fix(runtime): bound endpoint drain + route etcd lease loss through Runtime::shutdown()

lib/runtime/src/pipeline/network/ingress.rs

@@ -8,3 +8,96 @@ pub mod shared_tcp_endpoint;
 pub mod unified_server;
 
 use super::*;
+
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::time::Duration;
+use tokio::sync::Notify;
+
+/// Wait for inflight requests to drain, bounded by `timeout`. Returns the count
+/// still inflight when the wait ends: `0` on a clean drain, `>0` if the timeout
+/// fired (a stuck request must not wedge teardown). Shared by the NATS push and
+/// TCP request planes.
+async fn drain_inflight(
+    inflight: Arc<AtomicU64>,
+    notify: Arc<Notify>,
+    endpoint_name: &str,
+    timeout: Duration,
+) -> u64 {
+    let inflight_count = inflight.load(Ordering::SeqCst);
+    if inflight_count == 0 {
+        return 0;
+    }
+
+    tracing::info!(
+        endpoint_name,
+        inflight_count,
+        "Waiting for inflight requests to complete"
+    );
+
+    let wait = async {
+        while inflight.load(Ordering::SeqCst) > 0 {
+            notify.notified().await;
+        }
+    };
+
+    match tokio::time::timeout(timeout, wait).await {
+        Ok(()) => {
+            tracing::info!(endpoint_name, "All inflight requests completed");
+            0
+        }
+        Err(_) => {
+            let remaining = inflight.load(Ordering::SeqCst);
+            tracing::warn!(
+                endpoint_name,
+                timeout_secs = timeout.as_secs(),
+                remaining,
+                "Timed out waiting for inflight requests to drain; proceeding with shutdown"
+            );
+            remaining
+        }
+    }
+}
+
+#[cfg(test)]
+mod drain_tests {
+    use super::*;
+
+    /// Drain returns within the bound even if the request never completes.
+    #[tokio::test(start_paused = true)]
+    async fn drain_inflight_is_bounded_when_request_never_completes() {
+        let inflight = Arc::new(AtomicU64::new(1));
+        let notify = Arc::new(Notify::new());
+
+        let remaining = tokio::time::timeout(
+            Duration::from_secs(3600),
+            drain_inflight(inflight, notify, "test-endpoint", Duration::from_secs(5)),
+        )
+        .await
+        .expect("drain_inflight must be bounded; it hung past the outer guard");
+
+        assert_eq!(
+            remaining, 1,
+            "the stuck inflight request should still be counted as remaining"
+        );
+    }
+
+    /// A clean drain (request completes) returns 0 promptly.
+    #[tokio::test(start_paused = true)]
+    async fn drain_inflight_returns_zero_when_requests_complete() {
+        let inflight = Arc::new(AtomicU64::new(1));
+        let notify = Arc::new(Notify::new());
+
+        let inflight_clone = inflight.clone();
+        let notify_clone = notify.clone();
+        tokio::spawn(async move {
+            inflight_clone.fetch_sub(1, Ordering::SeqCst);
+            notify_clone.notify_one();
+        });
+
+        let remaining =
+            drain_inflight(inflight, notify, "test-endpoint", Duration::from_secs(5)).await;
+
+        assert_eq!(remaining, 0, "a completed request should drain cleanly");
+    }
+}

lib/runtime/src/pipeline/network/ingress/push_endpoint.rs

@@ -148,21 +148,13 @@ impl PushEndpoint {
 
         // await for all inflight requests to complete if graceful shutdown
         if self.graceful_shutdown {
-            let inflight_count = inflight.load(Ordering::SeqCst);
-            if inflight_count > 0 {
-                tracing::info!(
-                    endpoint_name = endpoint_name_local.as_str(),
-                    inflight_count = inflight_count,
-                    "Waiting for inflight NATS requests to complete"
-                );
-                while inflight.load(Ordering::SeqCst) > 0 {
-                    notify.notified().await;
-                }
-                tracing::info!(
-                    endpoint_name = endpoint_name_local.as_str(),
-                    "All inflight NATS requests completed"
-                );
-            }
+            super::drain_inflight(
+                inflight,
+                notify,
+                endpoint_name_local.as_str(),
+                crate::runtime::graceful_shutdown_timeout(),
+            )
+            .await;
         } else {
             tracing::info!(
                 endpoint_name = endpoint_name_local.as_str(),

lib/runtime/src/pipeline/network/ingress/shared_tcp_endpoint.rs

@@ -463,21 +463,13 @@ impl SharedTcpServer {
                 "Unregistered TCP endpoint handler"
             );
 
-            let inflight_count = handler.inflight.load(Ordering::SeqCst);
-            if inflight_count > 0 {
-                tracing::info!(
-                    endpoint_name = %endpoint_name,
-                    inflight_count = inflight_count,
-                    "Waiting for inflight TCP requests to complete"
-                );
-                while handler.inflight.load(Ordering::SeqCst) > 0 {
-                    handler.notify.notified().await;
-                }
-                tracing::info!(
-                    endpoint_name = %endpoint_name,
-                    "All inflight TCP requests completed"
-                );
-            }
+            super::drain_inflight(
+                handler.inflight.clone(),
+                handler.notify.clone(),
+                endpoint_name,
+                crate::runtime::graceful_shutdown_timeout(),
+            )
+            .await;
         }
     }
 

lib/runtime/src/runtime.rs

@@ -32,7 +32,7 @@ pub use tokio_util::sync::CancellationToken;
 
 const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS: u64 = 15 * 60;
 
-fn graceful_shutdown_timeout() -> Duration {
+pub(crate) fn graceful_shutdown_timeout() -> Duration {
     let timeout_secs = std::env::var(
         config::environment_names::runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS,
     )

lib/runtime/src/transports/etcd.rs

@@ -69,10 +69,10 @@ impl Client {
     /// If the lease expires, the [`Runtime`] will be shutdown.
     /// If the [`Runtime`] is shutdown, the lease will be revoked.
     pub async fn new(config: ClientOptions, runtime: Runtime) -> Result<Self> {
-        let token = runtime.primary_token();
+        let runtime_for_lease = runtime.clone();
 
         let ((connector, lease_id), rt) = build_in_runtime(
-            async move { Self::connect_with_startup_retry(&config, token).await },
+            async move { Self::connect_with_startup_retry(&config, runtime_for_lease).await },
             1,
         )
         .await?;
@@ -88,8 +88,9 @@ impl Client {
     /// Connect to etcd during startup, retrying with exponential backoff for up to 2 minutes.
     async fn connect_with_startup_retry(
         config: &ClientOptions,
-        token: CancellationToken,
+        runtime: Runtime,
     ) -> Result<(Arc<Connector>, u64)> {
+        let token = runtime.primary_token();
         let deadline = Instant::now() + STARTUP_CONNECT_TIMEOUT;
         let mut backoff = STARTUP_CONNECT_INITIAL_BACKOFF;
 
@@ -98,7 +99,7 @@ impl Client {
                 anyhow::bail!("etcd startup connection cancelled");
             }
 
-            let attempt = Self::connect_startup_attempt(config, &token).await;
+            let attempt = Self::connect_startup_attempt(config, &runtime).await;
 
             match attempt {
                 Ok(connection) => return Ok(connection),
@@ -134,13 +135,14 @@ impl Client {
 
     async fn connect_startup_attempt(
         config: &ClientOptions,
-        token: &CancellationToken,
+        runtime: &Runtime,
     ) -> Result<(Arc<Connector>, u64)> {
+        let token = runtime.primary_token();
         let connector =
             Connector::new(config.etcd_url.clone(), config.etcd_connect_options.clone()).await?;
 
         let lease_id = if config.attach_lease {
-            create_lease(connector.clone(), config.lease_ttl, token.clone())
+            create_lease(connector.clone(), config.lease_ttl, runtime.clone())
                 .await
                 .with_context(|| {
                     format!(

lib/runtime/src/transports/etcd/lease.rs

@@ -2,21 +2,24 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::connector::Connector;
+use crate::runtime::Runtime;
 use etcd_client::{LeaseKeepAliveStream, LeaseKeeper};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio_util::sync::CancellationToken;
 
-/// Create an etcd lease with the given TTL, attach it to the provided cancellation token,
+/// Create an etcd lease with the given TTL, tie its lifetime to the [`Runtime`],
 /// spawn a keep-alive task, and return the lease id (u64).
 ///
-/// Note: this function spawns a background task that maintains the lease until the token is
-/// cancelled or an unrecoverable error occurs.
+/// Note: this function spawns a background task that maintains the lease until the runtime is
+/// shut down or an unrecoverable error occurs. On an unrecoverable error the runtime is shut
+/// down, honoring the contract that a lost lease shuts the worker down.
 pub async fn create_lease(
     connector: Arc<Connector>,
     ttl: u64,
-    token: CancellationToken,
+    runtime: Runtime,
 ) -> anyhow::Result<u64> {
+    let token = runtime.primary_token();
     if token.is_cancelled() {
         anyhow::bail!("lease creation cancelled");
     }
@@ -47,7 +50,9 @@ pub async fn create_lease(
                     error = %e,
                     "Unable to maintain lease. Check etcd server status"
                 );
-                token.cancel();
+                // Phased shutdown (endpoint drain -> backend teardown), not a
+                // bare primary-token cancel, so teardown is ordered.
+                runtime.shutdown();
             }
         }
     });
@@ -204,3 +209,44 @@ async fn keep_alive_with_stream(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The lease-loss teardown path (`runtime.shutdown()`) must be phased:
+    /// Phase 1 cancels the endpoint token promptly, but Phase 3 (primary token
+    /// + backend teardown) waits for outstanding graceful tasks.
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn lease_loss_shutdown_is_phased() {
+        let runtime = Runtime::from_current().unwrap();
+
+        // Stands in for a serving endpoint; Phase 1 cancels the endpoint
+        // shutdown token, cascading to this child.
+        let endpoint_token = runtime.child_token();
+        let primary_token = runtime.primary_token();
+
+        // Hold a graceful task so Phase 2 cannot advance to Phase 3 yet.
+        let guard = runtime.graceful_shutdown_tracker().register_task();
+
+        runtime.shutdown();
+
+        // Phase 1 cancels the endpoint token promptly.
+        tokio::time::timeout(Duration::from_secs(5), endpoint_token.cancelled())
+            .await
+            .expect("Phase 1 must cancel the endpoint token on lease loss");
+
+        // While the graceful task is outstanding, Phase 3 must not have torn
+        // down the primary token. A bare-cancel regression fails this assert.
+        assert!(
+            !primary_token.is_cancelled(),
+            "primary token must not be cancelled while a graceful task is outstanding"
+        );
+
+        // Releasing the task lets Phase 3 proceed and cancel the primary token.
+        drop(guard);
+        tokio::time::timeout(Duration::from_secs(5), primary_token.cancelled())
+            .await
+            .expect("Phase 3 must cancel the primary token once graceful tasks complete");
+    }
+}