add metrics collection for logs, traces and metrics separately

add below metrics -
1. total logs collected by date
2. total size of logs collected by date
3. total metrics collected by date
4. total size of metrics collected by date
5. total traces collected by date
6. total size of traces collected by date

add all of these to pbilling

Author: nikhilsinhaparseable

src/event/format/json.rs

@@ -31,7 +31,9 @@ use std::{collections::HashMap, sync::Arc};
 use tracing::error;
 
 use super::EventFormat;
-use crate::{metadata::SchemaVersion, storage::StreamType, utils::arrow::get_field};
+use crate::{
+    handlers::TelemetryType, metadata::SchemaVersion, storage::StreamType, utils::arrow::get_field,
+};
 
 pub struct Event {
     pub json: Value,
@@ -147,6 +149,7 @@ impl EventFormat for Event {
         schema_version: SchemaVersion,
         stream_type: StreamType,
         p_custom_fields: &HashMap<String, String>,
+        telemetry_type: TelemetryType,
     ) -> Result<super::Event, anyhow::Error> {
         let custom_partition_values = match custom_partitions.as_ref() {
             Some(custom_partition) => {
@@ -179,6 +182,7 @@ impl EventFormat for Event {
             time_partition: None,
             custom_partition_values,
             stream_type,
+            telemetry_type,
         })
     }
 }

src/event/format/mod.rs

@@ -31,6 +31,7 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
 use crate::{
+    handlers::TelemetryType,
     metadata::SchemaVersion,
     storage::StreamType,
     utils::arrow::{add_parseable_fields, get_field},
@@ -220,6 +221,7 @@ pub trait EventFormat: Sized {
         schema_version: SchemaVersion,
         stream_type: StreamType,
         p_custom_fields: &HashMap<String, String>,
+        telemetry_type: TelemetryType,
     ) -> Result<Event, AnyError>;
 }
 

src/event/mod.rs

@@ -27,6 +27,7 @@ use std::sync::Arc;
 use self::error::EventError;
 use crate::{
     LOCK_EXPECT,
+    handlers::TelemetryType,
     metadata::update_stats,
     metrics::{increment_events_ingested_by_date, increment_events_ingested_size_by_date},
     parseable::{PARSEABLE, StagingError},
@@ -52,6 +53,7 @@ pub struct Event {
     pub time_partition: Option<String>,
     pub custom_partition_values: HashMap<String, String>,
     pub stream_type: StreamType,
+    pub telemetry_type: TelemetryType,
 }
 
 // Events holds the schema related to a each event for a single log stream
@@ -92,7 +94,7 @@ impl Event {
         // Track billing metrics for event ingestion
         let date_string = self.parsed_timestamp.date().to_string();
         increment_events_ingested_by_date(self.rb.num_rows() as u64, &date_string);
-        increment_events_ingested_size_by_date(self.origin_size, &date_string);
+        increment_events_ingested_size_by_date(self.origin_size, &date_string, self.telemetry_type);
 
         crate::livetail::LIVETAIL.process(&self.stream_name, &self.rb);
 

src/handlers/http/cluster/mod.rs

@@ -100,6 +100,11 @@ struct BillingMetricsCollector {
     pub total_input_llm_tokens_by_date: HashMap<String, HashMap<String, HashMap<String, u64>>>, // provider -> model -> date -> count
     pub total_output_llm_tokens_by_date: HashMap<String, HashMap<String, HashMap<String, u64>>>,
     pub total_metrics_collected_by_date: HashMap<String, u64>,
+    pub total_metrics_collected_size_by_date: HashMap<String, u64>,
+    pub total_logs_collected_by_date: HashMap<String, u64>,
+    pub total_logs_collected_size_by_date: HashMap<String, u64>,
+    pub total_traces_collected_by_date: HashMap<String, u64>,
+    pub total_traces_collected_size_by_date: HashMap<String, u64>,
     pub event_time: chrono::NaiveDateTime,
 }
 
@@ -202,6 +207,41 @@ impl BillingMetricsCollector {
                 &self.total_metrics_collected_by_date,
             );
         }
+        if !self.total_metrics_collected_size_by_date.is_empty() {
+            add_simple_metric(
+                events,
+                "total_metrics_collected_size",
+                &self.total_metrics_collected_size_by_date,
+            );
+        }
+        if !self.total_logs_collected_by_date.is_empty() {
+            add_simple_metric(
+                events,
+                "total_logs_collected",
+                &self.total_logs_collected_by_date,
+            );
+        }
+        if !self.total_logs_collected_size_by_date.is_empty() {
+            add_simple_metric(
+                events,
+                "total_logs_collected_size",
+                &self.total_logs_collected_size_by_date,
+            );
+        }
+        if !self.total_traces_collected_by_date.is_empty() {
+            add_simple_metric(
+                events,
+                "total_traces_collected",
+                &self.total_traces_collected_by_date,
+            );
+        }
+        if !self.total_traces_collected_size_by_date.is_empty() {
+            add_simple_metric(
+                events,
+                "total_traces_collected_size",
+                &self.total_traces_collected_size_by_date,
+            );
+        }
     }
 
     /// Add object store metrics (method-based) to the events vector
@@ -1273,6 +1313,11 @@ fn is_simple_metric(metric: &str) -> bool {
             | "parseable_total_files_scanned_in_query_by_date"
             | "parseable_total_bytes_scanned_in_query_by_date"
             | "parseable_total_metrics_collected_by_date"
+            | "parseable_total_metrics_collected_size_by_date"
+            | "parseable_total_logs_collected_by_date"
+            | "parseable_total_logs_collected_size_by_date"
+            | "parseable_total_traces_collected_by_date"
+            | "parseable_total_traces_collected_size_by_date"
     )
 }
 
@@ -1344,6 +1389,31 @@ fn process_simple_metric(
                     .total_metrics_collected_by_date
                     .insert(date.to_string(), value);
             }
+            "parseable_total_metrics_collected_size_by_date" => {
+                collector
+                    .total_metrics_collected_size_by_date
+                    .insert(date.to_string(), value);
+            }
+            "parseable_total_logs_collected_by_date" => {
+                collector
+                    .total_logs_collected_by_date
+                    .insert(date.to_string(), value);
+            }
+            "parseable_total_logs_collected_size_by_date" => {
+                collector
+                    .total_logs_collected_size_by_date
+                    .insert(date.to_string(), value);
+            }
+            "parseable_total_traces_collected_by_date" => {
+                collector
+                    .total_traces_collected_by_date
+                    .insert(date.to_string(), value);
+            }
+            "parseable_total_traces_collected_size_by_date" => {
+                collector
+                    .total_traces_collected_size_by_date
+                    .insert(date.to_string(), value);
+            }
             _ => {}
         }
     }

src/handlers/http/ingest.rs

@@ -125,7 +125,15 @@ pub async fn ingest(
         .add_update_log_source(&stream_name, log_source_entry)
         .await?;
 
-    flatten_and_push_logs(json, &stream_name, &log_source, &p_custom_fields, None).await?;
+    flatten_and_push_logs(
+        json,
+        &stream_name,
+        &log_source,
+        &p_custom_fields,
+        None,
+        telemetry_type,
+    )
+    .await?;
 
     Ok(HttpResponse::Ok().finish())
 }
@@ -149,6 +157,7 @@ pub async fn ingest_internal_stream(stream_name: String, body: Bytes) -> Result<
             SchemaVersion::V0,
             StreamType::Internal,
             &p_custom_fields,
+            TelemetryType::Logs,
         )?
         .process()?;
 
@@ -235,6 +244,7 @@ async fn process_otel_content(
     body: web::Bytes,
     stream_name: &str,
     log_source: &LogSource,
+    telemetry_type: TelemetryType,
 ) -> Result<(), PostError> {
     let p_custom_fields = get_custom_fields_from_header(req);
 
@@ -251,6 +261,7 @@ async fn process_otel_content(
                     log_source,
                     &p_custom_fields,
                     None,
+                    telemetry_type,
                 )
                 .await?;
             } else if content_type == CONTENT_TYPE_PROTOBUF {
@@ -289,7 +300,7 @@ pub async fn handle_otel_logs_ingestion(
     )
     .await?;
 
-    process_otel_content(&req, body, &stream_name, &log_source).await?;
+    process_otel_content(&req, body, &stream_name, &log_source, TelemetryType::Logs).await?;
 
     Ok(HttpResponse::Ok().finish())
 }
@@ -309,7 +320,14 @@ pub async fn handle_otel_metrics_ingestion(
     )
     .await?;
 
-    process_otel_content(&req, body, &stream_name, &log_source).await?;
+    process_otel_content(
+        &req,
+        body,
+        &stream_name,
+        &log_source,
+        TelemetryType::Metrics,
+    )
+    .await?;
 
     Ok(HttpResponse::Ok().finish())
 }
@@ -329,7 +347,7 @@ pub async fn handle_otel_traces_ingestion(
     )
     .await?;
 
-    process_otel_content(&req, body, &stream_name, &log_source).await?;
+    process_otel_content(&req, body, &stream_name, &log_source, TelemetryType::Traces).await?;
 
     Ok(HttpResponse::Ok().finish())
 }
@@ -396,7 +414,15 @@ pub async fn post_event(
     //return error if the stream log source is otel traces or otel metrics
     validate_stream_for_ingestion(&stream_name)?;
 
-    flatten_and_push_logs(json, &stream_name, &log_source, &p_custom_fields, None).await?;
+    flatten_and_push_logs(
+        json,
+        &stream_name,
+        &log_source,
+        &p_custom_fields,
+        None,
+        TelemetryType::Logs,
+    )
+    .await?;
 
     Ok(HttpResponse::Ok().finish())
 }
@@ -415,6 +441,7 @@ pub async fn push_logs_unchecked(
         is_first_event: true,                    // NOTE: Maybe should be false
         custom_partition_values: HashMap::new(), // should be an empty map for unchecked push
         stream_type: StreamType::UserDefined,
+        telemetry_type: TelemetryType::Logs,
     };
     unchecked_event.process_unchecked()?;
 

src/handlers/http/modal/utils/ingest_utils.rs

@@ -32,7 +32,7 @@ use crate::{
         format::{EventFormat, LogSource, json},
     },
     handlers::{
-        EXTRACT_LOG_KEY, LOG_SOURCE_KEY, STREAM_NAME_HEADER_KEY,
+        EXTRACT_LOG_KEY, LOG_SOURCE_KEY, STREAM_NAME_HEADER_KEY, TelemetryType,
         http::{
             ingest::PostError,
             kinesis::{Message, flatten_kinesis_logs},
@@ -54,6 +54,7 @@ pub async fn flatten_and_push_logs(
     log_source: &LogSource,
     p_custom_fields: &HashMap<String, String>,
     time_partition: Option<String>,
+    telemetry_type: TelemetryType,
 ) -> Result<(), PostError> {
     // Verify the dataset fields count
     verify_dataset_fields_count(stream_name)?;
@@ -70,6 +71,7 @@ pub async fn flatten_and_push_logs(
                 log_source,
                 p_custom_fields,
                 time_partition,
+                telemetry_type,
             )
             .await?;
         }
@@ -83,6 +85,7 @@ pub async fn flatten_and_push_logs(
                     log_source,
                     p_custom_fields,
                     time_partition.clone(),
+                    telemetry_type,
                 )
                 .await?;
             }
@@ -97,6 +100,7 @@ pub async fn flatten_and_push_logs(
                     log_source,
                     p_custom_fields,
                     time_partition.clone(),
+                    telemetry_type,
                 )
                 .await?;
             }
@@ -111,6 +115,7 @@ pub async fn flatten_and_push_logs(
                     log_source,
                     p_custom_fields,
                     time_partition.clone(),
+                    telemetry_type,
                 )
                 .await?;
             }
@@ -122,6 +127,7 @@ pub async fn flatten_and_push_logs(
                 log_source,
                 p_custom_fields,
                 time_partition,
+                telemetry_type,
             )
             .await?
         }
@@ -136,6 +142,7 @@ pub async fn push_logs(
     log_source: &LogSource,
     p_custom_fields: &HashMap<String, String>,
     time_partition: Option<String>,
+    telemetry_type: TelemetryType,
 ) -> Result<(), PostError> {
     let stream = PARSEABLE.get_stream(stream_name)?;
     let time_partition_limit = PARSEABLE
@@ -169,6 +176,7 @@ pub async fn push_logs(
                 schema_version,
                 StreamType::UserDefined,
                 p_custom_fields,
+                telemetry_type,
             )?
             .process()?;
     }