ntex-util 4.9.0

//! A one-shot, futures-aware channel.
use std::{future::Future, future::poll_fn, pin::Pin, task::Context, task::Poll};

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

/// Creates a new futures-aware, one-shot channel.
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let inner = Cell::new(Inner {
        value: None,
        rx_task: LocalWaker::new(),
    });
    let tx = Sender {
        inner: inner.clone(),
    };
    let rx = Receiver { inner };
    (tx, rx)
}

/// Represents the completion half of a oneshot through which the result of a
/// computation is signaled.
#[derive(Debug)]
pub struct Sender<T> {
    inner: Cell<Inner<T>>,
}

/// A future representing the completion of a computation happening elsewhere in
/// memory.
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Receiver<T> {
    inner: Cell<Inner<T>>,
}

// The channels do not ever project Pin to the inner T
impl<T> Unpin for Receiver<T> {}
impl<T> Unpin for Sender<T> {}

#[derive(Debug)]
struct Inner<T> {
    value: Option<T>,
    rx_task: LocalWaker,
}

impl<T> Sender<T> {
    /// Completes this oneshot with a successful result.
    ///
    /// This function will consume `self` and indicate to the other end, the
    /// `Receiver`, that the value provided is the result of the computation this
    /// represents.
    ///
    /// If the value is successfully enqueued for the remote end to receive,
    /// then `Ok(())` is returned. If the receiving end was dropped before
    /// this function was called, however, then `Err` is returned with the value
    /// provided.
    pub fn send(self, val: T) -> Result<(), T> {
        if self.inner.strong_count() == 2 {
            let waker = {
                let inner = self.inner.get_mut();
                inner.value = Some(val);
                inner.rx_task.take_local()
            };
            wake(waker);
            Ok(())
        } else {
            Err(val)
        }
    }

    /// Tests to see whether this `Sender`'s corresponding `Receiver`
    /// has gone away.
    pub fn is_canceled(&self) -> bool {
        self.inner.strong_count() == 1
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        let waker = self.inner.get_ref().rx_task.take_local();
        wake(waker);
    }
}

impl<T> Receiver<T> {
    /// Waits for the value.
    ///
    /// Returns [`Canceled`] if the sender is dropped without sending a value.
    pub async fn recv(&self) -> Result<T, Canceled> {
        poll_fn(|cx| self.poll_recv(cx)).await
    }

    /// Polls for the value.
    ///
    /// Returns [`Canceled`] if the sender is dropped without sending a value.
    pub fn poll_recv(&self, cx: &mut Context<'_>) -> Poll<Result<T, Canceled>> {
        if let Poll::Ready(res) = self.try_recv() {
            return Poll::Ready(res);
        }
        LocalWaker::register_in(|| &self.inner.get_ref().rx_task, cx.waker());

        // the waker callbacks may have sent the value or dropped the sender
        self.try_recv()
    }

    fn try_recv(&self) -> Poll<Result<T, Canceled>> {
        if let Some(val) = self.inner.get_mut().value.take() {
            Poll::Ready(Ok(val))
        } else if self.inner.strong_count() == 1 {
            Poll::Ready(Err(Canceled))
        } else {
            Poll::Pending
        }
    }
}

impl<T> Future for Receiver<T> {
    type Output = Result<T, Canceled>;

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

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

    /// STATEMENT: One-shot channels deliver once and expose cancellation when either endpoint drops.
    ///
    /// HOW:
    /// - Receive values through both future and `recv` APIs, including after a pending poll.
    /// - Drop each endpoint in fresh channels and check sender cancellation and receiver errors.
    #[ntex::test]
    async fn test_oneshot() {
        let (tx, rx) = channel();
        assert!(format!("{tx:?}").contains("Sender"));
        assert!(format!("{rx:?}").contains("Receiver"));

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

        let (tx, rx) = channel();
        tx.send("test").unwrap();
        assert_eq!(rx.recv().await.unwrap(), "test");

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

        let (tx, rx) = channel::<&'static str>();
        drop(tx);
        assert!(rx.await.is_err());

        let (tx, mut rx) = channel::<&'static str>();
        assert_eq!(lazy(|cx| Pin::new(&mut rx).poll(cx)).await, Poll::Pending);
        tx.send("test").unwrap();
        assert_eq!(rx.await.unwrap(), "test");

        let (tx, mut rx) = channel::<&'static str>();
        assert_eq!(lazy(|cx| Pin::new(&mut rx).poll(cx)).await, Poll::Pending);
        drop(tx);
        assert!(rx.await.is_err());
    }

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

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