From 7b9817359060ca70a61520830f4199c683aeef98 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Mon, 14 Sep 2026 15:40:01 +0800 Subject: [PATCH 1/9] feat(broadcast): add an SPMC specialization Summary Add `broadcast::spmc` with a non-cloneable sender whose publish methods take `&mut self`. Bounded waits for the slowest active subscription; unbounded never waits. Existing `broadcast::mpmc` is unchanged. --- CHANGELOG.md | 1 + asyncband/src/broadcast/mod.rs | 10 + asyncband/src/broadcast/spmc/bounded/mod.rs | 676 ++++++++++++++++ asyncband/src/broadcast/spmc/bounded/tests.rs | 83 ++ asyncband/src/broadcast/spmc/common.rs | 551 +++++++++++++ asyncband/src/broadcast/spmc/error.rs | 103 +++ asyncband/src/broadcast/spmc/mod.rs | 74 ++ asyncband/src/broadcast/spmc/unbounded/mod.rs | 421 ++++++++++ .../src/broadcast/spmc/unbounded/tests.rs | 70 ++ benchmarks/asyncband/broadcast/mod.rs | 1 + .../asyncband/broadcast/spmc/bounded.rs | 160 ++++ benchmarks/asyncband/broadcast/spmc/mod.rs | 19 + .../asyncband/broadcast/spmc/unbounded.rs | 191 +++++ benchmarks/ecosystem/broadcast/mod.rs | 1 + .../ecosystem/broadcast/spmc/adapters.rs | 310 ++++++++ .../ecosystem/broadcast/spmc/bounded.rs | 87 +++ benchmarks/ecosystem/broadcast/spmc/mod.rs | 21 + .../ecosystem/broadcast/spmc/support.rs | 232 ++++++ .../ecosystem/broadcast/spmc/unbounded.rs | 70 ++ .../tests/broadcast_spmc_bounded_test.rs | 732 ++++++++++++++++++ .../tests/broadcast_spmc_unbounded_test.rs | 499 ++++++++++++ tests-integration/tests/traits_test.rs | 17 + xtask/src/main.rs | 4 + 23 files changed, 4333 insertions(+) create mode 100644 asyncband/src/broadcast/spmc/bounded/mod.rs create mode 100644 asyncband/src/broadcast/spmc/bounded/tests.rs create mode 100644 asyncband/src/broadcast/spmc/common.rs create mode 100644 asyncband/src/broadcast/spmc/error.rs create mode 100644 asyncband/src/broadcast/spmc/mod.rs create mode 100644 asyncband/src/broadcast/spmc/unbounded/mod.rs create mode 100644 asyncband/src/broadcast/spmc/unbounded/tests.rs create mode 100644 benchmarks/asyncband/broadcast/spmc/bounded.rs create mode 100644 benchmarks/asyncband/broadcast/spmc/mod.rs create mode 100644 benchmarks/asyncband/broadcast/spmc/unbounded.rs create mode 100644 benchmarks/ecosystem/broadcast/spmc/adapters.rs create mode 100644 benchmarks/ecosystem/broadcast/spmc/bounded.rs create mode 100644 benchmarks/ecosystem/broadcast/spmc/mod.rs create mode 100644 benchmarks/ecosystem/broadcast/spmc/support.rs create mode 100644 benchmarks/ecosystem/broadcast/spmc/unbounded.rs create mode 100644 tests-integration/tests/broadcast_spmc_bounded_test.rs create mode 100644 tests-integration/tests/broadcast_spmc_unbounded_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 81e16b5f..e2834aed 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ All notable changes to this project will be documented in this file. * Add `event::AutoResetEvent`, a reusable signal that releases one waiter, retains at most one unassigned signal that can be cleared with `reset`, and transfers assigned signals when waits are cancelled. * Add `ManualResetEvent::try_wait` to check readiness without registering a waiter or consuming the set state. +* Add `broadcast::spmc`, a lossless single-producer broadcast family with a non-cloneable sender whose publish methods require exclusive access; bounded retains at most the requested capacity and makes the producer wait for the slowest active subscription, while unbounded never waits and lets the retained backlog grow. * 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`. diff --git a/asyncband/src/broadcast/mod.rs b/asyncband/src/broadcast/mod.rs index 81e56c55..0d9f5fab 100644 --- a/asyncband/src/broadcast/mod.rs +++ b/asyncband/src/broadcast/mod.rs @@ -16,5 +16,15 @@ // under the License. //! Broadcast channels grouped by producer topology. +//! +//! [`mpmc`] supports any number of concurrent producers. [`spmc`] is the single-producer +//! specialization: its sender is not [`Clone`] and publish methods require exclusive access +//! (`&mut self`). Choose `spmc` when the program has one publisher; choose `mpmc` when it does +//! not. Choose `spmc` for the exclusive-send API; do not assume it is faster on every path. +//! +//! Both topologies are fan-out broadcast: every accepted value is delivered to every active +//! subscription. That is a different delivery family from a competing crate-root `spmc` queue, +//! which would give each value to exactly one receiver. pub mod mpmc; +pub mod spmc; diff --git a/asyncband/src/broadcast/spmc/bounded/mod.rs b/asyncband/src/broadcast/spmc/bounded/mod.rs new file mode 100644 index 00000000..8ec3ccc4 --- /dev/null +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -0,0 +1,676 @@ +// 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 single-producer multi-consumer broadcast channel with a bounded buffer. +//! +//! Each message sent 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 the producer wait. The sender is not [`Clone`]; [`BoundedSender::send`] and +//! [`BoundedSender::try_send`] take `&mut self` so a second producer cannot exist at compile time. +//! +//! # 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 the producer, +//! 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. +//! +//! [`BoundedSender::send`] holds `&mut self` while it waits, so the sender cannot +//! [`subscribe`](BoundedSender::subscribe) until that send completes. Create the extra +//! subscription from a live receiver with [`BoundedReceiver::resubscribe`] instead. +//! +//! # Cancel safety +//! +//! Publication itself is one indivisible step, so cancelling a send can never leave a gap in the +//! committed order. Cancelling a pending `recv` does not advance the subscription cursor. +//! +//! # Examples +//! +//! Basic usage: +//! +//! ``` +//! use asyncband::broadcast::spmc; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (mut tx, mut rx1) = spmc::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::spmc; +//! use asyncband::broadcast::spmc::TrySendError; +//! +//! let (mut tx, mut rx1) = spmc::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::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::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::spmc; +/// +/// let (mut tx, mut rx) = spmc::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, + }); + let sender = BoundedSender { + shared: shared.clone(), + }; + let receiver = BoundedReceiver { shared, key }; + (sender, receiver) +} + +struct Shared { + /// Buffer, receiver cursors, parked receivers, and the single parked producer, all under one + /// lock. + inner: Mutex>, + /// `1` while the sender is alive, `0` after it is dropped. + senders: AtomicUsize, + /// The logical limit on the retained backlog. + capacity: usize, +} + +/// The sending side of a bounded broadcast channel. +/// +/// This handle is not [`Clone`]. Dropping it disconnects the channel. Each receiver may drain its +/// own buffered messages before observing disconnection. +pub struct BoundedSender { + shared: Arc>, +} + +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) { + self.shared.senders.store(0, Ordering::Release); + common::disconnect(&self.shared.inner); + } +} + +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::spmc; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (mut tx, mut rx) = spmc::bounded(1); + /// tx.send(10).await; + /// assert_eq!(rx.recv().await, Ok(10)); + /// # } + /// ``` + pub async fn send(&mut self, value: T) { + let value = match self.try_send(value) { + Ok(()) => return, + Err(TrySendError::Full(value)) => value, + }; + + struct SendState<'a, T> { + sender: &'a mut BoundedSender, + // Boxed once, out of the critical section, and reused by every retry. Dropped after + // `SendState::drop` has already released the producer slot, so a cancelled send + // unregisters before running the payload destructor. + value: Option>, + } + + impl Drop for SendState<'_, T> { + fn drop(&mut self) { + // Take the slot with the channel unlocked afterward so the replaced waker is + // dropped outside the lock. The payload in `value` is dropped only after this + // returns. + let waker = { + let mut inner = self.sender.shared.inner.lock(); + inner.producer.take() + }; + drop(waker); + } + } + + impl SendState<'_, T> { + fn poll_send(&mut self, cx: &mut Context<'_>) -> Poll<()> { + let msg = match self.value.take() { + Some(msg) => msg, + None => return Poll::Ready(()), + }; + + let mut inner = self.sender.shared.inner.lock(); + + if !inner.log.has_receivers() { + inner.log.publish_discarded(); + let retired_producer = inner.producer.take(); + let wakers = inner.waiters.drain(); + drop(inner); + wake_all(wakers); + drop(retired_producer); + drop(msg); + return Poll::Ready(()); + } + + if inner.log.retained() == self.sender.shared.capacity { + // Same critical section as the capacity check: a reclaim that lands between + // those two observations cannot skip this waiter. + let retired = inner.producer.replace(cx.waker().clone()); + drop(inner); + drop(retired); + self.value = Some(msg); + return Poll::Pending; + } + + inner.log.publish_retained(msg); + let retired_producer = inner.producer.take(); + let wakers = inner.waiters.drain(); + drop(inner); + // Wake receivers before dropping the retired producer waker: that waker is this + // send, so it must not run under the lock, but a panic in its Drop must not skip + // the receiver wake-ups either. + wake_all(wakers); + drop(retired_producer); + Poll::Ready(()) + } + } + + let mut send = SendState { + sender: self, + 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::spmc; + /// use asyncband::broadcast::spmc::TrySendError; + /// + /// let (mut tx, mut rx) = spmc::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(&mut 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. + self.publish(value, Arc::new).map_err(TrySendError::Full) + } + + /// 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. + /// + /// 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

(&mut 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::spmc; + /// + /// let (mut tx, mut rx) = spmc::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 the producer waits. + /// + /// 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::spmc; + /// + /// let (tx, _rx) = spmc::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. + /// + /// This cannot be called while [`send`](BoundedSender::send) is waiting, because that future + /// holds `&mut self`. Use [`BoundedReceiver::resubscribe`] from a live receiver instead. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc; + /// use asyncband::broadcast::spmc::TryRecvError; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (mut tx, _rx) = spmc::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 the producer. +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, producer) = { + let mut inner = self.shared.inner.lock(); + let reclaimed = inner.log.remove_receiver(self.key); + let drained_last = !inner.log.has_receivers(); + let producer = common::take_producer_on_reclaim(&mut inner, &reclaimed, drained_last); + (reclaimed, producer) + }; + + // Wake before dropping reclaimed payloads: a panicking destructor must not strand the + // producer on capacity this drop already freed. + common::wake_producer(producer); + drop(reclaimed); + } +} + +impl BoundedReceiver { + /// Receives the next value for this receiver. + /// + /// # Returns + /// + /// * `Ok(T)`: The next message. + /// * `Err(RecvError::Disconnected)`: The sender has 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::spmc; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (mut tx, mut rx) = spmc::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)`: The sender has been dropped and this receiver has no + /// remaining messages. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc; + /// + /// let (mut tx, mut rx) = spmc::bounded(4); + /// tx.try_send(10).unwrap(); + /// assert_eq!(rx.try_recv(), Ok(10)); + /// ``` + pub fn try_recv(&mut self) -> Result { + let ((msg, reclaimed), producer) = + common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; + + // Wake before taking the payload: `take_msg` runs `T::clone` and `T::drop`, and if either + // panics the slot this receive already freed would otherwise never reach the parked + // producer, stalling it permanently. + common::wake_producer(producer); + 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. + /// + /// This is also how to add a subscription while [`BoundedSender::send`] is waiting: that + /// future holds `&mut self` on the sender, so [`BoundedSender::subscribe`] cannot be called + /// until it completes. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc; + /// + /// let (mut tx, mut rx) = spmc::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)); + /// ``` + /// + /// Adding a subscription while the producer is waiting for capacity: + /// + /// ``` + /// use std::future::Future; + /// use std::task::Context; + /// use std::task::Poll; + /// use std::task::Waker; + /// + /// use asyncband::broadcast::spmc; + /// + /// let (mut tx, rx) = spmc::bounded(1); + /// tx.try_send(0).unwrap(); + /// + /// let mut send = Box::pin(tx.send(1)); + /// let mut cx = Context::from_waker(Waker::noop()); + /// assert!(matches!(send.as_mut().poll(&mut cx), Poll::Pending)); + /// + /// // `tx.subscribe()` would not compile here: `send` holds `&mut tx`. + /// let _late = rx.resubscribe(); + /// assert_eq!(rx.unread_message_count(), 1); + /// ``` + #[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::spmc; + /// + /// let (mut tx, mut rx) = spmc::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), producer) = 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(outcome)) => outcome, + }; + + // Wake before taking the payload, for the same reason as `try_recv`: a panicking + // `T::clone` must not strand the producer on a slot this receive already freed. + common::wake_producer(producer); + Poll::Ready(Ok(common::take_msg(msg, reclaimed))) + } +} diff --git a/asyncband/src/broadcast/spmc/bounded/tests.rs b/asyncband/src/broadcast/spmc/bounded/tests.rs new file mode 100644 index 00000000..ae110faa --- /dev/null +++ b/asyncband/src/broadcast/spmc/bounded/tests.rs @@ -0,0 +1,83 @@ +// 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_spmc_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 (mut 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 (mut 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 at_most_one_producer_can_wait() { + let (mut tx, mut rx) = bounded(1); + tx.try_send(0).unwrap(); + let shared = tx.shared.clone(); + + let mut cx = Context::from_waker(Waker::noop()); + let mut send = Box::pin(tx.send(1)); + assert!(send.as_mut().poll(&mut cx).is_pending()); + assert!(shared.inner.lock().producer.is_some()); + + assert_eq!(rx.try_recv(), Ok(0)); + assert!(send.as_mut().poll(&mut cx).is_ready()); + drop(send); + assert!(shared.inner.lock().producer.is_none()); +} diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs new file mode 100644 index 00000000..da5eb397 --- /dev/null +++ b/asyncband/src/broadcast/spmc/common.rs @@ -0,0 +1,551 @@ +// 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 SPMC broadcast +//! channels. +//! +//! Retention and reclaim match the MPMC broadcast backlog. `Inner` also holds the single parked +//! producer so a bounded `send` can check capacity and register its waker in one critical section. + +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 std::task::Waker; + +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); + +/// A receive plus the parked producer to wake if that receive freed capacity. +pub type RecvOutcome = (Received, Option); + +/// 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, parked receivers, and at most one parked producer, 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. +/// +/// The producer slot is the SPMC counterpart of MPMC's semaphore: `send(&mut self)` means at most +/// one waiting publisher, so occupancy is a single `Option`. Checking capacity and registering +/// that waker share this lock, which is what makes the wait path race-free without a recheck +/// protocol. +pub struct Inner { + pub log: Backlog, + pub waiters: WakerSet, + pub producer: Option, +} + +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(), + producer: None, + }); + (inner, key) + } +} + +/// Wakes every parked receiver so it can observe the channel's disconnected state. +/// +/// The sender's `Drop` calls this. +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, the prefix the receive released, and the +/// parked producer if that reclaim freed capacity. +/// +/// The caller owns what happens next: a bounded channel wakes the producer 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) => { + let producer = take_producer_on_reclaim(&mut inner, &received.1, false); + Ok((received, producer)) + } + 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. A reclaim that frees capacity takes the +/// producer waker in the same critical section. +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; + let producer = take_producer_on_reclaim(&mut inner, &received.1, false); + Poll::Ready(Ok((received, producer))) + } + 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 + } + } +} + +/// Takes the parked producer when this reclaim is one of the two sources that can free capacity. +/// +/// `buffer` shrinks only in `reclaim_consumed`, reachable from a receive that vacates the head +/// and from removing a subscription. Both must wake the producer. A new subscription starts at +/// the tail and never lowers `retained()`, so it never takes this path. +pub fn take_producer_on_reclaim( + inner: &mut Inner, + reclaimed: &Reclaimed, + drained_last: bool, +) -> Option { + if drained_last || !reclaimed.is_empty() { + inner.producer.take() + } else { + None + } +} + +/// Wakes the parked producer, if any, with the channel already unlocked. +pub fn wake_producer(producer: Option) { + if let Some(waker) = producer { + waker.wake(); + } +} + +/// 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/spmc/error.rs b/asyncband/src/broadcast/spmc/error.rs new file mode 100644 index 00000000..51fd9884 --- /dev/null +++ b/asyncband/src/broadcast/spmc/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::spmc::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 { + /// The sender has 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 the sender remains. + Empty, + /// The sender has 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/spmc/mod.rs b/asyncband/src/broadcast/spmc/mod.rs new file mode 100644 index 00000000..092b9488 --- /dev/null +++ b/asyncband/src/broadcast/spmc/mod.rs @@ -0,0 +1,74 @@ +// 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. + +//! Single-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 the +//! producer does when the slowest subscription stops reclaiming. [`bounded`] retains at most the +//! capacity it was built with and makes the producer wait for that subscription. [`unbounded`] +//! never waits to send and lets the retained backlog grow instead. +//! +//! The sender is not [`Clone`], and every publish method takes `&mut self`. That is the static +//! single-writer contract this topology adds over [`crate::broadcast::mpmc`]: there cannot be a +//! second producer, at compile time. `&Sender` can still be shared for [`subscribe`] and the +//! inspection methods. The family is offered for that exclusive-send API. Measured against +//! `async-broadcast` and `tokio::sync::broadcast`, tight bounded wait (capacity 1) is competitive; +//! it is not a general throughput upgrade over 1-producer `broadcast::mpmc` or over Tokio's +//! non-blocking ring. +//! +//! This is fan-out broadcast, not a competing queue: every accepted value is delivered to every +//! active subscription. A competitive `asyncband::spmc` queue would give each value to exactly one +//! receiver. +//! +//! # 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. +//! +//! # Subscribing while a bounded send is waiting +//! +//! [`BoundedSender::send`] holds `&mut self` for the duration of the wait, so the sender cannot +//! [`BoundedSender::subscribe`] until that send completes. Create the extra subscription from a +//! live receiver with [`BoundedReceiver::resubscribe`] instead; see the example on +//! [`BoundedReceiver::resubscribe`]. +//! +//! [`subscribe`]: BoundedSender::subscribe + +mod bounded; +mod common; +mod error; +mod unbounded; + +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/spmc/unbounded/mod.rs b/asyncband/src/broadcast/spmc/unbounded/mod.rs new file mode 100644 index 00000000..845bd2c0 --- /dev/null +++ b/asyncband/src/broadcast/spmc/unbounded/mod.rs @@ -0,0 +1,421 @@ +// 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. + +//! An unbounded fan-out channel with one sender and many receivers. +//! +//! A send publishes one value to every receiver that exists at that moment. Receivers advance +//! independently, and a receiver created later starts with the next value rather than replaying +//! earlier values. The sender is not [`Clone`]; [`UnboundedSender::send`] takes `&mut self` so a +//! second producer cannot exist at compile time. +//! +//! # Backlog and memory +//! +//! 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. Use [`bounded`] when +//! the producer should wait for the slowest receiver instead of growing the backlog. +//! +//! # Receivers +//! +//! [`UnboundedSender::subscribe`] and [`UnboundedReceiver::resubscribe`] add a receiver at the +//! current publication boundary. They do not copy another receiver's unread backlog. +//! +//! # Example +//! +//! ``` +//! use asyncband::broadcast::spmc::TryRecvError; +//! use asyncband::broadcast::spmc::unbounded; +//! +//! let (mut publisher, mut early) = unbounded(); +//! publisher.send("before subscription"); +//! +//! let mut late = publisher.subscribe(); +//! publisher.send("after subscription"); +//! +//! assert_eq!(early.try_recv(), Ok("before subscription")); +//! assert_eq!(early.try_recv(), Ok("after subscription")); +//! assert_eq!(late.try_recv(), Ok("after subscription")); +//! assert_eq!(late.try_recv(), Err(TryRecvError::Empty)); +//! ``` +//! +//! [`bounded`]: super::bounded + +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; + +use 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::WakerToken; + +#[cfg(test)] +mod tests; + +/// Creates an unbounded broadcast channel and its first receiver. +/// +/// The returned receiver is subscribed before any value can be published. Additional receivers can +/// be added with [`UnboundedSender::subscribe`]. +/// +/// # Examples +/// +/// ``` +/// use asyncband::broadcast::spmc::unbounded; +/// +/// let (mut publisher, mut receiver) = unbounded(); +/// publisher.send("ready"); +/// assert_eq!(receiver.try_recv(), Ok("ready")); +/// ``` +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { + let (inner, key) = Inner::with_first_subscription(Backlog::elastic()); + let shared = Arc::new(Shared { + inner, + senders: AtomicUsize::new(1), + }); + let sender = UnboundedSender { + shared: shared.clone(), + }; + let receiver = UnboundedReceiver { shared, key }; + (sender, receiver) +} + +struct Shared { + /// Buffer, receiver cursors, and parked receivers, all under a single lock. + inner: Mutex>, + /// `1` while the sender is alive, `0` after it is dropped. + senders: AtomicUsize, +} + +/// A publishing handle for an unbounded broadcast channel. +/// +/// This handle is not [`Clone`]. Once it is dropped, each receiver can drain the values already +/// published for it and then observes disconnection. +pub struct UnboundedSender { + shared: Arc>, +} + +impl fmt::Debug for UnboundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedSender").finish_non_exhaustive() + } +} + +impl Drop for UnboundedSender { + fn drop(&mut self) { + self.shared.senders.store(0, Ordering::Release); + common::disconnect(&self.shared.inner); + } +} + +impl UnboundedSender { + /// Publishes `msg` to every receiver currently subscribed. + /// + /// Sending has no backpressure. The channel retains the value until every eligible receiver + /// consumes it or is dropped. + /// + /// When no receivers exist, `msg` is discarded without entering the backlog. + /// + /// # Panics + /// + /// Panics if the channel has already published `u64::MAX` values. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc::unbounded; + /// + /// let (mut publisher, mut first) = unbounded(); + /// let mut second = publisher.subscribe(); + /// publisher.send("update"); + /// + /// assert_eq!(first.try_recv(), Ok("update")); + /// assert_eq!(second.try_recv(), Ok("update")); + /// ``` + pub fn send(&mut self, msg: T) { + let msg = Arc::new(msg); + + // Publishing and draining the wait set share one critical section, so a receiver can never + // observe an empty buffer and park after this message became visible. + let (unretained, wakers) = { + let mut inner = self.shared.inner.lock(); + 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. + /// + /// This is not an unread count for any particular receiver. A value remains included until the + /// last receiver eligible for it advances or is dropped. + /// + /// The result is an instantaneous observation and may become stale as other tasks send or + /// receive. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc::unbounded; + /// + /// let (mut publisher, mut fast) = unbounded(); + /// let mut slow = publisher.subscribe(); + /// publisher.send("update"); + /// assert_eq!(publisher.retained_message_count(), 1); + /// + /// assert_eq!(fast.try_recv(), Ok("update")); + /// assert_eq!(publisher.retained_message_count(), 1); + /// assert_eq!(slow.try_recv(), Ok("update")); + /// assert_eq!(publisher.retained_message_count(), 0); + /// ``` + pub fn retained_message_count(&self) -> usize { + self.shared.inner.lock().log.retained() + } + + /// Subscribes a new receiver for values published from this point forward. + /// + /// Values already in the backlog are not visible to the new receiver. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc::TryRecvError; + /// use asyncband::broadcast::spmc::unbounded; + /// + /// let (mut publisher, _) = unbounded(); + /// publisher.send("earlier"); + /// + /// let mut receiver = publisher.subscribe(); + /// assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); + /// publisher.send("later"); + /// assert_eq!(receiver.try_recv(), Ok("later")); + /// ``` + #[must_use = "the receiver is dropped immediately if it is not retained"] + pub fn subscribe(&self) -> UnboundedReceiver { + let key = self.shared.inner.lock().log.subscribe(); + UnboundedReceiver { + shared: self.shared.clone(), + key, + } + } +} + +/// An independent subscription to an unbounded broadcast channel. +/// +/// This receiver observes every value published after its subscription point and retains its own +/// position in the shared backlog. +pub struct UnboundedReceiver { + shared: Arc>, + key: SlotId, +} + +impl fmt::Debug for UnboundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() + } +} + +impl Drop for UnboundedReceiver { + fn drop(&mut self) { + let reclaimed = { + let mut inner = self.shared.inner.lock(); + inner.log.remove_receiver(self.key) + }; + drop(reclaimed); + } +} + +impl UnboundedReceiver { + /// Waits for this receiver's next value. + /// + /// Values already published for this receiver are returned before disconnection. + /// [`RecvError::Disconnected`] is returned only when no sender remains and this receiver's + /// backlog is empty. + /// + /// # Cancel safety + /// + /// Dropping a pending `recv` leaves this receiver's cursor unchanged. Its next call can still + /// return the same next value, so `recv` can be raced with other futures in a selection + /// construct. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc::RecvError; + /// use asyncband::broadcast::spmc::unbounded; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (mut publisher, mut receiver) = unbounded(); + /// publisher.send("final update"); + /// drop(publisher); + /// + /// assert_eq!(receiver.recv().await, Ok("final update")); + /// assert_eq!(receiver.recv().await, Err(RecvError::Disconnected)); + /// # } + /// ``` + pub async fn recv(&mut self) -> Result { + Recv { + receiver: self, + token: None, + } + .await + } + + /// Attempts to take this receiver's next value without waiting. + /// + /// [`TryRecvError::Empty`] means this receiver is currently caught up while a sender remains. + /// [`TryRecvError::Disconnected`] means no sender remains and this receiver has drained its + /// backlog. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc::TryRecvError; + /// use asyncband::broadcast::spmc::unbounded; + /// + /// let (mut publisher, mut receiver) = unbounded(); + /// assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); + /// + /// publisher.send("update"); + /// assert_eq!(receiver.try_recv(), Ok("update")); + /// drop(publisher); + /// assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected)); + /// ``` + pub fn try_recv(&mut self) -> Result { + let ((msg, reclaimed), _producer) = + common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; + Ok(common::take_msg(msg, reclaimed)) + } +} + +impl UnboundedReceiver { + /// 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. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc::TryRecvError; + /// use asyncband::broadcast::spmc::unbounded; + /// + /// let (mut publisher, mut original) = unbounded(); + /// publisher.send("pending for original"); + /// + /// let mut fresh = original.resubscribe(); + /// assert_eq!(fresh.try_recv(), Err(TryRecvError::Empty)); + /// publisher.send("visible to both"); + /// + /// assert_eq!(original.try_recv(), Ok("pending for original")); + /// assert_eq!(original.try_recv(), Ok("visible to both")); + /// assert_eq!(fresh.try_recv(), Ok("visible to both")); + /// ``` + #[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 this receiver's unread value count. + /// + /// Unlike [`UnboundedSender::retained_message_count`], this excludes values retained only for + /// other receivers. + /// + /// The result is an instantaneous observation and may become stale as other tasks publish + /// values. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::spmc::unbounded; + /// + /// let (mut publisher, mut receiver) = unbounded(); + /// publisher.send("first"); + /// publisher.send("second"); + /// assert_eq!(receiver.unread_message_count(), 2); + /// + /// assert_eq!(receiver.try_recv(), Ok("first")); + /// assert_eq!(receiver.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 UnboundedReceiver, + 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), _producer) = 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(outcome)) => outcome, + }; + + Poll::Ready(Ok(common::take_msg(msg, reclaimed))) + } +} diff --git a/asyncband/src/broadcast/spmc/unbounded/tests.rs b/asyncband/src/broadcast/spmc/unbounded/tests.rs new file mode 100644 index 00000000..86fc8e83 --- /dev/null +++ b/asyncband/src/broadcast/spmc/unbounded/tests.rs @@ -0,0 +1,70 @@ +// 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 super::*; +use crate::broadcast::spmc::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 (mut tx, _) = unbounded(); + tx.shared.inner.lock().log.set_tail(u64::MAX); + tx.send(()); +} + +#[test] +fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { + let (mut tx, mut rx) = unbounded(); + + let burst = MIN_RETAINED_CAPACITY * 16; + for i in 0..burst { + tx.send(i); + } + assert!(tx.shared.inner.lock().log.buffer_capacity() >= burst); + + for i in 0..burst { + assert_eq!(rx.try_recv(), Ok(i)); + } + + // 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().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().log.buffer_capacity() < burst); +} + +#[test] +fn repeated_bursts_keep_their_allocation() { + let (mut tx, mut rx) = unbounded(); + let burst = MIN_RETAINED_CAPACITY * 4; + + for _ in 0..4 { + for i in 0..burst { + tx.send(i); + } + for i in 0..burst { + assert_eq!(rx.try_recv(), Ok(i)); + } + } + + // Every cycle peaks at the same size, so the buffer must not rebuild its allocation each time. + assert!(tx.shared.inner.lock().log.buffer_capacity() >= burst); +} diff --git a/benchmarks/asyncband/broadcast/mod.rs b/benchmarks/asyncband/broadcast/mod.rs index b078f4ac..72736dd7 100644 --- a/benchmarks/asyncband/broadcast/mod.rs +++ b/benchmarks/asyncband/broadcast/mod.rs @@ -16,3 +16,4 @@ // under the License. mod mpmc; +mod spmc; diff --git a/benchmarks/asyncband/broadcast/spmc/bounded.rs b/benchmarks/asyncband/broadcast/spmc/bounded.rs new file mode 100644 index 00000000..fbb6d55c --- /dev/null +++ b/benchmarks/asyncband/broadcast/spmc/bounded.rs @@ -0,0 +1,160 @@ +// 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 and no parked producer left behind. 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::spmc; +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 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 (mut tx, rx) = spmc::bounded(CAPACITY); + drop(rx); + + bencher.bench_local(|| tx.try_send(black_box(1))); +} + +#[divan::bench] +fn try_send_and_try_recv(bencher: Bencher) { + let (mut tx, mut rx) = spmc::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 (mut tx, _rx) = spmc::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 (mut tx, rx) = spmc::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] +fn reclaim_wakes_the_blocked_producer(bencher: Bencher) { + let mut context = bench_context(); + + // Measures the whole backpressure cycle: park the single producer on a full channel, free one + // slot, and let it through. Each iteration ends with the producer parked and the same backlog, + // so the loop is stationary. + bencher + .with_inputs(|| { + let (mut tx, rx) = spmc::bounded(1); + tx.try_send(0).unwrap(); + (tx, rx) + }) + .bench_local_refs(|(tx, rx)| { + let mut send = Box::pin(tx.send(1)); + poll_pending(send.as_mut(), &mut context); + + black_box(rx.try_recv().unwrap()); + assert!(send.as_mut().poll(&mut context).is_ready()); + + // Drain the republished message so the next iteration starts from the same state. + black_box(rx.try_recv().unwrap()); + drop(send); + tx.try_send(0).unwrap(); + }); +} + +#[divan::bench] +fn cancel_blocked_send(bencher: Bencher) { + let mut context = bench_context(); + let (mut tx, _rx) = spmc::bounded(1); + tx.try_send(0).unwrap(); + + // Park the producer and immediately cancel it: measures registering and unlinking the 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 (mut tx, mut rx) = spmc::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_producer(bencher: Bencher, backlog: usize) { + bencher + .with_inputs(|| { + let (mut sender, mut fast) = spmc::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 send = Box::pin(async move { sender.send(backlog).await }); + poll_pending(send.as_mut(), &mut context); + (slow, fast, send) + }) + .bench_local_values(|(slow, fast, send)| { + // The fast subscription stays alive so this measures reclaim, not last-receiver exit. + // Preparing the backlog, parking the producer, and disposing of futures are outside + // timing. + drop(slow); + defer_input_drop((fast, send), ()) + }); +} diff --git a/benchmarks/asyncband/broadcast/spmc/mod.rs b/benchmarks/asyncband/broadcast/spmc/mod.rs new file mode 100644 index 00000000..e0ac8347 --- /dev/null +++ b/benchmarks/asyncband/broadcast/spmc/mod.rs @@ -0,0 +1,19 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod bounded; +mod unbounded; diff --git a/benchmarks/asyncband/broadcast/spmc/unbounded.rs b/benchmarks/asyncband/broadcast/spmc/unbounded.rs new file mode 100644 index 00000000..33aee514 --- /dev/null +++ b/benchmarks/asyncband/broadcast/spmc/unbounded.rs @@ -0,0 +1,191 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Every benchmark here must return the channel to a drained state on each iteration. Unlike the +// overflow policy, this channel has no capacity ceiling, so a timed loop that only sends would +// grow the retained backlog until the process runs out of memory. + +use std::fmt; +use std::pin::pin; + +use asyncband::broadcast::spmc; +use divan::Bencher; +use divan::black_box; + +use crate::support::bench_context; +use crate::support::poll_pending; +use crate::support::poll_pinned_ready; + +const RECEIVER_COUNTS: &[usize] = &[1, 8, 32]; + +/// A channel that peaked at `peak` receivers and currently has `live` of them. +/// +/// The two are measured separately because a dropped receiver leaves its slot behind: the reclaim +/// scan walks every slot the channel ever handed out, so a channel that shed receivers keeps +/// paying for the peak. Pairing each peak with a drained arena is what makes that visible. +#[derive(Clone, Copy)] +struct Fanout { + peak: usize, + live: usize, +} + +impl fmt::Display for Fanout { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "peak {} live {}", self.peak, self.live) + } +} + +const RECLAIM_FANOUTS: &[Fanout] = &[ + Fanout { peak: 1, live: 1 }, + Fanout { peak: 8, live: 8 }, + Fanout { peak: 8, live: 1 }, + Fanout { peak: 32, live: 32 }, + Fanout { peak: 32, live: 4 }, + Fanout { peak: 32, live: 1 }, + Fanout { + peak: 256, + live: 32, + }, + Fanout { peak: 256, live: 1 }, +]; + +#[divan::bench] +fn send_without_receivers(bencher: Bencher) { + let (mut sender, receiver) = spmc::unbounded::(); + drop(receiver); + bencher.bench_local(|| sender.send(black_box(1))); +} + +#[divan::bench] +fn try_recv_empty(bencher: Bencher) { + let (sender, mut receiver) = spmc::unbounded::(); + bencher.bench_local(|| black_box(receiver.try_recv())); + black_box(sender); +} + +// With the payload shared, each receive clones it and the second one reclaims the slot. +#[divan::bench] +fn send_and_try_recv_shared(bencher: Bencher) { + let (mut sender, mut first) = spmc::unbounded(); + let mut second = sender.subscribe(); + bencher.bench_local(|| { + sender.send(black_box(1usize)); + black_box(first.try_recv().unwrap()); + black_box(second.try_recv().unwrap()) + }); +} + +// The `usize` benchmarks above hide what a receive costs for a payload that owns memory: a clone +// there is an allocation, not a register move. +fn payload() -> String { + "x".repeat(64) +} + +#[divan::bench] +fn send_and_try_recv_owned(bencher: Bencher) { + let (mut sender, mut receiver) = spmc::unbounded(); + bencher.bench_local(|| { + sender.send(black_box(payload())); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn send_and_try_recv_owned_shared(bencher: Bencher) { + let (mut sender, mut first) = spmc::unbounded(); + let mut second = sender.subscribe(); + bencher.bench_local(|| { + sender.send(black_box(payload())); + black_box(first.try_recv().unwrap()); + black_box(second.try_recv().unwrap()) + }); +} + +// Measures the reclaim scan, which runs when the slowest cursor advances. Comparing a peak against +// the same peak drained down to fewer receivers shows what the slots left behind still cost. +#[divan::bench(args = RECLAIM_FANOUTS)] +fn drain_with_receivers(bencher: Bencher, fanout: Fanout) { + let (mut sender, receiver) = spmc::unbounded(); + drop(receiver); + let mut receivers = (0..fanout.peak) + .map(|_| sender.subscribe()) + .collect::>(); + // Dropping down to `live` leaves the arena holding a slot for every receiver that ever existed. + receivers.truncate(fanout.live); + + bencher.bench_local(|| { + sender.send(black_box(1usize)); + for receiver in &mut receivers { + black_box(receiver.try_recv().unwrap()); + } + }); +} + +#[divan::bench] +fn cancel_pending(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (sender, mut receiver) = spmc::unbounded::(); + { + let mut recv = pin!(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + } + black_box((sender, receiver)) + }); +} + +#[divan::bench] +fn deliver_to_waiter(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (mut sender, mut receiver) = spmc::unbounded(); + let mut recv = pin!(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + + sender.send(black_box(1usize)); + let value = poll_pinned_ready(recv.as_mut(), &mut context).unwrap(); + black_box(value) + }); +} + +#[divan::bench(args = RECEIVER_COUNTS)] +fn deliver_to_receiver_batch(bencher: Bencher, receiver_count: usize) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (mut sender, receiver) = spmc::unbounded(); + drop(receiver); + let mut receivers = (0..receiver_count) + .map(|_| sender.subscribe()) + .collect::>(); + let mut recvs = receivers + .iter_mut() + .map(|receiver| Box::pin(receiver.recv())) + .collect::>(); + for recv in &mut recvs { + poll_pending(recv.as_mut(), &mut context); + } + + sender.send(black_box(1usize)); + for mut recv in recvs { + let value = poll_pinned_ready(recv.as_mut(), &mut context).unwrap(); + black_box(value); + } + }); +} diff --git a/benchmarks/ecosystem/broadcast/mod.rs b/benchmarks/ecosystem/broadcast/mod.rs index b078f4ac..72736dd7 100644 --- a/benchmarks/ecosystem/broadcast/mod.rs +++ b/benchmarks/ecosystem/broadcast/mod.rs @@ -16,3 +16,4 @@ // under the License. mod mpmc; +mod spmc; diff --git a/benchmarks/ecosystem/broadcast/spmc/adapters.rs b/benchmarks/ecosystem/broadcast/spmc/adapters.rs new file mode 100644 index 00000000..e5405628 --- /dev/null +++ b/benchmarks/ecosystem/broadcast/spmc/adapters.rs @@ -0,0 +1,310 @@ +// 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::task::Context; + +use asyncband::blocking::FutureExt; + +use crate::support::poll_ready; + +pub struct Asyncband; +pub struct AsyncbandMpmc; +pub struct Tokio; +pub struct AsyncBroadcast; + +/// Single-producer broadcast: publish takes `&mut` so a non-`Clone` sender still fits. +pub trait BroadcastSpmc: Send + Sync + 'static { + type Sender: Send + 'static; + type Receiver: Send + 'static; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec); + fn send(sender: &mut Self::Sender, value: usize); + fn try_recv(receiver: &mut Self::Receiver) -> Option; + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize; +} + +impl BroadcastSpmc for Asyncband { + type Receiver = asyncband::broadcast::spmc::UnboundedReceiver; + type Sender = asyncband::broadcast::spmc::UnboundedSender; + + fn channel(_capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = asyncband::broadcast::spmc::unbounded(); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(sender.subscribe()); + } + (sender, receivers) + } + + fn send(sender: &mut Self::Sender, value: usize) { + sender.send(value); + } + + fn try_recv(receiver: &mut Self::Receiver) -> Option { + match receiver.try_recv() { + Ok(value) => Some(value), + Err(asyncband::broadcast::spmc::TryRecvError::Empty) => None, + Err(asyncband::broadcast::spmc::TryRecvError::Disconnected) => { + panic!("asyncband channel closed during benchmark") + } + } + } + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(receiver.recv(), context).unwrap() + } +} + +impl BroadcastSpmc for AsyncbandMpmc { + type Receiver = asyncband::broadcast::mpmc::UnboundedReceiver; + type Sender = asyncband::broadcast::mpmc::UnboundedSender; + + fn channel(_capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = asyncband::broadcast::mpmc::unbounded(); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(sender.subscribe()); + } + (sender, receivers) + } + + fn send(sender: &mut Self::Sender, value: usize) { + asyncband::broadcast::mpmc::UnboundedSender::send(sender, value); + } + + fn try_recv(receiver: &mut Self::Receiver) -> Option { + match receiver.try_recv() { + Ok(value) => Some(value), + Err(asyncband::broadcast::mpmc::TryRecvError::Empty) => None, + Err(asyncband::broadcast::mpmc::TryRecvError::Disconnected) => { + panic!("asyncband channel closed during benchmark") + } + } + } + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(receiver.recv(), context).unwrap() + } +} + +impl BroadcastSpmc for Tokio { + type Receiver = tokio::sync::broadcast::Receiver; + type Sender = tokio::sync::broadcast::Sender; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = tokio::sync::broadcast::channel(capacity); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(sender.subscribe()); + } + (sender, receivers) + } + + fn send(sender: &mut Self::Sender, value: usize) { + sender.send(value).unwrap(); + } + + fn try_recv(receiver: &mut Self::Receiver) -> Option { + match receiver.try_recv() { + Ok(value) => Some(value), + Err(tokio::sync::broadcast::error::TryRecvError::Empty) => None, + Err(error) => panic!("unexpected Tokio receive error: {error}"), + } + } + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(receiver.recv(), context).unwrap() + } +} + +impl BroadcastSpmc for AsyncBroadcast { + type Receiver = async_broadcast::Receiver; + type Sender = async_broadcast::Sender; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = async_broadcast::broadcast(capacity); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + let receiver = receivers[0].clone(); + receivers.push(receiver); + } + (sender, receivers) + } + + fn send(sender: &mut Self::Sender, value: usize) { + sender.try_broadcast(value).unwrap(); + } + + fn try_recv(receiver: &mut Self::Receiver) -> Option { + match receiver.try_recv() { + Ok(value) => Some(value), + Err(async_broadcast::TryRecvError::Empty) => None, + Err(error) => panic!("unexpected async-broadcast receive error: {error}"), + } + } + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + poll_ready(receiver.recv_direct(), context).unwrap() + } +} + +/// Lossless bounded broadcast with exclusive send. Tokio is omitted: it overwrites at capacity. +pub trait BoundedBroadcastSpmc: Send + Sync + 'static { + type Sender: Send + 'static; + type Receiver: Send + 'static; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec); + fn try_send(sender: &mut Self::Sender, value: usize); + fn send_async(sender: &mut Self::Sender, value: usize) -> impl Future + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; + + fn send_ready(sender: &mut Self::Sender, value: usize, context: &mut Context<'_>) { + poll_ready(Self::send_async(sender, value), context); + } + + fn send_blocking(sender: &mut 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 BoundedBroadcastSpmc for Asyncband { + type Receiver = asyncband::broadcast::spmc::BoundedReceiver; + type Sender = asyncband::broadcast::spmc::BoundedSender; + + fn channel(capacity: usize, receiver_count: usize) -> (Self::Sender, Vec) { + let (sender, receiver) = asyncband::broadcast::spmc::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: &mut Self::Sender, value: usize) { + sender.try_send(value).unwrap(); + } + + async fn send_async(sender: &mut 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::spmc::TryRecvError::Empty) => None, + Err(asyncband::broadcast::spmc::TryRecvError::Disconnected) => { + panic!("asyncband channel closed during benchmark") + } + } + } +} + +impl BoundedBroadcastSpmc for AsyncbandMpmc { + 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: &mut Self::Sender, value: usize) { + sender.try_send(value).unwrap(); + } + + async fn send_async(sender: &mut 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 BoundedBroadcastSpmc 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: &mut Self::Sender, value: usize) { + sender.try_broadcast(value).unwrap(); + } + + async fn send_async(sender: &mut 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/spmc/bounded.rs b/benchmarks/ecosystem/broadcast/spmc/bounded.rs new file mode 100644 index 00000000..7bf4290e --- /dev/null +++ b/benchmarks/ecosystem/broadcast/spmc/bounded.rs @@ -0,0 +1,87 @@ +// 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. + +// Lossless bounded path: producers wait at capacity. Tokio is absent because it overwrites. + +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use super::adapters::AsyncBroadcast; +use super::adapters::Asyncband; +use super::adapters::AsyncbandMpmc; +use super::adapters::BoundedBroadcastSpmc; +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; + +#[divan::bench(types = [Asyncband, AsyncbandMpmc, AsyncBroadcast], sample_size = 512)] +fn try_round_trip(bencher: Bencher) { + let (mut sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); + let mut receiver = receivers.pop().unwrap(); + + bencher.bench_local(|| { + C::try_send(&mut sender, black_box(usize::MAX)); + black_box(C::try_recv(&mut receiver).unwrap()) + }); +} + +#[divan::bench(types = [Asyncband, AsyncbandMpmc, AsyncBroadcast], sample_size = 512)] +fn ready_round_trip(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); + let mut receiver = receivers.pop().unwrap(); + + bencher.bench_local(|| { + C::send_ready(&mut sender, black_box(usize::MAX), &mut context); + black_box(C::recv_ready(&mut receiver, &mut context)) + }); +} + +#[divan::bench( + types = [Asyncband, AsyncbandMpmc, 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, AsyncbandMpmc, 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/spmc/mod.rs b/benchmarks/ecosystem/broadcast/spmc/mod.rs new file mode 100644 index 00000000..dd09282b --- /dev/null +++ b/benchmarks/ecosystem/broadcast/spmc/mod.rs @@ -0,0 +1,21 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod adapters; +mod bounded; +mod support; +mod unbounded; diff --git a/benchmarks/ecosystem/broadcast/spmc/support.rs b/benchmarks/ecosystem/broadcast/spmc/support.rs new file mode 100644 index 00000000..46adc0af --- /dev/null +++ b/benchmarks/ecosystem/broadcast/spmc/support.rs @@ -0,0 +1,232 @@ +// 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::fmt; +use std::sync::Arc; +use std::sync::Barrier; +use std::thread; +use std::thread::JoinHandle; + +use divan::black_box; +use tokio::runtime::Runtime; +use tokio::task::JoinSet; + +use super::adapters::BoundedBroadcastSpmc; +use super::adapters::BroadcastSpmc; + +pub const BATCH_MESSAGES: usize = 4096; +pub const RECEIVER_COUNTS: &[usize] = &[1, 2, 4, 8, 32]; +pub const ROUND_TRIP_CAPACITY: usize = 64; + +/// One-producer bounded workload. Capacity is the shared backlog; receivers are subscriptions. +#[derive(Clone, Copy)] +pub struct BoundedShape { + pub capacity: usize, + pub receivers: usize, +} + +impl BoundedShape { + const fn new(capacity: usize, receivers: usize) -> Self { + Self { + capacity, + receivers, + } + } +} + +impl fmt::Display for BoundedShape { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "cap {} 1 producer {} receivers", + self.capacity, self.receivers + ) + } +} + +/// Primary 1-producer shapes from the SPMC comparison matrix. +pub const BOUNDED_SHAPES: &[BoundedShape] = &[ + BoundedShape::new(1, 1), + BoundedShape::new(1, 8), + BoundedShape::new(1, 32), + BoundedShape::new(64, 1), + BoundedShape::new(64, 8), + BoundedShape::new(64, 32), +]; + +fn recv(receiver: &mut C::Receiver) -> usize { + C::try_recv(receiver).expect("the published benchmark batch must be ready") +} + +/// One producer publishes the batch, then every subscription drains it. +pub struct Fanout { + sender: C::Sender, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl Fanout { + pub fn new(receiver_count: usize) -> Self { + let (sender, receivers) = C::channel(BATCH_MESSAGES, receiver_count); + let start = Arc::new(Barrier::new(receiver_count + 1)); + let done = Arc::new(Barrier::new(receiver_count + 1)); + let workers = receivers + .into_iter() + .map(|mut receiver| { + let start = start.clone(); + let done = done.clone(); + thread::spawn(move || { + start.wait(); + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(recv::(&mut receiver)); + } + black_box(checksum); + done.wait(); + }) + }) + .collect(); + + Self { + sender, + start, + done, + workers, + } + } + + pub fn run(&mut self) { + for value in 0..BATCH_MESSAGES { + C::send(&mut self.sender, black_box(value)); + } + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for Fanout { + fn drop(&mut self) { + let panicking = thread::panicking(); + for worker in self.workers.drain(..) { + let result = worker.join(); + if !panicking { + result.expect("benchmark receiver panicked"); + } + } + } +} + +/// One producer and `receivers` subscriptions on native threads. All wait, then transfer. +pub struct BoundedConcurrent { + start: Arc, + workers: Vec>, +} + +impl BoundedConcurrent { + pub fn new(shape: BoundedShape) -> Self { + let BoundedShape { + capacity, + receivers, + } = shape; + let (mut sender, receivers) = C::channel(capacity, receivers); + let start = Arc::new(Barrier::new(receivers.len() + 2)); + let mut workers = Vec::with_capacity(receivers.len() + 1); + + 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); + })); + } + let producer_start = start.clone(); + workers.push(thread::spawn(move || { + producer_start.wait(); + for value in 0..BATCH_MESSAGES { + C::send_blocking(&mut 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 1-producer bounded workload on async tasks. +pub struct BoundedTasks { + start: Arc, + tasks: JoinSet<()>, +} + +impl BoundedTasks { + pub fn new(runtime: &Runtime, shape: BoundedShape) -> Self { + let BoundedShape { + capacity, + receivers, + } = shape; + let (mut sender, receivers) = C::channel(capacity, receivers); + let start = Arc::new(tokio::sync::Barrier::new(receivers.len() + 2)); + 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(), + ); + } + let start_producer = start.clone(); + tasks.spawn_on( + async move { + start_producer.wait().await; + for value in 0..BATCH_MESSAGES { + C::send_async(&mut 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/benchmarks/ecosystem/broadcast/spmc/unbounded.rs b/benchmarks/ecosystem/broadcast/spmc/unbounded.rs new file mode 100644 index 00000000..a7aafa01 --- /dev/null +++ b/benchmarks/ecosystem/broadcast/spmc/unbounded.rs @@ -0,0 +1,70 @@ +// 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. + +// One producer, lossless non-blocking path. Tokio and async-broadcast get room for the whole +// batch so this is not measuring overwrite or backpressure. + +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use super::adapters::AsyncBroadcast; +use super::adapters::Asyncband; +use super::adapters::AsyncbandMpmc; +use super::adapters::BroadcastSpmc; +use super::adapters::Tokio; +use super::support::BATCH_MESSAGES; +use super::support::Fanout; +use super::support::RECEIVER_COUNTS; +use super::support::ROUND_TRIP_CAPACITY; +use crate::support::bench_context; + +#[divan::bench(types = [Asyncband, AsyncbandMpmc, Tokio, AsyncBroadcast])] +fn try_round_trip(bencher: Bencher) { + let (mut sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); + let mut receiver = receivers.pop().unwrap(); + + bencher.bench_local(|| { + C::send(&mut sender, black_box(usize::MAX)); + black_box(C::try_recv(&mut receiver).unwrap()) + }); +} + +#[divan::bench(types = [Asyncband, AsyncbandMpmc, Tokio, AsyncBroadcast])] +fn ready_round_trip(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, mut receivers) = C::channel(ROUND_TRIP_CAPACITY, 1); + let mut receiver = receivers.pop().unwrap(); + + bencher.bench_local(|| { + C::send(&mut sender, black_box(usize::MAX)); + black_box(C::recv_ready(&mut receiver, &mut context)) + }); +} + +#[divan::bench( + types = [Asyncband, AsyncbandMpmc, Tokio, AsyncBroadcast], + args = RECEIVER_COUNTS, + sample_count = 20, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn fanout(bencher: Bencher, receiver_count: usize) { + bencher + .with_inputs(|| Fanout::::new(receiver_count)) + .bench_local_refs(Fanout::run); +} diff --git a/tests-integration/tests/broadcast_spmc_bounded_test.rs b/tests-integration/tests/broadcast_spmc_bounded_test.rs new file mode 100644 index 00000000..b86cee98 --- /dev/null +++ b/tests-integration/tests/broadcast_spmc_bounded_test.rs @@ -0,0 +1,732 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// The following MPMC cases are omitted on purpose: a non-cloneable `&mut self` sender makes them +// unrepresentable, not untested. +// +// * `bounded_concurrent_producers_commit_one_order_seen_by_every_receiver` — no concurrent +// producers; committed order is program order, covered by `publish_order_is_program_order`. +// * `wakes_blocked_senders_as_capacity_frees` — only one producer can wait; the singular case is +// `receive_that_vacates_the_head_wakes_the_blocked_sender` and +// `parked_recv_that_reclaims_wakes_the_blocked_sender`. +// * `dropping_the_last_receiver_wakes_every_blocked_sender` — singular case is +// `dropping_the_last_receiver_wakes_the_blocked_sender`. +// * `cancelled_notified_sender_passes_capacity_to_the_next_sender` — there is no next producer. +// * `a_large_reclaim_leaves_no_permit_slack` — SPMC has no semaphore and no permit slack. +// * `dropping_the_last_receiver_never_strands_a_racing_producer` — the multi-producer race is gone; +// the remaining drop-vs-wait race is `dropping_the_last_receiver_wakes_the_blocked_sender` +// together with `sends_never_block_once_all_receivers_are_gone`. + +use std::sync::Arc; +use std::sync::Mutex; +use std::task::Context; +use std::task::Wake; +use std::task::Waker; + +use asyncband::blocking::FutureExt; +use asyncband::broadcast::spmc::*; +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. +/// +/// The sender is not `Clone`, so the probe is a shared receiver handle. +struct Reentrant { + value: u64, + probe: Option>>>, +} + +impl Clone for Reentrant { + fn clone(&self) -> Self { + Self { + value: self.value, + probe: self.probe.clone(), + } + } +} + +impl Drop for Reentrant { + fn drop(&mut self) { + if let Some(probe) = &self.probe { + // `try_lock`: draining the probe itself may drop a payload while this mutex is held. + // Deadlocks if the channel still holds its lock while dropping reclaimed messages. + if let Ok(probe) = probe.try_lock() { + let _ = probe.unread_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 +// --------------------------------------------------------------------------------------------- + +#[test] +fn delivers_every_message_to_every_receiver() { + let (mut tx, mut rx1) = bounded(4); + let mut rx2 = tx.subscribe(); + + tx.try_send(10).unwrap(); + tx.try_send(20).unwrap(); + + assert_eq!(rx1.try_recv(), Ok(10)); + assert_eq!(rx1.try_recv(), Ok(20)); + assert_eq!(rx2.try_recv(), Ok(10)); + assert_eq!(rx2.try_recv(), Ok(20)); +} + +#[test] +fn slow_receiver_keeps_every_message_under_backpressure() { + let (mut 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 subscribe_starts_at_the_committed_tail() { + let (mut 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 resubscribe_keeps_the_original_receivers_backlog() { + let (mut 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 unread_message_count_tracks_each_receiver() { + let (mut 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); +} + +#[test] +fn publish_order_is_program_order() { + let (mut tx, mut rx1) = bounded(8); + let mut rx2 = tx.subscribe(); + + for value in 0..8 { + tx.try_send(value).unwrap(); + } + + for value in 0..8 { + assert_eq!(rx1.try_recv(), Ok(value)); + assert_eq!(rx2.try_recv(), Ok(value)); + } +} + +// --------------------------------------------------------------------------------------------- +// Strict capacity +// --------------------------------------------------------------------------------------------- + +#[test] +fn try_send_rejects_at_capacity_and_returns_the_value() { + let (mut 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 (mut 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 (mut 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 (mut 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_the_blocked_sender() { + let (mut 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_the_blocked_sender() { + let (mut 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 dropping_a_lagging_receiver_wakes_the_blocked_sender() { + let (mut 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!(rx1.unread_message_count(), 0); + assert!(poll_once(send.as_mut()).is_ready()); +} + +#[test] +fn dropping_the_last_receiver_wakes_the_blocked_sender() { + let (mut tx, rx) = bounded(2); + tx.try_send(0).unwrap(); + tx.try_send(1).unwrap(); + + let tracker = Arc::new(WakeCounter::default()); + let waker = Waker::from(tracker.clone()); + let mut send = Box::pin(tx.send(10)); + assert!( + send.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + + drop(rx); + + assert!( + tracker.count() > 0, + "blocked sender was never woken after the last receiver was dropped" + ); + assert!(poll_once(send.as_mut()).is_ready()); +} + +#[test] +fn sends_never_block_once_all_receivers_are_gone() { + let (mut 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()); + drop(send); + assert_eq!(tx.retained_message_count(), 0); +} + +#[test] +fn subscribing_while_the_producer_is_blocked_does_not_release_capacity() { + let (mut 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. The sender is + // exclusively borrowed by `send`, so the extra subscription comes from `resubscribe`. + let _late = rx.resubscribe(); + assert_eq!(rx.unread_message_count(), 1); + assert!(poll_once(send.as_mut()).is_pending()); +} + +// --------------------------------------------------------------------------------------------- +// Cancellation +// --------------------------------------------------------------------------------------------- + +#[test] +fn cancelled_send_publishes_nothing() { + let (mut 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_recv_releases_its_waker() { + let (mut 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); +} + +#[test] +fn dropping_a_woken_recv_keeps_another_receivers_waiter() { + let (mut 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); +} + +// --------------------------------------------------------------------------------------------- +// Disconnection +// --------------------------------------------------------------------------------------------- + +#[test] +fn recv_drains_buffered_messages_before_reporting_disconnection() { + let (mut tx, mut rx) = bounded(4); + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + drop(tx); + + assert_eq!(FutureExt::block_on(rx.recv()), Ok(1)); + assert_eq!(FutureExt::block_on(rx.recv()), Ok(2)); + assert_eq!(FutureExt::block_on(rx.recv()), Err(RecvError::Disconnected)); +} + +#[test] +fn recv_reports_disconnection_without_any_message() { + let (tx, mut rx) = bounded::(4); + drop(tx); + assert_eq!(FutureExt::block_on(rx.recv()), Err(RecvError::Disconnected)); +} + +#[test] +fn parked_recv_wakes_when_the_sender_drops() { + let (tx, mut rx) = bounded::(4); + + 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(), 1); + assert_eq!( + poll_once(recv.as_mut()), + std::task::Poll::Ready(Err(RecvError::Disconnected)) + ); +} + +// --------------------------------------------------------------------------------------------- +// Panic safety +// --------------------------------------------------------------------------------------------- + +#[test] +fn panicking_wake_does_not_strand_the_producer_after_a_large_reclaim() { + let (mut tx, mut fast) = bounded(8); + let slow = tx.subscribe(); + for value in 0..8 { + tx.try_send(value).unwrap(); + fast.try_recv().unwrap(); + } + + let tracker = Arc::new(WakeCounter::default()); + let waker = { + let tracker = tracker.clone(); + waker_on_wake(move || { + tracker.wake(); + panic!("producer wake panics"); + }) + }; + let mut send = Box::pin(tx.send(8)); + 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!(tracker.count(), 1); + assert!( + poll_once(send.as_mut()).is_ready(), + "a panicking producer wake must not strand the send" + ); + assert_eq!(fast.try_recv(), Ok(8)); +} + +#[test] +fn panicking_receiver_wake_still_wakes_the_rest() { + let (mut tx, mut rx1) = bounded(4); + let mut rx2 = tx.subscribe(); + let mut rx3 = tx.subscribe(); + + let trackers = [ + Arc::new(WakeCounter::default()), + Arc::new(WakeCounter::default()), + Arc::new(WakeCounter::default()), + ]; + let wakers = trackers + .iter() + .enumerate() + .map(|(index, tracker)| { + let tracker = tracker.clone(); + waker_on_wake(move || { + tracker.wake(); + assert_ne!(index, 0, "first receiver wake panics"); + }) + }) + .collect::>(); + + let mut recvs = [ + Box::pin(rx1.recv()), + Box::pin(rx2.recv()), + Box::pin(rx3.recv()), + ]; + for (recv, waker) in recvs.iter_mut().zip(&wakers) { + assert!( + recv.as_mut() + .poll(&mut Context::from_waker(waker)) + .is_pending() + ); + } + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + tx.try_send(1).unwrap(); + })); + assert!(result.is_err()); + for tracker in &trackers { + assert_eq!(tracker.count(), 1); + } +} + +#[test] +fn panicking_clone_leaves_the_channel_consistent() { + let (mut 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 (mut 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 message_destructors_run_outside_the_channel_lock() { + assert_completes_without_deadlock(|| { + let (mut tx, mut rx1) = bounded(8); + let rx2 = tx.subscribe(); + let probe = Arc::new(Mutex::new(tx.subscribe())); + + for value in 0..4 { + tx.try_send(Reentrant { + value, + probe: Some(probe.clone()), + }) + .unwrap(); + let _ = probe.lock().unwrap().try_recv(); + } + + // 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); + drop(probe); + + // With no receiver, both send paths discard the payload immediately. + tx.try_send(Reentrant { + value: 4, + probe: None, + }) + .unwrap(); + FutureExt::block_on(tx.send(Reentrant { + value: 5, + probe: None, + })); + }); +} + +#[test] +fn cancelling_a_blocked_send_drops_its_payload_outside_the_channel_lock() { + assert_completes_without_deadlock(|| { + let (mut tx, mut rx) = bounded(1); + tx.try_send(Reentrant { + value: 0, + probe: None, + }) + .unwrap(); + // Subscribe at the tail so the probe does not retain the buffered message. + let probe = Arc::new(Mutex::new(rx.resubscribe())); + + let mut send = Box::pin(tx.send(Reentrant { + value: 1, + probe: Some(probe.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, + probe: None, + }) + .unwrap(); + assert_eq!(rx.try_recv().unwrap().value, 2); + drop(probe); + }); +} diff --git a/tests-integration/tests/broadcast_spmc_unbounded_test.rs b/tests-integration/tests/broadcast_spmc_unbounded_test.rs new file mode 100644 index 00000000..c2825eed --- /dev/null +++ b/tests-integration/tests/broadcast_spmc_unbounded_test.rs @@ -0,0 +1,499 @@ +// 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. + +// The following MPMC cases are omitted on purpose: a non-cloneable `&mut self` sender makes them +// unrepresentable, not untested. +// +// * `concurrent_senders_deliver_every_message_to_every_receiver` — there is no second producer; +// committed order is program order, covered by `publish_order_is_program_order`. + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use asyncband::broadcast::spmc::*; +use tests_integration::WakeCounter; +use tests_integration::assert_completes_without_deadlock; + +/// A payload whose destructor re-enters the channel it was sent through. +/// +/// The sender is not `Clone`, so the probe is a shared receiver handle. `unread_message_count` +/// takes the same channel lock the destructor must not already hold. +struct Reentrant { + value: u64, + probe: Option>>>, +} + +impl Clone for Reentrant { + fn clone(&self) -> Self { + Self { + value: self.value, + probe: self.probe.clone(), + } + } +} + +impl Drop for Reentrant { + fn drop(&mut self) { + if let Some(probe) = &self.probe { + // `try_lock`: draining the probe itself may drop a payload while this mutex is held. + // Deadlocks if the channel still holds its lock while dropping reclaimed messages. + if let Ok(probe) = probe.try_lock() { + let _ = probe.unread_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, + } + } +} + +struct Rng(u64); + +impl Rng { + fn below(&mut self, n: u64) -> u64 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + x % n + } +} + +#[test] +fn try_recv_reports_empty_then_value_then_disconnected() { + let (mut tx, mut rx) = unbounded(); + + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + tx.send(10); + assert_eq!(rx.try_recv(), Ok(10)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + drop(tx); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[tokio::test] +async fn slow_receiver_keeps_every_message() { + let (mut tx, mut rx1) = unbounded(); + let mut rx2 = tx.subscribe(); + + for i in 0..1024 { + tx.send(i); + } + + // The fast receiver draining fully must not reclaim anything the slow one still needs. + for i in 0..1024 { + assert_eq!(rx1.recv().await, Ok(i)); + } + assert_eq!(tx.retained_message_count(), 1024); + + for i in 0..1024 { + assert_eq!(rx2.recv().await, Ok(i)); + } + assert_eq!(tx.retained_message_count(), 0); +} + +#[tokio::test] +async fn retained_message_count_tracks_the_slowest_receiver() { + let (mut tx, mut rx1) = unbounded(); + let mut rx2 = tx.subscribe(); + + tx.send(1); + tx.send(2); + assert_eq!(tx.retained_message_count(), 2); + + // Reclaiming waits for the slowest receiver, message by message. + assert_eq!(rx1.recv().await, Ok(1)); + assert_eq!(tx.retained_message_count(), 2); + assert_eq!(rx2.recv().await, Ok(1)); + assert_eq!(tx.retained_message_count(), 1); + + assert_eq!(rx1.recv().await, Ok(2)); + assert_eq!(tx.retained_message_count(), 1); + assert_eq!(rx2.recv().await, Ok(2)); + assert_eq!(tx.retained_message_count(), 0); +} + +#[tokio::test] +async fn dropping_a_lagging_receiver_releases_its_backlog() { + let (mut tx, mut rx1) = unbounded(); + let rx2 = tx.subscribe(); + + for i in 0..128 { + tx.send(i); + } + for i in 0..128 { + assert_eq!(rx1.recv().await, Ok(i)); + } + assert_eq!(tx.retained_message_count(), 128); + + drop(rx2); + assert_eq!(tx.retained_message_count(), 0); +} + +#[tokio::test] +async fn resubscribe_keeps_the_original_receivers_backlog() { + let (mut tx, mut rx) = unbounded(); + + tx.send(1); + tx.send(2); + + let mut rx2 = rx.resubscribe(); + assert_eq!(tx.retained_message_count(), 2); + + tx.send(3); + + assert_eq!(rx2.recv().await, Ok(3)); + assert_eq!(tx.retained_message_count(), 3); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Ok(3)); + assert_eq!(tx.retained_message_count(), 0); +} + +#[tokio::test] +async fn send_without_receivers_does_not_buffer() { + let (mut tx, rx) = unbounded(); + drop(rx); + + tx.send(1); + tx.send(2); + assert_eq!(tx.retained_message_count(), 0); + + let mut rx = tx.subscribe(); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + tx.send(3); + assert_eq!(rx.recv().await, Ok(3)); +} + +#[test] +fn unread_message_count_tracks_each_receiver() { + let (mut tx, mut rx1) = unbounded(); + assert_eq!(rx1.unread_message_count(), 0); + + tx.send(1); + tx.send(2); + assert_eq!(rx1.unread_message_count(), 2); + + let mut rx2 = tx.subscribe(); + assert_eq!(rx2.unread_message_count(), 0); + + tx.send(3); + assert_eq!(rx1.unread_message_count(), 3); + assert_eq!(rx2.unread_message_count(), 1); + + assert_eq!(rx2.try_recv(), Ok(3)); + assert_eq!(rx2.unread_message_count(), 0); + drop(rx2); + + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.unread_message_count(), 2); +} + +#[test] +fn sole_receiver_takes_messages_without_cloning() { + static CLONES: AtomicUsize = AtomicUsize::new(0); + + struct CountClone(u32); + + impl Clone for CountClone { + fn clone(&self) -> Self { + CLONES.fetch_add(1, Ordering::Relaxed); + Self(self.0) + } + } + + let (mut tx, mut rx) = unbounded(); + for i in 0..8 { + tx.send(CountClone(i)); + assert_eq!(rx.try_recv().unwrap().0, i); + } + assert_eq!(CLONES.load(Ordering::Relaxed), 0); + + // A second receiver means the payload is shared, so it has to be cloned again. + let mut second = tx.subscribe(); + tx.send(CountClone(8)); + assert_eq!(rx.try_recv().unwrap().0, 8); + assert_eq!(second.try_recv().unwrap().0, 8); + assert_eq!(CLONES.load(Ordering::Relaxed), 1); +} + +#[test] +fn panicking_clone_leaves_the_channel_consistent() { + let (mut tx, mut rx1) = unbounded(); + let mut rx2 = tx.subscribe(); + + tx.send(PanicOnClone { + value: 1, + panic: true, + }); + tx.send(PanicOnClone { + value: 2, + panic: false, + }); + + // Two receivers share the payload, so this receive has to clone it. + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + rx1.try_recv().map(|msg| msg.value) + })); + assert!(result.is_err()); + + // The failed receive still consumed the message for `rx1`, and left the channel usable for + // both receivers. + assert_eq!(rx1.try_recv().unwrap().value, 2); + assert_eq!(rx2.try_recv().unwrap().value, 1); + assert_eq!(rx2.try_recv().unwrap().value, 2); + assert_eq!(tx.retained_message_count(), 0); + assert_eq!(rx1.try_recv().unwrap_err(), TryRecvError::Empty); +} + +#[test] +fn message_destructors_run_outside_the_channel_lock() { + assert_completes_without_deadlock(|| { + let (mut tx, mut rx1) = unbounded(); + let rx2 = tx.subscribe(); + let probe = Arc::new(Mutex::new(tx.subscribe())); + + for value in 0..4 { + tx.send(Reentrant { + value, + probe: Some(probe.clone()), + }); + // Keep the probe at the tail so it does not retain the messages under test. + let _ = probe.lock().unwrap().try_recv(); + } + + // Reclaim through a receive, and then through a receiver drop. + assert_eq!(rx1.try_recv().unwrap().value, 0); + drop(rx2); + assert_eq!(rx1.try_recv().unwrap().value, 1); + drop(rx1); + drop(probe); + }); +} + +#[test] +fn send_wakes_a_parked_receiver_exactly_once() { + let (mut tx, mut rx) = unbounded(); + let tracker = Arc::new(WakeCounter::default()); + let waker = Waker::from(tracker.clone()); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + + tx.send(42); + + assert_eq!(tracker.count(), 1); + assert_eq!(recv.as_mut().poll(&mut context), Poll::Ready(Ok(42))); +} + +#[test] +fn cancelled_recv_releases_its_waker() { + let (mut tx, mut rx) = unbounded::<()>(); + let tracker = Arc::new(WakeCounter::default()); + let waker = Waker::from(tracker.clone()); + let baseline = Arc::strong_count(&tracker); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + + drop(recv); + assert_eq!(Arc::strong_count(&tracker), baseline); + + tx.send(()); + assert_eq!(tracker.count(), 0); + assert_eq!(rx.try_recv(), Ok(())); +} + +#[test] +fn dropping_a_woken_recv_keeps_another_receivers_waiter() { + let (mut tx, mut rx1) = unbounded::(); + 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.send(1); + 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.send(2); + + assert_eq!(second.count(), 1); +} + +#[test] +fn parked_recv_wakes_when_the_sender_drops() { + let (tx, mut rx) = unbounded::<()>(); + let tracker = Arc::new(WakeCounter::default()); + let waker = Waker::from(tracker.clone()); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + + drop(tx); + assert_eq!(tracker.count(), 1); + + drop(recv); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[test] +fn parked_recv_prefers_buffered_messages_over_disconnection() { + let (mut tx, mut rx) = unbounded(); + let tracker = Arc::new(WakeCounter::default()); + let waker = Waker::from(tracker); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + + tx.send(7); + drop(tx); + + assert_eq!(recv.as_mut().poll(&mut context), Poll::Ready(Ok(7))); + drop(recv); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[tokio::test] +async fn recv_drains_buffered_messages_before_reporting_disconnection() { + let (mut tx, mut rx) = unbounded(); + + tx.send(1); + tx.send(2); + drop(tx); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); +} + +#[tokio::test] +async fn recv_reports_disconnection_without_any_message() { + let (tx, mut rx) = unbounded::<()>(); + drop(tx); + assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); +} + +#[test] +fn publish_order_is_program_order() { + let (mut tx, mut rx1) = unbounded(); + let mut rx2 = tx.subscribe(); + + for value in 0..64 { + tx.send(value); + } + + for value in 0..64 { + assert_eq!(rx1.try_recv(), Ok(value)); + assert_eq!(rx2.try_recv(), Ok(value)); + } +} + +#[test] +fn randomized_operations_track_the_reference_model() { + for seed in 1..32u64 { + let mut rng = Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1); + let (mut tx, rx) = unbounded::(); + let mut tail = 0u64; + let mut model = vec![(rx, 0u64)]; + + for _ in 0..512 { + match rng.below(100) { + 0..=44 => { + tx.send(tail); + tail += 1; + } + 45..=79 if !model.is_empty() => { + let index = rng.below(model.len() as u64) as usize; + let (receiver, cursor) = &mut model[index]; + if *cursor < tail { + assert_eq!(receiver.try_recv(), Ok(*cursor), "seed {seed}"); + *cursor += 1; + } else { + assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty), "seed {seed}"); + } + } + 80..=89 => model.push((tx.subscribe(), tail)), + _ if !model.is_empty() => { + let index = rng.below(model.len() as u64) as usize; + model.swap_remove(index); + } + _ => {} + } + + let retained = model + .iter() + .map(|(_, cursor)| *cursor) + .min() + .map_or(0, |slowest| tail - slowest); + assert_eq!( + tx.retained_message_count(), + retained as usize, + "seed {seed}" + ); + for (receiver, cursor) in &model { + assert_eq!( + receiver.unread_message_count(), + (tail - cursor) as usize, + "seed {seed}" + ); + } + } + } +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 26205e34..2c50fc2b 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -111,6 +111,16 @@ fn public_types_are_send_and_sync() { assert_send_and_sync::(); assert_send_and_sync::(); assert_send_and_sync::>(); + // SPMC senders stay `Send + Sync` so `&Sender` can subscribe and inspect. They are not + // `Clone`; exclusive `&mut self` publish is the type-level single-writer contract (no + // trybuild). + 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::(); @@ -203,6 +213,13 @@ fn public_types_are_unpin() { assert_unpin::(); assert_unpin::(); assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::(); + assert_unpin::(); + assert_unpin::>(); assert_unpin::>(); assert_unpin::>(); assert_unpin::>(); diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 51fb9866..fcf77a3c 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -125,6 +125,10 @@ impl CommandMiri { "tests-integration", &["--test", "phaser_test"], )); + run_command(make_miri_cmd( + "tests-integration", + &["--test", "broadcast_spmc_bounded_test"], + )); } } From 1c35ef02695e8a7b264de465c8dbc29855cbc827 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Mon, 14 Sep 2026 16:48:37 +0800 Subject: [PATCH 2/9] perf(broadcast): drain SPMC slots without Arc or the publication lock Receivers clone T from a single-writer log without taking the waiter mutex. The last remaining reader takes the payload. --- CHANGELOG.md | 2 +- asyncband/src/broadcast/mod.rs | 5 +- asyncband/src/broadcast/spmc/bounded/mod.rs | 199 ++-- asyncband/src/broadcast/spmc/bounded/tests.rs | 10 +- asyncband/src/broadcast/spmc/common.rs | 865 ++++++++++-------- asyncband/src/broadcast/spmc/mod.rs | 6 +- asyncband/src/broadcast/spmc/unbounded/mod.rs | 93 +- .../src/broadcast/spmc/unbounded/tests.rs | 38 +- .../tests/broadcast_spmc_bounded_test.rs | 38 +- .../tests/broadcast_spmc_unbounded_test.rs | 41 +- 10 files changed, 696 insertions(+), 601 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e2834aed..4e3a85b7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,7 +27,7 @@ All notable changes to this project will be documented in this file. * Add `event::AutoResetEvent`, a reusable signal that releases one waiter, retains at most one unassigned signal that can be cleared with `reset`, and transfers assigned signals when waits are cancelled. * Add `ManualResetEvent::try_wait` to check readiness without registering a waiter or consuming the set state. -* Add `broadcast::spmc`, a lossless single-producer broadcast family with a non-cloneable sender whose publish methods require exclusive access; bounded retains at most the requested capacity and makes the producer wait for the slowest active subscription, while unbounded never waits and lets the retained backlog grow. +* Add `broadcast::spmc`, a lossless single-producer broadcast family with a non-cloneable sender whose publish methods require exclusive access; receivers drain already-published slots without taking the publication lock, bounded retains at most the requested capacity and makes the producer wait for the slowest active subscription, and unbounded never waits and lets the retained backlog grow. * 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`. diff --git a/asyncband/src/broadcast/mod.rs b/asyncband/src/broadcast/mod.rs index 0d9f5fab..d49f3e9b 100644 --- a/asyncband/src/broadcast/mod.rs +++ b/asyncband/src/broadcast/mod.rs @@ -19,8 +19,9 @@ //! //! [`mpmc`] supports any number of concurrent producers. [`spmc`] is the single-producer //! specialization: its sender is not [`Clone`] and publish methods require exclusive access -//! (`&mut self`). Choose `spmc` when the program has one publisher; choose `mpmc` when it does -//! not. Choose `spmc` for the exclusive-send API; do not assume it is faster on every path. +//! (`&mut self`). Receivers drain published slots without taking the publication lock, so a single +//! producer can fan out without serializing every subscription on that lock. Choose `spmc` when +//! the program has one publisher; choose `mpmc` when it does not. //! //! Both topologies are fan-out broadcast: every accepted value is delivered to every active //! subscription. That is a different delivery family from a competing crate-root `spmc` queue, diff --git a/asyncband/src/broadcast/spmc/bounded/mod.rs b/asyncband/src/broadcast/spmc/bounded/mod.rs index 8ec3ccc4..f013452f 100644 --- a/asyncband/src/broadcast/spmc/bounded/mod.rs +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -106,19 +106,17 @@ use std::future::Future; use std::future::poll_fn; 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::common::BoundedBuffer; +use super::common::Shared; +use super::common::SlotStore; 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::wake_all; use crate::internal::wakerset::WakerToken; @@ -152,35 +150,20 @@ pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver< "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, - }); + let shared = Arc::new(Shared::new(BoundedBuffer::new(capacity))); let sender = BoundedSender { shared: shared.clone(), }; - let receiver = BoundedReceiver { shared, key }; + let receiver = BoundedReceiver { shared, cursor: 0 }; (sender, receiver) } -struct Shared { - /// Buffer, receiver cursors, parked receivers, and the single parked producer, all under one - /// lock. - inner: Mutex>, - /// `1` while the sender is alive, `0` after it is dropped. - senders: AtomicUsize, - /// The logical limit on the retained backlog. - capacity: usize, -} - /// The sending side of a bounded broadcast channel. /// /// This handle is not [`Clone`]. Dropping it disconnects the channel. Each receiver may drain its /// own buffered messages before observing disconnection. pub struct BoundedSender { - shared: Arc>, + shared: Arc>>, } impl fmt::Debug for BoundedSender { @@ -192,7 +175,7 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { self.shared.senders.store(0, Ordering::Release); - common::disconnect(&self.shared.inner); + common::disconnect(&self.shared); } } @@ -238,7 +221,7 @@ impl BoundedSender { // Boxed once, out of the critical section, and reused by every retry. Dropped after // `SendState::drop` has already released the producer slot, so a cancelled send // unregisters before running the payload destructor. - value: Option>, + value: Option, } impl Drop for SendState<'_, T> { @@ -246,9 +229,13 @@ impl BoundedSender { // Take the slot with the channel unlocked afterward so the replaced waker is // dropped outside the lock. The payload in `value` is dropped only after this // returns. + self.sender + .shared + .producer_waiting + .store(0, Ordering::Release); let waker = { - let mut inner = self.sender.shared.inner.lock(); - inner.producer.take() + let mut state = self.sender.shared.state.lock(); + state.producer.take() }; drop(waker); } @@ -261,33 +248,59 @@ impl BoundedSender { None => return Poll::Ready(()), }; - let mut inner = self.sender.shared.inner.lock(); - - if !inner.log.has_receivers() { - inner.log.publish_discarded(); - let retired_producer = inner.producer.take(); - let wakers = inner.waiters.drain(); - drop(inner); + self.sender + .shared + .producer_waiting + .store(1, Ordering::Release); + let mut state = self.sender.shared.state.lock(); + + if state.receiver_count == 0 { + let next = Shared::>::next_tail( + self.sender.shared.tail.load(Ordering::Relaxed), + ); + let retired_producer = state.producer.take(); + self.sender + .shared + .producer_waiting + .store(0, Ordering::Release); + common::commit_discard( + &self.sender.shared.head, + &self.sender.shared.tail, + next, + ); + let wakers = state.waiters.drain(); + drop(state); wake_all(wakers); drop(retired_producer); drop(msg); return Poll::Ready(()); } - if inner.log.retained() == self.sender.shared.capacity { + let head = self.sender.shared.head.load(Ordering::Acquire); + let tail = self.sender.shared.tail.load(Ordering::Relaxed); + if tail - head >= self.sender.shared.buffer.cap as u64 { // Same critical section as the capacity check: a reclaim that lands between // those two observations cannot skip this waiter. - let retired = inner.producer.replace(cx.waker().clone()); - drop(inner); + let retired = state.producer.replace(cx.waker().clone()); + drop(state); drop(retired); self.value = Some(msg); return Poll::Pending; } - inner.log.publish_retained(msg); - let retired_producer = inner.producer.take(); - let wakers = inner.waiters.drain(); - drop(inner); + let next = Shared::>::next_tail(tail); + let n = state.receiver_count; + unsafe { + self.sender.shared.buffer.slot(tail).write(msg, n); + } + let retired_producer = state.producer.take(); + self.sender + .shared + .producer_waiting + .store(0, Ordering::Release); + common::commit_publish(&self.sender.shared.tail, next); + let wakers = state.waiters.drain(); + drop(state); // Wake receivers before dropping the retired producer waker: that waker is this // send, so it must not run under the lock, but a panic in its Drop must not skip // the receiver wake-ups either. @@ -299,7 +312,7 @@ impl BoundedSender { let mut send = SendState { sender: self, - value: Some(Arc::new(value)), + value: Some(value), }; poll_fn(|cx| send.poll_send(cx)).await } @@ -330,37 +343,43 @@ impl BoundedSender { /// tx.try_send(20).unwrap(); /// ``` pub fn try_send(&mut 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. - self.publish(value, Arc::new).map_err(TrySendError::Full) + self.publish(value).map_err(TrySendError::Full) } /// 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. - /// /// 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

(&mut self, payload: P, into_msg: impl FnOnce(P) -> Arc) -> Result<(), P> { + /// observe an empty buffer and park after this message became visible. The payload is moved + /// into the slot only after the capacity check, and `T::drop` for a discarded send runs after + /// the lock is released. + fn publish(&mut self, payload: T) -> Result<(), T> { let mut discarded = None; let wakers = { - let mut inner = self.shared.inner.lock(); + let mut state = self.shared.state.lock(); - if !inner.log.has_receivers() { + if state.receiver_count == 0 { // 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(); + let next = + Shared::>::next_tail(self.shared.tail.load(Ordering::Relaxed)); 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); + common::commit_discard(&self.shared.head, &self.shared.tail, next); } else { - inner.log.publish_retained(into_msg(payload)); - } + let head = self.shared.head.load(Ordering::Acquire); + let tail = self.shared.tail.load(Ordering::Relaxed); + if tail - head >= self.shared.buffer.cap as u64 { + // Nothing was published, so there is no wait set to drain. + return Err(payload); + } - inner.waiters.drain() + let next = Shared::>::next_tail(tail); + let n = state.receiver_count; + unsafe { + self.shared.buffer.slot(tail).write(payload, n); + } + common::commit_publish(&self.shared.tail, next); + } + state.waiters.drain() }; wake_all(wakers); @@ -390,7 +409,7 @@ impl BoundedSender { /// assert_eq!(tx.retained_message_count(), 0); /// ``` pub fn retained_message_count(&self) -> usize { - self.shared.inner.lock().log.retained() + self.shared.retained() } /// Returns the number of messages this channel retains before the producer waits. @@ -407,7 +426,7 @@ impl BoundedSender { /// assert_eq!(tx.capacity(), 8); /// ``` pub fn capacity(&self) -> usize { - self.shared.capacity + self.shared.buffer.cap } /// Creates a new receiver that starts receiving messages from the current tail of the channel. @@ -437,10 +456,10 @@ impl BoundedSender { /// ``` #[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(); + let cursor = self.shared.subscribe(); BoundedReceiver { shared: self.shared.clone(), - key, + cursor, } } } @@ -451,8 +470,8 @@ impl BoundedSender { /// that stops draining holds capacity for the whole channel, so dropping one that will not keep up /// is how a caller releases the producer. pub struct BoundedReceiver { - shared: Arc>, - key: SlotId, + shared: Arc>>, + cursor: u64, } impl fmt::Debug for BoundedReceiver { @@ -463,13 +482,7 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - let (reclaimed, producer) = { - let mut inner = self.shared.inner.lock(); - let reclaimed = inner.log.remove_receiver(self.key); - let drained_last = !inner.log.has_receivers(); - let producer = common::take_producer_on_reclaim(&mut inner, &reclaimed, drained_last); - (reclaimed, producer) - }; + let (reclaimed, producer) = common::drop_subscription(&self.shared, self.cursor); // Wake before dropping reclaimed payloads: a panicking destructor must not strand the // producer on capacity this drop already freed. @@ -532,14 +545,10 @@ impl BoundedReceiver { /// assert_eq!(rx.try_recv(), Ok(10)); /// ``` pub fn try_recv(&mut self) -> Result { - let ((msg, reclaimed), producer) = - common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; - - // Wake before taking the payload: `take_msg` runs `T::clone` and `T::drop`, and if either - // panics the slot this receive already freed would otherwise never reach the parked - // producer, stalling it permanently. + let consumed = common::try_receive(&self.shared, &mut self.cursor)?; + let producer = common::take_producer_on_reclaim(&self.shared, consumed.reclaimed, false); common::wake_producer(producer); - Ok(common::take_msg(msg, reclaimed)) + Ok(consumed.value) } } @@ -593,10 +602,10 @@ impl BoundedReceiver { /// ``` #[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(); + let cursor = self.shared.subscribe(); Self { shared: self.shared.clone(), - key, + cursor, } } @@ -625,7 +634,7 @@ impl BoundedReceiver { /// assert_eq!(rx.unread_message_count(), 1); /// ``` pub fn unread_message_count(&self) -> usize { - self.shared.inner.lock().log.unread(self.key) + self.shared.unread(self.cursor) } } @@ -641,12 +650,7 @@ impl Drop for Recv<'_, T> { return; } - common::unregister( - &self.receiver.shared.inner, - &self.receiver.shared.senders, - self.receiver.key, - &mut self.token, - ); + common::unregister(&self.receiver.shared, self.receiver.cursor, &mut self.token); } } @@ -656,21 +660,16 @@ impl Future for Recv<'_, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let Self { receiver, token } = self.get_mut(); - let ((msg, reclaimed), producer) = match common::poll_receive( - &receiver.shared.inner, - &receiver.shared.senders, - receiver.key, - token, - cx, - ) { + let consumed = match common::poll_receive(&receiver.shared, &mut receiver.cursor, token, cx) + { Poll::Pending => return Poll::Pending, Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), - Poll::Ready(Ok(outcome)) => outcome, + Poll::Ready(Ok(consumed)) => consumed, }; - // Wake before taking the payload, for the same reason as `try_recv`: a panicking - // `T::clone` must not strand the producer on a slot this receive already freed. + let producer = + common::take_producer_on_reclaim(&receiver.shared, consumed.reclaimed, false); common::wake_producer(producer); - Poll::Ready(Ok(common::take_msg(msg, reclaimed))) + Poll::Ready(Ok(consumed.value)) } } diff --git a/asyncband/src/broadcast/spmc/bounded/tests.rs b/asyncband/src/broadcast/spmc/bounded/tests.rs index ae110faa..49b0fbba 100644 --- a/asyncband/src/broadcast/spmc/bounded/tests.rs +++ b/asyncband/src/broadcast/spmc/bounded/tests.rs @@ -33,7 +33,7 @@ fn bounded_panics_on_zero_capacity() { fn send_panics_on_version_overflow() { // The receiver is dropped right away: the doctored counter would make its own drop overflow. let (mut tx, _) = bounded(1); - tx.shared.inner.lock().log.set_tail(u64::MAX); + tx.shared.set_tail(u64::MAX); let _ = tx.try_send(()); } @@ -41,7 +41,7 @@ fn send_panics_on_version_overflow() { fn buffer_is_preallocated_and_never_shrinks() { let capacity = 128; let (mut tx, mut rx) = bounded(capacity); - let allocated = tx.shared.inner.lock().log.buffer_capacity(); + let allocated = tx.shared.buffer.len(); assert!(allocated >= capacity); // Fill to capacity, drain completely, and repeat with a much smaller cycle. An elastic backlog @@ -56,7 +56,7 @@ fn buffer_is_preallocated_and_never_shrinks() { assert_eq!(rx.try_recv(), Ok(0)); assert_eq!(tx.retained_message_count(), 0); - assert_eq!(tx.shared.inner.lock().log.buffer_capacity(), allocated); + assert_eq!(tx.shared.buffer.len(), allocated); } #[test] @@ -74,10 +74,10 @@ fn at_most_one_producer_can_wait() { let mut cx = Context::from_waker(Waker::noop()); let mut send = Box::pin(tx.send(1)); assert!(send.as_mut().poll(&mut cx).is_pending()); - assert!(shared.inner.lock().producer.is_some()); + assert!(shared.state.lock().producer.is_some()); assert_eq!(rx.try_recv(), Ok(0)); assert!(send.as_mut().poll(&mut cx).is_ready()); drop(send); - assert!(shared.inner.lock().producer.is_none()); + assert!(shared.state.lock().producer.is_none()); } diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs index da5eb397..b2f8632a 100644 --- a/asyncband/src/broadcast/spmc/common.rs +++ b/asyncband/src/broadcast/spmc/common.rs @@ -15,15 +15,21 @@ // specific language governing permissions and limitations // under the License. -//! Storage, cursors, and the receive step shared by the bounded and unbounded SPMC broadcast +//! Slot log, waiters, and the receive step shared by the bounded and unbounded SPMC broadcast //! channels. //! -//! Retention and reclaim match the MPMC broadcast backlog. `Inner` also holds the single parked -//! producer so a bounded `send` can check capacity and register its waker in one critical section. - -use std::collections::VecDeque; -use std::mem; -use std::sync::Arc; +//! The producer is unique (`send` takes `&mut self`). Receivers drain already-published slots +//! without taking the waiter mutex: each slot carries a remaining-reader count, and each +//! subscription keeps its cursor locally. The mutex is only for subscribe/unsubscribe, parking, +//! and the producer's publish-and-drain critical section — which is what keeps a park from missing +//! a wake-up. + +use std::cell::UnsafeCell; +use std::mem::MaybeUninit; +use std::ptr; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicPtr; +use std::sync::atomic::AtomicU64; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -32,32 +38,30 @@ use std::task::Waker; 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; +/// Number of slots in one unbounded log chunk. +pub const CHUNK_LEN: usize = 64; -/// A received message together with the retained prefix that the receive released. +/// A received value together with whether this receive freed a retained slot. /// -/// 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); - -/// A receive plus the parked producer to wake if that receive freed capacity. -pub type RecvOutcome = (Received, Option); +/// A bounded channel wakes the producer from `reclaimed` before it returns `value`, so a panicking +/// `T::clone` cannot strand a waiter on capacity this receive already released. +pub struct Consumed { + pub value: T, + pub reclaimed: bool, +} -/// Messages removed from the shared buffer and waiting to be dropped after it is unlocked. +/// Messages removed from the shared log and waiting to be dropped after waiters are 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>, + first: Option, + rest: Vec, } impl Reclaimed { @@ -68,484 +72,547 @@ impl Reclaimed { } } - 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() + fn push(&mut self, msg: T) { + if self.first.is_none() { + self.first = Some(msg); + } else { + self.rest.push(msg); + } } } -/// 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, +/// One published value and the number of subscriptions that still have to consume it. +pub struct Slot { + msg: UnsafeCell>, + remaining: AtomicUsize, + /// `true` once the producer has written `msg` and until the last remaining reader takes it. + /// + /// Head only advances over a slot after this is cleared, so the producer cannot reuse the + /// memory while a reader is still cloning `T`. + occupied: AtomicBool, } -/// 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)`. +// SAFETY: The producer writes `msg` before publishing `tail`. Readers clone `T` only for versions +// they were counted in. The last remaining reader takes `msg` before clearing `occupied` and +// advancing `head`, which is what lets the producer reuse the slot. +unsafe impl Sync for Slot {} + +impl Slot { + fn empty() -> Self { + Self { + msg: UnsafeCell::new(MaybeUninit::uninit()), + remaining: AtomicUsize::new(0), + occupied: AtomicBool::new(false), + } + } + + /// Writes `msg` for `n` receivers. /// - /// 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, -} + /// The caller must publish `tail` with `Release` after this returns, and must be the unique + /// producer for this slot. + /// + /// # Safety + /// + /// `occupied` must be `false`. `n` must be the number of subscriptions that will consume this + /// version, and must be greater than zero. + pub unsafe fn write(&self, msg: T, n: usize) { + debug_assert!(n > 0); + debug_assert!(!self.occupied.load(Ordering::Relaxed)); + unsafe { + (*self.msg.get()).write(msg); + } + self.remaining.store(n, Ordering::Relaxed); + self.occupied.store(true, Ordering::Release); + } -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 }) + /// Clones the published value. + /// + /// # Safety + /// + /// This slot must currently hold a published message the caller is allowed to read. + unsafe fn clone_msg(&self) -> T + where + T: Clone, + { + unsafe { (*self.msg.get()).assume_init_ref().clone() } } - /// A backlog preallocated for `capacity` retained messages that never shrinks. - pub fn fixed(capacity: usize) -> Self { - Self::new(VecDeque::with_capacity(capacity), Retention::Fixed) + /// Takes the slot's value after the last remaining reader has consumed it. + /// + /// # Safety + /// + /// The caller must be the last remaining consumer of this slot. + unsafe fn take_msg(&self) -> T { + let msg = unsafe { (*self.msg.get()).assume_init_read() }; + self.occupied.store(false, Ordering::Release); + msg } +} + +/// Parked receivers, the live subscription count, and the single parked producer. +pub struct State { + pub waiters: WakerSet, + pub receiver_count: usize, + pub producer: Option, +} - fn new(buffer: VecDeque>, retention: Retention) -> Self { +/// Shared channel state: the slot log plus the waiter mutex. +pub struct Shared { + pub buffer: B, + pub head: AtomicU64, + pub tail: AtomicU64, + pub senders: AtomicUsize, + pub producer_waiting: AtomicUsize, + pub state: Mutex, +} + +impl Shared { + pub fn new(buffer: B) -> Self { Self { buffer, - head: 0, - head_receivers: 0, - tail: 0, - receivers: Arena::new(), - retention, + head: AtomicU64::new(0), + tail: AtomicU64::new(0), + senders: AtomicUsize::new(1), + producer_waiting: AtomicUsize::new(0), + state: Mutex::new(State { + waiters: WakerSet::new(), + receiver_count: 1, + producer: None, + }), } } - /// 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() + let tail = self.tail.load(Ordering::Acquire); + let head = self.head.load(Ordering::Acquire); + usize::try_from(tail - head).expect("retained broadcast message count exceeds usize") } - /// 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() + pub fn unread(&self, cursor: u64) -> usize { + let tail = self.tail.load(Ordering::Acquire); + debug_assert!(tail >= cursor); + usize::try_from(tail - cursor).expect("unread broadcast message count exceeds usize") } - /// 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") + /// Next committed version, panicking on overflow. + pub fn next_tail(tail: u64) -> u64 { + tail.checked_add(1) + .expect("broadcast channel version counter overflowed") } - /// 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"); + /// Registers a new subscription at the committed tail. + pub fn subscribe(&self) -> u64 { + let mut state = self.state.lock(); + state.receiver_count += 1; + self.tail.load(Ordering::Relaxed) } - /// 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 + #[cfg(test)] + pub fn set_tail(&self, tail: u64) { + self.tail.store(tail, Ordering::Relaxed); + self.head.store(tail, Ordering::Relaxed); + } +} + +/// A random-access published slot, addressed by its committed version. +pub trait SlotStore { + fn slot(&self, version: u64) -> &Slot; + + /// Moves the lookup start past chunks the live window has left behind. /// - /// 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); + /// Chunks stay allocated until the channel is dropped, so a receiver that still holds an old + /// pointer cannot observe a free. Skipping them keeps `slot` proportional to the live window + /// rather than to the lifetime message count. + fn sync_head(&self, _head: u64) {} +} + +/// Fixed ring used by the bounded channel. +pub struct BoundedBuffer { + slots: Box<[Slot]>, + pub cap: usize, +} + +impl BoundedBuffer { + pub fn new(capacity: usize) -> Self { + Self { + slots: (0..capacity).map(|_| Slot::empty()).collect(), + cap: capacity, } + } - self.publish_retained(msg); - None + #[cfg(test)] + pub fn len(&self) -> usize { + self.slots.len() } +} - /// 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()); +impl Drop for BoundedBuffer { + fn drop(&mut self) { + for slot in self.slots.iter() { + if slot.occupied.load(Ordering::Relaxed) { + drop(unsafe { slot.take_msg() }); + } } } +} - fn insert_receiver(&mut self, head: u64) -> SlotId { - if head == self.head { - self.head_receivers += 1; - } +impl SlotStore for BoundedBuffer { + #[inline] + fn slot(&self, version: u64) -> &Slot { + &self.slots[(version % self.cap as u64) as usize] + } +} + +/// One growable segment of the unbounded log. +pub struct Chunk { + slots: [Slot; CHUNK_LEN], + next: AtomicPtr>, + base: u64, +} - self.receivers.insert(head) +impl Chunk { + fn new(base: u64) -> Box { + Box::new(Self { + slots: std::array::from_fn(|_| Slot::empty()), + next: AtomicPtr::new(ptr::null_mut()), + base, + }) } +} - /// 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) +/// Linked chunks used by the unbounded channel. +/// +/// The producer appends chunks without moving earlier slots, so receivers can drain without a +/// publication lock. Fully consumed chunks are recycled onto the sender's spare list. +pub struct UnboundedBuffer { + /// First chunk ever allocated. Never moves; `Drop` walks from here. + root: AtomicPtr>, + /// First chunk that may still hold a live message. Lookup starts here. + head_chunk: AtomicPtr>, + tail_chunk: AtomicPtr>, + _marker: std::marker::PhantomData>, +} + +impl UnboundedBuffer { + pub fn new() -> Self { + let chunk = Box::into_raw(Chunk::new(0)); + Self { + root: AtomicPtr::new(chunk), + head_chunk: AtomicPtr::new(chunk), + tail_chunk: AtomicPtr::new(chunk), + _marker: std::marker::PhantomData, + } } - pub fn remove_receiver(&mut self, key: SlotId) -> Reclaimed { - let head = self.receivers.remove(key); + /// Ensures the chunk that holds `version` exists and returns that slot. + /// + /// The caller is the unique producer and must not publish `tail` past this version until this + /// returns. Fully consumed chunks stay allocated until the channel is dropped so a receiver + /// walking the list cannot observe a freed chunk. + pub fn slot_for_publish(&self, version: u64) -> &Slot { + loop { + let chunk = self.tail_chunk.load(Ordering::Acquire); + debug_assert!(!chunk.is_null()); + let current = unsafe { &*chunk }; + if version < current.base + CHUNK_LEN as u64 { + debug_assert!(version >= current.base); + return ¤t.slots[(version - current.base) as usize]; + } - if head == self.head { - self.release_head_receiver() - } else { - Reclaimed::empty() + let next_base = current.base + CHUNK_LEN as u64; + let raw = Box::into_raw(Chunk::new(next_base)); + current.next.store(raw, Ordering::Release); + self.tail_chunk.store(raw, Ordering::Release); } } - fn release_head_receiver(&mut self) -> Reclaimed { - self.head_receivers -= 1; + fn advance_head_chunk(&self, head: u64) { + loop { + let chunk = self.head_chunk.load(Ordering::Acquire); + let current = unsafe { &*chunk }; + let next = current.next.load(Ordering::Acquire); + if next.is_null() { + return; + } + if current.base + CHUNK_LEN as u64 > head { + return; + } + self.head_chunk.store(next, Ordering::Release); + } + } - if self.head_receivers == 0 { - self.reclaim_consumed() - } else { - Reclaimed::empty() + #[cfg(test)] + pub fn allocated_slots(&self) -> usize { + let mut n = 0; + let mut chunk = self.root.load(Ordering::Acquire); + while !chunk.is_null() { + n += CHUNK_LEN; + chunk = unsafe { (*chunk).next.load(Ordering::Acquire) }; } + n } +} - 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; +impl Drop for UnboundedBuffer { + fn drop(&mut self) { + let mut chunk = self.root.load(Ordering::Relaxed); + while !chunk.is_null() { + let boxed = unsafe { Box::from_raw(chunk) }; + for slot in &boxed.slots { + if slot.occupied.load(Ordering::Relaxed) { + drop(unsafe { slot.take_msg() }); + } } - let head = *cursor; - *cursor += 1; - head - }; + chunk = boxed.next.load(Ordering::Relaxed); + } + } +} - 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; +impl SlotStore for UnboundedBuffer { + #[inline] + fn slot(&self, version: u64) -> &Slot { + let mut chunk = self.head_chunk.load(Ordering::Acquire); + loop { + debug_assert!(!chunk.is_null()); + let current = unsafe { &*chunk }; + if version < current.base + CHUNK_LEN as u64 { + debug_assert!(version >= current.base); + return ¤t.slots[(version - current.base) as usize]; } + chunk = current.next.load(Ordering::Acquire); } + } - 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); + fn sync_head(&self, head: u64) { + self.advance_head_chunk(head); + } +} - self.head = next_head; - self.head_receivers = head_receivers; - self.shrink_buffer(); - reclaimed +/// Advances `head` over slots whose value has already been taken. +fn advance_head>(shared: &Shared) { + let mut h = shared.head.load(Ordering::Acquire); + loop { + let t = shared.tail.load(Ordering::Acquire); + if h >= t { + break; + } + if shared.buffer.slot(h).occupied.load(Ordering::Acquire) { + break; + } + match shared + .head + .compare_exchange_weak(h, h + 1, Ordering::AcqRel, Ordering::Acquire) + { + Ok(_) => h += 1, + Err(actual) => h = actual, + } } + shared.buffer.sync_head(h); +} - /// 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; +/// Consumes the message at `cursor` and advances the cursor. +/// +/// A subscription that is the last remaining reader takes the slot value without cloning. Any +/// other reader clones `T` with the waiter mutex not held. The cursor advances before that clone +/// so a panicking `T::clone` still consumes the subscription's claim; a drop guard releases the +/// remaining count if the clone unwinds. +pub fn consume>(shared: &Shared, cursor: &mut u64) -> Consumed { + let version = *cursor; + debug_assert!(version < shared.tail.load(Ordering::Acquire)); + let slot = shared.buffer.slot(version); + *cursor = version + 1; + + if slot.remaining.load(Ordering::Acquire) == 1 { + let value = unsafe { slot.take_msg() }; + slot.remaining.store(0, Ordering::Release); + advance_head(shared); + return Consumed { + value, + reclaimed: true, }; + } - if !self.buffer.is_empty() { - return; - } + struct RemainingGuard<'a, T, B: SlotStore> { + slot: &'a Slot, + shared: &'a Shared, + armed: bool, + } - 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)); + impl> Drop for RemainingGuard<'_, T, B> { + fn drop(&mut self) { + if !self.armed { + return; + } + if self.slot.remaining.fetch_sub(1, Ordering::AcqRel) == 1 { + drop(unsafe { self.slot.take_msg() }); + advance_head(self.shared); + } } } - #[cfg(test)] - pub fn buffer_capacity(&self) -> usize { - self.buffer.capacity() + let mut guard = RemainingGuard { + slot, + shared, + armed: true, + }; + let value = unsafe { slot.clone_msg() }; + let last = slot.remaining.fetch_sub(1, Ordering::AcqRel) == 1; + guard.armed = false; + if last { + drop(unsafe { slot.take_msg() }); + advance_head(shared); } + Consumed { + value, + reclaimed: last, + } +} - /// 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; +fn reclaim_range>(shared: &Shared, start: u64, end: u64) -> Reclaimed { + let mut reclaimed = Reclaimed::empty(); + let mut version = start; + while version < end { + let slot = shared.buffer.slot(version); + if slot.remaining.fetch_sub(1, Ordering::AcqRel) == 1 { + reclaimed.push(unsafe { slot.take_msg() }); + } + version += 1; + } + if !reclaimed.is_empty() { + advance_head(shared); } + reclaimed } -/// Buffer, receiver cursors, parked receivers, and at most one parked producer, 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. -/// -/// The producer slot is the SPMC counterpart of MPMC's semaphore: `send(&mut self)` means at most -/// one waiting publisher, so occupancy is a single `Option`. Checking capacity and registering -/// that waker share this lock, which is what makes the wait path race-free without a recheck -/// protocol. -pub struct Inner { - pub log: Backlog, - pub waiters: WakerSet, - pub producer: Option, +/// Takes the parked producer if a reclaim may have freed capacity. +pub fn take_producer_on_reclaim( + shared: &Shared, + reclaimed: bool, + drained_last: bool, +) -> Option { + if !drained_last && !reclaimed { + return None; + } + if !drained_last && shared.producer_waiting.load(Ordering::Acquire) == 0 { + return None; + } + + let mut state = shared.state.lock(); + state.producer.take() } -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(), - producer: None, - }); - (inner, key) +/// Wakes the parked producer, if any, with the channel already unlocked. +pub fn wake_producer(producer: Option) { + if let Some(waker) = producer { + waker.wake(); } } +/// Drops a subscription, reclaiming every unread slot it still held. +pub fn drop_subscription>( + shared: &Shared, + cursor: u64, +) -> (Reclaimed, Option) { + let mut state = shared.state.lock(); + state.receiver_count -= 1; + let last = state.receiver_count == 0; + let tail = shared.tail.load(Ordering::Acquire); + + if last { + // Take slot values before the producer can observe `receiver_count == 0` and reuse them. + let reclaimed = reclaim_range(shared, cursor, tail); + shared.producer_waiting.store(0, Ordering::Release); + let producer = state.producer.take(); + drop(state); + return (reclaimed, producer); + } + drop(state); + + let reclaimed = reclaim_range(shared, cursor, tail); + let producer = take_producer_on_reclaim(shared, !reclaimed.is_empty(), false); + (reclaimed, producer) +} + /// Wakes every parked receiver so it can observe the channel's disconnected state. -/// -/// The sender's `Drop` calls this. -pub fn disconnect(inner: &Mutex>) { +pub fn disconnect(shared: &Shared) { let wakers = { - let mut inner = inner.lock(); - inner.waiters.take_all() + let mut state = shared.state.lock(); + state.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. +pub fn unregister(shared: &Shared, cursor: u64, token: &mut Option) { + let mut state = shared.state.lock(); + if cursor < shared.tail.load(Ordering::Relaxed) || shared.senders.load(Ordering::Acquire) == 0 { *token = None; return; } - let waker = inner.waiters.unregister(token); - drop(inner); + let waker = state.waiters.unregister(token); + drop(state); drop(waker); } -/// Receives without waiting, yielding the message, the prefix the receive released, and the -/// parked producer if that reclaim freed capacity. -/// -/// The caller owns what happens next: a bounded channel wakes the producer 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) => { - let producer = take_producer_on_reclaim(&mut inner, &received.1, false); - Ok((received, producer)) - } - None if senders.load(Ordering::Acquire) == 0 => Err(TryRecvError::Disconnected), - None => Err(TryRecvError::Empty), +/// Receives without waiting. +pub fn try_receive>( + shared: &Shared, + cursor: &mut u64, +) -> Result, TryRecvError> { + if *cursor < shared.tail.load(Ordering::Acquire) { + return Ok(consume(shared, cursor)); + } + if shared.senders.load(Ordering::Acquire) == 0 { + Err(TryRecvError::Disconnected) + } else { + 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. A reclaim that frees capacity takes the -/// producer waker in the same critical section. -pub fn poll_receive( - inner: &Mutex>, - senders: &AtomicUsize, - key: SlotId, +/// The ready path does not take the waiter mutex. Parking rechecks `tail` under that mutex so a +/// publish that drains waiters cannot slip between the empty check and the registration. +pub fn poll_receive>( + shared: &Shared, + cursor: &mut u64, token: &mut Option, cx: &mut Context<'_>, -) -> Poll, RecvError>> { - let mut inner = inner.lock(); - match inner.log.receive(key) { - Some(received) => { - *token = None; - let producer = take_producer_on_reclaim(&mut inner, &received.1, false); - Poll::Ready(Ok((received, producer))) - } - 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 - } +) -> Poll, RecvError>> { + if *cursor < shared.tail.load(Ordering::Acquire) { + *token = None; + return Poll::Ready(Ok(consume(shared, cursor))); } -} -/// Takes the parked producer when this reclaim is one of the two sources that can free capacity. -/// -/// `buffer` shrinks only in `reclaim_consumed`, reachable from a receive that vacates the head -/// and from removing a subscription. Both must wake the producer. A new subscription starts at -/// the tail and never lowers `retained()`, so it never takes this path. -pub fn take_producer_on_reclaim( - inner: &mut Inner, - reclaimed: &Reclaimed, - drained_last: bool, -) -> Option { - if drained_last || !reclaimed.is_empty() { - inner.producer.take() - } else { - None + let mut state = shared.state.lock(); + if *cursor < shared.tail.load(Ordering::Acquire) { + *token = None; + drop(state); + return Poll::Ready(Ok(consume(shared, cursor))); } -} - -/// Wakes the parked producer, if any, with the channel already unlocked. -pub fn wake_producer(producer: Option) { - if let Some(waker) = producer { - waker.wake(); + if shared.senders.load(Ordering::Acquire) == 0 { + *token = None; + return Poll::Ready(Err(RecvError::Disconnected)); } -} -/// 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); + let retired_waker = state.waiters.register(token, cx.waker()); + drop(state); + drop(retired_waker); + Poll::Pending +} - if !sole_owner { - return (*msg).clone(); - } +/// Publishes `tail` after writing a slot. Caller holds `state` and drains waiters after this. +pub fn commit_publish(tail: &AtomicU64, next: u64) { + tail.store(next, Ordering::Release); +} - // 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()) +/// Advances `head` and `tail` together when nothing can read the message. +pub fn commit_discard(head: &AtomicU64, tail: &AtomicU64, next: u64) { + head.store(next, Ordering::Release); + tail.store(next, Ordering::Release); } diff --git a/asyncband/src/broadcast/spmc/mod.rs b/asyncband/src/broadcast/spmc/mod.rs index 092b9488..7afbddd0 100644 --- a/asyncband/src/broadcast/spmc/mod.rs +++ b/asyncband/src/broadcast/spmc/mod.rs @@ -26,10 +26,8 @@ //! The sender is not [`Clone`], and every publish method takes `&mut self`. That is the static //! single-writer contract this topology adds over [`crate::broadcast::mpmc`]: there cannot be a //! second producer, at compile time. `&Sender` can still be shared for [`subscribe`] and the -//! inspection methods. The family is offered for that exclusive-send API. Measured against -//! `async-broadcast` and `tokio::sync::broadcast`, tight bounded wait (capacity 1) is competitive; -//! it is not a general throughput upgrade over 1-producer `broadcast::mpmc` or over Tokio's -//! non-blocking ring. +//! inspection methods. Receivers drain already-published slots without taking the publication +//! lock, which is the throughput reason to pick this family when there is one producer. //! //! This is fan-out broadcast, not a competing queue: every accepted value is delivered to every //! active subscription. A competitive `asyncband::spmc` queue would give each value to exactly one diff --git a/asyncband/src/broadcast/spmc/unbounded/mod.rs b/asyncband/src/broadcast/spmc/unbounded/mod.rs index 845bd2c0..9c059eb2 100644 --- a/asyncband/src/broadcast/spmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/spmc/unbounded/mod.rs @@ -59,18 +59,15 @@ use std::fmt; use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; use super::common; -use super::common::Backlog; -use super::common::Inner; +use super::common::Shared; +use super::common::UnboundedBuffer; 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::WakerToken; @@ -92,31 +89,20 @@ mod tests; /// assert_eq!(receiver.try_recv(), Ok("ready")); /// ``` pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { - let (inner, key) = Inner::with_first_subscription(Backlog::elastic()); - let shared = Arc::new(Shared { - inner, - senders: AtomicUsize::new(1), - }); + let shared = Arc::new(Shared::new(UnboundedBuffer::new())); let sender = UnboundedSender { shared: shared.clone(), }; - let receiver = UnboundedReceiver { shared, key }; + let receiver = UnboundedReceiver { shared, cursor: 0 }; (sender, receiver) } -struct Shared { - /// Buffer, receiver cursors, and parked receivers, all under a single lock. - inner: Mutex>, - /// `1` while the sender is alive, `0` after it is dropped. - senders: AtomicUsize, -} - /// A publishing handle for an unbounded broadcast channel. /// /// This handle is not [`Clone`]. Once it is dropped, each receiver can drain the values already /// published for it and then observes disconnection. pub struct UnboundedSender { - shared: Arc>, + shared: Arc>>, } impl fmt::Debug for UnboundedSender { @@ -128,7 +114,7 @@ impl fmt::Debug for UnboundedSender { impl Drop for UnboundedSender { fn drop(&mut self) { self.shared.senders.store(0, Ordering::Release); - common::disconnect(&self.shared.inner); + common::disconnect(&self.shared); } } @@ -157,15 +143,25 @@ impl UnboundedSender { /// assert_eq!(second.try_recv(), Ok("update")); /// ``` pub fn send(&mut self, msg: T) { - let msg = Arc::new(msg); - // Publishing and draining the wait set share one critical section, so a receiver can never // observe an empty buffer and park after this message became visible. let (unretained, wakers) = { - let mut inner = self.shared.inner.lock(); - let unretained = inner.log.publish(msg); - let wakers = inner.waiters.drain(); - (unretained, wakers) + let mut state = self.shared.state.lock(); + let tail = self.shared.tail.load(Ordering::Relaxed); + let next = Shared::>::next_tail(tail); + let unretained = if state.receiver_count == 0 { + common::commit_discard(&self.shared.head, &self.shared.tail, next); + Some(msg) + } else { + let n = state.receiver_count; + let slot = self.shared.buffer.slot_for_publish(tail); + unsafe { + slot.write(msg, n); + } + common::commit_publish(&self.shared.tail, next); + None + }; + (unretained, state.waiters.drain()) }; // Notify all waiting receivers. An unsent message is dropped here too, once the lock is @@ -198,7 +194,7 @@ impl UnboundedSender { /// assert_eq!(publisher.retained_message_count(), 0); /// ``` pub fn retained_message_count(&self) -> usize { - self.shared.inner.lock().log.retained() + self.shared.retained() } /// Subscribes a new receiver for values published from this point forward. @@ -221,10 +217,10 @@ impl UnboundedSender { /// ``` #[must_use = "the receiver is dropped immediately if it is not retained"] pub fn subscribe(&self) -> UnboundedReceiver { - let key = self.shared.inner.lock().log.subscribe(); + let cursor = self.shared.subscribe(); UnboundedReceiver { shared: self.shared.clone(), - key, + cursor, } } } @@ -234,8 +230,8 @@ impl UnboundedSender { /// This receiver observes every value published after its subscription point and retains its own /// position in the shared backlog. pub struct UnboundedReceiver { - shared: Arc>, - key: SlotId, + shared: Arc>>, + cursor: u64, } impl fmt::Debug for UnboundedReceiver { @@ -246,10 +242,7 @@ impl fmt::Debug for UnboundedReceiver { impl Drop for UnboundedReceiver { fn drop(&mut self) { - let reclaimed = { - let mut inner = self.shared.inner.lock(); - inner.log.remove_receiver(self.key) - }; + let (reclaimed, _producer) = common::drop_subscription(&self.shared, self.cursor); drop(reclaimed); } } @@ -312,9 +305,7 @@ impl UnboundedReceiver { /// assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - let ((msg, reclaimed), _producer) = - common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; - Ok(common::take_msg(msg, reclaimed)) + Ok(common::try_receive(&self.shared, &mut self.cursor)?.value) } } @@ -344,10 +335,10 @@ impl UnboundedReceiver { /// ``` #[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(); + let cursor = self.shared.subscribe(); Self { shared: self.shared.clone(), - key, + cursor, } } @@ -373,7 +364,7 @@ impl UnboundedReceiver { /// assert_eq!(receiver.unread_message_count(), 1); /// ``` pub fn unread_message_count(&self) -> usize { - self.shared.inner.lock().log.unread(self.key) + self.shared.unread(self.cursor) } } @@ -389,12 +380,7 @@ impl Drop for Recv<'_, T> { return; } - common::unregister( - &self.receiver.shared.inner, - &self.receiver.shared.senders, - self.receiver.key, - &mut self.token, - ); + common::unregister(&self.receiver.shared, self.receiver.cursor, &mut self.token); } } @@ -404,18 +390,13 @@ impl Future for Recv<'_, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let Self { receiver, token } = self.get_mut(); - let ((msg, reclaimed), _producer) = match common::poll_receive( - &receiver.shared.inner, - &receiver.shared.senders, - receiver.key, - token, - cx, - ) { + let consumed = match common::poll_receive(&receiver.shared, &mut receiver.cursor, token, cx) + { Poll::Pending => return Poll::Pending, Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), - Poll::Ready(Ok(outcome)) => outcome, + Poll::Ready(Ok(consumed)) => consumed, }; - Poll::Ready(Ok(common::take_msg(msg, reclaimed))) + Poll::Ready(Ok(consumed.value)) } } diff --git a/asyncband/src/broadcast/spmc/unbounded/tests.rs b/asyncband/src/broadcast/spmc/unbounded/tests.rs index 86fc8e83..2ea9df9c 100644 --- a/asyncband/src/broadcast/spmc/unbounded/tests.rs +++ b/asyncband/src/broadcast/spmc/unbounded/tests.rs @@ -16,55 +16,33 @@ // under the License. use super::*; -use crate::broadcast::spmc::common::MIN_RETAINED_CAPACITY; +use crate::broadcast::spmc::common::CHUNK_LEN; #[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 (mut tx, _) = unbounded(); - tx.shared.inner.lock().log.set_tail(u64::MAX); + tx.shared.set_tail(u64::MAX); tx.send(()); } #[test] -fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { +fn chunks_grow_with_the_committed_log() { let (mut tx, mut rx) = unbounded(); - let burst = MIN_RETAINED_CAPACITY * 16; + let burst = CHUNK_LEN * 4; for i in 0..burst { tx.send(i); } - assert!(tx.shared.inner.lock().log.buffer_capacity() >= burst); + assert!(tx.shared.buffer.allocated_slots() >= burst); for i in 0..burst { assert_eq!(rx.try_recv(), Ok(i)); } - // Draining evaluates the cycle that just peaked, so the burst allocation is still held. + // Chunks stay allocated until the channel is dropped so receivers can walk them without a + // reclamation lock. assert_eq!(tx.retained_message_count(), 0); - 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().log.buffer_capacity() < burst); -} - -#[test] -fn repeated_bursts_keep_their_allocation() { - let (mut tx, mut rx) = unbounded(); - let burst = MIN_RETAINED_CAPACITY * 4; - - for _ in 0..4 { - for i in 0..burst { - tx.send(i); - } - for i in 0..burst { - assert_eq!(rx.try_recv(), Ok(i)); - } - } - - // Every cycle peaks at the same size, so the buffer must not rebuild its allocation each time. - assert!(tx.shared.inner.lock().log.buffer_capacity() >= burst); + assert!(tx.shared.buffer.allocated_slots() >= burst); } diff --git a/tests-integration/tests/broadcast_spmc_bounded_test.rs b/tests-integration/tests/broadcast_spmc_bounded_test.rs index b86cee98..e6e40f35 100644 --- a/tests-integration/tests/broadcast_spmc_bounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_bounded_test.rs @@ -36,6 +36,7 @@ use std::sync::Mutex; use std::task::Context; use std::task::Wake; use std::task::Waker; +use std::thread; use asyncband::blocking::FutureExt; use asyncband::broadcast::spmc::*; @@ -67,7 +68,8 @@ impl Drop for Reentrant { // `try_lock`: draining the probe itself may drop a payload while this mutex is held. // Deadlocks if the channel still holds its lock while dropping reclaimed messages. if let Ok(probe) = probe.try_lock() { - let _ = probe.unread_message_count(); + // `resubscribe` takes the waiter mutex; `unread_message_count` does not. + let _ = probe.resubscribe(); } } } @@ -201,6 +203,40 @@ fn unread_message_count_tracks_each_receiver() { assert_eq!(rx2.unread_message_count(), 2); } +#[test] +fn concurrent_receivers_keep_up_with_the_producer() { + const RECEIVERS: usize = 8; + const MESSAGES: usize = 256; + + let (mut tx, rx) = bounded(64); + let mut receivers = vec![rx]; + for _ in 1..RECEIVERS { + receivers.push(tx.subscribe()); + } + + let expected = (MESSAGES as u64 - 1) * MESSAGES as u64 / 2; + let handles: Vec<_> = receivers + .into_iter() + .map(|mut rx| { + thread::spawn(move || { + let mut sum = 0u64; + for _ in 0..MESSAGES { + sum += FutureExt::block_on(rx.recv()).unwrap(); + } + sum + }) + }) + .collect(); + + for value in 0..MESSAGES as u64 { + FutureExt::block_on(tx.send(value)); + } + + for handle in handles { + assert_eq!(handle.join().unwrap(), expected); + } +} + #[test] fn publish_order_is_program_order() { let (mut tx, mut rx1) = bounded(8); diff --git a/tests-integration/tests/broadcast_spmc_unbounded_test.rs b/tests-integration/tests/broadcast_spmc_unbounded_test.rs index c2825eed..a1b53469 100644 --- a/tests-integration/tests/broadcast_spmc_unbounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_unbounded_test.rs @@ -28,6 +28,7 @@ use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; use std::task::Waker; +use std::thread; use asyncband::broadcast::spmc::*; use tests_integration::WakeCounter; @@ -35,8 +36,7 @@ use tests_integration::assert_completes_without_deadlock; /// A payload whose destructor re-enters the channel it was sent through. /// -/// The sender is not `Clone`, so the probe is a shared receiver handle. `unread_message_count` -/// takes the same channel lock the destructor must not already hold. +/// The sender is not `Clone`, so the probe is a shared receiver handle. struct Reentrant { value: u64, probe: Option>>>, @@ -57,7 +57,8 @@ impl Drop for Reentrant { // `try_lock`: draining the probe itself may drop a payload while this mutex is held. // Deadlocks if the channel still holds its lock while dropping reclaimed messages. if let Ok(probe) = probe.try_lock() { - let _ = probe.unread_message_count(); + // `resubscribe` takes the waiter mutex; `unread_message_count` does not. + let _ = probe.resubscribe(); } } } @@ -430,6 +431,40 @@ async fn recv_reports_disconnection_without_any_message() { assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); } +#[test] +fn concurrent_receivers_drain_a_published_batch() { + const RECEIVERS: usize = 8; + const MESSAGES: usize = 256; + + let (mut tx, rx) = unbounded(); + let mut receivers = vec![rx]; + for _ in 1..RECEIVERS { + receivers.push(tx.subscribe()); + } + for value in 0..MESSAGES as u64 { + tx.send(value); + } + + let expected = (MESSAGES as u64 - 1) * MESSAGES as u64 / 2; + let handles: Vec<_> = receivers + .into_iter() + .map(|mut rx| { + thread::spawn(move || { + let mut sum = 0u64; + for _ in 0..MESSAGES { + sum += rx.try_recv().unwrap(); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + sum + }) + }) + .collect(); + + for handle in handles { + assert_eq!(handle.join().unwrap(), expected); + } +} + #[test] fn publish_order_is_program_order() { let (mut tx, mut rx1) = unbounded(); From c34cca12e8fc247e745efd032347cdd375332868 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Mon, 14 Sep 2026 17:12:12 +0800 Subject: [PATCH 3/9] perf(broadcast): park bounded SPMC waits on a shared condvar Native-thread send_blocking/recv_blocking wait on one epoch condvar instead of an async waker per task. --- CHANGELOG.md | 2 +- asyncband/src/broadcast/spmc/bounded/mod.rs | 65 +++++++++++++++++++ asyncband/src/broadcast/spmc/common.rs | 17 +++++ .../ecosystem/broadcast/spmc/adapters.rs | 8 +++ .../tests/broadcast_spmc_bounded_test.rs | 30 ++++++++- 5 files changed, 119 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4e3a85b7..f0d4f906 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,7 +27,7 @@ All notable changes to this project will be documented in this file. * Add `event::AutoResetEvent`, a reusable signal that releases one waiter, retains at most one unassigned signal that can be cleared with `reset`, and transfers assigned signals when waits are cancelled. * Add `ManualResetEvent::try_wait` to check readiness without registering a waiter or consuming the set state. -* Add `broadcast::spmc`, a lossless single-producer broadcast family with a non-cloneable sender whose publish methods require exclusive access; receivers drain already-published slots without taking the publication lock, bounded retains at most the requested capacity and makes the producer wait for the slowest active subscription, and unbounded never waits and lets the retained backlog grow. +* Add `broadcast::spmc`, a lossless single-producer broadcast family with a non-cloneable sender whose publish methods require exclusive access; receivers drain already-published slots without taking the publication lock, bounded retains at most the requested capacity and makes the producer wait for the slowest active subscription, unbounded never waits and lets the retained backlog grow, and bounded `send_blocking` / `recv_blocking` park native threads on a shared condvar. * 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`. diff --git a/asyncband/src/broadcast/spmc/bounded/mod.rs b/asyncband/src/broadcast/spmc/bounded/mod.rs index f013452f..7bfa8966 100644 --- a/asyncband/src/broadcast/spmc/bounded/mod.rs +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -271,6 +271,7 @@ impl BoundedSender { let wakers = state.waiters.drain(); drop(state); wake_all(wakers); + common::notify_blocking(self.sender.shared.as_ref()); drop(retired_producer); drop(msg); return Poll::Ready(()); @@ -305,6 +306,7 @@ impl BoundedSender { // send, so it must not run under the lock, but a panic in its Drop must not skip // the receiver wake-ups either. wake_all(wakers); + common::notify_blocking(self.sender.shared.as_ref()); drop(retired_producer); Poll::Ready(()) } @@ -346,6 +348,38 @@ impl BoundedSender { self.publish(value).map_err(TrySendError::Full) } + /// Broadcasts a value, parking the current thread while the channel is at capacity. + pub fn send_blocking(&mut self, mut value: T) { + loop { + match self.try_send(value) { + Ok(()) => return, + Err(TrySendError::Full(returned)) => { + value = returned; + let shared = self.shared.clone(); + shared.producer_waiting.store(1, Ordering::Release); + let mut epoch = shared.blocking.lock(); + let cap = shared.buffer.cap as u64; + if shared.tail.load(Ordering::Acquire) - shared.head.load(Ordering::Acquire) + < cap + { + drop(epoch); + continue; + } + let snap = *epoch; + while *epoch == snap + && shared.tail.load(Ordering::Acquire) - shared.head.load(Ordering::Acquire) + >= cap + { + epoch = shared + .blocking_cvar + .wait(epoch) + .unwrap_or_else(|e| e.into_inner()); + } + } + } + } + } + /// The publish step both send paths share. /// /// Publishing and draining the wait set share one critical section, so a receiver can never @@ -383,6 +417,7 @@ impl BoundedSender { }; wake_all(wakers); + common::notify_blocking(self.shared.as_ref()); drop(discarded); Ok(()) } @@ -550,6 +585,36 @@ impl BoundedReceiver { common::wake_producer(producer); Ok(consumed.value) } + + /// Receives the next value, parking the current thread while empty. + /// + /// Same result as [`recv`](Self::recv), without going through an async waker. Native-thread + /// waiters share one futex so a publish can wake every blocked receiver with a single notify. + pub fn recv_blocking(&mut self) -> Result { + loop { + match self.try_recv() { + Ok(value) => return Ok(value), + Err(TryRecvError::Disconnected) => return Err(RecvError::Disconnected), + Err(TryRecvError::Empty) => { + let shared = self.shared.clone(); + let mut epoch = shared.blocking.lock(); + if self.cursor < shared.tail.load(Ordering::Acquire) { + continue; + } + if shared.senders.load(Ordering::Acquire) == 0 { + return Err(RecvError::Disconnected); + } + let snap = *epoch; + while *epoch == snap && self.cursor >= shared.tail.load(Ordering::Acquire) { + epoch = shared + .blocking_cvar + .wait(epoch) + .unwrap_or_else(|e| e.into_inner()); + } + } + } + } + } } impl BoundedReceiver { diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs index b2f8632a..ce196795 100644 --- a/asyncband/src/broadcast/spmc/common.rs +++ b/asyncband/src/broadcast/spmc/common.rs @@ -27,6 +27,7 @@ use std::cell::UnsafeCell; use std::mem::MaybeUninit; use std::ptr; +use std::sync::Condvar; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicPtr; use std::sync::atomic::AtomicU64; @@ -168,6 +169,10 @@ pub struct Shared { pub senders: AtomicUsize, pub producer_waiting: AtomicUsize, pub state: Mutex, + /// Native-thread waiters. Publish and reclaim notify this condvar so a blocking receive does + /// not go through one async waker per parked task. + pub blocking: Mutex, + pub blocking_cvar: Condvar, } impl Shared { @@ -178,6 +183,8 @@ impl Shared { tail: AtomicU64::new(0), senders: AtomicUsize::new(1), producer_waiting: AtomicUsize::new(0), + blocking: Mutex::new(0), + blocking_cvar: Condvar::new(), state: Mutex::new(State { waiters: WakerSet::new(), receiver_count: 1, @@ -412,6 +419,9 @@ fn advance_head>(shared: &Shared) { } } shared.buffer.sync_head(h); + if shared.producer_waiting.load(Ordering::Acquire) != 0 { + notify_blocking(shared); + } } /// Consumes the message at `cursor` and advances the cursor. @@ -616,3 +626,10 @@ pub fn commit_discard(head: &AtomicU64, tail: &AtomicU64, next: u64) { head.store(next, Ordering::Release); tail.store(next, Ordering::Release); } + +/// Wakes native-thread waiters parked in `recv_blocking` / `send_blocking`. +pub fn notify_blocking(shared: &Shared) { + let mut epoch = shared.blocking.lock(); + *epoch = epoch.wrapping_add(1); + shared.blocking_cvar.notify_all(); +} diff --git a/benchmarks/ecosystem/broadcast/spmc/adapters.rs b/benchmarks/ecosystem/broadcast/spmc/adapters.rs index e5405628..707c2e89 100644 --- a/benchmarks/ecosystem/broadcast/spmc/adapters.rs +++ b/benchmarks/ecosystem/broadcast/spmc/adapters.rs @@ -222,6 +222,14 @@ impl BoundedBroadcastSpmc for Asyncband { receiver.recv().await.unwrap() } + fn send_blocking(sender: &mut Self::Sender, value: usize) { + sender.send_blocking(value); + } + + fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + receiver.recv_blocking().unwrap() + } + fn try_recv(receiver: &mut Self::Receiver) -> Option { match receiver.try_recv() { Ok(value) => Some(value), diff --git a/tests-integration/tests/broadcast_spmc_bounded_test.rs b/tests-integration/tests/broadcast_spmc_bounded_test.rs index e6e40f35..bf108dbf 100644 --- a/tests-integration/tests/broadcast_spmc_bounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_bounded_test.rs @@ -221,7 +221,7 @@ fn concurrent_receivers_keep_up_with_the_producer() { thread::spawn(move || { let mut sum = 0u64; for _ in 0..MESSAGES { - sum += FutureExt::block_on(rx.recv()).unwrap(); + sum += rx.recv_blocking().unwrap(); } sum }) @@ -229,7 +229,7 @@ fn concurrent_receivers_keep_up_with_the_producer() { .collect(); for value in 0..MESSAGES as u64 { - FutureExt::block_on(tx.send(value)); + tx.send_blocking(value); } for handle in handles { @@ -237,6 +237,32 @@ fn concurrent_receivers_keep_up_with_the_producer() { } } +#[test] +fn concurrent_blocking_wait_at_capacity_one() { + let (mut tx, rx) = bounded(1); + let mut rx2 = tx.subscribe(); + let handle = thread::spawn(move || { + let mut sum = 0u64; + for _ in 0..64 { + sum += rx2.recv_blocking().unwrap(); + } + sum + }); + let handle1 = thread::spawn(move || { + let mut sum = 0u64; + let mut rx = rx; + for _ in 0..64 { + sum += rx.recv_blocking().unwrap(); + } + sum + }); + for value in 0..64u64 { + tx.send_blocking(value); + } + assert_eq!(handle.join().unwrap(), 64 * 63 / 2); + assert_eq!(handle1.join().unwrap(), 64 * 63 / 2); +} + #[test] fn publish_order_is_program_order() { let (mut tx, mut rx1) = bounded(8); From b7799449cf6fe4bf81975af36207dbc4f5b526af Mon Sep 17 00:00:00 2001 From: onenewcode Date: Mon, 14 Sep 2026 17:34:47 +0800 Subject: [PATCH 4/9] perf(broadcast): publish unbounded SPMC without the waiter lock Unbounded send writes the slot and publishes tail without the waiter mutex when receivers exist. Fan-out at 32 receivers is now under Tokio's lossy ring. --- asyncband/src/broadcast/spmc/bounded/mod.rs | 42 +++++++----- asyncband/src/broadcast/spmc/common.rs | 68 +++++++++++++++---- asyncband/src/broadcast/spmc/unbounded/mod.rs | 43 ++++++++---- 3 files changed, 111 insertions(+), 42 deletions(-) diff --git a/asyncband/src/broadcast/spmc/bounded/mod.rs b/asyncband/src/broadcast/spmc/bounded/mod.rs index 7bfa8966..857ffa54 100644 --- a/asyncband/src/broadcast/spmc/bounded/mod.rs +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -357,24 +357,27 @@ impl BoundedSender { value = returned; let shared = self.shared.clone(); shared.producer_waiting.store(1, Ordering::Release); - let mut epoch = shared.blocking.lock(); + let snap = shared.epoch.load(Ordering::Acquire); let cap = shared.buffer.cap as u64; if shared.tail.load(Ordering::Acquire) - shared.head.load(Ordering::Acquire) < cap { - drop(epoch); continue; } - let snap = *epoch; - while *epoch == snap - && shared.tail.load(Ordering::Acquire) - shared.head.load(Ordering::Acquire) - >= cap + let guard = shared.blocking.lock(); + if shared.epoch.load(Ordering::Acquire) != snap + || shared.tail.load(Ordering::Acquire) - shared.head.load(Ordering::Acquire) + < cap { - epoch = shared - .blocking_cvar - .wait(epoch) - .unwrap_or_else(|e| e.into_inner()); + drop(guard); + continue; } + drop( + shared + .blocking_cvar + .wait(guard) + .unwrap_or_else(|e| e.into_inner()), + ); } } } @@ -597,20 +600,25 @@ impl BoundedReceiver { Err(TryRecvError::Disconnected) => return Err(RecvError::Disconnected), Err(TryRecvError::Empty) => { let shared = self.shared.clone(); - let mut epoch = shared.blocking.lock(); + let snap = shared.epoch.load(Ordering::Acquire); if self.cursor < shared.tail.load(Ordering::Acquire) { continue; } if shared.senders.load(Ordering::Acquire) == 0 { return Err(RecvError::Disconnected); } - let snap = *epoch; - while *epoch == snap && self.cursor >= shared.tail.load(Ordering::Acquire) { - epoch = shared - .blocking_cvar - .wait(epoch) - .unwrap_or_else(|e| e.into_inner()); + let guard = shared.blocking.lock(); + if shared.epoch.load(Ordering::Acquire) != snap + || self.cursor < shared.tail.load(Ordering::Acquire) + { + continue; } + drop( + shared + .blocking_cvar + .wait(guard) + .unwrap_or_else(|e| e.into_inner()), + ); } } } diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs index ce196795..53daff69 100644 --- a/asyncband/src/broadcast/spmc/common.rs +++ b/asyncband/src/broadcast/spmc/common.rs @@ -20,11 +20,11 @@ //! //! The producer is unique (`send` takes `&mut self`). Receivers drain already-published slots //! without taking the waiter mutex: each slot carries a remaining-reader count, and each -//! subscription keeps its cursor locally. The mutex is only for subscribe/unsubscribe, parking, -//! and the producer's publish-and-drain critical section — which is what keeps a park from missing -//! a wake-up. +//! subscription keeps its cursor locally. The mutex is for subscribe/unsubscribe and parking. +//! Unbounded send publishes without it and drains waiters only when a receiver has parked. use std::cell::UnsafeCell; +use std::hint; use std::mem::MaybeUninit; use std::ptr; use std::sync::Condvar; @@ -45,7 +45,7 @@ use crate::internal::wakerset::WakerSet; use crate::internal::wakerset::WakerToken; /// Number of slots in one unbounded log chunk. -pub const CHUNK_LEN: usize = 64; +pub const CHUNK_LEN: usize = 256; /// A received value together with whether this receive freed a retained slot. /// @@ -168,10 +168,15 @@ pub struct Shared { pub tail: AtomicU64, pub senders: AtomicUsize, pub producer_waiting: AtomicUsize, + /// Set around an unbounded lock-free publish so subscribe spins until `tail` is stable. + pub send_in_progress: AtomicBool, + pub receiver_count: AtomicUsize, + pub has_waiters: AtomicBool, pub state: Mutex, /// Native-thread waiters. Publish and reclaim notify this condvar so a blocking receive does /// not go through one async waker per parked task. - pub blocking: Mutex, + pub epoch: AtomicU64, + pub blocking: Mutex<()>, pub blocking_cvar: Condvar, } @@ -183,7 +188,11 @@ impl Shared { tail: AtomicU64::new(0), senders: AtomicUsize::new(1), producer_waiting: AtomicUsize::new(0), - blocking: Mutex::new(0), + send_in_progress: AtomicBool::new(false), + receiver_count: AtomicUsize::new(1), + has_waiters: AtomicBool::new(false), + epoch: AtomicU64::new(0), + blocking: Mutex::new(()), blocking_cvar: Condvar::new(), state: Mutex::new(State { waiters: WakerSet::new(), @@ -211,11 +220,26 @@ impl Shared { .expect("broadcast channel version counter overflowed") } + /// Waits until an unbounded lock-free send is not mid-publish, then takes `state`. + fn lock_idle_send(&self) -> std::sync::MutexGuard<'_, State> { + loop { + while self.send_in_progress.load(Ordering::Acquire) { + hint::spin_loop(); + } + let state = self.state.lock(); + if !self.send_in_progress.load(Ordering::Acquire) { + return state; + } + } + } + /// Registers a new subscription at the committed tail. pub fn subscribe(&self) -> u64 { - let mut state = self.state.lock(); + let mut state = self.lock_idle_send(); state.receiver_count += 1; - self.tail.load(Ordering::Relaxed) + self.receiver_count + .store(state.receiver_count, Ordering::Release); + self.tail.load(Ordering::Acquire) } #[cfg(test)] @@ -527,8 +551,11 @@ pub fn drop_subscription>( shared: &Shared, cursor: u64, ) -> (Reclaimed, Option) { - let mut state = shared.state.lock(); + let mut state = shared.lock_idle_send(); state.receiver_count -= 1; + shared + .receiver_count + .store(state.receiver_count, Ordering::Release); let last = state.receiver_count == 0; let tail = shared.tail.load(Ordering::Acquire); @@ -586,8 +613,8 @@ pub fn try_receive>( /// The one poll step behind `recv` on both channels. /// -/// The ready path does not take the waiter mutex. Parking rechecks `tail` under that mutex so a -/// publish that drains waiters cannot slip between the empty check and the registration. +/// The ready path does not take the waiter mutex. Parking stores `has_waiters` and then rechecks +/// `tail` so an unbounded send that published without this lock cannot leave the receiver parked. pub fn poll_receive>( shared: &Shared, cursor: &mut u64, @@ -611,6 +638,21 @@ pub fn poll_receive>( } let retired_waker = state.waiters.register(token, cx.waker()); + shared.has_waiters.store(true, Ordering::Release); + if *cursor < shared.tail.load(Ordering::Acquire) { + let waker = state.waiters.unregister(token); + drop(state); + drop(retired_waker); + drop(waker); + return Poll::Ready(Ok(consume(shared, cursor))); + } + if shared.senders.load(Ordering::Acquire) == 0 { + let waker = state.waiters.unregister(token); + drop(state); + drop(retired_waker); + drop(waker); + return Poll::Ready(Err(RecvError::Disconnected)); + } drop(state); drop(retired_waker); Poll::Pending @@ -629,7 +671,7 @@ pub fn commit_discard(head: &AtomicU64, tail: &AtomicU64, next: u64) { /// Wakes native-thread waiters parked in `recv_blocking` / `send_blocking`. pub fn notify_blocking(shared: &Shared) { - let mut epoch = shared.blocking.lock(); - *epoch = epoch.wrapping_add(1); + let _guard = shared.blocking.lock(); + shared.epoch.fetch_add(1, Ordering::Release); shared.blocking_cvar.notify_all(); } diff --git a/asyncband/src/broadcast/spmc/unbounded/mod.rs b/asyncband/src/broadcast/spmc/unbounded/mod.rs index 9c059eb2..afdc855c 100644 --- a/asyncband/src/broadcast/spmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/spmc/unbounded/mod.rs @@ -143,30 +143,49 @@ impl UnboundedSender { /// assert_eq!(second.try_recv(), Ok("update")); /// ``` pub fn send(&mut self, msg: T) { - // 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 (unretained, wakers) = { - let mut state = self.shared.state.lock(); + self.shared.send_in_progress.store(true, Ordering::Release); + let n = self.shared.receiver_count.load(Ordering::Acquire); + + let unretained = if n == 0 { + let state = self.shared.state.lock(); let tail = self.shared.tail.load(Ordering::Relaxed); let next = Shared::>::next_tail(tail); - let unretained = if state.receiver_count == 0 { + if state.receiver_count == 0 { common::commit_discard(&self.shared.head, &self.shared.tail, next); + drop(state); + self.shared.send_in_progress.store(false, Ordering::Release); Some(msg) } else { - let n = state.receiver_count; let slot = self.shared.buffer.slot_for_publish(tail); unsafe { - slot.write(msg, n); + slot.write(msg, state.receiver_count); } common::commit_publish(&self.shared.tail, next); + drop(state); + self.shared.send_in_progress.store(false, Ordering::Release); None - }; - (unretained, state.waiters.drain()) + } + } else { + let tail = self.shared.tail.load(Ordering::Relaxed); + let next = Shared::>::next_tail(tail); + let slot = self.shared.buffer.slot_for_publish(tail); + unsafe { + slot.write(msg, n); + } + common::commit_publish(&self.shared.tail, next); + self.shared.send_in_progress.store(false, Ordering::Release); + None }; - // Notify all waiting receivers. An unsent message is dropped here too, once the lock is - // released. - wake_all(wakers); + if self.shared.has_waiters.load(Ordering::Acquire) { + let wakers = { + let mut state = self.shared.state.lock(); + let wakers = state.waiters.drain(); + self.shared.has_waiters.store(false, Ordering::Release); + wakers + }; + wake_all(wakers); + } drop(unretained); } From a5697089354eb730fe325385cb229d1c8760dbf6 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Mon, 14 Sep 2026 17:47:12 +0800 Subject: [PATCH 5/9] fix(broadcast): wake blocking SPMC receivers when the sender drops disconnect() notifies the shared condvar, and recv_blocking rechecks disconnection after taking that lock so a parked native-thread receive cannot hang. --- asyncband/src/broadcast/spmc/bounded/mod.rs | 1 + asyncband/src/broadcast/spmc/common.rs | 1 + .../tests/broadcast_spmc_bounded_test.rs | 13 +++++++++++++ 3 files changed, 15 insertions(+) diff --git a/asyncband/src/broadcast/spmc/bounded/mod.rs b/asyncband/src/broadcast/spmc/bounded/mod.rs index 857ffa54..48285335 100644 --- a/asyncband/src/broadcast/spmc/bounded/mod.rs +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -610,6 +610,7 @@ impl BoundedReceiver { let guard = shared.blocking.lock(); if shared.epoch.load(Ordering::Acquire) != snap || self.cursor < shared.tail.load(Ordering::Acquire) + || shared.senders.load(Ordering::Acquire) == 0 { continue; } diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs index 53daff69..db0b168c 100644 --- a/asyncband/src/broadcast/spmc/common.rs +++ b/asyncband/src/broadcast/spmc/common.rs @@ -581,6 +581,7 @@ pub fn disconnect(shared: &Shared) { state.waiters.take_all() }; wake_all(wakers); + notify_blocking(shared); } /// Releases a cancelled receive's waker registration, dropping the waker unlocked. diff --git a/tests-integration/tests/broadcast_spmc_bounded_test.rs b/tests-integration/tests/broadcast_spmc_bounded_test.rs index bf108dbf..3ebb9635 100644 --- a/tests-integration/tests/broadcast_spmc_bounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_bounded_test.rs @@ -583,6 +583,19 @@ fn parked_recv_wakes_when_the_sender_drops() { ); } +#[test] +fn parked_recv_blocking_wakes_when_the_sender_drops() { + assert_completes_without_deadlock(|| { + let (tx, mut rx) = bounded::(4); + let parked = thread::spawn(move || rx.recv_blocking()); + // The worker parks on the empty channel. Dropping the sender must finish that receive + // with Disconnected; a missed condvar wake hangs inside assert_completes_without_deadlock. + thread::sleep(std::time::Duration::from_millis(50)); + drop(tx); + assert_eq!(parked.join().unwrap(), Err(RecvError::Disconnected)); + }); +} + // --------------------------------------------------------------------------------------------- // Panic safety // --------------------------------------------------------------------------------------------- From 820427b4792d35b5f2d61877a14a85dc0d42e39b Mon Sep 17 00:00:00 2001 From: onenewcode Date: Mon, 14 Sep 2026 18:04:15 +0800 Subject: [PATCH 6/9] refactor(broadcast): flatten SPMC send and blocking waits Drop the extra wait helper. Clear send_in_progress once, park without cloning the channel Arc, and recover condvar poison the same way as Mutex. --- asyncband/src/broadcast/spmc/bounded/mod.rs | 42 +++++++++---------- asyncband/src/broadcast/spmc/common.rs | 28 ++++++------- asyncband/src/broadcast/spmc/unbounded/mod.rs | 18 ++++---- 3 files changed, 42 insertions(+), 46 deletions(-) diff --git a/asyncband/src/broadcast/spmc/bounded/mod.rs b/asyncband/src/broadcast/spmc/bounded/mod.rs index 48285335..195537b8 100644 --- a/asyncband/src/broadcast/spmc/bounded/mod.rs +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -106,6 +106,7 @@ use std::future::Future; use std::future::poll_fn; use std::pin::Pin; use std::sync::Arc; +use std::sync::PoisonError; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; @@ -355,28 +356,28 @@ impl BoundedSender { Ok(()) => return, Err(TrySendError::Full(returned)) => { value = returned; - let shared = self.shared.clone(); - shared.producer_waiting.store(1, Ordering::Release); - let snap = shared.epoch.load(Ordering::Acquire); - let cap = shared.buffer.cap as u64; - if shared.tail.load(Ordering::Acquire) - shared.head.load(Ordering::Acquire) + self.shared.producer_waiting.store(1, Ordering::Release); + let snap = self.shared.epoch.load(Ordering::Acquire); + let cap = self.shared.buffer.cap as u64; + if self.shared.tail.load(Ordering::Acquire) + - self.shared.head.load(Ordering::Acquire) < cap { continue; } - let guard = shared.blocking.lock(); - if shared.epoch.load(Ordering::Acquire) != snap - || shared.tail.load(Ordering::Acquire) - shared.head.load(Ordering::Acquire) + let guard = self.shared.blocking.lock(); + if self.shared.epoch.load(Ordering::Acquire) != snap + || self.shared.tail.load(Ordering::Acquire) + - self.shared.head.load(Ordering::Acquire) < cap { - drop(guard); continue; } drop( - shared + self.shared .blocking_cvar .wait(guard) - .unwrap_or_else(|e| e.into_inner()), + .unwrap_or_else(PoisonError::into_inner), ); } } @@ -599,26 +600,25 @@ impl BoundedReceiver { Ok(value) => return Ok(value), Err(TryRecvError::Disconnected) => return Err(RecvError::Disconnected), Err(TryRecvError::Empty) => { - let shared = self.shared.clone(); - let snap = shared.epoch.load(Ordering::Acquire); - if self.cursor < shared.tail.load(Ordering::Acquire) { + let snap = self.shared.epoch.load(Ordering::Acquire); + if self.cursor < self.shared.tail.load(Ordering::Acquire) { continue; } - if shared.senders.load(Ordering::Acquire) == 0 { + if self.shared.senders.load(Ordering::Acquire) == 0 { return Err(RecvError::Disconnected); } - let guard = shared.blocking.lock(); - if shared.epoch.load(Ordering::Acquire) != snap - || self.cursor < shared.tail.load(Ordering::Acquire) - || shared.senders.load(Ordering::Acquire) == 0 + let guard = self.shared.blocking.lock(); + if self.shared.epoch.load(Ordering::Acquire) != snap + || self.cursor < self.shared.tail.load(Ordering::Acquire) + || self.shared.senders.load(Ordering::Acquire) == 0 { continue; } drop( - shared + self.shared .blocking_cvar .wait(guard) - .unwrap_or_else(|e| e.into_inner()), + .unwrap_or_else(PoisonError::into_inner), ); } } diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs index db0b168c..762b506d 100644 --- a/asyncband/src/broadcast/spmc/common.rs +++ b/asyncband/src/broadcast/spmc/common.rs @@ -28,6 +28,7 @@ use std::hint; use std::mem::MaybeUninit; use std::ptr; use std::sync::Condvar; +use std::sync::MutexGuard; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicPtr; use std::sync::atomic::AtomicU64; @@ -221,7 +222,7 @@ impl Shared { } /// Waits until an unbounded lock-free send is not mid-publish, then takes `state`. - fn lock_idle_send(&self) -> std::sync::MutexGuard<'_, State> { + fn lock_idle_send(&self) -> MutexGuard<'_, State> { loop { while self.send_in_progress.load(Ordering::Acquire) { hint::spin_loop(); @@ -640,26 +641,25 @@ pub fn poll_receive>( let retired_waker = state.waiters.register(token, cx.waker()); shared.has_waiters.store(true, Ordering::Release); - if *cursor < shared.tail.load(Ordering::Acquire) { - let waker = state.waiters.unregister(token); - drop(state); - drop(retired_waker); - drop(waker); - return Poll::Ready(Ok(consume(shared, cursor))); - } - if shared.senders.load(Ordering::Acquire) == 0 { - let waker = state.waiters.unregister(token); + if *cursor >= shared.tail.load(Ordering::Acquire) && shared.senders.load(Ordering::Acquire) != 0 + { drop(state); drop(retired_waker); - drop(waker); - return Poll::Ready(Err(RecvError::Disconnected)); + return Poll::Pending; } + + let waker = state.waiters.unregister(token); drop(state); drop(retired_waker); - Poll::Pending + drop(waker); + if *cursor < shared.tail.load(Ordering::Acquire) { + Poll::Ready(Ok(consume(shared, cursor))) + } else { + Poll::Ready(Err(RecvError::Disconnected)) + } } -/// Publishes `tail` after writing a slot. Caller holds `state` and drains waiters after this. +/// Publishes `tail` after writing a slot. pub fn commit_publish(tail: &AtomicU64, next: u64) { tail.store(next, Ordering::Release); } diff --git a/asyncband/src/broadcast/spmc/unbounded/mod.rs b/asyncband/src/broadcast/spmc/unbounded/mod.rs index afdc855c..f2b992fa 100644 --- a/asyncband/src/broadcast/spmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/spmc/unbounded/mod.rs @@ -152,37 +152,33 @@ impl UnboundedSender { let next = Shared::>::next_tail(tail); if state.receiver_count == 0 { common::commit_discard(&self.shared.head, &self.shared.tail, next); - drop(state); - self.shared.send_in_progress.store(false, Ordering::Release); Some(msg) } else { - let slot = self.shared.buffer.slot_for_publish(tail); unsafe { - slot.write(msg, state.receiver_count); + self.shared + .buffer + .slot_for_publish(tail) + .write(msg, state.receiver_count); } common::commit_publish(&self.shared.tail, next); - drop(state); - self.shared.send_in_progress.store(false, Ordering::Release); None } } else { let tail = self.shared.tail.load(Ordering::Relaxed); let next = Shared::>::next_tail(tail); - let slot = self.shared.buffer.slot_for_publish(tail); unsafe { - slot.write(msg, n); + self.shared.buffer.slot_for_publish(tail).write(msg, n); } common::commit_publish(&self.shared.tail, next); - self.shared.send_in_progress.store(false, Ordering::Release); None }; + self.shared.send_in_progress.store(false, Ordering::Release); if self.shared.has_waiters.load(Ordering::Acquire) { let wakers = { let mut state = self.shared.state.lock(); - let wakers = state.waiters.drain(); self.shared.has_waiters.store(false, Ordering::Release); - wakers + state.waiters.drain() }; wake_all(wakers); } From a2081b9017d153d3c888611dc5a192a488e8547f Mon Sep 17 00:00:00 2001 From: onenewcode Date: Tue, 15 Sep 2026 11:57:59 +0800 Subject: [PATCH 7/9] test(broadcast): cover exclusive send and subscription delivery Keep slot-log types family-private. Compile-fail examples show the sender is not Clone and publish takes &mut self. Cover unbounded fan-out, late subscribe, and slow-subscription growth. --- asyncband/src/broadcast/spmc/bounded/mod.rs | 14 ++++ asyncband/src/broadcast/spmc/common.rs | 78 ++++++++++--------- asyncband/src/broadcast/spmc/mod.rs | 4 +- asyncband/src/broadcast/spmc/unbounded/mod.rs | 14 ++++ .../tests/broadcast_spmc_bounded_test.rs | 15 ++++ .../tests/broadcast_spmc_unbounded_test.rs | 58 ++++++++++++++ 6 files changed, 147 insertions(+), 36 deletions(-) diff --git a/asyncband/src/broadcast/spmc/bounded/mod.rs b/asyncband/src/broadcast/spmc/bounded/mod.rs index 195537b8..c544da53 100644 --- a/asyncband/src/broadcast/spmc/bounded/mod.rs +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -163,6 +163,20 @@ pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver< /// /// This handle is not [`Clone`]. Dropping it disconnects the channel. Each receiver may drain its /// own buffered messages before observing disconnection. +/// +/// ```compile_fail +/// use asyncband::broadcast::spmc::bounded; +/// +/// let (tx, _rx) = bounded::(1); +/// let _ = tx.clone(); +/// ``` +/// +/// ```compile_fail +/// use asyncband::broadcast::spmc::bounded; +/// +/// let (tx, _rx) = bounded::(1); +/// tx.try_send(1); +/// ``` pub struct BoundedSender { shared: Arc>>, } diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs index 762b506d..41844706 100644 --- a/asyncband/src/broadcast/spmc/common.rs +++ b/asyncband/src/broadcast/spmc/common.rs @@ -18,6 +18,11 @@ //! Slot log, waiters, and the receive step shared by the bounded and unbounded SPMC broadcast //! channels. //! +//! This is family-private retention and sequencing machinery. The public retention contract is +//! the same as [`crate::broadcast::mpmc`]: lossless fan-out, bounded wait-at-capacity or unbounded +//! growth, and new subscriptions joining at the committed tail. Ring slots, chunk storage, and the +//! `head` / `tail` sequencer stay in this module. +//! //! The producer is unique (`send` takes `&mut self`). Receivers drain already-published slots //! without taking the waiter mutex: each slot carries a remaining-reader count, and each //! subscription keeps its cursor locally. The mutex is for subscribe/unsubscribe and parking. @@ -46,13 +51,13 @@ use crate::internal::wakerset::WakerSet; use crate::internal::wakerset::WakerToken; /// Number of slots in one unbounded log chunk. -pub const CHUNK_LEN: usize = 256; +pub(super) const CHUNK_LEN: usize = 256; /// A received value together with whether this receive freed a retained slot. /// /// A bounded channel wakes the producer from `reclaimed` before it returns `value`, so a panicking /// `T::clone` cannot strand a waiter on capacity this receive already released. -pub struct Consumed { +pub(super) struct Consumed { pub value: T, pub reclaimed: bool, } @@ -61,7 +66,7 @@ pub struct Consumed { /// /// 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 { +pub(super) struct Reclaimed { first: Option, rest: Vec, } @@ -74,7 +79,7 @@ impl Reclaimed { } } - pub fn is_empty(&self) -> bool { + pub(super) fn is_empty(&self) -> bool { self.first.is_none() } @@ -88,7 +93,7 @@ impl Reclaimed { } /// One published value and the number of subscriptions that still have to consume it. -pub struct Slot { +pub(super) struct Slot { msg: UnsafeCell>, remaining: AtomicUsize, /// `true` once the producer has written `msg` and until the last remaining reader takes it. @@ -121,7 +126,7 @@ impl Slot { /// /// `occupied` must be `false`. `n` must be the number of subscriptions that will consume this /// version, and must be greater than zero. - pub unsafe fn write(&self, msg: T, n: usize) { + pub(super) unsafe fn write(&self, msg: T, n: usize) { debug_assert!(n > 0); debug_assert!(!self.occupied.load(Ordering::Relaxed)); unsafe { @@ -156,14 +161,14 @@ impl Slot { } /// Parked receivers, the live subscription count, and the single parked producer. -pub struct State { +pub(super) struct State { pub waiters: WakerSet, pub receiver_count: usize, pub producer: Option, } /// Shared channel state: the slot log plus the waiter mutex. -pub struct Shared { +pub(super) struct Shared { pub buffer: B, pub head: AtomicU64, pub tail: AtomicU64, @@ -182,7 +187,7 @@ pub struct Shared { } impl Shared { - pub fn new(buffer: B) -> Self { + pub(super) fn new(buffer: B) -> Self { Self { buffer, head: AtomicU64::new(0), @@ -203,20 +208,20 @@ impl Shared { } } - pub fn retained(&self) -> usize { + pub(super) fn retained(&self) -> usize { let tail = self.tail.load(Ordering::Acquire); let head = self.head.load(Ordering::Acquire); usize::try_from(tail - head).expect("retained broadcast message count exceeds usize") } - pub fn unread(&self, cursor: u64) -> usize { + pub(super) fn unread(&self, cursor: u64) -> usize { let tail = self.tail.load(Ordering::Acquire); debug_assert!(tail >= cursor); usize::try_from(tail - cursor).expect("unread broadcast message count exceeds usize") } /// Next committed version, panicking on overflow. - pub fn next_tail(tail: u64) -> u64 { + pub(super) fn next_tail(tail: u64) -> u64 { tail.checked_add(1) .expect("broadcast channel version counter overflowed") } @@ -235,7 +240,7 @@ impl Shared { } /// Registers a new subscription at the committed tail. - pub fn subscribe(&self) -> u64 { + pub(super) fn subscribe(&self) -> u64 { let mut state = self.lock_idle_send(); state.receiver_count += 1; self.receiver_count @@ -244,14 +249,14 @@ impl Shared { } #[cfg(test)] - pub fn set_tail(&self, tail: u64) { + pub(super) fn set_tail(&self, tail: u64) { self.tail.store(tail, Ordering::Relaxed); self.head.store(tail, Ordering::Relaxed); } } /// A random-access published slot, addressed by its committed version. -pub trait SlotStore { +pub(super) trait SlotStore { fn slot(&self, version: u64) -> &Slot; /// Moves the lookup start past chunks the live window has left behind. @@ -263,13 +268,13 @@ pub trait SlotStore { } /// Fixed ring used by the bounded channel. -pub struct BoundedBuffer { +pub(super) struct BoundedBuffer { slots: Box<[Slot]>, pub cap: usize, } impl BoundedBuffer { - pub fn new(capacity: usize) -> Self { + pub(super) fn new(capacity: usize) -> Self { Self { slots: (0..capacity).map(|_| Slot::empty()).collect(), cap: capacity, @@ -277,7 +282,7 @@ impl BoundedBuffer { } #[cfg(test)] - pub fn len(&self) -> usize { + pub(super) fn len(&self) -> usize { self.slots.len() } } @@ -300,7 +305,7 @@ impl SlotStore for BoundedBuffer { } /// One growable segment of the unbounded log. -pub struct Chunk { +pub(super) struct Chunk { slots: [Slot; CHUNK_LEN], next: AtomicPtr>, base: u64, @@ -319,8 +324,8 @@ impl Chunk { /// Linked chunks used by the unbounded channel. /// /// The producer appends chunks without moving earlier slots, so receivers can drain without a -/// publication lock. Fully consumed chunks are recycled onto the sender's spare list. -pub struct UnboundedBuffer { +/// publication lock. Fully consumed chunks stay allocated until the channel is dropped. +pub(super) struct UnboundedBuffer { /// First chunk ever allocated. Never moves; `Drop` walks from here. root: AtomicPtr>, /// First chunk that may still hold a live message. Lookup starts here. @@ -330,7 +335,7 @@ pub struct UnboundedBuffer { } impl UnboundedBuffer { - pub fn new() -> Self { + pub(super) fn new() -> Self { let chunk = Box::into_raw(Chunk::new(0)); Self { root: AtomicPtr::new(chunk), @@ -345,7 +350,7 @@ impl UnboundedBuffer { /// The caller is the unique producer and must not publish `tail` past this version until this /// returns. Fully consumed chunks stay allocated until the channel is dropped so a receiver /// walking the list cannot observe a freed chunk. - pub fn slot_for_publish(&self, version: u64) -> &Slot { + pub(super) fn slot_for_publish(&self, version: u64) -> &Slot { loop { let chunk = self.tail_chunk.load(Ordering::Acquire); debug_assert!(!chunk.is_null()); @@ -378,7 +383,7 @@ impl UnboundedBuffer { } #[cfg(test)] - pub fn allocated_slots(&self) -> usize { + pub(super) fn allocated_slots(&self) -> usize { let mut n = 0; let mut chunk = self.root.load(Ordering::Acquire); while !chunk.is_null() { @@ -455,7 +460,10 @@ fn advance_head>(shared: &Shared) { /// other reader clones `T` with the waiter mutex not held. The cursor advances before that clone /// so a panicking `T::clone` still consumes the subscription's claim; a drop guard releases the /// remaining count if the clone unwinds. -pub fn consume>(shared: &Shared, cursor: &mut u64) -> Consumed { +pub(super) fn consume>( + shared: &Shared, + cursor: &mut u64, +) -> Consumed { let version = *cursor; debug_assert!(version < shared.tail.load(Ordering::Acquire)); let slot = shared.buffer.slot(version); @@ -524,7 +532,7 @@ fn reclaim_range>(shared: &Shared, start: u64, end: u64) - } /// Takes the parked producer if a reclaim may have freed capacity. -pub fn take_producer_on_reclaim( +pub(super) fn take_producer_on_reclaim( shared: &Shared, reclaimed: bool, drained_last: bool, @@ -541,14 +549,14 @@ pub fn take_producer_on_reclaim( } /// Wakes the parked producer, if any, with the channel already unlocked. -pub fn wake_producer(producer: Option) { +pub(super) fn wake_producer(producer: Option) { if let Some(waker) = producer { waker.wake(); } } /// Drops a subscription, reclaiming every unread slot it still held. -pub fn drop_subscription>( +pub(super) fn drop_subscription>( shared: &Shared, cursor: u64, ) -> (Reclaimed, Option) { @@ -576,7 +584,7 @@ pub fn drop_subscription>( } /// Wakes every parked receiver so it can observe the channel's disconnected state. -pub fn disconnect(shared: &Shared) { +pub(super) fn disconnect(shared: &Shared) { let wakers = { let mut state = shared.state.lock(); state.waiters.take_all() @@ -586,7 +594,7 @@ pub fn disconnect(shared: &Shared) { } /// Releases a cancelled receive's waker registration, dropping the waker unlocked. -pub fn unregister(shared: &Shared, cursor: u64, token: &mut Option) { +pub(super) fn unregister(shared: &Shared, cursor: u64, token: &mut Option) { let mut state = shared.state.lock(); if cursor < shared.tail.load(Ordering::Relaxed) || shared.senders.load(Ordering::Acquire) == 0 { *token = None; @@ -599,7 +607,7 @@ pub fn unregister(shared: &Shared, cursor: u64, token: &mut Option>( +pub(super) fn try_receive>( shared: &Shared, cursor: &mut u64, ) -> Result, TryRecvError> { @@ -617,7 +625,7 @@ pub fn try_receive>( /// /// The ready path does not take the waiter mutex. Parking stores `has_waiters` and then rechecks /// `tail` so an unbounded send that published without this lock cannot leave the receiver parked. -pub fn poll_receive>( +pub(super) fn poll_receive>( shared: &Shared, cursor: &mut u64, token: &mut Option, @@ -660,18 +668,18 @@ pub fn poll_receive>( } /// Publishes `tail` after writing a slot. -pub fn commit_publish(tail: &AtomicU64, next: u64) { +pub(super) fn commit_publish(tail: &AtomicU64, next: u64) { tail.store(next, Ordering::Release); } /// Advances `head` and `tail` together when nothing can read the message. -pub fn commit_discard(head: &AtomicU64, tail: &AtomicU64, next: u64) { +pub(super) fn commit_discard(head: &AtomicU64, tail: &AtomicU64, next: u64) { head.store(next, Ordering::Release); tail.store(next, Ordering::Release); } /// Wakes native-thread waiters parked in `recv_blocking` / `send_blocking`. -pub fn notify_blocking(shared: &Shared) { +pub(super) fn notify_blocking(shared: &Shared) { let _guard = shared.blocking.lock(); shared.epoch.fetch_add(1, Ordering::Release); shared.blocking_cvar.notify_all(); diff --git a/asyncband/src/broadcast/spmc/mod.rs b/asyncband/src/broadcast/spmc/mod.rs index 7afbddd0..811ff53a 100644 --- a/asyncband/src/broadcast/spmc/mod.rs +++ b/asyncband/src/broadcast/spmc/mod.rs @@ -21,7 +21,9 @@ //! that was active when it was accepted, so a receive never reports lag. They differ in what the //! producer does when the slowest subscription stops reclaiming. [`bounded`] retains at most the //! capacity it was built with and makes the producer wait for that subscription. [`unbounded`] -//! never waits to send and lets the retained backlog grow instead. +//! never waits to send and lets the retained backlog grow instead. That is the same public +//! retention contract as [`crate::broadcast::mpmc`]: no lag or overwrite, and a new subscription +//! starts at the committed tail. //! //! The sender is not [`Clone`], and every publish method takes `&mut self`. That is the static //! single-writer contract this topology adds over [`crate::broadcast::mpmc`]: there cannot be a diff --git a/asyncband/src/broadcast/spmc/unbounded/mod.rs b/asyncband/src/broadcast/spmc/unbounded/mod.rs index f2b992fa..70a26ee7 100644 --- a/asyncband/src/broadcast/spmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/spmc/unbounded/mod.rs @@ -101,6 +101,20 @@ pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { /// /// This handle is not [`Clone`]. Once it is dropped, each receiver can drain the values already /// published for it and then observes disconnection. +/// +/// ```compile_fail +/// use asyncband::broadcast::spmc::unbounded; +/// +/// let (tx, _rx) = unbounded::(); +/// let _ = tx.clone(); +/// ``` +/// +/// ```compile_fail +/// use asyncband::broadcast::spmc::unbounded; +/// +/// let (tx, _rx) = unbounded::(); +/// tx.send(1); +/// ``` pub struct UnboundedSender { shared: Arc>>, } diff --git a/tests-integration/tests/broadcast_spmc_bounded_test.rs b/tests-integration/tests/broadcast_spmc_bounded_test.rs index 3ebb9635..394e722a 100644 --- a/tests-integration/tests/broadcast_spmc_bounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_bounded_test.rs @@ -134,6 +134,20 @@ fn delivers_every_message_to_every_receiver() { assert_eq!(rx2.try_recv(), Ok(20)); } +#[test] +fn send_and_recv_deliver_every_accepted_value() { + let (mut tx, mut rx1) = bounded(4); + let mut rx2 = tx.subscribe(); + + FutureExt::block_on(tx.send(10)); + FutureExt::block_on(tx.send(20)); + + assert_eq!(FutureExt::block_on(rx1.recv()), Ok(10)); + assert_eq!(FutureExt::block_on(rx1.recv()), Ok(20)); + assert_eq!(FutureExt::block_on(rx2.recv()), Ok(10)); + assert_eq!(FutureExt::block_on(rx2.recv()), Ok(20)); +} + #[test] fn slow_receiver_keeps_every_message_under_backpressure() { let (mut tx, mut fast) = bounded(2); @@ -350,6 +364,7 @@ fn send_waits_while_the_slowest_subscription_holds_capacity() { assert_eq!(rx2.try_recv(), Ok(1)); assert!(poll_once(send.as_mut()).is_ready()); assert_eq!(rx1.try_recv(), Ok(2)); + assert_eq!(rx2.try_recv(), Ok(2)); } #[test] diff --git a/tests-integration/tests/broadcast_spmc_unbounded_test.rs b/tests-integration/tests/broadcast_spmc_unbounded_test.rs index a1b53469..04e3cac7 100644 --- a/tests-integration/tests/broadcast_spmc_unbounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_unbounded_test.rs @@ -96,6 +96,64 @@ impl Rng { } } +#[test] +fn delivers_every_message_to_every_receiver() { + let (mut tx, mut rx1) = unbounded(); + let mut rx2 = tx.subscribe(); + + tx.send(10); + tx.send(20); + + assert_eq!(rx1.try_recv(), Ok(10)); + assert_eq!(rx1.try_recv(), Ok(20)); + assert_eq!(rx2.try_recv(), Ok(10)); + assert_eq!(rx2.try_recv(), Ok(20)); +} + +#[tokio::test] +async fn send_and_recv_deliver_every_accepted_value() { + let (mut tx, mut rx1) = unbounded(); + let mut rx2 = tx.subscribe(); + + tx.send(10); + tx.send(20); + + assert_eq!(rx1.recv().await, Ok(10)); + assert_eq!(rx1.recv().await, Ok(20)); + assert_eq!(rx2.recv().await, Ok(10)); + assert_eq!(rx2.recv().await, Ok(20)); +} + +#[test] +fn subscribe_starts_at_the_committed_tail() { + let (mut tx, _rx) = unbounded(); + tx.send(1); + + let mut late = tx.subscribe(); + assert_eq!(late.try_recv(), Err(TryRecvError::Empty)); + + tx.send(2); + assert_eq!(late.try_recv(), Ok(2)); +} + +#[test] +fn send_never_waits_while_a_slow_subscription_lags() { + let (mut tx, mut fast) = unbounded(); + let slow = tx.subscribe(); + + for value in 0..1024 { + tx.send(value); + } + + assert_eq!(tx.retained_message_count(), 1024); + for value in 0..1024 { + assert_eq!(fast.try_recv(), Ok(value)); + } + assert_eq!(tx.retained_message_count(), 1024); + drop(slow); + assert_eq!(tx.retained_message_count(), 0); +} + #[test] fn try_recv_reports_empty_then_value_then_disconnected() { let (mut tx, mut rx) = unbounded(); From a9df3a737ce89319f9c94bf1f927a0d02f18b03e Mon Sep 17 00:00:00 2001 From: onenewcode Date: Wed, 16 Sep 2026 10:35:06 +0800 Subject: [PATCH 8/9] fix(broadcast): serialize SPMC publish with subscribe and parking Hold the waiter mutex across publication so a slot's remaining-reader count matches live subscriptions and a parking receiver cannot miss a send. Release unbounded chunk storage as head advances, and wake the producer if a panicking clone is the last reader. Drop native-thread send_blocking/recv_blocking; FutureExt::block_on covers that path. --- CHANGELOG.md | 2 +- asyncband/src/broadcast/spmc/bounded/mod.rs | 83 +----- asyncband/src/broadcast/spmc/common.rs | 247 +++++++++--------- asyncband/src/broadcast/spmc/mod.rs | 5 +- asyncband/src/broadcast/spmc/unbounded/mod.rs | 44 ++-- .../src/broadcast/spmc/unbounded/tests.rs | 15 +- .../ecosystem/broadcast/spmc/adapters.rs | 8 - .../tests/broadcast_spmc_bounded_test.rs | 77 +++++- .../tests/broadcast_spmc_unbounded_test.rs | 124 +++++++++ 9 files changed, 352 insertions(+), 253 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f0d4f906..aa333808 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,7 +27,7 @@ All notable changes to this project will be documented in this file. * Add `event::AutoResetEvent`, a reusable signal that releases one waiter, retains at most one unassigned signal that can be cleared with `reset`, and transfers assigned signals when waits are cancelled. * Add `ManualResetEvent::try_wait` to check readiness without registering a waiter or consuming the set state. -* Add `broadcast::spmc`, a lossless single-producer broadcast family with a non-cloneable sender whose publish methods require exclusive access; receivers drain already-published slots without taking the publication lock, bounded retains at most the requested capacity and makes the producer wait for the slowest active subscription, unbounded never waits and lets the retained backlog grow, and bounded `send_blocking` / `recv_blocking` park native threads on a shared condvar. +* Add `broadcast::spmc`, a lossless single-producer broadcast family with the same public surface and retention contract as `broadcast::mpmc` but a non-cloneable sender whose publish methods require exclusive access; receivers drain already-published slots without taking the publication lock, bounded retains at most the requested capacity and makes the producer wait for the slowest active subscription, and unbounded never waits and lets the retained backlog grow while releasing the storage the backlog has left behind. * 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`. diff --git a/asyncband/src/broadcast/spmc/bounded/mod.rs b/asyncband/src/broadcast/spmc/bounded/mod.rs index c544da53..988a6f6b 100644 --- a/asyncband/src/broadcast/spmc/bounded/mod.rs +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -106,8 +106,8 @@ use std::future::Future; use std::future::poll_fn; use std::pin::Pin; use std::sync::Arc; -use std::sync::PoisonError; use std::sync::atomic::Ordering; +use std::sync::atomic::fence; use std::task::Context; use std::task::Poll; @@ -267,6 +267,9 @@ impl BoundedSender { .shared .producer_waiting .store(1, Ordering::Release); + // Announce the intent to park before reading `head`. Pairs with the fence in + // `common::take_producer_on_reclaim`. + fence(Ordering::SeqCst); let mut state = self.sender.shared.state.lock(); if state.receiver_count == 0 { @@ -286,7 +289,6 @@ impl BoundedSender { let wakers = state.waiters.drain(); drop(state); wake_all(wakers); - common::notify_blocking(self.sender.shared.as_ref()); drop(retired_producer); drop(msg); return Poll::Ready(()); @@ -321,7 +323,6 @@ impl BoundedSender { // send, so it must not run under the lock, but a panic in its Drop must not skip // the receiver wake-ups either. wake_all(wakers); - common::notify_blocking(self.sender.shared.as_ref()); drop(retired_producer); Poll::Ready(()) } @@ -363,41 +364,6 @@ impl BoundedSender { self.publish(value).map_err(TrySendError::Full) } - /// Broadcasts a value, parking the current thread while the channel is at capacity. - pub fn send_blocking(&mut self, mut value: T) { - loop { - match self.try_send(value) { - Ok(()) => return, - Err(TrySendError::Full(returned)) => { - value = returned; - self.shared.producer_waiting.store(1, Ordering::Release); - let snap = self.shared.epoch.load(Ordering::Acquire); - let cap = self.shared.buffer.cap as u64; - if self.shared.tail.load(Ordering::Acquire) - - self.shared.head.load(Ordering::Acquire) - < cap - { - continue; - } - let guard = self.shared.blocking.lock(); - if self.shared.epoch.load(Ordering::Acquire) != snap - || self.shared.tail.load(Ordering::Acquire) - - self.shared.head.load(Ordering::Acquire) - < cap - { - continue; - } - drop( - self.shared - .blocking_cvar - .wait(guard) - .unwrap_or_else(PoisonError::into_inner), - ); - } - } - } - } - /// The publish step both send paths share. /// /// Publishing and draining the wait set share one critical section, so a receiver can never @@ -435,7 +401,6 @@ impl BoundedSender { }; wake_all(wakers); - common::notify_blocking(self.shared.as_ref()); drop(discarded); Ok(()) } @@ -599,45 +564,10 @@ impl BoundedReceiver { /// ``` pub fn try_recv(&mut self) -> Result { let consumed = common::try_receive(&self.shared, &mut self.cursor)?; - let producer = common::take_producer_on_reclaim(&self.shared, consumed.reclaimed, false); + let producer = common::take_producer_on_reclaim(&self.shared, consumed.reclaimed); common::wake_producer(producer); Ok(consumed.value) } - - /// Receives the next value, parking the current thread while empty. - /// - /// Same result as [`recv`](Self::recv), without going through an async waker. Native-thread - /// waiters share one futex so a publish can wake every blocked receiver with a single notify. - pub fn recv_blocking(&mut self) -> Result { - loop { - match self.try_recv() { - Ok(value) => return Ok(value), - Err(TryRecvError::Disconnected) => return Err(RecvError::Disconnected), - Err(TryRecvError::Empty) => { - let snap = self.shared.epoch.load(Ordering::Acquire); - if self.cursor < self.shared.tail.load(Ordering::Acquire) { - continue; - } - if self.shared.senders.load(Ordering::Acquire) == 0 { - return Err(RecvError::Disconnected); - } - let guard = self.shared.blocking.lock(); - if self.shared.epoch.load(Ordering::Acquire) != snap - || self.cursor < self.shared.tail.load(Ordering::Acquire) - || self.shared.senders.load(Ordering::Acquire) == 0 - { - continue; - } - drop( - self.shared - .blocking_cvar - .wait(guard) - .unwrap_or_else(PoisonError::into_inner), - ); - } - } - } - } } impl BoundedReceiver { @@ -755,8 +685,7 @@ impl Future for Recv<'_, T> { Poll::Ready(Ok(consumed)) => consumed, }; - let producer = - common::take_producer_on_reclaim(&receiver.shared, consumed.reclaimed, false); + let producer = common::take_producer_on_reclaim(&receiver.shared, consumed.reclaimed); common::wake_producer(producer); Poll::Ready(Ok(consumed.value)) } diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs index 41844706..33b80431 100644 --- a/asyncband/src/broadcast/spmc/common.rs +++ b/asyncband/src/broadcast/spmc/common.rs @@ -25,20 +25,19 @@ //! //! The producer is unique (`send` takes `&mut self`). Receivers drain already-published slots //! without taking the waiter mutex: each slot carries a remaining-reader count, and each -//! subscription keeps its cursor locally. The mutex is for subscribe/unsubscribe and parking. -//! Unbounded send publishes without it and drains waiters only when a receiver has parked. +//! subscription keeps its cursor locally. The mutex covers publication, subscribe/unsubscribe, and +//! parking, so a slot's remaining-reader count always matches its consumers and a parking receiver +//! cannot miss a publication. use std::cell::UnsafeCell; -use std::hint; use std::mem::MaybeUninit; use std::ptr; -use std::sync::Condvar; -use std::sync::MutexGuard; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicPtr; use std::sync::atomic::AtomicU64; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::sync::atomic::fence; use std::task::Context; use std::task::Poll; use std::task::Waker; @@ -98,8 +97,8 @@ pub(super) struct Slot { remaining: AtomicUsize, /// `true` once the producer has written `msg` and until the last remaining reader takes it. /// - /// Head only advances over a slot after this is cleared, so the producer cannot reuse the - /// memory while a reader is still cloning `T`. + /// `take_msg` clears it before `head` moves past the slot, so the producer cannot reuse the + /// memory while a reader is still cloning `T`. Teardown also reads it to find leftovers. occupied: AtomicBool, } @@ -174,16 +173,7 @@ pub(super) struct Shared { pub tail: AtomicU64, pub senders: AtomicUsize, pub producer_waiting: AtomicUsize, - /// Set around an unbounded lock-free publish so subscribe spins until `tail` is stable. - pub send_in_progress: AtomicBool, - pub receiver_count: AtomicUsize, - pub has_waiters: AtomicBool, pub state: Mutex, - /// Native-thread waiters. Publish and reclaim notify this condvar so a blocking receive does - /// not go through one async waker per parked task. - pub epoch: AtomicU64, - pub blocking: Mutex<()>, - pub blocking_cvar: Condvar, } impl Shared { @@ -194,12 +184,6 @@ impl Shared { tail: AtomicU64::new(0), senders: AtomicUsize::new(1), producer_waiting: AtomicUsize::new(0), - send_in_progress: AtomicBool::new(false), - receiver_count: AtomicUsize::new(1), - has_waiters: AtomicBool::new(false), - epoch: AtomicU64::new(0), - blocking: Mutex::new(()), - blocking_cvar: Condvar::new(), state: Mutex::new(State { waiters: WakerSet::new(), receiver_count: 1, @@ -226,25 +210,13 @@ impl Shared { .expect("broadcast channel version counter overflowed") } - /// Waits until an unbounded lock-free send is not mid-publish, then takes `state`. - fn lock_idle_send(&self) -> MutexGuard<'_, State> { - loop { - while self.send_in_progress.load(Ordering::Acquire) { - hint::spin_loop(); - } - let state = self.state.lock(); - if !self.send_in_progress.load(Ordering::Acquire) { - return state; - } - } - } - /// Registers a new subscription at the committed tail. + /// + /// Publication holds the same lock, so the returned cursor is this subscription's first + /// counted version. pub(super) fn subscribe(&self) -> u64 { - let mut state = self.lock_idle_send(); + let mut state = self.state.lock(); state.receiver_count += 1; - self.receiver_count - .store(state.receiver_count, Ordering::Release); self.tail.load(Ordering::Acquire) } @@ -259,11 +231,11 @@ impl Shared { pub(super) trait SlotStore { fn slot(&self, version: u64) -> &Slot; - /// Moves the lookup start past chunks the live window has left behind. + /// Moves the lookup start past storage the live window has left behind, releasing it. /// - /// Chunks stay allocated until the channel is dropped, so a receiver that still holds an old - /// pointer cannot observe a free. Skipping them keeps `slot` proportional to the live window - /// rather than to the lifetime message count. + /// Called after `head` advances. A fixed ring has nothing to release; the unbounded log frees + /// the storage of chunks now entirely below `head`, keeping both `slot` and the footprint + /// proportional to the live window rather than to the lifetime message count. fn sync_head(&self, _head: u64) {} } @@ -305,26 +277,63 @@ impl SlotStore for BoundedBuffer { } /// One growable segment of the unbounded log. +/// +/// The header is split from the storage. A lookup walks the `next` chain from +/// [`UnboundedBuffer::head_chunk`], possibly through segments the live window has left behind, so +/// headers live until the channel is dropped; `slots` is released once `head` passes the segment. pub(super) struct Chunk { - slots: [Slot; CHUNK_LEN], + slots: AtomicPtr<[Slot; CHUNK_LEN]>, next: AtomicPtr>, base: u64, } impl Chunk { fn new(base: u64) -> Box { + let slots = Box::into_raw(Box::new(std::array::from_fn(|_| Slot::empty()))); Box::new(Self { - slots: std::array::from_fn(|_| Slot::empty()), + slots: AtomicPtr::new(slots), next: AtomicPtr::new(ptr::null_mut()), base, }) } + + /// Returns this segment's slot for `version`. + /// + /// # Safety + /// + /// `version` must belong to this segment and be at or above the channel's `head`, so the + /// storage cannot have been released. + #[inline] + unsafe fn slot(&self, version: u64) -> &Slot { + let slots = self.slots.load(Ordering::Acquire); + debug_assert!(!slots.is_null()); + unsafe { &(*slots)[(version - self.base) as usize] } + } + + /// Releases this segment's message storage, dropping any message it still holds. + /// + /// Concurrent callers race on one swap, so the storage is freed exactly once. Only teardown + /// finds a message here: `head` passes a slot only after its value was taken, so + /// [`UnboundedBuffer::release_consumed_chunks`] never runs `T::drop`. + fn release(&self) { + let slots = self.slots.swap(ptr::null_mut(), Ordering::AcqRel); + if slots.is_null() { + return; + } + let slots = unsafe { Box::from_raw(slots) }; + for slot in slots.iter() { + if slot.occupied.load(Ordering::Relaxed) { + drop(unsafe { slot.take_msg() }); + } + } + } } /// Linked chunks used by the unbounded channel. /// /// The producer appends chunks without moving earlier slots, so receivers can drain without a -/// publication lock. Fully consumed chunks stay allocated until the channel is dropped. +/// publication lock. Storage is released as `head` advances past a chunk; the headers stay +/// allocated until the channel is dropped. pub(super) struct UnboundedBuffer { /// First chunk ever allocated. Never moves; `Drop` walks from here. root: AtomicPtr>, @@ -348,8 +357,7 @@ impl UnboundedBuffer { /// Ensures the chunk that holds `version` exists and returns that slot. /// /// The caller is the unique producer and must not publish `tail` past this version until this - /// returns. Fully consumed chunks stay allocated until the channel is dropped so a receiver - /// walking the list cannot observe a freed chunk. + /// returns. `version` is at or above `tail`, so its chunk still owns its storage. pub(super) fn slot_for_publish(&self, version: u64) -> &Slot { loop { let chunk = self.tail_chunk.load(Ordering::Acquire); @@ -357,7 +365,7 @@ impl UnboundedBuffer { let current = unsafe { &*chunk }; if version < current.base + CHUNK_LEN as u64 { debug_assert!(version >= current.base); - return ¤t.slots[(version - current.base) as usize]; + return unsafe { current.slot(version) }; } let next_base = current.base + CHUNK_LEN as u64; @@ -367,28 +375,36 @@ impl UnboundedBuffer { } } - fn advance_head_chunk(&self, head: u64) { + /// Releases the message storage of every chunk the live window has left behind. + /// + /// A receive only addresses versions at or above `head`, so a chunk entirely below it can + /// release its storage even while receivers walk past its header. Freeing the header instead + /// would race with that walk. + fn release_consumed_chunks(&self, head: u64) { loop { let chunk = self.head_chunk.load(Ordering::Acquire); + debug_assert!(!chunk.is_null()); let current = unsafe { &*chunk }; let next = current.next.load(Ordering::Acquire); - if next.is_null() { - return; - } - if current.base + CHUNK_LEN as u64 > head { + if next.is_null() || current.base + CHUNK_LEN as u64 > head { return; } + current.release(); self.head_chunk.store(next, Ordering::Release); } } + /// The number of message slots this buffer currently keeps allocated. #[cfg(test)] pub(super) fn allocated_slots(&self) -> usize { let mut n = 0; let mut chunk = self.root.load(Ordering::Acquire); while !chunk.is_null() { - n += CHUNK_LEN; - chunk = unsafe { (*chunk).next.load(Ordering::Acquire) }; + let current = unsafe { &*chunk }; + if !current.slots.load(Ordering::Acquire).is_null() { + n += CHUNK_LEN; + } + chunk = current.next.load(Ordering::Acquire); } n } @@ -399,11 +415,7 @@ impl Drop for UnboundedBuffer { let mut chunk = self.root.load(Ordering::Relaxed); while !chunk.is_null() { let boxed = unsafe { Box::from_raw(chunk) }; - for slot in &boxed.slots { - if slot.occupied.load(Ordering::Relaxed) { - drop(unsafe { slot.take_msg() }); - } - } + boxed.release(); chunk = boxed.next.load(Ordering::Relaxed); } } @@ -418,40 +430,27 @@ impl SlotStore for UnboundedBuffer { let current = unsafe { &*chunk }; if version < current.base + CHUNK_LEN as u64 { debug_assert!(version >= current.base); - return ¤t.slots[(version - current.base) as usize]; + return unsafe { current.slot(version) }; } chunk = current.next.load(Ordering::Acquire); } } fn sync_head(&self, head: u64) { - self.advance_head_chunk(head); + self.release_consumed_chunks(head); } } -/// Advances `head` over slots whose value has already been taken. -fn advance_head>(shared: &Shared) { - let mut h = shared.head.load(Ordering::Acquire); - loop { - let t = shared.tail.load(Ordering::Acquire); - if h >= t { - break; - } - if shared.buffer.slot(h).occupied.load(Ordering::Acquire) { - break; - } - match shared - .head - .compare_exchange_weak(h, h + 1, Ordering::AcqRel, Ordering::Acquire) - { - Ok(_) => h += 1, - Err(actual) => h = actual, - } - } - shared.buffer.sync_head(h); - if shared.producer_waiting.load(Ordering::Acquire) != 0 { - notify_blocking(shared); - } +/// Advances `head` past `version`, whose value the caller has just taken. +/// +/// Slots are released in version order: a subscription consumes in cursor order and a dropped one +/// reclaims in increasing order, so `remaining` can only reach zero at `version` once it has at +/// every earlier version. One `fetch_max` therefore does what a scan over the log would, and no +/// caller has to address a version the live window may already have passed. +fn release_head>(shared: &Shared, version: u64) { + let next = version + 1; + let head = shared.head.fetch_max(next, Ordering::AcqRel).max(next); + shared.buffer.sync_head(head); } /// Consumes the message at `cursor` and advances the cursor. @@ -472,7 +471,7 @@ pub(super) fn consume>( if slot.remaining.load(Ordering::Acquire) == 1 { let value = unsafe { slot.take_msg() }; slot.remaining.store(0, Ordering::Release); - advance_head(shared); + release_head(shared, version); return Consumed { value, reclaimed: true, @@ -482,6 +481,7 @@ pub(super) fn consume>( struct RemainingGuard<'a, T, B: SlotStore> { slot: &'a Slot, shared: &'a Shared, + version: u64, armed: bool, } @@ -492,7 +492,11 @@ pub(super) fn consume>( } if self.slot.remaining.fetch_sub(1, Ordering::AcqRel) == 1 { drop(unsafe { self.slot.take_msg() }); - advance_head(self.shared); + release_head(self.shared, self.version); + // Another reader may have advanced past this slot while the clone was running, so + // the panicking receive can be the one that frees capacity. The caller is + // unwinding and will not wake anyone, so wake from here. + wake_producer(take_producer_on_reclaim(self.shared, true)); } } } @@ -500,6 +504,7 @@ pub(super) fn consume>( let mut guard = RemainingGuard { slot, shared, + version, armed: true, }; let value = unsafe { slot.clone_msg() }; @@ -507,7 +512,7 @@ pub(super) fn consume>( guard.armed = false; if last { drop(unsafe { slot.take_msg() }); - advance_head(shared); + release_head(shared, version); } Consumed { value, @@ -517,30 +522,36 @@ pub(super) fn consume>( fn reclaim_range>(shared: &Shared, start: u64, end: u64) -> Reclaimed { let mut reclaimed = Reclaimed::empty(); + let mut released = None; let mut version = start; while version < end { let slot = shared.buffer.slot(version); if slot.remaining.fetch_sub(1, Ordering::AcqRel) == 1 { reclaimed.push(unsafe { slot.take_msg() }); + released = Some(version); } version += 1; } - if !reclaimed.is_empty() { - advance_head(shared); + if let Some(version) = released { + release_head(shared, version); } reclaimed } /// Takes the parked producer if a reclaim may have freed capacity. -pub(super) fn take_producer_on_reclaim( - shared: &Shared, - reclaimed: bool, - drained_last: bool, -) -> Option { - if !drained_last && !reclaimed { +/// +/// Skipping the lock keeps the common receive off the waiter mutex, but this reclaim stored `head` +/// before loading `producer_waiting` while a parking producer does the opposite. Release/acquire +/// does not order a store against a later load of another location, so without the fence — and its +/// counterpart in the send path — both sides can read stale values and the producer parks on +/// capacity this reclaim already released. +pub(super) fn take_producer_on_reclaim(shared: &Shared, reclaimed: bool) -> Option { + if !reclaimed { return None; } - if !drained_last && shared.producer_waiting.load(Ordering::Acquire) == 0 { + + fence(Ordering::SeqCst); + if shared.producer_waiting.load(Ordering::Acquire) == 0 { return None; } @@ -560,11 +571,8 @@ pub(super) fn drop_subscription>( shared: &Shared, cursor: u64, ) -> (Reclaimed, Option) { - let mut state = shared.lock_idle_send(); + let mut state = shared.state.lock(); state.receiver_count -= 1; - shared - .receiver_count - .store(state.receiver_count, Ordering::Release); let last = state.receiver_count == 0; let tail = shared.tail.load(Ordering::Acquire); @@ -579,7 +587,7 @@ pub(super) fn drop_subscription>( drop(state); let reclaimed = reclaim_range(shared, cursor, tail); - let producer = take_producer_on_reclaim(shared, !reclaimed.is_empty(), false); + let producer = take_producer_on_reclaim(shared, !reclaimed.is_empty()); (reclaimed, producer) } @@ -590,7 +598,6 @@ pub(super) fn disconnect(shared: &Shared) { state.waiters.take_all() }; wake_all(wakers); - notify_blocking(shared); } /// Releases a cancelled receive's waker registration, dropping the waker unlocked. @@ -623,8 +630,8 @@ pub(super) fn try_receive>( /// The one poll step behind `recv` on both channels. /// -/// The ready path does not take the waiter mutex. Parking stores `has_waiters` and then rechecks -/// `tail` so an unbounded send that published without this lock cannot leave the receiver parked. +/// The ready path does not take the waiter mutex. Parking takes it, and so do publication and the +/// sender's disconnect, so a registration made here is visible to whichever happens next. pub(super) fn poll_receive>( shared: &Shared, cursor: &mut u64, @@ -647,24 +654,12 @@ pub(super) fn poll_receive>( return Poll::Ready(Err(RecvError::Disconnected)); } + // A disconnect racing this registration still wakes it: the sender clears `senders` before + // taking this lock to drain the wait set. let retired_waker = state.waiters.register(token, cx.waker()); - shared.has_waiters.store(true, Ordering::Release); - if *cursor >= shared.tail.load(Ordering::Acquire) && shared.senders.load(Ordering::Acquire) != 0 - { - drop(state); - drop(retired_waker); - return Poll::Pending; - } - - let waker = state.waiters.unregister(token); drop(state); drop(retired_waker); - drop(waker); - if *cursor < shared.tail.load(Ordering::Acquire) { - Poll::Ready(Ok(consume(shared, cursor))) - } else { - Poll::Ready(Err(RecvError::Disconnected)) - } + Poll::Pending } /// Publishes `tail` after writing a slot. @@ -673,14 +668,10 @@ pub(super) fn commit_publish(tail: &AtomicU64, next: u64) { } /// Advances `head` and `tail` together when nothing can read the message. +/// +/// `tail` moves first so a concurrent `retained` never observes `head` ahead of it. `fetch_max` +/// keeps `head` monotonic against a reclaim that is releasing slots at the same time. pub(super) fn commit_discard(head: &AtomicU64, tail: &AtomicU64, next: u64) { - head.store(next, Ordering::Release); tail.store(next, Ordering::Release); -} - -/// Wakes native-thread waiters parked in `recv_blocking` / `send_blocking`. -pub(super) fn notify_blocking(shared: &Shared) { - let _guard = shared.blocking.lock(); - shared.epoch.fetch_add(1, Ordering::Release); - shared.blocking_cvar.notify_all(); + head.fetch_max(next, Ordering::AcqRel); } diff --git a/asyncband/src/broadcast/spmc/mod.rs b/asyncband/src/broadcast/spmc/mod.rs index 811ff53a..1f102b95 100644 --- a/asyncband/src/broadcast/spmc/mod.rs +++ b/asyncband/src/broadcast/spmc/mod.rs @@ -28,8 +28,9 @@ //! The sender is not [`Clone`], and every publish method takes `&mut self`. That is the static //! single-writer contract this topology adds over [`crate::broadcast::mpmc`]: there cannot be a //! second producer, at compile time. `&Sender` can still be shared for [`subscribe`] and the -//! inspection methods. Receivers drain already-published slots without taking the publication -//! lock, which is the throughput reason to pick this family when there is one producer. +//! inspection methods. Apart from that the public surface matches [`crate::broadcast::mpmc`] +//! method for method. Receivers drain already-published slots without taking the publication lock, +//! which is the throughput reason to pick this family when there is one producer. //! //! This is fan-out broadcast, not a competing queue: every accepted value is delivered to every //! active subscription. A competitive `asyncband::spmc` queue would give each value to exactly one diff --git a/asyncband/src/broadcast/spmc/unbounded/mod.rs b/asyncband/src/broadcast/spmc/unbounded/mod.rs index 70a26ee7..4ce9d9ea 100644 --- a/asyncband/src/broadcast/spmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/spmc/unbounded/mod.rs @@ -30,6 +30,11 @@ //! [`UnboundedSender::retained_message_count`] reports its current length. Use [`bounded`] when //! the producer should wait for the slowest receiver instead of growing the backlog. //! +//! Value storage is released as the backlog drains, so memory tracks the peak backlog rather than +//! the number of values ever published. The log keeps one small index entry per 256 values until +//! the channel is dropped: a receiver locates a value without any lock, and may still be walking +//! past entries the backlog has left behind. +//! //! # Receivers //! //! [`UnboundedSender::subscribe`] and [`UnboundedReceiver::resubscribe`] add a receiver at the @@ -157,16 +162,21 @@ impl UnboundedSender { /// assert_eq!(second.try_recv(), Ok("update")); /// ``` pub fn send(&mut self, msg: T) { - self.shared.send_in_progress.store(true, Ordering::Release); - let n = self.shared.receiver_count.load(Ordering::Acquire); - - let unretained = if n == 0 { - let state = self.shared.state.lock(); + let mut discarded = None; + let wakers = { + // Counting subscriptions, writing the slot, publishing `tail`, and draining the wait + // set share one critical section with subscribe and unsubscribe. That is what makes a + // slot's remaining-reader count match its consumers and keeps a parking receiver from + // missing this publication. + let mut state = self.shared.state.lock(); let tail = self.shared.tail.load(Ordering::Relaxed); let next = Shared::>::next_tail(tail); + if state.receiver_count == 0 { + // Nothing can read this message. It leaves the critical section with us and is + // dropped below, so `T::drop` never runs under the lock. + discarded = Some(msg); common::commit_discard(&self.shared.head, &self.shared.tail, next); - Some(msg) } else { unsafe { self.shared @@ -175,28 +185,12 @@ impl UnboundedSender { .write(msg, state.receiver_count); } common::commit_publish(&self.shared.tail, next); - None } - } else { - let tail = self.shared.tail.load(Ordering::Relaxed); - let next = Shared::>::next_tail(tail); - unsafe { - self.shared.buffer.slot_for_publish(tail).write(msg, n); - } - common::commit_publish(&self.shared.tail, next); - None + state.waiters.drain() }; - self.shared.send_in_progress.store(false, Ordering::Release); - if self.shared.has_waiters.load(Ordering::Acquire) { - let wakers = { - let mut state = self.shared.state.lock(); - self.shared.has_waiters.store(false, Ordering::Release); - state.waiters.drain() - }; - wake_all(wakers); - } - drop(unretained); + wake_all(wakers); + drop(discarded); } /// Returns the number of values in the shared backlog. diff --git a/asyncband/src/broadcast/spmc/unbounded/tests.rs b/asyncband/src/broadcast/spmc/unbounded/tests.rs index 2ea9df9c..714c7a06 100644 --- a/asyncband/src/broadcast/spmc/unbounded/tests.rs +++ b/asyncband/src/broadcast/spmc/unbounded/tests.rs @@ -28,7 +28,7 @@ fn send_panics_on_version_overflow() { } #[test] -fn chunks_grow_with_the_committed_log() { +fn chunk_storage_grows_and_drains_with_the_live_window() { let (mut tx, mut rx) = unbounded(); let burst = CHUNK_LEN * 4; @@ -41,8 +41,15 @@ fn chunks_grow_with_the_committed_log() { assert_eq!(rx.try_recv(), Ok(i)); } - // Chunks stay allocated until the channel is dropped so receivers can walk them without a - // reclamation lock. + // The footprint tracks the live window, not the lifetime message count: only the chunk + // holding the current window is left. assert_eq!(tx.retained_message_count(), 0); - assert!(tx.shared.buffer.allocated_slots() >= burst); + assert_eq!(tx.shared.buffer.allocated_slots(), CHUNK_LEN); + + // Released chunks keep their index entries, so lookups still reach live versions. + for i in 0..CHUNK_LEN { + tx.send(i); + assert_eq!(rx.try_recv(), Ok(i)); + } + assert_eq!(tx.shared.buffer.allocated_slots(), CHUNK_LEN); } diff --git a/benchmarks/ecosystem/broadcast/spmc/adapters.rs b/benchmarks/ecosystem/broadcast/spmc/adapters.rs index 707c2e89..e5405628 100644 --- a/benchmarks/ecosystem/broadcast/spmc/adapters.rs +++ b/benchmarks/ecosystem/broadcast/spmc/adapters.rs @@ -222,14 +222,6 @@ impl BoundedBroadcastSpmc for Asyncband { receiver.recv().await.unwrap() } - fn send_blocking(sender: &mut Self::Sender, value: usize) { - sender.send_blocking(value); - } - - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { - receiver.recv_blocking().unwrap() - } - fn try_recv(receiver: &mut Self::Receiver) -> Option { match receiver.try_recv() { Ok(value) => Some(value), diff --git a/tests-integration/tests/broadcast_spmc_bounded_test.rs b/tests-integration/tests/broadcast_spmc_bounded_test.rs index 394e722a..50daa8af 100644 --- a/tests-integration/tests/broadcast_spmc_bounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_bounded_test.rs @@ -31,6 +31,7 @@ // the remaining drop-vs-wait race is `dropping_the_last_receiver_wakes_the_blocked_sender` // together with `sends_never_block_once_all_receivers_are_gone`. +use std::future::Future; use std::sync::Arc; use std::sync::Mutex; use std::task::Context; @@ -94,6 +95,33 @@ impl Clone for PanicOnClone { } } +/// A payload whose clone lets another subscription consume the same slot, then panics. +/// +/// That makes the panicking receive the slot's last reader, so it releases the slot from the drop +/// guard inside the receive step — a path that unwinds before the caller can wake the producer. +struct RaceOnClone { + armed: bool, + probe: Option>>>, +} + +impl Clone for RaceOnClone { + fn clone(&self) -> Self { + if self.armed { + if let Some(probe) = &self.probe + && let Ok(mut probe) = probe.try_lock() + { + // Takes this slot's remaining count from 2 to 1 while we are still cloning. + let _ = probe.try_recv(); + } + panic!("panic while cloning a broadcast message"); + } + Self { + armed: self.armed, + probe: self.probe.clone(), + } + } +} + /// 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 @@ -235,7 +263,7 @@ fn concurrent_receivers_keep_up_with_the_producer() { thread::spawn(move || { let mut sum = 0u64; for _ in 0..MESSAGES { - sum += rx.recv_blocking().unwrap(); + sum += FutureExt::block_on(rx.recv()).unwrap(); } sum }) @@ -243,7 +271,7 @@ fn concurrent_receivers_keep_up_with_the_producer() { .collect(); for value in 0..MESSAGES as u64 { - tx.send_blocking(value); + FutureExt::block_on(tx.send(value)); } for handle in handles { @@ -258,7 +286,7 @@ fn concurrent_blocking_wait_at_capacity_one() { let handle = thread::spawn(move || { let mut sum = 0u64; for _ in 0..64 { - sum += rx2.recv_blocking().unwrap(); + sum += FutureExt::block_on(rx2.recv()).unwrap(); } sum }); @@ -266,12 +294,12 @@ fn concurrent_blocking_wait_at_capacity_one() { let mut sum = 0u64; let mut rx = rx; for _ in 0..64 { - sum += rx.recv_blocking().unwrap(); + sum += FutureExt::block_on(rx.recv()).unwrap(); } sum }); for value in 0..64u64 { - tx.send_blocking(value); + FutureExt::block_on(tx.send(value)); } assert_eq!(handle.join().unwrap(), 64 * 63 / 2); assert_eq!(handle1.join().unwrap(), 64 * 63 / 2); @@ -599,12 +627,12 @@ fn parked_recv_wakes_when_the_sender_drops() { } #[test] -fn parked_recv_blocking_wakes_when_the_sender_drops() { +fn parked_native_thread_recv_wakes_when_the_sender_drops() { assert_completes_without_deadlock(|| { let (tx, mut rx) = bounded::(4); - let parked = thread::spawn(move || rx.recv_blocking()); + let parked = thread::spawn(move || FutureExt::block_on(rx.recv())); // The worker parks on the empty channel. Dropping the sender must finish that receive - // with Disconnected; a missed condvar wake hangs inside assert_completes_without_deadlock. + // with Disconnected; a missed wake hangs inside assert_completes_without_deadlock. thread::sleep(std::time::Duration::from_millis(50)); drop(tx); assert_eq!(parked.join().unwrap(), Err(RecvError::Disconnected)); @@ -754,6 +782,39 @@ fn panicking_payload_destructor_still_releases_capacity() { ); } +#[test] +fn panicking_clone_that_becomes_the_last_reader_wakes_the_producer() { + let (mut tx, mut rx1) = bounded(1); + let rx2 = Arc::new(Mutex::new(tx.subscribe())); + + tx.try_send(RaceOnClone { + armed: true, + probe: Some(rx2.clone()), + }) + .unwrap(); + + let (waker, counter) = WakeCounter::new(); + let mut context = Context::from_waker(&waker); + let mut send = Box::pin(tx.send(RaceOnClone { + armed: false, + probe: None, + })); + assert!(send.as_mut().poll(&mut context).is_pending()); + assert_eq!(counter.count(), 0); + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _ = rx1.try_recv(); + })); + assert!(result.is_err(), "the clone was expected to panic"); + + assert!( + counter.count() > 0, + "a panicking clone must not strand a producer on capacity it already freed" + ); + drop(send); + assert_eq!(tx.retained_message_count(), 0); +} + #[test] fn message_destructors_run_outside_the_channel_lock() { assert_completes_without_deadlock(|| { diff --git a/tests-integration/tests/broadcast_spmc_unbounded_test.rs b/tests-integration/tests/broadcast_spmc_unbounded_test.rs index 04e3cac7..dbcbc5ad 100644 --- a/tests-integration/tests/broadcast_spmc_unbounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_unbounded_test.rs @@ -23,6 +23,8 @@ use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -523,6 +525,128 @@ fn concurrent_receivers_drain_a_published_batch() { } } +#[test] +fn subscription_churn_never_strands_the_backlog() { + const MESSAGES: u64 = 200_000; + + let (mut tx, rx) = unbounded::(); + let stop = Arc::new(AtomicBool::new(false)); + + // A publish has to agree with concurrent subscribe/unsubscribe on each version's consumer + // count. If it does not, a slot's remaining count never reaches zero and `head` stops. + let churn_stop = stop.clone(); + let churn = thread::spawn(move || { + while !churn_stop.load(Ordering::Relaxed) { + for _ in 0..64 { + drop(rx.resubscribe()); + } + } + rx + }); + + for value in 0..MESSAGES { + tx.send(value); + } + stop.store(true, Ordering::Relaxed); + let mut rx = churn.join().unwrap(); + + for value in 0..MESSAGES { + assert_eq!(rx.try_recv(), Ok(value)); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + assert_eq!(tx.retained_message_count(), 0); +} + +#[test] +fn publish_races_with_a_parking_receiver_without_losing_the_wakeup() { + const ROUNDS: u64 = 200_000; + + let (mut tx, mut rx) = unbounded::(); + let gate = Arc::new(AtomicU64::new(0)); + let published = Arc::new(AtomicU64::new(0)); + + // Start the send and the parking receive together, so the publish lands while the receiver is + // registering its waker. + let producer_gate = gate.clone(); + let producer_published = published.clone(); + let producer = thread::spawn(move || { + for round in 1..=ROUNDS { + while producer_gate.load(Ordering::Acquire) < round { + std::hint::spin_loop(); + } + tx.send(round); + producer_published.store(round, Ordering::Release); + } + tx + }); + + let (waker, counter) = WakeCounter::new(); + let mut context = Context::from_waker(&waker); + + for round in 1..=ROUNDS { + let woken = counter.count(); + let mut recv = Box::pin(rx.recv()); + gate.store(round, Ordering::Release); + if recv.as_mut().poll(&mut context).is_ready() { + continue; + } + + // The receive parked, so this round's send must have woken it. Otherwise the value sits + // unread with the receiver parked. + while published.load(Ordering::Acquire) < round { + std::hint::spin_loop(); + } + assert!( + counter.count() > woken, + "round {round}: publish left the receiver parked" + ); + drop(recv); + assert_eq!(rx.try_recv(), Ok(round)); + } + + drop(producer.join().unwrap()); +} + +#[test] +fn releasing_a_long_backlog_keeps_concurrent_receives_inside_the_live_window() { + const DRAINERS: usize = 32; + const BACKLOG: u64 = 20_000; + + // Dropping a receiver that pinned the whole prefix advances `head` many chunks in one step, + // while the other receivers are mid-receive. Locating a slot must not depend on a `head` + // snapshot that this jump has already invalidated. + for _ in 0..8 { + let (mut tx, rx) = unbounded::(); + let laggard = tx.subscribe(); + for value in 0..BACKLOG { + tx.send(value); + } + + let mut drainers = vec![rx]; + for _ in 1..DRAINERS { + drainers.push(tx.subscribe()); + } + let stop = Arc::new(AtomicBool::new(false)); + let handles: Vec<_> = drainers + .into_iter() + .map(|mut rx| { + let stop = stop.clone(); + thread::spawn(move || { + while !stop.load(Ordering::Relaxed) { + let _ = rx.try_recv(); + } + }) + }) + .collect(); + + drop(laggard); + stop.store(true, Ordering::Relaxed); + for handle in handles { + handle.join().unwrap(); + } + } +} + #[test] fn publish_order_is_program_order() { let (mut tx, mut rx1) = unbounded(); From 6f3ddc48ef2d52e244b94e3c33485ea606fc9334 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Wed, 16 Sep 2026 11:58:04 +0800 Subject: [PATCH 9/9] fix(broadcast): match MPMC drain and discard contracts in SPMC Reload tail after observing a dropped sender so try_recv cannot report Disconnected while a published value is still unread. Discard unbounded sends without advancing the version log so empty chunks are not allocated for messages that were never retained. --- asyncband/src/broadcast/spmc/common.rs | 14 +++++-- asyncband/src/broadcast/spmc/unbounded/mod.rs | 9 +++-- .../src/broadcast/spmc/unbounded/tests.rs | 30 +++++++++++++- .../tests/broadcast_spmc_bounded_test.rs | 39 +++++++++++++++++++ .../tests/broadcast_spmc_unbounded_test.rs | 32 +++++++++++++-- 5 files changed, 111 insertions(+), 13 deletions(-) diff --git a/asyncband/src/broadcast/spmc/common.rs b/asyncband/src/broadcast/spmc/common.rs index 33b80431..36cbc5c8 100644 --- a/asyncband/src/broadcast/spmc/common.rs +++ b/asyncband/src/broadcast/spmc/common.rs @@ -614,6 +614,11 @@ pub(super) fn unregister(shared: &Shared, cursor: u64, token: &mut Option< } /// Receives without waiting. +/// +/// `Disconnected` is only returned after a reload of `tail`. The sender stores `senders = 0` only +/// after its last publication, so seeing zero synchronizes with that store and cannot hide an +/// accepted message. Returning `Disconnected` while a value is still unread would violate the same +/// drain-before-disconnect contract as [`crate::broadcast::mpmc`]. pub(super) fn try_receive>( shared: &Shared, cursor: &mut u64, @@ -621,10 +626,13 @@ pub(super) fn try_receive>( if *cursor < shared.tail.load(Ordering::Acquire) { return Ok(consume(shared, cursor)); } - if shared.senders.load(Ordering::Acquire) == 0 { - Err(TryRecvError::Disconnected) + if shared.senders.load(Ordering::Acquire) != 0 { + return Err(TryRecvError::Empty); + } + if *cursor < shared.tail.load(Ordering::Acquire) { + Ok(consume(shared, cursor)) } else { - Err(TryRecvError::Empty) + Err(TryRecvError::Disconnected) } } diff --git a/asyncband/src/broadcast/spmc/unbounded/mod.rs b/asyncband/src/broadcast/spmc/unbounded/mod.rs index 4ce9d9ea..4f58c3ba 100644 --- a/asyncband/src/broadcast/spmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/spmc/unbounded/mod.rs @@ -169,15 +169,16 @@ impl UnboundedSender { // slot's remaining-reader count match its consumers and keeps a parking receiver from // missing this publication. let mut state = self.shared.state.lock(); - let tail = self.shared.tail.load(Ordering::Relaxed); - let next = Shared::>::next_tail(tail); if state.receiver_count == 0 { // Nothing can read this message. It leaves the critical section with us and is - // dropped below, so `T::drop` never runs under the lock. + // dropped below, so `T::drop` never runs under the lock. `head` / `tail` stay put: + // a discarded send is not a publication, so the version-addressed log must not + // allocate empty chunks for it. A later subscribe still joins at this tail. discarded = Some(msg); - common::commit_discard(&self.shared.head, &self.shared.tail, next); } else { + let tail = self.shared.tail.load(Ordering::Relaxed); + let next = Shared::>::next_tail(tail); unsafe { self.shared .buffer diff --git a/asyncband/src/broadcast/spmc/unbounded/tests.rs b/asyncband/src/broadcast/spmc/unbounded/tests.rs index 714c7a06..9741cbbe 100644 --- a/asyncband/src/broadcast/spmc/unbounded/tests.rs +++ b/asyncband/src/broadcast/spmc/unbounded/tests.rs @@ -15,15 +15,18 @@ // specific language governing permissions and limitations // under the License. +use super::TryRecvError; use super::*; use crate::broadcast::spmc::common::CHUNK_LEN; #[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 (mut tx, _) = unbounded(); + // Keep a live subscription at the doctored tail so drop does not walk the whole log, and so + // this send is a publication rather than a no-subscriber discard. + let (mut tx, mut rx) = unbounded(); tx.shared.set_tail(u64::MAX); + rx.cursor = u64::MAX; tx.send(()); } @@ -53,3 +56,26 @@ fn chunk_storage_grows_and_drains_with_the_live_window() { } assert_eq!(tx.shared.buffer.allocated_slots(), CHUNK_LEN); } + +#[test] +fn discarded_sends_do_not_grow_chunk_storage() { + let (mut tx, rx) = unbounded(); + drop(rx); + + for i in 0..CHUNK_LEN * 4 { + tx.send(i); + } + + // Discarded sends are not publications. The log must not allocate empty chunks for versions + // that were never retained, matching MPMC's empty-buffer discard path. + assert_eq!(tx.retained_message_count(), 0); + assert_eq!(tx.shared.buffer.allocated_slots(), CHUNK_LEN); + + let mut rx = tx.subscribe(); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + tx.send(0); + assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(tx.retained_message_count(), 0); + assert_eq!(tx.shared.buffer.allocated_slots(), CHUNK_LEN); +} diff --git a/tests-integration/tests/broadcast_spmc_bounded_test.rs b/tests-integration/tests/broadcast_spmc_bounded_test.rs index 50daa8af..c0460621 100644 --- a/tests-integration/tests/broadcast_spmc_bounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_bounded_test.rs @@ -495,6 +495,22 @@ fn sends_never_block_once_all_receivers_are_gone() { assert_eq!(tx.retained_message_count(), 0); } +#[test] +fn send_without_receivers_does_not_buffer() { + let (mut tx, rx) = bounded(4); + drop(rx); + + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + assert_eq!(tx.retained_message_count(), 0); + + let mut rx = tx.subscribe(); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + tx.try_send(3).unwrap(); + assert_eq!(rx.try_recv(), Ok(3)); +} + #[test] fn subscribing_while_the_producer_is_blocked_does_not_release_capacity() { let (mut tx, rx) = bounded(1); @@ -626,6 +642,29 @@ fn parked_recv_wakes_when_the_sender_drops() { ); } +#[test] +fn try_recv_does_not_report_disconnected_while_a_message_is_unread() { + // `try_recv` must not treat a sender drop as terminal until it has observed `tail`. MPMC + // gets that by receiving under the publication lock; SPMC reloads `tail` after `senders == 0`. + for _ in 0..10_000 { + let (mut tx, mut rx) = bounded(4); + let producer = thread::spawn(move || { + tx.try_send(1).unwrap(); + }); + loop { + match rx.try_recv() { + Ok(1) => break, + Err(TryRecvError::Empty) => std::hint::spin_loop(), + Err(TryRecvError::Disconnected) => { + panic!("try_recv dropped a published message on disconnect") + } + Ok(other) => panic!("unexpected value {other}"), + } + } + producer.join().unwrap(); + } +} + #[test] fn parked_native_thread_recv_wakes_when_the_sender_drops() { assert_completes_without_deadlock(|| { diff --git a/tests-integration/tests/broadcast_spmc_unbounded_test.rs b/tests-integration/tests/broadcast_spmc_unbounded_test.rs index dbcbc5ad..1d7a775e 100644 --- a/tests-integration/tests/broadcast_spmc_unbounded_test.rs +++ b/tests-integration/tests/broadcast_spmc_unbounded_test.rs @@ -255,15 +255,16 @@ async fn send_without_receivers_does_not_buffer() { let (mut tx, rx) = unbounded(); drop(rx); - tx.send(1); - tx.send(2); + for value in 0..1024 { + tx.send(value); + } assert_eq!(tx.retained_message_count(), 0); let mut rx = tx.subscribe(); assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - tx.send(3); - assert_eq!(rx.recv().await, Ok(3)); + tx.send(1024); + assert_eq!(rx.recv().await, Ok(1024)); } #[test] @@ -471,6 +472,29 @@ fn parked_recv_prefers_buffered_messages_over_disconnection() { assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); } +#[test] +fn try_recv_does_not_report_disconnected_while_a_message_is_unread() { + // `try_recv` must not treat a sender drop as terminal until it has observed `tail`. MPMC + // gets that by receiving under the publication lock; SPMC reloads `tail` after `senders == 0`. + for _ in 0..10_000 { + let (mut tx, mut rx) = unbounded(); + let producer = thread::spawn(move || { + tx.send(1); + }); + loop { + match rx.try_recv() { + Ok(1) => break, + Err(TryRecvError::Empty) => std::hint::spin_loop(), + Err(TryRecvError::Disconnected) => { + panic!("try_recv dropped a published message on disconnect") + } + Ok(other) => panic!("unexpected value {other}"), + } + } + producer.join().unwrap(); + } +} + #[tokio::test] async fn recv_drains_buffered_messages_before_reporting_disconnection() { let (mut tx, mut rx) = unbounded();