diff --git a/lib/runtime/src/pipeline/network/ingress/push_endpoint.rs b/lib/runtime/src/pipeline/network/ingress/push_endpoint.rs index f7d71c110dae..1796934e649f 100644 --- a/lib/runtime/src/pipeline/network/ingress/push_endpoint.rs +++ b/lib/runtime/src/pipeline/network/ingress/push_endpoint.rs @@ -58,11 +58,6 @@ impl PushEndpoint { let req = tokio::select! { biased; - // await on service request - req = endpoint.next() => { - req - } - // process shutdown _ = self.cancellation_token.cancelled() => { tracing::info!("PushEndpoint received cancellation signal, shutting down service"); @@ -71,6 +66,11 @@ impl PushEndpoint { } break; } + + // Cancellation wins even when the request queue stays ready. + req = endpoint.next() => { + req + } }; if let Some(req) = req { diff --git a/lib/runtime/tests/nats_shutdown.rs b/lib/runtime/tests/nats_shutdown.rs new file mode 100644 index 000000000000..90b67605bf76 --- /dev/null +++ b/lib/runtime/tests/nats_shutdown.rs @@ -0,0 +1,111 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Requires a NATS server at NATS_SERVER (default: nats://127.0.0.1:4222). +#![cfg(feature = "integration")] + +use std::sync::{ + Arc, + atomic::{AtomicU64, Ordering}, +}; +use std::time::Duration; + +use async_nats::service::ServiceExt; +use dynamo_runtime::pipeline::{ + PipelineError, + network::{PushWorkHandler, ingress::push_endpoint::PushEndpoint}, +}; +use dynamo_runtime::{SystemHealth, config::HealthStatus}; +use futures::StreamExt; +use parking_lot::Mutex; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +#[derive(Default)] +struct Handler { + calls: AtomicU64, + started: Notify, + release: Notify, +} + +#[async_trait::async_trait] +impl PushWorkHandler for Handler { + async fn handle_payload( + &self, + _: bytes::Bytes, + _: Option, + ) -> Result<(), PipelineError> { + self.calls.fetch_add(1, Ordering::SeqCst); + self.started.notify_one(); + self.release.notified().await; + Ok(()) + } + + fn add_metrics( + &self, + _: &dynamo_runtime::component::Endpoint, + _: Option<&[(&str, &str)]>, + ) -> anyhow::Result<()> { + Ok(()) + } +} + +#[tokio::test] +async fn cancellation_rejects_queued_requests_and_drains_accepted_work() { + tokio::time::timeout(Duration::from_secs(10), async { + let server = std::env::var(dynamo_runtime::config::environment_names::nats::NATS_SERVER) + .unwrap_or_else(|_| "nats://127.0.0.1:4222".into()); + let client = async_nats::connect(server).await.unwrap(); + let service = client.service_builder() + .start("shutdown_test", "0.0.1").await.unwrap(); + for accepted in [false, true] { + let subject = format!("shutdown.{}", uuid::Uuid::new_v4().simple()); + let endpoint = service.endpoint(&subject).await.unwrap(); + let handler = Arc::new(Handler::default()); + let cancellation = CancellationToken::new(); + let health = Arc::new(Mutex::new(SystemHealth::new( + HealthStatus::Ready, vec![subject.clone()], false, + "/health".into(), "/live".into(), + ))); + let push = PushEndpoint::builder() + .service_handler(handler.clone()) + .cancellation_token(cancellation.clone()) + .build().unwrap(); + let run = push.start( + endpoint, "test".into(), "worker".into(), subject.clone(), 1, health.clone(), + ); + // Receiving a later message on the same connection ensures the + // endpoint request has been delivered before we poll its receive loop. + let barrier_subject = client.new_inbox(); + let mut barrier = client.subscribe(barrier_subject.clone()).await.unwrap(); + client.publish_with_reply(subject.clone(), client.new_inbox(), "request".into()) + .await.unwrap(); + client.publish(barrier_subject, "barrier".into()).await.unwrap(); + barrier.next().await.unwrap(); + if !accepted { + cancellation.cancel(); + } + tokio::pin!(run); + if accepted { + tokio::select! { + _ = handler.started.notified() => {}, + result = &mut run => panic!("endpoint stopped before accepting work: {result:?}"), + } + cancellation.cancel(); + // Poll shutdown until it marks the endpoint not-ready and waits + // for the already accepted handler. Unsubscribing must not stop it. + tokio::select! { + _ = async { + while health.lock().get_endpoint_health_status(&subject) != Some(HealthStatus::NotReady) { + tokio::task::yield_now().await; + } + } => {}, + result = &mut run => panic!("endpoint stopped before draining: {result:?}"), + } + } + handler.release.notify_one(); + run.await.unwrap(); + assert_eq!(handler.calls.load(Ordering::SeqCst), u64::from(accepted)); + } + }).await.expect("NATS shutdown did not complete"); +}