dynamo_runtime/transports/event_plane/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Generic Event Plane for transport-agnostic pub/sub communication.
5
6mod codec;
7mod dynamic_subscriber;
8mod frame;
9mod nats_transport;
10mod traits;
11mod transport;
12pub mod zmq_transport;
13
14pub use codec::{Codec, MsgpackCodec};
15pub use dynamic_subscriber::DynamicSubscriber;
16pub use frame::{FRAME_HEADER_SIZE, FRAME_VERSION, Frame, FrameError, FrameHeader};
17pub use traits::{EventEnvelope, EventStream, TypedEventStream};
18pub use transport::{EventTransportRx, EventTransportTx, WireStream};
19pub use zmq_transport::{
20    ValidatedEnvelope, ValidatedZmqSource, ValidatedZmqSourceError, ZmqPubTransport,
21    ZmqSubTransport,
22};
23
24// Re-export transport kind from discovery for convenience
25pub use crate::discovery::{EventScope, EventTransportKind};
26
27use std::num::NonZeroUsize;
28use std::sync::Arc;
29use std::sync::atomic::{AtomicU64, Ordering};
30use std::time::{SystemTime, UNIX_EPOCH};
31
32use anyhow::Result;
33use bytes::Bytes;
34use futures::{Stream, StreamExt};
35use lru::LruCache;
36use rand::TryRngCore;
37use serde::Serialize;
38use serde::de::DeserializeOwned;
39use std::pin::Pin;
40use std::task::{Context, Poll};
41
42use crate::DistributedRuntime;
43use crate::component::{Component, Endpoint, Namespace};
44use crate::discovery::{
45    Discovery, DiscoveryInstance, DiscoveryQuery, DiscoverySpec, EventChannelQuery, EventTransport,
46    MAX_JSON_SAFE_PUBLISHER_ID,
47};
48use crate::protocols::EndpointId;
49use crate::traits::DistributedRuntimeProvider;
50use crate::utils::local_ip_for_advertise;
51
52// ============================================================================
53// Broker Resolution Logic
54// ============================================================================
55
56/// Broker endpoints for ZMQ broker mode
57#[derive(Debug, Clone)]
58struct BrokerEndpoints {
59    xsub_endpoints: Vec<String>,
60    xpub_endpoints: Vec<String>,
61}
62
63/// Whether the configured event plane selects direct per-publisher ZMQ sockets.
64///
65/// This intentionally answers only the topology question needed by specialized
66/// consumers. Broker discovery and connection errors remain owned by the normal
67/// [`EventSubscriber`] construction path.
68pub fn uses_direct_zmq(transport_kind: EventTransportKind) -> bool {
69    uses_direct_zmq_from_lookup(transport_kind, |key| std::env::var_os(key))
70}
71
72fn uses_direct_zmq_from_lookup(
73    transport_kind: EventTransportKind,
74    mut get_env: impl FnMut(&str) -> Option<std::ffi::OsString>,
75) -> bool {
76    if transport_kind != EventTransportKind::Zmq {
77        return false;
78    }
79
80    if get_env(crate::config::environment_names::zmq_broker::DYN_ZMQ_BROKER_URL).is_some() {
81        return false;
82    }
83
84    !get_env(crate::config::environment_names::zmq_broker::DYN_ZMQ_BROKER_ENABLED)
85        .is_some_and(|value| crate::config::is_truthy(&value.to_string_lossy()))
86}
87
88/// Resolve ZMQ broker endpoints from environment or discovery
89/// Returns None if broker mode is not configured (direct mode)
90async fn resolve_zmq_broker(
91    drt: &DistributedRuntime,
92    scope: &EventScope,
93) -> Result<Option<BrokerEndpoints>> {
94    // Priority 1: Explicit URL from DYN_ZMQ_BROKER_URL
95    if let Ok(broker_url) =
96        std::env::var(crate::config::environment_names::zmq_broker::DYN_ZMQ_BROKER_URL)
97    {
98        let (xsub_endpoints, xpub_endpoints) = parse_broker_url(&broker_url)?;
99        tracing::info!(
100            num_xsub = xsub_endpoints.len(),
101            num_xpub = xpub_endpoints.len(),
102            "Using explicit ZMQ broker URL"
103        );
104        return Ok(Some(BrokerEndpoints {
105            xsub_endpoints,
106            xpub_endpoints,
107        }));
108    }
109
110    // Priority 2: Discovery-based lookup if DYN_ZMQ_BROKER_ENABLED is truthy
111    if crate::config::env_is_truthy(
112        crate::config::environment_names::zmq_broker::DYN_ZMQ_BROKER_ENABLED,
113    ) {
114        let query = DiscoveryQuery::EventChannels(EventChannelQuery::component(
115            scope.namespace().to_string(),
116            "zmq_broker".to_string(),
117        ));
118
119        let instances = drt.discovery().list(query).await?;
120
121        // Collect all broker instances (multiple brokers for HA)
122        let mut xsub_endpoints = Vec::new();
123        let mut xpub_endpoints = Vec::new();
124
125        for instance in instances {
126            if let DiscoveryInstance::EventChannel { transport, .. } = instance
127                && let EventTransport::ZmqBroker {
128                    xsub_endpoints: xsubs,
129                    xpub_endpoints: xpubs,
130                } = transport
131            {
132                xsub_endpoints.extend(xsubs);
133                xpub_endpoints.extend(xpubs);
134            }
135        }
136
137        if xsub_endpoints.is_empty() {
138            anyhow::bail!(
139                "DYN_ZMQ_BROKER_ENABLED is set but no broker found in discovery for namespace '{}'",
140                scope.namespace()
141            );
142        }
143
144        tracing::info!(
145            num_brokers = xsub_endpoints.len(),
146            "Discovered ZMQ brokers from discovery plane"
147        );
148
149        return Ok(Some(BrokerEndpoints {
150            xsub_endpoints,
151            xpub_endpoints,
152        }));
153    }
154
155    // No broker configured - use direct mode
156    Ok(None)
157}
158
159/// Parse broker URL format: "xsub=tcp://host1:5555;tcp://host2:5555 , xpub=tcp://host1:5556;tcp://host2:5556"
160fn parse_broker_url(url: &str) -> Result<(Vec<String>, Vec<String>)> {
161    let parts: Vec<&str> = url.split(',').map(|s| s.trim()).collect();
162    if parts.len() != 2 {
163        anyhow::bail!(
164            "Invalid broker URL format. Expected 'xsub=<urls> , xpub=<urls>', got: {}",
165            url
166        );
167    }
168
169    let mut xsub_endpoints = Vec::new();
170    let mut xpub_endpoints = Vec::new();
171
172    for part in parts {
173        if let Some(urls_str) = part.strip_prefix("xsub=") {
174            xsub_endpoints = urls_str
175                .split(';')
176                .map(|s| s.trim().to_string())
177                .filter(|s| !s.is_empty())
178                .collect();
179        } else if let Some(urls_str) = part.strip_prefix("xpub=") {
180            xpub_endpoints = urls_str
181                .split(';')
182                .map(|s| s.trim().to_string())
183                .filter(|s| !s.is_empty())
184                .collect();
185        } else {
186            anyhow::bail!(
187                "Invalid broker URL part. Expected 'xsub=' or 'xpub=' prefix, got: {}",
188                part
189            );
190        }
191    }
192
193    if xsub_endpoints.is_empty() || xpub_endpoints.is_empty() {
194        anyhow::bail!(
195            "Broker URL must contain at least one xsub and one xpub endpoint. Got xsub={:?}, xpub={:?}",
196            xsub_endpoints,
197            xpub_endpoints
198        );
199    }
200
201    Ok((xsub_endpoints, xpub_endpoints))
202}
203
204/// Deduplicates events based on (publisher_id, sequence) tuple
205/// Required when connecting to multiple brokers in HA mode
206struct DeduplicatingStream {
207    inner: WireStream,
208    codec: Arc<Codec>,
209    seen_events: LruCache<(u64, u64), ()>, // (publisher_id, sequence) -> ()
210}
211
212impl DeduplicatingStream {
213    fn new(inner: WireStream, codec: Arc<Codec>, cache_size: usize) -> Self {
214        Self {
215            inner,
216            codec,
217            seen_events: LruCache::new(
218                NonZeroUsize::new(cache_size).expect("cache_size must be non-zero"),
219            ),
220        }
221    }
222}
223
224impl Stream for DeduplicatingStream {
225    type Item = Result<Bytes>;
226
227    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
228        loop {
229            match Pin::new(&mut self.inner).poll_next(cx) {
230                Poll::Ready(Some(Ok(bytes))) => {
231                    // Decode envelope to extract publisher_id and sequence
232                    match self.codec.decode_envelope_identity(&bytes) {
233                        Ok(key) => {
234                            // Check if we've seen this event before
235                            if self.seen_events.contains(&key) {
236                                // Duplicate - skip and continue loop
237                                tracing::debug!(
238                                    publisher_id = key.0,
239                                    sequence = key.1,
240                                    "Filtered duplicate event from multi-broker setup"
241                                );
242                                continue;
243                            }
244
245                            // New event - record and return
246                            self.seen_events.put(key, ());
247                            return Poll::Ready(Some(Ok(bytes)));
248                        }
249                        Err(e) => {
250                            tracing::warn!(error = %e, "Failed to decode envelope for deduplication");
251                            return Poll::Ready(Some(Err(e)));
252                        }
253                    }
254                }
255                Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
256                Poll::Ready(None) => return Poll::Ready(None),
257                Poll::Pending => return Poll::Pending,
258            }
259        }
260    }
261}
262
263/// Keep publisher IDs exactly representable in float64-backed JSON metadata.
264fn discovery_safe_publisher_id(random_id: u64) -> u64 {
265    random_id & MAX_JSON_SAFE_PUBLISHER_ID
266}
267
268/// Event publisher for a specific topic.
269pub struct EventPublisher {
270    transport_kind: EventTransportKind,
271    topic: String,
272    subject: String,
273    publisher_id: u64,
274    sequence: AtomicU64,
275    tx: Arc<dyn EventTransportTx>,
276    codec: Arc<Codec>,
277    runtime_handle: tokio::runtime::Handle,
278    // Keeps unregister work in graceful-shutdown Phase 2 when dropped before shutdown.
279    graceful_shutdown_tracker: Arc<crate::utils::GracefulShutdownTracker>,
280    /// Discovery client and registered instance for unregistration on drop
281    discovery_client: Option<Arc<dyn Discovery>>,
282    discovery_instance: Option<crate::discovery::DiscoveryInstance>,
283}
284
285impl EventPublisher {
286    /// Create a publisher for an endpoint-scoped topic.
287    pub async fn for_endpoint(endpoint: &Endpoint, topic: impl Into<String>) -> Result<Self> {
288        let transport_kind = endpoint.drt().default_event_transport_kind();
289        Self::for_endpoint_with_transport(endpoint, topic, transport_kind).await
290    }
291
292    /// Create an endpoint-scoped publisher with explicit transport.
293    pub async fn for_endpoint_with_transport(
294        endpoint: &Endpoint,
295        topic: impl Into<String>,
296        transport_kind: EventTransportKind,
297    ) -> Result<Self> {
298        Self::new_internal(
299            endpoint.drt(),
300            EventScope::Endpoint {
301                endpoint: endpoint.id(),
302            },
303            topic.into(),
304            transport_kind,
305        )
306        .await
307    }
308
309    /// Create a publisher for an endpoint identity without constructing a local endpoint handle.
310    pub async fn for_endpoint_id(
311        drt: &DistributedRuntime,
312        endpoint: &EndpointId,
313        topic: impl Into<String>,
314    ) -> Result<Self> {
315        let transport_kind = drt.default_event_transport_kind();
316        Self::for_endpoint_id_with_transport(drt, endpoint, topic, transport_kind).await
317    }
318
319    /// Create an endpoint-identity publisher with an explicit transport.
320    pub async fn for_endpoint_id_with_transport(
321        drt: &DistributedRuntime,
322        endpoint: &EndpointId,
323        topic: impl Into<String>,
324        transport_kind: EventTransportKind,
325    ) -> Result<Self> {
326        Self::new_internal(
327            drt,
328            EventScope::Endpoint {
329                endpoint: endpoint.clone(),
330            },
331            topic.into(),
332            transport_kind,
333        )
334        .await
335    }
336
337    /// Create a publisher for a component-scoped topic.
338    ///
339    /// The event transport is chosen automatically: if `DYN_EVENT_PLANE` is set that
340    /// value is used; otherwise the runtime's default is used (ZMQ for local backends
341    /// such as `file`/`mem`, NATS for distributed backends such as `etcd`/`kubernetes`).
342    /// Use [`for_component_with_transport`](Self::for_component_with_transport) to
343    /// override explicitly.
344    pub async fn for_component(comp: &Component, topic: impl Into<String>) -> Result<Self> {
345        let transport_kind = comp.drt().default_event_transport_kind();
346        Self::for_component_with_transport(comp, topic, transport_kind).await
347    }
348
349    /// Create a publisher with explicit transport.
350    pub async fn for_component_with_transport(
351        comp: &Component,
352        topic: impl Into<String>,
353        transport_kind: EventTransportKind,
354    ) -> Result<Self> {
355        let drt = comp.drt();
356        let scope = EventScope::Component {
357            namespace: comp.namespace().name(),
358            component: comp.name().to_string(),
359        };
360        Self::new_internal(drt, scope, topic.into(), transport_kind).await
361    }
362
363    /// Create a publisher for a namespace-scoped topic.
364    ///
365    /// The event transport is chosen automatically: if `DYN_EVENT_PLANE` is set that
366    /// value is used; otherwise the runtime's default is used (ZMQ for local backends
367    /// such as `file`/`mem`, NATS for distributed backends such as `etcd`/`kubernetes`).
368    /// Use [`for_namespace_with_transport`](Self::for_namespace_with_transport) to
369    /// override explicitly.
370    pub async fn for_namespace(ns: &Namespace, topic: impl Into<String>) -> Result<Self> {
371        let transport_kind = ns.drt().default_event_transport_kind();
372        Self::for_namespace_with_transport(ns, topic, transport_kind).await
373    }
374
375    /// Create a namespace publisher with explicit transport.
376    pub async fn for_namespace_with_transport(
377        ns: &Namespace,
378        topic: impl Into<String>,
379        transport_kind: EventTransportKind,
380    ) -> Result<Self> {
381        let drt = ns.drt();
382        let scope = EventScope::Namespace { name: ns.name() };
383        Self::new_internal(drt, scope, topic.into(), transport_kind).await
384    }
385
386    async fn new_internal(
387        drt: &DistributedRuntime,
388        scope: EventScope,
389        topic: String,
390        transport_kind: EventTransportKind,
391    ) -> Result<Self> {
392        // Publishers are discovery objects in their own right. A single process
393        // can host multiple publishers for the same scope/topic, each with its
394        // own ZMQ endpoint and sequence space, so the process ID is not unique
395        // enough here.
396        let publisher_id = discovery_safe_publisher_id(
397            rand::rngs::OsRng
398                .try_next_u64()
399                .map_err(|error| anyhow::anyhow!("failed to generate publisher ID: {error}"))?,
400        );
401        let discovery = Some(drt.discovery());
402        let runtime_handle = drt.runtime().secondary();
403        let subject = scope.subject(&topic);
404        let graceful_shutdown_tracker = drt.graceful_shutdown_tracker();
405
406        // Use Msgpack codec for all transports
407        enum TransportSetup {
408            Nats(Arc<dyn EventTransportTx>, Arc<Codec>),
409            ZmqDirect(Arc<dyn EventTransportTx>, Arc<Codec>, String), // includes public endpoint
410            ZmqBroker(Arc<dyn EventTransportTx>, Arc<Codec>),
411        }
412
413        let transport_setup = match transport_kind {
414            EventTransportKind::Nats => {
415                let transport = Arc::new(nats_transport::NatsTransport::new_publisher(
416                    drt.clone(),
417                    subject.clone(),
418                ));
419                let codec = Arc::new(Codec::Msgpack(MsgpackCodec));
420                TransportSetup::Nats(transport as Arc<dyn EventTransportTx>, codec)
421            }
422            EventTransportKind::Zmq => {
423                // Check for broker mode
424                if let Some(broker) = resolve_zmq_broker(drt, &scope).await? {
425                    // BROKER MODE: Connect to broker (single or multiple endpoints)
426                    let pub_transport = if broker.xsub_endpoints.len() == 1 {
427                        zmq_transport::ZmqPubTransport::connect(&broker.xsub_endpoints[0], &subject)
428                            .await?
429                    } else {
430                        zmq_transport::ZmqPubTransport::connect_multiple(
431                            &broker.xsub_endpoints,
432                            &subject,
433                        )
434                        .await?
435                    };
436
437                    let codec = Arc::new(Codec::Msgpack(MsgpackCodec));
438                    TransportSetup::ZmqBroker(
439                        Arc::new(pub_transport) as Arc<dyn EventTransportTx>,
440                        codec,
441                    )
442                } else {
443                    // DIRECT MODE: Bind PUB socket
444                    let (pub_transport, actual_bind_endpoint) = std::thread::spawn({
445                        let topic = topic.clone();
446                        move || {
447                            let rt = tokio::runtime::Builder::new_current_thread()
448                                .enable_all()
449                                .build()
450                                .expect("Failed to create Tokio runtime for ZMQ");
451
452                            rt.block_on(async move {
453                                zmq_transport::ZmqPubTransport::bind("tcp://0.0.0.0:0", &topic)
454                                    .await
455                                    .expect("Failed to bind ZMQ publisher")
456                            })
457                        }
458                    })
459                    .join()
460                    .expect("Failed to join ZMQ initialization thread");
461
462                    // Get local IP for public endpoint
463                    let actual_port: u16 = actual_bind_endpoint
464                        .rsplit(':')
465                        .next()
466                        .and_then(|s| s.parse().ok())
467                        .expect("Failed to parse port from bind endpoint");
468                    let local_ip = local_ip_for_advertise();
469                    let public_endpoint = format!("tcp://{}:{}", local_ip, actual_port);
470
471                    let codec = Arc::new(Codec::Msgpack(MsgpackCodec));
472                    TransportSetup::ZmqDirect(
473                        Arc::new(pub_transport) as Arc<dyn EventTransportTx>,
474                        codec,
475                        public_endpoint,
476                    )
477                }
478            }
479        };
480
481        // Extract transport and codec, and register if needed
482        let (tx, codec, discovery_instance) = match transport_setup {
483            TransportSetup::Nats(tx, codec) => {
484                let transport_config = EventTransport::nats(scope.subject_prefix());
485                let spec = DiscoverySpec::EventChannel {
486                    scope: scope.clone(),
487                    topic: topic.clone(),
488                    publisher_id,
489                    transport: transport_config,
490                };
491
492                let discovery_instance = drt.discovery().register(spec).await?;
493                tracing::info!(
494                    topic = %topic,
495                    transport = ?transport_kind,
496                    publisher_id = %publisher_id,
497                    "EventPublisher registered with discovery"
498                );
499                (tx, codec, Some(discovery_instance))
500            }
501            TransportSetup::ZmqDirect(tx, codec, public_endpoint) => {
502                let transport_config = EventTransport::zmq(public_endpoint);
503                let spec = DiscoverySpec::EventChannel {
504                    scope: scope.clone(),
505                    topic: topic.clone(),
506                    publisher_id,
507                    transport: transport_config,
508                };
509
510                let discovery_instance = drt.discovery().register(spec).await?;
511                tracing::info!(
512                    topic = %topic,
513                    transport = ?transport_kind,
514                    publisher_id = %publisher_id,
515                    "EventPublisher registered with discovery (direct mode)"
516                );
517                (tx, codec, Some(discovery_instance))
518            }
519            TransportSetup::ZmqBroker(tx, codec) => {
520                tracing::info!(
521                    topic = %topic,
522                    transport = ?transport_kind,
523                    "EventPublisher in broker mode - skipping discovery registration"
524                );
525                (tx, codec, None)
526            }
527        };
528
529        Ok(Self {
530            transport_kind,
531            topic,
532            subject,
533            publisher_id,
534            sequence: AtomicU64::new(0),
535            tx,
536            codec,
537            runtime_handle,
538            graceful_shutdown_tracker,
539            discovery_client: discovery,
540            discovery_instance,
541        })
542    }
543
544    /// Publish a serializable event.
545    pub async fn publish<T: Serialize + Send + Sync>(&self, event: &T) -> Result<()> {
546        let payload = self.codec.encode_payload(event)?;
547        self.publish_bytes_ref(payload.as_ref()).await
548    }
549
550    /// Publish raw bytes.
551    pub async fn publish_bytes(&self, bytes: Vec<u8>) -> Result<()> {
552        self.publish_bytes_ref(&bytes).await
553    }
554
555    /// Publish raw bytes without taking ownership of the payload buffer.
556    pub async fn publish_bytes_ref(&self, bytes: &[u8]) -> Result<()> {
557        let envelope_bytes = self.codec.encode_envelope_parts(
558            self.publisher_id,
559            self.sequence.fetch_add(1, Ordering::SeqCst),
560            current_timestamp_ms(),
561            &self.topic,
562            bytes,
563        )?;
564
565        self.tx.publish(&self.subject, envelope_bytes).await
566    }
567
568    /// Get the publisher ID.
569    pub fn publisher_id(&self) -> u64 {
570        self.publisher_id
571    }
572
573    /// Get the topic.
574    pub fn topic(&self) -> &str {
575        &self.topic
576    }
577
578    /// Get the transport kind.
579    pub fn transport_kind(&self) -> EventTransportKind {
580        self.transport_kind
581    }
582}
583
584impl Drop for EventPublisher {
585    fn drop(&mut self) {
586        // Unregister from discovery on drop
587        if let (Some(discovery), Some(instance)) =
588            (self.discovery_client.take(), self.discovery_instance.take())
589        {
590            let topic = self.topic.clone();
591            let publisher_id = instance.instance_id();
592            let runtime_handle = self.runtime_handle.clone();
593            let shutdown_guard = self.graceful_shutdown_tracker.register_task();
594
595            // Drop can run outside any Tokio context (notably via PyO3 finalizers), so use
596            // the runtime that created the publisher rather than the ambient thread state.
597            let spawn_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
598                runtime_handle.spawn(async move {
599                    let _shutdown_guard = shutdown_guard;
600                    match discovery.unregister(instance).await {
601                        Ok(()) => {
602                            tracing::info!(
603                                topic = %topic,
604                                publisher_id = %publisher_id,
605                                "EventPublisher unregistered from discovery"
606                            );
607                        }
608                        Err(e) => {
609                            tracing::warn!(
610                                topic = %topic,
611                                publisher_id = %publisher_id,
612                                error = %e,
613                                "Failed to unregister EventPublisher from discovery"
614                            );
615                        }
616                    }
617                });
618            }));
619
620            if spawn_result.is_err() {
621                tracing::warn!(
622                    topic = %self.topic,
623                    publisher_id = %publisher_id,
624                    "Skipping EventPublisher unregister during drop because the runtime is unavailable"
625                );
626            }
627        }
628    }
629}
630
631/// Event subscriber for a specific topic.
632pub struct EventSubscriber {
633    stream: EventStream,
634    #[allow(dead_code)]
635    scope: EventScope,
636    #[allow(dead_code)]
637    topic: String,
638    codec: Arc<Codec>,
639}
640
641impl EventSubscriber {
642    /// Create a subscriber for an endpoint-scoped topic.
643    pub async fn for_endpoint(endpoint: &Endpoint, topic: impl Into<String>) -> Result<Self> {
644        let transport_kind = endpoint.drt().default_event_transport_kind();
645        Self::for_endpoint_with_transport(endpoint, topic, transport_kind).await
646    }
647
648    /// Create an endpoint-scoped subscriber with explicit transport.
649    pub async fn for_endpoint_with_transport(
650        endpoint: &Endpoint,
651        topic: impl Into<String>,
652        transport_kind: EventTransportKind,
653    ) -> Result<Self> {
654        Self::for_endpoint_id_with_transport(endpoint.drt(), &endpoint.id(), topic, transport_kind)
655            .await
656    }
657
658    /// Create a subscriber for an endpoint identity without constructing a
659    /// local [`Endpoint`] handle.
660    pub async fn for_endpoint_id(
661        drt: &DistributedRuntime,
662        endpoint: &EndpointId,
663        topic: impl Into<String>,
664    ) -> Result<Self> {
665        let transport_kind = drt.default_event_transport_kind();
666        Self::for_endpoint_id_with_transport(drt, endpoint, topic, transport_kind).await
667    }
668
669    /// Create an endpoint-identity subscriber with explicit transport.
670    pub async fn for_endpoint_id_with_transport(
671        drt: &DistributedRuntime,
672        endpoint: &EndpointId,
673        topic: impl Into<String>,
674        transport_kind: EventTransportKind,
675    ) -> Result<Self> {
676        Self::new_internal(
677            drt,
678            EventScope::Endpoint {
679                endpoint: endpoint.clone(),
680            },
681            topic.into(),
682            transport_kind,
683        )
684        .await
685    }
686
687    /// Create a subscriber for a component-scoped topic.
688    ///
689    /// The event transport is chosen automatically: if `DYN_EVENT_PLANE` is set that
690    /// value is used; otherwise the runtime's default is used (ZMQ for local backends
691    /// such as `file`/`mem`, NATS for distributed backends such as `etcd`/`kubernetes`).
692    /// Use [`for_component_with_transport`](Self::for_component_with_transport) to
693    /// override explicitly.
694    pub async fn for_component(comp: &Component, topic: impl Into<String>) -> Result<Self> {
695        let transport_kind = comp.drt().default_event_transport_kind();
696        Self::for_component_with_transport(comp, topic, transport_kind).await
697    }
698
699    /// Create a subscriber with explicit transport.
700    pub async fn for_component_with_transport(
701        comp: &Component,
702        topic: impl Into<String>,
703        transport_kind: EventTransportKind,
704    ) -> Result<Self> {
705        let drt = comp.drt();
706        let scope = EventScope::Component {
707            namespace: comp.namespace().name(),
708            component: comp.name().to_string(),
709        };
710        Self::new_internal(drt, scope, topic.into(), transport_kind).await
711    }
712
713    /// Create a subscriber for a namespace-scoped topic.
714    ///
715    /// The event transport is chosen automatically: if `DYN_EVENT_PLANE` is set that
716    /// value is used; otherwise the runtime's default is used (ZMQ for local backends
717    /// such as `file`/`mem`, NATS for distributed backends such as `etcd`/`kubernetes`).
718    /// Use [`for_namespace_with_transport`](Self::for_namespace_with_transport) to
719    /// override explicitly.
720    pub async fn for_namespace(ns: &Namespace, topic: impl Into<String>) -> Result<Self> {
721        let transport_kind = ns.drt().default_event_transport_kind();
722        Self::for_namespace_with_transport(ns, topic, transport_kind).await
723    }
724
725    /// Create a namespace subscriber with explicit transport.
726    pub async fn for_namespace_with_transport(
727        ns: &Namespace,
728        topic: impl Into<String>,
729        transport_kind: EventTransportKind,
730    ) -> Result<Self> {
731        let drt = ns.drt();
732        let scope = EventScope::Namespace { name: ns.name() };
733        Self::new_internal(drt, scope, topic.into(), transport_kind).await
734    }
735
736    async fn new_internal(
737        drt: &DistributedRuntime,
738        scope: EventScope,
739        topic: String,
740        transport_kind: EventTransportKind,
741    ) -> Result<Self> {
742        let discovery = drt.discovery();
743        let routing_key = scope.subject(&topic);
744
745        // Use Msgpack codec for all transports
746        let (wire_stream, codec): (WireStream, Arc<Codec>) = match transport_kind {
747            EventTransportKind::Nats => {
748                let transport = nats_transport::NatsTransport::new(drt.clone());
749                let stream = transport.subscribe(&routing_key).await?;
750                let codec = Arc::new(Codec::Msgpack(MsgpackCodec));
751                (stream, codec)
752            }
753            EventTransportKind::Zmq => {
754                // Check for broker mode
755                if let Some(broker) = resolve_zmq_broker(drt, &scope).await? {
756                    // BROKER MODE: Connect to broker's XPUB (single or multiple endpoints)
757                    let codec = Arc::new(Codec::Msgpack(MsgpackCodec));
758
759                    let stream: WireStream = if broker.xpub_endpoints.len() == 1 {
760                        // Single broker - no deduplication needed
761                        let sub_transport = zmq_transport::ZmqSubTransport::connect_broker(
762                            &broker.xpub_endpoints[0],
763                            &routing_key,
764                        )
765                        .await?;
766                        sub_transport.subscribe(&routing_key).await?
767                    } else {
768                        // Multiple brokers - need deduplication
769                        let sub_transport =
770                            zmq_transport::ZmqSubTransport::connect_broker_multiple(
771                                &broker.xpub_endpoints,
772                                &routing_key,
773                            )
774                            .await?;
775                        let inner_stream = sub_transport.subscribe(&routing_key).await?;
776
777                        // Wrap with deduplication (default cache size: 100,000 entries)
778                        Box::pin(DeduplicatingStream::new(
779                            inner_stream,
780                            codec.clone(),
781                            100_000,
782                        ))
783                    };
784
785                    (stream, codec)
786                } else {
787                    // DIRECT MODE: Use dynamic subscriber to discover and connect to publishers
788                    let query = match &scope {
789                        EventScope::Namespace { name } => {
790                            crate::discovery::DiscoveryQuery::EventChannels(
791                                crate::discovery::EventChannelQuery::namespace_topic(
792                                    name.clone(),
793                                    topic.clone(),
794                                ),
795                            )
796                        }
797                        EventScope::Component {
798                            namespace,
799                            component,
800                        } => crate::discovery::DiscoveryQuery::EventChannels(
801                            crate::discovery::EventChannelQuery::topic(
802                                namespace.clone(),
803                                component.clone(),
804                                topic.clone(),
805                            ),
806                        ),
807                        EventScope::Endpoint { endpoint } => {
808                            crate::discovery::DiscoveryQuery::EventChannels(
809                                crate::discovery::EventChannelQuery::endpoint_topic(
810                                    endpoint.clone(),
811                                    topic.clone(),
812                                ),
813                            )
814                        }
815                    };
816
817                    let subscriber = Arc::new(DynamicSubscriber::with_cancel_token(
818                        discovery,
819                        query,
820                        topic.clone(),
821                        drt.primary_token().child_token(),
822                    ));
823
824                    let stream = subscriber.start_zmq().await?;
825                    let codec = Arc::new(Codec::Msgpack(MsgpackCodec));
826                    (stream, codec)
827                }
828            }
829        };
830
831        // Filter by topic and decode envelopes
832        let topic_filter = topic.clone();
833        let codec_for_stream = codec.clone();
834        let stream = wire_stream.filter_map(move |result| {
835            let codec = codec_for_stream.clone();
836            let topic_filter = topic_filter.clone();
837            async move {
838                match result {
839                    Ok(bytes) => match codec.decode_envelope(&bytes) {
840                        Ok(envelope) => {
841                            // Filter by topic for transports that don't support native filtering
842                            if envelope.topic == topic_filter {
843                                Some(Ok(envelope))
844                            } else {
845                                None
846                            }
847                        }
848                        Err(e) => Some(Err(e)),
849                    },
850                    Err(e) => Some(Err(e)),
851                }
852            }
853        });
854
855        tracing::info!(
856            topic = %topic,
857            transport = ?transport_kind,
858            "EventSubscriber created"
859        );
860
861        Ok(Self {
862            stream: Box::pin(stream),
863            scope,
864            topic,
865            codec,
866        })
867    }
868
869    /// Get the next event envelope.
870    pub async fn next(&mut self) -> Option<Result<EventEnvelope>> {
871        self.stream.next().await
872    }
873
874    /// Subscribe with automatic deserialization.
875    pub fn typed<T: DeserializeOwned + Send + 'static>(self) -> TypedEventSubscriber<T> {
876        TypedEventSubscriber {
877            stream: self.stream,
878            codec: self.codec,
879            _marker: std::marker::PhantomData,
880        }
881    }
882}
883
884/// Typed event subscriber that deserializes payloads.
885pub struct TypedEventSubscriber<T> {
886    stream: EventStream,
887    codec: Arc<Codec>,
888    _marker: std::marker::PhantomData<T>,
889}
890
891impl<T: DeserializeOwned + Send + 'static> TypedEventSubscriber<T> {
892    /// Get the next typed event with its envelope.
893    pub async fn next(&mut self) -> Option<Result<(EventEnvelope, T)>> {
894        std::future::poll_fn(|cx| self.poll_next(cx)).await
895    }
896
897    /// Poll for the next typed event.
898    pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<(EventEnvelope, T)>>> {
899        match self.stream.as_mut().poll_next(cx) {
900            Poll::Ready(Some(envelope)) => Poll::Ready(Some(match envelope {
901                Ok(env) => match self.codec.decode_payload(&env.payload) {
902                    Ok(typed) => Ok((env, typed)),
903                    Err(e) => Err(e),
904                },
905                Err(e) => Err(e),
906            })),
907            Poll::Ready(None) => Poll::Ready(None),
908            Poll::Pending => Poll::Pending,
909        }
910    }
911}
912
913/// Get current timestamp in milliseconds since Unix epoch.
914fn current_timestamp_ms() -> u64 {
915    SystemTime::now()
916        .duration_since(UNIX_EPOCH)
917        .map(|d| d.as_millis() as u64)
918        .unwrap_or(0)
919}
920
921#[cfg(test)]
922mod tests {
923    use super::*;
924    use crate::config::environment_names::zmq_broker as broker_env;
925
926    #[test]
927    fn publisher_ids_survive_a_json_number_round_trip() {
928        // This historical full-range ID is not JSON-safe. It was observed
929        // rounding to 13584172880116488000 after a discovery round trip.
930        let unsafe_id: u64 = 13_584_172_880_116_487_724;
931        assert!(unsafe_id > MAX_JSON_SAFE_PUBLISHER_ID);
932        assert_ne!(unsafe_id as f64 as u64, unsafe_id);
933
934        for random_id in [0, 1, u64::MAX, unsafe_id, 6_633_287_539_119_378] {
935            let publisher_id = discovery_safe_publisher_id(random_id);
936            assert!(
937                publisher_id <= MAX_JSON_SAFE_PUBLISHER_ID,
938                "publisher ID {publisher_id} exceeds the JSON-safe integer range"
939            );
940            assert_eq!(
941                publisher_id as f64 as u64, publisher_id,
942                "publisher ID {publisher_id} must survive an f64 round trip"
943            );
944        }
945    }
946
947    #[test]
948    fn direct_zmq_topology_selection_is_narrow() {
949        let lookup = |url: Option<&str>, enabled: Option<&str>| {
950            let url = url.map(std::ffi::OsString::from);
951            let enabled = enabled.map(std::ffi::OsString::from);
952            move |key: &str| match key {
953                broker_env::DYN_ZMQ_BROKER_URL => url.clone(),
954                broker_env::DYN_ZMQ_BROKER_ENABLED => enabled.clone(),
955                _ => None,
956            }
957        };
958
959        assert!(uses_direct_zmq_from_lookup(
960            EventTransportKind::Zmq,
961            lookup(None, None)
962        ));
963        assert!(!uses_direct_zmq_from_lookup(
964            EventTransportKind::Zmq,
965            lookup(Some("xsub=tcp://broker:5555,xpub=tcp://broker:5556"), None)
966        ));
967        assert!(!uses_direct_zmq_from_lookup(
968            EventTransportKind::Zmq,
969            lookup(None, Some("true"))
970        ));
971        assert!(uses_direct_zmq_from_lookup(
972            EventTransportKind::Zmq,
973            lookup(None, Some("false"))
974        ));
975        assert!(!uses_direct_zmq_from_lookup(
976            EventTransportKind::Nats,
977            lookup(None, None)
978        ));
979    }
980
981    #[tokio::test]
982    async fn direct_zmq_endpoint_scopes_are_isolated() {
983        temp_env::async_with_vars(
984            [
985                (broker_env::DYN_ZMQ_BROKER_URL, None::<&str>),
986                (broker_env::DYN_ZMQ_BROKER_ENABLED, None::<&str>),
987            ],
988            async {
989                let runtime = crate::Runtime::from_current().expect("create runtime handle");
990                let drt = DistributedRuntime::new(
991                    runtime,
992                    crate::distributed::DistributedConfig::process_local(),
993                )
994                .await
995                .expect("create distributed runtime");
996                let component = drt
997                    .namespace("endpoint-event-isolation-test")
998                    .expect("create namespace")
999                    .component("worker")
1000                    .expect("create component");
1001                let endpoint_a = component.endpoint("a");
1002                let endpoint_b = component.endpoint("b");
1003
1004                let publisher_a = EventPublisher::for_endpoint_with_transport(
1005                    &endpoint_a,
1006                    "events",
1007                    EventTransportKind::Zmq,
1008                )
1009                .await
1010                .expect("create endpoint A publisher");
1011                let publisher_b = EventPublisher::for_endpoint_with_transport(
1012                    &endpoint_b,
1013                    "events",
1014                    EventTransportKind::Zmq,
1015                )
1016                .await
1017                .expect("create endpoint B publisher");
1018                let mut subscriber_a = EventSubscriber::for_endpoint_with_transport(
1019                    &endpoint_a,
1020                    "events",
1021                    EventTransportKind::Zmq,
1022                )
1023                .await
1024                .expect("create endpoint A subscriber");
1025                let mut subscriber_b = EventSubscriber::for_endpoint_with_transport(
1026                    &endpoint_b,
1027                    "events",
1028                    EventTransportKind::Zmq,
1029                )
1030                .await
1031                .expect("create endpoint B subscriber");
1032
1033                let receive = async {
1034                    loop {
1035                        publisher_a
1036                            .publish_bytes(vec![0xa1])
1037                            .await
1038                            .expect("publish endpoint A event");
1039                        publisher_b
1040                            .publish_bytes(vec![0xb2])
1041                            .await
1042                            .expect("publish endpoint B event");
1043
1044                        let a = tokio::time::timeout(
1045                            std::time::Duration::from_millis(100),
1046                            subscriber_a.next(),
1047                        )
1048                        .await;
1049                        let b = tokio::time::timeout(
1050                            std::time::Duration::from_millis(100),
1051                            subscriber_b.next(),
1052                        )
1053                        .await;
1054                        if let (Ok(Some(Ok(a))), Ok(Some(Ok(b)))) = (a, b) {
1055                            assert_eq!(a.publisher_id, publisher_a.publisher_id());
1056                            assert_eq!(a.payload.as_ref(), &[0xa1]);
1057                            assert_eq!(b.publisher_id, publisher_b.publisher_id());
1058                            assert_eq!(b.payload.as_ref(), &[0xb2]);
1059                            break;
1060                        }
1061
1062                        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1063                    }
1064                };
1065
1066                tokio::time::timeout(std::time::Duration::from_secs(5), receive)
1067                    .await
1068                    .expect("each endpoint subscriber should receive only its own publisher");
1069            },
1070        )
1071        .await;
1072    }
1073
1074    #[tokio::test]
1075    async fn direct_zmq_publishers_in_one_endpoint_fan_into_one_subscriber() {
1076        temp_env::async_with_vars(
1077            [
1078                (broker_env::DYN_ZMQ_BROKER_URL, None::<&str>),
1079                (broker_env::DYN_ZMQ_BROKER_ENABLED, None::<&str>),
1080            ],
1081            async {
1082                let runtime = crate::Runtime::from_current().expect("create runtime handle");
1083                let drt = DistributedRuntime::new(
1084                    runtime,
1085                    crate::distributed::DistributedConfig::process_local(),
1086                )
1087                .await
1088                .expect("create distributed runtime");
1089                let endpoint = drt
1090                    .namespace("endpoint-event-fan-in-test")
1091                    .expect("create namespace")
1092                    .component("worker")
1093                    .expect("create component")
1094                    .endpoint("generate");
1095                let publisher_a = EventPublisher::for_endpoint_with_transport(
1096                    &endpoint,
1097                    "events",
1098                    EventTransportKind::Zmq,
1099                )
1100                .await
1101                .expect("create first publisher");
1102                let publisher_b = EventPublisher::for_endpoint_with_transport(
1103                    &endpoint,
1104                    "events",
1105                    EventTransportKind::Zmq,
1106                )
1107                .await
1108                .expect("create second publisher");
1109                let mut subscriber = EventSubscriber::for_endpoint_with_transport(
1110                    &endpoint,
1111                    "events",
1112                    EventTransportKind::Zmq,
1113                )
1114                .await
1115                .expect("create subscriber");
1116
1117                let receive = async {
1118                    let mut publisher_ids = std::collections::HashSet::new();
1119                    while publisher_ids.len() < 2 {
1120                        publisher_a.publish_bytes(vec![0xa1]).await.unwrap();
1121                        publisher_b.publish_bytes(vec![0xb2]).await.unwrap();
1122                        if let Ok(Some(Ok(envelope))) = tokio::time::timeout(
1123                            std::time::Duration::from_millis(100),
1124                            subscriber.next(),
1125                        )
1126                        .await
1127                        {
1128                            publisher_ids.insert(envelope.publisher_id);
1129                        }
1130                    }
1131                    assert_eq!(
1132                        publisher_ids,
1133                        std::collections::HashSet::from([
1134                            publisher_a.publisher_id(),
1135                            publisher_b.publisher_id(),
1136                        ])
1137                    );
1138                    for publisher_id in [publisher_a.publisher_id(), publisher_b.publisher_id()] {
1139                        assert!(
1140                            publisher_id <= MAX_JSON_SAFE_PUBLISHER_ID,
1141                            "publisher ID {publisher_id} exceeds the JSON-safe integer range"
1142                        );
1143                    }
1144                };
1145                tokio::time::timeout(std::time::Duration::from_secs(5), receive)
1146                    .await
1147                    .expect("subscriber should receive both endpoint publishers");
1148            },
1149        )
1150        .await;
1151    }
1152
1153    #[tokio::test]
1154    async fn same_topic_publishers_are_independent_across_recreation() {
1155        temp_env::async_with_vars(
1156            [
1157                (broker_env::DYN_ZMQ_BROKER_URL, None::<&str>),
1158                (broker_env::DYN_ZMQ_BROKER_ENABLED, None::<&str>),
1159            ],
1160            async {
1161                let runtime = crate::Runtime::from_current().expect("create runtime handle");
1162                let drt = DistributedRuntime::new(
1163                    runtime,
1164                    crate::distributed::DistributedConfig::process_local(),
1165                )
1166                .await
1167                .expect("create distributed runtime");
1168                let component = drt
1169                    .namespace("event-publisher-test")
1170                    .expect("create namespace")
1171                    .component("worker")
1172                    .expect("create component");
1173
1174                let publisher_a = EventPublisher::for_component_with_transport(
1175                    &component,
1176                    "events",
1177                    EventTransportKind::Zmq,
1178                )
1179                .await
1180                .expect("create first publisher");
1181                let publisher_b = EventPublisher::for_component_with_transport(
1182                    &component,
1183                    "events",
1184                    EventTransportKind::Zmq,
1185                )
1186                .await
1187                .expect("create second publisher");
1188                let publisher_a_id = publisher_a.publisher_id();
1189                let publisher_b_id = publisher_b.publisher_id();
1190
1191                assert_ne!(publisher_a_id, publisher_b_id);
1192
1193                let query = DiscoveryQuery::EventChannels(EventChannelQuery::topic(
1194                    "event-publisher-test",
1195                    "worker",
1196                    "events",
1197                ));
1198                let instances = drt
1199                    .discovery()
1200                    .list(query.clone())
1201                    .await
1202                    .expect("list event publishers");
1203                assert_eq!(instances.len(), 2);
1204                assert!(
1205                    instances
1206                        .iter()
1207                        .any(|instance| instance.instance_id() == publisher_a_id)
1208                );
1209                assert!(
1210                    instances
1211                        .iter()
1212                        .any(|instance| instance.instance_id() == publisher_b_id)
1213                );
1214
1215                let mut subscriber = EventSubscriber::for_component_with_transport(
1216                    &component,
1217                    "events",
1218                    EventTransportKind::Zmq,
1219                )
1220                .await
1221                .expect("create subscriber");
1222                let mut received_a = false;
1223                let mut received_b = false;
1224
1225                tokio::time::timeout(std::time::Duration::from_secs(5), async {
1226                    while !received_a || !received_b {
1227                        publisher_a
1228                            .publish_bytes(vec![0xa1])
1229                            .await
1230                            .expect("publish from first publisher");
1231                        publisher_b
1232                            .publish_bytes(vec![0xb2])
1233                            .await
1234                            .expect("publish from second publisher");
1235
1236                        if let Ok(Some(envelope)) = tokio::time::timeout(
1237                            std::time::Duration::from_millis(100),
1238                            subscriber.next(),
1239                        )
1240                        .await
1241                        {
1242                            let envelope = envelope.expect("receive event envelope");
1243                            if envelope.publisher_id == publisher_a_id {
1244                                assert_eq!(envelope.payload.as_ref(), &[0xa1]);
1245                                received_a = true;
1246                            } else if envelope.publisher_id == publisher_b_id {
1247                                assert_eq!(envelope.payload.as_ref(), &[0xb2]);
1248                                received_b = true;
1249                            } else {
1250                                panic!("event from unexpected publisher");
1251                            }
1252                        }
1253
1254                        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1255                    }
1256                })
1257                .await
1258                .expect("subscriber should receive events from both publishers");
1259
1260                drop(publisher_a);
1261                let publisher_a_recreated = EventPublisher::for_component_with_transport(
1262                    &component,
1263                    "events",
1264                    EventTransportKind::Zmq,
1265                )
1266                .await
1267                .expect("recreate first publisher");
1268                let publisher_a_recreated_id = publisher_a_recreated.publisher_id();
1269
1270                assert_ne!(publisher_a_recreated_id, publisher_a_id);
1271                assert_ne!(publisher_a_recreated_id, publisher_b_id);
1272                assert_eq!(
1273                    publisher_a_recreated.sequence.load(Ordering::SeqCst),
1274                    0,
1275                    "a recreated publisher starts a new sequence space"
1276                );
1277
1278                tokio::time::timeout(std::time::Duration::from_secs(1), async {
1279                    loop {
1280                        let instances = drt
1281                            .discovery()
1282                            .list(query.clone())
1283                            .await
1284                            .expect("list event publishers after recreation");
1285                        if instances.len() == 2
1286                            && instances
1287                                .iter()
1288                                .any(|instance| instance.instance_id() == publisher_b_id)
1289                            && instances
1290                                .iter()
1291                                .any(|instance| instance.instance_id() == publisher_a_recreated_id)
1292                        {
1293                            break;
1294                        }
1295                        tokio::task::yield_now().await;
1296                    }
1297                })
1298                .await
1299                .expect("old publisher should unregister without removing current publishers");
1300
1301                let mut received_b_after_recreation = false;
1302                let mut received_recreated_a = false;
1303
1304                tokio::time::timeout(std::time::Duration::from_secs(5), async {
1305                    while !received_b_after_recreation || !received_recreated_a {
1306                        publisher_b
1307                            .publish_bytes(vec![0xb3])
1308                            .await
1309                            .expect("publish from second publisher after recreation");
1310                        publisher_a_recreated
1311                            .publish_bytes(vec![0xa2])
1312                            .await
1313                            .expect("publish from recreated publisher");
1314
1315                        if let Ok(Some(envelope)) = tokio::time::timeout(
1316                            std::time::Duration::from_millis(100),
1317                            subscriber.next(),
1318                        )
1319                        .await
1320                        {
1321                            let envelope = envelope.expect("receive event envelope after drop");
1322                            if envelope.publisher_id == publisher_b_id
1323                                && envelope.payload.as_ref() == [0xb3]
1324                            {
1325                                received_b_after_recreation = true;
1326                            } else if envelope.publisher_id == publisher_a_recreated_id
1327                                && envelope.payload.as_ref() == [0xa2]
1328                            {
1329                                received_recreated_a = true;
1330                            }
1331                        }
1332
1333                        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1334                    }
1335                })
1336                .await
1337                .expect("subscriber should receive from surviving and recreated publishers");
1338            },
1339        )
1340        .await;
1341    }
1342
1343    #[tokio::test]
1344    async fn dropped_publisher_unregister_completes_within_graceful_shutdown() {
1345        temp_env::async_with_vars(
1346            [
1347                (broker_env::DYN_ZMQ_BROKER_URL, None::<&str>),
1348                (broker_env::DYN_ZMQ_BROKER_ENABLED, None::<&str>),
1349            ],
1350            async {
1351                let runtime = crate::Runtime::from_current().expect("create runtime handle");
1352                let drt = DistributedRuntime::new(
1353                    runtime,
1354                    crate::distributed::DistributedConfig::process_local(),
1355                )
1356                .await
1357                .expect("create distributed runtime");
1358                let component = drt
1359                    .namespace("event-publisher-shutdown-test")
1360                    .expect("create namespace")
1361                    .component("worker")
1362                    .expect("create component");
1363
1364                let publisher = EventPublisher::for_component_with_transport(
1365                    &component,
1366                    "events",
1367                    EventTransportKind::Zmq,
1368                )
1369                .await
1370                .expect("create publisher");
1371                let publisher_id = publisher.publisher_id();
1372
1373                let query = DiscoveryQuery::EventChannels(EventChannelQuery::topic(
1374                    "event-publisher-shutdown-test",
1375                    "worker",
1376                    "events",
1377                ));
1378                let instances = drt
1379                    .discovery()
1380                    .list(query.clone())
1381                    .await
1382                    .expect("list event publishers");
1383                assert_eq!(instances.len(), 1);
1384                assert_eq!(instances[0].instance_id(), publisher_id);
1385
1386                let tracker = drt.graceful_shutdown_tracker();
1387                assert_eq!(tracker.get_count(), 0);
1388
1389                let main_token = drt.runtime().primary_token();
1390                let endpoint_token = drt.runtime().child_token();
1391
1392                // Dropping the publisher schedules the async discovery unregister.
1393                // The drop path must synchronously take a graceful-shutdown guard so
1394                // that `Runtime::shutdown` Phase 2 waits for the unregister task.
1395                drop(publisher);
1396                assert_eq!(
1397                    tracker.get_count(),
1398                    1,
1399                    "dropping a publisher must register its unregister work with the graceful-shutdown tracker"
1400                );
1401
1402                drt.runtime().shutdown();
1403
1404                tokio::time::timeout(std::time::Duration::from_secs(5), main_token.cancelled())
1405                    .await
1406                    .expect("graceful shutdown should complete once the unregister task finishes");
1407
1408                assert!(endpoint_token.is_cancelled());
1409                assert_eq!(
1410                    tracker.get_count(),
1411                    0,
1412                    "the unregister task must release its graceful-shutdown guard"
1413                );
1414
1415                // Phase 3 (main token cancellation) only runs after Phase 2 drained
1416                // the tracker, so the dropped publisher must already be unregistered.
1417                let instances = drt
1418                    .discovery()
1419                    .list(query)
1420                    .await
1421                    .expect("list event publishers after shutdown");
1422                assert!(
1423                    instances.is_empty(),
1424                    "unregister must complete within the graceful-shutdown window"
1425                );
1426            },
1427        )
1428        .await;
1429    }
1430
1431    #[tokio::test]
1432    async fn runtime_cancellation_stops_retained_direct_zmq_subscriber() {
1433        temp_env::async_with_vars(
1434            [
1435                (broker_env::DYN_ZMQ_BROKER_URL, None::<&str>),
1436                (broker_env::DYN_ZMQ_BROKER_ENABLED, None::<&str>),
1437            ],
1438            async {
1439                let runtime = crate::Runtime::from_current().expect("create runtime handle");
1440                let drt = DistributedRuntime::new(
1441                    runtime,
1442                    crate::distributed::DistributedConfig::process_local(),
1443                )
1444                .await
1445                .expect("create distributed runtime");
1446                let component = drt
1447                    .namespace("event-subscriber-shutdown-test")
1448                    .expect("create namespace")
1449                    .component("worker")
1450                    .expect("create component");
1451
1452                let mut subscriber = EventSubscriber::for_component_with_transport(
1453                    &component,
1454                    "events",
1455                    EventTransportKind::Zmq,
1456                )
1457                .await
1458                .expect("create subscriber");
1459
1460                drt.primary_token().cancel();
1461
1462                let next =
1463                    tokio::time::timeout(std::time::Duration::from_secs(1), subscriber.next())
1464                        .await
1465                        .expect("runtime cancellation should stop the subscriber stream");
1466                assert!(next.is_none());
1467            },
1468        )
1469        .await;
1470    }
1471
1472    #[test]
1473    fn test_event_scope_subject_prefix() {
1474        let scopes = [
1475            (
1476                EventScope::Namespace {
1477                    name: "ns.one".to_string(),
1478                },
1479                "namespace.ns%2Eone",
1480            ),
1481            (
1482                EventScope::Component {
1483                    namespace: "ns.one".to_string(),
1484                    component: "worker/*".to_string(),
1485                },
1486                "namespace.ns%2Eone.component.worker%2F%2A",
1487            ),
1488            (
1489                EventScope::Endpoint {
1490                    endpoint: EndpointId {
1491                        namespace: "ns.one".to_string(),
1492                        component: "worker/*".to_string(),
1493                        name: "generate.>".to_string(),
1494                    },
1495                },
1496                "namespace.ns%2Eone.component.worker%2F%2A.endpoint.generate%2E%3E",
1497            ),
1498        ];
1499
1500        for (scope, expected_prefix) in scopes {
1501            assert_eq!(scope.subject_prefix(), expected_prefix);
1502            assert_eq!(
1503                scope.subject("kv.events/*"),
1504                format!("{expected_prefix}.kv%2Eevents%2F%2A")
1505            );
1506        }
1507    }
1508
1509    #[test]
1510    fn test_event_scope_accessors() {
1511        let ns_scope = EventScope::Namespace {
1512            name: "my-ns".to_string(),
1513        };
1514        assert_eq!(ns_scope.namespace(), "my-ns");
1515        assert_eq!(ns_scope.component(), None);
1516
1517        let comp_scope = EventScope::Component {
1518            namespace: "my-ns".to_string(),
1519            component: "my-comp".to_string(),
1520        };
1521        assert_eq!(comp_scope.namespace(), "my-ns");
1522        assert_eq!(comp_scope.component(), Some("my-comp"));
1523    }
1524
1525    #[test]
1526    fn test_event_envelope_serde() {
1527        let envelope = EventEnvelope {
1528            publisher_id: 42,
1529            sequence: 10,
1530            published_at: 1700000000000,
1531            topic: "test-topic".to_string(),
1532            payload: Bytes::from("test data"),
1533        };
1534
1535        let json = serde_json::to_string(&envelope).expect("serialize");
1536        let deserialized: EventEnvelope = serde_json::from_str(&json).expect("deserialize");
1537
1538        assert_eq!(deserialized.publisher_id, 42);
1539        assert_eq!(deserialized.sequence, 10);
1540        assert_eq!(deserialized.published_at, 1700000000000);
1541        assert_eq!(deserialized.topic, "test-topic");
1542        assert_eq!(deserialized.payload, Bytes::from("test data"));
1543    }
1544}
dynamo_runtime/transports/event_plane/mod.rs

dynamo_runtime/transports/event_plane/
mod.rs
