diff --git a/Cargo.lock b/Cargo.lock index fb5cb6c..63e9f20 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2373,6 +2373,16 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + [[package]] name = "signature" version = "2.2.0" @@ -2842,6 +2852,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2 0.6.4", "tokio-macros", "windows-sys 0.61.2", diff --git a/Cargo.toml b/Cargo.toml index 528f33b..68d83b8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,7 +32,7 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" chrono = { version = "0.4", features = ["serde"] } thiserror = "1" -tokio = { version = "1", features = ["sync", "macros", "rt-multi-thread", "net"] } +tokio = { version = "1", features = ["sync", "macros", "rt-multi-thread", "net", "signal"] } url = "2" uuid = { version = "1", features = ["v4"] } dashmap = "5" diff --git a/src/main.rs b/src/main.rs index 5acf03a..01a9617 100644 --- a/src/main.rs +++ b/src/main.rs @@ -37,6 +37,7 @@ mod native { use std::sync::Arc; use axum::extract::State; + use axum::http::StatusCode; use axum::response::IntoResponse; use axum::routing::{get, post}; use axum::{Json, Router}; @@ -55,6 +56,7 @@ mod native { cache: Arc, config_watcher: ConfigWatcher, config_version: u32, + shutdown_flag: Arc>, } /// Build the axum router with all application routes. @@ -72,8 +74,14 @@ mod native { } /// `GET /health` — returns a JSON health-check payload. - async fn health_handler() -> impl IntoResponse { - Json(json!({"status": "ok"})) + async fn health_handler(State(state): State) -> impl IntoResponse { + if *state.shutdown_flag.read().await { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"status": "shutting_down"})), + ); + } + (StatusCode::OK, Json(json!({"status": "ok"}))) } /// `GET /metrics` — returns Prometheus text-format metrics. @@ -161,6 +169,36 @@ mod native { } } + /// Listen for SIGTERM / SIGINT, then flip the shutdown flag so the + /// health endpoint reports 503 while `axum::serve` drains in-flight + /// connections. + async fn shutdown_signal(flag: Arc>) { + let ctrl_c = async { + tokio::signal::ctrl_c() + .await + .expect("failed to install Ctrl+C handler"); + }; + + #[cfg(unix)] + let terminate = async { + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("failed to install signal handler") + .recv() + .await; + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } + + *flag.write().await = true; + eprintln!("[proofstell] Shutdown signal received, starting graceful shutdown..."); + } + /// Bootstrap: load config, wire up services, and start the server. pub async fn run() -> anyhow::Result<()> { // ── Metrics ───────────────────────────────────────────────── @@ -263,12 +301,16 @@ mod native { }); // ── Router ────────────────────────────────────────────────── + let shutdown_flag = Arc::new(tokio::sync::RwLock::new(false)); + let webhook_for_shutdown = Arc::clone(&webhook); + let cache_for_shutdown = Arc::clone(&cache); let state = AppState { metrics: Arc::clone(&metrics), webhook, cache, config_watcher, config_version: AppConfig::version(), + shutdown_flag: Arc::clone(&shutdown_flag), }; let app = build_router(state); @@ -278,7 +320,15 @@ mod native { eprintln!("[proofstell] Config endpoints: GET /config/status, POST /config/reload"); let listener = tokio::net::TcpListener::bind(addr).await?; - axum::serve(listener, app).await?; + axum::serve(listener, app) + .with_graceful_shutdown(shutdown_signal(shutdown_flag)) + .await?; + + eprintln!("[proofstell] Draining webhook DLQ..."); + webhook_for_shutdown.drain_dlq().await; + eprintln!("[proofstell] Flushing cache..."); + drop(cache_for_shutdown); + eprintln!("[proofstell] Shutdown complete"); Ok(()) }