From 3b790301e1b31ded24825cc8cb1e24a169ee553a Mon Sep 17 00:00:00 2001 From: orthur2 Date: Sat, 6 Jun 2026 15:17:02 +0200 Subject: [PATCH 1/8] feat(broadcast): add unbounded policy --- CHANGELOG.md | 1 + README.md | 1 + asyncband/Cargo.toml | 1 + asyncband/src/channel/broadcast/mod.rs | 20 + .../src/channel/broadcast/unbounded/mod.rs | 683 ++++++++++++++++++ .../src/channel/broadcast/unbounded/tests.rs | 411 +++++++++++ asyncband/src/channel/mod.rs | 2 + asyncband/src/internal/arena.rs | 20 +- asyncband/src/internal/mod.rs | 5 +- asyncband/src/lib.rs | 4 +- benchmarks/Cargo.toml | 1 + tests-integration/Cargo.toml | 1 + tests-integration/tests/traits_test.rs | 9 + 13 files changed, 1152 insertions(+), 7 deletions(-) create mode 100644 asyncband/src/channel/broadcast/mod.rs create mode 100644 asyncband/src/channel/broadcast/unbounded/mod.rs create mode 100644 asyncband/src/channel/broadcast/unbounded/tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 78bc8ac..a65fea8 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. * Add opt-in bounded and unbounded runtime-agnostic object pools under `asyncband::pool`. * Add opt-in `asyncband::once::LazyCell` for values that own one asynchronous initializer and preserve its in-flight future across caller cancellation. diff --git a/README.md b/README.md index bcec6fe..566c6e6 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/) | `shutdown` | Coordinate shutdown signals and completion. | | Channels | [`oneshot`](https://docs.rs/asyncband/*/asyncband/oneshot/) | `oneshot` | Send one value between two tasks. | | | [`mpsc`](https://docs.rs/asyncband/*/asyncband/mpsc/) | `mpsc` | Send values from multiple producers through bounded or unbounded channels. | +| | [`broadcast::unbounded`](https://docs.rs/asyncband/*/asyncband/broadcast/unbounded/) | `broadcast` | Broadcast values and retain them until every active receiver consumes them. | | Resource reuse | [`pool`](https://docs.rs/asyncband/*/asyncband/pool/) | `pool` | Reuse objects through bounded or unbounded pool variants. | | Workload coordination | [`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/Cargo.toml b/asyncband/Cargo.toml index 3871529..aafe547 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -46,6 +46,7 @@ default = [] barrier = [] blocking = [] +broadcast = [] condvar = ["mutex"] latch = [] lazy-cell = ["mutex"] diff --git a/asyncband/src/channel/broadcast/mod.rs b/asyncband/src/channel/broadcast/mod.rs new file mode 100644 index 0000000..3a61b71 --- /dev/null +++ b/asyncband/src/channel/broadcast/mod.rs @@ -0,0 +1,20 @@ +// 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. + +//! Multi-producer, multi-consumer broadcast channels. + +pub mod unbounded; diff --git a/asyncband/src/channel/broadcast/unbounded/mod.rs b/asyncband/src/channel/broadcast/unbounded/mod.rs new file mode 100644 index 0000000..208b25a --- /dev/null +++ b/asyncband/src/channel/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::SlotId; +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) -> SlotId { + if head == self.head { + self.head_receivers += 1; + } + + self.receivers.insert(head) + } + + fn remove_receiver(&mut self, key: SlotId) -> Vec> { + let head = self.receivers.remove(key); + + if head == self.head { + self.release_head_receiver() + } else { + Vec::new() + } + } + + fn advance_receiver(&mut self, key: SlotId, 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: SlotId) -> 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: SlotId, +} + +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/channel/broadcast/unbounded/tests.rs b/asyncband/src/channel/broadcast/unbounded/tests.rs new file mode 100644 index 0000000..7c62175 --- /dev/null +++ b/asyncband/src/channel/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/channel/mod.rs b/asyncband/src/channel/mod.rs index 58689a6..bfe3be0 100644 --- a/asyncband/src/channel/mod.rs +++ b/asyncband/src/channel/mod.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +#[cfg(feature = "broadcast")] +pub mod broadcast; #[cfg(feature = "mpsc")] pub mod mpsc; #[cfg(feature = "oneshot")] diff --git a/asyncband/src/internal/arena.rs b/asyncband/src/internal/arena.rs index 98502aa..3158d31 100644 --- a/asyncband/src/internal/arena.rs +++ b/asyncband/src/internal/arena.rs @@ -136,6 +136,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, + }) + } + /// Removes the value stored at `id`. /// /// # Panics @@ -194,11 +209,6 @@ impl Arena { self.len = 0; values.into_iter() } - - #[cfg(test)] - pub fn len(&self) -> usize { - self.len - } } #[cfg(test)] diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index cdd33e3..2374f39 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -20,6 +20,7 @@ pub(crate) mod atomic_waker; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", feature = "mpsc", feature = "mutex", @@ -30,7 +31,7 @@ pub(crate) mod atomic_waker; // `WaitList` and `WaitSet` use different `Arena` operations. A single-primitive build therefore // leaves part of this shared API unused, while the all-feature build uses it. #[allow(dead_code)] -mod arena; +pub(crate) mod arena; #[cfg(any(feature = "latch", feature = "once", feature = "waitgroup"))] // `waitgroup` increments and decrements the countdown, while `latch` and `once` only decrement it. @@ -52,6 +53,7 @@ pub(crate) mod value_cell; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", feature = "mpsc", feature = "mutex", @@ -83,6 +85,7 @@ pub(crate) mod waitlist; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", feature = "once", feature = "waitgroup", diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 733275a..dae129a 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::LazyCell`], [`once::OnceMap`] | `once`, `once-cell`, `lazy-cell`, `once-map` | //! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | -//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`] | `oneshot`, `mpsc` | +//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::unbounded`] | `oneshot`, `mpsc`, `broadcast` | //! | Reuse objects | [`pool::bounded`], [`pool::unbounded`] | `pool` | //! | Coordinate workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | //! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | @@ -107,6 +107,8 @@ mod internal; pub mod barrier; #[cfg(feature = "blocking")] pub mod blocking; +#[cfg(feature = "broadcast")] +pub use self::channel::broadcast; #[cfg(feature = "condvar")] pub mod condvar; #[cfg(feature = "latch")] diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index c082f90..95eec00 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -27,6 +27,7 @@ async-channel = { workspace = true } asyncband = { workspace = true, features = [ "barrier", "blocking", + "broadcast", "condvar", "latch", "mpsc", diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 7ea541f..b66fda0 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -29,6 +29,7 @@ tokio = { workspace = true, features = ["full"] } asyncband = { workspace = true, features = [ "barrier", "blocking", + "broadcast", "condvar", "latch", "lazy-cell", diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index db52c62..b51d114 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -18,6 +18,7 @@ use std::cell::Cell; use asyncband::barrier::Barrier; +use asyncband::broadcast; use asyncband::condvar::Condvar; use asyncband::latch::Latch; use asyncband::mpsc; @@ -85,6 +86,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::>(); @@ -132,6 +137,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 11551ca085fba3fe6ccbae96105b8cfacf0722bf Mon Sep 17 00:00:00 2001 From: Orthur Date: Sat, 22 Aug 2026 11:13:42 -0400 Subject: [PATCH 2/8] perf(broadcast): serve unbounded receives from one critical section --- .../src/channel/broadcast/unbounded/mod.rs | 232 +++++++++++++----- 1 file changed, 177 insertions(+), 55 deletions(-) diff --git a/asyncband/src/channel/broadcast/unbounded/mod.rs b/asyncband/src/channel/broadcast/unbounded/mod.rs index 208b25a..5abae30 100644 --- a/asyncband/src/channel/broadcast/unbounded/mod.rs +++ b/asyncband/src/channel/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: SlotId, } +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))) } } From cb63504001952e3875fca3433afe66042e1f4ad2 Mon Sep 17 00:00:00 2001 From: Orthur Date: Sat, 22 Aug 2026 11:13:42 -0400 Subject: [PATCH 3/8] test(broadcast): move unbounded coverage into integration tests --- .../src/channel/broadcast/unbounded/tests.rs | 398 +---------- benchmarks/asyncband/broadcast/mod.rs | 18 + benchmarks/asyncband/broadcast/unbounded.rs | 356 ++++++++++ benchmarks/asyncband/main.rs | 1 + .../tests/broadcast_unbounded_test.rs | 654 ++++++++++++++++++ 5 files changed, 1057 insertions(+), 370 deletions(-) create mode 100644 benchmarks/asyncband/broadcast/mod.rs create mode 100644 benchmarks/asyncband/broadcast/unbounded.rs create mode 100644 tests-integration/tests/broadcast_unbounded_test.rs diff --git a/asyncband/src/channel/broadcast/unbounded/tests.rs b/asyncband/src/channel/broadcast/unbounded/tests.rs index 7c62175..a80a527 100644 --- a/asyncband/src/channel/broadcast/unbounded/tests.rs +++ b/asyncband/src/channel/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/asyncband/broadcast/mod.rs b/benchmarks/asyncband/broadcast/mod.rs new file mode 100644 index 0000000..78ef889 --- /dev/null +++ b/benchmarks/asyncband/broadcast/mod.rs @@ -0,0 +1,18 @@ +// 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. + +mod unbounded; diff --git a/benchmarks/asyncband/broadcast/unbounded.rs b/benchmarks/asyncband/broadcast/unbounded.rs new file mode 100644 index 0000000..74f9e15 --- /dev/null +++ b/benchmarks/asyncband/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/asyncband/main.rs b/benchmarks/asyncband/main.rs index cb8ac1a..c383615 100644 --- a/benchmarks/asyncband/main.rs +++ b/benchmarks/asyncband/main.rs @@ -17,6 +17,7 @@ mod barrier; mod blocking; +mod broadcast; 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}"); + } + } + } +} From 889b8913145c45329f7f3403494da88c694c652b Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 27 Aug 2026 01:33:36 +0800 Subject: [PATCH 4/8] refactor(broadcast): align unbounded MPMC API --- CHANGELOG.md | 2 +- README.md | 2 +- asyncband/src/channel/broadcast/mod.rs | 4 +- asyncband/src/channel/broadcast/mpmc/mod.rs | 26 +++ .../{unbounded/mod.rs => mpmc/unbounded.rs} | 149 +++++++----------- .../broadcast/{ => mpmc}/unbounded/tests.rs | 6 +- asyncband/src/lib.rs | 2 +- benchmarks/asyncband/broadcast/mod.rs | 2 +- benchmarks/asyncband/broadcast/mpmc/mod.rs | 18 +++ .../broadcast/{ => mpmc}/unbounded.rs | 36 ++--- ...st.rs => broadcast_mpmc_unbounded_test.rs} | 101 ++++-------- tests-integration/tests/traits_test.rs | 16 +- 12 files changed, 162 insertions(+), 202 deletions(-) create mode 100644 asyncband/src/channel/broadcast/mpmc/mod.rs rename asyncband/src/channel/broadcast/{unbounded/mod.rs => mpmc/unbounded.rs} (85%) rename asyncband/src/channel/broadcast/{ => mpmc}/unbounded/tests.rs (95%) create mode 100644 benchmarks/asyncband/broadcast/mpmc/mod.rs rename benchmarks/asyncband/broadcast/{ => mpmc}/unbounded.rs (91%) rename tests-integration/tests/{broadcast_unbounded_test.rs => broadcast_mpmc_unbounded_test.rs} (88%) diff --git a/CHANGELOG.md b/CHANGELOG.md index a65fea8..0457e74 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +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. +* Implement `broadcast::mpmc::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. * Add opt-in bounded and unbounded runtime-agnostic object pools under `asyncband::pool`. * Add opt-in `asyncband::once::LazyCell` for values that own one asynchronous initializer and preserve its in-flight future across caller cancellation. diff --git a/README.md b/README.md index 566c6e6..018e6ea 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/) | `shutdown` | Coordinate shutdown signals and completion. | | Channels | [`oneshot`](https://docs.rs/asyncband/*/asyncband/oneshot/) | `oneshot` | Send one value between two tasks. | | | [`mpsc`](https://docs.rs/asyncband/*/asyncband/mpsc/) | `mpsc` | Send values from multiple producers through bounded or unbounded channels. | -| | [`broadcast::unbounded`](https://docs.rs/asyncband/*/asyncband/broadcast/unbounded/) | `broadcast` | Broadcast values and retain them until every active receiver consumes them. | +| | [`broadcast::mpmc::unbounded`](https://docs.rs/asyncband/*/asyncband/broadcast/mpmc/fn.unbounded.html) | `broadcast` | Broadcast values from multiple producers and retain them until every active receiver consumes them. | | Resource reuse | [`pool`](https://docs.rs/asyncband/*/asyncband/pool/) | `pool` | Reuse objects through bounded or unbounded pool variants. | | Workload coordination | [`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/channel/broadcast/mod.rs b/asyncband/src/channel/broadcast/mod.rs index 3a61b71..81e56c5 100644 --- a/asyncband/src/channel/broadcast/mod.rs +++ b/asyncband/src/channel/broadcast/mod.rs @@ -15,6 +15,6 @@ // specific language governing permissions and limitations // under the License. -//! Multi-producer, multi-consumer broadcast channels. +//! Broadcast channels grouped by producer topology. -pub mod unbounded; +pub mod mpmc; diff --git a/asyncband/src/channel/broadcast/mpmc/mod.rs b/asyncband/src/channel/broadcast/mpmc/mod.rs new file mode 100644 index 0000000..9249df2 --- /dev/null +++ b/asyncband/src/channel/broadcast/mpmc/mod.rs @@ -0,0 +1,26 @@ +// 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. + +//! Multi-producer, multi-consumer broadcast channels. + +mod unbounded; + +pub use self::unbounded::RecvError; +pub use self::unbounded::TryRecvError; +pub use self::unbounded::UnboundedReceiver; +pub use self::unbounded::UnboundedSender; +pub use self::unbounded::unbounded; diff --git a/asyncband/src/channel/broadcast/unbounded/mod.rs b/asyncband/src/channel/broadcast/mpmc/unbounded.rs similarity index 85% rename from asyncband/src/channel/broadcast/unbounded/mod.rs rename to asyncband/src/channel/broadcast/mpmc/unbounded.rs index 5abae30..594a125 100644 --- a/asyncband/src/channel/broadcast/unbounded/mod.rs +++ b/asyncband/src/channel/broadcast/mpmc/unbounded.rs @@ -25,8 +25,8 @@ //! //! 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. +//! consumed them or the receiver is dropped. Use [`UnboundedSender::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 @@ -34,10 +34,8 @@ //! //! # Receivers //! -//! Each receiver has an independent cursor. Use [`Sender::subscribe`] to create a receiver that -//! 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. +//! Each receiver has an independent cursor. Use [`UnboundedSender::subscribe`] or +//! [`UnboundedReceiver::resubscribe`] to create a receiver that starts 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 @@ -49,11 +47,11 @@ //! Basic usage: //! //! ``` -//! use asyncband::broadcast::unbounded; +//! use asyncband::broadcast::mpmc; //! //! # #[tokio::main] //! # async fn main() { -//! let (tx, mut rx1) = unbounded::channel(); +//! let (tx, mut rx1) = mpmc::unbounded(); //! let mut rx2 = tx.subscribe(); //! //! tx.send(10); @@ -69,11 +67,11 @@ //! Slow receivers do not miss messages: //! //! ``` -//! use asyncband::broadcast::unbounded; +//! use asyncband::broadcast::mpmc; //! //! # #[tokio::main] //! # async fn main() { -//! let (tx, mut rx1) = unbounded::channel(); +//! let (tx, mut rx1) = mpmc::unbounded(); //! let mut rx2 = tx.subscribe(); //! //! tx.send(1); @@ -104,26 +102,26 @@ use std::task::Poll; use crate::internal::arena::Arena; use crate::internal::arena::SlotId; use crate::internal::mutex::Mutex; -use crate::internal::waitset::WaitRegistration; use crate::internal::waitset::WaitSet; +use crate::internal::waitset::WakerToken; #[cfg(test)] mod tests; /// Creates a new broadcast channel with an unbounded buffer. /// -/// See [module-level documentation](self) for broadcast channel semantics. +/// Every accepted value is retained until all active receivers consume it or are dropped. /// /// # Examples /// /// ``` -/// use asyncband::broadcast::unbounded; +/// use asyncband::broadcast::mpmc; /// -/// let (tx, mut rx) = unbounded::channel(); +/// let (tx, mut rx) = mpmc::unbounded(); /// tx.send(10); /// assert_eq!(rx.try_recv(), Ok(10)); /// ``` -pub fn channel() -> (Sender, Receiver) { +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let mut receivers = Arena::new(); let key = receivers.insert(0); let shared = Arc::new(Shared { @@ -138,14 +136,14 @@ pub fn channel() -> (Sender, Receiver) { }), senders: AtomicUsize::new(1), }); - let sender = Sender { + let sender = UnboundedSender { shared: shared.clone(), }; - let receiver = Receiver { shared, key }; + let receiver = UnboundedReceiver { shared, key }; (sender, receiver) } -/// Error returned by [`Receiver::recv`]. +/// Error returned by [`UnboundedReceiver::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. @@ -162,7 +160,7 @@ impl fmt::Display for RecvError { impl std::error::Error for RecvError {} -/// Error returned by [`Receiver::try_recv`]. +/// Error returned by [`UnboundedReceiver::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 @@ -204,7 +202,7 @@ struct Inner { receivers: Arena, /// The largest backlog retained since the buffer was last empty. peak_len: usize, - /// Receivers parked in [`Receiver::recv`]. + /// Receivers parked in [`UnboundedReceiver::recv`]. waiters: WaitSet, } @@ -341,11 +339,11 @@ struct Shared { /// /// The sender can be cloned to create multiple producers. When all senders are dropped, /// the channel is closed. -pub struct Sender { +pub struct UnboundedSender { shared: Arc>, } -impl Clone for Sender { +impl Clone for UnboundedSender { 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 @@ -357,13 +355,13 @@ impl Clone for Sender { } } -impl fmt::Debug for Sender { +impl fmt::Debug for UnboundedSender { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Sender").finish_non_exhaustive() + f.debug_struct("UnboundedSender").finish_non_exhaustive() } } -impl Drop for Sender { +impl Drop for UnboundedSender { fn drop(&mut self) { match self.shared.senders.fetch_sub(1, Ordering::AcqRel) { 1 => { @@ -381,7 +379,7 @@ impl Drop for Sender { } } -impl Sender { +impl UnboundedSender { /// Broadcasts a value to all active receivers. /// /// This operation does not wait for receiver capacity. If receivers fall behind, messages @@ -398,9 +396,9 @@ impl Sender { /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::mpmc; /// - /// let (tx, mut rx) = unbounded::channel(); + /// let (tx, mut rx) = mpmc::unbounded(); /// tx.send(10); /// assert_eq!(rx.try_recv(), Ok(10)); /// ``` @@ -447,9 +445,9 @@ impl Sender { /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::mpmc; /// - /// let (tx, mut rx) = unbounded::channel(); + /// let (tx, mut rx) = mpmc::unbounded(); /// tx.send(10); /// assert_eq!(tx.buffer_len(), 1); /// @@ -465,9 +463,9 @@ impl Sender { /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::mpmc; /// - /// let (tx, rx) = unbounded::channel::(); + /// let (tx, rx) = mpmc::unbounded::(); /// assert_eq!(tx.receiver_count(), 1); /// /// let rx2 = tx.subscribe(); @@ -486,12 +484,12 @@ impl Sender { /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; - /// use asyncband::broadcast::unbounded::TryRecvError; + /// use asyncband::broadcast::mpmc; + /// use asyncband::broadcast::mpmc::TryRecvError; /// /// # #[tokio::main] /// # async fn main() { - /// let (tx, _) = unbounded::channel(); + /// let (tx, _) = mpmc::unbounded(); /// tx.send(10); /// /// let mut rx = tx.subscribe(); @@ -500,67 +498,30 @@ impl Sender { /// assert_eq!(rx.recv().await, Ok(20)); /// # } /// ``` - pub fn subscribe(&self) -> Receiver { + pub fn subscribe(&self) -> UnboundedReceiver { 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 } + UnboundedReceiver { shared, key } } } /// 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 { +pub struct UnboundedReceiver { shared: Arc>, key: SlotId, } -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 { +impl fmt::Debug for UnboundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Receiver").finish_non_exhaustive() + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() } } -impl Drop for Receiver { +impl Drop for UnboundedReceiver { fn drop(&mut self) { let reclaimed = { let mut inner = self.shared.inner.lock(); @@ -570,7 +531,7 @@ impl Drop for Receiver { } } -impl Receiver { +impl UnboundedReceiver { /// Receives the next value for this receiver. /// /// # Returns @@ -588,11 +549,11 @@ impl Receiver { /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::mpmc; /// /// # #[tokio::main] /// # async fn main() { - /// let (tx, mut rx) = unbounded::channel(); + /// let (tx, mut rx) = mpmc::unbounded(); /// tx.send(10); /// assert_eq!(rx.recv().await, Ok(10)); /// # } @@ -617,9 +578,9 @@ impl Receiver { /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::mpmc; /// - /// let (tx, mut rx) = unbounded::channel(); + /// let (tx, mut rx) = mpmc::unbounded(); /// tx.send(10); /// assert_eq!(rx.try_recv(), Ok(10)); /// ``` @@ -654,7 +615,7 @@ fn take_msg(msg: Arc, reclaimed: Vec>) -> T { } } -impl Receiver { +impl UnboundedReceiver { 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 @@ -681,9 +642,9 @@ impl Receiver { /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::mpmc; /// - /// let (tx, mut rx) = unbounded::channel(); + /// let (tx, mut rx) = mpmc::unbounded(); /// tx.send(1); /// tx.send(2); /// @@ -702,15 +663,15 @@ impl Receiver { /// 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. + /// This count is specific to this receiver, unlike [`UnboundedSender::buffer_len`], which + /// reports the shared backlog retained by the slowest active receiver. /// /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::mpmc; /// - /// let (tx, mut rx) = unbounded::channel(); + /// let (tx, mut rx) = mpmc::unbounded(); /// assert_eq!(rx.len(), 0); /// /// tx.send(10); @@ -734,9 +695,9 @@ impl Receiver { /// # Examples /// /// ``` - /// use asyncband::broadcast::unbounded; + /// use asyncband::broadcast::mpmc; /// - /// let (tx, rx) = unbounded::channel(); + /// let (tx, rx) = mpmc::unbounded(); /// assert!(rx.is_empty()); /// /// tx.send(10); @@ -748,8 +709,8 @@ impl Receiver { } struct Recv<'a, T> { - receiver: &'a mut Receiver, - registration: Option, + receiver: &'a mut UnboundedReceiver, + registration: Option, } impl Drop for Recv<'_, T> { diff --git a/asyncband/src/channel/broadcast/unbounded/tests.rs b/asyncband/src/channel/broadcast/mpmc/unbounded/tests.rs similarity index 95% rename from asyncband/src/channel/broadcast/unbounded/tests.rs rename to asyncband/src/channel/broadcast/mpmc/unbounded/tests.rs index a80a527..e3f6746 100644 --- a/asyncband/src/channel/broadcast/unbounded/tests.rs +++ b/asyncband/src/channel/broadcast/mpmc/unbounded/tests.rs @@ -21,14 +21,14 @@ use super::*; #[should_panic(expected = "broadcast channel version counter overflowed")] fn send_panics_on_version_overflow() { // The receiver is dropped right away: the doctored counter would make its own drop overflow. - let (tx, _) = channel(); + let (tx, _) = unbounded(); tx.shared.inner.lock().tail = u64::MAX; tx.send(()); } #[test] fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); let burst = MIN_RETAINED_CAPACITY * 16; for i in 0..burst { @@ -52,7 +52,7 @@ fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { #[test] fn repeated_bursts_keep_their_allocation() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); let burst = MIN_RETAINED_CAPACITY * 4; for _ in 0..4 { diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index dae129a..aa9dbf4 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::LazyCell`], [`once::OnceMap`] | `once`, `once-cell`, `lazy-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::unbounded`] | `oneshot`, `mpsc`, `broadcast` | +//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::mpmc::unbounded`] | `oneshot`, `mpsc`, `broadcast` | //! | Reuse objects | [`pool::bounded`], [`pool::unbounded`] | `pool` | //! | Coordinate workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | //! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | diff --git a/benchmarks/asyncband/broadcast/mod.rs b/benchmarks/asyncband/broadcast/mod.rs index 78ef889..b078f4a 100644 --- a/benchmarks/asyncband/broadcast/mod.rs +++ b/benchmarks/asyncband/broadcast/mod.rs @@ -15,4 +15,4 @@ // specific language governing permissions and limitations // under the License. -mod unbounded; +mod mpmc; diff --git a/benchmarks/asyncband/broadcast/mpmc/mod.rs b/benchmarks/asyncband/broadcast/mpmc/mod.rs new file mode 100644 index 0000000..78ef889 --- /dev/null +++ b/benchmarks/asyncband/broadcast/mpmc/mod.rs @@ -0,0 +1,18 @@ +// 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. + +mod unbounded; diff --git a/benchmarks/asyncband/broadcast/unbounded.rs b/benchmarks/asyncband/broadcast/mpmc/unbounded.rs similarity index 91% rename from benchmarks/asyncband/broadcast/unbounded.rs rename to benchmarks/asyncband/broadcast/mpmc/unbounded.rs index 74f9e15..fa35c27 100644 --- a/benchmarks/asyncband/broadcast/unbounded.rs +++ b/benchmarks/asyncband/broadcast/mpmc/unbounded.rs @@ -26,13 +26,13 @@ use std::sync::Barrier; use std::thread; use std::thread::JoinHandle; -use asyncband::broadcast::unbounded; +use asyncband::broadcast::mpmc; use divan::Bencher; use divan::black_box; -use super::support::bench_context; -use super::support::poll_pending; -use super::support::poll_pinned_ready; +use crate::support::bench_context; +use crate::support::poll_pending; +use crate::support::poll_pinned_ready; const RECEIVER_COUNTS: &[usize] = &[1, 8, 32]; const CONCURRENCY_COUNTS: &[usize] = &[1, 2, 4, 8]; @@ -70,7 +70,7 @@ const RECLAIM_FANOUTS: &[Fanout] = &[ ]; struct ConcurrentSend { - receiver: unbounded::Receiver, + receiver: mpmc::UnboundedReceiver, start: Arc, done: Arc, workers: Vec>, @@ -78,7 +78,7 @@ struct ConcurrentSend { impl ConcurrentSend { fn new(sender_count: usize) -> Self { - let (sender, receiver) = unbounded::channel(); + let (sender, receiver) = mpmc::unbounded(); 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)); @@ -131,7 +131,7 @@ impl Drop for ConcurrentSend { } struct ConcurrentFanout { - sender: unbounded::Sender, + sender: mpmc::UnboundedSender, start: Arc, done: Arc, workers: Vec>, @@ -139,7 +139,7 @@ struct ConcurrentFanout { impl ConcurrentFanout { fn new(receiver_count: usize) -> Self { - let (sender, receiver) = unbounded::channel(); + let (sender, receiver) = mpmc::unbounded(); let mut receivers = Vec::with_capacity(receiver_count); receivers.push(receiver); for _ in 1..receiver_count { @@ -196,14 +196,14 @@ impl Drop for ConcurrentFanout { #[divan::bench] fn send_without_receivers(bencher: Bencher) { - let (sender, receiver) = unbounded::channel::(); + let (sender, receiver) = mpmc::unbounded::(); drop(receiver); bencher.bench_local(|| sender.send(black_box(1))); } #[divan::bench] fn try_recv_empty(bencher: Bencher) { - let (sender, mut receiver) = unbounded::channel::(); + let (sender, mut receiver) = mpmc::unbounded::(); bencher.bench_local(|| black_box(receiver.try_recv())); black_box(sender); } @@ -211,7 +211,7 @@ fn try_recv_empty(bencher: Bencher) { // 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(); + let (sender, mut receiver) = mpmc::unbounded(); bencher.bench_local(|| { sender.send(black_box(1usize)); black_box(receiver.try_recv().unwrap()) @@ -221,7 +221,7 @@ fn send_and_try_recv(bencher: Bencher) { // 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 (sender, mut first) = mpmc::unbounded(); let mut second = sender.subscribe(); bencher.bench_local(|| { sender.send(black_box(1usize)); @@ -238,7 +238,7 @@ fn payload() -> String { #[divan::bench] fn send_and_try_recv_owned(bencher: Bencher) { - let (sender, mut receiver) = unbounded::channel(); + let (sender, mut receiver) = mpmc::unbounded(); bencher.bench_local(|| { sender.send(black_box(payload())); black_box(receiver.try_recv().unwrap()) @@ -247,7 +247,7 @@ fn send_and_try_recv_owned(bencher: Bencher) { #[divan::bench] fn send_and_try_recv_owned_shared(bencher: Bencher) { - let (sender, mut first) = unbounded::channel(); + let (sender, mut first) = mpmc::unbounded(); let mut second = sender.subscribe(); bencher.bench_local(|| { sender.send(black_box(payload())); @@ -260,7 +260,7 @@ fn send_and_try_recv_owned_shared(bencher: Bencher) { // 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(); + let (sender, receiver) = mpmc::unbounded(); drop(receiver); let mut receivers = (0..fanout.peak) .map(|_| sender.subscribe()) @@ -305,7 +305,7 @@ fn cancel_pending(bencher: Bencher) { let mut context = bench_context(); bencher.bench_local(|| { - let (sender, mut receiver) = unbounded::channel::(); + let (sender, mut receiver) = mpmc::unbounded::(); { let mut recv = pin!(receiver.recv()); poll_pending(recv.as_mut(), &mut context); @@ -319,7 +319,7 @@ fn deliver_to_waiter(bencher: Bencher) { let mut context = bench_context(); bencher.bench_local(|| { - let (sender, mut receiver) = unbounded::channel(); + let (sender, mut receiver) = mpmc::unbounded(); let mut recv = pin!(receiver.recv()); poll_pending(recv.as_mut(), &mut context); @@ -334,7 +334,7 @@ fn deliver_to_receiver_batch(bencher: Bencher, receiver_count: usize) { let mut context = bench_context(); bencher.bench_local(|| { - let (sender, receiver) = unbounded::channel(); + let (sender, receiver) = mpmc::unbounded(); drop(receiver); let mut receivers = (0..receiver_count) .map(|_| sender.subscribe()) diff --git a/tests-integration/tests/broadcast_unbounded_test.rs b/tests-integration/tests/broadcast_mpmc_unbounded_test.rs similarity index 88% rename from tests-integration/tests/broadcast_unbounded_test.rs rename to tests-integration/tests/broadcast_mpmc_unbounded_test.rs index 9287a9f..7617ac4 100644 --- a/tests-integration/tests/broadcast_unbounded_test.rs +++ b/tests-integration/tests/broadcast_mpmc_unbounded_test.rs @@ -26,7 +26,7 @@ use std::thread; use std::time::Duration; use std::time::Instant; -use asyncband::broadcast::unbounded::*; +use asyncband::broadcast::mpmc::*; struct TrackWake(AtomicUsize); @@ -39,7 +39,7 @@ impl Wake for TrackWake { /// A payload whose destructor re-enters the channel it was sent through. struct Reentrant { value: u64, - channel: Option>, + channel: Option>, } impl Clone for Reentrant { @@ -95,7 +95,7 @@ impl Rng { #[tokio::test] async fn test_broadcast_basic() { - let (tx, mut rx1) = channel(); + let (tx, mut rx1) = unbounded(); let mut rx2 = tx.subscribe(); tx.send(10); @@ -109,7 +109,7 @@ async fn test_broadcast_basic() { #[tokio::test] async fn test_subscribe() { - let (tx, _rx) = channel(); + let (tx, _rx) = unbounded(); let mut rx = tx.subscribe(); tx.send(100); @@ -118,7 +118,7 @@ async fn test_subscribe() { #[tokio::test] async fn test_resubscribe() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); tx.send(1); tx.send(2); @@ -137,7 +137,7 @@ async fn test_resubscribe() { #[test] fn test_try_recv() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); // Empty assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); @@ -154,7 +154,7 @@ fn test_try_recv() { #[tokio::test] async fn test_slow_receiver_keeps_every_message() { - let (tx, mut rx1) = channel(); + let (tx, mut rx1) = unbounded(); let mut rx2 = tx.subscribe(); for i in 0..1024 { @@ -175,7 +175,7 @@ async fn test_slow_receiver_keeps_every_message() { #[tokio::test] async fn buffer_len_tracks_the_slowest_receiver() { - let (tx, mut rx1) = channel(); + let (tx, mut rx1) = unbounded(); let mut rx2 = tx.subscribe(); tx.send(1); @@ -196,7 +196,7 @@ async fn buffer_len_tracks_the_slowest_receiver() { #[tokio::test] async fn test_dropping_a_lagging_receiver_releases_its_backlog() { - let (tx, mut rx1) = channel(); + let (tx, mut rx1) = unbounded(); let rx2 = tx.subscribe(); for i in 0..128 { @@ -213,7 +213,7 @@ async fn test_dropping_a_lagging_receiver_releases_its_backlog() { #[tokio::test] async fn resubscribe_keeps_the_original_receivers_backlog() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); tx.send(1); tx.send(2); @@ -234,7 +234,7 @@ async fn resubscribe_keeps_the_original_receivers_backlog() { #[tokio::test] async fn send_without_receivers_does_not_buffer() { - let (tx, rx) = channel(); + let (tx, rx) = unbounded(); drop(rx); tx.send(1); @@ -250,7 +250,7 @@ async fn send_without_receivers_does_not_buffer() { #[test] fn receiver_count_and_len_track_each_receiver() { - let (tx, mut rx1) = channel(); + let (tx, mut rx1) = unbounded(); assert_eq!(tx.receiver_count(), 1); assert_eq!(rx1.len(), 0); assert!(rx1.is_empty()); @@ -278,44 +278,6 @@ fn receiver_count_and_len_track_each_receiver() { 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); @@ -329,7 +291,7 @@ fn sole_receiver_takes_messages_without_cloning() { } } - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); for i in 0..8 { tx.send(CountClone(i)); assert_eq!(rx.try_recv().unwrap().0, i); @@ -346,7 +308,7 @@ fn sole_receiver_takes_messages_without_cloning() { #[test] fn panicking_clone_leaves_the_channel_consistent() { - let (tx, mut rx1) = channel(); + let (tx, mut rx1) = unbounded(); let mut rx2 = tx.subscribe(); tx.send(PanicOnClone { @@ -379,7 +341,7 @@ fn message_destructors_run_outside_the_channel_lock() { let flag = finished.clone(); let worker = thread::spawn(move || { - let (tx, mut rx1) = channel(); + let (tx, mut rx1) = unbounded(); let rx2 = tx.subscribe(); for value in 0..4 { @@ -419,7 +381,7 @@ fn message_destructors_run_outside_the_channel_lock() { #[tokio::test] async fn test_wait_mechanism() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); let handle = tokio::spawn(async move { rx.recv().await }); @@ -431,7 +393,7 @@ async fn test_wait_mechanism() { #[test] fn send_wakes_a_parked_receiver_exactly_once() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); let waker = Waker::from(tracker.clone()); let mut context = Context::from_waker(&waker); @@ -447,7 +409,7 @@ fn send_wakes_a_parked_receiver_exactly_once() { #[test] fn cancelled_recv_releases_its_waker() { - let (tx, mut rx) = channel::<()>(); + let (tx, mut rx) = unbounded::<()>(); let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); let waker = Waker::from(tracker.clone()); let baseline = Arc::strong_count(&tracker); @@ -467,7 +429,7 @@ fn cancelled_recv_releases_its_waker() { #[test] fn dropping_a_woken_recv_keeps_another_receivers_waiter() { - let (tx, mut rx1) = channel::(); + let (tx, mut rx1) = unbounded::(); let mut rx2 = tx.subscribe(); let first = Arc::new(TrackWake(AtomicUsize::new(0))); let waker = Waker::from(first.clone()); @@ -495,7 +457,7 @@ fn dropping_a_woken_recv_keeps_another_receivers_waiter() { #[test] fn parked_recv_wakes_when_the_last_sender_drops() { - let (tx, mut rx) = channel::<()>(); + let (tx, mut rx) = unbounded::<()>(); let extra = tx.clone(); let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); let waker = Waker::from(tracker.clone()); @@ -516,7 +478,7 @@ fn parked_recv_wakes_when_the_last_sender_drops() { #[test] fn parked_recv_prefers_buffered_messages_over_disconnect() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); let waker = Waker::from(tracker); let mut context = Context::from_waker(&waker); @@ -534,7 +496,7 @@ fn parked_recv_prefers_buffered_messages_over_disconnect() { #[tokio::test] async fn recv_drains_buffered_messages_before_reporting_disconnect() { - let (tx, mut rx) = channel(); + let (tx, mut rx) = unbounded(); tx.send(1); tx.send(2); @@ -547,7 +509,7 @@ async fn recv_drains_buffered_messages_before_reporting_disconnect() { #[tokio::test] async fn recv_reports_disconnect_without_any_message() { - let (tx, mut rx) = channel::<()>(); + let (tx, mut rx) = unbounded::<()>(); drop(tx); assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); } @@ -557,17 +519,10 @@ 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 (tx, rx) = unbounded(); + let mut receivers = Vec::with_capacity(4); + receivers.push(rx); + receivers.extend((1..4).map(|_| tx.subscribe())); let senders = (0..SENDERS) .map(|worker| { @@ -610,7 +565,7 @@ fn concurrent_senders_deliver_every_message_to_every_receiver() { 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 (tx, rx) = unbounded::(); let mut tail = 0u64; let mut model = vec![(rx, 0u64)]; diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index b51d114..847de16 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -86,10 +86,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::(); + assert_send_and_sync::(); assert_send_and_sync::>(); assert_send_and_sync::>(); assert_send_and_sync::>(); @@ -137,10 +137,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::(); + assert_unpin::(); assert_unpin::>(); assert_unpin::>(); assert_unpin::>(); From c75ac73ab3043dff1dc18a67a28ccebd2f3d6727 Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 27 Aug 2026 10:37:47 +0800 Subject: [PATCH 5/8] move file Signed-off-by: tison --- .../src/channel/broadcast/mpmc/{unbounded.rs => unbounded/mod.rs} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename asyncband/src/channel/broadcast/mpmc/{unbounded.rs => unbounded/mod.rs} (100%) diff --git a/asyncband/src/channel/broadcast/mpmc/unbounded.rs b/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs similarity index 100% rename from asyncband/src/channel/broadcast/mpmc/unbounded.rs rename to asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs From c5bd0a7c3a8e0e77903d6815f827e6c66d00d0fa Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 27 Aug 2026 10:42:52 +0800 Subject: [PATCH 6/8] fixup Signed-off-by: tison --- .../src/channel/broadcast/mpmc/unbounded/mod.rs | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs index 594a125..c2d3b4e 100644 --- a/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs @@ -38,9 +38,9 @@ //! [`UnboundedReceiver::resubscribe`] to create a receiver that starts 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. +//! receiver. Only the receiver 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 //! @@ -265,7 +265,7 @@ impl Inner { 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. + // is `msg`. `take_msg` relies on this to recognize that it owns the payload. debug_assert!( reclaimed .first() @@ -609,10 +609,7 @@ fn take_msg(msg: Arc, reclaimed: Vec>) -> T { // 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(), - } + Arc::try_unwrap(msg).unwrap_or_else(|msg| (*msg).clone()) } impl UnboundedReceiver { From 53f0db4d13395dda4f23a5062b49cb3de4a8ef9f Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 27 Aug 2026 11:01:59 +0800 Subject: [PATCH 7/8] perf(broadcast): avoid single-message reclaim allocation --- .../channel/broadcast/mpmc/unbounded/mod.rs | 125 +++++++++++------- 1 file changed, 77 insertions(+), 48 deletions(-) diff --git a/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs index c2d3b4e..5eacfe1 100644 --- a/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/channel/broadcast/mpmc/unbounded/mod.rs @@ -206,6 +206,37 @@ struct Inner { waiters: WaitSet, } +/// Messages removed from the shared buffer and waiting to be dropped after it is unlocked. +/// +/// Keeping the first message out of the `Vec` avoids a heap allocation on the common path where +/// one receive reclaims exactly one message. +struct Reclaimed { + first: Option>, + rest: Vec>, +} + +impl Reclaimed { + fn empty() -> Self { + Self { + first: None, + rest: Vec::new(), + } + } + + fn first(&self) -> Option<&Arc> { + self.first.as_ref() + } + + fn is_empty(&self) -> bool { + self.first.is_none() + } + + fn drop_messages(self) { + let Self { first, rest } = self; + drop((first, rest)); + } +} + impl Inner { fn insert_receiver(&mut self, head: u64) -> SlotId { if head == self.head { @@ -215,69 +246,60 @@ impl Inner { self.receivers.insert(head) } - fn remove_receiver(&mut self, key: SlotId) -> Vec> { + fn remove_receiver(&mut self, key: SlotId) -> Reclaimed { let head = self.receivers.remove(key); if head == self.head { self.release_head_receiver() } else { - Vec::new() + Reclaimed::empty() } } - fn advance_receiver(&mut self, key: SlotId, 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> { + fn release_head_receiver(&mut self) -> Reclaimed { self.head_receivers -= 1; if self.head_receivers == 0 { self.reclaim_consumed() } else { - Vec::new() + Reclaimed::empty() } } - fn receive(&mut self, key: SlotId) -> Option<(Arc, Vec>)> { - let head = *self - .receivers - .get(key) - .expect("active broadcast receiver must be registered"); + fn receive(&mut self, key: SlotId) -> Option<(Arc, Reclaimed)> { + let head = { + let cursor = self + .receivers + .get_mut(key) + .expect("active broadcast receiver must be registered"); + if *cursor >= self.tail { + return None; + } + let head = *cursor; + *cursor += 1; + head + }; - 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); - // 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 recognize that it owns the payload. - debug_assert!( - reclaimed - .first() - .is_none_or(|first| Arc::ptr_eq(first, &msg)) - ); - Some((msg, reclaimed)) + debug_assert!(head >= self.head); + let offset = (head - self.head) as usize; + let msg = self.buffer[offset].clone(); + let reclaimed = if head == self.head { + self.release_head_receiver() } else { - None - } + Reclaimed::empty() + }; + // 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 recognize that it owns the payload. + debug_assert!( + reclaimed + .first() + .is_none_or(|first| Arc::ptr_eq(first, &msg)) + ); + Some((msg, reclaimed)) } - fn reclaim_consumed(&mut self) -> Vec> { + fn reclaim_consumed(&mut self) -> Reclaimed { let mut next_head = self.tail; let mut head_receivers = 0; @@ -293,8 +315,15 @@ impl Inner { 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(); + // Move reclaimed messages out so their Drop impls run after `inner` is unlocked. Keep the + // first one separate so the usual one-message reclaim does not allocate another buffer. + let first = if consumed == 0 { + None + } else { + self.buffer.pop_front() + }; + let rest = self.buffer.drain(..consumed.saturating_sub(1)).collect(); + let reclaimed = Reclaimed { first, rest }; self.head = next_head; self.head_receivers = head_receivers; @@ -599,9 +628,9 @@ impl UnboundedReceiver { /// 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 { +fn take_msg(msg: Arc, reclaimed: Reclaimed) -> T { let sole_owner = !reclaimed.is_empty(); - drop(reclaimed); + reclaimed.drop_messages(); if !sole_owner { return (*msg).clone(); @@ -613,7 +642,7 @@ fn take_msg(msg: Arc, reclaimed: Vec>) -> T { } impl UnboundedReceiver { - fn try_recv_shared(&mut self) -> Result<(Arc, Vec>), TryRecvError> { + fn try_recv_shared(&mut self) -> Result<(Arc, Reclaimed), 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. From ebca8ce6baa076308bc426ff701ac7643a76a869 Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 27 Aug 2026 11:02:04 +0800 Subject: [PATCH 8/8] test(broadcast): compare MPMC ecosystem performance --- Cargo.lock | 13 ++ Cargo.toml | 1 + benchmarks/Cargo.toml | 1 + benchmarks/ecosystem/broadcast/mod.rs | 18 ++ .../ecosystem/broadcast/mpmc/adapters.rs | 130 ++++++++++++++ benchmarks/ecosystem/broadcast/mpmc/mod.rs | 20 +++ .../ecosystem/broadcast/mpmc/support.rs | 159 ++++++++++++++++++ .../ecosystem/broadcast/mpmc/unbounded.rs | 86 ++++++++++ benchmarks/ecosystem/main.rs | 1 + 9 files changed, 429 insertions(+) create mode 100644 benchmarks/ecosystem/broadcast/mod.rs create mode 100644 benchmarks/ecosystem/broadcast/mpmc/adapters.rs create mode 100644 benchmarks/ecosystem/broadcast/mpmc/mod.rs create mode 100644 benchmarks/ecosystem/broadcast/mpmc/support.rs create mode 100644 benchmarks/ecosystem/broadcast/mpmc/unbounded.rs diff --git a/Cargo.lock b/Cargo.lock index 0c320cf..d293387 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -52,6 +52,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "async-broadcast" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "435a87a52755b8f27fcf321ac4f04b2802e337c8c4872923137471ec39c37532" +dependencies = [ + "event-listener", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + [[package]] name = "async-channel" version = "2.5.0" @@ -82,6 +94,7 @@ checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" name = "benchmarks" version = "0.0.0" dependencies = [ + "async-broadcast", "async-channel", "asyncband", "divan", diff --git a/Cargo.toml b/Cargo.toml index 872eae1..11086fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ asyncband = { path = "asyncband" } hashbrown = { version = "0.17.1", default-features = false } # Dev dependencies +async-broadcast = { version = "0.7.2" } async-channel = { version = "2.5.0" } cargo_metadata = { version = "0.23.1" } clap = { version = "4.6.5" } diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 95eec00..15a724f 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -23,6 +23,7 @@ edition.workspace = true rust-version.workspace = true [dev-dependencies] +async-broadcast = { workspace = true } async-channel = { workspace = true } asyncband = { workspace = true, features = [ "barrier", diff --git a/benchmarks/ecosystem/broadcast/mod.rs b/benchmarks/ecosystem/broadcast/mod.rs new file mode 100644 index 0000000..b078f4a --- /dev/null +++ b/benchmarks/ecosystem/broadcast/mod.rs @@ -0,0 +1,18 @@ +// 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. + +mod mpmc; diff --git a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs new file mode 100644 index 0000000..7b98870 --- /dev/null +++ b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs @@ -0,0 +1,130 @@ +// 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::task::Context; + +use crate::support::poll_ready; + +pub struct Asyncband; +pub struct Tokio; +pub struct AsyncBroadcast; + +pub trait BroadcastMpmc: Send + Sync + 'static { + type Sender: Clone + Send + 'static; + type Receiver: Send + 'static; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec); + fn send(sender: &Self::Sender, value: usize); + fn try_recv(receiver: &mut Self::Receiver) -> Option; + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize; +} + +impl BroadcastMpmc for Asyncband { + type Receiver = asyncband::broadcast::mpmc::UnboundedReceiver; + type Sender = asyncband::broadcast::mpmc::UnboundedSender; + + fn channel(_capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = asyncband::broadcast::mpmc::unbounded(); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(sender.subscribe()); + } + (sender, receivers) + } + + fn send(sender: &Self::Sender, value: usize) { + sender.send(value); + } + + fn try_recv(receiver: &mut Self::Receiver) -> Option { + match receiver.try_recv() { + Ok(value) => Some(value), + Err(asyncband::broadcast::mpmc::TryRecvError::Empty) => None, + Err(asyncband::broadcast::mpmc::TryRecvError::Disconnected) => { + panic!("asyncband channel closed during benchmark") + } + } + } + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(receiver.recv(), context).unwrap() + } +} + +impl BroadcastMpmc for Tokio { + type Receiver = tokio::sync::broadcast::Receiver; + type Sender = tokio::sync::broadcast::Sender; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = tokio::sync::broadcast::channel(capacity); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(sender.subscribe()); + } + (sender, receivers) + } + + fn send(sender: &Self::Sender, value: usize) { + sender.send(value).unwrap(); + } + + fn try_recv(receiver: &mut Self::Receiver) -> Option { + match receiver.try_recv() { + Ok(value) => Some(value), + Err(tokio::sync::broadcast::error::TryRecvError::Empty) => None, + Err(error) => panic!("unexpected Tokio receive error: {error}"), + } + } + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(receiver.recv(), context).unwrap() + } +} + +impl BroadcastMpmc for AsyncBroadcast { + type Receiver = async_broadcast::Receiver; + type Sender = async_broadcast::Sender; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = async_broadcast::broadcast(capacity); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + let receiver = receivers[0].clone(); + receivers.push(receiver); + } + (sender, receivers) + } + + fn send(sender: &Self::Sender, value: usize) { + sender.try_broadcast(value).unwrap(); + } + + fn try_recv(receiver: &mut Self::Receiver) -> Option { + match receiver.try_recv() { + Ok(value) => Some(value), + Err(async_broadcast::TryRecvError::Empty) => None, + Err(error) => panic!("unexpected async-broadcast receive error: {error}"), + } + } + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(receiver.recv_direct(), context).unwrap() + } +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/mod.rs b/benchmarks/ecosystem/broadcast/mpmc/mod.rs new file mode 100644 index 0000000..5e86ce6 --- /dev/null +++ b/benchmarks/ecosystem/broadcast/mpmc/mod.rs @@ -0,0 +1,20 @@ +// 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. + +mod adapters; +mod support; +mod unbounded; diff --git a/benchmarks/ecosystem/broadcast/mpmc/support.rs b/benchmarks/ecosystem/broadcast/mpmc/support.rs new file mode 100644 index 0000000..55c016a --- /dev/null +++ b/benchmarks/ecosystem/broadcast/mpmc/support.rs @@ -0,0 +1,159 @@ +// 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::marker::PhantomData; +use std::sync::Arc; +use std::sync::Barrier; +use std::thread; +use std::thread::JoinHandle; + +use divan::black_box; + +use super::adapters::BroadcastMpmc; + +pub const BATCH_MESSAGES: usize = 4096; +pub const PRODUCER_COUNTS: &[usize] = &[1, 2, 4, 8]; +pub const RECEIVER_COUNTS: &[usize] = &[1, 2, 4, 8, 32]; +pub const ROUND_TRIP_CAPACITY: usize = 64; + +fn recv(receiver: &mut C::Receiver) -> usize { + C::try_recv(receiver).expect("the published benchmark batch must be ready") +} + +pub struct ConcurrentSend { + receiver: C::Receiver, + start: Arc, + done: Arc, + workers: Vec>, + channel: PhantomData, +} + +impl ConcurrentSend { + pub fn new(producer_count: usize) -> Self { + assert_eq!(BATCH_MESSAGES % producer_count, 0); + let (sender, mut receivers) = C::channel(BATCH_MESSAGES, 1); + let receiver = receivers.pop().unwrap(); + let start = Arc::new(Barrier::new(producer_count + 1)); + let done = Arc::new(Barrier::new(producer_count + 1)); + let messages_per_producer = BATCH_MESSAGES / producer_count; + let workers = (0..producer_count) + .map(|producer| { + let sender = sender.clone(); + let start = start.clone(); + let done = done.clone(); + thread::spawn(move || { + start.wait(); + let first = producer * messages_per_producer; + for value in first..first + messages_per_producer { + C::send(&sender, black_box(value)); + } + done.wait(); + }) + }) + .collect(); + drop(sender); + + Self { + receiver, + start, + done, + workers, + channel: PhantomData, + } + } + + pub fn run(&mut self) -> usize { + self.start.wait(); + self.done.wait(); + + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(recv::(&mut self.receiver)); + } + black_box(checksum) + } +} + +impl Drop for ConcurrentSend { + fn drop(&mut self) { + let panicking = thread::panicking(); + for worker in self.workers.drain(..) { + let result = worker.join(); + if !panicking { + result.expect("benchmark producer panicked"); + } + } + } +} + +pub struct Fanout { + sender: C::Sender, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl Fanout { + pub fn new(receiver_count: usize) -> Self { + let (sender, receivers) = C::channel(BATCH_MESSAGES, receiver_count); + let start = Arc::new(Barrier::new(receiver_count + 1)); + let done = Arc::new(Barrier::new(receiver_count + 1)); + let workers = receivers + .into_iter() + .map(|mut receiver| { + let start = start.clone(); + let done = done.clone(); + thread::spawn(move || { + start.wait(); + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(recv::(&mut receiver)); + } + black_box(checksum); + done.wait(); + }) + }) + .collect(); + + Self { + sender, + start, + done, + workers, + } + } + + pub fn run(&mut self) { + for value in 0..BATCH_MESSAGES { + C::send(&self.sender, black_box(value)); + } + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for Fanout { + fn drop(&mut self) { + let panicking = thread::panicking(); + for worker in self.workers.drain(..) { + let result = worker.join(); + if !panicking { + result.expect("benchmark receiver panicked"); + } + } + } +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/unbounded.rs b/benchmarks/ecosystem/broadcast/mpmc/unbounded.rs new file mode 100644 index 0000000..71e1214 --- /dev/null +++ b/benchmarks/ecosystem/broadcast/mpmc/unbounded.rs @@ -0,0 +1,86 @@ +// 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. + +// Asyncband is the only unbounded channel in this comparison. Tokio broadcast overwrites messages +// at capacity, while async-broadcast applies backpressure by default. Every bounded peer gets room +// for the entire measured batch, so these workloads compare their common lossless, non-blocking +// path rather than their different lag and capacity policies. + +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use super::adapters::AsyncBroadcast; +use super::adapters::Asyncband; +use super::adapters::BroadcastMpmc; +use super::adapters::Tokio; +use super::support::BATCH_MESSAGES; +use super::support::ConcurrentSend; +use super::support::Fanout; +use super::support::PRODUCER_COUNTS; +use super::support::RECEIVER_COUNTS; +use super::support::ROUND_TRIP_CAPACITY; +use crate::support::bench_context; + +#[divan::bench(types = [Asyncband, Tokio, AsyncBroadcast])] +fn try_round_trip(bencher: Bencher) { + let (sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); + let mut receiver = receivers.pop().unwrap(); + + bencher.bench_local(|| { + C::send(&sender, black_box(usize::MAX)); + black_box(C::try_recv(&mut receiver).unwrap()) + }); +} + +#[divan::bench(types = [Asyncband, Tokio, AsyncBroadcast])] +fn ready_round_trip(bencher: Bencher) { + let mut context = bench_context(); + let (sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); + let mut receiver = receivers.pop().unwrap(); + + bencher.bench_local(|| { + C::send(&sender, black_box(usize::MAX)); + black_box(C::recv_ready(&mut receiver, &mut context)) + }); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncBroadcast], + args = PRODUCER_COUNTS, + sample_count = 20, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn concurrent_producers(bencher: Bencher, producer_count: usize) { + bencher + .with_inputs(|| ConcurrentSend::::new(producer_count)) + .bench_local_refs(ConcurrentSend::run); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncBroadcast], + args = RECEIVER_COUNTS, + sample_count = 20, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn fanout(bencher: Bencher, receiver_count: usize) { + bencher + .with_inputs(|| Fanout::::new(receiver_count)) + .bench_local_refs(Fanout::run); +} diff --git a/benchmarks/ecosystem/main.rs b/benchmarks/ecosystem/main.rs index ae6037b..6fd069a 100644 --- a/benchmarks/ecosystem/main.rs +++ b/benchmarks/ecosystem/main.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod broadcast; mod mpsc; #[allow(dead_code)]