Skip to main content

monocoque/zmq/
publisher.rs

1//! PUB socket implementation with worker pool architecture.
2
3use bytes::Bytes;
4use monocoque_core::monitor::{SocketEventSender, SocketMonitor, create_monitor};
5use monocoque_core::options::SocketOptions;
6use monocoque_core::rt::TcpListener;
7use monocoque_zmtp::SocketType;
8use monocoque_zmtp::publisher::PubSocket as InternalPub;
9use std::io;
10
11/// A PUB socket for broadcasting messages to multiple subscribers.
12///
13/// PubSocket uses a **worker pool architecture** to handle multiple subscribers efficiently:
14/// - Multiple OS threads (default: CPU core count)
15/// - Each worker runs its own compio runtime with io_uring
16/// - Round-robin subscriber distribution across workers
17/// - Zero-copy message broadcasting via `Arc<Bytes>`
18/// - Lock-free subscription management
19///
20/// ## Example
21///
22/// ```rust,no_run
23/// use monocoque::zmq::PubSocket;
24/// use bytes::Bytes;
25///
26/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
27/// let mut socket = PubSocket::bind("127.0.0.1:5555").await?;
28///
29/// // Accept subscribers (non-blocking with worker pool)
30/// socket.accept_subscriber().await?;
31///
32/// // Broadcast to all subscribers
33/// socket.send(vec![Bytes::from("topic"), Bytes::from("data")]).await?;
34/// # Ok(())
35/// # }
36/// ```
37pub struct PubSocket {
38    inner: InternalPub,
39    listener: TcpListener,
40    monitor: Option<SocketEventSender>,
41}
42
43impl PubSocket {
44    /// Bind to an address with default worker count (CPU cores).
45    pub async fn bind(addr: impl monocoque_core::rt::ToSocketAddrs) -> io::Result<Self> {
46        let listener = TcpListener::bind(addr).await?;
47        Ok(Self {
48            inner: InternalPub::new(),
49            listener,
50            monitor: None,
51        })
52    }
53
54    /// Bind with a specific number of worker threads.
55    pub async fn bind_with_workers(
56        addr: impl monocoque_core::rt::ToSocketAddrs,
57        worker_count: usize,
58    ) -> io::Result<Self> {
59        let listener = TcpListener::bind(addr).await?;
60        Ok(Self {
61            inner: InternalPub::with_workers(worker_count),
62            listener,
63            monitor: None,
64        })
65    }
66
67    /// Accept a new subscriber connection.
68    ///
69    /// Performs ZMTP handshake and assigns the subscriber to a worker thread.
70    /// Returns the subscriber ID.
71    pub async fn accept_subscriber(&mut self) -> io::Result<u64> {
72        self.inner.accept_subscriber(&self.listener).await
73    }
74
75    /// Broadcast a multipart message to all matching subscribers.
76    ///
77    /// Messages are distributed to all workers in parallel.
78    /// The first frame is typically used as a topic for subscription filtering.
79    pub async fn send(&mut self, msg: Vec<Bytes>) -> io::Result<()> {
80        self.inner.send(msg).await
81    }
82
83    /// Broadcast a message given as borrowed frames.
84    ///
85    /// Allocation-light counterpart to [`send`](Self::send): the shared message
86    /// is allocated only when it matches a subscription, so publishing from a
87    /// stack array (`send_frames(&[topic, payload])`) pays no per-message heap
88    /// allocation on the common drop path of a topic-filtered stream.
89    pub async fn send_frames(&mut self, frames: &[Bytes]) -> io::Result<()> {
90        self.inner.send_frames(frames).await
91    }
92
93    /// Get the number of active subscribers.
94    pub const fn subscriber_count(&self) -> usize {
95        self.inner.subscriber_count()
96    }
97
98    /// Get the local address this socket is bound to.
99    pub fn local_addr(&self) -> io::Result<std::net::SocketAddr> {
100        self.listener.local_addr()
101    }
102
103    /// Get the socket type.
104    ///
105    /// # ZeroMQ Compatibility
106    ///
107    /// Corresponds to `ZMQ_TYPE` (16) option.
108    #[inline]
109    pub const fn socket_type() -> SocketType {
110        SocketType::Pub
111    }
112
113    /// Enable monitoring for this socket.
114    ///
115    /// Returns a receiver for socket lifecycle events.
116    pub fn monitor(&mut self) -> SocketMonitor {
117        let (sender, receiver) = create_monitor();
118        self.monitor = Some(sender);
119        receiver
120    }
121
122    /// Get a mutable reference to this socket's options.
123    #[inline]
124    pub fn options_mut(&mut self) -> &mut SocketOptions {
125        self.inner.options_mut()
126    }
127
128    /// Number of messages dropped due to HWM backpressure.
129    #[inline]
130    pub fn drop_count(&self) -> u64 {
131        self.inner.drop_count()
132    }
133}