From 4ffb990603d51b684fdafe428621085dbbef3f98 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Mon, 31 Aug 2026 10:31:57 +0800 Subject: [PATCH] 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))); + }, + ); +}