From 49c418068e8e173739266bdc1c359cd8cba98ed2 Mon Sep 17 00:00:00 2001
From: azdle <azdle@azdle.net>
Date: Wed, 14 May 2025 13:13:00 -0500
Subject: [PATCH] add graceful shutdown

---
 src/server/mod.rs | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/src/server/mod.rs b/src/server/mod.rs
index d85a17f..c67d737 100644
--- a/src/server/mod.rs
+++ b/src/server/mod.rs
@@ -8,6 +8,7 @@ use std::future::{Future, IntoFuture};
 use std::net::SocketAddr;
 use std::pin::pin;
 use std::pin::Pin;
+use tokio::signal;
 use tracing::info;
 
 use crate::Conf;
@@ -45,7 +46,7 @@ impl ZeroToAxum {
 
         let listener = tokio::net::TcpListener::bind(&conf.listen).await.unwrap();
         let bound_addr = listener.local_addr().unwrap();
-        let server = axum::serve(listener, app);
+        let server = axum::serve(listener, app).with_graceful_shutdown(shutdown_signal());
 
         info!("server started, listening on {bound_addr:?}");
 
@@ -67,3 +68,31 @@ impl FromRef<AppState> for Key {
         state.key.clone()
     }
 }
+
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+
+        info!("got Ctrl+C, shutting down...");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+
+        info!("got SIGTERM, shutting down...");
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = terminate => {},
+    }
+}