starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-06 08:09:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: usage & billing changes (#1033)

Browse Source
This commit is contained in:
Ashish Kolhe
2023-06-26 17:32:35 +05:30
committed by GitHub
parent 4fb16a863f
commit b23751e4ce
19 changed files with 641 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
-1
View File
@@ -21,7 +21,6 @@ publish = false
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[features]
default = ["zo_functions"]
mimalloc = ["dep:mimalloc"]
jemalloc = ["dep:tikv-jemallocator"]
zo_functions = []
src/common/functions.rs
-1
View File
@@ -16,7 +16,6 @@ pub async fn get_all_transform_keys(org_id: &str) -> Vec<String> {
.collect()
}
#[cfg(feature = "zo_functions")]
pub fn init_vrl_runtime() -> vrl::compiler::runtime::Runtime {
vrl::compiler::runtime::Runtime::new(vrl::prelude::state::RuntimeState::default())
}
src/handler/http/request/search/mod.rs
+50
View File
@@ -26,8 +26,10 @@ use crate::common::json;
use crate::infra::config::CONFIG;
use crate::infra::{errors, metrics};
use crate::meta::http::HttpResponse as MetaHttpResponse;
use crate::meta::usage::{RequestStats, UsageType};
use crate::meta::{self, StreamType};
use crate::service::search as SearchService;
use crate::service::usage::report_usage_stats;
/** SearchStreamData*/
#[utoipa::path(
@@ -154,6 +156,22 @@ pub async fn search(
])
.inc();
res.set_local_took(start.elapsed().as_millis() as usize, took_wait);
let req_stats = RequestStats {
records: res.hits.len() as u64,
response_time: time,
size: res.scan_size as f64,
request_body: Some(req.query.sql),
};
let num_fn = req.query.query_fn.is_some() as u16;
report_usage_stats(
req_stats,
&org_id,
StreamType::Logs,
UsageType::Search,
num_fn,
)
.await;
Ok(HttpResponse::Ok().json(res))
}
Err(err) => {
@@ -427,6 +445,22 @@ pub async fn around(
])
.inc();
let req_stats = RequestStats {
records: resp.hits.len() as u64,
response_time: time,
size: resp.scan_size as f64,
request_body: Some(req.query.sql),
};
let num_fn = req.query.query_fn.is_some() as u16;
report_usage_stats(
req_stats,
&org_id,
StreamType::Logs,
UsageType::SearchAround,
num_fn,
)
.await;
Ok(HttpResponse::Ok().json(resp))
}
@@ -610,6 +644,22 @@ pub async fn values(
])
.inc();
let req_stats = RequestStats {
records: resp.hits.len() as u64,
response_time: time,
size: resp.scan_size as f64,
request_body: Some(req.query.sql),
};
let num_fn = req.query.query_fn.is_some() as u16;
report_usage_stats(
req_stats,
&org_id,
StreamType::Logs,
UsageType::SearchTopNValues,
num_fn,
)
.await;
Ok(HttpResponse::Ok().json(resp))
}
src/infra/config.rs
+14
View File
@@ -218,6 +218,20 @@ pub struct Common {
pub prometheus_enabled: bool,
#[env_config(name = "ZO_PRINT_KEY_CONFIG", default = false)]
pub print_key_config: bool,
#[env_config(
name = "USAGE_AUTH",
default = "cm9vdEBleGFtcGxlLmNvbTpDb21wbGV4cGFzcyMxMjM="
)]
pub usage_auth: String,
#[env_config(
name = "USAGE_EP",
default = "https://alpha1.gke.zinclabs.dev/api/atest/billings/_json"
)]
pub usage_url: String,
#[env_config(name = "USAGE_BATCH_SIZE", default = 5)]
pub usage_batch_size: usize,
#[env_config(name = "USAGE_REPORTING_ENABLED", default = false)]
pub usage_enabled: bool,
}
#[derive(EnvConfig)]
src/meta/mod.rs
+1
View File
@@ -30,6 +30,7 @@ pub mod stream;
pub mod syslog;
pub mod telemetry;
pub mod traces;
pub mod usage;
pub mod user;
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize, ToSchema)]
src/meta/usage.rs
+122
View File
@@ -0,0 +1,122 @@
use super::StreamType;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct UsageData {
pub event: UsageEvent,
pub day: u32,
pub hour: u32,
pub month: u32,
pub year: i32,
pub organization_identifier: String,
pub request_body: String,
pub size: f64,
pub unit: String,
pub user_email: String,
pub response_time: f64,
pub stream_type: StreamType,
pub num_records: u64,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum UsageEvent {
Ingestion,
Search,
Functions,
Other,
}
impl From<UsageType> for UsageEvent {
fn from(usage: UsageType) -> UsageEvent {
match usage {
UsageType::Bulk
| UsageType::Json
| UsageType::Multi
| UsageType::Traces
| UsageType::Metrics
| UsageType::KinesisFirehose
| UsageType::GCPSubscription
| UsageType::EnrichmentTable
| UsageType::Syslog
| UsageType::JsonMetrics => UsageEvent::Ingestion,
UsageType::Search | UsageType::SearchAround | UsageType::SearchTopNValues => {
UsageEvent::Search
}
UsageType::Functions => UsageEvent::Functions,
UsageType::Retention => UsageEvent::Other,
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum UsageType {
#[serde(rename = "logs/_bulk")]
Bulk,
#[serde(rename = "logs/_json")]
Json,
#[serde(rename = "logs/_multi")]
Multi,
#[serde(rename = "/traces")]
Traces,
#[serde(rename = "/v1/write")]
Metrics,
#[serde(rename = "/_search")]
Search,
#[serde(rename = "/_around")]
SearchAround,
#[serde(rename = "/_values")]
SearchTopNValues,
#[serde(rename = "functions")]
Functions,
#[serde(rename = "data_retention")]
Retention,
#[serde(rename = "/_kinesis_firehose")]
KinesisFirehose,
#[serde(rename = "/gcp/_sub")]
GCPSubscription,
#[serde(rename = "metrics/_json")]
JsonMetrics,
Syslog,
EnrichmentTable,
}
impl ToString for UsageType {
fn to_string(&self) -> String {
match self {
UsageType::Bulk => "logs/_bulk".to_owned(),
UsageType::Json => "logs/_json".to_owned(),
UsageType::JsonMetrics => "metrics/_json".to_owned(),
UsageType::Multi => "logs/_multi".to_owned(),
UsageType::Traces => "/traces".to_owned(),
UsageType::Metrics => "/v1/write".to_owned(),
UsageType::Search => "/_search".to_owned(),
UsageType::Functions => "functions".to_owned(),
UsageType::Retention => "data_retention".to_owned(),
UsageType::KinesisFirehose => "_kinesis_firehose".to_owned(),
UsageType::Syslog => "syslog".to_owned(),
UsageType::EnrichmentTable => "enrichment_table".to_owned(),
UsageType::SearchAround => "/_around".to_owned(),
UsageType::SearchTopNValues => "/_values".to_owned(),
UsageType::GCPSubscription => "/gcp/_sub".to_owned(),
}
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct RequestStats {
pub size: f64,
pub records: u64,
pub response_time: f64,
#[serde(default)]
pub request_body: Option<String>,
}
impl Default for RequestStats {
fn default() -> Self {
Self {
size: 0.0,
records: 0,
response_time: 0.0,
request_body: None,
}
}
}
src/service/enrichment_table/mod.rs
+16 -2
View File
@@ -23,19 +23,20 @@ use datafusion::arrow::datatypes::Schema;
use futures::{StreamExt, TryStreamExt};
use std::io::Error;
use crate::common::json;
use crate::infra::{
cache::stats,
cluster,
config::{CONFIG, STREAM_SCHEMAS},
};
use crate::meta::{self, http::HttpResponse as MetaHttpResponse, StreamType};
use crate::{common::json, meta::usage::UsageType};
use super::{
compact::delete,
db,
ingestion::{chk_schema_by_record, write_file},
schema::stream_schema_exists,
usage::report_usage_stats,
};
pub async fn save_enrichment_data(
@@ -44,6 +45,7 @@ pub async fn save_enrichment_data(
mut payload: Multipart,
thread_id: usize,
) -> Result<HttpResponse, Error> {
let start = std::time::Instant::now();
let mut hour_key = String::new();
let mut buf: AHashMap<String, Vec<String>> = AHashMap::new();
let stream_name = &crate::service::ingestion::format_stream_name(table_name);
@@ -142,13 +144,25 @@ pub async fn save_enrichment_data(
}
buf.insert(hour_key.clone(), records.clone());
write_file(
let mut stream_file_name = "".to_string();
let mut req_stats = write_file(
buf,
thread_id,
org_id,
stream_name,
&mut stream_file_name,
StreamType::EnrichmentTables,
);
req_stats.response_time = start.elapsed().as_secs_f64();
//metric + data usage
report_usage_stats(
req_stats,
org_id,
StreamType::Logs,
UsageType::EnrichmentTable,
0,
)
.await;
Ok(HttpResponse::Ok().json(MetaHttpResponse::error(
StatusCode::OK.into(),
src/service/ingestion/mod.rs
+57 -29
View File
@@ -12,42 +12,38 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use super::{db, triggers};
use crate::common::flatten;
use crate::infra::wal::get_or_create;
use ahash::AHashMap;
use arrow_schema::Schema;
use bytes::{BufMut, BytesMut};
use chrono::{TimeZone, Utc};
use datafusion::arrow::json::reader::infer_json_schema;
#[cfg(feature = "zo_functions")]
use std::collections::BTreeMap;
use std::io::BufReader;
#[cfg(feature = "zo_functions")]
use vector_enrichment::TableRegistry;
#[cfg(feature = "zo_functions")]
use vrl::compiler::{runtime::Runtime, CompilationResult, TargetValueRef};
#[cfg(feature = "zo_functions")]
use vrl::compiler::runtime::Runtime;
use vrl::compiler::{CompilationResult, TargetValueRef};
use vrl::prelude::state;
use super::{db, triggers};
#[cfg(feature = "zo_functions")]
use crate::common::flatten;
#[cfg(feature = "zo_functions")]
use crate::common::functions::get_vrl_compiler_config;
use crate::common::{
json::{self, Map, Value},
notification::send_notification,
};
#[cfg(feature = "zo_functions")]
use crate::infra::config::STREAM_FUNCTIONS;
use crate::infra::config::{CONFIG, STREAM_ALERTS};
use crate::infra::metrics;
#[cfg(feature = "zo_functions")]
use crate::meta::functions::{StreamTransform, VRLRuntimeConfig};
use crate::meta::usage::RequestStats;
use crate::meta::{
alert::{Alert, Trigger},
StreamType,
};
#[cfg(feature = "zo_functions")]
pub fn compile_vrl_function(func: &str, org_id: &str) -> Result<VRLRuntimeConfig, std::io::Error> {
if func.contains("get_env_var") {
return Err(std::io::Error::new(
@@ -80,7 +76,6 @@ pub fn compile_vrl_function(func: &str, org_id: &str) -> Result<VRLRuntimeConfig
}
}
#[cfg(feature = "zo_functions")]
pub fn apply_vrl_fn(runtime: &mut Runtime, vrl_runtime: &VRLRuntimeConfig, row: &Value) -> Value {
let mut metadata = vrl::value::Value::from(BTreeMap::new());
let mut target = TargetValueRef {
@@ -106,7 +101,6 @@ pub fn apply_vrl_fn(runtime: &mut Runtime, vrl_runtime: &VRLRuntimeConfig, row:
}
}
#[cfg(feature = "zo_functions")]
pub async fn get_stream_transforms<'a>(
org_id: &str,
stream_type: StreamType,
@@ -118,10 +112,10 @@ pub async fn get_stream_transforms<'a>(
if stream_transform_map.contains_key(&key) {
return;
}
let mut _local_tans: Vec<StreamTransform> = vec![];
(_local_tans, *stream_vrl_map) =
let mut _local_trans: Vec<StreamTransform> = vec![];
(_local_trans, *stream_vrl_map) =
crate::service::ingestion::register_stream_transforms(org_id, stream_type, stream_name);
stream_transform_map.insert(key, _local_tans);
stream_transform_map.insert(key, _local_trans);
}
pub async fn get_stream_partition_keys(
@@ -228,20 +222,19 @@ pub async fn send_ingest_notification(trigger: Trigger, alert: Alert) {
let _ = triggers::save_trigger(&trigger_to_save.alert_name, &trigger_to_save).await;
}
#[cfg(feature = "zo_functions")]
pub fn register_stream_transforms(
org_id: &str,
stream_type: StreamType,
stream_name: &str,
) -> (Vec<StreamTransform>, AHashMap<String, VRLRuntimeConfig>) {
let mut local_tans = vec![];
let mut local_trans = vec![];
let mut stream_vrl_map: AHashMap<String, VRLRuntimeConfig> = AHashMap::new();
let key = format!("{}/{}/{}", &org_id, stream_type, &stream_name);
if let Some(transforms) = STREAM_FUNCTIONS.get(&key) {
local_tans = (*transforms.list).to_vec();
local_tans.sort_by(|a, b| a.order.cmp(&b.order));
for trans in &local_tans {
local_trans = (*transforms.list).to_vec();
local_trans.sort_by(|a, b| a.order.cmp(&b.order));
for trans in &local_trans {
let func_key = format!("{}/{}", &stream_name, trans.transform.name);
if let Ok(vrl_runtime_config) = compile_vrl_function(&trans.transform.function, org_id)
{
@@ -255,19 +248,18 @@ pub fn register_stream_transforms(
}
}
(local_tans, stream_vrl_map)
(local_trans, stream_vrl_map)
}
#[cfg(feature = "zo_functions")]
pub fn apply_stream_transform<'a>(
local_tans: &Vec<StreamTransform>,
local_trans: &Vec<StreamTransform>,
value: &'a Value,
stream_vrl_map: &'a AHashMap<String, VRLRuntimeConfig>,
stream_name: &str,
runtime: &mut Runtime,
) -> Result<Value, anyhow::Error> {
let mut value = value.clone();
for trans in local_tans {
for trans in local_trans {
let func_key = format!("{stream_name}/{}", trans.transform.name);
if stream_vrl_map.contains_key(&func_key) && !value.is_null() {
let vrl_runtime = stream_vrl_map.get(&func_key).unwrap();
@@ -318,7 +310,7 @@ pub async fn chk_schema_by_record(
.unwrap();
}
pub fn write_file(
pub fn _write_file(
buf: AHashMap<String, Vec<String>>,
thread_id: usize,
org_id: &str,
@@ -368,11 +360,47 @@ pub fn write_file(
}
}
#[cfg(feature = "zo_functions")]
pub fn init_functions_runtime() -> Runtime {
crate::common::functions::init_vrl_runtime()
}
pub fn write_file(
buf: AHashMap<String, Vec<String>>,
thread_id: usize,
org_id: &str,
stream_name: &str,
stream_file_name: &mut String,
stream_type: StreamType,
) -> RequestStats {
let mut write_buf = BytesMut::new();
let mut req_stats = RequestStats::default();
for (key, entry) in buf {
if entry.is_empty() {
continue;
}
write_buf.clear();
for row in &entry {
write_buf.put(row.as_bytes());
write_buf.put("\n".as_bytes());
}
let file = get_or_create(
thread_id,
org_id,
stream_name,
stream_type,
&key,
CONFIG.common.wal_memory_mode_enabled,
);
if stream_file_name.is_empty() {
*stream_file_name = file.full_name();
}
file.write(write_buf.as_ref());
req_stats.size += write_buf.len() as f64 / (1024.0 * 1024.0);
req_stats.records += entry.len() as u64;
}
req_stats
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
@@ -434,7 +462,7 @@ mod tests {
}
#[actix_web::test]
#[cfg(feature = "zo_functions")]
async fn test_compile_vrl_function() {
let result = compile_vrl_function(
r#"if .country == "USA" {
src/service/logs/bulk.rs
+31 -15
View File
@@ -21,8 +21,9 @@ use std::io::{BufRead, BufReader};
use super::StreamMeta;
use crate::common::{flatten, json, time::parse_timestamp_micro_from_value};
use crate::infra::{cluster, config::CONFIG, metrics};
#[cfg(feature = "zo_functions")]
use crate::meta::functions::{StreamTransform, VRLRuntimeConfig};
use crate::meta::usage::{RequestStats, UsageType};
use crate::meta::{
alert::{Alert, Trigger},
ingestion::{
@@ -31,6 +32,7 @@ use crate::meta::{
},
StreamType,
};
use crate::service::usage::report_usage_stats;
use crate::service::{db, ingestion::write_file, schema::stream_schema_exists};
pub const TRANSFORM_FAILED: &str = "document_failed_transform";
@@ -60,14 +62,14 @@ pub async fn ingest(
let mut min_ts =
(Utc::now() + Duration::hours(CONFIG.limit.ingest_allowed_upto)).timestamp_micros();
#[cfg(feature = "zo_functions")]
let mut runtime = crate::service::ingestion::init_functions_runtime();
#[cfg(feature = "zo_functions")]
let mut stream_vrl_map: AHashMap<String, VRLRuntimeConfig> = AHashMap::new();
let mut stream_schema_map: AHashMap<String, Schema> = AHashMap::new();
let mut stream_data_map = AHashMap::new();
#[cfg(feature = "zo_functions")]
let mut stream_tansform_map: AHashMap<String, Vec<StreamTransform>> = AHashMap::new();
let mut stream_transform_map: AHashMap<String, Vec<StreamTransform>> = AHashMap::new();
let mut stream_partition_keys_map: AHashMap<String, (StreamSchemaChk, Vec<String>)> =
AHashMap::new();
let mut stream_alerts_map: AHashMap<String, Vec<Alert>> = AHashMap::new();
@@ -84,10 +86,9 @@ pub async fn ingest(
if line.is_empty() {
continue;
}
#[cfg(feature = "zo_functions")]
let value: json::Value = json::from_slice(line.as_bytes())?;
#[cfg(not(feature = "zo_functions"))]
let value: json::Value = json::from_slice(line.as_bytes())?;
if !next_line_is_data {
// check bulk operate
let ret = super::parse_bulk_index(&value);
@@ -98,12 +99,12 @@ pub async fn ingest(
next_line_is_data = true;
// Start Register Transfoms for stream
#[cfg(feature = "zo_functions")]
crate::service::ingestion::get_stream_transforms(
org_id,
StreamType::Logs,
&stream_name,
&mut stream_tansform_map,
&mut stream_transform_map,
&mut stream_vrl_map,
)
.await;
@@ -146,14 +147,13 @@ pub async fn ingest(
let buf = &mut stream_data.data;
//Start row based transform
#[cfg(feature = "zo_functions")]
let key = format!("{org_id}/{}/{stream_name}", StreamType::Logs);
//JSON Flattening
let mut value = flatten::flatten(&value)?;
#[cfg(feature = "zo_functions")]
if let Some(transforms) = stream_tansform_map.get(&key) {
if let Some(transforms) = stream_transform_map.get(&key) {
let mut ret_value = value.clone();
ret_value = crate::service::ingestion::apply_stream_transform(
transforms,
@@ -278,20 +278,25 @@ pub async fn ingest(
}
}
}
let mut final_req_stats = RequestStats::default();
for (stream_name, stream_data) in stream_data_map {
// check if we are allowed to ingest
if db::compact::delete::is_deleting_stream(org_id, &stream_name, StreamType::Logs, None) {
return Err(anyhow::anyhow!("stream [{stream_name}] is being deleted"));
}
// write to file
write_file(
let mut stream_file_name = "".to_string();
let req_stats = write_file(
stream_data.data,
thread_id,
org_id,
&stream_name,
&mut stream_file_name,
StreamType::Logs,
);
final_req_stats.size += req_stats.size;
final_req_stats.records += req_stats.records;
}
// only one trigger per request, as it updates etcd
@@ -300,6 +305,17 @@ pub async fn ingest(
}
let time = start.elapsed().as_secs_f64();
final_req_stats.response_time += time;
//metric + data usage
let fns_length: usize = stream_transform_map.values().map(|v| v.len()).sum();
report_usage_stats(
final_req_stats,
org_id,
StreamType::Logs,
UsageType::Bulk,
fns_length as u16,
)
.await;
metrics::HTTP_RESPONSE_TIME
.with_label_values(&[
"/api/org/ingest/logs/_bulk",
src/service/logs/gcs_pub_sub.rs
+26 -8
View File
@@ -6,6 +6,7 @@ use std::io::Read;
use crate::common::{flatten, json, time::parse_timestamp_micro_from_value};
use crate::infra::{cluster, config::CONFIG, metrics};
use crate::meta::usage::UsageType;
use crate::meta::{
alert::{Alert, Trigger},
ingestion::{
@@ -13,6 +14,7 @@ use crate::meta::{
},
StreamType,
};
use crate::service::usage::report_usage_stats;
use crate::service::{db, ingestion::write_file};
use super::StreamMeta;
@@ -38,7 +40,6 @@ pub async fn process(
)));
}
#[cfg(feature = "zo_functions")]
let mut runtime = crate::service::ingestion::init_functions_runtime();
let mut stream_schema_map: AHashMap<String, Schema> = AHashMap::new();
@@ -55,8 +56,8 @@ pub async fn process(
let mut trigger: Option<Trigger> = None;
// Start Register Transforms for stream
#[cfg(feature = "zo_functions")]
let (local_tans, stream_vrl_map) = crate::service::ingestion::register_stream_transforms(
let (local_trans, stream_vrl_map) = crate::service::ingestion::register_stream_transforms(
org_id,
StreamType::Logs,
stream_name,
@@ -100,15 +101,14 @@ pub async fn process(
value = flatten::flatten(&value)?;
// Start row based transform
#[cfg(feature = "zo_functions")]
let mut value = crate::service::ingestion::apply_stream_transform(
&local_tans,
&local_trans,
&value,
&stream_vrl_map,
stream_name,
&mut runtime,
)?;
#[cfg(feature = "zo_functions")]
if value.is_null() || !value.is_object() {
stream_status.status.failed += 1; // transform failed or dropped
}
@@ -164,9 +164,16 @@ pub async fn process(
timestamp: request.message.publish_time,
});
}
let mut stream_file_name = "".to_string();
// write to file
write_file(buf, thread_id, org_id, stream_name, StreamType::Logs);
let mut req_stats = write_file(
buf,
thread_id,
org_id,
stream_name,
&mut stream_file_name,
StreamType::Logs,
);
// only one trigger per request, as it updates etcd
super::evaluate_trigger(trigger, stream_alerts_map).await;
@@ -190,6 +197,17 @@ pub async fn process(
])
.inc();
req_stats.response_time = start.elapsed().as_secs_f64();
//metric + data usage
report_usage_stats(
req_stats,
org_id,
StreamType::Logs,
UsageType::GCPSubscription,
local_trans.len() as u16,
)
.await;
Ok(GCPIngestionResponse {
request_id: request.message.message_id,
timestamp: request.message.publish_time,
src/service/logs/json.rs
+31 -5
View File
@@ -20,11 +20,13 @@ use datafusion::arrow::datatypes::Schema;
use super::StreamMeta;
use crate::common::{flatten, json, time::parse_timestamp_micro_from_value};
use crate::infra::{cluster, config::CONFIG, metrics};
use crate::meta::usage::UsageType;
use crate::meta::{
alert::{Alert, Trigger},
ingestion::{IngestionResponse, StreamStatus},
StreamType,
};
use crate::service::usage::report_usage_stats;
use crate::service::{db, ingestion::write_file, schema::stream_schema_exists};
pub async fn ingest(
@@ -52,7 +54,6 @@ pub async fn ingest(
let mut min_ts =
(Utc::now() + Duration::hours(CONFIG.limit.ingest_allowed_upto)).timestamp_micros();
#[cfg(feature = "zo_functions")]
let mut runtime = crate::service::ingestion::init_functions_runtime();
let mut stream_schema_map: AHashMap<String, Schema> = AHashMap::new();
@@ -62,7 +63,7 @@ pub async fn ingest(
let mut trigger: Option<Trigger> = None;
// Start Register Transforms for stream
#[cfg(feature = "zo_functions")]
let (local_trans, stream_vrl_map) = crate::service::ingestion::register_stream_transforms(
org_id,
StreamType::Logs,
@@ -95,7 +96,6 @@ pub async fn ingest(
//JSON Flattening
let mut value = flatten::flatten(item)?;
#[cfg(feature = "zo_functions")]
if !local_trans.is_empty() {
value = crate::service::ingestion::apply_stream_transform(
&local_trans,
@@ -105,7 +105,7 @@ pub async fn ingest(
&mut runtime,
)?;
}
#[cfg(feature = "zo_functions")]
if value.is_null() || !value.is_object() {
stream_status.status.failed += 1; // transform failed or dropped
continue;
@@ -162,7 +162,22 @@ pub async fn ingest(
}
// write to file
write_file(buf, thread_id, org_id, stream_name, StreamType::Logs);
let mut stream_file_name = "".to_string();
let mut req_stats = write_file(
buf,
thread_id,
org_id,
stream_name,
&mut stream_file_name,
StreamType::Logs,
);
if stream_file_name.is_empty() {
return Ok(IngestionResponse::new(
http::StatusCode::OK.into(),
vec![stream_status],
));
}
// only one trigger per request, as it updates etcd
super::evaluate_trigger(trigger, stream_alerts_map).await;
@@ -187,6 +202,17 @@ pub async fn ingest(
])
.inc();
req_stats.response_time = start.elapsed().as_secs_f64();
//metric + data usage
report_usage_stats(
req_stats,
org_id,
StreamType::Logs,
UsageType::Json,
local_trans.len() as u16,
)
.await;
Ok(IngestionResponse::new(
http::StatusCode::OK.into(),
vec![stream_status],
src/service/logs/kinesis_firehose.rs
+43 -11
View File
@@ -18,11 +18,8 @@ use datafusion::arrow::datatypes::Schema;
use flate2::read::GzDecoder;
use std::io::Read;
use crate::common::{
flatten, json,
time::{parse_i64_to_timestamp_micros, parse_timestamp_micro_from_value},
};
use crate::infra::{cluster, config::CONFIG, metrics};
use crate::meta::usage::UsageType;
use crate::meta::{
alert::{Alert, Trigger},
ingestion::{
@@ -31,6 +28,13 @@ use crate::meta::{
StreamType,
};
use crate::service::{db, ingestion::write_file};
use crate::{
common::{
flatten, json,
time::{parse_i64_to_timestamp_micros, parse_timestamp_micro_from_value},
},
service::usage::report_usage_stats,
};
use super::StreamMeta;
@@ -52,7 +56,6 @@ pub async fn process(
return Err(anyhow::anyhow!("stream [{stream_name}] is being deleted"));
}
#[cfg(feature = "zo_functions")]
let mut runtime = crate::service::ingestion::init_functions_runtime();
let mut min_ts =
@@ -64,8 +67,8 @@ pub async fn process(
let mut trigger: Option<Trigger> = None;
// Start Register Transforms for stream
#[cfg(feature = "zo_functions")]
let (local_tans, stream_vrl_map) = crate::service::ingestion::register_stream_transforms(
let (local_trans, stream_vrl_map) = crate::service::ingestion::register_stream_transforms(
org_id,
StreamType::Logs,
stream_name,
@@ -172,15 +175,15 @@ pub async fn process(
value = flatten::flatten(&value)?;
// Start row based transform
#[cfg(feature = "zo_functions")]
let mut value = crate::service::ingestion::apply_stream_transform(
&local_tans,
&local_trans,
&value,
&stream_vrl_map,
stream_name,
&mut runtime,
)?;
#[cfg(feature = "zo_functions")]
if value.is_null() || !value.is_object() {
stream_status.status.failed += 1; // transform failed or dropped
continue;
@@ -229,11 +232,40 @@ pub async fn process(
}
// write to file
write_file(buf, thread_id, org_id, stream_name, StreamType::Logs);
let mut stream_file_name = "".to_string();
let mut req_stats = write_file(
buf,
thread_id,
org_id,
stream_name,
&mut stream_file_name,
StreamType::Logs,
);
if stream_file_name.is_empty() {
return Ok(KinesisFHIngestionResponse {
request_id: request.request_id,
error_message: Some(
json::to_string(&stream_status).unwrap_or("error processing request".to_string()),
),
timestamp: request.timestamp.unwrap_or(Utc::now().timestamp_micros()),
});
}
// only one trigger per request, as it updates etcd
super::evaluate_trigger(trigger, stream_alerts_map).await;
req_stats.response_time = start.elapsed().as_secs_f64();
//metric + data usage
report_usage_stats(
req_stats,
org_id,
StreamType::Logs,
UsageType::KinesisFirehose,
local_trans.len() as u16,
)
.await;
metrics::HTTP_RESPONSE_TIME
.with_label_values(&[
"/api/org/ingest/logs/_kinesis",
src/service/logs/multi.rs
+33 -5
View File
@@ -20,11 +20,13 @@ use std::io::{BufRead, BufReader};
use crate::common::{flatten, json, time::parse_timestamp_micro_from_value};
use crate::infra::{cluster, config::CONFIG, metrics};
use crate::meta::usage::UsageType;
use crate::meta::{
alert::{Alert, Trigger},
ingestion::{IngestionResponse, StreamStatus},
StreamType,
};
use crate::service::usage::report_usage_stats;
use crate::service::{db, ingestion::write_file, logs::StreamMeta, schema::stream_schema_exists};
pub async fn ingest(
@@ -48,7 +50,6 @@ pub async fn ingest(
if db::compact::delete::is_deleting_stream(org_id, stream_name, StreamType::Logs, None) {
return Err(anyhow::anyhow!("stream [{stream_name}] is being deleted"));
}
#[cfg(feature = "zo_functions")]
let mut runtime = crate::service::ingestion::init_functions_runtime();
let mut min_ts =
@@ -60,7 +61,7 @@ pub async fn ingest(
let mut trigger: Option<Trigger> = None;
// Start Register Transforms for stream
#[cfg(feature = "zo_functions")]
let (local_trans, stream_vrl_map) = crate::service::ingestion::register_stream_transforms(
org_id,
StreamType::Logs,
@@ -99,7 +100,7 @@ pub async fn ingest(
// JSON Flattening
value = flatten::flatten(&value)?;
// Start row based transform
#[cfg(feature = "zo_functions")]
if !local_trans.is_empty() {
value = crate::service::ingestion::apply_stream_transform(
&local_trans,
@@ -109,7 +110,7 @@ pub async fn ingest(
&mut runtime,
)?;
}
#[cfg(feature = "zo_functions")]
if value.is_null() || !value.is_object() {
stream_status.status.failed += 1; // transform failed or dropped
continue;
@@ -167,11 +168,38 @@ pub async fn ingest(
}
// write to file
write_file(buf, thread_id, org_id, stream_name, StreamType::Logs);
let mut stream_file_name = "".to_string();
let mut req_stats = write_file(
buf,
thread_id,
org_id,
stream_name,
&mut stream_file_name,
StreamType::Logs,
);
if stream_file_name.is_empty() {
return Ok(IngestionResponse::new(
http::StatusCode::OK.into(),
vec![stream_status],
));
}
// only one trigger per request, as it updates etcd
super::evaluate_trigger(trigger, stream_alerts_map).await;
req_stats.response_time = start.elapsed().as_secs_f64();
//metric + data usage
report_usage_stats(
req_stats,
org_id,
StreamType::Logs,
UsageType::Multi,
local_trans.len() as u16,
)
.await;
metrics::HTTP_RESPONSE_TIME
.with_label_values(&[
"/api/org/ingest/logs/_multi",
src/service/logs/syslog.rs
+16 -5
View File
@@ -84,7 +84,7 @@ pub async fn ingest(msg: &str, addr: SocketAddr) -> Result<HttpResponse, ()> {
let mut trigger: Option<Trigger> = None;
// Start Register Transforms for stream
/* #[cfg(feature = "zo_functions")]
/*
let (local_tans, _, stream_vrl_map) = crate::service::ingestion::register_stream_transforms(
org_id,
stream_name,
@@ -118,7 +118,7 @@ pub async fn ingest(msg: &str, addr: SocketAddr) -> Result<HttpResponse, ()> {
let mut value = message_to_value(parsed_msg);
value = flatten::flatten(&value).unwrap();
/* #[cfg(feature = "zo_functions")]
/*
let mut value = crate::service::ingestion::apply_stream_transform(
&local_tans,
&value,
@@ -128,7 +128,7 @@ pub async fn ingest(msg: &str, addr: SocketAddr) -> Result<HttpResponse, ()> {
stream_name,
&mut runtime,
);
#[cfg(feature = "zo_functions")]
if value.is_null() || !value.is_object() {
stream_status.status.failed += 1; // transform failed or dropped
} */
@@ -174,12 +174,23 @@ pub async fn ingest(msg: &str, addr: SocketAddr) -> Result<HttpResponse, ()> {
if local_trigger.is_some() {
trigger = Some(local_trigger.unwrap());
}
write_file(buf, thread_id, org_id, stream_name, StreamType::Logs);
let mut stream_file_name = "".to_string();
write_file(
buf,
thread_id,
org_id,
stream_name,
&mut stream_file_name,
StreamType::Logs,
);
// only one trigger per request, as it updates etcd
super::evaluate_trigger(trigger, stream_alerts_map).await;
/* req_stats.response_time = start.elapsed().as_secs_f64();
//metric + data usage
report_ingest_stats(&req_stats, org_id, StreamType::Logs, UsageEvent::Syslog).await; */
let time = start.elapsed().as_secs_f64();
metrics::HTTP_RESPONSE_TIME
.with_label_values(&[
src/service/metrics/json.rs
+20 -11
View File
@@ -17,16 +17,17 @@ use ahash::AHashMap;
use anyhow::{anyhow, Result};
use datafusion::arrow::{datatypes::Schema, json::reader::infer_json_schema};
use std::{collections::HashMap, io::BufReader};
#[cfg(feature = "zo_functions")]
use vrl::compiler::runtime::Runtime;
use crate::common::{flatten, json, time};
use crate::infra::{cluster, config::CONFIG, metrics};
use crate::meta::usage::{RequestStats, UsageType};
use crate::meta::{
ingestion::{IngestionResponse, StreamStatus},
prom::{Metadata, HASH_LABEL, METADATA_LABEL, NAME_LABEL, TYPE_LABEL, VALUE_LABEL},
StreamType,
};
use crate::service::usage::report_usage_stats;
use crate::service::{
db,
ingestion::{get_hour_key, write_file},
@@ -43,9 +44,7 @@ pub async fn ingest(org_id: &str, body: web::Bytes, thread_id: usize) -> Result<
return Err(anyhow::anyhow!("Quota exceeded for this organisation"));
}
#[cfg(feature = "zo_functions")]
let mut runtime = crate::service::ingestion::init_functions_runtime();
let mut stream_schema_map: AHashMap<String, Schema> = AHashMap::new();
let mut stream_status_map: AHashMap<String, StreamStatus> = AHashMap::new();
let mut stream_data_buf: AHashMap<String, AHashMap<String, Vec<String>>> = AHashMap::new();
@@ -70,11 +69,9 @@ pub async fn ingest(org_id: &str, body: web::Bytes, thread_id: usize) -> Result<
};
// apply functions
#[cfg(feature = "zo_functions")]
let mut record = json::Value::Object(record.to_owned());
#[cfg(feature = "zo_functions")]
apply_func(&mut runtime, org_id, &stream_name, &mut record)?;
#[cfg(feature = "zo_functions")]
let record = record.as_object_mut().unwrap();
// check timestamp & value
@@ -176,24 +173,28 @@ pub async fn ingest(org_id: &str, body: web::Bytes, thread_id: usize) -> Result<
.or_insert(StreamStatus::new(&stream_name));
stream_status.status.successful += 1;
}
let mut final_req_stats = RequestStats::default();
for (stream_name, stream_data) in stream_data_buf {
// check if we are allowed to ingest
if db::compact::delete::is_deleting_stream(org_id, &stream_name, StreamType::Metrics, None)
{
return Err(anyhow!("stream [{stream_name}] is being deleted"));
}
// write to file
write_file(
let mut stream_file_name = "".to_string();
let req_stats = write_file(
stream_data,
thread_id,
org_id,
&stream_name,
StreamType::Metrics,
&mut stream_file_name,
StreamType::Logs,
);
final_req_stats.size += req_stats.size;
final_req_stats.records += req_stats.records;
}
let time = start.elapsed().as_secs_f64();
final_req_stats.response_time += time;
metrics::HTTP_RESPONSE_TIME
.with_label_values(&[
"/api/org/ingest/metrics/_json",
@@ -212,6 +213,15 @@ pub async fn ingest(org_id: &str, body: web::Bytes, thread_id: usize) -> Result<
&StreamType::Metrics.to_string(),
])
.inc();
//let fns_length: usize = stream_transform_map.values().map(|v| v.len()).sum();
report_usage_stats(
final_req_stats,
org_id,
StreamType::Logs,
UsageType::JsonMetrics,
0,
)
.await;
Ok(IngestionResponse::new(
http::StatusCode::OK.into(),
@@ -219,7 +229,6 @@ pub async fn ingest(org_id: &str, body: web::Bytes, thread_id: usize) -> Result<
))
}
#[cfg(feature = "zo_functions")]
fn apply_func(
runtime: &mut Runtime,
org_id: &str,
src/service/metrics/prom.rs
+45 -14
View File
@@ -28,12 +28,15 @@ use crate::infra::{
errors::{Error, Result},
metrics,
};
use crate::meta::functions::StreamTransform;
use crate::meta::usage::{RequestStats, UsageType};
use crate::meta::{
self,
alert::{Alert, Trigger},
prom::{self, HASH_LABEL, METADATA_LABEL, NAME_LABEL, VALUE_LABEL},
StreamType,
};
use crate::service::usage::report_usage_stats;
use crate::service::{
db,
ingestion::{chk_schema_by_record, write_file},
@@ -71,6 +74,8 @@ pub async fn remote_write(
let mut metric_schema_map: AHashMap<String, Schema> = AHashMap::new();
let mut stream_alerts_map: AHashMap<String, Vec<Alert>> = AHashMap::new();
let mut stream_trigger_map: AHashMap<String, Trigger> = AHashMap::new();
let mut stream_transform_map: AHashMap<String, Vec<StreamTransform>> = AHashMap::new();
let mut final_req_stats = RequestStats::default();
let decoded = snap::raw::Decoder::new()
.decompress_vec(&body)
@@ -152,7 +157,10 @@ pub async fn remote_write(
value: sample_val,
};
let timestamp = parse_i64_to_timestamp_micros(sample.timestamp);
let mut timestamp = parse_i64_to_timestamp_micros(sample.timestamp);
if timestamp == 0 {
timestamp = Utc::now().timestamp_micros();
}
if timestamp < min_ts {
min_ts = timestamp;
}
@@ -210,34 +218,32 @@ pub async fn remote_write(
crate::service::ingestion::get_stream_alerts(key, &mut stream_alerts_map).await;
// End get stream alert
#[cfg(feature = "zo_functions")]
let mut runtime = crate::service::ingestion::init_functions_runtime();
// Start Register Transforms for stream
#[cfg(feature = "zo_functions")]
let (local_tans, stream_vrl_map) =
let (local_trans, stream_vrl_map) =
crate::service::ingestion::register_stream_transforms(
org_id,
StreamType::Metrics,
&metric_name,
);
stream_transform_map.insert(metric_name.to_owned(), local_trans.clone());
// End Register Transforms for stream
#[cfg(not(feature = "zo_functions"))]
let mut value: json::Value = json::to_value(&metric).unwrap();
#[cfg(feature = "zo_functions")]
let value: json::Value = json::to_value(&metric).unwrap();
// Start row based transform
#[cfg(feature = "zo_functions")]
let mut value = crate::service::ingestion::apply_stream_transform(
&local_tans,
value = crate::service::ingestion::apply_stream_transform(
&local_trans,
&value,
&stream_vrl_map,
&metric_name,
&mut runtime,
)?;
// End row based transform
// get json object
@@ -310,19 +316,33 @@ pub async fn remote_write(
}
for (stream_name, stream_data) in metric_data_map {
// write to file
let mut stream_file_name = "".to_string();
// check if we are allowed to ingest
if db::compact::delete::is_deleting_stream(org_id, &stream_name, StreamType::Metrics, None)
{
return Err(anyhow::anyhow!("stream [{stream_name}] is being deleted"));
}
// write to file
write_file(
let req_stats = write_file(
stream_data,
thread_id,
org_id,
&stream_name,
StreamType::Metrics,
&mut stream_file_name,
StreamType::Logs,
);
final_req_stats.size += req_stats.size;
final_req_stats.records += req_stats.records;
let _schema_exists = stream_schema_exists(
org_id,
&stream_name,
StreamType::Metrics,
&mut metric_schema_map,
)
.await;
}
// only one trigger per request, as it updates etcd
@@ -366,6 +386,17 @@ pub async fn remote_write(
&StreamType::Metrics.to_string(),
])
.inc();
final_req_stats.response_time += time;
//metric + data usage
let fns_length: usize = stream_transform_map.values().map(|v| v.len()).sum();
report_usage_stats(
final_req_stats,
org_id,
StreamType::Metrics,
UsageType::Metrics,
fns_length as u16,
)
.await;
Ok(())
}
src/service/mod.rs
+1
View File
@@ -34,6 +34,7 @@ pub mod stream;
pub mod syslogs_route;
pub mod traces;
pub mod triggers;
pub mod usage;
pub mod users;
// generate partition key for query
src/service/traces/otlp_http.rs
+41 -44
View File
@@ -14,7 +14,6 @@
use actix_web::{http, web, HttpResponse};
use ahash::AHashMap;
use bytes::{BufMut, BytesMut};
use chrono::{Duration, Utc};
use datafusion::arrow::datatypes::Schema;
use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
@@ -24,12 +23,11 @@ use std::{
io::{BufRead, BufReader, Error},
};
use crate::common::{flatten, json};
use crate::infra::{cluster, config::CONFIG, wal};
use crate::meta::{
alert::{Alert, Evaluate, Trigger},
http::HttpResponse as MetaHttpResponse,
traces::{Event, Span, SpanRefType},
usage::UsageType,
StreamType,
};
use crate::service::{
@@ -37,6 +35,14 @@ use crate::service::{
ingestion::{format_stream_name, get_partition_key_record},
schema::{add_stream_schema, stream_schema_exists},
};
use crate::{
common::{flatten, json},
service::ingestion::write_file,
};
use crate::{
infra::{cluster, config::CONFIG},
service::usage::report_usage_stats,
};
const PARENT_SPAN_ID: &str = "reference.parent_span_id";
const PARENT_TRACE_ID: &str = "reference.parent_trace_id";
@@ -74,6 +80,7 @@ pub async fn traces_json(
)));
}
let start = std::time::Instant::now();
let traces_stream_name = "default";
let mut trace_meta_coll: AHashMap<String, Vec<json::Map<String, json::Value>>> =
AHashMap::new();
@@ -107,9 +114,9 @@ pub async fn traces_json(
let mut trigger: Option<Trigger> = None;
/* // Start Register Transforms for stream
#[cfg(feature = "zo_functions")]
let (lua, mut runtime) = crate::service::ingestion::init_functions_runtime();
#[cfg(feature = "zo_functions")]
let (local_tans, stream_lua_map, stream_vrl_map) =
crate::service::ingestion::register_stream_transforms(
org_id,
@@ -269,7 +276,7 @@ pub async fn traces_json(
value = flatten::flatten(&value).unwrap();
/* // Start row based transform
#[cfg(feature = "zo_functions")]
let mut value = crate::service::ingestion::apply_stream_transform(
&local_tans,
&value,
@@ -279,7 +286,7 @@ pub async fn traces_json(
traces_stream_name,
&mut runtime,
);
#[cfg(feature = "zo_functions")]
if value.is_null() || !value.is_object() {
continue;
}
@@ -370,51 +377,41 @@ pub async fn traces_json(
)));
}
}
let mut write_buf = BytesMut::new();
for (key, entry) in data_buf {
if entry.is_empty() {
continue;
}
write_buf.clear();
for row in &entry {
write_buf.put(row.as_bytes());
write_buf.put("\n".as_bytes());
}
let file = wal::get_or_create(
thread_id,
org_id,
traces_stream_name,
StreamType::Traces,
&key,
false,
);
let traces_file_name = file.full_name();
let mut traces_file_name = "".to_string();
let mut req_stats = write_file(
data_buf,
thread_id,
org_id,
traces_stream_name,
&mut traces_file_name,
StreamType::Traces,
);
req_stats.response_time = start.elapsed().as_secs_f64();
//metric + data usage
report_usage_stats(req_stats, org_id, StreamType::Traces, UsageType::Traces, 0).await;
file.write(write_buf.as_ref());
let schema_exists = stream_schema_exists(
let schema_exists = stream_schema_exists(
org_id,
traces_stream_name,
StreamType::Traces,
&mut traces_schema_map,
)
.await;
if !schema_exists.has_fields && !traces_file_name.is_empty() {
let file = OpenOptions::new()
.read(true)
.open(&traces_file_name)
.unwrap();
add_stream_schema(
org_id,
traces_stream_name,
StreamType::Traces,
&file,
&mut traces_schema_map,
min_ts,
)
.await;
if !schema_exists.has_fields && !traces_file_name.is_empty() {
let file = OpenOptions::new()
.read(true)
.open(&traces_file_name)
.unwrap();
add_stream_schema(
org_id,
traces_stream_name,
StreamType::Traces,
&file,
&mut traces_schema_map,
min_ts,
)
.await;
}
}
// only one trigger per request, as it updates etcd
src/service/usage/mod.rs
+94
View File
@@ -0,0 +1,94 @@
use chrono::{Datelike, Timelike, Utc};
use once_cell::sync::Lazy;
use reqwest::Client;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::{
infra::{config::CONFIG, metrics},
meta::{
usage::{RequestStats, UsageData, UsageEvent, UsageType},
StreamType,
},
};
pub static USAGE_DATA: Lazy<Arc<RwLock<Vec<UsageData>>>> =
Lazy::new(|| Arc::new(RwLock::new(vec![])));
pub async fn report_usage_stats(
stats: RequestStats,
org_id: &str,
stream_type: StreamType,
event: UsageType,
num_functions: u16,
) {
if !CONFIG.common.usage_enabled {
return;
}
let local_stream_type = stream_type.to_string();
// metrics
metrics::HTTP_RESPONSE_TIME
.with_label_values(&[&event.to_string(), "200", org_id, "", &local_stream_type])
.observe(stats.response_time);
metrics::HTTP_INCOMING_REQUESTS
.with_label_values(&[&event.to_string(), "200", org_id, "", &local_stream_type])
.inc();
let request_body = stats.request_body.unwrap_or(event.to_string());
let now = Utc::now();
let mut usage = vec![UsageData {
event: event.into(),
day: now.day(),
hour: now.hour(),
month: now.month(),
year: now.year(),
organization_identifier: org_id.to_owned(),
request_body: request_body.to_owned(),
size: stats.size,
unit: "MB".to_owned(),
user_email: "".to_owned(),
response_time: stats.response_time,
num_records: stats.records,
stream_type,
}];
if num_functions > 0 {
usage.push(UsageData {
event: UsageEvent::Functions,
day: now.day(),
hour: now.hour(),
month: now.month(),
year: now.year(),
organization_identifier: org_id.to_owned(),
request_body,
size: stats.size,
unit: "MB".to_owned(),
user_email: "".to_owned(),
response_time: stats.response_time,
num_records: stats.records * num_functions as u64,
stream_type,
})
}
publish_usage(usage).await;
}
pub async fn publish_usage(mut usage: Vec<UsageData>) {
let mut usages = USAGE_DATA.write().await;
usages.append(&mut usage);
if usages.len() >= CONFIG.common.usage_batch_size {
let cl = Arc::new(Client::builder().build().unwrap());
let curr_usage = std::mem::take(&mut *usages);
let url = url::Url::parse(&CONFIG.common.usage_url).unwrap();
let auth = format!("Basic {}", &CONFIG.common.usage_auth);
let cl = Arc::clone(&cl);
tokio::task::spawn(async move {
let _ = cl
.post(url)
.header("Content-Type", "application/json")
.header(reqwest::header::AUTHORIZATION, auth)
.json(&curr_usage)
.send()
.await;
});
}
}