rskit_stream/

broadcaster.rs

//! Bounded fan-out broadcaster source.
//!
//! [`Broadcaster<T>`] turns "observe a backend" into a *bounded,
//! cancellable stream of typed events* fanned out to many independent subscribers.
//! Each subscriber owns a private bounded channel:
//! a subscriber that falls further behind than the buffer loses interim events (backpressure by drop)
//! but never blocks the broadcaster or its peers.
//! Disconnected subscribers are pruned lazily on the next broadcast.
//!
//! This is the canonical owner for the "watch a source → typed change stream" shape that recurs across config reloads,
//! service discovery, cache invalidation, and secret rotation.
//! Consumers that must not miss events should treat any received item as a signal to re-run their full load pipeline.

use std::fmt;
use std::pin::Pin;
use std::sync::Arc;

use futures::Stream;
use parking_lot::Mutex;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;

/// A bounded, owned, cancellable stream of broadcast events.
///
/// Yielded by [`Broadcaster::subscribe`]; terminates when the originating `CancellationToken` fires
/// or the broadcaster is dropped.
pub type BroadcastStream<T> = Pin<Box<dyn Stream<Item = T> + Send + 'static>>;

/// Default per-subscriber buffer used by [`Broadcaster::new`].
///
/// A subscriber lagging more than this many unconsumed events drops the overflow rather than stalling the broadcaster.
pub const DEFAULT_BROADCAST_BUFFER: usize = 64;

/// A bounded, cancellable fan-out broadcaster.
///
/// Cheaply cloneable; every clone shares the same subscriber set,
/// so an event broadcast through any clone reaches all live subscribers.
/// `T` must be `Clone` (each subscriber receives its own copy)
/// and `Send + 'static` to cross the per-subscriber channel.
pub struct Broadcaster<T> {
    senders: Arc<Mutex<Vec<mpsc::Sender<T>>>>,
    buffer: usize,
}

impl<T> Clone for Broadcaster<T> {
    fn clone(&self) -> Self {
        Self {
            senders: Arc::clone(&self.senders),
            buffer: self.buffer,
        }
    }
}

impl<T> fmt::Debug for Broadcaster<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Broadcaster")
            .field("buffer", &self.buffer)
            .field("subscribers", &self.subscriber_count())
            .finish_non_exhaustive()
    }
}

impl<T> Default for Broadcaster<T>
where
    T: Clone + Send + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Broadcaster<T>
where
    T: Clone + Send + 'static,
{
    /// Create a broadcaster with the [`DEFAULT_BROADCAST_BUFFER`] per-subscriber buffer.
    #[must_use]
    pub fn new() -> Self {
        Self::with_buffer(DEFAULT_BROADCAST_BUFFER)
    }

    /// Create a broadcaster with an explicit per-subscriber buffer.
    ///
    /// A `buffer` of zero is clamped to one so every subscriber can hold at least one in-flight event.
    #[must_use]
    pub fn with_buffer(buffer: usize) -> Self {
        Self {
            senders: Arc::new(Mutex::new(Vec::new())),
            buffer: buffer.max(1),
        }
    }
}

impl<T> Broadcaster<T> {
    /// The per-subscriber buffer capacity.
    #[must_use]
    pub const fn buffer(&self) -> usize {
        self.buffer
    }

    /// The number of currently registered subscribers.
    ///
    /// Disconnected subscribers are pruned lazily on [`broadcast`](Self::broadcast),
    /// so this count may transiently include subscribers whose streams have been dropped but not
    /// yet pruned.
    #[must_use]
    pub fn subscriber_count(&self) -> usize {
        self.senders.lock().len()
    }

    /// Register a new subscriber and return its bounded, cancellable stream.
    ///
    /// The returned stream yields events broadcast after subscription and terminates once `cancel` fires
    /// or the broadcaster is dropped.
    /// Pass a fresh `CancellationToken` when no external cancellation is required.
    pub fn subscribe(&self, cancel: CancellationToken) -> BroadcastStream<T>
    where
        T: Send + 'static,
    {
        use futures::StreamExt as _;

        let (tx, rx) = mpsc::channel(self.buffer);
        self.senders.lock().push(tx);
        Box::pin(ReceiverStream::new(rx).take_until(cancel.cancelled_owned()))
    }

    /// Fan `item` out to every live subscriber.
    ///
    /// Each subscriber receives an independent clone. A subscriber whose buffer is full is retained
    /// but skips this event (bounded drop backpressure);
    /// a subscriber whose stream has been dropped is pruned.
    pub fn broadcast(&self, item: &T)
    where
        T: Clone,
    {
        self.senders
            .lock()
            .retain(|tx| match tx.try_send(item.clone()) {
                Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => true,
                Err(mpsc::error::TrySendError::Closed(_)) => false,
            });
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::StreamExt as _;

    #[tokio::test]
    async fn delivers_to_all_subscribers() {
        let bc = Broadcaster::<u32>::new();
        let cancel = CancellationToken::new();
        let mut a = bc.subscribe(cancel.clone());
        let mut b = bc.subscribe(cancel.clone());

        bc.broadcast(&7);

        assert_eq!(a.next().await, Some(7));
        assert_eq!(b.next().await, Some(7));
    }

    #[tokio::test]
    async fn stream_terminates_on_cancel() {
        let bc = Broadcaster::<u32>::new();
        let cancel = CancellationToken::new();
        let mut sub = bc.subscribe(cancel.clone());

        bc.broadcast(&1);
        assert_eq!(sub.next().await, Some(1));

        cancel.cancel();
        assert_eq!(sub.next().await, None);
    }

    #[tokio::test]
    async fn dropped_subscriber_is_pruned() {
        let bc = Broadcaster::<u32>::new();
        let cancel = CancellationToken::new();
        let sub = bc.subscribe(cancel);
        assert_eq!(bc.subscriber_count(), 1);

        drop(sub);
        // Prune happens lazily on the next broadcast.
        bc.broadcast(&1);
        assert_eq!(bc.subscriber_count(), 0);
    }

    #[tokio::test]
    async fn full_subscriber_drops_overflow_without_blocking() {
        let bc = Broadcaster::<u32>::with_buffer(2);
        let cancel = CancellationToken::new();
        let mut sub = bc.subscribe(cancel.clone());

        // Send more than the buffer; overflow is dropped, not blocked.
        for i in 0..5 {
            bc.broadcast(&i);
        }

        let mut received = Vec::new();
        // Only the first `buffer` events are retained.
        while let Ok(Some(v)) =
            tokio::time::timeout(std::time::Duration::from_millis(50), sub.next()).await
        {
            received.push(v);
        }
        assert_eq!(received, vec![0, 1]);
    }

    #[tokio::test]
    async fn zero_buffer_is_clamped_to_one() {
        let bc = Broadcaster::<u32>::with_buffer(0);
        assert_eq!(bc.buffer(), 1);
    }

    #[tokio::test]
    async fn clones_share_subscriber_set() {
        let bc = Broadcaster::<u32>::new();
        let cancel = CancellationToken::new();
        let mut sub = bc.subscribe(cancel.clone());

        let clone = bc.clone();
        clone.broadcast(&42);

        assert_eq!(sub.next().await, Some(42));
    }

    #[test]
    fn default_and_debug_report_buffer_and_subscribers() {
        let bc = Broadcaster::<u32>::default();

        assert_eq!(bc.buffer(), DEFAULT_BROADCAST_BUFFER);
        assert!(format!("{bc:?}").contains("Broadcaster"));
    }
}