diff --git a/lib/runtime/src/transports/event_plane/dynamic_subscriber.rs b/lib/runtime/src/transports/event_plane/dynamic_subscriber.rs index ec82702459c3..e7af1682885f 100644 --- a/lib/runtime/src/transports/event_plane/dynamic_subscriber.rs +++ b/lib/runtime/src/transports/event_plane/dynamic_subscriber.rs @@ -32,11 +32,20 @@ pub struct DynamicSubscriber { impl DynamicSubscriber { pub fn new(discovery: Arc, query: DiscoveryQuery, topic: String) -> Self { + Self::with_cancel_token(discovery, query, topic, CancellationToken::new()) + } + + pub fn with_cancel_token( + discovery: Arc, + query: DiscoveryQuery, + topic: String, + cancel_token: CancellationToken, + ) -> Self { Self { discovery, query, topic, - cancel_token: CancellationToken::new(), + cancel_token, } } @@ -80,8 +89,12 @@ impl DynamicSubscriber { "Attempting to start discovery watch" ); - // Don't pass the cancel token to list_and_watch - we'll handle cancellation ourselves - let mut watch_stream = match discovery.list_and_watch(query.clone(), None).await { + // Pass cancellation through so the discovery backend can stop any + // task that it owns in addition to the consumer loop below. + let mut watch_stream = match discovery + .list_and_watch(query.clone(), Some(cancel_token.clone())) + .await + { Ok(stream) => { tracing::debug!("Successfully obtained discovery watch stream"); stream @@ -94,12 +107,21 @@ impl DynamicSubscriber { tracing::info!(?query, "Started dynamic discovery watch for ZMQ publishers"); - while let Some(event_result) = watch_stream.next().await { + loop { + let event_result = tokio::select! { + biased; + + _ = cancel_token.cancelled() => { + tracing::info!("Dynamic subscriber cancelled, stopping watch"); + break; + } + result = watch_stream.next() => match result { + Some(result) => result, + None => break, + }, + }; + tracing::debug!("Received discovery event: {:?}", event_result); - if cancel_token.is_cancelled() { - tracing::info!("Dynamic subscriber cancelled, stopping watch"); - break; - } match event_result { Ok(DiscoveryEvent::Added(instance)) => { @@ -305,6 +327,48 @@ impl Drop for DynamicSubscriber { #[cfg(test)] mod tests { use super::*; + use crate::discovery::{DiscoverySpec, DiscoveryStream, EventChannelQuery}; + use tokio::sync::Notify; + use tokio::time::{Duration, timeout}; + + struct CancellationAwareDiscovery { + backend_stopped: Arc, + } + + #[async_trait::async_trait] + impl Discovery for CancellationAwareDiscovery { + fn instance_id(&self) -> u64 { + 1 + } + + async fn register_internal(&self, _spec: DiscoverySpec) -> Result { + anyhow::bail!("register is not supported by this test discovery") + } + + async fn unregister(&self, _instance: DiscoveryInstance) -> Result<()> { + anyhow::bail!("unregister is not supported by this test discovery") + } + + async fn list(&self, _query: DiscoveryQuery) -> Result> { + Ok(Vec::new()) + } + + async fn list_and_watch( + &self, + _query: DiscoveryQuery, + cancel_token: Option, + ) -> Result { + let cancel_token = cancel_token + .ok_or_else(|| anyhow::anyhow!("dynamic subscriber must pass cancellation"))?; + let backend_stopped = Arc::clone(&self.backend_stopped); + tokio::spawn(async move { + cancel_token.cancelled().await; + backend_stopped.notify_one(); + }); + + Ok(Box::pin(futures::stream::pending())) + } + } fn event_channel(topic: &str, transport: EventTransport) -> DiscoveryInstance { DiscoveryInstance::EventChannel { @@ -338,4 +402,75 @@ mod tests { None ); } + + #[tokio::test] + async fn cancellation_stops_idle_discovery_watch() { + let backend_stopped = Arc::new(Notify::new()); + let discovery = Arc::new(CancellationAwareDiscovery { + backend_stopped: Arc::clone(&backend_stopped), + }); + let query = DiscoveryQuery::EventChannels(EventChannelQuery::topic( + "test-ns", + "test-component", + "kv-events", + )); + let subscriber = Arc::new(DynamicSubscriber::new( + discovery, + query, + "kv-events".to_string(), + )); + let mut stream = Arc::clone(&subscriber).start_zmq().await.unwrap(); + + tokio::task::yield_now().await; + subscriber.cancel(); + + let next = timeout(Duration::from_secs(1), stream.next()) + .await + .expect("subscriber stream should close promptly after cancellation"); + assert!(next.is_none()); + + timeout(Duration::from_secs(1), backend_stopped.notified()) + .await + .expect("discovery backend should receive cancellation"); + } + + #[tokio::test] + async fn dropping_returned_stream_cancels_idle_discovery_watch() { + let backend_stopped = Arc::new(Notify::new()); + let discovery = Arc::new(CancellationAwareDiscovery { + backend_stopped: Arc::clone(&backend_stopped), + }); + let query = DiscoveryQuery::EventChannels(EventChannelQuery::topic( + "test-ns", + "test-component", + "kv-events", + )); + let parent_token = CancellationToken::new(); + let subscriber = Arc::new(DynamicSubscriber::with_cancel_token( + discovery, + query, + "kv-events".to_string(), + parent_token.child_token(), + )); + let weak_subscriber = Arc::downgrade(&subscriber); + let stream = subscriber.start_zmq().await.unwrap(); + + assert!( + weak_subscriber.upgrade().is_some(), + "returned stream should retain the dynamic subscriber" + ); + drop(stream); + assert!( + weak_subscriber.upgrade().is_none(), + "dropping the returned stream should release the dynamic subscriber" + ); + assert!( + !parent_token.is_cancelled(), + "dropping a subscriber must not cancel its parent token" + ); + + timeout(Duration::from_secs(1), backend_stopped.notified()) + .await + .expect("dropping the stream should cancel the discovery backend"); + } } diff --git a/lib/runtime/src/transports/event_plane/mod.rs b/lib/runtime/src/transports/event_plane/mod.rs index 23b55af7853e..e82dfbe3a013 100644 --- a/lib/runtime/src/transports/event_plane/mod.rs +++ b/lib/runtime/src/transports/event_plane/mod.rs @@ -715,8 +715,12 @@ impl EventSubscriber { ), }; - let subscriber = - Arc::new(DynamicSubscriber::new(discovery, query, topic.clone())); + let subscriber = Arc::new(DynamicSubscriber::with_cancel_token( + discovery, + query, + topic.clone(), + drt.primary_token().child_token(), + )); let stream = subscriber.start_zmq().await?; let codec = Arc::new(Codec::Msgpack(MsgpackCodec)); @@ -1098,6 +1102,47 @@ mod tests { .await; } + #[tokio::test] + async fn runtime_cancellation_stops_retained_direct_zmq_subscriber() { + temp_env::async_with_vars( + [ + (broker_env::DYN_ZMQ_BROKER_URL, None::<&str>), + (broker_env::DYN_ZMQ_BROKER_ENABLED, None::<&str>), + ], + async { + let runtime = crate::Runtime::from_current().expect("create runtime handle"); + let drt = DistributedRuntime::new( + runtime, + crate::distributed::DistributedConfig::process_local(), + ) + .await + .expect("create distributed runtime"); + let component = drt + .namespace("event-subscriber-shutdown-test") + .expect("create namespace") + .component("worker") + .expect("create component"); + + let mut subscriber = EventSubscriber::for_component_with_transport( + &component, + "events", + EventTransportKind::Zmq, + ) + .await + .expect("create subscriber"); + + drt.primary_token().cancel(); + + let next = + tokio::time::timeout(std::time::Duration::from_secs(1), subscriber.next()) + .await + .expect("runtime cancellation should stop the subscriber stream"); + assert!(next.is_none()); + }, + ) + .await; + } + #[test] fn test_event_scope_subject_prefix() { let ns_scope = EventScope::Namespace { diff --git a/lib/runtime/src/transports/event_plane/zmq_transport.rs b/lib/runtime/src/transports/event_plane/zmq_transport.rs index eb91e83d5360..0877c14496c4 100644 --- a/lib/runtime/src/transports/event_plane/zmq_transport.rs +++ b/lib/runtime/src/transports/event_plane/zmq_transport.rs @@ -26,6 +26,7 @@ use tmq::{ subscribe::{Subscribe, subscribe}, }; use tokio::sync::{Mutex, broadcast}; +use tokio_util::task::AbortOnDropHandle; /// Returns the process-wide shared ZMQ context. /// @@ -207,7 +208,7 @@ impl EventTransportTx for ZmqPubTransport { /// Uses a background async reader to fan out frames to multiple local subscribers. pub struct ZmqSubTransport { broadcast_tx: broadcast::Sender, - _socket_pump_handle: tokio::task::JoinHandle<()>, + socket_pump_handle: Arc>, } impl ZmqSubTransport { @@ -230,7 +231,7 @@ impl ZmqSubTransport { Ok(Self { broadcast_tx, - _socket_pump_handle: pump_handle, + socket_pump_handle: Arc::new(AbortOnDropHandle::new(pump_handle)), }) } @@ -273,7 +274,7 @@ impl ZmqSubTransport { Ok(Self { broadcast_tx, - _socket_pump_handle: pump_handle, + socket_pump_handle: Arc::new(AbortOnDropHandle::new(pump_handle)), }) } @@ -353,8 +354,13 @@ impl ZmqSubTransport { impl EventTransportRx for ZmqSubTransport { async fn subscribe(&self, _subject: &str) -> Result { let mut receiver = self.broadcast_tx.subscribe(); + let socket_pump_handle = Arc::clone(&self.socket_pump_handle); let stream = stream! { + // Keep the socket pump alive after the transport is dropped. The + // final transport or subscription stream aborts the pump, which + // drops its owned ZMQ socket instead of detaching the task. + let _socket_pump_handle = socket_pump_handle; loop { match receiver.recv().await { Ok(payload) => yield Ok(payload), @@ -445,6 +451,10 @@ mod tests { .await .expect("Failed to create subscription"); + // Broker-mode callers retain only the returned stream. It must keep the + // socket pump alive after the transport itself leaves scope. + drop(subscriber); + tokio::time::sleep(Duration::from_millis(100)).await; let codec = MsgpackCodec; @@ -470,6 +480,33 @@ mod tests { assert_eq!(decoded.topic, topic); } + #[tokio::test] + async fn test_zmq_socket_pump_stops_with_last_owner() { + let endpoint = format!("inproc://dynamo-zmq-pump-lifetime-{}", std::process::id()); + let topic = "pump-lifetime"; + + let (_publisher, _) = ZmqPubTransport::bind(&endpoint, topic).await.unwrap(); + let subscriber = ZmqSubTransport::connect(&endpoint, topic).await.unwrap(); + let pump_handle = subscriber.socket_pump_handle.abort_handle(); + let stream = subscriber.subscribe(topic).await.unwrap(); + + drop(subscriber); + tokio::task::yield_now().await; + assert!( + !pump_handle.is_finished(), + "subscription stream should keep the socket pump alive" + ); + + drop(stream); + timeout(Duration::from_secs(1), async { + while !pump_handle.is_finished() { + tokio::task::yield_now().await; + } + }) + .await + .expect("socket pump should stop when its final owner is dropped"); + } + #[tokio::test] async fn test_zmq_multiple_messages() { let port = 25556;