diff --git a/crates/lance-context-master/src/main.rs b/crates/lance-context-master/src/main.rs index 611adeb1..6fa3c1a6 100644 --- a/crates/lance-context-master/src/main.rs +++ b/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"); } diff --git a/crates/lance-context-server/src/main.rs b/crates/lance-context-server/src/main.rs index 9ca9f4fc..5a9862a5 100644 --- a/crates/lance-context-server/src/main.rs +++ b/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"); } diff --git a/python/tests/conftest.py b/python/tests/conftest.py index 99d05500..6f6f8261 100644 --- a/python/tests/conftest.py +++ b/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})" + )