feat: parallelize flush+conversion per stream

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/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

@@ -41,6 +41,7 @@ use parquet::{
 };
 use rand::distributions::DistString;
 use relative_path::RelativePathBuf;
+use tokio::task::JoinSet;
 use tracing::{error, info, trace, warn};
 
 use crate::{
@@ -651,6 +652,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)]
@@ -717,21 +725,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/sync.rs

@@ -190,11 +190,15 @@ pub fn local_sync() -> (
                 select! {
                     // Spawns a flush+conversion task every `LOCAL_SYNC_INTERVAL` seconds
                     _ = sync_interval.tick() => {
-                        joinset.spawn(flush_and_convert());
+                        PARSEABLE.streams.flush_and_convert(&mut joinset, false)
                     },
                     // Joins and logs errors in spawned tasks
-                    Some(Err(e)) = joinset.join_next(), if !joinset.is_empty() => {
-                        error!("Issue joining flush+conversion: {e}")
+                    Some(res) = joinset.join_next(), if !joinset.is_empty() => {
+                        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!("Issue joining flush+conversion task: {err}"),
+                        }
                     }
                     res = &mut inbox_rx => {match res{
                         Ok(_) => break,
@@ -269,18 +273,3 @@ pub fn schedule_alert_task(
     });
     Ok(handle)
 }
-
-/// Asynchronously flushes all streams when called, then compacts them into parquet files ready to be pushed onto objectstore
-async fn flush_and_convert() {
-    trace!("Flushing Arrows to disk...");
-    PARSEABLE.flush_all_streams();
-
-    trace!("Converting Arrow to Parquet... ");
-    if let Err(e) = monitor_task_duration("arrow_conversion", Duration::from_secs(30), || async {
-        PARSEABLE.streams.prepare_parquet(false)
-    })
-    .await
-    {
-        warn!("failed to convert local arrow data to parquet. {e:?}");
-    }
-}