Merge remote-tracking branch 'origin/main'

Author: Devdutt Shenoi

src/handlers/http/health_check.rs

@@ -25,10 +25,7 @@ use actix_web::middleware::Next;
 use actix_web::{Error, HttpResponse};
 use lazy_static::lazy_static;
 use std::sync::Arc;
-use tokio::signal::ctrl_c;
-use tracing::info;
-
-use tokio::sync::{oneshot, Mutex};
+use tokio::sync::Mutex;
 
 // Create a global variable to store signal status
 lazy_static! {
@@ -53,46 +50,16 @@ pub async fn check_shutdown_middleware(
     }
 }
 
-pub async fn handle_signals(shutdown_signal: Arc<Mutex<Option<oneshot::Sender<()>>>>) {
-    #[cfg(windows)]
-    {
-        tokio::select! {
-            _ = ctrl_c() => {
-                info!("Received SIGINT signal at Readiness Probe Handler");
-                shutdown(shutdown_signal).await;
-            }
-        }
-    }
-    #[cfg(unix)]
-    {
-        use tokio::signal::unix::{signal, SignalKind};
-        let mut sigterm = signal(SignalKind::terminate()).unwrap();
-        tokio::select! {
-            _ = ctrl_c() => {
-                info!("Received SIGINT signal at Readiness Probe Handler");
-                shutdown(shutdown_signal).await;
-            },
-            _ = sigterm.recv() => {
-                info!("Received SIGTERM signal at Readiness Probe Handler");
-                shutdown(shutdown_signal).await;
-            }
-        }
-    }
-}
-
-async fn shutdown(shutdown_signal: Arc<Mutex<Option<oneshot::Sender<()>>>>) {
+// This function is called when the server is shutting down
+pub async fn shutdown() {
     // Set the shutdown flag to true
     let mut shutdown_flag = SIGNAL_RECEIVED.lock().await;
     *shutdown_flag = true;
 
     // Sync to local
     crate::event::STREAM_WRITERS.unset_all();
-
-    // Trigger graceful shutdown
-    if let Some(shutdown_sender) = shutdown_signal.lock().await.take() {
-        let _ = shutdown_sender.send(());
-    }
 }
+
 pub async fn readiness() -> HttpResponse {
     // Check the object store connection
     if CONFIG.storage().get_object_store().check().await.is_ok() {

src/handlers/http/modal/ingest_server.rs

@@ -51,6 +51,7 @@ use bytes::Bytes;
 use once_cell::sync::Lazy;
 use relative_path::RelativePathBuf;
 use serde_json::Value;
+use tokio::sync::oneshot;
 use tracing::error;
 
 /// ! have to use a guard before using it
@@ -97,7 +98,7 @@ impl ParseableServer for IngestServer {
     }
 
     /// configure the server and start an instance to ingest data
-    async fn init(&self) -> anyhow::Result<()> {
+    async fn init(&self, shutdown_rx: oneshot::Receiver<()>) -> anyhow::Result<()> {
         let prometheus = metrics::build_metrics_handler();
         CONFIG.storage().register_store_metrics(&prometheus);
 
@@ -114,7 +115,7 @@ impl ParseableServer for IngestServer {
         set_ingestor_metadata().await?;
 
         // Ingestors shouldn't have to deal with OpenId auth flow
-        let app = self.start(prometheus, None);
+        let app = self.start(shutdown_rx, prometheus, None);
 
         tokio::pin!(app);
         loop {

src/handlers/http/modal/mod.rs

@@ -38,7 +38,7 @@ use openid::Discovered;
 use serde::Deserialize;
 use serde::Serialize;
 use ssl_acceptor::get_ssl_acceptor;
-use tokio::sync::{oneshot, Mutex};
+use tokio::sync::oneshot;
 use tracing::{error, info, warn};
 
 use super::audit;
@@ -67,11 +67,12 @@ pub trait ParseableServer {
     async fn load_metadata(&self) -> anyhow::Result<Option<Bytes>>;
 
     /// code that describes starting and setup procedures for each type of server
-    async fn init(&self) -> anyhow::Result<()>;
+    async fn init(&self, shutdown_rx: oneshot::Receiver<()>) -> anyhow::Result<()>;
 
     /// configure the server
     async fn start(
         &self,
+        shutdown_rx: oneshot::Receiver<()>,
         prometheus: PrometheusMetrics,
         oidc_client: Option<crate::oidc::OpenidConfig>,
     ) -> anyhow::Result<()>
@@ -108,19 +109,6 @@ pub trait ParseableServer {
                 .wrap(cross_origin_config())
         };
 
-        // 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
-        let signal_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())
@@ -142,6 +130,8 @@ pub trait ParseableServer {
             // Wait for the shutdown signal
             let _ = shutdown_rx.await;
 
+            health_check::shutdown().await;
+
             // Perform S3 sync and wait for completion
             info!("Starting data sync to S3...");
             if let Err(e) = CONFIG.storage().get_object_store().sync(true).await {
@@ -158,11 +148,6 @@ pub trait ParseableServer {
         // 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 {
-            error!("Error in signal handler: {:?}", e);
-        }
-
         // Wait for the sync task to complete before exiting
         if let Err(e) = sync_task.await {
             error!("Error in sync task: {:?}", e);

src/handlers/http/modal/query_server.rs

@@ -34,6 +34,7 @@ use actix_web::web::{resource, ServiceConfig};
 use actix_web::{web, Scope};
 use async_trait::async_trait;
 use bytes::Bytes;
+use tokio::sync::oneshot;
 use tracing::{error, info};
 
 use crate::{option::CONFIG, ParseableServer};
@@ -85,7 +86,7 @@ impl ParseableServer for QueryServer {
     }
 
     /// initialize the server, run migrations as needed and start an instance
-    async fn init(&self) -> anyhow::Result<()> {
+    async fn init(&self, shutdown_rx: oneshot::Receiver<()>) -> anyhow::Result<()> {
         let prometheus = metrics::build_metrics_handler();
         CONFIG.storage().register_store_metrics(&prometheus);
 
@@ -121,7 +122,7 @@ impl ParseableServer for QueryServer {
             sync::object_store_sync().await;
 
         tokio::spawn(airplane::server());
-        let app = self.start(prometheus, CONFIG.options.openid());
+        let app = self.start(shutdown_rx, prometheus, CONFIG.options.openid());
 
         tokio::pin!(app);
         loop {

src/handlers/http/modal/server.rs

@@ -40,6 +40,7 @@ use actix_web::Scope;
 use actix_web_static_files::ResourceFiles;
 use async_trait::async_trait;
 use bytes::Bytes;
+use tokio::sync::oneshot;
 use tracing::error;
 
 use crate::{
@@ -94,7 +95,7 @@ impl ParseableServer for Server {
     }
 
     // configure the server and start an instance of the single server setup
-    async fn init(&self) -> anyhow::Result<()> {
+    async fn init(&self, shutdown_rx: oneshot::Receiver<()>) -> anyhow::Result<()> {
         let prometheus = metrics::build_metrics_handler();
         CONFIG.storage().register_store_metrics(&prometheus);
 
@@ -124,7 +125,7 @@ impl ParseableServer for Server {
         tokio::spawn(handlers::livetail::server());
         tokio::spawn(handlers::airplane::server());
 
-        let app = self.start(prometheus, CONFIG.options.openid());
+        let app = self.start(shutdown_rx, prometheus, CONFIG.options.openid());
 
         tokio::pin!(app);
 

src/main.rs

@@ -21,6 +21,9 @@ use parseable::{
     option::{Mode, CONFIG},
     rbac, storage, IngestServer, ParseableServer, QueryServer, Server,
 };
+use tokio::signal::ctrl_c;
+use tokio::sync::oneshot;
+use tracing::info;
 use tracing_subscriber::EnvFilter;
 
 #[cfg(any(
@@ -61,7 +64,37 @@ async fn main() -> anyhow::Result<()> {
         tokio::task::spawn(kafka::setup_integration());
     }
 
-    server.init().await?;
+    // Spawn a task to trigger graceful shutdown on appropriate signal
+    let (shutdown_trigger, shutdown_rx) = oneshot::channel::<()>();
+    tokio::spawn(async move {
+        block_until_shutdown_signal().await;
+
+        // Trigger graceful shutdown
+        println!("Received shutdown signal, notifying server to shut down...");
+        shutdown_trigger.send(()).unwrap();
+    });
+
+    server.init(shutdown_rx).await?;
 
     Ok(())
 }
+
+#[cfg(windows)]
+/// Asynchronously blocks until a shutdown signal is received
+pub async fn block_until_shutdown_signal() {
+    _ = ctrl_c().await;
+    info!("Received a CTRL+C event");
+}
+
+#[cfg(unix)]
+/// Asynchronously blocks until a shutdown signal is received
+pub async fn block_until_shutdown_signal() {
+    use tokio::signal::unix::{signal, SignalKind};
+    let mut sigterm =
+        signal(SignalKind::terminate()).expect("Failed to create SIGTERM signal handler");
+
+    tokio::select! {
+        _ = ctrl_c() => info!("Received SIGINT signal"),
+        _ = sigterm.recv() => info!("Received SIGTERM signal"),
+    }
+}