From 4ffb990603d51b684fdafe428621085dbbef3f98 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Mon, 31 Aug 2026 10:31:57 +0800 Subject: [PATCH 1/8] feat(broadcast): add a bounded MPMC broadcast channel Add `broadcast::mpmc::bounded`, a lossless bounded broadcast channel. Every accepted value stays readable by every subscription that was active when it was accepted, so a receive never reports lag. Capacity counts the shared backlog held by the slowest active subscription: `send` waits on it and `try_send` reports `TrySendError::Full` without taking the value. Nothing slides, overwrites, or is dropped to make room. Move the backlog, the cursors, and the one receive poll step both families share into a private `common` module, while each keeps its own publish path and `Recv` future so neither retention contract hides behind a shared abstraction. Ring, cursor, retention, and sequencing types stay private to `broadcast::mpmc`. Signed-off-by: onenewcode --- CHANGELOG.md | 1 + README.md | 2 +- asyncband/src/broadcast/mpmc/bounded/mod.rs | 732 ++++++++++++++++++ asyncband/src/broadcast/mpmc/bounded/tests.rs | 90 +++ asyncband/src/broadcast/mpmc/common.rs | 527 +++++++++++++ asyncband/src/broadcast/mpmc/error.rs | 103 +++ asyncband/src/broadcast/mpmc/mod.rs | 17 +- asyncband/src/broadcast/mpmc/unbounded/mod.rs | 402 ++-------- .../src/broadcast/mpmc/unbounded/tests.rs | 11 +- asyncband/src/internal/mod.rs | 9 +- asyncband/src/lib.rs | 2 +- .../asyncband/broadcast/mpmc/bounded.rs | 145 ++++ benchmarks/asyncband/broadcast/mpmc/mod.rs | 1 + .../ecosystem/broadcast/mpmc/adapters.rs | 108 +++ .../ecosystem/broadcast/mpmc/bounded.rs | 84 ++ benchmarks/ecosystem/broadcast/mpmc/mod.rs | 1 + .../ecosystem/broadcast/mpmc/support.rs | 150 ++++ .../tests/broadcast_mpmc_bounded_test.rs | 732 ++++++++++++++++++ tests-integration/tests/traits_test.rs | 6 + .../tests/waitset_reentrancy_test.rs | 17 + 20 files changed, 2771 insertions(+), 369 deletions(-) create mode 100644 asyncband/src/broadcast/mpmc/bounded/mod.rs create mode 100644 asyncband/src/broadcast/mpmc/bounded/tests.rs create mode 100644 asyncband/src/broadcast/mpmc/common.rs create mode 100644 asyncband/src/broadcast/mpmc/error.rs create mode 100644 benchmarks/asyncband/broadcast/mpmc/bounded.rs create mode 100644 benchmarks/ecosystem/broadcast/mpmc/bounded.rs create mode 100644 tests-integration/tests/broadcast_mpmc_bounded_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 3be44e7c..61c8dac4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ All notable changes to this project will be documented in this file. ### New features * Implement `broadcast::mpmc::unbounded`, an unbounded broadcast channel that retains messages until all active receivers consume them or are dropped. +* Implement `broadcast::mpmc::bounded`, a lossless bounded broadcast channel that retains at most the requested capacity and makes producers wait for the slowest active receiver. * Add an opt-in latest-state channel under `asyncband::watch`. * Add opt-in `asyncband::event::ManualResetEvent`, a reusable level-triggered signal that releases registered waits and remains ready for future waits until explicitly reset. * Add an opt-in shared one-shot completion primitive under `asyncband::completion` with a single-use completer, cloneable observers, a retained borrowed result, and observable abandonment. diff --git a/README.md b/README.md index 34e1e14c..4c360603 100644 --- a/README.md +++ b/README.md @@ -79,7 +79,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`Shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/struct.Shutdown.html) | `shutdown` | Coordinate shutdown signals and completion. | | Channels | [`oneshot`](https://docs.rs/asyncband/*/asyncband/oneshot/) | `oneshot` | Send one value from one sender to one receiver. | | | [`mpsc`](https://docs.rs/asyncband/*/asyncband/mpsc/) | `mpsc` | Send each value from multiple producers to one receiver through a bounded or unbounded queue. | -| | [`broadcast`](https://docs.rs/asyncband/*/asyncband/broadcast/) | `broadcast` | Broadcast values from one or more producers and retain them until every active receiver consumes them. | +| | [`broadcast`](https://docs.rs/asyncband/*/asyncband/broadcast/) | `broadcast` | Broadcast every value to all active receivers, with bounded backpressure or unbounded retention. | | | [`watch`](https://docs.rs/asyncband/*/asyncband/watch/) | `watch` | Publish the latest state to independently tracked receivers and coalesce intermediate updates. | | 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. | diff --git a/asyncband/src/broadcast/mpmc/bounded/mod.rs b/asyncband/src/broadcast/mpmc/bounded/mod.rs new file mode 100644 index 00000000..42066397 --- /dev/null +++ b/asyncband/src/broadcast/mpmc/bounded/mod.rs @@ -0,0 +1,732 @@ +// 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 a bounded buffer. +//! +//! This channel supports multiple senders and multiple receivers. Each message sent by any sender +//! is received by all active receivers. Nothing is ever displaced to make room, so a receive never +//! reports lag; instead the channel retains at most `capacity` messages and makes producers wait. +//! +//! # Capacity +//! +//! Capacity counts the *shared* backlog — the messages retained because the slowest active +//! receiver has not read them yet — not messages per receiver. Adding receivers therefore does not +//! consume capacity; falling behind does. +//! +//! Because the backlog is shared, a single receiver that stops draining stalls **every** producer +//! on the channel, however many other receivers are keeping up. That is what "the slowest +//! subscription exerts backpressure" means, and it is the trade a lossless bounded broadcast +//! makes. Drop a receiver that will not drain, and its backlog is released immediately. +//! +//! If no receivers are active the channel retains nothing, so a send never waits. +//! +//! # Receivers +//! +//! Each receiver has an independent cursor. Use [`BoundedSender::subscribe`] or +//! [`BoundedReceiver::resubscribe`] to create a receiver that starts at the current tail. A new +//! subscription never sees messages published before it existed. +//! +//! # Fairness +//! +//! Waiting producers are woken as capacity frees, but capacity is not reserved for them: a +//! producer calling [`BoundedSender::try_send`] can take a slot that a woken producer was about to +//! use, and that producer then waits again. Publication itself is one indivisible step, so +//! cancelling a send can never leave a gap in the committed order. +//! +//! # Examples +//! +//! Basic usage: +//! +//! ``` +//! use asyncband::broadcast::mpmc; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (tx, mut rx1) = mpmc::bounded(4); +//! let mut rx2 = tx.subscribe(); +//! +//! tx.send(10).await; +//! tx.send(20).await; +//! +//! 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)); +//! # } +//! ``` +//! +//! The slowest receiver holds the capacity: +//! +//! ``` +//! use asyncband::broadcast::mpmc; +//! use asyncband::broadcast::mpmc::TrySendError; +//! +//! let (tx, mut rx1) = mpmc::bounded(2); +//! let rx2 = tx.subscribe(); +//! +//! tx.try_send(1).unwrap(); +//! tx.try_send(2).unwrap(); +//! assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); +//! +//! // `rx1` draining is not enough: `rx2` has read neither message, so both stay retained. +//! assert_eq!(rx1.try_recv(), Ok(1)); +//! assert_eq!(tx.retained_message_count(), 2); +//! assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); +//! +//! // Dropping the lagging receiver releases the backlog only it was holding. `rx1` has still not +//! // read the second message, so that one stays. +//! drop(rx2); +//! assert_eq!(tx.retained_message_count(), 1); +//! tx.try_send(3).unwrap(); +//! ``` + +use std::fmt; +use std::future::Future; +use std::future::poll_fn; +use std::pin::Pin; +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 super::common; +use super::common::Backlog; +use super::common::Inner; +use super::error::RecvError; +use super::error::TryRecvError; +use super::error::TrySendError; +use crate::internal::arena::SlotId; +use crate::internal::mutex::Mutex; +use crate::internal::semaphore::Acquire; +use crate::internal::semaphore::Semaphore; +use crate::internal::waitset::WakerToken; +use crate::internal::waitset::wake_all; + +#[cfg(test)] +mod tests; + +/// Creates a new broadcast channel that retains at most `capacity` messages. +/// +/// Every accepted value stays readable by every receiver that was active when it was accepted. +/// Once `capacity` messages are retained, [`BoundedSender::send`] waits and +/// [`BoundedSender::try_send`] reports [`TrySendError::Full`] until the slowest active receiver +/// consumes a message or is dropped. +/// +/// # Panics +/// +/// Panics if `capacity` is zero. +/// +/// # Examples +/// +/// ``` +/// use asyncband::broadcast::mpmc; +/// +/// let (tx, mut rx) = mpmc::bounded(1); +/// tx.try_send(10).unwrap(); +/// assert_eq!(rx.try_recv(), Ok(10)); +/// ``` +#[track_caller] +pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { + assert!( + capacity > 0, + "broadcast bounded channel requires capacity > 0" + ); + + let (inner, key) = Inner::with_first_subscription(Backlog::fixed(capacity)); + let shared = Arc::new(Shared { + inner, + senders: AtomicUsize::new(1), + capacity, + tx_permits: Semaphore::new(0), + blocked_senders: AtomicUsize::new(0), + }); + let sender = BoundedSender { + shared: shared.clone(), + }; + let receiver = BoundedReceiver { shared, key }; + (sender, receiver) +} + +struct Shared { + /// Buffer, receiver cursors, and parked receivers, all under a single lock. + inner: Mutex>, + /// Number of active senders. + senders: AtomicUsize, + /// The logical limit on the retained backlog. + capacity: usize, + /// Producers parked in [`BoundedSender::send`]. + /// + /// A permit here is a wake-up hint, not a reserved slot: a woken producer rechecks the backlog + /// and parks again if another producer took the space first. The semaphore starts empty and + /// only ever grows when a reclaim finds someone waiting, so an idle channel accumulates none. + tx_permits: Semaphore, + /// How many producers are somewhere inside the waiting path of [`BoundedSender::send`]. + /// + /// An upper bound on the number of parked producers, and the only thing either release path + /// consults. It answers both questions a reclaim has — whether to wake anyone, and how many + /// permits are worth handing out — without taking the semaphore's lock. Reclaiming is far more + /// frequent than blocking — under fan-out every message is reclaimed, while a channel with + /// headroom never blocks at all — so paying an atomic load there instead of a lock acquisition + /// is what keeps an uncontended receive off the semaphore entirely. + blocked_senders: AtomicUsize, +} + +impl Shared { + /// Hands `freed` released slots back to producers parked in `send`. + /// + /// Capacity is `retained()`, which is `buffer.len()`. The buffer grows only in + /// `Backlog::publish_retained` and shrinks only in `Backlog::reclaim_consumed`, which is + /// reachable from exactly two places: a receive that vacates the last cursor at the backlog + /// head, and removing a subscription. Those are the only callers of this method, so no path + /// can free capacity without waking a producer. Subscribing cannot: a new cursor starts at the + /// tail and never lowers `retained()`. + /// + /// Callers must invoke this with the channel unlocked, and — on the receive path — before + /// touching the payload, since `common::take_msg` runs user code that may panic. + fn release_reclaimed(&self, freed: usize) { + // Release no more permits than there are producers to wake. A permit the semaphore cannot + // hand to a waiter is kept as slack, and the next producer to block has to burn it off one + // futile publish attempt — a channel lock apiece — at a time before it can park. Freeing a + // large prefix at once is not exotic: dropping a lagging subscription reclaims the whole + // backlog, which would otherwise leave nearly `capacity` permits behind. + // + // Capping cannot lose a wake-up, by the same argument that lets this read the count at all: + // a producer this load observes is one the release covers, and one it misses incremented + // after the load, which it does before taking the channel lock to recheck — so its recheck + // runs after the reclaim and finds the capacity itself. + let waiting = self.waiting_senders(); + if freed > 0 && waiting > 0 { + self.tx_permits.release_if_nonempty(freed.min(waiting)); + } + } + + /// Wakes every parked producer, however many slots came back. + /// + /// The last subscription leaving is not a reclaim of some number of slots — it removes the + /// limit itself, because a channel with no receivers discards instead of retaining. Releasing + /// only as many permits as that final reclaim freed would strand every producer beyond that + /// count, so this is the one release that must be unbounded. + fn release_all(&self) { + if self.waiting_senders() > 0 { + self.tx_permits.notify_all(); + } + } + + /// How many producers might be waiting, answered without touching the semaphore's lock. + /// + /// This cannot miss a wake-up. A producer increments the count before it ever takes the + /// channel lock to recheck capacity, and every caller here loads it after releasing that same + /// lock, so the mutex orders the two: either this load observes the producer, or the + /// producer's recheck runs after the change and finds the capacity itself. + fn waiting_senders(&self) -> usize { + self.blocked_senders.load(Ordering::Acquire) + } +} + +/// The sending side of a bounded broadcast channel. +/// +/// The sender can be cloned to create multiple producers. Dropping the final sender disconnects +/// the channel. Each receiver may drain its own buffered messages before observing disconnection. +pub struct BoundedSender { + shared: Arc>, +} + +impl Clone for BoundedSender { + fn clone(&self) -> Self { + // Relaxed is enough because this count publishes nothing on its own: receivers read it + // only to decide whether any sender remains, 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(), + } + } +} + +impl fmt::Debug for BoundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedSender").finish_non_exhaustive() + } +} + +impl Drop for BoundedSender { + fn drop(&mut self) { + match self.shared.senders.fetch_sub(1, Ordering::AcqRel) { + // Only parked receivers need waking. A parked producer borrows a live sender for the + // duration of its `send`, so the last sender cannot be dropping while one exists. + 1 => common::disconnect(&self.shared.inner), + _ => { + // there are still other senders left, do nothing + } + } + } +} + +impl BoundedSender { + /// Broadcasts a value to all active receivers, waiting for capacity if the channel is full. + /// + /// The wait ends when the slowest active receiver consumes a retained message or is dropped. + /// If no receivers are active, the message is dropped immediately and this returns without + /// waiting. + /// + /// # Cancel safety + /// + /// This method is cancel safe in the sense that matters for a lossless log: the value is + /// either published to every active receiver or not published at all. Publication happens in + /// one indivisible step, so a cancelled send cannot leave a reserved but unfilled position in + /// the committed order. A send cancelled before it published drops the value with the future. + /// + /// # 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::mpmc; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, mut rx) = mpmc::bounded(1); + /// tx.send(10).await; + /// assert_eq!(rx.recv().await, Ok(10)); + /// # } + /// ``` + pub async fn send(&self, value: T) { + let value = match self.try_send(value) { + Ok(()) => return, + Err(TrySendError::Full(value)) => value, + }; + + struct SendState<'a, T> { + sender: &'a BoundedSender, + // Declared before `value` so a cancelled send hands its registration back to the next + // waiting producer before running the payload's destructor. + acquire: Acquire<'a>, + // Boxed once, out of the critical section, and reused by every retry. + value: Option>, + } + + impl Drop for SendState<'_, T> { + fn drop(&mut self) { + // Runs before the fields, so the count drops while `acquire` is still queued. + // That ordering is safe in the direction that matters: the window can only make a + // receiver skip a wake-up this producer no longer wants, because the future being + // dropped is exactly the one leaving. Every other waiting producer still holds its + // own increment, so the count cannot reach zero while one of them needs waking. + self.sender + .shared + .blocked_senders + .fetch_sub(1, Ordering::Release); + } + } + + impl SendState<'_, T> { + fn poll_send(&mut self, cx: &mut Context<'_>) -> Poll<()> { + let mut msg = match self.value.take() { + Some(msg) => msg, + None => return Poll::Ready(()), + }; + + loop { + // Enqueue before rechecking. `release_if_nonempty` adds nothing when no + // producer is queued, so a reclaim landing between the recheck and the + // registration would otherwise drop its wake-up and park this producer for + // good. Registering first orders this producer's semaphore acquisition ahead + // of the reclaim's, so either the recheck sees the freed slot or the reclaim + // sees this waiter. + let poll = pin!(&mut self.acquire).poll(cx); + + msg = match self.sender.try_publish(msg) { + Ok(()) => return Poll::Ready(()), + Err(msg) => msg, + }; + + if poll.is_ready() { + self.acquire = self.sender.shared.tx_permits.poll_acquire(1); + } else { + self.value = Some(msg); + return Poll::Pending; + } + } + } + } + + // Announce this producer before it can recheck capacity, so a concurrent reclaim either + // sees it here or is seen by that recheck. + self.shared.blocked_senders.fetch_add(1, Ordering::Release); + let acquire = self.shared.tx_permits.poll_acquire(1); + let mut send = SendState { + sender: self, + acquire, + value: Some(Arc::new(value)), + }; + poll_fn(|cx| send.poll_send(cx)).await + } + + /// Attempts to broadcast a value to all active receivers without waiting. + /// + /// # Returns + /// + /// * `Ok(())`: The value was published, or discarded because no receivers are active. + /// * `Err(TrySendError::Full(value))`: The channel already retains `capacity` messages. The + /// value was not published and is returned unchanged. + /// + /// # Panics + /// + /// Panics if the internal message version counter overflows. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// use asyncband::broadcast::mpmc::TrySendError; + /// + /// let (tx, mut rx) = mpmc::bounded(1); + /// tx.try_send(10).unwrap(); + /// assert_eq!(tx.try_send(20), Err(TrySendError::Full(20))); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// tx.try_send(20).unwrap(); + /// ``` + pub fn try_send(&self, value: T) -> Result<(), TrySendError> { + // `Arc::new` runs inside the critical section, but only after the capacity check, so a + // rejected send never allocates. Unlike `T::clone` and `T::drop` it cannot run user code + // that reenters this channel, so it is safe to hold the lock across it. Hoisting it out + // measured no faster even with eight producers contending — the allocator's thread-local + // cache already makes it cheap — and it measured slower wherever sends block, because a + // rejected send would then allocate and free once before `send` boxes the value for real. + self.publish(value, Arc::new).map_err(TrySendError::Full) + } + + /// Publishes a message that is already boxed, handing it back if the channel is still full. + /// + /// This is the retry step of a waiting `send`, which boxes once with the channel unlocked and + /// then reuses that `Arc` for every attempt rather than reallocating per retry. + fn try_publish(&self, msg: Arc) -> Result<(), Arc> { + self.publish(msg, |msg| msg) + } + + /// The publish step both send paths share. + /// + /// `into_msg` is called only once this decides the message will actually be retained, which is + /// what lets `try_send` defer its allocation past the capacity check while `try_publish` hands + /// over an `Arc` it allocated with the channel unlocked. + /// + /// 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. + fn publish

(&self, payload: P, into_msg: impl FnOnce(P) -> Arc) -> Result<(), P> { + let mut discarded = None; + let wakers = { + let mut inner = self.shared.inner.lock(); + + if !inner.log.has_receivers() { + // Nothing can read this message. The payload leaves the critical section with us + // and is dropped below, so `T::drop` never runs under the lock. + inner.log.publish_discarded(); + discarded = Some(payload); + } else if inner.log.retained() == self.shared.capacity { + // Nothing was published, so there is no wait set to drain. + return Err(payload); + } else { + inner.log.publish_retained(into_msg(payload)); + } + + inner.waiters.drain() + }; + + wake_all(wakers); + drop(discarded); + Ok(()) + } + + /// Returns the number of messages currently retained by the channel. + /// + /// 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, and it is what this channel measures + /// against its [`capacity`](BoundedSender::capacity). + /// + /// The returned value is an instantaneous snapshot. It is suitable for diagnostics and soft + /// flow-control decisions, but concurrent sends and receives may change it immediately. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::bounded(4); + /// tx.try_send(10).unwrap(); + /// assert_eq!(tx.retained_message_count(), 1); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(tx.retained_message_count(), 0); + /// ``` + pub fn retained_message_count(&self) -> usize { + self.shared.inner.lock().log.retained() + } + + /// Returns the number of messages this channel retains before producers wait. + /// + /// This is the value passed to [`bounded`] and never changes. Pair it with + /// [`retained_message_count`](BoundedSender::retained_message_count) to compute headroom. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, _rx) = mpmc::bounded::(8); + /// assert_eq!(tx.capacity(), 8); + /// ``` + pub fn capacity(&self) -> usize { + self.shared.capacity + } + + /// Creates a new receiver that starts receiving messages from the current tail of the channel. + /// + /// Subscribing never consumes capacity: the new cursor starts at the tail, so it retains + /// nothing that was not already retained. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// use asyncband::broadcast::mpmc::TryRecvError; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, _rx) = mpmc::bounded(4); + /// tx.send(10).await; + /// + /// let mut rx = tx.subscribe(); + /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + /// tx.send(20).await; + /// assert_eq!(rx.recv().await, Ok(20)); + /// # } + /// ``` + pub fn subscribe(&self) -> BoundedReceiver { + let key = self.shared.inner.lock().log.subscribe(); + BoundedReceiver { + shared: self.shared.clone(), + key, + } + } +} + +/// A receiver for a bounded broadcast channel. +/// +/// Each receiver sees every message sent to the channel while the receiver is active. A receiver +/// that stops draining holds capacity for the whole channel, so dropping one that will not keep up +/// is how a caller releases producers. +pub struct BoundedReceiver { + shared: Arc>, + key: SlotId, +} + +impl fmt::Debug for BoundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedReceiver").finish_non_exhaustive() + } +} + +impl Drop for BoundedReceiver { + fn drop(&mut self) { + let (reclaimed, drained_last) = { + let mut inner = self.shared.inner.lock(); + let reclaimed = inner.log.remove_receiver(self.key); + let drained_last = !inner.log.has_receivers(); + (reclaimed, drained_last) + }; + + if drained_last { + self.shared.release_all(); + } else { + self.shared.release_reclaimed(reclaimed.len()); + } + + // Payload destructors run last, and unlocked. + drop(reclaimed); + } +} + +impl BoundedReceiver { + /// Receives the next value for this receiver. + /// + /// # Returns + /// + /// * `Ok(T)`: The next message. + /// * `Err(RecvError::Disconnected)`: All senders have been dropped and this receiver has no + /// remaining messages. + /// + /// # 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::mpmc; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, mut rx) = mpmc::bounded(4); + /// tx.send(10).await; + /// assert_eq!(rx.recv().await, Ok(10)); + /// # } + /// ``` + pub async fn recv(&mut self) -> Result { + Recv { + receiver: self, + token: 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 this receiver has no + /// remaining messages. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::bounded(4); + /// tx.try_send(10).unwrap(); + /// assert_eq!(rx.try_recv(), Ok(10)); + /// ``` + pub fn try_recv(&mut self) -> Result { + let (msg, reclaimed) = + common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; + + // Release before taking the payload: `take_msg` runs `T::clone` and `T::drop`, and if + // either panics the slots this receive already freed would otherwise never be handed to a + // parked producer, stalling it permanently. + self.shared.release_reclaimed(reclaimed.len()); + Ok(common::take_msg(msg, reclaimed)) + } +} + +impl BoundedReceiver { + /// 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::mpmc; + /// + /// let (tx, mut rx) = mpmc::bounded(4); + /// tx.try_send(1).unwrap(); + /// tx.try_send(2).unwrap(); + /// + /// let mut rx2 = rx.resubscribe(); + /// tx.try_send(3).unwrap(); + /// + /// assert_eq!(rx2.try_recv(), Ok(3)); + /// ``` + pub fn resubscribe(&self) -> Self { + let key = self.shared.inner.lock().log.subscribe(); + Self { + shared: self.shared.clone(), + key, + } + } + + /// Returns the number of messages this receiver can still read. + /// + /// This count is specific to this receiver, unlike + /// [`BoundedSender::retained_message_count`], which reports the shared backlog retained by the + /// slowest active receiver. + /// + /// The returned value is an instantaneous snapshot. It is suitable for detecting that this + /// receiver is falling behind, but concurrent sends may change it immediately. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::bounded(4); + /// assert_eq!(rx.unread_message_count(), 0); + /// + /// tx.try_send(10).unwrap(); + /// tx.try_send(20).unwrap(); + /// assert_eq!(rx.unread_message_count(), 2); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(rx.unread_message_count(), 1); + /// ``` + pub fn unread_message_count(&self) -> usize { + self.shared.inner.lock().log.unread(self.key) + } +} + +struct Recv<'a, T> { + receiver: &'a mut BoundedReceiver, + token: Option, +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + // Ready paths clear the token, so only a cancelled pending receive takes this lock. + if self.token.is_none() { + return; + } + + common::unregister(&self.receiver.shared.inner, &mut self.token); + } +} + +impl Future for Recv<'_, T> { + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let Self { receiver, token } = self.get_mut(); + + let (msg, reclaimed) = match common::poll_receive( + &receiver.shared.inner, + &receiver.shared.senders, + receiver.key, + token, + cx, + ) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + Poll::Ready(Ok(received)) => received, + }; + + // Release before taking the payload, for the same reason as `try_recv`: a panicking + // `T::clone` must not strand producers on slots this receive already freed. + receiver.shared.release_reclaimed(reclaimed.len()); + Poll::Ready(Ok(common::take_msg(msg, reclaimed))) + } +} diff --git a/asyncband/src/broadcast/mpmc/bounded/tests.rs b/asyncband/src/broadcast/mpmc/bounded/tests.rs new file mode 100644 index 00000000..594936ce --- /dev/null +++ b/asyncband/src/broadcast/mpmc/bounded/tests.rs @@ -0,0 +1,90 @@ +// 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. + +// These run under Miri via `cargo x miri`, so they stay single-threaded and small. Behavior +// reachable from the public API is covered in `tests-integration/broadcast_mpmc_bounded_test.rs`. + +use std::task::Waker; + +use super::*; + +#[test] +#[should_panic(expected = "broadcast bounded channel requires capacity > 0")] +fn bounded_panics_on_zero_capacity() { + let _ = bounded::<()>(0); +} + +#[test] +#[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, _) = bounded(1); + tx.shared.inner.lock().log.set_tail(u64::MAX); + let _ = tx.try_send(()); +} + +#[test] +fn buffer_is_preallocated_and_never_shrinks() { + let capacity = 128; + let (tx, mut rx) = bounded(capacity); + let allocated = tx.shared.inner.lock().log.buffer_capacity(); + assert!(allocated >= capacity); + + // Fill to capacity, drain completely, and repeat with a much smaller cycle. An elastic backlog + // would hand the allocation back after the small cycle; a fixed one must not. + for i in 0..capacity { + tx.try_send(i).unwrap(); + } + for i in 0..capacity { + assert_eq!(rx.try_recv(), Ok(i)); + } + tx.try_send(0).unwrap(); + assert_eq!(rx.try_recv(), Ok(0)); + + assert_eq!(tx.retained_message_count(), 0); + assert_eq!(tx.shared.inner.lock().log.buffer_capacity(), allocated); +} + +#[test] +fn capacity_reports_the_requested_value() { + let (tx, _rx) = bounded::(3); + assert_eq!(tx.capacity(), 3); +} + +#[test] +fn a_large_reclaim_leaves_no_permit_slack() { + // Dropping a lagging subscription frees the whole backlog in one step, far more slots than the + // single parked producer can use. Permits beyond that producer would sit in the semaphore, and + // the next send to block would burn each one on a publish attempt that cannot succeed. + let capacity = 64; + let (tx, mut fast) = bounded(capacity); + let lagging = tx.subscribe(); + for value in 0..capacity { + tx.try_send(value).unwrap(); + } + for _ in 0..capacity { + fast.try_recv().unwrap(); + } + + let mut cx = Context::from_waker(Waker::noop()); + let mut send = Box::pin(tx.send(capacity)); + assert!(send.as_mut().poll(&mut cx).is_pending()); + + drop(lagging); + assert!(send.as_mut().poll(&mut cx).is_ready()); + assert_eq!(tx.shared.tx_permits.available_permits(), 0); +} diff --git a/asyncband/src/broadcast/mpmc/common.rs b/asyncband/src/broadcast/mpmc/common.rs new file mode 100644 index 00000000..7e9241e2 --- /dev/null +++ b/asyncband/src/broadcast/mpmc/common.rs @@ -0,0 +1,527 @@ +// 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. + +//! Storage, cursors, and the receive step shared by the bounded and unbounded MPMC broadcast +//! channels. +//! +//! Both channels retain the same committed backlog and reclaim it the same way; they differ only +//! in what a producer does when that backlog is large. That difference stays in the two channel +//! modules, and so does each channel's own `Recv` future, so neither contract is hidden behind a +//! shared abstraction. What lives here is the state and the one poll step whose waker protocol is +//! subtle enough that a second copy would be a second thing to keep correct. + +use std::collections::VecDeque; +use std::mem; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; + +use super::error::RecvError; +use super::error::TryRecvError; +use crate::internal::arena::Arena; +use crate::internal::arena::SlotId; +use crate::internal::mutex::Mutex; +use crate::internal::waitset::WaitSet; +use crate::internal::waitset::WakerToken; +use crate::internal::waitset::wake_all; + +/// Retained capacity below which an elastic backlog is never shrunk back. +pub const MIN_RETAINED_CAPACITY: usize = 64; + +/// A received message together with the retained prefix that the receive released. +/// +/// The two travel together because the caller has to act on both with the channel unlocked, and a +/// bounded channel has to hand the released capacity back before it touches the payload. +pub type Received = (Arc, Reclaimed); + +/// 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. +pub struct Reclaimed { + first: Option>, + rest: Vec>, +} + +impl Reclaimed { + fn empty() -> Self { + Self { + first: None, + rest: vec![], + } + } + + fn first(&self) -> Option<&Arc> { + self.first.as_ref() + } + + pub fn is_empty(&self) -> bool { + self.first.is_none() + } + + /// The number of retained messages this reclaim released. + /// + /// A bounded channel turns this into the capacity it hands back to blocked producers. + pub fn len(&self) -> usize { + usize::from(self.first.is_some()) + self.rest.len() + } +} + +/// How a backlog manages the allocation behind its retained messages. +enum Retention { + /// Grow on demand, and return a burst allocation once a later cycle stays small. + Elastic { + /// The largest backlog retained since the buffer was last empty. + peak_len: usize, + }, + /// Allocated once for the requested capacity and never shrunk. + Fixed, +} + +/// The committed backlog: every message whose version falls in `[head, tail)`, plus one cursor for +/// each active subscription. +/// +/// This is the retention and sequencing machinery that stays private to the channel families. +/// `tail` is the single sequencer: it only advances while the channel lock is held, and a message +/// is placed in `buffer` in the same critical section, so a later publication can never become +/// visible ahead of an earlier one. +pub struct Backlog { + /// Messages whose versions are in the range `[head, tail)`. + /// + /// Each message is held behind an `Arc` so the receive path can move 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, + /// 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, + retention: Retention, +} + +impl Backlog { + /// A backlog that grows on demand and gives burst allocations back. + pub fn elastic() -> Self { + Self::new(VecDeque::new(), Retention::Elastic { peak_len: 0 }) + } + + /// A backlog preallocated for `capacity` retained messages that never shrinks. + pub fn fixed(capacity: usize) -> Self { + Self::new(VecDeque::with_capacity(capacity), Retention::Fixed) + } + + fn new(buffer: VecDeque>, retention: Retention) -> Self { + Self { + buffer, + head: 0, + head_receivers: 0, + tail: 0, + receivers: Arena::new(), + retention, + } + } + + /// The number of messages the channel currently retains. + /// + /// This is the shared backlog kept alive by the slowest active subscription, and it is what a + /// bounded channel measures its capacity against. + pub fn retained(&self) -> usize { + self.buffer.len() + } + + /// Whether any subscription is active. + /// + /// A channel with none retains nothing, so a bounded producer never waits on one. + pub fn has_receivers(&self) -> bool { + !self.receivers.is_empty() + } + + /// The number of messages the subscription registered as `key` can still read. + pub fn unread(&self, key: SlotId) -> usize { + let head = *self + .receivers + .get(key) + .expect("active broadcast receiver must be registered"); + usize::try_from(self.tail - head).expect("unread broadcast message count exceeds usize") + } + + /// Advances the committed tail. + /// + /// # Panics + /// + /// Panics if the message version counter overflows. + fn advance_tail(&mut self) { + self.tail = self + .tail + .checked_add(1) + .expect("broadcast channel version counter overflowed"); + } + + /// Advances the committed tail for a message no subscription can read. + /// + /// `head` moves with it 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. The caller keeps the payload and drops it after + /// releasing the channel lock. + pub fn publish_discarded(&mut self) { + debug_assert!(!self.has_receivers()); + debug_assert!(self.buffer.is_empty()); + debug_assert_eq!(self.head_receivers, 0); + self.advance_tail(); + self.head = self.tail; + } + + /// Advances the committed tail and retains `msg` for every currently active subscription. + /// + /// Returns the message when no subscription can read it, so the caller drops it after + /// releasing the channel lock rather than running `T::drop` inside the critical section. + /// + /// # Panics + /// + /// Panics if the message version counter overflows. + #[must_use = "drop the unretained message after releasing the channel lock"] + pub fn publish(&mut self, msg: Arc) -> Option> { + if !self.has_receivers() { + self.publish_discarded(); + return Some(msg); + } + + self.publish_retained(msg); + None + } + + /// Advances the committed tail and retains `msg`. + /// + /// The caller must already have established that a subscription is active, which is what a + /// bounded channel does anyway to decide between rejecting and discarding. + /// + /// # Panics + /// + /// Panics if the message version counter overflows. + pub fn publish_retained(&mut self, msg: Arc) { + debug_assert!(self.has_receivers()); + self.advance_tail(); + self.buffer.push_back(msg); + if let Retention::Elastic { peak_len } = &mut self.retention { + *peak_len = (*peak_len).max(self.buffer.len()); + } + } + + fn insert_receiver(&mut self, head: u64) -> SlotId { + if head == self.head { + self.head_receivers += 1; + } + + self.receivers.insert(head) + } + + /// Registers a new subscription at the committed tail. + /// + /// A new cursor never lowers `retained()`, so this can never release capacity. + pub fn subscribe(&mut self) -> SlotId { + let head = self.tail; + self.insert_receiver(head) + } + + pub fn remove_receiver(&mut self, key: SlotId) -> Reclaimed { + let head = self.receivers.remove(key); + + if head == self.head { + self.release_head_receiver() + } else { + Reclaimed::empty() + } + } + + fn release_head_receiver(&mut self) -> Reclaimed { + self.head_receivers -= 1; + + if self.head_receivers == 0 { + self.reclaim_consumed() + } else { + Reclaimed::empty() + } + } + + pub fn receive(&mut self, key: SlotId) -> Option> { + 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 + }; + + 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 { + 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)) + } + + /// Advances `head` to the slowest active cursor and hands the released prefix to the caller. + /// + /// `buffer` shrinks here and grows only in [`Backlog::publish`], so this is the one place + /// `retained()` can fall. A bounded channel therefore accounts for released capacity at + /// exactly the two call sites that reach this: [`Backlog::receive`] and + /// [`Backlog::remove_receiver`]. + fn reclaim_consumed(&mut self) -> Reclaimed { + 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 the channel 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() + }; + // Reclaiming exactly one message is the overwhelmingly common case — a cursor advances by + // one at a time — so skip building a `Drain` that would yield nothing. + let rest = if consumed > 1 { + self.buffer.drain(..consumed - 1).collect() + } else { + vec![] + }; + let reclaimed = Reclaimed { first, rest }; + debug_assert_eq!(reclaimed.len(), consumed); + + 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. + /// + /// A fixed backlog keeps the allocation it was built with, which is the whole point of asking + /// for a capacity up front. + fn shrink_buffer(&mut self) { + let Retention::Elastic { peak_len } = &mut self.retention else { + return; + }; + + if !self.buffer.is_empty() { + return; + } + + let peak = mem::take(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)); + } + } + + #[cfg(test)] + pub fn buffer_capacity(&self) -> usize { + self.buffer.capacity() + } + + /// Doctors the sequencer so a test can reach the overflow guard in `publish`. + #[cfg(test)] + pub fn set_tail(&mut self, tail: u64) { + self.tail = tail; + } +} + +/// Buffer, receiver cursors, and parked receivers, all under one lock. +/// +/// The wait set lives beside the backlog 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. +pub struct Inner { + pub log: Backlog, + pub waiters: WaitSet, +} + +impl Inner { + /// Wraps `log` in the channel lock and registers the subscription every constructor hands out + /// alongside its first sender. + pub fn with_first_subscription(mut log: Backlog) -> (Mutex, SlotId) { + let key = log.subscribe(); + let inner = Mutex::new(Self { + log, + waiters: WaitSet::new(), + }); + (inner, key) + } +} + +/// Wakes every parked receiver so it can observe the channel's disconnected state. +/// +/// Both families call this from the last sender's `Drop`. +pub fn disconnect(inner: &Mutex>) { + let wakers = { + let mut inner = inner.lock(); + inner.waiters.drain() + }; + wake_all(wakers); +} + +/// Releases a cancelled receive's waker registration, dropping the waker unlocked. +pub fn unregister(inner: &Mutex>, token: &mut Option) { + let waker = { + let mut inner = inner.lock(); + inner.waiters.unregister(token) + }; + drop(waker); +} + +/// Receives without waiting, yielding the message and the prefix the receive released. +/// +/// The caller owns what happens next: a bounded channel hands the released count back to blocked +/// producers before it touches the payload. +pub fn try_receive( + inner: &Mutex>, + senders: &AtomicUsize, + key: SlotId, +) -> Result, TryRecvError> { + // Check this receiver's cursor while holding `inner` before observing the sender count. + // 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 = inner.lock(); + match inner.log.receive(key) { + Some(received) => Ok(received), + None if senders.load(Ordering::Acquire) == 0 => Err(TryRecvError::Disconnected), + None => Err(TryRecvError::Empty), + } +} + +/// The one poll step behind `recv` on both channels. +/// +/// Buffered messages and repeated polls with the same task waker require no clone. If the pending +/// path needs a new waker, this releases the lock, clones, and repeats the full state check before +/// registration. Senders publish messages and drain waiters under the same lock, so the recheck +/// cannot miss a send, disconnection, or state change made by a reentrant clone callback. The loop +/// executes at most twice. +pub fn poll_receive( + inner: &Mutex>, + senders: &AtomicUsize, + key: SlotId, + token: &mut Option, + cx: &mut Context<'_>, +) -> Poll, RecvError>> { + let mut prepared_waker = None; + loop { + let mut guard = inner.lock(); + + match guard.log.receive(key) { + Some(received) => { + drop(guard); + drop(prepared_waker); + // Clearing the token without unregistering is safe, and it is what keeps the + // ready path off a second lock acquisition. A message can only become readable + // through a publish, and a publish drains the wait set in the same critical + // section that made the message visible — so any registration this future still + // held was already taken by that drain, and the token is stale. `Drop` reads the + // cleared token and skips its own lock for the same reason. + *token = None; + return Poll::Ready(Ok(received)); + } + None => { + if senders.load(Ordering::Acquire) == 0 { + *token = None; + drop(guard); + drop(prepared_waker); + return Poll::Ready(Err(RecvError::Disconnected)); + } + + if prepared_waker.is_none() && guard.waiters.will_wake(token, cx.waker()) { + return Poll::Pending; + } + let Some(waker) = prepared_waker.take() else { + drop(guard); + prepared_waker = Some(cx.waker().clone()); + continue; + }; + let retired_waker = guard.waiters.register(token, waker); + drop(guard); + drop(retired_waker); + return Poll::Pending; + } + } + } +} + +/// 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. +/// +/// This runs `T::clone` and `T::drop`, either of which may panic, so a bounded channel must +/// already have released the reclaimed capacity before calling it. +pub fn take_msg(msg: Arc, reclaimed: Reclaimed) -> 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. + Arc::try_unwrap(msg).unwrap_or_else(|msg| (*msg).clone()) +} diff --git a/asyncband/src/broadcast/mpmc/error.rs b/asyncband/src/broadcast/mpmc/error.rs new file mode 100644 index 00000000..f1cbce93 --- /dev/null +++ b/asyncband/src/broadcast/mpmc/error.rs @@ -0,0 +1,103 @@ +// 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::any::type_name; +use std::fmt; + +/// Error returned by [`BoundedSender::try_send`]. +/// +/// A bounded broadcast channel is lossless, so a publication that would exceed the requested +/// capacity is rejected rather than displacing a retained message. The message that could not be +/// sent can be retrieved again with [`TrySendError::into_inner`]. +/// +/// [`BoundedSender::try_send`]: crate::broadcast::mpmc::BoundedSender::try_send +#[derive(Clone, PartialEq, Eq)] +pub enum TrySendError { + /// The shared backlog is at capacity, so the message cannot be sent without waiting for the + /// slowest active receiver to release a retained message. + Full(T), +} + +impl TrySendError { + /// Gets a reference to the message that failed to be sent. + pub fn as_inner(&self) -> &T { + match self { + TrySendError::Full(msg) => msg, + } + } + + /// Consumes the error and returns the message that failed to be sent. + pub fn into_inner(self) -> T { + match self { + TrySendError::Full(msg) => msg, + } + } +} + +impl fmt::Display for TrySendError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + TrySendError::Full(_) => "sending on a full channel", + }) + } +} + +impl fmt::Debug for TrySendError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let ty = type_name::(); + match self { + TrySendError::Full(_) => write!(f, "TrySendError<{ty}>::Full(..)"), + } + } +} + +impl std::error::Error for TrySendError {} + +/// Error returned by `recv`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RecvError { + /// All senders have been dropped, and this receiver has no remaining messages. + Disconnected, +} + +impl fmt::Display for RecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("receiving on a disconnected channel") + } +} + +impl std::error::Error for RecvError {} + +/// Error returned by `try_recv`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TryRecvError { + /// No message is currently available, but at least one sender remains. + Empty, + /// All senders have been dropped, and this receiver has no remaining messages. + Disconnected, +} + +impl fmt::Display for TryRecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + TryRecvError::Empty => "receiving on an empty channel", + TryRecvError::Disconnected => "receiving on a disconnected channel", + }) + } +} + +impl std::error::Error for TryRecvError {} diff --git a/asyncband/src/broadcast/mpmc/mod.rs b/asyncband/src/broadcast/mpmc/mod.rs index 9249df21..db1a3682 100644 --- a/asyncband/src/broadcast/mpmc/mod.rs +++ b/asyncband/src/broadcast/mpmc/mod.rs @@ -16,11 +16,24 @@ // under the License. //! Multi-producer, multi-consumer broadcast channels. +//! +//! Both channels are lossless: every value a channel accepts stays readable by every subscription +//! that was active when it was accepted, so a receive never reports lag. They differ in what a +//! producer does when the slowest subscription stops reclaiming. [`bounded`] retains at most the +//! capacity it was built with and makes producers wait for that subscription. [`unbounded`] never +//! waits to send and lets the retained backlog grow instead. +mod bounded; +mod common; +mod error; mod unbounded; -pub use self::unbounded::RecvError; -pub use self::unbounded::TryRecvError; +pub use self::bounded::BoundedReceiver; +pub use self::bounded::BoundedSender; +pub use self::bounded::bounded; +pub use self::error::RecvError; +pub use self::error::TryRecvError; +pub use self::error::TrySendError; pub use self::unbounded::UnboundedReceiver; pub use self::unbounded::UnboundedSender; pub use self::unbounded::unbounded; diff --git a/asyncband/src/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/broadcast/mpmc/unbounded/mod.rs index 2924d488..1bd573a3 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -27,7 +27,8 @@ //! buffer to grow without bound, because messages are retained until every active receiver has //! consumed them or the receiver is dropped. Use //! [`UnboundedSender::retained_message_count`] to monitor the number of messages currently retained -//! by the channel. +//! by the channel. Use [`bounded`] instead when producers should wait for the slowest receiver +//! rather than let the backlog grow. //! //! 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 @@ -88,11 +89,11 @@ //! assert_eq!(tx.retained_message_count(), 0); //! # } //! ``` +//! +//! [`bounded`]: super::bounded -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; @@ -100,10 +101,13 @@ use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; -use crate::internal::arena::Arena; +use super::common; +use super::common::Backlog; +use super::common::Inner; +use super::error::RecvError; +use super::error::TryRecvError; use crate::internal::arena::SlotId; use crate::internal::mutex::Mutex; -use crate::internal::waitset::WaitSet; use crate::internal::waitset::WakerToken; use crate::internal::waitset::wake_all; @@ -124,18 +128,9 @@ mod tests; /// assert_eq!(rx.try_recv(), Ok(10)); /// ``` pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { - let mut receivers = Arena::new(); - let key = receivers.insert(0); + let (inner, key) = Inner::with_first_subscription(Backlog::elastic()); let shared = Arc::new(Shared { - inner: Mutex::new(Inner { - buffer: VecDeque::new(), - head: 0, - head_receivers: 1, - tail: 0, - receivers, - peak_len: 0, - waiters: WaitSet::new(), - }), + inner, senders: AtomicUsize::new(1), }); let sender = UnboundedSender { @@ -145,221 +140,8 @@ pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { (sender, receiver) } -/// Error returned by [`UnboundedReceiver::recv`]. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum RecvError { - /// All senders have been dropped, and this receiver has no remaining messages. - Disconnected, -} - -impl fmt::Display for RecvError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - RecvError::Disconnected => write!(f, "receiving on a disconnected channel"), - } - } -} - -impl std::error::Error for RecvError {} - -/// Error returned by [`UnboundedReceiver::try_recv`]. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum TryRecvError { - /// No message is currently available, but at least one sender remains. - Empty, - /// All senders have been dropped, and this receiver has no remaining messages. - 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 disconnected channel"), - } - } -} - -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 the receive path can move 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, - /// 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, - /// The largest backlog retained since the buffer was last empty. - peak_len: usize, - /// Receivers parked in [`UnboundedReceiver::recv`]. - 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![], - } - } - - 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 { - self.head_receivers += 1; - } - - self.receivers.insert(head) - } - - fn remove_receiver(&mut self, key: SlotId) -> Reclaimed { - let head = self.receivers.remove(key); - - if head == self.head { - self.release_head_receiver() - } else { - Reclaimed::empty() - } - } - - fn release_head_receiver(&mut self) -> Reclaimed { - self.head_receivers -= 1; - - if self.head_receivers == 0 { - self.reclaim_consumed() - } else { - Reclaimed::empty() - } - } - - 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 - }; - - 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 { - 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) -> Reclaimed { - 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. 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; - 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, @@ -394,14 +176,7 @@ impl fmt::Debug for UnboundedSender { impl Drop for UnboundedSender { fn drop(&mut self) { match self.shared.senders.fetch_sub(1, Ordering::AcqRel) { - 1 => { - // Wake every parked receiver so it can observe the channel's disconnected state. - let wakers = { - let mut inner = self.shared.inner.lock(); - inner.waiters.drain() - }; - wake_all(wakers); - } + 1 => common::disconnect(&self.shared.inner), _ => { // there are still other senders left, do nothing } @@ -437,32 +212,17 @@ impl UnboundedSender { // 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 (unretained, wakers) = { 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); - inner.peak_len = inner.peak_len.max(inner.buffer.len()); - } - - inner.waiters.drain() + let unretained = inner.log.publish(msg); + let wakers = inner.waiters.drain(); + (unretained, wakers) }; // Notify all waiting receivers. An unsent message is dropped here too, once the lock is // released. wake_all(wakers); + drop(unretained); } /// Returns the number of messages currently retained by the channel. @@ -486,7 +246,7 @@ impl UnboundedSender { /// assert_eq!(tx.retained_message_count(), 0); /// ``` pub fn retained_message_count(&self) -> usize { - self.shared.inner.lock().buffer.len() + self.shared.inner.lock().log.retained() } /// Creates a new receiver that starts receiving messages from the current tail of the channel. @@ -509,11 +269,11 @@ impl UnboundedSender { /// # } /// ``` 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(); - UnboundedReceiver { shared, key } + let key = self.shared.inner.lock().log.subscribe(); + UnboundedReceiver { + shared: self.shared.clone(), + key, + } } } @@ -535,7 +295,7 @@ impl Drop for UnboundedReceiver { fn drop(&mut self) { let reclaimed = { let mut inner = self.shared.inner.lock(); - inner.remove_receiver(self.key) + inner.log.remove_receiver(self.key) }; drop(reclaimed); } @@ -595,50 +355,13 @@ impl UnboundedReceiver { /// assert_eq!(rx.try_recv(), Ok(10)); /// ``` pub fn try_recv(&mut self) -> Result { - let (msg, reclaimed) = self.try_recv_shared()?; - 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: Reclaimed) -> T { - let sole_owner = !reclaimed.is_empty(); - reclaimed.drop_messages(); - - if !sole_owner { - return (*msg).clone(); + let (msg, reclaimed) = + common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; + Ok(common::take_msg(msg, reclaimed)) } - - // Another receiver can still hold an in-flight reference to the same message, so the clone - // remains the fallback. - Arc::try_unwrap(msg).unwrap_or_else(|msg| (*msg).clone()) } impl UnboundedReceiver { - 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. - 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. /// @@ -661,11 +384,11 @@ impl UnboundedReceiver { /// 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 } + let key = self.shared.inner.lock().log.subscribe(); + Self { + shared: self.shared.clone(), + key, + } } /// Returns the number of messages this receiver can still read. @@ -693,12 +416,7 @@ impl UnboundedReceiver { /// assert_eq!(rx.unread_message_count(), 1); /// ``` pub fn unread_message_count(&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") + self.shared.inner.lock().log.unread(self.key) } } @@ -714,11 +432,7 @@ impl Drop for Recv<'_, T> { return; } - let waker = { - let mut inner = self.receiver.shared.inner.lock(); - inner.waiters.unregister(&mut self.token) - }; - drop(waker); + common::unregister(&self.receiver.shared.inner, &mut self.token); } } @@ -728,44 +442,18 @@ impl Future for Recv<'_, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let Self { receiver, token } = self.get_mut(); - // Buffered messages and repeated polls with the same task waker require no clone. If the - // pending path needs a new waker, release the lock, clone, and repeat the full state check - // before registration. Senders publish messages and drain waiters under the same lock, so - // the recheck cannot miss a send, disconnection, or state change made by a reentrant clone - // callback. The loop executes at most twice. - let mut prepared_waker = None; - let received = loop { - let mut inner = receiver.shared.inner.lock(); - - match inner.receive(receiver.key) { - Some(received) => break received, - None => { - if receiver.shared.senders.load(Ordering::Acquire) == 0 { - *token = None; - drop(inner); - drop(prepared_waker); - return Poll::Ready(Err(RecvError::Disconnected)); - } - - if prepared_waker.is_none() && inner.waiters.will_wake(token, cx.waker()) { - return Poll::Pending; - } - let Some(waker) = prepared_waker.take() else { - drop(inner); - prepared_waker = Some(cx.waker().clone()); - continue; - }; - let retired_waker = inner.waiters.register(token, waker); - drop(inner); - drop(retired_waker); - return Poll::Pending; - } - } + let (msg, reclaimed) = match common::poll_receive( + &receiver.shared.inner, + &receiver.shared.senders, + receiver.key, + token, + cx, + ) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + Poll::Ready(Ok(received)) => received, }; - drop(prepared_waker); - let (msg, reclaimed) = received; - *token = None; - Poll::Ready(Ok(take_msg(msg, reclaimed))) + Poll::Ready(Ok(common::take_msg(msg, reclaimed))) } } diff --git a/asyncband/src/broadcast/mpmc/unbounded/tests.rs b/asyncband/src/broadcast/mpmc/unbounded/tests.rs index 933ee021..4bf330aa 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/tests.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/tests.rs @@ -16,13 +16,14 @@ // under the License. use super::*; +use crate::broadcast::mpmc::common::MIN_RETAINED_CAPACITY; #[test] #[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, _) = unbounded(); - tx.shared.inner.lock().tail = u64::MAX; + tx.shared.inner.lock().log.set_tail(u64::MAX); tx.send(()); } @@ -34,7 +35,7 @@ fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { for i in 0..burst { tx.send(i); } - assert!(tx.shared.inner.lock().buffer.capacity() >= burst); + assert!(tx.shared.inner.lock().log.buffer_capacity() >= burst); for i in 0..burst { assert_eq!(rx.try_recv(), Ok(i)); @@ -42,12 +43,12 @@ fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { // Draining evaluates the cycle that just peaked, so the burst allocation is still held. assert_eq!(tx.retained_message_count(), 0); - assert!(tx.shared.inner.lock().buffer.capacity() >= burst); + assert!(tx.shared.inner.lock().log.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); + assert!(tx.shared.inner.lock().log.buffer_capacity() < burst); } #[test] @@ -65,5 +66,5 @@ fn repeated_bursts_keep_their_allocation() { } // 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); + assert!(tx.shared.inner.lock().log.buffer_capacity() >= burst); } diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 143c60bd..5b11efba 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -64,18 +64,21 @@ pub(crate) mod value_cell; pub(crate) mod mutex; #[cfg(any( + feature = "broadcast", feature = "mpsc", feature = "mutex", feature = "rwlock", feature = "semaphore", ))] -// `mpsc` uses `poll_acquire`, `release_if_nonempty`, and `notify_all`; mutexes and rwlocks use -// `acquire`, `try_acquire`, and `release`; the public semaphore also uses the accounting methods. -// Each single-primitive build intentionally leaves the other groups unused. +// `broadcast` and `mpsc` park blocked producers with `poll_acquire`, `release_if_nonempty`, and +// `notify_all`; mutexes and rwlocks use `acquire`, `try_acquire`, and `release`; the public +// semaphore also uses the accounting methods. Each single-primitive build intentionally leaves the +// other groups unused. #[allow(dead_code)] pub(crate) mod semaphore; #[cfg(any( + feature = "broadcast", feature = "event", feature = "mpsc", feature = "mutex", diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index c17a8f43..e5e46246 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -71,7 +71,7 @@ //! | | [`Shutdown`](shutdown::Shutdown) | `shutdown` | Coordinate shutdown signals and completion. | //! | Channels | [`oneshot`] | `oneshot` | Send one value from one sender to one receiver. | //! | | [`mpsc`] | `mpsc` | Send each value from multiple producers to one receiver through a bounded or unbounded queue. | -//! | | [`broadcast`] | `broadcast` | Broadcast values from one or more producers and retain them until every active receiver consumes them. | +//! | | [`broadcast`] | `broadcast` | Broadcast every value to all active receivers, with bounded backpressure or unbounded retention. | //! | | [`watch`] | `watch` | Publish the latest state to independently tracked receivers and coalesce intermediate updates. | //! | Resource reuse | [`pool`] | `pool` | Reuse objects through bounded or unbounded pool variants. | //! | Workload coordination | [`Semaphore`](semaphore::Semaphore) | `semaphore` | Control concurrent access with permits. | diff --git a/benchmarks/asyncband/broadcast/mpmc/bounded.rs b/benchmarks/asyncband/broadcast/mpmc/bounded.rs new file mode 100644 index 00000000..c0800f83 --- /dev/null +++ b/benchmarks/asyncband/broadcast/mpmc/bounded.rs @@ -0,0 +1,145 @@ +// 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 steady state on each iteration: the retained +// backlog back where it started, no parked producer left behind, and no permit slack in the +// producer wait queue. Unlike the unbounded channel the hazard is not unbounded memory but a +// wedged timed loop — a send that never gets its capacity back would hang the bench, not slow it. + +use std::pin::pin; + +use asyncband::broadcast::mpmc; +use divan::Bencher; +use divan::black_box; + +use crate::support::bench_context; +use crate::support::poll_pending; +use crate::support::poll_pinned_ready; + +const RECEIVER_COUNTS: &[usize] = &[1, 8, 32]; +const BLOCKED_SENDER_COUNTS: &[usize] = &[1, 8, 32]; +const CAPACITY: usize = 64; + +#[divan::bench] +fn send_without_receivers(bencher: Bencher) { + // No subscription means nothing is retained, so this measures the discard path, which never + // allocates and never waits. + let (tx, rx) = mpmc::bounded(CAPACITY); + drop(rx); + + bencher.bench_local(|| tx.try_send(black_box(1))); +} + +#[divan::bench] +fn try_send_and_try_recv(bencher: Bencher) { + let (tx, mut rx) = mpmc::bounded(CAPACITY); + + bencher.bench_local(|| { + tx.try_send(black_box(1)).unwrap(); + black_box(rx.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn try_send_when_full(bencher: Bencher) { + let (tx, _rx) = mpmc::bounded(1); + tx.try_send(0).unwrap(); + + // The rejected value comes straight back, so the channel stays exactly as full as it started. + bencher.bench_local(|| black_box(tx.try_send(black_box(1))).is_err()); +} + +#[divan::bench(args = RECEIVER_COUNTS)] +fn try_send_and_drain_fanout(bencher: Bencher, receiver_count: usize) { + let (tx, rx) = mpmc::bounded(CAPACITY); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(rx); + for _ in 1..receiver_count { + receivers.push(tx.subscribe()); + } + + // One message in, every receiver drains it out: the last one to read pays the reclaim scan and + // the capacity release, and the channel is empty again for the next iteration. + bencher.bench_local(|| { + tx.try_send(black_box(1)).unwrap(); + for receiver in &mut receivers { + black_box(receiver.try_recv().unwrap()); + } + }); +} + +#[divan::bench(args = BLOCKED_SENDER_COUNTS)] +fn reclaim_wakes_blocked_senders(bencher: Bencher, sender_count: usize) { + let mut context = bench_context(); + + // Measures the whole backpressure cycle: park `sender_count` producers on a full channel, free + // one slot, and let exactly one of them through. Each iteration ends with the same number of + // producers parked and the same backlog, so the loop is stationary. + bencher + .with_inputs(|| { + let (tx, rx) = mpmc::bounded(1); + tx.try_send(0).unwrap(); + (tx, rx) + }) + .bench_local_refs(|(tx, rx)| { + let mut sends = (0..sender_count) + .map(|value| Box::pin(tx.send(value))) + .collect::>(); + for send in &mut sends { + poll_pending(send.as_mut(), &mut context); + } + + // Releasing one slot wakes the queue; one producer republishes and the rest re-park. + black_box(rx.try_recv().unwrap()); + for send in &mut sends { + if send.as_mut().poll(&mut context).is_ready() { + break; + } + } + + // Drain the republished message so the next iteration starts from the same state. + black_box(rx.try_recv().unwrap()); + drop(sends); + tx.try_send(0).unwrap(); + }); +} + +#[divan::bench] +fn cancel_blocked_send(bencher: Bencher) { + let mut context = bench_context(); + let (tx, _rx) = mpmc::bounded(1); + tx.try_send(0).unwrap(); + + // Park a producer and immediately cancel it: measures registering and unlinking one waiter. + bencher.bench_local(|| { + let send = pin!(tx.send(black_box(1))); + poll_pending(send, &mut context); + }); +} + +#[divan::bench] +fn deliver_to_waiting_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (tx, mut rx) = mpmc::bounded(CAPACITY); + + bencher.bench_local(|| { + let mut recv = pin!(rx.recv()); + poll_pending(recv.as_mut(), &mut context); + tx.try_send(black_box(1)).unwrap(); + black_box(poll_pinned_ready(recv, &mut context).unwrap()) + }); +} diff --git a/benchmarks/asyncband/broadcast/mpmc/mod.rs b/benchmarks/asyncband/broadcast/mpmc/mod.rs index 78ef889a..e0ac8347 100644 --- a/benchmarks/asyncband/broadcast/mpmc/mod.rs +++ b/benchmarks/asyncband/broadcast/mpmc/mod.rs @@ -15,4 +15,5 @@ // specific language governing permissions and limitations // under the License. +mod bounded; mod unbounded; diff --git a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs index 7b988701..0d08cf1f 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs @@ -128,3 +128,111 @@ impl BroadcastMpmc for AsyncBroadcast { poll_ready(receiver.recv_direct(), context).unwrap() } } + +/// A lossless bounded broadcast channel: every accepted value reaches every active subscription, +/// and a full channel makes producers wait rather than displacing anything. +/// +/// `tokio::sync::broadcast` deliberately has no implementation here — see the note in `bounded.rs`. +pub trait BoundedBroadcastMpmc: Send + Sync + 'static { + type Sender: Clone + Send + 'static; + type Receiver: Send + 'static; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec); + fn try_send(sender: &Self::Sender, value: usize); + fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>); + fn send_blocking(sender: &Self::Sender, value: usize); + fn try_recv(receiver: &mut Self::Receiver) -> Option; + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize; + fn recv_blocking(receiver: &mut Self::Receiver) -> usize; +} + +impl BoundedBroadcastMpmc for Asyncband { + type Receiver = asyncband::broadcast::mpmc::BoundedReceiver; + type Sender = asyncband::broadcast::mpmc::BoundedSender; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = asyncband::broadcast::mpmc::bounded(capacity); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(sender.subscribe()); + } + (sender, receivers) + } + + fn try_send(sender: &Self::Sender, value: usize) { + sender.try_send(value).unwrap(); + } + + fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + poll_ready(sender.send(value), context); + } + + fn send_blocking(sender: &Self::Sender, value: usize) { + pollster::block_on(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() + } + + fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + pollster::block_on(receiver.recv()).unwrap() + } +} + +impl BoundedBroadcastMpmc 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 try_send(sender: &Self::Sender, value: usize) { + sender.try_broadcast(value).unwrap(); + } + + fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + poll_ready(sender.broadcast_direct(value), context) + .expect("async-broadcast lost every receiver during benchmark"); + } + + fn send_blocking(sender: &Self::Sender, value: usize) { + pollster::block_on(sender.broadcast_direct(value)) + .expect("async-broadcast lost every receiver during benchmark"); + } + + 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() + } + + fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + pollster::block_on(receiver.recv_direct()).unwrap() + } +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/bounded.rs b/benchmarks/ecosystem/broadcast/mpmc/bounded.rs new file mode 100644 index 00000000..71540ebb --- /dev/null +++ b/benchmarks/ecosystem/broadcast/mpmc/bounded.rs @@ -0,0 +1,84 @@ +// 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 and async-broadcast are compared here because both are lossless and both make +// producers wait at capacity, so a small channel measures the same contract on each side. +// +// `tokio::sync::broadcast` is deliberately absent. It overwrites at capacity and reports `Lagged` +// rather than waiting, so it has no lossless bounded path to compare: it would be measuring the +// cheaper workload of dropping messages. It appears in `unbounded.rs` instead, where every peer is +// given room for the whole batch and the comparison is over their shared non-blocking path. +// +// `concurrent` sweeps capacity as well as producer and receiver counts, because the ratio of +// capacity to fanout is what decides how a bounded broadcast behaves. A backlog smaller than the +// fanout turns every message into a round trip — publish one value, wake every subscription, wait +// for the one slot to come back — while a roomy backlog lets both sides batch. The two regimes +// differ by an order of magnitude on every implementation measured, so reporting one capacity +// would describe half the channel. + +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use super::adapters::AsyncBroadcast; +use super::adapters::Asyncband; +use super::adapters::BoundedBroadcastMpmc; +use super::support::BATCH_MESSAGES; +use super::support::BOUNDED_SHAPES; +use super::support::BoundedConcurrent; +use super::support::BoundedShape; +use super::support::ROUND_TRIP_CAPACITY; +use crate::support::bench_context; + +// Send-then-receive pairing keeps at most one message retained, so these never reach capacity. +#[divan::bench(types = [Asyncband, 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::try_send(&sender, black_box(usize::MAX)); + black_box(C::try_recv(&mut receiver).unwrap()) + }); +} + +#[divan::bench(types = [Asyncband, 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_ready(&sender, black_box(usize::MAX), &mut context); + black_box(C::recv_ready(&mut receiver, &mut context)) + }); +} + +// `sample_size = 1` is required: `BoundedConcurrent` spawns its workers once and they exit after a +// single pass, so a second `run` on the same value would block forever. +#[divan::bench( + types = [Asyncband, AsyncBroadcast], + args = BOUNDED_SHAPES, + sample_count = 10, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn concurrent(bencher: Bencher, shape: BoundedShape) { + bencher + .with_inputs(|| BoundedConcurrent::::new(shape)) + .bench_local_refs(BoundedConcurrent::run); +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/mod.rs b/benchmarks/ecosystem/broadcast/mpmc/mod.rs index 5e86ce6e..dd09282b 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/mod.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/mod.rs @@ -16,5 +16,6 @@ // under the License. mod adapters; +mod bounded; mod support; mod unbounded; diff --git a/benchmarks/ecosystem/broadcast/mpmc/support.rs b/benchmarks/ecosystem/broadcast/mpmc/support.rs index 55c016a5..424263f4 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/support.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/support.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::fmt; use std::marker::PhantomData; use std::sync::Arc; use std::sync::Barrier; @@ -23,13 +24,69 @@ use std::thread::JoinHandle; use divan::black_box; +use super::adapters::BoundedBroadcastMpmc; 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]; +/// Capacity for the round-trip benches, which pair every send with a receive and so never fill +/// the channel. pub const ROUND_TRIP_CAPACITY: usize = 64; +/// One bounded workload: the channel capacity, how many producers publish, and how many +/// subscriptions read. +/// +/// Capacity is a dimension rather than a constant because it is the parameter that decides how a +/// bounded broadcast behaves. At `TIGHT` the backlog is a fraction of the fanout, so the run +/// degenerates into a per-message round trip: the producer publishes one value, every subscription +/// is woken to read it, and only then does a slot come back. At `ROOMY` the producer runs ahead +/// and both sides batch. Reporting only one of the two would describe half the channel. +#[derive(Clone, Copy)] +pub struct BoundedShape { + pub capacity: usize, + pub producers: usize, + pub receivers: usize, +} + +impl BoundedShape { + const TIGHT: usize = 64; + const ROOMY: usize = 1024; + + const fn new(capacity: usize, producers: usize, receivers: usize) -> Self { + Self { + capacity, + producers, + receivers, + } + } +} + +impl fmt::Display for BoundedShape { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "cap {} {} producers {} receivers", + self.capacity, self.producers, self.receivers + ) + } +} + +pub const BOUNDED_SHAPES: &[BoundedShape] = &[ + BoundedShape::new(BoundedShape::TIGHT, 1, 1), + BoundedShape::new(BoundedShape::TIGHT, 1, 8), + BoundedShape::new(BoundedShape::TIGHT, 1, 32), + BoundedShape::new(BoundedShape::TIGHT, 8, 1), + BoundedShape::new(BoundedShape::TIGHT, 8, 8), + BoundedShape::new(BoundedShape::TIGHT, 8, 32), + BoundedShape::new(BoundedShape::ROOMY, 1, 1), + BoundedShape::new(BoundedShape::ROOMY, 1, 8), + BoundedShape::new(BoundedShape::ROOMY, 1, 32), + BoundedShape::new(BoundedShape::ROOMY, 8, 1), + BoundedShape::new(BoundedShape::ROOMY, 8, 8), + BoundedShape::new(BoundedShape::ROOMY, 8, 32), +]; + fn recv(receiver: &mut C::Receiver) -> usize { C::try_recv(receiver).expect("the published benchmark batch must be ready") } @@ -157,3 +214,96 @@ impl Drop for Fanout { } } } + +/// Producers and receivers running concurrently against a channel far smaller than the batch. +/// +/// The unbounded fixtures above publish the whole batch before anyone drains it, which is only +/// safe because every peer is given room for the entire batch. That shape deadlocks a genuinely +/// bounded channel, so this one interleaves: every thread blocks, and the drain runs while the +/// producers are still publishing. +/// +/// This terminates. The run could only wedge if every producer and every receiver waited at the +/// same time, but a producer waits only while at least one message is retained, and a retained +/// message is by definition unread by the slowest receiver — so that receiver is runnable. The +/// counts balance exactly: the producers publish `BATCH_MESSAGES` between them and each receiver +/// consumes `BATCH_MESSAGES`, so no thread over- or under-runs. Every receiver is subscribed +/// before the first send, so every receiver sees every message. +/// +/// Benches using this must set `sample_size = 1`: the worker threads are spawned in `new` and exit +/// after one pass, so a second `run` on the same value would block forever. +pub struct BoundedConcurrent { + start: Arc, + done: Arc, + workers: Vec>, + channel: PhantomData, +} + +impl BoundedConcurrent { + pub fn new(shape: BoundedShape) -> Self { + let BoundedShape { + capacity, + producers, + receivers, + } = shape; + assert_eq!(BATCH_MESSAGES % producers, 0); + + let (sender, receivers) = C::channel(capacity, receivers); + let start = Arc::new(Barrier::new(producers + receivers.len() + 1)); + let done = Arc::new(Barrier::new(producers + receivers.len() + 1)); + let messages_per_producer = BATCH_MESSAGES / producers; + let mut workers = Vec::with_capacity(producers + receivers.len()); + + for mut receiver in receivers { + let start = start.clone(); + let done = done.clone(); + workers.push(thread::spawn(move || { + start.wait(); + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(C::recv_blocking(&mut receiver)); + } + black_box(checksum); + done.wait(); + })); + } + + for producer in 0..producers { + let sender = sender.clone(); + let start = start.clone(); + let done = done.clone(); + workers.push(thread::spawn(move || { + start.wait(); + let first = producer * messages_per_producer; + for value in first..first + messages_per_producer { + C::send_blocking(&sender, black_box(value)); + } + done.wait(); + })); + } + drop(sender); + + Self { + start, + done, + workers, + channel: PhantomData, + } + } + + pub fn run(&mut self) { + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for BoundedConcurrent { + fn drop(&mut self) { + let panicking = thread::panicking(); + for worker in self.workers.drain(..) { + let result = worker.join(); + if !panicking { + result.expect("bounded benchmark worker panicked"); + } + } + } +} diff --git a/tests-integration/tests/broadcast_mpmc_bounded_test.rs b/tests-integration/tests/broadcast_mpmc_bounded_test.rs new file mode 100644 index 00000000..bf06ef82 --- /dev/null +++ b/tests-integration/tests/broadcast_mpmc_bounded_test.rs @@ -0,0 +1,732 @@ +// 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::Wake; +use std::task::Waker; +use std::thread; +use std::time::Duration; + +use asyncband::broadcast::mpmc::*; +use tests_integration::poll_once; + +struct TrackWake(AtomicUsize); + +impl TrackWake { + fn new() -> Arc { + Arc::new(Self(AtomicUsize::new(0))) + } + + fn count(&self) -> usize { + self.0.load(Ordering::Relaxed) + } +} + +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.retained_message_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, + } + } +} + +/// A payload that panics while the channel drops a message it reclaimed. +/// +/// Clones disarm themselves, so only the copy the channel retains is dangerous. That lets a test +/// drain a receiver normally and still blow up inside the reclaim. +struct PanicOnDrop { + armed: bool, +} + +impl Clone for PanicOnDrop { + fn clone(&self) -> Self { + Self { armed: false } + } +} + +impl Drop for PanicOnDrop { + fn drop(&mut self) { + if self.armed { + panic!("panic while dropping a broadcast message"); + } + } +} + +// --------------------------------------------------------------------------------------------- +// Fanout and subscription +// --------------------------------------------------------------------------------------------- + +#[tokio::test] +async fn bounded_delivers_every_message_to_every_receiver() { + let (tx, mut rx1) = bounded(4); + let mut rx2 = tx.subscribe(); + + tx.send(10).await; + tx.send(20).await; + + 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)); +} + +#[test] +fn bounded_slow_receiver_keeps_every_message_under_backpressure() { + let (tx, mut fast) = bounded(2); + let mut slow = tx.subscribe(); + + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + + // The fast subscription draining does not release anything, because the slow one has read + // nothing — being bounded must not turn into dropping what the slow subscription still owes. + assert_eq!(fast.try_recv(), Ok(1)); + assert_eq!(fast.try_recv(), Ok(2)); + assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); + + // One read by the slow subscription frees exactly one slot. + assert_eq!(slow.try_recv(), Ok(1)); + tx.try_send(3).unwrap(); + + // Every value accepted while both were active reaches both, in order. + assert_eq!(slow.try_recv(), Ok(2)); + assert_eq!(slow.try_recv(), Ok(3)); + assert_eq!(fast.try_recv(), Ok(3)); + assert_eq!(tx.retained_message_count(), 0); +} + +#[test] +fn bounded_subscribe_starts_at_the_committed_tail() { + let (tx, _rx) = bounded(4); + tx.try_send(1).unwrap(); + + let mut late = tx.subscribe(); + assert_eq!(late.try_recv(), Err(TryRecvError::Empty)); + + tx.try_send(2).unwrap(); + assert_eq!(late.try_recv(), Ok(2)); +} + +#[test] +fn bounded_resubscribe_keeps_the_original_receivers_backlog() { + let (tx, mut rx) = bounded(4); + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + + let mut rx2 = rx.resubscribe(); + tx.try_send(3).unwrap(); + + assert_eq!(rx2.try_recv(), Ok(3)); + assert_eq!(rx.try_recv(), Ok(1)); + assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv(), Ok(3)); +} + +#[test] +fn bounded_unread_message_count_tracks_each_receiver() { + let (tx, mut rx1) = bounded(4); + let rx2 = tx.subscribe(); + + assert_eq!(rx1.unread_message_count(), 0); + + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + assert_eq!(rx1.unread_message_count(), 2); + assert_eq!(rx2.unread_message_count(), 2); + + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.unread_message_count(), 1); + assert_eq!(rx2.unread_message_count(), 2); +} + +// --------------------------------------------------------------------------------------------- +// Strict capacity +// --------------------------------------------------------------------------------------------- + +#[test] +fn try_send_rejects_at_capacity_and_returns_the_value() { + let (tx, mut rx) = bounded(2); + + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); + + // The rejected value is handed back untouched, and nothing was published. + assert_eq!(tx.retained_message_count(), 2); + assert_eq!(rx.try_recv(), Ok(1)); + assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); +} + +#[test] +fn capacity_counts_the_shared_backlog_not_receivers() { + let (tx, _rx) = bounded(2); + let _extra = (0..8).map(|_| tx.subscribe()).collect::>(); + + // Eight more subscriptions do not consume capacity; only unread messages do. + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); + assert_eq!(tx.capacity(), 2); + assert_eq!(tx.retained_message_count(), 2); +} + +#[test] +fn retained_message_count_tracks_the_slowest_receiver() { + let (tx, mut rx1) = bounded(4); + let mut rx2 = tx.subscribe(); + + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + assert_eq!(tx.retained_message_count(), 2); + + // Draining one receiver does not release what the other has not read. + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.try_recv(), Ok(2)); + assert_eq!(tx.retained_message_count(), 2); + + assert_eq!(rx2.try_recv(), Ok(1)); + assert_eq!(tx.retained_message_count(), 1); + assert_eq!(rx2.try_recv(), Ok(2)); + assert_eq!(tx.retained_message_count(), 0); +} + +// --------------------------------------------------------------------------------------------- +// Backpressure and capacity release +// --------------------------------------------------------------------------------------------- + +#[test] +fn send_waits_while_the_slowest_subscription_holds_capacity() { + let (tx, mut rx1) = bounded(1); + let mut rx2 = tx.subscribe(); + tx.try_send(1).unwrap(); + + let mut send = Box::pin(tx.send(2)); + assert!(poll_once(send.as_mut()).is_pending()); + + // The fast receiver draining is not enough while the slow one still retains the message. + assert_eq!(rx1.try_recv(), Ok(1)); + assert!(poll_once(send.as_mut()).is_pending()); + + assert_eq!(rx2.try_recv(), Ok(1)); + assert!(poll_once(send.as_mut()).is_ready()); + assert_eq!(rx1.try_recv(), Ok(2)); +} + +#[test] +fn receive_that_vacates_the_head_wakes_a_blocked_sender() { + let (tx, mut rx) = bounded(1); + tx.try_send(0).unwrap(); + + let tracker = TrackWake::new(); + let waker = Waker::from(tracker.clone()); + let mut send = Box::pin(tx.send(1)); + assert!( + send.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + assert_eq!(tracker.count(), 0); + + assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(tracker.count(), 1); + assert!(poll_once(send.as_mut()).is_ready()); +} + +#[test] +fn parked_recv_that_reclaims_wakes_a_blocked_sender() { + let (tx, mut rx) = bounded(1); + tx.try_send(0).unwrap(); + + let tracker = TrackWake::new(); + let waker = Waker::from(tracker.clone()); + let mut send = Box::pin(tx.send(1)); + assert!( + send.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + + // Reclaim through the `recv` future rather than `try_recv`: it is a separate call site, and a + // release wired into only one of them would strand this producer. + let mut recv = Box::pin(rx.recv()); + assert_eq!(poll_once(recv.as_mut()), std::task::Poll::Ready(Ok(0))); + drop(recv); + + assert_eq!(tracker.count(), 1); + assert!(poll_once(send.as_mut()).is_ready()); +} + +#[test] +fn wakes_blocked_senders_as_capacity_frees() { + let (tx, mut rx) = bounded(1); + tx.try_send(0).unwrap(); + + let mut first = Box::pin(tx.send(1)); + let mut second = Box::pin(tx.send(2)); + assert!(poll_once(first.as_mut()).is_pending()); + assert!(poll_once(second.as_mut()).is_pending()); + + // One freed slot admits exactly one producer. + assert_eq!(rx.try_recv(), Ok(0)); + assert!(poll_once(first.as_mut()).is_ready()); + assert!(poll_once(second.as_mut()).is_pending()); + + assert_eq!(rx.try_recv(), Ok(1)); + assert!(poll_once(second.as_mut()).is_ready()); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn dropping_a_lagging_receiver_wakes_blocked_senders() { + let (tx, mut rx1) = bounded(1); + let rx2 = tx.subscribe(); + tx.try_send(0).unwrap(); + + let mut send = Box::pin(tx.send(1)); + assert_eq!(rx1.try_recv(), Ok(0)); + assert!(poll_once(send.as_mut()).is_pending()); + + // `rx2` is the one holding the backlog; dropping it releases the slot. + drop(rx2); + assert_eq!(tx.retained_message_count(), 0); + assert!(poll_once(send.as_mut()).is_ready()); +} + +#[test] +fn dropping_the_last_receiver_wakes_every_blocked_sender() { + const BLOCKED: usize = 3; + + let (tx, rx) = bounded(2); + tx.try_send(0).unwrap(); + tx.try_send(1).unwrap(); + + // More blocked producers than the drop will reclaim slots. Once no receiver remains every + // send succeeds unconditionally, so waking only `reclaimed` of them would strand the rest. + let trackers = (0..BLOCKED).map(|_| TrackWake::new()).collect::>(); + let mut sends = (0..BLOCKED) + .map(|value| Box::pin(tx.send(10 + value as i32))) + .collect::>(); + + for (send, tracker) in sends.iter_mut().zip(&trackers) { + let waker = Waker::from(tracker.clone()); + assert!( + send.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + } + + drop(rx); + + for (index, tracker) in trackers.iter().enumerate() { + assert!( + tracker.count() > 0, + "blocked sender {index} was never woken after the last receiver was dropped" + ); + } + for send in &mut sends { + assert!(poll_once(send.as_mut()).is_ready()); + } +} + +#[test] +fn sends_never_block_once_all_receivers_are_gone() { + let (tx, rx) = bounded(1); + tx.try_send(0).unwrap(); + drop(rx); + + assert_eq!(tx.retained_message_count(), 0); + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + + let mut send = Box::pin(tx.send(3)); + assert!(poll_once(send.as_mut()).is_ready()); + assert_eq!(tx.retained_message_count(), 0); +} + +#[test] +fn subscribing_while_producers_are_blocked_does_not_release_capacity() { + let (tx, _rx) = bounded(1); + tx.try_send(0).unwrap(); + + let mut send = Box::pin(tx.send(1)); + assert!(poll_once(send.as_mut()).is_pending()); + + // A new cursor starts at the tail, so it cannot lower the retained backlog. + let _late = tx.subscribe(); + assert_eq!(tx.retained_message_count(), 1); + assert!(poll_once(send.as_mut()).is_pending()); +} + +// --------------------------------------------------------------------------------------------- +// Cancellation +// --------------------------------------------------------------------------------------------- + +#[test] +fn cancelled_send_publishes_nothing() { + let (tx, mut rx) = bounded(1); + tx.try_send(0).unwrap(); + + let mut send = Box::pin(tx.send(1)); + assert!(poll_once(send.as_mut()).is_pending()); + drop(send); + + // The cancelled value never entered the committed order, so the next receive sees only what + // was already published, and the one after it is a fresh send. + assert_eq!(rx.try_recv(), Ok(0)); + tx.try_send(2).unwrap(); + assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); +} + +#[test] +fn cancelled_notified_sender_passes_capacity_to_the_next_sender() { + let (tx, mut rx) = bounded(1); + tx.try_send(0).unwrap(); + + let mut first = Box::pin(tx.send(1)); + let mut second = Box::pin(tx.send(2)); + assert!(poll_once(first.as_mut()).is_pending()); + assert!(poll_once(second.as_mut()).is_pending()); + + assert_eq!(rx.try_recv(), Ok(0)); + drop(first); + + assert!(poll_once(second.as_mut()).is_ready()); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn cancelled_recv_releases_its_waker() { + let (tx, mut rx) = bounded(4); + + let tracker = TrackWake::new(); + let waker = Waker::from(tracker.clone()); + let baseline = Arc::strong_count(&tracker); + + let mut recv = Box::pin(rx.recv()); + assert!( + recv.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + + drop(recv); + assert_eq!(Arc::strong_count(&tracker), baseline); + + tx.try_send(1).unwrap(); + assert_eq!(tracker.count(), 0); +} + +// --------------------------------------------------------------------------------------------- +// Disconnection +// --------------------------------------------------------------------------------------------- + +#[tokio::test] +async fn bounded_recv_drains_buffered_messages_before_reporting_disconnection() { + let (tx, mut rx) = bounded(4); + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + 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 bounded_recv_reports_disconnection_without_any_message() { + let (tx, mut rx) = bounded::(4); + drop(tx); + assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); +} + +#[test] +fn bounded_parked_recv_wakes_when_the_last_sender_drops() { + let (tx, mut rx) = bounded::(4); + let second_tx = tx.clone(); + + let tracker = TrackWake::new(); + let waker = Waker::from(tracker.clone()); + let mut recv = Box::pin(rx.recv()); + assert!( + recv.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + + drop(tx); + assert_eq!(tracker.count(), 0); + + drop(second_tx); + assert_eq!(tracker.count(), 1); + assert_eq!( + poll_once(recv.as_mut()), + std::task::Poll::Ready(Err(RecvError::Disconnected)) + ); +} + +// --------------------------------------------------------------------------------------------- +// Panic safety +// --------------------------------------------------------------------------------------------- + +#[test] +fn bounded_panicking_clone_leaves_the_channel_consistent() { + let (tx, mut rx1) = bounded(4); + let mut rx2 = tx.subscribe(); + + tx.try_send(PanicOnClone { + value: 1, + panic: true, + }) + .unwrap(); + tx.try_send(PanicOnClone { + value: 2, + panic: false, + }) + .unwrap(); + + // 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.retained_message_count(), 0); + assert_eq!(rx1.try_recv().unwrap_err(), TryRecvError::Empty); +} + +#[test] +fn panicking_payload_destructor_still_releases_capacity() { + let (tx, mut rx1) = bounded(3); + let rx2 = tx.subscribe(); + + // Only the first retained message is armed: the reclaim drops the whole prefix, and a second + // panic while the first one unwinds would abort the process instead of failing the test. + for index in 0..3 { + tx.try_send(PanicOnDrop { armed: index == 0 }).unwrap(); + } + // `rx1` reads clones, which are disarmed; the armed originals stay retained for `rx2`. + for _ in 0..3 { + rx1.try_recv().unwrap(); + } + + let mut send = Box::pin(tx.send(PanicOnDrop { armed: false })); + assert!(poll_once(send.as_mut()).is_pending()); + + // Dropping `rx2` reclaims all three retained messages and their destructors panic. The + // capacity they released must already have reached the parked producer by then. + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(rx2))); + assert!(result.is_err()); + + assert!( + poll_once(send.as_mut()).is_ready(), + "a panicking payload destructor must not strand a producer on capacity it already freed" + ); +} + +#[test] +fn bounded_message_destructors_run_outside_the_channel_lock() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + + let worker = thread::spawn(move || { + let (tx, mut rx1) = bounded(8); + let rx2 = tx.subscribe(); + + for value in 0..4 { + tx.try_send(Reentrant { + value, + channel: Some(tx.clone()), + }) + .unwrap(); + } + + // Draining both receivers reclaims the prefix, whose destructors re-enter the channel. + for _ in 0..4 { + rx1.try_recv().unwrap(); + } + drop(rx2); + drop(rx1); + + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("reclaimed message destructors must not run while the channel is locked"); + worker.join().unwrap(); +} + +// --------------------------------------------------------------------------------------------- +// Concurrency +// --------------------------------------------------------------------------------------------- + +#[test] +fn dropping_the_last_receiver_never_strands_a_racing_producer() { + const ROUNDS: u64 = 150; + const PRODUCERS: u64 = 4; + + // The deterministic tests above drop the receiver at a fixed point. This races the drop + // against producers entering the waiting path, which is the window where the channel decides + // whether anybody needs waking. A missed wake-up here parks a producer forever, so the failure + // mode is a hang rather than a wrong value — hence the timeout instead of an assertion. + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + + let worker = thread::spawn(move || { + for round in 0..ROUNDS { + let (tx, rx) = bounded(1); + tx.try_send(0).unwrap(); + + let producers = (0..PRODUCERS) + .map(|producer| { + let tx = tx.clone(); + thread::spawn(move || pollster::block_on(tx.send(round * 10 + producer + 1))) + }) + .collect::>(); + + drop(rx); + + for producer in producers { + producer.join().unwrap(); + } + } + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(60)) + .expect("a producer was left waiting after the last receiver went away"); + worker.join().unwrap(); +} + +#[test] +fn bounded_concurrent_producers_commit_one_order_seen_by_every_receiver() { + const PRODUCERS: u64 = 4; + const PER_PRODUCER: u64 = 128; + const RECEIVERS: usize = 4; + const TOTAL: u64 = PRODUCERS * PER_PRODUCER; + + // Several producers publishing concurrently must still commit one contiguous order, and every + // subscription must observe that same order — not merely the same set. + // + // Capacity is far below the batch, so the producers really do block on the slowest receiver. + // This still terminates: the run could only wedge if every producer and every receiver waited + // at once, but a producer waits only while at least one message is retained, and a retained + // message is by definition unread by the slowest receiver — so that receiver is runnable. + let (tx, rx) = bounded(8); + let mut receivers = vec![rx]; + receivers.extend((1..RECEIVERS).map(|_| tx.subscribe())); + + let drains = receivers + .into_iter() + .map(|mut receiver| { + thread::spawn(move || { + let mut seen = Vec::with_capacity(TOTAL as usize); + for _ in 0..TOTAL { + seen.push(pollster::block_on(receiver.recv()).expect("sender dropped early")); + } + seen + }) + }) + .collect::>(); + + let producers = (0..PRODUCERS) + .map(|worker| { + let tx = tx.clone(); + thread::spawn(move || { + for value in 0..PER_PRODUCER { + pollster::block_on(tx.send(worker * PER_PRODUCER + value)); + } + }) + }) + .collect::>(); + + for producer in producers { + producer.join().unwrap(); + } + drop(tx); + + let orders = drains + .into_iter() + .map(|drain| drain.join().unwrap()) + .collect::>(); + + // Every subscription saw the identical sequence. + for (index, order) in orders.iter().enumerate().skip(1) { + assert_eq!( + order, &orders[0], + "subscription {index} observed a different committed order" + ); + } + + // And that sequence is every published value exactly once — no gap, no duplicate. + let mut sorted = orders[0].clone(); + sorted.sort_unstable(); + assert_eq!(sorted, (0..TOTAL).collect::>()); +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index d8683683..7d879a79 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -96,8 +96,11 @@ 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::>(); @@ -162,8 +165,11 @@ 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::>(); diff --git a/tests-integration/tests/waitset_reentrancy_test.rs b/tests-integration/tests/waitset_reentrancy_test.rs index a5867e59..53087674 100644 --- a/tests-integration/tests/waitset_reentrancy_test.rs +++ b/tests-integration/tests/waitset_reentrancy_test.rs @@ -213,3 +213,20 @@ fn broadcast_clones_wakers_outside_its_state_lock() { }, ); } + +#[test] +fn bounded_broadcast_clones_wakers_outside_its_state_lock() { + assert_completes_without_deadlock( + "waker clone callback deadlocked against the bounded broadcast lock", + || { + let (sender, mut receiver) = mpmc::bounded(1); + let callback_sender = sender.clone(); + let waker = waker_with_clone_callback(move || { + callback_sender.try_send(1).expect("channel has room"); + }); + let mut recv = Box::pin(receiver.recv()); + + assert_eq!(poll_with(recv.as_mut(), &waker), Poll::Ready(Ok(1))); + }, + ); +} From 4602e9b9791731cfc5f847ae57d6de268cdfa394 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 12 Sep 2026 11:10:09 +0800 Subject: [PATCH 2/8] fix(semaphore): finish permit releases after batched wake panics --- CHANGELOG.md | 1 + asyncband/src/internal/semaphore.rs | 96 +++++++++++++++++-- .../tests/broadcast_mpmc_bounded_test.rs | 49 ++++++++++ 3 files changed, 136 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 63f9328c..31998904 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,7 @@ All notable changes to this project will be documented in this file. ### Bug fixes +* Complete semaphore permit releases and notify all eligible waiters even if a wake callback panics during an earlier notification batch. * Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. * Notify all blocked bounded MPSC senders on receiver disconnection even when a buffered message destructor panics. * Avoid deadlocks when a bounded MPSC sender's waker clone callback receives from the same channel. diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index 5bb44a0c..96127dbb 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -26,6 +26,8 @@ use std::future::Future; use std::mem::MaybeUninit; +use std::panic; +use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::ptr; use std::sync::MutexGuard; @@ -86,19 +88,24 @@ impl WakeBatch { self.end == WAKE_BATCH_SIZE } - fn wake_all(&mut self) { - wake_all(std::iter::from_fn(|| { - if self.start == self.end { - return None; - } + fn wake_all(&mut self) -> std::thread::Result<()> { + let result = panic::catch_unwind(AssertUnwindSafe(|| { + wake_all(std::iter::from_fn(|| { + if self.start == self.end { + return None; + } - let index = self.start; - self.start += 1; - // SAFETY: `index` was within the initialized range before advancing `start`. - Some(unsafe { self.wakers[index].assume_init_read() }) + let index = self.start; + self.start += 1; + // SAFETY: `index` was within the initialized range before advancing `start`. + Some(unsafe { self.wakers[index].assume_init_read() }) + })); })); + // `wake_all` attempts every callback even after a panic, so the batch is empty in either + // case. Reset it before the next batch continues distributing the remaining permits. self.start = 0; self.end = 0; + result } } @@ -246,6 +253,7 @@ impl Semaphore { waiters: MutexGuard<'_, WaitList>, ) { let mut wakers = WakeBatch::new(); + let mut first_panic = None; let mut lock = Some(waiters); while rem > 0 { @@ -288,7 +296,17 @@ impl Semaphore { } drop(waiters); - wakers.wake_all(); + if let Err(payload) = wakers.wake_all() + && first_panic.is_none() + { + first_panic = Some(payload); + } + } + + // A callback must not prevent later batches from receiving permits that have already + // been released. Propagate its panic only after all accounting and notifications finish. + if let Some(payload) = first_panic { + panic::resume_unwind(payload); } } } @@ -515,4 +533,62 @@ mod tests { } assert_eq!(semaphore.waiters.lock().occupied_len(), 0); } + + #[test] + fn panicking_wakes_finish_later_batches_and_preserve_the_first_panic() { + const WAITER_COUNT: usize = WAKE_BATCH_SIZE * 2 + 1; + + struct TrackedWake { + count: AtomicUsize, + panic_message: Option<&'static str>, + } + + impl Wake for TrackedWake { + fn wake(self: Arc) { + self.count.fetch_add(1, Ordering::Relaxed); + if let Some(message) = self.panic_message { + panic::panic_any(message); + } + } + } + + let semaphore = Semaphore::new(0); + let trackers = (0..WAITER_COUNT) + .map(|index| { + Arc::new(TrackedWake { + count: AtomicUsize::new(0), + panic_message: match index { + 0 => Some("first wake panic"), + WAKE_BATCH_SIZE => Some("later wake panic"), + _ => None, + }, + }) + }) + .collect::>(); + let mut acquires = trackers + .iter() + .map(|tracker| { + let mut acquire = semaphore.poll_acquire(1); + assert!( + acquire + .poll_once(&Waker::from(tracker.clone())) + .is_pending() + ); + acquire + }) + .collect::>(); + + let payload = panic::catch_unwind(AssertUnwindSafe(|| semaphore.release(WAITER_COUNT + 2))) + .expect_err("the original wake panic must reach the caller"); + assert_eq!(payload.downcast_ref::<&str>(), Some(&"first wake panic")); + + for tracker in trackers { + assert_eq!(tracker.count.load(Ordering::Relaxed), 1); + } + for acquire in &mut acquires { + assert!(acquire.poll_once(Waker::noop()).is_ready()); + } + assert_eq!(semaphore.available_permits(), 2); + assert_eq!(semaphore.waiters.lock().occupied_len(), 0); + } } diff --git a/tests-integration/tests/broadcast_mpmc_bounded_test.rs b/tests-integration/tests/broadcast_mpmc_bounded_test.rs index edd64ea5..0d2da1c4 100644 --- a/tests-integration/tests/broadcast_mpmc_bounded_test.rs +++ b/tests-integration/tests/broadcast_mpmc_bounded_test.rs @@ -28,6 +28,7 @@ use std::time::Duration; use asyncband::blocking::FutureExt; use asyncband::broadcast::mpmc::*; use tests_integration::poll_once; +use tests_integration::waker_on_wake; struct TrackWake(AtomicUsize); @@ -560,6 +561,54 @@ fn bounded_parked_recv_wakes_when_the_last_sender_drops() { // Panic safety // --------------------------------------------------------------------------------------------- +#[test] +fn panicking_wake_does_not_strand_senders_after_a_large_reclaim() { + let (tx, mut fast) = bounded(40); + let slow = tx.subscribe(); + for value in 0..40 { + tx.try_send(value).unwrap(); + fast.try_recv().unwrap(); + } + + let trackers = (0..40).map(|_| TrackWake::new()).collect::>(); + let wakers = trackers + .iter() + .enumerate() + .map(|(index, tracker)| { + let tracker = tracker.clone(); + waker_on_wake(move || { + tracker.wake(); + assert_ne!(index, 0, "first sender wake panics"); + }) + }) + .collect::>(); + let mut sends = (40..80) + .map(|value| Box::pin(tx.send(value))) + .collect::>(); + for (send, waker) in sends.iter_mut().zip(&wakers) { + assert!( + send.as_mut() + .poll(&mut Context::from_waker(waker)) + .is_pending() + ); + } + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(slow))); + assert!(result.is_err()); + assert_eq!(tx.retained_message_count(), 0); + for tracker in trackers { + assert_eq!(tracker.count(), 1); + } + // Every send fits without another receive. All must have been notified, not merely made + // ready for a poll that an executor would otherwise have no reason to perform. + for send in &mut sends { + assert!(poll_once(send.as_mut()).is_ready()); + } + for value in 40..80 { + assert_eq!(fast.try_recv(), Ok(value)); + } +} + #[test] fn bounded_panicking_clone_leaves_the_channel_consistent() { let (tx, mut rx1) = bounded(4); From 4ee10707f839783e5a80b748f4f002f624bc8271 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 12 Sep 2026 11:10:09 +0800 Subject: [PATCH 3/8] test(broadcast): compare small capacities and async task scheduling --- .../ecosystem/broadcast/mpmc/adapters.rs | 42 +++++-- .../ecosystem/broadcast/mpmc/bounded.rs | 31 +++-- .../ecosystem/broadcast/mpmc/support.rs | 110 ++++++++++++++---- 3 files changed, 144 insertions(+), 39 deletions(-) diff --git a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs index 9369fa1f..21224281 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; use std::task::Context; use asyncband::blocking::FutureExt; @@ -136,11 +137,13 @@ impl BroadcastMpmc for AsyncBroadcast { /// /// `tokio::sync::broadcast` deliberately has no implementation here — see the note in `bounded.rs`. pub trait BoundedBroadcastMpmc: Send + Sync + 'static { - type Sender: Clone + Send + 'static; + type Sender: Clone + Send + Sync + 'static; type Receiver: Send + 'static; fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec); fn try_send(sender: &Self::Sender, value: usize); + fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>); fn send_blocking(sender: &Self::Sender, value: usize); fn try_recv(receiver: &mut Self::Receiver) -> Option; @@ -166,12 +169,20 @@ impl BoundedBroadcastMpmc for Asyncband { sender.try_send(value).unwrap(); } + async fn send_async(sender: &Self::Sender, value: usize) { + sender.send(value).await; + } + + async fn recv_async(receiver: &mut Self::Receiver) -> usize { + receiver.recv().await.unwrap() + } + fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { - poll_ready(sender.send(value), context); + poll_ready(Self::send_async(sender, value), context); } fn send_blocking(sender: &Self::Sender, value: usize) { - FutureExt::block_on(sender.send(value)); + FutureExt::block_on(Self::send_async(sender, value)); } fn try_recv(receiver: &mut Self::Receiver) -> Option { @@ -185,11 +196,11 @@ impl BoundedBroadcastMpmc for Asyncband { } fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { - poll_ready(receiver.recv(), context).unwrap() + poll_ready(Self::recv_async(receiver), context) } fn recv_blocking(receiver: &mut Self::Receiver) -> usize { - FutureExt::block_on(receiver.recv()).unwrap() + FutureExt::block_on(Self::recv_async(receiver)) } } @@ -212,14 +223,23 @@ impl BoundedBroadcastMpmc for AsyncBroadcast { sender.try_broadcast(value).unwrap(); } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { - poll_ready(sender.broadcast_direct(value), context) + async fn send_async(sender: &Self::Sender, value: usize) { + sender + .broadcast_direct(value) + .await .expect("async-broadcast lost every receiver during benchmark"); } + async fn recv_async(receiver: &mut Self::Receiver) -> usize { + receiver.recv_direct().await.unwrap() + } + + fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + poll_ready(Self::send_async(sender, value), context); + } + fn send_blocking(sender: &Self::Sender, value: usize) { - FutureExt::block_on(sender.broadcast_direct(value)) - .expect("async-broadcast lost every receiver during benchmark"); + FutureExt::block_on(Self::send_async(sender, value)); } fn try_recv(receiver: &mut Self::Receiver) -> Option { @@ -231,10 +251,10 @@ impl BoundedBroadcastMpmc for AsyncBroadcast { } fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { - poll_ready(receiver.recv_direct(), context).unwrap() + poll_ready(Self::recv_async(receiver), context) } fn recv_blocking(receiver: &mut Self::Receiver) -> usize { - FutureExt::block_on(receiver.recv_direct()).unwrap() + FutureExt::block_on(Self::recv_async(receiver)) } } diff --git a/benchmarks/ecosystem/broadcast/mpmc/bounded.rs b/benchmarks/ecosystem/broadcast/mpmc/bounded.rs index 71540ebb..db2f3d3e 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/bounded.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/bounded.rs @@ -23,12 +23,9 @@ // cheaper workload of dropping messages. It appears in `unbounded.rs` instead, where every peer is // given room for the whole batch and the comparison is over their shared non-blocking path. // -// `concurrent` sweeps capacity as well as producer and receiver counts, because the ratio of -// capacity to fanout is what decides how a bounded broadcast behaves. A backlog smaller than the -// fanout turns every message into a round trip — publish one value, wake every subscription, wait -// for the one slot to come back — while a roomy backlog lets both sides batch. The two regimes -// differ by an order of magnitude on every implementation measured, so reporting one capacity -// would describe half the channel. +// Sweep capacity and producer/subscription counts independently. Capacity one measures the +// per-message handoff; larger backlogs allow several messages to be outstanding. How effectively +// that headroom is used depends on scheduling and the slowest subscription, not just fanout. use divan::Bencher; use divan::black_box; @@ -41,11 +38,12 @@ use super::support::BATCH_MESSAGES; use super::support::BOUNDED_SHAPES; use super::support::BoundedConcurrent; use super::support::BoundedShape; +use super::support::BoundedTasks; use super::support::ROUND_TRIP_CAPACITY; use crate::support::bench_context; // Send-then-receive pairing keeps at most one message retained, so these never reach capacity. -#[divan::bench(types = [Asyncband, AsyncBroadcast])] +#[divan::bench(types = [Asyncband, AsyncBroadcast], sample_size = 512)] fn try_round_trip(bencher: Bencher) { let (sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); let mut receiver = receivers.pop().unwrap(); @@ -56,7 +54,7 @@ fn try_round_trip(bencher: Bencher) { }); } -#[divan::bench(types = [Asyncband, AsyncBroadcast])] +#[divan::bench(types = [Asyncband, AsyncBroadcast], sample_size = 512)] fn ready_round_trip(bencher: Bencher) { let mut context = bench_context(); let (sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); @@ -82,3 +80,20 @@ fn concurrent(bencher: Bencher, shape: BoundedShape) { .with_inputs(|| BoundedConcurrent::::new(shape)) .bench_local_refs(BoundedConcurrent::run); } + +#[divan::bench( + types = [Asyncband, AsyncBroadcast], + args = BOUNDED_SHAPES, + sample_count = 10, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn scheduled(bencher: Bencher, shape: BoundedShape) { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .build() + .unwrap(); + bencher + .with_inputs(|| BoundedTasks::new::(&runtime, shape)) + .bench_local_refs(|tasks| tasks.run(&runtime)); +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/support.rs b/benchmarks/ecosystem/broadcast/mpmc/support.rs index 424263f4..ccdc9780 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/support.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/support.rs @@ -23,6 +23,8 @@ use std::thread; use std::thread::JoinHandle; use divan::black_box; +use tokio::runtime::Runtime; +use tokio::task::JoinSet; use super::adapters::BoundedBroadcastMpmc; use super::adapters::BroadcastMpmc; @@ -37,11 +39,9 @@ pub const ROUND_TRIP_CAPACITY: usize = 64; /// One bounded workload: the channel capacity, how many producers publish, and how many /// subscriptions read. /// -/// Capacity is a dimension rather than a constant because it is the parameter that decides how a -/// bounded broadcast behaves. At `TIGHT` the backlog is a fraction of the fanout, so the run -/// degenerates into a per-message round trip: the producer publishes one value, every subscription -/// is woken to read it, and only then does a slot come back. At `ROOMY` the producer runs ahead -/// and both sides batch. Reporting only one of the two would describe half the channel. +/// Capacity bounds the shared backlog, independently of the subscription count. Capacity one +/// forces the next publication to wait until every subscription advances; larger capacities let +/// producers run ahead by that many messages, subject to scheduling and consumer progress. #[derive(Clone, Copy)] pub struct BoundedShape { pub capacity: usize, @@ -50,9 +50,6 @@ pub struct BoundedShape { } impl BoundedShape { - const TIGHT: usize = 64; - const ROOMY: usize = 1024; - const fn new(capacity: usize, producers: usize, receivers: usize) -> Self { Self { capacity, @@ -73,18 +70,30 @@ impl fmt::Display for BoundedShape { } pub const BOUNDED_SHAPES: &[BoundedShape] = &[ - BoundedShape::new(BoundedShape::TIGHT, 1, 1), - BoundedShape::new(BoundedShape::TIGHT, 1, 8), - BoundedShape::new(BoundedShape::TIGHT, 1, 32), - BoundedShape::new(BoundedShape::TIGHT, 8, 1), - BoundedShape::new(BoundedShape::TIGHT, 8, 8), - BoundedShape::new(BoundedShape::TIGHT, 8, 32), - BoundedShape::new(BoundedShape::ROOMY, 1, 1), - BoundedShape::new(BoundedShape::ROOMY, 1, 8), - BoundedShape::new(BoundedShape::ROOMY, 1, 32), - BoundedShape::new(BoundedShape::ROOMY, 8, 1), - BoundedShape::new(BoundedShape::ROOMY, 8, 8), - BoundedShape::new(BoundedShape::ROOMY, 8, 32), + BoundedShape::new(1, 1, 1), + BoundedShape::new(1, 1, 8), + BoundedShape::new(1, 8, 1), + BoundedShape::new(1, 8, 8), + BoundedShape::new(2, 1, 1), + BoundedShape::new(2, 1, 8), + BoundedShape::new(2, 8, 1), + BoundedShape::new(2, 8, 8), + BoundedShape::new(8, 1, 1), + BoundedShape::new(8, 1, 8), + BoundedShape::new(8, 8, 1), + BoundedShape::new(8, 8, 8), + BoundedShape::new(64, 1, 1), + BoundedShape::new(64, 1, 8), + BoundedShape::new(64, 1, 32), + BoundedShape::new(64, 8, 1), + BoundedShape::new(64, 8, 8), + BoundedShape::new(64, 8, 32), + BoundedShape::new(1024, 1, 1), + BoundedShape::new(1024, 1, 8), + BoundedShape::new(1024, 1, 32), + BoundedShape::new(1024, 8, 1), + BoundedShape::new(1024, 8, 8), + BoundedShape::new(1024, 8, 32), ]; fn recv(receiver: &mut C::Receiver) -> usize { @@ -307,3 +316,64 @@ impl Drop for BoundedConcurrent { } } } + +/// The same bounded workload on async tasks. Construction and task spawning happen outside the +/// timed section. Each fixture runs once, so its benchmark must use `sample_size = 1`. +pub struct BoundedTasks { + start: Arc, + tasks: JoinSet<()>, +} + +impl BoundedTasks { + pub fn new(runtime: &Runtime, shape: BoundedShape) -> Self { + let BoundedShape { + capacity, + producers, + receivers, + } = shape; + assert_eq!(BATCH_MESSAGES % producers, 0); + let (sender, receivers) = C::channel(capacity, receivers); + let start = Arc::new(tokio::sync::Barrier::new(producers + receivers.len() + 1)); + let mut tasks = JoinSet::new(); + + for mut receiver in receivers { + let start = start.clone(); + tasks.spawn_on( + async move { + start.wait().await; + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(C::recv_async(&mut receiver).await); + } + assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); + }, + runtime.handle(), + ); + } + for producer in 0..producers { + let sender = sender.clone(); + let start = start.clone(); + tasks.spawn_on( + async move { + start.wait().await; + let first = producer * (BATCH_MESSAGES / producers); + for value in first..first + BATCH_MESSAGES / producers { + C::send_async(&sender, black_box(value)).await; + } + }, + runtime.handle(), + ); + } + + Self { start, tasks } + } + + pub fn run(&mut self, runtime: &Runtime) { + runtime.block_on(async { + self.start.wait().await; + while let Some(result) = self.tasks.join_next().await { + result.expect("bounded benchmark task panicked"); + } + }); + } +} From 6dd3b2a5e07513c8f15eaa4a0cd7714caf2dff2f Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 12 Sep 2026 11:10:09 +0800 Subject: [PATCH 4/8] docs(broadcast): clarify delivery and subscription contracts --- asyncband/src/broadcast/mpmc/bounded/mod.rs | 12 +++++++++--- asyncband/src/broadcast/mpmc/mod.rs | 14 ++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/asyncband/src/broadcast/mpmc/bounded/mod.rs b/asyncband/src/broadcast/mpmc/bounded/mod.rs index 30d71522..cadb8037 100644 --- a/asyncband/src/broadcast/mpmc/bounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/bounded/mod.rs @@ -34,6 +34,10 @@ //! //! If no receivers are active the channel retains nothing, so a send never waits. //! +//! A successful receive releases its subscription's claim before returning the value; processing +//! that value afterward does not hold capacity. The capacity limit excludes pending sends and +//! values already handed to application code. +//! //! # Receivers //! //! Each receiver has an independent cursor. Use [`BoundedSender::subscribe`] or @@ -521,6 +525,7 @@ impl BoundedSender { /// assert_eq!(rx.recv().await, Ok(20)); /// # } /// ``` + #[must_use = "the receiver is dropped immediately if it is not retained"] pub fn subscribe(&self) -> BoundedReceiver { let key = self.shared.inner.lock().log.subscribe(); BoundedReceiver { @@ -635,9 +640,9 @@ impl BoundedReceiver { /// 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. + /// The new receiver skips every value already published, including the latest retained value. + /// The original receiver is unchanged and continues to retain its own backlog until it + /// consumes those messages or is dropped. /// /// # Examples /// @@ -653,6 +658,7 @@ impl BoundedReceiver { /// /// assert_eq!(rx2.try_recv(), Ok(3)); /// ``` + #[must_use = "the receiver is dropped immediately if it is not retained"] pub fn resubscribe(&self) -> Self { let key = self.shared.inner.lock().log.subscribe(); Self { diff --git a/asyncband/src/broadcast/mpmc/mod.rs b/asyncband/src/broadcast/mpmc/mod.rs index db1a3682..636f8958 100644 --- a/asyncband/src/broadcast/mpmc/mod.rs +++ b/asyncband/src/broadcast/mpmc/mod.rs @@ -22,6 +22,20 @@ //! producer does when the slowest subscription stops reclaiming. [`bounded`] retains at most the //! capacity it was built with and makes producers wait for that subscription. [`unbounded`] never //! waits to send and lets the retained backlog grow instead. +//! +//! # Delivery and processing +//! +//! A receive advances its subscription before returning the value. The channel tracks unread +//! messages, not application work: retaining a received value or processing it asynchronously +//! does not hold backlog capacity. There is no acknowledgement or processing-completion barrier. +//! If cloning a received value panics, that subscription has still advanced past the value. +//! +//! Sending with no subscriptions discards the value and succeeds. A later subscription starts +//! with future publications; it does not replay discarded or previously retained values. +//! +//! Operations briefly acquire internal mutexes. No mutex is held across an await point or while +//! cloning or dropping payloads. The `try_*` methods do not wait for messages or capacity, but may +//! wait to acquire a mutex. Sending a value does not wait for subscribers to receive or process it. mod bounded; mod common; From cad4333b92983e30424950e1357ec212de1a1f41 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 12 Sep 2026 16:01:46 +0800 Subject: [PATCH 5/8] fix(semaphore): preserve Rust 1.86 compatibility --- asyncband/src/internal/semaphore.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index 96127dbb..18ee9fd2 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -296,10 +296,8 @@ impl Semaphore { } drop(waiters); - if let Err(payload) = wakers.wake_all() - && first_panic.is_none() - { - first_panic = Some(payload); + if let Err(payload) = wakers.wake_all() { + first_panic.get_or_insert(payload); } } From ad8ca418923fde5f6d6868f31cf5dcaec1fce136 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 12 Sep 2026 17:12:28 +0800 Subject: [PATCH 6/8] refactor(semaphore): centralize panic-safe waking --- CHANGELOG.md | 2 +- asyncband/src/internal/semaphore.rs | 175 +++++++++------------------- 2 files changed, 54 insertions(+), 123 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 31998904..087014cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,7 +32,7 @@ All notable changes to this project will be documented in this file. ### Bug fixes -* Complete semaphore permit releases and notify all eligible waiters even if a wake callback panics during an earlier notification batch. +* Complete semaphore permit releases and notify all eligible waiters even if a wake callback panics. * Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. * Notify all blocked bounded MPSC senders on receiver disconnection even when a buffered message destructor panics. * Avoid deadlocks when a bounded MPSC sender's waker clone callback receives from the same channel. diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index 18ee9fd2..b32dbbed 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -20,16 +20,12 @@ // The Tokio-derived portions remain licensed under the MIT License. // Asyncband substantially replaced the waiter lifecycle with queue-owned WaitList nodes, supports // queue-head permit debt for exact reductions, has no closed state or reserved flag bits, and uses -// its own cancellation, detachment, and batched-waking machinery. +// its own cancellation, detachment, and waking machinery. // Upstream source: // https://github.com/tokio-rs/tokio/blob/bb9d57017e100985f86d8ca41ac105ee9140423e/tokio/src/sync/batch_semaphore.rs use std::future::Future; -use std::mem::MaybeUninit; -use std::panic; -use std::panic::AssertUnwindSafe; use std::pin::Pin; -use std::ptr; use std::sync::MutexGuard; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -58,68 +54,6 @@ struct WaitNode { waker: Option, } -const WAKE_BATCH_SIZE: usize = 32; - -/// The initialized entries in `wakers` are exactly `start..end`. -struct WakeBatch { - wakers: [MaybeUninit; WAKE_BATCH_SIZE], - start: usize, - end: usize, -} - -impl WakeBatch { - fn new() -> Self { - const UNINIT: MaybeUninit = MaybeUninit::uninit(); - Self { - wakers: [UNINIT; WAKE_BATCH_SIZE], - start: 0, - end: 0, - } - } - - fn push(&mut self, waker: Waker) { - debug_assert_eq!(self.start, 0); - debug_assert!(self.end < WAKE_BATCH_SIZE); - self.wakers[self.end].write(waker); - self.end += 1; - } - - fn is_full(&self) -> bool { - self.end == WAKE_BATCH_SIZE - } - - fn wake_all(&mut self) -> std::thread::Result<()> { - let result = panic::catch_unwind(AssertUnwindSafe(|| { - wake_all(std::iter::from_fn(|| { - if self.start == self.end { - return None; - } - - let index = self.start; - self.start += 1; - // SAFETY: `index` was within the initialized range before advancing `start`. - Some(unsafe { self.wakers[index].assume_init_read() }) - })); - })); - // `wake_all` attempts every callback even after a panic, so the batch is empty in either - // case. Reset it before the next batch continues distributing the remaining permits. - self.start = 0; - self.end = 0; - result - } -} - -impl Drop for WakeBatch { - fn drop(&mut self) { - let start = self.wakers[self.start..self.end] - .as_mut_ptr() - .cast::(); - let remaining = ptr::slice_from_raw_parts_mut(start, self.end - self.start); - // SAFETY: The initialized entries are exactly `start..end`. - unsafe { ptr::drop_in_place(remaining) }; - } -} - impl Semaphore { pub const fn new(permits: usize) -> Self { Self { @@ -244,68 +178,61 @@ impl Semaphore { } } drop(waiters); - crate::internal::wake_all(wakers.into_iter()); + wake_all(wakers.into_iter()); } fn insert_permits_with_lock( &self, mut rem: usize, - waiters: MutexGuard<'_, WaitList>, + mut waiters: MutexGuard<'_, WaitList>, ) { - let mut wakers = WakeBatch::new(); - let mut first_panic = None; - - let mut lock = Some(waiters); + // A single-waiter handoff should not allocate a wake buffer. + let mut first_waker = None; + let mut wakers = vec![]; while rem > 0 { - let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock()); - while !wakers.is_full() { - match waiters.unlink_first_waiter(|node| { - if node.permits <= rem { - rem -= node.permits; - node.permits = 0; - true - } else { - node.permits -= rem; - rem = 0; - false - } - }) { - None => break, - Some((id, waiter)) => { - let remove_now = waiter.waker.is_none(); - if let Some(waker) = waiter.waker.take() { + match waiters.unlink_first_waiter(|node| { + if node.permits <= rem { + rem -= node.permits; + node.permits = 0; + true + } else { + node.permits -= rem; + rem = 0; + false + } + }) { + None => break, + Some((id, waiter)) => { + let remove_now = waiter.waker.is_none(); + if let Some(waker) = waiter.waker.take() { + if first_waker.is_none() { + first_waker = Some(waker); + } else { wakers.push(waker); } - if remove_now { - waiters.remove_unlinked_waiter(id); - } + } + if remove_now { + waiters.remove_unlinked_waiter(id); } } } - - if rem > 0 && waiters.is_empty() { - // Holding `waiters` serializes all permit additions. Concurrent operations can - // only remove permits, so the count cannot grow between this check and fetch_add. - let current = self.permits.load(Ordering::Relaxed); - assert!( - current.checked_add(rem).is_some(), - "number of added permits ({rem}) would overflow usize::MAX (prev: {current})" - ); - self.permits.fetch_add(rem, Ordering::Release); - rem = 0; - } - - drop(waiters); - if let Err(payload) = wakers.wake_all() { - first_panic.get_or_insert(payload); - } } - // A callback must not prevent later batches from receiving permits that have already - // been released. Propagate its panic only after all accounting and notifications finish. - if let Some(payload) = first_panic { - panic::resume_unwind(payload); + if rem > 0 { + // Holding `waiters` serializes all permit additions. Concurrent operations can only + // remove permits, so the count cannot grow between this check and fetch_add. + let current = self.permits.load(Ordering::Relaxed); + assert!( + current.checked_add(rem).is_some(), + "number of added permits ({rem}) would overflow usize::MAX (prev: {current})" + ); + self.permits.fetch_add(rem, Ordering::Release); } + + // Finish all permit accounting before invoking callbacks. The shared helper attempts + // every wake even if one panics, and propagates the first panic afterward. + drop(waiters); + wake_all(first_waker.into_iter().chain(wakers)); } } @@ -479,6 +406,8 @@ fn acquired_or_enqueue( #[cfg(test)] mod tests { + use std::panic; + use std::panic::AssertUnwindSafe; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -508,8 +437,8 @@ mod tests { } #[test] - fn release_drains_more_than_one_wake_batch() { - const WAITER_COUNT: usize = WAKE_BATCH_SIZE + 3; + fn release_distributes_permits_to_all_waiters() { + const WAITER_COUNT: usize = 35; let semaphore = Semaphore::new(0); let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); @@ -533,8 +462,8 @@ mod tests { } #[test] - fn panicking_wakes_finish_later_batches_and_preserve_the_first_panic() { - const WAITER_COUNT: usize = WAKE_BATCH_SIZE * 2 + 1; + fn panicking_wakes_preserve_permits_and_the_first_panic() { + const WAITER_COUNT: usize = 65; struct TrackedWake { count: AtomicUsize, @@ -555,10 +484,12 @@ mod tests { .map(|index| { Arc::new(TrackedWake { count: AtomicUsize::new(0), - panic_message: match index { - 0 => Some("first wake panic"), - WAKE_BATCH_SIZE => Some("later wake panic"), - _ => None, + panic_message: if index == 0 { + Some("first wake panic") + } else if index == WAITER_COUNT / 2 { + Some("later wake panic") + } else { + None }, }) }) From 528e29f4dbc13767ae0bf56c7a6287c32b50e6e1 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 12 Sep 2026 23:57:27 +0800 Subject: [PATCH 7/8] refactor(semaphore): stream stack batches through panic-safe waking --- asyncband/src/internal/semaphore.rs | 150 ++++++++++++++++------ benchmarks/asyncband/semaphore/release.rs | 28 ++++ 2 files changed, 136 insertions(+), 42 deletions(-) diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index b32dbbed..ca4c3019 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -20,12 +20,14 @@ // The Tokio-derived portions remain licensed under the MIT License. // Asyncband substantially replaced the waiter lifecycle with queue-owned WaitList nodes, supports // queue-head permit debt for exact reductions, has no closed state or reserved flag bits, and uses -// its own cancellation, detachment, and waking machinery. +// its own cancellation, detachment, and batched-waking machinery. // Upstream source: // https://github.com/tokio-rs/tokio/blob/bb9d57017e100985f86d8ca41ac105ee9140423e/tokio/src/sync/batch_semaphore.rs use std::future::Future; +use std::mem::MaybeUninit; use std::pin::Pin; +use std::ptr; use std::sync::MutexGuard; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -54,6 +56,61 @@ struct WaitNode { waker: Option, } +const WAKE_BATCH_SIZE: usize = 32; + +/// The initialized entries in `wakers` are exactly `start..end`. +struct WakeBatch { + wakers: [MaybeUninit; WAKE_BATCH_SIZE], + start: usize, + end: usize, +} + +impl WakeBatch { + fn new() -> Self { + const UNINIT: MaybeUninit = MaybeUninit::uninit(); + Self { + wakers: [UNINIT; WAKE_BATCH_SIZE], + start: 0, + end: 0, + } + } + + fn push(&mut self, waker: Waker) { + debug_assert_eq!(self.start, 0); + debug_assert!(self.end < WAKE_BATCH_SIZE); + self.wakers[self.end].write(waker); + self.end += 1; + } + + fn is_full(&self) -> bool { + self.end == WAKE_BATCH_SIZE + } + + fn take_next(&mut self) -> Option { + if self.start == self.end { + self.start = 0; + self.end = 0; + return None; + } + + let index = self.start; + self.start += 1; + // SAFETY: `index` was within the initialized range before advancing `start`. + Some(unsafe { self.wakers[index].assume_init_read() }) + } +} + +impl Drop for WakeBatch { + fn drop(&mut self) { + let start = self.wakers[self.start..self.end] + .as_mut_ptr() + .cast::(); + let remaining = ptr::slice_from_raw_parts_mut(start, self.end - self.start); + // SAFETY: The initialized entries are exactly `start..end`. + unsafe { ptr::drop_in_place(remaining) }; + } +} + impl Semaphore { pub const fn new(permits: usize) -> Self { Self { @@ -184,55 +241,64 @@ impl Semaphore { fn insert_permits_with_lock( &self, mut rem: usize, - mut waiters: MutexGuard<'_, WaitList>, + waiters: MutexGuard<'_, WaitList>, ) { - // A single-waiter handoff should not allocate a wake buffer. - let mut first_waker = None; - let mut wakers = vec![]; - while rem > 0 { - match waiters.unlink_first_waiter(|node| { - if node.permits <= rem { - rem -= node.permits; - node.permits = 0; - true - } else { - node.permits -= rem; - rem = 0; - false + let mut batch = WakeBatch::new(); + let mut lock = Some(waiters); + + // One iterator covers the entire release. If a callback panics, `wake_all` keeps pulling + // batches during unwinding, so the remaining permits are still distributed and notified. + wake_all(std::iter::from_fn(|| { + loop { + if let Some(waker) = batch.take_next() { + return Some(waker); } - }) { - None => break, - Some((id, waiter)) => { - let remove_now = waiter.waker.is_none(); - if let Some(waker) = waiter.waker.take() { - if first_waker.is_none() { - first_waker = Some(waker); + if rem == 0 { + return None; + } + + let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock()); + while !batch.is_full() { + match waiters.unlink_first_waiter(|node| { + if node.permits <= rem { + rem -= node.permits; + node.permits = 0; + true } else { - wakers.push(waker); + node.permits -= rem; + rem = 0; + false + } + }) { + None => break, + Some((id, waiter)) => { + let remove_now = waiter.waker.is_none(); + if let Some(waker) = waiter.waker.take() { + batch.push(waker); + } + if remove_now { + waiters.remove_unlinked_waiter(id); + } } - } - if remove_now { - waiters.remove_unlinked_waiter(id); } } - } - } - if rem > 0 { - // Holding `waiters` serializes all permit additions. Concurrent operations can only - // remove permits, so the count cannot grow between this check and fetch_add. - let current = self.permits.load(Ordering::Relaxed); - assert!( - current.checked_add(rem).is_some(), - "number of added permits ({rem}) would overflow usize::MAX (prev: {current})" - ); - self.permits.fetch_add(rem, Ordering::Release); - } + if rem > 0 && waiters.is_empty() { + // Retire the remainder before the overflow check so unwinding cannot retry it. + let added = std::mem::take(&mut rem); + // The lock serializes additions; concurrent operations can only remove permits. + let current = self.permits.load(Ordering::Relaxed); + assert!( + current.checked_add(added).is_some(), + "number of added permits ({added}) would overflow usize::MAX (prev: {current})" + ); + self.permits.fetch_add(added, Ordering::Release); + } - // Finish all permit accounting before invoking callbacks. The shared helper attempts - // every wake even if one panics, and propagates the first panic afterward. - drop(waiters); - wake_all(first_waker.into_iter().chain(wakers)); + // Neither wake callbacks nor destruction of the taken waker run under this lock. + drop(waiters); + } + })); } } diff --git a/benchmarks/asyncband/semaphore/release.rs b/benchmarks/asyncband/semaphore/release.rs index d5423f43..4c400fa9 100644 --- a/benchmarks/asyncband/semaphore/release.rs +++ b/benchmarks/asyncband/semaphore/release.rs @@ -15,10 +15,16 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use asyncband::semaphore::Semaphore; use divan::Bencher; use divan::black_box; +use crate::support::bench_context; +use crate::support::defer_input_drop; +use crate::support::poll_pending; + #[divan::bench] fn fulfill_debt_repeatedly(bencher: Bencher) { const CYCLES: usize = 64; @@ -42,3 +48,25 @@ fn release(bencher: Bencher) { black_box(semaphore) }); } + +// The first two sizes distinguish a single handoff from fan-out; larger sizes measure bulk release. +#[divan::bench(args = [1, 2, 32, 256], sample_size = 64)] +fn release_to_waiters(bencher: Bencher, waiter_count: usize) { + bencher + .with_inputs(|| { + let semaphore = Arc::new(Semaphore::new(0)); + let mut context = bench_context(); + let mut waiters = (0..waiter_count) + .map(|_| Box::pin(semaphore.clone().acquire_owned(1))) + .collect::>(); + for waiter in &mut waiters { + poll_pending(waiter.as_mut(), &mut context); + } + (semaphore, waiters) + }) + .bench_local_values(|(semaphore, waiters)| { + // Only release and wake callbacks are timed; registration and future cleanup are not. + semaphore.release(black_box(waiter_count)); + defer_input_drop((semaphore, waiters), ()) + }); +} From a268c69aab72aaed6efc73787075be8e7ba439f2 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 12 Sep 2026 23:57:28 +0800 Subject: [PATCH 8/8] test(broadcast): simplify fixtures and cover bulk reclamation --- .../asyncband/broadcast/mpmc/bounded.rs | 31 +++ .../ecosystem/broadcast/mpmc/adapters.rs | 53 ++-- .../ecosystem/broadcast/mpmc/bounded.rs | 5 +- .../ecosystem/broadcast/mpmc/support.rs | 60 +---- .../tests/broadcast_mpmc_bounded_test.rs | 242 +++++++++--------- 5 files changed, 187 insertions(+), 204 deletions(-) diff --git a/benchmarks/asyncband/broadcast/mpmc/bounded.rs b/benchmarks/asyncband/broadcast/mpmc/bounded.rs index c0800f83..aabae2ef 100644 --- a/benchmarks/asyncband/broadcast/mpmc/bounded.rs +++ b/benchmarks/asyncband/broadcast/mpmc/bounded.rs @@ -27,6 +27,7 @@ use divan::Bencher; use divan::black_box; use crate::support::bench_context; +use crate::support::defer_input_drop; use crate::support::poll_pending; use crate::support::poll_pinned_ready; @@ -143,3 +144,33 @@ fn deliver_to_waiting_receiver(bencher: Bencher) { black_box(poll_pinned_ready(recv, &mut context).unwrap()) }); } + +#[divan::bench(args = [1, 2, 32, 256], sample_size = 64)] +fn drop_lagging_receiver_wakes_senders(bencher: Bencher, backlog: usize) { + bencher + .with_inputs(|| { + let (sender, mut fast) = mpmc::bounded(backlog); + let slow = sender.subscribe(); + for value in 0..backlog { + sender.try_send(value).unwrap(); + assert_eq!(fast.try_recv().unwrap(), value); + } + let mut context = bench_context(); + let mut sends = (0..backlog) + .map(|value| { + let sender = sender.clone(); + Box::pin(async move { sender.send(value).await }) + }) + .collect::>(); + for send in &mut sends { + poll_pending(send.as_mut(), &mut context); + } + (slow, fast, sends) + }) + .bench_local_values(|(slow, fast, sends)| { + // The fast subscription stays alive so this measures reclaim, not last-receiver exit. + // Preparing the backlog, parking senders, and disposing of futures are outside timing. + drop(slow); + defer_input_drop((fast, sends), ()) + }); +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs index 21224281..658707bb 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs @@ -144,11 +144,24 @@ pub trait BoundedBroadcastMpmc: Send + Sync + 'static { fn try_send(sender: &Self::Sender, value: usize); fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>); - fn send_blocking(sender: &Self::Sender, value: usize); + + fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + poll_ready(Self::send_async(sender, value), context); + } + + fn send_blocking(sender: &Self::Sender, value: usize) { + FutureExt::block_on(Self::send_async(sender, value)); + } + fn try_recv(receiver: &mut Self::Receiver) -> Option; - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize; - fn recv_blocking(receiver: &mut Self::Receiver) -> usize; + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(Self::recv_async(receiver), context) + } + + fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + FutureExt::block_on(Self::recv_async(receiver)) + } } impl BoundedBroadcastMpmc for Asyncband { @@ -177,14 +190,6 @@ impl BoundedBroadcastMpmc for Asyncband { receiver.recv().await.unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { - poll_ready(Self::send_async(sender, value), context); - } - - fn send_blocking(sender: &Self::Sender, value: usize) { - FutureExt::block_on(Self::send_async(sender, value)); - } - fn try_recv(receiver: &mut Self::Receiver) -> Option { match receiver.try_recv() { Ok(value) => Some(value), @@ -194,14 +199,6 @@ impl BoundedBroadcastMpmc for Asyncband { } } } - - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { - poll_ready(Self::recv_async(receiver), context) - } - - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { - FutureExt::block_on(Self::recv_async(receiver)) - } } impl BoundedBroadcastMpmc for AsyncBroadcast { @@ -234,14 +231,6 @@ impl BoundedBroadcastMpmc for AsyncBroadcast { receiver.recv_direct().await.unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { - poll_ready(Self::send_async(sender, value), context); - } - - fn send_blocking(sender: &Self::Sender, value: usize) { - FutureExt::block_on(Self::send_async(sender, value)); - } - fn try_recv(receiver: &mut Self::Receiver) -> Option { match receiver.try_recv() { Ok(value) => Some(value), @@ -249,12 +238,4 @@ impl BoundedBroadcastMpmc for AsyncBroadcast { Err(error) => panic!("unexpected async-broadcast receive error: {error}"), } } - - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { - poll_ready(Self::recv_async(receiver), context) - } - - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { - FutureExt::block_on(Self::recv_async(receiver)) - } } diff --git a/benchmarks/ecosystem/broadcast/mpmc/bounded.rs b/benchmarks/ecosystem/broadcast/mpmc/bounded.rs index db2f3d3e..d4cf7f9a 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/bounded.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/bounded.rs @@ -66,8 +66,7 @@ fn ready_round_trip(bencher: Bencher) { }); } -// `sample_size = 1` is required: `BoundedConcurrent` spawns its workers once and they exit after a -// single pass, so a second `run` on the same value would block forever. +// Keep one fixture per sample so workers from other fixtures are not alive during timing. #[divan::bench( types = [Asyncband, AsyncBroadcast], args = BOUNDED_SHAPES, @@ -77,7 +76,7 @@ fn ready_round_trip(bencher: Bencher) { )] fn concurrent(bencher: Bencher, shape: BoundedShape) { bencher - .with_inputs(|| BoundedConcurrent::::new(shape)) + .with_inputs(|| BoundedConcurrent::new::(shape)) .bench_local_refs(BoundedConcurrent::run); } diff --git a/benchmarks/ecosystem/broadcast/mpmc/support.rs b/benchmarks/ecosystem/broadcast/mpmc/support.rs index ccdc9780..cd6e83bd 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/support.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/support.rs @@ -224,101 +224,63 @@ impl Drop for Fanout { } } -/// Producers and receivers running concurrently against a channel far smaller than the batch. -/// -/// The unbounded fixtures above publish the whole batch before anyone drains it, which is only -/// safe because every peer is given room for the entire batch. That shape deadlocks a genuinely -/// bounded channel, so this one interleaves: every thread blocks, and the drain runs while the -/// producers are still publishing. -/// -/// This terminates. The run could only wedge if every producer and every receiver waited at the -/// same time, but a producer waits only while at least one message is retained, and a retained -/// message is by definition unread by the slowest receiver — so that receiver is runnable. The -/// counts balance exactly: the producers publish `BATCH_MESSAGES` between them and each receiver -/// consumes `BATCH_MESSAGES`, so no thread over- or under-runs. Every receiver is subscribed -/// before the first send, so every receiver sees every message. -/// -/// Benches using this must set `sample_size = 1`: the worker threads are spawned in `new` and exit -/// after one pass, so a second `run` on the same value would block forever. -pub struct BoundedConcurrent { +/// Concurrent producers and subscribers on native threads. Each subscriber drains the full batch. +/// Construction is outside timing; `run` includes barrier release, transfers, checksum validation, +/// and worker joins. +pub struct BoundedConcurrent { start: Arc, - done: Arc, workers: Vec>, - channel: PhantomData, } -impl BoundedConcurrent { - pub fn new(shape: BoundedShape) -> Self { +impl BoundedConcurrent { + pub fn new(shape: BoundedShape) -> Self { let BoundedShape { capacity, producers, receivers, } = shape; assert_eq!(BATCH_MESSAGES % producers, 0); - let (sender, receivers) = C::channel(capacity, receivers); let start = Arc::new(Barrier::new(producers + receivers.len() + 1)); - let done = Arc::new(Barrier::new(producers + receivers.len() + 1)); let messages_per_producer = BATCH_MESSAGES / producers; let mut workers = Vec::with_capacity(producers + receivers.len()); for mut receiver in receivers { let start = start.clone(); - let done = done.clone(); workers.push(thread::spawn(move || { start.wait(); let mut checksum = 0usize; for _ in 0..BATCH_MESSAGES { checksum = checksum.wrapping_add(C::recv_blocking(&mut receiver)); } - black_box(checksum); - done.wait(); + assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); })); } - for producer in 0..producers { let sender = sender.clone(); let start = start.clone(); - let done = done.clone(); workers.push(thread::spawn(move || { start.wait(); let first = producer * messages_per_producer; for value in first..first + messages_per_producer { C::send_blocking(&sender, black_box(value)); } - done.wait(); })); } - drop(sender); - Self { - start, - done, - workers, - channel: PhantomData, - } + Self { start, workers } } pub fn run(&mut self) { self.start.wait(); - self.done.wait(); - } -} - -impl Drop for BoundedConcurrent { - fn drop(&mut self) { - let panicking = thread::panicking(); for worker in self.workers.drain(..) { - let result = worker.join(); - if !panicking { - result.expect("bounded benchmark worker panicked"); - } + worker.join().expect("bounded benchmark worker panicked"); } } } -/// The same bounded workload on async tasks. Construction and task spawning happen outside the -/// timed section. Each fixture runs once, so its benchmark must use `sample_size = 1`. +/// The same bounded workload on async tasks. Construction and spawning are outside timing; `run` +/// includes barrier release, transfers, checksum validation, and task joins. pub struct BoundedTasks { start: Arc, tasks: JoinSet<()>, diff --git a/tests-integration/tests/broadcast_mpmc_bounded_test.rs b/tests-integration/tests/broadcast_mpmc_bounded_test.rs index 0d2da1c4..5acfd17b 100644 --- a/tests-integration/tests/broadcast_mpmc_bounded_test.rs +++ b/tests-integration/tests/broadcast_mpmc_bounded_test.rs @@ -17,37 +17,19 @@ use std::future::Future; use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; +use std::sync::Barrier; use std::task::Context; use std::task::Wake; use std::task::Waker; use std::thread; -use std::time::Duration; use asyncband::blocking::FutureExt; use asyncband::broadcast::mpmc::*; +use tests_integration::WakeCounter; +use tests_integration::assert_completes_without_deadlock; use tests_integration::poll_once; use tests_integration::waker_on_wake; -struct TrackWake(AtomicUsize); - -impl TrackWake { - fn new() -> Arc { - Arc::new(Self(AtomicUsize::new(0))) - } - - fn count(&self) -> usize { - self.0.load(Ordering::Relaxed) - } -} - -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, @@ -279,7 +261,7 @@ fn receive_that_vacates_the_head_wakes_a_blocked_sender() { let (tx, mut rx) = bounded(1); tx.try_send(0).unwrap(); - let tracker = TrackWake::new(); + let tracker = Arc::new(WakeCounter::default()); let waker = Waker::from(tracker.clone()); let mut send = Box::pin(tx.send(1)); assert!( @@ -299,7 +281,7 @@ fn parked_recv_that_reclaims_wakes_a_blocked_sender() { let (tx, mut rx) = bounded(1); tx.try_send(0).unwrap(); - let tracker = TrackWake::new(); + let tracker = Arc::new(WakeCounter::default()); let waker = Waker::from(tracker.clone()); let mut send = Box::pin(tx.send(1)); assert!( @@ -364,7 +346,9 @@ fn dropping_the_last_receiver_wakes_every_blocked_sender() { // More blocked producers than the drop will reclaim slots. Once no receiver remains every // send succeeds unconditionally, so waking only `reclaimed` of them would strand the rest. - let trackers = (0..BLOCKED).map(|_| TrackWake::new()).collect::>(); + let trackers = (0..BLOCKED) + .map(|_| Arc::new(WakeCounter::default())) + .collect::>(); let mut sends = (0..BLOCKED) .map(|value| Box::pin(tx.send(10 + value as i32))) .collect::>(); @@ -462,7 +446,7 @@ fn cancelled_notified_sender_passes_capacity_to_the_next_sender() { fn cancelled_recv_releases_its_waker() { let (tx, mut rx) = bounded(4); - let tracker = TrackWake::new(); + let tracker = Arc::new(WakeCounter::default()); let waker = Waker::from(tracker.clone()); let baseline = Arc::strong_count(&tracker); @@ -489,7 +473,7 @@ fn cancelled_recv_releases_its_waker() { fn dropping_a_woken_recv_keeps_another_receivers_waiter() { let (tx, mut rx1) = bounded::(2); let mut rx2 = tx.subscribe(); - let first = TrackWake::new(); + let first = Arc::new(WakeCounter::default()); let waker = Waker::from(first.clone()); let mut context = Context::from_waker(&waker); let mut recv1 = Box::pin(rx1.recv()); @@ -500,7 +484,7 @@ fn dropping_a_woken_recv_keeps_another_receivers_waiter() { assert_eq!(first.count(), 1); assert_eq!(rx2.try_recv(), Ok(1)); - let second = TrackWake::new(); + let second = Arc::new(WakeCounter::default()); let waker = Waker::from(second.clone()); let mut context = Context::from_waker(&waker); let mut recv2 = Box::pin(rx2.recv()); @@ -537,7 +521,7 @@ fn bounded_parked_recv_wakes_when_the_last_sender_drops() { let (tx, mut rx) = bounded::(4); let second_tx = tx.clone(); - let tracker = TrackWake::new(); + let tracker = Arc::new(WakeCounter::default()); let waker = Waker::from(tracker.clone()); let mut recv = Box::pin(rx.recv()); assert!( @@ -570,7 +554,9 @@ fn panicking_wake_does_not_strand_senders_after_a_large_reclaim() { fast.try_recv().unwrap(); } - let trackers = (0..40).map(|_| TrackWake::new()).collect::>(); + let trackers = (0..40) + .map(|_| Arc::new(WakeCounter::default())) + .collect::>(); let wakers = trackers .iter() .enumerate() @@ -671,9 +657,7 @@ fn panicking_payload_destructor_still_releases_capacity() { #[test] fn bounded_message_destructors_run_outside_the_channel_lock() { - let (finished_tx, finished_rx) = std::sync::mpsc::channel(); - - let worker = thread::spawn(move || { + assert_completes_without_deadlock(|| { let (tx, mut rx1) = bounded(8); let rx2 = tx.subscribe(); @@ -685,20 +669,51 @@ fn bounded_message_destructors_run_outside_the_channel_lock() { .unwrap(); } - // Draining both receivers reclaims the prefix, whose destructors re-enter the channel. - for _ in 0..4 { - rx1.try_recv().unwrap(); - } + // Reclaim through a receive, and then through receiver drops. + assert_eq!(rx1.try_recv().unwrap().value, 0); drop(rx2); + assert_eq!(rx1.try_recv().unwrap().value, 1); drop(rx1); - finished_tx.send(()).unwrap(); + // With no receiver, both send paths discard the payload immediately. + tx.try_send(Reentrant { + value: 4, + channel: Some(tx.clone()), + }) + .unwrap(); + FutureExt::block_on(tx.send(Reentrant { + value: 5, + channel: Some(tx.clone()), + })); }); +} - finished_rx - .recv_timeout(Duration::from_secs(10)) - .expect("reclaimed message destructors must not run while the channel is locked"); - worker.join().unwrap(); +#[test] +fn cancelling_a_blocked_send_drops_its_payload_outside_the_channel_lock() { + assert_completes_without_deadlock(|| { + let (tx, mut rx) = bounded(1); + tx.try_send(Reentrant { + value: 0, + channel: None, + }) + .unwrap(); + + let mut send = Box::pin(tx.send(Reentrant { + value: 1, + channel: Some(tx.clone()), + })); + assert!(poll_once(send.as_mut()).is_pending()); + drop(send); + + assert_eq!(rx.try_recv().unwrap().value, 0); + assert!(matches!(rx.try_recv(), Err(TryRecvError::Empty))); + tx.try_send(Reentrant { + value: 2, + channel: None, + }) + .unwrap(); + assert_eq!(rx.try_recv().unwrap().value, 2); + }); } // --------------------------------------------------------------------------------------------- @@ -710,101 +725,96 @@ fn dropping_the_last_receiver_never_strands_a_racing_producer() { const ROUNDS: u64 = 150; const PRODUCERS: u64 = 4; - // The deterministic tests above drop the receiver at a fixed point. This races the drop - // against producers entering the waiting path, which is the window where the channel decides - // whether anybody needs waking. A missed wake-up here parks a producer forever, so the failure - // mode is a hang rather than a wrong value — hence the timeout instead of an assertion. - let (finished_tx, finished_rx) = std::sync::mpsc::channel(); - - let worker = thread::spawn(move || { - for round in 0..ROUNDS { + // Bound each race independently of the total time spent starting threads across all rounds. + for round in 0..ROUNDS { + assert_completes_without_deadlock(move || { let (tx, rx) = bounded(1); tx.try_send(0).unwrap(); + let start = Arc::new(Barrier::new(PRODUCERS as usize + 1)); let producers = (0..PRODUCERS) .map(|producer| { let tx = tx.clone(); - thread::spawn(move || FutureExt::block_on(tx.send(round * 10 + producer + 1))) + let start = start.clone(); + thread::spawn(move || { + start.wait(); + FutureExt::block_on(tx.send(round * 10 + producer + 1)); + }) }) .collect::>(); + // Race removal against senders entering the wait path after all workers are ready. + start.wait(); drop(rx); for producer in producers { producer.join().unwrap(); } - } - finished_tx.send(()).unwrap(); - }); - - finished_rx - .recv_timeout(Duration::from_secs(60)) - .expect("a producer was left waiting after the last receiver went away"); - worker.join().unwrap(); + }); + } } #[test] fn bounded_concurrent_producers_commit_one_order_seen_by_every_receiver() { - const PRODUCERS: u64 = 4; - const PER_PRODUCER: u64 = 128; - const RECEIVERS: usize = 4; - const TOTAL: u64 = PRODUCERS * PER_PRODUCER; - - // Several producers publishing concurrently must still commit one contiguous order, and every - // subscription must observe that same order — not merely the same set. - // - // Capacity is far below the batch, so the producers really do block on the slowest receiver. - // This still terminates: the run could only wedge if every producer and every receiver waited - // at once, but a producer waits only while at least one message is retained, and a retained - // message is by definition unread by the slowest receiver — so that receiver is runnable. - let (tx, rx) = bounded(8); - let mut receivers = vec![rx]; - receivers.extend((1..RECEIVERS).map(|_| tx.subscribe())); - - let drains = receivers - .into_iter() - .map(|mut receiver| { - thread::spawn(move || { - let mut seen = Vec::with_capacity(TOTAL as usize); - for _ in 0..TOTAL { - seen.push(FutureExt::block_on(receiver.recv()).expect("sender dropped early")); - } - seen + assert_completes_without_deadlock(|| { + const PRODUCERS: u64 = 4; + const PER_PRODUCER: u64 = 128; + const RECEIVERS: usize = 4; + const TOTAL: u64 = PRODUCERS * PER_PRODUCER; + + // Capacity is below the message count, exercising backpressure while every subscription + // checks the same committed order rather than merely the same set of messages. + let (tx, rx) = bounded(8); + let mut receivers = vec![rx]; + receivers.extend((1..RECEIVERS).map(|_| tx.subscribe())); + + let drains = receivers + .into_iter() + .map(|mut receiver| { + thread::spawn(move || { + let mut seen = Vec::with_capacity(TOTAL as usize); + for _ in 0..TOTAL { + seen.push( + FutureExt::block_on(receiver.recv()).expect("sender dropped early"), + ); + } + seen + }) }) - }) - .collect::>(); - - let producers = (0..PRODUCERS) - .map(|worker| { - let tx = tx.clone(); - thread::spawn(move || { - for value in 0..PER_PRODUCER { - FutureExt::block_on(tx.send(worker * PER_PRODUCER + value)); - } + .collect::>(); + + let producers = (0..PRODUCERS) + .map(|worker| { + let tx = tx.clone(); + thread::spawn(move || { + for value in 0..PER_PRODUCER { + FutureExt::block_on(tx.send(worker * PER_PRODUCER + value)); + } + }) }) - }) - .collect::>(); + .collect::>(); - for producer in producers { - producer.join().unwrap(); - } - drop(tx); - - let orders = drains - .into_iter() - .map(|drain| drain.join().unwrap()) - .collect::>(); - - // Every subscription saw the identical sequence. - for (index, order) in orders.iter().enumerate().skip(1) { - assert_eq!( - order, &orders[0], - "subscription {index} observed a different committed order" - ); - } + for producer in producers { + producer.join().unwrap(); + } + drop(tx); + + let orders = drains + .into_iter() + .map(|drain| drain.join().unwrap()) + .collect::>(); + + // Every subscription saw the identical sequence. + for (index, order) in orders.iter().enumerate().skip(1) { + assert_eq!( + order, &orders[0], + "subscription {index} observed a different committed order" + ); + } - // And that sequence is every published value exactly once — no gap, no duplicate. - let mut sorted = orders[0].clone(); - sorted.sort_unstable(); - assert_eq!(sorted, (0..TOTAL).collect::>()); + // And that sequence is every published value exactly once — no gap, no duplicate. + let mut sorted = orders[0].clone(); + sorted.sort_unstable(); + assert_eq!(sorted, (0..TOTAL).collect::>()); + }); }