fix: commit schema bug (#1666)

* fix: commit schema bug

* fix: flush pending batches to disk, schema mutex

pending batches by default get pushed to disk near instantly

schema writer is behind a mutex to prevent incorrect schemas

* fix: deepsource and coderabbit

Author: parmesant

src/event/mod.rs

@@ -86,19 +86,22 @@ impl Event {
             }
         }
 
+        let stream = PARSEABLE.get_or_create_stream(&self.stream_name, &self.tenant_id);
+
         if self.is_first_event {
             commit_schema(&self.stream_name, self.rb.schema(), &self.tenant_id)?;
+            if !stream.get_static_schema_flag() {
+                stream.stage_schema_file(stream.get_schema().as_ref().clone())?;
+            }
         }
 
-        PARSEABLE
-            .get_or_create_stream(&self.stream_name, &self.tenant_id)
-            .push(
-                &key,
-                &self.rb,
-                self.parsed_timestamp,
-                &self.custom_partition_values,
-                self.stream_type,
-            )?;
+        stream.push(
+            &key,
+            &self.rb,
+            self.parsed_timestamp,
+            &self.custom_partition_values,
+            self.stream_type,
+        )?;
 
         let tenant = self.tenant_id.as_deref().unwrap_or(DEFAULT_TENANT);
         update_stats(

src/parseable/staging/mod.rs

@@ -40,4 +40,6 @@ pub enum StagingError {
     TenantNotFound(#[from] TenantNotFound),
     #[error("{0}")]
     PoisonError(#[from] PoisonError<String>),
+    #[error("JSON Error {0}")]
+    Json(#[from] serde_json::Error),
 }

src/parseable/staging/writer.rs

@@ -29,8 +29,9 @@ use arrow_array::RecordBatch;
 use arrow_ipc::writer::StreamWriter;
 use arrow_schema::Schema;
 use arrow_select::concat::concat_batches;
-use chrono::Utc;
+use chrono::{TimeDelta, Utc};
 use itertools::Itertools;
+use once_cell::sync::Lazy;
 use rand::distributions::{Alphanumeric, DistString};
 use tracing::error;
 
@@ -41,7 +42,16 @@ use crate::{
 
 use super::StagingError;
 
-const DISK_WRITE_BATCH_ROWS: usize = 32_768;
+const DISK_WRITE_BATCH_MAX_AGE_SECS_VAR: &str = "DISK_WRITE_BATCH_MAX_AGE_SECS";
+static DISK_WRITE_BATCH_MAX_AGE_SECS: Lazy<i64> = Lazy::new(|| {
+    if let Ok(var) = std::env::var(DISK_WRITE_BATCH_MAX_AGE_SECS_VAR)
+        && let Ok(var) = var.parse::<i64>()
+    {
+        var
+    } else {
+        1
+    }
+});
 
 #[derive(Default)]
 pub struct Writer {
@@ -58,13 +68,18 @@ impl Writer {
         rb: &RecordBatch,
         file_path: PathBuf,
         range: TimeRange,
+        batch_rows: usize,
     ) -> Result<(), StagingError> {
+        let now = Utc::now();
         let pending = self.disk_pending.entry(filename.clone()).or_default();
         pending.rows += rb.num_rows();
         pending.range.get_or_insert(range);
+        pending.first_seen.get_or_insert(now);
         pending.batches.push(rb.clone());
 
-        if pending.rows >= DISK_WRITE_BATCH_ROWS {
+        let should_flush = pending.rows >= batch_rows.max(1)
+            || pending.is_older_than(now, TimeDelta::seconds(*DISK_WRITE_BATCH_MAX_AGE_SECS));
+        if should_flush {
             self.flush_pending_disk(&filename, file_path)?;
         }
 
@@ -104,8 +119,12 @@ impl Writer {
 
         let mut flushable_pending = HashMap::new();
         let old_pending = std::mem::take(&mut self.disk_pending);
+        let now = Utc::now();
         for (filename, pending) in old_pending {
-            if !forced && pending.is_current() {
+            if !forced
+                && pending.is_current()
+                && !pending.is_older_than(now, TimeDelta::seconds(*DISK_WRITE_BATCH_MAX_AGE_SECS))
+            {
                 self.disk_pending.insert(filename, pending);
             } else {
                 flushable_pending.insert(filename, pending);
@@ -164,6 +183,7 @@ struct PendingDiskBatch {
     rows: usize,
     batches: Vec<RecordBatch>,
     range: Option<TimeRange>,
+    first_seen: Option<chrono::DateTime<Utc>>,
 }
 
 impl PendingDiskBatch {
@@ -172,6 +192,11 @@ impl PendingDiskBatch {
             .as_ref()
             .is_some_and(|range| range.contains(Utc::now()))
     }
+
+    fn is_older_than(&self, now: chrono::DateTime<Utc>, age: TimeDelta) -> bool {
+        self.first_seen
+            .is_some_and(|first_seen| now.signed_duration_since(first_seen) >= age)
+    }
 }
 
 pub struct DiskWriter {

src/parseable/streams.rs

@@ -23,6 +23,7 @@ use arrow_schema::{Field, Fields, Schema};
 use chrono::{NaiveDateTime, Timelike, Utc};
 use derive_more::derive::{Deref, DerefMut};
 use itertools::Itertools;
+use once_cell::sync::Lazy;
 use parquet::{
     arrow::ArrowWriter,
     basic::Encoding,
@@ -69,6 +70,17 @@ use super::{
     staging::{StagingError, reader::MergedReverseRecordReader, writer::Writer},
 };
 
+const DISK_WRITE_BATCH_ROWS_VAR: &str = "DISK_WRITE_BATCH_ROWS";
+static DISK_WRITE_BATCH_ROWS: Lazy<usize> = Lazy::new(|| {
+    if let Ok(var) = std::env::var(DISK_WRITE_BATCH_ROWS_VAR)
+        && let Ok(var) = var.parse::<usize>()
+    {
+        var
+    } else {
+        1
+    }
+});
+
 const INPROCESS_DIR_PREFIX: &str = "processing_";
 const METRIC_ROW_GROUP_PREP_IN_FLIGHT: usize = 1;
 
@@ -115,6 +127,7 @@ pub struct Stream {
     pub data_path: PathBuf,
     pub options: Arc<Options>,
     pub writer: Mutex<Writer>,
+    schema_writer: Mutex<()>,
     pub ingestor_id: Option<String>,
 }
 impl Stream {
@@ -134,6 +147,7 @@ impl Stream {
             data_path,
             options,
             writer: Mutex::new(Writer::default()),
+            schema_writer: Mutex::new(()),
             ingestor_id,
         })
     }
@@ -181,7 +195,7 @@ impl Stream {
             );
             let file_path = self.data_path.join(&filename);
 
-            guard.push_disk(filename, record, file_path, range)?;
+            guard.push_disk(filename, record, file_path, range, *DISK_WRITE_BATCH_ROWS)?;
         }
 
         guard.mem.push(schema_key, record)?;
@@ -435,32 +449,22 @@ impl Stream {
             .collect()
     }
 
-    pub fn get_schemas_if_present(&self) -> Option<Vec<Schema>> {
-        let Ok(dir) = self.data_path.read_dir() else {
-            return None;
-        };
+    pub fn get_schemas_if_present(&self) -> Result<Vec<Schema>, StagingError> {
+        let dir = self.data_path.read_dir()?;
 
         let mut schemas: Vec<Schema> = Vec::new();
 
         for file in dir.flatten() {
             if let Some(ext) = file.path().extension()
                 && ext.eq("schema")
             {
-                let file = File::open(file.path()).expect("Schema File should exist");
+                let file = File::open(file.path())?;
 
-                let schema = match serde_json::from_reader(file) {
-                    Ok(schema) => schema,
-                    Err(_) => continue,
-                };
-                schemas.push(schema);
+                schemas.push(serde_json::from_reader(file)?);
             }
         }
 
-        if !schemas.is_empty() {
-            Some(schemas)
-        } else {
-            None
-        }
+        Ok(schemas)
     }
 
     /// Converts arrow files in staging into parquet files, does so only for past minutes when run with `!shutdown_signal`
@@ -497,28 +501,48 @@ impl Stream {
         // check if there is already a schema file in staging pertaining to this stream
         // if yes, then merge them and save
 
-        if let Some(mut schema) = schema {
+        if let Some(schema) = schema {
             let static_schema_flag = self.get_static_schema_flag();
             if !static_schema_flag {
-                // schema is dynamic, read from staging and merge if present
+                self.stage_schema_file(schema)?;
+            }
+        }
+
+        Ok(())
+    }
 
-                // need to add something before .schema to make the file have an extension of type `schema`
-                let path = RelativePathBuf::from_iter([format!("{}.schema", self.stream_name)])
-                    .to_path(&self.data_path);
+    pub fn stage_schema_file(&self, mut schema: Schema) -> Result<(), StagingError> {
+        let _schema_writer = self.schema_writer.lock().map_err(|poisoned| {
+            StagingError::PoisonError(PoisonError::new(format!(
+                "Schema writer lock poisoned while staging schema for stream {} - {}",
+                self.stream_name, poisoned
+            )))
+        })?;
 
-                let staging_schemas = self.get_schemas_if_present();
-                if let Some(mut staging_schemas) = staging_schemas {
-                    staging_schemas.push(schema);
-                    schema = Schema::try_merge(staging_schemas)?;
-                }
+        // schema is dynamic, read from staging and merge if present
+        fs::create_dir_all(&self.data_path)?;
 
-                // save the merged schema on staging disk
-                // the path should be stream/.ingestor.{id}.schema
-                info!("writing schema to path - {path:?}");
-                write(path, to_bytes(&schema))?;
-            }
+        // need to add something before .schema to make the file have an extension of type `schema`
+        let file_name = self.ingestor_id.as_ref().map_or_else(
+            || ".schema".to_owned(),
+            |id| format!(".ingestor.{id}.schema"),
+        );
+        let path = RelativePathBuf::from_iter([file_name]).to_path(&self.data_path);
+        let tmp_path = path.with_extension("schema.tmp");
+
+        let staging_schemas = self.get_schemas_if_present()?;
+        if !staging_schemas.is_empty() {
+            let mut staging_schemas = staging_schemas;
+            staging_schemas.push(schema);
+            schema = Schema::try_merge(staging_schemas)?;
         }
 
+        // save the merged schema on staging disk
+        // the path should be stream/.ingestor.{id}.schema
+        info!("writing schema to path - {path:?}");
+        write(&tmp_path, to_bytes(&schema))?;
+        fs::rename(tmp_path, path)?;
+
         Ok(())
     }