unsync 0.2.0

//! An unsynchronized single-producer, single-consumer channel.
//!
//! You might also know this simply as a "queue", but I'm sticking with a
//! uniform naming scheme here so give me a break.
//!
//! This allocates storage internally to maintain shared state between the
//! [Sender] and [Receiver].

use core::error::Error;
use core::fmt::{self, Debug, Display, Formatter};
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll, Waker};

use alloc::collections::VecDeque;

use crate::bi_rc::BiRc;

/// Error raised when sending a message over the queue.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct SendError<T>(pub T);

impl<T> Display for SendError<T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "channel disconnected")
    }
}

impl<T> Error for SendError<T> where T: Debug {}

/// Interior shared state.
///
/// Note that we maintain two sets of waker to avoid having to clone the waker
/// associated with the channel unecessarily through the [`Waker::will_wake`]
/// optimization. This is done because it's presumed that the channel will be
/// re-used.
struct Shared<T> {
    /// Waker to wake once sending is available.
    tx: Option<Waker>,
    /// Waker to wake once receiving is available.
    rx: Option<Waker>,
    /// The buffer of messages for this channel.
    buf: VecDeque<T>,
    /// Indicates if the channel is unbounded.
    unbounded: bool,
}

impl<T> Shared<T> {
    /// Test if the current channel is at capacity.
    fn at_capacity(&self) -> bool {
        !self.unbounded && self.buf.capacity() == self.buf.len()
    }
}

/// Sender end of the channel created through [`channel`].
pub struct Sender<T> {
    inner: BiRc<Shared<T>>,
}

impl<T> Sender<T> {
    /// Try to send a message on this channel without blocking.
    ///
    /// This will succeed if there is sufficient capacity to send, but fail
    /// otherwise.
    ///
    /// Note: don't attempt to use this as an optimization over [`Sender::send`]
    /// since it already performs this operation internally as needed.
    ///
    /// # Examples
    ///
    /// ```rust
    /// #[tokio::main(flavor = "current_thread")]
    /// # async fn main() {
    /// let (mut tx, mut rx) = unsync::spsc::channel(3);
    /// assert!(tx.try_send(1).is_ok());
    /// assert!(tx.try_send(2).is_ok());
    /// assert!(tx.try_send(3).is_ok());
    /// assert!(tx.try_send(4).is_err());
    ///
    /// let first = rx.recv().await;
    /// assert_eq!(first, Some(1));
    ///
    /// assert!(tx.try_send(5).is_ok());
    /// assert!(tx.try_send(6).is_err());
    ///
    /// let mut collected = Vec::new();
    ///
    /// // Drop sender so that the channel "ends".
    /// drop(tx);
    ///
    /// while let Some(value) = rx.recv().await {
    ///     collected.push(value);
    /// }
    ///
    /// assert_eq!(collected, vec![2, 3, 5]);
    /// # }
    /// ```
    pub fn try_send(&mut self, value: T) -> Result<(), SendError<T>> {
        unsafe {
            let (inner, both_present) = self.inner.get_mut_unchecked();

            if !both_present || inner.at_capacity() {
                return Err(SendError(value));
            }

            inner.buf.push_back(value);

            if let Some(waker) = &inner.rx {
                waker.wake_by_ref();
            };

            Ok(())
        }
    }

    /// Send a message on this channel.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tokio::task;
    ///
    /// # #[tokio::main(flavor = "current_thread")] async fn main() -> Result<(), task::JoinError> {
    /// let (mut tx, mut rx) = unsync::spsc::channel(1);
    ///
    /// let local = task::LocalSet::new();
    ///
    /// let collected = local.run_until(async move {
    ///     let collect = task::spawn_local(async move {
    ///         let mut out = Vec::new();
    ///
    ///         while let Some(value) = rx.recv().await {
    ///             out.push(value);
    ///         }
    ///
    ///         out
    ///     });
    ///
    ///     let sender = task::spawn_local(async move {
    ///         for n in 0..10 {
    ///             let result = tx.send(n).await;
    ///         }
    ///     });
    ///
    ///     collect.await
    /// }).await?;
    ///
    /// assert_eq!(collected, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
    /// # Ok(()) }
    /// ```
    pub async fn send(&mut self, value: T) -> Result<(), SendError<T>> {
        Send {
            inner: &self.inner,
            value: Some(value),
        }
        .await
    }
}

/// Future returned when sending a value through [Sender::send].
struct Send<'a, T> {
    inner: &'a BiRc<Shared<T>>,
    value: Option<T>,
}

impl<T> Future for Send<'_, T> {
    type Output = Result<(), SendError<T>>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = Pin::into_inner(self);

        unsafe {
            let (inner, both_present) = this.inner.get_mut_unchecked();

            if !both_present {
                inner.tx = None;
                let value = this.value.take().expect("future already completed");
                return Poll::Ready(Err(SendError(value)));
            }

            // If we are at capacity, register ourselves as an interested waker
            // and move on.
            if inner.at_capacity() {
                if !matches!(&inner.tx, Some(w) if w.will_wake(cx.waker())) {
                    inner.tx = Some(cx.waker().clone());
                }

                return Poll::Pending;
            };

            inner
                .buf
                .push_back(this.value.take().expect("future already completed"));

            if let Some(waker) = &inner.rx {
                waker.wake_by_ref();
            };

            Poll::Ready(Ok(()))
        }
    }
}

impl<T> Unpin for Send<'_, T> {}

/// Receiver end of the channel created through [channel].
pub struct Receiver<T> {
    inner: BiRc<Shared<T>>,
}

impl<T> Receiver<T> {
    /// Receive a message on this channel.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tokio::task;
    ///
    /// # #[tokio::main(flavor = "current_thread")] async fn main() -> Result<(), task::JoinError> {
    /// let (mut tx, mut rx) = unsync::spsc::channel(1);
    ///
    /// let local = task::LocalSet::new();
    ///
    /// let collected = local.run_until(async move {
    ///     let collect = task::spawn_local(async move {
    ///         let mut out = Vec::new();
    ///
    ///         while let Some(value) = rx.recv().await {
    ///             out.push(value);
    ///         }
    ///
    ///         out
    ///     });
    ///
    ///     let sender = task::spawn_local(async move {
    ///         for n in 0..10 {
    ///             let result = tx.send(n).await;
    ///         }
    ///     });
    ///
    ///     collect.await
    /// }).await?;
    ///
    /// assert_eq!(collected, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
    /// # Ok(()) }
    /// ```
    pub async fn recv(&mut self) -> Option<T> {
        Recv { inner: &self.inner }.await
    }
}

struct Recv<'a, T> {
    inner: &'a BiRc<Shared<T>>,
}

impl<T> Future for Recv<'_, T> {
    type Output = Option<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = Pin::into_inner(self);

        unsafe {
            let (inner, both_present) = this.inner.get_mut_unchecked();

            if let Some(value) = inner.buf.pop_front() {
                // Popping freed capacity, so wake a sender waiting to send.
                if let Some(tx) = &inner.tx {
                    tx.wake_by_ref();
                }

                return Poll::Ready(Some(value));
            }

            if !both_present {
                inner.rx = None;
                return Poll::Ready(None);
            }

            if !matches!(&inner.rx, Some(w) if w.will_wake(cx.waker())) {
                inner.rx = Some(cx.waker().clone())
            }

            if let Some(tx) = &inner.tx {
                tx.wake_by_ref();
            }

            Poll::Pending
        }
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        unsafe {
            if let Some(waker) = self.inner.get_mut_unchecked().0.rx.take() {
                waker.wake();
            }
        }
    }
}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        unsafe {
            if let Some(waker) = self.inner.get_mut_unchecked().0.tx.take() {
                waker.wake();
            }
        }
    }
}

/// Setup a spsc with the given capacity.
///
/// Any sender is capable of sending without blocking up until `capacity` number
/// of elements have been buffered.
///
/// # Panics
///
/// Panics if `capacity` is set to `0`.
pub fn channel<T>(capacity: usize) -> (Sender<T>, Receiver<T>) {
    assert!(capacity > 0, "capacity cannot be 0");

    let (a, b) = BiRc::new(Shared {
        tx: None,
        rx: None,
        buf: VecDeque::with_capacity(capacity),
        unbounded: false,
    });

    let rx = Receiver { inner: a };
    let tx = Sender { inner: b };

    (tx, rx)
}

/// Setup a spsc with an unbounded capacity.
///
/// Sending through this channel will never block.
pub fn unbounded<T>() -> (Sender<T>, Receiver<T>) {
    let (a, b) = BiRc::new(Shared {
        tx: None,
        rx: None,
        buf: VecDeque::new(),
        unbounded: true,
    });

    let rx = Receiver { inner: a };
    let tx = Sender { inner: b };

    (tx, rx)
}

#[cfg(test)]
mod tests {
    use core::future::Future;
    use core::sync::atomic::{AtomicBool, Ordering};
    use core::task::{Context, Poll, Waker};

    use alloc::boxed::Box;
    use alloc::sync::Arc;
    use alloc::task::Wake;

    use super::channel;
    use crate::utils::noop_cx;

    struct Flag(AtomicBool);

    impl Wake for Flag {
        fn wake(self: Arc<Self>) {
            self.0.store(true, Ordering::SeqCst);
        }

        fn wake_by_ref(self: &Arc<Self>) {
            self.0.store(true, Ordering::SeqCst);
        }
    }

    // A `recv` re-polled with a new waker (e.g. moved between tasks) must register
    // it. Otherwise a later send wakes the old waker and the receiver never runs.
    #[test]
    fn recv_updates_changed_waker() {
        let (mut tx, mut rx) = channel::<u32>(1);

        let first = Arc::new(Flag(AtomicBool::new(false)));
        let second = Arc::new(Flag(AtomicBool::new(false)));
        let w1 = Waker::from(first.clone());
        let w2 = Waker::from(second.clone());

        let mut fut = Box::pin(rx.recv());
        assert!(
            fut.as_mut()
                .poll(&mut Context::from_waker(&w1))
                .is_pending()
        );
        assert!(
            fut.as_mut()
                .poll(&mut Context::from_waker(&w2))
                .is_pending()
        );

        tx.try_send(1).unwrap();

        assert!(
            second.0.load(Ordering::SeqCst),
            "latest waker was not woken"
        );
    }

    // A sender blocked on a full channel must be woken when the receiver pops a
    // value (freeing capacity), not only the next time it finds the queue empty.
    #[test]
    fn recv_wakes_blocked_sender_on_pop() {
        let (mut tx, mut rx) = channel::<u32>(1);
        tx.try_send(1).unwrap();

        let sender = Arc::new(Flag(AtomicBool::new(false)));
        let ws = Waker::from(sender.clone());

        // The channel is full, so this send parks and registers `ws`.
        let mut send = Box::pin(tx.send(2));
        assert!(
            send.as_mut()
                .poll(&mut Context::from_waker(&ws))
                .is_pending()
        );

        // Popping a value frees capacity and must wake the parked sender.
        let mut recv = Box::pin(rx.recv());
        assert_eq!(recv.as_mut().poll(&mut noop_cx()), Poll::Ready(Some(1)));

        assert!(
            sender.0.load(Ordering::SeqCst),
            "blocked sender was not woken"
        );
    }
}