refactor: use middlewares to reject requests when shutting down (#974)

Author: Devdutt Shenoi

server/Cargo.toml

@@ -25,7 +25,7 @@ tower-http = { version = "0.6.1", features = ["cors"] }
 
 ### actix dependencies
 actix-web-httpauth = "0.8"
-actix-web = { version = "4.5.1", features = ["rustls-0_22"] }
+actix-web = { version = "4.9.0", features = ["rustls-0_22"] }
 actix-cors = "0.7.0"
 actix-web-prometheus = { version = "0.1" }
 actix-web-static-files = "4.0"

server/src/handlers/http/health_check.rs

@@ -17,8 +17,12 @@
  */
 
 use crate::option::CONFIG;
+use actix_web::body::MessageBody;
+use actix_web::dev::{ServiceRequest, ServiceResponse};
+use actix_web::error::ErrorServiceUnavailable;
 use actix_web::http::StatusCode;
-use actix_web::HttpResponse;
+use actix_web::middleware::Next;
+use actix_web::{Error, HttpResponse};
 use lazy_static::lazy_static;
 use std::sync::Arc;
 use tokio::signal::unix::{signal, SignalKind};
@@ -34,9 +38,21 @@ 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();
+pub async fn check_shutdown_middleware(
+    req: ServiceRequest,
+    next: Next<impl MessageBody>,
+) -> Result<ServiceResponse<impl MessageBody>, Error> {
+    // Acquire the shutdown flag to check if the server is shutting down.
+    if *SIGNAL_RECEIVED.lock().await {
+        // Return 503 Service Unavailable if the server is shutting down.
+        Err(ErrorServiceUnavailable("Server is shutting down"))
+    } else {
+        // Continue processing the request if the server is not shutting down.
+        next.call(req).await
+    }
+}
 
+pub async fn handle_signals(shutdown_signal: Arc<Mutex<Option<oneshot::Sender<()>>>>) {
     let mut sigterm =
         signal(SignalKind::terminate()).expect("Failed to set up SIGTERM signal handler");
     log::info!("Signal handler task started");
@@ -47,7 +63,7 @@ pub async fn handle_signals(shutdown_signal: Arc<Mutex<Option<oneshot::Sender<()
             log::info!("Received SIGTERM signal at Readiness Probe Handler");
 
             // Set the shutdown flag to true
-            let mut shutdown_flag = signal_received.lock().await;
+            let mut shutdown_flag = SIGNAL_RECEIVED.lock().await;
             *shutdown_flag = true;
 
             // Trigger graceful shutdown
@@ -77,12 +93,6 @@ pub async fn handle_signals(shutdown_signal: Arc<Mutex<Option<oneshot::Sender<()
 }
 
 pub async fn readiness() -> HttpResponse {
-    // 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);
-    }
-
     // Check the object store connection
     if CONFIG.storage().get_object_store().check().await.is_ok() {
         HttpResponse::new(StatusCode::OK)

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

@@ -54,6 +54,7 @@ use crate::{
     option::CONFIG,
 };
 use actix_web::body::MessageBody;
+use actix_web::middleware::from_fn;
 use actix_web::web::resource;
 use actix_web::Scope;
 use actix_web::{web, App, HttpServer};
@@ -97,6 +98,7 @@ impl ParseableServer for IngestServer {
             App::new()
                 .wrap(prometheus.clone())
                 .configure(|config| IngestServer::configure_routes(config, None))
+                .wrap(from_fn(health_check::check_shutdown_middleware))
                 .wrap(actix_web::middleware::Logger::default())
                 .wrap(actix_web::middleware::Compress::default())
                 .wrap(cross_origin_config())

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

@@ -29,6 +29,7 @@ use crate::sync;
 use crate::users::dashboards::DASHBOARDS;
 use crate::users::filters::FILTERS;
 use crate::{analytics, banner, metrics, migration, rbac, storage};
+use actix_web::middleware::from_fn;
 use actix_web::web::{resource, ServiceConfig};
 use actix_web::{web, Scope};
 use actix_web::{App, HttpServer};
@@ -74,6 +75,7 @@ impl ParseableServer for QueryServer {
             App::new()
                 .wrap(prometheus.clone())
                 .configure(|config| QueryServer::configure_routes(config, oidc_client.clone()))
+                .wrap(from_fn(health_check::check_shutdown_middleware))
                 .wrap(actix_web::middleware::Logger::default())
                 .wrap(actix_web::middleware::Compress::default())
                 .wrap(cross_origin_config())

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

@@ -38,6 +38,7 @@ use crate::storage;
 use crate::sync;
 use crate::users::dashboards::DASHBOARDS;
 use crate::users::filters::FILTERS;
+use actix_web::middleware::from_fn;
 use std::sync::Arc;
 use tokio::sync::{oneshot, Mutex};
 
@@ -89,6 +90,7 @@ impl ParseableServer for Server {
             App::new()
                 .wrap(prometheus.clone())
                 .configure(|cfg| Server::configure_routes(cfg, oidc_client.clone()))
+                .wrap(from_fn(health_check::check_shutdown_middleware))
                 .wrap(actix_web::middleware::Logger::default())
                 .wrap(actix_web::middleware::Compress::default())
                 .wrap(cross_origin_config())