ntex_util/channel/

mpsc.rs

//! A multi-producer, single-consumer, futures-aware, FIFO queue.
use std::collections::VecDeque;
use std::future::poll_fn;
use std::{fmt, panic::UnwindSafe, pin::Pin, task::Context, task::Poll};

use futures_core::{FusedStream, Stream};

use super::cell::{Cell, wake};
use crate::task::LocalWaker;

/// Creates an unbounded in-memory channel.
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let shared = Cell::new(Shared {
        has_receiver: true,
        buffer: VecDeque::new(),
        blocked_recv: LocalWaker::new(),
        closed: LocalWaker::new(),
    });
    let sender = Sender {
        shared: shared.clone(),
    };
    let receiver = Receiver { shared };
    (sender, receiver)
}

#[derive(Debug)]
struct Shared<T> {
    buffer: VecDeque<T>,
    blocked_recv: LocalWaker,
    closed: LocalWaker,
    has_receiver: bool,
}

fn close<T>(shared: &Cell<Shared<T>>) {
    let (recv, closed) = {
        let shared = shared.get_mut();
        shared.has_receiver = false;
        (shared.blocked_recv.take_local(), shared.closed.take_local())
    };
    wake(recv);
    wake(closed);
}

/// The transmission end of a channel.
///
/// This is created by the `channel` function.
#[derive(Debug)]
pub struct Sender<T> {
    shared: Cell<Shared<T>>,
}

impl<T> Unpin for Sender<T> {}

impl<T> Sender<T> {
    /// Sends the provided message along this channel.
    pub fn send(&self, item: T) -> Result<(), SendError<T>> {
        let waker = {
            let shared = self.shared.get_mut();
            if !shared.has_receiver {
                return Err(SendError(item)); // receiver was dropped
            }
            shared.buffer.push_back(item);
            shared.blocked_recv.take_local()
        };
        wake(waker);
        Ok(())
    }

    /// Closes the channel for every sender.
    ///
    /// This prevents any further messages from being sent on the channel, by
    /// this sender or any of its clones, while still enabling the receiver to
    /// drain messages that are buffered.
    pub fn close(&self) {
        close(&self.shared);
    }

    /// Returns `true` if the channel is closed or the receiver has been dropped.
    pub fn is_closed(&self) -> bool {
        self.shared.strong_count() == 1 || !self.shared.get_ref().has_receiver
    }

    /// Polls whether the channel is closed or the receiver has been dropped.
    ///
    /// Only the task from the most recent call is woken, across this sender
    /// and all of its clones.
    pub fn poll_closed(&self, cx: &mut Context<'_>) -> Poll<()> {
        if !self.shared.get_ref().has_receiver {
            return Poll::Ready(());
        }
        LocalWaker::register_in(|| &self.shared.get_ref().closed, cx.waker());

        // the waker callbacks may have closed the channel
        if self.shared.get_ref().has_receiver {
            Poll::Pending
        } else {
            Poll::Ready(())
        }
    }

    /// Waits until the channel is closed or the receiver has been dropped.
    ///
    /// Only one task may wait at a time, across this sender and all of its
    /// clones; see [`poll_closed`](Self::poll_closed).
    pub async fn closed(&self) {
        poll_fn(|cx| self.poll_closed(cx)).await;
    }
}

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

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        let count = self.shared.strong_count();
        let waker = {
            let shared = self.shared.get_mut();

            // check is last sender is about to drop
            if shared.has_receiver && count == 2 {
                // the channel is closed for good, and the receiver's
                // stream has ended
                shared.has_receiver = false;
                shared.blocked_recv.take_local()
            } else {
                None
            }
        };
        wake(waker);
    }
}

/// The receiving end of a channel which implements the `Stream` trait.
///
/// This is created by the `channel` function.
///
/// # Single waiter
///
/// The receiver stores one waker. [`recv`](Self::recv),
/// [`poll_recv`](Self::poll_recv) and [`Stream::poll_next`] take `&self`, but
/// only the task that polled most recently is woken. Wait for messages from
/// one task at a time; an earlier concurrent waiter may stay pending even
/// while messages are buffered.
#[derive(Debug)]
pub struct Receiver<T> {
    shared: Cell<Shared<T>>,
}

impl<T> Receiver<T> {
    /// Creates an additional sender for this channel.
    ///
    /// Closing is permanent: once the channel is closed, or every sender has
    /// been dropped, the returned sender is closed too and its `send()` fails.
    pub fn sender(&self) -> Sender<T> {
        Sender {
            shared: self.shared.clone(),
        }
    }

    /// Closes the receiving half of a channel, without dropping it.
    ///
    /// This prevents any further messages from being sent on the channel
    /// while still enabling the receiver to drain messages that are buffered.
    pub fn close(&self) {
        close(&self.shared);
    }

    /// Returns whether this channel is closed.
    pub fn is_closed(&self) -> bool {
        self.shared.strong_count() == 1 || !self.shared.get_ref().has_receiver
    }

    /// Waits for the next message.
    ///
    /// Returns `None` once the channel is closed or every sender has been
    /// dropped, and every buffered message has been received.
    ///
    /// Only one task may wait at a time; see [single waiter](Self#single-waiter).
    pub async fn recv(&self) -> Option<T> {
        poll_fn(|cx| self.poll_recv(cx)).await
    }

    /// Polls for the next message.
    ///
    /// Returns `Ready(None)` once the channel is closed or every sender has
    /// been dropped, and every buffered message has been received.
    ///
    /// A pending poll replaces the previously registered waker; see
    /// [single waiter](Self#single-waiter).
    pub fn poll_recv(&self, cx: &mut Context<'_>) -> Poll<Option<T>> {
        if let Poll::Ready(msg) = self.try_recv() {
            return Poll::Ready(msg);
        }
        LocalWaker::register_in(|| &self.shared.get_ref().blocked_recv, cx.waker());

        // the waker callbacks may have sent a message or closed the channel
        self.try_recv()
    }

    fn try_recv(&self) -> Poll<Option<T>> {
        let shared = self.shared.get_mut();
        if let Some(msg) = shared.buffer.pop_front() {
            Poll::Ready(Some(msg))
        } else if !shared.has_receiver || self.shared.strong_count() == 1 {
            // The channel is closed or all senders have been dropped, and the
            // buffer is drained, so end the stream.
            Poll::Ready(None)
        } else {
            Poll::Pending
        }
    }
}

impl<T> Unpin for Receiver<T> {}

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

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

impl<T> FusedStream for Receiver<T> {
    /// Returns `true` once the channel is closed and every buffered message
    /// has been received.
    fn is_terminated(&self) -> bool {
        self.is_closed() && self.shared.get_ref().buffer.is_empty()
    }
}

impl<T> UnwindSafe for Receiver<T> {}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        let (buffer, waker) = {
            let shared = self.shared.get_mut();
            shared.has_receiver = false;
            (
                std::mem::take(&mut shared.buffer),
                shared.closed.take_local(),
            )
        };
        wake(waker);

        // queued messages may send on this channel from their `Drop`
        drop(buffer);
    }
}

/// Error type for sending, used when the receiving end of a channel is
/// dropped
pub struct SendError<T>(T);

impl<T> std::error::Error for SendError<T> {}

impl<T> fmt::Debug for SendError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_tuple("SendError").field(&"...").finish()
    }
}

impl<T> fmt::Display for SendError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "send failed because receiver is gone")
    }
}

impl<T> SendError<T> {
    /// Returns the message that was attempted to be sent but failed.
    pub fn into_inner(self) -> T {
        self.0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::task::hooks::{ON_CLONE, hook_waker};
    use crate::{future::lazy, future::stream_recv};

    /// STATEMENT: MPSC sender clones share delivery and closure state.
    ///
    /// HOW:
    /// - Send from original and cloned handles and verify an empty live channel remains pending.
    /// - Close or drop endpoints, check failed sends, and recover the value from `SendError`.
    #[ntex::test]
    async fn test_mpsc() {
        let (tx, mut rx) = channel();
        assert!(format!("{tx:?}").contains("Sender"));
        assert!(format!("{rx:?}").contains("Receiver"));

        tx.send("test").unwrap();
        assert_eq!(stream_recv(&mut rx).await.unwrap(), "test");

        let tx2 = tx.clone();
        tx2.send("test2").unwrap();
        assert_eq!(stream_recv(&mut rx).await.unwrap(), "test2");

        assert_eq!(
            lazy(|cx| Pin::new(&mut rx).poll_next(cx)).await,
            Poll::Pending
        );
        drop(tx2);
        assert_eq!(
            lazy(|cx| Pin::new(&mut rx).poll_next(cx)).await,
            Poll::Pending
        );
        drop(tx);

        let (tx, mut rx) = channel::<String>();
        tx.close();
        assert_eq!(stream_recv(&mut rx).await, None);

        let (tx, rx) = channel();
        tx.send("test").unwrap();
        drop(rx);
        assert!(tx.send("test").is_err());

        let (tx, _) = channel();
        let tx2 = tx.clone();
        tx.close();
        assert!(tx.send("test").is_err());
        assert!(tx2.send("test").is_err());

        let err = SendError("test");
        assert!(format!("{err:?}").contains("SendError"));
        assert!(format!("{err}").contains("send failed because receiver is gone"));
        assert_eq!(err.into_inner(), "test");
    }

    /// STATEMENT: Explicit closure wakes receivers, and closure by dropping the last
    /// sender is permanent.
    ///
    /// HOW:
    /// - Close from each endpoint and check closed, terminated, and registered-waker state.
    /// - Drop the last sender, then create one through the receiver and check the channel
    ///   stays closed and terminated, `send()` fails, and `recv()` returns `None`.
    #[ntex::test]
    async fn test_close() {
        let (tx, rx) = channel::<()>();
        assert!(!tx.is_closed());
        assert!(!rx.is_closed());
        assert!(!rx.is_terminated());

        tx.close();
        assert!(tx.is_closed());
        assert!(rx.is_closed());
        assert!(rx.is_terminated());

        let (tx, rx) = channel::<()>();
        assert_eq!(lazy(|cx| rx.poll_recv(cx)).await, Poll::Pending);
        assert!(rx.shared.get_ref().blocked_recv.is_set());
        rx.close();
        assert!(tx.is_closed());
        assert!(!rx.shared.get_ref().blocked_recv.is_set());
        assert_eq!(lazy(|cx| rx.poll_recv(cx)).await, Poll::Ready(None));

        let (tx, rx) = channel::<()>();
        drop(tx);
        assert!(rx.is_closed());
        assert!(rx.is_terminated());
        let tx = rx.sender();
        assert!(tx.is_closed());
        assert!(rx.is_closed());
        assert!(rx.is_terminated());
        assert!(tx.send(()).is_err());
        assert_eq!(rx.recv().await, None);
    }

    /// STATEMENT: Sender closure waiters complete when either endpoint closes the channel.
    ///
    /// HOW:
    /// - Register closure waiters before receiver drop, receiver close, and sender close.
    /// - Verify stored wakers are consumed and closure futures complete for cloned handles too.
    #[ntex::test]
    async fn test_poll_closed() {
        let (tx, rx) = channel::<()>();
        assert_eq!(lazy(|cx| tx.poll_closed(cx)).await, Poll::Pending);
        assert!(tx.shared.get_ref().closed.is_set());
        drop(rx);
        assert!(!tx.shared.get_ref().closed.is_set());
        assert_eq!(lazy(|cx| tx.poll_closed(cx)).await, Poll::Ready(()));
        tx.closed().await;

        let (tx, rx) = channel::<()>();
        assert_eq!(lazy(|cx| tx.poll_closed(cx)).await, Poll::Pending);
        rx.close();
        assert!(!tx.shared.get_ref().closed.is_set());
        tx.closed().await;

        let (tx, _rx) = channel::<()>();
        let tx2 = tx.clone();
        assert_eq!(lazy(|cx| tx2.poll_closed(cx)).await, Poll::Pending);
        tx.close();
        assert!(!tx.shared.get_ref().closed.is_set());
        tx2.closed().await;
    }

    /// STATEMENT: A receiver and the closure state shared by sender clones each wake
    /// only the most recently registered task, and a superseded waiter still
    /// observes the result when polled again.
    ///
    /// HOW:
    /// - Poll `poll_recv` with two distinct counted wakers, then send two messages.
    /// - Check only the second waker was woken once, and re-polling with the first
    ///   returns the first message.
    /// - Poll `poll_closed` on two sender clones with distinct wakers, drop the
    ///   receiver, and check only the second waker was woken while both polls are ready.
    #[test]
    fn single_waiter_wakes_latest_task_only() {
        use crate::task::hooks::counted_waker;

        let (tx, rx) = channel::<u32>();
        let (first, first_waker) = counted_waker();
        let (second, second_waker) = counted_waker();
        let mut first_cx = Context::from_waker(&first_waker);
        let mut second_cx = Context::from_waker(&second_waker);

        assert_eq!(rx.poll_recv(&mut first_cx), Poll::Pending);
        assert_eq!(rx.poll_recv(&mut second_cx), Poll::Pending);
        tx.send(1).unwrap();
        tx.send(2).unwrap();
        assert_eq!((first.get(), second.get()), (0, 1));
        assert_eq!(rx.poll_recv(&mut first_cx), Poll::Ready(Some(1)));

        let tx2 = tx.clone();
        let (first, first_waker) = counted_waker();
        let (second, second_waker) = counted_waker();
        let mut first_cx = Context::from_waker(&first_waker);
        let mut second_cx = Context::from_waker(&second_waker);

        assert_eq!(tx.poll_closed(&mut first_cx), Poll::Pending);
        assert_eq!(tx2.poll_closed(&mut second_cx), Poll::Pending);
        drop(rx);
        assert_eq!((first.get(), second.get()), (0, 1));
        assert_eq!(tx.poll_closed(&mut first_cx), Poll::Ready(()));
        assert_eq!(tx2.poll_closed(&mut second_cx), Poll::Ready(()));
    }

    /// STATEMENT: Dropping an MPSC receiver safely rejects reentrant sends from message destructors.
    ///
    /// HOW:
    /// - Queue four values whose destructors attempt another send and count all drops.
    /// - Drop the receiver and verify eight drops and a closed sender.
    #[test]
    fn test_drop_receiver_reentrant_send() {
        use std::rc::Rc;

        struct Msg(Option<Sender<Msg>>, Rc<std::cell::Cell<usize>>);

        impl Drop for Msg {
            fn drop(&mut self) {
                self.1.set(self.1.get() + 1);
                if let Some(tx) = self.0.take() {
                    assert!(tx.send(Msg(None, self.1.clone())).is_err());
                }
            }
        }

        let drops = Rc::new(std::cell::Cell::new(0));
        let (tx, rx) = channel();
        for _ in 0..4 {
            tx.send(Msg(Some(tx.clone()), drops.clone())).unwrap();
        }
        drop(rx);
        assert_eq!(drops.get(), 8);
        assert!(tx.is_closed());
    }

    /// STATEMENT: A closed MPSC stream is terminated only after its buffered messages are drained.
    ///
    /// HOW:
    /// - Queue values, drop the sender, and check ordered delivery before termination.
    /// - Repeat with explicit receiver closure and verify the queued value remains readable.
    #[ntex::test]
    async fn test_fused_drains_buffer() {
        let (tx, mut rx) = channel();
        tx.send(1).unwrap();
        tx.send(2).unwrap();
        drop(tx);

        assert!(!rx.is_terminated());
        assert_eq!(stream_recv(&mut rx).await, Some(1));
        assert_eq!(stream_recv(&mut rx).await, Some(2));
        assert!(rx.is_terminated());
        assert_eq!(stream_recv(&mut rx).await, None);

        let (tx, mut rx) = channel();
        tx.send(1).unwrap();
        rx.close();
        assert!(!rx.is_terminated());
        assert_eq!(stream_recv(&mut rx).await, Some(1));
        assert!(rx.is_terminated());
    }

    /// STATEMENT: Direct MPSC receive returns a sent value and then ends after sender drop.
    ///
    /// HOW:
    /// - Send and receive one value, drop the sender, and verify the next receive returns `None`.
    #[ntex::test]
    async fn test_mpsc_recv() {
        let (tx, rx) = channel();
        tx.send(1).unwrap();
        assert_eq!(rx.recv().await, Some(1));
        drop(tx);
        assert_eq!(rx.recv().await, None);
    }

    /// STATEMENT: `poll_recv` returns a message sent by its waker's clone callback instead of missing it.
    ///
    /// HOW:
    /// - Set an `ON_CLONE` hook that sends `1` through a sender clone.
    /// - Poll the empty receiver with `hook_waker` and check it returns `Ready(Some(1))`.
    #[test]
    fn miri_poll_recv_with_reentrant_clone() {
        let (tx, rx) = channel();
        let tx2 = tx.clone();
        ON_CLONE.set(Some(Box::new(move || tx2.send(1).unwrap())));

        let waker = hook_waker();
        let mut cx = Context::from_waker(&waker);
        assert_eq!(rx.poll_recv(&mut cx), Poll::Ready(Some(1)));
        drop(tx);
    }
}