From 8e1f1f8bf64d85b18557eb703111e81c294e1c70 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 24 Aug 2026 05:32:05 +0800 Subject: [PATCH] feat(channel): implement terminal channel taxonomy --- CHANGELOG.md | 6 +- HISTORY.md | 2 + README.md | 8 +- asyncband/Cargo.toml | 6 + asyncband/src/channel/broadcast/internal.rs | 644 ++++++++++++++++++ asyncband/src/channel/broadcast/mod.rs | 36 + asyncband/src/channel/broadcast/mpmc.rs | 247 +++++++ asyncband/src/channel/broadcast/spmc.rs | 245 +++++++ asyncband/src/{mpsc => channel}/error.rs | 67 +- asyncband/src/channel/mod.rs | 43 ++ asyncband/src/channel/mpmc.rs | 192 ++++++ asyncband/src/{ => channel}/mpsc/bounded.rs | 8 +- asyncband/src/{ => channel}/mpsc/mod.rs | 10 +- asyncband/src/{ => channel}/mpsc/unbounded.rs | 6 +- asyncband/src/{ => channel}/oneshot/mod.rs | 0 .../src/{ => channel}/oneshot/receiver.rs | 18 +- asyncband/src/{ => channel}/oneshot/sender.rs | 16 +- asyncband/src/{ => channel}/oneshot/tests.rs | 0 asyncband/src/channel/queue.rs | 384 +++++++++++ asyncband/src/channel/spmc.rs | 197 ++++++ asyncband/src/channel/spsc.rs | 187 +++++ asyncband/src/channel/watch.rs | 288 ++++++++ asyncband/src/internal/arena.rs | 14 +- asyncband/src/internal/mod.rs | 23 +- asyncband/src/internal/waitset.rs | 6 + asyncband/src/lib.rs | 22 +- benchmarks/Cargo.toml | 5 + benchmarks/broadcast.rs | 274 ++++++++ benchmarks/main.rs | 3 + benchmarks/mpsc.rs | 46 ++ benchmarks/queue.rs | 84 +++ benchmarks/watch.rs | 42 ++ tests-integration/Cargo.toml | 5 + tests-integration/tests/broadcast_test.rs | 547 +++++++++++++++ .../tests/queue_topologies_test.rs | 339 +++++++++ tests-integration/tests/support/mod.rs | 53 ++ tests-integration/tests/traits_test.rs | 79 +++ tests-integration/tests/watch_test.rs | 228 +++++++ 38 files changed, 4328 insertions(+), 52 deletions(-) create mode 100644 asyncband/src/channel/broadcast/internal.rs create mode 100644 asyncband/src/channel/broadcast/mod.rs create mode 100644 asyncband/src/channel/broadcast/mpmc.rs create mode 100644 asyncband/src/channel/broadcast/spmc.rs rename asyncband/src/{mpsc => channel}/error.rs (80%) create mode 100644 asyncband/src/channel/mod.rs create mode 100644 asyncband/src/channel/mpmc.rs rename asyncband/src/{ => channel}/mpsc/bounded.rs (98%) rename asyncband/src/{ => channel}/mpsc/mod.rs (88%) rename asyncband/src/{ => channel}/mpsc/unbounded.rs (99%) rename asyncband/src/{ => channel}/oneshot/mod.rs (100%) rename asyncband/src/{ => channel}/oneshot/receiver.rs (98%) rename asyncband/src/{ => channel}/oneshot/sender.rs (97%) rename asyncband/src/{ => channel}/oneshot/tests.rs (100%) create mode 100644 asyncband/src/channel/queue.rs create mode 100644 asyncband/src/channel/spmc.rs create mode 100644 asyncband/src/channel/spsc.rs create mode 100644 asyncband/src/channel/watch.rs create mode 100644 benchmarks/broadcast.rs create mode 100644 benchmarks/queue.rs create mode 100644 benchmarks/watch.rs create mode 100644 tests-integration/tests/broadcast_test.rs create mode 100644 tests-integration/tests/queue_topologies_test.rs create mode 100644 tests-integration/tests/support/mod.rs create mode 100644 tests-integration/tests/watch_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a0d21b9..e7c6e48 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,13 +10,17 @@ All notable changes to this project will be documented in this file. * 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 SPSC, SPMC, and MPMC competing queues with topology-specific endpoint capabilities. +* Add opt-in lossless bounded and unbounded SPMC and MPMC broadcast channels. +* Add an opt-in latest-state watch channel. +* Add a `channel` umbrella feature that enables every channel API while keeping their public paths at the crate root. ### Breaking changes * Gate all exported primitives behind opt-in Cargo features and enable no features by default; downstream dependencies must explicitly enable the APIs they use. * Remove `admission::FairShare` and its `admission` Cargo feature from the feature set. * Remove the `asyncband::atomicbox` module and its `AtomicBox` and `AtomicOptionBox` types from the public API. -* Remove the lossy `broadcast::overflow` channel and its `broadcast` Cargo feature; future broadcast APIs will use explicit bounded and unbounded lossless semantics. +* Remove the lossy `broadcast::overflow` API; the `broadcast` Cargo feature now selects explicit bounded and unbounded lossless channels. * Remove `Semaphore::try_acquire_and_forget`, `Semaphore::acquire_and_forget`, `Semaphore::try_acquire_owned_and_forget`, and `Semaphore::acquire_owned_and_forget`; acquire a permit and call its `forget` method instead. * Rename `oneshot::Sender::is_closed` and `oneshot::Receiver::is_closed` to `is_disconnected`. * Replace `Semaphore::forget` with `Semaphore::drain_permits` and `Semaphore::forget_exact` with `Semaphore::reduce_permits`; permit-level `forget` methods are unchanged. diff --git a/HISTORY.md b/HISTORY.md index 9ffeb03..347e31b 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -9,6 +9,7 @@ Asyncband collects composable, runtime-agnostic concurrency building blocks info - `condvar::Condvar` is inspired by [`std::sync::Condvar`](https://doc.rust-lang.org/std/sync/struct.Condvar.html) and [`async_std::sync::Condvar`](https://docs.rs/async-std/latest/async_std/sync/struct.Condvar.html), with a fair FIFO waiter queue and standard non-buffered notification semantics. - `latch::Latch` is inspired by [`latches`](https://github.com/mirromutth/latches), with a different implementation based on the internal `CountdownState` primitive. - `mutex::Mutex` is derived from [`tokio::sync::Mutex`](https://docs.rs/tokio/latest/tokio/sync/struct.Mutex.html). +- The cloneable competing-receiver topology of `spmc` and `mpmc` is informed by [`flume`](https://github.com/zesterer/flume), with an independent runtime-agnostic implementation built on Asyncband's waiter arena. - `once::OnceCell` is derived from [`tokio::sync::OnceCell`](https://docs.rs/tokio/latest/tokio/sync/struct.OnceCell.html), but uses Asyncband's semaphore implementation. - `once::OnceMap` is inspired by [`uv-once-map`](https://github.com/astral-sh/uv/tree/main/crates/uv-once-map), with a redesigned interface and implementation. - `oneshot::channel` is derived from the [`oneshot`](https://github.com/faern/oneshot) crate, with significant simplifications because Asyncband does not provide synchronized receive operations. @@ -16,3 +17,4 @@ Asyncband collects composable, runtime-agnostic concurrency building blocks info - `rwlock::RwLock` is derived from [`tokio::sync::RwLock`](https://docs.rs/tokio/latest/tokio/sync/struct.RwLock.html), but accepts any `NonZeroUsize` as `max_readers` instead of Tokio's restricted range. - `semaphore::Semaphore` is derived from [`tokio::sync::Semaphore`](https://docs.rs/tokio/latest/tokio/sync/struct.Semaphore.html), but omits `close`, avoids Tokio's fixed maximum-permit constant, and adds operations such as `reduce_permits` for Asyncband's use cases. - `waitgroup::WaitGroup` is inspired by [`waitgroup-rs`](https://github.com/laizy/waitgroup-rs), with a different API and an implementation based on the internal `CountdownState` primitive. +- `watch` is inspired by [`tokio::sync::watch`](https://docs.rs/tokio/latest/tokio/sync/watch/), but returns owned `Arc` snapshots instead of runtime-specific borrow guards. diff --git a/README.md b/README.md index 44ecbb6..6db909a 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,7 @@ async fn increment() { } ``` -Public paths stay direct—such as `asyncband::mutex`, `asyncband::pool`, and `asyncband::once::OnceCell`—while Cargo features keep unused implementations out of the build. +Public paths stay direct—such as `asyncband::mutex`, `asyncband::mpsc`, and `asyncband::once::OnceCell`—while Cargo features keep unused implementations out of the build. The `channel` feature enables every channel API without adding an `asyncband::channel` namespace. ## API map @@ -71,8 +71,14 @@ Public paths stay direct—such as `asyncband::mutex`, `asyncband::pool`, and `a | | [`WaitGroup`](https://docs.rs/asyncband/*/asyncband/waitgroup/struct.WaitGroup.html) | `waitgroup` | Wait for a dynamic group of tasks to finish. | | | [`shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/) | `shutdown` | Coordinate shutdown signals and completion. | | Channels | [`oneshot::channel`](https://docs.rs/asyncband/*/asyncband/oneshot/fn.channel.html) | `oneshot` | Send one value between two tasks. | +| | [`spsc`](https://docs.rs/asyncband/*/asyncband/spsc/) | `spsc` | Queue each value for one producer and one receiver. | | | [`mpsc::bounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.bounded.html) | `mpsc` | Send values from multiple producers through a bounded channel. | | | [`mpsc::unbounded`](https://docs.rs/asyncband/*/asyncband/mpsc/fn.unbounded.html) | `mpsc` | Send values from multiple producers through an unbounded channel. | +| | [`spmc`](https://docs.rs/asyncband/*/asyncband/spmc/) | `spmc` | Let multiple receivers compete for values from one producer. | +| | [`mpmc`](https://docs.rs/asyncband/*/asyncband/mpmc/) | `mpmc` | Let multiple producers and receivers share a competing queue. | +| | [`broadcast::spmc`](https://docs.rs/asyncband/*/asyncband/broadcast/spmc/) | `broadcast` | Broadcast every value from one producer to every subscription. | +| | [`broadcast::mpmc`](https://docs.rs/asyncband/*/asyncband/broadcast/mpmc/) | `broadcast` | Broadcast one committed order from concurrent producers. | +| | [`watch`](https://docs.rs/asyncband/*/asyncband/watch/) | `watch` | Retain the latest state and coalesce intermediate updates. | | Resource reuse | [`pool::bounded`](https://docs.rs/asyncband/*/asyncband/pool/bounded/) | `pool` | Reuse managed objects up to a configured capacity. | | | [`pool::unbounded`](https://docs.rs/asyncband/*/asyncband/pool/unbounded/) | `pool` | Reuse manually supplied or manager-created objects. | | Workload coordination | [`Semaphore`](https://docs.rs/asyncband/*/asyncband/semaphore/struct.Semaphore.html) | `semaphore` | Control concurrent access with permits. | diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index af8e642..af34f69 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -46,8 +46,11 @@ default = [] barrier = [] blocking = [] +broadcast = [] +channel = ["broadcast", "mpmc", "mpsc", "oneshot", "spmc", "spsc", "watch"] condvar = ["mutex"] latch = [] +mpmc = [] mpsc = [] mutex = [] once = ["semaphore"] @@ -59,7 +62,10 @@ rwlock = [] semaphore = [] shutdown = ["latch", "waitgroup"] singleflight = ["dep:hashbrown", "once-cell"] +spmc = [] +spsc = [] waitgroup = [] +watch = [] [dependencies] hashbrown = { workspace = true, default-features = false, features = [ diff --git a/asyncband/src/channel/broadcast/internal.rs b/asyncband/src/channel/broadcast/internal.rs new file mode 100644 index 0000000..0e07f0b --- /dev/null +++ b/asyncband/src/channel/broadcast/internal.rs @@ -0,0 +1,644 @@ +// 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::collections::VecDeque; +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use crate::channel::error::RecvError; +use crate::channel::error::SendError; +use crate::channel::error::TryRecvError; +use crate::channel::error::TrySendError; +use crate::internal::arena::Arena; +use crate::internal::arena::ArenaKey; +use crate::internal::arena::ArenaValues; +use crate::internal::mutex::Mutex; +use crate::internal::waitset::WaitRegistration; +use crate::internal::waitset::WaitSet; + +type Received = (Arc, Reclaimed); +type TryReceived = ( + Result, TryRecvError>, + Option>, +); + +struct Reclaimed { + first: Option>, + // Boxing the uncommon multi-value tail keeps the zero/one-value hot path to two words. + #[allow(clippy::box_collection)] + _rest: Option>>>, +} + +impl Reclaimed { + fn none() -> Self { + Self { + first: None, + _rest: None, + } + } + + fn is_empty(&self) -> bool { + self.first.is_none() + } + + fn first(&self) -> Option<&Arc> { + self.first.as_ref() + } +} + +#[derive(Clone, Copy)] +enum Retention { + Bounded(usize), + Unbounded, +} + +pub(super) fn bounded(capacity: usize) -> (Sender, Receiver) { + debug_assert!(capacity > 0); + channel(Retention::Bounded(capacity)) +} + +pub(super) fn unbounded() -> (Sender, Receiver) { + channel(Retention::Unbounded) +} + +fn channel(retention: Retention) -> (Sender, Receiver) { + let mut receivers = Receivers::new(); + let key = receivers.insert(0); + let buffer = match retention { + Retention::Bounded(capacity) => VecDeque::with_capacity(capacity), + Retention::Unbounded => VecDeque::new(), + }; + let shared = Arc::new(Shared { + retention, + inner: Mutex::new(Inner { + buffer, + head: 0, + head_receivers: 1, + tail: 0, + receivers, + send_waiters: WaitSet::new(), + recv_waiters: WaitSet::new(), + }), + senders: AtomicUsize::new(1), + }); + ( + Sender { + shared: shared.clone(), + }, + Receiver { + shared, + key, + cursor: 0, + }, + ) +} + +struct Shared { + retention: Retention, + inner: Mutex>, + senders: AtomicUsize, +} + +struct Inner { + /// Messages with sequence numbers in `[head, tail)`. + /// + /// `Arc` lets cursor bookkeeping finish under the lock while user-defined `Clone` and `Drop` + /// implementations run after unlocking. + buffer: VecDeque>, + head: u64, + /// Number of live subscriptions currently equal to `head`. + /// + /// Only the last one to advance scans all live cursors and reclaims the common prefix. + head_receivers: usize, + tail: u64, + receivers: Receivers, + send_waiters: WaitSet, + recv_waiters: WaitSet, +} + +/// Stable receiver keys backed by a dense cursor list. +/// +/// The dense list keeps reclaim scans proportional to the number of live subscriptions instead of +/// the arena's historical high-water mark. Arena slots map stable endpoint keys into that list. +struct Receivers { + slots: Arena, + active: Vec, +} + +struct ReceiverCursor { + key: ArenaKey, + sequence: u64, +} + +impl Receivers { + fn new() -> Self { + Self { + slots: Arena::new(), + active: Vec::new(), + } + } + + fn insert(&mut self, sequence: u64) -> ArenaKey { + let active_index = self.active.len(); + let key = self.slots.insert(active_index); + self.active.push(ReceiverCursor { key, sequence }); + key + } + + fn set_sequence(&mut self, key: ArenaKey, sequence: u64) { + let active_index = *self + .slots + .get(key) + .expect("active broadcast receiver must be registered"); + let receiver = &mut self.active[active_index]; + debug_assert_eq!(receiver.sequence + 1, sequence); + receiver.sequence = sequence; + } + + fn remove(&mut self, key: ArenaKey) -> u64 { + let active_index = self.slots.remove(key); + let removed = self.active.swap_remove(active_index); + debug_assert_eq!(removed.key, key); + if let Some(moved) = self.active.get(active_index) { + *self + .slots + .get_mut(moved.key) + .expect("active broadcast receiver must be registered") = active_index; + } + removed.sequence + } + + fn len(&self) -> usize { + self.active.len() + } + + fn is_empty(&self) -> bool { + self.active.is_empty() + } + + fn sequences(&self) -> impl Iterator + '_ { + self.active.iter().map(|receiver| receiver.sequence) + } +} + +impl Inner { + fn insert_receiver(&mut self, cursor: u64) -> ArenaKey { + if cursor == self.head { + self.head_receivers += 1; + } + self.receivers.insert(cursor) + } + + fn remove_receiver(&mut self, key: ArenaKey) -> Reclaimed { + let cursor = self.receivers.remove(key); + if cursor == self.head { + self.release_head_receiver() + } else { + Reclaimed::none() + } + } + + fn receive(&mut self, key: ArenaKey, cursor: u64) -> Option<(Arc, Reclaimed)> { + if cursor == self.tail { + return None; + } + + debug_assert!(cursor >= self.head); + let offset = usize::try_from(cursor - self.head) + .expect("retained broadcast message count exceeds usize"); + let value = self.buffer[offset].clone(); + self.receivers.set_sequence(key, cursor + 1); + let reclaimed = if cursor == self.head { + self.release_head_receiver() + } else { + Reclaimed::none() + }; + debug_assert!( + reclaimed + .first() + .is_none_or(|first| Arc::ptr_eq(first, &value)) + ); + Some((value, reclaimed)) + } + + fn release_head_receiver(&mut self) -> Reclaimed { + self.head_receivers -= 1; + if self.head_receivers == 0 { + self.reclaim_consumed() + } else { + Reclaimed::none() + } + } + + fn reclaim_consumed(&mut self) -> Reclaimed { + let mut next_head = self.tail; + let mut head_receivers = 0; + for cursor in self.receivers.sequences() { + if cursor < next_head { + next_head = cursor; + head_receivers = 1; + } else if cursor == next_head { + head_receivers += 1; + } + } + + let consumed = usize::try_from(next_head - self.head) + .expect("retained broadcast message count exceeds usize"); + let first = (consumed > 0).then(|| { + self.buffer + .pop_front() + .expect("a reclaimed broadcast range must be buffered") + }); + let rest = if consumed > 1 { + Some(Box::new(self.buffer.drain(..consumed - 1).collect())) + } else { + None + }; + self.head = next_head; + self.head_receivers = head_receivers; + Reclaimed { first, _rest: rest } + } +} + +pub(super) struct Sender { + shared: Arc>, +} + +impl Clone for Sender { + fn clone(&self) -> Self { + let mut senders = self.shared.senders.load(Ordering::Relaxed); + loop { + let next = senders + .checked_add(1) + .expect("broadcast sender count overflow"); + match self.shared.senders.compare_exchange_weak( + senders, + next, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(actual) => senders = actual, + } + } + Self { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for Sender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Sender").finish_non_exhaustive() + } +} + +impl Drop for Sender { + fn drop(&mut self) { + if self.shared.senders.fetch_sub(1, Ordering::AcqRel) == 1 { + let wakers = { + let mut inner = self.shared.inner.lock(); + inner.recv_waiters.take_wakers() + }; + wake_all(wakers); + } + } +} + +impl Sender { + pub(super) async fn send(&self, value: T) -> Result<(), SendError> { + Send { + sender: self, + value: Some(value), + registration: None, + } + .await + } + + pub(super) fn try_send(&self, value: T) -> Result<(), TrySendError> { + let wakers = { + let mut inner = self.shared.inner.lock(); + if inner.receivers.is_empty() { + return Err(TrySendError::Disconnected(value)); + } + let Retention::Bounded(capacity) = self.shared.retention else { + unreachable!("try_send is only used by bounded broadcast endpoints") + }; + if inner.buffer.len() == capacity { + return Err(TrySendError::Full(value)); + } + append(&mut inner, value); + (!inner.recv_waiters.is_empty()).then(|| inner.recv_waiters.take_wakers()) + }; + if let Some(wakers) = wakers { + wake_all(wakers); + } + Ok(()) + } + + pub(super) fn send_unbounded(&self, value: T) -> Result<(), SendError> { + let wakers = { + let mut inner = self.shared.inner.lock(); + if inner.receivers.is_empty() { + return Err(SendError::new(value)); + } + debug_assert!(matches!(self.shared.retention, Retention::Unbounded)); + append(&mut inner, value); + (!inner.recv_waiters.is_empty()).then(|| inner.recv_waiters.take_wakers()) + }; + if let Some(wakers) = wakers { + wake_all(wakers); + } + Ok(()) + } + + pub(super) fn subscribe(&self) -> Receiver { + let (key, cursor) = { + let mut inner = self.shared.inner.lock(); + let cursor = inner.tail; + (inner.insert_receiver(cursor), cursor) + }; + Receiver { + shared: self.shared.clone(), + key, + cursor, + } + } + + pub(super) fn receiver_count(&self) -> usize { + self.shared.inner.lock().receivers.len() + } + + pub(super) fn buffer_len(&self) -> usize { + self.shared.inner.lock().buffer.len() + } + + fn cancel_send(&self, registration: &mut Option) { + let waker = { + let mut inner = self.shared.inner.lock(); + inner.send_waiters.unregister_waker(registration) + }; + drop(waker); + } +} + +fn append(inner: &mut Inner, value: T) { + inner.tail = inner + .tail + .checked_add(1) + .expect("broadcast sequence overflow"); + inner.buffer.push_back(Arc::new(value)); +} + +pub(super) struct Receiver { + shared: Arc>, + key: ArenaKey, + // Keep the owning endpoint's hot cursor local; the registry copy exists only for gating. + cursor: u64, +} + +impl fmt::Debug for Receiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Receiver").finish_non_exhaustive() + } +} + +impl Drop for Receiver { + fn drop(&mut self) { + let (reclaimed, wakers) = { + let mut inner = self.shared.inner.lock(); + let reclaimed = inner.remove_receiver(self.key); + let wakers = (!reclaimed.is_empty() && !inner.send_waiters.is_empty()) + .then(|| inner.send_waiters.take_wakers()); + (reclaimed, wakers) + }; + drop(reclaimed); + if let Some(wakers) = wakers { + wake_all(wakers); + } + } +} + +impl Receiver { + pub(super) async fn recv(&mut self) -> Result { + Recv { + receiver: self, + registration: None, + } + .await + } + + pub(super) fn try_recv(&mut self) -> Result { + let (result, wakers) = self.try_recv_shared(); + if let Some(wakers) = wakers { + wake_all(wakers); + } + let (value, reclaimed) = result?; + Ok(take_value(value, reclaimed)) + } +} + +impl Receiver { + fn try_recv_shared(&mut self) -> TryReceived { + let mut inner = self.shared.inner.lock(); + if let Some((value, reclaimed)) = inner.receive(self.key, self.cursor) { + self.cursor += 1; + let wakers = (!reclaimed.is_empty() && !inner.send_waiters.is_empty()) + .then(|| inner.send_waiters.take_wakers()); + return (Ok((value, reclaimed)), wakers); + } + if self.shared.senders.load(Ordering::Acquire) == 0 { + (Err(TryRecvError::Disconnected), None) + } else { + (Err(TryRecvError::Empty), None) + } + } + + pub(super) fn len(&self) -> usize { + let inner = self.shared.inner.lock(); + usize::try_from(inner.tail - self.cursor) + .expect("unread broadcast message count exceeds usize") + } + + pub(super) fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub(super) fn is_disconnected(&self) -> bool { + self.shared.senders.load(Ordering::Acquire) == 0 + } + + fn cancel_recv(&self, registration: &mut Option) { + let waker = { + let mut inner = self.shared.inner.lock(); + inner.recv_waiters.unregister_waker(registration) + }; + drop(waker); + } +} + +fn take_value(value: Arc, reclaimed: Reclaimed) -> T { + // Reclaiming this sequence means `value` becomes uniquely owned after the drained buffer + // references are dropped. The common single-subscription path can therefore move T out. + let reclaimed_value = !reclaimed.is_empty(); + drop(reclaimed); + if reclaimed_value { + match Arc::try_unwrap(value) { + Ok(value) => value, + Err(value) => (*value).clone(), + } + } else { + (*value).clone() + } +} + +struct Send<'a, T> { + sender: &'a Sender, + value: Option, + registration: Option, +} + +impl Unpin for Send<'_, T> {} + +impl Future for Send<'_, T> { + type Output = Result<(), SendError>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let (poll, retired_waker, wake_receivers) = { + let mut inner = this.sender.shared.inner.lock(); + if inner.receivers.is_empty() { + let retired_waker = inner.send_waiters.unregister_waker(&mut this.registration); + ( + Poll::Ready(Err(SendError::new( + this.value + .take() + .expect("an incomplete send owns its value"), + ))), + retired_waker, + None, + ) + } else { + let Retention::Bounded(capacity) = this.sender.shared.retention else { + unreachable!("async send is only used by bounded broadcast endpoints") + }; + if inner.buffer.len() == capacity { + let retired_waker = inner + .send_waiters + .register_waker(&mut this.registration, cx); + (Poll::Pending, retired_waker, None) + } else { + let retired_waker = inner.send_waiters.unregister_waker(&mut this.registration); + append( + &mut inner, + this.value + .take() + .expect("an incomplete send owns its value"), + ); + let wake_receivers = + (!inner.recv_waiters.is_empty()).then(|| inner.recv_waiters.take_wakers()); + (Poll::Ready(Ok(())), retired_waker, wake_receivers) + } + } + }; + drop(retired_waker); + if let Some(wakers) = wake_receivers { + wake_all(wakers); + } + poll + } +} + +impl Drop for Send<'_, T> { + fn drop(&mut self) { + if self.registration.is_some() { + self.sender.cancel_send(&mut self.registration); + } + } +} + +struct Recv<'a, T> { + receiver: &'a mut Receiver, + registration: Option, +} + +impl Future for Recv<'_, T> { + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let (poll, retired_waker, wake_senders) = { + let mut inner = this.receiver.shared.inner.lock(); + if let Some((value, reclaimed)) = inner.receive(this.receiver.key, this.receiver.cursor) + { + let retired_waker = inner.recv_waiters.unregister_waker(&mut this.registration); + this.receiver.cursor += 1; + let wake_senders = (!reclaimed.is_empty() && !inner.send_waiters.is_empty()) + .then(|| inner.send_waiters.take_wakers()); + ( + Poll::Ready(Ok((value, reclaimed))), + retired_waker, + wake_senders, + ) + } else if this.receiver.shared.senders.load(Ordering::Acquire) == 0 { + let retired_waker = inner.recv_waiters.unregister_waker(&mut this.registration); + ( + Poll::Ready(Err(RecvError::Disconnected)), + retired_waker, + None, + ) + } else { + let retired_waker = inner + .recv_waiters + .register_waker(&mut this.registration, cx); + (Poll::Pending, retired_waker, None) + } + }; + drop(retired_waker); + if let Some(wakers) = wake_senders { + wake_all(wakers); + } + match poll { + Poll::Ready(Ok((value, reclaimed))) => Poll::Ready(Ok(take_value(value, reclaimed))), + Poll::Ready(Err(error)) => Poll::Ready(Err(error)), + Poll::Pending => Poll::Pending, + } + } +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + if self.registration.is_some() { + self.receiver.cancel_recv(&mut self.registration); + } + } +} + +fn wake_all(wakers: ArenaValues) { + // Every parked subscription observes a publication. Backpressured senders also wake as a set + // so cancellation of one selected future cannot strand newly available capacity. + for waker in wakers { + waker.wake(); + } +} diff --git a/asyncband/src/channel/broadcast/mod.rs b/asyncband/src/channel/broadcast/mod.rs new file mode 100644 index 0000000..10b7b2a --- /dev/null +++ b/asyncband/src/channel/broadcast/mod.rs @@ -0,0 +1,36 @@ +// 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 broadcast channels. +//! +//! Every active subscription observes every accepted value. Bounded retention applies +//! backpressure at the slowest subscription, while unbounded retention grows until subscriptions +//! advance or are dropped. +//! +//! Unlike a competing `asyncband::spmc` or `asyncband::mpmc` queue, each subscription has +//! independent receive progress. New subscriptions start at the committed tail and observe future +//! publications only. Both retention modes are lossless; there is no lag or overwrite result. +//! +//! A bounded channel enforces the requested capacity as a strict logical limit and backpressures +//! senders until the slowest subscription advances or drops. An unbounded channel grows subject to +//! process memory and reclaims a prefix only after every active subscription advances past it. +//! Pending sends and receives are cancel safe and do not publish or advance a subscription. + +mod internal; + +pub mod mpmc; +pub mod spmc; diff --git a/asyncband/src/channel/broadcast/mpmc.rs b/asyncband/src/channel/broadcast/mpmc.rs new file mode 100644 index 0000000..177a566 --- /dev/null +++ b/asyncband/src/channel/broadcast/mpmc.rs @@ -0,0 +1,247 @@ +// 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 lossless broadcast channels. +//! +//! ``` +//! use asyncband::broadcast::mpmc; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (tx, mut first) = mpmc::bounded(2); +//! let mut second = tx.subscribe(); +//! tx.send("event").await.unwrap(); +//! +//! assert_eq!(first.recv().await, Ok("event")); +//! assert_eq!(second.recv().await, Ok("event")); +//! # } +//! ``` +//! +//! Subscriptions are intentionally non-cloneable. Call [`BoundedSender::subscribe`] or +//! [`UnboundedSender::subscribe`] to create one that visibly starts at the current tail. +//! +//! ```compile_fail +//! let (_, receiver) = asyncband::broadcast::mpmc::unbounded::(); +//! let _ = receiver.clone(); +//! ``` + +use std::fmt; + +pub use crate::channel::error::RecvError; +pub use crate::channel::error::SendError; +pub use crate::channel::error::TryRecvError; +pub use crate::channel::error::TrySendError; + +/// Creates a bounded MPMC broadcast channel. +/// +/// The slowest active subscription gates senders once `capacity` values are retained. +/// +/// # Panics +/// +/// Panics if `capacity` is zero. +#[track_caller] +pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { + assert!( + capacity > 0, + "bounded broadcast channel requires capacity > 0" + ); + let (sender, receiver) = super::internal::bounded(capacity); + ( + BoundedSender { inner: sender }, + BoundedReceiver { inner: receiver }, + ) +} + +/// Creates an unbounded MPMC broadcast channel. +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { + let (sender, receiver) = super::internal::unbounded(); + ( + UnboundedSender { inner: sender }, + UnboundedReceiver { inner: receiver }, + ) +} + +/// A sending endpoint of a bounded MPMC broadcast channel. +pub struct BoundedSender { + inner: super::internal::Sender, +} + +impl Clone for BoundedSender { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl fmt::Debug for BoundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedSender").finish_non_exhaustive() + } +} + +impl BoundedSender { + /// Broadcasts a value, waiting while the retention buffer is full. + pub async fn send(&self, value: T) -> Result<(), SendError> { + self.inner.send(value).await + } + + /// Attempts to broadcast a value without waiting. + pub fn try_send(&self, value: T) -> Result<(), TrySendError> { + self.inner.try_send(value) + } + + /// Creates a subscription starting at the current committed tail. + pub fn subscribe(&self) -> BoundedReceiver { + BoundedReceiver { + inner: self.inner.subscribe(), + } + } + + /// Returns the number of active subscriptions. + pub fn receiver_count(&self) -> usize { + self.inner.receiver_count() + } + + /// Returns the number of values retained for the slowest subscription. + pub fn buffer_len(&self) -> usize { + self.inner.buffer_len() + } +} + +/// A subscription to a bounded MPMC broadcast channel. +pub struct BoundedReceiver { + inner: super::internal::Receiver, +} + +impl fmt::Debug for BoundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedReceiver").finish_non_exhaustive() + } +} + +impl BoundedReceiver { + /// Receives the next value for this subscription. + pub async fn recv(&mut self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&mut self) -> Result { + self.inner.try_recv() + } +} + +impl BoundedReceiver { + /// Returns the number of values currently available to this subscription. + pub fn len(&self) -> usize { + self.inner.len() + } + + /// Returns whether no value is currently available to this subscription. + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Returns whether every sender has been dropped. + pub fn is_disconnected(&self) -> bool { + self.inner.is_disconnected() + } +} + +/// A sending endpoint of an unbounded MPMC broadcast channel. +pub struct UnboundedSender { + inner: super::internal::Sender, +} + +impl Clone for UnboundedSender { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl fmt::Debug for UnboundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedSender").finish_non_exhaustive() + } +} + +impl UnboundedSender { + /// Broadcasts a value without waiting. + pub fn send(&self, value: T) -> Result<(), SendError> { + self.inner.send_unbounded(value) + } + + /// Creates a subscription starting at the current committed tail. + pub fn subscribe(&self) -> UnboundedReceiver { + UnboundedReceiver { + inner: self.inner.subscribe(), + } + } + + /// Returns the number of active subscriptions. + pub fn receiver_count(&self) -> usize { + self.inner.receiver_count() + } + + /// Returns the number of values retained for the slowest subscription. + pub fn buffer_len(&self) -> usize { + self.inner.buffer_len() + } +} + +/// A subscription to an unbounded MPMC broadcast channel. +pub struct UnboundedReceiver { + inner: super::internal::Receiver, +} + +impl fmt::Debug for UnboundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() + } +} + +impl UnboundedReceiver { + /// Receives the next value for this subscription. + pub async fn recv(&mut self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&mut self) -> Result { + self.inner.try_recv() + } +} + +impl UnboundedReceiver { + /// Returns the number of values currently available to this subscription. + pub fn len(&self) -> usize { + self.inner.len() + } + + /// Returns whether no value is currently available to this subscription. + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Returns whether every sender has been dropped. + pub fn is_disconnected(&self) -> bool { + self.inner.is_disconnected() + } +} diff --git a/asyncband/src/channel/broadcast/spmc.rs b/asyncband/src/channel/broadcast/spmc.rs new file mode 100644 index 0000000..1ce11d7 --- /dev/null +++ b/asyncband/src/channel/broadcast/spmc.rs @@ -0,0 +1,245 @@ +// 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 lossless broadcast channels. +//! +//! ``` +//! use asyncband::broadcast::spmc; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (mut tx, mut first) = spmc::bounded(2); +//! let mut second = tx.subscribe(); +//! tx.send("event").await.unwrap(); +//! +//! assert_eq!(first.recv().await, Ok("event")); +//! assert_eq!(second.recv().await, Ok("event")); +//! # } +//! ``` +//! +//! The single-producer contract is enforced by the sender type: +//! +//! ```compile_fail +//! let (tx, _) = asyncband::broadcast::spmc::unbounded::(); +//! let _ = tx.clone(); +//! ``` +//! +//! ```compile_fail +//! fn assert_sync() {} +//! assert_sync::>(); +//! ``` + +use std::cell::Cell; +use std::fmt; +use std::marker::PhantomData; + +pub use crate::channel::error::RecvError; +pub use crate::channel::error::SendError; +pub use crate::channel::error::TryRecvError; +pub use crate::channel::error::TrySendError; + +/// Creates a bounded SPMC broadcast channel. +/// +/// The slowest active subscription gates the sender once `capacity` values are retained. +/// +/// # Panics +/// +/// Panics if `capacity` is zero. +#[track_caller] +pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { + assert!( + capacity > 0, + "bounded broadcast channel requires capacity > 0" + ); + let (sender, receiver) = super::internal::bounded(capacity); + ( + BoundedSender { + inner: sender, + not_sync: PhantomData, + }, + BoundedReceiver { inner: receiver }, + ) +} + +/// Creates an unbounded SPMC broadcast channel. +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { + let (sender, receiver) = super::internal::unbounded(); + ( + UnboundedSender { + inner: sender, + not_sync: PhantomData, + }, + UnboundedReceiver { inner: receiver }, + ) +} + +/// The sending endpoint of a bounded SPMC broadcast channel. +pub struct BoundedSender { + inner: super::internal::Sender, + not_sync: PhantomData>, +} + +impl fmt::Debug for BoundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedSender").finish_non_exhaustive() + } +} + +impl BoundedSender { + /// Broadcasts a value, waiting while the retention buffer is full. + pub async fn send(&mut self, value: T) -> Result<(), SendError> { + self.inner.send(value).await + } + + /// Attempts to broadcast a value without waiting. + pub fn try_send(&mut self, value: T) -> Result<(), TrySendError> { + self.inner.try_send(value) + } + + /// Creates a subscription starting at the current committed tail. + pub fn subscribe(&self) -> BoundedReceiver { + BoundedReceiver { + inner: self.inner.subscribe(), + } + } + + /// Returns the number of active subscriptions. + pub fn receiver_count(&self) -> usize { + self.inner.receiver_count() + } + + /// Returns the number of values retained for the slowest subscription. + pub fn buffer_len(&self) -> usize { + self.inner.buffer_len() + } +} + +/// A subscription to a bounded SPMC broadcast channel. +pub struct BoundedReceiver { + inner: super::internal::Receiver, +} + +impl fmt::Debug for BoundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedReceiver").finish_non_exhaustive() + } +} + +impl BoundedReceiver { + /// Receives the next value for this subscription. + pub async fn recv(&mut self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&mut self) -> Result { + self.inner.try_recv() + } +} + +impl BoundedReceiver { + /// Returns the number of values currently available to this subscription. + pub fn len(&self) -> usize { + self.inner.len() + } + + /// Returns whether no value is currently available to this subscription. + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Returns whether every sender has been dropped. + pub fn is_disconnected(&self) -> bool { + self.inner.is_disconnected() + } +} + +/// The sending endpoint of an unbounded SPMC broadcast channel. +pub struct UnboundedSender { + inner: super::internal::Sender, + not_sync: PhantomData>, +} + +impl fmt::Debug for UnboundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedSender").finish_non_exhaustive() + } +} + +impl UnboundedSender { + /// Broadcasts a value without waiting. + pub fn send(&mut self, value: T) -> Result<(), SendError> { + self.inner.send_unbounded(value) + } + + /// Creates a subscription starting at the current committed tail. + pub fn subscribe(&self) -> UnboundedReceiver { + UnboundedReceiver { + inner: self.inner.subscribe(), + } + } + + /// Returns the number of active subscriptions. + pub fn receiver_count(&self) -> usize { + self.inner.receiver_count() + } + + /// Returns the number of values retained for the slowest subscription. + pub fn buffer_len(&self) -> usize { + self.inner.buffer_len() + } +} + +/// A subscription to an unbounded SPMC broadcast channel. +pub struct UnboundedReceiver { + inner: super::internal::Receiver, +} + +impl fmt::Debug for UnboundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() + } +} + +impl UnboundedReceiver { + /// Receives the next value for this subscription. + pub async fn recv(&mut self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&mut self) -> Result { + self.inner.try_recv() + } +} + +impl UnboundedReceiver { + /// Returns the number of values currently available to this subscription. + pub fn len(&self) -> usize { + self.inner.len() + } + + /// Returns whether no value is currently available to this subscription. + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Returns whether every sender has been dropped. + pub fn is_disconnected(&self) -> bool { + self.inner.is_disconnected() + } +} diff --git a/asyncband/src/mpsc/error.rs b/asyncband/src/channel/error.rs similarity index 80% rename from asyncband/src/mpsc/error.rs rename to asyncband/src/channel/error.rs index cf5d649..6a465bf 100644 --- a/asyncband/src/mpsc/error.rs +++ b/asyncband/src/channel/error.rs @@ -20,17 +20,8 @@ use std::fmt; /// An error returned when trying to send on a closed channel. /// -/// Returned from [`UnboundedSender::send`] or [`BoundedSender::send`] if the -/// corresponding [`UnboundedReceiver`] or [`BoundedReceiver`] has already been -/// dropped. -/// /// The message that could not be sent can be retrieved again with /// [`SendError::into_inner`]. -/// -/// [`UnboundedSender::send`]: crate::mpsc::UnboundedSender::send -/// [`BoundedSender::send`]: crate::mpsc::BoundedSender::send -/// [`UnboundedReceiver`]: crate::mpsc::UnboundedReceiver -/// [`BoundedReceiver`]: crate::mpsc::BoundedReceiver #[derive(Clone, PartialEq, Eq)] pub struct SendError(T); @@ -46,7 +37,7 @@ impl SendError { } /// Creates a new `SendError` with the given message. - pub(super) fn new(msg: T) -> SendError { + pub(crate) fn new(msg: T) -> SendError { SendError(msg) } } @@ -66,6 +57,13 @@ impl fmt::Debug for SendError { impl std::error::Error for SendError {} /// Error returned by `try_send`. +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", +))] #[derive(Clone, PartialEq, Eq)] pub enum TrySendError { /// The channel is full, so data may not be sent at this time, but the receiver has not yet @@ -75,6 +73,13 @@ pub enum TrySendError { Disconnected(T), } +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", +))] impl TrySendError { /// Gets a reference to the message that failed to be sent. pub fn as_inner(&self) -> &T { @@ -91,6 +96,13 @@ impl TrySendError { } } +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", +))] impl fmt::Display for TrySendError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(match self { @@ -100,6 +112,13 @@ impl fmt::Display for TrySendError { } } +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", +))] impl fmt::Debug for TrySendError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let ty = type_name::(); @@ -110,6 +129,13 @@ impl fmt::Debug for TrySendError { } } +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", +))] impl std::error::Error for TrySendError {} /// Error returned by `recv`. @@ -128,6 +154,13 @@ impl fmt::Display for RecvError { impl std::error::Error for RecvError {} /// Error returned by `try_recv`. +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", +))] #[derive(Debug, Clone, PartialEq, Eq)] pub enum TryRecvError { /// This channel is currently empty, but the sender(s) have not yet disconnected, so data may @@ -137,6 +170,13 @@ pub enum TryRecvError { Disconnected, } +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", +))] impl fmt::Display for TryRecvError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(match self { @@ -146,4 +186,11 @@ impl fmt::Display for TryRecvError { } } +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", +))] impl std::error::Error for TryRecvError {} diff --git a/asyncband/src/channel/mod.rs b/asyncband/src/channel/mod.rs new file mode 100644 index 0000000..4bec501 --- /dev/null +++ b/asyncband/src/channel/mod.rs @@ -0,0 +1,43 @@ +// 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. + +#[cfg(any( + feature = "broadcast", + feature = "mpmc", + feature = "mpsc", + feature = "spmc", + feature = "spsc", + feature = "watch", +))] +mod error; + +#[cfg(feature = "broadcast")] +pub mod broadcast; +#[cfg(feature = "mpmc")] +pub mod mpmc; +#[cfg(any(feature = "mpsc", feature = "spsc"))] +pub mod mpsc; +#[cfg(feature = "oneshot")] +pub mod oneshot; +#[cfg(any(feature = "mpmc", feature = "spmc"))] +mod queue; +#[cfg(feature = "spmc")] +pub mod spmc; +#[cfg(feature = "spsc")] +pub mod spsc; +#[cfg(feature = "watch")] +pub mod watch; diff --git a/asyncband/src/channel/mpmc.rs b/asyncband/src/channel/mpmc.rs new file mode 100644 index 0000000..bf9beab --- /dev/null +++ b/asyncband/src/channel/mpmc.rs @@ -0,0 +1,192 @@ +// 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 competing queues. +//! +//! Both endpoints are cloneable and may be used concurrently. Every accepted value is delivered +//! to exactly one receiver. +//! +//! A bounded queue retains at most its requested capacity and waits rather than dropping values. +//! An unbounded queue sends synchronously and grows subject to process memory. Accepted values +//! drain before disconnection is reported, while a send after the last receiver disconnects +//! returns its value. +//! +//! Pending `send` and `recv` operations are cancel safe: dropping their futures does not transfer +//! a value. +//! +//! ``` +//! use asyncband::mpmc; +//! +//! let (tx, rx) = mpmc::unbounded(); +//! let other_tx = tx.clone(); +//! let other_rx = rx.clone(); +//! tx.send(1).unwrap(); +//! other_tx.send(2).unwrap(); +//! +//! assert_eq!(rx.try_recv(), Ok(1)); +//! assert_eq!(other_rx.try_recv(), Ok(2)); +//! ``` + +use std::fmt; + +pub use super::error::RecvError; +pub use super::error::SendError; +pub use super::error::TryRecvError; +pub use super::error::TrySendError; + +/// Creates a bounded MPMC queue with the given capacity. +/// +/// # Panics +/// +/// Panics if `capacity` is zero. +#[track_caller] +pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { + assert!(capacity > 0, "mpmc bounded channel requires capacity > 0"); + let (sender, receiver) = super::queue::bounded(capacity); + ( + BoundedSender { inner: sender }, + BoundedReceiver { inner: receiver }, + ) +} + +/// Creates an unbounded MPMC queue. +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { + let (sender, receiver) = super::queue::unbounded(); + ( + UnboundedSender { inner: sender }, + UnboundedReceiver { inner: receiver }, + ) +} + +/// A sending endpoint of a bounded MPMC queue. +pub struct BoundedSender { + inner: super::queue::Sender, +} + +impl Clone for BoundedSender { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl fmt::Debug for BoundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedSender").finish_non_exhaustive() + } +} + +impl BoundedSender { + /// Sends a value, waiting for capacity when the queue is full. + pub async fn send(&self, value: T) -> Result<(), SendError> { + self.inner.send(value).await + } + + /// Attempts to send a value without waiting. + pub fn try_send(&self, value: T) -> Result<(), TrySendError> { + self.inner.try_send(value) + } +} + +/// A receiving endpoint of a bounded MPMC queue. +pub struct BoundedReceiver { + inner: super::queue::Receiver, +} + +impl Clone for BoundedReceiver { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl fmt::Debug for BoundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedReceiver").finish_non_exhaustive() + } +} + +impl BoundedReceiver { + /// Receives the next value, waiting while the connected queue is empty. + pub async fn recv(&self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&self) -> Result { + self.inner.try_recv() + } +} + +/// A sending endpoint of an unbounded MPMC queue. +pub struct UnboundedSender { + inner: super::queue::Sender, +} + +impl Clone for UnboundedSender { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl fmt::Debug for UnboundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedSender").finish_non_exhaustive() + } +} + +impl UnboundedSender { + /// Sends a value without waiting. + pub fn send(&self, value: T) -> Result<(), SendError> { + self.inner.send_unbounded(value) + } +} + +/// A receiving endpoint of an unbounded MPMC queue. +pub struct UnboundedReceiver { + inner: super::queue::Receiver, +} + +impl Clone for UnboundedReceiver { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl fmt::Debug for UnboundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() + } +} + +impl UnboundedReceiver { + /// Receives the next value, waiting while the connected queue is empty. + pub async fn recv(&self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&self) -> Result { + self.inner.try_recv() + } +} diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/channel/mpsc/bounded.rs similarity index 98% rename from asyncband/src/mpsc/bounded.rs rename to asyncband/src/channel/mpsc/bounded.rs index 60f0da0..08aefd4 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/channel/mpsc/bounded.rs @@ -28,13 +28,13 @@ use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use super::RecvError; +use super::SendError; +use super::TryRecvError; +use super::TrySendError; use crate::internal::atomic_waker::AtomicWaker; use crate::internal::semaphore::Acquire; use crate::internal::semaphore::Semaphore; -use crate::mpsc::RecvError; -use crate::mpsc::SendError; -use crate::mpsc::TryRecvError; -use crate::mpsc::error::TrySendError; /// Creates a bounded mpsc channel for communicating between asynchronous /// tasks with backpressure. diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/channel/mpsc/mod.rs similarity index 88% rename from asyncband/src/mpsc/mod.rs rename to asyncband/src/channel/mpsc/mod.rs index 87c7c8f..2134bfc 100644 --- a/asyncband/src/mpsc/mod.rs +++ b/asyncband/src/channel/mpsc/mod.rs @@ -18,16 +18,16 @@ //! A multi-producer, single-consumer queue for sending values between asynchronous tasks. mod bounded; -mod error; mod unbounded; pub use bounded::BoundedReceiver; pub use bounded::BoundedSender; pub use bounded::bounded; -pub use error::RecvError; -pub use error::SendError; -pub use error::TryRecvError; -pub use error::TrySendError; pub use unbounded::UnboundedReceiver; pub use unbounded::UnboundedSender; pub use unbounded::unbounded; + +pub use super::error::RecvError; +pub use super::error::SendError; +pub use super::error::TryRecvError; +pub use super::error::TrySendError; diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/channel/mpsc/unbounded.rs similarity index 99% rename from asyncband/src/mpsc/unbounded.rs rename to asyncband/src/channel/mpsc/unbounded.rs index eb446be..3ee892a 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/channel/mpsc/unbounded.rs @@ -26,10 +26,10 @@ use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use super::RecvError; +use super::SendError; +use super::TryRecvError; use crate::internal::atomic_waker::AtomicWaker; -use crate::mpsc::RecvError; -use crate::mpsc::SendError; -use crate::mpsc::TryRecvError; /// Creates an unbounded mpsc channel for communicating between asynchronous /// tasks without backpressure. diff --git a/asyncband/src/oneshot/mod.rs b/asyncband/src/channel/oneshot/mod.rs similarity index 100% rename from asyncband/src/oneshot/mod.rs rename to asyncband/src/channel/oneshot/mod.rs diff --git a/asyncband/src/oneshot/receiver.rs b/asyncband/src/channel/oneshot/receiver.rs similarity index 98% rename from asyncband/src/oneshot/receiver.rs rename to asyncband/src/channel/oneshot/receiver.rs index 29f913b..7d232dc 100644 --- a/asyncband/src/oneshot/receiver.rs +++ b/asyncband/src/channel/oneshot/receiver.rs @@ -24,16 +24,16 @@ use std::sync::atomic::fence; use std::task::Context; use std::task::Poll; -use crate::oneshot::AWAKING; -use crate::oneshot::Channel; -use crate::oneshot::DISCONNECTED; -use crate::oneshot::EMPTY; -use crate::oneshot::MESSAGE; -use crate::oneshot::RECEIVING; +use super::AWAKING; +use super::Channel; +use super::DISCONNECTED; +use super::EMPTY; +use super::MESSAGE; +use super::RECEIVING; #[cfg(doc)] -use crate::oneshot::Sender; -use crate::oneshot::deallocate_empty_channel; -use crate::oneshot::drop_message_and_deallocate_channel; +use super::Sender; +use super::deallocate_empty_channel; +use super::drop_message_and_deallocate_channel; /// Receives a value from the associated [`Sender`]. pub struct Receiver { diff --git a/asyncband/src/oneshot/sender.rs b/asyncband/src/channel/oneshot/sender.rs similarity index 97% rename from asyncband/src/oneshot/sender.rs rename to asyncband/src/channel/oneshot/sender.rs index adf45ee..58a712c 100644 --- a/asyncband/src/oneshot/sender.rs +++ b/asyncband/src/channel/oneshot/sender.rs @@ -22,15 +22,15 @@ use std::ptr::NonNull; use std::sync::atomic::Ordering; use std::sync::atomic::fence; -use crate::oneshot::Channel; -use crate::oneshot::DISCONNECTED; -use crate::oneshot::EMPTY; -use crate::oneshot::MESSAGE; -use crate::oneshot::RECEIVING; +use super::Channel; +use super::DISCONNECTED; +use super::EMPTY; +use super::MESSAGE; +use super::RECEIVING; #[cfg(doc)] -use crate::oneshot::Receiver; -use crate::oneshot::deallocate_empty_channel; -use crate::oneshot::drop_message_and_deallocate_channel; +use super::Receiver; +use super::deallocate_empty_channel; +use super::drop_message_and_deallocate_channel; /// Sends a value to the associated [`Receiver`]. pub struct Sender { diff --git a/asyncband/src/oneshot/tests.rs b/asyncband/src/channel/oneshot/tests.rs similarity index 100% rename from asyncband/src/oneshot/tests.rs rename to asyncband/src/channel/oneshot/tests.rs diff --git a/asyncband/src/channel/queue.rs b/asyncband/src/channel/queue.rs new file mode 100644 index 0000000..cadbd8f --- /dev/null +++ b/asyncband/src/channel/queue.rs @@ -0,0 +1,384 @@ +// 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. + +//! Shared storage for competing multi-consumer queues. + +use std::collections::VecDeque; +use std::fmt; +use std::future::Future; +use std::mem; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use super::error::RecvError; +use super::error::SendError; +use super::error::TryRecvError; +use super::error::TrySendError; +use crate::internal::arena::ArenaValues; +use crate::internal::mutex::Mutex; +use crate::internal::waitset::WaitRegistration; +use crate::internal::waitset::WaitSet; + +#[derive(Clone, Copy)] +enum Capacity { + Bounded(usize), + Unbounded, +} + +pub(super) fn bounded(capacity: usize) -> (Sender, Receiver) { + debug_assert!(capacity > 0); + channel(Capacity::Bounded(capacity)) +} + +pub(super) fn unbounded() -> (Sender, Receiver) { + channel(Capacity::Unbounded) +} + +fn channel(capacity: Capacity) -> (Sender, Receiver) { + let queue = match capacity { + Capacity::Bounded(capacity) => VecDeque::with_capacity(capacity), + Capacity::Unbounded => VecDeque::new(), + }; + let shared = Arc::new(Shared { + capacity, + state: Mutex::new(State { + queue, + senders: 1, + receivers: 1, + send_waiters: WaitSet::new(), + recv_waiters: WaitSet::new(), + }), + }); + ( + Sender { + shared: shared.clone(), + }, + Receiver { shared }, + ) +} + +struct Shared { + capacity: Capacity, + state: Mutex>, +} + +struct State { + queue: VecDeque, + senders: usize, + receivers: usize, + send_waiters: WaitSet, + recv_waiters: WaitSet, +} + +pub(super) struct Sender { + shared: Arc>, +} + +impl Clone for Sender { + fn clone(&self) -> Self { + let mut state = self.shared.state.lock(); + state.senders = state + .senders + .checked_add(1) + .expect("queue sender count overflow"); + drop(state); + Self { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for Sender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Sender").finish_non_exhaustive() + } +} + +impl Drop for Sender { + fn drop(&mut self) { + let wakers = { + let mut state = self.shared.state.lock(); + state.senders -= 1; + (state.senders == 0).then(|| state.recv_waiters.take_wakers()) + }; + if let Some(wakers) = wakers { + wake_all(wakers); + } + } +} + +impl Sender { + pub(super) async fn send(&self, value: T) -> Result<(), SendError> { + Send { + sender: self, + value: Some(value), + registration: None, + } + .await + } + + pub(super) fn try_send(&self, value: T) -> Result<(), TrySendError> { + let wakers = { + let mut state = self.shared.state.lock(); + if state.receivers == 0 { + return Err(TrySendError::Disconnected(value)); + } + let Capacity::Bounded(capacity) = self.shared.capacity else { + unreachable!("try_send is only used by bounded queue endpoints") + }; + if state.queue.len() == capacity { + return Err(TrySendError::Full(value)); + } + state.queue.push_back(value); + (!state.recv_waiters.is_empty()).then(|| state.recv_waiters.take_wakers()) + }; + if let Some(wakers) = wakers { + wake_all(wakers); + } + Ok(()) + } + + pub(super) fn send_unbounded(&self, value: T) -> Result<(), SendError> { + let wakers = { + let mut state = self.shared.state.lock(); + if state.receivers == 0 { + return Err(SendError::new(value)); + } + debug_assert!(matches!(self.shared.capacity, Capacity::Unbounded)); + state.queue.push_back(value); + (!state.recv_waiters.is_empty()).then(|| state.recv_waiters.take_wakers()) + }; + if let Some(wakers) = wakers { + wake_all(wakers); + } + Ok(()) + } + + fn cancel_send(&self, registration: &mut Option) { + let waker = { + let mut state = self.shared.state.lock(); + state.send_waiters.unregister_waker(registration) + }; + drop(waker); + } +} + +pub(super) struct Receiver { + shared: Arc>, +} + +impl Clone for Receiver { + fn clone(&self) -> Self { + let mut state = self.shared.state.lock(); + state.receivers = state + .receivers + .checked_add(1) + .expect("queue receiver count overflow"); + drop(state); + Self { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for Receiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Receiver").finish_non_exhaustive() + } +} + +impl Drop for Receiver { + fn drop(&mut self) { + let (wakers, queued) = { + let mut state = self.shared.state.lock(); + state.receivers -= 1; + if state.receivers == 0 { + ( + Some(state.send_waiters.take_wakers()), + mem::take(&mut state.queue), + ) + } else { + (None, VecDeque::new()) + } + }; + drop(queued); + if let Some(wakers) = wakers { + wake_all(wakers); + } + } +} + +impl Receiver { + pub(super) async fn recv(&self) -> Result { + Recv { + receiver: self, + registration: None, + } + .await + } + + pub(super) fn try_recv(&self) -> Result { + let (result, wakers) = { + let mut state = self.shared.state.lock(); + if let Some(value) = state.queue.pop_front() { + let wakers = (matches!(self.shared.capacity, Capacity::Bounded(_)) + && !state.send_waiters.is_empty()) + .then(|| state.send_waiters.take_wakers()); + (Ok(value), wakers) + } else if state.senders == 0 { + (Err(TryRecvError::Disconnected), None) + } else { + (Err(TryRecvError::Empty), None) + } + }; + if let Some(wakers) = wakers { + wake_all(wakers); + } + result + } + + fn cancel_recv(&self, registration: &mut Option) { + let waker = { + let mut state = self.shared.state.lock(); + state.recv_waiters.unregister_waker(registration) + }; + drop(waker); + } +} + +struct Send<'a, T> { + sender: &'a Sender, + value: Option, + registration: Option, +} + +impl Unpin for Send<'_, T> {} + +impl Future for Send<'_, T> { + type Output = Result<(), SendError>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let (poll, retired_waker, wake_receivers) = { + let mut state = this.sender.shared.state.lock(); + if state.receivers == 0 { + let retired_waker = state.send_waiters.unregister_waker(&mut this.registration); + ( + Poll::Ready(Err(SendError::new( + this.value + .take() + .expect("an incomplete send owns its value"), + ))), + retired_waker, + None, + ) + } else { + let Capacity::Bounded(capacity) = this.sender.shared.capacity else { + unreachable!("async send is only used by bounded queue endpoints") + }; + if state.queue.len() == capacity { + let retired_waker = state + .send_waiters + .register_waker(&mut this.registration, cx); + (Poll::Pending, retired_waker, None) + } else { + let retired_waker = state.send_waiters.unregister_waker(&mut this.registration); + state.queue.push_back( + this.value + .take() + .expect("an incomplete send owns its value"), + ); + let wake_receivers = + (!state.recv_waiters.is_empty()).then(|| state.recv_waiters.take_wakers()); + (Poll::Ready(Ok(())), retired_waker, wake_receivers) + } + } + }; + drop(retired_waker); + if let Some(wakers) = wake_receivers { + wake_all(wakers); + } + poll + } +} + +impl Drop for Send<'_, T> { + fn drop(&mut self) { + if self.registration.is_some() { + self.sender.cancel_send(&mut self.registration); + } + } +} + +struct Recv<'a, T> { + receiver: &'a Receiver, + registration: Option, +} + +impl Future for Recv<'_, T> { + type Output = Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let (poll, retired_waker, wake_senders) = { + let mut state = this.receiver.shared.state.lock(); + if let Some(value) = state.queue.pop_front() { + let retired_waker = state.recv_waiters.unregister_waker(&mut this.registration); + let wake_senders = (matches!(this.receiver.shared.capacity, Capacity::Bounded(_)) + && !state.send_waiters.is_empty()) + .then(|| state.send_waiters.take_wakers()); + (Poll::Ready(Ok(value)), retired_waker, wake_senders) + } else if state.senders == 0 { + let retired_waker = state.recv_waiters.unregister_waker(&mut this.registration); + ( + Poll::Ready(Err(RecvError::Disconnected)), + retired_waker, + None, + ) + } else { + let retired_waker = state + .recv_waiters + .register_waker(&mut this.registration, cx); + (Poll::Pending, retired_waker, None) + } + }; + drop(retired_waker); + if let Some(wakers) = wake_senders { + wake_all(wakers); + } + poll + } +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + if self.registration.is_some() { + self.receiver.cancel_recv(&mut self.registration); + } + } +} + +fn wake_all(wakers: ArenaValues) { + // A wake is not a reservation. Waking every contender prevents available work or capacity + // from being stranded if the first selected future is cancelled before it polls again. + for waker in wakers { + waker.wake(); + } +} diff --git a/asyncband/src/channel/spmc.rs b/asyncband/src/channel/spmc.rs new file mode 100644 index 0000000..b2a16eb --- /dev/null +++ b/asyncband/src/channel/spmc.rs @@ -0,0 +1,197 @@ +// 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 competing queues. +//! +//! The sender is non-cloneable and requires exclusive access. Receivers are cloneable and may +//! receive concurrently; every accepted value is delivered to exactly one receiver. +//! +//! A bounded queue retains at most its requested capacity and waits rather than dropping values. +//! An unbounded queue sends synchronously and grows subject to process memory. Accepted values +//! drain before disconnection is reported, while a send after the last receiver disconnects +//! returns its value. +//! +//! Pending `send` and `recv` operations are cancel safe: dropping their futures does not transfer +//! a value. +//! +//! ``` +//! use asyncband::spmc; +//! +//! let (mut tx, first) = spmc::unbounded(); +//! let second = first.clone(); +//! tx.send(1).unwrap(); +//! tx.send(2).unwrap(); +//! +//! assert_eq!(first.try_recv(), Ok(1)); +//! assert_eq!(second.try_recv(), Ok(2)); +//! ``` +//! +//! The single-producer contract is enforced by the sender type: +//! +//! ```compile_fail +//! let (tx, _) = asyncband::spmc::unbounded::<()>(); +//! let _ = tx.clone(); +//! ``` +//! +//! ```compile_fail +//! fn assert_sync() {} +//! assert_sync::>(); +//! ``` + +use std::cell::Cell; +use std::fmt; +use std::marker::PhantomData; + +pub use super::error::RecvError; +pub use super::error::SendError; +pub use super::error::TryRecvError; +pub use super::error::TrySendError; + +/// Creates a bounded SPMC queue with the given capacity. +/// +/// # Panics +/// +/// Panics if `capacity` is zero. +#[track_caller] +pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { + assert!(capacity > 0, "spmc bounded channel requires capacity > 0"); + let (sender, receiver) = super::queue::bounded(capacity); + ( + BoundedSender { + inner: sender, + not_sync: PhantomData, + }, + BoundedReceiver { inner: receiver }, + ) +} + +/// Creates an unbounded SPMC queue. +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { + let (sender, receiver) = super::queue::unbounded(); + ( + UnboundedSender { + inner: sender, + not_sync: PhantomData, + }, + UnboundedReceiver { inner: receiver }, + ) +} + +/// The sending endpoint of a bounded SPMC queue. +pub struct BoundedSender { + inner: super::queue::Sender, + not_sync: PhantomData>, +} + +impl fmt::Debug for BoundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedSender").finish_non_exhaustive() + } +} + +impl BoundedSender { + /// Sends a value, waiting for capacity when the queue is full. + pub async fn send(&mut self, value: T) -> Result<(), SendError> { + self.inner.send(value).await + } + + /// Attempts to send a value without waiting. + pub fn try_send(&mut self, value: T) -> Result<(), TrySendError> { + self.inner.try_send(value) + } +} + +/// A receiving endpoint of a bounded SPMC queue. +pub struct BoundedReceiver { + inner: super::queue::Receiver, +} + +impl Clone for BoundedReceiver { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl fmt::Debug for BoundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedReceiver").finish_non_exhaustive() + } +} + +impl BoundedReceiver { + /// Receives the next value, waiting while the connected queue is empty. + pub async fn recv(&self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&self) -> Result { + self.inner.try_recv() + } +} + +/// The sending endpoint of an unbounded SPMC queue. +pub struct UnboundedSender { + inner: super::queue::Sender, + not_sync: PhantomData>, +} + +impl fmt::Debug for UnboundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedSender").finish_non_exhaustive() + } +} + +impl UnboundedSender { + /// Sends a value without waiting. + pub fn send(&mut self, value: T) -> Result<(), SendError> { + self.inner.send_unbounded(value) + } +} + +/// A receiving endpoint of an unbounded SPMC queue. +pub struct UnboundedReceiver { + inner: super::queue::Receiver, +} + +impl Clone for UnboundedReceiver { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl fmt::Debug for UnboundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() + } +} + +impl UnboundedReceiver { + /// Receives the next value, waiting while the connected queue is empty. + pub async fn recv(&self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&self) -> Result { + self.inner.try_recv() + } +} diff --git a/asyncband/src/channel/spsc.rs b/asyncband/src/channel/spsc.rs new file mode 100644 index 0000000..1b245ef --- /dev/null +++ b/asyncband/src/channel/spsc.rs @@ -0,0 +1,187 @@ +// 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, single-consumer queues. +//! +//! Both endpoints are non-cloneable. Sending and receiving require exclusive access, preserving +//! the topology's static single-writer and single-reader guarantees. +//! +//! A bounded queue retains at most its requested capacity and waits rather than dropping values. +//! An unbounded queue sends synchronously and grows subject to process memory. Accepted values +//! drain before disconnection is reported, while a send after receiver disconnection returns its +//! value. +//! +//! Pending `send` and `recv` operations are cancel safe: dropping their futures does not transfer +//! a value. +//! +//! ``` +//! use asyncband::spsc; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (mut tx, mut rx) = spsc::bounded(2); +//! tx.send("event").await.unwrap(); +//! assert_eq!(rx.recv().await, Ok("event")); +//! # } +//! ``` +//! +//! The endpoints deliberately cannot be cloned or shared by reference across threads: +//! +//! ```compile_fail +//! let (tx, _) = asyncband::spsc::unbounded::<()>(); +//! let _ = tx.clone(); +//! ``` +//! +//! ```compile_fail +//! fn assert_sync() {} +//! assert_sync::>(); +//! ``` + +use std::cell::Cell; +use std::fmt; +use std::marker::PhantomData; + +pub use super::error::RecvError; +pub use super::error::SendError; +pub use super::error::TryRecvError; +pub use super::error::TrySendError; + +/// Creates a bounded SPSC queue with the given capacity. +/// +/// # Panics +/// +/// Panics if `capacity` is zero. +#[track_caller] +pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { + let (sender, receiver) = super::mpsc::bounded(capacity); + ( + BoundedSender { + inner: sender, + not_sync: PhantomData, + }, + BoundedReceiver { + inner: receiver, + not_sync: PhantomData, + }, + ) +} + +/// Creates an unbounded SPSC queue. +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { + let (sender, receiver) = super::mpsc::unbounded(); + ( + UnboundedSender { + inner: sender, + not_sync: PhantomData, + }, + UnboundedReceiver { + inner: receiver, + not_sync: PhantomData, + }, + ) +} + +/// The sending endpoint of a bounded SPSC queue. +pub struct BoundedSender { + inner: super::mpsc::BoundedSender, + not_sync: PhantomData>, +} + +impl fmt::Debug for BoundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedSender").finish_non_exhaustive() + } +} + +impl BoundedSender { + /// Sends a value, waiting for capacity when the queue is full. + pub async fn send(&mut self, value: T) -> Result<(), SendError> { + self.inner.send(value).await + } + + /// Attempts to send a value without waiting. + pub fn try_send(&mut self, value: T) -> Result<(), TrySendError> { + self.inner.try_send(value) + } +} + +/// The receiving endpoint of a bounded SPSC queue. +pub struct BoundedReceiver { + inner: super::mpsc::BoundedReceiver, + not_sync: PhantomData>, +} + +impl fmt::Debug for BoundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedReceiver").finish_non_exhaustive() + } +} + +impl BoundedReceiver { + /// Receives the next value, waiting while the connected queue is empty. + pub async fn recv(&mut self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&mut self) -> Result { + self.inner.try_recv() + } +} + +/// The sending endpoint of an unbounded SPSC queue. +pub struct UnboundedSender { + inner: super::mpsc::UnboundedSender, + not_sync: PhantomData>, +} + +impl fmt::Debug for UnboundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedSender").finish_non_exhaustive() + } +} + +impl UnboundedSender { + /// Sends a value without waiting. + pub fn send(&mut self, value: T) -> Result<(), SendError> { + self.inner.send(value) + } +} + +/// The receiving endpoint of an unbounded SPSC queue. +pub struct UnboundedReceiver { + inner: super::mpsc::UnboundedReceiver, + not_sync: PhantomData>, +} + +impl fmt::Debug for UnboundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() + } +} + +impl UnboundedReceiver { + /// Receives the next value, waiting while the connected queue is empty. + pub async fn recv(&mut self) -> Result { + self.inner.recv().await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv(&mut self) -> Result { + self.inner.try_recv() + } +} diff --git a/asyncband/src/channel/watch.rs b/asyncband/src/channel/watch.rs new file mode 100644 index 0000000..d103d3a --- /dev/null +++ b/asyncband/src/channel/watch.rs @@ -0,0 +1,288 @@ +// 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 channel that retains and distributes the latest state. +//! +//! Every receiver tracks the last version it observed. Intermediate updates are coalesced, so a +//! slow receiver sees the newest value rather than a backlog of every update. +//! +//! The receiver returned by [`channel`] considers the initial value observed. [`Sender::subscribe`] +//! likewise starts at the current version, while cloning a receiver preserves that receiver's +//! observed version. An unseen final update remains available after every sender disconnects. +//! Dropping a pending [`Receiver::changed`] future does not mark an update observed. +//! +//! ``` +//! use asyncband::watch; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (tx, mut rx) = watch::channel(0); +//! tx.send(1).unwrap(); +//! tx.send(2).unwrap(); +//! +//! assert_eq!(*rx.changed().await.unwrap(), 2); +//! assert_eq!(rx.has_changed(), Ok(false)); +//! # } +//! ``` + +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +pub use super::error::RecvError; +pub use super::error::SendError; +use crate::internal::mutex::Mutex; +use crate::internal::waitset::WaitRegistration; +use crate::internal::waitset::WaitSet; + +/// Creates a watch channel with an initial value. +pub fn channel(initial: T) -> (Sender, Receiver) { + let shared = Arc::new(Shared { + state: Mutex::new(State { + value: Arc::new(initial), + version: 0, + senders: 1, + receivers: 1, + waiters: WaitSet::new(), + }), + }); + ( + Sender { + shared: shared.clone(), + }, + Receiver { shared, seen: 0 }, + ) +} + +struct Shared { + state: Mutex>, +} + +struct State { + value: Arc, + version: u64, + senders: usize, + receivers: usize, + waiters: WaitSet, +} + +/// A sending endpoint of a watch channel. +pub struct Sender { + shared: Arc>, +} + +impl Clone for Sender { + fn clone(&self) -> Self { + let mut state = self.shared.state.lock(); + state.senders = state + .senders + .checked_add(1) + .expect("watch sender count overflow"); + drop(state); + Self { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for Sender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Sender").finish_non_exhaustive() + } +} + +impl Drop for Sender { + fn drop(&mut self) { + let wakers = { + let mut state = self.shared.state.lock(); + state.senders -= 1; + (state.senders == 0).then(|| state.waiters.take_wakers()) + }; + if let Some(wakers) = wakers { + for waker in wakers { + waker.wake(); + } + } + } +} + +impl Sender { + /// Publishes a new current value. + /// + /// Returns the value if no receivers remain. + pub fn send(&self, value: T) -> Result<(), SendError> { + let (wakers, replaced) = { + let mut state = self.shared.state.lock(); + if state.receivers == 0 { + return Err(SendError::new(value)); + } + let version = state + .version + .checked_add(1) + .expect("watch version overflow"); + let replaced = std::mem::replace(&mut state.value, Arc::new(value)); + state.version = version; + let wakers = (!state.waiters.is_empty()).then(|| state.waiters.take_wakers()); + (wakers, replaced) + }; + if let Some(wakers) = wakers { + for waker in wakers { + waker.wake(); + } + } + drop(replaced); + Ok(()) + } + + /// Creates a receiver that considers the current value already observed. + pub fn subscribe(&self) -> Receiver { + let mut state = self.shared.state.lock(); + state.receivers = state + .receivers + .checked_add(1) + .expect("watch receiver count overflow"); + let seen = state.version; + drop(state); + Receiver { + shared: self.shared.clone(), + seen, + } + } + + /// Returns the number of active receivers. + pub fn receiver_count(&self) -> usize { + self.shared.state.lock().receivers + } +} + +/// A receiving endpoint of a watch channel. +pub struct Receiver { + shared: Arc>, + seen: u64, +} + +impl Clone for Receiver { + fn clone(&self) -> Self { + let mut state = self.shared.state.lock(); + state.receivers = state + .receivers + .checked_add(1) + .expect("watch receiver count overflow"); + drop(state); + Self { + shared: self.shared.clone(), + seen: self.seen, + } + } +} + +impl fmt::Debug for Receiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Receiver").finish_non_exhaustive() + } +} + +impl Drop for Receiver { + fn drop(&mut self) { + self.shared.state.lock().receivers -= 1; + } +} + +impl Receiver { + /// Returns the current value without marking its version observed. + pub fn borrow(&self) -> Arc { + self.shared.state.lock().value.clone() + } + + /// Returns the current value and marks its version observed. + pub fn borrow_and_update(&mut self) -> Arc { + let state = self.shared.state.lock(); + self.seen = state.version; + state.value.clone() + } + + /// Returns whether a version newer than the last observed version exists. + pub fn has_changed(&self) -> Result { + let state = self.shared.state.lock(); + if state.version != self.seen { + Ok(true) + } else if state.senders == 0 { + Err(RecvError::Disconnected) + } else { + Ok(false) + } + } + + /// Waits for a newer version and returns the latest value. + pub async fn changed(&mut self) -> Result, RecvError> { + Changed { + receiver: self, + registration: None, + } + .await + } + + /// Returns whether every sender has been dropped. + pub fn is_disconnected(&self) -> bool { + self.shared.state.lock().senders == 0 + } +} + +struct Changed<'a, T> { + receiver: &'a mut Receiver, + registration: Option, +} + +impl Future for Changed<'_, T> { + type Output = Result, RecvError>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let (poll, retired_waker) = { + let mut state = this.receiver.shared.state.lock(); + if state.version != this.receiver.seen { + let retired_waker = state.waiters.unregister_waker(&mut this.registration); + this.receiver.seen = state.version; + (Poll::Ready(Ok(state.value.clone())), retired_waker) + } else if state.senders == 0 { + let retired_waker = state.waiters.unregister_waker(&mut this.registration); + (Poll::Ready(Err(RecvError::Disconnected)), retired_waker) + } else { + let retired_waker = state.waiters.register_waker(&mut this.registration, cx); + (Poll::Pending, retired_waker) + } + }; + drop(retired_waker); + poll + } +} + +impl Drop for Changed<'_, T> { + fn drop(&mut self) { + if self.registration.is_none() { + return; + } + let waker = { + let mut state = self.receiver.shared.state.lock(); + state.waiters.unregister_waker(&mut self.registration) + }; + drop(waker); + } +} diff --git a/asyncband/src/internal/arena.rs b/asyncband/src/internal/arena.rs index d35ca30..855d0c0 100644 --- a/asyncband/src/internal/arena.rs +++ b/asyncband/src/internal/arena.rs @@ -117,6 +117,15 @@ impl Arena { } } + #[cfg(test)] + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + pub fn remove(&mut self, key: ArenaKey) -> T { let index = key.0; let slot = self @@ -160,11 +169,6 @@ impl Arena { self.len = 0; values } - - #[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 248f9ee..edc5104 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -15,16 +15,21 @@ // specific language governing permissions and limitations // under the License. -#[cfg(feature = "mpsc")] +#[cfg(any(feature = "mpsc", feature = "spsc"))] pub(crate) mod atomic_waker; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", + feature = "mpmc", feature = "mpsc", feature = "mutex", feature = "rwlock", feature = "semaphore", + feature = "spmc", + feature = "spsc", + feature = "watch", feature = "waitgroup", ))] // `WaitList` and `WaitSet` use different `Arena` operations. A single-primitive build therefore @@ -44,11 +49,16 @@ pub(crate) mod once_table; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", + feature = "mpmc", feature = "mpsc", feature = "mutex", feature = "rwlock", feature = "semaphore", + feature = "spmc", + feature = "spsc", + feature = "watch", feature = "waitgroup", ))] pub(crate) mod mutex; @@ -58,9 +68,11 @@ pub(crate) mod mutex; feature = "mutex", feature = "rwlock", feature = "semaphore", + feature = "spsc", ))] -// `mpsc` uses `poll_acquire`, `release_if_nonempty`, and `notify_all`; mutexes and rwlocks use -// `acquire`, `try_acquire`, and `release`; the public semaphore also uses the accounting methods. +// The MPSC backend, also reused by SPSC, uses `poll_acquire`, `release_if_nonempty`, and +// `notify_all`; mutexes and rwlocks use `acquire`, `try_acquire`, and `release`; the public +// semaphore also uses the accounting methods. // Each single-primitive build intentionally leaves the other groups unused. #[allow(dead_code)] pub(crate) mod semaphore; @@ -70,13 +82,18 @@ pub(crate) mod semaphore; feature = "mutex", feature = "rwlock", feature = "semaphore", + feature = "spsc", ))] pub(crate) mod waitlist; #[cfg(any( feature = "barrier", + feature = "broadcast", feature = "latch", + feature = "mpmc", feature = "once", + feature = "spmc", + feature = "watch", feature = "waitgroup", ))] // `barrier` constructs a wait set with `with_capacity`, while countdown-based primitives use diff --git a/asyncband/src/internal/waitset.rs b/asyncband/src/internal/waitset.rs index 669e99d..bc67415 100644 --- a/asyncband/src/internal/waitset.rs +++ b/asyncband/src/internal/waitset.rs @@ -57,6 +57,12 @@ impl WaitSet { } } + /// Returns whether no task is currently registered. + #[inline] + pub fn is_empty(&self) -> bool { + self.waiters.is_empty() + } + /// Takes all registered wakers without waking them. #[inline] pub fn take_wakers(&mut self) -> ArenaValues { diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 1c3ffa5..3932e88 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -56,11 +56,16 @@ //! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | //! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::OnceMap`] | `once`, `once-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` | +//! | Transfer one value | [`oneshot::channel`] | `oneshot` | +//! | Compete for queued values | [`spsc`], [`mpsc`], [`spmc`], [`mpmc`] | matching module name | +//! | Broadcast every value | [`broadcast::spmc`], [`broadcast::mpmc`] | `broadcast` | +//! | Observe latest state | [`watch`] | `watch` | //! | Reuse objects | [`pool::bounded`], [`pool::unbounded`] | `pool` | //! | Coordinate workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | //! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | //! +//! Enable `channel` to select every channel feature without adding a public `channel` path. +//! //! # Scope and runtime model //! //! The project is not limited to small or stateless primitives. Stateful tools such as @@ -100,24 +105,31 @@ //! //! While incubation status is not necessarily a reflection of the completeness or stability of the //! code, it does indicate that the project has yet to be fully endorsed by the ASF. +mod channel; mod internal; #[cfg(feature = "barrier")] 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")] pub mod latch; +#[cfg(feature = "mpmc")] +pub use self::channel::mpmc; #[cfg(feature = "mpsc")] -pub mod mpsc; +pub use self::channel::mpsc; #[cfg(feature = "mutex")] pub mod mutex; #[cfg(any(feature = "once", feature = "once-cell", feature = "once-map"))] pub mod once; #[cfg(feature = "oneshot")] -pub mod oneshot; +pub use self::channel::oneshot; +#[cfg(feature = "spsc")] +pub use self::channel::spsc; #[cfg(feature = "pool")] pub mod pool; #[cfg(feature = "rwlock")] @@ -128,8 +140,12 @@ pub mod semaphore; pub mod shutdown; #[cfg(feature = "singleflight")] pub mod singleflight; +#[cfg(feature = "spmc")] +pub use self::channel::spmc; #[cfg(feature = "waitgroup")] pub mod waitgroup; +#[cfg(feature = "watch")] +pub use self::channel::watch; #[cfg(all(test, any(feature = "once-map", feature = "singleflight")))] mod test_support; diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 6905629..37932a8 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -26,8 +26,10 @@ rust-version.workspace = true asyncband = { workspace = true, features = [ "barrier", "blocking", + "broadcast", "condvar", "latch", + "mpmc", "mpsc", "mutex", "once", @@ -39,7 +41,10 @@ asyncband = { workspace = true, features = [ "semaphore", "shutdown", "singleflight", + "spmc", + "spsc", "waitgroup", + "watch", ] } divan = { workspace = true } diff --git a/benchmarks/broadcast.rs b/benchmarks/broadcast.rs new file mode 100644 index 0000000..b99a45d --- /dev/null +++ b/benchmarks/broadcast.rs @@ -0,0 +1,274 @@ +// 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 asyncband::broadcast::mpmc; +use divan::Bencher; +use divan::black_box; + +use super::support::bench_context; +use super::support::poll_pending; +use super::support::poll_pinned_ready; + +const RECEIVER_COUNTS: &[usize] = &[1, 8, 32]; +const SENDER_COUNTS: &[usize] = &[1, 2, 4, 8]; +const CONCURRENT_BATCH_SIZE: usize = 4096; + +#[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: 1 }, + Fanout { + peak: 256, + live: 32, + }, + Fanout { peak: 256, live: 1 }, +]; + +struct ConcurrentSend { + receiver: mpmc::UnboundedReceiver, + start: Arc, + done: Arc, + workers: Vec>, +} + +struct ConcurrentBoundedSend { + _receiver: mpmc::BoundedReceiver, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl ConcurrentBoundedSend { + fn new(sender_count: usize) -> Self { + let (sender, receiver) = mpmc::bounded(CONCURRENT_BATCH_SIZE); + 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.try_send(black_box(value)).unwrap(); + } + done.wait(); + })); + } + drop(sender); + ready.wait(); + + Self { + _receiver: receiver, + start, + done, + workers, + } + } + + fn run(&mut self) { + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for ConcurrentBoundedSend { + fn drop(&mut self) { + for worker in self.workers.drain(..) { + worker.join().unwrap(); + } + } +} + +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)).unwrap(); + } + done.wait(); + })); + } + drop(sender); + ready.wait(); + + Self { + receiver, + start, + done, + workers, + } + } + + fn run(&mut self) { + self.start.wait(); + self.done.wait(); + for _ in 0..CONCURRENT_BATCH_SIZE { + black_box(self.receiver.try_recv().unwrap()); + } + } +} + +impl Drop for ConcurrentSend { + fn drop(&mut self) { + for worker in self.workers.drain(..) { + worker.join().unwrap(); + } + } +} + +#[divan::bench] +fn bounded_round_trip(bencher: Bencher) { + let (sender, mut receiver) = mpmc::bounded(64); + bencher.bench_local(|| { + sender.try_send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn unbounded_round_trip(bencher: Bencher) { + let (sender, mut receiver) = mpmc::unbounded(); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn unbounded_round_trip_two_subscriptions(bencher: Bencher) { + let (sender, mut first) = mpmc::unbounded(); + let mut second = sender.subscribe(); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(first.try_recv().unwrap()); + black_box(second.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn deliver_to_waiting_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (sender, mut receiver) = mpmc::unbounded(); + + bencher.bench_local(|| { + let mut recv = Box::pin(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + sender.send(black_box(1usize)).unwrap(); + black_box(poll_pinned_ready(recv.as_mut(), &mut context).unwrap()) + }); +} + +#[divan::bench(args = RECEIVER_COUNTS)] +fn fanout_round_trip(bencher: Bencher, receiver_count: usize) { + 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()); + } + + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + for receiver in &mut receivers { + black_box(receiver.try_recv().unwrap()); + } + }); +} + +#[divan::bench(args = RECLAIM_FANOUTS)] +fn reclaim_after_receiver_high_water(bencher: Bencher, fanout: Fanout) { + let (sender, initial) = mpmc::unbounded(); + drop(initial); + let mut receivers = (0..fanout.peak) + .map(|_| sender.subscribe()) + .collect::>(); + receivers.truncate(fanout.live); + + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + for receiver in &mut receivers { + black_box(receiver.try_recv().unwrap()); + } + }); +} + +#[divan::bench( + args = SENDER_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 = SENDER_COUNTS, + sample_count = 50, + sample_size = 1, + counters = [CONCURRENT_BATCH_SIZE] +)] +fn concurrent_bounded_send(bencher: Bencher, sender_count: usize) { + bencher + .with_inputs(|| ConcurrentBoundedSend::new(sender_count)) + .bench_local_refs(ConcurrentBoundedSend::run); +} diff --git a/benchmarks/main.rs b/benchmarks/main.rs index cb8ac1a..68fbc57 100644 --- a/benchmarks/main.rs +++ b/benchmarks/main.rs @@ -17,6 +17,7 @@ mod barrier; mod blocking; +mod broadcast; mod condvar; mod latch; mod mpsc; @@ -25,12 +26,14 @@ mod once; mod once_map; mod oneshot; mod pool; +mod queue; mod rwlock; mod semaphore; mod shutdown; mod singleflight; mod support; mod waitgroup; +mod watch; fn main() { divan::main(); diff --git a/benchmarks/mpsc.rs b/benchmarks/mpsc.rs index f86a7ca..486997e 100644 --- a/benchmarks/mpsc.rs +++ b/benchmarks/mpsc.rs @@ -22,8 +22,54 @@ use divan::black_box; use super::support::bench_context; use super::support::poll_pending; use super::support::poll_pinned_ready; +use super::support::poll_ready; const SENDER_COUNTS: &[usize] = &[1, 8, 32]; +const CAPACITIES: &[usize] = &[1, 64]; + +#[divan::bench(args = CAPACITIES)] +fn bounded_try_round_trip(bencher: Bencher, capacity: usize) { + let (sender, mut receiver) = mpsc::bounded(capacity); + + bencher.bench_local(|| { + sender.try_send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench(args = CAPACITIES)] +fn bounded_ready_send_round_trip(bencher: Bencher, capacity: usize) { + let mut context = bench_context(); + let (sender, mut receiver) = mpsc::bounded(capacity); + + bencher.bench_local(|| { + poll_ready(sender.send(black_box(1usize)), &mut context).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn unbounded_round_trip(bencher: Bencher) { + let (sender, mut receiver) = mpsc::unbounded(); + + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn deliver_to_waiting_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (sender, mut receiver) = mpsc::unbounded(); + + bencher.bench_local(|| { + let mut recv = Box::pin(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + sender.send(black_box(1usize)).unwrap(); + black_box(poll_pinned_ready(recv.as_mut(), &mut context).unwrap()) + }); +} #[divan::bench] fn reregister_pending_receiver(bencher: Bencher) { diff --git a/benchmarks/queue.rs b/benchmarks/queue.rs new file mode 100644 index 0000000..4429b46 --- /dev/null +++ b/benchmarks/queue.rs @@ -0,0 +1,84 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use asyncband::mpmc; +use asyncband::spmc; +use asyncband::spsc; +use divan::Bencher; +use divan::black_box; + +use super::support::bench_context; +use super::support::poll_pending; +use super::support::poll_pinned_ready; + +#[divan::bench] +fn spsc_bounded_round_trip(bencher: Bencher) { + let (mut sender, mut receiver) = spsc::bounded(64); + bencher.bench_local(|| { + sender.try_send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn spsc_unbounded_round_trip(bencher: Bencher) { + let (mut sender, mut receiver) = spsc::unbounded(); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn spmc_unbounded_round_trip(bencher: Bencher) { + let (mut sender, receiver) = spmc::unbounded(); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn mpmc_bounded_round_trip(bencher: Bencher) { + let (sender, receiver) = mpmc::bounded(64); + bencher.bench_local(|| { + sender.try_send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn mpmc_unbounded_round_trip(bencher: Bencher) { + let (sender, receiver) = mpmc::unbounded(); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench] +fn mpmc_deliver_to_waiting_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (sender, receiver) = mpmc::unbounded(); + + bencher.bench_local(|| { + let mut recv = Box::pin(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + sender.send(black_box(1usize)).unwrap(); + black_box(poll_pinned_ready(recv.as_mut(), &mut context).unwrap()) + }); +} diff --git a/benchmarks/watch.rs b/benchmarks/watch.rs new file mode 100644 index 0000000..bb31644 --- /dev/null +++ b/benchmarks/watch.rs @@ -0,0 +1,42 @@ +// 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 asyncband::watch; +use divan::Bencher; +use divan::black_box; + +use super::support::bench_context; +use super::support::poll_ready; + +#[divan::bench] +fn send_and_borrow(bencher: Bencher) { + let (sender, receiver) = watch::channel(0usize); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(receiver.borrow()) + }); +} + +#[divan::bench] +fn send_and_changed(bencher: Bencher) { + let mut context = bench_context(); + let (sender, mut receiver) = watch::channel(0usize); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(poll_ready(receiver.changed(), &mut context).unwrap()) + }); +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 27c0594..f3d8d72 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -29,8 +29,10 @@ tokio = { workspace = true, features = ["full"] } asyncband = { workspace = true, features = [ "barrier", "blocking", + "broadcast", "condvar", "latch", + "mpmc", "mpsc", "mutex", "once", @@ -42,7 +44,10 @@ asyncband = { workspace = true, features = [ "semaphore", "shutdown", "singleflight", + "spmc", + "spsc", "waitgroup", + "watch", ] } pollster = { workspace = true, features = ["macro"] } tokio-test = { workspace = true } diff --git a/tests-integration/tests/broadcast_test.rs b/tests-integration/tests/broadcast_test.rs new file mode 100644 index 0000000..6e03104 --- /dev/null +++ b/tests-integration/tests/broadcast_test.rs @@ -0,0 +1,547 @@ +// 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::Poll; +use std::task::Waker; +use std::time::Duration; + +use asyncband::broadcast::mpmc; +use asyncband::broadcast::spmc; +use tokio_test::assert_pending; +use tokio_test::assert_ready_eq; +use tokio_test::task::spawn; + +mod support; +use support::TrackWake; +use support::callback_waker; +use support::poll_with_waker; + +#[test] +fn bounded_capacity_is_strict_and_gated_by_the_slowest_subscription() { + let (tx, mut first) = mpmc::bounded(3); + let mut second = tx.subscribe(); + + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + tx.try_send(3).unwrap(); + assert_eq!(tx.try_send(4), Err(mpmc::TrySendError::Full(4))); + + assert_eq!(first.try_recv(), Ok(1)); + assert_eq!(tx.try_send(4), Err(mpmc::TrySendError::Full(4))); + assert_eq!(second.try_recv(), Ok(1)); + tx.try_send(4).unwrap(); +} + +#[test] +fn bounded_send_wakes_only_after_retention_is_reclaimed() { + let (tx, mut first) = mpmc::bounded(1); + let mut second = tx.subscribe(); + tx.try_send(1).unwrap(); + + let mut send = spawn(tx.send(2)); + assert_pending!(send.poll()); + assert_eq!(first.try_recv(), Ok(1)); + assert_pending!(send.poll()); + assert_eq!(second.try_recv(), Ok(1)); + assert_ready_eq!(send.poll(), Ok(())); + + assert_eq!(first.try_recv(), Ok(2)); + assert_eq!(second.try_recv(), Ok(2)); +} + +#[test] +fn dropping_the_slowest_subscription_releases_bounded_capacity() { + let (tx, mut fast) = mpmc::bounded(1); + let slow = tx.subscribe(); + tx.try_send(1).unwrap(); + assert_eq!(fast.try_recv(), Ok(1)); + + let mut send = spawn(tx.send(2)); + assert_pending!(send.poll()); + drop(slow); + assert_ready_eq!(send.poll(), Ok(())); + assert_eq!(fast.try_recv(), Ok(2)); +} + +#[test] +fn cancelling_a_bounded_send_does_not_publish() { + let (tx, mut rx) = mpmc::bounded(1); + tx.try_send(1).unwrap(); + let mut cancelled = spawn(tx.send(2)); + assert_pending!(cancelled.poll()); + drop(cancelled); + + assert_eq!(rx.try_recv(), Ok(1)); + assert_eq!(rx.try_recv(), Err(mpmc::TryRecvError::Empty)); + tx.try_send(3).unwrap(); + assert_eq!(rx.try_recv(), Ok(3)); +} + +#[test] +fn unbounded_retains_until_every_subscription_advances() { + let (tx, mut first) = mpmc::unbounded(); + let mut second = tx.subscribe(); + tx.send(1).unwrap(); + tx.send(2).unwrap(); + + assert_eq!(first.try_recv(), Ok(1)); + assert_eq!(first.try_recv(), Ok(2)); + assert_eq!(tx.buffer_len(), 2); + assert_eq!(second.try_recv(), Ok(1)); + assert_eq!(tx.buffer_len(), 1); + assert_eq!(second.try_recv(), Ok(2)); + assert_eq!(tx.buffer_len(), 0); +} + +#[test] +fn subscription_counts_and_backlogs_track_each_receiver() { + let (tx, mut first) = mpmc::unbounded(); + assert_eq!(tx.receiver_count(), 1); + assert!(first.is_empty()); + + tx.send(1).unwrap(); + tx.send(2).unwrap(); + let mut second = tx.subscribe(); + tx.send(3).unwrap(); + assert_eq!(tx.receiver_count(), 2); + assert_eq!(first.len(), 3); + assert_eq!(second.len(), 1); + + assert_eq!(second.try_recv(), Ok(3)); + assert!(second.is_empty()); + drop(second); + assert_eq!(tx.receiver_count(), 1); + assert_eq!(first.try_recv(), Ok(1)); + assert_eq!(first.len(), 2); +} + +#[test] +fn a_sender_can_create_a_new_subscription_after_all_previous_ones_drop() { + let (tx, rx) = mpmc::unbounded(); + drop(rx); + assert_eq!(tx.send(1).unwrap_err().into_inner(), 1); + assert_eq!(tx.buffer_len(), 0); + + let mut replacement = tx.subscribe(); + tx.send(2).unwrap(); + assert_eq!(replacement.try_recv(), Ok(2)); +} + +#[test] +fn new_subscriptions_start_at_the_committed_tail() { + let (tx, mut first) = mpmc::unbounded(); + tx.send(1).unwrap(); + let mut second = tx.subscribe(); + tx.send(2).unwrap(); + + assert_eq!(first.try_recv(), Ok(1)); + assert_eq!(first.try_recv(), Ok(2)); + assert_eq!(second.try_recv(), Ok(2)); + assert_eq!(second.try_recv(), Err(mpmc::TryRecvError::Empty)); +} + +#[tokio::test] +async fn accepted_values_drain_before_disconnection() { + let (mut tx, mut rx) = spmc::unbounded(); + tx.send(1).unwrap(); + tx.send(2).unwrap(); + drop(tx); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Err(spmc::RecvError::Disconnected)); +} + +#[test] +fn sends_return_the_value_after_the_last_subscription_drops() { + let (tx, rx) = mpmc::bounded(1); + drop(rx); + assert_eq!(tx.try_send(1), Err(mpmc::TrySendError::Disconnected(1))); + + let (tx, rx) = mpmc::unbounded(); + drop(rx); + assert_eq!(tx.send(2).unwrap_err().into_inner(), 2); +} + +#[test] +fn bounded_capacity_must_be_positive() { + assert!(std::panic::catch_unwind(|| mpmc::bounded::<()>(0)).is_err()); + assert!(std::panic::catch_unwind(|| spmc::bounded::<()>(0)).is_err()); +} + +#[test] +fn concurrent_publishers_expose_one_order_to_every_subscription() { + const PRODUCERS: usize = 4; + const VALUES_PER_PRODUCER: usize = 1_000; + + let (tx, mut first) = mpmc::unbounded(); + let mut second = tx.subscribe(); + std::thread::scope(|scope| { + for producer in 0..PRODUCERS { + let sender = tx.clone(); + scope.spawn(move || { + let start = producer * VALUES_PER_PRODUCER; + for value in start..start + VALUES_PER_PRODUCER { + sender.send(value).unwrap(); + } + }); + } + }); + + let expected = PRODUCERS * VALUES_PER_PRODUCER; + let first = (0..expected) + .map(|_| first.try_recv().unwrap()) + .collect::>(); + let second = (0..expected) + .map(|_| second.try_recv().unwrap()) + .collect::>(); + assert_eq!(first, second); +} + +#[test] +fn concurrent_publishers_deliver_every_value_to_every_subscription() { + const PRODUCERS: usize = 4; + const SUBSCRIPTIONS: usize = 4; + const VALUES_PER_PRODUCER: usize = 500; + + let (tx, first) = mpmc::unbounded(); + let mut receivers = vec![first]; + for _ in 1..SUBSCRIPTIONS { + receivers.push(tx.subscribe()); + } + + let drains = receivers + .into_iter() + .map(|mut receiver| { + std::thread::spawn(move || { + let mut values = Vec::new(); + while let Ok(value) = pollster::block_on(receiver.recv()) { + values.push(value); + } + values + }) + }) + .collect::>(); + + std::thread::scope(|scope| { + for producer in 0..PRODUCERS { + let sender = tx.clone(); + scope.spawn(move || { + let start = producer * VALUES_PER_PRODUCER; + for value in start..start + VALUES_PER_PRODUCER { + sender.send(value).unwrap(); + } + }); + } + }); + drop(tx); + + let expected = (0..PRODUCERS * VALUES_PER_PRODUCER).collect::>(); + for drain in drains { + let mut values = drain.join().unwrap(); + values.sort_unstable(); + assert_eq!(values, expected); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn bounded_concurrent_publishers_make_progress_in_one_shared_order() { + const PRODUCERS: usize = 4; + const VALUES_PER_PRODUCER: usize = 500; + let expected = PRODUCERS * VALUES_PER_PRODUCER; + + let (tx, mut first) = mpmc::bounded(7); + let mut second = tx.subscribe(); + let first_drain = tokio::spawn(async move { + let mut values = Vec::with_capacity(expected); + for _ in 0..expected { + values.push(first.recv().await.unwrap()); + } + values + }); + let second_drain = tokio::spawn(async move { + let mut values = Vec::with_capacity(expected); + for _ in 0..expected { + values.push(second.recv().await.unwrap()); + } + values + }); + + let producers = (0..PRODUCERS) + .map(|producer| { + let sender = tx.clone(); + tokio::spawn(async move { + let start = producer * VALUES_PER_PRODUCER; + for value in start..start + VALUES_PER_PRODUCER { + sender.send(value).await.unwrap(); + } + }) + }) + .collect::>(); + drop(tx); + for producer in producers { + producer.await.unwrap(); + } + + let first = first_drain.await.unwrap(); + let second = second_drain.await.unwrap(); + assert_eq!(first, second); + let mut values = first; + values.sort_unstable(); + assert_eq!(values, (0..expected).collect::>()); +} + +#[derive(Debug)] +struct CloneProbe { + value: usize, + clones: Arc, +} + +impl Clone for CloneProbe { + fn clone(&self) -> Self { + self.clones.fetch_add(1, Ordering::Relaxed); + Self { + value: self.value, + clones: self.clones.clone(), + } + } +} + +#[test] +fn single_subscription_receive_moves_the_payload_without_cloning() { + let clones = Arc::new(AtomicUsize::new(0)); + let (tx, mut rx) = mpmc::unbounded(); + tx.send(CloneProbe { + value: 7, + clones: clones.clone(), + }) + .unwrap(); + + assert_eq!(rx.try_recv().unwrap().value, 7); + assert_eq!(clones.load(Ordering::Relaxed), 0); +} + +struct ReentrantDrop(Option>); + +impl Clone for ReentrantDrop { + fn clone(&self) -> Self { + Self(None) + } +} + +impl Drop for ReentrantDrop { + fn drop(&mut self) { + if let Some(callback) = self.0.take() { + callback(); + } + } +} + +#[test] +fn reclaimed_payloads_are_dropped_after_unlocking_the_channel() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let (tx, mut fast) = mpmc::unbounded(); + let slow = tx.subscribe(); + let callback_sender = tx.clone(); + tx.send(ReentrantDrop(Some(Box::new(move || { + let _ = callback_sender.receiver_count(); + })))) + .unwrap(); + + drop(fast.try_recv().unwrap()); + drop(slow); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("payload destructor deadlocked against the broadcast lock"); +} + +#[derive(Debug)] +struct PanicOnClone { + value: usize, + panic: bool, +} + +impl Clone for PanicOnClone { + fn clone(&self) -> Self { + assert!(!self.panic, "panic while cloning a broadcast value"); + Self { + value: self.value, + panic: self.panic, + } + } +} + +#[test] +fn panicking_payload_clone_leaves_the_channel_consistent() { + let (tx, mut first) = mpmc::unbounded(); + let mut second = tx.subscribe(); + tx.send(PanicOnClone { + value: 1, + panic: true, + }) + .unwrap(); + tx.send(PanicOnClone { + value: 2, + panic: false, + }) + .unwrap(); + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| first.try_recv())); + assert!(result.is_err()); + assert_eq!(first.try_recv().unwrap().value, 2); + assert_eq!(second.try_recv().unwrap().value, 1); + assert_eq!(second.try_recv().unwrap().value, 2); + assert_eq!(tx.buffer_len(), 0); +} + +#[test] +fn cancelled_broadcast_operations_release_their_wakers() { + let (tx, mut rx) = mpmc::bounded::(1); + tx.try_send(0).unwrap(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let baseline = Arc::strong_count(&tracker); + let mut send = Box::pin(tx.send(1)); + assert!(poll_with_waker(send.as_mut(), &waker).is_pending()); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + drop(send); + assert_eq!(Arc::strong_count(&tracker), baseline); + assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(tracker.0.load(Ordering::Relaxed), 0); + + let mut recv = Box::pin(rx.recv()); + assert!(poll_with_waker(recv.as_mut(), &waker).is_pending()); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + drop(recv); + assert_eq!(Arc::strong_count(&tracker), baseline); + tx.try_send(2).unwrap(); + assert_eq!(tracker.0.load(Ordering::Relaxed), 0); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn broadcast_wakers_run_after_unlocking_the_channel() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let (tx, mut rx) = mpmc::unbounded(); + let callback_sender = tx.clone(); + let waker = callback_waker(move || { + callback_sender.send(2).unwrap(); + }); + let mut recv = Box::pin(rx.recv()); + + assert!(poll_with_waker(recv.as_mut(), &waker).is_pending()); + tx.send(1).unwrap(); + assert_eq!(poll_with_waker(recv.as_mut(), &waker), Poll::Ready(Ok(1))); + drop(recv); + assert_eq!(rx.try_recv(), Ok(2)); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("waker callback deadlocked against the broadcast lock"); +} + +#[test] +fn parked_subscription_wakes_when_the_last_sender_drops() { + let (tx, mut rx) = mpmc::unbounded::<()>(); + let other = tx.clone(); + let mut recv = spawn(rx.recv()); + assert_pending!(recv.poll()); + + drop(tx); + assert!(!recv.is_woken()); + drop(other); + assert!(recv.is_woken()); + assert_ready_eq!(recv.poll(), Err(mpmc::RecvError::Disconnected)); +} + +struct Rng(u64); + +impl Rng { + fn below(&mut self, upper: u64) -> u64 { + let mut value = self.0; + value ^= value << 13; + value ^= value >> 7; + value ^= value << 17; + self.0 = value; + value % upper + } +} + +#[test] +fn randomized_subscription_operations_match_a_reference_model() { + for seed in 1..32u64 { + let mut rng = Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1); + let (tx, rx) = mpmc::unbounded::(); + let mut tail = 0u64; + let mut model = vec![(rx, 0u64)]; + + for _ in 0..512 { + match rng.below(100) { + 0..=44 => { + if model.is_empty() { + assert_eq!(tx.send(tail).unwrap_err().into_inner(), tail); + } else { + tx.send(tail).unwrap(); + 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(mpmc::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/queue_topologies_test.rs b/tests-integration/tests/queue_topologies_test.rs new file mode 100644 index 0000000..675e470 --- /dev/null +++ b/tests-integration/tests/queue_topologies_test.rs @@ -0,0 +1,339 @@ +// 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::Poll; +use std::task::Waker; +use std::time::Duration; + +use asyncband::mpmc; +use asyncband::spmc; +use asyncband::spsc; +use tokio_test::assert_pending; +use tokio_test::assert_ready_eq; +use tokio_test::task::spawn; + +mod support; +use support::TrackWake; +use support::callback_waker; +use support::poll_with_waker; + +#[tokio::test] +async fn spsc_preserves_fifo_and_disconnection() { + let (mut tx, mut rx) = spsc::bounded(2); + tx.send(1).await.unwrap(); + tx.send(2).await.unwrap(); + drop(tx); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await, Err(spsc::RecvError::Disconnected)); +} + +#[test] +fn bounded_queue_capacity_must_be_positive() { + assert!(std::panic::catch_unwind(|| spsc::bounded::<()>(0)).is_err()); + assert!(std::panic::catch_unwind(|| spmc::bounded::<()>(0)).is_err()); + assert!(std::panic::catch_unwind(|| mpmc::bounded::<()>(0)).is_err()); +} + +#[tokio::test] +async fn cancelling_a_bounded_send_does_not_enqueue_its_value() { + let (mut tx, rx) = spmc::bounded(1); + tx.send(1).await.unwrap(); + + let mut cancelled = spawn(tx.send(2)); + assert_pending!(cancelled.poll()); + drop(cancelled); + + assert_eq!(rx.recv().await, Ok(1)); + assert_eq!(rx.try_recv(), Err(spmc::TryRecvError::Empty)); + tx.try_send(3).unwrap(); + assert_eq!(rx.recv().await, Ok(3)); +} + +#[test] +fn bounded_send_wakes_after_receive_frees_capacity() { + let (tx, rx) = mpmc::bounded(1); + tx.try_send(1).unwrap(); + + let mut send = spawn(tx.send(2)); + assert_pending!(send.poll()); + assert_eq!(rx.try_recv(), Ok(1)); + assert_ready_eq!(send.poll(), Ok(())); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn freeing_capacity_wakes_every_sender_for_cancellation_transfer() { + let (tx, rx) = mpmc::bounded(1); + tx.try_send(0).unwrap(); + let mut first = spawn(tx.send(1)); + let mut second = spawn(tx.send(2)); + assert_pending!(first.poll()); + assert_pending!(second.poll()); + + assert_eq!(rx.try_recv(), Ok(0)); + assert!(first.is_woken()); + assert!(second.is_woken()); + drop(first); + assert_ready_eq!(second.poll(), Ok(())); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn sending_wakes_every_receiver_for_cancellation_transfer() { + let (tx, rx) = mpmc::unbounded(); + let other = rx.clone(); + let mut first = spawn(rx.recv()); + let mut second = spawn(other.recv()); + assert_pending!(first.poll()); + assert_pending!(second.poll()); + + tx.send(1).unwrap(); + assert!(first.is_woken()); + assert!(second.is_woken()); + drop(first); + assert_ready_eq!(second.poll(), Ok(1)); +} + +#[test] +fn cancelled_queue_operations_release_their_wakers() { + let (tx, rx) = mpmc::bounded::(1); + tx.try_send(0).unwrap(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let baseline = Arc::strong_count(&tracker); + let mut send = Box::pin(tx.send(1)); + assert!(poll_with_waker(send.as_mut(), &waker).is_pending()); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + drop(send); + assert_eq!(Arc::strong_count(&tracker), baseline); + assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(tracker.0.load(Ordering::Relaxed), 0); + + let mut recv = Box::pin(rx.recv()); + assert!(poll_with_waker(recv.as_mut(), &waker).is_pending()); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + drop(recv); + assert_eq!(Arc::strong_count(&tracker), baseline); + tx.try_send(2).unwrap(); + assert_eq!(tracker.0.load(Ordering::Relaxed), 0); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn queue_wakers_run_after_unlocking_the_queue() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let (tx, rx) = mpmc::unbounded(); + let callback_sender = tx.clone(); + let waker = callback_waker(move || { + callback_sender.send(2).unwrap(); + }); + let mut recv = Box::pin(rx.recv()); + + assert!(poll_with_waker(recv.as_mut(), &waker).is_pending()); + tx.send(1).unwrap(); + assert_eq!(poll_with_waker(recv.as_mut(), &waker), Poll::Ready(Ok(1))); + assert_eq!(rx.try_recv(), Ok(2)); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("waker callback deadlocked against the queue lock"); +} + +#[test] +fn dropping_the_last_receiver_wakes_a_backpressured_sender() { + let (tx, rx) = mpmc::bounded(1); + tx.try_send(1).unwrap(); + let mut send = spawn(tx.send(2)); + assert_pending!(send.poll()); + + drop(rx); + assert!(send.is_woken()); + let Poll::Ready(Err(error)) = send.poll() else { + panic!("send should fail after the last receiver drops"); + }; + assert_eq!(error.into_inner(), 2); +} + +#[test] +fn dropping_the_last_sender_wakes_a_parked_receiver() { + let (tx, rx) = mpmc::unbounded::<()>(); + let other = tx.clone(); + let mut recv = spawn(rx.recv()); + assert_pending!(recv.poll()); + + drop(tx); + assert!(!recv.is_woken()); + drop(other); + assert!(recv.is_woken()); + assert_ready_eq!(recv.poll(), Err(mpmc::RecvError::Disconnected)); +} + +#[test] +fn spmc_receivers_compete_for_each_value() { + let (mut tx, rx1) = spmc::unbounded(); + let rx2 = rx1.clone(); + tx.send(1).unwrap(); + tx.send(2).unwrap(); + + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx2.try_recv(), Ok(2)); + assert_eq!(rx1.try_recv(), Err(spmc::TryRecvError::Empty)); + assert_eq!(rx2.try_recv(), Err(spmc::TryRecvError::Empty)); +} + +#[test] +fn dropping_last_receiver_returns_and_drops_queued_values() { + let (mut tx, rx) = spmc::unbounded(); + tx.send(String::from("queued")).unwrap(); + drop(rx); + + let error = tx.send(String::from("unsent")).unwrap_err(); + assert_eq!(error.into_inner(), "unsent"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn mpmc_delivers_every_value_exactly_once_under_contention() { + const PRODUCERS: usize = 4; + const CONSUMERS: usize = 4; + const VALUES_PER_PRODUCER: usize = 2_000; + + let (tx, rx) = mpmc::unbounded(); + let mut consumers = Vec::new(); + for _ in 0..CONSUMERS { + let receiver = rx.clone(); + consumers.push(tokio::spawn(async move { + let mut values = Vec::new(); + while let Ok(value) = receiver.recv().await { + values.push(value); + } + values + })); + } + drop(rx); + + let mut producers = Vec::new(); + for producer in 0..PRODUCERS { + let sender = tx.clone(); + producers.push(tokio::spawn(async move { + let start = producer * VALUES_PER_PRODUCER; + for value in start..start + VALUES_PER_PRODUCER { + sender.send(value).unwrap(); + } + })); + } + drop(tx); + + for producer in producers { + producer.await.unwrap(); + } + + let mut received = Vec::new(); + for consumer in consumers { + received.extend(consumer.await.unwrap()); + } + received.sort_unstable(); + + assert_eq!( + received, + (0..PRODUCERS * VALUES_PER_PRODUCER).collect::>() + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn bounded_mpmc_makes_progress_under_contention() { + const PRODUCERS: usize = 4; + const CONSUMERS: usize = 4; + const VALUES_PER_PRODUCER: usize = 500; + + let (tx, rx) = mpmc::bounded(7); + let consumers = (0..CONSUMERS) + .map(|_| { + let receiver = rx.clone(); + tokio::spawn(async move { + let mut values = Vec::new(); + while let Ok(value) = receiver.recv().await { + values.push(value); + } + values + }) + }) + .collect::>(); + drop(rx); + + let producers = (0..PRODUCERS) + .map(|producer| { + let sender = tx.clone(); + tokio::spawn(async move { + let start = producer * VALUES_PER_PRODUCER; + for value in start..start + VALUES_PER_PRODUCER { + sender.send(value).await.unwrap(); + } + }) + }) + .collect::>(); + drop(tx); + + for producer in producers { + producer.await.unwrap(); + } + let mut received = Vec::new(); + for consumer in consumers { + received.extend(consumer.await.unwrap()); + } + received.sort_unstable(); + assert_eq!( + received, + (0..PRODUCERS * VALUES_PER_PRODUCER).collect::>() + ); +} + +struct ReentrantDrop(Option>); + +impl Drop for ReentrantDrop { + fn drop(&mut self) { + if let Some(callback) = self.0.take() { + callback(); + } + } +} + +#[test] +fn queued_payloads_are_dropped_after_unlocking_the_queue() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let (tx, rx) = mpmc::unbounded(); + let callback_sender = tx.clone(); + tx.send(ReentrantDrop(Some(Box::new(move || { + assert!(callback_sender.send(ReentrantDrop(None)).is_err()); + })))) + .unwrap(); + drop(rx); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("payload destructor deadlocked against the queue lock"); +} diff --git a/tests-integration/tests/support/mod.rs b/tests-integration/tests/support/mod.rs new file mode 100644 index 0000000..58a52f2 --- /dev/null +++ b/tests-integration/tests/support/mod.rs @@ -0,0 +1,53 @@ +// 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::pin::Pin; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Wake; +use std::task::Waker; + +pub struct TrackWake(pub AtomicUsize); + +impl Wake for TrackWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +pub fn poll_with_waker(future: Pin<&mut F>, waker: &Waker) -> Poll { + future.poll(&mut Context::from_waker(waker)) +} + +struct CallbackWake(Mutex>>); + +impl Wake for CallbackWake { + fn wake(self: Arc) { + if let Some(callback) = self.0.lock().unwrap().take() { + callback(); + } + } +} + +pub fn callback_waker(callback: impl FnOnce() + Send + 'static) -> Waker { + Waker::from(Arc::new(CallbackWake(Mutex::new(Some(Box::new(callback)))))) +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 4419543..aa51973 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -15,11 +15,14 @@ // specific language governing permissions and limitations // under the License. +use std::any::TypeId; use std::cell::Cell; use asyncband::barrier::Barrier; +use asyncband::broadcast; use asyncband::condvar::Condvar; use asyncband::latch::Latch; +use asyncband::mpmc; use asyncband::mpsc; use asyncband::mutex::Mutex; use asyncband::mutex::MutexGuard; @@ -39,8 +42,11 @@ use asyncband::shutdown::ShutdownRecv; use asyncband::shutdown::ShutdownSend; use asyncband::shutdown::ShutdownWatch; use asyncband::singleflight; +use asyncband::spmc; +use asyncband::spsc; use asyncband::waitgroup::Wait; use asyncband::waitgroup::WaitGroup; +use asyncband::watch; struct PoolManager; @@ -95,6 +101,20 @@ fn public_types_are_send_and_sync() { assert_send_and_sync::>(); assert_send_and_sync::>(); assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); } #[test] @@ -105,6 +125,14 @@ fn movable_public_types_are_send() { assert_send::>(); assert_send::>(); assert_send::>>(); + assert_send::>(); + assert_send::>(); + assert_send::>(); + assert_send::>(); + assert_send::>(); + assert_send::>(); + assert_send::>(); + assert_send::>(); } #[test] @@ -142,6 +170,28 @@ 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::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); } #[test] @@ -152,3 +202,32 @@ fn unbounded_manual_manager_traits_do_not_depend_on_the_object() { assert_copy::>(); assert_debug::>(); } + +#[test] +fn channel_endpoints_are_nominal_types() { + assert_ne!( + TypeId::of::>(), + TypeId::of::>() + ); + assert_ne!( + TypeId::of::>(), + TypeId::of::>() + ); + assert_ne!( + TypeId::of::>(), + TypeId::of::>() + ); + + assert_eq!( + TypeId::of::>(), + TypeId::of::>() + ); + assert_eq!( + TypeId::of::(), + TypeId::of::() + ); + assert_eq!( + TypeId::of::(), + TypeId::of::() + ); +} diff --git a/tests-integration/tests/watch_test.rs b/tests-integration/tests/watch_test.rs new file mode 100644 index 0000000..c823d99 --- /dev/null +++ b/tests-integration/tests/watch_test.rs @@ -0,0 +1,228 @@ +// 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::Poll; +use std::task::Waker; +use std::time::Duration; + +use asyncband::watch; +use tokio_test::assert_pending; +use tokio_test::assert_ready; +use tokio_test::task::spawn; + +mod support; +use support::TrackWake; +use support::callback_waker; +use support::poll_with_waker; + +#[test] +fn watch_coalesces_updates_to_the_latest_value() { + let (tx, mut rx) = watch::channel(0); + assert_eq!(*rx.borrow(), 0); + assert_eq!(rx.has_changed(), Ok(false)); + + tx.send(1).unwrap(); + tx.send(2).unwrap(); + assert_eq!(rx.has_changed(), Ok(true)); + + let mut changed = spawn(rx.changed()); + let value = assert_ready!(changed.poll()).unwrap(); + assert_eq!(*value, 2); + drop(changed); + assert_eq!(rx.has_changed(), Ok(false)); +} + +#[test] +fn borrow_does_not_mark_a_version_observed() { + let (tx, mut rx) = watch::channel(0); + tx.send(1).unwrap(); + + assert_eq!(*rx.borrow(), 1); + assert_eq!(rx.has_changed(), Ok(true)); + assert_eq!(*rx.borrow_and_update(), 1); + assert_eq!(rx.has_changed(), Ok(false)); +} + +#[test] +fn subscriptions_start_with_the_current_value_observed() { + let (tx, _rx) = watch::channel(0); + tx.send(1).unwrap(); + let mut subscribed = tx.subscribe(); + assert_eq!(*subscribed.borrow(), 1); + assert_eq!(subscribed.has_changed(), Ok(false)); + + tx.send(2).unwrap(); + assert_eq!(*pollster::block_on(subscribed.changed()).unwrap(), 2); +} + +#[test] +fn cloned_receiver_preserves_its_observed_version() { + let (tx, mut first) = watch::channel(0); + tx.send(1).unwrap(); + let mut second = first.clone(); + + assert_eq!(*pollster::block_on(first.changed()).unwrap(), 1); + assert_eq!(*pollster::block_on(second.changed()).unwrap(), 1); +} + +#[test] +fn changed_drains_the_last_update_before_disconnection() { + let (tx, mut rx) = watch::channel(0); + tx.send(1).unwrap(); + drop(tx); + + assert_eq!(*pollster::block_on(rx.changed()).unwrap(), 1); + assert_eq!( + pollster::block_on(rx.changed()), + Err(watch::RecvError::Disconnected) + ); +} + +#[test] +fn cancelling_changed_does_not_consume_the_next_update() { + let (tx, mut rx) = watch::channel(0); + let mut cancelled = spawn(rx.changed()); + assert_pending!(cancelled.poll()); + drop(cancelled); + + tx.send(1).unwrap(); + assert_eq!(*pollster::block_on(rx.changed()).unwrap(), 1); +} + +#[test] +fn send_returns_the_value_when_no_receivers_remain() { + let (tx, rx) = watch::channel(String::from("initial")); + drop(rx); + + assert_eq!( + tx.send(String::from("unsent")).unwrap_err().into_inner(), + "unsent" + ); +} + +#[test] +fn sender_tracks_receivers_and_can_subscribe_after_disconnection() { + let (tx, rx) = watch::channel(0); + assert_eq!(tx.receiver_count(), 1); + drop(rx); + assert_eq!(tx.receiver_count(), 0); + assert_eq!(tx.send(1).unwrap_err().into_inner(), 1); + + let replacement = tx.subscribe(); + assert_eq!(tx.receiver_count(), 1); + assert_eq!(*replacement.borrow(), 0); + tx.send(2).unwrap(); + assert_eq!(*replacement.borrow(), 2); +} + +#[test] +fn changed_wakes_on_update_and_on_final_sender_drop() { + let (tx, mut rx) = watch::channel(0); + let other = tx.clone(); + let mut changed = spawn(rx.changed()); + assert_pending!(changed.poll()); + tx.send(1).unwrap(); + assert!(changed.is_woken()); + assert_eq!(*assert_ready!(changed.poll()).unwrap(), 1); + drop(changed); + + let mut changed = spawn(rx.changed()); + assert_pending!(changed.poll()); + drop(tx); + assert!(!changed.is_woken()); + drop(other); + assert!(changed.is_woken()); + assert_eq!( + assert_ready!(changed.poll()), + Err(watch::RecvError::Disconnected) + ); +} + +#[test] +fn cancelled_changed_releases_its_waker() { + let (tx, mut rx) = watch::channel(0); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let baseline = Arc::strong_count(&tracker); + let mut changed = Box::pin(rx.changed()); + + assert!(poll_with_waker(changed.as_mut(), &waker).is_pending()); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + drop(changed); + assert_eq!(Arc::strong_count(&tracker), baseline); + + tx.send(1).unwrap(); + assert_eq!(tracker.0.load(Ordering::Relaxed), 0); + assert_eq!(*pollster::block_on(rx.changed()).unwrap(), 1); +} + +#[test] +fn watch_wakers_run_after_unlocking_the_channel() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let (tx, mut rx) = watch::channel(0); + let callback_sender = tx.clone(); + let waker = callback_waker(move || { + callback_sender.send(2).unwrap(); + }); + let mut changed = Box::pin(rx.changed()); + + assert!(poll_with_waker(changed.as_mut(), &waker).is_pending()); + tx.send(1).unwrap(); + let Poll::Ready(Ok(value)) = poll_with_waker(changed.as_mut(), &waker) else { + panic!("watch update was not ready"); + }; + assert_eq!(*value, 2); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("waker callback deadlocked against the watch lock"); +} + +struct ReentrantDrop(Option>); + +impl Drop for ReentrantDrop { + fn drop(&mut self) { + if let Some(callback) = self.0.take() { + callback(); + } + } +} + +#[test] +fn replaced_values_are_dropped_after_unlocking_the_channel() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let (tx, _rx) = watch::channel(ReentrantDrop(None)); + let callback_sender = tx.clone(); + tx.send(ReentrantDrop(Some(Box::new(move || { + let _ = callback_sender.receiver_count(); + })))) + .unwrap(); + tx.send(ReentrantDrop(None)).unwrap(); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("value destructor deadlocked against the watch lock"); +}