rust_zero_core/
pubsub.rs

1use std::{
2    collections::HashMap,
3    hash::Hash,
4    sync::{Arc, Mutex},
5};
6
7use tokio::sync::broadcast;
8
9/// A typed, topic-based in-process message broker.
10///
11/// Slow subscribers receive [`broadcast::error::RecvError::Lagged`] rather than blocking
12/// publishers, matching the fail-fast behavior expected from service event buses.
13pub struct Broker<Topic, Message> {
14    capacity: usize,
15    topics: Arc<Mutex<HashMap<Topic, broadcast::Sender<Message>>>>,
16}
17
18impl<Topic, Message> Clone for Broker<Topic, Message> {
19    fn clone(&self) -> Self {
20        Self {
21            capacity: self.capacity,
22            topics: Arc::clone(&self.topics),
23        }
24    }
25}
26
27impl<Topic, Message> Broker<Topic, Message>
28where
29    Topic: Clone + Eq + Hash,
30    Message: Clone,
31{
32    pub fn new(capacity: usize) -> Self {
33        assert!(capacity > 0, "broker capacity must be greater than zero");
34        Self {
35            capacity,
36            topics: Arc::new(Mutex::new(HashMap::new())),
37        }
38    }
39
40    pub fn subscribe(&self, topic: Topic) -> Subscription<Message> {
41        let mut topics = self.topics.lock().expect("broker lock poisoned");
42        let sender = topics.entry(topic).or_insert_with(|| {
43            let (sender, _) = broadcast::channel(self.capacity);
44            sender
45        });
46        Subscription {
47            receiver: sender.subscribe(),
48        }
49    }
50
51    /// Publishes a message and returns the number of active subscribers.
52    pub fn publish(&self, topic: Topic, message: Message) -> usize {
53        let mut topics = self.topics.lock().expect("broker lock poisoned");
54        let sender = topics.entry(topic).or_insert_with(|| {
55            let (sender, _) = broadcast::channel(self.capacity);
56            sender
57        });
58        sender.send(message).unwrap_or(0)
59    }
60
61    pub fn subscriber_count(&self, topic: &Topic) -> usize {
62        self.topics
63            .lock()
64            .expect("broker lock poisoned")
65            .get(topic)
66            .map_or(0, broadcast::Sender::receiver_count)
67    }
68}
69
70/// A subscription to one broker topic.
71pub struct Subscription<Message> {
72    receiver: broadcast::Receiver<Message>,
73}
74
75impl<Message> Subscription<Message>
76where
77    Message: Clone,
78{
79    pub async fn recv(&mut self) -> Result<Message, broadcast::error::RecvError> {
80        self.receiver.recv().await
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[tokio::test]
89    async fn fans_messages_out_to_topic_subscribers() {
90        let broker = Broker::new(8);
91        let mut first = broker.subscribe("orders");
92        let mut second = broker.subscribe("orders");
93        let mut unrelated = broker.subscribe("users");
94
95        assert_eq!(broker.publish("orders", 42), 2);
96        assert_eq!(first.recv().await.unwrap(), 42);
97        assert_eq!(second.recv().await.unwrap(), 42);
98        assert!(unrelated.receiver.try_recv().is_err());
99    }
100}
rust_zero_core/pubsub.rs

rust_zero_core/
pubsub.rs
