fix(server): handle graceful process shutdown (#294)

* fix(server): handle graceful process shutdown

Signed-off-by: nachiketb <nachiketb@nvidia.com>

* refactor(server): isolate platform shutdown signals

Signed-off-by: nachiketb <nachiketb@nvidia.com>

* refactor(server): await shutdown without spawning

Signed-off-by: nachiketb <nachiketb@nvidia.com>

---------

Signed-off-by: nachiketb <nachiketb@nvidia.com>

Author: nachiketb-nvidia

crates/switchyard-py/src/server_bindings.rs

@@ -46,6 +46,7 @@ impl PyServer {
                     addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port),
                     backlog: DEFAULT_LISTEN_BACKLOG,
                     dry_run: false,
+                    shutdown_timeout: Duration::from_secs(2),
                     tls: None,
                 },
             )

crates/switchyard-server/README.md

@@ -58,6 +58,9 @@ export API_KEY="..."
 cargo run -p switchyard-server -- --config routes.toml
 ```
 
+Ctrl+C and Unix `SIGTERM` stop new connections and allow active requests to drain for up to
+`--shutdown-timeout` (30 seconds by default) before they are terminated.
+
 The server logs exactly one structured terminal event per LLM request: successful responses at
 `INFO`, 4xx responses at `WARN`, and 5xx responses at `ERROR`. Set
 `RUST_LOG=switchyard_server=debug,libsy=debug` to include routing decisions and nested failure

crates/switchyard-server/src/cli.rs

@@ -9,8 +9,8 @@ use std::path::PathBuf;
 use clap::Parser;
 use switchyard_server::config::load_server_state;
 use switchyard_server::{
-    DEFAULT_LISTEN_BACKLOG, ServerError, ServerResult, ServerRunOptions, ServerState, TlsOptions,
-    run_server,
+    DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT, DEFAULT_LISTEN_BACKLOG, ServerError, ServerResult,
+    ServerRunOptions, ServerState, TlsOptions, run_server,
 };
 
 const DEFAULT_HOST: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED);
@@ -40,6 +40,10 @@ pub(crate) struct ServerArgs {
     #[arg(long, default_value_t = DEFAULT_LISTEN_BACKLOG)]
     backlog: u32,
 
+    /// Maximum time active requests may drain during shutdown.
+    #[arg(long, default_value_t = humantime::Duration::from(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT))]
+    shutdown_timeout: humantime::Duration,
+
     /// Validate the algorithm and client configuration without binding a socket.
     #[arg(long)]
     dry_run: bool,
@@ -85,6 +89,7 @@ impl ServerArgs {
             addr: SocketAddr::new(self.host, self.port),
             backlog: self.backlog,
             dry_run: self.dry_run,
+            shutdown_timeout: self.shutdown_timeout.into(),
             tls,
         };
         Ok((state, options))

crates/switchyard-server/src/lib.rs

@@ -8,6 +8,7 @@ mod metrics;
 mod observability;
 mod response;
 mod routing_log;
+mod shutdown;
 mod sse;
 mod stats;
 mod usage_metrics;
@@ -50,6 +51,9 @@ pub use observability::{flush_observability, initialize_observability};
 /// Default TCP listen backlog used by the Rust server.
 pub const DEFAULT_LISTEN_BACKLOG: u32 = 65_535;
 
+/// Default time allowed for active requests to finish during shutdown.
+pub const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30);
+
 /// Maximum buffered JSON request size accepted by the LLM endpoints.
 pub const DEFAULT_MAX_REQUEST_BODY_BYTES: usize = 32 * 1024 * 1024;
 
@@ -214,6 +218,8 @@ pub struct ServerRunOptions {
     pub backlog: u32,
     /// Validate runtime construction without binding a socket.
     pub dry_run: bool,
+    /// Maximum time active requests may drain after shutdown begins.
+    pub shutdown_timeout: Duration,
     /// TLS certificate configuration, when HTTPS is enabled.
     pub tls: Option<TlsOptions>,
 }
@@ -242,7 +248,7 @@ pub async fn run_server(state: ServerState, options: ServerRunOptions) -> Server
 
     let server = BoundServer::bind(state, options)?;
     println!("{}", server.startup_banner(std::io::stdout().is_terminal()));
-    server.serve(shutdown_signal()).await
+    server.serve(shutdown::signal()).await
 }
 
 /// A configured server with its listening socket already bound.
@@ -276,10 +282,11 @@ impl BoundServer {
         self,
         shutdown: impl Future<Output = ()> + Send + 'static,
     ) -> ServerResult<()> {
+        let shutdown_timeout = self.options.shutdown_timeout;
         if let Some(tls) = self.options.tls {
-            serve_tls(self.listener, self.router, tls, shutdown).await
+            serve_tls(self.listener, self.router, tls, shutdown_timeout, shutdown).await
         } else {
-            serve(self.listener, self.router, shutdown).await
+            serve(self.listener, self.router, shutdown_timeout, shutdown).await
         }
     }
 
@@ -292,6 +299,7 @@ async fn serve_tls(
     listener: TcpListener,
     router: Router,
     tls: TlsOptions,
+    shutdown_timeout: Duration,
     shutdown: impl Future<Output = ()> + Send + 'static,
 ) -> ServerResult<()> {
     if let Err(error) = rustls::crypto::aws_lc_rs::default_provider().install_default() {
@@ -301,32 +309,49 @@ async fn serve_tls(
     let config = RustlsConfig::from_pem_file(tls.cert, tls.key)
         .await
         .map_err(server_io_error)?;
-    let handle = axum_server::Handle::new();
-
-    let shutdown_handle = handle.clone();
-    tokio::spawn(async move {
-        shutdown.await;
-        shutdown_handle.graceful_shutdown(Some(Duration::from_secs(2)));
-    });
-
     let std_listener = listener.into_std().map_err(server_io_error)?;
-    axum_server::from_tcp_rustls(std_listener, config)
-        .map_err(server_io_error)?
-        .handle(handle)
-        .serve(router.into_make_service())
-        .await
-        .map_err(server_io_error)
+    let server = axum_server::from_tcp_rustls(std_listener, config).map_err(server_io_error)?;
+    let handle = axum_server::Handle::new();
+    let server = server
+        .handle(handle.clone())
+        .serve(router.into_make_service());
+    serve_until_shutdown(server, handle, shutdown_timeout, shutdown).await
 }
 
 async fn serve(
     listener: TcpListener,
     router: Router,
+    shutdown_timeout: Duration,
     shutdown: impl Future<Output = ()> + Send + 'static,
 ) -> ServerResult<()> {
-    axum::serve(listener, router)
-        .with_graceful_shutdown(shutdown)
-        .await
-        .map_err(server_io_error)
+    let std_listener = listener.into_std().map_err(server_io_error)?;
+    let server = axum_server::from_tcp(std_listener).map_err(server_io_error)?;
+    let handle = axum_server::Handle::new();
+    let server = server
+        .handle(handle.clone())
+        .serve(router.into_make_service());
+    serve_until_shutdown(server, handle, shutdown_timeout, shutdown).await
+}
+
+/// Runs the server until it exits or shutdown begins, then drains active requests.
+async fn serve_until_shutdown(
+    server: impl Future<Output = std::io::Result<()>>,
+    handle: axum_server::Handle<SocketAddr>,
+    timeout: Duration,
+    shutdown: impl Future<Output = ()> + Send + 'static,
+) -> ServerResult<()> {
+    tokio::pin!(server);
+    tokio::select! {
+        result = &mut server => result.map_err(server_io_error),
+        _ = shutdown => {
+            tracing::info!(
+                ?timeout,
+                "shutdown signal received; draining active requests"
+            );
+            handle.graceful_shutdown(Some(timeout));
+            server.await.map_err(server_io_error)
+        }
+    }
 }
 
 /// Ingress timestamp for one request, taken before any body is read.
@@ -411,16 +436,6 @@ fn server_io_error(error: std::io::Error) -> ServerError {
     ServerError::new(error.to_string())
 }
 
-async fn shutdown_signal() {
-    if let Err(error) = tokio::signal::ctrl_c().await {
-        tracing::warn!(
-            error = %error,
-            "ctrl-c shutdown signal unavailable; continuing without shutdown trigger"
-        );
-        std::future::pending::<()>().await;
-    }
-}
-
 async fn openai_chat_completions(
     State(state): State<ServerState>,
     Extension(started): Extension<RequestStart>,
@@ -1079,8 +1094,111 @@ fn endpoint_listing(has_routing_log: bool) -> String {
 
 #[cfg(test)]
 mod tests {
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+    use tokio::sync::{Notify, oneshot};
+
     use super::*;
 
+    #[derive(Clone)]
+    struct ShutdownTestState {
+        started: Arc<Notify>,
+        release: Arc<Notify>,
+    }
+
+    struct ShutdownTestServer {
+        state: ShutdownTestState,
+        shutdown: oneshot::Sender<()>,
+        server: task::JoinHandle<ServerResult<()>>,
+        request: task::JoinHandle<std::io::Result<Vec<u8>>>,
+    }
+
+    async fn blocked_request(State(state): State<ShutdownTestState>) -> &'static str {
+        state.started.notify_one();
+        state.release.notified().await;
+        "done"
+    }
+
+    async fn raw_request(addr: SocketAddr) -> std::io::Result<Vec<u8>> {
+        let mut stream = tokio::net::TcpStream::connect(addr).await?;
+        stream
+            .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")
+            .await?;
+        let mut response = Vec::new();
+        stream.read_to_end(&mut response).await?;
+        Ok(response)
+    }
+
+    fn shutdown_test_server(shutdown_timeout: Duration) -> ShutdownTestServer {
+        let state = ShutdownTestState {
+            started: Arc::new(Notify::new()),
+            release: Arc::new(Notify::new()),
+        };
+        let router = Router::new()
+            .route("/", get(blocked_request))
+            .with_state(state.clone());
+        let listener = bind_tcp_listener("127.0.0.1:0".parse().expect("valid address"), 16)
+            .expect("listener binds");
+        let addr = listener.local_addr().expect("listener has an address");
+        let (shutdown, shutdown_receiver) = oneshot::channel();
+        let server = tokio::spawn(serve(listener, router, shutdown_timeout, async move {
+            let _ = shutdown_receiver.await;
+        }));
+        let request = tokio::spawn(raw_request(addr));
+        ShutdownTestServer {
+            state,
+            shutdown,
+            server,
+            request,
+        }
+    }
+
+    // Active requests may finish within the grace period, while stuck requests are bounded.
+    #[tokio::test]
+    async fn shutdown_drains_until_configured_deadline() {
+        let ShutdownTestServer {
+            state,
+            shutdown,
+            mut server,
+            request,
+        } = shutdown_test_server(Duration::from_secs(1));
+        state.started.notified().await;
+        shutdown.send(()).expect("server receives shutdown");
+        assert!(
+            tokio::time::timeout(Duration::from_millis(25), &mut server)
+                .await
+                .is_err(),
+            "server must wait for the active request"
+        );
+        state.release.notify_one();
+        tokio::time::timeout(Duration::from_secs(1), server)
+            .await
+            .expect("server stops after request drains")
+            .expect("server task completes")
+            .expect("server exits cleanly");
+        let response = request
+            .await
+            .expect("request task completes")
+            .expect("request succeeds");
+        assert!(response.windows(8).any(|part| part == b"200 OK\r\n"));
+        assert!(response.ends_with(b"done"));
+
+        let ShutdownTestServer {
+            state,
+            shutdown,
+            server,
+            request,
+        } = shutdown_test_server(Duration::from_millis(25));
+        state.started.notified().await;
+        shutdown.send(()).expect("server receives shutdown");
+        tokio::time::timeout(Duration::from_secs(1), server)
+            .await
+            .expect("shutdown deadline is enforced")
+            .expect("server task completes")
+            .expect("server exits cleanly");
+        state.release.notify_one();
+        request.abort();
+    }
+
     // Terminal request severity follows HTTP status instead of error-path bookkeeping.
     #[test]
     fn request_log_level_follows_http_status() {

crates/switchyard-server/src/shutdown.rs

@@ -0,0 +1,52 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Platform-specific process shutdown signals.
+
+/// Waits for the platform's normal process termination signal.
+pub(crate) async fn signal() {
+    platform::signal().await;
+}
+
+async fn ctrl_c() {
+    if let Err(error) = tokio::signal::ctrl_c().await {
+        tracing::warn!(
+            error = %error,
+            "ctrl-c shutdown signal unavailable; continuing without shutdown trigger"
+        );
+        std::future::pending::<()>().await;
+    }
+}
+
+#[cfg(unix)]
+mod platform {
+    pub(super) async fn signal() {
+        tokio::select! {
+            _ = super::ctrl_c() => {},
+            _ = terminate() => {},
+        }
+    }
+
+    async fn terminate() {
+        let mut signal =
+            match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
+                Ok(signal) => signal,
+                Err(error) => {
+                    tracing::warn!(
+                        error = %error,
+                        "SIGTERM shutdown signal unavailable; continuing without SIGTERM trigger"
+                    );
+                    std::future::pending::<()>().await;
+                    return;
+                }
+            };
+        signal.recv().await;
+    }
+}
+
+#[cfg(not(unix))]
+mod platform {
+    pub(super) async fn signal() {
+        super::ctrl_c().await;
+    }
+}