add graceful shutdown

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 => {},
+    }
+}