From 50326c205f30fcffa3ada13a4b886b01a7df3c4a Mon Sep 17 00:00:00 2001
From: NetworkCats <networkcatlol@gmail.com>
Date: Wed, 4 Mar 2026 04:20:35 +0800
Subject: [PATCH] feat: add graceful shutdown with SIGTERM handling and updater
 cancellation

---
 Cargo.lock     | 20 ++++++++++++++++++++
 Cargo.toml     |  1 +
 Dockerfile     |  2 ++
 src/main.rs    | 36 +++++++++++++++++++++++++++++++-----
 src/updater.rs | 17 +++++++++++++----
 5 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3f0284f..cfb0fc7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -329,6 +329,12 @@ version = "0.3.32"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
 
+[[package]]
+name = "futures-sink"
+version = "0.3.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893"
+
 [[package]]
 name = "futures-task"
 version = "0.3.32"
@@ -1104,6 +1110,7 @@ dependencies = [
  "serde_json",
  "static-files",
  "tokio",
+ "tokio-util",
  "tower",
  "tracing",
  "tracing-subscriber",
@@ -1434,6 +1441,19 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "tokio-util"
+version = "0.7.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098"
+dependencies = [
+ "bytes",
+ "futures-core",
+ "futures-sink",
+ "pin-project-lite",
+ "tokio",
+]
+
 [[package]]
 name = "tower"
 version = "0.5.3"
diff --git a/Cargo.toml b/Cargo.toml
index a6bf030..c488897 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ reqwest = { version = "0.13", default-features = false, features = ["rustls"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 dotenvy = "0.15"
+tokio-util = "0.7"
 static-files = "0.3.1"
 
 [build-dependencies]
diff --git a/Dockerfile b/Dockerfile
index fcb9038..3eb4aca 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -25,4 +25,6 @@ RUN chown -R rustyip:rustyip /app
 
 USER rustyip
 
+STOPSIGNAL SIGTERM
+
 ENTRYPOINT ["rustyip"]
diff --git a/src/main.rs b/src/main.rs
index 23cc7b6..fc6388b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,7 @@
 use std::path::PathBuf;
 
-use tracing::error;
+use tokio_util::sync::CancellationToken;
+use tracing::{error, info};
 use tracing_subscriber::EnvFilter;
 
 use rustyip::{config, db, handlers, routes, updater};
@@ -32,12 +33,14 @@ async fn main() {
     };
 
     let shared_db = db::new_shared(reader);
+    let cancel = CancellationToken::new();
 
-    updater::spawn_updater(
+    let updater_handle = updater::spawn_updater(
         shared_db.clone(),
         db_path,
         config.db_update_url,
         config.db_update_interval_hours,
+        cancel.clone(),
     );
 
     let state = handlers::AppState {
@@ -56,10 +59,33 @@ async fn main() {
         .with_graceful_shutdown(shutdown_signal())
         .await
         .expect("server error");
+
+    info!("server stopped, cleaning up background tasks");
+    cancel.cancel();
+    let _ = updater_handle.await;
+    info!("shutdown complete");
 }
 
 async fn shutdown_signal() {
-    tokio::signal::ctrl_c()
-        .await
-        .expect("failed to listen for ctrl+c");
+    let ctrl_c = async {
+        tokio::signal::ctrl_c()
+            .await
+            .expect("failed to listen for ctrl+c");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
+            .expect("failed to listen for SIGTERM")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        () = ctrl_c => info!("received SIGINT, shutting down"),
+        () = terminate => info!("received SIGTERM, shutting down"),
+    }
 }
diff --git a/src/updater.rs b/src/updater.rs
index 0dc3f4d..3a89c94 100644
--- a/src/updater.rs
+++ b/src/updater.rs
@@ -4,7 +4,9 @@ use std::time::Duration;
 
 use tokio::fs;
 use tokio::io::AsyncWriteExt;
-use tracing::error;
+use tokio::task::JoinHandle;
+use tokio_util::sync::CancellationToken;
+use tracing::{error, info};
 
 use crate::db::{DbReader, SharedDb, load_db};
 
@@ -55,16 +57,23 @@ pub fn spawn_updater(
     db_path: PathBuf,
     update_url: String,
     interval_hours: u64,
-) {
+    cancel: CancellationToken,
+) -> JoinHandle<()> {
     tokio::spawn(async move {
         let interval = Duration::from_secs(interval_hours * 3600);
         loop {
-            tokio::time::sleep(interval).await;
+            tokio::select! {
+                () = tokio::time::sleep(interval) => {}
+                () = cancel.cancelled() => {
+                    info!("updater task cancelled, stopping");
+                    break;
+                }
+            }
             if let Err(e) = update_db(&shared_db, &db_path, &update_url).await {
                 error!("database update failed: {e}");
             }
         }
-    });
+    })
 }
 
 async fn update_db(