axum with server graceful shutdown

be-rust-axum/rust-template/Cargo.toml

@@ -10,7 +10,7 @@ license = "MIT OR Apache-2.0"
 [dependencies]
 # one can add more dependencies via cargo to Cargo.toml as shown next: cargo add axum -F axum/http2
 axum = { version = "0.7", features = ["http2"] }
-tokio = { version = "1.35", features = ["rt-multi-thread", "macros"] }
+tokio = { version = "1.35", features = ["rt-multi-thread", "macros", "signal"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 http = "1.0"

be-rust-axum/rust-template/src/api/router.rs

@@ -1,6 +1,7 @@
 use axum::{Json, Router};
 use http::StatusCode;
 use std::net::SocketAddr;
+use tokio::signal::unix::{signal, SignalKind};
 use tracing::info;
 
 use crate::api::routes;
@@ -17,6 +18,17 @@ pub fn app() -> Router {
     .fallback(fallback)
 }
 
+/// The serve function that starts the Axum server
+pub async fn serve(address: SocketAddr) {
+  let listener = tokio::net::TcpListener::bind(address).await.unwrap();
+
+  info!("Server listening on {}", &address);
+  axum::serve(listener, app().into_make_service())
+		.with_graceful_shutdown(shutdown_signal())
+		.await
+    .expect("Failed to start server");
+}
+
 /// The fallback function when a non configured endpoint is reached
 async fn fallback() -> (StatusCode, Json<Status>) {
   (
@@ -27,12 +39,22 @@ async fn fallback() -> (StatusCode, Json<Status>) {
   )
 }
 
-/// The serve function that starts the Axum server
-pub async fn serve(address: SocketAddr) {
-  let listener = tokio::net::TcpListener::bind(address).await.unwrap();
-
-  info!("Server listening on {}", &address);
-  axum::serve(listener, app().into_make_service())
-    .await
-    .expect("Failed to start server");
+/// Setup example of the OS signal handlers to catch for proper server graceful shutdown.
+async fn shutdown_signal() {
+  let interrupt = async {
+      signal(SignalKind::interrupt())
+          .expect("failed to install signal handler")
+          .recv()
+          .await;
+  };
+  let terminate = async {
+      signal(SignalKind::terminate())
+          .expect("failed to install signal handler")
+          .recv()
+          .await;
+  };
+  tokio::select! {
+      _ = interrupt => {},
+      _ = terminate => {},
+  }
 }