fix: bug in sending data more than 4MB

Eshanatnight · Eshanatnight · commit 321b0beae4f0 · 2024-05-09T13:28:55.000+05:30
Need to update the `push_logs_unchecked` function as event processing
has changed
Need to clean up the `do_get` function for airplane
diff --git a/server/src/event.rs b/server/src/event.rs
@@ -28,7 +28,7 @@ use std::sync::Arc;
 
 use self::error::EventError;
 pub use self::writer::STREAM_WRITERS;
-use crate::metadata;
+use crate::{handlers::http::ingest::PostError, metadata};
 use chrono::NaiveDateTime;
 
 pub const DEFAULT_TIMESTAMP_KEY: &str = "p_timestamp";
@@ -86,12 +86,18 @@ impl Event {
         Ok(())
     }
 
-    pub fn process_unchecked(&self) -> Result<(), EventError> {
+    pub fn process_unchecked(self) -> Result<Self, PostError> {
         let key = get_schema_key(&self.rb.schema().fields);
 
-        Self::process_event(&self.stream_name, &key, self.rb.clone(), self.parsed_timestamp)?;
+        Self::process_event(
+            &self.stream_name,
+            &key,
+            self.rb.clone(),
+            self.parsed_timestamp,
+        )
+        .map_err(PostError::Event)?;
 
-        Ok(())
+        Ok(self)
     }
 
     pub fn clear(&self, stream_name: &str) {
diff --git a/server/src/event/writer.rs b/server/src/event/writer.rs
@@ -22,7 +22,7 @@ mod mem_writer;
 
 use std::{
     collections::HashMap,
-    sync::{Arc, Mutex, RwLock},
+    sync::{Arc, Mutex, RwLock, RwLockWriteGuard},
 };
 
 use crate::{
@@ -67,7 +67,6 @@ impl Writer {
         Ok(())
     }
 
-
     fn push_mem(&mut self, schema_key: &str, rb: RecordBatch) -> Result<(), StreamWriterError> {
         self.mem.push(schema_key, rb);
         Ok(())
@@ -90,61 +89,97 @@ impl WriterTable {
 
         match hashmap_guard.get(stream_name) {
             Some(stream_writer) => {
-                if CONFIG.parseable.mode != Mode::Query {
-                    stream_writer.lock().unwrap().push(
-                        stream_name,
-                        schema_key,
-                        record,
+                self.handle_existing_writer(
+                    stream_writer,
+                    stream_name,
+                    schema_key,
+                    record,
                     parsed_timestamp,
                 )?;
-                } else {
-                    stream_writer
-                        .lock()
-                        .unwrap()
-                        .push_mem(stream_name, record)?;
-                }
             }
             None => {
                 drop(hashmap_guard);
-                let mut map = self.write().unwrap();
+                let map = self.write().unwrap();
                 // check for race condition
                 // if map contains entry then just
-                if let Some(writer) = map.get(stream_name) {
-                    if CONFIG.parseable.mode != Mode::Query {
-                        writer.lock().unwrap().push(
-                            stream_name,
-                            schema_key,
-                            record,
+                self.handle_missing_writer(
+                    map,
+                    stream_name,
+                    schema_key,
+                    record,
+                    parsed_timestamp,
+                )?;
+            }
+        };
+        Ok(())
+    }
+
+    fn handle_existing_writer(
+        &self,
+        stream_writer: &Mutex<Writer>,
+        stream_name: &str,
+        schema_key: &str,
+        record: RecordBatch,
+        parsed_timestamp: NaiveDateTime,
+    ) -> Result<(), StreamWriterError> {
+        if CONFIG.parseable.mode != Mode::Query {
+            stream_writer.lock().unwrap().push(
+                stream_name,
+                schema_key,
+                record,
+                parsed_timestamp,
+            )?;
+        } else {
+            stream_writer
+                .lock()
+                .unwrap()
+                .push_mem(stream_name, record)?;
+        }
+
+        Ok(())
+    }
+
+    fn handle_missing_writer(
+        &self,
+        mut map: RwLockWriteGuard<HashMap<String, Mutex<Writer>>>,
+        stream_name: &str,
+        schema_key: &str,
+        record: RecordBatch,
+        parsed_timestamp: NaiveDateTime,
+    ) -> Result<(), StreamWriterError> {
+        match map.get(stream_name) {
+            Some(writer) => {
+                if CONFIG.parseable.mode != Mode::Query {
+                    writer.lock().unwrap().push(
+                        stream_name,
+                        schema_key,
+                        record,
                         parsed_timestamp,
                     )?;
-                    } else {
-                        writer.lock().unwrap().push_mem(stream_name, record)?;
-                    }
-                } else if CONFIG.parseable.mode != Mode::Query {
+                } else {
+                    writer.lock().unwrap().push_mem(stream_name, record)?;
+                }
+            }
+            None => {
+                if CONFIG.parseable.mode != Mode::Query {
                     let mut writer = Writer::default();
                     writer.push(stream_name, schema_key, record, parsed_timestamp)?;
                     map.insert(stream_name.to_owned(), Mutex::new(writer));
                 } else {
                     let mut writer = Writer::default();
-                    writer.push(stream_name, schema_key, record, parsed_timestamp)?;
+                    writer.push_mem(schema_key, record)?;
                     map.insert(stream_name.to_owned(), Mutex::new(writer));
                 }
             }
-        };
+        }
         Ok(())
     }
 
     pub fn clear(&self, stream_name: &str) {
         let map = self.write().unwrap();
         if let Some(writer) = map.get(stream_name) {
             let w = &mut writer.lock().unwrap().mem;
-            dbg!(&w.read_buffer.len());
-            dbg!(&w.mutable_buffer.inner.len());
-
-            dbg!(&w.mutable_buffer.inner);
             w.clear();
-            dbg!(&w.read_buffer.len());
-            dbg!(&w.mutable_buffer.inner.len());
         }
     }
 
diff --git a/server/src/handlers/airplane.rs b/server/src/handlers/airplane.rs
@@ -1,10 +1,9 @@
 use arrow_array::RecordBatch;
+use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use arrow_flight::PollInfo;
 use arrow_schema::{ArrowError, Schema};
-use arrow_select::concat::concat_batches;
 use chrono::Utc;
-use crossterm::event;
 use datafusion::common::tree_node::TreeNode;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -19,7 +18,6 @@ use tonic_web::GrpcWebLayer;
 use crate::event::commit_schema;
 use crate::handlers::http::cluster::get_ingestor_info;
 use crate::handlers::http::fetch_schema;
-use crate::handlers::http::ingest::push_logs_unchecked;
 use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::option::{Mode, CONFIG};
 
@@ -28,14 +26,14 @@ use crate::handlers::livetail::cross_origin_config;
 use crate::handlers::http::query::{authorize_and_set_filter_tags, into_query};
 use crate::query::{TableScanVisitor, QUERY_SESSION};
 use crate::storage::object_storage::commit_schema_to_storage;
-use crate::utils::arrow::flight::{get_query_from_ticket, run_do_get_rpc};
+use crate::utils::arrow::flight::{append_temporary_events, get_query_from_ticket, run_do_get_rpc};
 use arrow_flight::{
     flight_service_server::FlightService, Action, ActionType, Criteria, Empty, FlightData,
     FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, SchemaAsIpc,
     SchemaResult, Ticket,
 };
-use arrow_ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
-use futures::stream::BoxStream;
+use arrow_ipc::writer::IpcWriteOptions;
+use futures::{stream, TryStreamExt};
 use tonic::{Request, Response, Status, Streaming};
 
 use crate::handlers::livetail::extract_session_key;
@@ -50,13 +48,13 @@ pub struct AirServiceImpl {}
 
 #[tonic::async_trait]
 impl FlightService for AirServiceImpl {
-    type HandshakeStream = BoxStream<'static, Result<HandshakeResponse, Status>>;
-    type ListFlightsStream = BoxStream<'static, Result<FlightInfo, Status>>;
-    type DoGetStream = BoxStream<'static, Result<FlightData, Status>>;
-    type DoPutStream = BoxStream<'static, Result<PutResult, Status>>;
-    type DoActionStream = BoxStream<'static, Result<arrow_flight::Result, Status>>;
-    type ListActionsStream = BoxStream<'static, Result<ActionType, Status>>;
-    type DoExchangeStream = BoxStream<'static, Result<FlightData, Status>>;
+    type HandshakeStream = stream::BoxStream<'static, Result<HandshakeResponse, Status>>;
+    type ListFlightsStream = stream::BoxStream<'static, Result<FlightInfo, Status>>;
+    type DoGetStream = stream::BoxStream<'static, Result<FlightData, Status>>;
+    type DoPutStream = stream::BoxStream<'static, Result<PutResult, Status>>;
+    type DoActionStream = stream::BoxStream<'static, Result<arrow_flight::Result, Status>>;
+    type ListActionsStream = stream::BoxStream<'static, Result<ActionType, Status>>;
+    type DoExchangeStream = stream::BoxStream<'static, Result<FlightData, Status>>;
 
     async fn handshake(
         &self,
@@ -175,24 +173,27 @@ impl FlightService for AirServiceImpl {
                 let mut batches = run_do_get_rpc(im, sql.clone()).await?;
                 minute_result.append(&mut batches);
             }
-            let mr = minute_result.iter().map(|rb| rb).collect::<Vec<_>>();
-            let schema = STREAM_INFO
-                .schema(&stream_name)
-                .map_err(|err| Status::failed_precondition(format!("Metadata Error: {}", err)))?;
-            let rb = concat_batches(&schema, mr)
-                .map_err(|err| Status::failed_precondition(format!("ArrowError: {}", err)))?;
-
-            let event = push_logs_unchecked(rb, &stream_name)
-                .await
-                .map_err(|err| Status::internal(err.to_string()))?;
-            let mut events = vec![];
-            for batch in minute_result {
-                events.push(
-                    push_logs_unchecked(batch, &stream_name)
-                        .await
-                        .map_err(|err| Status::internal(err.to_string()))?,
-                );
-            }
+            let mr = minute_result.iter().collect::<Vec<_>>();
+            dbg!(&mr.len());
+            let event = append_temporary_events(&stream_name, mr).await?;
+
+            // let schema = STREAM_INFO
+            //     .schema(&stream_name)
+            //     .map_err(|err| Status::failed_precondition(format!("Metadata Error: {}", err)))?;
+            // let rb = concat_batches(&schema, mr)
+            //     .map_err(|err| Status::failed_precondition(format!("ArrowError: {}", err)))?;
+            //
+            // let event = push_logs_unchecked(rb, &stream_name)
+            //     .await
+            //     .map_err(|err| Status::internal(err.to_string()))?;
+            // let mut events = vec![];
+            // for batch in minute_result {
+            //     events.push(
+            //         push_logs_unchecked(batch, &stream_name)
+            //             .await
+            //             .map_err(|err| Status::internal(err.to_string()))?,
+            //     );
+            // }
             Some(event)
         } else {
             None
@@ -210,28 +211,39 @@ impl FlightService for AirServiceImpl {
             .await
             .map_err(|err| Status::internal(err.to_string()))
             .unwrap();
+        if results.len() > 1 {
+            dbg!(&results.len());
+        }
 
         let schemas = results
             .iter()
             .map(|batch| batch.schema())
             .map(|s| s.as_ref().clone())
             .collect::<Vec<_>>();
-
         let schema = Schema::try_merge(schemas).map_err(|err| Status::internal(err.to_string()))?;
-        let options = datafusion::arrow::ipc::writer::IpcWriteOptions::default();
-        let schema_flight_data = SchemaAsIpc::new(&schema, &options);
-
-        let mut flights = vec![FlightData::from(schema_flight_data)];
-        let encoder = IpcDataGenerator::default();
-        let mut tracker = DictionaryTracker::new(false);
-        for batch in &results {
-            let (flight_dictionaries, flight_batch) = encoder
-                .encoded_batch(batch, &mut tracker, &options)
-                .map_err(|e| Status::internal(e.to_string()))?;
-            flights.extend(flight_dictionaries.into_iter().map(Into::into));
-            flights.push(flight_batch.into());
-        }
-        let output = futures::stream::iter(flights.into_iter().map(Ok));
+        let input_stream = futures::stream::iter(results.into_iter().map(Ok));
+
+        let flight_data_stream = FlightDataEncoderBuilder::new()
+            .with_max_flight_data_size(usize::MAX)
+            .with_schema(schema.into())
+            .build(input_stream);
+
+        let flight_data_stream = flight_data_stream.map_err(|err| Status::unknown(err.to_string()));
+
+        // let options = datafusion::arrow::ipc::writer::IpcWriteOptions::default();
+        // let schema_flight_data = SchemaAsIpc::new(&schema, &options);
+        //
+        // let mut flights = vec![FlightData::from(schema_flight_data)];
+        // let encoder = IpcDataGenerator::default();
+        // let mut tracker = DictionaryTracker::new(false);
+        // for batch in &results {
+        //     let (flight_dictionaries, flight_batch) = encoder
+        //         .encoded_batch(batch, &mut tracker, &options)
+        //         .map_err(|e| Status::internal(e.to_string()))?;
+        //     flights.extend(flight_dictionaries.into_iter().map(Into::into));
+        //     flights.push(flight_batch.into());
+        // }
+        // let output = futures::stream::iter(flights.into_iter().map(Ok));
         if let Some(events) = events {
             events.clear(&stream_name);
             // for event in events {
@@ -244,7 +256,9 @@ impl FlightService for AirServiceImpl {
             .with_label_values(&[&format!("flight-query-{}", stream_name)])
             .observe(time);
 
-        Ok(Response::new(Box::pin(output) as Self::DoGetStream))
+        Ok(Response::new(
+            Box::pin(flight_data_stream) as Self::DoGetStream
+        ))
     }
 
     async fn do_put(
@@ -297,9 +311,8 @@ pub fn server() -> impl Future<Output = Result<(), Box<dyn std::error::Error + S
     let svc = FlightServiceServer::new(service)
         .max_encoding_message_size(usize::MAX)
         .max_decoding_message_size(usize::MAX)
-        .send_compressed(CompressionEncoding::Gzip);
-
-    dbg!(&svc);
+        .send_compressed(CompressionEncoding::Zstd)
+        .accept_compressed(CompressionEncoding::Zstd);
 
     let cors = cross_origin_config();
 
diff --git a/server/src/handlers/http/ingest.rs b/server/src/handlers/http/ingest.rs
@@ -102,17 +102,16 @@ pub async fn push_logs_unchecked(
     batches: RecordBatch,
     stream_name: &str,
 ) -> Result<event::Event, PostError> {
-    todo!("timepartition fix");
-
-    // let event = event::Event {
-    //     rb: batches,
-    //     stream_name: stream_name.to_string(),
-    //     origin_format: "json",
-    //     origin_size: 0,
-    //     is_first_event: true,
-    // };
-    // event.process_unchecked()?;
-    // Ok(event)
+    event::Event {
+        rb: batches,
+        stream_name: stream_name.to_string(),
+        origin_format: "json",
+        origin_size: 0,
+        parsed_timestamp: Utc::now().naive_utc(),
+        time_partition: None,
+        is_first_event: true, // NOTE: Maybe should be false
+    }
+    .process_unchecked()
 }
 
 async fn push_logs(stream_name: String, req: HttpRequest, body: Bytes) -> Result<(), PostError> {
diff --git a/server/src/handlers/http/query.rs b/server/src/handlers/http/query.rs
@@ -99,6 +99,9 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Respon
     let time = Instant::now();
 
     let (records, fields) = query.execute(table_name.clone()).await?;
+    if records.len() > 1 {
+        dbg!(&records.len());
+    }
     let response = QueryResponse {
         records,
         fields,
diff --git a/server/src/utils/arrow/flight.rs b/server/src/utils/arrow/flight.rs
