diff --git a/CHANGELOG.md b/CHANGELOG.md index 81e16b5f..aa333808 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 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/mod.rs b/asyncband/src/broadcast/mod.rs index 81e56c55..d49f3e9b 100644 --- a/asyncband/src/broadcast/mod.rs +++ b/asyncband/src/broadcast/mod.rs @@ -16,5 +16,16 @@ // 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`). 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, +//! 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..988a6f6b --- /dev/null +++ b/asyncband/src/broadcast/spmc/bounded/mod.rs @@ -0,0 +1,692 @@ +// 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::Ordering; +use std::sync::atomic::fence; +use std::task::Context; +use std::task::Poll; + +use super::common; +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::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 shared = Arc::new(Shared::new(BoundedBuffer::new(capacity))); + let sender = BoundedSender { + shared: shared.clone(), + }; + let receiver = BoundedReceiver { shared, cursor: 0 }; + (sender, receiver) +} + +/// 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. +/// +/// ```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>>, +} + +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); + } +} + +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. + self.sender + .shared + .producer_waiting + .store(0, Ordering::Release); + let waker = { + let mut state = self.sender.shared.state.lock(); + state.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(()), + }; + + self.sender + .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 { + 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(()); + } + + 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 = state.producer.replace(cx.waker().clone()); + drop(state); + drop(retired); + self.value = Some(msg); + return Poll::Pending; + } + + 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. + wake_all(wakers); + drop(retired_producer); + Poll::Ready(()) + } + } + + let mut send = SendState { + sender: self, + value: Some(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> { + self.publish(value).map_err(TrySendError::Full) + } + + /// The publish step both send paths share. + /// + /// 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. 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 state = self.shared.state.lock(); + + 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. + let next = + Shared::>::next_tail(self.shared.tail.load(Ordering::Relaxed)); + discarded = Some(payload); + common::commit_discard(&self.shared.head, &self.shared.tail, next); + } else { + 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); + } + + 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); + 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.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.buffer.cap + } + + /// 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 cursor = self.shared.subscribe(); + BoundedReceiver { + shared: self.shared.clone(), + cursor, + } + } +} + +/// 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>>, + cursor: u64, +} + +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) = 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. + 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 consumed = common::try_receive(&self.shared, &mut self.cursor)?; + let producer = common::take_producer_on_reclaim(&self.shared, consumed.reclaimed); + common::wake_producer(producer); + Ok(consumed.value) + } +} + +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 cursor = self.shared.subscribe(); + Self { + shared: self.shared.clone(), + cursor, + } + } + + /// 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.unread(self.cursor) + } +} + +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, self.receiver.cursor, &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 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(consumed)) => consumed, + }; + + 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/bounded/tests.rs b/asyncband/src/broadcast/spmc/bounded/tests.rs new file mode 100644 index 00000000..49b0fbba --- /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.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.buffer.len(); + 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.buffer.len(), 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.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.state.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..36cbc5c8 --- /dev/null +++ b/asyncband/src/broadcast/spmc/common.rs @@ -0,0 +1,685 @@ +// 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. + +//! 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 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::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::sync::atomic::fence; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use super::error::RecvError; +use super::error::TryRecvError; +use crate::internal::mutex::Mutex; +use crate::internal::wake_all; +use crate::internal::wakerset::WakerSet; +use crate::internal::wakerset::WakerToken; + +/// Number of slots in one unbounded log chunk. +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(super) struct Consumed { + pub value: T, + pub reclaimed: bool, +} + +/// 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(super) struct Reclaimed { + first: Option, + rest: Vec, +} + +impl Reclaimed { + fn empty() -> Self { + Self { + first: None, + rest: vec![], + } + } + + pub(super) fn is_empty(&self) -> bool { + self.first.is_none() + } + + fn push(&mut self, msg: T) { + if self.first.is_none() { + self.first = Some(msg); + } else { + self.rest.push(msg); + } + } +} + +/// One published value and the number of subscriptions that still have to consume it. +pub(super) struct Slot { + msg: UnsafeCell>, + remaining: AtomicUsize, + /// `true` once the producer has written `msg` and until the last remaining reader takes it. + /// + /// `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, +} + +// 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. + /// + /// 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(super) 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); + } + + /// 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() } + } + + /// 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(super) struct State { + pub waiters: WakerSet, + pub receiver_count: usize, + pub producer: Option, +} + +/// Shared channel state: the slot log plus the waiter mutex. +pub(super) struct Shared { + pub buffer: B, + pub head: AtomicU64, + pub tail: AtomicU64, + pub senders: AtomicUsize, + pub producer_waiting: AtomicUsize, + pub state: Mutex, +} + +impl Shared { + pub(super) fn new(buffer: B) -> Self { + Self { + buffer, + 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, + }), + } + } + + 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(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(super) fn next_tail(tail: u64) -> u64 { + tail.checked_add(1) + .expect("broadcast channel version counter overflowed") + } + + /// 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.state.lock(); + state.receiver_count += 1; + self.tail.load(Ordering::Acquire) + } + + #[cfg(test)] + 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(super) trait SlotStore { + fn slot(&self, version: u64) -> &Slot; + + /// Moves the lookup start past storage the live window has left behind, releasing it. + /// + /// 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) {} +} + +/// Fixed ring used by the bounded channel. +pub(super) struct BoundedBuffer { + slots: Box<[Slot]>, + pub cap: usize, +} + +impl BoundedBuffer { + pub(super) fn new(capacity: usize) -> Self { + Self { + slots: (0..capacity).map(|_| Slot::empty()).collect(), + cap: capacity, + } + } + + #[cfg(test)] + pub(super) fn len(&self) -> usize { + self.slots.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() }); + } + } + } +} + +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. +/// +/// 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: 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: 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. 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>, + /// 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(super) 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, + } + } + + /// 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. `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); + debug_assert!(!chunk.is_null()); + let current = unsafe { &*chunk }; + if version < current.base + CHUNK_LEN as u64 { + debug_assert!(version >= current.base); + return unsafe { current.slot(version) }; + } + + 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); + } + } + + /// 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() || 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() { + let current = unsafe { &*chunk }; + if !current.slots.load(Ordering::Acquire).is_null() { + n += CHUNK_LEN; + } + chunk = current.next.load(Ordering::Acquire); + } + n + } +} + +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) }; + boxed.release(); + chunk = boxed.next.load(Ordering::Relaxed); + } + } +} + +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 unsafe { current.slot(version) }; + } + chunk = current.next.load(Ordering::Acquire); + } + } + + fn sync_head(&self, head: u64) { + self.release_consumed_chunks(head); + } +} + +/// 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. +/// +/// 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(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); + *cursor = version + 1; + + if slot.remaining.load(Ordering::Acquire) == 1 { + let value = unsafe { slot.take_msg() }; + slot.remaining.store(0, Ordering::Release); + release_head(shared, version); + return Consumed { + value, + reclaimed: true, + }; + } + + struct RemainingGuard<'a, T, B: SlotStore> { + slot: &'a Slot, + shared: &'a Shared, + version: u64, + armed: bool, + } + + 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() }); + 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)); + } + } + } + + let mut guard = RemainingGuard { + slot, + shared, + version, + 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() }); + release_head(shared, version); + } + Consumed { + value, + reclaimed: last, + } +} + +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 let Some(version) = released { + release_head(shared, version); + } + reclaimed +} + +/// Takes the parked producer if a reclaim may have freed capacity. +/// +/// 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; + } + + fence(Ordering::SeqCst); + if shared.producer_waiting.load(Ordering::Acquire) == 0 { + return None; + } + + let mut state = shared.state.lock(); + state.producer.take() +} + +/// Wakes the parked producer, if any, with the channel already unlocked. +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(super) 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()); + (reclaimed, producer) +} + +/// Wakes every parked receiver so it can observe the channel's disconnected state. +pub(super) fn disconnect(shared: &Shared) { + let wakers = { + 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(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; + return; + } + + let waker = state.waiters.unregister(token); + drop(state); + drop(waker); +} + +/// 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, +) -> Result, TryRecvError> { + if *cursor < shared.tail.load(Ordering::Acquire) { + return Ok(consume(shared, cursor)); + } + 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::Disconnected) + } +} + +/// The one poll step behind `recv` on both channels. +/// +/// 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, + token: &mut Option, + cx: &mut Context<'_>, +) -> Poll, RecvError>> { + if *cursor < shared.tail.load(Ordering::Acquire) { + *token = None; + return Poll::Ready(Ok(consume(shared, cursor))); + } + + let mut state = shared.state.lock(); + if *cursor < shared.tail.load(Ordering::Acquire) { + *token = None; + drop(state); + return Poll::Ready(Ok(consume(shared, cursor))); + } + if shared.senders.load(Ordering::Acquire) == 0 { + *token = None; + 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()); + drop(state); + drop(retired_waker); + Poll::Pending +} + +/// Publishes `tail` after writing a slot. +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. +/// +/// `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) { + tail.store(next, Ordering::Release); + head.fetch_max(next, Ordering::AcqRel); +} 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..1f102b95 --- /dev/null +++ b/asyncband/src/broadcast/spmc/mod.rs @@ -0,0 +1,75 @@ +// 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. 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 +//! second producer, at compile time. `&Sender` can still be shared for [`subscribe`] and the +//! 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 +//! 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..4f58c3ba --- /dev/null +++ b/asyncband/src/broadcast/spmc/unbounded/mod.rs @@ -0,0 +1,426 @@ +// 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. +//! +//! 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 +//! 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::Ordering; +use std::task::Context; +use std::task::Poll; + +use super::common; +use super::common::Shared; +use super::common::UnboundedBuffer; +use super::error::RecvError; +use super::error::TryRecvError; +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 shared = Arc::new(Shared::new(UnboundedBuffer::new())); + let sender = UnboundedSender { + shared: shared.clone(), + }; + let receiver = UnboundedReceiver { shared, cursor: 0 }; + (sender, receiver) +} + +/// 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. +/// +/// ```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>>, +} + +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); + } +} + +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 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(); + + 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. `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); + } 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, state.receiver_count); + } + common::commit_publish(&self.shared.tail, next); + } + state.waiters.drain() + }; + + wake_all(wakers); + drop(discarded); + } + + /// 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.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 cursor = self.shared.subscribe(); + UnboundedReceiver { + shared: self.shared.clone(), + cursor, + } + } +} + +/// 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>>, + cursor: u64, +} + +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, _producer) = common::drop_subscription(&self.shared, self.cursor); + 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 { + Ok(common::try_receive(&self.shared, &mut self.cursor)?.value) + } +} + +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 cursor = self.shared.subscribe(); + Self { + shared: self.shared.clone(), + cursor, + } + } + + /// 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.unread(self.cursor) + } +} + +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, self.receiver.cursor, &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 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(consumed)) => consumed, + }; + + Poll::Ready(Ok(consumed.value)) + } +} diff --git a/asyncband/src/broadcast/spmc/unbounded/tests.rs b/asyncband/src/broadcast/spmc/unbounded/tests.rs new file mode 100644 index 00000000..9741cbbe --- /dev/null +++ b/asyncband/src/broadcast/spmc/unbounded/tests.rs @@ -0,0 +1,81 @@ +// 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::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() { + // 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(()); +} + +#[test] +fn chunk_storage_grows_and_drains_with_the_live_window() { + let (mut tx, mut rx) = unbounded(); + + let burst = CHUNK_LEN * 4; + for i in 0..burst { + tx.send(i); + } + assert!(tx.shared.buffer.allocated_slots() >= burst); + + for i in 0..burst { + assert_eq!(rx.try_recv(), Ok(i)); + } + + // 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_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); +} + +#[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/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..c0460621 --- /dev/null +++ b/tests-integration/tests/broadcast_spmc_bounded_test.rs @@ -0,0 +1,922 @@ +// 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::future::Future; +use std::sync::Arc; +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::*; +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() { + // `resubscribe` takes the waiter mutex; `unread_message_count` does not. + let _ = probe.resubscribe(); + } + } + } +} + +/// 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 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 +/// 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 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); + 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 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 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 += FutureExt::block_on(rx2.recv()).unwrap(); + } + sum + }); + let handle1 = thread::spawn(move || { + let mut sum = 0u64; + let mut rx = rx; + for _ in 0..64 { + sum += FutureExt::block_on(rx.recv()).unwrap(); + } + sum + }); + for value in 0..64u64 { + FutureExt::block_on(tx.send(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); + 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)); + assert_eq!(rx2.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 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); + 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)) + ); +} + +#[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(|| { + let (tx, mut rx) = bounded::(4); + 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 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 +// --------------------------------------------------------------------------------------------- + +#[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 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(|| { + 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..1d7a775e --- /dev/null +++ b/tests-integration/tests/broadcast_spmc_unbounded_test.rs @@ -0,0 +1,740 @@ +// 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::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::AtomicUsize; +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; +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. +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() { + // `resubscribe` takes the waiter mutex; `unread_message_count` does not. + let _ = probe.resubscribe(); + } + } + } +} + +/// 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 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(); + + 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); + + 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(1024); + assert_eq!(rx.recv().await, Ok(1024)); +} + +#[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)); +} + +#[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(); + + 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 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 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(); + 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"], + )); } }