Add graceful shutdown.

src/main.rs

@@ -1,10 +1,14 @@
 #![forbid(unsafe_code)]
+use std::time::Duration;
+
 use axum::http::StatusCode;
 use axum::routing::get;
 use axum::routing::post;
 use axum::Json;
 use axum::Router;
 use serde::Serialize;
+use tokio::signal;
+use tower_http::timeout::TimeoutLayer;
 use tower_http::trace::TraceLayer;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
@@ -27,14 +31,44 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let app = Router::new()
         .route("/health", get(health))
         .route("/hook", post(hook))
-        .layer(TraceLayer::new_for_http());
+        .layer((
+            TraceLayer::new_for_http(),
+            // Add a timeout layer so graceful shutdown can't wait forever.
+            TimeoutLayer::new(Duration::from_secs(600)),
+        ));
 
     let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
     tracing::info!("listening on {}", listener.local_addr().unwrap());
-    axum::serve(listener, app).await?;
+    axum::serve(listener, app)
+        .with_graceful_shutdown(shutdown_signal())
+        .await?;
     Ok(())
 }
 
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(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 => {},
+    }
+}
+
 async fn health() -> (StatusCode, Json<HealthResponse>) {
     (StatusCode::OK, Json(HealthResponse { ok: true }))
 }