Handling SIGTERM for Distributed and Standalone Mode

Dockerfile

@@ -34,4 +34,4 @@ WORKDIR /parseable
 COPY --from=builder /bin/sh /bin/sh
 COPY --from=builder /parseable/target/release/parseable /usr/bin/parseable
 
-CMD ["parseable"]
+CMD ["/usr/bin/parseable"]

Dockerfile.debug

@@ -32,4 +32,4 @@ WORKDIR /parseable
 
 COPY --from=builder /parseable/target/debug/parseable /usr/bin/parseable
 
-CMD ["parseable"]
+CMD ["/usr/bin/parseable"]

server/Cargo.toml

@@ -36,6 +36,7 @@ anyhow = { version = "1.0", features = ["backtrace"] }
 argon2 = "0.5.0"
 async-trait = "0.1"
 base64 = "0.22.0"
+lazy_static = "1.4"
 bytes = "1.4"
 byteorder = "1.4.3"
 bzip2 = { version = "*", features = ["static"] }

server/src/handlers/http/health_check.rs

@@ -1,34 +1,74 @@
-/*
- * Parseable Server (C) 2022 - 2024 Parseable, Inc.
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU Affero General Public License as
- * published by the Free Software Foundation, either version 3 of the
- * License, or (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU Affero General Public License for more details.
- *
- * You should have received a copy of the GNU Affero General Public License
- * along with this program.  If not, see <http://www.gnu.org/licenses/>.
- *
- */
-
+use crate::option::CONFIG;
 use actix_web::http::StatusCode;
 use actix_web::HttpResponse;
+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};
 
-use crate::option::CONFIG;
+// Create a global variable to store signal status
+lazy_static! {
+    static ref SIGNAL_RECEIVED: Arc<Mutex<bool>> = Arc::new(Mutex::new(false));
+}
 
 pub async fn liveness() -> HttpResponse {
     HttpResponse::new(StatusCode::OK)
 }
 
+pub async fn handle_signals(shutdown_signal: Arc<Mutex<Option<oneshot::Sender<()>>>>) {
+    let signal_received = SIGNAL_RECEIVED.clone();
+
+    let mut sigterm =
+        signal(SignalKind::terminate()).expect("Failed to set up SIGTERM signal handler");
+    log::info!("Signal handler task started");
+
+    // Block until SIGTERM is received
+    match sigterm.recv().await {
+        Some(_) => {
+            log::info!("Received SIGTERM signal at Readiness Probe Handler");
+
+            // Set the shutdown flag to true
+            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);
+            }
+
+            log::info!("Local and S3 Sync done, handler SIGTERM completed.");
+        }
+        None => {
+            println!("Signal handler received None, indicating an error or end of stream");
+        }
+    }
+
+    eprintln!("Signal handler task completed");
+}
+
 pub async fn readiness() -> HttpResponse {
-    if CONFIG.storage().get_object_store().check().await.is_ok() {
-        return HttpResponse::new(StatusCode::OK);
+    // Check if the application has received a shutdown signal
+    let shutdown_flag = SIGNAL_RECEIVED.lock().await;
+    if *shutdown_flag {
+        return HttpResponse::new(StatusCode::SERVICE_UNAVAILABLE);
     }
 
-    HttpResponse::new(StatusCode::SERVICE_UNAVAILABLE)
+    // Check the object store connection
+    if CONFIG.storage().get_object_store().check().await.is_ok() {
+        HttpResponse::new(StatusCode::OK)
+    } else {
+        HttpResponse::new(StatusCode::SERVICE_UNAVAILABLE)
+    }
 }

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

@@ -18,6 +18,7 @@
 use crate::analytics;
 use crate::banner;
 use crate::handlers::airplane;
+use crate::handlers::http::health_check;
 use crate::handlers::http::ingest;
 use crate::handlers::http::logstream;
 use crate::handlers::http::middleware::RouteExt;
@@ -35,6 +36,8 @@ use crate::storage::ObjectStorageError;
 use crate::storage::PARSEABLE_ROOT_DIRECTORY;
 use crate::sync;
 
+use std::sync::Arc;
+
 use super::server::Server;
 use super::ssl_acceptor::get_ssl_acceptor;
 use super::IngestorMetadata;
@@ -56,6 +59,7 @@ use bytes::Bytes;
 use once_cell::sync::Lazy;
 use relative_path::RelativePathBuf;
 use serde_json::Value;
+use tokio::sync::{oneshot, Mutex};
 
 /// ! have to use a guard before using it
 pub static INGESTOR_META: Lazy<IngestorMetadata> =
@@ -91,17 +95,47 @@ impl ParseableServer for IngestServer {
                 .wrap(cross_origin_config())
         };
 
-        // concurrent workers equal to number of logical cores
-        let http_server = HttpServer::new(create_app_fn).workers(num_cpus::get());
+        // Create a channel to trigger server shutdown
+        let (shutdown_trigger, shutdown_rx) = oneshot::channel::<()>();
+        let server_shutdown_signal = Arc::new(Mutex::new(Some(shutdown_trigger)));
+
+        // Clone the shutdown signal for the signal handler
+        let shutdown_signal = server_shutdown_signal.clone();
+
+        // Spawn the signal handler task
+        tokio::spawn(async move {
+            health_check::handle_signals(shutdown_signal).await;
+            println!("Received shutdown signal, notifying server to shut down...");
+        });
 
-        if let Some(config) = ssl {
+        // Create the HTTP server
+        let http_server = HttpServer::new(create_app_fn)
+            .workers(num_cpus::get())
+            .shutdown_timeout(60);
+
+        // Start the server with or without TLS
+        let srv = if let Some(config) = ssl {
             http_server
                 .bind_rustls_0_22(&CONFIG.parseable.address, config)?
                 .run()
-                .await?;
         } else {
-            http_server.bind(&CONFIG.parseable.address)?.run().await?;
-        }
+            http_server.bind(&CONFIG.parseable.address)?.run()
+        };
+
+        // Graceful shutdown handling
+        let srv_handle = srv.handle();
+
+        tokio::spawn(async move {
+            // Wait for the shutdown signal
+            shutdown_rx.await.ok();
+
+            // 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?;
 
         Ok(())
     }

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

@@ -18,6 +18,7 @@
 
 use crate::handlers::airplane;
 use crate::handlers::http::cluster::{self, init_cluster_metrics_schedular};
+use crate::handlers::http::health_check;
 use crate::handlers::http::logstream::create_internal_stream_if_not_exists;
 use crate::handlers::http::middleware::RouteExt;
 use crate::handlers::http::{base_path, cross_origin_config, API_BASE_PATH, API_VERSION};
@@ -32,6 +33,7 @@ use actix_web::web::ServiceConfig;
 use actix_web::{App, HttpServer};
 use async_trait::async_trait;
 use std::sync::Arc;
+use tokio::sync::{oneshot, Mutex};
 
 use crate::option::CONFIG;
 
@@ -74,16 +76,47 @@ impl ParseableServer for QueryServer {
                 .wrap(cross_origin_config())
         };
 
-        // concurrent workers equal to number of cores on the cpu
-        let http_server = HttpServer::new(create_app_fn).workers(num_cpus::get());
-        if let Some(config) = ssl {
+        // Create a channel to trigger server shutdown
+        let (shutdown_trigger, shutdown_rx) = oneshot::channel::<()>();
+        let server_shutdown_signal = Arc::new(Mutex::new(Some(shutdown_trigger)));
+
+        // Clone the shutdown signal for the signal handler
+        let shutdown_signal = server_shutdown_signal.clone();
+
+        // Spawn the signal handler task
+        tokio::spawn(async move {
+            health_check::handle_signals(shutdown_signal).await;
+            println!("Received shutdown signal, notifying server to shut down...");
+        });
+
+        // Create the HTTP server
+        let http_server = HttpServer::new(create_app_fn)
+            .workers(num_cpus::get())
+            .shutdown_timeout(60);
+
+        // Start the server with or without TLS
+        let srv = if let Some(config) = ssl {
             http_server
                 .bind_rustls_0_22(&CONFIG.parseable.address, config)?
                 .run()
-                .await?;
         } else {
-            http_server.bind(&CONFIG.parseable.address)?.run().await?;
-        }
+            http_server.bind(&CONFIG.parseable.address)?.run()
+        };
+
+        // Graceful shutdown handling
+        let srv_handle = srv.handle();
+
+        tokio::spawn(async move {
+            // Wait for the shutdown signal
+            shutdown_rx.await.ok();
+
+            // 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?;
 
         Ok(())
     }

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

@@ -37,6 +37,7 @@ use crate::sync;
 use crate::users::dashboards::DASHBOARDS;
 use crate::users::filters::FILTERS;
 use std::sync::Arc;
+use tokio::sync::{oneshot, Mutex};
 
 use actix_web::web::resource;
 use actix_web::Resource;
@@ -96,16 +97,47 @@ impl ParseableServer for Server {
             &CONFIG.parseable.tls_key_path,
         )?;
 
-        // concurrent workers equal to number of cores on the cpu
-        let http_server = HttpServer::new(create_app_fn).workers(num_cpus::get());
-        if let Some(config) = ssl {
+        // Create a channel to trigger server shutdown
+        let (shutdown_trigger, shutdown_rx) = oneshot::channel::<()>();
+        let server_shutdown_signal = Arc::new(Mutex::new(Some(shutdown_trigger)));
+
+        // Clone the shutdown signal for the signal handler
+        let shutdown_signal = server_shutdown_signal.clone();
+
+        // Spawn the signal handler task
+        tokio::spawn(async move {
+            health_check::handle_signals(shutdown_signal).await;
+            println!("Received shutdown signal, notifying server to shut down...");
+        });
+
+        // Create the HTTP server
+        let http_server = HttpServer::new(create_app_fn)
+            .workers(num_cpus::get())
+            .shutdown_timeout(60);
+
+        // Start the server with or without TLS
+        let srv = if let Some(config) = ssl {
             http_server
                 .bind_rustls_0_22(&CONFIG.parseable.address, config)?
                 .run()
-                .await?;
         } else {
-            http_server.bind(&CONFIG.parseable.address)?.run().await?;
-        }
+            http_server.bind(&CONFIG.parseable.address)?.run()
+        };
+
+        // Graceful shutdown handling
+        let srv_handle = srv.handle();
+
+        tokio::spawn(async move {
+            // Wait for the shutdown signal
+            shutdown_rx.await.ok();
+
+            // 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?;
 
         Ok(())
     }
@@ -548,6 +580,7 @@ impl Server {
         let app = self.start(prometheus, CONFIG.parseable.openid.clone());
 
         tokio::pin!(app);
+
         loop {
             tokio::select! {
                 e = &mut app => {
@@ -574,6 +607,7 @@ impl Server {
                     }
                     (remote_sync_handler, remote_sync_outbox, remote_sync_inbox) = sync::object_store_sync().await;
                 }
+
             };
         }
     }