From c4f599e39489b06b5297864cc7150e48ded0d84c Mon Sep 17 00:00:00 2001
From: Joe Ardent <code@ardent.nebcorp.com>
Date: Fri, 8 Mar 2024 22:40:23 -0800
Subject: [PATCH] add graceful shutdown handler

---
 Cargo.lock  | 10 ++++++++++
 Cargo.toml  |  2 +-
 src/main.rs | 30 +++++++++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 39daee3..c3041dd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1481,6 +1481,15 @@ dependencies = [
  "digest",
 ]
 
+[[package]]
+name = "signal-hook-registry"
+version = "1.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "signature"
 version = "2.2.0"
@@ -1888,6 +1897,7 @@ dependencies = [
  "mio",
  "num_cpus",
  "pin-project-lite",
+ "signal-hook-registry",
  "socket2",
  "tokio-macros",
  "windows-sys 0.48.0",
diff --git a/Cargo.toml b/Cargo.toml
index 19604fc..bc79a60 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,7 +20,7 @@ serde_json = { version = "1", default-features = false }
 sqlx = { version = "0.7", default-features = false, features = ["runtime-tokio", "sqlite", "tls-none", "migrate", "macros"] }
 thiserror = { version = "1" }
 time = { version = "0.3", default-features = false }
-tokio = { version = "1", default-features = false, features = ["rt-multi-thread"] }
+tokio = { version = "1", default-features = false, features = ["rt-multi-thread", "signal"] }
 tower-http = { version = "0.5", default-features = false, features = ["fs"] }
 tower-sessions = { version = "0.11", default-features = false, features = ["axum-core"] }
 tower-sessions-sqlx-store = { version = "0.11", default-features = false, features = ["sqlite"] }
diff --git a/src/main.rs b/src/main.rs
index fca7b5c..8ac260f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -55,7 +55,10 @@ async fn main() {
         .with_state(pool)
         .into_make_service();
     let listener = mklistener().await;
-    axum::serve(listener, app).await.unwrap();
+    axum::serve(listener, app)
+        .with_graceful_shutdown(shutdown_signal())
+        .await
+        .unwrap();
 }
 
 //-************************************************************************
@@ -114,3 +117,28 @@ where
             .route(path.trim_end_matches('/'), method_router)
     }
 }
+
+async fn shutdown_signal() {
+    use tokio::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 => {},
+    }
+}