From b69a77b63faff5ec6db9f267294ee80d814830e3 Mon Sep 17 00:00:00 2001 From: Devdutt Shenoi Date: Fri, 25 Oct 2024 12:07:28 +0530 Subject: [PATCH] refactor: use middlewares to reject requests when shutting down --- server/Cargo.toml | 2 +- server/src/handlers/http/health_check.rs | 30 ++++++++++++------- .../src/handlers/http/modal/ingest_server.rs | 2 ++ .../src/handlers/http/modal/query_server.rs | 2 ++ server/src/handlers/http/modal/server.rs | 2 ++ 5 files changed, 27 insertions(+), 11 deletions(-) diff --git a/server/Cargo.toml b/server/Cargo.toml index 806909759..04aec2f47 100644 --- a/server/Cargo.toml +++ b/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" diff --git a/server/src/handlers/http/health_check.rs b/server/src/handlers/http/health_check.rs index de4a01ba9..d59f3bec0 100644 --- a/server/src/handlers/http/health_check.rs +++ b/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>>>) { - let signal_received = SIGNAL_RECEIVED.clone(); +pub async fn check_shutdown_middleware( + req: ServiceRequest, + next: Next, +) -> Result, 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>>>) { 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 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) diff --git a/server/src/handlers/http/modal/ingest_server.rs b/server/src/handlers/http/modal/ingest_server.rs index 23a0dcc55..27fd99426 100644 --- a/server/src/handlers/http/modal/ingest_server.rs +++ b/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()) diff --git a/server/src/handlers/http/modal/query_server.rs b/server/src/handlers/http/modal/query_server.rs index 34f6d46a8..394853f30 100644 --- a/server/src/handlers/http/modal/query_server.rs +++ b/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()) diff --git a/server/src/handlers/http/modal/server.rs b/server/src/handlers/http/modal/server.rs index ef8467f67..51195609b 100644 --- a/server/src/handlers/http/modal/server.rs +++ b/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())