From 3e1a936e25e600aacb2462d93b08cfa362f2653f Mon Sep 17 00:00:00 2001
From: KeyCode17 <m.daffa.karyudi@gmail.com>
Date: Mon, 18 May 2026 06:04:05 +0700
Subject: [PATCH] feat(server): drain in-flight solves on SIGTERM / Ctrl-C
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wire `axum::serve(...).with_graceful_shutdown(shutdown_signal())`.
The signal handler resolves on the first of:

- SIGINT (Ctrl-C from a TTY)
- SIGTERM (sent by systemd, Kubernetes, `kill`) — Unix only; on
  other targets the SIGTERM branch is a pending future so SIGINT
  alone still wins.

After either signal we log a single "draining in-flight solves"
line; axum stops accepting new connections, lets the existing
`/v1/solve` requests finish, and the binary exits with a clean
"px-server stopped cleanly" log. Errors installing the handlers
log a warning but do not abort the process — the worst case
degrades to the previous "hard kill" behaviour.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 px-server/src/main.rs | 42 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/px-server/src/main.rs b/px-server/src/main.rs
index 114e737..587931b 100644
--- a/px-server/src/main.rs
+++ b/px-server/src/main.rs
@@ -59,10 +59,50 @@ async fn main() -> Result<()> {
         .await
         .with_context(|| format!("bind {bind}"))?;
     tracing::info!(%bind, "px-server listening");
-    axum::serve(listener, app).await.context("axum serve")?;
+    axum::serve(listener, app)
+        .with_graceful_shutdown(shutdown_signal())
+        .await
+        .context("axum serve")?;
+    tracing::info!("px-server stopped cleanly");
     Ok(())
 }
 
+/// Resolves when the process should begin a graceful shutdown.
+///
+/// Listens for SIGINT (Ctrl-C from a TTY) and, on Unix, SIGTERM (sent by
+/// systemd, Kubernetes, and `kill`). The first signal wins. Axum's
+/// `with_graceful_shutdown` then stops accepting new connections and lets
+/// in-flight `/v1/solve` requests finish before returning.
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        if let Err(e) = tokio::signal::ctrl_c().await {
+            tracing::warn!(error = %e, "ctrl_c signal handler install failed");
+        }
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        use tokio::signal::unix::{SignalKind, signal};
+        match signal(SignalKind::terminate()) {
+            Ok(mut s) => {
+                s.recv().await;
+            }
+            Err(e) => {
+                tracing::warn!(error = %e, "SIGTERM handler install failed");
+                std::future::pending::<()>().await;
+            }
+        }
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => tracing::info!("SIGINT received, draining in-flight solves"),
+        _ = terminate => tracing::info!("SIGTERM received, draining in-flight solves"),
+    }
+}
+
 /// Collect the set of domains that should route through the Cloudflare
 /// (Camoufox) handler. Sources, unioned and lowercased:
 ///