[feat] - Graceful shutdown

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"
+wiremock = "0.6.2"

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"

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"

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"

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"

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");
   }

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...");
+      },
+  }
 }

src/util/task.rs

@@ -17,18 +17,20 @@ pub async fn join_all(tasks: Vec<Task>) -> 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(()),
   }
 }