moirai_async/sync/

broadcast.rs

//! Broadcast channel for one-to-many async communication
//!
//! Provides broadcast channel implementation that allows one sender to
//! broadcast messages to multiple receivers with SLAP-compliant design.

#![expect(
    clippy::unwrap_used,
    reason = "ratchet MOIRAI-UNWRAP-1: pre-existing debt"
)]

use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

use super::subscribers::{SubscriberRegistry, wake_drained};

/// Broadcast channel for one-to-many communication
pub struct Broadcast<T> {
    _phantom: std::marker::PhantomData<T>,
}

struct BroadcastState<T> {
    messages: VecDeque<(u64, T)>,
    sequence: u64,
    closed: bool,
    /// One slot per receiver, shared with `Watch` (see `subscribers`). The
    /// cursor is the last sequence that receiver consumed, which doubles as the
    /// input to the retention boundary in `send`.
    subscribers: SubscriberRegistry<u64>,
    capacity: usize,
}

impl<T: Clone + Send + 'static> Broadcast<T> {
    /// Create a new broadcast channel with the given capacity
    /// Returns (sender, receiver) tuple per channel pattern conventions
    #[allow(clippy::new_ret_no_self)] // Standard channel pattern per Rust Book Ch.16
    pub fn new(capacity: usize) -> (BroadcastSender<T>, BroadcastReceiver<T>) {
        let state = Arc::new(Mutex::new(BroadcastState {
            messages: VecDeque::new(),
            sequence: 0,
            closed: false,
            subscribers: SubscriberRegistry::with_initial(0),
            capacity,
        }));

        let sender = BroadcastSender {
            state: state.clone(),
        };

        let receiver = BroadcastReceiver {
            state: state.clone(),
            id: 0,
            position: 0,
        };

        (sender, receiver)
    }
}

/// Sender half of broadcast channel
pub struct BroadcastSender<T> {
    state: Arc<Mutex<BroadcastState<T>>>,
}

impl<T: Clone> BroadcastSender<T> {
    /// Send a message to all receivers
    pub fn send(&self, message: T) -> Result<usize, BroadcastError> {
        let (receiver_count, wakers) = {
            let mut state = self.state.lock().unwrap();
            if state.closed {
                return Err(BroadcastError::Closed);
            }

            // Add new message first, then reclaim messages already read by every
            // receiver.  The retention boundary is the minimum position across
            // all live receivers; messages with sequence <= that boundary have
            // been consumed by everyone and can be dropped.
            state.sequence += 1;
            let sequence = state.sequence;
            state.messages.push_back((sequence, message));

            // Publication is done; take the registered wakers. They observe
            // either the appended message or the explicit lag state retained by
            // the capacity contract. Woken only after the lock is released —
            // see `drain_wakers`.
            let receiver_count = state.subscribers.len();
            let wakers = state.subscribers.drain_wakers();

            // `min` over the cursors, copied out before the mutation below so
            // the immutable borrow of the registry ends here.
            let min_position = state
                .subscribers
                .cursors()
                .copied()
                .min()
                .unwrap_or(sequence);
            while state.messages.len() > state.capacity
                || state
                    .messages
                    .front()
                    .is_some_and(|(seq, _)| *seq <= min_position)
            {
                state.messages.pop_front();
            }

            (receiver_count, wakers)
        };
        wake_drained(wakers);

        Ok(receiver_count)
    }

    /// Get the number of active receivers
    pub fn receiver_count(&self) -> usize {
        self.state.lock().unwrap().subscribers.len()
    }
}

impl<T> Drop for BroadcastSender<T> {
    fn drop(&mut self) {
        let wakers = {
            let mut state = self.state.lock().unwrap();
            state.closed = true;
            state.subscribers.drain_wakers()
        };
        wake_drained(wakers);
    }
}

/// Receiver half of broadcast channel
pub struct BroadcastReceiver<T> {
    state: Arc<Mutex<BroadcastState<T>>>,
    id: u64,
    position: u64,
}

impl<T: Clone> BroadcastReceiver<T> {
    /// Receive the next message
    pub fn recv(&mut self) -> BroadcastRecv<'_, T> {
        BroadcastRecv { receiver: self }
    }

    /// Try to receive a message immediately
    pub fn try_recv(&mut self) -> Result<T, BroadcastError> {
        let mut state = self.state.lock().unwrap();

        if state.messages.is_empty() {
            if state.closed {
                return Err(BroadcastError::Closed);
            }
            return Err(BroadcastError::Empty);
        }

        // Lagging check
        let oldest_seq = state.messages.front().unwrap().0;
        if self.position + 1 < oldest_seq {
            self.position = oldest_seq - 1;
            if let Some(subscriber) = state.subscribers.get_mut(self.id) {
                subscriber.cursor = self.position;
            }
            return Err(BroadcastError::Lagged);
        }

        // Sequences are dense (each `send` appends `sequence + 1` and only
        // `pop_front` removes), so the next unread message sits at a directly
        // computable offset from the front — no scan. After the lag check,
        // `position + 1 >= oldest_seq`, and `position <= sequence` bounds the
        // offset by the queue length, so the conversion cannot truncate.
        let offset = usize::try_from(self.position + 1 - oldest_seq)
            .expect("invariant: unread offset is bounded by the message queue length");
        let found = state.messages.get(offset).map(|(seq, message)| {
            debug_assert_eq!(*seq, self.position + 1, "broadcast sequences must be dense");
            self.position = *seq;
            message.clone()
        });
        if let Some(message) = found {
            if let Some(subscriber) = state.subscribers.get_mut(self.id) {
                subscriber.cursor = self.position;
            }
            return Ok(message);
        }

        if state.closed {
            Err(BroadcastError::Closed)
        } else {
            Err(BroadcastError::Empty)
        }
    }

    /// Clone this receiver to create a new independent receiver
    pub fn resubscribe(&self) -> BroadcastReceiver<T> {
        let mut state = self.state.lock().unwrap();
        let current_sequence = state.sequence;
        let id = state.subscribers.register(current_sequence);

        BroadcastReceiver {
            state: self.state.clone(),
            id,
            position: current_sequence,
        }
    }

    /// Poll to receive a message, registering waker if empty.
    pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, BroadcastError>> {
        let mut state = self.state.lock().unwrap();

        if state.messages.is_empty() {
            if state.closed {
                return Poll::Ready(Err(BroadcastError::Closed));
            }
            if let Some(subscriber) = state.subscribers.get_mut(self.id) {
                subscriber.waker = Some(cx.waker().clone());
            }
            return Poll::Pending;
        }

        // Lagging check
        let oldest_seq = state.messages.front().unwrap().0;
        if self.position + 1 < oldest_seq {
            self.position = oldest_seq - 1;
            if let Some(subscriber) = state.subscribers.get_mut(self.id) {
                subscriber.cursor = self.position;
            }
            return Poll::Ready(Err(BroadcastError::Lagged));
        }

        // Dense-sequence direct index; see `try_recv` for the derivation.
        let offset = usize::try_from(self.position + 1 - oldest_seq)
            .expect("invariant: unread offset is bounded by the message queue length");
        let found_msg = state.messages.get(offset).map(|(seq, message)| {
            debug_assert_eq!(*seq, self.position + 1, "broadcast sequences must be dense");
            self.position = *seq;
            (*seq, message.clone())
        });

        if let Some((_, message)) = found_msg {
            if let Some(subscriber) = state.subscribers.get_mut(self.id) {
                subscriber.cursor = self.position;
            }
            Poll::Ready(Ok(message))
        } else if state.closed {
            Poll::Ready(Err(BroadcastError::Closed))
        } else {
            if let Some(subscriber) = state.subscribers.get_mut(self.id) {
                subscriber.waker = Some(cx.waker().clone());
            }
            Poll::Pending
        }
    }
}

impl<T: Clone> Clone for BroadcastReceiver<T> {
    fn clone(&self) -> Self {
        self.resubscribe()
    }
}

impl<T> Drop for BroadcastReceiver<T> {
    fn drop(&mut self) {
        if let Ok(mut state) = self.state.lock() {
            state.subscribers.remove(self.id);
        }
    }
}

/// Future for receiving from broadcast channel
pub struct BroadcastRecv<'a, T> {
    receiver: &'a mut BroadcastReceiver<T>,
}

impl<'a, T: Clone> Future for BroadcastRecv<'a, T> {
    type Output = Result<T, BroadcastError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.receiver.poll_recv(cx)
    }
}

impl<'a, T> Drop for BroadcastRecv<'a, T> {
    fn drop(&mut self) {
        // Clear the waker registered by `poll_recv` so a cancelled recv future
        // does not leave a stale waker that the next `send` would spuriously wake
        // (and retain until overwritten). This mirrors `WatchChanged::drop`. The
        // future holds `&mut BroadcastReceiver`, so it is the unique registrant
        // for this receiver id — clearing here cannot drop another future's waker.
        if let Ok(mut state) = self.receiver.state.lock() {
            let id = self.receiver.id;
            if let Some(subscriber) = state.subscribers.get_mut(id) {
                subscriber.waker = None;
            }
        }
    }
}

/// Error types for broadcast channel operations
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BroadcastError {
    /// Channel is empty
    Empty,
    /// Channel has been closed
    Closed,
    /// Message was lost due to channel overflow
    Lagged,
}

impl std::fmt::Display for BroadcastError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BroadcastError::Empty => write!(f, "broadcast channel is empty"),
            BroadcastError::Closed => write!(f, "broadcast channel is closed"),
            BroadcastError::Lagged => write!(f, "broadcast channel lagged"),
        }
    }
}

impl std::error::Error for BroadcastError {}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::task::{Wake, Waker};

    struct CountingWake(Arc<AtomicUsize>);
    impl Wake for CountingWake {
        fn wake(self: Arc<Self>) {
            self.0.fetch_add(1, Ordering::Release);
        }
        fn wake_by_ref(self: &Arc<Self>) {
            self.0.fetch_add(1, Ordering::Release);
        }
    }

    #[test]
    fn cancelled_recv_clears_waker_and_is_not_spuriously_woken() {
        let (tx, mut rx) = Broadcast::<u32>::new(8);
        let count = Arc::new(AtomicUsize::new(0));
        let waker = Waker::from(Arc::new(CountingWake(Arc::clone(&count))));
        let mut cx = Context::from_waker(&waker);

        {
            let mut fut = rx.recv();
            assert!(Pin::new(&mut fut).poll(&mut cx).is_pending());
            // `fut` is dropped here; its Drop must clear the registered waker.
        }

        // The cancelled future's waker must not fire on the next send.
        tx.send(42).expect("send must succeed");
        assert_eq!(
            count.load(Ordering::Acquire),
            0,
            "a cancelled recv future must not be spuriously woken"
        );

        // The message is still deliverable to a fresh recv on the same receiver.
        assert_eq!(rx.try_recv(), Ok(42));
    }

    #[test]
    fn live_recv_is_woken_on_send() {
        // Control: a still-live registered waker IS woken by send.
        let (tx, mut rx) = Broadcast::<u32>::new(8);
        let count = Arc::new(AtomicUsize::new(0));
        let waker = Waker::from(Arc::new(CountingWake(Arc::clone(&count))));
        let mut cx = Context::from_waker(&waker);

        let mut fut = rx.recv();
        assert!(Pin::new(&mut fut).poll(&mut cx).is_pending());
        tx.send(7).expect("send must succeed");
        assert_eq!(
            count.load(Ordering::Acquire),
            1,
            "a live recv future must be woken by send"
        );
        // Keep `fut` alive across the send so its waker stays registered.
        drop(fut);
    }

    #[test]
    fn messages_read_by_every_receiver_are_reclaimed_on_next_send() {
        let (tx, mut first) = Broadcast::<u32>::new(8);
        let mut second = first.resubscribe();

        tx.send(10).expect("first send must succeed");
        tx.send(20).expect("second send must succeed");
        assert_eq!(first.try_recv(), Ok(10));
        assert_eq!(second.try_recv(), Ok(10));
        assert_eq!(first.try_recv(), Ok(20));
        assert_eq!(second.try_recv(), Ok(20));

        tx.send(30).expect("third send must succeed");
        let state = tx.state.lock().expect("broadcast state must not poison");
        assert_eq!(
            state.messages.iter().copied().collect::<Vec<_>>(),
            vec![(3, 30)],
            "the next send must retain only the new unread message"
        );
    }
}