Skip to main content

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}