From 229f504d30847603801beb29edc4a55df91f0575 Mon Sep 17 00:00:00 2001 From: zhongdaor Date: Fri, 26 Jun 2026 21:39:56 -0700 Subject: [PATCH 1/4] fix(runtime): give event publishers unique discovery identities Signed-off-by: zhongdaor --- lib/runtime/src/discovery/kube.rs | 18 +- lib/runtime/src/discovery/kv_store.rs | 10 +- lib/runtime/src/discovery/mod.rs | 24 ++- .../event_plane/dynamic_subscriber.rs | 91 +++++++-- lib/runtime/src/transports/event_plane/mod.rs | 175 +++++++++++++++++- 5 files changed, 283 insertions(+), 35 deletions(-) diff --git a/lib/runtime/src/discovery/kube.rs b/lib/runtime/src/discovery/kube.rs index 611f2c2862c1..87f353bb9fc8 100644 --- a/lib/runtime/src/discovery/kube.rs +++ b/lib/runtime/src/discovery/kube.rs @@ -129,13 +129,15 @@ impl Discovery for KubeDiscoveryClient { } async fn register_internal(&self, spec: DiscoverySpec) -> Result { - let instance_id = self.instance_id(); - let instance = spec.with_instance_id(instance_id); + let owner_instance_id = self.instance_id(); + let instance = spec.with_instance_id(owner_instance_id); + let registered_instance_id = instance.instance_id(); tracing::debug!( - "Registering instance: {:?} with instance_id={:x}", + "Registering instance: {:?} with instance_id={:x}, owner_instance_id={:x}", instance, - instance_id + registered_instance_id, + owner_instance_id ); // Write to local metadata and persist to CR @@ -152,7 +154,7 @@ impl Discovery for KubeDiscoveryClient { inst.namespace, inst.component, inst.endpoint, - instance_id + registered_instance_id ); metadata.register_endpoint(instance.clone())?; } @@ -167,7 +169,7 @@ impl Discovery for KubeDiscoveryClient { namespace, component, endpoint, - instance_id + registered_instance_id ); metadata.register_model_card(instance.clone())?; } @@ -182,7 +184,7 @@ impl Discovery for KubeDiscoveryClient { namespace, component, topic, - instance_id + registered_instance_id ); metadata.register_event_channel(instance.clone())?; } @@ -214,7 +216,7 @@ impl Discovery for KubeDiscoveryClient { } async fn unregister(&self, instance: DiscoveryInstance) -> Result<()> { - let instance_id = self.instance_id(); + let instance_id = instance.instance_id(); // Write to local metadata and persist to CR // IMPORTANT: Hold the write lock across the CR write to prevent race conditions diff --git a/lib/runtime/src/discovery/kv_store.rs b/lib/runtime/src/discovery/kv_store.rs index 528b20ff4c9f..8fceb106367a 100644 --- a/lib/runtime/src/discovery/kv_store.rs +++ b/lib/runtime/src/discovery/kv_store.rs @@ -416,8 +416,9 @@ impl Discovery for KVStoreDiscovery { } async fn register_internal(&self, spec: DiscoverySpec) -> Result { - let instance_id = self.instance_id(); - let instance = spec.with_instance_id(instance_id); + let owner_instance_id = self.instance_id(); + let instance = spec.with_instance_id(owner_instance_id); + let registered_instance_id = instance.instance_id(); let (bucket_name, key_path) = match &instance { DiscoveryInstance::Endpoint(inst) => { @@ -528,8 +529,9 @@ impl Discovery for KVStoreDiscovery { // Use revision 0 for initial registration let outcome = bucket.insert(&key, instance_json.into(), 0).await?; tracing::debug!( - "KVStoreDiscovery::register: Successfully registered instance_id={}, key={}, outcome={:?}", - instance_id, + "KVStoreDiscovery::register: Successfully registered instance_id={}, owner_instance_id={}, key={}, outcome={:?}", + registered_instance_id, + owner_instance_id, key_path, outcome ); diff --git a/lib/runtime/src/discovery/mod.rs b/lib/runtime/src/discovery/mod.rs index 91ec2b76c911..e72c4cfe55fe 100644 --- a/lib/runtime/src/discovery/mod.rs +++ b/lib/runtime/src/discovery/mod.rs @@ -350,6 +350,11 @@ pub enum DiscoverySpec { component: String, /// Topic name for this channel (e.g., "kv-events", "kv-metrics") topic: String, + /// Unique identity of this publisher incarnation. + /// + /// A process can host multiple publishers for the same topic, so event + /// channels cannot use the process-level discovery instance ID. + publisher_id: u64, /// Event transport type (NATS subject prefix or ZMQ endpoint) transport: EventTransport, }, @@ -392,8 +397,12 @@ impl DiscoverySpec { }) } - /// Attaches an instance ID to create a DiscoveryInstance - pub fn with_instance_id(self, instance_id: u64) -> DiscoveryInstance { + /// Creates a discovery instance owned by the given process. + /// + /// Endpoint and model instances use the process-level ID. Event channels + /// use their publisher-level ID because one process may own multiple + /// publishers for the same topic. + pub fn with_instance_id(self, owner_instance_id: u64) -> DiscoveryInstance { match self { Self::Endpoint { namespace, @@ -405,7 +414,7 @@ impl DiscoverySpec { namespace, component, endpoint, - instance_id, + instance_id: owner_instance_id, transport, device_type, }), @@ -419,7 +428,7 @@ impl DiscoverySpec { namespace, component, endpoint, - instance_id, + instance_id: owner_instance_id, card_json, model_suffix, }, @@ -427,12 +436,13 @@ impl DiscoverySpec { namespace, component, topic, + publisher_id, transport, } => DiscoveryInstance::EventChannel { namespace, component, topic, - instance_id, + instance_id: publisher_id, transport, }, } @@ -789,7 +799,9 @@ fn find_conflicting_model_name( #[async_trait] pub trait Discovery: Send + Sync { /// Returns a unique identifier for this worker (e.g lease id if using etcd or generated id for memory store) - /// Discovery objects created by this worker will be associated with this id. + /// Endpoint and model objects created by this worker use this ID. Event + /// channels use a publisher-level ID because a worker can own more than one + /// publisher for the same topic. fn instance_id(&self) -> u64; /// Registers an object in the discovery plane with the instance id diff --git a/lib/runtime/src/transports/event_plane/dynamic_subscriber.rs b/lib/runtime/src/transports/event_plane/dynamic_subscriber.rs index 1d8bfc8d7ab6..ec82702459c3 100644 --- a/lib/runtime/src/transports/event_plane/dynamic_subscriber.rs +++ b/lib/runtime/src/transports/event_plane/dynamic_subscriber.rs @@ -58,8 +58,9 @@ impl DynamicSubscriber { let (event_tx, event_rx) = mpsc::channel::(channel_cap); // Track active endpoint connections with instance ID to endpoint mapping - let active_endpoints: Arc>> = - Arc::new(RwLock::new(HashMap::new())); + let active_endpoints: Arc< + RwLock>, + > = Arc::new(RwLock::new(HashMap::new())); // Clone self for the spawned task let subscriber_clone = Arc::clone(&self); @@ -103,19 +104,19 @@ impl DynamicSubscriber { match event_result { Ok(DiscoveryEvent::Added(instance)) => { tracing::info!(instance = ?instance, "Discovery Added event received"); - let instance_id = instance.instance_id().to_string(); + let instance_id = instance.id(); // Extract ZMQ endpoint from the instance - if let Some(endpoint) = Self::extract_zmq_endpoint(&instance) { + if let Some(endpoint) = Self::extract_zmq_endpoint(&instance, &zmq_topic) { let mut endpoints_guard = endpoints.write().await; // Skip if instance already tracked if endpoints_guard.contains_key(&instance_id) { - tracing::debug!(endpoint = %endpoint, instance_id = %instance_id, "Already connected to ZMQ publisher"); + tracing::debug!(endpoint = %endpoint, ?instance_id, "Already connected to ZMQ publisher"); continue; } - tracing::info!(endpoint = %endpoint, instance_id = %instance_id, "Connecting to new ZMQ publisher"); + tracing::info!(endpoint = %endpoint, ?instance_id, "Connecting to new ZMQ publisher"); // Create cancellation token for this endpoint's stream let endpoint_cancel = CancellationToken::new(); @@ -151,25 +152,44 @@ impl DynamicSubscriber { endpoints_clone.write().await.remove(&instance_id_clone); }); } else { - tracing::warn!( + tracing::debug!( instance = ?instance, - "Discovery Added event did not contain a ZMQ endpoint" + expected_topic = %zmq_topic, + "Discovery event is not a matching ZMQ publisher" ); } } Ok(DiscoveryEvent::Removed(instance_id)) => { - let id_str = instance_id.instance_id().to_string(); + let is_expected_topic = matches!( + &instance_id, + DiscoveryInstanceId::EventChannel(channel_id) + if channel_id.topic == zmq_topic + ); + if !is_expected_topic { + tracing::debug!( + ?instance_id, + expected_topic = %zmq_topic, + "Ignoring removal for unrelated event channel" + ); + continue; + } + tracing::info!( - instance_id = %id_str, + ?instance_id, "ZMQ publisher removed from discovery, cancelling endpoint stream" ); // Cancel the endpoint's stream via its CancellationToken - if let Some((_endpoint, cancel)) = endpoints.write().await.remove(&id_str) { + if let Some((_endpoint, cancel)) = + endpoints.write().await.remove(&instance_id) + { cancel.cancel(); - tracing::info!(instance_id = %id_str, "Cancelled endpoint stream"); + tracing::info!(?instance_id, "Cancelled endpoint stream"); } else { - tracing::warn!(instance_id = %id_str, "No active endpoint found for removed stream instance"); + tracing::debug!( + ?instance_id, + "No active endpoint found for removed stream instance" + ); } } Err(e) => { @@ -201,8 +221,11 @@ impl DynamicSubscriber { } /// Extract ZMQ endpoint from a discovery instance. - fn extract_zmq_endpoint(instance: &DiscoveryInstance) -> Option { - if let DiscoveryInstance::EventChannel { transport, .. } = instance + fn extract_zmq_endpoint(instance: &DiscoveryInstance, expected_topic: &str) -> Option { + if let DiscoveryInstance::EventChannel { + topic, transport, .. + } = instance + && topic == expected_topic && let EventTransport::Zmq { endpoint } = transport { return Some(endpoint.clone()); @@ -278,3 +301,41 @@ impl Drop for DynamicSubscriber { self.cancel_token.cancel(); } } + +#[cfg(test)] +mod tests { + use super::*; + + fn event_channel(topic: &str, transport: EventTransport) -> DiscoveryInstance { + DiscoveryInstance::EventChannel { + namespace: "test-ns".to_string(), + component: "test-component".to_string(), + topic: topic.to_string(), + instance_id: 1, + transport, + } + } + + #[test] + fn extracts_only_matching_zmq_topic() { + let matching = event_channel("kv-events", EventTransport::zmq("tcp://127.0.0.1:1")); + let wrong_topic = event_channel("kv-metrics", EventTransport::zmq("tcp://127.0.0.1:2")); + let wrong_transport = event_channel( + "kv-events", + EventTransport::nats("namespace.test-ns.component.test-component"), + ); + + assert_eq!( + DynamicSubscriber::extract_zmq_endpoint(&matching, "kv-events").as_deref(), + Some("tcp://127.0.0.1:1") + ); + assert_eq!( + DynamicSubscriber::extract_zmq_endpoint(&wrong_topic, "kv-events"), + None + ); + assert_eq!( + DynamicSubscriber::extract_zmq_endpoint(&wrong_transport, "kv-events"), + None + ); + } +} diff --git a/lib/runtime/src/transports/event_plane/mod.rs b/lib/runtime/src/transports/event_plane/mod.rs index 07f937f160a0..1e7c140bfb3d 100644 --- a/lib/runtime/src/transports/event_plane/mod.rs +++ b/lib/runtime/src/transports/event_plane/mod.rs @@ -30,6 +30,7 @@ use anyhow::Result; use bytes::Bytes; use futures::{Stream, StreamExt}; use lru::LruCache; +use rand::TryRngCore; use serde::Serialize; use serde::de::DeserializeOwned; use std::pin::Pin; @@ -285,6 +286,8 @@ pub struct EventPublisher { tx: Arc, codec: Arc, runtime_handle: tokio::runtime::Handle, + // Keeps unregister work in graceful-shutdown Phase 2 when dropped before shutdown. + graceful_shutdown_tracker: Arc, /// Discovery client and registered instance for unregistration on drop discovery_client: Option>, discovery_instance: Option, @@ -346,10 +349,17 @@ impl EventPublisher { topic: String, transport_kind: EventTransportKind, ) -> Result { - let publisher_id = drt.discovery().instance_id(); + // Publishers are discovery objects in their own right. A single process + // can host multiple publishers for the same scope/topic, each with its + // own ZMQ endpoint and sequence space, so the process ID is not unique + // enough here. + let publisher_id = rand::rngs::OsRng + .try_next_u64() + .map_err(|error| anyhow::anyhow!("failed to generate publisher ID: {error}"))?; let discovery = Some(drt.discovery()); let runtime_handle = drt.runtime().secondary(); let subject = format!("{}.{}", scope.subject_prefix(), topic); + let graceful_shutdown_tracker = drt.graceful_shutdown_tracker(); // Use Msgpack codec for all transports enum TransportSetup { @@ -434,6 +444,7 @@ impl EventPublisher { namespace: scope.namespace().to_string(), component: scope.component().unwrap_or("").to_string(), topic: topic.clone(), + publisher_id, transport: transport_config, }; @@ -452,6 +463,7 @@ impl EventPublisher { namespace: scope.namespace().to_string(), component: scope.component().unwrap_or("").to_string(), topic: topic.clone(), + publisher_id, transport: transport_config, }; @@ -483,6 +495,7 @@ impl EventPublisher { tx, codec, runtime_handle, + graceful_shutdown_tracker, discovery_client: discovery, discovery_instance, }) @@ -537,11 +550,13 @@ impl Drop for EventPublisher { let topic = self.topic.clone(); let instance_id = instance.instance_id(); let runtime_handle = self.runtime_handle.clone(); + let shutdown_guard = self.graceful_shutdown_tracker.register_task(); // Drop can run outside any Tokio context (notably via PyO3 finalizers), so use // the runtime that created the publisher rather than the ambient thread state. let spawn_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || { runtime_handle.spawn(async move { + let _shutdown_guard = shutdown_guard; match discovery.unregister(instance).await { Ok(()) => { tracing::info!( @@ -806,7 +821,163 @@ fn current_timestamp_ms() -> u64 { #[cfg(test)] mod tests { use super::*; - use crate::config::environment_names::event_plane as env; + use crate::config::environment_names::zmq_broker as broker_env; + + #[tokio::test] + async fn same_topic_publishers_are_discovered_and_delivered_independently() { + 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-publisher-test") + .expect("create namespace") + .component("worker") + .expect("create component"); + + let publisher_a = EventPublisher::for_component_with_transport( + &component, + "events", + EventTransportKind::Zmq, + ) + .await + .expect("create first publisher"); + let publisher_b = EventPublisher::for_component_with_transport( + &component, + "events", + EventTransportKind::Zmq, + ) + .await + .expect("create second publisher"); + let publisher_a_id = publisher_a.publisher_id(); + let publisher_b_id = publisher_b.publisher_id(); + + assert_ne!(publisher_a_id, publisher_b_id); + + let query = DiscoveryQuery::EventChannels(EventChannelQuery::topic( + "event-publisher-test", + "worker", + "events", + )); + let instances = drt + .discovery() + .list(query.clone()) + .await + .expect("list event publishers"); + assert_eq!(instances.len(), 2); + assert!( + instances + .iter() + .any(|instance| instance.instance_id() == publisher_a_id) + ); + assert!( + instances + .iter() + .any(|instance| instance.instance_id() == publisher_b_id) + ); + + let mut subscriber = EventSubscriber::for_component_with_transport( + &component, + "events", + EventTransportKind::Zmq, + ) + .await + .expect("create subscriber"); + let mut received_a = false; + let mut received_b = false; + + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while !received_a || !received_b { + publisher_a + .publish_bytes(vec![0xa1]) + .await + .expect("publish from first publisher"); + publisher_b + .publish_bytes(vec![0xb2]) + .await + .expect("publish from second publisher"); + + if let Ok(Some(envelope)) = tokio::time::timeout( + std::time::Duration::from_millis(100), + subscriber.next(), + ) + .await + { + let envelope = envelope.expect("receive event envelope"); + if envelope.publisher_id == publisher_a_id { + assert_eq!(envelope.payload.as_ref(), &[0xa1]); + received_a = true; + } else if envelope.publisher_id == publisher_b_id { + assert_eq!(envelope.payload.as_ref(), &[0xb2]); + received_b = true; + } else { + panic!("event from unexpected publisher"); + } + } + + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("subscriber should receive events from both publishers"); + + drop(publisher_a); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + let instances = drt + .discovery() + .list(query.clone()) + .await + .expect("list event publishers after drop"); + if instances.len() == 1 { + assert_eq!(instances[0].instance_id(), publisher_b_id); + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("first publisher should unregister without removing the second"); + + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + publisher_b + .publish_bytes(vec![0xb3]) + .await + .expect("publish from remaining publisher"); + + if let Ok(Some(envelope)) = tokio::time::timeout( + std::time::Duration::from_millis(100), + subscriber.next(), + ) + .await + { + let envelope = envelope.expect("receive event envelope after drop"); + if envelope.publisher_id == publisher_b_id + && envelope.payload.as_ref() == [0xb3] + { + break; + } + } + + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("remaining publisher should stay connected after peer removal"); + }, + ) + .await; + } #[test] fn test_event_scope_subject_prefix() { From 02d77dcca59b9ab13662831c8052b80c2e6b5beb Mon Sep 17 00:00:00 2001 From: zhongdaor Date: Fri, 26 Jun 2026 22:22:01 -0700 Subject: [PATCH 2/4] test(router): use default event plane in mocker e2e Signed-off-by: zhongdaor --- tests/router/test_router_e2e_with_mockers.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index 89c734deaa7e..2d3f0acc9c7e 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -61,25 +61,6 @@ COUNTER_WORKER_SCRIPT = os.path.join(os.path.dirname(__file__), "counter_worker.py") -@pytest.fixture(autouse=True) -def _pin_nats_event_plane_for_mocker(request, monkeypatch): - """Pin the NATS event plane for etcd-backed mocker tests. - - The mock engine publishes KV cache events instantly -- before the in-process - router's ZMQ subscription has connected (ZMQ slow-joiner) -- so on the (now - default) ZMQ event plane the router observes zero events and the routing - assertions fail. Real engines start slowly enough to avoid this, so the - vLLM/SGLang router e2e tests cover the ZMQ default; only the fast mocker needs - NATS here. file-backed variants keep the ZMQ default, and an explicitly set - DYN_EVENT_PLANE (e.g. via durable_kv_events) is left untouched. - """ - callspec = getattr(request.node, "callspec", None) - store_backend = callspec.params.get("store_backend", "etcd") if callspec else "etcd" - if store_backend != "file" and not os.environ.get("DYN_EVENT_PLANE"): - monkeypatch.setenv("DYN_EVENT_PLANE", "nats") - yield - - pytestmark = [ pytest.mark.pre_merge, pytest.mark.gpu_0, From 5cb2d1bf8257b9102b37c1bc6f999c1001612638 Mon Sep 17 00:00:00 2001 From: zhongdaor Date: Mon, 29 Jun 2026 18:36:05 -0700 Subject: [PATCH 3/4] refactor(runtime): clarify publisher discovery identities Signed-off-by: zhongdaor --- lib/runtime/src/discovery/kube.rs | 16 ++--- lib/runtime/src/discovery/kv_store.rs | 10 ++- lib/runtime/src/discovery/mock.rs | 2 +- lib/runtime/src/discovery/mod.rs | 19 +++-- lib/runtime/src/transports/event_plane/mod.rs | 72 ++++++++++++++----- 5 files changed, 77 insertions(+), 42 deletions(-) diff --git a/lib/runtime/src/discovery/kube.rs b/lib/runtime/src/discovery/kube.rs index 87f353bb9fc8..bc09b9897c62 100644 --- a/lib/runtime/src/discovery/kube.rs +++ b/lib/runtime/src/discovery/kube.rs @@ -129,15 +129,13 @@ impl Discovery for KubeDiscoveryClient { } async fn register_internal(&self, spec: DiscoverySpec) -> Result { - let owner_instance_id = self.instance_id(); - let instance = spec.with_instance_id(owner_instance_id); - let registered_instance_id = instance.instance_id(); + let instance = spec.into_instance(self.instance_id()); + let instance_id = instance.instance_id(); tracing::debug!( - "Registering instance: {:?} with instance_id={:x}, owner_instance_id={:x}", + "Registering discovery instance: {:?}, instance_id={:x}", instance, - registered_instance_id, - owner_instance_id + instance_id ); // Write to local metadata and persist to CR @@ -154,7 +152,7 @@ impl Discovery for KubeDiscoveryClient { inst.namespace, inst.component, inst.endpoint, - registered_instance_id + instance_id ); metadata.register_endpoint(instance.clone())?; } @@ -169,7 +167,7 @@ impl Discovery for KubeDiscoveryClient { namespace, component, endpoint, - registered_instance_id + instance_id ); metadata.register_model_card(instance.clone())?; } @@ -184,7 +182,7 @@ impl Discovery for KubeDiscoveryClient { namespace, component, topic, - registered_instance_id + instance_id ); metadata.register_event_channel(instance.clone())?; } diff --git a/lib/runtime/src/discovery/kv_store.rs b/lib/runtime/src/discovery/kv_store.rs index 8fceb106367a..eae35ac080ab 100644 --- a/lib/runtime/src/discovery/kv_store.rs +++ b/lib/runtime/src/discovery/kv_store.rs @@ -416,9 +416,8 @@ impl Discovery for KVStoreDiscovery { } async fn register_internal(&self, spec: DiscoverySpec) -> Result { - let owner_instance_id = self.instance_id(); - let instance = spec.with_instance_id(owner_instance_id); - let registered_instance_id = instance.instance_id(); + let instance = spec.into_instance(self.instance_id()); + let instance_id = instance.instance_id(); let (bucket_name, key_path) = match &instance { DiscoveryInstance::Endpoint(inst) => { @@ -529,9 +528,8 @@ impl Discovery for KVStoreDiscovery { // Use revision 0 for initial registration let outcome = bucket.insert(&key, instance_json.into(), 0).await?; tracing::debug!( - "KVStoreDiscovery::register: Successfully registered instance_id={}, owner_instance_id={}, key={}, outcome={:?}", - registered_instance_id, - owner_instance_id, + "KVStoreDiscovery::register: Registration insert completed instance_id={}, key={}, outcome={:?}", + instance_id, key_path, outcome ); diff --git a/lib/runtime/src/discovery/mock.rs b/lib/runtime/src/discovery/mock.rs index 9788cd223366..8d0b522d8a6e 100644 --- a/lib/runtime/src/discovery/mock.rs +++ b/lib/runtime/src/discovery/mock.rs @@ -158,7 +158,7 @@ impl Discovery for MockDiscovery { } async fn register_internal(&self, spec: DiscoverySpec) -> Result { - let instance = spec.with_instance_id(self.instance_id); + let instance = spec.into_instance(self.instance_id); self.registry .instances diff --git a/lib/runtime/src/discovery/mod.rs b/lib/runtime/src/discovery/mod.rs index e72c4cfe55fe..d453402ed2fb 100644 --- a/lib/runtime/src/discovery/mod.rs +++ b/lib/runtime/src/discovery/mod.rs @@ -397,12 +397,12 @@ impl DiscoverySpec { }) } - /// Creates a discovery instance owned by the given process. + /// Converts this registration spec into a discovery instance. /// - /// Endpoint and model instances use the process-level ID. Event channels - /// use their publisher-level ID because one process may own multiple - /// publishers for the same topic. - pub fn with_instance_id(self, owner_instance_id: u64) -> DiscoveryInstance { + /// Endpoint and model specs use `default_instance_id`, normally the + /// discovery client's process-level ID. Event channel specs already carry + /// a publisher-level ID, so they use that instead. + pub fn into_instance(self, default_instance_id: u64) -> DiscoveryInstance { match self { Self::Endpoint { namespace, @@ -414,7 +414,7 @@ impl DiscoverySpec { namespace, component, endpoint, - instance_id: owner_instance_id, + instance_id: default_instance_id, transport, device_type, }), @@ -428,7 +428,7 @@ impl DiscoverySpec { namespace, component, endpoint, - instance_id: owner_instance_id, + instance_id: default_instance_id, card_json, model_suffix, }, @@ -447,6 +447,11 @@ impl DiscoverySpec { }, } } + + /// Compatibility alias for [`DiscoverySpec::into_instance`]. + pub fn with_instance_id(self, default_instance_id: u64) -> DiscoveryInstance { + self.into_instance(default_instance_id) + } } /// Registered instances in the discovery plane diff --git a/lib/runtime/src/transports/event_plane/mod.rs b/lib/runtime/src/transports/event_plane/mod.rs index 1e7c140bfb3d..c5545749348b 100644 --- a/lib/runtime/src/transports/event_plane/mod.rs +++ b/lib/runtime/src/transports/event_plane/mod.rs @@ -448,14 +448,14 @@ impl EventPublisher { transport: transport_config, }; - let registered_instance = drt.discovery().register(spec).await?; + let discovery_instance = drt.discovery().register(spec).await?; tracing::info!( topic = %topic, transport = ?transport_kind, - instance_id = %registered_instance.instance_id(), + publisher_id = %publisher_id, "EventPublisher registered with discovery" ); - (tx, codec, Some(registered_instance)) + (tx, codec, Some(discovery_instance)) } TransportSetup::ZmqDirect(tx, codec, public_endpoint) => { let transport_config = EventTransport::zmq(public_endpoint); @@ -467,14 +467,14 @@ impl EventPublisher { transport: transport_config, }; - let registered_instance = drt.discovery().register(spec).await?; + let discovery_instance = drt.discovery().register(spec).await?; tracing::info!( topic = %topic, transport = ?transport_kind, - instance_id = %registered_instance.instance_id(), + publisher_id = %publisher_id, "EventPublisher registered with discovery (direct mode)" ); - (tx, codec, Some(registered_instance)) + (tx, codec, Some(discovery_instance)) } TransportSetup::ZmqBroker(tx, codec) => { tracing::info!( @@ -548,7 +548,7 @@ impl Drop for EventPublisher { (self.discovery_client.take(), self.discovery_instance.take()) { let topic = self.topic.clone(); - let instance_id = instance.instance_id(); + let publisher_id = instance.instance_id(); let runtime_handle = self.runtime_handle.clone(); let shutdown_guard = self.graceful_shutdown_tracker.register_task(); @@ -561,14 +561,14 @@ impl Drop for EventPublisher { Ok(()) => { tracing::info!( topic = %topic, - instance_id = %instance_id, + publisher_id = %publisher_id, "EventPublisher unregistered from discovery" ); } Err(e) => { tracing::warn!( topic = %topic, - instance_id = %instance_id, + publisher_id = %publisher_id, error = %e, "Failed to unregister EventPublisher from discovery" ); @@ -580,7 +580,7 @@ impl Drop for EventPublisher { if spawn_result.is_err() { tracing::warn!( topic = %self.topic, - instance_id = %instance_id, + publisher_id = %publisher_id, "Skipping EventPublisher unregister during drop because the runtime is unavailable" ); } @@ -824,7 +824,7 @@ mod tests { use crate::config::environment_names::zmq_broker as broker_env; #[tokio::test] - async fn same_topic_publishers_are_discovered_and_delivered_independently() { + async fn same_topic_publishers_are_independent_across_recreation() { temp_env::async_with_vars( [ (broker_env::DYN_ZMQ_BROKER_URL, None::<&str>), @@ -931,29 +931,59 @@ mod tests { .expect("subscriber should receive events from both publishers"); drop(publisher_a); + let publisher_a_recreated = EventPublisher::for_component_with_transport( + &component, + "events", + EventTransportKind::Zmq, + ) + .await + .expect("recreate first publisher"); + let publisher_a_recreated_id = publisher_a_recreated.publisher_id(); + + assert_ne!(publisher_a_recreated_id, publisher_a_id); + assert_ne!(publisher_a_recreated_id, publisher_b_id); + assert_eq!( + publisher_a_recreated.sequence.load(Ordering::SeqCst), + 0, + "a recreated publisher starts a new sequence space" + ); + tokio::time::timeout(std::time::Duration::from_secs(1), async { loop { let instances = drt .discovery() .list(query.clone()) .await - .expect("list event publishers after drop"); - if instances.len() == 1 { - assert_eq!(instances[0].instance_id(), publisher_b_id); + .expect("list event publishers after recreation"); + if instances.len() == 2 + && instances + .iter() + .any(|instance| instance.instance_id() == publisher_b_id) + && instances + .iter() + .any(|instance| instance.instance_id() == publisher_a_recreated_id) + { break; } tokio::task::yield_now().await; } }) .await - .expect("first publisher should unregister without removing the second"); + .expect("old publisher should unregister without removing current publishers"); + + let mut received_b_after_recreation = false; + let mut received_recreated_a = false; tokio::time::timeout(std::time::Duration::from_secs(5), async { - loop { + while !received_b_after_recreation || !received_recreated_a { publisher_b .publish_bytes(vec![0xb3]) .await - .expect("publish from remaining publisher"); + .expect("publish from second publisher after recreation"); + publisher_a_recreated + .publish_bytes(vec![0xa2]) + .await + .expect("publish from recreated publisher"); if let Ok(Some(envelope)) = tokio::time::timeout( std::time::Duration::from_millis(100), @@ -965,7 +995,11 @@ mod tests { if envelope.publisher_id == publisher_b_id && envelope.payload.as_ref() == [0xb3] { - break; + received_b_after_recreation = true; + } else if envelope.publisher_id == publisher_a_recreated_id + && envelope.payload.as_ref() == [0xa2] + { + received_recreated_a = true; } } @@ -973,7 +1007,7 @@ mod tests { } }) .await - .expect("remaining publisher should stay connected after peer removal"); + .expect("subscriber should receive from surviving and recreated publishers"); }, ) .await; From 3a977ca6c906e080420fb33bb6f2890b8f1d8c29 Mon Sep 17 00:00:00 2001 From: zhongdaor Date: Thu, 2 Jul 2026 18:49:59 -0700 Subject: [PATCH 4/4] test(runtime): cover EventPublisher graceful-shutdown unregister path Dropping an EventPublisher schedules an async discovery unregister and holds a GracefulShutdownTracker guard so Runtime::shutdown Phase 2 waits for it. Add a regression test that drops a publisher, runs the real shutdown sequence, and asserts the guard is taken synchronously on drop, released once unregister finishes, and that the publisher is gone from discovery before Phase 3 cancels the main token. Co-Authored-By: Claude Fable 5 --- lib/runtime/src/transports/event_plane/mod.rs | 88 +++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/lib/runtime/src/transports/event_plane/mod.rs b/lib/runtime/src/transports/event_plane/mod.rs index c5545749348b..95ed01c1b5e9 100644 --- a/lib/runtime/src/transports/event_plane/mod.rs +++ b/lib/runtime/src/transports/event_plane/mod.rs @@ -1013,6 +1013,94 @@ mod tests { .await; } + #[tokio::test] + async fn dropped_publisher_unregister_completes_within_graceful_shutdown() { + 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-publisher-shutdown-test") + .expect("create namespace") + .component("worker") + .expect("create component"); + + let publisher = EventPublisher::for_component_with_transport( + &component, + "events", + EventTransportKind::Zmq, + ) + .await + .expect("create publisher"); + let publisher_id = publisher.publisher_id(); + + let query = DiscoveryQuery::EventChannels(EventChannelQuery::topic( + "event-publisher-shutdown-test", + "worker", + "events", + )); + let instances = drt + .discovery() + .list(query.clone()) + .await + .expect("list event publishers"); + assert_eq!(instances.len(), 1); + assert_eq!(instances[0].instance_id(), publisher_id); + + let tracker = drt.graceful_shutdown_tracker(); + assert_eq!(tracker.get_count(), 0); + + let main_token = drt.runtime().primary_token(); + let endpoint_token = drt.runtime().child_token(); + + // Dropping the publisher schedules the async discovery unregister. + // The drop path must synchronously take a graceful-shutdown guard so + // that `Runtime::shutdown` Phase 2 waits for the unregister task. + drop(publisher); + assert_eq!( + tracker.get_count(), + 1, + "dropping a publisher must register its unregister work with the graceful-shutdown tracker" + ); + + drt.runtime().shutdown(); + + tokio::time::timeout(std::time::Duration::from_secs(5), main_token.cancelled()) + .await + .expect("graceful shutdown should complete once the unregister task finishes"); + + assert!(endpoint_token.is_cancelled()); + assert_eq!( + tracker.get_count(), + 0, + "the unregister task must release its graceful-shutdown guard" + ); + + // Phase 3 (main token cancellation) only runs after Phase 2 drained + // the tracker, so the dropped publisher must already be unregistered. + let instances = drt + .discovery() + .list(query) + .await + .expect("list event publishers after shutdown"); + assert!( + instances.is_empty(), + "unregister must complete within the graceful-shutdown window" + ); + }, + ) + .await; + } + #[test] fn test_event_scope_subject_prefix() { let ns_scope = EventScope::Namespace {