From 40f4e6c1e69cf442ccebc5707539014b4e834546 Mon Sep 17 00:00:00 2001 From: orthur2 Date: Sat, 6 Jun 2026 15:17:02 +0200 Subject: [PATCH 1/3] feat(broadcast): add unbounded policy --- CHANGELOG.md | 1 + README.md | 1 + asyncband/src/broadcast/mod.rs | 2 + asyncband/src/broadcast/unbounded/mod.rs | 683 +++++++++++++++++++++ asyncband/src/broadcast/unbounded/tests.rs | 411 +++++++++++++ asyncband/src/internal/arena.rs | 20 +- asyncband/src/lib.rs | 2 +- tests-integration/tests/traits_test.rs | 8 + 8 files changed, 1122 insertions(+), 6 deletions(-) create mode 100644 asyncband/src/broadcast/unbounded/mod.rs create mode 100644 asyncband/src/broadcast/unbounded/tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 0335243..47b999b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file. ### New features +* Implement `broadcast::unbounded`, an unbounded broadcast channel that retains messages until all active receivers consume them or are dropped. * Add an opt-in `asyncband::blocking::FutureExt` bridge with `block_on` and `wait_timeout` methods for waiting on runtime-agnostic futures from synchronous code. ### Breaking changes diff --git a/README.md b/README.md index a3def5b..ff38a46 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ The crate enables no primitives by default. Categories describe each primitive's | | [`mpsc::bounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.bounded.html) | `mpsc` | Send values from multiple producers through a bounded channel. | | | [`mpsc::unbounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.unbounded.html) | `mpsc` | Send values from multiple producers through an unbounded channel. | | | [`broadcast::overflow`](https://docs.rs/asyncband/*/asyncband/broadcast/overflow/) | `broadcast` | Broadcast values and report when slow receivers miss overwritten items. | +| | [`broadcast::unbounded`](https://docs.rs/asyncband/*/asyncband/broadcast/unbounded/) | `broadcast` | Broadcast values and retain them until every active receiver consumes them. | | Workload control | [`Semaphore`](https://docs.rs/asyncband/*/asyncband/semaphore/struct.Semaphore.html) | `semaphore` | Control concurrent access with permits. | | | [`Group`](https://docs.rs/asyncband/*/asyncband/singleflight/struct.Group.html) | `singleflight` | Coalesce concurrent calls for the same key. | diff --git a/asyncband/src/broadcast/mod.rs b/asyncband/src/broadcast/mod.rs index 96b8d1b..63da2a1 100644 --- a/asyncband/src/broadcast/mod.rs +++ b/asyncband/src/broadcast/mod.rs @@ -20,5 +20,7 @@ //! This module provides broadcast channels in one of the following policies: //! //! * [`overflow`]: when the channel is full, the oldest messages are overwritten. +//! * [`unbounded`]: messages are retained until every active receiver consumes them or is dropped. pub mod overflow; +pub mod unbounded; diff --git a/asyncband/src/broadcast/unbounded/mod.rs b/asyncband/src/broadcast/unbounded/mod.rs new file mode 100644 index 0000000..0a2d21d --- /dev/null +++ b/asyncband/src/broadcast/unbounded/mod.rs @@ -0,0 +1,683 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A multi-producer multi-consumer broadcast channel with an unbounded buffer. +//! +//! This channel supports multiple senders and multiple receivers. Each message sent by any +//! sender is received by all active receivers. If a receiver falls behind, messages are buffered +//! until the receiver consumes them or is dropped. +//! +//! # Memory usage +//! +//! This channel does not impose a capacity limit. A slow or stalled receiver can cause the +//! buffer to grow without bound, because messages are retained until every active receiver has +//! consumed them or the receiver is dropped. Use [`Sender::buffer_len`] to monitor the number of +//! messages currently retained by the shared buffer. +//! +//! # Receivers +//! +//! Each receiver has an independent cursor. Use [`Sender::subscribe`] to create a receiver that +//! starts at the current tail of the channel, or [`Receiver::resubscribe`] to skip this receiver's +//! backlog and start a new receiver at the current tail. +//! +//! # Examples +//! +//! Basic usage: +//! +//! ``` +//! use asyncband::broadcast::unbounded; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (tx, mut rx1) = unbounded::channel(); +//! let mut rx2 = tx.subscribe(); +//! +//! tx.send(10); +//! tx.send(20); +//! +//! assert_eq!(rx1.recv().await, Ok(10)); +//! assert_eq!(rx1.recv().await, Ok(20)); +//! assert_eq!(rx2.recv().await, Ok(10)); +//! assert_eq!(rx2.recv().await, Ok(20)); +//! # } +//! ``` +//! +//! Slow receivers do not miss messages: +//! +//! ``` +//! use asyncband::broadcast::unbounded; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (tx, mut rx) = unbounded::channel(); +//! +//! tx.send(1); +//! tx.send(2); +//! tx.send(3); +//! +//! assert_eq!(rx.recv().await, Ok(1)); +//! assert_eq!(rx.recv().await, Ok(2)); +//! assert_eq!(rx.recv().await, Ok(3)); +//! # } +//! ``` + +use std::collections::VecDeque; +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; + +use crate::internal::arena::Arena; +use crate::internal::arena::ArenaKey; +use crate::internal::mutex::Mutex; +use crate::internal::waitset::WaitRegistration; +use crate::internal::waitset::WaitSet; + +#[cfg(test)] +mod tests; + +/// Creates a new broadcast channel with an unbounded buffer. +/// +/// See [module-level documentation](self) for broadcast channel semantics. +/// +/// # Examples +/// +/// ``` +/// use asyncband::broadcast::unbounded; +/// +/// let (tx, mut rx) = unbounded::channel(); +/// tx.send(10); +/// assert_eq!(rx.try_recv(), Ok(10)); +/// ``` +pub fn channel() -> (Sender, Receiver) { + let mut receivers = Arena::new(); + let key = receivers.insert(0); + let shared = Arc::new(Shared { + inner: Mutex::new(Inner { + buffer: VecDeque::new(), + head: 0, + head_receivers: 1, + tail: 0, + receivers, + }), + senders: AtomicUsize::new(1), + waiters: Mutex::new(WaitSet::new()), + }); + let sender = Sender { + shared: shared.clone(), + }; + let receiver = Receiver { shared, key }; + (sender, receiver) +} + +/// Error returned by [`Receiver::recv`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RecvError { + /// The sender has become disconnected, and there will never be any more data received on it. + Disconnected, +} + +impl fmt::Display for RecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + RecvError::Disconnected => write!(f, "receiving on a closed channel"), + } + } +} + +impl std::error::Error for RecvError {} + +/// Error returned by [`Receiver::try_recv`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TryRecvError { + /// This channel is currently empty, but the sender(s) have not yet disconnected, so data may + /// yet become available. + Empty, + /// The sender has become disconnected, and there will never be any more data received on it. + Disconnected, +} + +impl fmt::Display for TryRecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + TryRecvError::Empty => write!(f, "receiving on an empty channel"), + TryRecvError::Disconnected => write!(f, "receiving on a closed channel"), + } + } +} + +impl std::error::Error for TryRecvError {} + +struct Inner { + /// Messages whose versions are in the range `[head, tail)`. + buffer: VecDeque>, + /// The version of the first message in `buffer`. + head: u64, + /// The number of active receivers whose cursor equals `head`. + head_receivers: usize, + /// The next message version to assign. + tail: u64, + /// Cursor for each active receiver. + receivers: Arena, +} + +impl Inner { + fn insert_receiver(&mut self, head: u64) -> ArenaKey { + if head == self.head { + self.head_receivers += 1; + } + + self.receivers.insert(head) + } + + fn remove_receiver(&mut self, key: ArenaKey) -> Vec> { + let head = self.receivers.remove(key); + + if head == self.head { + self.release_head_receiver() + } else { + Vec::new() + } + } + + fn advance_receiver(&mut self, key: ArenaKey, next_head: u64) -> Vec> { + let head = *self + .receivers + .get(key) + .expect("active broadcast receiver must be registered"); + *self + .receivers + .get_mut(key) + .expect("active broadcast receiver must be registered") = next_head; + + if head == self.head { + self.release_head_receiver() + } else { + Vec::new() + } + } + + fn release_head_receiver(&mut self) -> Vec> { + self.head_receivers -= 1; + + if self.head_receivers == 0 { + self.reclaim_consumed() + } else { + Vec::new() + } + } + + fn receive(&mut self, key: ArenaKey) -> Option<(Arc, Vec>)> { + let head = *self + .receivers + .get(key) + .expect("active broadcast receiver must be registered"); + + if head < self.tail { + debug_assert!(head >= self.head); + let offset = (head - self.head) as usize; + let msg = self.buffer[offset].clone(); + let reclaimed = self.advance_receiver(key, head + 1); + Some((msg, reclaimed)) + } else { + None + } + } + + fn reclaim_consumed(&mut self) -> Vec> { + let mut next_head = self.tail; + let mut head_receivers = 0; + + for head in self.receivers.values() { + if *head < next_head { + next_head = *head; + head_receivers = 1; + } else if *head == next_head { + head_receivers += 1; + } + } + + debug_assert!(next_head >= self.head); + let consumed = usize::try_from(next_head - self.head) + .expect("retained broadcast message count exceeds usize"); + // Move reclaimed messages out so their Drop impls run after `inner` is unlocked. + let reclaimed = self.buffer.drain(..consumed).collect(); + + self.head = next_head; + self.head_receivers = head_receivers; + reclaimed + } +} + +struct Shared { + inner: Mutex>, + /// Number of active senders. + senders: AtomicUsize, + /// Waiters (receivers) waiting for new messages. + waiters: Mutex, +} + +/// A sender handle to the broadcast channel. +/// +/// The sender can be cloned to create multiple producers. When all senders are dropped, +/// the channel is closed. +pub struct Sender { + shared: Arc>, +} + +impl Clone for Sender { + fn clone(&self) -> Self { + self.shared.senders.fetch_add(1, Ordering::Relaxed); + Self { + shared: self.shared.clone(), + } + } +} + +impl Drop for Sender { + fn drop(&mut self) { + match self.shared.senders.fetch_sub(1, Ordering::AcqRel) { + 1 => { + // If this is the last sender, we need to wake up the receiver so it can + // observe the disconnected state. + let wakers = { + let mut waiters = self.shared.waiters.lock(); + waiters.take_wakers() + }; + for waker in wakers { + waker.wake(); + } + } + _ => { + // there are still other senders left, do nothing + } + } + } +} + +impl Sender { + /// Broadcasts a value to all active receivers. + /// + /// This operation does not wait for receiver capacity. If receivers fall behind, messages + /// remain buffered until all active receivers have consumed them or the lagging receivers + /// are dropped. + /// + /// If no receivers are active, the message is dropped immediately. + /// + /// # Panics + /// + /// Panics if the internal message version counter overflows. After `u64::MAX` successful sends + /// on one channel instance, the next send panics. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// let (tx, mut rx) = unbounded::channel(); + /// tx.send(10); + /// assert_eq!(rx.try_recv(), Ok(10)); + /// ``` + pub fn send(&self, msg: T) { + let msg = Arc::new(msg); + + { + let mut inner = self.shared.inner.lock(); + inner.tail = inner + .tail + .checked_add(1) + .expect("broadcast channel version counter overflowed"); + + if inner.receivers.is_empty() { + // No receivers means no one will read this message; advance `head` so the + // invariant that `buffer` covers versions `[head, tail)` still holds without + // buffering anything. The buffer is already drained when the last receiver was + // dropped, so there is nothing to clear here. + debug_assert!(inner.buffer.is_empty()); + debug_assert_eq!(inner.head_receivers, 0); + inner.head = inner.tail; + } else { + inner.buffer.push_back(msg); + } + } + + // Notify all waiting receivers. + let wakers = { + let mut waiters = self.shared.waiters.lock(); + waiters.take_wakers() + }; + for waker in wakers { + waker.wake(); + } + } + + /// Returns the number of messages currently retained by the shared buffer. + /// + /// This is not the number of messages any single receiver can still read. It is the shared + /// backlog kept alive by the slowest active receiver. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// let (tx, mut rx) = unbounded::channel(); + /// tx.send(10); + /// assert_eq!(tx.buffer_len(), 1); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(tx.buffer_len(), 0); + /// ``` + pub fn buffer_len(&self) -> usize { + self.shared.inner.lock().buffer.len() + } + + /// Returns the number of active receivers. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// let (tx, rx) = unbounded::channel::(); + /// assert_eq!(tx.receiver_count(), 1); + /// + /// let rx2 = tx.subscribe(); + /// assert_eq!(tx.receiver_count(), 2); + /// + /// drop(rx); + /// drop(rx2); + /// assert_eq!(tx.receiver_count(), 0); + /// ``` + pub fn receiver_count(&self) -> usize { + self.shared.inner.lock().receivers.len() + } + + /// Creates a new receiver that starts receiving messages from the current tail of the channel. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::unbounded::TryRecvError; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, _) = unbounded::channel(); + /// tx.send(10); + /// + /// let mut rx = tx.subscribe(); + /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + /// tx.send(20); + /// assert_eq!(rx.recv().await, Ok(20)); + /// # } + /// ``` + pub fn subscribe(&self) -> Receiver { + let mut inner = self.shared.inner.lock(); + let head = inner.tail; + let key = inner.insert_receiver(head); + let shared = self.shared.clone(); + Receiver { shared, key } + } +} + +/// A receiver handle to the broadcast channel. +/// +/// Each receiver sees every message sent to the channel while the receiver is active. +pub struct Receiver { + shared: Arc>, + key: ArenaKey, +} + +impl Drop for Receiver { + fn drop(&mut self) { + let reclaimed = { + let mut inner = self.shared.inner.lock(); + inner.remove_receiver(self.key) + }; + drop(reclaimed); + } +} + +impl Receiver { + /// Receives the next value for this receiver. + /// + /// # Returns + /// + /// * `Ok(T)`: The next message. + /// * `Err(RecvError::Disconnected)`: All senders have been dropped and no more messages are + /// available. + /// + /// # Cancel safety + /// + /// This method is cancel safe. If `recv` is used as the event in a `select` statement and some + /// other branch completes first, it is guaranteed that no messages were received on this + /// channel. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, mut rx) = unbounded::channel(); + /// tx.send(10); + /// assert_eq!(rx.recv().await, Ok(10)); + /// # } + /// ``` + pub async fn recv(&mut self) -> Result { + Recv { + receiver: self, + registration: None, + } + .await + } + + /// Attempts to receive the next value for this receiver without blocking. + /// + /// # Returns + /// + /// * `Ok(T)`: The next message. + /// * `Err(TryRecvError::Empty)`: No message is currently available. + /// * `Err(TryRecvError::Disconnected)`: All senders have been dropped and no more messages are + /// available. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// let (tx, mut rx) = unbounded::channel(); + /// tx.send(10); + /// assert_eq!(rx.try_recv(), Ok(10)); + /// ``` + pub fn try_recv(&mut self) -> Result { + let (msg, reclaimed) = self.try_recv_shared()?; + drop(reclaimed); + Ok((*msg).clone()) + } +} + +impl Receiver { + fn try_recv_shared(&mut self) -> Result<(Arc, Vec>), TryRecvError> { + // Check this receiver's cursor while holding `inner` before observing `senders`. Senders + // append messages under the same lock before they can be dropped, so an empty result here + // means this receiver has no unread buffered message. + let mut inner = self.shared.inner.lock(); + if let Some(received) = inner.receive(self.key) { + return Ok(received); + } + + if self.shared.senders.load(Ordering::Acquire) == 0 { + Err(TryRecvError::Disconnected) + } else { + Err(TryRecvError::Empty) + } + } + + /// Re-subscribes to the channel, returning a new receiver that starts receiving messages from + /// the *current* tail of the channel. + /// + /// This is useful if the receiver wants to jump to the latest message, skipping everything in + /// between. The original receiver is unchanged and continues to retain its own backlog until + /// it consumes those messages or is dropped. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// let (tx, mut rx) = unbounded::channel(); + /// tx.send(1); + /// tx.send(2); + /// + /// let mut rx2 = rx.resubscribe(); + /// tx.send(3); + /// + /// assert_eq!(rx2.try_recv(), Ok(3)); + /// ``` + pub fn resubscribe(&self) -> Self { + let mut inner = self.shared.inner.lock(); + let head = inner.tail; + let key = inner.insert_receiver(head); + let shared = self.shared.clone(); + Self { shared, key } + } + + /// Returns the number of messages this receiver can still read. + /// + /// This count is specific to this receiver, unlike [`Sender::buffer_len`], which reports the + /// shared backlog retained by the slowest active receiver. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// let (tx, mut rx) = unbounded::channel(); + /// assert_eq!(rx.len(), 0); + /// + /// tx.send(10); + /// tx.send(20); + /// assert_eq!(rx.len(), 2); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(rx.len(), 1); + /// ``` + pub fn len(&self) -> usize { + let inner = self.shared.inner.lock(); + let head = *inner + .receivers + .get(self.key) + .expect("active broadcast receiver must be registered"); + usize::try_from(inner.tail - head).expect("unread broadcast message count exceeds usize") + } + + /// Returns `true` if this receiver has no currently available messages. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// let (tx, rx) = unbounded::channel(); + /// assert!(rx.is_empty()); + /// + /// tx.send(10); + /// assert!(!rx.is_empty()); + /// ``` + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +struct Recv<'a, T> { + receiver: &'a mut Receiver, + registration: Option, +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + // Ready paths clear the registration, so only a cancelled pending receive takes this lock. + if self.registration.is_none() { + return; + } + + let waker = { + let mut waiters = self.receiver.shared.waiters.lock(); + waiters.unregister_waker(&mut self.registration) + }; + drop(waker); + } +} + +impl Future for Recv<'_, T> { + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let Self { + receiver, + registration, + } = self.get_mut(); + + match receiver.try_recv_shared() { + Ok((msg, reclaimed)) => { + *registration = None; + drop(reclaimed); + return Poll::Ready(Ok((*msg).clone())); + } + Err(TryRecvError::Disconnected) => { + *registration = None; + return Poll::Ready(Err(RecvError::Disconnected)); + } + Err(TryRecvError::Empty) => {} + } + + let received = { + let mut waiters = receiver.shared.waiters.lock(); + let mut inner = receiver.shared.inner.lock(); + + if let Some(received) = inner.receive(receiver.key) { + received + } else { + // A sender may have disconnected after the first `try_recv` returned `Empty`. + // Check again before registering the waker so `recv` does not miss the final wake. + if receiver.shared.senders.load(Ordering::Acquire) == 0 { + *registration = None; + return Poll::Ready(Err(RecvError::Disconnected)); + } + + // Register Waker + let waker = waiters.register_waker(registration, cx); + drop(waiters); + drop(inner); + drop(waker); + return Poll::Pending; + } + }; + + let (msg, reclaimed) = received; + *registration = None; + drop(reclaimed); + Poll::Ready(Ok((*msg).clone())) + } +} diff --git a/asyncband/src/broadcast/unbounded/tests.rs b/asyncband/src/broadcast/unbounded/tests.rs new file mode 100644 index 0000000..7c62175 --- /dev/null +++ b/asyncband/src/broadcast/unbounded/tests.rs @@ -0,0 +1,411 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Wake; +use std::task::Waker; + +use super::*; + +struct TrackWake(Arc); + +impl Wake for TrackWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +fn count_waker() -> (Waker, Arc) { + let count = Arc::new(AtomicUsize::new(0)); + let waker = Waker::from(Arc::new(TrackWake(count.clone()))); + (waker, count) +} + +#[tokio::test] +async fn test_broadcast_basic() { + let (tx, mut rx1) = channel(); + let mut rx2 = tx.subscribe(); + + tx.send(10); + tx.send(20); + + assert_eq!(rx1.recv().await, Ok(10)); + assert_eq!(rx1.recv().await, Ok(20)); + assert_eq!(rx2.recv().await, Ok(10)); + assert_eq!(rx2.recv().await, Ok(20)); +} + +#[tokio::test] +async fn test_broadcast_slow_receiver() { + let (tx, mut rx) = channel(); + + tx.send(1); + tx.send(2); + tx.send(3); + tx.send(4); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Ok(3)); + assert_eq!(rx.recv().await, Ok(4)); +} + +#[tokio::test] +async fn test_broadcast_closed() { + let (tx, mut rx) = channel::<()>(); + drop(tx); + assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); +} + +#[tokio::test] +async fn test_broadcast_closed_after_buffered_messages() { + let (tx, mut rx) = channel(); + + tx.send(1); + tx.send(2); + drop(tx); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); +} + +#[test] +fn test_wait_mechanism() { + let (tx, mut rx) = channel(); + let (waker, wake_count) = count_waker(); + let mut cx = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut cx).is_pending()); + + tx.send(42); + + assert_eq!(wake_count.load(Ordering::Relaxed), 1); + assert_eq!(recv.as_mut().poll(&mut cx), Poll::Ready(Ok(42))); +} + +#[tokio::test] +async fn test_recv_cancellation_removes_waiter() { + let (tx, mut rx) = channel::(); + let (waker, wake_count) = count_waker(); + let mut cx = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut cx).is_pending()); + + drop(recv); + tx.send(1); + + assert_eq!(wake_count.load(Ordering::Relaxed), 0); + assert_eq!(rx.try_recv(), Ok(1)); +} + +#[tokio::test] +async fn test_dropped_woken_recv_does_not_remove_reused_waiter() { + let (tx, mut rx1) = channel::(); + let mut rx2 = tx.subscribe(); + let (waker1, wake_count1) = count_waker(); + let mut cx1 = Context::from_waker(&waker1); + let mut recv1 = Box::pin(rx1.recv()); + + assert!(recv1.as_mut().poll(&mut cx1).is_pending()); + + tx.send(1); + assert_eq!(wake_count1.load(Ordering::Relaxed), 1); + assert_eq!(rx2.try_recv(), Ok(1)); + + let (waker2, wake_count2) = count_waker(); + let mut cx2 = Context::from_waker(&waker2); + let mut recv2 = Box::pin(rx2.recv()); + assert!(recv2.as_mut().poll(&mut cx2).is_pending()); + + drop(recv1); + tx.send(2); + + assert_eq!(wake_count2.load(Ordering::Relaxed), 1); + drop(recv2); +} + +#[tokio::test] +async fn test_subscribe() { + let (tx, _rx) = channel(); + let mut rx = tx.subscribe(); + + tx.send(100); + assert_eq!(rx.recv().await, Ok(100)); +} + +#[tokio::test] +async fn test_receiver_count_and_len() { + let (tx, mut rx1) = channel(); + assert_eq!(tx.receiver_count(), 1); + assert_eq!(rx1.len(), 0); + assert!(rx1.is_empty()); + + tx.send(1); + tx.send(2); + assert_eq!(rx1.len(), 2); + assert!(!rx1.is_empty()); + + let mut rx2 = tx.subscribe(); + assert_eq!(tx.receiver_count(), 2); + assert_eq!(rx2.len(), 0); + assert!(rx2.is_empty()); + + tx.send(3); + assert_eq!(rx1.len(), 3); + assert_eq!(rx2.len(), 1); + + assert_eq!(rx2.try_recv(), Ok(3)); + assert_eq!(rx2.len(), 0); + drop(rx2); + assert_eq!(tx.receiver_count(), 1); + + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.len(), 2); +} + +#[tokio::test] +async fn test_resubscribe() { + let (tx, mut rx) = channel(); + + tx.send(1); + tx.send(2); + + let mut rx2 = rx.resubscribe(); + + // rx sees 1, 2 + // rx2 sees nothing yet (starts at tail=2) + + tx.send(3); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx2.recv().await, Ok(3)); +} + +#[tokio::test] +async fn test_try_recv() { + let (tx, mut rx) = channel(); + + // Empty + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + // Success + tx.send(10); + assert_eq!(rx.try_recv(), Ok(10)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + // Closed + drop(tx); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[tokio::test] +async fn test_consumed_messages_are_reclaimed() { + let (tx, mut rx1) = channel(); + let mut rx2 = tx.subscribe(); + + tx.send(1); + tx.send(2); + assert_eq!(tx.buffer_len(), 2); + + assert_eq!(rx1.recv().await, Ok(1)); + assert_eq!(tx.buffer_len(), 2); + + assert_eq!(rx2.recv().await, Ok(1)); + assert_eq!(tx.buffer_len(), 1); + + assert_eq!(rx1.recv().await, Ok(2)); + assert_eq!(tx.buffer_len(), 1); + + assert_eq!(rx2.recv().await, Ok(2)); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn test_drop_receiver_reclaims_messages() { + let (tx, mut rx1) = channel(); + let rx2 = tx.subscribe(); + + tx.send(1); + tx.send(2); + tx.send(3); + + assert_eq!(rx1.recv().await, Ok(1)); + assert_eq!(rx1.recv().await, Ok(2)); + assert_eq!(rx1.recv().await, Ok(3)); + assert_eq!(tx.buffer_len(), 3); + + drop(rx2); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn test_buffer_len_tracks_shared_backlog() { + let (tx, mut rx1) = channel(); + let mut rx2 = tx.subscribe(); + + tx.send(1); + tx.send(2); + assert_eq!(tx.buffer_len(), 2); + + assert_eq!(rx1.recv().await, Ok(1)); + assert_eq!(rx1.recv().await, Ok(2)); + assert_eq!(tx.buffer_len(), 2); + + assert_eq!(rx2.recv().await, Ok(1)); + assert_eq!(tx.buffer_len(), 1); + assert_eq!(rx2.recv().await, Ok(2)); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn test_resubscribe_keeps_original_receiver_backlog() { + let (tx, mut rx) = channel(); + + tx.send(1); + tx.send(2); + + let mut rx2 = rx.resubscribe(); + assert_eq!(tx.buffer_len(), 2); + + tx.send(3); + + assert_eq!(rx2.recv().await, Ok(3)); + assert_eq!(tx.buffer_len(), 3); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Ok(3)); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +#[should_panic(expected = "broadcast channel version counter overflowed")] +async fn test_send_panics_on_version_overflow() { + let (tx, _) = channel(); + tx.shared.inner.lock().tail = u64::MAX; + tx.send(()); +} + +#[tokio::test] +async fn test_send_without_receivers_does_not_buffer() { + let (tx, rx) = channel(); + drop(rx); + + tx.send(1); + tx.send(2); + assert_eq!(tx.buffer_len(), 0); + + let mut rx = tx.subscribe(); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + tx.send(3); + assert_eq!(rx.recv().await, Ok(3)); +} + +#[tokio::test] +async fn test_multi_senders_concurrent() { + let (tx, mut rx) = channel(); + let tx1 = tx.clone(); + let tx2 = tx.clone(); + + let handle1 = tokio::spawn(async move { + for i in 0..10 { + tx1.send(i); + } + }); + + let handle2 = tokio::spawn(async move { + for i in 10..20 { + tx2.send(i); + } + }); + + // Main tx can also send + for i in 20..30 { + tx.send(i); + } + + handle1.await.unwrap(); + handle2.await.unwrap(); + drop(tx); + + let mut received = Vec::new(); + while let Ok(n) = rx.recv().await { + received.push(n); + } + received.sort(); + + let expected = (0..30).collect::>(); + assert_eq!(received, expected); +} + +#[tokio::test] +async fn test_multi_senders_multiple_receivers_receive_all() { + let (tx, mut rx1) = channel(); + let mut rx2 = tx.subscribe(); + let mut rx3 = tx.subscribe(); + let tx1 = tx.clone(); + let tx2 = tx.clone(); + + let handle1 = tokio::spawn(async move { + for i in 0..10 { + tx1.send(i); + } + }); + let handle2 = tokio::spawn(async move { + for i in 10..20 { + tx2.send(i); + } + }); + + for i in 20..30 { + tx.send(i); + } + + handle1.await.unwrap(); + handle2.await.unwrap(); + drop(tx); + + let received1 = drain(&mut rx1).await; + let received2 = drain(&mut rx2).await; + let received3 = drain(&mut rx3).await; + + assert_eq!(received1, received2); + assert_eq!(received1, received3); + + let mut sorted = received1; + sorted.sort(); + let expected = (0..30).collect::>(); + assert_eq!(sorted, expected); +} + +async fn drain(rx: &mut Receiver) -> Vec { + let mut received = Vec::new(); + while let Ok(n) = rx.recv().await { + received.push(n); + } + received +} diff --git a/asyncband/src/internal/arena.rs b/asyncband/src/internal/arena.rs index d35ca30..7c9ff2f 100644 --- a/asyncband/src/internal/arena.rs +++ b/asyncband/src/internal/arena.rs @@ -117,6 +117,21 @@ impl Arena { } } + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + pub fn values(&self) -> impl Iterator { + self.slots.iter().filter_map(|slot| match slot { + Slot::Occupied(value) => Some(value), + Slot::Vacant(_) => None, + }) + } + pub fn remove(&mut self, key: ArenaKey) -> T { let index = key.0; let slot = self @@ -160,11 +175,6 @@ impl Arena { self.len = 0; values } - - #[cfg(test)] - pub fn len(&self) -> usize { - self.len - } } #[cfg(test)] diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 3d810e9..e8ffdc8 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -56,7 +56,7 @@ //! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | //! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::OnceMap`] | `once`, `once-cell`, `once-map` | //! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | -//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::overflow`] | `oneshot`, `mpsc`, `broadcast` | +//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::overflow`], [`broadcast::unbounded`] | `oneshot`, `mpsc`, `broadcast` | //! | Control workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | //! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | //! diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index e894246..8f06027 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -64,6 +64,10 @@ fn public_types_are_send_and_sync() { assert_send_and_sync::>(); assert_send_and_sync::(); assert_send_and_sync::(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::(); + assert_send_and_sync::(); assert_send_and_sync::>(); assert_send_and_sync::>(); assert_send_and_sync::>(); @@ -108,6 +112,10 @@ fn public_types_are_unpin() { assert_unpin::>(); assert_unpin::(); assert_unpin::(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::(); + assert_unpin::(); assert_unpin::>(); assert_unpin::>(); assert_unpin::>(); From fcd4f18b3eaa099d95d2290686decf58c25e1289 Mon Sep 17 00:00:00 2001 From: Orthur Date: Sat, 22 Aug 2026 11:13:42 -0400 Subject: [PATCH 2/3] perf(broadcast): serve unbounded receives from one critical section --- README.md | 36 ++-- asyncband/src/broadcast/unbounded/mod.rs | 232 +++++++++++++++++------ asyncband/src/lib.rs | 14 +- 3 files changed, 202 insertions(+), 80 deletions(-) diff --git a/README.md b/README.md index ff38a46..0201b9b 100644 --- a/README.md +++ b/README.md @@ -32,25 +32,25 @@ Asyncband is a runtime-agnostic library providing essential synchronization prim The crate enables no primitives by default. Categories describe each primitive's primary purpose and do not add another module level, so public paths remain concise, such as `asyncband::mutex` and `asyncband::once::OnceCell`. -| Category | Primitive | Feature | Purpose | -| ----------------------- | ------------------------------------------------------------------------------------ | -------------- | ----------------------------------------------------------------------- | -| Shared state | [`Mutex`](https://docs.rs/asyncband/*/asyncband/mutex/struct.Mutex.html) | `mutex` | Protect shared data with asynchronous mutual exclusion. | -| | [`RwLock`](https://docs.rs/asyncband/*/asyncband/rwlock/struct.RwLock.html) | `rwlock` | Allow multiple readers or one writer. | -| | [`Condvar`](https://docs.rs/asyncband/*/asyncband/condvar/struct.Condvar.html) | `condvar` | Wait for notifications while releasing a mutex. | -| One-time initialization | [`Once`](https://docs.rs/asyncband/*/asyncband/once/struct.Once.html) | `once` | Run asynchronous initialization exactly once. | -| | [`OnceCell`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceCell.html) | `once-cell` | Initialize and store one asynchronous value. | -| | [`OnceMap`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceMap.html) | `once-map` | Initialize and store one value per key. | -| Task coordination | [`Barrier`](https://docs.rs/asyncband/*/asyncband/barrier/struct.Barrier.html) | `barrier` | Wait until all participants reach a synchronization point. | -| | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a one-way countdown completes. | -| | [`WaitGroup`](https://docs.rs/asyncband/*/asyncband/waitgroup/struct.WaitGroup.html) | `waitgroup` | Wait for a dynamic group of tasks to finish. | -| | [`shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/) | `shutdown` | Coordinate shutdown signals and completion. | -| Channels | [`oneshot::channel`](https://docs.rs/asyncband/*/asyncband/oneshot/fn.channel.html) | `oneshot` | Send one value between two tasks. | -| | [`mpsc::bounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.bounded.html) | `mpsc` | Send values from multiple producers through a bounded channel. | -| | [`mpsc::unbounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.unbounded.html) | `mpsc` | Send values from multiple producers through an unbounded channel. | -| | [`broadcast::overflow`](https://docs.rs/asyncband/*/asyncband/broadcast/overflow/) | `broadcast` | Broadcast values and report when slow receivers miss overwritten items. | +| Category | Primitive | Feature | Purpose | +| ----------------------- | ------------------------------------------------------------------------------------ | -------------- | --------------------------------------------------------------------------- | +| Shared state | [`Mutex`](https://docs.rs/asyncband/*/asyncband/mutex/struct.Mutex.html) | `mutex` | Protect shared data with asynchronous mutual exclusion. | +| | [`RwLock`](https://docs.rs/asyncband/*/asyncband/rwlock/struct.RwLock.html) | `rwlock` | Allow multiple readers or one writer. | +| | [`Condvar`](https://docs.rs/asyncband/*/asyncband/condvar/struct.Condvar.html) | `condvar` | Wait for notifications while releasing a mutex. | +| One-time initialization | [`Once`](https://docs.rs/asyncband/*/asyncband/once/struct.Once.html) | `once` | Run asynchronous initialization exactly once. | +| | [`OnceCell`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceCell.html) | `once-cell` | Initialize and store one asynchronous value. | +| | [`OnceMap`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceMap.html) | `once-map` | Initialize and store one value per key. | +| Task coordination | [`Barrier`](https://docs.rs/asyncband/*/asyncband/barrier/struct.Barrier.html) | `barrier` | Wait until all participants reach a synchronization point. | +| | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a one-way countdown completes. | +| | [`WaitGroup`](https://docs.rs/asyncband/*/asyncband/waitgroup/struct.WaitGroup.html) | `waitgroup` | Wait for a dynamic group of tasks to finish. | +| | [`shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/) | `shutdown` | Coordinate shutdown signals and completion. | +| Channels | [`oneshot::channel`](https://docs.rs/asyncband/*/asyncband/oneshot/fn.channel.html) | `oneshot` | Send one value between two tasks. | +| | [`mpsc::bounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.bounded.html) | `mpsc` | Send values from multiple producers through a bounded channel. | +| | [`mpsc::unbounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.unbounded.html) | `mpsc` | Send values from multiple producers through an unbounded channel. | +| | [`broadcast::overflow`](https://docs.rs/asyncband/*/asyncband/broadcast/overflow/) | `broadcast` | Broadcast values and report when slow receivers miss overwritten items. | | | [`broadcast::unbounded`](https://docs.rs/asyncband/*/asyncband/broadcast/unbounded/) | `broadcast` | Broadcast values and retain them until every active receiver consumes them. | -| Workload control | [`Semaphore`](https://docs.rs/asyncband/*/asyncband/semaphore/struct.Semaphore.html) | `semaphore` | Control concurrent access with permits. | -| | [`Group`](https://docs.rs/asyncband/*/asyncband/singleflight/struct.Group.html) | `singleflight` | Coalesce concurrent calls for the same key. | +| Workload control | [`Semaphore`](https://docs.rs/asyncband/*/asyncband/semaphore/struct.Semaphore.html) | `semaphore` | Control concurrent access with permits. | +| | [`Group`](https://docs.rs/asyncband/*/asyncband/singleflight/struct.Group.html) | `singleflight` | Coalesce concurrent calls for the same key. | ## Installation diff --git a/asyncband/src/broadcast/unbounded/mod.rs b/asyncband/src/broadcast/unbounded/mod.rs index 0a2d21d..7906487 100644 --- a/asyncband/src/broadcast/unbounded/mod.rs +++ b/asyncband/src/broadcast/unbounded/mod.rs @@ -28,11 +28,21 @@ //! consumed them or the receiver is dropped. Use [`Sender::buffer_len`] to monitor the number of //! messages currently retained by the shared buffer. //! +//! The buffer keeps the capacity a steady workload needs, so a channel that repeatedly fills and +//! drains does not reallocate. Capacity grown for a one-off burst is released once a later cycle +//! drains completely without needing it. +//! //! # Receivers //! //! Each receiver has an independent cursor. Use [`Sender::subscribe`] to create a receiver that -//! starts at the current tail of the channel, or [`Receiver::resubscribe`] to skip this receiver's -//! backlog and start a new receiver at the current tail. +//! starts at the current tail of the channel, [`Receiver::clone`] to create one that shares this +//! receiver's unread backlog, or [`Receiver::resubscribe`] to skip this receiver's backlog and +//! start a new receiver at the current tail. +//! +//! Messages are reclaimed once the slowest receiver moves past them, which scans one slot per +//! receiver. Only the receive that advances the slowest cursor pays for that scan, and the channel +//! keeps a slot for every receiver it hands out, so the cost follows the largest number of +//! receivers that were ever active at once rather than the number active now. //! //! # Examples //! @@ -63,21 +73,27 @@ //! //! # #[tokio::main] //! # async fn main() { -//! let (tx, mut rx) = unbounded::channel(); +//! let (tx, mut rx1) = unbounded::channel(); +//! let mut rx2 = tx.subscribe(); //! //! tx.send(1); //! tx.send(2); -//! tx.send(3); //! -//! assert_eq!(rx.recv().await, Ok(1)); -//! assert_eq!(rx.recv().await, Ok(2)); -//! assert_eq!(rx.recv().await, Ok(3)); +//! // One receiver draining the channel does not discard what the other has not read yet. +//! assert_eq!(rx1.recv().await, Ok(1)); +//! assert_eq!(rx1.recv().await, Ok(2)); +//! assert_eq!(tx.buffer_len(), 2); +//! +//! assert_eq!(rx2.recv().await, Ok(1)); +//! assert_eq!(rx2.recv().await, Ok(2)); +//! assert_eq!(tx.buffer_len(), 0); //! # } //! ``` use std::collections::VecDeque; use std::fmt; use std::future::Future; +use std::mem; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicUsize; @@ -117,9 +133,10 @@ pub fn channel() -> (Sender, Receiver) { head_receivers: 1, tail: 0, receivers, + peak_len: 0, + waiters: WaitSet::new(), }), senders: AtomicUsize::new(1), - waiters: Mutex::new(WaitSet::new()), }); let sender = Sender { shared: shared.clone(), @@ -166,8 +183,16 @@ impl fmt::Display for TryRecvError { impl std::error::Error for TryRecvError {} +/// Retained capacity below which the shared buffer is never shrunk back. +const MIN_RETAINED_CAPACITY: usize = 64; + struct Inner { /// Messages whose versions are in the range `[head, tail)`. + /// + /// Each message is held behind an `Arc` so a receive can hand the payload out of the critical + /// section. Cloning the `Arc` under the lock keeps `T::clone` — and, for reclaimed messages, + /// `T::drop` — outside it, which matters because both are arbitrary user code that may call + /// back into this channel. buffer: VecDeque>, /// The version of the first message in `buffer`. head: u64, @@ -177,6 +202,10 @@ struct Inner { tail: u64, /// Cursor for each active receiver. receivers: Arena, + /// The largest backlog retained since the buffer was last empty. + peak_len: usize, + /// Receivers parked in [`Receiver::recv`]. + waiters: WaitSet, } impl Inner { @@ -236,6 +265,14 @@ impl Inner { let offset = (head - self.head) as usize; let msg = self.buffer[offset].clone(); let reclaimed = self.advance_receiver(key, head + 1); + // A reclaim triggered by this receive always begins with this receiver's own message: + // the reclaim path runs only for a cursor sitting at `head`, so the first slot drained + // is `msg`. `take_msg` relies on this to recognise that it owns the payload. + debug_assert!( + reclaimed + .first() + .is_none_or(|first| Arc::ptr_eq(first, &msg)) + ); Some((msg, reclaimed)) } else { None @@ -263,16 +300,41 @@ impl Inner { self.head = next_head; self.head_receivers = head_receivers; + self.shrink_buffer(); reclaimed } + + /// Returns the allocation grown for a stalled receiver once that backlog is behind us. + /// + /// Without this, a single burst pins its peak allocation for the lifetime of the channel. + /// The decision is deliberately made only when the buffer drains completely, and against the + /// peak of the cycle that just ended rather than the current length: a channel that repeatedly + /// fills and drains keeps a peak as large as its bursts, so it holds its allocation instead of + /// reallocating on every cycle. Only once a full cycle stays small does the buffer give the + /// memory back. + fn shrink_buffer(&mut self) { + if !self.buffer.is_empty() { + return; + } + + let peak = mem::take(&mut self.peak_len); + let capacity = self.buffer.capacity(); + if capacity > MIN_RETAINED_CAPACITY && peak <= capacity / 4 { + self.buffer.shrink_to(MIN_RETAINED_CAPACITY.max(peak * 2)); + } + } } struct Shared { + /// Buffer, receiver cursors, and parked receivers, all under a single lock. + /// + /// The wait set lives here rather than beside it so that publishing a message and draining the + /// waiters happen in one critical section. That is what makes the park path race-free: a + /// receiver that finds no message and then registers still holds this lock, so a concurrent + /// `send` cannot slip between the two steps and skip the wake-up. inner: Mutex>, /// Number of active senders. senders: AtomicUsize, - /// Waiters (receivers) waiting for new messages. - waiters: Mutex, } /// A sender handle to the broadcast channel. @@ -285,6 +347,9 @@ pub struct Sender { impl Clone for Sender { fn clone(&self) -> Self { + // Relaxed is enough because this count publishes nothing on its own: receivers read it + // only to decide whether the channel is closed, and every message it could hide is + // published under `inner`, which a receiver holds before it observes the count. self.shared.senders.fetch_add(1, Ordering::Relaxed); Self { shared: self.shared.clone(), @@ -292,16 +357,19 @@ impl Clone for Sender { } } +impl fmt::Debug for Sender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Sender").finish_non_exhaustive() + } +} + impl Drop for Sender { fn drop(&mut self) { match self.shared.senders.fetch_sub(1, Ordering::AcqRel) { 1 => { // If this is the last sender, we need to wake up the receiver so it can // observe the disconnected state. - let wakers = { - let mut waiters = self.shared.waiters.lock(); - waiters.take_wakers() - }; + let wakers = self.shared.inner.lock().waiters.take_wakers(); for waker in wakers { waker.wake(); } @@ -339,7 +407,9 @@ impl Sender { pub fn send(&self, msg: T) { let msg = Arc::new(msg); - { + // Publishing and draining the wait set share one critical section, so a receiver can never + // observe an empty buffer and park after this message became visible. + let wakers = { let mut inner = self.shared.inner.lock(); inner.tail = inner .tail @@ -356,14 +426,14 @@ impl Sender { inner.head = inner.tail; } else { inner.buffer.push_back(msg); + inner.peak_len = inner.peak_len.max(inner.buffer.len()); } - } - // Notify all waiting receivers. - let wakers = { - let mut waiters = self.shared.waiters.lock(); - waiters.take_wakers() + inner.waiters.take_wakers() }; + + // Notify all waiting receivers. An unsent message is dropped here too, once the lock is + // released. for waker in wakers { waker.wake(); } @@ -442,11 +512,54 @@ impl Sender { /// A receiver handle to the broadcast channel. /// /// Each receiver sees every message sent to the channel while the receiver is active. +/// +/// Cloning a receiver creates one that shares this receiver's unread backlog, while +/// [`Receiver::resubscribe`] creates one that starts at the current tail instead. pub struct Receiver { shared: Arc>, key: ArenaKey, } +impl Clone for Receiver { + /// Creates a receiver that starts from this receiver's current position. + /// + /// The clone reads this receiver's unread backlog and every later message. Use + /// [`Receiver::resubscribe`] instead to start at the current tail and skip the backlog. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::unbounded; + /// + /// let (tx, mut rx) = unbounded::channel(); + /// tx.send(1); + /// + /// let mut clone = rx.clone(); + /// assert_eq!(rx.try_recv(), Ok(1)); + /// assert_eq!(clone.try_recv(), Ok(1)); + /// ``` + fn clone(&self) -> Self { + let key = { + let mut inner = self.shared.inner.lock(); + let head = *inner + .receivers + .get(self.key) + .expect("active broadcast receiver must be registered"); + inner.insert_receiver(head) + }; + Self { + shared: self.shared.clone(), + key, + } + } +} + +impl fmt::Debug for Receiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Receiver").finish_non_exhaustive() + } +} + impl Drop for Receiver { fn drop(&mut self) { let reclaimed = { @@ -512,8 +625,32 @@ impl Receiver { /// ``` pub fn try_recv(&mut self) -> Result { let (msg, reclaimed) = self.try_recv_shared()?; - drop(reclaimed); - Ok((*msg).clone()) + Ok(take_msg(msg, reclaimed)) + } +} + +/// Drops the reclaimed backlog, then yields the received message, both with the channel unlocked. +/// +/// A non-empty backlog means this receive drained `msg` from the buffer, so once the backlog is +/// dropped this receive holds the only reference and the payload can be moved out instead of +/// cloned. A channel with a single receiver therefore never clones a payload. +/// +/// Ownership is decided from that bookkeeping rather than by probing the reference count. An +/// [`Arc::try_unwrap`] on every receive would fail under fan-out, and its failed compare-exchange +/// writes to a cache line that every receiver draining the message shares. +fn take_msg(msg: Arc, reclaimed: Vec>) -> T { + let sole_owner = !reclaimed.is_empty(); + drop(reclaimed); + + if !sole_owner { + return (*msg).clone(); + } + + // Another receiver can still hold an in-flight reference to the same message, so the clone + // remains the fallback. + match Arc::try_unwrap(msg) { + Ok(msg) => msg, + Err(msg) => (*msg).clone(), } } @@ -623,8 +760,8 @@ impl Drop for Recv<'_, T> { } let waker = { - let mut waiters = self.receiver.shared.waiters.lock(); - waiters.unregister_waker(&mut self.registration) + let mut inner = self.receiver.shared.inner.lock(); + inner.waiters.unregister_waker(&mut self.registration) }; drop(waker); } @@ -639,45 +776,30 @@ impl Future for Recv<'_, T> { registration, } = self.get_mut(); - match receiver.try_recv_shared() { - Ok((msg, reclaimed)) => { - *registration = None; - drop(reclaimed); - return Poll::Ready(Ok((*msg).clone())); - } - Err(TryRecvError::Disconnected) => { - *registration = None; - return Poll::Ready(Err(RecvError::Disconnected)); - } - Err(TryRecvError::Empty) => {} - } - + // One critical section decides between all three outcomes. Senders append messages and + // drain the wait set under this same lock, so registering here cannot miss a wake-up and + // cannot observe a closed channel that still has a message for this receiver. let received = { - let mut waiters = receiver.shared.waiters.lock(); let mut inner = receiver.shared.inner.lock(); - if let Some(received) = inner.receive(receiver.key) { - received - } else { - // A sender may have disconnected after the first `try_recv` returned `Empty`. - // Check again before registering the waker so `recv` does not miss the final wake. - if receiver.shared.senders.load(Ordering::Acquire) == 0 { - *registration = None; - return Poll::Ready(Err(RecvError::Disconnected)); + match inner.receive(receiver.key) { + Some(received) => received, + None => { + if receiver.shared.senders.load(Ordering::Acquire) == 0 { + *registration = None; + return Poll::Ready(Err(RecvError::Disconnected)); + } + + let waker = inner.waiters.register_waker(registration, cx); + drop(inner); + drop(waker); + return Poll::Pending; } - - // Register Waker - let waker = waiters.register_waker(registration, cx); - drop(waiters); - drop(inner); - drop(waker); - return Poll::Pending; } }; let (msg, reclaimed) = received; *registration = None; - drop(reclaimed); - Poll::Ready(Ok((*msg).clone())) + Poll::Ready(Ok(take_msg(msg, reclaimed))) } } diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index e8ffdc8..0555ec1 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -51,14 +51,14 @@ //! //! # API guide //! -//! | Use case | APIs | Cargo features | -//! | -------------------------- | ------------------------------------------------------------------------------------- | ------------------------------------------- | -//! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | -//! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::OnceMap`] | `once`, `once-cell`, `once-map` | -//! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | +//! | Use case | APIs | Cargo features | +//! | -------------------------- | --------------------------------------------------------------------------------------------------------------- | ------------------------------------------- | +//! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | +//! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::OnceMap`] | `once`, `once-cell`, `once-map` | +//! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | //! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::overflow`], [`broadcast::unbounded`] | `oneshot`, `mpsc`, `broadcast` | -//! | Control workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | -//! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | +//! | Control workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | +//! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | //! //! # Runtime and blocking model //! From 54355f843779809f08a9301e71655e49c17e83a8 Mon Sep 17 00:00:00 2001 From: Orthur Date: Sat, 22 Aug 2026 11:13:42 -0400 Subject: [PATCH 3/3] test(broadcast): move unbounded coverage into integration tests --- asyncband/src/broadcast/unbounded/tests.rs | 398 +---------- benchmarks/broadcast_unbounded.rs | 356 ++++++++++ benchmarks/main.rs | 1 + .../tests/broadcast_unbounded_test.rs | 654 ++++++++++++++++++ 4 files changed, 1039 insertions(+), 370 deletions(-) create mode 100644 benchmarks/broadcast_unbounded.rs create mode 100644 tests-integration/tests/broadcast_unbounded_test.rs diff --git a/asyncband/src/broadcast/unbounded/tests.rs b/asyncband/src/broadcast/unbounded/tests.rs index 7c62175..a80a527 100644 --- a/asyncband/src/broadcast/unbounded/tests.rs +++ b/asyncband/src/broadcast/unbounded/tests.rs @@ -15,397 +15,55 @@ // specific language governing permissions and limitations // under the License. -use std::future::Future; -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use std::task::Wake; -use std::task::Waker; - use super::*; -struct TrackWake(Arc); - -impl Wake for TrackWake { - fn wake(self: Arc) { - self.0.fetch_add(1, Ordering::Relaxed); - } -} - -fn count_waker() -> (Waker, Arc) { - let count = Arc::new(AtomicUsize::new(0)); - let waker = Waker::from(Arc::new(TrackWake(count.clone()))); - (waker, count) -} - -#[tokio::test] -async fn test_broadcast_basic() { - let (tx, mut rx1) = channel(); - let mut rx2 = tx.subscribe(); - - tx.send(10); - tx.send(20); - - assert_eq!(rx1.recv().await, Ok(10)); - assert_eq!(rx1.recv().await, Ok(20)); - assert_eq!(rx2.recv().await, Ok(10)); - assert_eq!(rx2.recv().await, Ok(20)); -} - -#[tokio::test] -async fn test_broadcast_slow_receiver() { - let (tx, mut rx) = channel(); - - tx.send(1); - tx.send(2); - tx.send(3); - tx.send(4); - - assert_eq!(rx.recv().await, Ok(1)); - assert_eq!(rx.recv().await, Ok(2)); - assert_eq!(rx.recv().await, Ok(3)); - assert_eq!(rx.recv().await, Ok(4)); -} - -#[tokio::test] -async fn test_broadcast_closed() { - let (tx, mut rx) = channel::<()>(); - drop(tx); - assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); -} - -#[tokio::test] -async fn test_broadcast_closed_after_buffered_messages() { - let (tx, mut rx) = channel(); - - tx.send(1); - tx.send(2); - drop(tx); - - assert_eq!(rx.recv().await, Ok(1)); - assert_eq!(rx.recv().await, Ok(2)); - assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); -} - #[test] -fn test_wait_mechanism() { - let (tx, mut rx) = channel(); - let (waker, wake_count) = count_waker(); - let mut cx = Context::from_waker(&waker); - let mut recv = Box::pin(rx.recv()); - - assert!(recv.as_mut().poll(&mut cx).is_pending()); - - tx.send(42); - - assert_eq!(wake_count.load(Ordering::Relaxed), 1); - assert_eq!(recv.as_mut().poll(&mut cx), Poll::Ready(Ok(42))); -} - -#[tokio::test] -async fn test_recv_cancellation_removes_waiter() { - let (tx, mut rx) = channel::(); - let (waker, wake_count) = count_waker(); - let mut cx = Context::from_waker(&waker); - let mut recv = Box::pin(rx.recv()); - - assert!(recv.as_mut().poll(&mut cx).is_pending()); - - drop(recv); - tx.send(1); - - assert_eq!(wake_count.load(Ordering::Relaxed), 0); - assert_eq!(rx.try_recv(), Ok(1)); -} - -#[tokio::test] -async fn test_dropped_woken_recv_does_not_remove_reused_waiter() { - let (tx, mut rx1) = channel::(); - let mut rx2 = tx.subscribe(); - let (waker1, wake_count1) = count_waker(); - let mut cx1 = Context::from_waker(&waker1); - let mut recv1 = Box::pin(rx1.recv()); - - assert!(recv1.as_mut().poll(&mut cx1).is_pending()); - - tx.send(1); - assert_eq!(wake_count1.load(Ordering::Relaxed), 1); - assert_eq!(rx2.try_recv(), Ok(1)); - - let (waker2, wake_count2) = count_waker(); - let mut cx2 = Context::from_waker(&waker2); - let mut recv2 = Box::pin(rx2.recv()); - assert!(recv2.as_mut().poll(&mut cx2).is_pending()); - - drop(recv1); - tx.send(2); - - assert_eq!(wake_count2.load(Ordering::Relaxed), 1); - drop(recv2); -} - -#[tokio::test] -async fn test_subscribe() { - let (tx, _rx) = channel(); - let mut rx = tx.subscribe(); - - tx.send(100); - assert_eq!(rx.recv().await, Ok(100)); -} - -#[tokio::test] -async fn test_receiver_count_and_len() { - let (tx, mut rx1) = channel(); - assert_eq!(tx.receiver_count(), 1); - assert_eq!(rx1.len(), 0); - assert!(rx1.is_empty()); - - tx.send(1); - tx.send(2); - assert_eq!(rx1.len(), 2); - assert!(!rx1.is_empty()); - - let mut rx2 = tx.subscribe(); - assert_eq!(tx.receiver_count(), 2); - assert_eq!(rx2.len(), 0); - assert!(rx2.is_empty()); - - tx.send(3); - assert_eq!(rx1.len(), 3); - assert_eq!(rx2.len(), 1); - - assert_eq!(rx2.try_recv(), Ok(3)); - assert_eq!(rx2.len(), 0); - drop(rx2); - assert_eq!(tx.receiver_count(), 1); - - assert_eq!(rx1.try_recv(), Ok(1)); - assert_eq!(rx1.len(), 2); -} - -#[tokio::test] -async fn test_resubscribe() { - let (tx, mut rx) = channel(); - - tx.send(1); - tx.send(2); - - let mut rx2 = rx.resubscribe(); - - // rx sees 1, 2 - // rx2 sees nothing yet (starts at tail=2) - - tx.send(3); - - assert_eq!(rx.recv().await, Ok(1)); - assert_eq!(rx.recv().await, Ok(2)); - assert_eq!(rx2.recv().await, Ok(3)); -} - -#[tokio::test] -async fn test_try_recv() { - let (tx, mut rx) = channel(); - - // Empty - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - - // Success - tx.send(10); - assert_eq!(rx.try_recv(), Ok(10)); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - - // Closed - drop(tx); - assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); -} - -#[tokio::test] -async fn test_consumed_messages_are_reclaimed() { - let (tx, mut rx1) = channel(); - let mut rx2 = tx.subscribe(); - - tx.send(1); - tx.send(2); - assert_eq!(tx.buffer_len(), 2); - - assert_eq!(rx1.recv().await, Ok(1)); - assert_eq!(tx.buffer_len(), 2); - - assert_eq!(rx2.recv().await, Ok(1)); - assert_eq!(tx.buffer_len(), 1); - - assert_eq!(rx1.recv().await, Ok(2)); - assert_eq!(tx.buffer_len(), 1); - - assert_eq!(rx2.recv().await, Ok(2)); - assert_eq!(tx.buffer_len(), 0); -} - -#[tokio::test] -async fn test_drop_receiver_reclaims_messages() { - let (tx, mut rx1) = channel(); - let rx2 = tx.subscribe(); - - tx.send(1); - tx.send(2); - tx.send(3); - - assert_eq!(rx1.recv().await, Ok(1)); - assert_eq!(rx1.recv().await, Ok(2)); - assert_eq!(rx1.recv().await, Ok(3)); - assert_eq!(tx.buffer_len(), 3); - - drop(rx2); - assert_eq!(tx.buffer_len(), 0); -} - -#[tokio::test] -async fn test_buffer_len_tracks_shared_backlog() { - let (tx, mut rx1) = channel(); - let mut rx2 = tx.subscribe(); - - tx.send(1); - tx.send(2); - assert_eq!(tx.buffer_len(), 2); - - assert_eq!(rx1.recv().await, Ok(1)); - assert_eq!(rx1.recv().await, Ok(2)); - assert_eq!(tx.buffer_len(), 2); - - assert_eq!(rx2.recv().await, Ok(1)); - assert_eq!(tx.buffer_len(), 1); - assert_eq!(rx2.recv().await, Ok(2)); - assert_eq!(tx.buffer_len(), 0); -} - -#[tokio::test] -async fn test_resubscribe_keeps_original_receiver_backlog() { - let (tx, mut rx) = channel(); - - tx.send(1); - tx.send(2); - - let mut rx2 = rx.resubscribe(); - assert_eq!(tx.buffer_len(), 2); - - tx.send(3); - - assert_eq!(rx2.recv().await, Ok(3)); - assert_eq!(tx.buffer_len(), 3); - - assert_eq!(rx.recv().await, Ok(1)); - assert_eq!(rx.recv().await, Ok(2)); - assert_eq!(rx.recv().await, Ok(3)); - assert_eq!(tx.buffer_len(), 0); -} - -#[tokio::test] #[should_panic(expected = "broadcast channel version counter overflowed")] -async fn test_send_panics_on_version_overflow() { +fn send_panics_on_version_overflow() { + // The receiver is dropped right away: the doctored counter would make its own drop overflow. let (tx, _) = channel(); tx.shared.inner.lock().tail = u64::MAX; tx.send(()); } -#[tokio::test] -async fn test_send_without_receivers_does_not_buffer() { - let (tx, rx) = channel(); - drop(rx); - - tx.send(1); - tx.send(2); - assert_eq!(tx.buffer_len(), 0); - - let mut rx = tx.subscribe(); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - - tx.send(3); - assert_eq!(rx.recv().await, Ok(3)); -} - -#[tokio::test] -async fn test_multi_senders_concurrent() { +#[test] +fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { let (tx, mut rx) = channel(); - let tx1 = tx.clone(); - let tx2 = tx.clone(); - let handle1 = tokio::spawn(async move { - for i in 0..10 { - tx1.send(i); - } - }); - - let handle2 = tokio::spawn(async move { - for i in 10..20 { - tx2.send(i); - } - }); - - // Main tx can also send - for i in 20..30 { + let burst = MIN_RETAINED_CAPACITY * 16; + for i in 0..burst { tx.send(i); } + assert!(tx.shared.inner.lock().buffer.capacity() >= burst); - handle1.await.unwrap(); - handle2.await.unwrap(); - drop(tx); - - let mut received = Vec::new(); - while let Ok(n) = rx.recv().await { - received.push(n); + for i in 0..burst { + assert_eq!(rx.try_recv(), Ok(i)); } - received.sort(); - let expected = (0..30).collect::>(); - assert_eq!(received, expected); + // Draining evaluates the cycle that just peaked, so the burst allocation is still held. + assert_eq!(tx.buffer_len(), 0); + assert!(tx.shared.inner.lock().buffer.capacity() >= burst); + + // The next cycle stays small, which is what releases the memory. + tx.send(0); + assert_eq!(rx.try_recv(), Ok(0)); + assert!(tx.shared.inner.lock().buffer.capacity() < burst); } -#[tokio::test] -async fn test_multi_senders_multiple_receivers_receive_all() { - let (tx, mut rx1) = channel(); - let mut rx2 = tx.subscribe(); - let mut rx3 = tx.subscribe(); - let tx1 = tx.clone(); - let tx2 = tx.clone(); +#[test] +fn repeated_bursts_keep_their_allocation() { + let (tx, mut rx) = channel(); + let burst = MIN_RETAINED_CAPACITY * 4; - let handle1 = tokio::spawn(async move { - for i in 0..10 { - tx1.send(i); + for _ in 0..4 { + for i in 0..burst { + tx.send(i); } - }); - let handle2 = tokio::spawn(async move { - for i in 10..20 { - tx2.send(i); + for i in 0..burst { + assert_eq!(rx.try_recv(), Ok(i)); } - }); - - for i in 20..30 { - tx.send(i); } - handle1.await.unwrap(); - handle2.await.unwrap(); - drop(tx); - - let received1 = drain(&mut rx1).await; - let received2 = drain(&mut rx2).await; - let received3 = drain(&mut rx3).await; - - assert_eq!(received1, received2); - assert_eq!(received1, received3); - - let mut sorted = received1; - sorted.sort(); - let expected = (0..30).collect::>(); - assert_eq!(sorted, expected); -} - -async fn drain(rx: &mut Receiver) -> Vec { - let mut received = Vec::new(); - while let Ok(n) = rx.recv().await { - received.push(n); - } - received + // Every cycle peaks at the same size, so the buffer must not rebuild its allocation each time. + assert!(tx.shared.inner.lock().buffer.capacity() >= burst); } diff --git a/benchmarks/broadcast_unbounded.rs b/benchmarks/broadcast_unbounded.rs new file mode 100644 index 0000000..74f9e15 --- /dev/null +++ b/benchmarks/broadcast_unbounded.rs @@ -0,0 +1,356 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Every benchmark here must return the channel to a drained state on each iteration. Unlike the +// overflow policy, this channel has no capacity ceiling, so a timed loop that only sends would +// grow the retained backlog until the process runs out of memory. + +use std::fmt; +use std::pin::pin; +use std::sync::Arc; +use std::sync::Barrier; +use std::thread; +use std::thread::JoinHandle; + +use asyncband::broadcast::unbounded; +use divan::Bencher; +use divan::black_box; + +use super::support::bench_context; +use super::support::poll_pending; +use super::support::poll_pinned_ready; + +const RECEIVER_COUNTS: &[usize] = &[1, 8, 32]; +const CONCURRENCY_COUNTS: &[usize] = &[1, 2, 4, 8]; +const CONCURRENT_BATCH_SIZE: usize = 4096; + +/// A channel that peaked at `peak` receivers and currently has `live` of them. +/// +/// The two are measured separately because a dropped receiver leaves its slot behind: the reclaim +/// scan walks every slot the channel ever handed out, so a channel that shed receivers keeps +/// paying for the peak. Pairing each peak with a drained arena is what makes that visible. +#[derive(Clone, Copy)] +struct Fanout { + peak: usize, + live: usize, +} + +impl fmt::Display for Fanout { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "peak {} live {}", self.peak, self.live) + } +} + +const RECLAIM_FANOUTS: &[Fanout] = &[ + Fanout { peak: 1, live: 1 }, + Fanout { peak: 8, live: 8 }, + Fanout { peak: 8, live: 1 }, + Fanout { peak: 32, live: 32 }, + Fanout { peak: 32, live: 4 }, + Fanout { peak: 32, live: 1 }, + Fanout { + peak: 256, + live: 32, + }, + Fanout { peak: 256, live: 1 }, +]; + +struct ConcurrentSend { + receiver: unbounded::Receiver, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl ConcurrentSend { + fn new(sender_count: usize) -> Self { + let (sender, receiver) = unbounded::channel(); + let ready = Arc::new(Barrier::new(sender_count + 1)); + let start = Arc::new(Barrier::new(sender_count + 1)); + let done = Arc::new(Barrier::new(sender_count + 1)); + let sends_per_worker = CONCURRENT_BATCH_SIZE / sender_count; + let mut workers = Vec::with_capacity(sender_count); + + for worker_index in 0..sender_count { + let sender = sender.clone(); + let ready = ready.clone(); + let start = start.clone(); + let done = done.clone(); + workers.push(thread::spawn(move || { + ready.wait(); + start.wait(); + let first = worker_index * sends_per_worker; + for value in first..first + sends_per_worker { + sender.send(black_box(value)); + } + done.wait(); + })); + } + drop(sender); + ready.wait(); + + Self { + receiver, + start, + done, + workers, + } + } + + // The drain is inside the measured region on purpose: it is what keeps the backlog bounded + // across samples, and reclaiming the batch is part of the cost of an unbounded send. + fn run(&mut self) { + self.start.wait(); + self.done.wait(); + while let Ok(value) = self.receiver.try_recv() { + black_box(value); + } + } +} + +impl Drop for ConcurrentSend { + fn drop(&mut self) { + for worker in self.workers.drain(..) { + worker.join().unwrap(); + } + } +} + +struct ConcurrentFanout { + sender: unbounded::Sender, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl ConcurrentFanout { + fn new(receiver_count: usize) -> Self { + let (sender, receiver) = unbounded::channel(); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(sender.subscribe()); + } + + let ready = Arc::new(Barrier::new(receiver_count + 1)); + let start = Arc::new(Barrier::new(receiver_count + 1)); + let done = Arc::new(Barrier::new(receiver_count + 1)); + let mut workers = Vec::with_capacity(receiver_count); + + for mut receiver in receivers { + let ready = ready.clone(); + let start = start.clone(); + let done = done.clone(); + workers.push(thread::spawn(move || { + ready.wait(); + start.wait(); + let result = (0..CONCURRENT_BATCH_SIZE).try_for_each(|_| { + receiver.try_recv().map(|value| { + black_box(value); + }) + }); + done.wait(); + result.unwrap(); + })); + } + ready.wait(); + + Self { + sender, + start, + done, + workers, + } + } + + fn run(&mut self) { + for value in 0..CONCURRENT_BATCH_SIZE { + self.sender.send(black_box(value)); + } + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for ConcurrentFanout { + fn drop(&mut self) { + for worker in self.workers.drain(..) { + worker.join().unwrap(); + } + } +} + +#[divan::bench] +fn send_without_receivers(bencher: Bencher) { + let (sender, receiver) = unbounded::channel::(); + drop(receiver); + bencher.bench_local(|| sender.send(black_box(1))); +} + +#[divan::bench] +fn try_recv_empty(bencher: Bencher) { + let (sender, mut receiver) = unbounded::channel::(); + bencher.bench_local(|| black_box(receiver.try_recv())); + black_box(sender); +} + +// A sole receiver takes ownership of the payload, so this path never clones the message. +#[divan::bench] +fn send_and_try_recv(bencher: Bencher) { + let (sender, mut receiver) = unbounded::channel(); + bencher.bench_local(|| { + sender.send(black_box(1usize)); + black_box(receiver.try_recv().unwrap()) + }); +} + +// With the payload shared, each receive clones it and the second one reclaims the slot. +#[divan::bench] +fn send_and_try_recv_shared(bencher: Bencher) { + let (sender, mut first) = unbounded::channel(); + let mut second = sender.subscribe(); + bencher.bench_local(|| { + sender.send(black_box(1usize)); + black_box(first.try_recv().unwrap()); + black_box(second.try_recv().unwrap()) + }); +} + +// The `usize` benchmarks above hide what a receive costs for a payload that owns memory: a clone +// there is an allocation, not a register move. +fn payload() -> String { + "x".repeat(64) +} + +#[divan::bench] +fn send_and_try_recv_owned(bencher: Bencher) { + let (sender, mut receiver) = unbounded::channel(); + bencher.bench_local(|| { + sender.send(black_box(payload())); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn send_and_try_recv_owned_shared(bencher: Bencher) { + let (sender, mut first) = unbounded::channel(); + let mut second = sender.subscribe(); + bencher.bench_local(|| { + sender.send(black_box(payload())); + black_box(first.try_recv().unwrap()); + black_box(second.try_recv().unwrap()) + }); +} + +// Measures the reclaim scan, which runs when the slowest cursor advances. Comparing a peak against +// the same peak drained down to fewer receivers shows what the slots left behind still cost. +#[divan::bench(args = RECLAIM_FANOUTS)] +fn drain_with_receivers(bencher: Bencher, fanout: Fanout) { + let (sender, receiver) = unbounded::channel(); + drop(receiver); + let mut receivers = (0..fanout.peak) + .map(|_| sender.subscribe()) + .collect::>(); + // Dropping down to `live` leaves the arena holding a slot for every receiver that ever existed. + receivers.truncate(fanout.live); + + bencher.bench_local(|| { + sender.send(black_box(1usize)); + for receiver in &mut receivers { + black_box(receiver.try_recv().unwrap()); + } + }); +} + +#[divan::bench( + args = CONCURRENCY_COUNTS, + sample_count = 50, + sample_size = 1, + counters = [CONCURRENT_BATCH_SIZE] +)] +fn concurrent_send_and_drain(bencher: Bencher, sender_count: usize) { + bencher + .with_inputs(|| ConcurrentSend::new(sender_count)) + .bench_local_refs(ConcurrentSend::run); +} + +#[divan::bench( + args = CONCURRENCY_COUNTS, + sample_count = 50, + sample_size = 1, + counters = [CONCURRENT_BATCH_SIZE] +)] +fn concurrent_fanout(bencher: Bencher, receiver_count: usize) { + bencher + .with_inputs(|| ConcurrentFanout::new(receiver_count)) + .bench_local_refs(ConcurrentFanout::run); +} + +#[divan::bench] +fn cancel_pending(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (sender, mut receiver) = unbounded::channel::(); + { + let mut recv = pin!(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + } + black_box((sender, receiver)) + }); +} + +#[divan::bench] +fn deliver_to_waiter(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (sender, mut receiver) = unbounded::channel(); + let mut recv = pin!(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + + sender.send(black_box(1usize)); + let value = poll_pinned_ready(recv.as_mut(), &mut context).unwrap(); + black_box(value) + }); +} + +#[divan::bench(args = RECEIVER_COUNTS)] +fn deliver_to_receiver_batch(bencher: Bencher, receiver_count: usize) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (sender, receiver) = unbounded::channel(); + drop(receiver); + let mut receivers = (0..receiver_count) + .map(|_| sender.subscribe()) + .collect::>(); + let mut recvs = receivers + .iter_mut() + .map(|receiver| Box::pin(receiver.recv())) + .collect::>(); + for recv in &mut recvs { + poll_pending(recv.as_mut(), &mut context); + } + + sender.send(black_box(1usize)); + for mut recv in recvs { + let value = poll_pinned_ready(recv.as_mut(), &mut context).unwrap(); + black_box(value); + } + }); +} diff --git a/benchmarks/main.rs b/benchmarks/main.rs index 43c5d32..31e9cae 100644 --- a/benchmarks/main.rs +++ b/benchmarks/main.rs @@ -18,6 +18,7 @@ mod barrier; mod blocking; mod broadcast; +mod broadcast_unbounded; mod condvar; mod latch; mod mpsc; diff --git a/tests-integration/tests/broadcast_unbounded_test.rs b/tests-integration/tests/broadcast_unbounded_test.rs new file mode 100644 index 0000000..9287a9f --- /dev/null +++ b/tests-integration/tests/broadcast_unbounded_test.rs @@ -0,0 +1,654 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Wake; +use std::task::Waker; +use std::thread; +use std::time::Duration; +use std::time::Instant; + +use asyncband::broadcast::unbounded::*; + +struct TrackWake(AtomicUsize); + +impl Wake for TrackWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +/// A payload whose destructor re-enters the channel it was sent through. +struct Reentrant { + value: u64, + channel: Option>, +} + +impl Clone for Reentrant { + fn clone(&self) -> Self { + Self { + value: self.value, + channel: self.channel.clone(), + } + } +} + +impl Drop for Reentrant { + fn drop(&mut self) { + if let Some(channel) = &self.channel { + // Deadlocks if the channel still holds its lock while dropping reclaimed messages. + let _ = channel.buffer_len(); + let _ = channel.receiver_count(); + } + } +} + +/// A payload that panics while a shared receive clones it. +#[derive(Debug)] +struct PanicOnClone { + value: u64, + panic: bool, +} + +impl Clone for PanicOnClone { + fn clone(&self) -> Self { + if self.panic { + panic!("panic while cloning a broadcast message"); + } + Self { + value: self.value, + panic: self.panic, + } + } +} + +struct Rng(u64); + +impl Rng { + fn below(&mut self, n: u64) -> u64 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + x % n + } +} + +#[tokio::test] +async fn test_broadcast_basic() { + let (tx, mut rx1) = channel(); + let mut rx2 = tx.subscribe(); + + tx.send(10); + tx.send(20); + + assert_eq!(rx1.recv().await, Ok(10)); + assert_eq!(rx1.recv().await, Ok(20)); + assert_eq!(rx2.recv().await, Ok(10)); + assert_eq!(rx2.recv().await, Ok(20)); +} + +#[tokio::test] +async fn test_subscribe() { + let (tx, _rx) = channel(); + let mut rx = tx.subscribe(); + + tx.send(100); + assert_eq!(rx.recv().await, Ok(100)); +} + +#[tokio::test] +async fn test_resubscribe() { + let (tx, mut rx) = channel(); + + tx.send(1); + tx.send(2); + + let mut rx2 = rx.resubscribe(); + + // rx sees 1, 2 + // rx2 sees nothing yet (starts at tail=2) + + tx.send(3); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx2.recv().await, Ok(3)); +} + +#[test] +fn test_try_recv() { + let (tx, mut rx) = channel(); + + // Empty + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + // Success + tx.send(10); + assert_eq!(rx.try_recv(), Ok(10)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + // Closed + drop(tx); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[tokio::test] +async fn test_slow_receiver_keeps_every_message() { + let (tx, mut rx1) = channel(); + let mut rx2 = tx.subscribe(); + + for i in 0..1024 { + tx.send(i); + } + + // The fast receiver draining fully must not reclaim anything the slow one still needs. + for i in 0..1024 { + assert_eq!(rx1.recv().await, Ok(i)); + } + assert_eq!(tx.buffer_len(), 1024); + + for i in 0..1024 { + assert_eq!(rx2.recv().await, Ok(i)); + } + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn buffer_len_tracks_the_slowest_receiver() { + let (tx, mut rx1) = channel(); + let mut rx2 = tx.subscribe(); + + tx.send(1); + tx.send(2); + assert_eq!(tx.buffer_len(), 2); + + // Reclaiming waits for the slowest receiver, message by message. + assert_eq!(rx1.recv().await, Ok(1)); + assert_eq!(tx.buffer_len(), 2); + assert_eq!(rx2.recv().await, Ok(1)); + assert_eq!(tx.buffer_len(), 1); + + assert_eq!(rx1.recv().await, Ok(2)); + assert_eq!(tx.buffer_len(), 1); + assert_eq!(rx2.recv().await, Ok(2)); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn test_dropping_a_lagging_receiver_releases_its_backlog() { + let (tx, mut rx1) = channel(); + let rx2 = tx.subscribe(); + + for i in 0..128 { + tx.send(i); + } + for i in 0..128 { + assert_eq!(rx1.recv().await, Ok(i)); + } + assert_eq!(tx.buffer_len(), 128); + + drop(rx2); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn resubscribe_keeps_the_original_receivers_backlog() { + let (tx, mut rx) = channel(); + + tx.send(1); + tx.send(2); + + let mut rx2 = rx.resubscribe(); + assert_eq!(tx.buffer_len(), 2); + + tx.send(3); + + assert_eq!(rx2.recv().await, Ok(3)); + assert_eq!(tx.buffer_len(), 3); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Ok(3)); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn send_without_receivers_does_not_buffer() { + let (tx, rx) = channel(); + drop(rx); + + tx.send(1); + tx.send(2); + assert_eq!(tx.buffer_len(), 0); + + let mut rx = tx.subscribe(); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + tx.send(3); + assert_eq!(rx.recv().await, Ok(3)); +} + +#[test] +fn receiver_count_and_len_track_each_receiver() { + let (tx, mut rx1) = channel(); + assert_eq!(tx.receiver_count(), 1); + assert_eq!(rx1.len(), 0); + assert!(rx1.is_empty()); + + tx.send(1); + tx.send(2); + assert_eq!(rx1.len(), 2); + assert!(!rx1.is_empty()); + + let mut rx2 = tx.subscribe(); + assert_eq!(tx.receiver_count(), 2); + assert_eq!(rx2.len(), 0); + assert!(rx2.is_empty()); + + tx.send(3); + assert_eq!(rx1.len(), 3); + assert_eq!(rx2.len(), 1); + + assert_eq!(rx2.try_recv(), Ok(3)); + assert_eq!(rx2.len(), 0); + drop(rx2); + assert_eq!(tx.receiver_count(), 1); + + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.len(), 2); +} + +#[tokio::test] +async fn clone_shares_the_current_position() { + let (tx, mut rx) = channel(); + + tx.send(1); + tx.send(2); + assert_eq!(rx.recv().await, Ok(1)); + + // The clone inherits the unread backlog, unlike `resubscribe`. + let mut clone = rx.clone(); + let mut fresh = rx.resubscribe(); + assert_eq!(tx.receiver_count(), 3); + + tx.send(3); + + assert_eq!(clone.recv().await, Ok(2)); + assert_eq!(clone.recv().await, Ok(3)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Ok(3)); + assert_eq!(fresh.recv().await, Ok(3)); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn clone_at_head_keeps_the_backlog_alive() { + let (tx, mut rx) = channel(); + + tx.send(1); + let clone = rx.clone(); + + assert_eq!(rx.recv().await, Ok(1)); + // The clone still sits at `head`, so the message must not be reclaimed yet. + assert_eq!(tx.buffer_len(), 1); + + drop(clone); + assert_eq!(tx.buffer_len(), 0); +} + +#[test] +fn sole_receiver_takes_messages_without_cloning() { + static CLONES: AtomicUsize = AtomicUsize::new(0); + + struct CountClone(u32); + + impl Clone for CountClone { + fn clone(&self) -> Self { + CLONES.fetch_add(1, Ordering::Relaxed); + Self(self.0) + } + } + + let (tx, mut rx) = channel(); + for i in 0..8 { + tx.send(CountClone(i)); + assert_eq!(rx.try_recv().unwrap().0, i); + } + assert_eq!(CLONES.load(Ordering::Relaxed), 0); + + // A second receiver means the payload is shared, so it has to be cloned again. + let mut second = tx.subscribe(); + tx.send(CountClone(8)); + assert_eq!(rx.try_recv().unwrap().0, 8); + assert_eq!(second.try_recv().unwrap().0, 8); + assert_eq!(CLONES.load(Ordering::Relaxed), 1); +} + +#[test] +fn panicking_clone_leaves_the_channel_consistent() { + let (tx, mut rx1) = channel(); + let mut rx2 = tx.subscribe(); + + tx.send(PanicOnClone { + value: 1, + panic: true, + }); + tx.send(PanicOnClone { + value: 2, + panic: false, + }); + + // Two receivers share the payload, so this receive has to clone it. + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + rx1.try_recv().map(|msg| msg.value) + })); + assert!(result.is_err()); + + // The failed receive still consumed the message for `rx1`, and left the channel usable for + // both receivers. + assert_eq!(rx1.try_recv().unwrap().value, 2); + assert_eq!(rx2.try_recv().unwrap().value, 1); + assert_eq!(rx2.try_recv().unwrap().value, 2); + assert_eq!(tx.buffer_len(), 0); + assert_eq!(rx1.try_recv().unwrap_err(), TryRecvError::Empty); +} + +#[test] +fn message_destructors_run_outside_the_channel_lock() { + let finished = Arc::new(AtomicUsize::new(0)); + let flag = finished.clone(); + + let worker = thread::spawn(move || { + let (tx, mut rx1) = channel(); + let rx2 = tx.subscribe(); + + for value in 0..4 { + tx.send(Reentrant { + value, + channel: Some(tx.clone()), + }); + } + + // Reclaim through a receive, and then through a receiver drop. + assert_eq!(rx1.try_recv().unwrap().value, 0); + drop(rx2); + assert_eq!(rx1.try_recv().unwrap().value, 1); + drop(rx1); + + // With no receiver left, `send` drops the message itself; that must be unlocked too. + tx.send(Reentrant { + value: 4, + channel: Some(tx.clone()), + }); + drop(tx); + + flag.store(1, Ordering::SeqCst); + }); + + let deadline = Instant::now() + Duration::from_secs(10); + while Instant::now() < deadline && finished.load(Ordering::SeqCst) == 0 { + thread::sleep(Duration::from_millis(10)); + } + assert_eq!( + finished.load(Ordering::SeqCst), + 1, + "a message destructor deadlocked against the channel lock" + ); + worker.join().unwrap(); +} + +#[tokio::test] +async fn test_wait_mechanism() { + let (tx, mut rx) = channel(); + + let handle = tokio::spawn(async move { rx.recv().await }); + + tokio::time::sleep(Duration::from_millis(100)).await; + tx.send(42); + + assert_eq!(handle.await.unwrap(), Ok(42)); +} + +#[test] +fn send_wakes_a_parked_receiver_exactly_once() { + let (tx, mut rx) = channel(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + + tx.send(42); + + assert_eq!(tracker.0.load(Ordering::Relaxed), 1); + assert_eq!(recv.as_mut().poll(&mut context), Poll::Ready(Ok(42))); +} + +#[test] +fn cancelled_recv_releases_its_waker() { + let (tx, mut rx) = channel::<()>(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let baseline = Arc::strong_count(&tracker); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + + drop(recv); + assert_eq!(Arc::strong_count(&tracker), baseline); + + tx.send(()); + assert_eq!(tracker.0.load(Ordering::Relaxed), 0); + assert_eq!(rx.try_recv(), Ok(())); +} + +#[test] +fn dropping_a_woken_recv_keeps_another_receivers_waiter() { + let (tx, mut rx1) = channel::(); + let mut rx2 = tx.subscribe(); + let first = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(first.clone()); + let mut context = Context::from_waker(&waker); + let mut recv1 = Box::pin(rx1.recv()); + + assert!(recv1.as_mut().poll(&mut context).is_pending()); + + tx.send(1); + assert_eq!(first.0.load(Ordering::Relaxed), 1); + assert_eq!(rx2.try_recv(), Ok(1)); + + let second = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(second.clone()); + let mut context = Context::from_waker(&waker); + let mut recv2 = Box::pin(rx2.recv()); + assert!(recv2.as_mut().poll(&mut context).is_pending()); + + // `recv1` was already woken, so dropping it must not release the slot `recv2` now owns. + drop(recv1); + tx.send(2); + + assert_eq!(second.0.load(Ordering::Relaxed), 1); +} + +#[test] +fn parked_recv_wakes_when_the_last_sender_drops() { + let (tx, mut rx) = channel::<()>(); + let extra = tx.clone(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + + drop(tx); + assert_eq!(tracker.0.load(Ordering::Relaxed), 0); + + drop(extra); + assert_eq!(tracker.0.load(Ordering::Relaxed), 1); + + drop(recv); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[test] +fn parked_recv_prefers_buffered_messages_over_disconnect() { + let (tx, mut rx) = channel(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + + tx.send(7); + drop(tx); + + assert_eq!(recv.as_mut().poll(&mut context), Poll::Ready(Ok(7))); + drop(recv); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[tokio::test] +async fn recv_drains_buffered_messages_before_reporting_disconnect() { + let (tx, mut rx) = channel(); + + tx.send(1); + tx.send(2); + drop(tx); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); +} + +#[tokio::test] +async fn recv_reports_disconnect_without_any_message() { + let (tx, mut rx) = channel::<()>(); + drop(tx); + assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); +} + +#[test] +fn concurrent_senders_deliver_every_message_to_every_receiver() { + const SENDERS: u64 = 4; + const PER_SENDER: u64 = 512; + + let (tx, rx) = channel(); + let receivers = (0..4) + .map(|index| { + if index == 0 { + rx.clone() + } else { + tx.subscribe() + } + }) + .collect::>(); + drop(rx); + + let senders = (0..SENDERS) + .map(|worker| { + let tx = tx.clone(); + thread::spawn(move || { + for value in 0..PER_SENDER { + tx.send(worker * PER_SENDER + value); + } + }) + }) + .collect::>(); + + let drains = receivers + .into_iter() + .map(|mut receiver| { + thread::spawn(move || { + let mut seen = Vec::new(); + while let Ok(value) = pollster::block_on(receiver.recv()) { + seen.push(value); + } + seen + }) + }) + .collect::>(); + + for sender in senders { + sender.join().unwrap(); + } + drop(tx); + + let expected = (0..SENDERS * PER_SENDER).collect::>(); + for drain in drains { + let mut seen = drain.join().unwrap(); + seen.sort_unstable(); + assert_eq!(seen, expected); + } +} + +#[test] +fn randomized_operations_track_the_reference_model() { + for seed in 1..32u64 { + let mut rng = Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1); + let (tx, rx) = channel::(); + let mut tail = 0u64; + let mut model = vec![(rx, 0u64)]; + + for _ in 0..512 { + match rng.below(100) { + 0..=44 => { + tx.send(tail); + tail += 1; + } + 45..=79 if !model.is_empty() => { + let index = rng.below(model.len() as u64) as usize; + let (receiver, cursor) = &mut model[index]; + if *cursor < tail { + assert_eq!(receiver.try_recv(), Ok(*cursor), "seed {seed}"); + *cursor += 1; + } else { + assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty), "seed {seed}"); + } + } + 80..=89 => model.push((tx.subscribe(), tail)), + _ if !model.is_empty() => { + let index = rng.below(model.len() as u64) as usize; + model.swap_remove(index); + } + _ => {} + } + + assert_eq!(tx.receiver_count(), model.len(), "seed {seed}"); + let retained = model + .iter() + .map(|(_, cursor)| *cursor) + .min() + .map_or(0, |slowest| tail - slowest); + assert_eq!(tx.buffer_len(), retained as usize, "seed {seed}"); + for (receiver, cursor) in &model { + assert_eq!(receiver.len(), (tail - cursor) as usize, "seed {seed}"); + assert_eq!(receiver.is_empty(), *cursor == tail, "seed {seed}"); + } + } + } +}