Skip to main content

max / makenotwork

14.8 KB · 443 lines History Blame Raw
1 //! In-process fan-out.
2 //!
3 //! One [`tokio::sync::broadcast`] channel per room, created when the first
4 //! listener arrives and pruned when the last one leaves. This is correct for the
5 //! deployment both consumers actually have (a single process on a single box) and
6 //! is knowingly the wrong shape for more than one node. Going multi-node means
7 //! putting Postgres `LISTEN/NOTIFY` or a message bus behind [`Hub::publish`]; it
8 //! is a second entry on the same list as the in-process rate limiter, not a
9 //! surprise.
10 //!
11 //! # Why the connection caps are load-bearing
12 //!
13 //! Both units run under `MemoryMax=512M`, a hard cgroup cap. Each listener holds
14 //! a task, a buffer, and a broadcast receiver, and an OOM does not degrade chat,
15 //! it restarts the whole site. So the caps here are a availability control, not
16 //! tidiness, and they are enforced before any allocation the connection would
17 //! own.
18 //!
19 //! # Divergence from the SyncKit SSE pattern
20 //!
21 //! This is modeled on `server/src/routes/synckit/subscribe.rs` and deviates from
22 //! it twice, both deliberately:
23 //!
24 //! 1. **Pruning is atomic.** SyncKit reads `receiver_count()`, drops the guard,
25 //! then calls `remove()`. Between those two steps a new subscriber can take
26 //! the entry, and the `remove()` then drops a channel that has a live
27 //! listener; the next publish lazily creates a *different* channel and the
28 //! listener silently receives nothing. SyncKit tolerates this because its
29 //! clients re-pull on their own schedule. A chat client would just sit in a
30 //! dead room. [`DashMap::remove_if`] runs the predicate under the shard lock,
31 //! which closes the race and still avoids the same-shard deadlock that made
32 //! SyncKit split the read and the remove in the first place.
33 //! 2. **Lag is surfaced, not swallowed.** SyncKit discards `Lagged`. Here it
34 //! becomes [`ChatEvent::Gap`] so the client re-fetches from its cursor,
35 //! because a room with a silent hole in it is worse than a visible one.
36
37 use std::sync::Arc;
38 use std::sync::atomic::{AtomicUsize, Ordering};
39
40 use dashmap::DashMap;
41 use tokio::sync::broadcast;
42
43 use crate::error::{ChatError, LimitScope};
44 use crate::event::ChatEvent;
45 use crate::ids::{RoomId, UserId};
46
47 /// Caps and buffer sizes. Every consumer sets these explicitly; there is no
48 /// unbounded configuration.
49 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
50 pub struct HubLimits {
51 /// Concurrent listeners one user may hold. A few, for multiple tabs. Not a
52 /// spam control, a fairness one: without it a single client can consume the
53 /// global budget.
54 pub max_connections_per_user: usize,
55 /// Concurrent listeners across every room in the process. This is the number
56 /// that keeps the box alive.
57 pub max_connections_total: usize,
58 /// Messages buffered per room before a slow listener starts missing events
59 /// and gets a [`ChatEvent::Gap`].
60 pub room_buffer: usize,
61 }
62
63 impl Default for HubLimits {
64 fn default() -> Self {
65 Self {
66 max_connections_per_user: 4,
67 max_connections_total: 2_000,
68 room_buffer: 256,
69 }
70 }
71 }
72
73 struct Inner {
74 rooms: DashMap<RoomId, broadcast::Sender<ChatEvent>>,
75 per_user: DashMap<UserId, AtomicUsize>,
76 total: AtomicUsize,
77 limits: HubLimits,
78 }
79
80 /// Per-room fan-out with connection accounting.
81 ///
82 /// Cheap to clone; both consumers hang one off application state.
83 #[derive(Clone)]
84 pub struct Hub {
85 inner: Arc<Inner>,
86 }
87
88 impl Hub {
89 pub fn new(limits: HubLimits) -> Self {
90 Self {
91 inner: Arc::new(Inner {
92 rooms: DashMap::new(),
93 per_user: DashMap::new(),
94 total: AtomicUsize::new(0),
95 limits,
96 }),
97 }
98 }
99
100 /// Send to everyone listening to `room`, returning how many received it.
101 ///
102 /// Zero is normal and not an error: an empty room is the common case, and a
103 /// message still lands in storage regardless of who is connected.
104 pub fn publish(&self, room: RoomId, event: ChatEvent) -> usize {
105 self.inner
106 .rooms
107 .get(&room)
108 .and_then(|tx| tx.send(event).ok())
109 .unwrap_or(0)
110 }
111
112 /// Start listening to `room`.
113 ///
114 /// # Errors
115 ///
116 /// [`ChatError::ConnectionLimit`] when the user or the process is at capacity.
117 /// The caller has already passed authz by this point, so this is a resource
118 /// refusal and should read as one: it is a 503, not a 403.
119 pub fn subscribe(&self, room: RoomId, user: UserId) -> Result<Subscription, ChatError> {
120 // The permit is acquired first and owns the rollback. There is no window
121 // in which a counter is incremented but nothing is responsible for
122 // decrementing it, including if anything below panics.
123 let permit = ConnectionPermit::acquire(self.inner.clone(), user)?;
124
125 let rx = self
126 .inner
127 .rooms
128 .entry(room)
129 .or_insert_with(|| broadcast::channel(self.inner.limits.room_buffer).0)
130 .subscribe();
131
132 Ok(Subscription { rx, room, permit })
133 }
134
135 /// Rooms currently holding a channel. Diagnostics only.
136 pub fn room_count(&self) -> usize {
137 self.inner.rooms.len()
138 }
139
140 /// Listeners across every room. Worth exposing in a health payload: it is the
141 /// number that predicts the memory cap.
142 pub fn connection_count(&self) -> usize {
143 self.inner.total.load(Ordering::Acquire)
144 }
145
146 /// Listeners in one room.
147 pub fn room_connection_count(&self, room: RoomId) -> usize {
148 self.inner
149 .rooms
150 .get(&room)
151 .map_or(0, |tx| tx.receiver_count())
152 }
153 }
154
155 /// Holds one connection's slot in the global and per-user counters, releasing
156 /// both on drop.
157 struct ConnectionPermit {
158 inner: Arc<Inner>,
159 user: UserId,
160 }
161
162 impl ConnectionPermit {
163 fn acquire(inner: Arc<Inner>, user: UserId) -> Result<Self, ChatError> {
164 // Global first: it is a single atomic, so the cheap check rejects the
165 // overload case without touching the per-user map at all.
166 let prev = inner.total.fetch_add(1, Ordering::AcqRel);
167 if prev >= inner.limits.max_connections_total {
168 inner.total.fetch_sub(1, Ordering::AcqRel);
169 return Err(ChatError::ConnectionLimit(LimitScope::Global));
170 }
171
172 let over_user_limit = {
173 let entry = inner
174 .per_user
175 .entry(user)
176 .or_insert_with(|| AtomicUsize::new(0));
177 let prev = entry.value().fetch_add(1, Ordering::AcqRel);
178 if prev >= inner.limits.max_connections_per_user {
179 entry.value().fetch_sub(1, Ordering::AcqRel);
180 true
181 } else {
182 false
183 }
184 };
185
186 if over_user_limit {
187 inner.total.fetch_sub(1, Ordering::AcqRel);
188 // The entry may now be at zero and prunable, but leaving it costs one
189 // small map slot and the next successful subscribe reuses it. Do not
190 // remove here: this path is reached under contention, which is
191 // exactly when an extra shard lock is least welcome.
192 return Err(ChatError::ConnectionLimit(LimitScope::User));
193 }
194
195 Ok(Self { inner, user })
196 }
197 }
198
199 impl Drop for ConnectionPermit {
200 fn drop(&mut self) {
201 self.inner.total.fetch_sub(1, Ordering::AcqRel);
202
203 if let Some(entry) = self.inner.per_user.get(&self.user) {
204 entry.value().fetch_sub(1, Ordering::AcqRel);
205 }
206 // Predicate runs under the shard lock, so a user reconnecting in this
207 // instant either keeps their entry or gets a fresh one. Never both.
208 self.inner
209 .per_user
210 .remove_if(&self.user, |_, count| count.load(Ordering::Acquire) == 0);
211 }
212 }
213
214 /// A live listener on one room.
215 ///
216 /// Dropping it releases the connection slot and prunes the room channel if this
217 /// was the last listener.
218 pub struct Subscription {
219 rx: broadcast::Receiver<ChatEvent>,
220 room: RoomId,
221 permit: ConnectionPermit,
222 }
223
224 impl Subscription {
225 /// Next event, or `None` once the room is gone and nothing more will arrive.
226 ///
227 /// A listener that fell behind gets [`ChatEvent::Gap`] rather than silence.
228 /// The events it missed are not recoverable from the channel by design: the
229 /// client re-fetches from its cursor, which is the same path a reconnect
230 /// after a deploy takes.
231 pub async fn recv(&mut self) -> Option<ChatEvent> {
232 match self.rx.recv().await {
233 Ok(event) => Some(event),
234 Err(broadcast::error::RecvError::Lagged(_)) => Some(ChatEvent::Gap),
235 Err(broadcast::error::RecvError::Closed) => None,
236 }
237 }
238
239 pub fn room(&self) -> RoomId {
240 self.room
241 }
242 }
243
244 impl Drop for Subscription {
245 fn drop(&mut self) {
246 // `self.rx` is still alive here (fields drop after this body), so this
247 // subscription is still counted. Compare against 1, not 0.
248 //
249 // The predicate runs under the shard lock, so a subscriber arriving in
250 // this instant either keeps this channel or creates a fresh one after the
251 // removal. It cannot end up holding a receiver on a sender that has just
252 // been dropped from the map, which is the silent-dead-room failure.
253 self.permit
254 .inner
255 .rooms
256 .remove_if(&self.room, |_, tx| tx.receiver_count() <= 1);
257 }
258 }
259
260 #[cfg(test)]
261 mod tests {
262 use super::*;
263 use crate::ids::MessageId;
264 use uuid::Uuid;
265
266 fn hub(limits: HubLimits) -> Hub {
267 Hub::new(limits)
268 }
269
270 fn room() -> RoomId {
271 RoomId(Uuid::new_v4())
272 }
273
274 fn user() -> UserId {
275 UserId(Uuid::new_v4())
276 }
277
278 fn delete(id: i64) -> ChatEvent {
279 ChatEvent::Delete { id: MessageId(id) }
280 }
281
282 #[tokio::test]
283 async fn delivers_to_every_listener_in_the_room() {
284 let hub = hub(HubLimits::default());
285 let r = room();
286 let mut a = hub.subscribe(r, user()).unwrap();
287 let mut b = hub.subscribe(r, user()).unwrap();
288
289 assert_eq!(hub.publish(r, delete(1)), 2);
290 assert_eq!(a.recv().await, Some(delete(1)));
291 assert_eq!(b.recv().await, Some(delete(1)));
292 }
293
294 #[tokio::test]
295 async fn rooms_are_isolated() {
296 let hub = hub(HubLimits::default());
297 let (r1, r2) = (room(), room());
298 let mut listener = hub.subscribe(r1, user()).unwrap();
299
300 assert_eq!(hub.publish(r2, delete(1)), 0);
301 assert_eq!(hub.publish(r1, delete(2)), 1);
302 assert_eq!(listener.recv().await, Some(delete(2)));
303 }
304
305 #[test]
306 fn publishing_to_an_empty_room_is_not_an_error() {
307 let hub = hub(HubLimits::default());
308 assert_eq!(hub.publish(room(), delete(1)), 0);
309 // and does not conjure a channel for a room nobody is listening to
310 assert_eq!(hub.room_count(), 0);
311 }
312
313 #[test]
314 fn room_channel_is_pruned_when_the_last_listener_leaves() {
315 let hub = hub(HubLimits::default());
316 let r = room();
317
318 let a = hub.subscribe(r, user()).unwrap();
319 let b = hub.subscribe(r, user()).unwrap();
320 assert_eq!(hub.room_count(), 1);
321
322 drop(a);
323 assert_eq!(hub.room_count(), 1, "one listener remains");
324
325 drop(b);
326 assert_eq!(hub.room_count(), 0, "last listener pruned the room");
327 }
328
329 #[tokio::test]
330 async fn a_room_recreated_after_pruning_still_delivers() {
331 // The failure this guards: pruning a channel that a new subscriber has
332 // just taken, leaving them attached to an orphaned sender.
333 let hub = hub(HubLimits::default());
334 let r = room();
335
336 let first = hub.subscribe(r, user()).unwrap();
337 drop(first);
338 assert_eq!(hub.room_count(), 0);
339
340 let mut second = hub.subscribe(r, user()).unwrap();
341 assert_eq!(hub.publish(r, delete(9)), 1);
342 assert_eq!(second.recv().await, Some(delete(9)));
343 }
344
345 #[test]
346 fn per_user_limit_is_enforced_and_released() {
347 let hub = hub(HubLimits {
348 max_connections_per_user: 2,
349 ..HubLimits::default()
350 });
351 let r = room();
352 let u = user();
353
354 let a = hub.subscribe(r, u).unwrap();
355 let b = hub.subscribe(r, u).unwrap();
356 assert!(matches!(
357 hub.subscribe(r, u),
358 Err(ChatError::ConnectionLimit(LimitScope::User))
359 ));
360
361 // A different user is unaffected by another's exhausted budget.
362 let other = hub.subscribe(r, user()).unwrap();
363
364 drop(a);
365 let _c = hub.subscribe(r, u).expect("slot released on drop");
366 drop((b, other));
367 }
368
369 #[test]
370 fn global_limit_is_enforced_and_released() {
371 let hub = hub(HubLimits {
372 max_connections_total: 2,
373 ..HubLimits::default()
374 });
375 let r = room();
376
377 let a = hub.subscribe(r, user()).unwrap();
378 let b = hub.subscribe(r, user()).unwrap();
379 assert!(matches!(
380 hub.subscribe(r, user()),
381 Err(ChatError::ConnectionLimit(LimitScope::Global))
382 ));
383
384 drop(a);
385 assert_eq!(hub.connection_count(), 1);
386 let _c = hub.subscribe(r, user()).expect("slot released on drop");
387 drop(b);
388 }
389
390 #[test]
391 fn a_rejected_subscribe_leaks_no_slot() {
392 let hub = hub(HubLimits {
393 max_connections_per_user: 1,
394 ..HubLimits::default()
395 });
396 let r = room();
397 let u = user();
398
399 let held = hub.subscribe(r, u).unwrap();
400 for _ in 0..10 {
401 assert!(hub.subscribe(r, u).is_err());
402 }
403 // The global counter must not have drifted upward on the rejected path.
404 assert_eq!(hub.connection_count(), 1);
405 drop(held);
406 assert_eq!(hub.connection_count(), 0);
407 }
408
409 #[tokio::test]
410 async fn a_listener_that_falls_behind_gets_a_gap_not_silence() {
411 let hub = hub(HubLimits {
412 room_buffer: 2,
413 ..HubLimits::default()
414 });
415 let r = room();
416 let mut slow = hub.subscribe(r, user()).unwrap();
417
418 for i in 0..10 {
419 hub.publish(r, delete(i));
420 }
421
422 assert_eq!(slow.recv().await, Some(ChatEvent::Gap));
423 // and the stream keeps working afterward
424 hub.publish(r, delete(99));
425 let next = slow.recv().await;
426 assert!(next.is_some(), "stream survives a gap");
427 }
428
429 #[test]
430 fn counters_return_to_zero_when_everything_drops() {
431 let hub = hub(HubLimits::default());
432 let r = room();
433 let subs: Vec<_> = (0..5).map(|_| hub.subscribe(r, user()).unwrap()).collect();
434 assert_eq!(hub.connection_count(), 5);
435 assert_eq!(hub.room_connection_count(r), 5);
436
437 drop(subs);
438 assert_eq!(hub.connection_count(), 0);
439 assert_eq!(hub.room_count(), 0);
440 assert_eq!(hub.room_connection_count(r), 0);
441 }
442 }
443