fix(server): handle graceful process shutdown - #294
Conversation
Signed-off-by: nachiketb <nachiketb@nvidia.com>
WalkthroughChangesGraceful shutdown
Estimated code review effort: 4 (Complex) | ~45 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
crates/switchyard-server/src/lib.rs (2)
311-341: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winSchedule the shutdown task after the fallible listener setup.
In both
serveandserve_tls,schedule_shutdownspawns the task before the fallible calls that follow it. Iflistener.into_std()fails, or ifaxum_server::from_tcp_rustls/from_tcpfails, the function returns early through?and never reachesshutdown_task.abort(). The spawned task then stays alive for the process lifetime, awaiting a shutdown future that will callgraceful_shutdownon a handle whose server never started.Move the listener conversion and builder construction ahead of
schedule_shutdownso every error path returns before a task is spawned.Also add a short comment on
schedule_shutdown. It carries the shutdown lifecycle contract for both transports, and the guideline asks for comments on non-obvious private helpers and important async and lifecycle behavior.♻️ Proposed reordering for the non-TLS path
async fn serve( listener: TcpListener, router: Router, shutdown_timeout: Duration, shutdown: impl Future<Output = ()> + Send + 'static, ) -> ServerResult<()> { - let handle = axum_server::Handle::new(); - let shutdown_task = schedule_shutdown(handle.clone(), shutdown_timeout, shutdown); let std_listener = listener.into_std().map_err(server_io_error)?; - let result = axum_server::from_tcp(std_listener) - .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 shutdown_task = schedule_shutdown(handle.clone(), shutdown_timeout, shutdown); + let result = server .handle(handle) .serve(router.into_make_service()) .await .map_err(server_io_error); shutdown_task.abort(); result }Apply the same reordering in
serve_tls, movingRustlsConfig::from_pem_file,listener.into_std(), andaxum_server::from_tcp_rustlsbeforeschedule_shutdown.+/// Aborts nothing on its own: the caller must abort the returned task once the +/// server stops, otherwise the task outlives the server it was created for. fn schedule_shutdown(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/switchyard-server/src/lib.rs` around lines 311 - 341, In both serve and serve_tls, complete all fallible configuration, listener conversion, and axum_server builder construction before calling schedule_shutdown, then create the shutdown task only after those steps succeed and retain the existing abort after serving. Add a brief comment on schedule_shutdown documenting its shared transport shutdown lifecycle contract and async behavior.Source: Coding guidelines
1196-1240: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWiden the timing margins and split the two scenarios.
Both scenarios depend on real wall-clock timing under
#[tokio::test]. Two margins are tight enough to flake on a loaded CI machine:
- Line 1207: a 25 ms probe asserts the server has not finished. A scheduling stall that delays the shutdown task past 25 ms is not distinguishable from correct draining, so the assertion can fail spuriously.
- Line 1230: a 25 ms grace period is close to the runtime's own scheduling jitter, so the second scenario can pass without proving the deadline was enforced.
Increase the separation between the grace period and the probe window. Use a multi-second grace period with a sub-second probe in the first scenario, and a grace period of a few hundred milliseconds in the second.
Also split the function into two tests, one per scenario. A failure then names the behavior that broke instead of one shared test name.
♻️ Proposed margin change
- } = shutdown_test_server(Duration::from_secs(1)); + } = shutdown_test_server(Duration::from_secs(10)); state.started.notified().await; shutdown.send(()).expect("server receives shutdown"); assert!( - tokio::time::timeout(Duration::from_millis(25), &mut server) + tokio::time::timeout(Duration::from_millis(250), &mut server) .await .is_err(), "server must wait for the active request" );- } = shutdown_test_server(Duration::from_millis(25)); + } = shutdown_test_server(Duration::from_millis(250));🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/switchyard-server/src/lib.rs` around lines 1196 - 1240, Update the shutdown tests by splitting shutdown_drains_until_configured_deadline into two independently named tests, one covering request draining and one covering deadline enforcement. In the draining test, use a multi-second configured grace period and a sub-second timeout for the pre-release probe; in the deadline test, use a grace period of a few hundred milliseconds while retaining the bounded server completion assertion. Preserve each scenario’s existing request, release, and response assertions.crates/switchyard-py/src/server_bindings.rs (1)
49-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive the drain deadline from
DEFAULT_SHUTDOWN_TIMEOUT_SECS.Line 21 already defines
DEFAULT_SHUTDOWN_TIMEOUT_SECS: f64 = 2.0, andcloseand__exit__use it as the caller's wait budget. Line 49 repeats the same two-second value as a separate literal. The two values now encode one contract in two places, so a later change to one leaves the other stale.Note also that the two deadlines are equal.
closestarts its wait, the server then begins draining, and the server drain can reach its own deadline at the same momentclosegives up. A drain deadline strictly shorter than the caller's wait budget makesclosedeterministic. If a behavior change is out of scope for this PR, keep the value and only remove the duplication.♻️ Proposed deduplication
- shutdown_timeout: Duration::from_secs(2), + shutdown_timeout: Duration::from_secs_f64(DEFAULT_SHUTDOWN_TIMEOUT_SECS),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/switchyard-py/src/server_bindings.rs` at line 49, Update the server shutdown configuration near `shutdown_timeout` to derive its duration from `DEFAULT_SHUTDOWN_TIMEOUT_SECS` instead of duplicating the two-second literal. Preserve the current timeout value and behavior; only remove the duplicated constant usage unless an existing shorter drain-timeout convention is already defined.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@crates/switchyard-py/src/server_bindings.rs`:
- Line 49: Update the server shutdown configuration near `shutdown_timeout` to
derive its duration from `DEFAULT_SHUTDOWN_TIMEOUT_SECS` instead of duplicating
the two-second literal. Preserve the current timeout value and behavior; only
remove the duplicated constant usage unless an existing shorter drain-timeout
convention is already defined.
In `@crates/switchyard-server/src/lib.rs`:
- Around line 311-341: In both serve and serve_tls, complete all fallible
configuration, listener conversion, and axum_server builder construction before
calling schedule_shutdown, then create the shutdown task only after those steps
succeed and retain the existing abort after serving. Add a brief comment on
schedule_shutdown documenting its shared transport shutdown lifecycle contract
and async behavior.
- Around line 1196-1240: Update the shutdown tests by splitting
shutdown_drains_until_configured_deadline into two independently named tests,
one covering request draining and one covering deadline enforcement. In the
draining test, use a multi-second configured grace period and a sub-second
timeout for the pre-release probe; in the deadline test, use a grace period of a
few hundred milliseconds while retaining the bounded server completion
assertion. Preserve each scenario’s existing request, release, and response
assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a62ca03b-ec5a-49ee-bbf9-fa1ad5b4b113
📒 Files selected for processing (4)
crates/switchyard-py/src/server_bindings.rscrates/switchyard-server/README.mdcrates/switchyard-server/src/cli.rscrates/switchyard-server/src/lib.rs
Signed-off-by: nachiketb <nachiketb@nvidia.com>
Signed-off-by: nachiketb <nachiketb@nvidia.com>
2 participants
What
--shutdown-timeout, defaulting to 30 secondsWhy
Docker and Kubernetes terminate processes with SIGTERM. The server previously handled only Ctrl+C, plain HTTP could drain forever, and TLS allowed only two seconds. That could either drop telemetry and active requests immediately or block shutdown indefinitely.
How
Both transports now use the existing
axum-servergraceful-shutdown handle. The signal future remains a small Tokio wrapper because SIGTERM is Unix-specific; embedded callers can continue supplying their own shutdown future and deadline.What to review
shutdown_signalBoundServer::serveValidation
cargo test -p switchyard-servercargo clippy -p switchyard-server --all-targets -- -D warningscargo check -p switchyard-pySummary by CodeRabbit
New Features
--shutdown-timeout, defaulting to 30 seconds.SIGTERM, allowing active requests to complete before stopping.Bug Fixes
Documentation