From b23751e4cebf11b7631b0dfd9fd5e265f922eac2 Mon Sep 17 00:00:00 2001 From: Ashish Kolhe <35160958+oasisk@users.noreply.github.com> Date: Mon, 26 Jun 2023 17:32:35 +0530 Subject: [PATCH] feat: usage & billing changes (#1033) --- Cargo.toml | 1 - src/common/functions.rs | 1 - src/handler/http/request/search/mod.rs | 50 ++++++++++ src/infra/config.rs | 14 +++ src/meta/mod.rs | 1 + src/meta/usage.rs | 122 +++++++++++++++++++++++++ src/service/enrichment_table/mod.rs | 18 +++- src/service/ingestion/mod.rs | 86 +++++++++++------ src/service/logs/bulk.rs | 46 +++++++--- src/service/logs/gcs_pub_sub.rs | 34 +++++-- src/service/logs/json.rs | 36 +++++++- src/service/logs/kinesis_firehose.rs | 54 ++++++++--- src/service/logs/multi.rs | 38 +++++++- src/service/logs/syslog.rs | 21 ++++- src/service/metrics/json.rs | 31 ++++--- src/service/metrics/prom.rs | 59 +++++++++--- src/service/mod.rs | 1 + src/service/traces/otlp_http.rs | 85 +++++++++-------- src/service/usage/mod.rs | 94 +++++++++++++++++++ 19 files changed, 641 insertions(+), 151 deletions(-) create mode 100644 src/meta/usage.rs create mode 100644 src/service/usage/mod.rs diff --git a/Cargo.toml b/Cargo.toml index b8e6511905..4cd7aa98b8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 = [] diff --git a/src/common/functions.rs b/src/common/functions.rs index 87bb186989..e5475be41c 100644 --- a/src/common/functions.rs +++ b/src/common/functions.rs @@ -16,7 +16,6 @@ pub async fn get_all_transform_keys(org_id: &str) -> Vec { .collect() } -#[cfg(feature = "zo_functions")] pub fn init_vrl_runtime() -> vrl::compiler::runtime::Runtime { vrl::compiler::runtime::Runtime::new(vrl::prelude::state::RuntimeState::default()) } diff --git a/src/handler/http/request/search/mod.rs b/src/handler/http/request/search/mod.rs index 74d47000b2..b68148de8f 100644 --- a/src/handler/http/request/search/mod.rs +++ b/src/handler/http/request/search/mod.rs @@ -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)) } diff --git a/src/infra/config.rs b/src/infra/config.rs index 611a41fcd3..f33c79be35 100644 --- a/src/infra/config.rs +++ b/src/infra/config.rs @@ -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)] diff --git a/src/meta/mod.rs b/src/meta/mod.rs index daaa6eea65..8b3a7894c9 100644 --- a/src/meta/mod.rs +++ b/src/meta/mod.rs @@ -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)] diff --git a/src/meta/usage.rs b/src/meta/usage.rs new file mode 100644 index 0000000000..286f543792 --- /dev/null +++ b/src/meta/usage.rs @@ -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 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, +} +impl Default for RequestStats { + fn default() -> Self { + Self { + size: 0.0, + records: 0, + response_time: 0.0, + request_body: None, + } + } +} diff --git a/src/service/enrichment_table/mod.rs b/src/service/enrichment_table/mod.rs index 7d74d5829f..54f8fa9a0d 100644 --- a/src/service/enrichment_table/mod.rs +++ b/src/service/enrichment_table/mod.rs @@ -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 { + let start = std::time::Instant::now(); let mut hour_key = String::new(); let mut buf: AHashMap> = 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(), diff --git a/src/service/ingestion/mod.rs b/src/service/ingestion/mod.rs index 2c2818540f..75fe9f266c 100644 --- a/src/service/ingestion/mod.rs +++ b/src/service/ingestion/mod.rs @@ -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 { 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 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 = vec![]; - (_local_tans, *stream_vrl_map) = + let mut _local_trans: Vec = 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, AHashMap) { - let mut local_tans = vec![]; + let mut local_trans = vec![]; let mut stream_vrl_map: AHashMap = 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, + local_trans: &Vec, value: &'a Value, stream_vrl_map: &'a AHashMap, stream_name: &str, runtime: &mut Runtime, ) -> Result { 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>, 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>, + 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" { diff --git a/src/service/logs/bulk.rs b/src/service/logs/bulk.rs index 1bef835e4f..14601c4307 100644 --- a/src/service/logs/bulk.rs +++ b/src/service/logs/bulk.rs @@ -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 = AHashMap::new(); let mut stream_schema_map: AHashMap = AHashMap::new(); let mut stream_data_map = AHashMap::new(); - #[cfg(feature = "zo_functions")] - let mut stream_tansform_map: AHashMap> = AHashMap::new(); + + let mut stream_transform_map: AHashMap> = AHashMap::new(); let mut stream_partition_keys_map: AHashMap)> = AHashMap::new(); let mut stream_alerts_map: AHashMap> = 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", diff --git a/src/service/logs/gcs_pub_sub.rs b/src/service/logs/gcs_pub_sub.rs index d732e7bf50..98ed945ba6 100644 --- a/src/service/logs/gcs_pub_sub.rs +++ b/src/service/logs/gcs_pub_sub.rs @@ -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 = AHashMap::new(); @@ -55,8 +56,8 @@ pub async fn process( let mut trigger: Option = 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, diff --git a/src/service/logs/json.rs b/src/service/logs/json.rs index 9347e6e2a8..403c69cd20 100644 --- a/src/service/logs/json.rs +++ b/src/service/logs/json.rs @@ -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 = AHashMap::new(); @@ -62,7 +63,7 @@ pub async fn ingest( let mut trigger: Option = 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], diff --git a/src/service/logs/kinesis_firehose.rs b/src/service/logs/kinesis_firehose.rs index 72b15187a3..05ed3d67ca 100644 --- a/src/service/logs/kinesis_firehose.rs +++ b/src/service/logs/kinesis_firehose.rs @@ -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 = 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", diff --git a/src/service/logs/multi.rs b/src/service/logs/multi.rs index 9ffd1a712e..25e089c5da 100644 --- a/src/service/logs/multi.rs +++ b/src/service/logs/multi.rs @@ -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 = 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", diff --git a/src/service/logs/syslog.rs b/src/service/logs/syslog.rs index 002dea30a1..9c71a120a6 100644 --- a/src/service/logs/syslog.rs +++ b/src/service/logs/syslog.rs @@ -84,7 +84,7 @@ pub async fn ingest(msg: &str, addr: SocketAddr) -> Result { let mut trigger: Option = 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 { 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 { 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 { 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(&[ diff --git a/src/service/metrics/json.rs b/src/service/metrics/json.rs index 7bcdd8f5f4..e99cae1703 100644 --- a/src/service/metrics/json.rs +++ b/src/service/metrics/json.rs @@ -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 = AHashMap::new(); let mut stream_status_map: AHashMap = AHashMap::new(); let mut stream_data_buf: AHashMap>> = 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, diff --git a/src/service/metrics/prom.rs b/src/service/metrics/prom.rs index f64fa9b6ea..c7ba99c0f1 100644 --- a/src/service/metrics/prom.rs +++ b/src/service/metrics/prom.rs @@ -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 = AHashMap::new(); let mut stream_alerts_map: AHashMap> = AHashMap::new(); let mut stream_trigger_map: AHashMap = AHashMap::new(); + let mut stream_transform_map: AHashMap> = 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(()) } diff --git a/src/service/mod.rs b/src/service/mod.rs index 470f453a5d..cff66334e8 100644 --- a/src/service/mod.rs +++ b/src/service/mod.rs @@ -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 diff --git a/src/service/traces/otlp_http.rs b/src/service/traces/otlp_http.rs index 3bfdca08c2..00e1620dec 100644 --- a/src/service/traces/otlp_http.rs +++ b/src/service/traces/otlp_http.rs @@ -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>> = AHashMap::new(); @@ -107,9 +114,9 @@ pub async fn traces_json( let mut trigger: Option = 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 diff --git a/src/service/usage/mod.rs b/src/service/usage/mod.rs new file mode 100644 index 0000000000..3b66bbb106 --- /dev/null +++ b/src/service/usage/mod.rs @@ -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>>> = + 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) { + 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; + }); + } +}