Merge branch 'main' into time-created

Author: Devdutt Shenoi

src/handlers/http/health_check.rs

@@ -27,7 +27,8 @@ use actix_web::{
     HttpResponse,
 };
 use http::StatusCode;
-use tokio::sync::Mutex;
+use tokio::{sync::Mutex, task::JoinSet};
+use tracing::{error, info, warn};
 
 use crate::parseable::PARSEABLE;
 
@@ -60,8 +61,29 @@ pub async fn shutdown() {
     let mut shutdown_flag = SIGNAL_RECEIVED.lock().await;
     *shutdown_flag = true;
 
+    let mut joinset = JoinSet::new();
+
     // Sync staging
-    PARSEABLE.flush_all_streams();
+    PARSEABLE.streams.flush_and_convert(&mut joinset, true);
+
+    while let Some(res) = joinset.join_next().await {
+        match res {
+            Ok(Ok(_)) => info!("Successfully converted arrow files to parquet."),
+            Ok(Err(err)) => warn!("Failed to convert arrow files to parquet. {err:?}"),
+            Err(err) => error!("Failed to join async task: {err}"),
+        }
+    }
+
+    if let Err(e) = PARSEABLE
+        .storage
+        .get_object_store()
+        .upload_files_from_staging()
+        .await
+    {
+        warn!("Failed to sync local data with object store. {:?}", e);
+    } else {
+        info!("Successfully synced all data to S3.");
+    }
 }
 
 pub async fn readiness() -> HttpResponse {

src/handlers/http/modal/mod.rs

@@ -136,26 +136,6 @@ pub trait ParseableServer {
 
             health_check::shutdown().await;
 
-            // Perform S3 sync and wait for completion
-            info!("Starting data sync to S3...");
-
-            if let Err(e) = PARSEABLE.streams.prepare_parquet(true) {
-                warn!("Failed to convert arrow files to parquet. {:?}", e);
-            } else {
-                info!("Successfully converted arrow files to parquet.");
-            }
-
-            if let Err(e) = PARSEABLE
-                .storage
-                .get_object_store()
-                .upload_files_from_staging()
-                .await
-            {
-                warn!("Failed to sync local data with object store. {:?}", e);
-            } else {
-                info!("Successfully synced all data to S3.");
-            }
-
             // Initiate graceful shutdown
             info!("Graceful shutdown of HTTP server triggered");
             srv_handle.stop(true).await;

src/lib.rs

@@ -59,8 +59,15 @@ use reqwest::{Client, ClientBuilder};
 // It is very unlikely that panic will occur when dealing with locks.
 pub const LOCK_EXPECT: &str = "Thread shouldn't panic while holding a lock";
 
-pub const STORAGE_CONVERSION_INTERVAL: u64 = 60;
-pub const STORAGE_UPLOAD_INTERVAL: u64 = 30;
+/// Describes the duration at the end of which in-memory buffers are flushed,
+/// arrows files are "finished" and compacted into parquet files.
+pub const LOCAL_SYNC_INTERVAL: Duration = Duration::from_secs(60);
+
+/// Duration used to configure prefix generation.
+pub const OBJECT_STORE_DATA_GRANULARITY: u32 = LOCAL_SYNC_INTERVAL.as_secs() as u32 / 60;
+
+/// Describes the duration at the end of which parquets are pushed into objectstore.
+pub const STORAGE_UPLOAD_INTERVAL: Duration = Duration::from_secs(30);
 
 // A single HTTP client for all outgoing HTTP requests from the parseable server
 static HTTP_CLIENT: Lazy<Client> = Lazy::new(|| {

src/parseable/mod.rs

@@ -179,16 +179,6 @@ impl Parseable {
                     .unwrap_or_default())
     }
 
-    /// Writes all streams in staging onto disk, awaiting conversion into parquet.
-    /// Deletes all in memory recordbatches, freeing up rows in mem-writer.
-    pub fn flush_all_streams(&self) {
-        let streams = self.streams.read().unwrap();
-
-        for staging in streams.values() {
-            staging.flush()
-        }
-    }
-
     // validate the storage, if the proper path for staging directory is provided
     // if the proper data directory is provided, or s3 bucket is provided etc
     pub async fn validate_storage(&self) -> Result<Option<Bytes>, ObjectStorageError> {

src/parseable/streams.rs

@@ -42,6 +42,7 @@ use parquet::{
 };
 use rand::distributions::DistString;
 use relative_path::RelativePathBuf;
+use tokio::task::JoinSet;
 use tracing::{error, info, trace, warn};
 
 use crate::{
@@ -50,11 +51,9 @@ use crate::{
     metadata::{LogStreamMetadata, SchemaVersion},
     metrics,
     option::Mode,
-    storage::{
-        object_storage::to_bytes, retention::Retention, StreamType, OBJECT_STORE_DATA_GRANULARITY,
-    },
+    storage::{object_storage::to_bytes, retention::Retention, StreamType},
     utils::minute_to_slot,
-    LOCK_EXPECT,
+    LOCK_EXPECT, OBJECT_STORE_DATA_GRANULARITY,
 };
 
 use super::{
@@ -446,21 +445,27 @@ impl Stream {
                 .set(0);
         }
 
-        // warn!("staging files-\n{staging_files:?}\n");
-        for (parquet_path, arrow_files) in staging_files {
-            metrics::STAGING_FILES
-                .with_label_values(&[&self.stream_name])
-                .set(arrow_files.len() as i64);
-
-            for file in &arrow_files {
-                let file_size = file.metadata().unwrap().len();
-                let file_type = file.extension().unwrap().to_str().unwrap();
+        //find sum of arrow files in staging directory for a stream
+        let total_arrow_files = staging_files.values().map(|v| v.len()).sum::<usize>();
+        metrics::STAGING_FILES
+            .with_label_values(&[&self.stream_name])
+            .set(total_arrow_files as i64);
 
-                metrics::STORAGE_SIZE
-                    .with_label_values(&["staging", &self.stream_name, file_type])
-                    .add(file_size as i64);
-            }
+        //find sum of file sizes of all arrow files in staging_files
+        let total_arrow_files_size = staging_files
+            .values()
+            .map(|v| {
+                v.iter()
+                    .map(|file| file.metadata().unwrap().len())
+                    .sum::<u64>()
+            })
+            .sum::<u64>();
+        metrics::STORAGE_SIZE
+            .with_label_values(&["staging", &self.stream_name, "arrows"])
+            .set(total_arrow_files_size as i64);
 
+        // warn!("staging files-\n{staging_files:?}\n");
+        for (parquet_path, arrow_files) in staging_files {
             let record_reader = MergedReverseRecordReader::try_new(&arrow_files);
             if record_reader.readers.is_empty() {
                 continue;
@@ -496,6 +501,7 @@ impl Stream {
                         "Couldn't rename part file: {part_path:?} -> {parquet_path:?}, error = {e}"
                     );
                 }
+
                 for file in arrow_files {
                     // warn!("file-\n{file:?}\n");
                     let file_size = file.metadata().unwrap().len();
@@ -655,6 +661,13 @@ impl Stream {
     pub fn get_stream_type(&self) -> StreamType {
         self.metadata.read().expect(LOCK_EXPECT).stream_type
     }
+
+    /// First flushes arrows onto disk and then converts the arrow into parquet
+    pub fn flush_and_convert(&self, shutdown_signal: bool) -> Result<(), StagingError> {
+        self.flush();
+
+        self.prepare_parquet(shutdown_signal)
+    }
 }
 
 #[derive(Deref, DerefMut, Default)]
@@ -721,21 +734,22 @@ impl Streams {
             .collect()
     }
 
-    /// Convert arrow files into parquet, preparing it for upload
-    pub fn prepare_parquet(&self, shutdown_signal: bool) -> Result<(), StagingError> {
+    /// Asynchronously flushes arrows and compacts into parquet data on all streams in staging,
+    /// so that it is ready to be pushed onto objectstore.
+    pub fn flush_and_convert(
+        &self,
+        joinset: &mut JoinSet<Result<(), StagingError>>,
+        shutdown_signal: bool,
+    ) {
         let streams: Vec<Arc<Stream>> = self
             .read()
             .expect(LOCK_EXPECT)
             .values()
             .map(Arc::clone)
             .collect();
         for stream in streams {
-            stream
-                .prepare_parquet(shutdown_signal)
-                .inspect_err(|err| error!("Failed to run conversion task {err:?}"))?;
+            joinset.spawn(async move { stream.flush_and_convert(shutdown_signal) });
         }
-
-        Ok(())
     }
 }
 

src/query/listing_table_builder.rs

@@ -32,9 +32,8 @@ use itertools::Itertools;
 use object_store::{path::Path, ObjectMeta, ObjectStore};
 
 use crate::{
-    event::DEFAULT_TIMESTAMP_KEY,
-    storage::{ObjectStorage, OBJECT_STORE_DATA_GRANULARITY},
-    utils::time::TimeRange,
+    event::DEFAULT_TIMESTAMP_KEY, storage::ObjectStorage, utils::time::TimeRange,
+    OBJECT_STORE_DATA_GRANULARITY,
 };
 
 use super::PartialTimeFilter;

src/storage/mod.rs

@@ -57,14 +57,6 @@ pub const SCHEMA_FILE_NAME: &str = ".schema";
 pub const ALERTS_ROOT_DIRECTORY: &str = ".alerts";
 pub const MANIFEST_FILE: &str = "manifest.json";
 
-/// local sync interval to move data.records to /tmp dir of that stream.
-/// 60 sec is a reasonable value.
-pub const LOCAL_SYNC_INTERVAL: u64 = 60;
-
-/// duration used to configure prefix in objectstore and local disk structure
-/// used for storage. Defaults to 1 min.
-pub const OBJECT_STORE_DATA_GRANULARITY: u32 = (LOCAL_SYNC_INTERVAL as u32) / 60;
-
 // max concurrent request allowed for datafusion object store
 const MAX_OBJECT_STORE_REQUESTS: usize = 1000;
 

src/storage/object_storage.rs

@@ -36,7 +36,8 @@ use datafusion::{datasource::listing::ListingTableUrl, execution::runtime_env::R
 use once_cell::sync::OnceCell;
 use relative_path::RelativePath;
 use relative_path::RelativePathBuf;
-use tracing::{debug, error, warn};
+use tracing::info;
+use tracing::{error, warn};
 use ulid::Ulid;
 
 use crate::alerts::AlertConfig;
@@ -712,13 +713,13 @@ pub trait ObjectStorage: Debug + Send + Sync + 'static {
     }
 
     async fn upload_files_from_staging(&self) -> Result<(), ObjectStorageError> {
-        if !Path::new(&PARSEABLE.options.staging_dir()).exists() {
+        if !PARSEABLE.options.staging_dir().exists() {
             return Ok(());
         }
 
         // get all streams
         for stream_name in PARSEABLE.streams.list() {
-            debug!("Starting object_store_sync for stream- {stream_name}");
+            info!("Starting object_store_sync for stream- {stream_name}");
 
             let stream = PARSEABLE.get_or_create_stream(&stream_name);
             let custom_partition = stream.get_custom_partition();