diff --git a/CHANGELOG.md b/CHANGELOG.md index 0335243..8585812 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,10 +9,15 @@ All notable changes to this project will be documented in this file. ### New features * 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 SPSC, MPSC, SPMC, and MPMC queues with rendezvous, bounded, and unbounded capacities. +* Add overflow, backpressure, and unbounded broadcast retention policies. +* Add coalescing watch channels and single-producer and multi-producer Disruptor-style multicast sequencers. +* Add `sync`, `channel`, and `coordination` public module and Cargo feature groups. ### 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. +* Group all channel families under `asyncband::channel` and remove the previous root-level oneshot, MPSC, and broadcast implementations. * 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. * Rename `oneshot::Sender::is_closed` and `oneshot::Receiver::is_closed` to `is_disconnected`. diff --git a/Cargo.lock b/Cargo.lock index c72436f..92a7cd0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -57,7 +57,9 @@ name = "asyncband" version = "0.6.7" dependencies = [ "hashbrown", + "pollster", "tokio", + "tokio-test", ] [[package]] diff --git a/README.md b/README.md index a3def5b..be095b9 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ Asyncband is a runtime-agnostic library providing essential synchronization prim ## Available primitives -The crate enables no primitives by default. Categories describe each primitive's primary purpose and do not add another module level, so public paths remain concise, such as `asyncband::mutex` and `asyncband::once::OnceCell`. +The crate enables no primitives by default. Synchronization primitives retain their existing root paths, while transfer APIs are grouped under `asyncband::channel`. The `sync`, `channel`, and `coordination` umbrella features enable their corresponding leaf features. | Category | Primitive | Feature | Purpose | | ----------------------- | ------------------------------------------------------------------------------------ | -------------- | ----------------------------------------------------------------------- | @@ -44,13 +44,16 @@ The crate enables no primitives by default. Categories describe each primitive's | | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a one-way countdown completes. | | | [`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. | -| | [`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. | -| | [`broadcast::overflow`](https://docs.rs/asyncband/*/asyncband/broadcast/overflow/) | `broadcast` | Broadcast values and report when slow receivers miss overwritten items. | +| Channels | [`channel::oneshot`](https://docs.rs/asyncband/*/asyncband/channel/oneshot/) | `oneshot` | Send one value between two tasks. | +| | [`channel::{spsc,mpsc,spmc,mpmc}`](https://docs.rs/asyncband/*/asyncband/channel/) | `queue` | Send each value to one competing receiver. | +| | [`channel::broadcast`](https://docs.rs/asyncband/*/asyncband/channel/broadcast/) | `broadcast` | Select overflow, backpressure, or unbounded multicast retention. | +| | [`channel::watch`](https://docs.rs/asyncband/*/asyncband/channel/watch/) | `watch` | Distribute the latest state and coalesce intermediate versions. | +| | [`channel::disruptor`](https://docs.rs/asyncband/*/asyncband/channel/disruptor/) | `disruptor` | Publish through bounded sequenced multicast rings. | | Workload control | [`Semaphore`](https://docs.rs/asyncband/*/asyncband/semaphore/struct.Semaphore.html) | `semaphore` | Control concurrent access with permits. | | | [`Group`](https://docs.rs/asyncband/*/asyncband/singleflight/struct.Group.html) | `singleflight` | Coalesce concurrent calls for the same key. | +See the [channel design for 0.7](docs/channel-design.md) for the topology matrix, overload policies, cancellation rules, and Disruptor invariants. + ## Installation Add the dependency to your `Cargo.toml` via: @@ -99,7 +102,7 @@ All synchronization primitives in this library are runtime-agnostic, meaning the ## Thread Safety -Asyncband primitives and guards implement `Send` and `Sync` only when the protected or transferred value satisfies the necessary bounds. In particular, owned read guards that may move destruction to another thread require the protected value to be `Send` as well as `Sync`. See each type's documentation for its exact bounds. +Asyncband primitives and guards implement `Send` and `Sync` only when the protected or transferred value satisfies the necessary bounds. In particular, owned read guards that may move destruction to another thread require the protected value to be `Send` as well as `Sync`. A channel endpoint is cloneable and `Sync` only when its topology supports multiple producers or consumers; single-side endpoints require exclusive access. See each type's documentation for its exact bounds. ## Minimum Supported Rust Version (MSRV) diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index c36f13f..64ae05d 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -38,22 +38,40 @@ rustdoc-args = ["--cfg", "docsrs"] [features] default = [] +channel = ["broadcast", "disruptor", "oneshot", "queue", "watch"] +coordination = ["shutdown", "singleflight"] +full = ["blocking", "channel", "coordination", "sync"] +sync = [ + "barrier", + "condvar", + "latch", + "mutex", + "once", + "once-cell", + "once-map", + "rwlock", + "semaphore", + "waitgroup", +] + barrier = [] blocking = [] broadcast = [] condvar = ["mutex"] +disruptor = [] latch = [] -mpsc = [] mutex = [] once = ["semaphore"] once-cell = ["semaphore"] once-map = ["dep:hashbrown", "once-cell"] oneshot = [] +queue = [] rwlock = [] semaphore = [] shutdown = ["latch", "waitgroup"] singleflight = ["dep:hashbrown", "once-cell"] waitgroup = [] +watch = [] [dependencies] hashbrown = { workspace = true, default-features = false, features = [ @@ -61,7 +79,9 @@ hashbrown = { workspace = true, default-features = false, features = [ ], optional = true } [dev-dependencies] +pollster.workspace = true tokio = { workspace = true, features = ["full"] } +tokio-test.workspace = true [lints] workspace = true diff --git a/asyncband/src/broadcast/overflow/mod.rs b/asyncband/src/broadcast/overflow/mod.rs deleted file mode 100644 index 176379a..0000000 --- a/asyncband/src/broadcast/overflow/mod.rs +++ /dev/null @@ -1,548 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! A multi-producer multi-consumer broadcast channel. -//! -//! This channel supports multiple senders and multiple receivers. Each message sent by any -//! sender is received by all receivers. If a receiver falls behind, it may miss messages, -//! which is reported via [`RecvError::Lagged`]. -//! -//! # Examples -//! -//! Basic usage: -//! -//! ``` -//! use asyncband::broadcast::overflow; -//! -//! # #[tokio::main] -//! # async fn main() { -//! let (tx, mut rx1) = overflow::channel(16); -//! let mut rx2 = tx.subscribe(); -//! -//! tx.send(10); -//! tx.send(20); -//! -//! assert_eq!(rx1.recv().await, Ok(10)); -//! assert_eq!(rx1.recv().await, Ok(20)); -//! assert_eq!(rx2.recv().await, Ok(10)); -//! assert_eq!(rx2.recv().await, Ok(20)); -//! # } -//! ``` -//! -//! Handling lag: -//! -//! ``` -//! use asyncband::broadcast::overflow; -//! use asyncband::broadcast::overflow::RecvError; -//! -//! # #[tokio::main] -//! # async fn main() { -//! let (tx, mut rx) = overflow::channel(2); -//! -//! tx.send(1); -//! tx.send(2); -//! tx.send(3); // overwrites the oldest message (1) -//! -//! assert_eq!(rx.recv().await, Err(RecvError::Lagged(1))); -//! assert_eq!(rx.recv().await, Ok(2)); -//! assert_eq!(rx.recv().await, Ok(3)); -//! # } -//! ``` - -use std::fmt; -use std::future::Future; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use std::task::Waker; - -use crate::internal::mutex::Mutex; -use crate::internal::rwlock::RwLock; -use crate::internal::waitset::WaitRegistration; -use crate::internal::waitset::WaitSet; - -#[cfg(test)] -mod tests; - -/// Creates a new broadcast channel with the given hint `capacity`. The actual capacity may be -/// greater than the provided `capacity`. -/// -/// See [module-level documentation](self) for broadcast channel semantics. -/// -/// # Panics -/// -/// Panics if `capacity` is 0. -/// -/// # Examples -/// -/// ``` -/// use asyncband::broadcast::overflow; -/// -/// let (tx, mut rx) = overflow::channel(16); -/// tx.send(10); -/// assert_eq!(rx.try_recv(), Ok(10)); -/// ``` -pub fn channel(capacity: usize) -> (Sender, Receiver) { - assert!(capacity > 0, "capacity must be greater than 0"); - - let capacity = capacity.next_power_of_two(); - let mask = capacity - 1; - - let mut buffer = Vec::with_capacity(capacity); - for _ in 0..capacity { - buffer.push(RwLock::new(Slot { - msg: None, - version: 0, - })); - } - - let shared = Arc::new(Shared { - buffer: buffer.into_boxed_slice(), - capacity, - mask, - tail: AtomicU64::new(0), - state: Mutex::new(State { - waiters: WaitSet::new(), - }), - senders: AtomicUsize::new(1), - }); - let sender = Sender { - shared: shared.clone(), - }; - let receiver = Receiver { shared, head: 0 }; - (sender, receiver) -} - -/// Error returned by [`Receiver::recv`]. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum RecvError { - /// The receiver lagged too far behind. - /// - /// The count is the number of messages skipped. The receiver's internal cursor has been - /// advanced to the oldest available message. - Lagged(u64), - /// The sender has become disconnected, and there will never be any more data received on it. - Disconnected, -} - -impl fmt::Display for RecvError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - RecvError::Lagged(n) => write!(f, "receiver has been lagged by {n}"), - RecvError::Disconnected => write!(f, "receiving on a closed channel"), - } - } -} - -impl std::error::Error for RecvError {} - -/// Error returned by [`Receiver::try_recv`]. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum TryRecvError { - /// This channel is currently empty, but the sender(s) have not yet disconnected, so data may - /// yet become available. - Empty, - /// The receiver lagged too far behind. - /// - /// The count is the number of messages skipped. The receiver's internal cursor has been - /// advanced to the oldest available message. - Lagged(u64), - /// The sender has become disconnected, and there will never be any more data received on it. - Disconnected, -} - -impl fmt::Display for TryRecvError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - TryRecvError::Empty => write!(f, "receiving on an empty channel"), - TryRecvError::Lagged(n) => write!(f, "receiver has been lagged by {n}"), - TryRecvError::Disconnected => write!(f, "receiving on a closed channel"), - } - } -} - -impl std::error::Error for TryRecvError {} - -#[derive(Debug)] -struct Slot { - /// The message. `None` if the slot is empty (initial state only). - msg: Option, - /// The absolute version of the message in this slot. - version: u64, -} - -struct Shared { - buffer: Box<[RwLock>]>, - capacity: usize, - mask: usize, - /// The next sequence after the contiguous prefix of fully published slots. - tail: AtomicU64, - /// Serializes senders and makes publishing atomic with draining or registering waiters. - state: Mutex, - /// Number of active senders. - senders: AtomicUsize, -} - -struct State { - /// Receivers waiting for a new message. - waiters: WaitSet, -} - -fn wake_waiters(wakers: impl IntoIterator) { - for waker in wakers { - waker.wake(); - } -} - -/// A sender handle to the broadcast channel. -/// -/// The sender can be cloned to create multiple producers. When all senders are dropped, -/// the channel is closed. -pub struct Sender { - shared: Arc>, -} - -impl Clone for Sender { - fn clone(&self) -> Self { - self.shared.senders.fetch_add(1, Ordering::Release); - Self { - shared: self.shared.clone(), - } - } -} - -impl Drop for Sender { - fn drop(&mut self) { - match self.shared.senders.fetch_sub(1, Ordering::AcqRel) { - 1 => { - // If this is the last sender, we need to wake up the receiver so it can - // observe the disconnected state. - let wakers = self.shared.state.lock().waiters.take_wakers(); - wake_waiters(wakers); - } - _ => { - // there are still other senders left, do nothing - } - } - } -} - -impl Sender { - /// Broadcasts a value to all active receivers. - /// - /// This operation is non-blocking. If the channel buffer is full, the oldest message - /// in the buffer is overwritten. Any receiver that was waiting for that overwritten - /// message will receive a [`RecvError::Lagged`] error on its next call to `recv`. - /// - /// # Examples - /// - /// ``` - /// use asyncband::broadcast::overflow; - /// - /// let (tx, mut rx) = overflow::channel(16); - /// tx.send(10); - /// assert_eq!(rx.try_recv(), Ok(10)); - /// ``` - pub fn send(&self, msg: T) { - let wakers = { - let mut state = self.shared.state.lock(); - let tail = self.shared.tail.load(Ordering::Relaxed); - let idx = (tail as usize) & self.shared.mask; - - let mut slot = self.shared.buffer[idx].write(); - slot.msg = Some(msg); - slot.version = tail; - - // Publish the completed slot before releasing its write lock. A receiver that sees the - // new tail either held the old slot lock first or waits until the new value is - // complete. - self.shared - .tail - .store(tail.wrapping_add(1), Ordering::Release); - drop(slot); - - state.waiters.take_wakers() - }; - - // Notify all waiting receivers. - wake_waiters(wakers); - } - - /// Creates a new receiver that starts receiving messages from the current tail of the channel. - /// - /// # Examples - /// - /// ``` - /// use asyncband::broadcast::overflow; - /// use asyncband::broadcast::overflow::TryRecvError; - /// - /// # #[tokio::main] - /// # async fn main() { - /// let (tx, _) = overflow::channel(16); - /// tx.send(10); - /// - /// let mut rx = tx.subscribe(); - /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - /// tx.send(20); - /// assert_eq!(rx.recv().await, Ok(20)); - /// # } - /// ``` - pub fn subscribe(&self) -> Receiver { - // Receiver starts at the current tail. - let head = self.shared.tail.load(Ordering::Acquire); - let shared = self.shared.clone(); - Receiver { shared, head } - } -} - -/// A receiver handle to the broadcast channel. -/// -/// The receiver can be cloned to create multiple consumers. Each receiver sees every -/// message sent to the channel (unless it lags behind). -pub struct Receiver { - shared: Arc>, - head: u64, -} - -impl Clone for Receiver { - fn clone(&self) -> Self { - Self { - shared: self.shared.clone(), - head: self.head, - } - } -} - -impl Receiver { - /// Receives the next value for this receiver. - /// - /// # Returns - /// - /// * `Ok(T)`: The next message. - /// * `Err(RecvError::Lagged(u64))`: The receiver lagged behind. The internal cursor is advanced - /// to the oldest available message. The count indicates how many messages were skipped. - /// * `Err(RecvError::Disconnected)`: All senders have been dropped and no more messages are - /// available. - /// - /// # Examples - /// - /// ``` - /// use asyncband::broadcast::overflow; - /// - /// # #[tokio::main] - /// # async fn main() { - /// let (tx, mut rx) = overflow::channel(16); - /// tx.send(10); - /// assert_eq!(rx.recv().await, Ok(10)); - /// # } - /// ``` - pub async fn recv(&mut self) -> Result { - Recv { - receiver: self, - registration: None, - } - .await - } - - /// Attempts to receive the next value for this receiver without blocking. - /// - /// # Returns - /// - /// * `Ok(T)`: The next message. - /// * `Err(TryRecvError::Empty)`: No message is currently available. - /// * `Err(TryRecvError::Lagged(u64))`: The receiver lagged behind. The internal cursor is - /// advanced to the oldest available message. The count indicates how many messages were - /// skipped. - /// * `Err(TryRecvError::Disconnected)`: All senders have been dropped and no more messages are - /// available. - /// - /// # Examples - /// - /// ``` - /// use asyncband::broadcast::overflow; - /// - /// let (tx, mut rx) = overflow::channel(16); - /// tx.send(10); - /// assert_eq!(rx.try_recv(), Ok(10)); - /// ``` - pub fn try_recv(&mut self) -> Result { - let shared = &self.shared; - let cap = shared.capacity as u64; - - loop { - let tail = shared.tail.load(Ordering::Acquire); - let head = self.head; - - // diff represents how far behind the head is from the tail. - let diff = tail.wrapping_sub(head); - - // 1. Check for Lag - if diff > cap { - let missed = diff - cap; - self.head = tail.wrapping_sub(cap); - return Err(TryRecvError::Lagged(missed)); - } - - // 2. Check if a message is available - if diff > 0 { - let idx = (head as usize) & shared.mask; - let slot = shared.buffer[idx].read(); - - if slot.version == head { - return if let Some(msg) = &slot.msg { - self.head = head.wrapping_add(1); - Ok(msg.clone()) - } else { - Err(TryRecvError::Empty) - }; - } - - drop(slot); - - // The slot may have been overwritten after the first tail snapshot. Publication - // happens while holding the slot write lock, so a fresh tail now includes that - // overwrite and produces an accurate lag count. - let tail = shared.tail.load(Ordering::Acquire); - let diff = tail.wrapping_sub(head); - if diff > cap { - let missed = diff - cap; - self.head = tail.wrapping_sub(cap); - return Err(TryRecvError::Lagged(missed)); - } - - return Err(TryRecvError::Empty); - } - - // 3. No message available (diff == 0). Check for Closed. - if shared.senders.load(Ordering::Acquire) == 0 { - // Observing the final sender drop synchronizes with all preceding sends, but the - // first tail snapshot predates that acquire. Reload it before declaring the - // channel drained so a published final message cannot be hidden by closure. - if shared.tail.load(Ordering::Acquire) != head { - continue; - } - return Err(TryRecvError::Disconnected); - } - - return Err(TryRecvError::Empty); - } - } -} - -impl Receiver { - /// Re-subscribes to the channel, returning a new receiver that starts receiving messages - /// from the *current* tail of the channel. - /// - /// This is useful if the receiver has lagged too far behind and wants to jump to the latest - /// message, skipping everything in between. - /// - /// # Examples - /// - /// ``` - /// use asyncband::broadcast::overflow; - /// - /// let (tx, mut rx) = overflow::channel(2); - /// tx.send(1); - /// tx.send(2); - /// - /// let mut rx2 = rx.resubscribe(); - /// tx.send(3); - /// - /// assert_eq!(rx2.try_recv(), Ok(3)); - /// ``` - pub fn resubscribe(&self) -> Self { - // Resubscribe starts at the current tail. - let head = self.shared.tail.load(Ordering::Acquire); - let shared = self.shared.clone(); - Self { shared, head } - } -} - -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 Self { - receiver, - registration, - } = self.get_mut(); - - loop { - // Senders publish data or closure before draining the current wake epoch. Once a - // result is observable, Drop does not need to lock the waiter set again. - match receiver.try_recv() { - Ok(val) => { - *registration = None; - return Poll::Ready(Ok(val)); - } - Err(TryRecvError::Lagged(n)) => { - *registration = None; - return Poll::Ready(Err(RecvError::Lagged(n))); - } - Err(TryRecvError::Disconnected) => { - *registration = None; - return Poll::Ready(Err(RecvError::Disconnected)); - } - Err(TryRecvError::Empty) => {} - } - - let shared = &receiver.shared; - let mut state = shared.state.lock(); - - // Double check tail to avoid race conditions. - if shared.tail.load(Ordering::Acquire) != receiver.head { - // New message arrived while acquiring the lock. Retry. - drop(state); - continue; - } - - // Check for Closed - // Use Acquire to ensure we see all writes before the sender dropped. - if shared.senders.load(Ordering::Acquire) == 0 { - *registration = None; - return Poll::Ready(Err(RecvError::Disconnected)); - } - - // Register Waker - let replaced_waker = state.waiters.register_waker(registration, cx); - drop(state); - drop(replaced_waker); - return Poll::Pending; - } - } -} - -impl Drop for Recv<'_, T> { - fn drop(&mut self) { - if self.registration.is_some() { - let removed_waker = { - let mut state = self.receiver.shared.state.lock(); - state.waiters.unregister_waker(&mut self.registration) - }; - drop(removed_waker); - } - } -} diff --git a/asyncband/src/broadcast/overflow/tests.rs b/asyncband/src/broadcast/overflow/tests.rs deleted file mode 100644 index da49b0c..0000000 --- a/asyncband/src/broadcast/overflow/tests.rs +++ /dev/null @@ -1,86 +0,0 @@ -// 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::atomic::Ordering; - -use super::*; - -// These tests stay next to the implementation because they inspect private state. - -#[tokio::test] -async fn sequence_number_wraparound() { - let (tx, mut rx) = channel(4); - let mut rx2 = rx.clone(); - - let boundary = u64::MAX - 2; - tx.shared.tail.store(boundary, Ordering::Release); - rx.head = boundary; - - tx.send(1); - assert_eq!(rx.recv().await, Ok(1)); - - for value in 2..=8 { - tx.send(value); - } - - assert_eq!(rx.recv().await, Err(RecvError::Lagged(3))); - for value in 5..=8 { - assert_eq!(rx.recv().await, Ok(value)); - } - - assert_eq!(rx2.recv().await, Err(RecvError::Lagged(1))); - for value in 5..=8 { - assert_eq!(rx2.recv().await, Ok(value)); - } -} - -#[tokio::test] -async fn sequence_number_wraparound_exactly_overwritten() { - let (tx, mut rx) = channel(4); - let mut rx2 = rx.clone(); - - let boundary = u64::MAX - 2; - tx.shared.tail.store(boundary, Ordering::Release); - rx.head = boundary; - - tx.send(1); - assert_eq!(rx.recv().await, Ok(1)); - - for value in 2..=5 { - tx.send(value); - } - - assert_eq!(rx.recv().await, Ok(2)); - // Wrapping the complete u64 space creates an ABA ambiguity. At 10^9 messages per second this - // takes roughly 584 years, so the implementation accepts it in favor of cheaper arithmetic. - assert_eq!(rx2.recv().await, Ok(4)); -} - -#[test] -fn capacity_is_rounded_to_a_power_of_two() { - let (tx, _) = channel::<()>(3); - assert_eq!(tx.shared.capacity, 4); - assert_eq!(tx.shared.mask, 3); - - let (tx, _) = channel::<()>(4); - assert_eq!(tx.shared.capacity, 4); - assert_eq!(tx.shared.mask, 3); - - let (tx, _) = channel::<()>(5); - assert_eq!(tx.shared.capacity, 8); - assert_eq!(tx.shared.mask, 7); -} diff --git a/asyncband/src/channel/broadcast/mod.rs b/asyncband/src/channel/broadcast/mod.rs new file mode 100644 index 0000000..e4b3a08 --- /dev/null +++ b/asyncband/src/channel/broadcast/mod.rs @@ -0,0 +1,746 @@ +// 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 channels where every receiver observes every retained value. +//! +//! Retention is selected by constructor rather than hidden behind a generic configuration: +//! +//! * [overflow] overwrites the oldest value and reports lag to slow receivers. +//! * [backpressure] gates producers on the slowest receiver. +//! * [unbounded] grows as needed and reclaims values observed by every receiver. +//! +//! ~~~ +//! use std::num::NonZeroUsize; +//! +//! use asyncband::channel::broadcast; +//! +//! let (tx, mut rx) = broadcast::overflow::channel(NonZeroUsize::new(16).unwrap()); +//! tx.send("event").unwrap(); +//! assert_eq!(rx.try_recv(), Ok("event")); +//! ~~~ + +use std::any::type_name; +use std::collections::HashMap; +use std::collections::VecDeque; +use std::fmt; +use std::future::Future; +use std::marker::PhantomData; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use crate::channel::SendError; +use crate::channel::TrySendError; +use crate::channel::wait::WaitQueue; +use crate::channel::wait::wake_all; +use crate::internal::mutex::Mutex; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum InternalRecvError { + Lagged(u64), + Disconnected, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum InternalTryRecvError { + Empty, + Lagged(u64), + Disconnected, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Retention { + Overflow(usize), + Backpressure(usize), + Unbounded, +} + +/// The overflow retention policy. +#[doc(hidden)] +pub struct Overflow; + +/// The backpressure retention policy. +#[doc(hidden)] +pub struct Backpressure; + +/// The unbounded retention policy. +#[doc(hidden)] +pub struct Unbounded; + +/// A sending endpoint whose concrete retention policy is chosen by its constructor module. +#[doc(hidden)] +pub struct Sender { + shared: Arc>, + policy: PhantomData, +} + +/// A receiving endpoint whose concrete retention policy is chosen by its constructor module. +#[doc(hidden)] +pub struct Receiver { + shared: Arc>, + id: u64, + cursor: u64, + policy: PhantomData, +} + +struct Shared { + retention: Retention, + state: Mutex>, +} + +struct State { + log: VecDeque>, + base: u64, + tail: u64, + senders: usize, + receivers: HashMap, + next_receiver_id: u64, + recv_waiters: WaitQueue, + send_waiters: WaitQueue, +} + +fn channel(retention: Retention) -> (Sender, Receiver) { + let mut receivers = HashMap::new(); + receivers.insert(0, 0); + let shared = Arc::new(Shared { + retention, + state: Mutex::new(State { + log: VecDeque::new(), + base: 0, + tail: 0, + senders: 1, + receivers, + next_receiver_id: 1, + recv_waiters: WaitQueue::default(), + send_waiters: WaitQueue::default(), + }), + }); + ( + Sender { + shared: shared.clone(), + policy: PhantomData, + }, + Receiver { + shared, + id: 0, + cursor: 0, + policy: PhantomData, + }, + ) +} + +impl Clone for Sender { + fn clone(&self) -> Self { + let mut state = self.shared.state.lock(); + state.senders = state + .senders + .checked_add(1) + .expect("broadcast sender count overflow"); + drop(state); + Self { + shared: self.shared.clone(), + policy: PhantomData, + } + } +} + +impl fmt::Debug for Sender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Sender") + .field("capacity", &self.capacity()) + .field("len", &self.len()) + .field("receiver_count", &self.receiver_count()) + .finish() + } +} + +impl Sender { + fn try_send_internal(&self, value: T) -> Result> { + let (receivers, wakers, replaced) = { + let mut state = self.shared.state.lock(); + if state.receivers.is_empty() { + return Err(TrySendError::Disconnected(value)); + } + if matches!( + self.shared.retention, + Retention::Backpressure(capacity) if state.log.len() >= capacity + ) { + return Err(TrySendError::Full(value)); + } + let replaced = state.push(self.shared.retention, value); + ( + state.receivers.len(), + state.recv_waiters.take_all(), + replaced, + ) + }; + wake_all(wakers); + drop(replaced); + Ok(receivers) + } + + fn send_nonblocking(&self, value: T) -> Result> { + match self.try_send_internal(value) { + Ok(receivers) => Ok(receivers), + Err(TrySendError::Disconnected(value)) => Err(SendError::new(value)), + Err(TrySendError::Full(_)) => { + unreachable!("nonblocking broadcast retention cannot become full") + } + } + } + + /// Creates a receiver that starts after the latest published value. + pub fn subscribe(&self) -> Receiver { + let mut state = self.shared.state.lock(); + let id = state.allocate_receiver_id(); + let cursor = state.tail; + state.receivers.insert(id, cursor); + drop(state); + Receiver { + shared: self.shared.clone(), + id, + cursor, + policy: PhantomData, + } + } + + /// Returns the configured capacity, or None for unbounded retention. + pub fn capacity(&self) -> Option { + match self.shared.retention { + Retention::Overflow(capacity) | Retention::Backpressure(capacity) => Some(capacity), + Retention::Unbounded => None, + } + } + + /// Returns the number of values currently retained. + pub fn len(&self) -> usize { + self.shared.state.lock().log.len() + } + + /// Returns true if no value is currently retained. + pub fn is_empty(&self) -> bool { + self.shared.state.lock().log.is_empty() + } + + /// Returns the number of active receivers. + pub fn receiver_count(&self) -> usize { + self.shared.state.lock().receivers.len() + } +} + +impl Sender { + /// Sends a value, overwriting the oldest retained value when necessary. + pub fn send(&self, value: T) -> Result> { + self.send_nonblocking(value) + } + + /// Attempts to send a value without waiting. + /// + /// This policy never returns [`TrySendError::Full`]. + pub fn try_send(&self, value: T) -> Result> { + self.try_send_internal(value) + } +} + +impl Sender { + /// Sends a value, waiting until every receiver leaves room in the retained window. + pub async fn send(&self, value: T) -> Result> { + Send { + shared: &self.shared, + value: Some(value), + waiter: None, + completed: false, + } + .await + } + + /// Attempts to send a value without waiting. + pub fn try_send(&self, value: T) -> Result> { + self.try_send_internal(value) + } +} + +impl Sender { + /// Sends a value without waiting. + pub fn send(&self, value: T) -> Result> { + self.send_nonblocking(value) + } + + /// Attempts to send a value without waiting. + /// + /// This policy never returns [`TrySendError::Full`]. + pub fn try_send(&self, value: T) -> Result> { + self.try_send_internal(value) + } +} + +impl Drop for Sender { + fn drop(&mut self) { + let wakers = { + let mut state = self.shared.state.lock(); + state.senders -= 1; + if state.senders == 0 { + state.recv_waiters.take_all() + } else { + Vec::new() + } + }; + wake_all(wakers); + } +} + +impl Clone for Receiver { + fn clone(&self) -> Self { + let mut state = self.shared.state.lock(); + let id = state.allocate_receiver_id(); + state.receivers.insert(id, self.cursor); + drop(state); + Self { + shared: self.shared.clone(), + id, + cursor: self.cursor, + policy: PhantomData, + } + } +} + +impl fmt::Debug for Receiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Receiver") + .field("cursor", &self.cursor) + .field("disconnected", &self.is_disconnected()) + .finish() + } +} + +impl Receiver { + async fn recv_internal(&mut self) -> Result { + Recv { + receiver: self, + waiter: None, + completed: false, + } + .await + } + + fn try_recv_internal(&mut self) -> Result { + self.try_recv_arc().map(|value| (*value).clone()) + } +} + +impl Receiver { + /// Receives the next retained value. + pub async fn recv(&mut self) -> Result { + self.recv_internal().await.map_err(|error| match error { + InternalRecvError::Lagged(count) => overflow::RecvError::Lagged(count), + InternalRecvError::Disconnected => overflow::RecvError::Disconnected, + }) + } + + /// Attempts to receive the next retained value without waiting. + pub fn try_recv(&mut self) -> Result { + self.try_recv_internal().map_err(|error| match error { + InternalTryRecvError::Empty => overflow::TryRecvError::Empty, + InternalTryRecvError::Lagged(count) => overflow::TryRecvError::Lagged(count), + InternalTryRecvError::Disconnected => overflow::TryRecvError::Disconnected, + }) + } +} + +macro_rules! impl_lossless_receiver { + ($policy:ty) => { + impl Receiver { + /// Receives the next retained value. + pub async fn recv(&mut self) -> Result { + self.recv_internal().await.map_err(|error| match error { + InternalRecvError::Disconnected => crate::channel::RecvError::Disconnected, + InternalRecvError::Lagged(_) => { + unreachable!("a lossless broadcast receiver cannot lag") + } + }) + } + + /// Attempts to receive the next retained value without waiting. + pub fn try_recv(&mut self) -> Result { + self.try_recv_internal().map_err(|error| match error { + InternalTryRecvError::Empty => crate::channel::TryRecvError::Empty, + InternalTryRecvError::Disconnected => { + crate::channel::TryRecvError::Disconnected + } + InternalTryRecvError::Lagged(_) => { + unreachable!("a lossless broadcast receiver cannot lag") + } + }) + } + } + }; +} + +impl_lossless_receiver!(Backpressure); +impl_lossless_receiver!(Unbounded); + +impl Receiver { + /// Creates a receiver that starts after the latest published value. + pub fn resubscribe(&self) -> Self { + let mut state = self.shared.state.lock(); + let id = state.allocate_receiver_id(); + let cursor = state.tail; + state.receivers.insert(id, cursor); + drop(state); + Self { + shared: self.shared.clone(), + id, + cursor, + policy: PhantomData, + } + } + + /// Returns true if every sender has been dropped. + pub fn is_disconnected(&self) -> bool { + self.shared.state.lock().senders == 0 + } + + fn try_recv_arc(&mut self) -> Result, InternalTryRecvError> { + let (result, wakers, reclaimed) = { + let mut state = self.shared.state.lock(); + if self.cursor < state.base { + let missed = state.base - self.cursor; + self.cursor = state.base; + state.receivers.insert(self.id, self.cursor); + ( + Err(InternalTryRecvError::Lagged(missed)), + Vec::new(), + Vec::new(), + ) + } else if self.cursor < state.tail { + let index = usize::try_from(self.cursor - state.base) + .expect("the retained broadcast range fits in memory"); + let value = state.log[index].clone(); + self.cursor += 1; + state.receivers.insert(self.id, self.cursor); + let reclaimed = state.reclaim(self.shared.retention); + (Ok(value), state.send_waiters.take_all(), reclaimed) + } else if state.senders == 0 { + ( + Err(InternalTryRecvError::Disconnected), + Vec::new(), + Vec::new(), + ) + } else { + (Err(InternalTryRecvError::Empty), Vec::new(), Vec::new()) + } + }; + wake_all(wakers); + drop(reclaimed); + result + } +} + +impl Drop for Receiver { + fn drop(&mut self) { + let (wakers, reclaimed) = { + let mut state = self.shared.state.lock(); + state.receivers.remove(&self.id); + let reclaimed = state.reclaim(self.shared.retention); + (state.send_waiters.take_all(), reclaimed) + }; + wake_all(wakers); + drop(reclaimed); + } +} + +impl State { + fn push(&mut self, retention: Retention, value: T) -> Option> { + let next_tail = self + .tail + .checked_add(1) + .expect("broadcast sequence overflow"); + let mut replaced = None; + if let Retention::Overflow(capacity) = retention { + if self.log.len() == capacity { + replaced = self.log.pop_front(); + self.base += 1; + } + } + self.log.push_back(Arc::new(value)); + self.tail = next_tail; + replaced + } + + fn reclaim(&mut self, retention: Retention) -> Vec> { + if matches!(retention, Retention::Overflow(_)) { + return Vec::new(); + } + let retain_from = self.receivers.values().copied().min().unwrap_or(self.tail); + let mut reclaimed = Vec::new(); + while self.base < retain_from { + reclaimed.push( + self.log + .pop_front() + .expect("the retained broadcast prefix exists"), + ); + self.base += 1; + } + reclaimed + } + + fn allocate_receiver_id(&mut self) -> u64 { + loop { + let id = self.next_receiver_id; + self.next_receiver_id = self.next_receiver_id.wrapping_add(1); + if !self.receivers.contains_key(&id) { + return id; + } + } + } +} + +struct Send<'a, T> { + shared: &'a Shared, + value: Option, + waiter: Option, + completed: bool, +} + +impl Unpin for Send<'_, T> {} + +impl Future for Send<'_, T> { + type Output = Result>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = &mut *self; + let mut recv_wakers = Vec::new(); + let mut retired_wakers = Vec::new(); + let mut replaced = None; + let poll = { + let mut state = this.shared.state.lock(); + if state.receivers.is_empty() { + retired_wakers.extend(state.send_waiters.remove(&mut this.waiter)); + Poll::Ready(Err(SendError::new( + this.value + .take() + .expect("an incomplete send owns its value"), + ))) + } else if matches!( + this.shared.retention, + Retention::Backpressure(capacity) if state.log.len() >= capacity + ) { + retired_wakers.extend(state.send_waiters.register(&mut this.waiter, cx.waker())); + Poll::Pending + } else { + retired_wakers.extend(state.send_waiters.remove(&mut this.waiter)); + replaced = state.push( + this.shared.retention, + this.value + .take() + .expect("an incomplete send owns its value"), + ); + recv_wakers = state.recv_waiters.take_all(); + Poll::Ready(Ok(state.receivers.len())) + } + }; + drop(retired_wakers); + wake_all(recv_wakers); + drop(replaced); + if poll.is_ready() { + this.completed = true; + } + poll + } +} + +impl Drop for Send<'_, T> { + fn drop(&mut self) { + if !self.completed { + let retired_waker = { + let mut state = self.shared.state.lock(); + state.send_waiters.remove(&mut self.waiter) + }; + drop(retired_waker); + } + } +} + +struct Recv<'a, T, Policy> { + receiver: &'a mut Receiver, + waiter: Option, + completed: bool, +} + +impl Future for Recv<'_, T, Policy> { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = &mut *self; + let result = this.receiver.try_recv_arc(); + let poll = match result { + Ok(value) => { + this.waiter = None; + Poll::Ready(Ok((*value).clone())) + } + Err(InternalTryRecvError::Lagged(count)) => { + this.waiter = None; + Poll::Ready(Err(InternalRecvError::Lagged(count))) + } + Err(InternalTryRecvError::Disconnected) => { + this.waiter = None; + Poll::Ready(Err(InternalRecvError::Disconnected)) + } + Err(InternalTryRecvError::Empty) => { + let mut state = this.receiver.shared.state.lock(); + let retired_waker = if this.receiver.cursor < state.tail || state.senders == 0 { + this.waiter = None; + drop(state); + cx.waker().wake_by_ref(); + None + } else { + state.recv_waiters.register(&mut this.waiter, cx.waker()) + }; + drop(retired_waker); + Poll::Pending + } + }; + if poll.is_ready() { + this.completed = true; + } + poll + } +} + +impl Drop for Recv<'_, T, Policy> { + fn drop(&mut self) { + if !self.completed { + let retired_waker = { + let mut state = self.receiver.shared.state.lock(); + state.recv_waiters.remove(&mut self.waiter) + }; + drop(retired_waker); + } + } +} + +/// Bounded broadcast that overwrites the oldest value for slow receivers. +pub mod overflow { + use std::fmt; + use std::num::NonZeroUsize; + + use super::Retention; + pub use crate::channel::SendError; + pub use crate::channel::TrySendError; + + /// An error returned when an overflow receiver cannot produce the next value. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub enum RecvError { + /// The receiver missed the contained number of values. + Lagged(u64), + /// Every sender has been dropped and no retained value remains. + Disconnected, + } + + impl fmt::Display for RecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Lagged(count) => write!(f, "receiver lagged by {count} values"), + Self::Disconnected => f.write_str("receiving on a closed broadcast channel"), + } + } + } + + impl std::error::Error for RecvError {} + + /// An error returned by a non-waiting overflow receive. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub enum TryRecvError { + /// No value is currently available, but the channel is connected. + Empty, + /// The receiver missed the contained number of values. + Lagged(u64), + /// Every sender has been dropped and no retained value remains. + Disconnected, + } + + impl fmt::Display for TryRecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Empty => f.write_str("receiving on an empty broadcast channel"), + Self::Lagged(count) => write!(f, "receiver lagged by {count} values"), + Self::Disconnected => f.write_str("receiving on a closed broadcast channel"), + } + } + } + + impl std::error::Error for TryRecvError {} + + /// A sender for overflow retention. + pub type Sender = super::Sender; + /// A receiver for overflow retention. + pub type Receiver = super::Receiver; + + /// Creates a bounded overflow broadcast channel. + pub fn channel(capacity: NonZeroUsize) -> (Sender, Receiver) { + super::channel::(Retention::Overflow(capacity.get())) + } +} + +/// Bounded broadcast that waits for the slowest receiver. +pub mod backpressure { + use std::num::NonZeroUsize; + + use super::Retention; + pub use crate::channel::RecvError; + pub use crate::channel::SendError; + pub use crate::channel::TryRecvError; + pub use crate::channel::TrySendError; + + /// A sender for backpressure retention. + pub type Sender = super::Sender; + /// A receiver for backpressure retention. + pub type Receiver = super::Receiver; + + /// Creates a bounded backpressure broadcast channel. + pub fn channel(capacity: NonZeroUsize) -> (Sender, Receiver) { + super::channel::(Retention::Backpressure(capacity.get())) + } +} + +/// Unbounded broadcast that reclaims values after every receiver advances. +pub mod unbounded { + use super::Retention; + pub use crate::channel::RecvError; + pub use crate::channel::SendError; + pub use crate::channel::TryRecvError; + pub use crate::channel::TrySendError; + + /// A sender for unbounded retention. + pub type Sender = super::Sender; + /// A receiver for unbounded retention. + pub type Receiver = super::Receiver; + + /// Creates an unbounded broadcast channel. + pub fn channel() -> (Sender, Receiver) { + super::channel::(Retention::Unbounded) + } +} + +impl fmt::Debug for Send<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Send") + .field("value", &format_args!("{}(..)", type_name::())) + .finish() + } +} diff --git a/asyncband/src/channel/disruptor/mod.rs b/asyncband/src/channel/disruptor/mod.rs new file mode 100644 index 0000000..ef9ec50 --- /dev/null +++ b/asyncband/src/channel/disruptor/mod.rs @@ -0,0 +1,596 @@ +// 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. + +//! Bounded multicast rings modeled after the LMAX Disruptor sequencer. +//! +//! Unlike a competing-consumer MPMC queue, every subscriber observes every published sequence. +//! Producers reserve ring sequences, write their slots, and publish only the highest contiguous +//! range. Subscriber cursors gate wrap-around, so unread slots are never overwritten. +//! +//! This async-oriented implementation parks tasks with wakers. It intentionally does not expose +//! busy-spin or blocking wait strategies. +//! +//! ~~~ +//! use asyncband::channel::disruptor; +//! +//! let capacity = disruptor::Capacity::new(16).unwrap(); +//! let (mut publisher, mut subscriber) = disruptor::single_producer::channel(capacity); +//! assert_eq!(publisher.try_publish("event"), Ok(0)); +//! assert_eq!(subscriber.try_recv(), Ok((0, "event"))); +//! ~~~ + +use std::cell::Cell; +use std::collections::HashMap; +use std::fmt; +use std::future::Future; +use std::marker::PhantomData; +use std::num::NonZeroUsize; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +pub use crate::channel::RecvError; +pub use crate::channel::SendError; +pub use crate::channel::TryRecvError; +pub use crate::channel::TrySendError; +use crate::channel::wait::WaitQueue; +use crate::channel::wait::wake_all; +use crate::internal::mutex::Mutex; + +/// Marker for the single-producer sequencer. +#[doc(hidden)] +pub struct SingleProducer(PhantomData>); + +/// Marker for the multi-producer sequencer. +#[doc(hidden)] +pub struct MultiProducer; + +/// A validated non-zero power-of-two ring capacity. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Capacity(NonZeroUsize); + +impl Capacity { + /// Creates a capacity when the value is a non-zero power of two. + pub const fn new(value: usize) -> Option { + if value.is_power_of_two() { + match NonZeroUsize::new(value) { + Some(value) => Some(Self(value)), + None => None, + } + } else { + None + } + } + + /// Returns the capacity as a usize. + pub const fn get(self) -> usize { + self.0.get() + } +} + +/// A ring publisher parameterized by its producer sequencer. +#[doc(hidden)] +pub struct Publisher { + shared: Arc>, + marker: PhantomData, +} + +/// A multicast subscriber that gates ring wrap-around. +#[doc(hidden)] +pub struct Subscriber { + shared: Arc>, + id: u64, + cursor: u64, +} + +struct Shared { + capacity: usize, + mask: usize, + slots: Box<[Mutex>]>, + state: Mutex, +} + +struct Slot { + sequence: Option, + value: Option>, +} + +struct State { + next_claim: u64, + published: u64, + available: Box<[Option]>, + publishers: usize, + subscribers: HashMap, + next_subscriber_id: u64, + publish_waiters: WaitQueue, + recv_waiters: WaitQueue, +} + +type TryRecvArc = Result<(u64, Arc), TryRecvError>; + +fn channel_with_capacity(capacity: Capacity) -> (Publisher, Subscriber) { + let capacity = capacity.get(); + let mut subscribers = HashMap::new(); + subscribers.insert(0, 0); + let slots = (0..capacity) + .map(|_| { + Mutex::new(Slot { + sequence: None, + value: None, + }) + }) + .collect::>() + .into_boxed_slice(); + let shared = Arc::new(Shared { + capacity, + mask: capacity - 1, + slots, + state: Mutex::new(State { + next_claim: 0, + published: 0, + available: vec![None; capacity].into_boxed_slice(), + publishers: 1, + subscribers, + next_subscriber_id: 1, + publish_waiters: WaitQueue::default(), + recv_waiters: WaitQueue::default(), + }), + }); + ( + Publisher { + shared: shared.clone(), + marker: PhantomData, + }, + Subscriber { + shared, + id: 0, + cursor: 0, + }, + ) +} + +impl fmt::Debug for Publisher { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Publisher") + .field("capacity", &self.shared.capacity) + .field("cursor", &self.cursor()) + .field("remaining_capacity", &self.remaining_capacity()) + .finish() + } +} + +impl Clone for Publisher { + fn clone(&self) -> Self { + let mut state = self.shared.state.lock(); + state.publishers = state + .publishers + .checked_add(1) + .expect("disruptor publisher count overflow"); + drop(state); + Self { + shared: self.shared.clone(), + marker: PhantomData, + } + } +} + +impl Publisher { + /// Creates a subscriber starting at the highest contiguous published sequence. + pub fn subscribe(&self) -> Subscriber { + let mut state = self.shared.state.lock(); + let id = state.allocate_subscriber_id(); + let cursor = state.published; + state.subscribers.insert(id, cursor); + drop(state); + Subscriber { + shared: self.shared.clone(), + id, + cursor, + } + } + + /// Returns the next sequence after the highest contiguous publication. + pub fn cursor(&self) -> u64 { + self.shared.state.lock().published + } + + /// Returns the number of sequences that may currently be reserved. + pub fn remaining_capacity(&self) -> usize { + let state = self.shared.state.lock(); + state.remaining_capacity(self.shared.capacity) + } +} + +macro_rules! publish_methods { + ($mode:ty, $this:ident, $($receiver:tt)+) => { + impl Publisher { + /// Publishes a value, waiting until every subscriber has released its ring slot. + pub async fn publish($($receiver)+, value: T) -> Result> { + Publish { + shared: &$this.shared, + value: Some(value), + waiter: None, + completed: false, + } + .await + } + + /// Attempts to reserve and publish a value without waiting. + pub fn try_publish($($receiver)+, value: T) -> Result> { + $this.shared.try_publish(value) + } + } + }; +} + +publish_methods!(SingleProducer, self, &mut self); +publish_methods!(MultiProducer, self, &self); + +impl Drop for Publisher { + fn drop(&mut self) { + let wakers = { + let mut state = self.shared.state.lock(); + state.publishers -= 1; + if state.publishers == 0 { + state.recv_waiters.take_all() + } else { + Vec::new() + } + }; + wake_all(wakers); + } +} + +impl Clone for Subscriber { + fn clone(&self) -> Self { + let mut state = self.shared.state.lock(); + let id = state.allocate_subscriber_id(); + state.subscribers.insert(id, self.cursor); + drop(state); + Self { + shared: self.shared.clone(), + id, + cursor: self.cursor, + } + } +} + +impl fmt::Debug for Subscriber { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Subscriber") + .field("cursor", &self.cursor) + .field("disconnected", &self.is_disconnected()) + .finish() + } +} + +impl Subscriber { + /// Receives the next published sequence and value. + pub async fn recv(&mut self) -> Result<(u64, T), RecvError> { + Recv { + subscriber: self, + waiter: None, + completed: false, + } + .await + } + + /// Attempts to receive the next published sequence and value without waiting. + pub fn try_recv(&mut self) -> Result<(u64, T), TryRecvError> { + self.try_recv_arc() + .map(|(sequence, value)| (sequence, (*value).clone())) + } +} + +impl Subscriber { + /// Returns the next sequence this subscriber will receive. + pub fn cursor(&self) -> u64 { + self.cursor + } + + /// Returns true if every publisher has been dropped. + pub fn is_disconnected(&self) -> bool { + self.shared.state.lock().publishers == 0 + } + + fn try_recv_arc(&mut self) -> TryRecvArc { + let (result, wakers) = { + let mut state = self.shared.state.lock(); + if self.cursor < state.published { + let sequence = self.cursor; + let slot = self.shared.slots[sequence as usize & self.shared.mask].lock(); + assert_eq!( + slot.sequence, + Some(sequence), + "a published disruptor sequence must occupy its gated slot" + ); + let value = slot + .value + .as_ref() + .expect("a published disruptor slot contains a value") + .clone(); + drop(slot); + self.cursor += 1; + state.subscribers.insert(self.id, self.cursor); + (Ok((sequence, value)), state.publish_waiters.take_all()) + } else if state.publishers == 0 { + (Err(TryRecvError::Disconnected), Vec::new()) + } else { + (Err(TryRecvError::Empty), Vec::new()) + } + }; + wake_all(wakers); + result + } +} + +impl Drop for Subscriber { + fn drop(&mut self) { + let wakers = { + let mut state = self.shared.state.lock(); + state.subscribers.remove(&self.id); + state.publish_waiters.take_all() + }; + wake_all(wakers); + } +} + +impl Shared { + fn try_publish(&self, value: T) -> Result> { + let sequence = { + let mut state = self.state.lock(); + if state.subscribers.is_empty() { + return Err(TrySendError::Disconnected(value)); + } + if state.remaining_capacity(self.capacity) == 0 { + return Err(TrySendError::Full(value)); + } + state.claim() + }; + self.finish_publish(sequence, value); + Ok(sequence) + } + + fn finish_publish(&self, sequence: u64, value: T) { + let index = sequence as usize & self.mask; + let replaced = { + let mut slot = self.slots[index].lock(); + let replaced = slot.value.replace(Arc::new(value)); + slot.sequence = Some(sequence); + replaced + }; + + let wakers = { + let mut state = self.state.lock(); + state.available[index] = Some(sequence); + let before = state.published; + while state.available[state.published as usize & self.mask] == Some(state.published) { + state.published += 1; + } + if state.published != before { + state.recv_waiters.take_all() + } else { + Vec::new() + } + }; + wake_all(wakers); + drop(replaced); + } +} + +impl State { + fn claim(&mut self) -> u64 { + let sequence = self.next_claim; + self.next_claim = self + .next_claim + .checked_add(1) + .expect("disruptor sequence overflow"); + sequence + } + + fn remaining_capacity(&self, capacity: usize) -> usize { + let gating = self + .subscribers + .values() + .copied() + .min() + .unwrap_or(self.next_claim); + let used = usize::try_from(self.next_claim - gating) + .expect("the gated disruptor range fits in memory"); + capacity - used + } + + fn allocate_subscriber_id(&mut self) -> u64 { + loop { + let id = self.next_subscriber_id; + self.next_subscriber_id = self.next_subscriber_id.wrapping_add(1); + if !self.subscribers.contains_key(&id) { + return id; + } + } + } +} + +struct Publish<'a, T> { + shared: &'a Shared, + value: Option, + waiter: Option, + completed: bool, +} + +impl Unpin for Publish<'_, T> {} + +impl Future for Publish<'_, T> { + type Output = Result>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = &mut *self; + let mut retired_wakers = Vec::new(); + let sequence = { + let mut state = this.shared.state.lock(); + if state.subscribers.is_empty() { + retired_wakers.extend(state.publish_waiters.remove(&mut this.waiter)); + this.completed = true; + return Poll::Ready(Err(SendError::new( + this.value + .take() + .expect("an incomplete publication owns its value"), + ))); + } + if state.remaining_capacity(this.shared.capacity) == 0 { + retired_wakers.extend(state.publish_waiters.register(&mut this.waiter, cx.waker())); + return Poll::Pending; + } + retired_wakers.extend(state.publish_waiters.remove(&mut this.waiter)); + state.claim() + }; + + drop(retired_wakers); + this.shared.finish_publish( + sequence, + this.value + .take() + .expect("an incomplete publication owns its value"), + ); + this.completed = true; + Poll::Ready(Ok(sequence)) + } +} + +impl Drop for Publish<'_, T> { + fn drop(&mut self) { + if !self.completed { + let retired_waker = { + let mut state = self.shared.state.lock(); + state.publish_waiters.remove(&mut self.waiter) + }; + drop(retired_waker); + } + } +} + +struct Recv<'a, T> { + subscriber: &'a mut Subscriber, + waiter: Option, + completed: bool, +} + +impl Future for Recv<'_, T> { + type Output = Result<(u64, T), RecvError>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = &mut *self; + let result = this.subscriber.try_recv_arc(); + let poll = match result { + Ok((sequence, value)) => { + this.waiter = None; + Poll::Ready(Ok((sequence, (*value).clone()))) + } + Err(TryRecvError::Disconnected) => { + this.waiter = None; + Poll::Ready(Err(RecvError::Disconnected)) + } + Err(TryRecvError::Empty) => { + let mut state = this.subscriber.shared.state.lock(); + let retired_waker = + if this.subscriber.cursor < state.published || state.publishers == 0 { + this.waiter = None; + drop(state); + cx.waker().wake_by_ref(); + None + } else { + state.recv_waiters.register(&mut this.waiter, cx.waker()) + }; + drop(retired_waker); + Poll::Pending + } + }; + if poll.is_ready() { + this.completed = true; + } + poll + } +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + if !self.completed { + let retired_waker = { + let mut state = self.subscriber.shared.state.lock(); + state.recv_waiters.remove(&mut self.waiter) + }; + drop(retired_waker); + } + } +} + +/// A ring with one statically enforced publisher. +pub mod single_producer { + /// The single publisher endpoint. + pub type Publisher = super::Publisher; + /// A multicast subscriber endpoint. + pub type Subscriber = super::Subscriber; + + /// Creates a single-producer Disruptor ring. + pub fn channel(capacity: super::Capacity) -> (Publisher, Subscriber) { + super::channel_with_capacity(capacity) + } +} + +/// A ring with cloneable concurrent publishers. +pub mod multi_producer { + /// A concurrent publisher endpoint. + pub type Publisher = super::Publisher; + /// A multicast subscriber endpoint. + pub type Subscriber = super::Subscriber; + + /// Creates a multi-producer Disruptor ring. + pub fn channel(capacity: super::Capacity) -> (Publisher, Subscriber) { + super::channel_with_capacity(capacity) + } +} + +#[cfg(test)] +mod tests { + use super::MultiProducer; + use super::channel_with_capacity; + use crate::channel::TryRecvError; + + #[test] + fn multi_producer_only_exposes_contiguous_publications() { + let capacity = super::Capacity::new(4).unwrap(); + let (publisher, mut subscriber) = channel_with_capacity::<_, MultiProducer>(capacity); + let first = { + let mut state = publisher.shared.state.lock(); + state.claim() + }; + let second = { + let mut state = publisher.shared.state.lock(); + state.claim() + }; + + publisher.shared.finish_publish(second, 2); + assert_eq!(publisher.cursor(), 0); + assert_eq!(subscriber.try_recv(), Err(TryRecvError::Empty)); + + publisher.shared.finish_publish(first, 1); + assert_eq!(publisher.cursor(), 2); + assert_eq!(subscriber.try_recv(), Ok((0, 1))); + assert_eq!(subscriber.try_recv(), Ok((1, 2))); + } +} diff --git a/asyncband/src/mpsc/error.rs b/asyncband/src/channel/error.rs similarity index 50% rename from asyncband/src/mpsc/error.rs rename to asyncband/src/channel/error.rs index cf5d649..645538a 100644 --- a/asyncband/src/mpsc/error.rs +++ b/asyncband/src/channel/error.rs @@ -18,36 +18,42 @@ use std::any::type_name; 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 +/// Selects which buffered value an explicit lossy send replaces. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FullBehavior { + /// Replace the oldest buffered value. + DropOldest, + /// Replace the newest buffered value. + DropNewest, +} + +/// The result of an explicit lossy send. +#[derive(Debug, Clone, PartialEq, Eq)] +#[must_use] +pub enum SendOutcome { + /// The value was sent without replacing another value. + Sent, + /// The value was sent and the contained buffered value was replaced. + Replaced(T), +} + +/// An error returned when all receivers have been dropped. #[derive(Clone, PartialEq, Eq)] pub struct SendError(T); impl SendError { - /// Get a reference to the message that failed to be sent. + /// Returns a reference to the value that was not sent. pub fn as_inner(&self) -> &T { &self.0 } - /// Consumes the error and returns the message that failed to be sent. + /// Returns the value that was not sent. pub fn into_inner(self) -> T { self.0 } - /// Creates a new `SendError` with the given message. - pub(super) fn new(msg: T) -> SendError { - SendError(msg) + pub(crate) fn new(value: T) -> Self { + Self(value) } } @@ -65,57 +71,57 @@ impl fmt::Debug for SendError { impl std::error::Error for SendError {} -/// Error returned by `try_send`. +/// An error returned by a non-waiting send. #[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 - /// disconnected. + /// The channel is full but still connected. Full(T), - /// The receiver has become disconnected, and there will never be any more data sent on it. + /// All receivers have been dropped. Disconnected(T), } impl TrySendError { - /// Gets a reference to the message that failed to be sent. + /// Returns a reference to the value that was not sent. pub fn as_inner(&self) -> &T { match self { - TrySendError::Full(msg) | TrySendError::Disconnected(msg) => msg, + Self::Full(value) | Self::Disconnected(value) => value, } } - /// Consumes the error and returns the message that failed to be sent. + /// Returns the value that was not sent. pub fn into_inner(self) -> T { match self { - TrySendError::Full(msg) | TrySendError::Disconnected(msg) => msg, + Self::Full(value) | Self::Disconnected(value) => value, } } } impl fmt::Display for TrySendError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(match self { - TrySendError::Full(_) => "sending on a full channel", - TrySendError::Disconnected(_) => "sending on a closed channel", - }) + match self { + Self::Full(_) => f.write_str("sending on a full channel"), + Self::Disconnected(_) => f.write_str("sending on a closed channel"), + } } } impl fmt::Debug for TrySendError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let ty = type_name::(); match self { - TrySendError::Full(_) => write!(f, "TrySendError<{ty}>::Full(..)"), - TrySendError::Disconnected(_) => write!(f, "TrySendError<{ty}>::Disconnected(..)"), + Self::Full(_) => write!(f, "TrySendError<{}>::Full(..)", type_name::()), + Self::Disconnected(_) => { + write!(f, "TrySendError<{}>::Disconnected(..)", type_name::()) + } } } } impl std::error::Error for TrySendError {} -/// Error returned by `recv`. -#[derive(Debug, Clone, PartialEq, Eq)] +/// An error returned when a channel is closed and drained. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RecvError { - /// The sender has become disconnected, and there will never be any more data received on it. + /// All senders have been dropped and no buffered value remains. Disconnected, } @@ -127,22 +133,21 @@ impl fmt::Display for RecvError { impl std::error::Error for RecvError {} -/// Error returned by `try_recv`. -#[derive(Debug, Clone, PartialEq, Eq)] +/// An error returned by a non-waiting receive. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum TryRecvError { - /// This channel is currently empty, but the sender(s) have not yet disconnected, so data may - /// yet become available. + /// No value is currently available, but the channel is still connected. Empty, - /// The sender has become disconnected, and there will never be any more data received on it. + /// All senders have been dropped and no buffered value remains. Disconnected, } impl fmt::Display for TryRecvError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(match self { - TryRecvError::Empty => "receiving on an empty channel", - TryRecvError::Disconnected => "receiving on a closed channel", - }) + match self { + Self::Empty => f.write_str("receiving on an empty channel"), + Self::Disconnected => f.write_str("receiving on a closed channel"), + } } } diff --git a/asyncband/src/channel/mod.rs b/asyncband/src/channel/mod.rs new file mode 100644 index 0000000..c7f2710 --- /dev/null +++ b/asyncband/src/channel/mod.rs @@ -0,0 +1,85 @@ +// 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. + +//! Runtime-agnostic channels for asynchronous tasks. +//! +//! The module separates queue topology from delivery policy: +//! +//! * [oneshot] transfers one value once. +//! * [spsc], [mpsc], [spmc], and [mpmc] are competing-consumer queues. +//! * [broadcast] provides multicast retention policies. +//! * [watch] keeps only the latest value. +//! * [disruptor] provides bounded multicast rings with explicit sequencing. +//! +//! # Example +//! +//! ~~~ +//! use std::num::NonZeroUsize; +//! +//! use asyncband::channel::mpsc; +//! +//! let (tx, mut rx) = mpsc::bounded(NonZeroUsize::new(16).unwrap()); +//! pollster::block_on(async { +//! tx.send("hello").await.unwrap(); +//! assert_eq!(rx.recv().await, Ok("hello")); +//! }); +//! ~~~ + +#[cfg(feature = "broadcast")] +pub mod broadcast; +#[cfg(feature = "disruptor")] +pub mod disruptor; +mod error; +#[cfg(feature = "queue")] +pub mod mpmc; +#[cfg(feature = "queue")] +pub mod mpsc; +#[cfg(feature = "oneshot")] +pub mod oneshot; +#[doc(hidden)] +#[cfg(feature = "queue")] +pub mod queue; +#[cfg(feature = "queue")] +pub mod spmc; +#[cfg(feature = "queue")] +pub mod spsc; +#[cfg(any( + feature = "broadcast", + feature = "disruptor", + feature = "queue", + feature = "watch", +))] +mod wait; +#[cfg(feature = "watch")] +pub mod watch; + +pub use error::FullBehavior; +pub use error::RecvError; +pub use error::SendError; +pub use error::SendOutcome; +pub use error::TryRecvError; +pub use error::TrySendError; + +#[cfg(all( + test, + feature = "broadcast", + feature = "disruptor", + feature = "oneshot", + feature = "queue", + feature = "watch", +))] +mod tests; diff --git a/asyncband/src/channel/mpmc.rs b/asyncband/src/channel/mpmc.rs new file mode 100644 index 0000000..c691544 --- /dev/null +++ b/asyncband/src/channel/mpmc.rs @@ -0,0 +1,52 @@ +// 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 queues. +//! +//! Both endpoints are cloneable and may be shared concurrently. Receivers compete for values, so +//! each value is delivered to exactly one receiver. + +use std::num::NonZeroUsize; + +pub use super::FullBehavior; +pub use super::RecvError; +pub use super::SendError; +pub use super::SendOutcome; +pub use super::TryRecvError; +pub use super::TrySendError; +use super::queue; + +/// A sending endpoint of an MPMC queue. +pub type Sender = queue::Sender; + +/// A receiving endpoint of an MPMC queue. +pub type Receiver = queue::Receiver; + +/// Creates an unbuffered MPMC rendezvous channel. +pub fn rendezvous() -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Rendezvous) +} + +/// Creates a bounded MPMC queue. +pub fn bounded(capacity: NonZeroUsize) -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Bounded(capacity.get())) +} + +/// Creates an unbounded MPMC queue. +pub fn unbounded() -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Unbounded) +} diff --git a/asyncband/src/channel/mpsc.rs b/asyncband/src/channel/mpsc.rs new file mode 100644 index 0000000..185cf28 --- /dev/null +++ b/asyncband/src/channel/mpsc.rs @@ -0,0 +1,52 @@ +// 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, single-consumer queues. +//! +//! Senders are cloneable and may be shared concurrently. The receiver is non-cloneable, +//! non-Sync, and requires mutable access. + +use std::num::NonZeroUsize; + +pub use super::FullBehavior; +pub use super::RecvError; +pub use super::SendError; +pub use super::SendOutcome; +pub use super::TryRecvError; +pub use super::TrySendError; +use super::queue; + +/// A sending endpoint of an MPSC queue. +pub type Sender = queue::Sender; + +/// The receiving endpoint of an MPSC queue. +pub type Receiver = queue::Receiver; + +/// Creates an unbuffered MPSC rendezvous channel. +pub fn rendezvous() -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Rendezvous) +} + +/// Creates a bounded MPSC queue. +pub fn bounded(capacity: NonZeroUsize) -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Bounded(capacity.get())) +} + +/// Creates an unbounded MPSC queue. +pub fn unbounded() -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Unbounded) +} diff --git a/asyncband/src/channel/oneshot.rs b/asyncband/src/channel/oneshot.rs new file mode 100644 index 0000000..3d9e05e --- /dev/null +++ b/asyncband/src/channel/oneshot.rs @@ -0,0 +1,224 @@ +// 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 transfers at most one value. + +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +pub use crate::channel::RecvError; +pub use crate::channel::SendError; +pub use crate::channel::TryRecvError; +use crate::internal::mutex::Mutex; + +/// Creates a one-shot channel. +pub fn channel() -> (Sender, Receiver) { + let shared = Arc::new(Shared { + state: Mutex::new(State { + value: None, + sender_alive: true, + receiver_alive: true, + receiver_waker: None, + }), + }); + ( + Sender { + shared: Some(shared.clone()), + }, + Receiver { + shared: Some(shared), + }, + ) +} + +struct Shared { + state: Mutex>, +} + +struct State { + value: Option, + sender_alive: bool, + receiver_alive: bool, + receiver_waker: Option, +} + +/// The sending half of a one-shot channel. +pub struct Sender { + shared: Option>>, +} + +impl fmt::Debug for Sender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Sender") + .field("completed", &self.shared.is_none()) + .finish() + } +} + +impl Sender { + /// Sends the channel's value. + pub fn send(mut self, value: T) -> Result<(), SendError> { + let shared = self + .shared + .take() + .expect("a one-shot sender can only complete once"); + let waker = { + let mut state = shared.state.lock(); + if !state.receiver_alive { + return Err(SendError::new(value)); + } + state.value = Some(value); + state.sender_alive = false; + state.receiver_waker.take() + }; + if let Some(waker) = waker { + waker.wake(); + } + Ok(()) + } + + /// Returns true if the receiver has been dropped. + pub fn is_disconnected(&self) -> bool { + self.shared + .as_ref() + .is_none_or(|shared| !shared.state.lock().receiver_alive) + } +} + +impl Drop for Sender { + fn drop(&mut self) { + let Some(shared) = self.shared.take() else { + return; + }; + let waker = { + let mut state = shared.state.lock(); + state.sender_alive = false; + state.receiver_waker.take() + }; + if let Some(waker) = waker { + waker.wake(); + } + } +} + +/// The receiving half of a one-shot channel. +pub struct Receiver { + shared: Option>>, +} + +impl fmt::Debug for Receiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Receiver") + .field("completed", &self.shared.is_none()) + .finish() + } +} + +impl Receiver { + /// Receives the value, waiting for the sender when necessary. + pub async fn recv(self) -> Result { + self.await + } + + /// Attempts to receive the value without waiting. + pub fn try_recv(&mut self) -> Result { + let Some(shared) = self.shared.as_ref() else { + return Err(TryRecvError::Disconnected); + }; + let result = { + let mut state = shared.state.lock(); + if let Some(value) = state.value.take() { + state.receiver_alive = false; + Ok(value) + } else if !state.sender_alive { + state.receiver_alive = false; + Err(TryRecvError::Disconnected) + } else { + Err(TryRecvError::Empty) + } + }; + if !matches!(result, Err(TryRecvError::Empty)) { + self.shared = None; + } + result + } + + /// Returns true if no value can still be received. + pub fn is_disconnected(&self) -> bool { + self.shared.as_ref().is_none_or(|shared| { + let state = shared.state.lock(); + !state.sender_alive && state.value.is_none() + }) + } +} + +impl Future for Receiver { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let Some(shared) = self.shared.as_ref() else { + return Poll::Ready(Err(RecvError::Disconnected)); + }; + let mut retired_waker = None; + let result = { + let mut state = shared.state.lock(); + if let Some(value) = state.value.take() { + state.receiver_alive = false; + Some(Ok(value)) + } else if !state.sender_alive { + state.receiver_alive = false; + Some(Err(RecvError::Disconnected)) + } else { + if state + .receiver_waker + .as_ref() + .is_none_or(|waker| !waker.will_wake(cx.waker())) + { + retired_waker = state.receiver_waker.replace(cx.waker().clone()); + } + None + } + }; + drop(retired_waker); + match result { + Some(result) => { + self.shared = None; + Poll::Ready(result) + } + None => Poll::Pending, + } + } +} + +impl Drop for Receiver { + fn drop(&mut self) { + let Some(shared) = self.shared.take() else { + return; + }; + let retired_waker = { + let mut state = shared.state.lock(); + state.receiver_alive = false; + state.receiver_waker.take() + }; + drop(retired_waker); + } +} diff --git a/asyncband/src/channel/queue.rs b/asyncband/src/channel/queue.rs new file mode 100644 index 0000000..6d60b09 --- /dev/null +++ b/asyncband/src/channel/queue.rs @@ -0,0 +1,706 @@ +// 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::cell::Cell; +use std::collections::VecDeque; +use std::fmt; +use std::future::Future; +use std::marker::PhantomData; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use crate::channel::FullBehavior; +use crate::channel::RecvError; +use crate::channel::SendError; +use crate::channel::SendOutcome; +use crate::channel::TryRecvError; +use crate::channel::TrySendError; +use crate::channel::wait::WaitQueue; +use crate::channel::wait::wake_all; +use crate::internal::mutex::Mutex; + +/// Marker for an endpoint that cannot be cloned or shared concurrently. +#[doc(hidden)] +pub struct Single(PhantomData>); + +/// Marker for a cloneable, concurrently shared endpoint. +#[doc(hidden)] +pub struct Multiple; + +#[derive(Debug, Clone, Copy)] +pub(super) enum QueueKind { + Rendezvous, + Bounded(usize), + Unbounded, +} + +/// A sending endpoint parameterized by its producer and consumer cardinalities. +#[doc(hidden)] +pub struct Sender { + core: Arc>, + producer: PhantomData, + consumer: PhantomData Consumer>, +} + +/// A receiving endpoint parameterized by its producer and consumer cardinalities. +#[doc(hidden)] +pub struct Receiver { + core: Arc>, + producer: PhantomData Producer>, + consumer: PhantomData, +} + +pub(super) fn channel( + kind: QueueKind, +) -> ( + Sender, + Receiver, +) { + debug_assert!(!matches!(kind, QueueKind::Bounded(0))); + let core = Arc::new(Core::new(kind)); + ( + Sender { + core: core.clone(), + producer: PhantomData, + consumer: PhantomData, + }, + Receiver { + core, + producer: PhantomData, + consumer: PhantomData, + }, + ) +} + +impl fmt::Debug for Sender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Sender") + .field("capacity", &self.capacity()) + .field("len", &self.len()) + .field("disconnected", &self.is_disconnected()) + .finish() + } +} + +impl fmt::Debug for Receiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Receiver") + .field("capacity", &self.capacity()) + .field("len", &self.len()) + .field("disconnected", &self.is_disconnected()) + .finish() + } +} + +impl Clone for Sender { + fn clone(&self) -> Self { + self.core.add_sender(); + Self { + core: self.core.clone(), + producer: PhantomData, + consumer: PhantomData, + } + } +} + +impl Clone for Receiver { + fn clone(&self) -> Self { + self.core.add_receiver(); + Self { + core: self.core.clone(), + producer: PhantomData, + consumer: PhantomData, + } + } +} + +impl Drop for Sender { + fn drop(&mut self) { + self.core.drop_sender(); + } +} + +impl Drop for Receiver { + fn drop(&mut self) { + self.core.drop_receiver(); + } +} + +macro_rules! sender_methods { + ($mode:ty, $this:ident, $($receiver:tt)+) => { + impl Sender { + /// Sends a value, waiting for capacity when the channel is full. + pub async fn send($($receiver)+, value: T) -> Result<(), SendError> { + Send { + core: &$this.core, + value: Some(value), + waiter: None, + rendezvous_id: None, + completed: false, + } + .await + } + + /// Attempts to send a value without waiting. + pub fn try_send($($receiver)+, value: T) -> Result<(), TrySendError> { + $this.core.try_send(value).map(|_| ()) + } + + /// Sends a value by explicitly replacing one buffered value when full. + /// + /// Rendezvous channels cannot replace a value and return Full unless a receiver is + /// already waiting. Unbounded channels never replace a value. + pub fn force_send( + $($receiver)+, + value: T, + behavior: FullBehavior, + ) -> Result, TrySendError> { + $this.core.force_send(value, behavior) + } + } + }; +} + +sender_methods!(Single, self, &mut self); +sender_methods!(Multiple, self, &self); + +macro_rules! receiver_methods { + ($mode:ty, $this:ident, $($receiver:tt)+) => { + impl Receiver { + /// Receives the next value, waiting while the connected channel is empty. + pub async fn recv($($receiver)+) -> Result { + Recv { + core: &$this.core, + waiter: None, + completed: false, + } + .await + } + + /// Attempts to receive the next value without waiting. + pub fn try_recv($($receiver)+) -> Result { + $this.core.try_recv() + } + } + }; +} + +receiver_methods!(Single, self, &mut self); +receiver_methods!(Multiple, self, &self); + +impl Sender { + /// Returns the configured buffer capacity, or None for an unbounded channel. + pub fn capacity(&self) -> Option { + self.core.capacity() + } + + /// Returns the number of accepted values waiting in the buffer. + pub fn len(&self) -> usize { + self.core.len() + } + + /// Returns true if no accepted value is waiting in the buffer. + pub fn is_empty(&self) -> bool { + self.core.len() == 0 + } + + /// Returns true if no receivers remain. + pub fn is_disconnected(&self) -> bool { + self.core.receiver_count() == 0 + } +} + +impl Receiver { + /// Returns the configured buffer capacity, or None for an unbounded channel. + pub fn capacity(&self) -> Option { + self.core.capacity() + } + + /// Returns the number of accepted values waiting in the buffer. + pub fn len(&self) -> usize { + self.core.len() + } + + /// Returns true if no accepted value is waiting in the buffer. + pub fn is_empty(&self) -> bool { + self.core.len() == 0 + } + + /// Returns true if no senders remain. + pub fn is_disconnected(&self) -> bool { + self.core.sender_count() == 0 + } +} + +struct Core { + kind: QueueKind, + state: Mutex>, +} + +struct State { + queue: VecDeque, + senders: usize, + receivers: usize, + send_waiters: WaitQueue, + recv_waiters: WaitQueue, + rendezvous_sends: VecDeque>, + next_rendezvous_id: u64, +} + +struct PendingSend { + id: u64, + value: T, + waker: Waker, +} + +impl Core { + fn new(kind: QueueKind) -> Self { + Self { + kind, + state: Mutex::new(State { + queue: VecDeque::new(), + senders: 1, + receivers: 1, + send_waiters: WaitQueue::default(), + recv_waiters: WaitQueue::default(), + rendezvous_sends: VecDeque::new(), + next_rendezvous_id: 0, + }), + } + } + + fn capacity(&self) -> Option { + match self.kind { + QueueKind::Rendezvous => Some(0), + QueueKind::Bounded(capacity) => Some(capacity), + QueueKind::Unbounded => None, + } + } + + fn len(&self) -> usize { + self.state.lock().queue.len() + } + + fn sender_count(&self) -> usize { + self.state.lock().senders + } + + fn receiver_count(&self) -> usize { + self.state.lock().receivers + } + + fn add_sender(&self) { + let mut state = self.state.lock(); + state.senders = state + .senders + .checked_add(1) + .expect("channel sender count overflow"); + } + + fn add_receiver(&self) { + let mut state = self.state.lock(); + state.receivers = state + .receivers + .checked_add(1) + .expect("channel receiver count overflow"); + } + + fn drop_sender(&self) { + let wakers = { + let mut state = self.state.lock(); + state.senders -= 1; + if state.senders == 0 { + state.recv_waiters.take_all() + } else { + Vec::new() + } + }; + wake_all(wakers); + } + + fn drop_receiver(&self) { + let wakers = { + let mut state = self.state.lock(); + state.receivers -= 1; + if state.receivers == 0 { + let mut wakers = state.send_waiters.take_all(); + wakers.extend( + state + .rendezvous_sends + .iter() + .map(|pending| pending.waker.clone()), + ); + wakers + } else { + Vec::new() + } + }; + wake_all(wakers); + } + + fn try_send(&self, value: T) -> Result, TrySendError> { + let (result, wakers) = { + let mut state = self.state.lock(); + if state.receivers == 0 { + return Err(TrySendError::Disconnected(value)); + } + + match self.kind { + QueueKind::Rendezvous => { + let wakers = state.recv_waiters.take_all(); + if wakers.is_empty() { + return Err(TrySendError::Full(value)); + } + state.queue.push_back(value); + (Ok(SendOutcome::Sent), wakers) + } + QueueKind::Bounded(capacity) if state.queue.len() >= capacity => { + return Err(TrySendError::Full(value)); + } + QueueKind::Bounded(_) | QueueKind::Unbounded => { + state.queue.push_back(value); + (Ok(SendOutcome::Sent), state.recv_waiters.take_all()) + } + } + }; + + wake_all(wakers); + result + } + + fn force_send( + &self, + value: T, + behavior: FullBehavior, + ) -> Result, TrySendError> { + let (result, wakers) = { + let mut state = self.state.lock(); + if state.receivers == 0 { + return Err(TrySendError::Disconnected(value)); + } + + match self.kind { + QueueKind::Rendezvous => { + let wakers = state.recv_waiters.take_all(); + if wakers.is_empty() { + return Err(TrySendError::Full(value)); + } + state.queue.push_back(value); + (Ok(SendOutcome::Sent), wakers) + } + QueueKind::Unbounded => { + state.queue.push_back(value); + (Ok(SendOutcome::Sent), state.recv_waiters.take_all()) + } + QueueKind::Bounded(capacity) if state.queue.len() < capacity => { + state.queue.push_back(value); + (Ok(SendOutcome::Sent), state.recv_waiters.take_all()) + } + QueueKind::Bounded(_) => { + let replaced = match behavior { + FullBehavior::DropOldest => state.queue.pop_front(), + FullBehavior::DropNewest => state.queue.pop_back(), + } + .expect("a full bounded channel has a buffered value"); + state.queue.push_back(value); + ( + Ok(SendOutcome::Replaced(replaced)), + state.recv_waiters.take_all(), + ) + } + } + }; + + wake_all(wakers); + result + } + + fn try_recv(&self) -> Result { + let (result, wakers) = { + let mut state = self.state.lock(); + + if let Some(value) = state.queue.pop_front() { + let waker = match self.kind { + QueueKind::Rendezvous => Vec::new(), + QueueKind::Bounded(_) | QueueKind::Unbounded => state.send_waiters.take_all(), + }; + (Ok(value), waker) + } else if matches!(self.kind, QueueKind::Rendezvous) { + if let Some(pending) = state.rendezvous_sends.pop_front() { + (Ok(pending.value), vec![pending.waker]) + } else if state.senders == 0 { + (Err(TryRecvError::Disconnected), Vec::new()) + } else { + (Err(TryRecvError::Empty), Vec::new()) + } + } else if state.senders == 0 { + (Err(TryRecvError::Disconnected), Vec::new()) + } else { + (Err(TryRecvError::Empty), Vec::new()) + } + }; + + wake_all(wakers); + result + } + + fn poll_send( + &self, + value: &mut Option, + waiter: &mut Option, + rendezvous_id: &mut Option, + cx: &mut Context<'_>, + ) -> Poll>> { + let mut wake_receivers = Vec::new(); + let mut retired_wakers = Vec::new(); + let poll = { + let mut state = self.state.lock(); + + if matches!(self.kind, QueueKind::Rendezvous) { + if let Some(id) = *rendezvous_id { + if state + .rendezvous_sends + .iter() + .all(|pending| pending.id != id) + { + return Poll::Ready(Ok(())); + } + } + } + + if state.receivers == 0 { + retired_wakers.extend(state.send_waiters.remove(waiter)); + if let Some(id) = rendezvous_id.take() { + let index = state + .rendezvous_sends + .iter() + .position(|pending| pending.id == id) + .expect("a pending rendezvous send owns its value"); + let pending = state + .rendezvous_sends + .remove(index) + .expect("the rendezvous send index exists"); + let PendingSend { + value: pending_value, + waker, + .. + } = pending; + retired_wakers.push(waker); + *value = Some(pending_value); + } + return Poll::Ready(Err(SendError::new( + value.take().expect("an incomplete send owns its value"), + ))); + } + + match self.kind { + QueueKind::Rendezvous => { + if let Some(id) = *rendezvous_id { + if let Some(pending) = state + .rendezvous_sends + .iter_mut() + .find(|pending| pending.id == id) + { + if !pending.waker.will_wake(cx.waker()) { + retired_wakers.push(std::mem::replace( + &mut pending.waker, + cx.waker().clone(), + )); + } + Poll::Pending + } else { + Poll::Ready(Ok(())) + } + } else { + let id = state.allocate_rendezvous_id(); + state.rendezvous_sends.push_back(PendingSend { + id, + value: value.take().expect("an incomplete send owns its value"), + waker: cx.waker().clone(), + }); + *rendezvous_id = Some(id); + wake_receivers = state.recv_waiters.take_all(); + Poll::Pending + } + } + QueueKind::Bounded(capacity) if state.queue.len() >= capacity => { + retired_wakers.extend(state.send_waiters.register(waiter, cx.waker())); + Poll::Pending + } + QueueKind::Bounded(_) | QueueKind::Unbounded => { + retired_wakers.extend(state.send_waiters.remove(waiter)); + state + .queue + .push_back(value.take().expect("an incomplete send owns its value")); + wake_receivers = state.recv_waiters.take_all(); + Poll::Ready(Ok(())) + } + } + }; + + drop(retired_wakers); + wake_all(wake_receivers); + poll + } + + fn poll_recv( + &self, + waiter: &mut Option, + cx: &mut Context<'_>, + ) -> Poll> { + let mut wake_senders = Vec::new(); + let mut retired_wakers = Vec::new(); + let poll = { + let mut state = self.state.lock(); + + if let Some(value) = state.queue.pop_front() { + retired_wakers.extend(state.recv_waiters.remove(waiter)); + if !matches!(self.kind, QueueKind::Rendezvous) { + wake_senders = state.send_waiters.take_all(); + } + Poll::Ready(Ok(value)) + } else if matches!(self.kind, QueueKind::Rendezvous) { + if let Some(pending) = state.rendezvous_sends.pop_front() { + retired_wakers.extend(state.recv_waiters.remove(waiter)); + wake_senders.push(pending.waker); + Poll::Ready(Ok(pending.value)) + } else if state.senders == 0 { + retired_wakers.extend(state.recv_waiters.remove(waiter)); + Poll::Ready(Err(RecvError::Disconnected)) + } else { + retired_wakers.extend(state.recv_waiters.register(waiter, cx.waker())); + Poll::Pending + } + } else if state.senders == 0 { + retired_wakers.extend(state.recv_waiters.remove(waiter)); + Poll::Ready(Err(RecvError::Disconnected)) + } else { + retired_wakers.extend(state.recv_waiters.register(waiter, cx.waker())); + Poll::Pending + } + }; + + drop(retired_wakers); + wake_all(wake_senders); + poll + } + + fn cancel_send(&self, waiter: &mut Option, rendezvous_id: &mut Option) { + let (retired_waker, pending) = { + let mut state = self.state.lock(); + let retired_waker = state.send_waiters.remove(waiter); + let pending = if let Some(id) = rendezvous_id.take() { + state + .rendezvous_sends + .iter() + .position(|pending| pending.id == id) + .and_then(|index| state.rendezvous_sends.remove(index)) + } else { + None + }; + (retired_waker, pending) + }; + drop(retired_waker); + drop(pending); + } + + fn cancel_recv(&self, waiter: &mut Option) { + let retired_waker = { + let mut state = self.state.lock(); + state.recv_waiters.remove(waiter) + }; + drop(retired_waker); + } +} + +impl State { + fn allocate_rendezvous_id(&mut self) -> u64 { + loop { + let id = self.next_rendezvous_id; + self.next_rendezvous_id = self.next_rendezvous_id.wrapping_add(1); + if self.rendezvous_sends.iter().all(|pending| pending.id != id) { + return id; + } + } + } +} + +struct Send<'a, T> { + core: &'a Core, + value: Option, + waiter: Option, + rendezvous_id: Option, + completed: bool, +} + +impl Unpin for Send<'_, T> {} + +impl Future for Send<'_, T> { + type Output = Result<(), SendError>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = &mut *self; + let poll = this.core.poll_send( + &mut this.value, + &mut this.waiter, + &mut this.rendezvous_id, + cx, + ); + if poll.is_ready() { + this.completed = true; + } + poll + } +} + +impl Drop for Send<'_, T> { + fn drop(&mut self) { + if !self.completed { + self.core + .cancel_send(&mut self.waiter, &mut self.rendezvous_id); + } + } +} + +struct Recv<'a, T> { + core: &'a Core, + waiter: Option, + completed: bool, +} + +impl Future for Recv<'_, T> { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let poll = self.core.poll_recv(&mut self.waiter, cx); + if poll.is_ready() { + self.completed = true; + } + poll + } +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + if !self.completed { + self.core.cancel_recv(&mut self.waiter); + } + } +} diff --git a/asyncband/src/channel/spmc.rs b/asyncband/src/channel/spmc.rs new file mode 100644 index 0000000..e45f248 --- /dev/null +++ b/asyncband/src/channel/spmc.rs @@ -0,0 +1,52 @@ +// 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 queues. +//! +//! The sender is non-cloneable, non-Sync, and requires mutable access. Receivers are cloneable and +//! compete for values, so each value is delivered to exactly one receiver. + +use std::num::NonZeroUsize; + +pub use super::FullBehavior; +pub use super::RecvError; +pub use super::SendError; +pub use super::SendOutcome; +pub use super::TryRecvError; +pub use super::TrySendError; +use super::queue; + +/// The sending endpoint of an SPMC queue. +pub type Sender = queue::Sender; + +/// A receiving endpoint of an SPMC queue. +pub type Receiver = queue::Receiver; + +/// Creates an unbuffered SPMC rendezvous channel. +pub fn rendezvous() -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Rendezvous) +} + +/// Creates a bounded SPMC queue. +pub fn bounded(capacity: NonZeroUsize) -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Bounded(capacity.get())) +} + +/// Creates an unbounded SPMC queue. +pub fn unbounded() -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Unbounded) +} diff --git a/asyncband/src/channel/spsc.rs b/asyncband/src/channel/spsc.rs new file mode 100644 index 0000000..bfa273f --- /dev/null +++ b/asyncband/src/channel/spsc.rs @@ -0,0 +1,58 @@ +// 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 intentionally non-cloneable and non-Sync. Sending and receiving require +//! mutable access, making the single-writer and single-reader guarantees visible to the compiler. +//! +//! ~~~compile_fail +//! fn require_sync() {} +//! +//! require_sync::>(); +//! ~~~ + +use std::num::NonZeroUsize; + +pub use super::FullBehavior; +pub use super::RecvError; +pub use super::SendError; +pub use super::SendOutcome; +pub use super::TryRecvError; +pub use super::TrySendError; +use super::queue; + +/// The sending endpoint of an SPSC queue. +pub type Sender = queue::Sender; + +/// The receiving endpoint of an SPSC queue. +pub type Receiver = queue::Receiver; + +/// Creates an unbuffered SPSC rendezvous channel. +pub fn rendezvous() -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Rendezvous) +} + +/// Creates a bounded SPSC queue. +pub fn bounded(capacity: NonZeroUsize) -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Bounded(capacity.get())) +} + +/// Creates an unbounded SPSC queue. +pub fn unbounded() -> (Sender, Receiver) { + queue::channel(queue::QueueKind::Unbounded) +} diff --git a/asyncband/src/channel/tests.rs b/asyncband/src/channel/tests.rs new file mode 100644 index 0000000..51c166d --- /dev/null +++ b/asyncband/src/channel/tests.rs @@ -0,0 +1,523 @@ +// 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::num::NonZeroUsize; +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::Wake; +use std::task::Waker; + +use super::FullBehavior; +use super::SendOutcome; +use super::TryRecvError; +use super::TrySendError; +use super::broadcast; +use super::disruptor; +use super::mpmc; +use super::mpsc; +use super::oneshot; +use super::spmc; +use super::spsc; +use super::watch; +use crate::test_support::poll_once; + +fn capacity(value: usize) -> NonZeroUsize { + NonZeroUsize::new(value).unwrap() +} + +fn ring_capacity(value: usize) -> disruptor::Capacity { + disruptor::Capacity::new(value).unwrap() +} + +fn poll_with_waker(future: Pin<&mut F>, waker: &Waker) -> Poll { + future.poll(&mut Context::from_waker(waker)) +} + +#[derive(Debug)] +struct WakeProbe { + wakes: AtomicUsize, +} + +impl Wake for WakeProbe { + fn wake(self: Arc) { + self.wakes.fetch_add(1, Ordering::Relaxed); + } +} + +#[test] +fn topology_endpoints_have_expected_positive_auto_traits() { + use std::any::TypeId; + + fn assert_send() {} + fn assert_send_sync() {} + + assert_send::>(); + assert_send::>(); + assert_send_sync::>(); + assert_send::>(); + assert_send::>(); + assert_send_sync::>(); + assert_send_sync::>(); + assert_send_sync::>(); + assert_send::>(); + assert_send_sync::>(); + assert_send_sync::>(); + assert_send_sync::>(); + assert_send_sync::>(); + assert_send_sync::>(); + + assert_ne!( + TypeId::of::>(), + TypeId::of::>() + ); + assert_ne!( + TypeId::of::>(), + TypeId::of::>() + ); + assert_ne!( + TypeId::of::>(), + TypeId::of::>() + ); + assert_ne!( + TypeId::of::>(), + TypeId::of::>() + ); +} + +#[test] +fn oneshot_transfers_or_returns_the_value() { + let (tx, mut rx) = oneshot::channel(); + tx.send(42).unwrap(); + assert_eq!(rx.try_recv(), Ok(42)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); + + let (tx, rx) = oneshot::channel(); + drop(rx); + assert_eq!(tx.send(7).unwrap_err().into_inner(), 7); + + let (tx, rx) = oneshot::channel::(); + drop(tx); + assert!(rx.is_disconnected()); + assert_eq!( + pollster::block_on(rx.recv()), + Err(super::RecvError::Disconnected) + ); +} + +#[test] +fn bounded_queue_is_fifo_and_loss_is_explicit() { + let (mut tx, mut rx) = spsc::bounded(capacity(2)); + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); + assert_eq!( + tx.force_send(3, FullBehavior::DropOldest), + Ok(SendOutcome::Replaced(1)) + ); + assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv(), Ok(3)); + + tx.try_send(4).unwrap(); + tx.try_send(5).unwrap(); + assert_eq!( + tx.force_send(6, FullBehavior::DropNewest), + Ok(SendOutcome::Replaced(5)) + ); + assert_eq!(rx.try_recv(), Ok(4)); + assert_eq!(rx.try_recv(), Ok(6)); +} + +#[test] +fn rendezvous_send_completes_after_receive() { + let (mut tx, mut rx) = spsc::rendezvous(); + assert_eq!(tx.try_send(1), Err(TrySendError::Full(1))); + + let mut send = Box::pin(tx.send(2)); + assert_eq!(poll_once(send.as_mut()), Poll::Pending); + assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(poll_once(send.as_mut()), Poll::Ready(Ok(()))); +} + +#[test] +fn accepted_rendezvous_send_stays_successful_after_receiver_drop() { + let (mut tx, mut rx) = spsc::rendezvous(); + let mut send = tokio_test::task::spawn(tx.send(1)); + assert_eq!(send.poll(), Poll::Pending); + assert_eq!(rx.try_recv(), Ok(1)); + drop(rx); + assert!(send.is_woken()); + assert_eq!(send.poll(), Poll::Ready(Ok(()))); +} + +#[test] +fn dropping_receiver_returns_pending_rendezvous_value() { + let (mut tx, rx) = spsc::rendezvous(); + let mut send = tokio_test::task::spawn(tx.send(1)); + assert_eq!(send.poll(), Poll::Pending); + drop(rx); + assert!(send.is_woken()); + let Poll::Ready(Err(error)) = send.poll() else { + panic!("pending rendezvous send should observe disconnection"); + }; + assert_eq!(error.into_inner(), 1); +} + +#[test] +fn cancelled_rendezvous_send_does_not_leave_a_value() { + let (mut tx, mut rx) = spsc::rendezvous(); + let mut send = Box::pin(tx.send(1)); + assert_eq!(poll_once(send.as_mut()), Poll::Pending); + drop(send); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + + let mut recv = Box::pin(rx.recv()); + assert_eq!(poll_once(recv.as_mut()), Poll::Pending); + assert_eq!(tx.try_send(2), Ok(())); + assert_eq!(poll_once(recv.as_mut()), Poll::Ready(Ok(2))); +} + +#[test] +fn cloned_queue_endpoints_follow_topology() { + let (tx, mut rx) = mpsc::unbounded(); + let tx2 = tx.clone(); + tx.try_send(1).unwrap(); + tx2.try_send(2).unwrap(); + assert_eq!(rx.try_recv(), Ok(1)); + assert_eq!(rx.try_recv(), Ok(2)); + + let (mut tx, rx) = spmc::unbounded(); + let rx2 = rx.clone(); + tx.try_send(3).unwrap(); + tx.try_send(4).unwrap(); + assert_eq!(rx.try_recv(), Ok(3)); + assert_eq!(rx2.try_recv(), Ok(4)); +} + +#[test] +fn freeing_capacity_wakes_all_senders_for_cancellation_transfer() { + let (tx, rx) = mpmc::bounded(capacity(1)); + tx.try_send(0).unwrap(); + let mut first = tokio_test::task::spawn(tx.send(1)); + let mut second = tokio_test::task::spawn(tx.send(2)); + assert_eq!(first.poll(), Poll::Pending); + assert_eq!(second.poll(), Poll::Pending); + + assert_eq!(rx.try_recv(), Ok(0)); + assert!(first.is_woken()); + assert!(second.is_woken()); + drop(first); + assert_eq!(second.poll(), Poll::Ready(Ok(()))); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn sending_wakes_all_receivers_for_cancellation_transfer() { + let (tx, rx) = mpmc::unbounded(); + let rx2 = rx.clone(); + let mut first = tokio_test::task::spawn(rx.recv()); + let mut second = tokio_test::task::spawn(rx2.recv()); + assert_eq!(first.poll(), Poll::Pending); + assert_eq!(second.poll(), Poll::Pending); + + tx.try_send(1).unwrap(); + assert!(first.is_woken()); + assert!(second.is_woken()); + drop(first); + assert_eq!(second.poll(), Poll::Ready(Ok(1))); +} + +#[test] +fn cancelled_receive_deregisters_its_waker() { + let (_tx, rx) = mpmc::unbounded::(); + let probe = Arc::new(WakeProbe { + wakes: AtomicUsize::new(0), + }); + let waker = Waker::from(probe.clone()); + let mut recv = Box::pin(rx.recv()); + assert_eq!(poll_with_waker(recv.as_mut(), &waker), Poll::Pending); + assert_eq!(Arc::strong_count(&probe), 3); + drop(recv); + assert_eq!(Arc::strong_count(&probe), 2); +} + +#[test] +fn mpmc_supports_concurrent_producers_and_consumers() { + let (tx, rx) = mpmc::unbounded(); + let values = Arc::new(std::sync::Mutex::new(Vec::new())); + + std::thread::scope(|scope| { + let mut producers = Vec::new(); + for producer in 0..4 { + let tx = tx.clone(); + producers.push(scope.spawn(move || { + for value in 0..100 { + tx.try_send(producer * 100 + value).unwrap(); + } + })); + } + + let mut consumers = Vec::new(); + for _ in 0..3 { + let rx = rx.clone(); + let values = values.clone(); + consumers.push(scope.spawn(move || { + while let Ok(value) = pollster::block_on(rx.recv()) { + values.lock().unwrap().push(value); + } + })); + } + drop(rx); + + for producer in producers { + producer.join().unwrap(); + } + drop(tx); + for consumer in consumers { + consumer.join().unwrap(); + } + }); + + let mut values = Arc::into_inner(values).unwrap().into_inner().unwrap(); + values.sort_unstable(); + assert_eq!(values, (0..400).collect::>()); +} + +#[test] +fn overflow_broadcast_reports_exact_lag() { + let (tx, mut rx) = broadcast::overflow::channel(capacity(2)); + tx.send(1).unwrap(); + tx.send(2).unwrap(); + tx.send(3).unwrap(); + assert_eq!( + rx.try_recv(), + Err(broadcast::overflow::TryRecvError::Lagged(1)) + ); + assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv(), Ok(3)); +} + +#[test] +fn concurrent_broadcast_senders_publish_one_shared_order() { + let (tx, mut rx1) = broadcast::overflow::channel(capacity(512)); + let mut rx2 = rx1.clone(); + std::thread::scope(|scope| { + let mut senders = Vec::new(); + for producer in 0..4 { + let tx = tx.clone(); + senders.push(scope.spawn(move || { + for value in 0..100 { + tx.try_send(producer * 100 + value).unwrap(); + } + })); + } + for sender in senders { + sender.join().unwrap(); + } + }); + + let first = (0..400) + .map(|_| rx1.try_recv().unwrap()) + .collect::>(); + let second = (0..400) + .map(|_| rx2.try_recv().unwrap()) + .collect::>(); + assert_eq!(first, second); + let mut values = first; + values.sort_unstable(); + assert_eq!(values, (0..400).collect::>()); +} + +#[test] +fn concurrent_overflow_broadcast_keeps_one_committed_suffix() { + let (tx, mut rx1) = broadcast::overflow::channel(capacity(8)); + let mut rx2 = rx1.clone(); + std::thread::scope(|scope| { + for producer in 0..4 { + let tx = tx.clone(); + scope.spawn(move || { + for value in 0..100 { + tx.send(producer * 100 + value).unwrap(); + } + }); + } + }); + + assert_eq!( + rx1.try_recv(), + Err(broadcast::overflow::TryRecvError::Lagged(392)) + ); + assert_eq!( + rx2.try_recv(), + Err(broadcast::overflow::TryRecvError::Lagged(392)) + ); + let first = (0..8).map(|_| rx1.try_recv().unwrap()).collect::>(); + let second = (0..8).map(|_| rx2.try_recv().unwrap()).collect::>(); + assert_eq!(first, second); + let mut values = first; + values.sort_unstable(); + values.dedup(); + assert_eq!(values.len(), 8); +} + +#[test] +fn cancelled_broadcast_receive_deregisters_its_waker() { + let (_tx, mut rx) = broadcast::overflow::channel::(capacity(1)); + let probe = Arc::new(WakeProbe { + wakes: AtomicUsize::new(0), + }); + let waker = Waker::from(probe.clone()); + let mut recv = Box::pin(rx.recv()); + assert_eq!(poll_with_waker(recv.as_mut(), &waker), Poll::Pending); + assert_eq!(Arc::strong_count(&probe), 3); + drop(recv); + assert_eq!(Arc::strong_count(&probe), 2); +} + +#[test] +fn backpressure_broadcast_is_gated_by_slowest_receiver() { + let (tx, mut rx1) = broadcast::backpressure::channel(capacity(2)); + let mut rx2 = rx1.clone(); + tx.try_send(1).unwrap(); + tx.try_send(2).unwrap(); + + let mut send = Box::pin(tx.send(3)); + assert_eq!(poll_once(send.as_mut()), Poll::Pending); + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(poll_once(send.as_mut()), Poll::Pending); + assert_eq!(rx2.try_recv(), Ok(1)); + assert_eq!(poll_once(send.as_mut()), Poll::Ready(Ok(2))); + + assert_eq!(rx1.try_recv(), Ok(2)); + assert_eq!(rx1.try_recv(), Ok(3)); + assert_eq!(rx2.try_recv(), Ok(2)); + assert_eq!(rx2.try_recv(), Ok(3)); +} + +#[test] +fn cancelled_backpressure_broadcast_send_deregisters_its_waker() { + let (tx, mut rx) = broadcast::backpressure::channel(capacity(1)); + tx.try_send(1).unwrap(); + let probe = Arc::new(WakeProbe { + wakes: AtomicUsize::new(0), + }); + let waker = Waker::from(probe.clone()); + let mut send = Box::pin(tx.send(2)); + assert_eq!(poll_with_waker(send.as_mut(), &waker), Poll::Pending); + assert_eq!(Arc::strong_count(&probe), 3); + drop(send); + assert_eq!(Arc::strong_count(&probe), 2); + assert_eq!(rx.try_recv(), Ok(1)); + assert_eq!(tx.try_send(3), Ok(1)); +} + +#[test] +fn unbounded_broadcast_reclaims_after_all_receivers_advance() { + let (tx, mut rx1) = broadcast::unbounded::channel(); + let mut rx2 = rx1.clone(); + tx.send(1).unwrap(); + tx.send(2).unwrap(); + assert_eq!(tx.len(), 2); + + assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(tx.len(), 2); + assert_eq!(rx2.try_recv(), Ok(1)); + assert_eq!(tx.len(), 1); + assert_eq!(rx1.try_recv(), Ok(2)); + assert_eq!(rx2.try_recv(), Ok(2)); + assert_eq!(tx.len(), 0); +} + +#[test] +fn watch_coalesces_versions() { + 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)); + assert_eq!(*pollster::block_on(rx.changed()).unwrap(), 2); + assert_eq!(rx.has_changed(), Ok(false)); +} + +#[test] +fn disruptor_multicasts_and_gates_wraparound() { + let (mut publisher, mut subscriber1) = disruptor::single_producer::channel(ring_capacity(2)); + let mut subscriber2 = subscriber1.clone(); + assert_eq!(publisher.try_publish(10), Ok(0)); + assert_eq!(publisher.try_publish(20), Ok(1)); + assert_eq!(publisher.try_publish(30), Err(TrySendError::Full(30))); + + assert_eq!(subscriber1.try_recv(), Ok((0, 10))); + assert_eq!(publisher.try_publish(30), Err(TrySendError::Full(30))); + assert_eq!(subscriber2.try_recv(), Ok((0, 10))); + assert_eq!(publisher.try_publish(30), Ok(2)); + + assert_eq!(subscriber1.try_recv(), Ok((1, 20))); + assert_eq!(subscriber1.try_recv(), Ok((2, 30))); + assert_eq!(subscriber2.try_recv(), Ok((1, 20))); + assert_eq!(subscriber2.try_recv(), Ok((2, 30))); +} + +#[test] +fn multi_producer_disruptor_assigns_one_sequence_per_publication() { + let (publisher, mut subscriber) = disruptor::multi_producer::channel(ring_capacity(4)); + let publisher2 = publisher.clone(); + assert_eq!(publisher.try_publish("a"), Ok(0)); + assert_eq!(publisher2.try_publish("b"), Ok(1)); + assert_eq!(subscriber.try_recv(), Ok((0, "a"))); + assert_eq!(subscriber.try_recv(), Ok((1, "b"))); +} + +#[test] +fn concurrent_disruptor_publishers_form_a_contiguous_sequence() { + let (publisher, mut subscriber) = disruptor::multi_producer::channel(ring_capacity(512)); + std::thread::scope(|scope| { + let mut publishers = Vec::new(); + for producer in 0..4 { + let publisher = publisher.clone(); + publishers.push(scope.spawn(move || { + for value in 0..100 { + publisher.try_publish(producer * 100 + value).unwrap(); + } + })); + } + for publisher in publishers { + publisher.join().unwrap(); + } + }); + + let mut values = Vec::new(); + for sequence in 0..400 { + let (actual, value) = subscriber.try_recv().unwrap(); + assert_eq!(actual, sequence); + values.push(value); + } + values.sort_unstable(); + assert_eq!(values, (0..400).collect::>()); +} + +#[test] +fn disruptor_validates_power_of_two_capacity() { + assert_eq!(disruptor::Capacity::new(0), None); + assert_eq!(disruptor::Capacity::new(3), None); + assert_eq!(disruptor::Capacity::new(4).unwrap().get(), 4); +} diff --git a/asyncband/src/channel/wait.rs b/asyncband/src/channel/wait.rs new file mode 100644 index 0000000..7dfc621 --- /dev/null +++ b/asyncband/src/channel/wait.rs @@ -0,0 +1,71 @@ +// 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::task::Waker; + +#[derive(Debug, Default)] +pub(super) struct WaitQueue { + next_id: u64, + entries: VecDeque<(u64, Waker)>, +} + +impl WaitQueue { + pub(super) fn register(&mut self, id: &mut Option, waker: &Waker) -> Option { + if let Some(id) = id { + if let Some((_, registered)) = self.entries.iter_mut().find(|(entry, _)| entry == id) { + if !registered.will_wake(waker) { + return Some(std::mem::replace(registered, waker.clone())); + } + return None; + } + } + + let new_id = self.allocate_id(); + self.entries.push_back((new_id, waker.clone())); + *id = Some(new_id); + None + } + + pub(super) fn remove(&mut self, id: &mut Option) -> Option { + let id = id.take()?; + if let Some(index) = self.entries.iter().position(|(entry, _)| *entry == id) { + return self.entries.remove(index).map(|(_, waker)| waker); + } + None + } + + pub(super) fn take_all(&mut self) -> Vec { + self.entries.drain(..).map(|(_, waker)| waker).collect() + } + + fn allocate_id(&mut self) -> u64 { + loop { + let id = self.next_id; + self.next_id = self.next_id.wrapping_add(1); + if self.entries.iter().all(|(entry, _)| *entry != id) { + return id; + } + } + } +} + +pub(super) fn wake_all(wakers: Vec) { + for waker in wakers { + waker.wake(); + } +} diff --git a/asyncband/src/channel/watch.rs b/asyncband/src/channel/watch.rs new file mode 100644 index 0000000..a0c1864 --- /dev/null +++ b/asyncband/src/channel/watch.rs @@ -0,0 +1,284 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A multi-producer, multi-consumer channel that retains only the latest value. +//! +//! Watch is a coalescing state channel rather than a queue. A slow receiver observes the latest +//! version and does not receive an error for skipped intermediate versions. +//! +//! ~~~ +//! use asyncband::channel::watch; +//! +//! let (tx, mut rx) = watch::channel(0); +//! tx.send(1).unwrap(); +//! assert_eq!(*pollster::block_on(rx.changed()).unwrap(), 1); +//! ~~~ + +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 crate::channel::RecvError; +pub use crate::channel::SendError; +use crate::channel::wait::WaitQueue; +use crate::channel::wait::wake_all; +use crate::internal::mutex::Mutex; + +/// 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: WaitQueue::default(), + }), + }); + ( + Sender { + shared: shared.clone(), + }, + Receiver { shared, seen: 0 }, + ) +} + +struct Shared { + state: Mutex>, +} + +struct State { + value: Arc, + version: u64, + senders: usize, + receivers: usize, + waiters: WaitQueue, +} + +/// 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") + .field("disconnected", &self.is_disconnected()) + .finish() + } +} + +impl Sender { + /// Publishes a new current value and returns the number of receivers. + pub fn send(&self, value: T) -> Result> { + let (receivers, wakers, replaced) = { + let mut state = self.shared.state.lock(); + if state.receivers == 0 { + return Err(SendError::new(value)); + } + let next_version = state + .version + .checked_add(1) + .expect("watch version overflow"); + let replaced = std::mem::replace(&mut state.value, Arc::new(value)); + state.version = next_version; + (state.receivers, state.waiters.take_all(), replaced) + }; + wake_all(wakers); + drop(replaced); + Ok(receivers) + } + + /// 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 true if no receivers remain. + pub fn is_disconnected(&self) -> bool { + self.shared.state.lock().receivers == 0 + } +} + +impl Drop for Sender { + fn drop(&mut self) { + let wakers = { + let mut state = self.shared.state.lock(); + state.senders -= 1; + if state.senders == 0 { + state.waiters.take_all() + } else { + Vec::new() + } + }; + wake_all(wakers); + } +} + +/// 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") + .field("seen", &self.seen) + .field("disconnected", &self.is_disconnected()) + .finish() + } +} + +impl Receiver { + /// Returns the current value without marking it 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 new version and returns the latest value. + pub async fn changed(&mut self) -> Result, RecvError> { + Changed { + receiver: self, + waiter: None, + completed: false, + } + .await + } + + /// Returns true if no senders remain. + pub fn is_disconnected(&self) -> bool { + self.shared.state.lock().senders == 0 + } +} + +impl Drop for Receiver { + fn drop(&mut self) { + let mut state = self.shared.state.lock(); + state.receivers -= 1; + } +} + +struct Changed<'a, T> { + receiver: &'a mut Receiver, + waiter: Option, + completed: bool, +} + +impl Future for Changed<'_, T> { + type Output = Result, RecvError>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = &mut *self; + let mut state = this.receiver.shared.state.lock(); + let poll = if state.version != this.receiver.seen { + let retired_waker = state.waiters.remove(&mut this.waiter); + this.receiver.seen = state.version; + let value = state.value.clone(); + drop(state); + drop(retired_waker); + Poll::Ready(Ok(value)) + } else if state.senders == 0 { + let retired_waker = state.waiters.remove(&mut this.waiter); + drop(state); + drop(retired_waker); + Poll::Ready(Err(RecvError::Disconnected)) + } else { + let retired_waker = state.waiters.register(&mut this.waiter, cx.waker()); + drop(state); + drop(retired_waker); + Poll::Pending + }; + if poll.is_ready() { + this.completed = true; + } + poll + } +} + +impl Drop for Changed<'_, T> { + fn drop(&mut self) { + if !self.completed { + let retired_waker = { + let mut state = self.receiver.shared.state.lock(); + state.waiters.remove(&mut self.waiter) + }; + drop(retired_waker); + } + } +} diff --git a/asyncband/src/broadcast/mod.rs b/asyncband/src/coordination.rs similarity index 56% rename from asyncband/src/broadcast/mod.rs rename to asyncband/src/coordination.rs index 96b8d1b..41dbb6e 100644 --- a/asyncband/src/broadcast/mod.rs +++ b/asyncband/src/coordination.rs @@ -15,10 +15,18 @@ // specific language governing permissions and limitations // under the License. -//! A multi-producer multi-consumer broadcast channel. +//! Higher-level coordination protocols built from synchronization primitives. //! -//! This module provides broadcast channels in one of the following policies: +//! Coordination protocols express application-level relationships rather than direct access to a +//! protected value: //! -//! * [`overflow`]: when the channel is full, the oldest messages are overwritten. +//! * [shutdown] coordinates shutdown initiation, participation, and observation. +//! * [singleflight] coalesces concurrent work for the same key. +//! +//! Keeping these protocols outside [crate::sync] leaves the synchronization group focused on +//! mutexes, notifications, permits, and one-time initialization. -pub mod overflow; +#[cfg(feature = "shutdown")] +pub use crate::shutdown; +#[cfg(feature = "singleflight")] +pub use crate::singleflight; diff --git a/asyncband/src/internal/atomic_option_box.rs b/asyncband/src/internal/atomic_option_box.rs deleted file mode 100644 index 26d9227..0000000 --- a/asyncband/src/internal/atomic_option_box.rs +++ /dev/null @@ -1,167 +0,0 @@ -// 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. - -// This is derived from https://github.com/jorendorff/atomicbox/blob/07756444/src/atomic_option_box.rs. - -use std::marker::PhantomData; -use std::ptr; -use std::sync::atomic::AtomicPtr; -use std::sync::atomic::Ordering; - -/// Internal owning atomic pointer used to store an optional receiver waker. -pub(crate) struct AtomicOptionBox { - /// Pointer to a `T` value in the heap, representing `Some(t)`; - /// or a null pointer for `None`. - ptr: AtomicPtr, - - /// This effectively makes `AtomicOptionBox` non-`Send` and non-`Sync` - /// if `T` is non-`Send`. - phantom: PhantomData>, -} - -/// Mark `AtomicOptionBox` as safe to share across threads. -/// -/// This is safe because shared access to an `AtomicOptionBox` does not -/// provide shared access to any `T` value. However, it does provide the -/// ability to get a `Box` from another thread, so `T: Send` is required. -unsafe impl Sync for AtomicOptionBox where T: Send {} - -fn into_ptr(value: Option>) -> *mut T { - match value { - Some(box_value) => Box::into_raw(box_value), - None => ptr::null_mut(), - } -} - -// SAFETY: The caller must ensure that `ptr` was obtained from `Box::into_raw` or is null. -unsafe fn from_ptr(ptr: *mut T) -> Option> { - if ptr.is_null() { - None - } else { - Some(unsafe { Box::from_raw(ptr) }) - } -} - -impl AtomicOptionBox { - /// Creates a new `AtomicOptionBox` with no value. - pub(crate) const fn none() -> Self { - Self { - ptr: AtomicPtr::new(ptr::null_mut()), - phantom: PhantomData, - } - } - - fn swap(&self, other: Option>) -> Option> { - let order = match other { - Some(_) => Ordering::AcqRel, - None => Ordering::Acquire, - }; - let new_ptr = into_ptr(other); - let old_ptr = self.ptr.swap(new_ptr, order); - unsafe { from_ptr(old_ptr) } - } - - /// Stores `other` and drops the previous value. - pub(crate) fn store(&self, other: Option>) { - drop(self.swap(other)); - } - - /// Replaces the value with `None` and returns the previous value. - pub(crate) fn take(&self) -> Option> { - self.swap(None) - } -} - -impl Drop for AtomicOptionBox { - /// Dropping an `AtomicOptionBox` drops the final `Box` value (if any) stored in it. - fn drop(&mut self) { - let ptr = *self.ptr.get_mut(); - unsafe { drop(from_ptr(ptr)) } - } -} - -#[cfg(test)] -mod tests { - use core::sync::atomic::Ordering; - use std::sync::Arc; - use std::sync::atomic::AtomicUsize; - - use super::*; - - #[test] - fn atomic_option_box_swap_works() { - let b = AtomicOptionBox::none(); - let bis = Box::new("bis"); - assert_eq!(b.swap(Some(bis)), None); - assert_eq!(b.swap(None), Some(Box::new("bis"))); - } - - #[test] - fn atomic_option_box_store_works() { - let b = AtomicOptionBox::none(); - let bis = Box::new("bis"); - b.store(Some(bis)); - assert_eq!(b.take(), Some(Box::new("bis"))); - assert_eq!(b.take(), None); - } - - #[test] - fn atomic_option_box_pointer_identity() { - let box1 = Box::new(1); - let p1 = &*box1 as *const i32; - let atom = AtomicOptionBox::none(); - atom.store(Some(box1)); - - let box2 = Box::new(2); - let p2 = &*box2 as *const i32; - assert_ne!(p2, p1); - - let box3 = atom.swap(Some(box2)).unwrap(); // box1 out, box2 in - let p3 = &*box3 as *const i32; - assert_eq!(p3, p1); // box3 is box1 - - let box4 = atom.swap(None).unwrap(); // box2 out, None in - let p4 = &*box4 as *const i32; - assert_eq!(p4, p2); // box4 is box2 - } - - #[test] - fn stored_values_are_dropped() { - struct K(Arc, usize); - - impl Drop for K { - fn drop(&mut self) { - self.0.fetch_add(self.1, Ordering::Relaxed); - } - } - - let n = Arc::new(AtomicUsize::new(0)); - { - let ab = AtomicOptionBox::none(); - ab.store(Some(Box::new(K(n.clone(), 5)))); - assert_eq!(n.load(Ordering::Relaxed), 0); - let first = ab.swap(None); - assert_eq!(n.load(Ordering::Relaxed), 0); - drop(first); - assert_eq!(n.load(Ordering::Relaxed), 5); - let second = ab.swap(Some(Box::new(K(n.clone(), 13)))); - assert!(second.is_none()); - assert_eq!(n.load(Ordering::Relaxed), 5); - } - assert_eq!(n.load(Ordering::Relaxed), 5 + 13); - } -} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 995db3c..2264be7 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -15,14 +15,9 @@ // specific language governing permissions and limitations // under the License. -#[cfg(feature = "mpsc")] -pub(crate) mod atomic_option_box; - #[cfg(any( feature = "barrier", - feature = "broadcast", feature = "latch", - feature = "mpsc", feature = "mutex", feature = "rwlock", feature = "semaphore", @@ -46,46 +41,35 @@ pub(crate) mod once_table; #[cfg(any( feature = "barrier", feature = "broadcast", + feature = "disruptor", feature = "latch", - feature = "mpsc", feature = "mutex", + feature = "oneshot", + feature = "queue", feature = "rwlock", feature = "semaphore", + feature = "watch", feature = "waitgroup", ))] pub(crate) mod mutex; -#[cfg(feature = "broadcast")] -pub(crate) mod rwlock; - -#[cfg(any( - feature = "mpsc", - feature = "mutex", - feature = "rwlock", - feature = "semaphore", -))] -// `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. -// Each single-primitive build intentionally leaves the other groups unused. +#[cfg(any(feature = "mutex", feature = "rwlock", feature = "semaphore"))] +// 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; -#[cfg(any( - feature = "mpsc", - feature = "mutex", - feature = "rwlock", - feature = "semaphore", -))] +#[cfg(any(feature = "mutex", feature = "rwlock", feature = "semaphore"))] pub(crate) mod waitlist; #[cfg(any( feature = "barrier", - feature = "broadcast", feature = "latch", feature = "once", feature = "waitgroup", ))] -// `barrier` constructs a wait set with `with_capacity`, while broadcast and countdown-based -// primitives use `new`. One constructor is therefore unused in every single-primitive build. +// `barrier` constructs a wait set with `with_capacity`, while countdown-based primitives use +// `new`. One constructor is therefore unused in every single-primitive build. #[allow(dead_code)] pub(crate) mod waitset; diff --git a/asyncband/src/internal/rwlock.rs b/asyncband/src/internal/rwlock.rs deleted file mode 100644 index 6d24dea..0000000 --- a/asyncband/src/internal/rwlock.rs +++ /dev/null @@ -1,56 +0,0 @@ -// 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::PoisonError; - -pub struct RwLock(std::sync::RwLock); - -impl RwLock { - pub const fn new(t: T) -> Self { - Self(std::sync::RwLock::new(t)) - } -} - -impl RwLock { - pub fn read(&self) -> std::sync::RwLockReadGuard<'_, T> { - self.0.read().unwrap_or_else(PoisonError::into_inner) - } - - pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, T> { - self.0.write().unwrap_or_else(PoisonError::into_inner) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use crate::internal::rwlock::RwLock; - - #[test] - fn test_poison_rwlock() { - let rwlock = Arc::new(RwLock::new(42)); - let r = rwlock.clone(); - let handle = std::thread::spawn(move || { - let _guard = r.write(); - panic!("poison"); - }); - let _ = handle.join(); - assert_eq!(*rwlock.read(), 42); - assert_eq!(*rwlock.write(), 42); - } -} diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index 5394941..d4d0d21 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -120,16 +120,6 @@ impl Semaphore { fut.await } - /// Returns a future that is resolved when acquired `n` permits from the semaphore. - pub fn poll_acquire(&self, n: usize) -> Acquire<'_> { - Acquire { - permits: n, - index: None, - semaphore: self, - done: false, - } - } - /// Adds `n` permits to the semaphore. pub fn release(&self, n: usize) { if n != 0 { @@ -137,41 +127,6 @@ impl Semaphore { } } - /// Adds `n` permits to the semaphore if there is any waiter. - pub fn release_if_nonempty(&self, n: usize) { - let waiters = self.waiters.lock(); - if !waiters.is_empty() { - self.insert_permits_with_lock(n, waiters); - } - } - - /// Adds as many permits until there is no waiter. - pub fn notify_all(&self) { - let mut waiters = self.waiters.lock(); - let mut wakers = Vec::new(); - loop { - match waiters.unlink_first_waiter(|node| { - node.permits = 0; - true - }) { - None => break, - Some((id, waiter)) => { - let remove_now = waiter.waker.is_none(); - if let Some(waker) = waiter.waker.take() { - wakers.push(waker); - } - if remove_now { - waiters.remove_unlinked_waiter(id); - } - } - } - } - drop(waiters); - for w in wakers.drain(..) { - w.wake(); - } - } - fn insert_permits_with_lock( &self, mut rem: usize, diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 3d810e9..3e1df62 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -56,7 +56,7 @@ //! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | //! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::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`], [`broadcast::overflow`] | `oneshot`, `mpsc`, `broadcast` | +//! | Send values | [`channel`] | `channel` or a channel policy feature | //! | Control workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | //! | Wait from synchronous code | [`blocking::FutureExt`] | `blocking` | //! @@ -93,20 +93,24 @@ mod internal; pub mod barrier; #[cfg(feature = "blocking")] pub mod blocking; -#[cfg(feature = "broadcast")] -pub mod broadcast; +#[cfg(any( + feature = "broadcast", + feature = "disruptor", + feature = "oneshot", + feature = "queue", + feature = "watch", +))] +pub mod channel; #[cfg(feature = "condvar")] pub mod condvar; +#[cfg(any(feature = "shutdown", feature = "singleflight"))] +pub mod coordination; #[cfg(feature = "latch")] pub mod latch; -#[cfg(feature = "mpsc")] -pub mod 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; #[cfg(feature = "rwlock")] pub mod rwlock; #[cfg(feature = "semaphore")] @@ -115,8 +119,34 @@ pub mod semaphore; pub mod shutdown; #[cfg(feature = "singleflight")] pub mod singleflight; +#[cfg(any( + feature = "barrier", + feature = "condvar", + feature = "latch", + feature = "mutex", + feature = "once", + feature = "once-cell", + feature = "once-map", + feature = "rwlock", + feature = "semaphore", + feature = "waitgroup", +))] +pub mod sync; #[cfg(feature = "waitgroup")] pub mod waitgroup; -#[cfg(all(test, any(feature = "once-map", feature = "singleflight")))] +#[cfg(all( + test, + any( + feature = "once-map", + feature = "singleflight", + all( + feature = "broadcast", + feature = "disruptor", + feature = "oneshot", + feature = "queue", + feature = "watch", + ), + ) +))] mod test_support; diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs deleted file mode 100644 index 0e90bbf..0000000 --- a/asyncband/src/mpsc/bounded.rs +++ /dev/null @@ -1,361 +0,0 @@ -// 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 bounded multi-producer, single-consumer queue for sending values between asynchronous -//! tasks with backpressure control. - -use std::fmt; -use std::future::Future; -use std::future::poll_fn; -use std::pin::pin; -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use std::task::Waker; - -use crate::internal::atomic_option_box::AtomicOptionBox; -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. -/// -/// A `send` on this channel will wait if the buffer of the channel is full until a -/// `recv` is called on the receiver, which will consume the message and -/// free up space in the buffer. -#[track_caller] -pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { - assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); - let state = Arc::new(BoundedState { - senders: AtomicUsize::new(1), - tx_permits: Semaphore::new(0), - rx_task: AtomicOptionBox::none(), - }); - let (sender, receiver) = std::sync::mpsc::sync_channel(buffer); - let sender = BoundedSender { - state: state.clone(), - sender: Some(sender), - }; - let receiver = BoundedReceiver { - state: state.clone(), - receiver: Some(receiver), - }; - (sender, receiver) -} - -struct BoundedState { - senders: AtomicUsize, - tx_permits: Semaphore, - rx_task: AtomicOptionBox, -} - -/// Send values to the associated [`BoundedReceiver`]. -/// -/// Instances are created by the [`bounded`] function. -pub struct BoundedSender { - state: Arc, - sender: Option>, -} - -impl Clone for BoundedSender { - fn clone(&self) -> Self { - self.state.senders.fetch_add(1, Ordering::Release); - BoundedSender { - state: self.state.clone(), - sender: self.sender.clone(), - } - } -} - -impl fmt::Debug for BoundedSender { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("BoundedSender").finish_non_exhaustive() - } -} - -impl Drop for BoundedSender { - fn drop(&mut self) { - // drop the sender; this closes the channel if it is the last sender - drop(self.sender.take()); - - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { - 1 => { - // If this is the last sender, we need to wake up the receiver so it can - // observe the disconnected state. - if let Some(waker) = self.state.rx_task.take() { - waker.wake(); - } - } - _ => { - // there are still other senders left, do nothing - } - } - } -} - -impl BoundedSender { - /// Attempts to send a message to the associated receiver. - /// - /// This method will wait if the buffer of the channel is full until a `recv` is called on the - /// receiver, which will consume the message and free up space in the buffer. - /// - /// If the receiver has been dropped, this function returns an error. The error includes - /// the value passed to `send`. - pub async fn send(&self, value: T) -> Result<(), SendError> { - let value = match self.try_send(value) { - Ok(()) => return Ok(()), - Err(TrySendError::Disconnected(value)) => return Err(SendError::new(value)), - Err(TrySendError::Full(value)) => value, - }; - - struct SendState<'a, T> { - sender: &'a BoundedSender, - value: Option, - acquire: Acquire<'a>, - } - - impl SendState<'_, T> { - fn poll_send(&mut self, cx: &mut Context<'_>) -> Poll>> { - let mut value = match self.value.take() { - Some(value) => value, - None => return Poll::Ready(Ok(())), - }; - - loop { - let poll = pin!(&mut self.acquire).poll(cx); - - value = match self.sender.try_send(value) { - Ok(()) => return Poll::Ready(Ok(())), - Err(TrySendError::Disconnected(value)) => { - return Poll::Ready(Err(SendError::new(value))); - } - Err(TrySendError::Full(value)) => value, - }; - - if poll.is_ready() { - self.acquire = self.sender.state.tx_permits.poll_acquire(1); - } else { - self.value = Some(value); - return Poll::Pending; - } - } - } - } - - let acquire = self.state.tx_permits.poll_acquire(1); - let mut send = SendState { - sender: self, - value: Some(value), - acquire, - }; - poll_fn(|cx| send.poll_send(cx)).await - } - - /// Attempts to send a message to the associated receiver without waiting. - /// - /// This method returns the [`Full`] error if the buffer of the channel is full. - /// - /// This method returns the [`Disconnected`] error if the channel is currently empty, and there - /// are no outstanding [receivers]. - /// - /// [`Full`]: TrySendError::Full - /// [`Disconnected`]: TrySendError::Disconnected - /// [receivers]: BoundedReceiver - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc::TrySendError; - /// use asyncband::mpsc::bounded; - /// let (tx, mut rx) = bounded::(1); - /// - /// tx.try_send(1).unwrap(); - /// assert_eq!(tx.try_send(2), Err(TrySendError::Full(2))); - /// - /// drop(rx); - /// assert_eq!(tx.try_send(3), Err(TrySendError::Disconnected(3))); - /// # } - /// ``` - pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - // SAFETY: The sender is guaranteed to be non-null before dropped. - let sender = self.sender.as_ref().unwrap(); - match sender.try_send(value) { - Ok(()) => { - if let Some(waker) = self.state.rx_task.take() { - waker.wake(); - } - - Ok(()) - } - Err(std::sync::mpsc::TrySendError::Full(value)) => Err(TrySendError::Full(value)), - Err(std::sync::mpsc::TrySendError::Disconnected(value)) => { - Err(TrySendError::Disconnected(value)) - } - } - } -} - -/// Receives values from the associated [`BoundedSender`]. -/// -/// Instances are created by the [`bounded`] function. -pub struct BoundedReceiver { - state: Arc, - receiver: Option>, -} - -/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`. -/// That is, `BoundedReceiver` can only be accessed by one thread at a time. -unsafe impl Sync for BoundedReceiver {} - -impl fmt::Debug for BoundedReceiver { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("BoundedReceiver").finish_non_exhaustive() - } -} - -impl Drop for BoundedReceiver { - fn drop(&mut self) { - drop(self.receiver.take()); - self.state.tx_permits.notify_all(); - } -} - -impl BoundedReceiver { - /// Tries to receive the next value for this receiver and frees up a space in the buffer if - /// successful. - /// - /// This method returns the [`Empty`] error if the channel is currently - /// empty, but there are still outstanding [senders]. - /// - /// This method returns the [`Disconnected`] error if the channel is - /// currently empty, and there are no outstanding [senders]. - /// - /// [`Empty`]: TryRecvError::Empty - /// [`Disconnected`]: TryRecvError::Disconnected - /// [senders]: BoundedSender - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc; - /// use asyncband::mpsc::TryRecvError; - /// let (tx, mut rx) = mpsc::bounded(2); - /// - /// tx.send("hello").await.unwrap(); - /// - /// assert_eq!(Ok("hello"), rx.try_recv()); - /// assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); - /// - /// tx.send("hello").await.unwrap(); - /// drop(tx); - /// - /// assert_eq!(Ok("hello"), rx.try_recv()); - /// assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); - /// # } - /// ``` - pub fn try_recv(&mut self) -> Result { - // SAFETY: The receiver is guaranteed to be non-null before dropped. - let receiver = self.receiver.as_ref().unwrap(); - match receiver.try_recv() { - Ok(v) => { - self.state.tx_permits.release_if_nonempty(1); - Ok(v) - } - Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected), - Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty), - } - } - - /// Receives the next value for this receiver and frees up a space in the buffer if successful. - /// - /// This method returns `Err(RecvError::Disconnected)` if the channel has been closed and there - /// are no remaining messages in the channel's buffer. This indicates that no further values - /// can ever be received from this `Receiver`. The channel is closed when all senders have been - /// dropped. - /// - /// If there are no messages in the channel's buffer, but the channel has not yet been closed, - /// this method will sleep until a message is sent or the channel is closed. - /// - /// # Cancel safety - /// - /// This method is cancel safe. If `recv` is used as the event in a `select` statement - /// and some other branch completes first, it is guaranteed that no messages were received - /// on this channel. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc; - /// let (tx, mut rx) = mpsc::bounded(1); - /// - /// tokio::spawn(async move { - /// tx.send("hello").await.unwrap(); - /// }); - /// - /// assert_eq!(Ok("hello"), rx.recv().await); - /// assert_eq!(Err(mpsc::RecvError::Disconnected), rx.recv().await); - /// # } - /// ``` - /// - /// Values are buffered if the channel has enough capacity: - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc; - /// let (tx, mut rx) = mpsc::bounded(2); - /// - /// tx.send("hello").await.unwrap(); - /// tx.send("world").await.unwrap(); - /// - /// assert_eq!(Ok("hello"), rx.recv().await); - /// assert_eq!(Ok("world"), rx.recv().await); - /// # } - /// ``` - pub async fn recv(&mut self) -> Result { - poll_fn(|cx| self.poll_recv(cx)).await - } - - fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => { - let waker = Some(Box::new(cx.waker().clone())); - self.state.rx_task.store(waker); - - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => Poll::Pending, - } - } - } - } -} diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs deleted file mode 100644 index 87c7c8f..0000000 --- a/asyncband/src/mpsc/mod.rs +++ /dev/null @@ -1,33 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! A multi-producer, 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; diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs deleted file mode 100644 index 06be655..0000000 --- a/asyncband/src/mpsc/unbounded.rs +++ /dev/null @@ -1,260 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! An unbounded multi-producer, single-consumer queue for sending values between asynchronous -//! tasks. - -use std::fmt; -use std::future::poll_fn; -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::internal::atomic_option_box::AtomicOptionBox; -use crate::mpsc::RecvError; -use crate::mpsc::SendError; -use crate::mpsc::TryRecvError; - -/// Creates an unbounded mpsc channel for communicating between asynchronous -/// tasks without backpressure. -/// -/// A `send` on this channel will always succeed as long as the receiver is alive. -/// If the receiver falls behind, messages will be arbitrarily buffered. -/// -/// Note that the amount of available system memory is an implicit bound to -/// the channel. Using an `unbounded` channel has the ability of causing the -/// process to run out of memory. In this case, the process will be aborted. -pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { - let state = Arc::new(UnboundedState { - senders: AtomicUsize::new(1), - rx_task: AtomicOptionBox::none(), - }); - let (sender, receiver) = std::sync::mpsc::channel(); - let sender = UnboundedSender { - state: state.clone(), - sender: Some(sender), - }; - let receiver = UnboundedReceiver { - state: state.clone(), - receiver, - }; - (sender, receiver) -} - -struct UnboundedState { - senders: AtomicUsize, - rx_task: AtomicOptionBox, -} - -/// Send values to the associated [`UnboundedReceiver`]. -/// -/// Instances are created by the [`unbounded`] function. -pub struct UnboundedSender { - state: Arc, - sender: Option>, -} - -impl Clone for UnboundedSender { - fn clone(&self) -> Self { - self.state.senders.fetch_add(1, Ordering::Release); - UnboundedSender { - state: self.state.clone(), - sender: self.sender.clone(), - } - } -} - -impl fmt::Debug for UnboundedSender { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("UnboundedSender").finish_non_exhaustive() - } -} - -impl Drop for UnboundedSender { - fn drop(&mut self) { - // drop the sender; this closes the channel if it is the last sender - drop(self.sender.take()); - - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { - 1 => { - // If this is the last sender, we need to wake up the receiver so it can - // observe the disconnected state. - if let Some(waker) = self.state.rx_task.take() { - waker.wake(); - } - } - _ => { - // there are still other senders left, do nothing - } - } - } -} - -impl UnboundedSender { - /// Attempts to send a message without blocking. - /// - /// This method is not marked async because sending a message to an unbounded channel - /// never requires any form of waiting. Because of this, the `send` method can be - /// used in both synchronous and asynchronous code without problems. - /// - /// If the receiver has been dropped, this function returns an error. The error includes - /// the value passed to `send`. - pub fn send(&self, value: T) -> Result<(), SendError> { - // SAFETY: The sender is guaranteed to be non-null before dropped. - let sender = self.sender.as_ref().unwrap(); - sender.send(value).map_err(|err| SendError::new(err.0))?; - - if let Some(waker) = self.state.rx_task.take() { - waker.wake(); - } - - Ok(()) - } -} - -/// Receive values from the associated [`UnboundedSender`]. -/// -/// Instances are created by the [`unbounded`] function. -pub struct UnboundedReceiver { - state: Arc, - receiver: std::sync::mpsc::Receiver, -} - -/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`. -/// That is, `UnboundedReceiver` can only be accessed by one thread at a time. -unsafe impl Sync for UnboundedReceiver {} - -impl fmt::Debug for UnboundedReceiver { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("UnboundedReceiver").finish_non_exhaustive() - } -} - -impl UnboundedReceiver { - /// Tries to receive the next value for this receiver. - /// - /// This method returns the [`Empty`] error if the channel is currently - /// empty, but there are still outstanding [senders]. - /// - /// This method returns the [`Disconnected`] error if the channel is - /// currently empty, and there are no outstanding [senders]. - /// - /// [`Empty`]: TryRecvError::Empty - /// [`Disconnected`]: TryRecvError::Disconnected - /// [senders]: UnboundedSender - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc; - /// use asyncband::mpsc::TryRecvError; - /// let (tx, mut rx) = mpsc::unbounded(); - /// - /// tx.send("hello").unwrap(); - /// - /// assert_eq!(Ok("hello"), rx.try_recv()); - /// assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); - /// - /// tx.send("hello").unwrap(); - /// drop(tx); - /// - /// assert_eq!(Ok("hello"), rx.try_recv()); - /// assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); - /// # } - /// ``` - pub fn try_recv(&mut self) -> Result { - match self.receiver.try_recv() { - Ok(v) => Ok(v), - Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected), - Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty), - } - } - - /// Receives the next value for this receiver. - /// - /// This method returns `Err(RecvError::Disconnected)` if the channel has been closed and there - /// are no remaining messages in the channel's buffer. This indicates that no further values - /// can ever be received from this `Receiver`. The channel is closed when all senders have been - /// dropped. - /// - /// If there are no messages in the channel's buffer, but the channel has not yet been closed, - /// this method will sleep until a message is sent or the channel is closed. - /// - /// # Cancel safety - /// - /// This method is cancel safe. If `recv` is used as the event in a `select` statement - /// and some other branch completes first, it is guaranteed that no messages were received - /// on this channel. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc; - /// let (tx, mut rx) = mpsc::unbounded(); - /// - /// tokio::spawn(async move { - /// tx.send("hello").unwrap(); - /// }); - /// - /// assert_eq!(Ok("hello"), rx.recv().await); - /// assert_eq!(Err(mpsc::RecvError::Disconnected), rx.recv().await); - /// # } - /// ``` - /// - /// Values are buffered: - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc; - /// let (tx, mut rx) = mpsc::unbounded(); - /// - /// tx.send("hello").unwrap(); - /// tx.send("world").unwrap(); - /// - /// assert_eq!(Ok("hello"), rx.recv().await); - /// assert_eq!(Ok("world"), rx.recv().await); - /// # } - /// ``` - pub async fn recv(&mut self) -> Result { - poll_fn(|cx| self.poll_recv(cx)).await - } - - fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => { - let waker = Some(Box::new(cx.waker().clone())); - self.state.rx_task.store(waker); - - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => Poll::Pending, - } - } - } - } -} diff --git a/asyncband/src/oneshot/mod.rs b/asyncband/src/oneshot/mod.rs deleted file mode 100644 index 4ee05fa..0000000 --- a/asyncband/src/oneshot/mod.rs +++ /dev/null @@ -1,374 +0,0 @@ -// 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. - -// This implementation is derived from the `oneshot` crate [1], with significant simplifications -// because this crate supports only asynchronous receive operations. -// -// [1] https://github.com/faern/oneshot/blob/83fd0864/src/lib.rs - -//! A one-shot channel is used for sending a single message between asynchronous tasks. The -//! [`channel`] function is used to create a [`Sender`] and [`Receiver`] pair that form the channel. -//! -//! The sender is used by the producer to send the value. The receiver is used by the consumer -//! to receive the value. -//! -//! The sender and receiver can be used by separate tasks. -//! -//! Since [`Sender::send`] is not async, it can be used anywhere. This includes sending between -//! two runtimes, and using it from non-async code. -//! -//! # Examples -//! -//! ``` -//! # #[tokio::main] -//! # async fn main() { -//! use asyncband::oneshot; -//! -//! let (tx, rx) = oneshot::channel(); -//! -//! tokio::spawn(async move { -//! if let Err(_) = tx.send(3) { -//! println!("the receiver dropped"); -//! } -//! }); -//! -//! match rx.await { -//! Ok(v) => println!("got = {:?}", v), -//! Err(_) => println!("the sender dropped"), -//! } -//! # } -//! ``` -//! -//! If the sender is dropped without sending, the receiver will fail with [`RecvError`]: -//! -//! ``` -//! # #[tokio::main] -//! # async fn main() { -//! use asyncband::oneshot; -//! -//! let (tx, rx) = oneshot::channel::(); -//! -//! tokio::spawn(async move { drop(tx) }); -//! -//! match rx.await { -//! Ok(_) => panic!("This doesn't happen"), -//! Err(_) => println!("the sender dropped"), -//! } -//! # } -//! ``` -//! -//! If the receiver is dropped before receiving, the sender will fail with [`SendError`]: -//! -//! ``` -//! use asyncband::oneshot; -//! -//! let (tx, rx) = oneshot::channel::(); -//! -//! drop(rx); -//! -//! match tx.send(42) { -//! Ok(_) => panic!("This doesn't happen"), -//! Err(_) => println!("the receiver dropped"), -//! } -//! ``` - -mod receiver; -mod sender; - -use std::cell::UnsafeCell; -use std::mem::MaybeUninit; -use std::ptr::NonNull; -use std::sync::atomic::AtomicU8; -use std::sync::atomic::Ordering; -use std::sync::atomic::fence; -use std::task::Poll; -use std::task::Waker; - -pub use self::receiver::Receiver; -pub use self::receiver::Recv; -pub use self::receiver::RecvError; -pub use self::receiver::TryRecvError; -pub use self::sender::SendError; -pub use self::sender::Sender; - -#[cfg(test)] -mod tests; - -/// Creates a new oneshot channel and returns the [`Sender`] and [`Receiver`]. -pub fn channel() -> (Sender, Receiver) { - let channel_ptr = NonNull::from(Box::leak(Box::new(Channel::new()))); - (Sender::new(channel_ptr), Receiver::new(channel_ptr)) -} - -const EMPTY: u8 = 0b011; -const RECEIVING: u8 = 0b000; -const AWAKING: u8 = 0b001; -const MESSAGE: u8 = 0b100; -const DISCONNECTED: u8 = 0b010; - -/// Shared storage and state machine for a oneshot channel. -/// -/// The atomic state publishes access to the `message` and `waker` slots. Initializing a slot does -/// not by itself transfer ownership; the corresponding release operation does, and an acquire -/// operation is required before the other endpoint accesses it. -/// -/// * `EMPTY`: no message or waker is published. The sender may be initializing the message, and the -/// receiver may temporarily own a reclaimed waker, but neither slot is available to the other -/// endpoint. -/// * `RECEIVING`: the receiver has published an initialized waker. It may reclaim the waker by -/// returning to `EMPTY`, or the sender may move to `AWAKING` and take ownership of it. The sender -/// retains ownership of any message that it has not yet published. -/// * `AWAKING`: the sender exclusively owns the published waker and any unpublished message while -/// it publishes either a message or a disconnect. The receiver must not access either slot; -/// cancellation may only transfer allocation cleanup to the sender by moving to `DISCONNECTED`. -/// * `MESSAGE`: the sender has published an initialized message and no longer accesses the channel. -/// The receiver owns the message and the allocation. -/// * `DISCONNECTED`: no message can subsequently be received. The transition that reaches or -/// observes this state determines which endpoint owns any remaining message, waker, and -/// allocation cleanup. -/// -/// The state is no longer meaningful after an operation obtains exclusive ownership of the whole -/// allocation, such as when `send` creates a `SendError`. -struct Channel { - state: AtomicU8, - message: UnsafeCell>, - waker: UnsafeCell>, -} - -impl Channel { - fn new() -> Self { - Self { - state: AtomicU8::new(EMPTY), - message: UnsafeCell::new(MaybeUninit::uninit()), - waker: UnsafeCell::new(MaybeUninit::uninit()), - } - } - - /// Returns a shared reference to the initialized message. - /// - /// # Safety - /// - /// The message must be initialized and remain initialized and immutably accessible for the - /// returned reference's lifetime. - #[inline(always)] - unsafe fn message(&self) -> &T { - unsafe { - let slot = &*self.message.get(); - slot.assume_init_ref() - } - } - - /// Moves the initialized message out of its slot. - /// - /// # Safety - /// - /// The caller must exclusively own an initialized message and must not subsequently read or - /// drop the slot as initialized unless it is initialized again. - #[inline(always)] - unsafe fn take_message(&self) -> T { - unsafe { - let slot = &*self.message.get(); - slot.assume_init_read() - } - } - - /// Initializes the message slot. - /// - /// # Safety - /// - /// The message slot must be uninitialized and exclusively accessible to the caller. The caller - /// must publish the initialized message before another thread accesses it. - #[inline(always)] - unsafe fn write_message(&self, message: T) { - unsafe { - let slot = &mut *self.message.get(); - slot.as_mut_ptr().write(message); - } - } - - /// Drops the initialized message in place. - /// - /// # Safety - /// - /// The caller must exclusively own an initialized message. The slot must not subsequently be - /// read or dropped as initialized unless it is initialized again. - #[inline(always)] - unsafe fn drop_message(&self) { - unsafe { - let slot = &mut *self.message.get(); - slot.assume_init_drop(); - } - } - - /// Stores and publishes a receiver waker, resolving a raced terminal state immediately. - /// - /// # Safety - /// - /// * The `waker` field must not contain an initialized waker when calling this method. - /// * The `state` must not be in the `RECEIVING` or `AWAKING` state when calling this method. - /// * No other receiver operation may access the waker slot concurrently. - unsafe fn register_waker(&self, waker: Waker) -> Poll> { - // SAFETY: The sender cannot access the waker until the state becomes RECEIVING. - unsafe { - let slot = &mut *self.waker.get(); - slot.as_mut_ptr().write(waker); - } - - // ORDERING: On success, Release publishes the initialized waker. Failure only observes the - // current state; the MESSAGE branch performs its own conditional Acquire below, while the - // DISCONNECTED branch does not access sender-owned data. - match self - .state - .compare_exchange(EMPTY, RECEIVING, Ordering::Release, Ordering::Relaxed) - { - // The waker is registered for the sender to take and wake. - Ok(_) => Poll::Pending, - // The sender sent the message while we prepared to await. - // We take the message and mark the channel disconnected. - Err(MESSAGE) => { - // SAFETY: We wrote a waker above. The sender cannot have observed the RECEIVING - // state, so it has not accessed the waker. We must drop it. - unsafe { self.drop_waker() }; - - // ORDERING: The sender has completed, so this receiver-only terminal update does - // not publish data to another thread. - self.state.store(DISCONNECTED, Ordering::Relaxed); - - // ORDERING: The failed CAS read MESSAGE from the sender's Release publication. This - // conditional Acquire makes the initialized message visible before it is taken. - fence(Ordering::Acquire); - - // SAFETY: The MESSAGE state tells us there is a correctly initialized message, - // and the fence above synchronizes with that write. - Poll::Ready(Ok(unsafe { self.take_message() })) - } - // The sender was dropped before sending anything while we prepared to await. - Err(DISCONNECTED) => { - // SAFETY: We wrote a waker above. The sender cannot have observed the RECEIVING - // state, so it has not accessed the waker. We must drop it. - unsafe { self.drop_waker() }; - Poll::Ready(Err(RecvError::Disconnected)) - } - Err(state) => unreachable!("unexpected channel state: {}", state), - } - } - - /// Drops the initialized waker in place. - /// - /// # Safety - /// - /// The caller must exclusively own an initialized waker. The slot must not subsequently be - /// read or dropped as initialized unless it is initialized again. - #[inline(always)] - unsafe fn drop_waker(&self) { - unsafe { - let slot = &mut *self.waker.get(); - slot.assume_init_drop(); - } - } - - /// Moves the initialized waker out of its slot. - /// - /// # Safety - /// - /// The caller must exclusively own an initialized waker. The slot must not subsequently be - /// read or dropped as initialized unless it is initialized again. - #[inline(always)] - unsafe fn take_waker(&self) -> Waker { - unsafe { - let slot = &*self.waker.get(); - slot.assume_init_read() - } - } - - /// Finishes the sender-owned `AWAKING` state by taking the receiver waker and publishing the - /// final channel state. - /// - /// Returns the waker and whether the receiver still owns allocation cleanup. If the receiver - /// cancelled from `AWAKING`, this returns `false` and transfers cleanup to the caller. - /// - /// # Safety - /// - /// * `final_state` must be `MESSAGE` or `DISCONNECTED`. - /// * The caller must have just observed `RECEIVING` with an atomic read-modify-write that - /// changed the state to `AWAKING`. This gives the caller exclusive ownership of the - /// initialized waker and provides the atomic read paired with the acquire fence in this - /// method. - /// * When publishing `MESSAGE`, the caller must own an initialized message that precedes the - /// release operation in this method. - #[inline(always)] - unsafe fn finish_sender_awakening(&self, final_state: u8) -> (Waker, bool) { - debug_assert!(matches!(final_state, MESSAGE | DISCONNECTED)); - - // ORDERING: The caller's Release RMW read RECEIVING with a Relaxed load. Acquire - // synchronizes that read with the receiver's Release publication before taking the waker. - fence(Ordering::Acquire); - - // SAFETY: The caller's RECEIVING-to-AWAKING transition transferred exclusive ownership of - // the initialized waker to the sender. - let waker = unsafe { self.take_waker() }; - - // ORDERING: Release publishes the message or disconnect when this replaces AWAKING. The - // RMW's load half is Relaxed; if it reads a receiver-written DISCONNECTED, the conditional - // Acquire below completes the reverse allocation-ownership handoff. - let previous_state = self.state.swap(final_state, Ordering::Release); - if matches!(previous_state, AWAKING) { - (waker, true) - } else { - // The receiver has been dropped. - debug_assert_eq!(previous_state, DISCONNECTED); - - // ORDERING: The swap read DISCONNECTED from the receiver's Release cancellation. - // Acquire makes every preceding receiver access happen before sender-side reclamation. - fence(Ordering::Acquire); - - (waker, false) - } - } -} - -/// Deallocates a channel whose slots no longer contain values that need to be dropped. -/// -/// # Safety -/// -/// `channel_ptr` must retain the provenance of the live allocation created by `channel`. The caller -/// must exclusively own allocation cleanup, neither slot may contain a value that still needs to be -/// dropped, and no access through any pointer or reference may follow this call. -unsafe fn deallocate_empty_channel(channel_ptr: NonNull>) { - // SAFETY: The caller transfers exclusive ownership of the original allocation to this function, - // so this is the only Box reconstructed from the pointer. - unsafe { drop(Box::from_raw(channel_ptr.as_ptr())) }; -} - -/// Drops the initialized message and then deallocates the channel. -/// -/// # Safety -/// -/// `channel_ptr` must retain the provenance of the live allocation created by `channel`. The caller -/// must exclusively own the initialized message and allocation cleanup, the waker slot must not -/// contain a value that still needs to be dropped, and no access through any pointer or reference -/// may follow this call. -unsafe fn drop_message_and_deallocate_channel(channel_ptr: NonNull>) { - // SAFETY: The caller transfers exclusive allocation ownership to this function, so this is the - // only Box reconstructed from the pointer. - let channel = unsafe { Box::from_raw(channel_ptr.as_ptr()) }; - - // SAFETY: The caller guarantees that the message is initialized and exclusively owned. The Box - // deallocates the channel on normal return and during unwinding if `T::drop` panics. Since the - // message is stored in MaybeUninit, dropping the Box will not drop it a second time. - unsafe { channel.drop_message() }; -} diff --git a/asyncband/src/oneshot/receiver.rs b/asyncband/src/oneshot/receiver.rs deleted file mode 100644 index 29f913b..0000000 --- a/asyncband/src/oneshot/receiver.rs +++ /dev/null @@ -1,422 +0,0 @@ -// 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::mem::ManuallyDrop; -use std::pin::Pin; -use std::ptr::NonNull; -use std::sync::atomic::Ordering; -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; -#[cfg(doc)] -use crate::oneshot::Sender; -use crate::oneshot::deallocate_empty_channel; -use crate::oneshot::drop_message_and_deallocate_channel; - -/// Receives a value from the associated [`Sender`]. -pub struct Receiver { - channel_ptr: NonNull>, -} - -impl fmt::Debug for Receiver { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Receiver").finish_non_exhaustive() - } -} - -unsafe impl Send for Receiver {} - -// Receiver must not be `Sync`: receive operations taking `&self` assume that no other receive -// operation runs concurrently. - -impl Unpin for Receiver {} - -impl IntoFuture for Receiver { - type Output = Result; - - type IntoFuture = Recv; - - fn into_future(self) -> Self::IntoFuture { - // `Recv` takes over receiver-side cleanup, so `Receiver::drop` must not run afterward. - let receiver = ManuallyDrop::new(self); - let channel_ptr = receiver.channel_ptr; - Recv { channel_ptr } - } -} - -impl Receiver { - /// Returns `true` if the channel is disconnected. - /// - /// This occurs when the associated [`Sender`] is dropped without sending a message, or after - /// the message is received. - /// - /// If `true` is returned, all future receive operations are guaranteed to return an error. - pub fn is_disconnected(&self) -> bool { - // SAFETY: The existence of `self` guarantees that the receiver is still alive. If the - // sender was dropped, it observed the live receiver and left allocation cleanup to it, so - // `channel_ptr` remains valid. - let channel = unsafe { self.channel_ptr.as_ref() }; - - // ORDERING: Relaxed is sufficient to enforce the method's contract. - // - // Once true has been observed, it will remain true. However, if false is observed, - // the sender might have just disconnected but this thread has not observed it yet. - matches!(channel.state.load(Ordering::Relaxed), DISCONNECTED) - } - - /// Returns true if there is a message in the channel, ready to be received. - /// - /// If `true` is returned, the next call to receive the message is guaranteed to return - /// the message immediately. - pub fn has_message(&self) -> bool { - // SAFETY: The existence of `self` guarantees that the receiver is still alive. If the - // sender was dropped, it observed the live receiver and left allocation cleanup to it, so - // `channel_ptr` remains valid. - let channel = unsafe { self.channel_ptr.as_ref() }; - - // ORDERING: This method only observes the atomic state. MESSAGE is terminal for the sender, - // and receiver operations cannot run concurrently, so atomic coherence preserves this - // observation for the next receive. Accessing the message synchronizes separately. - matches!(channel.state.load(Ordering::Relaxed), MESSAGE) - } - - /// Checks if there is a message in the channel without blocking. Returns: - /// - /// * `Ok(message)` if there was a message in the channel. - /// * `Err(TryRecvError::Empty)` if the [`Sender`] is alive, but has not yet sent a message. - /// * `Err(TryRecvError::Disconnected)` if the [`Sender`] was dropped before sending anything or - /// if the message has already been extracted by a previous `try_recv` call. - /// - /// If a message is returned, the channel is disconnected and any subsequent receive operation - /// using this receiver will return an error: [`TryRecvError::Disconnected`] for `try_recv`, - /// or [`RecvError::Disconnected`] for [`recv`](Receiver::into_future). - pub fn try_recv(&self) -> Result { - // SAFETY: The channel will not be freed while this method is still running. - let channel = unsafe { self.channel_ptr.as_ref() }; - - // ORDERING: Relaxed is fine since the only branch that needs synchronization is MESSAGE, - // and that branch has its own synchronization. - match channel.state.load(Ordering::Relaxed) { - MESSAGE => { - // It is okay to break up the load and store since once we are in the MESSAGE state, - // the sender no longer modifies the state - // - // ORDERING: The sender has completed, so this receiver-only terminal update does - // not publish data to another thread. - channel.state.store(DISCONNECTED, Ordering::Relaxed); - - // ORDERING: The preceding Relaxed load read MESSAGE from the sender's Release - // publication. This conditional Acquire makes the message visible before it is - // taken. - fence(Ordering::Acquire); - - // SAFETY: we are in the MESSAGE state so the message is present and synchronized. - Ok(unsafe { channel.take_message() }) - } - EMPTY => Err(TryRecvError::Empty), - DISCONNECTED => Err(TryRecvError::Disconnected), - state => unreachable!("unexpected channel state: {}", state), - } - } - - pub(super) fn new(channel_ptr: NonNull>) -> Self { - Self { channel_ptr } - } -} - -impl Drop for Receiver { - fn drop(&mut self) { - // SAFETY: The live receiver guarantees that a dropped sender left allocation cleanup to - // this side. - let channel = unsafe { self.channel_ptr.as_ref() }; - - // Set the channel state to disconnected and read what state the channel was in. - // - // ORDERING: This is a bidirectional ownership handoff. Release publishes the receiver's - // last access when the sender must reclaim the allocation; Acquire receives a - // sender-published message or disconnect before receiver-side cleanup. - match channel.state.swap(DISCONNECTED, Ordering::AcqRel) { - // The sender has not sent anything, nor is it dropped. The sender is responsible for - // deallocating the channel. - EMPTY => {} - // The sender already sent something. We must drop it, and free the channel. - MESSAGE => { - // SAFETY: The MESSAGE state plus acquire ordering guarantees the sender has - // written a message and that it has a happens-before relationship with this drop. - // In addition, the acquire ordering above synchronizes with the sender's final - // write of the state, so we can safely deallocate the channel. - unsafe { drop_message_and_deallocate_channel(self.channel_ptr) }; - } - // The sender was already dropped. We are responsible for freeing the channel. - DISCONNECTED => { - // SAFETY: If the sender published DISCONNECTED, the swap's Acquire half makes its - // preceding accesses happen before reclamation. If this receiver previously wrote - // DISCONNECTED after taking the message, no cross-thread synchronization is needed. - unsafe { deallocate_empty_channel(self.channel_ptr) }; - } - // NOTE: the receiver, unless transformed into a future, will never see the RECEIVING or - // AWAKING states, so we can ignore them here. - state => unreachable!("unexpected channel state: {}", state), - } - } -} - -/// A future that completes when the message is sent from the associated [`Sender`], or the -/// [`Sender`] is dropped before sending a message. -pub struct Recv { - channel_ptr: NonNull>, -} - -unsafe impl Send for Recv {} - -impl fmt::Debug for Recv { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Recv").finish_non_exhaustive() - } -} - -impl Future for Recv { - type Output = Result; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - // SAFETY: The existence of `self` guarantees that the receiver is still alive. If the - // sender was dropped, it observed the live receiver and left allocation cleanup to it, so - // `channel_ptr` remains valid. - let channel = unsafe { self.channel_ptr.as_ref() }; - - // ORDERING: This load only selects a state-machine branch. Branches that access a published - // message or waker perform their own Acquire operation. - match channel.state.load(Ordering::Relaxed) { - // The sender is alive but has not sent anything yet. - EMPTY => { - let waker = cx.waker().clone(); - // SAFETY: EMPTY means no waker is initialized or owned by the sender. - unsafe { channel.register_waker(waker) } - } - // The sender sent the message. - MESSAGE => { - // ORDERING: The sender has completed, so this receiver-only terminal update does - // not publish data to another thread. - channel.state.store(DISCONNECTED, Ordering::Relaxed); - - // ORDERING: The preceding Relaxed load read MESSAGE from the sender's Release - // publication. This conditional Acquire makes the message visible before it is - // taken. - fence(Ordering::Acquire); - - // SAFETY: we are in the MESSAGE state and have synchronized with the sender. - Poll::Ready(Ok(unsafe { channel.take_message() })) - } - // We were polled again while waiting for the sender. Replace the waker with the new - // one. - RECEIVING => { - // ORDERING: On success, Acquire synchronizes with the Release that published the - // stored waker before this poll reclaims it. Failure does not access that waker. - match channel.state.compare_exchange( - RECEIVING, - EMPTY, - Ordering::Acquire, - Ordering::Relaxed, - ) { - // The state is EMPTY again. - Ok(_) => { - let waker = cx.waker().clone(); - - // SAFETY: The successful exchange makes the state EMPTY, so the sender - // cannot take the stored waker. The acquire ordering synchronizes with the - // waker write. - unsafe { channel.drop_waker() }; - - // SAFETY: The old waker was dropped while the state was EMPTY, so no waker - // remains initialized or owned by the sender. - unsafe { channel.register_waker(waker) } - } - // The sender sent the message while we prepared to replace the waker. - // We take the message and mark the channel disconnected. - // The sender has already taken the waker. - Err(MESSAGE) => { - // ORDERING: The sender has completed, so this receiver-only terminal update - // does not publish data to another thread. - channel.state.store(DISCONNECTED, Ordering::Relaxed); - - // ORDERING: The failed CAS read MESSAGE from the sender's Release - // publication. This conditional Acquire makes the message visible before it - // is taken. - fence(Ordering::Acquire); - - // SAFETY: The state tells us the sender has initialized the message, and - // the fence above synchronizes with that write. - Poll::Ready(Ok(unsafe { channel.take_message() })) - } - // The sender started awakening us while we prepared to replace the waker. - Err(AWAKING) => { - cx.waker().wake_by_ref(); - Poll::Pending - } - // The sender was dropped before sending anything while we prepared to park. - // The sender has taken the waker already. - Err(DISCONNECTED) => Poll::Ready(Err(RecvError::Disconnected)), - Err(state) => unreachable!("unexpected channel state: {}", state), - } - } - // The sender is publishing the final state and owns the stored waker. Schedule this - // poll's potentially different waker and return without waiting for the - // sender to make progress. - AWAKING => { - cx.waker().wake_by_ref(); - Poll::Pending - } - // The sender was dropped before sending anything. - DISCONNECTED => Poll::Ready(Err(RecvError::Disconnected)), - state => unreachable!("unexpected channel state: {}", state), - } - } -} - -impl Drop for Recv { - fn drop(&mut self) { - // SAFETY: The live receiver guarantees that a dropped sender left allocation cleanup to - // this side. - let channel = unsafe { self.channel_ptr.as_ref() }; - - loop { - // ORDERING: Acquire synchronizes terminal sender publications before cleanup and the - // receiver's earlier waker publication before reclaiming it. EMPTY and AWAKING need - // only the atomic state observation, but using one load keeps the cleanup - // paths fence-free. - match channel.state.load(Ordering::Acquire) { - // The sender has not sent anything, nor is it dropped. Mark the receiver as - // dropped; the sender is responsible for deallocating the channel. - EMPTY => { - if channel - .state - .compare_exchange(EMPTY, DISCONNECTED, Ordering::Release, Ordering::Relaxed) - .is_ok() - { - break; - } - } - // The sender already sent something. We must drop it, and free the channel. - MESSAGE => { - // SAFETY: The MESSAGE state plus acquire ordering guarantees the sender has - // written a message and that it has a happens-before relationship with this - // drop. The same load orders sender accesses before allocation reclamation. - unsafe { drop_message_and_deallocate_channel(self.channel_ptr) }; - break; - } - // This receiver was previously polled, but was not polled to completion. Move away - // from RECEIVING before dropping the waker so the sender cannot take the same - // waker. - // - // A successful exchange creates a short EMPTY window before the next iteration can - // mark DISCONNECTED. This branch owns and drops the stored waker first. A sender - // that observes EMPTY does not touch the waker. It either stores MESSAGE and - // leaves the message and allocation to this loop, or stores DISCONNECTED and - // leaves the allocation to this loop. If this loop marks DISCONNECTED first, the - // sender observes DISCONNECTED and owns any send error cleanup. - RECEIVING => { - if channel - .state - .compare_exchange(RECEIVING, EMPTY, Ordering::Relaxed, Ordering::Relaxed) - .is_ok() - { - // SAFETY: The successful exchange makes the state EMPTY, so the sender - // cannot take the stored waker. The preceding Acquire load synchronized - // with its publication. No transition can recreate RECEIVING, so the - // successful Relaxed CAS still claims that same waker. - unsafe { channel.drop_waker() }; - } - } - // The sender owns the waker. Transfer allocation cleanup to it instead of waiting - // for it to publish the terminal state. - AWAKING => { - if channel - .state - .compare_exchange( - AWAKING, - DISCONNECTED, - Ordering::Release, - Ordering::Relaxed, - ) - .is_ok() - { - break; - } - } - // The sender was already dropped, or this future was previously polled to - // completion. We are responsible for freeing the channel. - DISCONNECTED => { - // SAFETY: If the sender published DISCONNECTED, the Acquire load makes its - // preceding accesses happen before reclamation. If this future wrote - // DISCONNECTED after taking the message, no cross-thread synchronization is - // needed. - unsafe { deallocate_empty_channel(self.channel_ptr) }; - break; - } - state => unreachable!("unexpected channel state: {}", state), - } - } - } -} - -/// Error returned by [`Receiver::try_recv`]. -#[derive(Debug, Clone, Eq, PartialEq)] -pub enum TryRecvError { - /// This channel is currently empty, but the sender has not yet disconnected, so data may yet - /// become available. - Empty, - /// The sender has become disconnected, and there will never be any more data received on it. - Disconnected, -} - -impl fmt::Display for TryRecvError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(match self { - TryRecvError::Empty => "receiving on an empty channel", - TryRecvError::Disconnected => "receiving on a closed channel", - }) - } -} - -impl std::error::Error for TryRecvError {} - -/// An error returned when awaiting the message via [`Receiver`]. -/// -/// This error indicates that the corresponding [`Sender`] was dropped before sending any message. -/// Note that if a message was already received (e.g., via [`Receiver::try_recv`]), subsequent -/// `try_recv` calls will return [`TryRecvError::Disconnected`] instead. -#[derive(Debug, Clone, Eq, PartialEq)] -pub enum RecvError { - /// The sender has become disconnected, and there will never be any more data received on it. - Disconnected, -} - -impl fmt::Display for RecvError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("receiving on a closed channel") - } -} - -impl std::error::Error for RecvError {} diff --git a/asyncband/src/oneshot/sender.rs b/asyncband/src/oneshot/sender.rs deleted file mode 100644 index adf45ee..0000000 --- a/asyncband/src/oneshot/sender.rs +++ /dev/null @@ -1,274 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::any::type_name; -use std::fmt; -use std::mem::ManuallyDrop; -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; -#[cfg(doc)] -use crate::oneshot::Receiver; -use crate::oneshot::deallocate_empty_channel; -use crate::oneshot::drop_message_and_deallocate_channel; - -/// Sends a value to the associated [`Receiver`]. -pub struct Sender { - channel_ptr: NonNull>, -} - -impl fmt::Debug for Sender { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Sender").finish_non_exhaustive() - } -} - -unsafe impl Send for Sender {} -unsafe impl Sync for Sender {} - -impl Sender { - /// Attempts to send a value on this channel, returning an error containing the message if it - /// could not be sent. - pub fn send(self, message: T) -> Result<(), SendError> { - // `send` takes over endpoint cleanup, so `Sender::drop` must not run afterward. - let sender = ManuallyDrop::new(self); - let channel_ptr = sender.channel_ptr; - - // SAFETY: The channel exists on the heap for the entire duration of this method, and we - // only ever acquire shared references to it. Note that if the receiver disconnects it - // does not free the channel. - let channel = unsafe { channel_ptr.as_ref() }; - - // Write the message into the channel on the heap. - // - // SAFETY: The receiver only ever accesses this memory location if we are in the MESSAGE - // state, and since we are responsible for setting that state, we can guarantee that we have - // exclusive access to this memory location to perform this write. - unsafe { channel.write_message(message) }; - - // Publish the message directly, or begin awakening a receiving task: - // - // * EMPTY + 1 = MESSAGE - // * RECEIVING + 1 = AWAKING - // * DISCONNECTED + 1 = EMPTY (invalid), however this state is never observed - // - // ORDERING: Release publishes the message directly for EMPTY-to-MESSAGE and orders it - // before the waiting path's final publication. The RMW's load half is Relaxed, so - // branches that consume receiver-published resources use an Acquire fence. - match channel.state.fetch_add(1, Ordering::Release) { - // The receiver is alive and has not started waiting. Send done. - EMPTY => Ok(()), - // The receiver is waiting. Wake it up so it can return the message. - RECEIVING => { - // SAFETY: fetch_add observed RECEIVING and changed it to AWAKING, transferring - // exclusive ownership of the published waker to this sender. The message was - // initialized before the RMW and is ready to publish. - let (waker, receiver_owns_allocation) = - unsafe { channel.finish_sender_awakening(MESSAGE) }; - if receiver_owns_allocation { - waker.wake(); - } else { - // The send remains successful because this sender owned the waker before the - // receiver cancelled. - // - // SAFETY: Receiver cancellation transferred message and allocation cleanup to - // this sender. The original pointer provenance may therefore be reclaimed as a - // Box; the message is initialized and the waker has been moved out. - unsafe { drop_message_and_deallocate_channel(channel_ptr) }; - } - Ok(()) - } - // The receiver was already dropped. The error is responsible for freeing the channel. - // - // SAFETY: The acquire fence in this arm synchronizes with the receiver's write of the - // DISCONNECTED state. Since the receiver will no longer access `channel_ptr`, the error - // takes exclusive ownership of the channel's resources. - // Moreover, since we just placed the message in the channel, the channel contains a - // valid message. - DISCONNECTED => { - // ORDERING: The RMW read DISCONNECTED from the receiver's Release endpoint drop. - // This Acquire completes the ownership handoff before SendError accesses the - // allocation. - fence(Ordering::Acquire); - Err(SendError { channel_ptr }) - } - state => unreachable!("unexpected channel state: {}", state), - } - } - - /// Returns `true` if the channel is disconnected. - /// - /// This occurs when the associated receiving endpoint is dropped. - /// - /// If `true` is returned, a future call to [`send`](Sender::send) is guaranteed to return an - /// error. - pub fn is_disconnected(&self) -> bool { - // SAFETY: The channel exists on the heap for the entire duration of this method, and we - // only ever acquire shared references to it. Note that if the receiver disconnects it - // does not free the channel. - let channel = unsafe { self.channel_ptr.as_ref() }; - - // ORDERING: Relaxed is sufficient for the method's contract: if this returns true, a - // future call to send is guaranteed to return an error. - // - // Once true has been observed, it will remain true. However, if false is observed, - // the receiver might have just disconnected but this thread has not observed it yet. - matches!(channel.state.load(Ordering::Relaxed), DISCONNECTED) - } - - pub(super) fn new(channel_ptr: NonNull>) -> Self { - Self { channel_ptr } - } -} - -impl Drop for Sender { - fn drop(&mut self) { - // SAFETY: The receiver only ever frees the channel if we are in the MESSAGE or - // DISCONNECTED states. - // - // * If we are in the MESSAGE state, then `send` suppressed `Sender::drop`, so we should not - // be in this function call. - // * If we are in the DISCONNECTED state, then the receiver either received the message, - // making this statement unreachable, or was dropped and observed that our side was still - // alive, and thus didn't free the channel. - let channel = unsafe { self.channel_ptr.as_ref() }; - - // Disconnect directly, or begin awakening a receiving task: - // - // * EMPTY ^ 001 = DISCONNECTED - // * RECEIVING ^ 001 = AWAKING - // * DISCONNECTED ^ 001 = EMPTY (invalid), but this state is never observed - // - // ORDERING: Release publishes a direct disconnect and orders it before the waiting path's - // final publication. The RMW's load half is Relaxed, so branches that consume - // receiver-published resources use an Acquire fence. - match channel.state.fetch_xor(0b001, Ordering::Release) { - // The receiver is not waiting, nor is it dropped. The receiver is responsible for - // deallocating the channel. - EMPTY => {} - // The receiver is waiting. Wake it up so it can detect that the channel disconnected. - RECEIVING => { - // SAFETY: fetch_xor observed RECEIVING and changed it to AWAKING, transferring - // exclusive ownership of the published waker to this sender. Sender::drop has no - // message to publish. - let (waker, receiver_owns_allocation) = - unsafe { channel.finish_sender_awakening(DISCONNECTED) }; - if receiver_owns_allocation { - waker.wake(); - } else { - // SAFETY: Receiver cancellation transferred allocation cleanup to this sender. - // Both slots are uninitialized, and the original pointer provenance may be - // reclaimed as a Box. - unsafe { deallocate_empty_channel(self.channel_ptr) }; - } - } - // The receiver was already dropped. We are responsible for freeing the channel. - DISCONNECTED => { - // ORDERING: The RMW read DISCONNECTED from the receiver's Release endpoint drop. - // Acquire makes all preceding receiver accesses happen before deallocation. - fence(Ordering::Acquire); - // SAFETY: when the receiver switches the state to DISCONNECTED they have received - // the message or will no longer be trying to receive the message, and have - // observed that the sender is still alive, meaning that we are responsible for - // freeing the channel allocation. The acquire ordering above synchronizes with - // the receiver's final write of the state. - unsafe { deallocate_empty_channel(self.channel_ptr) }; - } - state => unreachable!("unexpected channel state: {}", state), - } - } -} - -/// An error returned when trying to send on a closed channel. Returned from -/// [`Sender::send`] if the corresponding [`Receiver`] has already been dropped. -/// -/// The message that could not be sent can be retrieved again with [`SendError::into_inner`]. -pub struct SendError { - channel_ptr: NonNull>, -} - -// SAFETY: SendError exclusively owns the channel allocation and its initialized message. If the -// message is Send, the error may transfer that ownership to another thread. -unsafe impl Send for SendError {} - -// SAFETY: SendError retains exclusive ownership while shared references only expose `&T`, which -// may cross threads when T is Sync. -unsafe impl Sync for SendError {} - -impl SendError { - /// Get a reference to the message that failed to be sent. - pub fn as_inner(&self) -> &T { - // SAFETY: SendError exclusively owns the allocation and its initialized message. - unsafe { self.channel_ptr.as_ref().message() } - } - - /// Consumes the error and returns the message that failed to be sent. - pub fn into_inner(self) -> T { - // The returned message and this method take over cleanup, so `SendError::drop` must not - // run. - let error = ManuallyDrop::new(self); - let channel_ptr = error.channel_ptr; - - // SAFETY: SendError exclusively owns the allocation. - let channel: &Channel = unsafe { channel_ptr.as_ref() }; - - // SAFETY: The send path initialized the message before constructing SendError. - let message = unsafe { channel.take_message() }; - - // SAFETY: SendError exclusively owns the allocation, so its original pointer provenance may - // be reclaimed as a Box. The message has been moved out and no waker remains initialized. - unsafe { deallocate_empty_channel(channel_ptr) }; - - message - } -} - -impl Drop for SendError { - fn drop(&mut self) { - // SAFETY: SendError exclusively owns the initialized message and allocation, so its - // original pointer provenance may be reclaimed as a Box. No waker remains in the - // channel. - unsafe { drop_message_and_deallocate_channel(self.channel_ptr) }; - } -} - -impl fmt::Display for SendError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("sending on a closed channel") - } -} - -impl fmt::Debug for SendError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "SendError<{}>(..)", type_name::()) - } -} - -impl std::error::Error for SendError {} - -#[cfg(test)] -impl Sender { - pub(super) fn channel_ptr(&self) -> NonNull> { - self.channel_ptr - } -} diff --git a/asyncband/src/oneshot/tests.rs b/asyncband/src/oneshot/tests.rs deleted file mode 100644 index 84d8b28..0000000 --- a/asyncband/src/oneshot/tests.rs +++ /dev/null @@ -1,242 +0,0 @@ -// 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::future::IntoFuture; -use std::mem; -use std::pin::Pin; -use std::sync::atomic::Ordering; -use std::sync::mpsc; -use std::task::Context; -use std::task::Poll; -use std::time::Duration; - -use self::support::DropProbe; -use self::support::WakerProbe; -use self::support::spawn_named; -use crate::oneshot; - -// These tests stay next to the implementation because they inspect private state. - -#[test] -fn poll_returns_while_sender_owns_waker() { - let (sender, receiver) = oneshot::channel(); - let channel_ptr = sender.channel_ptr(); - mem::forget(sender); - let mut receiver = receiver.into_future(); - - let (stored_waker, stored_probe) = WakerProbe::new(); - let mut stored_context = Context::from_waker(&stored_waker); - assert_eq!( - Pin::new(&mut receiver).poll(&mut stored_context), - Poll::Pending - ); - - let channel = unsafe { channel_ptr.as_ref() }; - unsafe { channel.write_message(1234u128) }; - // Pause the synthetic sender in AWAKING immediately after it takes the stored waker. - assert_eq!( - channel.state.fetch_add(1, Ordering::Release), - super::RECEIVING - ); - - let (started_tx, started_rx) = mpsc::sync_channel(0); - let (result_tx, result_rx) = mpsc::sync_channel(1); - let poll_thread = spawn_named("receiver", move || { - let (current_waker, current_probe) = WakerProbe::new(); - let mut current_context = Context::from_waker(¤t_waker); - started_tx.send(()).unwrap(); - let result = Pin::new(&mut receiver).poll(&mut current_context); - result_tx.send((receiver, current_probe, result)).unwrap(); - }); - - started_rx.recv().unwrap(); - let result = result_rx.recv_timeout(Duration::from_secs(5)); - let returned_before_publish = result.is_ok(); - - // SAFETY: fetch_add observed RECEIVING and changed it to AWAKING after the message and waker - // were initialized. - let (sender_waker, receiver_owns_allocation) = - unsafe { channel.finish_sender_awakening(super::MESSAGE) }; - assert!(receiver_owns_allocation); - sender_waker.wake(); - assert_eq!(stored_probe.wake_count(), 1); - - let (mut receiver, current_probe, result) = match result { - Ok(result) => result, - Err(mpsc::RecvTimeoutError::Timeout) => result_rx - .recv_timeout(Duration::from_secs(5)) - .expect("receiver remained blocked after the sender published MESSAGE"), - Err(mpsc::RecvTimeoutError::Disconnected) => panic!("receiver thread exited unexpectedly"), - }; - poll_thread.join().unwrap(); - assert!( - returned_before_publish, - "poll blocked while the sender owned the waker" - ); - assert_eq!(result, Poll::Pending); - assert_eq!(current_probe.wake_count(), 1); - - let (current_waker, _) = WakerProbe::new(); - let mut current_context = Context::from_waker(¤t_waker); - assert_eq!( - Pin::new(&mut receiver).poll(&mut current_context), - Poll::Ready(Ok(1234)) - ); -} - -#[test] -fn drop_transfers_cleanup_while_sender_owns_waker() { - let (sender, receiver) = oneshot::channel(); - let channel_ptr = sender.channel_ptr(); - mem::forget(sender); - let mut receiver = receiver.into_future(); - let (message, message_drop_count) = DropProbe::new(1234u128); - - let (stored_waker, stored_probe) = WakerProbe::new(); - let mut context = Context::from_waker(&stored_waker); - assert!(matches!( - Pin::new(&mut receiver).poll(&mut context), - Poll::Pending - )); - - let channel = unsafe { channel_ptr.as_ref() }; - unsafe { channel.write_message(message) }; - // Pause the synthetic sender in AWAKING immediately after it takes the stored waker. - assert_eq!( - channel.state.fetch_add(1, Ordering::Release), - super::RECEIVING - ); - - let (started_tx, started_rx) = mpsc::sync_channel(0); - let (done_tx, done_rx) = mpsc::sync_channel(1); - let drop_thread = spawn_named("receiver", move || { - started_tx.send(()).unwrap(); - drop(receiver); - done_tx.send(()).unwrap(); - }); - - started_rx.recv().unwrap(); - let result = done_rx.recv_timeout(Duration::from_secs(5)); - let returned_before_publish = result.is_ok(); - - // SAFETY: fetch_add observed RECEIVING and changed it to AWAKING after the message and waker - // were initialized. - let (sender_waker, receiver_owns_allocation) = - unsafe { channel.finish_sender_awakening(super::MESSAGE) }; - if receiver_owns_allocation { - sender_waker.wake(); - } else { - // SAFETY: Receiver cancellation transferred allocation cleanup to the synthetic sender, so - // the original pointer provenance may be reclaimed as a Box. The message is initialized and - // the waker was moved out above. - unsafe { super::drop_message_and_deallocate_channel(channel_ptr) }; - drop(sender_waker); - } - - match result { - Ok(()) => {} - Err(mpsc::RecvTimeoutError::Timeout) => done_rx - .recv_timeout(Duration::from_secs(5)) - .expect("receiver remained blocked after the sender published MESSAGE"), - Err(mpsc::RecvTimeoutError::Disconnected) => panic!("receiver thread exited unexpectedly"), - } - drop_thread.join().unwrap(); - assert!( - returned_before_publish, - "drop blocked while the sender owned the waker" - ); - assert!(!receiver_owns_allocation); - assert_eq!(message_drop_count.load(Ordering::Relaxed), 1); - assert_eq!(WakerProbe::live_waker_count(&stored_probe), 1); -} - -mod support { - use std::sync::Arc; - use std::sync::atomic::AtomicU32; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - use std::task::Wake; - use std::task::Waker; - - pub(super) struct DropProbe { - drop_count: Arc, - _value: T, - } - - impl DropProbe { - pub(super) fn new(value: T) -> (Self, Arc) { - let drop_count = Arc::new(AtomicUsize::new(0)); - ( - Self { - drop_count: drop_count.clone(), - _value: value, - }, - drop_count, - ) - } - } - - impl Drop for DropProbe { - fn drop(&mut self) { - self.drop_count.fetch_add(1, Ordering::Relaxed); - } - } - - #[derive(Default)] - pub(super) struct WakerProbe { - wake_count: AtomicU32, - } - - impl WakerProbe { - pub(super) fn new() -> (Waker, Arc) { - let probe = Arc::new(Self::default()); - (Waker::from(probe.clone()), probe) - } - - pub(super) fn live_waker_count(this: &Arc) -> usize { - // The returned probe owns one strong reference; every other reference belongs to a - // live Waker created from it. - Arc::strong_count(this) - 1 - } - - pub(super) fn wake_count(&self) -> u32 { - self.wake_count.load(Ordering::Relaxed) - } - } - - impl Wake for WakerProbe { - fn wake(self: Arc) { - self.wake_count.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wake_count.fetch_add(1, Ordering::Relaxed); - } - } - - pub(super) fn spawn_named(name: &str, f: F) -> std::thread::JoinHandle - where - F: FnOnce() -> T + Send + 'static, - T: Send + 'static, - { - std::thread::Builder::new() - .name(name.to_owned()) - .spawn(f) - .unwrap() - } -} diff --git a/asyncband/src/sync.rs b/asyncband/src/sync.rs new file mode 100644 index 0000000..0a6cbe1 --- /dev/null +++ b/asyncband/src/sync.rs @@ -0,0 +1,75 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Mutual exclusion, notification, initialization, and task synchronization primitives. + +#[cfg(feature = "barrier")] +pub use crate::barrier::Barrier; +#[cfg(feature = "barrier")] +pub use crate::barrier::BarrierWaitResult; +#[cfg(feature = "condvar")] +pub use crate::condvar::Condvar; +#[cfg(feature = "latch")] +pub use crate::latch::Latch; +#[cfg(feature = "latch")] +pub use crate::latch::LatchWait; +#[cfg(feature = "latch")] +pub use crate::latch::OwnedLatchWait; +#[cfg(feature = "mutex")] +pub use crate::mutex::MappedMutexGuard; +#[cfg(feature = "mutex")] +pub use crate::mutex::Mutex; +#[cfg(feature = "mutex")] +pub use crate::mutex::MutexGuard; +#[cfg(feature = "mutex")] +pub use crate::mutex::OwnedMappedMutexGuard; +#[cfg(feature = "mutex")] +pub use crate::mutex::OwnedMutexGuard; +#[cfg(feature = "once")] +pub use crate::once::Once; +#[cfg(feature = "once-cell")] +pub use crate::once::OnceCell; +#[cfg(feature = "once-map")] +pub use crate::once::OnceMap; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::MappedRwLockReadGuard; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::MappedRwLockWriteGuard; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::OwnedMappedRwLockReadGuard; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::OwnedMappedRwLockWriteGuard; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::OwnedRwLockReadGuard; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::OwnedRwLockWriteGuard; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::RwLock; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::RwLockReadGuard; +#[cfg(feature = "rwlock")] +pub use crate::rwlock::RwLockWriteGuard; +#[cfg(feature = "semaphore")] +pub use crate::semaphore::OwnedSemaphorePermit; +#[cfg(feature = "semaphore")] +pub use crate::semaphore::Semaphore; +#[cfg(feature = "semaphore")] +pub use crate::semaphore::SemaphorePermit; +#[cfg(feature = "waitgroup")] +pub use crate::waitgroup::Wait; +#[cfg(feature = "waitgroup")] +pub use crate::waitgroup::WaitGroup; diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 15b6faa..9354861 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -26,15 +26,13 @@ rust-version.workspace = true asyncband = { workspace = true, features = [ "barrier", "blocking", - "broadcast", + "channel", "condvar", "latch", - "mpsc", "mutex", "once", "once-cell", "once-map", - "oneshot", "rwlock", "semaphore", "shutdown", diff --git a/benchmarks/broadcast.rs b/benchmarks/broadcast.rs index 87d9870..2142954 100644 --- a/benchmarks/broadcast.rs +++ b/benchmarks/broadcast.rs @@ -15,13 +15,14 @@ // specific language governing permissions and limitations // under the License. +use std::num::NonZeroUsize; use std::pin::pin; use std::sync::Arc; use std::sync::Barrier; use std::thread; use std::thread::JoinHandle; -use asyncband::broadcast::overflow; +use asyncband::channel::broadcast::overflow; use divan::Bencher; use divan::black_box; @@ -42,7 +43,8 @@ struct ConcurrentSend { impl ConcurrentSend { fn new(sender_count: usize) -> Self { - let (sender, receiver) = overflow::channel(CONCURRENT_BATCH_SIZE); + let (sender, receiver) = + overflow::channel(NonZeroUsize::new(CONCURRENT_BATCH_SIZE).unwrap()); 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)); @@ -59,7 +61,7 @@ impl ConcurrentSend { start.wait(); let first = worker_index * sends_per_worker; for value in first..first + sends_per_worker { - sender.send(black_box(value)); + sender.send(black_box(value)).unwrap(); } done.wait(); })); @@ -98,7 +100,8 @@ struct ConcurrentFanout { impl ConcurrentFanout { fn new(receiver_count: usize) -> Self { - let (sender, receiver) = overflow::channel(CONCURRENT_BATCH_SIZE); + let (sender, receiver) = + overflow::channel(NonZeroUsize::new(CONCURRENT_BATCH_SIZE).unwrap()); let mut receivers = Vec::with_capacity(receiver_count); receivers.push(receiver); for _ in 1..receiver_count { @@ -138,7 +141,7 @@ impl ConcurrentFanout { fn run(&mut self) { for value in 0..CONCURRENT_BATCH_SIZE { - self.sender.send(black_box(value)); + self.sender.send(black_box(value)).unwrap(); } self.start.wait(); self.done.wait(); @@ -155,23 +158,23 @@ impl Drop for ConcurrentFanout { #[divan::bench] fn send_overwrite(bencher: Bencher) { - let (sender, receiver) = overflow::channel::(1); - bencher.bench_local(|| sender.send(black_box(1))); + let (sender, receiver) = overflow::channel::(NonZeroUsize::new(1).unwrap()); + bencher.bench_local(|| sender.send(black_box(1)).unwrap()); black_box(receiver); } #[divan::bench] fn try_recv_empty(bencher: Bencher) { - let (sender, mut receiver) = overflow::channel::(1); + let (sender, mut receiver) = overflow::channel::(NonZeroUsize::new(1).unwrap()); bencher.bench_local(|| black_box(receiver.try_recv())); black_box(sender); } #[divan::bench] fn send_and_try_recv(bencher: Bencher) { - let (sender, mut receiver) = overflow::channel(1); + let (sender, mut receiver) = overflow::channel(NonZeroUsize::new(1).unwrap()); bencher.bench_local(|| { - sender.send(black_box(1)); + sender.send(black_box(1)).unwrap(); black_box(receiver.try_recv().unwrap()) }); } @@ -205,7 +208,7 @@ fn cancel_pending(bencher: Bencher) { let mut context = bench_context(); bencher.bench_local(|| { - let (sender, mut receiver) = overflow::channel::(1); + let (sender, mut receiver) = overflow::channel::(NonZeroUsize::new(1).unwrap()); { let mut recv = pin!(receiver.recv()); poll_pending(recv.as_mut(), &mut context); @@ -219,11 +222,11 @@ fn deliver_to_waiter(bencher: Bencher) { let mut context = bench_context(); bencher.bench_local(|| { - let (sender, mut receiver) = overflow::channel(1); + let (sender, mut receiver) = overflow::channel(NonZeroUsize::new(1).unwrap()); let mut recv = pin!(receiver.recv()); poll_pending(recv.as_mut(), &mut context); - sender.send(black_box(1usize)); + sender.send(black_box(1usize)).unwrap(); let value = poll_pinned_ready(recv.as_mut(), &mut context).unwrap(); black_box(value) }); @@ -234,7 +237,7 @@ fn deliver_to_receiver_batch(bencher: Bencher, receiver_count: usize) { let mut context = bench_context(); bencher.bench_local(|| { - let (sender, receiver) = overflow::channel(1); + let (sender, receiver) = overflow::channel(NonZeroUsize::new(1).unwrap()); let mut receivers = (0..receiver_count) .map(|_| receiver.resubscribe()) .collect::>(); @@ -246,7 +249,7 @@ fn deliver_to_receiver_batch(bencher: Bencher, receiver_count: usize) { poll_pending(recv.as_mut(), &mut context); } - sender.send(black_box(1usize)); + sender.send(black_box(1usize)).unwrap(); for mut recv in recvs { let value = poll_pinned_ready(recv.as_mut(), &mut context).unwrap(); black_box(value); diff --git a/benchmarks/mpsc.rs b/benchmarks/mpsc.rs index 8764734..f797759 100644 --- a/benchmarks/mpsc.rs +++ b/benchmarks/mpsc.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. -use asyncband::mpsc; +use std::num::NonZeroUsize; + +use asyncband::channel::mpsc; use divan::Bencher; use divan::black_box; @@ -28,7 +30,7 @@ const SENDER_COUNTS: &[usize] = &[1, 8, 32]; #[divan::bench(args = SENDER_COUNTS)] fn cancel_backpressured_senders(bencher: Bencher, sender_count: usize) { let mut context = bench_context(); - let (sender, mut receiver) = mpsc::bounded(1); + let (sender, mut receiver) = mpsc::bounded(NonZeroUsize::new(1).unwrap()); let senders = (0..sender_count) .map(|_| sender.clone()) .collect::>(); @@ -52,7 +54,7 @@ fn cancel_backpressured_senders(bencher: Bencher, sender_count: usize) { #[divan::bench(args = SENDER_COUNTS)] fn drain_backpressured_senders(bencher: Bencher, sender_count: usize) { let mut context = bench_context(); - let (sender, mut receiver) = mpsc::bounded(1); + let (sender, mut receiver) = mpsc::bounded(NonZeroUsize::new(1).unwrap()); let senders = (0..sender_count) .map(|_| sender.clone()) .collect::>(); diff --git a/benchmarks/oneshot.rs b/benchmarks/oneshot.rs index 1b4029b..2b2c3e4 100644 --- a/benchmarks/oneshot.rs +++ b/benchmarks/oneshot.rs @@ -22,7 +22,7 @@ use std::task::Context; use std::task::Poll; use std::task::Waker; -use asyncband::oneshot; +use asyncband::channel::oneshot; use divan::Bencher; use divan::black_box; diff --git a/docs/channel-design.md b/docs/channel-design.md new file mode 100644 index 0000000..8e2ee0e --- /dev/null +++ b/docs/channel-design.md @@ -0,0 +1,210 @@ +# Channel design for 0.7 + +## Status + +This document describes the API and semantic reference implementation in the channel redesign draft. It is intentionally a correctness-first baseline, not a claim that one synchronization strategy wins every workload. + +The old 0.6 channel implementations are replaced rather than used as architectural constraints. The new channel families do not depend on the internal atomic pointer slot retained for legacy implementation needs. + +## Design dimensions + +Producer and consumer counts are only one part of a channel contract. The public family must also make delivery, capacity, retention, overload, and waiting semantics discoverable. + +| Dimension | Choices represented in the draft | +| --- | --- | +| Delivery | one value once, competing consumers, multicast log, coalesced latest state | +| Producers | single, multiple | +| Consumers | single, multiple | +| Capacity | rendezvous, bounded, unbounded | +| Full behavior | wait, reject, explicitly replace oldest, explicitly replace newest | +| Multicast retention | overwrite and report lag, gate on slowest receiver, grow and reclaim | +| Waiting | immediate try operation or runtime-agnostic async task parking | + +An MPMC queue and an MPMC broadcast are therefore separate families: queue receivers compete for each value, while broadcast receivers each observe every retained value. + +## Public shape + +All channel APIs live under `asyncband::channel`. + +| Family | Constructors or variants | Producers | Consumers | Delivery | +| --- | --- | ---: | ---: | --- | +| oneshot | channel | 1 | 1 | one value once | +| spsc | rendezvous, bounded, unbounded | 1 | 1 | each value once | +| mpsc | rendezvous, bounded, unbounded | many | 1 | each value once | +| spmc | rendezvous, bounded, unbounded | 1 | many | each value once | +| mpmc | rendezvous, bounded, unbounded | many | many | each value once | +| broadcast::overflow | bounded | many | many | every retained value; slow receivers report Lagged | +| broadcast::backpressure | bounded | many | many | every value; slowest receiver gates producers | +| broadcast::unbounded | unbounded | many | many | every value; prefix reclaimed after all receivers advance | +| watch | one retained state | many | many | latest version only | +| disruptor::single_producer | power-of-two bounded ring | 1 | many | every contiguous published sequence | +| disruptor::multi_producer | power-of-two bounded ring | many | many | every contiguous published sequence | + +Single-producer and single-consumer endpoint types are non-cloneable and non-Sync. Their operations require mutable access, so a caller cannot accidentally turn an SPSC endpoint into a locally concurrent producer or consumer. Multiple endpoints are cloneable and Sync. Every endpoint remains Send when its value type permits it. + +The four competing-consumer queue modules are aliases over one semantic core. The aliases are not merely names: both endpoint cardinalities are part of their nominal type, while the local endpoint cardinality determines Clone and Sync behavior and whether an operation requires mutable or shared access. + +## Immediate, waiting, and blocking operations + +Every operation has one nonblocking state transition: + +| Situation | API | +| --- | --- | +| Send immediately or report Full or Disconnected | send, try_send, or try_publish | +| Receive immediately or report Empty or Disconnected | try_recv | +| Wait for capacity | send().await or publish().await | +| Wait for a value | recv().await or changed().await | + +Async methods park only the current task by registering its Waker. Asyncband does not add thread-blocking channel methods. A synchronous caller can park its thread around the same Future with the optional `asyncband::blocking` adapter, following the existing runtime-agnostic policy in the README. + +This split also gives integrations a direct base for future Stream and Sink adapters without making either trait part of the core API. + +## Queue capacity and overload + +Rendezvous has logical capacity zero. An async send owns its value until a receiver accepts it, and cancellation removes an unaccepted handoff. A nonwaiting try_send succeeds only when a receiver is already registered; once accepted, the handoff remains committed even if that particular receive future is subsequently cancelled. + +Bounded queues use send().await for backpressure and try_send for rejection. Loss is never an invisible constructor setting. force_send must be called explicitly with FullBehavior::DropOldest or FullBehavior::DropNewest, and SendOutcome::Replaced returns the displaced value to the caller. + +Bounded queue and broadcast constructors accept NonZeroUsize because rendezvous is a separate constructor and zero has no valid buffered interpretation. Disruptor constructors accept a validated Capacity that also makes the power-of-two ring requirement unrepresentable after construction. + +Unbounded queues accept while receivers exist and reclaim values as consumers pop them. They are logically unbounded, not memory-safe under an indefinitely faster producer. + +The strategy set is informed by [.NET System.Threading.Channels full modes](https://learn.microsoft.com/en-us/dotnet/api/system.threading.channels.boundedchannelfullmode), which distinguishes Wait, DropNewest, DropOldest, and DropWrite. The Asyncband draft maps Wait to send().await, maps rejection and caller-controlled drop-write to try_send, and makes replacement observable through force_send. Crossbeam's [ArrayQueue](https://docs.rs/crossbeam-queue/latest/crossbeam_queue/struct.ArrayQueue.html) similarly separates normal push failure from an explicit force_push overwrite operation. + +[Flume](https://docs.rs/flume/latest/flume/) provides the ecosystem precedent for one MPMC API spanning rendezvous, bounded, and unbounded queues. [Postage](https://docs.rs/postage/latest/postage/) makes the equally important semantic split between a competing-consumer dispatch queue, lossless broadcast, MPSC, oneshot, and watch. The Asyncband taxonomy combines those lessons while retaining explicit endpoint cardinalities. + +## Queue state and cancellation + +The ordinary queue core uses one short std mutex around a VecDeque, endpoint counts, pending rendezvous sends, and send or receive waiter queues. User Wakers are always invoked after releasing the state lock. + +Waiter identities are monotonically generated tokens rather than reusable slab indices. Cancelling a pending future removes only the registration bearing its exact identity, so an old future cannot remove a newer waiter after a wake-and-reuse cycle. Normal completion does not require an extra cleanup lock after its waiter has already been drained. + +The mutex baseline has three useful properties for an async runtime-agnostic crate: the state transition is auditable, a poll never spins waiting for another thread, and no endpoint pays an atomic read-modify-write merely to discover that it must park. Lock-free specialization remains possible behind the same topology types if measurements justify it. + +Queue value order is FIFO, but waiter scheduling does not promise strict fairness. Capacity and data transitions wake all eligible waiters so cancellation of the first scheduled task cannot strand a permit or buffered value; the executor then determines which competing task wins the state transition. + +## Implementation strategies surveyed + +The public contract and the storage algorithm are deliberately separate decisions. Existing libraries make different, internally coherent choices: + +| Concern | Established choices | Choice in this draft | +| --- | --- | --- | +| Queue type surface | Crossbeam and Flume use one endpoint type for rendezvous, bounded, and unbounded queues; Tokio gives bounded and unbounded MPSC distinct endpoint types | one endpoint type per producer/consumer topology, with capacity selected by constructor | +| Buffered storage | mutex-protected deque; fixed ring with per-slot state; segmented linked blocks | mutex-protected VecDeque reference core | +| Rendezvous | zero-capacity queue with paired waiters or a dedicated handoff state | dedicated pending-send handoff state inside the queue core | +| Producer progress | serialized critical section; single-writer cursor; multi-writer CAS reservation plus publication tracking | serialized state transition for every topology | +| Waiting | thread blocking, task parking, spinning, yielding, sleeping, or phased backoff | task parking through Wakers only | +| Multicast publication | serialize append; or reserve independently and expose only a contiguous published prefix | serialized append for broadcast; explicit reservation and contiguous publication for Disruptor | + +[Crossbeam channels](https://docs.rs/crossbeam-channel/latest/crossbeam_channel/) and [Flume](https://docs.rs/flume/latest/flume/) show that one endpoint type can consistently span zero, bounded, and unbounded capacities. [Tokio MPSC](https://docs.rs/tokio/latest/tokio/sync/mpsc/) instead uses distinct bounded and unbounded endpoints; that lets unbounded send be synchronous, and its implementation uses a lock-free linked list of fixed-size blocks. The current Asyncband draft chooses the smaller common capacity surface: `try_send` is always synchronous, while `send` is uniformly async and immediately becomes ready for unbounded queues. A policy-typed capacity axis remains a viable follow-up if the ability to remove impossible methods and make unbounded `send` synchronous is worth the extra public types. + +A fixed ring with per-slot sequence generations is a strong candidate for bounded SPSC and MPMC specialization; a segmented list avoids reallocating a monolithic unbounded buffer; a single-producer cursor can eliminate producer-side contention; and a CAS claim cursor can scale multiple producers. None is a free substitution. Each needs its own proof for cancellation, destruction, publication gaps, ABA or generation reuse, wrap-around, and wake registration. The draft therefore first fixes the observable protocol and leaves these as benchmark-driven internal specializations. + +The Broadcast policy is nominal in the public API even though the reference implementation shares log and cursor machinery. This prevents accidentally exchanging overflow and lossless endpoints and permits policy-specific send and error APIs. Splitting the three storage backends remains possible without changing callers if profiling or a simpler invariant warrants it. + +## Broadcast retention + +Broadcast uses one committed tail under the same state lock as its log and waiter metadata. A sender appends the complete value before advancing the tail, so a receiver never mistakes reservation for publication. This directly closes the publication-hole and stale-writer class tracked by the Broadcast correctness issue. + +The three retention policies share the sequenced log and receiver-cursor machinery because their state invariants are the same: + +| Policy | Full behavior | Slow receiver behavior | Reclamation | +| --- | --- | --- | --- | +| overflow | remove oldest and publish new value | next receive returns exact Lagged count, then resumes at oldest retained value | bounded ring-like prefix | +| backpressure | try_send returns Full; send().await parks | no loss | minimum receiver cursor gates producers | +| unbounded | always accepts while connected | no loss | prefix removed after every receiver advances or drops | + +DropNewest is not offered for multicast. Replacing the newest retained sequence can mutate a value after some receivers have already observed that sequence while other receivers have not, violating the single-value-per-sequence contract. Dropping the incoming broadcast would also make send success ambiguous. A caller that wants coalescing should use watch instead. + +The retention choice is part of the endpoint's nominal type rather than runtime configuration. Overflow and unbounded `send` calls are synchronous because those policies never wait for capacity. Backpressure alone exposes an async `send`, alongside `try_send`. Only overflow receive errors contain `Lagged`; the two lossless policies reuse the common channel receive errors, so callers do not need to handle an impossible state. + +New subscriptions start at the committed tail. Cloning a receiver preserves its cursor. Sending with no receivers returns the original value, rather than retaining values for a hypothetical future subscriber. + +These policies implement the three levels proposed in [issue #95](https://github.com/apache/asyncband/issues/95) and match the distinction between lagging broadcast and backpressured broadcast discussed in [issue #88](https://github.com/apache/asyncband/issues/88). Tokio's [broadcast channel](https://docs.rs/tokio/latest/tokio/sync/broadcast/) is the primary reference for bounded lag reporting. + +[async-broadcast](https://docs.rs/async-broadcast/latest/async_broadcast/) demonstrates that bounded backpressure and opt-in oldest-value overflow can share one multicast abstraction. Its overflow send reports the removed value, reinforcing the rule that deliberate loss should be observable rather than silently configured. + +## Watch + +Watch is included because latest-state distribution is not a special case of queue or broadcast retention. Each send publishes a new version and replaces the previous Arc-backed value. changed().await coalesces intermediate versions and returns the latest value. A receiver can borrow without marking the version observed or borrow_and_update to advance explicitly. + +The contract follows the broad shape of [Tokio watch](https://docs.rs/tokio/latest/tokio/sync/watch/) while returning Arc values so the channel itself does not require T: Clone. + +## Disruptor-style sequencers + +The Disruptor modules are bounded multicast logs, not MPMC work queues. Their essential contract follows the [LMAX Disruptor user guide](https://lmax-exchange.github.io/disruptor/user-guide/): + +1. A producer reserves a monotonically increasing sequence. +2. It writes the corresponding preallocated ring slot. +3. It marks that sequence available. +4. Consumers see only the highest contiguous published prefix. +5. The minimum subscriber sequence gates wrap-around so an unread slot is never overwritten. + +The multi-producer implementation deliberately permits reservations to finish out of order. If sequence N+1 is ready before N, availability records N+1 but the published cursor does not advance. Publishing N then scans the per-slot availability generations and exposes both sequences together. This is the distinction that an atomic reservation cursor alone cannot provide. + +The single-producer type enforces the single-writer rule and uses the same semantic core, without cloning or concurrent access to the publisher. The multi-producer type is cloneable and allows concurrent reservations. Both currently serialize reservation metadata with the state mutex; no atomic pointer utility or standard atomic is required for the reference implementation. + +This differs from the Java implementation's CAS-oriented [MultiProducerSequencer](https://github.com/LMAX-Exchange/disruptor/blob/c871ca49826a6be7ada6957f6fbafcfecf7b1f87/src/main/java/com/lmax/disruptor/MultiProducerSequencer.java), but preserves its observable claim, availability, contiguous publication, and gating rules. A lock-free backend would be an optimization of this contract, not a different public channel. + +Only task-parking waits are built in. Busy spin, yield, phased backoff, blocking conditions, batch translation, mutable preallocated event factories, and consumer dependency graphs are intentionally outside this first API. Busy waiting is a scheduler and deployment choice that a runtime-agnostic async primitive should not perform inside Future::poll. + +LMAX exposes [blocking, timeout-blocking, sleeping, yielding, busy-spin, and phased-backoff wait strategies](https://lmax-exchange.github.io/disruptor/javadoc/com.lmax.disruptor/com/lmax/disruptor/WaitStrategy.html). In Asyncband, Waker registration is the task-level analogue of the conservative blocking strategy. Timeouts compose around the returned Future, while spin, yield, sleep, and phased policies belong in a dedicated thread-driven event processor rather than in the channel Future itself. If such an event processor is added later, its wait strategy should be separate from the sequencer so the same ring correctness contract remains usable by ordinary async tasks. + +## Disconnection and error vocabulary + +Channel modules re-export a common SendError, TrySendError, RecvError, and TryRecvError when their state space matches. Only overflow broadcast adds Lagged to its receive errors. Endpoint predicates use is_disconnected, matching the Disconnected variants and the 0.7 naming decision in [issue #141](https://github.com/apache/asyncband/issues/141). + +Buffered channels drain accepted values before returning Disconnected. Sending fails as soon as no receiver remains and returns ownership of the unsent value. + +## Correctness invariants + +- Accepted queue values are received at most once, and FIFO order is preserved at the shared queue boundary. +- A rendezvous send future does not complete before a receiver accepts its value. +- Every published broadcast sequence identifies one immutable value. +- Broadcast tail never advances past an incomplete value. +- Backpressured multicast producers never wrap or reclaim past the minimum active receiver cursor. +- A Disruptor cursor denotes a contiguous published prefix, not the largest reservation. +- No Waker is invoked while an internal channel lock is held. +- Cancelling a pending send or receive deregisters only that future's waiter. +- Dropping the last opposite endpoint wakes every task that can now observe Disconnected. + +## 0.7 crate grouping + +The draft introduces three top-level groups while leaving non-channel root modules available during review. The formerly public atomic pointer utilities and admission policy have already been removed on `main` and are not reintroduced here: + +| Group | Contents | +| --- | --- | +| `asyncband::sync` | Barrier, Condvar, Latch, Mutex, Once, RwLock, Semaphore, WaitGroup, and guards | +| `asyncband::channel` | all transfer, queue, multicast, state, and sequenced-ring channels | +| `asyncband::coordination` | shutdown and singleflight protocols | + +For the breaking release, the remaining synchronization and coordination root-level duplicates can be reviewed separately. This draft removes the superseded root-level channel paths. + +A matching additive Cargo feature layout can be introduced after the module names settle: + +~~~toml +[features] +default = [] +full = ["blocking", "sync", "channel", "coordination"] +sync = ["barrier", "condvar", "latch", "mutex", "once", "rwlock", "semaphore", "waitgroup"] +channel = ["oneshot", "queue", "broadcast", "watch", "disruptor"] +coordination = ["shutdown", "singleflight"] +~~~ + +Umbrella features only group leaf features; leaf features encode implementation dependencies. CI should cover no-default, each umbrella, default, and all-features rather than the exponential set of every leaf combination. This follows Cargo's [additive feature model](https://doc.rust-lang.org/cargo/reference/features.html). + +The feature grouping is applied in this rebased draft because `main` now enables no primitive modules by default and validates every individual feature in CI. + +## Benchmark gates for later specialization + +The reference implementation should be benchmarked by operation and contention shape before replacing its synchronization mechanism: + +- uncontended try send and try receive for each topology; +- producer contention for MPSC and MPMC; +- consumer contention and fairness for SPMC and MPMC; +- cancellation-heavy async waits; +- broadcast fan-out with fast and slow receivers; +- Disruptor single-producer and multi-producer throughput at several subscriber counts; +- tail latency under task oversubscription, where busy-spin designs commonly regress. + +An optimized backend must retain the invariants above and demonstrate a material workload benefit. In particular, replacing a mutex with an atomic claim cursor is incomplete unless publication gaps, per-slot generations, wrap gating, cancellation, and Waker registration all remain correct. diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 104171c..d42fdbf 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -29,15 +29,13 @@ tokio = { workspace = true, features = ["full"] } asyncband = { workspace = true, features = [ "barrier", "blocking", - "broadcast", + "channel", "condvar", "latch", - "mpsc", "mutex", "once", "once-cell", "once-map", - "oneshot", "rwlock", "semaphore", "shutdown", diff --git a/tests-integration/tests/broadcast_test.rs b/tests-integration/tests/broadcast_test.rs deleted file mode 100644 index 56ab69d..0000000 --- a/tests-integration/tests/broadcast_test.rs +++ /dev/null @@ -1,350 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::future::Future; -use std::sync::Arc; -use std::sync::Barrier; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Wake; -use std::task::Waker; -use std::thread; - -use asyncband::broadcast::overflow::*; - -struct TrackWake(AtomicUsize); - -impl Wake for TrackWake { - fn wake(self: Arc) { - self.0.fetch_add(1, Ordering::Relaxed); - } -} - -#[derive(Debug)] -struct PanicOnDrop { - value: u64, - panic: bool, - panicked: Arc, -} - -impl Clone for PanicOnDrop { - fn clone(&self) -> Self { - Self { - value: self.value, - panic: false, - panicked: self.panicked.clone(), - } - } -} - -impl Drop for PanicOnDrop { - fn drop(&mut self) { - if self.panic && !self.panicked.swap(true, Ordering::Relaxed) { - panic!("panic while replacing a broadcast slot"); - } - } -} - -#[tokio::test] -async fn test_broadcast_basic() { - let (tx, mut rx1) = channel(10); - let mut rx2 = rx1.clone(); - - tx.send(10); - tx.send(20); - - assert_eq!(rx1.recv().await, Ok(10)); - assert_eq!(rx1.recv().await, Ok(20)); - assert_eq!(rx2.recv().await, Ok(10)); - assert_eq!(rx2.recv().await, Ok(20)); -} - -#[tokio::test] -async fn test_broadcast_lagged() { - let (tx, mut rx) = channel(2); - - tx.send(1); - tx.send(2); - tx.send(3); - - // Overwrites 1. Rx lagged by 1 (missed msg '1'). - // Rx should return Lagged(1) and catch up to 2 (oldest valid). - assert_eq!(rx.recv().await, Err(RecvError::Lagged(1))); - assert_eq!(rx.recv().await, Ok(2)); - assert_eq!(rx.recv().await, Ok(3)); -} - -#[tokio::test] -async fn test_broadcast_lagged_multi() { - let (tx, mut rx) = channel(2); - - tx.send(1); - tx.send(2); - tx.send(3); - tx.send(4); - - // Overwrites 1 and 2. Missed 2 messages. - assert_eq!(rx.recv().await, Err(RecvError::Lagged(2))); - assert_eq!(rx.recv().await, Ok(3)); - assert_eq!(rx.recv().await, Ok(4)); -} - -#[tokio::test] -async fn test_broadcast_closed() { - let (tx, mut rx) = channel::<()>(10); - drop(tx); - assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); -} - -#[tokio::test] -async fn test_wait_mechanism() { - let (tx, mut rx) = channel(10); - - let handle = tokio::spawn(async move { rx.recv().await }); - - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - tx.send(42); - - assert_eq!(handle.await.unwrap(), Ok(42)); -} - -#[test] -fn cancelled_recv_releases_its_waker() { - let (tx, mut rx) = channel::<()>(1); - let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); - let waker = Waker::from(tracker.clone()); - let baseline = Arc::strong_count(&tracker); - let mut context = Context::from_waker(&waker); - let mut recv = Box::pin(rx.recv()); - - assert!(recv.as_mut().poll(&mut context).is_pending()); - assert_eq!(Arc::strong_count(&tracker), baseline + 1); - - drop(recv); - assert_eq!(Arc::strong_count(&tracker), baseline); - - tx.send(()); - assert_eq!(tracker.0.load(Ordering::Relaxed), 0); - assert_eq!(rx.try_recv(), Ok(())); -} - -#[tokio::test] -async fn test_subscribe() { - let (tx, _rx) = channel(10); - let mut rx = tx.subscribe(); - - tx.send(100); - assert_eq!(rx.recv().await, Ok(100)); -} - -#[tokio::test] -async fn test_resubscribe() { - let (tx, mut rx) = channel(2); - - tx.send(1); - tx.send(2); - - let mut rx2 = rx.resubscribe(); - - // rx sees 1, 2 - // rx2 sees nothing yet (starts at tail=2) - - tx.send(3); - - assert_eq!(rx.recv().await, Err(RecvError::Lagged(1))); - assert_eq!(rx.recv().await, Ok(2)); - assert_eq!(rx2.recv().await, Ok(3)); -} - -#[tokio::test] -async fn test_try_recv() { - let (tx, mut rx) = channel(16); - - // Empty - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - - // Success - tx.send(10); - assert_eq!(rx.try_recv(), Ok(10)); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - - // Closed - drop(tx); - assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); -} - -#[tokio::test] -async fn test_try_recv_lagged() { - let (tx, mut rx) = channel(2); - tx.send(1); - tx.send(2); - tx.send(3); - - assert_eq!(rx.try_recv(), Err(TryRecvError::Lagged(1))); - assert_eq!(rx.try_recv(), Ok(2)); - assert_eq!(rx.try_recv(), Ok(3)); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); -} - -#[test] -fn panicking_send_does_not_publish_an_unwritten_slot() { - let panicked = Arc::new(AtomicBool::new(false)); - let (tx, mut rx) = channel(1); - tx.send(PanicOnDrop { - value: 1, - panic: true, - panicked: panicked.clone(), - }); - - let received = rx.try_recv().unwrap(); - assert_eq!(received.value, 1); - drop(received); - - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - tx.send(PanicOnDrop { - value: 2, - panic: false, - panicked: panicked.clone(), - }); - })); - assert!(result.is_err()); - assert!(panicked.load(Ordering::Relaxed)); - - assert!(matches!(rx.try_recv(), Err(TryRecvError::Empty))); - tx.send(PanicOnDrop { - value: 3, - panic: false, - panicked: panicked.clone(), - }); - assert_eq!(rx.try_recv().unwrap().value, 3); - - drop(tx); - assert!(matches!(rx.try_recv(), Err(TryRecvError::Disconnected))); -} - -#[test] -fn concurrent_overwrite_preserves_sequence_and_lag_count() { - const MESSAGE_COUNT: u64 = 200_000; - - let (tx, mut rx) = channel(2); - let producer = thread::spawn(move || { - for value in 0..MESSAGE_COUNT { - tx.send(value); - } - }); - - let mut next = 0_u64; - loop { - match rx.try_recv() { - Ok(value) => { - assert_eq!(value, next); - next = next.wrapping_add(1); - } - Err(TryRecvError::Lagged(missed)) => { - assert!(missed > 0); - next = next.wrapping_add(missed); - } - Err(TryRecvError::Empty) => thread::yield_now(), - Err(TryRecvError::Disconnected) => break, - } - } - - producer.join().unwrap(); - assert_eq!(next, MESSAGE_COUNT); -} - -#[test] -fn concurrent_receivers_observe_the_same_sequence() { - const MESSAGE_COUNT: usize = 4096; - const RECEIVER_COUNT: usize = 8; - - let (tx, receiver) = channel(MESSAGE_COUNT); - let mut receivers = Vec::with_capacity(RECEIVER_COUNT); - receivers.push(receiver); - for _ in 1..RECEIVER_COUNT { - receivers.push(receivers[0].clone()); - } - - let ready = Arc::new(Barrier::new(RECEIVER_COUNT + 1)); - let workers = receivers - .into_iter() - .map(|mut receiver| { - let ready = ready.clone(); - thread::spawn(move || { - ready.wait(); - let mut received = Vec::with_capacity(MESSAGE_COUNT); - loop { - match receiver.try_recv() { - Ok(value) => received.push(value), - Err(TryRecvError::Empty) => thread::yield_now(), - Err(TryRecvError::Disconnected) => return received, - Err(TryRecvError::Lagged(missed)) => { - panic!("receiver unexpectedly lagged by {missed}") - } - } - } - }) - }) - .collect::>(); - - ready.wait(); - for value in 0..MESSAGE_COUNT { - tx.send(value); - } - drop(tx); - - let expected = (0..MESSAGE_COUNT).collect::>(); - for worker in workers { - assert_eq!(worker.join().unwrap(), expected); - } -} - -#[tokio::test] -async fn test_multi_senders_concurrent() { - let (tx, mut rx) = channel(100); - let tx1 = tx.clone(); - let tx2 = tx.clone(); - - tokio::spawn(async move { - for i in 0..10 { - tx1.send(i); - } - }); - - tokio::spawn(async move { - for i in 10..20 { - tx2.send(i); - } - }); - - // Main tx can also send - for i in 20..30 { - tx.send(i); - } - drop(tx); - - let mut received = Vec::new(); - while let Ok(n) = rx.recv().await { - received.push(n); - } - received.sort(); - - let expected = (0..30).collect::>(); - assert_eq!(received, expected); -} diff --git a/tests-integration/tests/composition_test.rs b/tests-integration/tests/composition_test.rs index b7de067..06afd37 100644 --- a/tests-integration/tests/composition_test.rs +++ b/tests-integration/tests/composition_test.rs @@ -21,8 +21,8 @@ use std::time::Duration; use asyncband::barrier::Barrier; use asyncband::blocking::FutureExt as _; +use asyncband::channel::oneshot; use asyncband::mutex::Mutex; -use asyncband::oneshot; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn public_primitives_compose_across_modules() { diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs deleted file mode 100644 index fab703e..0000000 --- a/tests-integration/tests/mpsc_test.rs +++ /dev/null @@ -1,290 +0,0 @@ -// 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::time::Instant; - -use asyncband::mpsc; -use asyncband::mpsc::RecvError; -use asyncband::mpsc::TryRecvError; -use asyncband::mpsc::TrySendError; -use tests_integration::test_runtime; -use tokio_test::assert_ok; - -#[test] -fn test_unbounded_pressure() { - let n = 1024 * 1024; - let (tx, mut rx) = mpsc::unbounded(); - - test_runtime().block_on(async move { - let start = Instant::now(); - tokio::spawn(async move { - for i in 0..n { - tx.send(i).unwrap(); - } - }); - - for i in 0..n { - assert_eq!(rx.recv().await, Ok(i)); - } - println!("Elapsed: {:?}", start.elapsed()); - }); -} - -#[test] -fn test_unbounded_sum() { - let (tx, mut rx) = mpsc::unbounded(); - - test_runtime().block_on(async move { - for i in 0..100 { - let tx = tx.clone(); - tokio::spawn(async move { - tx.send(i).unwrap(); - }); - } - drop(tx); - - let mut sum = 0; - while let Ok(i) = rx.recv().await { - sum += i; - } - assert_eq!(sum, 4950); - }); -} - -#[tokio::test] -async fn select_streams() { - let (tx1, mut rx1) = mpsc::unbounded::(); - let (tx2, mut rx2) = mpsc::unbounded::(); - let (tx3, mut rx3) = mpsc::bounded(1); - let (tx4, mut rx4) = mpsc::bounded(1); - - tokio::spawn(async move { - assert_ok!(tx2.send(1)); - tokio::task::yield_now().await; - - assert_ok!(tx1.send(2)); - tokio::task::yield_now().await; - - assert_ok!(tx2.send(3)); - tokio::task::yield_now().await; - - assert_ok!(tx3.send(4).await); - tokio::task::yield_now().await; - - assert_ok!(tx4.send(5).await); - tokio::task::yield_now().await; - - assert_ok!(tx3.send(6).await); - tokio::task::yield_now().await; - - drop((tx1, tx2)); - }); - - let mut rem = true; - let mut msgs = vec![]; - let mut rx1_closed = false; - let mut rx2_closed = false; - let mut rx3_closed = false; - let mut rx4_closed = false; - - while rem { - rem = !(rx1_closed && rx2_closed && rx3_closed && rx4_closed); - - tokio::select! { - result = rx1.recv(), if !rx1_closed => { - match result { - Ok(x) => msgs.push(x), - Err(RecvError::Disconnected) => rx1_closed = true, - } - } - result = rx2.recv(), if !rx2_closed => { - match result { - Ok(y) => msgs.push(y), - Err(RecvError::Disconnected) => rx2_closed = true, - } - } - result = rx3.recv(), if !rx3_closed => { - match result { - Ok(z) => msgs.push(z), - Err(RecvError::Disconnected) => rx3_closed = true, - } - } - result = rx4.recv(), if !rx4_closed => { - match result { - Ok(w) => msgs.push(w), - Err(RecvError::Disconnected) => rx4_closed = true, - } - } - else => { - rx1_closed = true; - rx2_closed = true; - rx3_closed = true; - rx4_closed = true; - } - } - } - - msgs.sort_unstable(); - assert_eq!(&msgs[..], &[1, 2, 3, 4, 5, 6]); -} - -#[tokio::test] -async fn send_recv_unbounded() { - let (tx, mut rx) = mpsc::unbounded::(); - - // Using `try_send` - assert_ok!(tx.send(1)); - assert_ok!(tx.send(2)); - - assert_eq!(rx.recv().await, Ok(1)); - assert_eq!(rx.recv().await, Ok(2)); - - drop(tx); - - assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); -} - -#[tokio::test] -async fn async_send_recv_unbounded() { - let (tx, mut rx) = mpsc::unbounded(); - - tokio::spawn(async move { - assert_ok!(tx.send(1)); - assert_ok!(tx.send(2)); - }); - - assert_eq!(Ok(1), rx.recv().await); - assert_eq!(Ok(2), rx.recv().await); - assert_eq!(Err(RecvError::Disconnected), rx.recv().await); -} - -#[test] -fn try_recv_unbounded() { - for num in 0..100 { - let (tx, mut rx) = mpsc::unbounded(); - - for i in 0..num { - tx.send(i).unwrap(); - } - - for i in 0..num { - assert_eq!(rx.try_recv(), Ok(i)); - } - - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - drop(tx); - assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); - } -} - -#[test] -fn try_recv_close_while_empty_unbounded() { - let (tx, mut rx) = mpsc::unbounded::<()>(); - - assert_eq!(Err(TryRecvError::Empty), rx.try_recv()); - drop(tx); - assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); -} - -#[tokio::test] -async fn send_recv_bounded() { - let (tx, mut rx) = mpsc::bounded(1); - - tx.send(1).await.unwrap(); - assert_eq!(rx.recv().await, Ok(1)); - - drop(tx); - assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); -} - -#[tokio::test] -async fn async_send_recv_bounded() { - let (tx, mut rx) = mpsc::bounded(1); - - tx.send(1).await.unwrap(); - // This will block until the receiver is ready to receive. - tokio::spawn(async move { - tx.send(2).await.unwrap(); - }); - - assert_eq!(Ok(1), rx.recv().await); - assert_eq!(Ok(2), rx.recv().await); - assert_eq!(Err(RecvError::Disconnected), rx.recv().await); -} - -#[test] -fn try_send_recv_bounded() { - for num in 1..101 { - let (tx, mut rx) = mpsc::bounded(num); - - for i in 0..num { - tx.try_send(i).unwrap(); - } - - assert_eq!(tx.try_send(num), Err(TrySendError::Full(num))); - - for i in 0..num { - assert_eq!(rx.try_recv(), Ok(i)); - } - - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - drop(tx); - assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); - } -} - -#[tokio::test] -async fn try_send_after_close_bounded() { - let (tx, rx) = mpsc::bounded(1); - - tx.try_send(1).unwrap(); - drop(rx); - - assert_eq!(tx.try_send(3), Err(TrySendError::Disconnected(3))); -} - -#[tokio::test] -async fn send_after_close_bounded() { - let (tx, mut rx) = mpsc::bounded(1); - - tx.send(1).await.unwrap(); - assert_eq!(rx.recv().await, Ok(1)); - - drop(rx); - let error = tx.send(2).await.unwrap_err(); - assert_eq!(error.into_inner(), 2); -} - -#[test] -fn test_bounded_pressure() { - let n = 1024 * 1024; - let (tx, mut rx) = mpsc::bounded(1024); - - test_runtime().block_on(async move { - let start = Instant::now(); - tokio::spawn(async move { - for i in 0..n { - tx.send(i).await.unwrap(); - } - }); - - for i in 0..n { - assert_eq!(rx.recv().await, Ok(i)); - } - println!("Elapsed: {:?}", start.elapsed()); - }); -} diff --git a/tests-integration/tests/oneshot_test/main.rs b/tests-integration/tests/oneshot_test/main.rs deleted file mode 100644 index 0cb6152..0000000 --- a/tests-integration/tests/oneshot_test/main.rs +++ /dev/null @@ -1,385 +0,0 @@ -// 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::future::IntoFuture; -use std::pin::Pin; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; - -use asyncband::oneshot; -use asyncband::oneshot::TryRecvError; - -use self::support::DropProbe; -use self::support::WakerProbe; -use self::support::spawn_named; -use self::support::spin_until; - -mod support; - -#[test] -fn send_before_await() { - let (sender, receiver) = oneshot::channel(); - assert!(!receiver.has_message()); - assert!(sender.send(19i128).is_ok()); - assert!(receiver.has_message()); - assert_eq!(pollster::block_on(receiver), Ok(19i128)); -} - -#[test] -fn await_with_dropped_sender() { - let (sender, receiver) = oneshot::channel::(); - assert!(!receiver.is_disconnected()); - drop(sender); - assert!(receiver.is_disconnected()); - assert_eq!( - pollster::block_on(receiver), - Err(oneshot::RecvError::Disconnected) - ); -} - -#[test] -fn try_recv_success_then_disconnected() { - let (tx, rx) = oneshot::channel::(); - tx.send(10).unwrap(); - - assert!(!rx.is_disconnected()); - assert_eq!(rx.try_recv(), Ok(10)); - assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); - assert!(rx.is_disconnected()); - assert!(!rx.has_message()); - assert_eq!( - pollster::block_on(rx.into_future()), - Err(oneshot::RecvError::Disconnected) - ); -} - -#[test] -fn try_recv_distinguishes_empty_from_disconnected() { - let (tx, rx) = oneshot::channel::<()>(); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - drop(tx); - assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); -} - -#[test] -fn send_error_preserves_message_until_consumed() { - let (sender, receiver) = oneshot::channel(); - let (message, message_drop_count) = DropProbe::new(17u128); - - assert!(!sender.is_disconnected()); - drop(receiver); - assert!(sender.is_disconnected()); - - let error = sender.send(message).unwrap_err(); - assert_eq!(message_drop_count.load(Ordering::Relaxed), 0); - assert_eq!(*error.as_inner().value(), 17); - - let message = error.into_inner(); - assert_eq!(message_drop_count.load(Ordering::Relaxed), 0); - drop(message); - assert_eq!(message_drop_count.load(Ordering::Relaxed), 1); -} - -#[test] -fn dropping_send_error_drops_message() { - let (sender, receiver) = oneshot::channel(); - let (message, message_drop_count) = DropProbe::new(()); - - drop(receiver); - drop(sender.send(message).unwrap_err()); - - assert_eq!(message_drop_count.load(Ordering::Relaxed), 1); -} - -#[test] -fn dropping_receiver_after_send_drops_message() { - let (sender, receiver) = oneshot::channel(); - let (message, message_drop_count) = DropProbe::new(()); - - sender.send(message).unwrap(); - assert_eq!(message_drop_count.load(Ordering::Relaxed), 0); - drop(receiver); - - assert_eq!(message_drop_count.load(Ordering::Relaxed), 1); -} - -#[test] -fn dropping_unpolled_recv_closes_channel() { - let (sender, receiver) = oneshot::channel::(); - let receiver = receiver.into_future(); - - drop(receiver); - - assert!(sender.is_disconnected()); - assert_eq!(sender.send(17).unwrap_err().into_inner(), 17); -} - -#[test] -fn dropping_recv_after_send_drops_message() { - let (sender, receiver) = oneshot::channel(); - let (message, message_drop_count) = DropProbe::new(()); - let receiver = receiver.into_future(); - - sender.send(message).unwrap(); - assert_eq!(message_drop_count.load(Ordering::Relaxed), 0); - drop(receiver); - - assert_eq!(message_drop_count.load(Ordering::Relaxed), 1); -} - -#[test] -fn poll_then_send() { - let (sender, receiver) = oneshot::channel::(); - let mut receiver = receiver.into_future(); - - let (waker, waker_probe) = WakerProbe::new(); - let mut context = Context::from_waker(&waker); - - assert_eq!(Pin::new(&mut receiver).poll(&mut context), Poll::Pending); - assert_eq!(WakerProbe::live_waker_count(&waker_probe), 2); - assert_eq!(waker_probe.wake_count(), 0); - - sender.send(1234).unwrap(); - assert_eq!(WakerProbe::live_waker_count(&waker_probe), 1); - assert_eq!(waker_probe.wake_count(), 1); - - assert_eq!( - Pin::new(&mut receiver).poll(&mut context), - Poll::Ready(Ok(1234)) - ); - assert_eq!(WakerProbe::live_waker_count(&waker_probe), 1); - assert_eq!(waker_probe.wake_count(), 1); -} - -#[test] -fn poll_then_drop_sender() { - let (sender, receiver) = oneshot::channel::(); - let mut receiver = receiver.into_future(); - - let (waker, waker_probe) = WakerProbe::new(); - let mut context = Context::from_waker(&waker); - - assert_eq!(Pin::new(&mut receiver).poll(&mut context), Poll::Pending); - assert_eq!(WakerProbe::live_waker_count(&waker_probe), 2); - assert_eq!(waker_probe.wake_count(), 0); - - drop(sender); - assert_eq!(WakerProbe::live_waker_count(&waker_probe), 1); - assert_eq!(waker_probe.wake_count(), 1); - - assert_eq!( - Pin::new(&mut receiver).poll(&mut context), - Poll::Ready(Err(oneshot::RecvError::Disconnected)) - ); - assert_eq!(WakerProbe::live_waker_count(&waker_probe), 1); - assert_eq!(waker_probe.wake_count(), 1); -} - -#[test] -fn poll_with_different_wakers() { - let (sender, receiver) = oneshot::channel::(); - let mut receiver = receiver.into_future(); - - let (waker1, waker_probe1) = WakerProbe::new(); - let mut context1 = Context::from_waker(&waker1); - - assert_eq!(Pin::new(&mut receiver).poll(&mut context1), Poll::Pending); - assert_eq!(WakerProbe::live_waker_count(&waker_probe1), 2); - assert_eq!(waker_probe1.wake_count(), 0); - - let (waker2, waker_probe2) = WakerProbe::new(); - let mut context2 = Context::from_waker(&waker2); - - assert_eq!(Pin::new(&mut receiver).poll(&mut context2), Poll::Pending); - assert_eq!(WakerProbe::live_waker_count(&waker_probe1), 1); - assert_eq!(waker_probe1.wake_count(), 0); - assert_eq!(WakerProbe::live_waker_count(&waker_probe2), 2); - assert_eq!(waker_probe2.wake_count(), 0); - - sender.send(1234).unwrap(); - assert_eq!(WakerProbe::live_waker_count(&waker_probe1), 1); - assert_eq!(waker_probe1.wake_count(), 0); - assert_eq!(WakerProbe::live_waker_count(&waker_probe2), 1); - assert_eq!(waker_probe2.wake_count(), 1); -} - -#[test] -fn poll_with_different_wakers_across_threads() { - let (sender, receiver) = oneshot::channel::(); - let mut receiver = receiver.into_future(); - - let (waker1, waker_probe1) = WakerProbe::new(); - let mut context1 = Context::from_waker(&waker1); - - assert_eq!(Pin::new(&mut receiver).poll(&mut context1), Poll::Pending); - assert_eq!(WakerProbe::live_waker_count(&waker_probe1), 2); - assert_eq!(waker_probe1.wake_count(), 0); - - let receiver_thread = spawn_named("receiver", move || { - let (waker2, waker_probe2) = WakerProbe::new(); - let mut context2 = Context::from_waker(&waker2); - - assert_eq!(Pin::new(&mut receiver).poll(&mut context2), Poll::Pending); - assert_eq!(WakerProbe::live_waker_count(&waker_probe2), 2); - assert_eq!(waker_probe2.wake_count(), 0); - - drop(receiver); - assert_eq!(WakerProbe::live_waker_count(&waker_probe2), 1); - }); - - receiver_thread.join().unwrap(); - assert_eq!(WakerProbe::live_waker_count(&waker_probe1), 1); - assert!(sender.is_disconnected()); -} - -#[test] -fn drop_pending_receiver_closes_channel_and_drops_waker() { - let (sender, receiver) = oneshot::channel::(); - let mut receiver = receiver.into_future(); - - let (waker, waker_probe) = WakerProbe::new(); - let mut context = Context::from_waker(&waker); - - assert_eq!(Pin::new(&mut receiver).poll(&mut context), Poll::Pending); - assert_eq!(WakerProbe::live_waker_count(&waker_probe), 2); - assert_eq!(waker_probe.wake_count(), 0); - - drop(receiver); - assert_eq!(WakerProbe::live_waker_count(&waker_probe), 1); - assert_eq!(waker_probe.wake_count(), 0); - assert!(sender.is_disconnected()); - - let error = sender.send(1234).unwrap_err(); - assert_eq!(*error.as_inner(), 1234); -} - -#[test] -fn poll_then_drop_receiver_during_send() { - let (sender, receiver) = oneshot::channel(); - let (message, message_drop_count) = DropProbe::new(1234u128); - let mut receiver = receiver.into_future(); - - let (waker, _waker_probe) = WakerProbe::new(); - let mut context = Context::from_waker(&waker); - - assert!(matches!( - Pin::new(&mut receiver).poll(&mut context), - Poll::Pending - )); - - let sender_thread = spawn_named("sender", move || sender.send(message)); - drop(receiver); - - // Whether send or receiver drop wins, exactly one side owns and drops the message. - drop(sender_thread.join().unwrap()); - assert_eq!(message_drop_count.load(Ordering::Relaxed), 1); -} - -#[test] -fn concurrent_send_and_try_recv_to_completion() { - let (sender, receiver) = oneshot::channel::(); - - let receiver_thread = spawn_named("receiver", move || { - spin_until("message from sender", || match receiver.try_recv() { - Ok(999) => true, - Ok(value) => panic!("unexpected value: {value}"), - Err(TryRecvError::Empty) => false, - Err(TryRecvError::Disconnected) => panic!("unexpected disconnect"), - }); - }); - - let sender_thread = spawn_named("sender", move || { - sender.send(999).unwrap(); - }); - - receiver_thread.join().unwrap(); - sender_thread.join().unwrap(); -} - -#[test] -fn concurrent_drop_sender_and_try_recv_to_completion() { - let (sender, receiver) = oneshot::channel::(); - - let receiver_thread = spawn_named("receiver", move || { - spin_until("sender disconnect", || match receiver.try_recv() { - Ok(value) => panic!("unexpected value: {value}"), - Err(TryRecvError::Empty) => false, - Err(TryRecvError::Disconnected) => true, - }); - }); - - let sender_thread = spawn_named("sender", move || { - drop(sender); - }); - - receiver_thread.join().unwrap(); - sender_thread.join().unwrap(); -} - -#[test] -fn concurrent_send_and_poll_to_completion() { - let (sender, receiver) = oneshot::channel::(); - - let receiver_thread = spawn_named("receiver", move || { - let mut receiver = receiver.into_future(); - let (waker, _waker_probe) = WakerProbe::new(); - let mut context = Context::from_waker(&waker); - - spin_until("poll ready with message", || { - match Pin::new(&mut receiver).poll(&mut context) { - Poll::Ready(Ok(999)) => true, - Poll::Ready(result) => panic!("unexpected result: {result:?}"), - Poll::Pending => false, - } - }); - }); - - let sender_thread = spawn_named("sender", move || { - sender.send(999).unwrap(); - }); - - receiver_thread.join().unwrap(); - sender_thread.join().unwrap(); -} - -#[test] -fn concurrent_drop_sender_and_poll_to_completion() { - let (sender, receiver) = oneshot::channel::(); - - let receiver_thread = spawn_named("receiver", move || { - let mut receiver = receiver.into_future(); - let (waker, _waker_probe) = WakerProbe::new(); - let mut context = Context::from_waker(&waker); - - spin_until("poll ready with disconnect", || { - match Pin::new(&mut receiver).poll(&mut context) { - Poll::Ready(Err(oneshot::RecvError::Disconnected)) => true, - Poll::Ready(result) => panic!("unexpected result: {result:?}"), - Poll::Pending => false, - } - }); - }); - - let sender_thread = spawn_named("sender", move || { - drop(sender); - }); - - receiver_thread.join().unwrap(); - sender_thread.join().unwrap(); -} diff --git a/tests-integration/tests/oneshot_test/support.rs b/tests-integration/tests/oneshot_test/support.rs deleted file mode 100644 index 821930c..0000000 --- a/tests-integration/tests/oneshot_test/support.rs +++ /dev/null @@ -1,121 +0,0 @@ -// 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::hint::spin_loop; -use std::sync::Arc; -use std::sync::atomic::AtomicU32; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Wake; -use std::task::Waker; -use std::time::Duration; -use std::time::Instant; - -pub(super) struct DropProbe { - drop_count: Arc, - value: T, -} - -impl DropProbe { - pub(super) fn new(value: T) -> (Self, Arc) { - let drop_count = Arc::new(AtomicUsize::new(0)); - ( - Self { - drop_count: drop_count.clone(), - value, - }, - drop_count, - ) - } - - pub(super) fn value(&self) -> &T { - &self.value - } -} - -impl Drop for DropProbe { - fn drop(&mut self) { - self.drop_count.fetch_add(1, Ordering::Relaxed); - } -} - -#[derive(Default)] -pub(super) struct WakerProbe { - wake_count: AtomicU32, -} - -impl WakerProbe { - pub(super) fn new() -> (Waker, Arc) { - let probe = Arc::new(Self::default()); - (Waker::from(probe.clone()), probe) - } - - pub(super) fn live_waker_count(this: &Arc) -> usize { - // The returned probe owns one strong reference; every other reference belongs to a live - // Waker created from it. - Arc::strong_count(this) - 1 - } - - pub(super) fn wake_count(&self) -> u32 { - self.wake_count.load(Ordering::Relaxed) - } -} - -impl Wake for WakerProbe { - fn wake(self: Arc) { - self.wake_count.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wake_count.fetch_add(1, Ordering::Relaxed); - } -} - -pub(super) fn spawn_named(name: &str, f: F) -> std::thread::JoinHandle -where - F: FnOnce() -> T + Send + 'static, - T: Send + 'static, -{ - std::thread::Builder::new() - .name(name.to_owned()) - .spawn(f) - .unwrap() -} - -pub(super) fn spin_until(label: &str, mut f: F) -where - F: FnMut() -> bool, -{ - let deadline = Instant::now() + Duration::from_secs(5); - let mut spins = 0usize; - - loop { - if f() { - break; - } - - assert!(Instant::now() < deadline, "timed out waiting for {label}"); - - if spins % 64 == 0 { - std::thread::yield_now(); - } else { - spin_loop(); - } - - spins += 1; - } -} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index e894246..53cb3fc 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -16,16 +16,16 @@ // under the License. use asyncband::barrier::Barrier; -use asyncband::broadcast; +use asyncband::channel::broadcast; +use asyncband::channel::mpsc; +use asyncband::channel::oneshot; use asyncband::condvar::Condvar; use asyncband::latch::Latch; -use asyncband::mpsc; use asyncband::mutex::Mutex; use asyncband::mutex::MutexGuard; use asyncband::once::Once; use asyncband::once::OnceCell; use asyncband::once::OnceMap; -use asyncband::oneshot; use asyncband::rwlock::OwnedRwLockReadGuard; use asyncband::rwlock::RwLock; use asyncband::rwlock::RwLockReadGuard; @@ -67,10 +67,7 @@ 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::>(); } #[test] @@ -79,7 +76,7 @@ fn movable_public_types_are_send() { assert_send::>>(); assert_send::>(); - assert_send::>(); + assert_send::>(); } #[test] @@ -111,10 +108,7 @@ 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::>(); }