diff --git a/CHANGELOG.md b/CHANGELOG.md index ef8a4f69..087014cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,12 +25,14 @@ All notable changes to this project will be documented in this file. ### New features +* 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 opt-in bounded and unbounded `asyncband::mpmc` queues with cloneable producers and competing consumers, delivering each accepted value to exactly one receiver while a receiver remains. * Add an opt-in runtime-agnostic `Phaser` with shared observer handles, dynamic RAII participants registered individually or in batches through an owning iterator, `u64` phase numbers, split arrival/wait with cancellation-resilient retries, and a `close` operation that releases unfinished waits with `Closed`. * Add bounded MPSC `reserve` and `try_reserve` methods returning a `Permit`, allowing callers to wait for capacity before constructing a message; pending sends and reservations receive capacity in wait-queue order, and unused permits release capacity without claiming message order. ### Bug fixes +* Complete semaphore permit releases and notify all eligible waiters even if a wake callback panics. * Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. * Notify all blocked bounded MPSC senders on receiver disconnection even when a buffered message destructor panics. * Avoid deadlocks when a bounded MPSC sender's waker clone callback receives from the same channel. diff --git a/README.md b/README.md index 287ac380..ed3dae37 100644 --- a/README.md +++ b/README.md @@ -102,7 +102,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`oneshot`](https://docs.rs/asyncband/*/asyncband/oneshot/) | `oneshot` | Send one value from one sender to one receiver. | | | [`mpmc`](https://docs.rs/asyncband/*/asyncband/mpmc/) | `mpmc` | Distribute each value to exactly one of multiple competing receivers. | | | [`mpsc`](https://docs.rs/asyncband/*/asyncband/mpsc/) | `mpsc` | Send each value from multiple producers to one receiver with bounded backpressure or an unbounded queue. | -| | [`broadcast`](https://docs.rs/asyncband/*/asyncband/broadcast/) | `broadcast` | Deliver every value to receivers active at send time; retain an unbounded backlog until each consumes or drops. | +| | [`broadcast`](https://docs.rs/asyncband/*/asyncband/broadcast/) | `broadcast` | Deliver every value to active receivers with bounded backpressure or unbounded retention. | | | [`watch`](https://docs.rs/asyncband/*/asyncband/watch/) | `watch` | Publish cloneable latest state from one or more senders; receivers independently coalesce intermediate updates. | | Object reuse | [`pool`](https://docs.rs/asyncband/*/asyncband/pool/) | `pool` | Reuse objects through bounded or unbounded pool variants. | | Sync interop | [`FutureExt`](https://docs.rs/asyncband/*/asyncband/blocking/trait.FutureExt.html) | `blocking` | Drive one runtime-agnostic future from a blocking thread. | diff --git a/asyncband/src/broadcast/mpmc/bounded/mod.rs b/asyncband/src/broadcast/mpmc/bounded/mod.rs new file mode 100644 index 00000000..cadb8037 --- /dev/null +++ b/asyncband/src/broadcast/mpmc/bounded/mod.rs @@ -0,0 +1,743 @@ +// 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. +//! +//! A successful receive releases its subscription's claim before returning the value; processing +//! that value afterward does not hold capacity. The capacity limit excludes pending sends and +//! values already handed to application code. +//! +//! # Receivers +//! +//! Each receiver has an independent cursor. Use [`BoundedSender::subscribe`] or +//! [`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::wake_all; +use crate::internal::wakerset::WakerToken; + +#[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)); + /// # } + /// ``` + #[must_use = "the receiver is dropped immediately if it is not retained"] + pub fn subscribe(&self) -> BoundedReceiver { + let key = self.shared.inner.lock().log.subscribe(); + BoundedReceiver { + 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. + /// + /// The new receiver skips every value already published, including the latest retained value. + /// The original receiver is unchanged and continues to retain its own backlog until it + /// consumes those messages or is dropped. + /// + /// # Examples + /// + /// ``` + /// 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)); + /// ``` + #[must_use = "the receiver is dropped immediately if it is not retained"] + pub fn resubscribe(&self) -> Self { + let key = self.shared.inner.lock().log.subscribe(); + Self { + 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, + &self.receiver.shared.senders, + self.receiver.key, + &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..94f0b1eb --- /dev/null +++ b/asyncband/src/broadcast/mpmc/common.rs @@ -0,0 +1,513 @@ +// 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::wake_all; +use crate::internal::wakerset::WakerSet; +use crate::internal::wakerset::WakerToken; + +/// 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: WakerSet, +} + +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: WakerSet::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.take_all() + }; + wake_all(wakers); +} + +/// Releases a cancelled receive's waker registration, dropping the waker unlocked. +pub fn unregister( + inner: &Mutex>, + senders: &AtomicUsize, + key: SlotId, + token: &mut Option, +) { + let mut inner = inner.lock(); + if inner.log.unread(key) != 0 || senders.load(Ordering::Acquire) == 0 { + // Publication or disconnection detached this registration under the channel lock. + *token = None; + return; + } + + let waker = inner.waiters.unregister(token); + drop(inner); + 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. +/// +/// Checking the backlog and registering a waker under the same lock prevents a publication from +/// landing between those steps. Publication and disconnection detach all registrations, so their +/// ready paths clear the token without unregistering it. +pub fn poll_receive( + inner: &Mutex>, + senders: &AtomicUsize, + key: SlotId, + token: &mut Option, + cx: &mut Context<'_>, +) -> Poll, RecvError>> { + let mut inner = inner.lock(); + match inner.log.receive(key) { + Some(received) => { + *token = None; + Poll::Ready(Ok(received)) + } + None => { + if senders.load(Ordering::Acquire) == 0 { + *token = None; + return Poll::Ready(Err(RecvError::Disconnected)); + } + + let retired_waker = inner.waiters.register(token, cx.waker()); + drop(inner); + drop(retired_waker); + 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..636f8958 100644 --- a/asyncband/src/broadcast/mpmc/mod.rs +++ b/asyncband/src/broadcast/mpmc/mod.rs @@ -16,11 +16,38 @@ // 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. +//! +//! # Delivery and processing +//! +//! A receive advances its subscription before returning the value. The channel tracks unread +//! messages, not application work: retaining a received value or processing it asynchronously +//! does not hold backlog capacity. There is no acknowledgement or processing-completion barrier. +//! If cloning a received value panics, that subscription has still advanced past the value. +//! +//! Sending with no subscriptions discards the value and succeeds. A later subscription starts +//! with future publications; it does not replay discarded or previously retained values. +//! +//! Operations briefly acquire internal mutexes. No mutex is held across an await point or while +//! cloning or dropping payloads. The `try_*` methods do not wait for messages or capacity, but may +//! wait to acquire a mutex. Sending a value does not wait for subscribers to receive or process it. +mod bounded; +mod common; +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 3ac8e8f4..934374d3 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -26,7 +26,8 @@ //! Published values remain in the shared backlog until every receiver that was eligible for them //! has advanced past them or been dropped. Because sending has no capacity limit, one stalled //! receiver can make that backlog exhaust available memory. -//! [`UnboundedSender::retained_message_count`] reports its current length. +//! [`UnboundedSender::retained_message_count`] reports its current length. Use [`bounded`] when +//! producers should wait for the slowest receiver instead of growing the backlog. //! //! # Receivers //! @@ -50,11 +51,11 @@ //! assert_eq!(late.try_recv(), Ok("after subscription")); //! assert_eq!(late.try_recv(), Err(TryRecvError::Empty)); //! ``` +//! +//! [`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; @@ -62,11 +63,14 @@ 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::wake_all; -use crate::internal::wakerset::WakerSet; use crate::internal::wakerset::WakerToken; #[cfg(test)] @@ -87,18 +91,9 @@ mod tests; /// assert_eq!(receiver.try_recv(), Ok("ready")); /// ``` 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: WakerSet::new(), - }), + inner, senders: AtomicUsize::new(1), }); let sender = UnboundedSender { @@ -108,219 +103,8 @@ pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { (sender, receiver) } -/// A receive operation reached the end of its subscription. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum RecvError { - /// No sender remains and this receiver has consumed its entire backlog. - 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 {} - -/// A non-blocking receive did not yield a value. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum TryRecvError { - /// This receiver is caught up, but a sender can still publish more values. - Empty, - /// No sender remains and this receiver has consumed its entire backlog. - 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 {} - -/// 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: WakerSet, -} - -/// 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, @@ -355,14 +139,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.take_all() - }; - wake_all(wakers); - } + 1 => common::disconnect(&self.shared.inner), _ => { // there are still other senders left, do nothing } @@ -399,32 +176,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 values in the shared backlog. @@ -451,7 +213,7 @@ impl UnboundedSender { /// assert_eq!(publisher.retained_message_count(), 0); /// ``` pub fn retained_message_count(&self) -> usize { - self.shared.inner.lock().buffer.len() + self.shared.inner.lock().log.retained() } /// Subscribes a new receiver for values published from this point forward. @@ -474,11 +236,11 @@ impl UnboundedSender { /// ``` #[must_use = "the receiver is dropped immediately if it is not retained"] 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, + } } } @@ -501,7 +263,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); } @@ -565,51 +327,15 @@ impl UnboundedReceiver { /// assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected)); /// ``` 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) - } - } - - /// Creates another receiver at the current publication boundary. + /// Re-subscribes to the channel, returning a new receiver that starts receiving messages from + /// the *current* tail of the channel. /// /// The new receiver skips this receiver's unread backlog. The original receiver remains at its /// current position and continues retaining those values until it consumes them or is dropped. @@ -633,11 +359,11 @@ impl UnboundedReceiver { /// ``` #[must_use = "the receiver is dropped immediately if it is not retained"] 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 this receiver's unread value count. @@ -662,12 +388,7 @@ impl UnboundedReceiver { /// assert_eq!(receiver.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) } } @@ -683,21 +404,12 @@ impl Drop for Recv<'_, T> { return; } - let mut inner = self.receiver.shared.inner.lock(); - let cursor = *inner - .receivers - .get(self.receiver.key) - .expect("active broadcast receiver must be registered"); - if cursor != inner.tail || self.receiver.shared.senders.load(Ordering::Acquire) == 0 { - // A publisher or the final sender owns this registration or has already detached it - // under the channel lock. - self.token = None; - return; - } - - let waker = inner.waiters.unregister(&mut self.token); - drop(inner); - drop(waker); + common::unregister( + &self.receiver.shared.inner, + &self.receiver.shared.senders, + self.receiver.key, + &mut self.token, + ); } } @@ -707,26 +419,18 @@ impl Future for Recv<'_, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let Self { receiver, token } = self.get_mut(); - let received = { - let mut inner = receiver.shared.inner.lock(); - match inner.receive(receiver.key) { - Some(received) => received, - None => { - if receiver.shared.senders.load(Ordering::Acquire) == 0 { - *token = None; - return Poll::Ready(Err(RecvError::Disconnected)); - } - - let retired_waker = inner.waiters.register(token, cx.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, }; - 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 7b8cb352..7998e25a 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -98,17 +98,20 @@ pub(crate) mod value_cell; pub(crate) mod mutex; #[cfg(any( + feature = "broadcast", feature = "mpmc", feature = "mutex", feature = "rwlock", feature = "semaphore", ))] -// MPMC uses waiter notifications; mutexes and rwlocks use acquire/release operations; the public -// semaphore also exposes permit accounting. Single-primitive builds leave part of this API unused. +// Broadcast and MPMC use waiter notifications; mutexes and rwlocks use acquire/release operations; +// the public semaphore also exposes permit accounting. Single-primitive builds leave part of this +// API unused. #[allow(dead_code)] pub(crate) mod semaphore; #[cfg(any( + feature = "broadcast", feature = "event", feature = "mpmc", feature = "mpsc", diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index 8c2de0c1..ca4c3019 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -86,19 +86,17 @@ impl WakeBatch { self.end == WAKE_BATCH_SIZE } - fn wake_all(&mut self) { - wake_all(std::iter::from_fn(|| { - if self.start == self.end { - return None; - } + fn take_next(&mut self) -> Option { + if self.start == self.end { + self.start = 0; + self.end = 0; + return None; + } - let index = self.start; - self.start += 1; - // SAFETY: `index` was within the initialized range before advancing `start`. - Some(unsafe { self.wakers[index].assume_init_read() }) - })); - self.start = 0; - self.end = 0; + let index = self.start; + self.start += 1; + // SAFETY: `index` was within the initialized range before advancing `start`. + Some(unsafe { self.wakers[index].assume_init_read() }) } } @@ -206,7 +204,7 @@ impl Semaphore { } /// Adds `n` permits to the semaphore if there is any waiter. - #[cfg(feature = "mpmc")] + #[cfg(any(feature = "broadcast", feature = "mpmc"))] pub fn release_if_nonempty(&self, n: usize) { let waiters = self.waiters.lock(); if !waiters.is_empty() { @@ -215,7 +213,7 @@ impl Semaphore { } /// Adds as many permits until there is no waiter. - #[cfg(feature = "mpmc")] + #[cfg(any(feature = "broadcast", feature = "mpmc"))] pub fn notify_all(&self) { let mut waiters = self.waiters.lock(); let mut wakers = vec![]; @@ -237,7 +235,7 @@ impl Semaphore { } } drop(waiters); - crate::internal::wake_all(wakers.into_iter()); + wake_all(wakers.into_iter()); } fn insert_permits_with_lock( @@ -245,51 +243,62 @@ impl Semaphore { mut rem: usize, waiters: MutexGuard<'_, WaitList>, ) { - let mut wakers = WakeBatch::new(); - + let mut batch = WakeBatch::new(); let mut lock = Some(waiters); - while rem > 0 { - let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock()); - while !wakers.is_full() { - match waiters.unlink_first_waiter(|node| { - if node.permits <= rem { - rem -= node.permits; - node.permits = 0; - true - } else { - node.permits -= rem; - rem = 0; - false - } - }) { - None => break, - Some((id, waiter)) => { - let remove_now = waiter.waker.is_none(); - if let Some(waker) = waiter.waker.take() { - wakers.push(waker); + + // One iterator covers the entire release. If a callback panics, `wake_all` keeps pulling + // batches during unwinding, so the remaining permits are still distributed and notified. + wake_all(std::iter::from_fn(|| { + loop { + if let Some(waker) = batch.take_next() { + return Some(waker); + } + if rem == 0 { + return None; + } + + let mut waiters = lock.take().unwrap_or_else(|| self.waiters.lock()); + while !batch.is_full() { + match waiters.unlink_first_waiter(|node| { + if node.permits <= rem { + rem -= node.permits; + node.permits = 0; + true + } else { + node.permits -= rem; + rem = 0; + false } - if remove_now { - waiters.remove_unlinked_waiter(id); + }) { + None => break, + Some((id, waiter)) => { + let remove_now = waiter.waker.is_none(); + if let Some(waker) = waiter.waker.take() { + batch.push(waker); + } + if remove_now { + waiters.remove_unlinked_waiter(id); + } } } } - } - if rem > 0 && waiters.is_empty() { - // Holding `waiters` serializes all permit additions. Concurrent operations can - // only remove permits, so the count cannot grow between this check and fetch_add. - let current = self.permits.load(Ordering::Relaxed); - assert!( - current.checked_add(rem).is_some(), - "number of added permits ({rem}) would overflow usize::MAX (prev: {current})" - ); - self.permits.fetch_add(rem, Ordering::Release); - rem = 0; - } + if rem > 0 && waiters.is_empty() { + // Retire the remainder before the overflow check so unwinding cannot retry it. + let added = std::mem::take(&mut rem); + // The lock serializes additions; concurrent operations can only remove permits. + let current = self.permits.load(Ordering::Relaxed); + assert!( + current.checked_add(added).is_some(), + "number of added permits ({added}) would overflow usize::MAX (prev: {current})" + ); + self.permits.fetch_add(added, Ordering::Release); + } - drop(waiters); - wakers.wake_all(); - } + // Neither wake callbacks nor destruction of the taken waker run under this lock. + drop(waiters); + } + })); } } @@ -463,6 +472,8 @@ fn acquired_or_enqueue( #[cfg(test)] mod tests { + use std::panic; + use std::panic::AssertUnwindSafe; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -492,8 +503,8 @@ mod tests { } #[test] - fn release_drains_more_than_one_wake_batch() { - const WAITER_COUNT: usize = WAKE_BATCH_SIZE + 3; + fn release_distributes_permits_to_all_waiters() { + const WAITER_COUNT: usize = 35; let semaphore = Semaphore::new(0); let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); @@ -515,4 +526,64 @@ mod tests { } assert_eq!(semaphore.waiters.lock().occupied_len(), 0); } + + #[test] + fn panicking_wakes_preserve_permits_and_the_first_panic() { + const WAITER_COUNT: usize = 65; + + struct TrackedWake { + count: AtomicUsize, + panic_message: Option<&'static str>, + } + + impl Wake for TrackedWake { + fn wake(self: Arc) { + self.count.fetch_add(1, Ordering::Relaxed); + if let Some(message) = self.panic_message { + panic::panic_any(message); + } + } + } + + let semaphore = Semaphore::new(0); + let trackers = (0..WAITER_COUNT) + .map(|index| { + Arc::new(TrackedWake { + count: AtomicUsize::new(0), + panic_message: if index == 0 { + Some("first wake panic") + } else if index == WAITER_COUNT / 2 { + Some("later wake panic") + } else { + None + }, + }) + }) + .collect::>(); + let mut acquires = trackers + .iter() + .map(|tracker| { + let mut acquire = semaphore.poll_acquire(1); + assert!( + acquire + .poll_once(&Waker::from(tracker.clone())) + .is_pending() + ); + acquire + }) + .collect::>(); + + let payload = panic::catch_unwind(AssertUnwindSafe(|| semaphore.release(WAITER_COUNT + 2))) + .expect_err("the original wake panic must reach the caller"); + assert_eq!(payload.downcast_ref::<&str>(), Some(&"first wake panic")); + + for tracker in trackers { + assert_eq!(tracker.count.load(Ordering::Relaxed), 1); + } + for acquire in &mut acquires { + assert!(acquire.poll_once(Waker::noop()).is_ready()); + } + assert_eq!(semaphore.available_permits(), 2); + assert_eq!(semaphore.waiters.lock().occupied_len(), 0); + } } diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index ae99158c..afd6fafa 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -75,7 +75,7 @@ //! | | [`oneshot`] | `oneshot` | Send one value from one sender to one receiver. | //! | | [`mpmc`] | `mpmc` | Distribute each value to exactly one of multiple competing receivers. | //! | | [`mpsc`] | `mpsc` | Send each value from multiple producers to one receiver with bounded backpressure or an unbounded queue. | -//! | | [`broadcast`] | `broadcast` | Deliver every value to receivers active at send time; retain an unbounded backlog until each consumes or drops. | +//! | | [`broadcast`] | `broadcast` | Deliver every value to active receivers with bounded backpressure or unbounded retention. | //! | | [`watch`] | `watch` | Publish cloneable latest state from one or more senders; receivers independently coalesce intermediate updates. | //! | Object reuse | [`pool`] | `pool` | Reuse objects through bounded or unbounded pool variants. | //! | Sync interop | [`FutureExt`](blocking::FutureExt) | `blocking` | Drive one runtime-agnostic future from a blocking thread. | diff --git a/benchmarks/asyncband/broadcast/mpmc/bounded.rs b/benchmarks/asyncband/broadcast/mpmc/bounded.rs new file mode 100644 index 00000000..aabae2ef --- /dev/null +++ b/benchmarks/asyncband/broadcast/mpmc/bounded.rs @@ -0,0 +1,176 @@ +// 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::defer_input_drop; +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()) + }); +} + +#[divan::bench(args = [1, 2, 32, 256], sample_size = 64)] +fn drop_lagging_receiver_wakes_senders(bencher: Bencher, backlog: usize) { + bencher + .with_inputs(|| { + let (sender, mut fast) = mpmc::bounded(backlog); + let slow = sender.subscribe(); + for value in 0..backlog { + sender.try_send(value).unwrap(); + assert_eq!(fast.try_recv().unwrap(), value); + } + let mut context = bench_context(); + let mut sends = (0..backlog) + .map(|value| { + let sender = sender.clone(); + Box::pin(async move { sender.send(value).await }) + }) + .collect::>(); + for send in &mut sends { + poll_pending(send.as_mut(), &mut context); + } + (slow, fast, sends) + }) + .bench_local_values(|(slow, fast, sends)| { + // The fast subscription stays alive so this measures reclaim, not last-receiver exit. + // Preparing the backlog, parking senders, and disposing of futures are outside timing. + drop(slow); + defer_input_drop((fast, sends), ()) + }); +} diff --git a/benchmarks/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/asyncband/semaphore/release.rs b/benchmarks/asyncband/semaphore/release.rs index d5423f43..4c400fa9 100644 --- a/benchmarks/asyncband/semaphore/release.rs +++ b/benchmarks/asyncband/semaphore/release.rs @@ -15,10 +15,16 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use asyncband::semaphore::Semaphore; use divan::Bencher; use divan::black_box; +use crate::support::bench_context; +use crate::support::defer_input_drop; +use crate::support::poll_pending; + #[divan::bench] fn fulfill_debt_repeatedly(bencher: Bencher) { const CYCLES: usize = 64; @@ -42,3 +48,25 @@ fn release(bencher: Bencher) { black_box(semaphore) }); } + +// The first two sizes distinguish a single handoff from fan-out; larger sizes measure bulk release. +#[divan::bench(args = [1, 2, 32, 256], sample_size = 64)] +fn release_to_waiters(bencher: Bencher, waiter_count: usize) { + bencher + .with_inputs(|| { + let semaphore = Arc::new(Semaphore::new(0)); + let mut context = bench_context(); + let mut waiters = (0..waiter_count) + .map(|_| Box::pin(semaphore.clone().acquire_owned(1))) + .collect::>(); + for waiter in &mut waiters { + poll_pending(waiter.as_mut(), &mut context); + } + (semaphore, waiters) + }) + .bench_local_values(|(semaphore, waiters)| { + // Only release and wake callbacks are timed; registration and future cleanup are not. + semaphore.release(black_box(waiter_count)); + defer_input_drop((semaphore, waiters), ()) + }); +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs index 7b988701..658707bb 100644 --- a/benchmarks/ecosystem/broadcast/mpmc/adapters.rs +++ b/benchmarks/ecosystem/broadcast/mpmc/adapters.rs @@ -15,8 +15,11 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; use std::task::Context; +use asyncband::blocking::FutureExt; + use crate::support::poll_ready; pub struct Asyncband; @@ -128,3 +131,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 + Sync + 'static; + type Receiver: Send + 'static; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec); + fn try_send(sender: &Self::Sender, value: usize); + fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; + + fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + poll_ready(Self::send_async(sender, value), context); + } + + fn send_blocking(sender: &Self::Sender, value: usize) { + FutureExt::block_on(Self::send_async(sender, value)); + } + + fn try_recv(receiver: &mut Self::Receiver) -> Option; + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(Self::recv_async(receiver), context) + } + + fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + FutureExt::block_on(Self::recv_async(receiver)) + } +} + +impl BoundedBroadcastMpmc for Asyncband { + 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(); + } + + async fn send_async(sender: &Self::Sender, value: usize) { + sender.send(value).await; + } + + async fn recv_async(receiver: &mut Self::Receiver) -> usize { + receiver.recv().await.unwrap() + } + + fn 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") + } + } + } +} + +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(); + } + + async fn send_async(sender: &Self::Sender, value: usize) { + sender + .broadcast_direct(value) + .await + .expect("async-broadcast lost every receiver during benchmark"); + } + + async fn recv_async(receiver: &mut Self::Receiver) -> usize { + receiver.recv_direct().await.unwrap() + } + + fn 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}"), + } + } +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/bounded.rs b/benchmarks/ecosystem/broadcast/mpmc/bounded.rs new file mode 100644 index 00000000..d4cf7f9a --- /dev/null +++ b/benchmarks/ecosystem/broadcast/mpmc/bounded.rs @@ -0,0 +1,98 @@ +// 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. +// +// Sweep capacity and producer/subscription counts independently. Capacity one measures the +// per-message handoff; larger backlogs allow several messages to be outstanding. How effectively +// that headroom is used depends on scheduling and the slowest subscription, not just fanout. + +use divan::Bencher; +use divan::black_box; +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::BoundedTasks; +use super::support::ROUND_TRIP_CAPACITY; +use crate::support::bench_context; + +// Send-then-receive pairing keeps at most one message retained, so these never reach capacity. +#[divan::bench(types = [Asyncband, AsyncBroadcast], sample_size = 512)] +fn try_round_trip(bencher: Bencher) { + let (sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); + let mut receiver = receivers.pop().unwrap(); + + bencher.bench_local(|| { + C::try_send(&sender, black_box(usize::MAX)); + black_box(C::try_recv(&mut receiver).unwrap()) + }); +} + +#[divan::bench(types = [Asyncband, AsyncBroadcast], sample_size = 512)] +fn ready_round_trip(bencher: Bencher) { + let mut context = bench_context(); + let (sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); + 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)) + }); +} + +// Keep one fixture per sample so workers from other fixtures are not alive during timing. +#[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); +} + +#[divan::bench( + types = [Asyncband, AsyncBroadcast], + args = BOUNDED_SHAPES, + sample_count = 10, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn scheduled(bencher: Bencher, shape: BoundedShape) { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .build() + .unwrap(); + bencher + .with_inputs(|| BoundedTasks::new::(&runtime, shape)) + .bench_local_refs(|tasks| tasks.run(&runtime)); +} diff --git a/benchmarks/ecosystem/broadcast/mpmc/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..cd6e83bd 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; @@ -22,14 +23,79 @@ use std::thread; use std::thread::JoinHandle; use divan::black_box; +use tokio::runtime::Runtime; +use tokio::task::JoinSet; +use super::adapters::BoundedBroadcastMpmc; use super::adapters::BroadcastMpmc; 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 bounds the shared backlog, independently of the subscription count. Capacity one +/// forces the next publication to wait until every subscription advances; larger capacities let +/// producers run ahead by that many messages, subject to scheduling and consumer progress. +#[derive(Clone, Copy)] +pub struct BoundedShape { + pub capacity: usize, + pub producers: usize, + pub receivers: usize, +} + +impl BoundedShape { + 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(1, 1, 1), + BoundedShape::new(1, 1, 8), + BoundedShape::new(1, 8, 1), + BoundedShape::new(1, 8, 8), + BoundedShape::new(2, 1, 1), + BoundedShape::new(2, 1, 8), + BoundedShape::new(2, 8, 1), + BoundedShape::new(2, 8, 8), + BoundedShape::new(8, 1, 1), + BoundedShape::new(8, 1, 8), + BoundedShape::new(8, 8, 1), + BoundedShape::new(8, 8, 8), + BoundedShape::new(64, 1, 1), + BoundedShape::new(64, 1, 8), + BoundedShape::new(64, 1, 32), + BoundedShape::new(64, 8, 1), + BoundedShape::new(64, 8, 8), + BoundedShape::new(64, 8, 32), + BoundedShape::new(1024, 1, 1), + BoundedShape::new(1024, 1, 8), + BoundedShape::new(1024, 1, 32), + BoundedShape::new(1024, 8, 1), + BoundedShape::new(1024, 8, 8), + BoundedShape::new(1024, 8, 32), +]; + fn recv(receiver: &mut C::Receiver) -> usize { C::try_recv(receiver).expect("the published benchmark batch must be ready") } @@ -157,3 +223,119 @@ impl Drop for Fanout { } } } + +/// Concurrent producers and subscribers on native threads. Each subscriber drains the full batch. +/// Construction is outside timing; `run` includes barrier release, transfers, checksum validation, +/// and worker joins. +pub struct BoundedConcurrent { + start: Arc, + workers: Vec>, +} + +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 messages_per_producer = BATCH_MESSAGES / producers; + let mut workers = Vec::with_capacity(producers + receivers.len()); + + for mut receiver in receivers { + let start = start.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)); + } + assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); + })); + } + for producer in 0..producers { + let sender = sender.clone(); + let start = start.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)); + } + })); + } + + Self { start, workers } + } + + pub fn run(&mut self) { + self.start.wait(); + for worker in self.workers.drain(..) { + worker.join().expect("bounded benchmark worker panicked"); + } + } +} + +/// The same bounded workload on async tasks. Construction and spawning are outside timing; `run` +/// includes barrier release, transfers, checksum validation, and task joins. +pub struct BoundedTasks { + start: Arc, + tasks: JoinSet<()>, +} + +impl BoundedTasks { + pub fn new(runtime: &Runtime, shape: BoundedShape) -> Self { + let BoundedShape { + capacity, + producers, + receivers, + } = shape; + assert_eq!(BATCH_MESSAGES % producers, 0); + let (sender, receivers) = C::channel(capacity, receivers); + let start = Arc::new(tokio::sync::Barrier::new(producers + receivers.len() + 1)); + let mut tasks = JoinSet::new(); + + for mut receiver in receivers { + let start = start.clone(); + tasks.spawn_on( + async move { + start.wait().await; + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(C::recv_async(&mut receiver).await); + } + assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); + }, + runtime.handle(), + ); + } + for producer in 0..producers { + let sender = sender.clone(); + let start = start.clone(); + tasks.spawn_on( + async move { + start.wait().await; + let first = producer * (BATCH_MESSAGES / producers); + for value in first..first + BATCH_MESSAGES / producers { + C::send_async(&sender, black_box(value)).await; + } + }, + runtime.handle(), + ); + } + + Self { start, tasks } + } + + pub fn run(&mut self, runtime: &Runtime) { + runtime.block_on(async { + self.start.wait().await; + while let Some(result) = self.tasks.join_next().await { + result.expect("bounded benchmark task panicked"); + } + }); + } +} 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..5acfd17b --- /dev/null +++ b/tests-integration/tests/broadcast_mpmc_bounded_test.rs @@ -0,0 +1,820 @@ +// 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::Barrier; +use std::task::Context; +use std::task::Wake; +use std::task::Waker; +use std::thread; + +use asyncband::blocking::FutureExt; +use asyncband::broadcast::mpmc::*; +use tests_integration::WakeCounter; +use tests_integration::assert_completes_without_deadlock; +use tests_integration::poll_once; +use tests_integration::waker_on_wake; + +/// 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 = Arc::new(WakeCounter::default()); + 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 = Arc::new(WakeCounter::default()); + 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(|_| Arc::new(WakeCounter::default())) + .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 = Arc::new(WakeCounter::default()); + 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 +// --------------------------------------------------------------------------------------------- + +#[test] +fn dropping_a_woken_recv_keeps_another_receivers_waiter() { + let (tx, mut rx1) = bounded::(2); + let mut rx2 = tx.subscribe(); + let first = Arc::new(WakeCounter::default()); + let waker = Waker::from(first.clone()); + let mut context = Context::from_waker(&waker); + let mut recv1 = Box::pin(rx1.recv()); + + assert!(recv1.as_mut().poll(&mut context).is_pending()); + + tx.try_send(1).unwrap(); + assert_eq!(first.count(), 1); + assert_eq!(rx2.try_recv(), Ok(1)); + + let second = Arc::new(WakeCounter::default()); + let waker = Waker::from(second.clone()); + let mut context = Context::from_waker(&waker); + let mut recv2 = Box::pin(rx2.recv()); + assert!(recv2.as_mut().poll(&mut context).is_pending()); + + // `recv1` was already woken, so dropping it must not release the slot `recv2` now owns. + drop(recv1); + tx.try_send(2).unwrap(); + + assert_eq!(second.count(), 1); +} + +#[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 = Arc::new(WakeCounter::default()); + 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 panicking_wake_does_not_strand_senders_after_a_large_reclaim() { + let (tx, mut fast) = bounded(40); + let slow = tx.subscribe(); + for value in 0..40 { + tx.try_send(value).unwrap(); + fast.try_recv().unwrap(); + } + + let trackers = (0..40) + .map(|_| Arc::new(WakeCounter::default())) + .collect::>(); + let wakers = trackers + .iter() + .enumerate() + .map(|(index, tracker)| { + let tracker = tracker.clone(); + waker_on_wake(move || { + tracker.wake(); + assert_ne!(index, 0, "first sender wake panics"); + }) + }) + .collect::>(); + let mut sends = (40..80) + .map(|value| Box::pin(tx.send(value))) + .collect::>(); + for (send, waker) in sends.iter_mut().zip(&wakers) { + assert!( + send.as_mut() + .poll(&mut Context::from_waker(waker)) + .is_pending() + ); + } + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(slow))); + assert!(result.is_err()); + assert_eq!(tx.retained_message_count(), 0); + for tracker in trackers { + assert_eq!(tracker.count(), 1); + } + // Every send fits without another receive. All must have been notified, not merely made + // ready for a poll that an executor would otherwise have no reason to perform. + for send in &mut sends { + assert!(poll_once(send.as_mut()).is_ready()); + } + for value in 40..80 { + assert_eq!(fast.try_recv(), Ok(value)); + } +} + +#[test] +fn bounded_panicking_clone_leaves_the_channel_consistent() { + let (tx, mut rx1) = bounded(4); + 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() { + assert_completes_without_deadlock(|| { + 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(); + } + + // Reclaim through a receive, and then through receiver drops. + assert_eq!(rx1.try_recv().unwrap().value, 0); + drop(rx2); + assert_eq!(rx1.try_recv().unwrap().value, 1); + drop(rx1); + + // With no receiver, both send paths discard the payload immediately. + tx.try_send(Reentrant { + value: 4, + channel: Some(tx.clone()), + }) + .unwrap(); + FutureExt::block_on(tx.send(Reentrant { + value: 5, + channel: Some(tx.clone()), + })); + }); +} + +#[test] +fn cancelling_a_blocked_send_drops_its_payload_outside_the_channel_lock() { + assert_completes_without_deadlock(|| { + let (tx, mut rx) = bounded(1); + tx.try_send(Reentrant { + value: 0, + channel: None, + }) + .unwrap(); + + let mut send = Box::pin(tx.send(Reentrant { + value: 1, + channel: Some(tx.clone()), + })); + assert!(poll_once(send.as_mut()).is_pending()); + drop(send); + + assert_eq!(rx.try_recv().unwrap().value, 0); + assert!(matches!(rx.try_recv(), Err(TryRecvError::Empty))); + tx.try_send(Reentrant { + value: 2, + channel: None, + }) + .unwrap(); + assert_eq!(rx.try_recv().unwrap().value, 2); + }); +} + +// --------------------------------------------------------------------------------------------- +// Concurrency +// --------------------------------------------------------------------------------------------- + +#[test] +fn dropping_the_last_receiver_never_strands_a_racing_producer() { + const ROUNDS: u64 = 150; + const PRODUCERS: u64 = 4; + + // Bound each race independently of the total time spent starting threads across all rounds. + for round in 0..ROUNDS { + assert_completes_without_deadlock(move || { + let (tx, rx) = bounded(1); + tx.try_send(0).unwrap(); + let start = Arc::new(Barrier::new(PRODUCERS as usize + 1)); + + let producers = (0..PRODUCERS) + .map(|producer| { + let tx = tx.clone(); + let start = start.clone(); + thread::spawn(move || { + start.wait(); + FutureExt::block_on(tx.send(round * 10 + producer + 1)); + }) + }) + .collect::>(); + + // Race removal against senders entering the wait path after all workers are ready. + start.wait(); + drop(rx); + + for producer in producers { + producer.join().unwrap(); + } + }); + } +} + +#[test] +fn bounded_concurrent_producers_commit_one_order_seen_by_every_receiver() { + assert_completes_without_deadlock(|| { + const PRODUCERS: u64 = 4; + const PER_PRODUCER: u64 = 128; + const RECEIVERS: usize = 4; + const TOTAL: u64 = PRODUCERS * PER_PRODUCER; + + // Capacity is below the message count, exercising backpressure while every subscription + // checks the same committed order rather than merely the same set of messages. + let (tx, rx) = bounded(8); + let mut receivers = vec![rx]; + receivers.extend((1..RECEIVERS).map(|_| tx.subscribe())); + + let drains = receivers + .into_iter() + .map(|mut receiver| { + thread::spawn(move || { + let mut seen = Vec::with_capacity(TOTAL as usize); + for _ in 0..TOTAL { + seen.push( + FutureExt::block_on(receiver.recv()).expect("sender dropped early"), + ); + } + seen + }) + }) + .collect::>(); + + let producers = (0..PRODUCERS) + .map(|worker| { + let tx = tx.clone(); + thread::spawn(move || { + for value in 0..PER_PRODUCER { + FutureExt::block_on(tx.send(worker * PER_PRODUCER + value)); + } + }) + }) + .collect::>(); + + 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 f421a63d..00d44dcc 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -104,8 +104,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::(); @@ -192,8 +195,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::>();