From 4ea81c49d8d31b9d8fc314c251af0d56cf724ccc Mon Sep 17 00:00:00 2001 From: Yijun Zhao Date: Wed, 21 Oct 2026 16:39:50 +0800 Subject: [PATCH 1/4] perf(mpmc): back bounded queues with a ring buffer Replace the bounded queue's VecDeque behind one mutex with a preallocated ring of stamped slots adapted from crossbeam-channel's array flavor. Sends, receives, and reservations take no lock unless a sender or receiver has to wait, while wait-queue order, exact capacity, permits, cancellation, and disconnection behave as before. Storage for the full capacity is allocated when the queue is created. The capacity that permits and grants hold, and a bit marking queued senders, live in the tail cursor beside the position, so every exchange on it compares the reservation state it was computed from. While senders are queued, a send with a permit grants the capacity that receivers free until it finishes, so a single receiver does not have to wake every queued sender itself. Positions are 64 bits, and capacity is limited to 2,097,151, so a stalled exchange can be fooled only after at least 2^40 sends. A 32-bit cursor would leave as few as two lap bits, so the crate now requires 64-bit atomics and supports x86, x86-64, ARMv6 or later, and AArch64. Copy CachePadded from crossbeam-utils and record the incorporated crossbeam portions in LICENSE, licenserc.toml, and .asf.yaml. Add benchmarks against crossbeam-channel: a spinning batch that isolates queue contention from wakeups, a capacity sweep, the 1P/8C topology, and a reservation micro-benchmark. Split the bounded adapter trait so channels without async operations can join the thread workloads. --- .asf.yaml | 1 + CHANGELOG.md | 9 + Cargo.lock | 10 + Cargo.toml | 1 + LICENSE | 16 + README.md | 4 + asyncband/src/internal/cache_padded.rs | 181 ++ asyncband/src/internal/mod.rs | 4 + asyncband/src/internal/waker_batch.rs | 4 +- asyncband/src/lib.rs | 12 + asyncband/src/mpmc/bounded.rs | 13 +- asyncband/src/mpmc/mod.rs | 4 +- asyncband/src/mpmc/queue.rs | 328 +--- asyncband/src/mpmc/ring.rs | 1483 +++++++++++++++++ asyncband/src/mpmc/unbounded.rs | 7 +- asyncband/src/mpmc/waiters.rs | 65 + benchmarks/Cargo.toml | 1 + benchmarks/benches/ecosystem/mpmc/bounded.rs | 39 +- benchmarks/benches/primitives/mpmc/bounded.rs | 13 + benchmarks/src/channels/adapters.rs | 158 +- benchmarks/src/channels/mpmc.rs | 41 + licenserc.toml | 1 + .../tests/mpmc_test/concurrency.rs | 95 ++ 23 files changed, 2136 insertions(+), 354 deletions(-) create mode 100644 asyncband/src/internal/cache_padded.rs create mode 100644 asyncband/src/mpmc/ring.rs create mode 100644 asyncband/src/mpmc/waiters.rs diff --git a/.asf.yaml b/.asf.yaml index 19a37e74..94c1e74b 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -41,6 +41,7 @@ project: source_excludes_lightweight: - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/blocking/executor.rs" - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/blocking/parker.rs" + - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/internal/cache_padded.rs" - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/pool/bounded.rs" - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/pool/common.rs" - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/pool/mod.rs" diff --git a/CHANGELOG.md b/CHANGELOG.md index f36264e9..571f0db6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,15 @@ All notable changes to this project will be documented in this file. ## Unreleased +### Notable changes + +* Asyncband now requires targets with 64-bit atomics, such as x86, x86-64, ARMv6 or later, and AArch64, and fails to compile on targets without them, such as ARMv5TE, 32-bit PowerPC, and 32-bit MIPS. +* 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. + ## 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..8a1cfe37 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 requires 64-bit atomics and fails to compile on targets without them, such as ARMv5TE, 32-bit PowerPC, and 32-bit MIPS. + ## 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/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..6df013ca 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -54,6 +54,9 @@ pub fn register_waker(slot: &mut Option, waker: &Waker) -> Option #[allow(dead_code)] pub(crate) mod arena; +#[cfg(feature = "mpmc")] +pub(crate) mod cache_padded; + #[cfg(any(feature = "latch", feature = "once"))] pub(crate) mod countdown; @@ -116,6 +119,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..aedf196b 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -112,6 +112,11 @@ //! 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 requires 64-bit atomics +//! and fails to compile on targets without them, such as ARMv5TE, 32-bit PowerPC, and 32-bit MIPS. +//! //! # Disclaimer //! //! Apache Asyncband (Incubating) is an effort undergoing incubation at the Apache Software @@ -123,6 +128,13 @@ //! //! 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. +// The bounded MPMC queue packs a slot position, its held capacity, and a lap count into one 64-bit +// cursor. Narrower cursors leave too few lap bits to rule out a stale claim in practice. +#[cfg(not(target_has_atomic = "64"))] +compile_error!( + "asyncband requires a target with 64-bit atomics, such as x86, x86-64, ARMv6 or later, or AArch64" +); + 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..ddb92fa4 --- /dev/null +++ b/asyncband/src/mpmc/ring.rs @@ -0,0 +1,1483 @@ +// 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::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, which is why the crate requires 64-bit atomics. +type AtomicIndex = std::sync::atomic::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/licenserc.toml b/licenserc.toml index 6b883f0b..3a32f0c0 100644 --- a/licenserc.toml +++ b/licenserc.toml @@ -23,6 +23,7 @@ excludes = [ # These third-party-derived files retain the applicable upstream license and copyright notices. "asyncband/src/blocking/executor.rs", "asyncband/src/blocking/parker.rs", + "asyncband/src/internal/cache_padded.rs", "asyncband/src/pool/bounded.rs", "asyncband/src/pool/common.rs", "asyncband/src/pool/mod.rs", 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::>() + ); +} From 0d1ba7496884368408c7f486d6a9626562d83103 Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 22 Oct 2026 07:49:54 +0800 Subject: [PATCH 2/4] Apply batched suggestions from code review Co-authored-by: tison --- .asf.yaml | 1 - licenserc.toml | 1 - 2 files changed, 2 deletions(-) diff --git a/.asf.yaml b/.asf.yaml index 94c1e74b..19a37e74 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -41,7 +41,6 @@ project: source_excludes_lightweight: - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/blocking/executor.rs" - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/blocking/parker.rs" - - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/internal/cache_padded.rs" - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/pool/bounded.rs" - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/pool/common.rs" - "/apache-asyncband-*-incubating-src.tar.gz/*/asyncband/src/pool/mod.rs" diff --git a/licenserc.toml b/licenserc.toml index 3a32f0c0..6b883f0b 100644 --- a/licenserc.toml +++ b/licenserc.toml @@ -23,7 +23,6 @@ excludes = [ # These third-party-derived files retain the applicable upstream license and copyright notices. "asyncband/src/blocking/executor.rs", "asyncband/src/blocking/parker.rs", - "asyncband/src/internal/cache_padded.rs", "asyncband/src/pool/bounded.rs", "asyncband/src/pool/common.rs", "asyncband/src/pool/mod.rs", From 3220f33feafb155240ab13c66759faad54529135 Mon Sep 17 00:00:00 2001 From: tison Date: Thu, 22 Oct 2026 07:51:16 +0800 Subject: [PATCH 3/4] Apply suggestion from @tisonkun --- asyncband/src/lib.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index aedf196b..e9ad2620 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -128,6 +128,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. + // The bounded MPMC queue packs a slot position, its held capacity, and a lap count into one 64-bit // cursor. Narrower cursors leave too few lap bits to rule out a stale claim in practice. #[cfg(not(target_has_atomic = "64"))] From 3cfaaad0d54772840cf96cadc2b2969669080636 Mon Sep 17 00:00:00 2001 From: Yijun Zhao Date: Thu, 22 Oct 2026 10:58:50 +0800 Subject: [PATCH 4/4] fix(mpmc): fall back to locked cursors without 64-bit atomics Keep the bounded MPMC ring's 64-bit positions on every target instead of refusing to compile without 64-bit atomics. Where they are missing, such as ARMv5TE and 32-bit PowerPC and MIPS, each cursor and stamp is a u64 behind a lock with the same API subset, which is slower but runs the same algorithm. Targets with 64-bit atomics use std's AtomicU64 unchanged. --- CHANGELOG.md | 3 +- README.md | 2 +- asyncband/src/internal/atomic_u64.rs | 168 +++++++++++++++++++++++++++ asyncband/src/internal/mod.rs | 3 + asyncband/src/lib.rs | 12 +- asyncband/src/mpmc/ring.rs | 6 +- 6 files changed, 180 insertions(+), 14 deletions(-) create mode 100644 asyncband/src/internal/atomic_u64.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 571f0db6..13d5213c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,12 +25,11 @@ All notable changes to this project will be documented in this file. ### Notable changes -* Asyncband now requires targets with 64-bit atomics, such as x86, x86-64, ARMv6 or later, and AArch64, and fails to compile on targets without them, such as ARMv5TE, 32-bit PowerPC, and 32-bit MIPS. * 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. +* 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 diff --git a/README.md b/README.md index 8a1cfe37..04ce53e4 100644 --- a/README.md +++ b/README.md @@ -148,7 +148,7 @@ Asyncband types implement `Send` and `Sync` only when the protected, transferred ## Platform support -Asyncband supports x86, x86-64, ARMv6 or later, and AArch64 targets. It requires 64-bit atomics and fails to compile on targets without them, such as ARMv5TE, 32-bit PowerPC, and 32-bit MIPS. +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) 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/mod.rs b/asyncband/src/internal/mod.rs index 6df013ca..54e4a2bf 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -54,6 +54,9 @@ 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; diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index e9ad2620..9207bb19 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -114,8 +114,9 @@ //! //! # Platform support //! -//! Asyncband supports x86, x86-64, ARMv6 or later, and AArch64 targets. It requires 64-bit atomics -//! and fails to compile on targets without them, such as ARMv5TE, 32-bit PowerPC, and 32-bit MIPS. +//! 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 //! @@ -129,13 +130,6 @@ //! 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. -// The bounded MPMC queue packs a slot position, its held capacity, and a lap count into one 64-bit -// cursor. Narrower cursors leave too few lap bits to rule out a stale claim in practice. -#[cfg(not(target_has_atomic = "64"))] -compile_error!( - "asyncband requires a target with 64-bit atomics, such as x86, x86-64, ARMv6 or later, or AArch64" -); - mod internal; #[cfg(feature = "barrier")] diff --git a/asyncband/src/mpmc/ring.rs b/asyncband/src/mpmc/ring.rs index ddb92fa4..e6635889 100644 --- a/asyncband/src/mpmc/ring.rs +++ b/asyncband/src/mpmc/ring.rs @@ -42,6 +42,7 @@ 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; @@ -54,8 +55,9 @@ use crate::mpmc::TrySendError; use crate::mpmc::waiters::RecvWaiter; use crate::mpmc::waiters::remove_waiter; -// Positions are 64 bits on every target, which is why the crate requires 64-bit atomics. -type AtomicIndex = std::sync::atomic::AtomicU64; +// 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