Handle SIGTERM gracefully in server and master

crates/lance-context-master/src/main.rs

@@ -71,14 +71,34 @@ async fn main() {
 
     let listener = TcpListener::bind(&addr).await.unwrap();
     axum::serve(listener, app)
-        .with_graceful_shutdown(shutdown_signal())
+        .with_graceful_shutdown(wait_for_shutdown_signal())
         .await
         .unwrap();
 }
 
-async fn shutdown_signal() {
+#[cfg(unix)]
+async fn wait_for_shutdown_signal() {
+    use tokio::signal::unix::{signal as unix_signal, SignalKind};
+
+    // Kubernetes and other process managers stop containers with SIGTERM.
+    let mut sigterm =
+        unix_signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
+
+    tokio::select! {
+        result = tokio::signal::ctrl_c() => {
+            result.expect("failed to install Ctrl+C handler");
+            tracing::info!(signal = "SIGINT", "Shutting down");
+        }
+        _ = sigterm.recv() => {
+            tracing::info!(signal = "SIGTERM", "Shutting down");
+        }
+    }
+}
+
+#[cfg(not(unix))]
+async fn wait_for_shutdown_signal() {
     tokio::signal::ctrl_c()
         .await
         .expect("failed to install Ctrl+C handler");
-    tracing::info!("Shutting down");
+    tracing::info!(signal = "Ctrl+C", "Shutting down");
 }

crates/lance-context-server/src/main.rs

@@ -81,19 +81,38 @@ async fn main() {
 
     let listener = TcpListener::bind(&addr).await.unwrap();
     axum::serve(listener, app)
-        .with_graceful_shutdown(shutdown_signal())
+        .with_graceful_shutdown(wait_for_shutdown_signal())
         .await
         .unwrap();
 
-    // Connections have drained. Deterministically close resident rollout writers
-    // (whose `ShardWriter` background tasks need an explicit `close().await`)
+    // Connections have drained. Deterministically close every resident writer
     // before the runtime tears down.
     state.shutdown().await;
 }
 
-async fn shutdown_signal() {
+#[cfg(unix)]
+async fn wait_for_shutdown_signal() {
+    use tokio::signal::unix::{signal as unix_signal, SignalKind};
+
+    // Kubernetes and other process managers stop containers with SIGTERM.
+    let mut sigterm =
+        unix_signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
+
+    tokio::select! {
+        result = tokio::signal::ctrl_c() => {
+            result.expect("failed to install Ctrl+C handler");
+            tracing::info!(signal = "SIGINT", "Shutting down");
+        }
+        _ = sigterm.recv() => {
+            tracing::info!(signal = "SIGTERM", "Shutting down");
+        }
+    }
+}
+
+#[cfg(not(unix))]
+async fn wait_for_shutdown_signal() {
     tokio::signal::ctrl_c()
         .await
         .expect("failed to install Ctrl+C handler");
-    tracing::info!("Shutting down");
+    tracing::info!(signal = "Ctrl+C", "Shutting down");
 }

python/tests/conftest.py

@@ -83,3 +83,10 @@ def server() -> Iterator[str]:
                     proc.wait(timeout=10)
                 log.seek(0)
                 print(log.read().decode(errors="replace"))
+                if os.name == "posix":
+                    # Popen.terminate sends SIGTERM on POSIX. A clean exit
+                    # proves the server reached its graceful-shutdown path.
+                    assert proc.returncode == 0, (
+                        "server did not handle SIGTERM gracefully "
+                        f"(exit code {proc.returncode})"
+                    )