async_observe/

lib.rs

#![cfg_attr(all(doc, not(doctest)), doc = include_str!("../README.md"))]
#![cfg_attr(
    any(not(doc), doctest),
    doc = "Async single-producer, multi-consumer channel that only retains the last sent value"
)]

use {
    event_listener::Event,
    futures_lite::{Stream, stream},
    std::{
        error, fmt,
        sync::{
            Arc, RwLock, RwLockReadGuard, RwLockWriteGuard,
            atomic::{AtomicUsize, Ordering},
        },
    },
};

/// Creates a new observer channel, returning the sender and receiver halves.
///
/// All values sent by [`Sender`] will become visible to the [`Receiver`]
/// handles. Only the last value sent is made available to the [`Receiver`]
/// half. All intermediate values are dropped.
///
/// # Examples
///
/// This example prints numbers from 0 to 9:
///
/// ```
/// # fn f() {
/// use {
///     futures_lite::future,
///     std::{thread, time::Duration},
/// };
///
/// let (tx, mut rx) = async_observe::channel(0);
///
/// // Perform computations in another thread
/// thread::spawn(move || {
///     for n in 1..10 {
///         thread::sleep(Duration::from_secs(1));
///
///         // Send a new value without blocking the thread.
///         // If sending fails, it means the sender was dropped.
///         // In that case stop the computation.
///         if tx.send(n).is_err() {
///             break;
///         }
///     }
/// });
///
/// // Print the initial value (0)
/// let n = rx.observe(|n| *n);
/// println!("{n}");
///
/// future::block_on(async {
///     // Print the value whenever it changes
///     while let Ok(n) = rx.recv().await {
///         println!("{n}");
///     }
/// });
/// # }
/// ```
///
/// In this example a new thread is spawned, but you can also use the channel
/// to update values from a future/async task.
///
/// Note that this channel does not have a message queue - it only stores the
/// latest update. Therefore, it does *not* guarantee that you will observe
/// every intermediate value. If you need to observe each change, use a
/// message-queue channel such as [`async-channel`].
///
/// [`async-channel`]: https://docs.rs/async-channel/latest/async_channel
pub fn channel<T>(init: T) -> (Sender<T>, Receiver<T>) {
    let shared = Arc::new(Shared {
        value: RwLock::new(init),
        state: State::new(),
        rx_count: AtomicUsize::new(1),
        changed: Event::new(),
        all_receivers_dropped: Event::new(),
    });

    let tx = Sender {
        shared: shared.clone(),
    };

    let rx = Receiver {
        shared,
        last_version: 0,
    };

    (tx, rx)
}

/// Sends values to the associated [`Receiver`]s.
///
/// Created by the [`channel`] function.
#[derive(Debug)]
pub struct Sender<T> {
    /// The inner shared state.
    shared: Arc<Shared<T>>,
}

impl<T> Sender<T> {
    /// Sends a new value to the channel and notifies all receivers.
    ///
    /// # Examples
    ///
    /// ```
    /// let (tx, rx) = async_observe::channel(0);
    /// assert_eq!(rx.observe(|n| *n), 0);
    ///
    /// // Send a new value
    /// tx.send(1);
    ///
    /// // Now the receiver can see it
    /// assert_eq!(rx.observe(|n| *n), 1);
    /// ```
    ///
    /// To wait until the value is updated, use the receiver's async methods
    /// [`changed`](Receiver::changed) or [`recv`](Receiver::recv).
    pub fn send(&self, value: T) -> Result<(), SendError<T>> {
        if self.shared.rx_count.load(Ordering::Relaxed) == 0 {
            // All receivers have been dropped
            return Err(SendError(value));
        }

        // Replace the value
        *self.shared.write_value() = value;
        self.shared.state.increment_version();

        // Notify all receivers
        self.shared.changed.notify(usize::MAX);

        Ok(())
    }

    /// Waits until all receivers are dropped.
    ///
    /// # Examples
    ///
    /// A producer can wait until no consumers is interested in its updates
    /// anymore and then stop working.
    ///
    /// ```
    /// # futures_lite::future::block_on(async {
    /// # use futures_lite::future::yield_now as wait_long_time;
    /// use futures_lite::future;
    ///
    /// let (tx, mut rx) = async_observe::channel(0);
    ///
    /// // The producer runs concurrently and waits until the channel
    /// // is closed, then cancels the main future
    /// let producer = future::or(
    ///     async {
    ///         let mut n = 0;
    ///         loop {
    ///             wait_long_time().await;
    ///
    ///             // The producer could check if the channel is closed on send,
    ///             // but computing a new value might take a long time.
    ///             // Instead, we cancel the whole working future.
    ///             _ = tx.send(n);
    ///             n += 1;
    ///         }
    ///     },
    ///     tx.closed(),
    /// );
    ///
    /// let consumer = async move {
    /// //                   ^^^^
    /// // Note: rx is moved into this future,
    /// // so it will be dropped when the loop ends.
    /// // Alternatively, you could call `drop(rx);` explicitly.
    ///
    ///     while let Ok(n) = rx.recv().await {
    ///         // After receiving number 5,
    ///         // the consumer is no longer interested.
    ///         if n == 5 {
    ///             break;
    ///         }
    ///     }
    /// };
    ///
    /// future::zip(producer, consumer).await;
    /// # });
    /// ```
    ///
    /// The receiver is dropped after getting the number 5. Since there are no
    /// other receivers, this makes `tx.closed()` finish and the sender stops
    /// sending new values.
    pub async fn closed(&self) {
        if self.shared.rx_count.load(Ordering::Relaxed) == 0 {
            return;
        }

        // In order to avoid a notification loss, we first request a notification,
        // **then** check the current `rx_count`.
        // If `rx_count` is 0, the notification request is dropped.
        event_listener::listener!(self.shared.all_receivers_dropped => listener);

        if self.shared.rx_count.load(Ordering::Relaxed) == 0 {
            return;
        }

        listener.await;

        debug_assert_eq!(self.shared.rx_count.load(Ordering::Relaxed), 0);
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        self.shared.state.close();
        self.shared.changed.notify(usize::MAX);
    }
}

/// Error produced when sending a value fails.
#[derive(PartialEq, Eq)]
pub struct SendError<T>(pub T);

impl<T> fmt::Display for SendError<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("sending on a closed channel")
    }
}

impl<T> fmt::Debug for SendError<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("sending on a closed channel")
    }
}

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

/// Receives values from the associated [`Sender`].
///
/// Created by the [`channel`] function.
#[derive(Debug)]
pub struct Receiver<T> {
    /// The inner shared state.
    shared: Arc<Shared<T>>,

    /// Last observed version.
    last_version: usize,
}

impl<T> Receiver<T> {
    /// Observes the latest value sent to the channel.
    ///
    /// This method takes a closure and calls it with a reference to the value.
    /// While the closure is running [`send`](Sender::send) calls are blocked.
    /// Because of this, the closure should run only as long as needed.
    /// A common pattern is to copy or clone the value inside the closure, then
    /// return and work with the copy outside.
    ///
    /// You can observe the value at any time, but usually you want to wait
    /// until it changes. For that, use the [`changed`](Receiver::changed)
    /// async method.
    ///
    /// # Examples
    ///
    /// ```
    /// # futures_lite::future::block_on(async {
    /// let (tx, mut rx) = async_observe::channel(0);
    ///
    /// // Send a new value
    /// tx.send(1);
    ///
    /// // Wait until the value changes
    /// rx.changed().await?;
    ///
    /// // Now we can read the new value
    /// let n = rx.observe(|n| *n);
    /// assert_eq!(n, 1);
    /// # Ok::<_, async_observe::RecvError>(())
    /// # });
    /// ```
    ///
    /// If the value type implements `Clone`, you can use
    /// [`recv`](Receiver::recv) instead, which waits for a change and returns
    /// the new value.
    ///
    /// ```
    /// # futures_lite::future::block_on(async {
    /// let (tx, mut rx) = async_observe::channel(0);
    ///
    /// // Send a new value
    /// tx.send(1);
    ///
    /// // Wait until the value changes and read it
    /// let n = rx.recv().await?;
    /// assert_eq!(n, 1);
    /// # Ok::<_, async_observe::RecvError>(())
    /// # });
    /// ```
    ///
    /// # Possible deadlock
    ///
    /// Calling [`send`](Sender::send) inside the closure will deadlock:
    ///
    /// ```no_run
    /// let (tx, rx) = async_observe::channel(0);
    /// rx.observe(|n| {
    ///     _ = tx.send(n + 1);
    /// });
    /// ```
    pub fn observe<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&T) -> R,
    {
        f(&self.shared.read_value())
    }

    /// Waits for the value to change.
    ///
    /// Call [`observe`](Receiver::observe) to read the new value.
    pub async fn changed(&mut self) -> Result<(), RecvError> {
        if self
            .shared
            .state
            .version_changed(&mut self.last_version)
            .ok_or(RecvError)?
        {
            return Ok(());
        }

        // In order to avoid a notification loss, we first request a notification,
        // **then** check the current value's version.
        // If a new version exists, the notification request is dropped.
        event_listener::listener!(self.shared.changed => listener);

        if self
            .shared
            .state
            .version_changed(&mut self.last_version)
            .ok_or(RecvError)?
        {
            return Ok(());
        }

        listener.await;

        let changed = self
            .shared
            .state
            .version_changed(&mut self.last_version)
            .ok_or(RecvError)?;

        debug_assert!(changed);
        Ok(())
    }

    /// Waits for the value to change and then returns a clone of it.
    ///
    /// # Examples
    ///
    /// ```
    /// # futures_lite::future::block_on(async {
    /// let (tx, mut rx) = async_observe::channel(0);
    ///
    /// // Send a new value
    /// tx.send(1);
    ///
    /// // Wait until the value changes and read it
    /// let n = rx.recv().await?;
    /// assert_eq!(n, 1);
    /// # Ok::<_, async_observe::RecvError>(())
    /// # });
    /// ```
    pub async fn recv(&mut self) -> Result<T, RecvError>
    where
        T: Clone,
    {
        self.changed().await?;
        Ok(self.observe(T::clone))
    }

    /// Creates a [stream](Stream) from the receiver.
    ///
    /// The stream ends when the [`Sender`] is dropped.
    ///
    /// # Examples
    ///
    /// ```
    /// # futures_lite::future::block_on(async {
    /// # use futures_lite::future::yield_now as wait_long_time;
    /// use futures_lite::{StreamExt, future};
    ///
    /// let (tx, rx) = async_observe::channel(0);
    ///
    /// let producer = async move {
    /// //                   ^^^^
    /// // Move tx into the future so it is dropped after the loop ends.
    /// // Dropping the sender is important so the receiver can
    /// // see it and stop the stream.
    ///
    ///     for n in 1..10 {
    ///         wait_long_time().await;
    ///         _ = tx.send(n);
    ///     }
    /// };
    ///
    /// // Create a stream from the receiver
    /// let consumer = rx
    ///     .into_stream()
    ///     .for_each(|n| println!("{n}"));
    ///
    /// future::zip(producer, consumer).await;
    /// # });
    /// ```
    ///
    /// # Return type
    ///
    /// Due to implementation details, the stream currently does not have
    /// a named type. For example, you cannot store this stream as a struct
    /// field (without a generic) or use it as an associated type in a trait
    /// implementation. As long as you don't need that, just use the returned
    /// anonymous type directly.
    ///
    /// However, if you need a named type, one solution is to box the stream:
    ///
    /// ```
    /// use {
    ///     async_observe::Receiver,
    ///     futures_lite::{StreamExt, stream::Boxed},
    /// };
    ///
    /// fn named_stream(rx: Receiver<u32>) -> Boxed<u32> {
    ///     rx.into_stream().boxed()
    /// }
    /// ```
    pub fn into_stream(self) -> impl Stream<Item = T>
    where
        T: Clone,
    {
        stream::unfold(self, async |mut me| {
            let value = me.recv().await.ok()?;
            Some((value, me))
        })
    }
}

impl<T> Clone for Receiver<T> {
    fn clone(&self) -> Self {
        self.shared.rx_count.fetch_add(1, Ordering::Relaxed);
        Self {
            shared: self.shared.clone(),
            last_version: self.last_version,
        }
    }
}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        if self.shared.rx_count.fetch_sub(1, Ordering::Relaxed) == 1 {
            // Notify all senders.
            // Even though `Sender` is not `Clone`, it can still wait for
            // the channel to close from multiple places, since the `closed`
            // method takes `&self`.
            self.shared.all_receivers_dropped.notify(usize::MAX);
        }
    }
}

/// Error produced when receiving a change notification.
#[derive(PartialEq, Eq)]
pub struct RecvError;

impl fmt::Display for RecvError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("receiving on a closed channel")
    }
}

impl fmt::Debug for RecvError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("receiving on a closed channel")
    }
}

impl error::Error for RecvError {}

#[derive(Debug)]
struct Shared<T> {
    /// The most recent value.
    value: RwLock<T>,

    /// The current state.
    state: State,

    /// Tracks the number of `Receiver` instances.
    rx_count: AtomicUsize,

    /// Event when the value has changed or the `Sender` has been dropped.
    changed: Event,

    /// Event when all `Receiver`s have been dropped.
    all_receivers_dropped: Event,
}

impl<T> Shared<T> {
    fn read_value(&self) -> RwLockReadGuard<'_, T> {
        // The `RwLock` has no poisoned state
        match self.value.read() {
            Ok(guard) => guard,
            Err(e) => e.into_inner(),
        }
    }

    fn write_value(&self) -> RwLockWriteGuard<'_, T> {
        // The `RwLock` has no poisoned state
        match self.value.write() {
            Ok(guard) => guard,
            Err(e) => e.into_inner(),
        }
    }
}

#[derive(Debug)]
struct State(AtomicUsize);

impl State {
    /// Using 2 as the version step preserves the `CLOSED_BIT`.
    const VERSION_STEP: usize = 2;

    /// The least significant bit signifies a closed channel.
    const CLOSED_BIT: usize = 1;

    fn new() -> Self {
        Self(AtomicUsize::new(0))
    }

    fn increment_version(&self) {
        self.0.fetch_add(Self::VERSION_STEP, Ordering::Release);
    }

    fn version_changed(&self, last_version: &mut usize) -> Option<bool> {
        let state = self.0.load(Ordering::Acquire);
        let new_version = state & !Self::CLOSED_BIT;

        if *last_version != new_version {
            *last_version = new_version;
            return Some(true);
        }

        if Self::CLOSED_BIT == state & Self::CLOSED_BIT {
            return None;
        }

        Some(false)
    }

    fn close(&self) {
        self.0.fetch_or(Self::CLOSED_BIT, Ordering::Release);
    }
}