refactor: extract byte_size during json deserialization

Author: Devdutt Shenoi

src/connectors/kafka/processor.rs

@@ -36,7 +36,7 @@ use super::{config::BufferConfig, ConsumerRecord, StreamConsumer, TopicPartition
 pub struct ParseableSinkProcessor;
 
 impl ParseableSinkProcessor {
-    async fn process_event_from_chunk(&self, records: &[ConsumerRecord]) -> anyhow::Result<u64> {
+    async fn process_event_from_chunk(&self, records: &[ConsumerRecord]) -> anyhow::Result<usize> {
         let stream_name = records
             .first()
             .map(|r| r.topic.as_str())
@@ -47,10 +47,10 @@ impl ParseableSinkProcessor {
             .await?;
 
         let mut json_vec = Vec::with_capacity(records.len());
-        let mut total_payload_size = 0u64;
+        let mut total_payload_size = 0;
 
         for record in records.iter().filter_map(|r| r.payload.as_ref()) {
-            total_payload_size += record.len() as u64;
+            total_payload_size += record.len();
             if let Ok(value) = serde_json::from_slice::<Value>(record) {
                 json_vec.push(value);
             }
@@ -60,6 +60,7 @@ impl ParseableSinkProcessor {
             .get_or_create_stream(stream_name)
             .push_logs(
                 Value::Array(json_vec),
+                total_payload_size,
                 &LogSource::Custom("Kafka".to_owned()),
             )
             .await?;

src/event/format/json.rs

@@ -225,7 +225,7 @@ impl EventFormat for Event {
     /// Converts a JSON event into a Parseable Event
     fn into_event(
         self,
-        origin_size: u64,
+        origin_size: usize,
         stream: &Stream,
         log_source: &LogSource,
     ) -> anyhow::Result<super::Event> {

src/event/format/mod.rs

@@ -167,7 +167,7 @@ pub trait EventFormat: Sized {
 
     fn into_event(
         self,
-        origin_size: u64,
+        origin_size: usize,
         stream: &Stream,
         log_source: &LogSource,
     ) -> Result<Event, AnyError>;

src/event/mod.rs

@@ -39,7 +39,7 @@ pub struct PartitionEvent {
 
 pub struct Event {
     pub origin_format: &'static str,
-    pub origin_size: u64,
+    pub origin_size: usize,
     pub is_first_event: bool,
     pub time_partition: Option<String>,
     pub partitions: HashMap<String, PartitionEvent>,

src/handlers/http/cluster/mod.rs

@@ -37,7 +37,7 @@ use tracing::{error, info, warn};
 use url::Url;
 use utils::{check_liveness, to_url_string, IngestionStats, QueriedStats, StorageStats};
 
-use crate::handlers::http::ingest::ingest_internal_stream;
+use crate::event::format::LogSource;
 use crate::metrics::prom_utils::Metrics;
 use crate::parseable::PARSEABLE;
 use crate::rbac::role::model::DefaultPrivilege;
@@ -774,29 +774,29 @@ pub fn init_cluster_metrics_schedular() -> Result<(), PostError> {
     scheduler
         .every(CLUSTER_METRICS_INTERVAL_SECONDS)
         .run(move || async {
+            let internal_stream = PARSEABLE.get_or_create_stream(INTERNAL_STREAM_NAME);
             let result: Result<(), PostError> = async {
                 let cluster_metrics = fetch_cluster_metrics().await;
-                if let Ok(metrics) = cluster_metrics {
-                    if !metrics.is_empty() {
-                        info!("Cluster metrics fetched successfully from all ingestors");
-                        if let Ok(metrics_bytes) = serde_json::to_vec(&metrics) {
-                            if matches!(
-                                ingest_internal_stream(
-                                    INTERNAL_STREAM_NAME.to_string(),
-                                    bytes::Bytes::from(metrics_bytes),
-                                )
-                                .await,
-                                Ok(())
-                            ) {
-                                info!("Cluster metrics successfully ingested into internal stream");
-                            } else {
-                                error!("Failed to ingest cluster metrics into internal stream");
-                            }
-                        } else {
-                            error!("Failed to serialize cluster metrics");
-                        }
+                let Ok(metrics) = cluster_metrics else {
+                    return Ok(());
+                };
+                if !metrics.is_empty() {
+                    info!("Cluster metrics fetched successfully from all ingestors");
+                    let json = serde_json::to_value(&metrics).expect("should be json serializable");
+                    let byte_size = serde_json::to_vec(&metrics).unwrap().len();
+
+                    if matches!(
+                        internal_stream
+                            .push_logs(json, byte_size, &LogSource::Pmeta)
+                            .await,
+                        Ok(())
+                    ) {
+                        info!("Cluster metrics successfully ingested into internal stream");
+                    } else {
+                        error!("Failed to ingest cluster metrics into internal stream");
                     }
                 }
+
                 Ok(())
             }
             .await;

src/handlers/http/cluster/utils.rs

@@ -16,11 +16,23 @@
  *
  */
 
-use crate::{handlers::http::base_path_without_preceding_slash, HTTP_CLIENT};
-use actix_web::http::header;
+use std::{future::Future, pin::Pin};
+
+use crate::{
+    handlers::http::{base_path_without_preceding_slash, MAX_EVENT_PAYLOAD_SIZE},
+    HTTP_CLIENT,
+};
+use actix_web::{
+    dev::Payload,
+    error::{ErrorPayloadTooLarge, JsonPayloadError},
+    http::header,
+    FromRequest, HttpRequest,
+};
+use bytes::BytesMut;
 use chrono::{DateTime, Utc};
+use futures::StreamExt;
 use itertools::Itertools;
-use serde::{Deserialize, Serialize};
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use tracing::error;
 use url::Url;
 
@@ -248,3 +260,49 @@ pub fn to_url_string(str: String) -> String {
 
     format!("http://{}/", str)
 }
+
+pub struct JsonWithSize<T> {
+    pub json: T,
+    pub byte_size: usize,
+}
+
+impl<T: DeserializeOwned + 'static> FromRequest for JsonWithSize<T> {
+    type Error = actix_web::error::Error;
+    type Future = Pin<Box<dyn Future<Output = Result<Self, Self::Error>>>>;
+
+    fn from_request(_: &HttpRequest, payload: &mut Payload) -> Self::Future {
+        let limit = MAX_EVENT_PAYLOAD_SIZE;
+
+        // Take ownership of payload for async processing
+        let mut payload = payload.take();
+
+        Box::pin(async move {
+            // Buffer to collect all bytes
+            let mut body = BytesMut::new();
+            let mut byte_size = 0;
+
+            // Collect all bytes from the payload stream
+            while let Some(chunk) = payload.next().await {
+                let chunk = chunk?;
+                byte_size += chunk.len();
+
+                // Check the size limit
+                if byte_size > limit {
+                    return Err(ErrorPayloadTooLarge(byte_size).into());
+                }
+
+                // Extend our buffer with the chunk
+                body.extend_from_slice(&chunk);
+            }
+
+            // Convert the collected bytes to Bytes
+            let bytes = body.freeze();
+
+            // Deserialize the JSON payload
+            let json = serde_json::from_slice::<T>(&bytes)
+                .map_err(|e| JsonPayloadError::Deserialize(e))?;
+
+            Ok(JsonWithSize { json, byte_size })
+        })
+    }
+}

src/handlers/http/ingest.rs

@@ -18,10 +18,9 @@
 
 use std::collections::HashMap;
 
-use actix_web::web::{Json, Path};
+use actix_web::web::Path;
 use actix_web::{http::header::ContentType, HttpRequest, HttpResponse};
 use arrow_array::RecordBatch;
-use bytes::Bytes;
 use chrono::Utc;
 use http::StatusCode;
 use serde_json::Value;
@@ -36,14 +35,18 @@ use crate::storage::{ObjectStorageError, StreamType};
 use crate::utils::header_parsing::ParseHeaderError;
 use crate::utils::json::flatten::JsonFlattenError;
 
+use super::cluster::utils::JsonWithSize;
 use super::logstream::error::{CreateStreamError, StreamError};
 use super::users::dashboards::DashboardError;
 use super::users::filters::FiltersError;
 
 // Handler for POST /api/v1/ingest
 // ingests events by extracting stream name from header
 // creates if stream does not exist
-pub async fn ingest(req: HttpRequest, Json(json): Json<Value>) -> Result<HttpResponse, PostError> {
+pub async fn ingest(
+    req: HttpRequest,
+    JsonWithSize { json, byte_size }: JsonWithSize<Value>,
+) -> Result<HttpResponse, PostError> {
     let Some(stream_name) = req.headers().get(STREAM_NAME_HEADER_KEY) else {
         return Err(PostError::Header(ParseHeaderError::MissingStreamName));
     };
@@ -72,29 +75,18 @@ pub async fn ingest(req: HttpRequest, Json(json): Json<Value>) -> Result<HttpRes
 
     PARSEABLE
         .get_or_create_stream(&stream_name)
-        .push_logs(json, &log_source)
+        .push_logs(json, byte_size, &log_source)
         .await?;
 
     Ok(HttpResponse::Ok().finish())
 }
 
-pub async fn ingest_internal_stream(stream_name: String, body: Bytes) -> Result<(), PostError> {
-    let json: Value = serde_json::from_slice(&body)?;
-
-    PARSEABLE
-        .get_stream(&stream_name)?
-        .push_logs(json, &LogSource::Pmeta)
-        .await?;
-
-    Ok(())
-}
-
 // Handler for POST /v1/logs to ingest OTEL logs
 // ingests events by extracting stream name from header
 // creates if stream does not exist
 pub async fn handle_otel_logs_ingestion(
     req: HttpRequest,
-    Json(json): Json<Value>,
+    JsonWithSize { json, byte_size }: JsonWithSize<Value>,
 ) -> Result<HttpResponse, PostError> {
     let Some(stream_name) = req.headers().get(STREAM_NAME_HEADER_KEY) else {
         return Err(PostError::Header(ParseHeaderError::MissingStreamName));
@@ -115,7 +107,7 @@ pub async fn handle_otel_logs_ingestion(
 
     PARSEABLE
         .get_or_create_stream(&stream_name)
-        .push_logs(json, &log_source)
+        .push_logs(json, byte_size, &log_source)
         .await?;
 
     Ok(HttpResponse::Ok().finish())
@@ -126,7 +118,7 @@ pub async fn handle_otel_logs_ingestion(
 // creates if stream does not exist
 pub async fn handle_otel_metrics_ingestion(
     req: HttpRequest,
-    Json(json): Json<Value>,
+    JsonWithSize { json, byte_size }: JsonWithSize<Value>,
 ) -> Result<HttpResponse, PostError> {
     let Some(stream_name) = req.headers().get(STREAM_NAME_HEADER_KEY) else {
         return Err(PostError::Header(ParseHeaderError::MissingStreamName));
@@ -149,7 +141,7 @@ pub async fn handle_otel_metrics_ingestion(
 
     PARSEABLE
         .get_or_create_stream(&stream_name)
-        .push_logs(json, &log_source)
+        .push_logs(json, byte_size, &log_source)
         .await?;
 
     Ok(HttpResponse::Ok().finish())
@@ -160,7 +152,7 @@ pub async fn handle_otel_metrics_ingestion(
 // creates if stream does not exist
 pub async fn handle_otel_traces_ingestion(
     req: HttpRequest,
-    Json(json): Json<Value>,
+    JsonWithSize { json, byte_size }: JsonWithSize<Value>,
 ) -> Result<HttpResponse, PostError> {
     let Some(stream_name) = req.headers().get(STREAM_NAME_HEADER_KEY) else {
         return Err(PostError::Header(ParseHeaderError::MissingStreamName));
@@ -180,7 +172,7 @@ pub async fn handle_otel_traces_ingestion(
 
     PARSEABLE
         .get_or_create_stream(&stream_name)
-        .push_logs(json, &log_source)
+        .push_logs(json, byte_size, &log_source)
         .await?;
 
     Ok(HttpResponse::Ok().finish())
@@ -192,7 +184,7 @@ pub async fn handle_otel_traces_ingestion(
 pub async fn post_event(
     req: HttpRequest,
     stream_name: Path<String>,
-    Json(json): Json<Value>,
+    JsonWithSize { json, byte_size }: JsonWithSize<Value>,
 ) -> Result<HttpResponse, PostError> {
     let stream_name = stream_name.into_inner();
 
@@ -232,7 +224,7 @@ pub async fn post_event(
 
     PARSEABLE
         .get_or_create_stream(&stream_name)
-        .push_logs(json, &log_source)
+        .push_logs(json, byte_size, &log_source)
         .await?;
 
     Ok(HttpResponse::Ok().finish())

src/metadata.rs

@@ -35,7 +35,7 @@ use crate::storage::StreamType;
 pub fn update_stats(
     stream_name: &str,
     origin: &'static str,
-    size: u64,
+    size: usize,
     num_rows: usize,
     parsed_date: NaiveDate,
 ) {

src/parseable/streams.rs

@@ -113,9 +113,12 @@ impl Stream {
         })
     }
 
-    pub async fn push_logs(&self, json: Value, log_source: &LogSource) -> anyhow::Result<()> {
-        let origin_size = serde_json::to_vec(&json).unwrap().len() as u64; // string length need not be the same as byte length
-
+    pub async fn push_logs(
+        &self,
+        json: Value,
+        origin_size: usize,
+        log_source: &LogSource,
+    ) -> anyhow::Result<()> {
         json::Event::new(json)
             .into_event(origin_size, self, log_source)?
             .process(self)?;