diff --git a/src/main.rs b/src/main.rs index c7681fd..40dcc28 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,9 +15,11 @@ use magnetar_calckey_model::{CacheConnectorConfig, CalckeyCache, CalckeyModel, C use miette::{miette, IntoDiagnostic}; use std::net::SocketAddr; use std::sync::Arc; +use tokio::signal; use tower_http::cors::{Any, CorsLayer}; use tower_http::trace::TraceLayer; use tracing::info; +use tracing::log::error; use tracing_subscriber::EnvFilter; #[tokio::main] @@ -85,6 +87,33 @@ async fn main() -> miette::Result<()> { info!("Serving on: {addr}"); axum::Server::bind(&addr) .serve(app.into_make_service()) + .with_graceful_shutdown(shutdown_signal()) .await .map_err(|e| miette!("Error running server: {}", e)) } + +async fn shutdown_signal() { + let ctrl_c = async { + if let Err(e) = signal::ctrl_c().await { + error!("Ctrl+C signal handler error: {}", e); + } + }; + + #[cfg(unix)] + let terminate = async { + signal::unix::signal(signal::unix::SignalKind::terminate()) + .expect("SIGTERM handler error") + .recv() + .await; + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } + + info!("Shutting down..."); +}