sigterm fixes

Author: AdheipSingh

server/src/handlers/http/health_check.rs

@@ -50,23 +50,13 @@ pub async fn handle_signals(shutdown_signal: Arc<Mutex<Option<oneshot::Sender<()
             let mut shutdown_flag = signal_received.lock().await;
             *shutdown_flag = true;
 
-            // Trigger graceful shutdown
-            if let Some(shutdown_sender) = shutdown_signal.lock().await.take() {
-                let _ = shutdown_sender.send(());
-            }
-
-            // Delay to allow readiness probe to return SERVICE_UNAVAILABLE
-            let _ = sleep(Duration::from_secs(20)).await;
-
             // Sync to local
             crate::event::STREAM_WRITERS.unset_all();
 
-            // Sync to S3
-            if let Err(e) = CONFIG.storage().get_object_store().sync().await {
-                log::warn!("Failed to sync local data with object store. {:?}", e);
+            // Trigger graceful shutdown
+            if let Some(shutdown_sender) = shutdown_signal.lock().await.take() {
+                let _ = shutdown_sender.send(());
             }
-
-            log::info!("Local and S3 Sync done, handler SIGTERM completed.");
         }
         None => {
             log::info!("Signal handler received None, indicating an error or end of stream");

server/src/handlers/http/modal/server.rs

@@ -40,6 +40,7 @@ use crate::users::dashboards::DASHBOARDS;
 use crate::users::filters::FILTERS;
 use std::sync::Arc;
 use tokio::sync::{oneshot, Mutex};
+use tokio::time::{sleep, Duration};
 
 use actix_web::web::resource;
 use actix_web::Resource;
@@ -64,7 +65,6 @@ use super::generate;
 use super::ssl_acceptor::get_ssl_acceptor;
 use super::OpenIdClient;
 use super::ParseableServer;
-
 #[derive(Default)]
 pub struct Server;
 
@@ -108,9 +108,9 @@ impl ParseableServer for Server {
         let shutdown_signal = server_shutdown_signal.clone();
 
         // Spawn the signal handler task
-        tokio::spawn(async move {
+        let signal_task = tokio::spawn(async move {
             health_check::handle_signals(shutdown_signal).await;
-            println!("Received shutdown signal, notifying server to shut down...");
+            log::info!("Received shutdown signal, notifying server to shut down...");
         });
 
         // Create the HTTP server
@@ -129,18 +129,41 @@ impl ParseableServer for Server {
 
         // Graceful shutdown handling
         let srv_handle = srv.handle();
-
-        tokio::spawn(async move {
+        
+        let sync_task = tokio::spawn(async move {
             // Wait for the shutdown signal
-            shutdown_rx.await.ok();
+            let _ = shutdown_rx.await;
+
+            // Perform S3 sync and wait for completion
+            log::info!("Starting data sync to S3...");
+            if let Err(e) = CONFIG.storage().get_object_store().sync(true).await {
+                log::warn!("Failed to sync local data with object store. {:?}", e);
+            } else {
+                log::info!("Successfully synced all data to S3.");
+            }
 
             // Initiate graceful shutdown
             log::info!("Graceful shutdown of HTTP server triggered");
             srv_handle.stop(true).await;
         });
 
-        // Await the server to run and handle shutdown
-        srv.await?;
+        // Await the HTTP server to run
+        let server_result = srv.await;
+
+        // Await the signal handler to ensure proper cleanup
+        if let Err(e) = signal_task.await {
+            log::error!("Error in signal handler: {:?}", e);
+        }
+
+        // Wait for the sync task to complete before exiting
+        if let Err(e) = sync_task.await {
+            log::error!("Error in sync task: {:?}", e);
+        } else {
+            log::info!("Sync task completed successfully.");
+        }
+
+        // Return the result of the server
+        server_result?;
 
         Ok(())
     }

server/src/main.rs

@@ -52,7 +52,7 @@ use option::{Mode, CONFIG};
 use crate::handlers::http::modal::{
     ingest_server::IngestServer, query_server::QueryServer, server::Server,
 };
-pub const STORAGE_UPLOAD_INTERVAL: u32 = 60;
+pub const STORAGE_UPLOAD_INTERVAL: u32 = 600;
 
 #[actix_web::main]
 async fn main() -> anyhow::Result<()> {

server/src/storage/object_storage.rs

@@ -426,16 +426,16 @@ pub trait ObjectStorage: Sync + 'static {
             .await
     }
 
-    async fn sync(&self) -> Result<(), ObjectStorageError> {
+    async fn sync(&self, shutdown_signal: bool) -> Result<(), ObjectStorageError> {
         if !Path::new(&CONFIG.staging_dir()).exists() {
             return Ok(());
         }
-
+    
         let streams = STREAM_INFO.list_streams();
-
+    
         let cache_manager = LocalCacheManager::global();
         let mut cache_updates: HashMap<&String, Vec<_>> = HashMap::new();
-
+    
         for stream in &streams {
             let cache_enabled = STREAM_INFO
                 .get_cache_enabled(stream)
@@ -452,9 +452,10 @@ pub trait ObjectStorage: Sync + 'static {
                 &dir,
                 time_partition,
                 custom_partition.clone(),
+                shutdown_signal,
             )
             .map_err(|err| ObjectStorageError::UnhandledError(Box::new(err)))?;
-
+    
             if let Some(schema) = schema {
                 let static_schema_flag = STREAM_INFO
                     .get_static_schema_flag(stream)
@@ -463,14 +464,18 @@ pub trait ObjectStorage: Sync + 'static {
                     commit_schema_to_storage(stream, schema).await?;
                 }
             }
+    
             let parquet_files = dir.parquet_files();
-
             for file in parquet_files {
                 let filename = file
                     .file_name()
                     .expect("only parquet files are returned by iterator")
                     .to_str()
                     .expect("filename is valid string");
+    
+                // Log the filename being processed
+                log::debug!("Processing file: {}", filename);
+    
                 let mut file_date_part = filename.split('.').collect::<Vec<&str>>()[0];
                 file_date_part = file_date_part.split('=').collect::<Vec<&str>>()[1];
                 let compressed_size = file.metadata().map_or(0, |meta| meta.len());
@@ -484,7 +489,7 @@ pub trait ObjectStorage: Sync + 'static {
                     .with_label_values(&["data", stream, "parquet"])
                     .add(compressed_size as i64);
                 let mut file_suffix = str::replacen(filename, ".", "/", 3);
-
+    
                 let custom_partition_clone = custom_partition.clone();
                 if custom_partition_clone.is_some() {
                     let custom_partition_fields = custom_partition_clone.unwrap();
@@ -493,8 +498,15 @@ pub trait ObjectStorage: Sync + 'static {
                     file_suffix =
                         str::replacen(filename, ".", "/", 3 + custom_partition_list.len());
                 }
+    
                 let stream_relative_path = format!("{stream}/{file_suffix}");
-                self.upload_file(&stream_relative_path, &file).await?;
+    
+                // Try uploading the file, handle potential errors without breaking the loop
+                if let Err(e) = self.upload_file(&stream_relative_path, &file).await {
+                    log::error!("Failed to upload file {}: {:?}", filename, e);
+                    continue; // Skip to the next file
+                }
+    
                 let absolute_path = self
                     .absolute_url(RelativePath::from_path(&stream_relative_path).unwrap())
                     .to_string();
@@ -512,28 +524,30 @@ pub trait ObjectStorage: Sync + 'static {
                 }
             }
         }
-
+    
+        // Cache management logic
         if let Some(manager) = cache_manager {
             let cache_updates = cache_updates
                 .into_iter()
                 .map(|(key, value)| (key.to_owned(), value))
                 .collect_vec();
-
+    
             tokio::spawn(async move {
                 for (stream, files) in cache_updates {
                     for (storage_path, file) in files {
-                        manager
+                        if let Err(e) = manager
                             .move_to_cache(&stream, storage_path, file.to_owned())
-                            .await
-                            .unwrap()
+                            .await {
+                            log::error!("Failed to move file to cache: {:?}", e);
+                        }
                     }
                 }
             });
         }
-
+    
         Ok(())
     }
-
+    
     // pick a better name
     fn get_bucket_name(&self) -> String;
 }

server/src/storage/staging.rs

@@ -155,17 +155,21 @@ impl StorageDir {
         &self,
         exclude: NaiveDateTime,
         stream: &str,
+        shutdown_signal: bool,
     ) -> HashMap<PathBuf, Vec<PathBuf>> {
         let mut grouped_arrow_file: HashMap<PathBuf, Vec<PathBuf>> = HashMap::new();
         let mut arrow_files = self.arrow_files();
-        arrow_files.retain(|path| {
-            !path
-                .file_name()
-                .unwrap()
-                .to_str()
-                .unwrap()
-                .starts_with(&exclude.format("%Y%m%dT%H%M").to_string())
-        });
+
+        if !shutdown_signal {
+            arrow_files.retain(|path| {
+                path.file_name()
+                    .unwrap()
+                    .to_str()
+                    .unwrap()
+                    .starts_with(&exclude.format("%Y%m%dT%H%M").to_string())
+            });
+        }
+
         let random_string =
             rand::distributions::Alphanumeric.sample_string(&mut rand::thread_rng(), 15);
         for arrow_file_path in arrow_files {
@@ -223,11 +227,12 @@ pub fn convert_disk_files_to_parquet(
     dir: &StorageDir,
     time_partition: Option<String>,
     custom_partition: Option<String>,
+    shutdown_signal: bool,
 ) -> Result<Option<Schema>, MoveDataError> {
     let mut schemas = Vec::new();
 
     let time = chrono::Utc::now().naive_utc();
-    let staging_files = dir.arrow_files_grouped_exclude_time(time, stream);
+    let staging_files = dir.arrow_files_grouped_exclude_time(time, stream, shutdown_signal);
     if staging_files.is_empty() {
         metrics::STAGING_FILES.with_label_values(&[stream]).set(0);
         metrics::STORAGE_SIZE

server/src/sync.rs

@@ -40,7 +40,7 @@ pub async fn object_store_sync() -> (
                 .every(STORAGE_UPLOAD_INTERVAL.seconds())
                 .plus(5u32.seconds())
                 .run(|| async {
-                    if let Err(e) = CONFIG.storage().get_object_store().sync().await {
+                    if let Err(e) = CONFIG.storage().get_object_store().sync(false).await {
                         log::warn!("failed to sync local data with object store. {:?}", e);
                     }
                 });