From 93da12ddd5a11db08eb2f99e6f59e96c6b34b00d Mon Sep 17 00:00:00 2001 From: Michael Feil <63565275+michaelfeil@users.noreply.github.com> Date: Sun, 14 Jun 2026 15:42:59 -0700 Subject: [PATCH 1/3] fix(runtime): bound graceful shutdown drain Signed-off-by: Michael Feil <63565275+michaelfeil@users.noreply.github.com> --- docs/fault-tolerance/graceful-shutdown.md | 7 +- lib/runtime/src/config/environment_names.rs | 6 ++ lib/runtime/src/runtime.rs | 77 +++++++++++++++++++-- 3 files changed, 85 insertions(+), 5 deletions(-) diff --git a/docs/fault-tolerance/graceful-shutdown.md b/docs/fault-tolerance/graceful-shutdown.md index 65244b6838ce..d7a4cfa5a6da 100644 --- a/docs/fault-tolerance/graceful-shutdown.md +++ b/docs/fault-tolerance/graceful-shutdown.md @@ -45,6 +45,11 @@ The `graceful_shutdown()` function: 5. Waits for in-flight requests (based on `graceful_shutdown` per endpoint) 6. Returns to allow cleanup to proceed +The aggregate wait in `runtime.shutdown()` is bounded by +`DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS`, which defaults to 900 seconds +(15 minutes). If endpoint draining exceeds this timeout, Dynamo logs the +remaining graceful endpoint count and proceeds with runtime teardown. + ## Endpoint Draining After the grace period, `runtime.shutdown()` invalidates endpoints so no new requests are accepted. The behavior for in-flight requests depends on the `graceful_shutdown` parameter when serving the endpoint. @@ -64,7 +69,7 @@ generate_endpoint.serve_endpoint( | `graceful_shutdown` | Behavior | |---------------------|----------| -| `True` | Wait for all in-flight requests to complete before returning | +| `True` | Wait for all in-flight requests to complete, bounded by `DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS` | | `False` | Return immediately without waiting for requests | ### Component-Specific Behavior diff --git a/lib/runtime/src/config/environment_names.rs b/lib/runtime/src/config/environment_names.rs index 5463921436f1..1e5a4c21b6ed 100644 --- a/lib/runtime/src/config/environment_names.rs +++ b/lib/runtime/src/config/environment_names.rs @@ -70,6 +70,10 @@ pub mod runtime { /// Maximum number of blocking threads for Tokio runtime pub const DYN_RUNTIME_MAX_BLOCKING_THREADS: &str = "DYN_RUNTIME_MAX_BLOCKING_THREADS"; + /// Maximum time to wait for graceful endpoint drain during runtime shutdown. + pub const DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS: &str = + "DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS"; + /// Enable Tokio task poll-time histogram (calls enable_metrics_poll_time_histogram on builder). /// Set to "1", "true", or "yes" to enable. Adds ~2× overhead of Instant::now() per task poll. pub const DYN_ENABLE_POLL_HISTOGRAM: &str = "DYN_ENABLE_POLL_HISTOGRAM"; @@ -654,6 +658,7 @@ mod tests { // Runtime runtime::DYN_RUNTIME_NUM_WORKER_THREADS, runtime::DYN_RUNTIME_MAX_BLOCKING_THREADS, + runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, runtime::system::DYN_SYSTEM_ENABLED, runtime::system::DYN_SYSTEM_HOST, runtime::system::DYN_SYSTEM_PORT, @@ -791,6 +796,7 @@ mod tests { fn test_naming_conventions() { // Dynamo-specific vars should start with DYN_ assert!(runtime::DYN_RUNTIME_NUM_WORKER_THREADS.starts_with("DYN_")); + assert!(runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS.starts_with("DYN_")); assert!(runtime::system::DYN_SYSTEM_ENABLED.starts_with("DYN_")); assert!(kvbm::DYN_KVBM_METRICS.starts_with("DYN_")); assert!(worker::DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT.starts_with("DYN_")); diff --git a/lib/runtime/src/runtime.rs b/lib/runtime/src/runtime.rs index dffdd78a492c..332b887719d5 100644 --- a/lib/runtime/src/runtime.rs +++ b/lib/runtime/src/runtime.rs @@ -24,11 +24,25 @@ use once_cell::sync::OnceCell; use std::{ mem::ManuallyDrop, sync::{Arc, atomic::Ordering}, + time::Duration, }; use tokio::{signal, sync::Mutex, task::JoinHandle}; pub use tokio_util::sync::CancellationToken; +const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS: u64 = 15 * 60; + +fn graceful_shutdown_timeout() -> Duration { + let timeout_secs = std::env::var( + config::environment_names::runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, + ) + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS); + + Duration::from_secs(timeout_secs) +} + /// Types of Tokio runtimes that can be used to construct a Dynamo [Runtime]. #[derive(Clone, Debug)] enum RuntimeType { @@ -330,18 +344,73 @@ impl Runtime { tracing::info!("Active graceful endpoints: {count}"); if count != 0 { - tracker.wait_for_completion().await; + let timeout = graceful_shutdown_timeout(); + if tokio::time::timeout(timeout, tracker.wait_for_completion()) + .await + .is_err() + { + let remaining = tracker.get_count(); + tracing::error!( + unified_model_logs = true, + timeout_secs = timeout.as_secs(), + remaining_endpoints = remaining, + "Graceful endpoint shutdown timed out; proceeding with runtime teardown" + ); + } } // Phase 3: Now connections will be disconnected to backend services (e.g. NATS/ETCD) by cancelling the main token - tracing::info!( - "Phase 3: All endpoints ended gracefully. Connections to backend services will now be disconnected" - ); + tracing::info!("Phase 3: Connections to backend services will now be disconnected"); main_token.cancel(); }); } } +#[cfg(test)] +mod tests { + use super::*; + use crate::config::environment_names::runtime as env_runtime; + + #[test] + fn graceful_shutdown_timeout_defaults_to_15_minutes() { + temp_env::with_var( + env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, + None::<&str>, + || { + assert_eq!( + graceful_shutdown_timeout(), + Duration::from_secs(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS) + ); + }, + ); + } + + #[test] + fn graceful_shutdown_timeout_uses_env_override() { + temp_env::with_var( + env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, + Some("42"), + || { + assert_eq!(graceful_shutdown_timeout(), Duration::from_secs(42)); + }, + ); + } + + #[test] + fn graceful_shutdown_timeout_defaults_for_invalid_env() { + temp_env::with_var( + env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, + Some("not-a-number"), + || { + assert_eq!( + graceful_shutdown_timeout(), + Duration::from_secs(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS) + ); + }, + ); + } +} + impl RuntimeType { /// Get [`tokio::runtime::Handle`] to runtime pub fn handle(&self) -> tokio::runtime::Handle { From b3f5f4af8dbcbb03afa63893b30cfe4f57f6afe6 Mon Sep 17 00:00:00 2001 From: Michael Feil <63565275+michaelfeil@users.noreply.github.com> Date: Sun, 14 Jun 2026 15:49:30 -0700 Subject: [PATCH 2/3] test(runtime): cover graceful shutdown timeout Signed-off-by: Michael Feil <63565275+michaelfeil@users.noreply.github.com> --- lib/runtime/src/runtime.rs | 112 ++++++++++++++++++++++--------------- 1 file changed, 67 insertions(+), 45 deletions(-) diff --git a/lib/runtime/src/runtime.rs b/lib/runtime/src/runtime.rs index 332b887719d5..e745cae890a1 100644 --- a/lib/runtime/src/runtime.rs +++ b/lib/runtime/src/runtime.rs @@ -366,51 +366,6 @@ impl Runtime { } } -#[cfg(test)] -mod tests { - use super::*; - use crate::config::environment_names::runtime as env_runtime; - - #[test] - fn graceful_shutdown_timeout_defaults_to_15_minutes() { - temp_env::with_var( - env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, - None::<&str>, - || { - assert_eq!( - graceful_shutdown_timeout(), - Duration::from_secs(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS) - ); - }, - ); - } - - #[test] - fn graceful_shutdown_timeout_uses_env_override() { - temp_env::with_var( - env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, - Some("42"), - || { - assert_eq!(graceful_shutdown_timeout(), Duration::from_secs(42)); - }, - ); - } - - #[test] - fn graceful_shutdown_timeout_defaults_for_invalid_env() { - temp_env::with_var( - env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, - Some("not-a-number"), - || { - assert_eq!( - graceful_shutdown_timeout(), - Duration::from_secs(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS) - ); - }, - ); - } -} - impl RuntimeType { /// Get [`tokio::runtime::Handle`] to runtime pub fn handle(&self) -> tokio::runtime::Handle { @@ -460,3 +415,70 @@ impl Drop for RuntimeType { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::environment_names::runtime as env_runtime; + + #[test] + fn graceful_shutdown_timeout_uses_env_override() { + temp_env::with_var( + env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, + Some("42"), + || { + assert_eq!(graceful_shutdown_timeout(), Duration::from_secs(42)); + }, + ); + } + + #[test] + fn graceful_shutdown_timeout_defaults_for_invalid_env() { + temp_env::with_var( + env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, + Some("not-a-number"), + || { + assert_eq!( + graceful_shutdown_timeout(), + Duration::from_secs(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS) + ); + }, + ); + } + + #[tokio::test(start_paused = true)] + async fn shutdown_cancels_main_token_after_graceful_timeout() { + temp_env::async_with_vars( + [( + env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, + Some("5"), + )], + async { + let runtime = Runtime::from_current().unwrap(); + let tracker = runtime.graceful_shutdown_tracker(); + let _guard = tracker.register_task(); + let main_token = runtime.primary_token(); + let endpoint_token = runtime.child_token(); + + runtime.shutdown(); + tokio::task::yield_now().await; + + assert!(endpoint_token.is_cancelled()); + assert!(!main_token.is_cancelled()); + assert_eq!(tracker.get_count(), 1); + + tokio::time::advance(Duration::from_secs(4)).await; + tokio::task::yield_now().await; + + assert!(!main_token.is_cancelled()); + + tokio::time::advance(Duration::from_secs(1)).await; + tokio::task::yield_now().await; + + assert!(main_token.is_cancelled()); + assert_eq!(tracker.get_count(), 1); + }, + ) + .await; + } +} From 95f16346e2df063af41eb0f34d71c18d159aeabd Mon Sep 17 00:00:00 2001 From: Michael Feil <63565275+michaelfeil@users.noreply.github.com> Date: Mon, 15 Jun 2026 08:58:04 -0700 Subject: [PATCH 3/3] test(runtime): address graceful shutdown review feedback Signed-off-by: Michael Feil <63565275+michaelfeil@users.noreply.github.com> --- lib/runtime/src/runtime.rs | 26 -------------------------- 1 file changed, 26 deletions(-) diff --git a/lib/runtime/src/runtime.rs b/lib/runtime/src/runtime.rs index e745cae890a1..aaee61ea9939 100644 --- a/lib/runtime/src/runtime.rs +++ b/lib/runtime/src/runtime.rs @@ -351,7 +351,6 @@ impl Runtime { { let remaining = tracker.get_count(); tracing::error!( - unified_model_logs = true, timeout_secs = timeout.as_secs(), remaining_endpoints = remaining, "Graceful endpoint shutdown timed out; proceeding with runtime teardown" @@ -421,31 +420,6 @@ mod tests { use super::*; use crate::config::environment_names::runtime as env_runtime; - #[test] - fn graceful_shutdown_timeout_uses_env_override() { - temp_env::with_var( - env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, - Some("42"), - || { - assert_eq!(graceful_shutdown_timeout(), Duration::from_secs(42)); - }, - ); - } - - #[test] - fn graceful_shutdown_timeout_defaults_for_invalid_env() { - temp_env::with_var( - env_runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, - Some("not-a-number"), - || { - assert_eq!( - graceful_shutdown_timeout(), - Duration::from_secs(DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS) - ); - }, - ); - } - #[tokio::test(start_paused = true)] async fn shutdown_cancels_main_token_after_graceful_timeout() { temp_env::async_with_vars(