sigterm on ingest and query

Author: AdheipSingh

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

@@ -110,9 +110,9 @@ impl ParseableServer for IngestServer {
         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
@@ -131,18 +131,41 @@ impl ParseableServer for IngestServer {
 
         // 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/handlers/http/modal/query_server.rs

@@ -87,8 +87,9 @@ impl ParseableServer for QueryServer {
         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;
+            log::info!("Received shutdown signal, notifying server to shut down...");
         });
 
         // Create the HTTP server
@@ -107,18 +108,41 @@ impl ParseableServer for QueryServer {
 
         // 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(())
     }