ullun/
hub.rs

1//! Publish/subscribe rooms
2//!
3//! A [`Hub`] fans messages out to everyone subscribed to a named room: chat
4//! rooms, live dashboards, per-document editing sessions. Subscribers that fall
5//! behind do not slow the others down: they are told how many messages they
6//! missed and carry on with the newest ones.
7//!
8//! ```
9//! # tokio_test_block(async {
10//! use ullun::hub::{Delivery, Hub};
11//!
12//! let hub: Hub<String> = Hub::new(64);
13//! let mut alice = hub.subscribe("lobby");
14//! let mut bob = hub.subscribe("lobby");
15//! assert_eq!(hub.publish("lobby", "hello".to_string()), 2);
16//!
17//! assert_eq!(alice.recv().await, Some(Delivery::Message("hello".to_string())));
18//! assert_eq!(bob.recv().await, Some(Delivery::Message("hello".to_string())));
19//! assert_eq!(hub.members("lobby"), 2);
20//! drop(bob);
21//! assert_eq!(hub.members("lobby"), 1);
22//! # });
23//! # fn tokio_test_block<F: std::future::Future>(f: F) {
24//! #     tokio::runtime::Builder::new_current_thread().build().unwrap().block_on(f);
25//! # }
26//! ```
27//!
28//! Rooms exist while they have subscribers and disappear with the last one, so
29//! a hub cannot be filled with empty rooms. Messages are not stored: keep
30//! history yourself if late joiners need it.
31
32use std::collections::HashMap;
33use std::sync::{Arc, Mutex, MutexGuard};
34use tokio::sync::broadcast;
35
36struct Room<T> {
37    tx: broadcast::Sender<T>,
38    subscribers: usize,
39}
40
41struct Shared<T> {
42    rooms: Mutex<HashMap<String, Room<T>>>,
43    capacity: usize,
44}
45
46/// Named rooms with fan-out; cheap to clone, all clones share the rooms
47pub struct Hub<T> {
48    shared: Arc<Shared<T>>,
49}
50
51impl<T> Clone for Hub<T> {
52    fn clone(&self) -> Self {
53        Self {
54            shared: self.shared.clone(),
55        }
56    }
57}
58
59/// What a subscription yields
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum Delivery<T> {
62    Message(T),
63    /// The subscriber was too slow and this many messages were dropped for it
64    Missed(u64),
65}
66
67/// A member of a room; leaves it when dropped
68pub struct Subscription<T> {
69    rx: broadcast::Receiver<T>,
70    room: String,
71    hub: Hub<T>,
72}
73
74impl<T: Clone + Send + 'static> Hub<T> {
75    /// `capacity`: messages buffered per room for slow subscribers (at least 1)
76    pub fn new(capacity: usize) -> Self {
77        Self {
78            shared: Arc::new(Shared {
79                rooms: Mutex::new(HashMap::new()),
80                capacity: capacity.max(1),
81            }),
82        }
83    }
84
85    fn lock(&self) -> MutexGuard<'_, HashMap<String, Room<T>>> {
86        self.shared.rooms.lock().unwrap_or_else(|e| e.into_inner())
87    }
88
89    /// Join `room`, creating it if needed. Only messages published after this call arrive.
90    pub fn subscribe(&self, room: &str) -> Subscription<T> {
91        let mut rooms = self.lock();
92        let entry = rooms.entry(room.to_string()).or_insert_with(|| Room {
93            tx: broadcast::channel(self.shared.capacity).0,
94            subscribers: 0,
95        });
96        entry.subscribers += 1;
97        Subscription {
98            rx: entry.tx.subscribe(),
99            room: room.to_string(),
100            hub: self.clone(),
101        }
102    }
103
104    /// Send to everyone currently in `room`; returns how many received it
105    /// (0 for a room nobody is in)
106    pub fn publish(&self, room: &str, message: T) -> usize {
107        let rooms = self.lock();
108        rooms
109            .get(room)
110            .map(|r| r.tx.send(message).unwrap_or(0))
111            .unwrap_or(0)
112    }
113
114    /// Subscribers in `room`
115    pub fn members(&self, room: &str) -> usize {
116        self.lock().get(room).map_or(0, |r| r.subscribers)
117    }
118
119    /// Rooms with at least one member, with their member counts, by name
120    pub fn rooms(&self) -> Vec<(String, usize)> {
121        let mut rooms: Vec<_> = self
122            .lock()
123            .iter()
124            .map(|(name, r)| (name.clone(), r.subscribers))
125            .collect();
126        rooms.sort();
127        rooms
128    }
129}
130
131impl<T: Clone + Send + 'static> Subscription<T> {
132    /// The next message, or how many were missed. `None` only if the hub is gone.
133    pub async fn recv(&mut self) -> Option<Delivery<T>> {
134        match self.rx.recv().await {
135            Ok(m) => Some(Delivery::Message(m)),
136            Err(broadcast::error::RecvError::Lagged(n)) => Some(Delivery::Missed(n)),
137            Err(broadcast::error::RecvError::Closed) => None,
138        }
139    }
140
141    pub fn room(&self) -> &str {
142        &self.room
143    }
144}
145
146impl<T> Drop for Subscription<T> {
147    fn drop(&mut self) {
148        let mut rooms = self
149            .hub
150            .shared
151            .rooms
152            .lock()
153            .unwrap_or_else(|e| e.into_inner());
154        if let Some(r) = rooms.get_mut(&self.room) {
155            r.subscribers = r.subscribers.saturating_sub(1);
156            if r.subscribers == 0 {
157                rooms.remove(&self.room);
158            }
159        }
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[tokio::test]
168    async fn rooms_are_isolated_and_counted() {
169        let hub: Hub<u32> = Hub::new(8);
170        let mut a1 = hub.subscribe("a");
171        let mut a2 = hub.subscribe("a");
172        let mut b = hub.subscribe("b");
173
174        assert_eq!(hub.publish("a", 1), 2);
175        assert_eq!(hub.publish("b", 2), 1);
176        assert_eq!(hub.publish("nobody", 3), 0);
177
178        assert_eq!(a1.recv().await, Some(Delivery::Message(1)));
179        assert_eq!(a2.recv().await, Some(Delivery::Message(1)));
180        assert_eq!(b.recv().await, Some(Delivery::Message(2)));
181        assert_eq!(hub.rooms(), [("a".to_string(), 2), ("b".to_string(), 1)]);
182        assert_eq!(a1.room(), "a");
183    }
184
185    #[tokio::test]
186    async fn late_joiners_only_see_new_messages_and_empty_rooms_vanish() {
187        let hub: Hub<&'static str> = Hub::new(8);
188        let early = hub.subscribe("r");
189        hub.publish("r", "before");
190        let mut late = hub.subscribe("r");
191        hub.publish("r", "after");
192        assert_eq!(late.recv().await, Some(Delivery::Message("after")));
193
194        drop(early);
195        assert_eq!(hub.members("r"), 1);
196        drop(late);
197        assert_eq!(hub.members("r"), 0);
198        assert!(
199            hub.rooms().is_empty(),
200            "the last member takes the room with it"
201        );
202        assert_eq!(hub.publish("r", "void"), 0);
203    }
204
205    #[tokio::test]
206    async fn slow_subscribers_learn_what_they_missed_and_do_not_block_others() {
207        let hub: Hub<u32> = Hub::new(2);
208        let mut slow = hub.subscribe("r");
209        let mut fast = hub.subscribe("r");
210        for i in 1..=5 {
211            hub.publish("r", i);
212            // the fast one keeps up
213            assert_eq!(fast.recv().await, Some(Delivery::Message(i)));
214        }
215        // the slow one lost the oldest messages, then gets the newest two
216        assert_eq!(slow.recv().await, Some(Delivery::Missed(3)));
217        assert_eq!(slow.recv().await, Some(Delivery::Message(4)));
218        assert_eq!(slow.recv().await, Some(Delivery::Message(5)));
219    }
220
221    #[tokio::test]
222    async fn clones_share_rooms_and_zero_capacity_is_bumped() {
223        let hub: Hub<u8> = Hub::new(0);
224        let other = hub.clone();
225        let mut sub = other.subscribe("x");
226        assert_eq!(hub.publish("x", 9), 1);
227        assert_eq!(sub.recv().await, Some(Delivery::Message(9)));
228    }
229}
ullun/hub.rs

ullun/
hub.rs
