tempest_rt/sync/

mpsc.rs

//! Bounded multi-producer, single-consumer channel.
//!
//! Create a channel with [`bounded`]. The [`BoundedSender`] can be cloned for multiple producers;
//! there is exactly one [`Receiver`]. Sending blocks when the channel is full, providing
//! natural backpressure.

use std::{
    cell::RefCell,
    collections::VecDeque,
    future::poll_fn,
    pin::Pin,
    rc::Rc,
    task::{Context, Poll, Waker},
};

use derive_more::{Display, Error};
use futures::Stream;

struct Inner<T> {
    queue: VecDeque<T>,
    /// Is the consumer still alive?
    rx_alive: bool,
    /// Stores the waker of the waiting consumer that is waiting for produced values.
    rx_waker: Option<Waker>,
    /// Stores any wakers of consumers that are waiting on the channel to make space.
    tx_wakers: Vec<Waker>,
}

/// Sending half of a bounded channel. Cloneable for multiple producers.
pub struct BoundedSender<T> {
    inner: Rc<RefCell<Inner<T>>>,
}

impl<T> Clone for BoundedSender<T> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

/// Receiving half of a bounded channel.
pub struct Receiver<T> {
    inner: Rc<RefCell<Inner<T>>>,
}

/// Creates a bounded channel with the given capacity.
///
/// Panics if `cap` is 0.
pub fn bounded<T>(cap: usize) -> (BoundedSender<T>, Receiver<T>) {
    assert_ne!(cap, 0, "a bounded channel with capacity 0 does not work");
    let mut queue = VecDeque::new();
    queue.reserve_exact(cap);
    let inner = Rc::new(RefCell::new(Inner {
        queue,
        rx_alive: true,
        rx_waker: None,
        tx_wakers: Vec::new(),
    }));
    let tx = BoundedSender {
        inner: inner.clone(),
    };
    let rx = Receiver { inner };
    (tx, rx)
}

/// Error returned by [`BoundedSender::send`] when the receiver has been dropped.
#[derive(Debug, Display, Error)]
#[display("receiver has been dropped")]
pub struct SendError<T>(pub T);

/// Error returned by [`BoundedSender::try_send`].
#[derive(Debug, Display, Error)]
pub enum TrySendError<T> {
    /// The channel is at capacity.
    #[display("channel is full")]
    Full(T),
    /// The receiver has been dropped.
    #[display("receiver has been dropped")]
    Closed(T),
}

impl<T> BoundedSender<T> {
    /// Sends `val`, parking until space is available.
    ///
    /// Returns `Err` if the receiver has been dropped.
    pub async fn send(&mut self, val: T) -> Result<(), SendError<T>> {
        let mut val = Some(val);
        poll_fn(|cx| {
            if let Some(waker) = self.inner.borrow_mut().rx_waker.take() {
                waker.wake();
            }
            match self.try_send(val.take().unwrap()) {
                Ok(()) => Poll::Ready(Ok(())),
                Err(TrySendError::Full(v)) => {
                    val = Some(v);
                    self.inner.borrow_mut().tx_wakers.push(cx.waker().clone());
                    Poll::Pending
                }
                Err(TrySendError::Closed(v)) => Poll::Ready(Err(SendError(v))),
            }
        })
        .await
    }

    /// Sends `val` without waiting.
    ///
    /// Returns `Err(Full)` if the channel is at capacity, or `Err(Closed)` if the receiver
    /// has been dropped.
    pub fn try_send(&mut self, val: T) -> Result<(), TrySendError<T>> {
        let mut inner = self.inner.borrow_mut();
        if !inner.rx_alive {
            Err(TrySendError::Closed(val))
        } else if inner.queue.len() == inner.queue.capacity() {
            Err(TrySendError::Full(val))
        } else {
            inner.queue.push_back(val);
            Ok(())
        }
    }
}

/// Error returned by [`Receiver::recv`] when all senders have been dropped.
#[derive(Debug, Display, Error, PartialEq, Eq)]
#[display("all senders have been dropped")]
pub struct RecvError;

/// Error returned by [`Receiver::try_recv`].
#[derive(Debug, Display, Error, PartialEq, Eq)]
pub enum TryRecvError {
    /// The queue is empty but senders are still alive.
    #[display("channel is empty")]
    Empty,
    /// All senders have been dropped and the queue is empty.
    #[display("all senders have been dropped")]
    Closed,
}

impl<T> Receiver<T> {
    pub(crate) fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
        match self.try_recv() {
            Ok(val) => {
                for waker in self.inner.borrow_mut().tx_wakers.drain(..) {
                    waker.wake();
                }
                Poll::Ready(Ok(val))
            }
            Err(TryRecvError::Empty) => {
                self.inner.borrow_mut().rx_waker = Some(cx.waker().clone());
                Poll::Pending
            }
            Err(TryRecvError::Closed) => Poll::Ready(Err(RecvError)),
        }
    }

    /// Receives the next value, parking until one arrives.
    ///
    /// Returns `Err` if all senders have been dropped and the queue is empty.
    pub async fn recv(&mut self) -> Result<T, RecvError> {
        poll_fn(|cx| self.poll_recv(cx)).await
    }

    /// Receives without waiting.
    ///
    /// Returns `Err(Empty)` if the queue is empty, or `Err(Closed)` if all senders have been
    /// dropped and the queue is empty.
    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
        if let Some(val) = self.inner.borrow_mut().queue.pop_front() {
            return Ok(val);
        }
        if Rc::strong_count(&self.inner) == 1 {
            Err(TryRecvError::Closed)
        } else {
            Err(TryRecvError::Empty)
        }
    }
}

impl<T> Stream for Receiver<T> {
    type Item = T;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.poll_recv(cx).map(|r| r.ok())
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.inner.borrow().queue.len(), None)
    }
}

impl<T> Drop for BoundedSender<T> {
    fn drop(&mut self) {
        // When the last sender is dropped (strong_count == 2: this sender + receiver),
        // wake the receiver so it can observe the Closed state.
        if Rc::strong_count(&self.inner) == 2 {
            if let Some(waker) = self.inner.borrow_mut().rx_waker.take() {
                waker.wake();
            }
        }
    }
}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        self.inner.borrow_mut().rx_alive = false;
    }
}

#[cfg(test)]
mod tests {
    use std::{
        pin::{Pin, pin},
        task::Poll,
    };

    use futures::StreamExt;
    use tempest_io::VirtualIo;

    use crate::{block_on, spawn};

    use super::*;

    // -- try_send --

    #[test]
    fn test_try_send_one() {
        block_on(VirtualIo::default(), async {
            let (mut tx, _rx) = bounded(1);
            assert!(tx.try_send(42).is_ok());
        });
    }

    #[test]
    fn test_try_send_exactly_full() {
        block_on(VirtualIo::default(), async {
            let (mut tx, _rx) = bounded(2);
            assert!(tx.try_send(1).is_ok());
            assert!(tx.try_send(2).is_ok());
        });
    }

    #[test]
    fn test_try_send_over_full() {
        block_on(VirtualIo::default(), async {
            let (mut tx, _rx) = bounded(1);
            tx.try_send(1).unwrap();
            match tx.try_send(99) {
                Err(TrySendError::Full(v)) => assert_eq!(v, 99),
                _ => panic!("expected Full"),
            }
        });
    }

    #[test]
    fn test_try_send_closed() {
        block_on(VirtualIo::default(), async {
            let (mut tx, rx) = bounded::<i32>(1);
            drop(rx);
            match tx.try_send(99) {
                Err(TrySendError::Closed(v)) => assert_eq!(v, 99),
                _ => panic!("expected Closed"),
            }
        });
    }

    // -- send --

    #[test]
    fn test_send_one() {
        block_on(VirtualIo::default(), async {
            let (mut tx, _rx) = bounded(1);
            tx.send(42).await.unwrap();
        });
    }

    #[test]
    fn test_send_exactly_full() {
        block_on(VirtualIo::default(), async {
            let (mut tx, _rx) = bounded(2);
            tx.send(1).await.unwrap();
            tx.send(2).await.unwrap();
        });
    }

    #[test]
    fn test_send_pending_when_full() {
        block_on(VirtualIo::default(), async {
            let (mut tx, _rx) = bounded(1);
            tx.try_send(1).unwrap();

            let waker = std::task::Waker::noop();
            let mut cx = std::task::Context::from_waker(&waker);
            let mut fut = pin!(tx.send(2));
            assert!(matches!(fut.as_mut().poll(&mut cx), Poll::Pending));
        });
    }

    #[test]
    fn test_send_when_full_eventually_resolves() {
        block_on(VirtualIo::default(), async {
            let (mut tx, mut rx) = bounded(1);
            tx.try_send(1).unwrap();
            // this will block, until the rx frees up space
            let mut handle = spawn(async move { tx.send(2).await.unwrap() });

            // the task handle should be pending, until we make space in the rx
            let waker = std::task::Waker::noop();
            let mut cx = std::task::Context::from_waker(&waker);
            assert!(matches!(Pin::new(&mut handle).poll(&mut cx), Poll::Pending));

            assert_eq!(rx.recv().await, Ok(1));
            assert_eq!(rx.recv().await, Ok(2));
            // channel now empty and receiver dropped
            assert!(rx.try_recv().is_err())
        });
    }

    #[test]
    fn test_recv_when_empty_eventually_resolves() {
        block_on(VirtualIo::default(), async {
            let (mut tx, mut rx) = bounded(1);
            spawn(async move {
                assert_eq!(rx.recv().await, Ok(1));
                assert_eq!(rx.recv().await, Ok(2));
            });

            tx.send(1).await.unwrap();
            assert!(matches!(tx.try_send(2), Err(TrySendError::Full(2))));
            tx.send(2).await.unwrap();
        });
    }

    #[test]
    fn test_send_closed() {
        block_on(VirtualIo::default(), async {
            let (mut tx, rx) = bounded::<i32>(1);
            drop(rx);
            match tx.send(99).await {
                Err(SendError(v)) => assert_eq!(v, 99),
                Ok(()) => panic!("expected Err"),
            }
        });
    }

    // -- try_recv --

    #[test]
    fn test_try_recv_one() {
        block_on(VirtualIo::default(), async {
            let (mut tx, mut rx) = bounded(1);
            tx.try_send(42).unwrap();
            assert_eq!(rx.try_recv().unwrap(), 42);
        });
    }

    #[test]
    fn test_try_recv_in_order() {
        block_on(VirtualIo::default(), async {
            let (mut tx, mut rx) = bounded(3);
            tx.try_send(1).unwrap();
            tx.try_send(2).unwrap();
            tx.try_send(3).unwrap();
            assert_eq!(rx.try_recv().unwrap(), 1);
            assert_eq!(rx.try_recv().unwrap(), 2);
            assert_eq!(rx.try_recv().unwrap(), 3);
        });
    }

    #[test]
    fn test_try_recv_empty() {
        block_on(VirtualIo::default(), async {
            let (_tx, mut rx) = bounded::<i32>(1);
            assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
        });
    }

    #[test]
    fn test_try_recv_closed() {
        block_on(VirtualIo::default(), async {
            let (tx, mut rx) = bounded::<i32>(1);
            drop(tx);
            assert_eq!(rx.try_recv(), Err(TryRecvError::Closed));
        });
    }

    // -- recv --

    #[test]
    fn test_recv_one() {
        block_on(VirtualIo::default(), async {
            let (mut tx, mut rx) = bounded(1);
            tx.send(42).await.unwrap();
            assert_eq!(rx.recv().await.unwrap(), 42);
        });
    }

    #[test]
    fn test_recv_in_order() {
        block_on(VirtualIo::default(), async {
            let (mut tx, mut rx) = bounded(3);
            tx.send(1).await.unwrap();
            tx.send(2).await.unwrap();
            tx.send(3).await.unwrap();
            assert_eq!(rx.recv().await.unwrap(), 1);
            assert_eq!(rx.recv().await.unwrap(), 2);
            assert_eq!(rx.recv().await.unwrap(), 3);
        });
    }

    #[test]
    fn test_recv_pending_when_empty() {
        block_on(VirtualIo::default(), async {
            let (_tx, mut rx) = bounded::<i32>(1);

            let waker = std::task::Waker::noop();
            let mut cx = std::task::Context::from_waker(&waker);
            let mut fut = pin!(rx.recv());
            assert!(matches!(fut.as_mut().poll(&mut cx), Poll::Pending));
        });
    }

    #[test]
    fn test_recv_closed() {
        block_on(VirtualIo::default(), async {
            let (tx, mut rx) = bounded::<i32>(1);
            drop(tx);
            assert_eq!(rx.recv().await, Err(RecvError));
        });
    }

    #[test]
    fn test_recv_woken_when_last_sender_dropped() {
        // Receiver is blocked waiting on an empty channel. Dropping the last sender
        // must wake the receiver so it can observe Closed without deadlocking.
        block_on(VirtualIo::default(), async {
            let (tx, mut rx) = bounded::<i32>(1);
            spawn(async move {
                drop(tx);
            });
            assert_eq!(rx.recv().await, Err(RecvError));
        });
    }

    #[test]
    fn test_recv_woken_when_last_of_multiple_senders_dropped() {
        block_on(VirtualIo::default(), async {
            let (tx, mut rx) = bounded::<i32>(1);
            let tx2 = tx.clone();
            spawn(async move {
                drop(tx);
                drop(tx2);
            });
            assert_eq!(rx.recv().await, Err(RecvError));
        });
    }

    // -- stream implementation of recv --

    #[test]
    fn test_stream_recv() {
        const ITEMS: &[i32; 3] = &[1, 2, 3];
        block_on(VirtualIo::default(), async {
            let (mut tx, rx) = bounded::<i32>(1);
            spawn(async move {
                for &item in ITEMS {
                    tx.send(item).await.unwrap();
                }
            });

            let result: Vec<_> = rx.collect().await;
            assert_eq!(result, ITEMS);
        })
    }
}