monocoque_zmtp/
xsub.rs

1//! XSUB (Extended Subscriber) socket implementation
2//!
3//! XSUB extends SUB by sending subscription messages upstream to publishers,
4//! enabling subscription forwarding in message brokers and dynamic subscription
5//! management.
6//!
7//! # Use Cases
8//!
9//! - **Message brokers**: Forward subscriptions from frontend to backend
10//! - **Cascading pub/sub**: Build subscription trees across network boundaries
11//! - **Dynamic subscriptions**: Programmatically manage topic interests
12//!
13//! # Pattern
14//!
15//! ```text
16//! XSUB ──subscribe("topic.a")──> Publisher
17//!      <──────data("topic.a")───
18//! XSUB ──subscribe("topic.b")──> Publisher
19//!      <──────data("topic.b")───
20//! ```
21
22use crate::base::SocketBase;
23use bytes::Bytes;
24use compio_io::{AsyncRead, AsyncWrite};
25use monocoque_core::endpoint::Endpoint;
26use monocoque_core::options::SocketOptions;
27use monocoque_core::rt::TcpStream;
28use monocoque_core::subscription::{SubscriptionEvent, SubscriptionTrie};
29use smallvec::SmallVec;
30use std::io;
31use tracing::{debug, trace};
32
33use crate::handshake::perform_handshake_with_options;
34use crate::session::SocketType;
35
36/// XSUB (Extended Subscriber) socket.
37///
38/// Receives data messages and can send subscription messages upstream.
39///
40/// # Features
41///
42/// - **Dynamic subscriptions**: Subscribe/unsubscribe at runtime
43/// - **Subscription forwarding**: Forward subscriptions in proxies
44/// - **Verbose unsubscribe**: Optionally send explicit unsubscribe messages
45///
46/// # Examples
47///
48/// ```no_run
49/// use monocoque_zmtp::xsub::XSubSocket;
50/// use bytes::Bytes;
51///
52/// # async fn example() -> std::io::Result<()> {
53/// let mut xsub = XSubSocket::connect("127.0.0.1:5555").await?;
54///     
55///     // Subscribe to topics
56///     xsub.subscribe("topic.").await?;
57///
58///     // Receive messages
59///     if let Some(msg) = xsub.recv().await? {
60///         println!("Received: {:?}", msg);
61///     }
62///
63///     // Unsubscribe
64/// xsub.unsubscribe("topic.").await?;
65///
66/// # Ok(())
67/// # }
68/// ```
69pub struct XSubSocket<S = TcpStream>
70where
71    S: AsyncRead + AsyncWrite + Unpin,
72{
73    /// Base socket infrastructure
74    base: SocketBase<S>,
75    /// Local subscription tracking (XSUB manages subscriptions locally)
76    subscriptions: SubscriptionTrie,
77}
78
79impl<S> XSubSocket<S>
80where
81    S: AsyncRead + AsyncWrite + Unpin,
82{
83    /// Create a new XSUB socket from a stream.
84    pub async fn new(stream: S) -> io::Result<Self> {
85        Self::with_options(stream, SocketOptions::default()).await
86    }
87
88    /// Create a new XSUB socket with custom configuration and options.
89    pub async fn with_options(mut stream: S, options: SocketOptions) -> io::Result<Self> {
90        debug!("[XSUB] Creating new XSUB socket");
91
92        // Perform ZMTP handshake
93        debug!("[XSUB] Performing ZMTP handshake...");
94        let handshake_result = perform_handshake_with_options(
95            &mut stream,
96            SocketType::Xsub,
97            options.routing_id.as_deref(),
98            Some(options.handshake_timeout),
99            &options,
100        )
101        .await
102        .map_err(|e| io::Error::other(format!("Handshake failed: {}", e)))?;
103
104        debug!(
105            peer_socket_type = ?handshake_result.peer_socket_type,
106            "[XSUB] Handshake complete"
107        );
108
109        let mut base = SocketBase::new(stream, SocketType::Xsub, options);
110        base.curve_cipher = handshake_result.curve_cipher;
111        Ok(Self {
112            base,
113            subscriptions: SubscriptionTrie::new(),
114        })
115    }
116
117    /// Subscribe to messages with the given prefix.
118    ///
119    /// Sends a subscription message upstream to the publisher.
120    ///
121    /// # Examples
122    ///
123    /// ```no_run
124    /// # use monocoque_zmtp::xsub::XSubSocket;
125    /// # async fn example(mut xsub: XSubSocket) -> std::io::Result<()> {
126    /// // Subscribe to all messages starting with "topic."
127    /// xsub.subscribe("topic.").await?;
128    ///
129    /// // Subscribe to all messages (empty prefix)
130    /// xsub.subscribe("").await?;
131    /// # Ok(())
132    /// # }
133    /// ```
134    pub async fn subscribe(&mut self, prefix: impl Into<Bytes>) -> io::Result<()> {
135        let prefix = prefix.into();
136        trace!("[XSUB] Subscribing to: {:?}", prefix);
137
138        // Send subscription message upstream first, then record the prefix locally.
139        self.send_subscription_event_prefix(0x01, &prefix).await?;
140        self.subscriptions.subscribe(prefix);
141
142        Ok(())
143    }
144
145    /// Unsubscribe from messages with the given prefix.
146    ///
147    /// Optionally sends an unsubscribe message upstream (if verbose mode enabled).
148    ///
149    /// # Examples
150    ///
151    /// ```no_run
152    /// # use monocoque_zmtp::xsub::XSubSocket;
153    /// # use bytes::Bytes;
154    /// # async fn example(mut xsub: XSubSocket) -> std::io::Result<()> {
155    /// let prefix = Bytes::from_static(b"topic.");
156    /// xsub.unsubscribe(prefix).await?;
157    /// # Ok(())
158    /// # }
159    /// ```
160    pub async fn unsubscribe(&mut self, prefix: impl Into<Bytes>) -> io::Result<()> {
161        let prefix = prefix.into();
162        trace!("[XSUB] Unsubscribing from: {:?}", prefix);
163
164        self.subscriptions.unsubscribe(&prefix);
165
166        // Send unsubscribe message if verbose mode enabled
167        if self.base.options.xsub_verbose_unsubs {
168            self.send_subscription_event_prefix(0x00, &prefix).await?;
169        }
170
171        Ok(())
172    }
173
174    /// Send a raw subscription event upstream (for proxies).
175    ///
176    /// This allows forwarding subscription messages in broker patterns.
177    ///
178    /// # Examples
179    ///
180    /// ```no_run
181    /// # use monocoque_zmtp::xsub::XSubSocket;
182    /// # use monocoque_core::subscription::SubscriptionEvent;
183    /// # use bytes::Bytes;
184    /// # async fn example(mut xsub: XSubSocket) -> std::io::Result<()> {
185    /// let event = SubscriptionEvent::Subscribe(Bytes::from("topic"));
186    /// xsub.send_subscription_event(event).await?;
187    /// # Ok(())
188    /// # }
189    /// ```
190    pub async fn send_subscription_event(&mut self, event: SubscriptionEvent) -> io::Result<()> {
191        let (cmd, prefix) = match &event {
192            SubscriptionEvent::Subscribe(prefix) => (0x01, prefix.as_ref()),
193            SubscriptionEvent::Unsubscribe(prefix) => (0x00, prefix.as_ref()),
194        };
195
196        self.send_subscription_event_prefix(cmd, prefix).await
197    }
198
199    async fn send_subscription_event_prefix(&mut self, cmd: u8, prefix: &[u8]) -> io::Result<()> {
200        use bytes::BytesMut;
201        use compio_buf::BufResult;
202        use compio_io::AsyncWriteExt;
203
204        trace!(
205            "[XSUB] Sending subscription event ({} bytes)",
206            1 + prefix.len()
207        );
208
209        let mut raw = BytesMut::with_capacity(1 + prefix.len());
210        raw.extend_from_slice(&[cmd]);
211        raw.extend_from_slice(prefix);
212        let raw = raw.freeze();
213
214        // Encrypt if CURVE is active; otherwise plain ZMTP frame.
215        let mut wire = BytesMut::with_capacity(raw.len() + 9);
216        if let Some(ref mut cipher) = self.base.curve_cipher {
217            let body = cipher
218                .encrypt_frame(&raw, false)
219                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
220            crate::base::append_zmtp_cmd_frame(&mut wire, &body);
221        } else {
222            crate::codec::encode_multipart(&[raw], &mut wire);
223        }
224        let wire = wire.freeze();
225
226        let stream =
227            self.base.stream.as_mut().ok_or_else(|| {
228                io::Error::new(io::ErrorKind::NotConnected, "Socket not connected")
229            })?;
230
231        let BufResult(result, _) = stream.write_all(wire).await;
232        result?;
233
234        trace!("[XSUB] Subscription event sent successfully");
235        Ok(())
236    }
237
238    /// Receive a data message (non-blocking).
239    ///
240    /// Returns `None` if no message is available.
241    ///
242    /// # Examples
243    ///
244    /// ```no_run
245    /// # use monocoque_zmtp::xsub::XSubSocket;
246    /// # async fn example(mut xsub: XSubSocket) -> std::io::Result<()> {
247    /// if let Some(msg) = xsub.recv().await? {
248    ///     for frame in msg {
249    ///         println!("Frame: {:?}", frame);
250    ///     }
251    /// }
252    /// # Ok(())
253    /// # }
254    /// ```
255    pub async fn recv(&mut self) -> io::Result<Option<Vec<Bytes>>> {
256        let mut frames: SmallVec<[Bytes; 4]> = SmallVec::new();
257
258        loop {
259            loop {
260                match self.base.process_frame()? {
261                    crate::base::FrameResult::NeedMore => break,
262                    crate::base::FrameResult::CommandHandled => {
263                        if !self.base.send_buffer.is_empty() {
264                            self.base.flush_send_buffer().await?;
265                        }
266                    }
267                    crate::base::FrameResult::Data(more, payload) => {
268                        frames.push(payload);
269                        if !more {
270                            trace!("[XSUB] Received {} frames", frames.len());
271                            return Ok(Some(frames.into_vec()));
272                        }
273                    }
274                }
275            }
276
277            let n = self.base.read_raw().await?;
278            if n == 0 {
279                trace!("[XSUB] Connection closed");
280                return Ok(None);
281            }
282            if self.base.check_heartbeat()? {
283                self.base.flush_send_buffer().await?;
284            }
285        }
286    }
287
288    /// Get the number of active subscriptions.
289    pub fn subscription_count(&self) -> usize {
290        self.subscriptions.len()
291    }
292
293    /// Check if subscribed to a specific topic.
294    pub fn is_subscribed(&self, topic: &[u8]) -> bool {
295        self.subscriptions.matches(topic)
296    }
297
298    /// Get all subscriptions.
299    pub fn subscriptions(&self) -> Vec<monocoque_core::subscription::Subscription> {
300        self.subscriptions.subscriptions()
301    }
302
303    /// Get the socket type.
304    pub const fn socket_type(&self) -> SocketType {
305        SocketType::Xsub
306    }
307
308    /// Get the endpoint this socket is connected/bound to, if available.
309    ///
310    /// Returns `None` if the socket was created from a raw stream.
311    ///
312    /// # ZeroMQ Compatibility
313    ///
314    /// Corresponds to `ZMQ_LAST_ENDPOINT` (32) option.
315    #[inline]
316    pub fn last_endpoint(&self) -> Option<&Endpoint> {
317        self.base.last_endpoint()
318    }
319
320    /// Check if the last received message has more frames coming.
321    ///
322    /// Returns `true` if there are more frames in the current multipart message.
323    ///
324    /// # ZeroMQ Compatibility
325    ///
326    /// Corresponds to `ZMQ_RCVMORE` (13) option.
327    #[inline]
328    pub fn has_more(&self) -> bool {
329        self.base.has_more()
330    }
331
332    /// Get the event state of the socket.
333    ///
334    /// Returns a bitmask indicating ready-to-receive and ready-to-send states.
335    ///
336    /// # Returns
337    ///
338    /// - `1` (POLLIN) - Socket is ready to receive
339    /// - `2` (POLLOUT) - Socket is ready to send
340    /// - `3` (POLLIN | POLLOUT) - Socket is ready for both
341    ///
342    /// # ZeroMQ Compatibility
343    ///
344    /// Corresponds to `ZMQ_EVENTS` (15) option.
345    #[inline]
346    pub fn events(&self) -> u32 {
347        self.base.events()
348    }
349}
350
351impl XSubSocket<TcpStream> {
352    /// Connect to a publisher, storing the endpoint for automatic reconnection.
353    ///
354    /// # Examples
355    ///
356    /// ```no_run
357    /// # use monocoque_zmtp::xsub::XSubSocket;
358    /// # async fn example() -> std::io::Result<()> {
359    /// let xsub = XSubSocket::connect("127.0.0.1:5555").await?;
360    /// # Ok(())
361    /// # }
362    /// ```
363    pub async fn connect(addr: &str) -> io::Result<Self> {
364        Self::connect_with_options(addr, SocketOptions::default()).await
365    }
366
367    /// Connect with custom socket options, storing the endpoint for automatic reconnection.
368    pub async fn connect_with_options(addr: &str, options: SocketOptions) -> io::Result<Self> {
369        let stream = TcpStream::connect(addr).await?;
370        let peer_addr = stream.peer_addr()?;
371
372        // Enable TCP_NODELAY (and keepalive) on the outbound connection, matching
373        // SUB's connect path. One-time setsockopt at connect, off the hot path.
374        crate::utils::configure_tcp_stream(&stream, &options, "XSUB")?;
375
376        let mut stream = stream;
377        let handshake_result = crate::handshake::perform_handshake_with_options(
378            &mut stream,
379            crate::session::SocketType::Xsub,
380            options.routing_id.as_deref(),
381            Some(options.handshake_timeout),
382            &options,
383        )
384        .await
385        .map_err(|e| io::Error::other(format!("Handshake failed: {}", e)))?;
386
387        debug!(
388            peer_identity = ?handshake_result.peer_identity,
389            peer_socket_type = ?handshake_result.peer_socket_type,
390            "[XSUB] Connected to {} (endpoint stored for reconnection)",
391            peer_addr
392        );
393
394        let endpoint = monocoque_core::endpoint::Endpoint::Tcp(peer_addr);
395        let mut base = crate::base::SocketBase::with_endpoint(
396            stream,
397            crate::session::SocketType::Xsub,
398            endpoint,
399            options,
400        );
401        base.curve_cipher = handshake_result.curve_cipher;
402        Ok(Self {
403            base,
404            subscriptions: SubscriptionTrie::new(),
405        })
406    }
407
408    /// Check if the socket is currently connected.
409    #[inline]
410    pub fn is_connected(&self) -> bool {
411        self.base.is_connected()
412    }
413
414    /// Try to reconnect to the stored endpoint, re-sending all active subscriptions.
415    pub async fn try_reconnect(&mut self) -> io::Result<()> {
416        self.base
417            .try_reconnect(crate::session::SocketType::Xsub)
418            .await?;
419        // Re-send all subscriptions to the fresh connection
420        let prefixes: Vec<bytes::Bytes> = self
421            .subscriptions
422            .subscriptions()
423            .iter()
424            .map(|s| s.prefix.clone())
425            .collect();
426        for prefix in prefixes {
427            self.send_subscription_event(
428                monocoque_core::subscription::SubscriptionEvent::Subscribe(prefix),
429            )
430            .await?;
431        }
432        Ok(())
433    }
434
435    /// Receive a message with automatic reconnection on EOF or network error.
436    ///
437    /// Respects `max_reconnect_attempts` - returns `NotConnected` when exhausted.
438    pub async fn recv_with_reconnect(&mut self) -> io::Result<Option<Vec<bytes::Bytes>>> {
439        let max = self.base.options.max_reconnect_attempts;
440        let mut attempts = 0u32;
441
442        loop {
443            if self.base.stream.is_none() {
444                if let Some(limit) = max
445                    && attempts >= limit
446                {
447                    return Err(io::Error::new(
448                        io::ErrorKind::NotConnected,
449                        format!("Max {} reconnection attempts exceeded", limit),
450                    ));
451                }
452                attempts += 1;
453                trace!(
454                    "[XSUB] Stream disconnected, reconnecting (attempt {})",
455                    attempts
456                );
457                self.try_reconnect().await?;
458            }
459
460            match self.recv().await {
461                Ok(Some(msg)) => return Ok(Some(msg)),
462                Ok(None) => {
463                    debug!("[XSUB] EOF on recv, will reconnect");
464                }
465                Err(e) => {
466                    if self.base.stream.is_none()
467                        || matches!(
468                            e.kind(),
469                            io::ErrorKind::ConnectionReset
470                                | io::ErrorKind::ConnectionAborted
471                                | io::ErrorKind::BrokenPipe
472                                | io::ErrorKind::UnexpectedEof
473                        )
474                    {
475                        debug!("[XSUB] Connection error on recv ({}), will reconnect", e);
476                        self.base.stream = None;
477                    } else {
478                        return Err(e);
479                    }
480                }
481            }
482        }
483    }
484}
485
486#[cfg(test)]
487mod tests {
488    use super::*;
489
490    /// XSUB's outbound TCP connection must have TCP_NODELAY set, matching SUB.
491    /// Drives a real connect against an XPUB peer and reads TCP_NODELAY off the
492    /// live XSUB socket fd. One-time setsockopt at connect - off the hot path.
493    #[cfg(unix)]
494    #[test]
495    fn xsub_connect_sets_tcp_nodelay() {
496        use monocoque_core::rt::LocalRuntime;
497        use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
498        use std::sync::mpsc;
499        use std::thread;
500
501        fn fd_nodelay(fd: RawFd) -> bool {
502            let sock = unsafe { socket2::Socket::from_raw_fd(fd) };
503            let nd = sock.nodelay().expect("query TCP_NODELAY");
504            std::mem::forget(sock); // borrowed fd - do not close it
505            nd
506        }
507
508        let (port_tx, port_rx) = mpsc::channel::<u16>();
509        let (nd_tx, nd_rx) = mpsc::channel::<bool>();
510        let (done_tx, done_rx) = mpsc::channel::<()>();
511
512        // XPUB peer: bind, announce the port, accept the XSUB, hold it open.
513        let server = thread::spawn(move || {
514            let rt = LocalRuntime::new().unwrap();
515            rt.block_on(async move {
516                let mut xpub = crate::xpub::XPubSocket::bind("127.0.0.1:0").await.unwrap();
517                port_tx.send(xpub.local_addr().unwrap().port()).unwrap();
518                xpub.accept().await.unwrap();
519                done_rx.recv().unwrap();
520            });
521        });
522
523        let port = port_rx.recv().unwrap();
524        let client = thread::spawn(move || {
525            let rt = LocalRuntime::new().unwrap();
526            rt.block_on(async move {
527                let xsub = XSubSocket::connect(&format!("127.0.0.1:{port}"))
528                    .await
529                    .unwrap();
530                let fd = xsub.base.stream.as_ref().unwrap().as_raw_fd();
531                nd_tx.send(fd_nodelay(fd)).unwrap();
532                done_tx.send(()).unwrap();
533            });
534        });
535
536        let nodelay = nd_rx.recv().unwrap();
537        client.join().unwrap();
538        server.join().unwrap();
539        assert!(
540            nodelay,
541            "XSUB connect must set TCP_NODELAY on the outbound socket",
542        );
543    }
544
545    #[test]
546    fn test_subscription_tracking() {
547        use monocoque_core::rt::LocalRuntime as Runtime;
548
549        Runtime::new().unwrap().block_on(async {
550            // Mock stream for testing
551            // In real tests, use actual TCP connection
552        });
553    }
554
555    #[test]
556    fn test_subscription_event_creation() {
557        let event = SubscriptionEvent::Subscribe(Bytes::from_static(b"topic"));
558        let msg = event.to_message();
559        assert_eq!(msg[0], 0x01);
560        assert_eq!(&msg[1..], b"topic");
561    }
562}
563
564crate::impl_socket_trait!(XSubSocket<S>, SocketType::Xsub);
monocoque_zmtp/xsub.rs

monocoque_zmtp/
xsub.rs
