diff --git a/CHANGELOG.md b/CHANGELOG.md index 78bc8ac..0457e74 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file. ### New features +* Implement `broadcast::mpmc::unbounded`, an unbounded broadcast channel that retains messages until all active receivers consume them or are dropped. * Add an opt-in `asyncband::blocking::FutureExt` bridge with `block_on` and `wait_timeout` methods for waiting on runtime-agnostic futures from synchronous code. * Add opt-in bounded and unbounded runtime-agnostic object pools under `asyncband::pool`. * Add opt-in `asyncband::once::LazyCell` for values that own one asynchronous initializer and preserve its in-flight future across caller cancellation. diff --git a/README.md b/README.md index bcec6fe..018e6ea 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/) | `shutdown` | Coordinate shutdown signals and completion. | | Channels | [`oneshot`](https://docs.rs/asyncband/*/asyncband/oneshot/) | `oneshot` | Send one value between two tasks. | | | [`mpsc`](https://docs.rs/asyncband/*/asyncband/mpsc/) | `mpsc` | Send values from multiple producers through bounded or unbounded channels. | +| | [`broadcast::mpmc::unbounded`](https://docs.rs/asyncband/*/asyncband/broadcast/mpmc/fn.unbounded.html) | `broadcast` | Broadcast values from multiple producers and retain them until every active receiver consumes them. | | Resource reuse | [`pool`](https://docs.rs/asyncband/*/asyncband/pool/) | `pool` | Reuse objects through bounded or unbounded pool variants. | | Workload coordination | [`Semaphore`](https://docs.rs/asyncband/*/asyncband/semaphore/struct.Semaphore.html) | `semaphore` | Control concurrent access with permits. | | | [`Group`](https://docs.rs/asyncband/*/asyncband/singleflight/struct.Group.html) | `singleflight` | Coalesce concurrent calls for the same key. | diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index 3871529..aafe547 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -46,6 +46,7 @@ default = [] barrier = [] blocking = [] +broadcast = [] condvar = ["mutex"] latch = [] lazy-cell = ["mutex"] diff --git a/asyncband/src/channel/broadcast/mod.rs b/asyncband/src/channel/broadcast/mod.rs new file mode 100644 index 0000000..81e56c5 --- /dev/null +++ b/asyncband/src/channel/broadcast/mod.rs @@ -0,0 +1,20 @@ +// 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. + +//! Broadcast channels grouped by producer topology. + +pub mod mpmc; diff --git a/asyncband/src/channel/broadcast/mpmc/mod.rs b/asyncband/src/channel/broadcast/mpmc/mod.rs new file mode 100644 index 0000000..9249df2 --- /dev/null +++ b/asyncband/src/channel/broadcast/mpmc/mod.rs @@ -0,0 +1,26 @@ +// 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. + +//! Multi-producer, multi-consumer broadcast channels. + +mod unbounded; + +pub use self::unbounded::RecvError; +pub use self::unbounded::TryRecvError; +pub use self::unbounded::UnboundedReceiver; +pub use self::unbounded::UnboundedSender; +pub use self::unbounded::unbounded; diff --git a/asyncband/src/channel/broadcast/mpmc/unbounded.rs b/asyncband/src/channel/broadcast/mpmc/unbounded.rs new file mode 100644 index 0000000..594a125 --- /dev/null +++ b/asyncband/src/channel/broadcast/mpmc/unbounded.rs @@ -0,0 +1,766 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A multi-producer multi-consumer broadcast channel with an unbounded buffer. +//! +//! This channel supports multiple senders and multiple receivers. Each message sent by any +//! sender is received by all active receivers. If a receiver falls behind, messages are buffered +//! until the receiver consumes them or is dropped. +//! +//! # Memory usage +//! +//! This channel does not impose a capacity limit. A slow or stalled receiver can cause the +//! buffer to grow without bound, because messages are retained until every active receiver has +//! consumed them or the receiver is dropped. Use [`UnboundedSender::buffer_len`] to monitor the +//! number of messages currently retained by the shared buffer. +//! +//! The buffer keeps the capacity a steady workload needs, so a channel that repeatedly fills and +//! drains does not reallocate. Capacity grown for a one-off burst is released once a later cycle +//! drains completely without needing it. +//! +//! # Receivers +//! +//! Each receiver has an independent cursor. Use [`UnboundedSender::subscribe`] or +//! [`UnboundedReceiver::resubscribe`] to create a receiver that starts at the current tail. +//! +//! Messages are reclaimed once the slowest receiver moves past them, which scans one slot per +//! receiver. Only the receive that advances the slowest cursor pays for that scan, and the channel +//! keeps a slot for every receiver it hands out, so the cost follows the largest number of +//! receivers that were ever active at once rather than the number active now. +//! +//! # Examples +//! +//! Basic usage: +//! +//! ``` +//! use asyncband::broadcast::mpmc; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (tx, mut rx1) = mpmc::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)); +//! # } +//! ``` +//! +//! Slow receivers do not miss messages: +//! +//! ``` +//! use asyncband::broadcast::mpmc; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (tx, mut rx1) = mpmc::unbounded(); +//! let mut rx2 = tx.subscribe(); +//! +//! tx.send(1); +//! tx.send(2); +//! +//! // One receiver draining the channel does not discard what the other has not read yet. +//! assert_eq!(rx1.recv().await, Ok(1)); +//! assert_eq!(rx1.recv().await, Ok(2)); +//! assert_eq!(tx.buffer_len(), 2); +//! +//! assert_eq!(rx2.recv().await, Ok(1)); +//! assert_eq!(rx2.recv().await, Ok(2)); +//! assert_eq!(tx.buffer_len(), 0); +//! # } +//! ``` + +use std::collections::VecDeque; +use std::fmt; +use std::future::Future; +use std::mem; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; + +use crate::internal::arena::Arena; +use crate::internal::arena::SlotId; +use crate::internal::mutex::Mutex; +use crate::internal::waitset::WaitSet; +use crate::internal::waitset::WakerToken; + +#[cfg(test)] +mod tests; + +/// Creates a new broadcast channel with an unbounded buffer. +/// +/// Every accepted value is retained until all active receivers consume it or are dropped. +/// +/// # Examples +/// +/// ``` +/// use asyncband::broadcast::mpmc; +/// +/// let (tx, mut rx) = mpmc::unbounded(); +/// tx.send(10); +/// assert_eq!(rx.try_recv(), Ok(10)); +/// ``` +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { + let mut receivers = Arena::new(); + let key = receivers.insert(0); + let shared = Arc::new(Shared { + inner: Mutex::new(Inner { + buffer: VecDeque::new(), + head: 0, + head_receivers: 1, + tail: 0, + receivers, + peak_len: 0, + waiters: WaitSet::new(), + }), + senders: AtomicUsize::new(1), + }); + let sender = UnboundedSender { + shared: shared.clone(), + }; + let receiver = UnboundedReceiver { shared, key }; + (sender, receiver) +} + +/// Error returned by [`UnboundedReceiver::recv`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RecvError { + /// The sender has become disconnected, and there will never be any more data received on it. + Disconnected, +} + +impl fmt::Display for RecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + RecvError::Disconnected => write!(f, "receiving on a closed channel"), + } + } +} + +impl std::error::Error for RecvError {} + +/// Error returned by [`UnboundedReceiver::try_recv`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TryRecvError { + /// This channel is currently empty, but the sender(s) have not yet disconnected, so data may + /// yet become available. + Empty, + /// The sender has become disconnected, and there will never be any more data received on it. + Disconnected, +} + +impl fmt::Display for TryRecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + TryRecvError::Empty => write!(f, "receiving on an empty channel"), + TryRecvError::Disconnected => write!(f, "receiving on a closed channel"), + } + } +} + +impl std::error::Error for TryRecvError {} + +/// Retained capacity below which the shared buffer is never shrunk back. +const MIN_RETAINED_CAPACITY: usize = 64; + +struct Inner { + /// Messages whose versions are in the range `[head, tail)`. + /// + /// Each message is held behind an `Arc` so a receive can hand the payload out of the critical + /// section. Cloning the `Arc` under the lock keeps `T::clone` — and, for reclaimed messages, + /// `T::drop` — outside it, which matters because both are arbitrary user code that may call + /// back into this channel. + buffer: VecDeque>, + /// The version of the first message in `buffer`. + head: u64, + /// The number of active receivers whose cursor equals `head`. + head_receivers: usize, + /// The next message version to assign. + tail: u64, + /// Cursor for each active receiver. + receivers: Arena, + /// The largest backlog retained since the buffer was last empty. + peak_len: usize, + /// Receivers parked in [`UnboundedReceiver::recv`]. + waiters: WaitSet, +} + +impl Inner { + fn insert_receiver(&mut self, head: u64) -> SlotId { + if head == self.head { + self.head_receivers += 1; + } + + self.receivers.insert(head) + } + + fn remove_receiver(&mut self, key: SlotId) -> Vec> { + let head = self.receivers.remove(key); + + if head == self.head { + self.release_head_receiver() + } else { + Vec::new() + } + } + + fn advance_receiver(&mut self, key: SlotId, next_head: u64) -> Vec> { + let head = *self + .receivers + .get(key) + .expect("active broadcast receiver must be registered"); + *self + .receivers + .get_mut(key) + .expect("active broadcast receiver must be registered") = next_head; + + if head == self.head { + self.release_head_receiver() + } else { + Vec::new() + } + } + + fn release_head_receiver(&mut self) -> Vec> { + self.head_receivers -= 1; + + if self.head_receivers == 0 { + self.reclaim_consumed() + } else { + Vec::new() + } + } + + fn receive(&mut self, key: SlotId) -> Option<(Arc, Vec>)> { + let head = *self + .receivers + .get(key) + .expect("active broadcast receiver must be registered"); + + if head < self.tail { + debug_assert!(head >= self.head); + let offset = (head - self.head) as usize; + let msg = self.buffer[offset].clone(); + let reclaimed = self.advance_receiver(key, head + 1); + // A reclaim triggered by this receive always begins with this receiver's own message: + // the reclaim path runs only for a cursor sitting at `head`, so the first slot drained + // is `msg`. `take_msg` relies on this to recognise that it owns the payload. + debug_assert!( + reclaimed + .first() + .is_none_or(|first| Arc::ptr_eq(first, &msg)) + ); + Some((msg, reclaimed)) + } else { + None + } + } + + fn reclaim_consumed(&mut self) -> Vec> { + let mut next_head = self.tail; + let mut head_receivers = 0; + + for head in self.receivers.values() { + if *head < next_head { + next_head = *head; + head_receivers = 1; + } else if *head == next_head { + head_receivers += 1; + } + } + + debug_assert!(next_head >= self.head); + let consumed = usize::try_from(next_head - self.head) + .expect("retained broadcast message count exceeds usize"); + // Move reclaimed messages out so their Drop impls run after `inner` is unlocked. + let reclaimed = self.buffer.drain(..consumed).collect(); + + self.head = next_head; + self.head_receivers = head_receivers; + self.shrink_buffer(); + reclaimed + } + + /// Returns the allocation grown for a stalled receiver once that backlog is behind us. + /// + /// Without this, a single burst pins its peak allocation for the lifetime of the channel. + /// The decision is deliberately made only when the buffer drains completely, and against the + /// peak of the cycle that just ended rather than the current length: a channel that repeatedly + /// fills and drains keeps a peak as large as its bursts, so it holds its allocation instead of + /// reallocating on every cycle. Only once a full cycle stays small does the buffer give the + /// memory back. + fn shrink_buffer(&mut self) { + if !self.buffer.is_empty() { + return; + } + + let peak = mem::take(&mut self.peak_len); + let capacity = self.buffer.capacity(); + if capacity > MIN_RETAINED_CAPACITY && peak <= capacity / 4 { + self.buffer.shrink_to(MIN_RETAINED_CAPACITY.max(peak * 2)); + } + } +} + +struct Shared { + /// Buffer, receiver cursors, and parked receivers, all under a single lock. + /// + /// The wait set lives here rather than beside it so that publishing a message and draining the + /// waiters happen in one critical section. That is what makes the park path race-free: a + /// receiver that finds no message and then registers still holds this lock, so a concurrent + /// `send` cannot slip between the two steps and skip the wake-up. + inner: Mutex>, + /// Number of active senders. + senders: AtomicUsize, +} + +/// A sender handle to the broadcast channel. +/// +/// The sender can be cloned to create multiple producers. When all senders are dropped, +/// the channel is closed. +pub struct UnboundedSender { + shared: Arc>, +} + +impl Clone for UnboundedSender { + fn clone(&self) -> Self { + // Relaxed is enough because this count publishes nothing on its own: receivers read it + // only to decide whether the channel is closed, and every message it could hide is + // published under `inner`, which a receiver holds before it observes the count. + self.shared.senders.fetch_add(1, Ordering::Relaxed); + Self { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for 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) { + match self.shared.senders.fetch_sub(1, Ordering::AcqRel) { + 1 => { + // If this is the last sender, we need to wake up the receiver so it can + // observe the disconnected state. + let wakers = self.shared.inner.lock().waiters.take_wakers(); + for waker in wakers { + waker.wake(); + } + } + _ => { + // there are still other senders left, do nothing + } + } + } +} + +impl UnboundedSender { + /// Broadcasts a value to all active receivers. + /// + /// This operation does not wait for receiver capacity. If receivers fall behind, messages + /// remain buffered until all active receivers have consumed them or the lagging receivers + /// are dropped. + /// + /// If no receivers are active, the message is dropped immediately. + /// + /// # Panics + /// + /// Panics if the internal message version counter overflows. After `u64::MAX` successful sends + /// on one channel instance, the next send panics. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::unbounded(); + /// tx.send(10); + /// assert_eq!(rx.try_recv(), Ok(10)); + /// ``` + pub fn send(&self, msg: T) { + let msg = Arc::new(msg); + + // Publishing and draining the wait set share one critical section, so a receiver can never + // observe an empty buffer and park after this message became visible. + let wakers = { + let mut inner = self.shared.inner.lock(); + inner.tail = inner + .tail + .checked_add(1) + .expect("broadcast channel version counter overflowed"); + + if inner.receivers.is_empty() { + // No receivers means no one will read this message; advance `head` so the + // invariant that `buffer` covers versions `[head, tail)` still holds without + // buffering anything. The buffer is already drained when the last receiver was + // dropped, so there is nothing to clear here. + debug_assert!(inner.buffer.is_empty()); + debug_assert_eq!(inner.head_receivers, 0); + inner.head = inner.tail; + } else { + inner.buffer.push_back(msg); + inner.peak_len = inner.peak_len.max(inner.buffer.len()); + } + + inner.waiters.take_wakers() + }; + + // Notify all waiting receivers. An unsent message is dropped here too, once the lock is + // released. + for waker in wakers { + waker.wake(); + } + } + + /// Returns the number of messages currently retained by the shared buffer. + /// + /// 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. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::unbounded(); + /// tx.send(10); + /// assert_eq!(tx.buffer_len(), 1); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(tx.buffer_len(), 0); + /// ``` + pub fn buffer_len(&self) -> usize { + self.shared.inner.lock().buffer.len() + } + + /// Returns the number of active receivers. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, rx) = mpmc::unbounded::(); + /// assert_eq!(tx.receiver_count(), 1); + /// + /// let rx2 = tx.subscribe(); + /// assert_eq!(tx.receiver_count(), 2); + /// + /// drop(rx); + /// drop(rx2); + /// assert_eq!(tx.receiver_count(), 0); + /// ``` + pub fn receiver_count(&self) -> usize { + self.shared.inner.lock().receivers.len() + } + + /// Creates a new receiver that starts receiving messages from the current tail of the channel. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// use asyncband::broadcast::mpmc::TryRecvError; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, _) = mpmc::unbounded(); + /// tx.send(10); + /// + /// let mut rx = tx.subscribe(); + /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + /// tx.send(20); + /// assert_eq!(rx.recv().await, Ok(20)); + /// # } + /// ``` + pub fn subscribe(&self) -> UnboundedReceiver { + let mut inner = self.shared.inner.lock(); + let head = inner.tail; + let key = inner.insert_receiver(head); + let shared = self.shared.clone(); + UnboundedReceiver { shared, key } + } +} + +/// A receiver handle to the broadcast channel. +/// +/// Each receiver sees every message sent to the channel while the receiver is active. +pub struct UnboundedReceiver { + shared: Arc>, + key: SlotId, +} + +impl fmt::Debug for UnboundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() + } +} + +impl Drop for UnboundedReceiver { + fn drop(&mut self) { + let reclaimed = { + let mut inner = self.shared.inner.lock(); + inner.remove_receiver(self.key) + }; + drop(reclaimed); + } +} + +impl UnboundedReceiver { + /// Receives the next value for this receiver. + /// + /// # Returns + /// + /// * `Ok(T)`: The next message. + /// * `Err(RecvError::Disconnected)`: All senders have been dropped and no more messages are + /// available. + /// + /// # Cancel safety + /// + /// This method is cancel safe. If `recv` is used as the event in a `select` statement and some + /// other branch completes first, it is guaranteed that no messages were received on this + /// channel. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, mut rx) = mpmc::unbounded(); + /// tx.send(10); + /// assert_eq!(rx.recv().await, Ok(10)); + /// # } + /// ``` + pub async fn recv(&mut self) -> Result { + Recv { + receiver: self, + registration: None, + } + .await + } + + /// Attempts to receive the next value for this receiver without blocking. + /// + /// # Returns + /// + /// * `Ok(T)`: The next message. + /// * `Err(TryRecvError::Empty)`: No message is currently available. + /// * `Err(TryRecvError::Disconnected)`: All senders have been dropped and no more messages are + /// available. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::unbounded(); + /// tx.send(10); + /// assert_eq!(rx.try_recv(), Ok(10)); + /// ``` + pub fn try_recv(&mut self) -> Result { + let (msg, reclaimed) = self.try_recv_shared()?; + Ok(take_msg(msg, reclaimed)) + } +} + +/// Drops the reclaimed backlog, then yields the received message, both with the channel unlocked. +/// +/// A non-empty backlog means this receive drained `msg` from the buffer, so once the backlog is +/// dropped this receive holds the only reference and the payload can be moved out instead of +/// cloned. A channel with a single receiver therefore never clones a payload. +/// +/// Ownership is decided from that bookkeeping rather than by probing the reference count. An +/// [`Arc::try_unwrap`] on every receive would fail under fan-out, and its failed compare-exchange +/// writes to a cache line that every receiver draining the message shares. +fn take_msg(msg: Arc, reclaimed: Vec>) -> T { + let sole_owner = !reclaimed.is_empty(); + drop(reclaimed); + + if !sole_owner { + return (*msg).clone(); + } + + // Another receiver can still hold an in-flight reference to the same message, so the clone + // remains the fallback. + match Arc::try_unwrap(msg) { + Ok(msg) => msg, + Err(msg) => (*msg).clone(), + } +} + +impl UnboundedReceiver { + fn try_recv_shared(&mut self) -> Result<(Arc, Vec>), TryRecvError> { + // Check this receiver's cursor while holding `inner` before observing `senders`. Senders + // append messages under the same lock before they can be dropped, so an empty result here + // means this receiver has no unread buffered message. + let mut inner = self.shared.inner.lock(); + if let Some(received) = inner.receive(self.key) { + return Ok(received); + } + + if self.shared.senders.load(Ordering::Acquire) == 0 { + Err(TryRecvError::Disconnected) + } else { + Err(TryRecvError::Empty) + } + } + + /// Re-subscribes to the channel, returning a new receiver that starts receiving messages from + /// the *current* tail of the channel. + /// + /// This is useful if the receiver wants to jump to the latest message, skipping everything in + /// between. The original receiver is unchanged and continues to retain its own backlog until + /// it consumes those messages or is dropped. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::unbounded(); + /// tx.send(1); + /// tx.send(2); + /// + /// let mut rx2 = rx.resubscribe(); + /// tx.send(3); + /// + /// assert_eq!(rx2.try_recv(), Ok(3)); + /// ``` + pub fn resubscribe(&self) -> Self { + let mut inner = self.shared.inner.lock(); + let head = inner.tail; + let key = inner.insert_receiver(head); + let shared = self.shared.clone(); + Self { shared, key } + } + + /// Returns the number of messages this receiver can still read. + /// + /// This count is specific to this receiver, unlike [`UnboundedSender::buffer_len`], which + /// reports the shared backlog retained by the slowest active receiver. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::unbounded(); + /// assert_eq!(rx.len(), 0); + /// + /// tx.send(10); + /// tx.send(20); + /// assert_eq!(rx.len(), 2); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(rx.len(), 1); + /// ``` + pub fn len(&self) -> usize { + let inner = self.shared.inner.lock(); + let head = *inner + .receivers + .get(self.key) + .expect("active broadcast receiver must be registered"); + usize::try_from(inner.tail - head).expect("unread broadcast message count exceeds usize") + } + + /// Returns `true` if this receiver has no currently available messages. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, rx) = mpmc::unbounded(); + /// assert!(rx.is_empty()); + /// + /// tx.send(10); + /// assert!(!rx.is_empty()); + /// ``` + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +struct Recv<'a, T> { + receiver: &'a mut UnboundedReceiver, + registration: Option, +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + // Ready paths clear the registration, so only a cancelled pending receive takes this lock. + if self.registration.is_none() { + return; + } + + let waker = { + let mut inner = self.receiver.shared.inner.lock(); + inner.waiters.unregister_waker(&mut self.registration) + }; + drop(waker); + } +} + +impl Future for Recv<'_, T> { + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let Self { + receiver, + registration, + } = self.get_mut(); + + // One critical section decides between all three outcomes. Senders append messages and + // drain the wait set under this same lock, so registering here cannot miss a wake-up and + // cannot observe a closed channel that still has a message for this receiver. + let received = { + let mut inner = receiver.shared.inner.lock(); + + match inner.receive(receiver.key) { + Some(received) => received, + None => { + if receiver.shared.senders.load(Ordering::Acquire) == 0 { + *registration = None; + return Poll::Ready(Err(RecvError::Disconnected)); + } + + let waker = inner.waiters.register_waker(registration, cx); + drop(inner); + drop(waker); + return Poll::Pending; + } + } + }; + + let (msg, reclaimed) = received; + *registration = None; + Poll::Ready(Ok(take_msg(msg, reclaimed))) + } +} diff --git a/asyncband/src/channel/broadcast/mpmc/unbounded/tests.rs b/asyncband/src/channel/broadcast/mpmc/unbounded/tests.rs new file mode 100644 index 0000000..e3f6746 --- /dev/null +++ b/asyncband/src/channel/broadcast/mpmc/unbounded/tests.rs @@ -0,0 +1,69 @@ +// 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::*; + +#[test] +#[should_panic(expected = "broadcast channel version counter overflowed")] +fn send_panics_on_version_overflow() { + // The receiver is dropped right away: the doctored counter would make its own drop overflow. + let (tx, _) = unbounded(); + tx.shared.inner.lock().tail = u64::MAX; + tx.send(()); +} + +#[test] +fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { + let (tx, mut rx) = unbounded(); + + let burst = MIN_RETAINED_CAPACITY * 16; + for i in 0..burst { + tx.send(i); + } + assert!(tx.shared.inner.lock().buffer.capacity() >= burst); + + for i in 0..burst { + assert_eq!(rx.try_recv(), Ok(i)); + } + + // Draining evaluates the cycle that just peaked, so the burst allocation is still held. + assert_eq!(tx.buffer_len(), 0); + assert!(tx.shared.inner.lock().buffer.capacity() >= burst); + + // The next cycle stays small, which is what releases the memory. + tx.send(0); + assert_eq!(rx.try_recv(), Ok(0)); + assert!(tx.shared.inner.lock().buffer.capacity() < burst); +} + +#[test] +fn repeated_bursts_keep_their_allocation() { + let (tx, mut rx) = unbounded(); + let burst = MIN_RETAINED_CAPACITY * 4; + + for _ in 0..4 { + for i in 0..burst { + tx.send(i); + } + for i in 0..burst { + assert_eq!(rx.try_recv(), Ok(i)); + } + } + + // Every cycle peaks at the same size, so the buffer must not rebuild its allocation each time. + assert!(tx.shared.inner.lock().buffer.capacity() >= burst); +} diff --git a/asyncband/src/channel/mod.rs b/asyncband/src/channel/mod.rs index 58689a6..bfe3be0 100644 --- a/asyncband/src/channel/mod.rs +++ b/asyncband/src/channel/mod.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +#[cfg(feature = "broadcast")] +pub mod broadcast; #[cfg(feature = "mpsc")] pub mod mpsc; #[cfg(feature = "oneshot")] diff --git a/asyncband/src/internal/arena.rs b/asyncband/src/internal/arena.rs index 98502aa..3158d31 100644 --- a/asyncband/src/internal/arena.rs +++ b/asyncband/src/internal/arena.rs @@ -136,6 +136,21 @@ impl Arena { } } + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + pub fn values(&self) -> impl Iterator { + self.slots.iter().filter_map(|slot| match slot { + Slot::Occupied(value) => Some(value), + Slot::Vacant { .. } => None, + }) + } + /// Removes the value stored at `id`. /// /// # Panics @@ -194,11 +209,6 @@ impl Arena { self.len = 0; values.into_iter() } - - #[cfg(test)] - pub fn len(&self) -> usize { - self.len - } } #[cfg(test)] diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index cdd33e3..2374f39 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -20,6 +20,7 @@ pub(crate) mod atomic_waker; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", feature = "mpsc", feature = "mutex", @@ -30,7 +31,7 @@ pub(crate) mod atomic_waker; // `WaitList` and `WaitSet` use different `Arena` operations. A single-primitive build therefore // leaves part of this shared API unused, while the all-feature build uses it. #[allow(dead_code)] -mod arena; +pub(crate) mod arena; #[cfg(any(feature = "latch", feature = "once", feature = "waitgroup"))] // `waitgroup` increments and decrements the countdown, while `latch` and `once` only decrement it. @@ -52,6 +53,7 @@ pub(crate) mod value_cell; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", feature = "mpsc", feature = "mutex", @@ -83,6 +85,7 @@ pub(crate) mod waitlist; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", feature = "once", feature = "waitgroup", diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 733275a..aa9dbf4 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -56,7 +56,7 @@ //! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | //! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::LazyCell`], [`once::OnceMap`] | `once`, `once-cell`, `lazy-cell`, `once-map` | //! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | -//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`] | `oneshot`, `mpsc` | +//! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::mpmc::unbounded`] | `oneshot`, `mpsc`, `broadcast` | //! | Reuse objects | [`pool::bounded`], [`pool::unbounded`] | `pool` | //! | Coordinate workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | //! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | @@ -107,6 +107,8 @@ mod internal; pub mod barrier; #[cfg(feature = "blocking")] pub mod blocking; +#[cfg(feature = "broadcast")] +pub use self::channel::broadcast; #[cfg(feature = "condvar")] pub mod condvar; #[cfg(feature = "latch")] diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index c082f90..95eec00 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -27,6 +27,7 @@ async-channel = { workspace = true } asyncband = { workspace = true, features = [ "barrier", "blocking", + "broadcast", "condvar", "latch", "mpsc", diff --git a/benchmarks/asyncband/broadcast/mod.rs b/benchmarks/asyncband/broadcast/mod.rs new file mode 100644 index 0000000..b078f4a --- /dev/null +++ b/benchmarks/asyncband/broadcast/mod.rs @@ -0,0 +1,18 @@ +// 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 mpmc; diff --git a/benchmarks/asyncband/broadcast/mpmc/mod.rs b/benchmarks/asyncband/broadcast/mpmc/mod.rs new file mode 100644 index 0000000..78ef889 --- /dev/null +++ b/benchmarks/asyncband/broadcast/mpmc/mod.rs @@ -0,0 +1,18 @@ +// 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 unbounded; diff --git a/benchmarks/asyncband/broadcast/mpmc/unbounded.rs b/benchmarks/asyncband/broadcast/mpmc/unbounded.rs new file mode 100644 index 0000000..fa35c27 --- /dev/null +++ b/benchmarks/asyncband/broadcast/mpmc/unbounded.rs @@ -0,0 +1,356 @@ +// 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 std::sync::Arc; +use std::sync::Barrier; +use std::thread; +use std::thread::JoinHandle; + +use asyncband::broadcast::mpmc; +use divan::Bencher; +use divan::black_box; + +use crate::support::bench_context; +use crate::support::poll_pending; +use crate::support::poll_pinned_ready; + +const RECEIVER_COUNTS: &[usize] = &[1, 8, 32]; +const CONCURRENCY_COUNTS: &[usize] = &[1, 2, 4, 8]; +const CONCURRENT_BATCH_SIZE: usize = 4096; + +/// 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 }, +]; + +struct ConcurrentSend { + receiver: mpmc::UnboundedReceiver, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl ConcurrentSend { + fn new(sender_count: usize) -> Self { + let (sender, receiver) = mpmc::unbounded(); + let ready = Arc::new(Barrier::new(sender_count + 1)); + let start = Arc::new(Barrier::new(sender_count + 1)); + let done = Arc::new(Barrier::new(sender_count + 1)); + let sends_per_worker = CONCURRENT_BATCH_SIZE / sender_count; + let mut workers = Vec::with_capacity(sender_count); + + for worker_index in 0..sender_count { + let sender = sender.clone(); + let ready = ready.clone(); + let start = start.clone(); + let done = done.clone(); + workers.push(thread::spawn(move || { + ready.wait(); + start.wait(); + let first = worker_index * sends_per_worker; + for value in first..first + sends_per_worker { + sender.send(black_box(value)); + } + done.wait(); + })); + } + drop(sender); + ready.wait(); + + Self { + receiver, + start, + done, + workers, + } + } + + // The drain is inside the measured region on purpose: it is what keeps the backlog bounded + // across samples, and reclaiming the batch is part of the cost of an unbounded send. + fn run(&mut self) { + self.start.wait(); + self.done.wait(); + while let Ok(value) = self.receiver.try_recv() { + black_box(value); + } + } +} + +impl Drop for ConcurrentSend { + fn drop(&mut self) { + for worker in self.workers.drain(..) { + worker.join().unwrap(); + } + } +} + +struct ConcurrentFanout { + sender: mpmc::UnboundedSender, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl ConcurrentFanout { + fn new(receiver_count: usize) -> Self { + let (sender, receiver) = mpmc::unbounded(); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(sender.subscribe()); + } + + let ready = Arc::new(Barrier::new(receiver_count + 1)); + let start = Arc::new(Barrier::new(receiver_count + 1)); + let done = Arc::new(Barrier::new(receiver_count + 1)); + let mut workers = Vec::with_capacity(receiver_count); + + for mut receiver in receivers { + let ready = ready.clone(); + let start = start.clone(); + let done = done.clone(); + workers.push(thread::spawn(move || { + ready.wait(); + start.wait(); + let result = (0..CONCURRENT_BATCH_SIZE).try_for_each(|_| { + receiver.try_recv().map(|value| { + black_box(value); + }) + }); + done.wait(); + result.unwrap(); + })); + } + ready.wait(); + + Self { + sender, + start, + done, + workers, + } + } + + fn run(&mut self) { + for value in 0..CONCURRENT_BATCH_SIZE { + self.sender.send(black_box(value)); + } + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for ConcurrentFanout { + fn drop(&mut self) { + for worker in self.workers.drain(..) { + worker.join().unwrap(); + } + } +} + +#[divan::bench] +fn send_without_receivers(bencher: Bencher) { + let (sender, receiver) = mpmc::unbounded::(); + drop(receiver); + bencher.bench_local(|| sender.send(black_box(1))); +} + +#[divan::bench] +fn try_recv_empty(bencher: Bencher) { + let (sender, mut receiver) = mpmc::unbounded::(); + bencher.bench_local(|| black_box(receiver.try_recv())); + black_box(sender); +} + +// A sole receiver takes ownership of the payload, so this path never clones the message. +#[divan::bench] +fn send_and_try_recv(bencher: Bencher) { + let (sender, mut receiver) = mpmc::unbounded(); + bencher.bench_local(|| { + sender.send(black_box(1usize)); + black_box(receiver.try_recv().unwrap()) + }); +} + +// 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 (sender, mut first) = mpmc::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 (sender, mut receiver) = mpmc::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 (sender, mut first) = mpmc::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 (sender, receiver) = mpmc::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( + args = CONCURRENCY_COUNTS, + sample_count = 50, + sample_size = 1, + counters = [CONCURRENT_BATCH_SIZE] +)] +fn concurrent_send_and_drain(bencher: Bencher, sender_count: usize) { + bencher + .with_inputs(|| ConcurrentSend::new(sender_count)) + .bench_local_refs(ConcurrentSend::run); +} + +#[divan::bench( + args = CONCURRENCY_COUNTS, + sample_count = 50, + sample_size = 1, + counters = [CONCURRENT_BATCH_SIZE] +)] +fn concurrent_fanout(bencher: Bencher, receiver_count: usize) { + bencher + .with_inputs(|| ConcurrentFanout::new(receiver_count)) + .bench_local_refs(ConcurrentFanout::run); +} + +#[divan::bench] +fn cancel_pending(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (sender, mut receiver) = mpmc::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 (sender, mut receiver) = mpmc::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 (sender, receiver) = mpmc::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/asyncband/main.rs b/benchmarks/asyncband/main.rs index cb8ac1a..c383615 100644 --- a/benchmarks/asyncband/main.rs +++ b/benchmarks/asyncband/main.rs @@ -17,6 +17,7 @@ mod barrier; mod blocking; +mod broadcast; mod condvar; mod latch; mod mpsc; diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 7ea541f..b66fda0 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -29,6 +29,7 @@ tokio = { workspace = true, features = ["full"] } asyncband = { workspace = true, features = [ "barrier", "blocking", + "broadcast", "condvar", "latch", "lazy-cell", diff --git a/tests-integration/tests/broadcast_mpmc_unbounded_test.rs b/tests-integration/tests/broadcast_mpmc_unbounded_test.rs new file mode 100644 index 0000000..7617ac4 --- /dev/null +++ b/tests-integration/tests/broadcast_mpmc_unbounded_test.rs @@ -0,0 +1,609 @@ +// 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::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Wake; +use std::task::Waker; +use std::thread; +use std::time::Duration; +use std::time::Instant; + +use asyncband::broadcast::mpmc::*; + +struct TrackWake(AtomicUsize); + +impl Wake for TrackWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +/// A payload whose destructor re-enters the channel it was sent through. +struct Reentrant { + value: u64, + channel: Option>, +} + +impl Clone for Reentrant { + fn clone(&self) -> Self { + Self { + value: self.value, + channel: self.channel.clone(), + } + } +} + +impl Drop for Reentrant { + fn drop(&mut self) { + if let Some(channel) = &self.channel { + // Deadlocks if the channel still holds its lock while dropping reclaimed messages. + let _ = channel.buffer_len(); + let _ = channel.receiver_count(); + } + } +} + +/// A payload that panics while a shared receive clones it. +#[derive(Debug)] +struct PanicOnClone { + value: u64, + panic: bool, +} + +impl Clone for PanicOnClone { + fn clone(&self) -> Self { + if self.panic { + panic!("panic while cloning a broadcast message"); + } + Self { + value: self.value, + panic: self.panic, + } + } +} + +struct Rng(u64); + +impl Rng { + fn below(&mut self, n: u64) -> u64 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + x % n + } +} + +#[tokio::test] +async fn test_broadcast_basic() { + let (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)); +} + +#[tokio::test] +async fn test_subscribe() { + let (tx, _rx) = unbounded(); + let mut rx = tx.subscribe(); + + tx.send(100); + assert_eq!(rx.recv().await, Ok(100)); +} + +#[tokio::test] +async fn test_resubscribe() { + let (tx, mut rx) = unbounded(); + + tx.send(1); + tx.send(2); + + let mut rx2 = rx.resubscribe(); + + // rx sees 1, 2 + // rx2 sees nothing yet (starts at tail=2) + + tx.send(3); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx2.recv().await, Ok(3)); +} + +#[test] +fn test_try_recv() { + let (tx, mut rx) = unbounded(); + + // Empty + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + // Success + tx.send(10); + assert_eq!(rx.try_recv(), Ok(10)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + // Closed + drop(tx); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[tokio::test] +async fn test_slow_receiver_keeps_every_message() { + let (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.buffer_len(), 1024); + + for i in 0..1024 { + assert_eq!(rx2.recv().await, Ok(i)); + } + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn buffer_len_tracks_the_slowest_receiver() { + let (tx, mut rx1) = unbounded(); + let mut rx2 = tx.subscribe(); + + tx.send(1); + tx.send(2); + assert_eq!(tx.buffer_len(), 2); + + // Reclaiming waits for the slowest receiver, message by message. + assert_eq!(rx1.recv().await, Ok(1)); + assert_eq!(tx.buffer_len(), 2); + assert_eq!(rx2.recv().await, Ok(1)); + assert_eq!(tx.buffer_len(), 1); + + assert_eq!(rx1.recv().await, Ok(2)); + assert_eq!(tx.buffer_len(), 1); + assert_eq!(rx2.recv().await, Ok(2)); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn test_dropping_a_lagging_receiver_releases_its_backlog() { + let (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.buffer_len(), 128); + + drop(rx2); + assert_eq!(tx.buffer_len(), 0); +} + +#[tokio::test] +async fn resubscribe_keeps_the_original_receivers_backlog() { + let (tx, mut rx) = unbounded(); + + tx.send(1); + tx.send(2); + + let mut rx2 = rx.resubscribe(); + assert_eq!(tx.buffer_len(), 2); + + tx.send(3); + + assert_eq!(rx2.recv().await, Ok(3)); + assert_eq!(tx.buffer_len(), 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.buffer_len(), 0); +} + +#[tokio::test] +async fn send_without_receivers_does_not_buffer() { + let (tx, rx) = unbounded(); + drop(rx); + + tx.send(1); + tx.send(2); + assert_eq!(tx.buffer_len(), 0); + + let mut rx = tx.subscribe(); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + tx.send(3); + assert_eq!(rx.recv().await, Ok(3)); +} + +#[test] +fn receiver_count_and_len_track_each_receiver() { + let (tx, mut rx1) = unbounded(); + assert_eq!(tx.receiver_count(), 1); + assert_eq!(rx1.len(), 0); + assert!(rx1.is_empty()); + + tx.send(1); + tx.send(2); + assert_eq!(rx1.len(), 2); + assert!(!rx1.is_empty()); + + let mut rx2 = tx.subscribe(); + assert_eq!(tx.receiver_count(), 2); + assert_eq!(rx2.len(), 0); + assert!(rx2.is_empty()); + + tx.send(3); + assert_eq!(rx1.len(), 3); + assert_eq!(rx2.len(), 1); + + assert_eq!(rx2.try_recv(), Ok(3)); + assert_eq!(rx2.len(), 0); + drop(rx2); + assert_eq!(tx.receiver_count(), 1); + + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.len(), 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 (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 (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.buffer_len(), 0); + assert_eq!(rx1.try_recv().unwrap_err(), TryRecvError::Empty); +} + +#[test] +fn message_destructors_run_outside_the_channel_lock() { + let finished = Arc::new(AtomicUsize::new(0)); + let flag = finished.clone(); + + let worker = thread::spawn(move || { + let (tx, mut rx1) = unbounded(); + let rx2 = tx.subscribe(); + + for value in 0..4 { + tx.send(Reentrant { + value, + channel: Some(tx.clone()), + }); + } + + // 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); + + // With no receiver left, `send` drops the message itself; that must be unlocked too. + tx.send(Reentrant { + value: 4, + channel: Some(tx.clone()), + }); + drop(tx); + + flag.store(1, Ordering::SeqCst); + }); + + let deadline = Instant::now() + Duration::from_secs(10); + while Instant::now() < deadline && finished.load(Ordering::SeqCst) == 0 { + thread::sleep(Duration::from_millis(10)); + } + assert_eq!( + finished.load(Ordering::SeqCst), + 1, + "a message destructor deadlocked against the channel lock" + ); + worker.join().unwrap(); +} + +#[tokio::test] +async fn test_wait_mechanism() { + let (tx, mut rx) = unbounded(); + + let handle = tokio::spawn(async move { rx.recv().await }); + + tokio::time::sleep(Duration::from_millis(100)).await; + tx.send(42); + + assert_eq!(handle.await.unwrap(), Ok(42)); +} + +#[test] +fn send_wakes_a_parked_receiver_exactly_once() { + let (tx, mut rx) = unbounded(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + 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.0.load(Ordering::Relaxed), 1); + assert_eq!(recv.as_mut().poll(&mut context), Poll::Ready(Ok(42))); +} + +#[test] +fn cancelled_recv_releases_its_waker() { + let (tx, mut rx) = unbounded::<()>(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + 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.0.load(Ordering::Relaxed), 0); + assert_eq!(rx.try_recv(), Ok(())); +} + +#[test] +fn dropping_a_woken_recv_keeps_another_receivers_waiter() { + let (tx, mut rx1) = unbounded::(); + let mut rx2 = tx.subscribe(); + let first = Arc::new(TrackWake(AtomicUsize::new(0))); + 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.0.load(Ordering::Relaxed), 1); + assert_eq!(rx2.try_recv(), Ok(1)); + + let second = Arc::new(TrackWake(AtomicUsize::new(0))); + 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.0.load(Ordering::Relaxed), 1); +} + +#[test] +fn parked_recv_wakes_when_the_last_sender_drops() { + let (tx, mut rx) = unbounded::<()>(); + let extra = tx.clone(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + 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.0.load(Ordering::Relaxed), 0); + + drop(extra); + assert_eq!(tracker.0.load(Ordering::Relaxed), 1); + + drop(recv); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[test] +fn parked_recv_prefers_buffered_messages_over_disconnect() { + let (tx, mut rx) = unbounded(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + + tx.send(7); + drop(tx); + + assert_eq!(recv.as_mut().poll(&mut context), Poll::Ready(Ok(7))); + drop(recv); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); +} + +#[tokio::test] +async fn recv_drains_buffered_messages_before_reporting_disconnect() { + let (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_disconnect_without_any_message() { + let (tx, mut rx) = unbounded::<()>(); + drop(tx); + assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); +} + +#[test] +fn concurrent_senders_deliver_every_message_to_every_receiver() { + const SENDERS: u64 = 4; + const PER_SENDER: u64 = 512; + + let (tx, rx) = unbounded(); + let mut receivers = Vec::with_capacity(4); + receivers.push(rx); + receivers.extend((1..4).map(|_| tx.subscribe())); + + let senders = (0..SENDERS) + .map(|worker| { + let tx = tx.clone(); + thread::spawn(move || { + for value in 0..PER_SENDER { + tx.send(worker * PER_SENDER + value); + } + }) + }) + .collect::>(); + + let drains = receivers + .into_iter() + .map(|mut receiver| { + thread::spawn(move || { + let mut seen = Vec::new(); + while let Ok(value) = pollster::block_on(receiver.recv()) { + seen.push(value); + } + seen + }) + }) + .collect::>(); + + for sender in senders { + sender.join().unwrap(); + } + drop(tx); + + let expected = (0..SENDERS * PER_SENDER).collect::>(); + for drain in drains { + let mut seen = drain.join().unwrap(); + seen.sort_unstable(); + assert_eq!(seen, expected); + } +} + +#[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 (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); + } + _ => {} + } + + assert_eq!(tx.receiver_count(), model.len(), "seed {seed}"); + let retained = model + .iter() + .map(|(_, cursor)| *cursor) + .min() + .map_or(0, |slowest| tail - slowest); + assert_eq!(tx.buffer_len(), retained as usize, "seed {seed}"); + for (receiver, cursor) in &model { + assert_eq!(receiver.len(), (tail - cursor) as usize, "seed {seed}"); + assert_eq!(receiver.is_empty(), *cursor == tail, "seed {seed}"); + } + } + } +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index db52c62..847de16 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -18,6 +18,7 @@ use std::cell::Cell; use asyncband::barrier::Barrier; +use asyncband::broadcast; use asyncband::condvar::Condvar; use asyncband::latch::Latch; use asyncband::mpsc; @@ -85,6 +86,10 @@ fn public_types_are_send_and_sync() { assert_send_and_sync::>(); assert_send_and_sync::>(); assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::(); + assert_send_and_sync::(); assert_send_and_sync::>(); assert_send_and_sync::>(); assert_send_and_sync::>(); @@ -132,6 +137,10 @@ 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::>();