diff --git a/Cargo.toml b/Cargo.toml index 85f7f59..f524a3c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,6 +60,8 @@ tera = "1.20.0" test-context = "0.3.0" thiserror = "2.0.6" tokio = { version = "1.42.0", features = ["full"] } +tower = { version = "0.5.1", features = ["full"] } +tower-http = { version = "0.6.2", features = ["full"] } tracing = { version = "0.1.41", features = ["attributes"] } tracing-appender = "0.2.3" tracing-bunyan-formatter = "0.3.10" @@ -76,4 +78,4 @@ uuid = { version = "1.11.0", features = ["v4", "serde"] } tokio-tungstenite = "0.24.0" garde = { version = "0.20.0", features = ["full"] } regex = "1.11.1" -wiremock = "0.6.2" \ No newline at end of file +wiremock = "0.6.2" diff --git a/settings/base.toml b/settings/base.toml index b419aeb..f4296df 100644 --- a/settings/base.toml +++ b/settings/base.toml @@ -1,6 +1,7 @@ [server] addr = "0.0.0.0" port = 8_080 +grace_shutdown_secs = 10 [db] host = "127.0.0.1" diff --git a/settings/dev.toml b/settings/dev.toml index 8519e7b..a219e1c 100644 --- a/settings/dev.toml +++ b/settings/dev.toml @@ -3,6 +3,7 @@ profile = "dev" [server] addr = "127.0.0.1" port = 8_080 +grace_shutdown_secs = 10 [db] host = "127.0.0.1" diff --git a/settings/prod.toml b/settings/prod.toml index fe8fc86..0a13b21 100644 --- a/settings/prod.toml +++ b/settings/prod.toml @@ -3,6 +3,7 @@ profile = "prod" [server] addr = "0.0.0.0" port = 8_080 +grace_shutdown_secs = 30 [db] host = "127.0.0.1" diff --git a/settings/test.toml b/settings/test.toml index 6f8af9b..05ea43d 100644 --- a/settings/test.toml +++ b/settings/test.toml @@ -3,6 +3,7 @@ profile = "test" [server] addr = "127.0.0.1" port = 0 +grace_shutdown_secs = 1 [db] host = "127.0.0.1" diff --git a/src/configure/server.rs b/src/configure/server.rs index 6eaf085..d2223da 100644 --- a/src/configure/server.rs +++ b/src/configure/server.rs @@ -6,6 +6,7 @@ use serde::Deserialize; pub struct ServerConfig { pub addr: String, pub port: u16, + pub grace_shutdown_secs: i64, } impl ServerConfig { @@ -31,6 +32,7 @@ pub mod tests { let config = ServerConfig { addr: "127.0.0.1".to_string(), port: 1024, + grace_shutdown_secs: 10, }; assert_eq!(config.get_http_addr(), "http://127.0.0.1:1024"); } diff --git a/src/server/mod.rs b/src/server/mod.rs index 86acefe..339a329 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -2,6 +2,9 @@ use self::state::AppState; use crate::configure::AppConfig; use crate::error::AppResult; use crate::router::create_router_app; +use tower_http::timeout::TimeoutLayer; +use tower_http::trace::TraceLayer; +use tracing::info; pub mod state; pub mod worker; @@ -20,8 +23,46 @@ impl AppServer { } pub async fn run(self) -> AppResult<()> { - let router = create_router_app(self.state); - axum::serve(self.tcp, router).await?; + let shutdown = self.grace_shutdown_time(); + let router = create_router_app(self.state) + .layer(TraceLayer::new_for_http()) // Visibility of the request and response, change as needed. + .layer(TimeoutLayer::new(shutdown)); // Graceful shutdown for hosting services requries N time to complete the request before shutting down. + + axum::serve(self.tcp, router) + .with_graceful_shutdown(shutdown_signal()) + .await?; Ok(()) } + + fn grace_shutdown_time(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.state.config.server.grace_shutdown_secs as u64) + } +} + +async fn shutdown_signal() { + let ctrl_c = async { + tokio::signal::ctrl_c() + .await + .expect("Failed to install CTRL+C signal handler"); + }; + + #[cfg(unix)] + let terminate = async { + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("Failed to install SIGTERM signal handler") + .recv() + .await; + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => { + info!("Received Ctrl-C signal. Shutting down..."); + }, + _ = terminate => { + info!("Received SIGTERM signal. Shutting down..."); + }, + } } diff --git a/src/util/task.rs b/src/util/task.rs index 7293d0c..f54f366 100644 --- a/src/util/task.rs +++ b/src/util/task.rs @@ -17,18 +17,20 @@ pub async fn join_all(tasks: Vec) -> AppResult { tokio::spawn(async { if let Err(e) = task.await { if let Some(sender) = sender { - sender - .send(e) - .await - .unwrap_or_else(|_| unreachable!("This channel never closed.")); + let _ = sender.send(e).await; } else { error!("A task failed: {e}."); } } }); } + + // Explicitly drop the sender to close the channel. + drop(sender); + + // Return Ok(()) if all futures are completed without error. match receiver.recv().await { Some(err) => Err(err), - None => unreachable!("This channel never closed."), + None => Ok(()), } }