fix(event-plane): stop retired ZMQ subscriber pumps

lib/runtime/src/transports/event_plane/dynamic_subscriber.rs

@@ -32,11 +32,20 @@ pub struct DynamicSubscriber {
 
 impl DynamicSubscriber {
     pub fn new(discovery: Arc<dyn Discovery>, query: DiscoveryQuery, topic: String) -> Self {
+        Self::with_cancel_token(discovery, query, topic, CancellationToken::new())
+    }
+
+    pub fn with_cancel_token(
+        discovery: Arc<dyn Discovery>,
+        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<Notify>,
+    }
+
+    #[async_trait::async_trait]
+    impl Discovery for CancellationAwareDiscovery {
+        fn instance_id(&self) -> u64 {
+            1
+        }
+
+        async fn register_internal(&self, _spec: DiscoverySpec) -> Result<DiscoveryInstance> {
+            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<Vec<DiscoveryInstance>> {
+            Ok(Vec::new())
+        }
+
+        async fn list_and_watch(
+            &self,
+            _query: DiscoveryQuery,
+            cancel_token: Option<CancellationToken>,
+        ) -> Result<DiscoveryStream> {
+            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");
+    }
 }

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 {

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<Bytes>,
-    _socket_pump_handle: tokio::task::JoinHandle<()>,
+    socket_pump_handle: Arc<AbortOnDropHandle<()>>,
 }
 
 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<WireStream> {
         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;