forest/utils/publisher.rs
1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use parking_lot::Mutex;
5use std::sync::Arc;
6
7/// A non-blocking fan-out publisher.
8///
9/// Each subscriber gets its own [`flume`] queue and cloning shares the subscriber
10/// registry. [`Self::subscribe`] gives a subscriber an unbounded, lossless queue;
11/// [`Self::subscribe_bounded`] gives a bounded queue that drops new events for that
12/// subscriber alone once it is full (use it for best-effort consumers that must not be
13/// able to grow memory without bound, e.g. ones fed by untrusted clients). Either way the
14/// producer never blocks and a slow subscriber never stalls the others.
15pub struct Publisher<T>(Arc<Mutex<Vec<flume::Sender<T>>>>);
16
17impl<T> Clone for Publisher<T> {
18 fn clone(&self) -> Self {
19 Self(self.0.clone())
20 }
21}
22
23impl<T> Default for Publisher<T> {
24 fn default() -> Self {
25 Self(Arc::new(Mutex::new(Vec::new())))
26 }
27}
28
29impl<T: Clone> Publisher<T> {
30 /// Registers a new subscriber with an unbounded, lossless queue and returns its receiver.
31 pub fn subscribe(&self) -> flume::Receiver<T> {
32 let (tx, rx) = flume::unbounded();
33 self.0.lock().push(tx);
34 rx
35 }
36
37 /// Registers a new subscriber with a bounded queue of capacity `cap`. When the subscriber
38 /// falls `cap` events behind, the newest events are dropped for it alone (it keeps the
39 /// oldest `cap`; the producer and other subscribers are unaffected).
40 pub fn subscribe_bounded(&self, cap: usize) -> flume::Receiver<T> {
41 let (tx, rx) = flume::bounded(cap);
42 self.0.lock().push(tx);
43 rx
44 }
45
46 /// Delivers `msg` to every subscriber. Never blocks: for a bounded subscriber that is
47 /// full the event is dropped for that subscriber; a subscriber whose receiver is gone
48 /// is pruned.
49 pub fn publish(&self, msg: T) {
50 self.0.lock().retain(|tx| {
51 !matches!(
52 tx.try_send(msg.clone()),
53 Err(flume::TrySendError::Disconnected(_))
54 )
55 });
56 }
57
58 /// Cheap check for whether any subscriber is registered. Does not prune, so it may
59 /// briefly report `true` after the last receiver is gone (until the next [`Self::publish`]
60 /// prunes it), but never reports `false` while a live subscriber exists.
61 pub fn has_subscribers(&self) -> bool {
62 !self.0.lock().is_empty()
63 }
64}
65
66#[cfg(test)]
67mod tests {
68 use super::*;
69 use itertools::Itertools as _;
70
71 #[test]
72 fn publisher_is_lossless_under_lag() {
73 let publisher = Publisher::default();
74 let rx1 = publisher.subscribe();
75 let rx2 = publisher.subscribe();
76
77 // Far more than any bounded channel would hold; nothing is drained meanwhile.
78 const N: u32 = 10_000;
79 for i in 0..N {
80 publisher.publish(i);
81 }
82
83 for rx in [&rx1, &rx2] {
84 for expected in 0..N {
85 assert_eq!(rx.recv().unwrap(), expected);
86 }
87 assert!(rx.try_recv().is_err());
88 }
89 }
90
91 #[test]
92 fn publisher_prunes_dropped_subscribers() {
93 let publisher = Publisher::<u32>::default();
94 let rx_live = publisher.subscribe();
95 let rx_dead = publisher.subscribe();
96 assert!(publisher.has_subscribers());
97
98 drop(rx_dead);
99 // Publishing prunes the dead sender while still delivering to the live one.
100 publisher.publish(7);
101 assert!(publisher.has_subscribers());
102 assert_eq!(rx_live.recv().unwrap(), 7);
103
104 drop(rx_live);
105 publisher.publish(8);
106 assert!(!publisher.has_subscribers());
107 }
108
109 #[test]
110 fn publisher_bounded_subscriber_drops_without_blocking_others() {
111 let publisher = Publisher::default();
112 let unbounded = publisher.subscribe();
113 let bounded = publisher.subscribe_bounded(2);
114
115 // Publishing well past the bound must not block and must not affect the unbounded sub.
116 for i in 0..10 {
117 publisher.publish(i);
118 }
119
120 // Bounded subscriber kept only up to its capacity; the excess was dropped for it alone.
121 let bounded_items = bounded.try_iter().collect_vec();
122 assert_eq!(bounded_items, vec![0, 1]);
123
124 // Unbounded subscriber still received everything, in order.
125 let unbounded_items = unbounded.try_iter().collect_vec();
126 assert_eq!(unbounded_items, (0..10).collect_vec());
127 }
128}