diff --git a/CHANGELOG.md b/CHANGELOG.md index f36264e9..13d5213c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,14 @@ All notable changes to this project will be documented in this file. ## Unreleased +### Notable changes + +* Bounded MPMC queues now panic when created with a capacity greater than 2,097,151. + +### Improvements + +* Bounded MPMC sends, receives, and reservations no longer take a lock unless a sender or receiver is waiting, which speeds up contended workloads on threads and async runtimes. Creating a bounded queue now allocates storage for its full capacity, and handoffs where every send waits for a receiver, such as through a capacity-1 queue, are slightly slower. On targets without 64-bit atomics, bounded queues protect their cursors with locks, which can be slower than before. + ## v0.7.4 ### Notable changes diff --git a/Cargo.lock b/Cargo.lock index 07175ce9..7ced7491 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -103,6 +103,7 @@ dependencies = [ "async-broadcast", "async-channel", "asyncband", + "crossbeam-channel", "divan", "flume", "tokio", @@ -260,6 +261,15 @@ dependencies = [ "url", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.22" diff --git a/Cargo.toml b/Cargo.toml index 8689fca6..f906782e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ async-broadcast = { version = "0.7.2" } async-channel = { version = "2.5.0" } cargo_metadata = { version = "0.23.1" } clap = { version = "4.6.5" } +crossbeam-channel = { version = "0.5.17" } divan = { version = "0.1.21" } flume = { version = "0.12.0", default-features = false } semver = { version = "1.0.28" } diff --git a/LICENSE b/LICENSE index 643e18a4..692b2368 100644 --- a/LICENSE +++ b/LICENSE @@ -416,6 +416,22 @@ the following copyright notice: Copyright 2014-2020 The Rust Project Developers +Portions of asyncband/src/internal/cache_padded.rs are adapted from +crossbeam-utils 0.8.23. Portions of the slot-stamp algorithm and backoff helper +in asyncband/src/mpmc/ring.rs are adapted from crossbeam-channel 0.5.17 and +crossbeam-utils 0.8.23. Both are taken from the following exact revision and +source paths: + + https://github.com/crossbeam-rs/crossbeam/blob/2920c984290229ab4e0ca0452ef09e48a82063f3/crossbeam-utils/src/cache_padded.rs + https://github.com/crossbeam-rs/crossbeam/blob/2920c984290229ab4e0ca0452ef09e48a82063f3/crossbeam-utils/src/backoff.rs + https://github.com/crossbeam-rs/crossbeam/blob/2920c984290229ab4e0ca0452ef09e48a82063f3/crossbeam-channel/src/flavors/array.rs + +Crossbeam is licensed under Apache-2.0 or MIT. Apache Asyncband uses the +Apache-2.0 option for the incorporated portions. The upstream project carries +the following copyright notice: + + Copyright (c) 2019 The Crossbeam Project Developers + Portions of the following files are ported from Fastpool 1.1.1. Each local path is followed by its upstream source path or paths: diff --git a/README.md b/README.md index 17ee4426..04ce53e4 100644 --- a/README.md +++ b/README.md @@ -146,6 +146,10 @@ The `blocking` module is a lightweight, thread-parking single-future executor, n Asyncband types implement `Send` and `Sync` only when the protected, transferred, or managed value satisfies the necessary bounds. See each API's documentation for its exact contract. +## Platform support + +Asyncband supports x86, x86-64, ARMv6 or later, and AArch64 targets. It also compiles on targets without 64-bit atomics, such as ARMv5TE, 32-bit PowerPC, and 32-bit MIPS, where bounded MPMC queues protect their cursors with locks and run slower. + ## Minimum Supported Rust Version (MSRV) This crate is built against the latest stable release, and its minimum supported rustc version is 1.86.0. diff --git a/asyncband/src/internal/atomic_u64.rs b/asyncband/src/internal/atomic_u64.rs new file mode 100644 index 00000000..e7fe429f --- /dev/null +++ b/asyncband/src/internal/atomic_u64.rs @@ -0,0 +1,168 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[cfg(target_has_atomic = "64")] +pub use std::sync::atomic::AtomicU64; + +#[cfg(not(target_has_atomic = "64"))] +pub use self::locked::AtomicU64; + +#[cfg(any(test, not(target_has_atomic = "64")))] +mod locked { + use std::sync::atomic::Ordering; + + use crate::internal::mutex::Mutex; + + /// A `u64` with the subset of the `AtomicU64` API that the bounded MPMC ring uses, for targets + /// without native 64-bit atomics. + /// + /// Each operation runs as a critical section on the value's lock, so all operations on a value + /// are totally ordered, and each happens after every operation ordered before it. That provides + /// the acquire, release, and sequentially consistent orderings the ring requests, so the + /// orderings passed in are ignored. + #[derive(Debug)] + pub struct AtomicU64(Mutex); + + impl AtomicU64 { + pub const fn new(value: u64) -> Self { + Self(Mutex::new(value)) + } + + pub fn load(&self, _: Ordering) -> u64 { + *self.0.lock() + } + + pub fn store(&self, value: u64, _: Ordering) { + *self.0.lock() = value; + } + + pub fn compare_exchange( + &self, + current: u64, + new: u64, + _: Ordering, + _: Ordering, + ) -> Result { + let mut value = self.0.lock(); + if *value == current { + *value = new; + Ok(current) + } else { + Err(*value) + } + } + + /// Behaves like `compare_exchange`; it never fails spuriously. + pub fn compare_exchange_weak( + &self, + current: u64, + new: u64, + success: Ordering, + failure: Ordering, + ) -> Result { + self.compare_exchange(current, new, success, failure) + } + + pub fn fetch_add(&self, delta: u64, _: Ordering) -> u64 { + self.update(|value| value.wrapping_add(delta)) + } + + pub fn fetch_sub(&self, delta: u64, _: Ordering) -> u64 { + self.update(|value| value.wrapping_sub(delta)) + } + + pub fn fetch_or(&self, bits: u64, _: Ordering) -> u64 { + self.update(|value| value | bits) + } + + pub fn fetch_and(&self, bits: u64, _: Ordering) -> u64 { + self.update(|value| value & bits) + } + + pub fn get_mut(&mut self) -> &mut u64 { + self.0.get_mut() + } + + /// Replaces the value with `update(value)` and returns the previous value. + fn update(&self, update: impl FnOnce(u64) -> u64) -> u64 { + let mut value = self.0.lock(); + let previous = *value; + *value = update(previous); + previous + } + } + + #[cfg(test)] + mod tests { + use std::thread; + + use super::*; + + #[test] + fn operations_return_the_previous_value_and_wrap() { + let mut atomic = AtomicU64::new(u64::MAX); + assert_eq!(atomic.fetch_add(2, Ordering::SeqCst), u64::MAX); + assert_eq!(atomic.fetch_sub(3, Ordering::SeqCst), 1); + assert_eq!(atomic.load(Ordering::SeqCst), u64::MAX - 1); + assert_eq!(atomic.fetch_and(0xff, Ordering::SeqCst), u64::MAX - 1); + assert_eq!(atomic.fetch_or(1 << 40, Ordering::SeqCst), 0xfe); + assert_eq!( + atomic.compare_exchange(0, 7, Ordering::SeqCst, Ordering::SeqCst), + Err((1 << 40) | 0xfe) + ); + assert_eq!( + atomic.compare_exchange_weak( + (1 << 40) | 0xfe, + 7, + Ordering::SeqCst, + Ordering::SeqCst + ), + Ok((1 << 40) | 0xfe) + ); + atomic.store(9, Ordering::SeqCst); + *atomic.get_mut() += 1; + assert_eq!(atomic.load(Ordering::SeqCst), 10); + } + + #[test] + fn concurrent_updates_are_not_lost() { + const THREADS: u64 = 4; + const UPDATES: u64 = if cfg!(miri) { 20 } else { 10_000 }; + + let atomic = AtomicU64::new(0); + thread::scope(|scope| { + for _ in 0..THREADS { + scope.spawn(|| { + for _ in 0..UPDATES { + atomic.fetch_add(3, Ordering::SeqCst); + let mut current = atomic.load(Ordering::SeqCst); + while let Err(actual) = atomic.compare_exchange_weak( + current, + current - 1, + Ordering::SeqCst, + Ordering::SeqCst, + ) { + current = actual; + } + } + }); + } + }); + assert_eq!(atomic.load(Ordering::SeqCst), THREADS * UPDATES * 2); + } + } +} diff --git a/asyncband/src/internal/cache_padded.rs b/asyncband/src/internal/cache_padded.rs new file mode 100644 index 00000000..9a05937a --- /dev/null +++ b/asyncband/src/internal/cache_padded.rs @@ -0,0 +1,181 @@ +// 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 file contains code adapted from crossbeam-utils 0.8.23. +// Copyright (c) 2019 The Crossbeam Project Developers +// Asyncband uses the upstream project's Apache-2.0 license option for the incorporated code. +// The incorporated code has been modified for use in Apache Asyncband: examples, conversions, and +// explicit Send and Sync implementations are omitted, and imports follow Asyncband's style. +// Upstream source: +// https://github.com/crossbeam-rs/crossbeam/blob/2920c984290229ab4e0ca0452ef09e48a82063f3/crossbeam-utils/src/cache_padded.rs + +use std::fmt; +use std::ops::Deref; +use std::ops::DerefMut; + +/// Pads and aligns a value to the length of a cache line. +/// +/// In concurrent programming, sometimes it is desirable to make sure commonly accessed pieces of +/// data are not placed into the same cache line. Updating an atomic value invalidates the whole +/// cache line it belongs to, which makes the next access to the same cache line slower for other +/// CPU cores. Use `CachePadded` to ensure updating one piece of data doesn't invalidate other +/// cached data. +/// +/// # Size and alignment +/// +/// Cache lines are assumed to be N bytes long, depending on the architecture: +/// +/// * On x86-64, aarch64, and powerpc64, N = 128. +/// * On arm, mips, mips64, sparc, and hexagon, N = 32. +/// * On m68k, N = 16. +/// * On s390x, N = 256. +/// * On all others, N = 64. +/// +/// Note that N is just a reasonable guess and is not guaranteed to match the actual cache line +/// length of the machine the program is running on. On modern Intel architectures, spatial +/// prefetcher is pulling pairs of 64-byte cache lines at a time, so we pessimistically assume that +/// cache lines are 128 bytes long. +/// +/// The size of `CachePadded` is the smallest multiple of N bytes large enough to accommodate +/// a value of type `T`. +/// +/// The alignment of `CachePadded` is the maximum of N bytes and the alignment of `T`. +/// +/// # Layout +/// +/// This type is `#[repr(C)]`, so a pointer to `CachePadded` has the same address as the pointer +/// to the underlying `T`. +#[derive(Clone, Copy, Default, Hash, PartialEq, Eq)] +// Starting from Intel's Sandy Bridge, spatial prefetcher is now pulling pairs of 64-byte cache +// lines at a time, so we have to align to 128 bytes rather than 64. +// +// Sources: +// - https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf +// - https://github.com/facebook/folly/blob/1b5288e6eea6df074758f877c849b6e73bbb9fbb/folly/lang/Align.h#L107 +// +// aarch64/arm64ec's big.LITTLE architecture has asymmetric cores and "big" cores have 128-byte +// cache line size. +// +// Sources: +// - https://www.mono-project.com/news/2016/09/12/arm64-icache/ +// +// powerpc64 has 128-byte cache line size. +// +// Sources: +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_ppc64x.go#L9 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/powerpc/include/asm/cache.h#L26 +#[cfg_attr( + any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "arm64ec", + target_arch = "powerpc64", + ), + repr(align(128)) +)] +// arm, mips, mips64, sparc, and hexagon have 32-byte cache line size. +// +// Sources: +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_arm.go#L7 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips.go#L7 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mipsle.go#L7 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_mips64x.go#L9 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/sparc/include/asm/cache.h#L17 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/hexagon/include/asm/cache.h#L12 +#[cfg_attr( + any( + target_arch = "arm", + target_arch = "mips", + target_arch = "mips32r6", + target_arch = "mips64", + target_arch = "mips64r6", + target_arch = "sparc", + target_arch = "hexagon", + ), + repr(align(32)) +)] +// m68k has 16-byte cache line size. +// +// Sources: +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/m68k/include/asm/cache.h#L9 +#[cfg_attr(target_arch = "m68k", repr(align(16)))] +// s390x has 256-byte cache line size. +// +// Sources: +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_s390x.go#L7 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/s390/include/asm/cache.h#L13 +#[cfg_attr(target_arch = "s390x", repr(align(256)))] +// x86, wasm, riscv, and sparc64 have 64-byte cache line size. +// +// Sources: +// - https://github.com/golang/go/blob/dda2991c2ea0c5914714469c4defc2562a907230/src/internal/cpu/cpu_x86.go#L9 +// - https://github.com/golang/go/blob/3dd58676054223962cd915bb0934d1f9f489d4d2/src/internal/cpu/cpu_wasm.go#L7 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/riscv/include/asm/cache.h#L10 +// - https://github.com/torvalds/linux/blob/3516bd729358a2a9b090c1905bd2a3fa926e24c6/arch/sparc/include/asm/cache.h#L19 +// +// All others are assumed to have 64-byte cache line size. +#[cfg_attr( + not(any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "arm64ec", + target_arch = "powerpc64", + target_arch = "arm", + target_arch = "mips", + target_arch = "mips32r6", + target_arch = "mips64", + target_arch = "mips64r6", + target_arch = "sparc", + target_arch = "hexagon", + target_arch = "m68k", + target_arch = "s390x", + )), + repr(align(64)) +)] +#[repr(C)] +pub struct CachePadded { + value: T, +} + +impl CachePadded { + /// Pads and aligns a value to the length of a cache line. + pub const fn new(t: T) -> CachePadded { + CachePadded:: { value: t } + } +} + +impl Deref for CachePadded { + type Target = T; + + fn deref(&self) -> &T { + &self.value + } +} + +impl DerefMut for CachePadded { + fn deref_mut(&mut self) -> &mut T { + &mut self.value + } +} + +impl fmt::Debug for CachePadded { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CachePadded") + .field("value", &self.value) + .finish() + } +} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index ccaa11ac..54e4a2bf 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -54,6 +54,12 @@ pub fn register_waker(slot: &mut Option, waker: &Waker) -> Option #[allow(dead_code)] pub(crate) mod arena; +#[cfg(feature = "mpmc")] +pub(crate) mod atomic_u64; + +#[cfg(feature = "mpmc")] +pub(crate) mod cache_padded; + #[cfg(any(feature = "latch", feature = "once"))] pub(crate) mod countdown; @@ -116,6 +122,7 @@ pub(crate) mod waitlist; feature = "completion", feature = "event", feature = "latch", + feature = "mpmc", feature = "mutex", feature = "once", feature = "phaser", diff --git a/asyncband/src/internal/waker_batch.rs b/asyncband/src/internal/waker_batch.rs index af27db78..25e9aa05 100644 --- a/asyncband/src/internal/waker_batch.rs +++ b/asyncband/src/internal/waker_batch.rs @@ -75,7 +75,9 @@ impl WakerBatch { // SAFETY: This prefix was initialized by `push` and is no longer owned by the batch. unsafe { slot.assume_init_read() }.wake(); } - self.spilled.drain(..).for_each(Waker::wake); + if !self.spilled.is_empty() { + self.spilled.drain(..).for_each(Waker::wake); + } } } diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index f2003d85..9207bb19 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -112,6 +112,12 @@ //! Asyncband types implement `Send` and `Sync` only when their protected, transferred, or managed //! values satisfy the required bounds. Consult each API's documentation for its exact contract. //! +//! # Platform support +//! +//! Asyncband supports x86, x86-64, ARMv6 or later, and AArch64 targets. It also compiles on targets +//! without 64-bit atomics, such as ARMv5TE, 32-bit PowerPC, and 32-bit MIPS, where bounded MPMC +//! queues protect their cursors with locks and run slower. +//! //! # Disclaimer //! //! Apache Asyncband (Incubating) is an effort undergoing incubation at the Apache Software @@ -123,6 +129,7 @@ //! //! While incubation status is not necessarily a reflection of the completeness or stability of the //! code, it does indicate that the project has yet to be fully endorsed by the ASF. + mod internal; #[cfg(feature = "barrier")] diff --git a/asyncband/src/mpmc/bounded.rs b/asyncband/src/mpmc/bounded.rs index 84f19e23..e1d58d3b 100644 --- a/asyncband/src/mpmc/bounded.rs +++ b/asyncband/src/mpmc/bounded.rs @@ -22,25 +22,26 @@ use crate::mpmc::RecvError; use crate::mpmc::SendError; use crate::mpmc::TryRecvError; use crate::mpmc::TrySendError; -use crate::mpmc::queue::Permit; -use crate::mpmc::queue::Shared; +use crate::mpmc::ring::Permit; +use crate::mpmc::ring::Shared; /// Creates a bounded multi-producer, multi-consumer queue. /// /// Queued values, held permits, and capacity granted to waiting senders occupy at most `capacity` /// slots. Pending `send` and `reserve` operations receive capacity in wait-queue order. Sending -/// waits for a receiver to free capacity when none is available. +/// waits for a receiver to free capacity when none is available. Storage for `capacity` values is +/// allocated when the queue is created. /// /// The `try_*` methods do not wait for capacity or messages, but may briefly block on an internal -/// mutex. +/// mutex or spin while another sender or receiver finishes with the same slot. /// /// # Panics /// -/// Panics if `capacity` is zero. +/// Panics if `capacity` is zero or greater than 2,097,151. #[track_caller] pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { assert!(capacity > 0, "mpmc bounded queue requires capacity > 0"); - let shared = Arc::new(Shared::bounded(capacity)); + let shared = Arc::new(Shared::new(capacity)); ( BoundedSender { shared: shared.clone(), diff --git a/asyncband/src/mpmc/mod.rs b/asyncband/src/mpmc/mod.rs index ba484949..e7be2a6d 100644 --- a/asyncband/src/mpmc/mod.rs +++ b/asyncband/src/mpmc/mod.rs @@ -25,7 +25,9 @@ mod bounded; mod error; mod queue; +mod ring; mod unbounded; +mod waiters; pub use self::bounded::BoundedReceiver; pub use self::bounded::BoundedSender; @@ -34,7 +36,7 @@ pub use self::error::RecvError; pub use self::error::SendError; pub use self::error::TryRecvError; pub use self::error::TrySendError; -pub use self::queue::Permit; +pub use self::ring::Permit; pub use self::unbounded::UnboundedReceiver; pub use self::unbounded::UnboundedSender; pub use self::unbounded::unbounded; diff --git a/asyncband/src/mpmc/queue.rs b/asyncband/src/mpmc/queue.rs index 02bb9d7a..47731453 100644 --- a/asyncband/src/mpmc/queue.rs +++ b/asyncband/src/mpmc/queue.rs @@ -16,12 +16,10 @@ // under the License. use std::collections::VecDeque; -use std::fmt; use std::future::poll_fn; use std::mem; use std::task::Context; use std::task::Poll; -use std::task::Waker; use crate::internal::mutex::Mutex; use crate::internal::waitlist::WaitList; @@ -29,68 +27,28 @@ use crate::internal::waitlist::WaiterId; use crate::mpmc::RecvError; use crate::mpmc::SendError; use crate::mpmc::TryRecvError; -use crate::mpmc::TrySendError; +use crate::mpmc::waiters::RecvWaiter; +use crate::mpmc::waiters::remove_waiter; +/// The unbounded queue's shared state. pub struct Shared { state: Mutex>, } -/// Values, endpoint counts, and both waiter queues share one lock, so each transition and the +/// Values, endpoint counts, and the receiver wait queue share one lock, so each transition and the /// waiter it selects are decided together. Wake callbacks and waker or value destructors run /// outside the lock, because they may reenter the queue. struct State { values: VecDeque, - // None is an unbounded queue. In a bounded queue, each slot belongs to available, - // a queued value, a Permit, or a detached waiter with a grant. - available: Option, senders: usize, receivers: usize, recv_waiters: WaitList, - send_waiters: WaitList, } impl State { - fn has_capacity(&self) -> bool { - self.available.is_none_or(|available| available != 0) - } - - fn acquire(&mut self) -> Result<(), TrySendError<()>> { - if self.receivers == 0 { - return Err(TrySendError::Disconnected(())); - } - if !self.has_capacity() { - return Err(TrySendError::Full(())); - } - if let Some(available) = &mut self.available { - *available -= 1; - } - Ok(()) - } - - fn release(&mut self) -> Option { - if self.receivers == 0 { - return None; - } - let Some(available) = &mut self.available else { - return None; - }; - if let Some(waker) = self.send_waiters.grant_one() { - return Some(waker); - } - *available += 1; - None - } - - /// Queues a value and selects a waiting receiver's waker. - fn push(&mut self, value: T) -> Option { - self.values.push_back(value); - self.recv_waiters.notify_one() - } - - /// Takes the next value and grants its capacity to the first waiting operation. - fn pop(&mut self) -> Result<(T, Option), TryRecvError> { + fn pop(&mut self) -> Result { if let Some(value) = self.values.pop_front() { - Ok((value, self.release())) + Ok(value) } else if self.senders == 0 { Err(TryRecvError::Disconnected) } else { @@ -99,107 +57,14 @@ impl State { } } -/// A receiver notification makes the task runnable without reserving a value. -enum RecvWaiter { - Waiting(Waker), - Notified, -} - -/// A grant assigns capacity to a detached waiter node until its future claims or releases it. -enum SendWaiter { - Waiting(Waker), - Granted, -} - -fn remove_waiter(waiters: &mut WaitList, id: WaiterId) -> W { - // Unlinking is idempotent, so detached waiters are removed the same way as linked ones. - waiters.unlink_waiter(id, |_| true); - waiters.remove_unlinked_waiter(id) -} - -impl WaitList { - fn notify_one(&mut self) -> Option { - let (_, waiter) = self.unlink_first_waiter(|_| true)?; - let RecvWaiter::Waiting(waker) = mem::replace(waiter, RecvWaiter::Notified) else { - unreachable!("only waiting operations remain linked"); - }; - Some(waker) - } - - /// Registers a pending operation's waker or refreshes an existing registration. - /// - /// If the future finds no value after notification, its waiter rejoins the queue. - #[must_use = "drop the replaced waker after releasing the queue lock"] - fn register(&mut self, id: &mut Option, current: &Waker) -> Option { - if let Some(queued) = *id { - if let RecvWaiter::Waiting(waker) = self.waiter_mut(queued) { - if waker.will_wake(current) { - return None; - } - return Some(mem::replace(waker, current.clone())); - } - } - if let Some(notified) = id.take() { - // The notification already took this node's waker, so nothing is retired. - remove_waiter(self, notified); - } - *id = Some(self.push_back(RecvWaiter::Waiting(current.clone()))); - None - } -} - -impl WaitList { - /// Registers an operation waiting for capacity or refreshes its waker without losing its place. - #[must_use = "drop the replaced waker after releasing the queue lock"] - fn register_waiter(&mut self, id: &mut Option, current: &Waker) -> Option { - if let Some(queued) = *id { - let SendWaiter::Waiting(waker) = self.waiter_mut(queued) else { - unreachable!("a granted waiter must be claimed before registration"); - }; - if waker.will_wake(current) { - return None; - } - return Some(mem::replace(waker, current.clone())); - } - *id = Some(self.push_back(SendWaiter::Waiting(current.clone()))); - None - } - - fn grant_one(&mut self) -> Option { - let (_, waiter) = self.unlink_first_waiter(|_| true)?; - let SendWaiter::Waiting(waker) = mem::replace(waiter, SendWaiter::Granted) else { - unreachable!("only waiting operations remain linked"); - }; - Some(waker) - } - - fn take_waiting_waker(&mut self) -> Option { - let (id, _) = self.unlink_first_waiter(|_| true)?; - let SendWaiter::Waiting(waker) = self.remove_unlinked_waiter(id) else { - unreachable!("only waiting operations remain linked"); - }; - Some(waker) - } -} - impl Shared { - pub fn bounded(capacity: usize) -> Self { - Self::new(Some(capacity)) - } - pub fn unbounded() -> Self { - Self::new(None) - } - - fn new(capacity: Option) -> Self { Self { state: Mutex::new(State { values: VecDeque::new(), - available: capacity, senders: 1, receivers: 1, recv_waiters: WaitList::new(), - send_waiters: WaitList::new(), }), } } @@ -229,35 +94,25 @@ impl Shared { } pub fn drop_receiver(&self) { - let (discarded, mut waiters) = { + let discarded = { let mut state = self.state.lock(); state.receivers -= 1; if state.receivers != 0 { return; } - ( - mem::take(&mut state.values), - mem::replace(&mut state.send_waiters, WaitList::new()), - ) + mem::take(&mut state.values) }; - // Notify blocked senders before destroying buffered values, whose destructors may panic. - while let Some(waker) = waiters.take_waiting_waker() { - waker.wake(); - } drop(discarded); } - pub fn try_send(&self, value: T) -> Result<(), TrySendError> { + pub fn send(&self, value: T) -> Result<(), SendError> { let waker = { let mut state = self.state.lock(); - match state.acquire() { - Ok(()) => {} - Err(TrySendError::Full(())) => return Err(TrySendError::Full(value)), - Err(TrySendError::Disconnected(())) => { - return Err(TrySendError::Disconnected(value)); - } + if state.receivers == 0 { + return Err(SendError::new(value)); } - state.push(value) + state.values.push_back(value); + state.recv_waiters.notify_one() }; if let Some(waker) = waker { waker.wake(); @@ -265,37 +120,8 @@ impl Shared { Ok(()) } - 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, - }; - match self.reserve().await { - Ok(permit) => permit.send(value), - Err(_) => Err(SendError::new(value)), - } - } - - pub fn try_reserve(&self) -> Result, TrySendError<()>> { - self.state.lock().acquire()?; - Ok(Permit { shared: self }) - } - - pub async fn reserve(&self) -> Result, SendError<()>> { - let mut reserve = Reserve { - shared: self, - waiter: None, - }; - poll_fn(|cx| reserve.poll(cx)).await - } - pub fn try_recv(&self) -> Result { - let (value, waker) = self.state.lock().pop()?; - if let Some(waker) = waker { - waker.wake(); - } - Ok(value) + self.state.lock().pop() } pub async fn recv(&self) -> Result { @@ -312,122 +138,6 @@ impl Shared { } } -/// Capacity reserved for one value on a bounded MPMC queue. -/// -/// A permit does not keep any receiver alive or claim message order. Dropping it without sending -/// passes its capacity to the next waiting sender or makes it available again. -#[must_use = "dropping the permit releases its reserved capacity"] -pub struct Permit<'a, T> { - shared: &'a Shared, -} - -impl fmt::Debug for Permit<'_, T> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Permit").finish_non_exhaustive() - } -} - -impl Permit<'_, T> { - /// Enqueues a value using the reserved capacity without waiting for space. - /// - /// If the last receiver has been dropped, the error returns ownership of the value. - pub fn send(self, value: T) -> Result<(), SendError> { - let waker = { - let mut state = self.shared.state.lock(); - if state.receivers == 0 { - return Err(SendError::new(value)); - } - let waker = state.push(value); - // The queued value now owns this capacity. - mem::forget(self); - waker - }; - if let Some(waker) = waker { - waker.wake(); - } - Ok(()) - } -} - -impl Drop for Permit<'_, T> { - fn drop(&mut self) { - let waker = self.shared.state.lock().release(); - if let Some(waker) = waker { - waker.wake(); - } - } -} - -struct Reserve<'a, T> { - shared: &'a Shared, - waiter: Option, -} - -impl<'a, T> Reserve<'a, T> { - fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { - let mut state = self.shared.state.lock(); - // The last receiver may have dropped while this reservation was waiting or granted. - if state.receivers == 0 { - self.waiter = None; - return Poll::Ready(Err(SendError::new(()))); - } - // A previous poll queued this reservation; a released slot may now belong to it. - if let Some(id) = self.waiter { - if matches!(state.send_waiters.waiter_mut(id), SendWaiter::Granted) { - let retired = state.send_waiters.remove_unlinked_waiter(id); - self.waiter = None; - drop(state); - drop(retired); - return Poll::Ready(Ok(Permit { - shared: self.shared, - })); - } - } else if state.has_capacity() { - // This reservation has not queued yet and can claim unassigned capacity immediately. - if let Some(available) = &mut state.available { - *available -= 1; - } - drop(state); - return Poll::Ready(Ok(Permit { - shared: self.shared, - })); - } - // Otherwise, wait for a slot or refresh the waker of the queued reservation. - let retired = state - .send_waiters - .register_waiter(&mut self.waiter, cx.waker()); - drop(state); - drop(retired); - Poll::Pending - } -} - -impl Drop for Reserve<'_, T> { - fn drop(&mut self) { - let Some(id) = self.waiter.take() else { - return; - }; - let (retired, waker) = { - let mut state = self.shared.state.lock(); - if state.receivers == 0 { - return; - } - let retired = remove_waiter(&mut state.send_waiters, id); - // A grant already owns a slot, so returning it does not require a capacity check. - let waker = if matches!(retired, SendWaiter::Granted) { - state.release() - } else { - None - }; - (retired, waker) - }; - if let Some(waker) = waker { - waker.wake(); - } - drop(retired); - } -} - struct Recv<'a, T> { shared: &'a Shared, waiter: Option, @@ -441,7 +151,7 @@ impl Recv<'_, T> { self.waiter = None; } let outcome = match state.pop() { - Ok(popped) => Ok(popped), + Ok(value) => Ok(value), Err(TryRecvError::Disconnected) => Err(RecvError::Disconnected), Err(TryRecvError::Empty) => { let retired = state.recv_waiters.register(&mut self.waiter, cx.waker()); @@ -455,14 +165,8 @@ impl Recv<'_, T> { .take() .map(|id| remove_waiter(&mut state.recv_waiters, id)); drop(state); - let result = outcome.map(|(value, waker)| { - if let Some(waker) = waker { - waker.wake(); - } - value - }); drop(retired); - Poll::Ready(result) + Poll::Ready(outcome) } } diff --git a/asyncband/src/mpmc/ring.rs b/asyncband/src/mpmc/ring.rs new file mode 100644 index 00000000..e6635889 --- /dev/null +++ b/asyncband/src/mpmc/ring.rs @@ -0,0 +1,1485 @@ +// 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. + +// Portions of the slot-stamp algorithm are adapted from the array flavor of crossbeam-channel +// 0.5.17, which implements Dmitry Vyukov's bounded MPMC queue, and the backoff helper follows +// crossbeam-utils 0.8.23. +// Copyright (c) 2019 The Crossbeam Project Developers +// Asyncband uses the upstream project's Apache-2.0 license option for the incorporated portions. +// Asyncband adds reservations that hold capacity without occupying a slot, grants capacity to +// waiting senders in queue order, and parks receivers in its own wait list. +// Upstream sources: +// https://github.com/crossbeam-rs/crossbeam/blob/2920c984290229ab4e0ca0452ef09e48a82063f3/crossbeam-channel/src/flavors/array.rs +// https://github.com/crossbeam-rs/crossbeam/blob/2920c984290229ab4e0ca0452ef09e48a82063f3/crossbeam-utils/src/backoff.rs + +use std::cell::UnsafeCell; +use std::fmt; +use std::future::poll_fn; +use std::hint; +use std::marker::PhantomPinned; +use std::mem; +use std::mem::MaybeUninit; +use std::ptr::NonNull; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; +use std::thread; + +use crate::internal::atomic_u64::AtomicU64; +use crate::internal::cache_padded::CachePadded; +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaitList; +use crate::internal::waitlist::WaiterId; +use crate::internal::waker_batch::WakerBatch; +use crate::mpmc::RecvError; +use crate::mpmc::SendError; +use crate::mpmc::TryRecvError; +use crate::mpmc::TrySendError; +use crate::mpmc::waiters::RecvWaiter; +use crate::mpmc::waiters::remove_waiter; + +// Positions are 64 bits on every target. Targets without native 64-bit atomics protect each cursor +// and stamp with a lock instead, which is slower but keeps the same algorithm. +type AtomicIndex = AtomicU64; +type Index = u64; + +/// The largest capacity whose tail keeps enough lap bits that an exchange prepared from an earlier +/// tail can succeed only after at least 2^40 sends. +const MAX_CAPACITY: usize = (1 << 21) - 1; + +/// The bounded queue's shared state. +/// +/// Values move through a ring of stamped slots without a lock. A position packs a slot index into +/// its low bits and a lap count into its high bits. A slot's stamp equals the position that may +/// write it next, or that position plus one once its value is published. Senders advance `tail` and +/// receivers advance `head`, so each cursor is written by one side only. +/// +/// Between the index and the lap, `tail` carries `mark_bit` once either side disconnects, the +/// capacity that permits and grants to waiting senders hold without a slot, and `queued_bit` while +/// senders wait. Every exchange on `tail` therefore compares the reservation state it was computed +/// from. The other values it depends on, `head` and the slot stamps, only advance, so reading them +/// early can only understate free capacity. An exchange prepared from an earlier tail succeeds only +/// once the lap count wraps, which takes 2^54 sends at a capacity of 64 and at least 2^40 sends up +/// to `MAX_CAPACITY`. +pub struct Shared { + head: CachePadded, + tail: CachePadded, + send_waiters: CachePadded, + recv_waiters: CachePadded, + slots: Box<[Slot]>, + mark_bit: Index, + /// The shift of the reserved count, which fills the bits between `mark_bit` and `queued_bit`. + count_shift: u32, + queued_bit: Index, + one_lap: Index, + sender_count: AtomicUsize, + receiver_count: AtomicUsize, +} + +// SAFETY: Values move into and out of slots by value and are never shared by reference. Each stamp +// hands its slot to exactly one sender or receiver at a time, so sharing the queue only requires +// `T: Send`. +unsafe impl Sync for Shared {} + +struct Slot { + stamp: AtomicIndex, + value: UnsafeCell>, +} + +struct SendWaiters { + /// Mirrors `queued_bit`, so receivers check for queued senders without loading `tail`, which + /// senders write on every send. + queued: AtomicBool, + /// Counts sends with a reservation that began while senders were queued and have not yet + /// checked for capacity to grant. Receivers leave grants to these sends. + handoffs: AtomicUsize, + state: Mutex, +} + +struct SendState { + waiters: WaitList, + /// Set when the last receiver disconnects. The waiter storage is replaced then, so earlier + /// waiter IDs no longer refer to it. + closed: bool, +} + +/// A sender waiting for capacity. A grant removes the node and sets the flag it points to. +struct SendWaiter { + waker: Waker, + /// Points into the pinned `Reserve` that owns this node. It is written only under the + /// wait-queue lock, and the `Reserve` removes its node under that lock before it is dropped. + granted: NonNull, +} + +// SAFETY: `granted` is dereferenced only under the wait-queue lock while its `Reserve` is alive, +// and `AtomicBool` is `Sync`. +unsafe impl Send for SendWaiter {} + +struct RecvWaiters { + /// Whether a receiver is linked in `state`, so senders can skip the lock otherwise. + linked: AtomicBool, + state: Mutex, +} + +struct RecvState { + waiters: WaitList, + /// Set when the last sender disconnects. The waiter storage is replaced then, so earlier + /// waiter IDs no longer refer to it. + closed: bool, +} + +/// Why a send without a reservation did not claim a slot. +enum Unclaimed { + /// No unreserved capacity remains, or queued senders own the capacity that frees next. + Full(T), + Disconnected(T), +} + +/// What a receiver would find without taking a value. +#[derive(Clone, Copy, PartialEq, Eq)] +enum Peek { + /// A value is published, or a sender has claimed a slot and is writing it. + Ready, + Empty, + Disconnected, +} + +/// Returns `mark_bit`, the unit of the reserved count, `queued_bit`, and `one_lap` for a capacity. +/// +/// From the low bits up, a tail holds the slot index, `mark_bit`, a reserved count as wide as the +/// index, `queued_bit`, and the lap count. +fn tail_layout(capacity: usize) -> (Index, Index, Index, Index) { + assert!( + capacity <= MAX_CAPACITY, + "mpmc bounded queue capacity is too large" + ); + let mark_bit = (capacity as Index + 1).next_power_of_two(); + let count_unit = mark_bit << 1; + let queued_bit = count_unit * mark_bit; + (mark_bit, count_unit, queued_bit, queued_bit << 1) +} + +impl Shared { + pub fn new(capacity: usize) -> Self { + let (mark_bit, count_unit, queued_bit, one_lap) = tail_layout(capacity); + let slots = (0..capacity) + .map(|index| Slot { + stamp: AtomicIndex::new(index as Index), + value: UnsafeCell::new(MaybeUninit::uninit()), + }) + .collect(); + Self { + head: CachePadded::new(AtomicIndex::new(0)), + tail: CachePadded::new(AtomicIndex::new(0)), + send_waiters: CachePadded::new(SendWaiters { + queued: AtomicBool::new(false), + handoffs: AtomicUsize::new(0), + state: Mutex::new(SendState { + waiters: WaitList::new(), + closed: false, + }), + }), + recv_waiters: CachePadded::new(RecvWaiters { + linked: AtomicBool::new(false), + state: Mutex::new(RecvState { + waiters: WaitList::new(), + closed: false, + }), + }), + slots, + mark_bit, + count_shift: count_unit.trailing_zeros(), + queued_bit, + one_lap, + sender_count: AtomicUsize::new(1), + receiver_count: AtomicUsize::new(1), + } + } + + pub fn clone_sender(&self) { + self.sender_count.fetch_add(1, Ordering::Relaxed); + } + + pub fn drop_sender(&self) { + if self.sender_count.fetch_sub(1, Ordering::AcqRel) != 1 { + return; + } + self.tail.fetch_or(self.mark_bit, Ordering::SeqCst); + let mut waiters = { + let mut state = self.recv_waiters.state.lock(); + state.closed = true; + self.recv_waiters.linked.store(false, Ordering::SeqCst); + // Disconnection invalidates every receiver waiter ID. Move the storage out so both + // notification and reclamation happen without holding the lock. + mem::replace(&mut state.waiters, WaitList::new()) + }; + while let Some(waker) = waiters.notify_one() { + waker.wake(); + } + } + + pub fn clone_receiver(&self) { + self.receiver_count.fetch_add(1, Ordering::Relaxed); + } + + pub fn drop_receiver(&self) { + // AcqRel orders every other receiver's slot accesses before the discard below. + if self.receiver_count.fetch_sub(1, Ordering::AcqRel) != 1 { + return; + } + let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst); + let mut waiters = { + let mut state = self.send_waiters.state.lock(); + state.closed = true; + // As in `drop_sender`, waiter IDs no longer refer to the replaced storage. + mem::replace(&mut state.waiters, WaitList::new()) + }; + // Notify blocked senders before destroying buffered values, whose destructors may panic. + // Granted senders were already woken, and their nodes removed. + while let Some((id, _)) = waiters.unlink_first_waiter(|_| true) { + waiters.remove_unlinked_waiter(id).waker.wake(); + } + drop(waiters); + self.discard_values(self.position(tail)); + } + + pub fn try_send(&self, value: T) -> Result<(), TrySendError> { + match self.push_unreserved(value) { + Ok(()) => { + self.notify_receiver(); + Ok(()) + } + Err(Unclaimed::Full(value)) => Err(TrySendError::Full(value)), + Err(Unclaimed::Disconnected(value)) => Err(TrySendError::Disconnected(value)), + } + } + + 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, + }; + match self.reserve().await { + Ok(permit) => permit.send(value), + Err(_) => Err(SendError::new(value)), + } + } + + pub fn try_reserve(&self) -> Result, TrySendError<()>> { + self.reserve_unqueued()?; + Ok(Permit { shared: self }) + } + + pub async fn reserve(&self) -> Result, SendError<()>> { + // The async state pins `reserve`, which a queued node points into. + let mut reserve = Reserve::new(self); + poll_fn(|cx| reserve.poll(cx)).await + } + + pub fn try_recv(&self) -> Result { + self.pop() + } + + pub async fn recv(&self) -> Result { + match self.pop() { + Ok(value) => return Ok(value), + Err(TryRecvError::Disconnected) => return Err(RecvError::Disconnected), + Err(TryRecvError::Empty) => {} + } + let mut recv = Recv { + shared: self, + waiter: None, + }; + poll_fn(|cx| recv.poll(cx)).await + } + + /// Reports whether either side has disconnected. The tail exchanges are authoritative. + fn is_disconnected(&self) -> bool { + self.tail.load(Ordering::Relaxed) & self.mark_bit != 0 + } + + /// Strips the mark bit and the reservation state from a tail. + fn position(&self, tail: Index) -> Index { + tail & !(self.one_lap - self.mark_bit) + } + + /// Replaces the position in a tail, keeping its mark bit and reservation state. + fn with_position(&self, tail: Index, position: Index) -> Index { + (tail & (self.one_lap - self.mark_bit)) | position + } + + /// Returns the capacity that permits and grants hold, according to a tail. + fn reserved(&self, tail: Index) -> usize { + ((tail & (self.queued_bit - 1)) >> self.count_shift) as usize + } + + /// Returns one unit of the reserved count in `tail`. + fn count_unit(&self) -> Index { + 1 << self.count_shift + } + + fn slot(&self, position: Index) -> &Slot { + &self.slots[(position & (self.mark_bit - 1)) as usize] + } + + fn next_position(&self, position: Index) -> Index { + let index = position & (self.mark_bit - 1); + if index + 1 < self.slots.len() as Index { + position + 1 + } else { + // Wrap to the first slot of the next lap. + (position & !(self.one_lap - 1)).wrapping_add(self.one_lap) + } + } + + /// Counts the values between `head` and the tail position. + /// + /// Returns `None` if `head` was read after receivers passed the tail position. + fn queued_len(&self, tail: Index, head: Index) -> Option { + let capacity = self.slots.len() as Index; + let laps = tail.wrapping_sub(head).wrapping_add(capacity) / self.one_lap; + let len = (laps * capacity + (tail & (self.mark_bit - 1))) + .checked_sub(head & (self.mark_bit - 1))?; + (laps <= 1 && len <= capacity).then_some(len as usize) + } + + /// Reserves one unit of capacity for a sender while none is queued. + fn reserve_unqueued(&self) -> Result<(), TrySendError<()>> { + let mut backoff = Backoff::new(); + let mut tail = self.tail.load(Ordering::Relaxed); + loop { + if tail & self.mark_bit != 0 { + return Err(TrySendError::Disconnected(())); + } + if tail & self.queued_bit != 0 { + // Queued senders own released capacity first. + return Err(TrySendError::Full(())); + } + // Reading `head` after `tail` can only overstate the queued values. + match self.queued_len(self.position(tail), self.head.load(Ordering::SeqCst)) { + Some(len) if len + self.reserved(tail) < self.slots.len() => {} + Some(_) => return Err(TrySendError::Full(())), + None => { + // Receivers passed this tail, so the exchange below would fail. + backoff.spin(); + tail = self.tail.load(Ordering::Relaxed); + continue; + } + } + match self.tail.compare_exchange_weak( + tail, + tail + self.count_unit(), + Ordering::SeqCst, + Ordering::Relaxed, + ) { + Ok(_) => return Ok(()), + Err(current) => { + tail = current; + backoff.spin(); + } + } + } + } + + /// Reports whether capacity remains for a grant, with the wait-queue lock held while senders + /// are queued. + fn can_grant(&self) -> bool { + // While senders are queued, only grants under this lock add to the reserved count, and + // sends without a reservation cannot claim slots, so capacity seen free stays free. + loop { + let tail = self.tail.load(Ordering::SeqCst); + if let Some(len) = + self.queued_len(self.position(tail), self.head.load(Ordering::SeqCst)) + { + return len + self.reserved(tail) < self.slots.len(); + } + // Receivers passed this tail after it was loaded. + hint::spin_loop(); + } + } + + /// Queues a sender, then grants it capacity if some is free. + fn enqueue_sender( + &self, + state: &mut SendState, + waiter: &mut Option, + waker: &Waker, + granted: NonNull, + wakers: &mut WakerBatch, + ) { + if state.waiters.is_empty() { + // From here on, sends and reservations without a place in the queue leave freed + // capacity to it. Those prepared from an earlier tail fail their exchange. + self.send_waiters.queued.store(true, Ordering::SeqCst); + self.tail.fetch_or(self.queued_bit, Ordering::SeqCst); + } + *waiter = Some(state.waiters.push_back(SendWaiter { + waker: waker.clone(), + granted, + })); + // Pairs with `grant_after_receive` and `release_reservation`: either they observe queued + // senders, or this check observes the capacity they freed. + self.grant_queued(state, wakers); + } + + /// Grants free capacity to queued senders in order, and clears `queued_bit` once none remain. + fn grant_queued(&self, state: &mut SendState, wakers: &mut WakerBatch) { + while !state.waiters.is_empty() && self.can_grant() { + let (id, _) = state + .waiters + .unlink_first_waiter(|_| true) + .expect("the wait queue is not empty"); + let SendWaiter { waker, granted } = state.waiters.remove_unlinked_waiter(id); + // Count the grant before the sender can observe it and spend it. Granting the last + // queued sender clears `queued_bit` in the same update. + let update = if state.waiters.is_empty() { + self.send_waiters.queued.store(false, Ordering::Relaxed); + self.count_unit().wrapping_sub(self.queued_bit) + } else { + self.count_unit() + }; + self.tail.fetch_add(update, Ordering::SeqCst); + // SAFETY: The node was queued, so its `Reserve` is alive, and it cannot be dropped + // while this lock is held. + unsafe { granted.as_ref() }.store(true, Ordering::Release); + wakers.push(waker); + } + if state.waiters.is_empty() && self.tail.load(Ordering::Relaxed) & self.queued_bit != 0 { + // A cancelled waiter emptied the queue. + self.send_waiters.queued.store(false, Ordering::Relaxed); + self.tail.fetch_and(!self.queued_bit, Ordering::SeqCst); + } + } + + /// Runs `grant_queued` under the wait-queue lock and wakes the granted senders. + fn grant_queued_locked(&self) { + let mut wakers = WakerBatch::new(); + { + let mut state = self.send_waiters.state.lock(); + if !state.closed { + self.grant_queued(&mut state, &mut wakers); + } + } + wakers.wake_all(); + } + + /// Hands capacity freed by a receive to the first queued sender. + fn grant_after_receive(&self) { + // The head exchange in `pop` and these loads are sequentially consistent, as are the + // `queued` store and the capacity check in `enqueue_sender`, and the end of a handoff and + // the capacity check in `finish_handoff`. A handoff in progress grants the freed capacity. + let senders = &self.send_waiters; + if senders.queued.load(Ordering::SeqCst) && senders.handoffs.load(Ordering::SeqCst) == 0 { + self.grant_queued_locked(); + } + } + + /// Ends a handoff that a send with a reservation began while senders were queued, then grants + /// capacity that receivers freed meanwhile. + fn finish_handoff(&self) { + let senders = &self.send_waiters; + senders.handoffs.fetch_sub(1, Ordering::SeqCst); + if senders.queued.load(Ordering::SeqCst) { + self.grant_queued_locked(); + } + } + + /// Returns capacity held by an unused permit to queued senders, or makes it available again. + fn release_reservation(&self) { + // Either this update observes `queued_bit`, or the capacity check after setting it + // observes this update. + let tail = self.tail.fetch_sub(self.count_unit(), Ordering::SeqCst); + if tail & self.queued_bit != 0 { + self.grant_queued_locked(); + } + } + + /// Pushes a value using capacity a reservation holds, then notifies a waiting receiver. + /// + /// Returns the value if the last receiver disconnected first. + fn publish_reserved(&self, value: T) -> Result<(), T> { + // While senders are queued, this send takes over granting capacity that receivers free + // until it finishes. Otherwise, a receiver that grants and wakes every queued sender itself + // can fall behind them, which keeps them queued and makes every send wait. + let handoff = self.send_waiters.queued.load(Ordering::SeqCst); + if handoff { + self.send_waiters.handoffs.fetch_add(1, Ordering::SeqCst); + } + let result = self.push_reserved(value); + if result.is_ok() { + self.notify_receiver(); + } + if handoff { + self.finish_handoff(); + } + result + } + + /// Claims a free slot without a reservation, leaving room for reserved capacity. + fn push_unreserved(&self, value: T) -> Result<(), Unclaimed> { + let mut backoff = Backoff::new(); + let mut tail = self.tail.load(Ordering::Relaxed); + loop { + if tail & self.mark_bit != 0 { + return Err(Unclaimed::Disconnected(value)); + } + if tail & self.queued_bit != 0 { + // Queued senders own released capacity first. + return Err(Unclaimed::Full(value)); + } + let position = self.position(tail); + let slot = self.slot(position); + let stamp = slot.stamp.load(Ordering::Acquire); + if stamp == position { + let reserved = self.reserved(tail); + if reserved != 0 { + match self.queued_len(position, self.head.load(Ordering::SeqCst)) { + Some(len) if len + reserved >= self.slots.len() => { + return Err(Unclaimed::Full(value)); + } + Some(_) => {} + None => { + // Receivers passed this tail, so the exchange below would fail. + backoff.spin(); + tail = self.tail.load(Ordering::Relaxed); + continue; + } + } + } + // The exchange fails if the reservation state changed after the load. + match self.tail.compare_exchange_weak( + tail, + self.with_position(tail, self.next_position(position)), + Ordering::SeqCst, + Ordering::Relaxed, + ) { + Ok(_) => { + // SAFETY: The exchange claimed this free slot. No receiver reads it until + // the stamp below publishes the value. + unsafe { slot.value.get().write(MaybeUninit::new(value)) }; + slot.stamp.store(position + 1, Ordering::Release); + return Ok(()); + } + Err(current) => { + tail = current; + backoff.spin(); + } + } + } else if stamp.wrapping_add(self.one_lap) == position + 1 { + // This slot still holds the previous lap's value. A receiver that has not stored + // the head exchange yet is still receiving, so reporting full linearizes before it. + if self.head.load(Ordering::SeqCst).wrapping_add(self.one_lap) == position { + return Err(Unclaimed::Full(value)); + } + // A receiver has claimed that value and is still reading it. + backoff.snooze(); + tail = self.tail.load(Ordering::Relaxed); + } else { + // Another sender advanced `tail` after it was loaded. + backoff.spin(); + tail = self.tail.load(Ordering::Relaxed); + } + } + } + + /// Claims a slot with capacity a reservation holds, returning the reservation in the same + /// exchange. + fn push_reserved(&self, value: T) -> Result<(), T> { + let mut backoff = Backoff::new(); + let mut tail = self.tail.load(Ordering::Relaxed); + loop { + if tail & self.mark_bit != 0 { + return Err(value); + } + let position = self.position(tail); + let slot = self.slot(position); + let stamp = slot.stamp.load(Ordering::Acquire); + if stamp == position { + // The queued value takes over the capacity, so the free capacity is unchanged. + match self.tail.compare_exchange_weak( + tail, + self.with_position(tail, self.next_position(position)) - self.count_unit(), + Ordering::SeqCst, + Ordering::Relaxed, + ) { + Ok(_) => { + // SAFETY: As in `push_unreserved`. + unsafe { slot.value.get().write(MaybeUninit::new(value)) }; + slot.stamp.store(position + 1, Ordering::Release); + return Ok(()); + } + Err(current) => { + tail = current; + backoff.spin(); + } + } + } else if stamp.wrapping_add(self.one_lap) == position + 1 { + // A receiver is still reading this slot's previous value. The reservation rules + // out a previous value that no receiver has claimed. + backoff.snooze(); + tail = self.tail.load(Ordering::Relaxed); + } else { + // Another sender advanced `tail` after it was loaded. + backoff.spin(); + tail = self.tail.load(Ordering::Relaxed); + } + } + } + + fn pop(&self) -> Result { + let mut backoff = Backoff::new(); + let mut head = self.head.load(Ordering::Relaxed); + loop { + let slot = self.slot(head); + let stamp = slot.stamp.load(Ordering::Acquire); + if stamp == head + 1 { + match self.head.compare_exchange_weak( + head, + self.next_position(head), + Ordering::SeqCst, + Ordering::Relaxed, + ) { + Ok(_) => { + // SAFETY: The exchange claimed this published slot. No sender writes it + // again until the stamp below frees it for the next lap. + let value = unsafe { slot.value.get().read().assume_init() }; + slot.stamp + .store(head.wrapping_add(self.one_lap), Ordering::Release); + self.grant_after_receive(); + return Ok(value); + } + Err(current) => { + head = current; + backoff.spin(); + } + } + } else if stamp == head { + // Reporting empty before a sender's tail exchange is visible linearizes before + // that send. + let tail = self.tail.load(Ordering::SeqCst); + if self.position(tail) == head { + return Err(if tail & self.mark_bit == 0 { + TryRecvError::Empty + } else { + TryRecvError::Disconnected + }); + } + // A sender has claimed this slot and is still writing it. + backoff.snooze(); + head = self.head.load(Ordering::Relaxed); + } else if stamp.wrapping_add(self.one_lap) == head + 1 { + // Another receiver is still reading this slot's previous value. + backoff.snooze(); + head = self.head.load(Ordering::Relaxed); + } else { + // Another receiver advanced `head` after it was loaded. + backoff.spin(); + head = self.head.load(Ordering::Relaxed); + } + } + } + + fn peek(&self) -> Peek { + let head = self.head.load(Ordering::SeqCst); + let tail = self.tail.load(Ordering::SeqCst); + if self.position(tail) != head { + Peek::Ready + } else if tail & self.mark_bit == 0 { + Peek::Empty + } else { + Peek::Disconnected + } + } + + fn notify_receiver(&self) { + // The tail exchange in a push and this load are sequentially consistent, as are a + // receiver's registration and its recheck in `Recv::poll`. Either this load observes the + // registration, or the recheck observes the claimed slot. + if !self.recv_waiters.linked.load(Ordering::SeqCst) { + return; + } + let waker = { + let mut state = self.recv_waiters.state.lock(); + let waker = state.waiters.notify_one(); + self.recv_waiters + .linked + .store(!state.waiters.is_empty(), Ordering::SeqCst); + waker + }; + if let Some(waker) = waker { + waker.wake(); + } + } + + /// Drops buffered values once the last receiver has disconnected the queue. + /// + /// `tail` is the final tail position. Senders that claimed a slot before the disconnection may + /// still be writing it. + fn discard_values(&self, tail: Index) { + // If a destructor panics, the guard drops the remaining values while unwinding. + struct Discard<'a, T> { + shared: &'a Shared, + tail: Index, + } + + impl Drop for Discard<'_, T> { + fn drop(&mut self) { + while self.shared.drop_next_value(self.tail) {} + } + } + + let discard = Discard { shared: self, tail }; + while self.drop_next_value(tail) {} + drop(discard); + } + + /// Drops the value at `head` once it is published, and returns whether one remained. + /// + /// Only the last receiver calls this, so `head` cannot change concurrently. + fn drop_next_value(&self, tail: Index) -> bool { + let head = self.head.load(Ordering::Relaxed); + if head == tail { + return false; + } + let slot = self.slot(head); + let mut backoff = Backoff::new(); + while slot.stamp.load(Ordering::Acquire) != head + 1 { + backoff.snooze(); + } + // SAFETY: The stamp shows a published value, and the last receiver is its only reader. + let value = unsafe { slot.value.get().read().assume_init() }; + slot.stamp + .store(head.wrapping_add(self.one_lap), Ordering::Relaxed); + self.head.store(self.next_position(head), Ordering::Relaxed); + drop(value); + true + } +} + +impl Drop for Shared { + fn drop(&mut self) { + // The last receiver discards buffered values; this only reclaims values it could not. + let tail = *self.tail.get_mut(); + let tail = self.position(tail); + while self.drop_next_value(tail) {} + } +} + +/// Capacity reserved for one value on a bounded MPMC queue. +/// +/// A permit does not keep any receiver alive or claim message order. Dropping it without sending +/// passes its capacity to the next waiting sender or makes it available again. +#[must_use = "dropping the permit releases its reserved capacity"] +pub struct Permit<'a, T> { + shared: &'a Shared, +} + +impl fmt::Debug for Permit<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Permit").finish_non_exhaustive() + } +} + +impl Permit<'_, T> { + /// Enqueues a value using the reserved capacity without waiting for space. + /// + /// If the last receiver has been dropped, the error returns ownership of the value. + pub fn send(self, value: T) -> Result<(), SendError> { + let shared = self.shared; + mem::forget(self); + // Claiming the slot returns the reservation. Once the last receiver has disconnected, the + // capacity no longer matters. + shared.publish_reserved(value).map_err(SendError::new) + } +} + +impl Drop for Permit<'_, T> { + fn drop(&mut self) { + self.shared.release_reservation(); + } +} + +struct Reserve<'a, T> { + shared: &'a Shared, + waiter: Option, + /// Set by a grant, which also removes the queued node. + granted: AtomicBool, + /// The queued node points at `granted`, so this future must not move while queued. + _pinned: PhantomPinned, +} + +impl<'a, T> Reserve<'a, T> { + fn new(shared: &'a Shared) -> Self { + Self { + shared, + waiter: None, + granted: AtomicBool::new(false), + _pinned: PhantomPinned, + } + } + + /// Takes a grant, whose capacity passes to the returned permit. + fn take_grant(&mut self) -> Permit<'a, T> { + // The grant already removed the node. + self.waiter = None; + self.granted.store(false, Ordering::Relaxed); + Permit { + shared: self.shared, + } + } + + fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { + let shared = self.shared; + // A queued node or a grant stays until `drop` handles it under the lock, because a grant + // may still write to `granted` before the last receiver closes the wait queue. + if shared.is_disconnected() { + return Poll::Ready(Err(SendError::new(()))); + } + // Pairs with the release store in `grant_queued`, so claiming a grant needs no lock. + if self.granted.load(Ordering::Acquire) { + return Poll::Ready(Ok(self.take_grant())); + } + let mut wakers = WakerBatch::new(); + let (outcome, retired) = { + let mut state = shared.send_waiters.state.lock(); + if state.closed { + // The closed storage that held this waiter has been replaced. + self.waiter = None; + return Poll::Ready(Err(SendError::new(()))); + } + let mut retired = None; + let outcome = match self.waiter { + // A grant may have arrived after the first check. + _ if self.granted.load(Ordering::Relaxed) => Some(self.take_grant()), + // Reserving free capacity directly avoids marking senders queued, which would turn + // concurrent sends away. Trying only under the lock keeps waiting senders from + // contending on the cursors. A disconnection fails this attempt, and closing the + // queue then wakes this sender. + None if state.waiters.is_empty() && shared.reserve_unqueued().is_ok() => { + Some(Permit { shared }) + } + None => { + let granted = NonNull::from(&self.granted); + shared.enqueue_sender( + &mut state, + &mut self.waiter, + cx.waker(), + granted, + &mut wakers, + ); + // The enqueue grants this sender if capacity was already free. + self.granted + .load(Ordering::Relaxed) + .then(|| self.take_grant()) + } + Some(id) => { + let waiter = state.waiters.waiter_mut(id); + if !waiter.waker.will_wake(cx.waker()) { + retired = Some(mem::replace(&mut waiter.waker, cx.waker().clone())); + } + None + } + }; + (outcome, retired) + }; + wakers.wake_all(); + drop(retired); + match outcome { + Some(permit) => Poll::Ready(Ok(permit)), + None => Poll::Pending, + } + } +} + +impl Drop for Reserve<'_, T> { + fn drop(&mut self) { + if self.waiter.is_none() && !self.granted.load(Ordering::Relaxed) { + return; + } + let shared = self.shared; + let mut wakers = WakerBatch::new(); + let retired = { + let mut state = shared.send_waiters.state.lock(); + if state.closed { + return; + } + let retired = if self.granted.load(Ordering::Relaxed) { + // Return the granted capacity, which passes to the next queued sender below or + // becomes available again. + shared.tail.fetch_sub(shared.count_unit(), Ordering::SeqCst); + None + } else { + self.waiter + .take() + .map(|id| remove_waiter(&mut state.waiters, id)) + }; + shared.grant_queued(&mut state, &mut wakers); + retired + }; + wakers.wake_all(); + drop(retired); + } +} + +struct Recv<'a, T> { + shared: &'a Shared, + waiter: Option, +} + +impl Recv<'_, T> { + fn poll(&mut self, cx: &mut Context<'_>) -> Poll> { + loop { + match self.shared.peek() { + Peek::Ready => {} + Peek::Disconnected => { + self.unregister(); + return Poll::Ready(Err(RecvError::Disconnected)); + } + Peek::Empty => { + if !self.register(cx.waker()) { + // The last sender has disconnected, which the next peek observes. + continue; + } + if self.shared.peek() == Peek::Empty { + return Poll::Pending; + } + } + } + // Leave the wait queue before taking a value. Otherwise, a sender could assign its + // notification to this receiver after it took an earlier value. + self.unregister(); + match self.shared.pop() { + Ok(value) => return Poll::Ready(Ok(value)), + Err(TryRecvError::Disconnected) => { + return Poll::Ready(Err(RecvError::Disconnected)); + } + // Another receiver took the value first. + Err(TryRecvError::Empty) => {} + } + } + } + + /// Registers this receiver's waker, or refreshes it without losing its place. + /// + /// Returns `false` once the last sender has disconnected. + fn register(&mut self, waker: &Waker) -> bool { + let receivers = &self.shared.recv_waiters; + let retired = { + let mut state = receivers.state.lock(); + if state.closed { + self.waiter = None; + return false; + } + let retired = state.waiters.register(&mut self.waiter, waker); + receivers.linked.store(true, Ordering::SeqCst); + retired + }; + drop(retired); + true + } + + /// Leaves the wait queue, consuming any notification this receiver received. + fn unregister(&mut self) { + let Some(id) = self.waiter.take() else { + return; + }; + let receivers = &self.shared.recv_waiters; + let retired = { + let mut state = receivers.state.lock(); + if state.closed { + return; + } + let retired = remove_waiter(&mut state.waiters, id); + receivers + .linked + .store(!state.waiters.is_empty(), Ordering::SeqCst); + retired + }; + drop(retired); + } +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + let Some(id) = self.waiter.take() else { + return; + }; + let receivers = &self.shared.recv_waiters; + let (retired, waker) = { + let mut state = receivers.state.lock(); + if state.closed { + return; + } + let retired = remove_waiter(&mut state.waiters, id); + // Hand an unconsumed notification to the next receiver while a value still waits. + let waker = + if matches!(retired, RecvWaiter::Notified) && self.shared.peek() == Peek::Ready { + state.waiters.notify_one() + } else { + None + }; + receivers + .linked + .store(!state.waiters.is_empty(), Ordering::SeqCst); + (retired, waker) + }; + if let Some(waker) = waker { + waker.wake(); + } + drop(retired); + } +} + +/// Spins briefly, then yields, while another thread finishes a short operation on a slot or cursor. +struct Backoff { + spins: u32, + snoozes: u32, +} + +impl Backoff { + const SPIN_LIMIT: u32 = 6; + + fn new() -> Self { + Self { + spins: 0, + snoozes: 0, + } + } + + /// Backs off after losing a race on a cursor, without yielding. + fn spin(&mut self) { + for _ in 0..1 << self.spins { + hint::spin_loop(); + } + self.spins = (self.spins + 1).min(Self::SPIN_LIMIT); + } + + /// Waits for another thread to finish writing or reading a slot. + fn snooze(&mut self) { + if self.snoozes <= Self::SPIN_LIMIT { + for _ in 0..1 << self.snoozes { + hint::spin_loop(); + } + self.snoozes += 1; + } else { + thread::yield_now(); + } + } +} + +#[cfg(test)] +mod tests { + use std::panic; + use std::panic::AssertUnwindSafe; + use std::pin::pin; + use std::sync::Arc; + use std::task::Wake; + + use super::*; + + struct CountWakes(AtomicUsize); + + impl Wake for CountWakes { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + fn wake_counter() -> (Waker, Arc) { + let wakes = Arc::new(CountWakes(AtomicUsize::new(0))); + (Waker::from(wakes.clone()), wakes) + } + + /// Returns the capacity held by permits and grants, and whether senders are queued. + fn reservation_state(shared: &Shared) -> (usize, bool) { + let tail = shared.tail.load(Ordering::Relaxed); + (shared.reserved(tail), tail & shared.queued_bit != 0) + } + + struct CountDrops(Arc); + + impl Drop for CountDrops { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[test] + fn positions_wrap_across_laps_in_fifo_order() { + let shared = Shared::new(3); + let mut next = 0; + for lap in 0..20 { + // Leave a varying number of values queued so every slot wraps at a different offset. + for _ in 0..lap % 3 + 1 { + shared.try_send(next).unwrap(); + next += 1; + } + for expected in next - (lap % 3 + 1)..next { + assert_eq!(shared.try_recv(), Ok(expected)); + } + } + for value in 0..3 { + shared.try_send(value).unwrap(); + } + assert_eq!(shared.try_send(3), Err(TrySendError::Full(3))); + } + + #[test] + #[should_panic(expected = "mpmc bounded queue capacity is too large")] + fn capacity_above_the_maximum_panics() { + let _ = Shared::<()>::new(MAX_CAPACITY + 1); + } + + #[test] + fn every_allowed_capacity_needs_2_40_sends_to_wrap_the_tail() { + // A capacity just above a power of two leaves the fewest sends per wrap. + for capacity in [ + 1, + 2, + 63, + 64, + 65, + 1 << 16, + (1 << 20) - 1, + 1 << 20, + MAX_CAPACITY, + ] { + let (.., one_lap) = tail_layout(capacity); + let laps = 1u128 << (Index::BITS - one_lap.trailing_zeros()); + assert!(laps * capacity as u128 >= 1 << 40, "capacity {capacity}"); + } + } + + #[test] + fn last_receiver_drops_buffered_values_across_the_wrap() { + let drops = Arc::new(AtomicUsize::new(0)); + let shared = Shared::new(2); + shared.try_send(CountDrops(drops.clone())).unwrap(); + shared.try_send(CountDrops(drops.clone())).unwrap(); + drop(shared.try_recv().unwrap()); + shared.try_send(CountDrops(drops.clone())).unwrap(); + assert_eq!(drops.load(Ordering::Relaxed), 1); + + shared.drop_receiver(); + assert_eq!(drops.load(Ordering::Relaxed), 3); + assert!(matches!( + shared.try_send(CountDrops(drops.clone())), + Err(TrySendError::Disconnected(_)) + )); + assert_eq!(drops.load(Ordering::Relaxed), 4); + } + + #[test] + fn discard_continues_after_a_destructor_panics() { + struct PanicOnDrop; + + impl Drop for PanicOnDrop { + fn drop(&mut self) { + panic!("value destructor panicked"); + } + } + + let drops = Arc::new(AtomicUsize::new(0)); + let shared = Shared::new(3); + shared.try_send((None, Some(PanicOnDrop))).unwrap(); + shared + .try_send((Some(CountDrops(drops.clone())), None)) + .unwrap(); + shared + .try_send((Some(CountDrops(drops.clone())), None)) + .unwrap(); + + let result = panic::catch_unwind(AssertUnwindSafe(|| shared.drop_receiver())); + assert!(result.is_err()); + assert_eq!(drops.load(Ordering::Relaxed), 2); + drop(shared); + assert_eq!(drops.load(Ordering::Relaxed), 2); + } + + #[test] + fn receiver_that_takes_a_silently_published_value_leaves_the_wait_queue() { + let shared = Shared::new(2); + shared.clone_receiver(); + let (first_waker, first_wakes) = wake_counter(); + let (second_waker, second_wakes) = wake_counter(); + let mut first = Recv { + shared: &shared, + waiter: None, + }; + let mut second = Recv { + shared: &shared, + waiter: None, + }; + assert!( + first + .poll(&mut Context::from_waker(&first_waker)) + .is_pending() + ); + assert!( + second + .poll(&mut Context::from_waker(&second_waker)) + .is_pending() + ); + + // A sender that read the wait-queue flag before either registration publishes silently. + assert!(matches!(shared.push_unreserved(1), Ok(()))); + // A spurious poll finds that value. Its receiver must leave the queue before taking it, so + // the next notification goes to the receiver that is still waiting. + assert_eq!( + first.poll(&mut Context::from_waker(&first_waker)), + Poll::Ready(Ok(1)) + ); + shared.try_send(2).unwrap(); + assert_eq!(first_wakes.0.load(Ordering::Relaxed), 0); + assert_eq!(second_wakes.0.load(Ordering::Relaxed), 1); + assert_eq!( + second.poll(&mut Context::from_waker(&second_waker)), + Poll::Ready(Ok(2)) + ); + } + + #[test] + fn cancelled_reservation_keeps_capacity_exact() { + let shared = Shared::new(1); + let permit = shared.try_reserve().unwrap(); + let (waker, wakes) = wake_counter(); + let mut reserve = pin!(shared.reserve()); + assert!( + reserve + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(permit); + assert_eq!(wakes.0.load(Ordering::Relaxed), 1); + let Poll::Ready(Ok(granted)) = reserve.as_mut().poll(&mut Context::from_waker(&waker)) + else { + panic!("the released capacity must be granted to the waiting reservation"); + }; + assert!(matches!(shared.try_send(1), Err(TrySendError::Full(1)))); + granted.send(2).unwrap(); + assert_eq!(shared.try_recv(), Ok(2)); + shared.try_send(3).unwrap(); + assert_eq!(shared.try_recv(), Ok(3)); + } + + #[test] + fn reservations_and_grants_keep_capacity_exact_across_laps() { + let shared = Shared::new(2); + let (waker, wakes) = wake_counter(); + for round in 0..6 { + let permit = shared.try_reserve().unwrap(); + shared.try_send(round).unwrap(); + assert!(matches!(shared.try_send(9), Err(TrySendError::Full(9)))); + let mut reserve = pin!(shared.reserve()); + assert!( + reserve + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + // The queued reservation owns the capacity the next receive frees. + assert_eq!(shared.try_recv(), Ok(round)); + assert!(matches!(shared.try_send(9), Err(TrySendError::Full(9)))); + let Poll::Ready(Ok(granted)) = reserve.as_mut().poll(&mut Context::from_waker(&waker)) + else { + panic!("the freed capacity must be granted to the queued reservation"); + }; + drop(permit); + granted.send(round + 10).unwrap(); + assert_eq!(reservation_state(&shared), (0, false)); + assert_eq!(shared.try_recv(), Ok(round + 10)); + } + assert_eq!(wakes.0.load(Ordering::Relaxed), 6); + for value in 0..2 { + shared.try_send(value).unwrap(); + } + assert!(matches!(shared.try_send(2), Err(TrySendError::Full(2)))); + } + + #[test] + fn send_in_a_handoff_grants_capacity_that_a_receive_freed() { + let shared = Shared::new(1); + shared.try_send(1).unwrap(); + let (waker, wakes) = wake_counter(); + let mut reserve = pin!(shared.reserve()); + assert!( + reserve + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + // A send with a reservation is between its claim and its grant check. + shared.send_waiters.handoffs.fetch_add(1, Ordering::SeqCst); + assert_eq!(shared.try_recv(), Ok(1)); + assert_eq!( + wakes.0.load(Ordering::Relaxed), + 0, + "the receive left the grant" + ); + shared.finish_handoff(); + assert_eq!(wakes.0.load(Ordering::Relaxed), 1); + let Poll::Ready(Ok(granted)) = reserve.as_mut().poll(&mut Context::from_waker(&waker)) + else { + panic!("the handoff must grant the capacity the receive freed"); + }; + granted.send(2).unwrap(); + assert_eq!(shared.try_recv(), Ok(2)); + assert_eq!(reservation_state(&shared), (0, false)); + } + + #[test] + fn stalled_send_cannot_claim_capacity_that_a_later_permit_holds() { + // Enough reservations to wrap any counter narrower than 25 bits run while the send stalls. + const RESERVATIONS: usize = if cfg!(miri) { 4 } else { 1 << 24 }; + + let shared = Shared::new(1); + // The send loads the tail and finds the only slot free, then stalls before its exchange. + let stale = shared.tail.load(Ordering::Relaxed); + for _ in 0..RESERVATIONS { + drop(shared.try_reserve().unwrap()); + } + let permit = shared.try_reserve().unwrap(); + let position = shared.position(stale); + let claim = shared.tail.compare_exchange( + stale, + shared.with_position(stale, shared.next_position(position)), + Ordering::SeqCst, + Ordering::Relaxed, + ); + assert!(claim.is_err(), "the permit holds the only capacity"); + permit.send(1).unwrap(); + assert_eq!(shared.try_recv(), Ok(1)); + assert_eq!(reservation_state(&shared), (0, false)); + } + + #[test] + fn threads_mixing_reservations_and_direct_sends_deliver_every_value_once() { + // Alternating permits and direct sends keeps reservations changing while slots are claimed. + const THREADS: usize = 3; + const VALUES: usize = if cfg!(miri) { 16 } else { 20_000 }; + + let shared = Shared::new(2); + let received = thread::scope(|scope| { + for producer in 0..THREADS { + let shared = &shared; + scope.spawn(move || { + for value in producer * VALUES..(producer + 1) * VALUES { + if value % 2 == 0 { + let mut value = value; + while let Err(error) = shared.try_send(value) { + let TrySendError::Full(returned) = error else { + unreachable!("receivers outlive the producers"); + }; + value = returned; + thread::yield_now(); + } + } else { + let permit = loop { + match shared.try_reserve() { + Ok(permit) => break permit, + Err(TrySendError::Full(())) => thread::yield_now(), + Err(TrySendError::Disconnected(())) => unreachable!(), + } + }; + permit.send(value).unwrap(); + } + } + }); + } + let consumers = (0..THREADS) + .map(|_| { + let shared = &shared; + scope.spawn(move || { + let mut values = Vec::new(); + while values.len() < VALUES { + match shared.try_recv() { + Ok(value) => values.push(value), + Err(TryRecvError::Empty) => thread::yield_now(), + Err(TryRecvError::Disconnected) => unreachable!(), + } + } + values + }) + }) + .collect::>(); + consumers + .into_iter() + .flat_map(|consumer| consumer.join().unwrap()) + .collect::>() + }); + let mut received = received; + received.sort_unstable(); + assert_eq!(received, (0..THREADS * VALUES).collect::>()); + assert_eq!(reservation_state(&shared), (0, false)); + } + + #[test] + fn threads_deliver_every_value_exactly_once_through_a_small_ring() { + const THREADS: usize = 3; + const VALUES: usize = if cfg!(miri) { 20 } else { 20_000 }; + + let shared = Shared::new(2); + let received = thread::scope(|scope| { + for producer in 0..THREADS { + let shared = &shared; + scope.spawn(move || { + for value in producer * VALUES..(producer + 1) * VALUES { + let mut value = value; + loop { + match shared.try_send(value) { + Ok(()) => break, + Err(TrySendError::Full(returned)) => { + value = returned; + thread::yield_now(); + } + Err(TrySendError::Disconnected(_)) => unreachable!(), + } + } + } + }); + } + let consumers = (0..THREADS) + .map(|_| { + let shared = &shared; + scope.spawn(move || { + let mut values = Vec::new(); + while values.len() < VALUES { + match shared.try_recv() { + Ok(value) => values.push(value), + Err(TryRecvError::Empty) => thread::yield_now(), + Err(TryRecvError::Disconnected) => unreachable!(), + } + } + values + }) + }) + .collect::>(); + consumers + .into_iter() + .flat_map(|consumer| consumer.join().unwrap()) + .collect::>() + }); + let mut received = received; + received.sort_unstable(); + assert_eq!(received, (0..THREADS * VALUES).collect::>()); + } +} diff --git a/asyncband/src/mpmc/unbounded.rs b/asyncband/src/mpmc/unbounded.rs index 1a9e6bc8..be3c4265 100644 --- a/asyncband/src/mpmc/unbounded.rs +++ b/asyncband/src/mpmc/unbounded.rs @@ -21,7 +21,6 @@ use std::sync::Arc; use crate::mpmc::RecvError; use crate::mpmc::SendError; use crate::mpmc::TryRecvError; -use crate::mpmc::TrySendError; use crate::mpmc::queue::Shared; /// Creates an unbounded multi-producer, multi-consumer queue. @@ -73,11 +72,7 @@ impl UnboundedSender { /// /// If all receivers have been dropped, the value is returned in [`SendError`]. pub fn send(&self, value: T) -> Result<(), SendError> { - match self.shared.try_send(value) { - Ok(()) => Ok(()), - Err(TrySendError::Disconnected(value)) => Err(SendError::new(value)), - Err(TrySendError::Full(_)) => unreachable!("unbounded queue cannot be full"), - } + self.shared.send(value) } } diff --git a/asyncband/src/mpmc/waiters.rs b/asyncband/src/mpmc/waiters.rs new file mode 100644 index 00000000..38c91c59 --- /dev/null +++ b/asyncband/src/mpmc/waiters.rs @@ -0,0 +1,65 @@ +// 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::mem; +use std::task::Waker; + +use crate::internal::waitlist::WaitList; +use crate::internal::waitlist::WaiterId; + +/// A receiver notification makes the task runnable without reserving a value. +pub enum RecvWaiter { + Waiting(Waker), + Notified, +} + +pub fn remove_waiter(waiters: &mut WaitList, id: WaiterId) -> W { + // Unlinking is idempotent, so detached waiters are removed the same way as linked ones. + waiters.unlink_waiter(id, |_| true); + waiters.remove_unlinked_waiter(id) +} + +impl WaitList { + pub fn notify_one(&mut self) -> Option { + let (_, waiter) = self.unlink_first_waiter(|_| true)?; + let RecvWaiter::Waiting(waker) = mem::replace(waiter, RecvWaiter::Notified) else { + unreachable!("only waiting operations remain linked"); + }; + Some(waker) + } + + /// Registers a pending operation's waker or refreshes an existing registration. + /// + /// If the future finds no value after notification, its waiter rejoins the queue. + #[must_use = "drop the replaced waker after releasing the queue lock"] + pub fn register(&mut self, id: &mut Option, current: &Waker) -> Option { + if let Some(queued) = *id { + if let RecvWaiter::Waiting(waker) = self.waiter_mut(queued) { + if waker.will_wake(current) { + return None; + } + return Some(mem::replace(waker, current.clone())); + } + } + if let Some(notified) = id.take() { + // The notification already took this node's waker, so nothing is retired. + remove_waiter(self, notified); + } + *id = Some(self.push_back(RecvWaiter::Waiting(current.clone()))); + None + } +} diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index adbc3ad0..360efbe0 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -49,6 +49,7 @@ asyncband = { workspace = true, features = [ "waitgroup", "watch", ] } +crossbeam-channel = { workspace = true } divan = { workspace = true } flume = { workspace = true, features = ["async"] } tokio = { workspace = true, features = ["rt-multi-thread", "sync"] } diff --git a/benchmarks/benches/ecosystem/mpmc/bounded.rs b/benchmarks/benches/ecosystem/mpmc/bounded.rs index 6cc84fe2..65655a53 100644 --- a/benchmarks/benches/ecosystem/mpmc/bounded.rs +++ b/benchmarks/benches/ecosystem/mpmc/bounded.rs @@ -17,9 +17,11 @@ use benchmarks::channels::BATCH_MESSAGES; use benchmarks::channels::BOUNDED_CAPACITY; +use benchmarks::channels::adapters::AsyncBoundedMpmc; use benchmarks::channels::adapters::AsyncChannel; use benchmarks::channels::adapters::Bounded; use benchmarks::channels::adapters::BoundedMpmc; +use benchmarks::channels::adapters::Crossbeam; use benchmarks::channels::adapters::Flume; use benchmarks::channels::adapters::Mpmc; use benchmarks::channels::mpmc::TOPOLOGIES; @@ -31,7 +33,7 @@ use divan::Bencher; use divan::counter::ItemsCount; #[divan::bench( - types = [Mpmc, AsyncChannel, Flume], + types = [Mpmc, AsyncChannel, Flume, Crossbeam], args = TOPOLOGIES, sample_count = 20, sample_size = 1, @@ -43,6 +45,39 @@ fn thread_batch(bencher: Bencher, topology: Topology) { .bench_local_refs(|batch| batch.run()); } +/// Capacity 1 makes every send wait for a receiver, while 1024 rarely fills. +#[divan::bench( + types = [Mpmc, AsyncChannel, Flume, Crossbeam], + consts = [1, 1024], + args = TOPOLOGIES, + sample_count = 20, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn thread_batch_capacity( + bencher: Bencher, + topology: Topology, +) { + bencher + .with_inputs(|| ThreadBatch::new_bounded::(CAPACITY, topology)) + .bench_local_refs(|batch| batch.run()); +} + +/// Endpoints spin on the nonblocking operations, so this measures queue contention without +/// parking or wakeups. +#[divan::bench( + types = [Mpmc, AsyncChannel, Flume, Crossbeam], + args = TOPOLOGIES, + sample_count = 20, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn spin_batch(bencher: Bencher, topology: Topology) { + bencher + .with_inputs(|| ThreadBatch::new_bounded_spinning::(BOUNDED_CAPACITY, topology)) + .bench_local_refs(|batch| batch.run()); +} + #[divan::bench( types = [Mpmc, AsyncChannel, Flume], consts = [0, 4], @@ -51,7 +86,7 @@ fn thread_batch(bencher: Bencher, topology: Topology) { sample_size = 1, counter = ItemsCount::new(BATCH_MESSAGES), )] -fn task_batch(bencher: Bencher, topology: Topology) { +fn task_batch(bencher: Bencher, topology: Topology) { let runtime = runtime(WORKERS); bencher .with_inputs(|| TaskBatch::new::>(&runtime, topology)) diff --git a/benchmarks/benches/primitives/mpmc/bounded.rs b/benchmarks/benches/primitives/mpmc/bounded.rs index 5c3e52ac..3a024896 100644 --- a/benchmarks/benches/primitives/mpmc/bounded.rs +++ b/benchmarks/benches/primitives/mpmc/bounded.rs @@ -46,6 +46,19 @@ fn send_then_recv(bencher: Bencher) { }); } +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn try_reserve_then_send(bencher: Bencher) { + let (sender, receiver) = mpmc::bounded(1); + bencher.bench_local(|| { + sender + .try_reserve() + .unwrap() + .send(black_box(1usize)) + .unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + #[divan::bench(sample_size = FAST_SAMPLE_SIZE)] fn wake_blocked_sender(bencher: Bencher) { let mut context = bench_context(); diff --git a/benchmarks/src/channels/adapters.rs b/benchmarks/src/channels/adapters.rs index 488fc021..b5a8b13d 100644 --- a/benchmarks/src/channels/adapters.rs +++ b/benchmarks/src/channels/adapters.rs @@ -17,6 +17,7 @@ use std::future::Future; use std::marker::PhantomData; +use std::task::Poll; use asyncband::blocking::FutureExt; use asyncband::mpmc; @@ -28,6 +29,7 @@ pub struct Mpmc; pub struct Spmc; pub struct AsyncChannel; pub struct Flume; +pub struct Crossbeam; pub struct Bounded(PhantomData); pub struct Unbounded(PhantomData); @@ -76,7 +78,7 @@ impl Channel for Unbounded { } } -impl Channel for Bounded { +impl Channel for Bounded { type Sender = C::Sender; type Receiver = C::Receiver; @@ -110,21 +112,26 @@ impl Channel for Unbounded { } } +/// A bounded MPMC queue that dedicated threads can use. pub trait BoundedMpmc: Send + Sync + 'static { type Sender: Clone + Send + Sync + 'static; type Receiver: Clone + Send + Sync + 'static; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver); - fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; - fn recv_async(receiver: &Self::Receiver) -> impl Future> + Send; + fn send(sender: &Self::Sender, value: usize); + fn recv(receiver: &Self::Receiver) -> Option; - fn send(sender: &Self::Sender, value: usize) { - Self::send_async(sender, value).block_on(); - } + /// Sends without waiting, returning the value while the queue is full. + fn try_send(sender: &Self::Sender, value: usize) -> Result<(), usize>; - fn recv(receiver: &Self::Receiver) -> Option { - Self::recv_async(receiver).block_on() - } + /// Receives without waiting. `Pending` means empty and `Ready(None)` means disconnected. + fn try_recv(receiver: &Self::Receiver) -> Poll>; +} + +/// A bounded MPMC queue that asynchronous tasks can also use. +pub trait AsyncBoundedMpmc: BoundedMpmc { + fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; + fn recv_async(receiver: &Self::Receiver) -> impl Future> + Send; } pub trait UnboundedMpmc: Send + Sync + 'static { @@ -148,6 +155,32 @@ impl BoundedMpmc for Mpmc { mpmc::bounded(capacity) } + fn send(sender: &Self::Sender, value: usize) { + ::send_async(sender, value).block_on(); + } + + fn recv(receiver: &Self::Receiver) -> Option { + ::recv_async(receiver).block_on() + } + + fn try_send(sender: &Self::Sender, value: usize) -> Result<(), usize> { + match sender.try_send(value) { + Ok(()) => Ok(()), + Err(mpmc::TrySendError::Full(value)) => Err(value), + Err(mpmc::TrySendError::Disconnected(_)) => panic!("benchmark receiver disconnected"), + } + } + + fn try_recv(receiver: &Self::Receiver) -> Poll> { + match receiver.try_recv() { + Ok(value) => Poll::Ready(Some(value)), + Err(mpmc::TryRecvError::Empty) => Poll::Pending, + Err(mpmc::TryRecvError::Disconnected) => Poll::Ready(None), + } + } +} + +impl AsyncBoundedMpmc for Mpmc { async fn send_async(sender: &Self::Sender, value: usize) { sender .send(value) @@ -185,25 +218,45 @@ impl BoundedMpmc for AsyncChannel { async_channel::bounded(capacity) } - async fn send_async(sender: &Self::Sender, value: usize) { + fn send(sender: &Self::Sender, value: usize) { sender - .send(value) - .await + .send_blocking(value) .expect("benchmark sender disconnected"); } - async fn recv_async(receiver: &Self::Receiver) -> Option { - receiver.recv().await.ok() + fn recv(receiver: &Self::Receiver) -> Option { + receiver.recv_blocking().ok() } - fn send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: usize) -> Result<(), usize> { + match sender.try_send(value) { + Ok(()) => Ok(()), + Err(async_channel::TrySendError::Full(value)) => Err(value), + Err(async_channel::TrySendError::Closed(_)) => { + panic!("benchmark receiver disconnected") + } + } + } + + fn try_recv(receiver: &Self::Receiver) -> Poll> { + match receiver.try_recv() { + Ok(value) => Poll::Ready(Some(value)), + Err(async_channel::TryRecvError::Empty) => Poll::Pending, + Err(async_channel::TryRecvError::Closed) => Poll::Ready(None), + } + } +} + +impl AsyncBoundedMpmc for AsyncChannel { + async fn send_async(sender: &Self::Sender, value: usize) { sender - .send_blocking(value) + .send(value) + .await .expect("benchmark sender disconnected"); } - fn recv(receiver: &Self::Receiver) -> Option { - receiver.recv_blocking().ok() + async fn recv_async(receiver: &Self::Receiver) -> Option { + receiver.recv().await.ok() } } @@ -238,6 +291,32 @@ impl BoundedMpmc for Flume { flume::bounded(capacity) } + fn send(sender: &Self::Sender, value: usize) { + sender.send(value).expect("benchmark sender disconnected"); + } + + fn recv(receiver: &Self::Receiver) -> Option { + receiver.recv().ok() + } + + fn try_send(sender: &Self::Sender, value: usize) -> Result<(), usize> { + match sender.try_send(value) { + Ok(()) => Ok(()), + Err(flume::TrySendError::Full(value)) => Err(value), + Err(flume::TrySendError::Disconnected(_)) => panic!("benchmark receiver disconnected"), + } + } + + fn try_recv(receiver: &Self::Receiver) -> Poll> { + match receiver.try_recv() { + Ok(value) => Poll::Ready(Some(value)), + Err(flume::TryRecvError::Empty) => Poll::Pending, + Err(flume::TryRecvError::Disconnected) => Poll::Ready(None), + } + } +} + +impl AsyncBoundedMpmc for Flume { async fn send_async(sender: &Self::Sender, value: usize) { sender .send_async(value) @@ -248,14 +327,6 @@ impl BoundedMpmc for Flume { async fn recv_async(receiver: &Self::Receiver) -> Option { receiver.recv_async().await.ok() } - - fn send(sender: &Self::Sender, value: usize) { - sender.send(value).expect("benchmark sender disconnected"); - } - - fn recv(receiver: &Self::Receiver) -> Option { - receiver.recv().ok() - } } impl UnboundedMpmc for Flume { @@ -278,3 +349,38 @@ impl UnboundedMpmc for Flume { receiver.recv().ok() } } + +impl BoundedMpmc for Crossbeam { + type Receiver = crossbeam_channel::Receiver; + type Sender = crossbeam_channel::Sender; + + fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { + crossbeam_channel::bounded(capacity) + } + + fn send(sender: &Self::Sender, value: usize) { + sender.send(value).expect("benchmark receiver disconnected"); + } + + fn recv(receiver: &Self::Receiver) -> Option { + receiver.recv().ok() + } + + fn try_send(sender: &Self::Sender, value: usize) -> Result<(), usize> { + match sender.try_send(value) { + Ok(()) => Ok(()), + Err(crossbeam_channel::TrySendError::Full(value)) => Err(value), + Err(crossbeam_channel::TrySendError::Disconnected(_)) => { + panic!("benchmark receiver disconnected") + } + } + } + + fn try_recv(receiver: &Self::Receiver) -> Poll> { + match receiver.try_recv() { + Ok(value) => Poll::Ready(Some(value)), + Err(crossbeam_channel::TryRecvError::Empty) => Poll::Pending, + Err(crossbeam_channel::TryRecvError::Disconnected) => Poll::Ready(None), + } + } +} diff --git a/benchmarks/src/channels/mpmc.rs b/benchmarks/src/channels/mpmc.rs index 8348058d..98ef2af3 100644 --- a/benchmarks/src/channels/mpmc.rs +++ b/benchmarks/src/channels/mpmc.rs @@ -16,8 +16,10 @@ // under the License. use std::fmt; +use std::hint; use std::sync::Arc; use std::sync::Barrier; +use std::task::Poll; use std::thread; use std::thread::JoinHandle; @@ -55,6 +57,10 @@ pub const TOPOLOGIES: &[Topology] = &[ producers: 1, consumers: 4, }, + Topology { + producers: 1, + consumers: 8, + }, Topology { producers: 2, consumers: 2, @@ -157,6 +163,13 @@ impl ThreadBatch { Self::new(sender, receiver, topology, C::send, C::recv) } + /// Uses only nonblocking operations, so endpoints spin instead of parking. This isolates queue + /// contention from wakeup costs. + pub fn new_bounded_spinning(capacity: usize, topology: Topology) -> Self { + let (sender, receiver) = C::channel(capacity); + Self::new(sender, receiver, topology, spin_send::, spin_recv::) + } + pub fn new_unbounded(topology: Topology) -> Self { let (sender, receiver) = C::channel(); Self::new(sender, receiver, topology, C::send, C::recv) @@ -230,3 +243,31 @@ impl ThreadBatch { black_box((count, checksum)) } } + +fn spin_send(sender: &C::Sender, mut value: usize) { + let mut spins = 0u32; + while let Err(returned) = C::try_send(sender, value) { + value = returned; + backoff(&mut spins); + } +} + +fn spin_recv(receiver: &C::Receiver) -> Option { + let mut spins = 0u32; + loop { + match C::try_recv(receiver) { + Poll::Ready(value) => return value, + Poll::Pending => backoff(&mut spins), + } + } +} + +/// Spins briefly, then yields, so oversubscribed runs still let the other side progress. +fn backoff(spins: &mut u32) { + if *spins < 64 { + hint::spin_loop(); + *spins += 1; + } else { + thread::yield_now(); + } +} diff --git a/tests-integration/tests/mpmc_test/concurrency.rs b/tests-integration/tests/mpmc_test/concurrency.rs index dbcbede8..33d9c32b 100644 --- a/tests-integration/tests/mpmc_test/concurrency.rs +++ b/tests-integration/tests/mpmc_test/concurrency.rs @@ -16,8 +16,10 @@ // under the License. use std::ops::Range; +use std::thread; use std::time::Duration; +use asyncband::blocking::FutureExt; use asyncband::mpmc; use tokio::task::JoinHandle; @@ -105,3 +107,96 @@ async fn unbounded_values_are_delivered_exactly_once_under_contention() { assert_delivered_exactly_once(receiver, producers).await; } + +#[test] +fn bounded_blocked_threads_deliver_exactly_once_through_a_small_ring() { + // Two slots keep every slot wrapping while senders wait for grants and receivers for values. + const THREADS: usize = 3; + const VALUES_PER_THREAD: usize = if cfg!(miri) { 16 } else { 4096 }; + + let (sender, receiver) = mpmc::bounded(2); + let mut received = thread::scope(|scope| { + for producer in 0..THREADS { + let sender = sender.clone(); + scope.spawn(move || { + let first = producer * VALUES_PER_THREAD; + for value in first..first + VALUES_PER_THREAD { + sender.send(value).block_on().unwrap(); + } + }); + } + drop(sender); + let consumers = (0..THREADS) + .map(|_| { + let receiver = receiver.clone(); + scope.spawn(move || { + let mut values = vec![]; + while let Ok(value) = receiver.recv().block_on() { + values.push(value); + } + values + }) + }) + .collect::>(); + drop(receiver); + consumers + .into_iter() + .flat_map(|consumer| consumer.join().unwrap()) + .collect::>() + }); + received.sort_unstable(); + assert_eq!( + received, + (0..THREADS * VALUES_PER_THREAD).collect::>() + ); +} + +#[test] +fn bounded_threads_releasing_unused_permits_deliver_exactly_once() { + // Dropping permits while other senders wait passes their capacity along the wait queue. + const THREADS: usize = 3; + const VALUES_PER_THREAD: usize = if cfg!(miri) { 12 } else { 4096 }; + + let (sender, receiver) = mpmc::bounded(2); + let mut received = thread::scope(|scope| { + for producer in 0..THREADS { + let sender = sender.clone(); + scope.spawn(move || { + let first = producer * VALUES_PER_THREAD; + for value in first..first + VALUES_PER_THREAD { + match value % 3 { + 0 => sender.send(value).block_on().unwrap(), + 1 => sender.reserve().block_on().unwrap().send(value).unwrap(), + _ => { + drop(sender.reserve().block_on().unwrap()); + sender.send(value).block_on().unwrap(); + } + } + } + }); + } + drop(sender); + let consumers = (0..THREADS) + .map(|_| { + let receiver = receiver.clone(); + scope.spawn(move || { + let mut values = vec![]; + while let Ok(value) = receiver.recv().block_on() { + values.push(value); + } + values + }) + }) + .collect::>(); + drop(receiver); + consumers + .into_iter() + .flat_map(|consumer| consumer.join().unwrap()) + .collect::>() + }); + received.sort_unstable(); + assert_eq!( + received, + (0..THREADS * VALUES_PER_THREAD).collect::>() + ); +}