Fix: Sync Files on Sigterm (#969)

Co-authored-by: Nikhil Sinha <[email protected]>
parseablehq · Oct 25, 2024 · b3f4090 · b3f4090
1 parent d2df652
commit b3f4090
Show file tree

Hide file tree

Showing 9 changed files with 151 additions and 86 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/helm/templates/ingestor-statefulset.yaml b/helm/templates/ingestor-statefulset.yaml
@@ -103,4 +103,4 @@ spec:
         requests:
           storage: {{ .Values.parseable.persistence.ingestor.size | quote }}
   {{- end }}
-{{- end }}
+{{- end }}
diff --git a/server/src/handlers/http/health_check.rs b/server/src/handlers/http/health_check.rs
@@ -27,11 +27,10 @@ use lazy_static::lazy_static;
 use std::sync::Arc;
 use tokio::signal::unix::{signal, SignalKind};
 use tokio::sync::{oneshot, Mutex};
-use tokio::time::{sleep, Duration};
 
 // Create a global variable to store signal status
 lazy_static! {
-    static ref SIGNAL_RECEIVED: Arc<Mutex<bool>> = Arc::new(Mutex::new(false));
+    pub static ref SIGNAL_RECEIVED: Arc<Mutex<bool>> = Arc::new(Mutex::new(false));
 }
 
 pub async fn liveness() -> HttpResponse {
@@ -66,23 +65,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");

diff --git a/server/src/handlers/http/modal/ingest_server.rs b/server/src/handlers/http/modal/ingest_server.rs
@@ -112,9 +112,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
@@ -134,17 +134,40 @@ 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(())
     }

diff --git a/server/src/handlers/http/modal/query_server.rs b/server/src/handlers/http/modal/query_server.rs
@@ -89,8 +89,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
@@ -110,17 +111,40 @@ 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(())
     }

diff --git a/server/src/handlers/http/modal/server.rs b/server/src/handlers/http/modal/server.rs
@@ -65,7 +65,6 @@ use super::generate;
 use super::ssl_acceptor::get_ssl_acceptor;
 use super::OpenIdClient;
 use super::ParseableServer;
-
 #[derive(Default)]
 pub struct Server;
 
@@ -110,9 +109,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
@@ -132,17 +131,40 @@ 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(())
     }