diff --git a/CHANGELOG.md b/CHANGELOG.md index 7113dc6..3762c21 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,7 @@ All notable changes to this project will be documented in this file. ### Improvements * Finish releasing buffered bounded MPSC messages even if one message destructor panics. -* Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. +* Improve unbounded MPSC throughput for ready-message and contended asynchronous workloads; reclaim consumed storage incrementally and bound empty-buffer retention independently of previous peak occupancy. ## v0.7.2 diff --git a/Cargo.lock b/Cargo.lock index c6baa36..a684042 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -103,6 +103,7 @@ dependencies = [ "async-broadcast", "async-channel", "asyncband", + "crossbeam-channel", "divan", "flume", "pollster", @@ -261,6 +262,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 05e963c..6d2e664 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.16" } divan = { version = "0.1.21" } flume = { version = "0.12.0", default-features = false } pollster = { version = "1.0.1" } diff --git a/LICENSE b/LICENSE index 3bd8900..c4e9421 100644 --- a/LICENSE +++ b/LICENSE @@ -436,3 +436,15 @@ composition. Fastpool is licensed under Apache-2.0, and its source carries the following copyright notice: Copyright 2025 FastLabs Developers + +Portions of asyncband/src/mpsc/unbounded/queue.rs are adapted from the +segmented list in crossbeam-channel 0.5.16 at the following exact revision: + + https://github.com/crossbeam-rs/crossbeam/blob/9b56303b8aa9ff8ec5bbebb9d2da05e034977889/crossbeam-channel/src/flavors/list.rs + +The single-consumer implementation replaces multi-reader reclamation with FIFO +block ownership, bounds slots by payload size, and separates receiver notification +from storage. Crossbeam is licensed under Apache-2.0 or MIT; Asyncband uses the +Apache-2.0 option. The upstream distribution carries this copyright notice: + + Copyright (c) 2019 The Crossbeam Project Developers diff --git a/asyncband/src/mpsc/unbounded/buffer.rs b/asyncband/src/mpsc/unbounded/buffer.rs deleted file mode 100644 index ee195a9..0000000 --- a/asyncband/src/mpsc/unbounded/buffer.rs +++ /dev/null @@ -1,132 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::VecDeque; -use std::mem; - -// Bound the inline storage retained by a partial batch. Boxed payloads belong to individual -// messages, not these backing allocations. Empty buffers are reused without retaining peak size. -pub const SEGMENT_BYTES: usize = 32 * 1024; - -pub struct Buffer { - writable: VecDeque, - sealed: VecDeque>, -} - -impl Buffer { - pub fn new() -> Self { - Self { - writable: VecDeque::new(), - sealed: VecDeque::new(), - } - } - - fn segment_capacity() -> usize { - if size_of::() == 0 { - return usize::MAX; - } - let limit = (SEGMENT_BYTES / size_of::()).max(1); - // Power-of-two limits let VecDeque grow naturally without exceeding the segment budget. - 1 << (usize::BITS - 1 - limit.leading_zeros()) - } - - pub fn push(&mut self, value: T) { - if self.writable.len() == Self::segment_capacity() { - let next = VecDeque::with_capacity(Self::segment_capacity()); - let sealed = mem::replace(&mut self.writable, next); - self.sealed.push_back(sealed); - } - self.writable.push_back(value); - } - - pub fn refill(&mut self, batch: &mut VecDeque) { - debug_assert!(batch.is_empty()); - if let Some(sealed) = self.sealed.pop_front() { - *batch = sealed; - if self.sealed.is_empty() && self.sealed.capacity() * size_of::>() > 1024 { - self.sealed = VecDeque::new(); - } - } else if !self.writable.is_empty() { - mem::swap(batch, &mut self.writable); - } - } -} - -pub fn pop_batch(batch: &mut VecDeque) -> T { - if batch.len() == 1 && batch.capacity().saturating_mul(size_of::()) > SEGMENT_BYTES { - // Retire the allocation on the last value, outside the shared lock. Keep this as a tail - // expression to avoid intermediate storage for large inline values. - mem::take(batch).pop_front() - } else { - batch.pop_front() - } - .expect("receiver batch must not be empty") -} - -#[cfg(test)] -mod tests { - use std::collections::VecDeque; - - use super::Buffer; - use super::SEGMENT_BYTES; - use super::pop_batch; - - fn allocated_bytes(buffer: &Buffer, batch: &VecDeque) -> usize { - let slots = batch.capacity() - + buffer.writable.capacity() - + buffer.sealed.iter().map(VecDeque::capacity).sum::(); - slots * size_of::() - } - - fn receive(buffer: &mut Buffer, batch: &mut VecDeque) -> T { - if batch.is_empty() { - buffer.refill(batch); - } - pop_batch(batch) - } - - #[test] - fn a_partial_drain_reclaims_segments_and_preserves_new_sends() { - let mut buffer = Buffer::new(); - let mut batch = VecDeque::new(); - for value in 0..1024usize { - buffer.push([value; 128]); - } - let peak = allocated_bytes(&buffer, &batch); - for value in 0..512 { - assert_eq!(receive(&mut buffer, &mut batch), [value; 128]); - } - assert!(allocated_bytes(&buffer, &batch) <= peak * 3 / 4); - // This value must stay behind both the current batch and the sealed segments. - buffer.push([1024; 128]); - for value in 512..=1024 { - assert_eq!(receive(&mut buffer, &mut batch), [value; 128]); - } - assert!(allocated_bytes(&buffer, &batch) <= 2 * SEGMENT_BYTES); - buffer.refill(&mut batch); - assert!(batch.is_empty()); - } - - #[test] - fn oversized_inline_values_release_the_allocation_on_the_last_receive() { - let mut buffer = Buffer::new(); - let mut batch = VecDeque::new(); - buffer.push([7u8; SEGMENT_BYTES + 1]); - assert_eq!(receive(&mut buffer, &mut batch), [7u8; SEGMENT_BYTES + 1]); - assert_eq!(allocated_bytes(&buffer, &batch), 0); - } -} diff --git a/asyncband/src/mpsc/unbounded/mod.rs b/asyncband/src/mpsc/unbounded/mod.rs index be2b087..ce868d4 100644 --- a/asyncband/src/mpsc/unbounded/mod.rs +++ b/asyncband/src/mpsc/unbounded/mod.rs @@ -19,12 +19,15 @@ //! tasks. use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::task::Waker; -use self::buffer::Buffer; +use self::queue::Queue; use crate::internal::mutex::Mutex; -mod buffer; +mod queue; mod receiver; mod sender; @@ -43,28 +46,92 @@ pub use self::sender::UnboundedSender; /// Storage is reclaimed incrementally as messages are received. A bounded amount of empty /// storage may be retained for reuse, independently of the channel's previous peak occupancy. /// -/// Operations briefly acquire an internal mutex; no lock is held across an await point or while -/// invoking waker callbacks or message destructors. Sending and trying to receive may wait to -/// acquire this mutex, but never wait for capacity or new messages. +/// Sending and receiving may briefly wait for an in-progress producer or an internal mutex, but +/// never wait for capacity or new messages in `send` or `try_recv`. No lock is held across an +/// await point or while invoking waker callbacks or message destructors. pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { - let shared = Arc::new(Mutex::new(State { - buffer: Buffer::new(), - senders: 1, - receiver: true, - recv_waker: None, - })); + let (queue, consumer) = Queue::new(); + let shared = Arc::new(State { + queue, + senders: CachePadded(AtomicUsize::new(1)), + recv: ReceiverWake { + waiting: CachePadded(AtomicBool::new(false)), + waker: Mutex::new(None), + }, + }); ( UnboundedSender::new(shared.clone()), - UnboundedReceiver::new(shared), + UnboundedReceiver::new(shared, consumer), ) } -// Queue contents, endpoint liveness, and wake registration share one lock. Only the receiver -// accesses its current batch; refilling that batch preserves the order of concurrent sends. +// Separate producer reservations, sender counts, and the receiver's mostly-read waiting flag. +#[cfg_attr( + any( + target_arch = "aarch64", + target_arch = "arm64ec", + target_arch = "x86_64", + target_arch = "powerpc64" + ), + repr(align(128)) +)] +#[cfg_attr(target_arch = "s390x", repr(align(256)))] +#[cfg_attr( + not(any( + target_arch = "aarch64", + target_arch = "arm64ec", + target_arch = "x86_64", + target_arch = "powerpc64", + target_arch = "s390x" + )), + repr(align(64)) +)] +struct CachePadded(T); + struct State { - buffer: Buffer, - senders: usize, - // True while the receiving endpoint is alive. - receiver: bool, - recv_waker: Option, + queue: Queue, + senders: CachePadded, + recv: ReceiverWake, +} + +struct ReceiverWake { + waiting: CachePadded, + waker: Mutex>, +} + +impl ReceiverWake { + fn register(&self, waker: &Waker) { + // Clone, replacement destruction, and wake may reenter this channel. + let waker = waker.clone(); + let mut slot = self.waker.lock(); + let old = slot.replace(waker); + // Publication and this flag share an SC order: after register -> recheck, either the + // receiver sees the value or its producer observes this registration and wakes it. + self.waiting.0.store(true, Ordering::SeqCst); + drop(slot); + drop(old); + } + + fn take(&self) -> Option { + let mut slot = self.waker.lock(); + let waker = slot.take(); + // Clear under the registration lock so a delayed notifier cannot erase a newer wait. + self.waiting.0.store(false, Ordering::SeqCst); + waker + } + + fn wake(&self) { + if !self.waiting.0.load(Ordering::SeqCst) + || self + .waiting + .0 + .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst) + .is_err() + { + return; + } + if let Some(waker) = self.take() { + waker.wake(); + } + } } diff --git a/asyncband/src/mpsc/unbounded/queue.rs b/asyncband/src/mpsc/unbounded/queue.rs new file mode 100644 index 0000000..a8cbf9d --- /dev/null +++ b/asyncband/src/mpsc/unbounded/queue.rs @@ -0,0 +1,375 @@ +// Portions of this queue are adapted from crossbeam-channel 0.5.16. +// Copyright (c) 2019 The Crossbeam Project Developers +// Asyncband uses the Apache-2.0 license option for the incorporated code. +// The incorporated code has been modified for use in Apache Asyncband. +// Upstream source: +// https://github.com/crossbeam-rs/crossbeam/blob/9b56303b8aa9ff8ec5bbebb9d2da05e034977889/crossbeam-channel/src/flavors/list.rs + +//! Producer reservations own distinct slots. Publication transfers each value to the unique +//! consumer; FIFO consumption delays block reclamation until every producer using it is done. + +use std::cell::UnsafeCell; +use std::hint; +use std::marker::PhantomData; +use std::mem::MaybeUninit; +use std::panic::RefUnwindSafe; +use std::panic::UnwindSafe; +use std::ptr; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicPtr; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::thread; + +use super::CachePadded; +use crate::internal::mutex::Mutex; + +const BLOCK_BYTES: usize = 32 * 1024; +const SHIFT: usize = 1; +const STEP: usize = 1 << SHIFT; +const CLOSED: usize = 1; + +struct Slot { + value: UnsafeCell>, + ready: AtomicBool, +} + +impl Slot { + fn new() -> Self { + Self { + value: UnsafeCell::new(MaybeUninit::uninit()), + ready: AtomicBool::new(false), + } + } +} + +struct Block { + next: AtomicPtr>, + slots: Box<[Slot]>, +} + +impl Block { + // A power-of-two lap preserves slot identity when the wrapping position overflows. + // Small messages keep the original 31-slot layout; larger messages reduce the allocation. + const CAPACITY: usize = { + let slots = BLOCK_BYTES / size_of::>(); + let mut lap = 2; + while lap < 32 && lap * 2 - 1 <= slots { + lap *= 2; + } + lap - 1 + }; + const LAP: usize = Self::CAPACITY + 1; + + fn new() -> Box { + Box::new(Self { + next: AtomicPtr::new(ptr::null_mut()), + slots: (0..Self::CAPACITY).map(|_| Slot::new()).collect(), + }) + } +} + +struct Position { + index: AtomicUsize, + block: AtomicPtr>, +} + +/// Producer-owned half of an unbounded MPSC queue. +pub struct Queue { + tail: CachePadded>, + initialize: Mutex<()>, + first: AtomicPtr>, + _marker: PhantomData, +} + +// SAFETY: each producer reserves a distinct slot through `tail.index`. A value placed in a slot is +// only accessed by the single consumer after the producer publishes `ready` with release semantics. +unsafe impl Send for Queue {} +// SAFETY: the producer algorithm coordinates all shared mutation through atomics. The consumer is +// separate and may only read a slot after acquiring its `ready` publication. +unsafe impl Sync for Queue {} + +/// Consumer-owned position in an unbounded MPSC queue. +pub struct Consumer { + index: usize, + block: *mut Block, + _marker: PhantomData, +} + +// SAFETY: moving the unique consumer moves exclusive ownership of its position. Queue values cross +// the thread boundary only when `T: Send`. +unsafe impl Send for Consumer {} +// SAFETY: shared references cannot pop or close the consumer because both operations require +// exclusive access. Sharing an idle consumer therefore exposes neither its position nor `T`. +unsafe impl Sync for Consumer {} + +/// Result of attempting to pop one queue slot. +pub enum Pop { + /// A published value was removed. + Value(T), + /// No producer has reserved the next slot. + Empty, +} + +impl Queue { + /// Creates the producer queue and its unique consumer position. + pub fn new() -> (Self, Consumer) { + let block = ptr::null_mut(); + let queue = Self { + tail: CachePadded(Position { + index: AtomicUsize::new(0), + block: AtomicPtr::new(block), + }), + initialize: Mutex::new(()), + first: AtomicPtr::new(ptr::null_mut()), + _marker: PhantomData, + }; + let consumer = Consumer { + index: 0, + block, + _marker: PhantomData, + }; + (queue, consumer) + } + + /// Appends a value, or returns it if the consumer has closed the queue. + pub fn push(&self, value: T) -> Result<(), T> { + let mut backoff = Backoff::new(); + let mut tail = self.tail.0.index.load(Ordering::Acquire); + let mut block = self.tail.0.block.load(Ordering::Acquire); + let mut next_block = None; + + loop { + if tail & CLOSED != 0 { + return Err(value); + } + + if block.is_null() { + // Serialize first allocation with close. Ordinary sends never take this mutex. + let _guard = self.initialize.lock(); + tail = self.tail.0.index.load(Ordering::Acquire); + if tail & CLOSED != 0 { + return Err(value); + } + block = self.tail.0.block.load(Ordering::Acquire); + if block.is_null() { + block = Box::into_raw(Block::new()); + self.first.store(block, Ordering::Release); + self.tail.0.block.store(block, Ordering::Release); + } + } + let offset = (tail >> SHIFT) % Block::::LAP; + if offset == Block::::CAPACITY { + backoff.snooze(); + tail = self.tail.0.index.load(Ordering::Acquire); + block = self.tail.0.block.load(Ordering::Acquire); + continue; + } + + // The producer that reserves a block's final slot also installs its successor. Doing + // the allocation before the reservation keeps the boundary transition short. + if offset + 1 == Block::::CAPACITY && next_block.is_none() { + next_block = Some(Block::new()); + } + + let new_tail = tail.wrapping_add(STEP); + match self.tail.0.index.compare_exchange_weak( + tail, + new_tail, + Ordering::SeqCst, + Ordering::Acquire, + ) { + Ok(_) => { + // SAFETY: a successful reservation gives this producer exclusive ownership of + // `block.slots[offset]`. The receiver cannot reclaim the block until this slot + // publishes `ready` and is read in FIFO order. + unsafe { + if offset + 1 == Block::::CAPACITY { + let next = Box::into_raw(next_block.take().unwrap()); + (*block).next.store(next, Ordering::Release); + self.tail.0.block.store(next, Ordering::Release); + + // The reserved sentinel index keeps other producers and close cleanup + // from entering the next block before both pointers are installed. + self.tail.0.index.fetch_add(STEP, Ordering::Release); + } + + let slots = &(*block).slots; + let slot = slots.get_unchecked(offset); + (*slot.value.get()).write(value); + // This publication joins the notification gate's sequentially consistent + // order. If a producer checks the gate before the receiver arms it, the + // receiver's subsequent recheck must observe this completed publication. + slot.ready.store(true, Ordering::SeqCst); + } + return Ok(()); + } + Err(observed) => { + tail = observed; + block = self.tail.0.block.load(Ordering::Acquire); + backoff.spin(); + } + } + } + } + + /// Closes the producer side, waits for already-reserved slots, and discards buffered values. + pub fn close(&self, consumer: &mut Consumer) { + let mut backoff = Backoff::new(); + let mut tail = { + let _guard = self.initialize.lock(); + self.tail.0.index.fetch_or(CLOSED, Ordering::SeqCst) | CLOSED + }; + + // A producer at the sentinel owns the block transition. It reserved before close and must + // finish installing the next block before cleanup can traverse it. + while (tail >> SHIFT) % Block::::LAP == Block::::CAPACITY { + backoff.snooze(); + tail = self.tail.0.index.load(Ordering::Acquire); + } + + let mut cleanup = Cleanup { + queue: self, + consumer, + complete: false, + }; + cleanup.drain(); + cleanup.complete = true; + } +} + +impl Consumer { + /// Removes the next value, waiting only for a producer that already reserved its slot. + pub fn pop(&mut self, queue: &Queue) -> Pop { + if self.block.is_null() { + if self.index >> SHIFT == queue.tail.0.index.load(Ordering::SeqCst) >> SHIFT { + return Pop::Empty; + } + self.block = queue.first.load(Ordering::Acquire); + } + let offset = (self.index >> SHIFT) % Block::::LAP; + debug_assert!(offset < Block::::CAPACITY); + + // SAFETY: the consumer alone reads slots and frees blocks. A ready slot transfers the + // initialized value. Reading ready first keeps a draining consumer off the producer's + // tail cache line; the tail is needed only to distinguish empty from unpublished slots. + unsafe { + let slots = &(*self.block).slots; + let slot = slots.get_unchecked(offset); + if !slot.ready.load(Ordering::SeqCst) { + if self.index >> SHIFT == queue.tail.0.index.load(Ordering::SeqCst) >> SHIFT { + return Pop::Empty; + } + // A later completed send must not be hidden behind an Empty result. Wait for + // the earlier reservation to publish, as FIFO order requires. + let mut backoff = Backoff::new(); + while !slot.ready.load(Ordering::Acquire) { + backoff.snooze(); + } + } + + let value = (*slot.value.get()).assume_init_read(); + let new_index = self.index.wrapping_add(STEP); + + if offset + 1 == Block::::CAPACITY { + let old = self.block; + let next = (*old).next.load(Ordering::Acquire); + debug_assert!(!next.is_null()); + self.block = next; + self.index = new_index.wrapping_add(STEP); + + // Observing the last slot's publication also observes the producer's earlier next + // pointer publication. FIFO consumption means every producer using `old` has + // finished publishing before the consumer reaches this point. + drop(Box::from_raw(old)); + } else { + self.index = new_index; + } + + Pop::Value(value) + } + } + + fn finish(&mut self, queue: &Queue) { + // An initialized queue may still be empty if close won before the first reservation. + if self.block.is_null() { + self.block = queue.first.load(Ordering::Acquire); + } + if !self.block.is_null() { + // SAFETY: close prevents new reservations and every reserved value was consumed. + unsafe { drop(Box::from_raw(self.block)) }; + } + self.block = ptr::null_mut(); + queue.tail.0.block.store(ptr::null_mut(), Ordering::Release); + queue.first.store(ptr::null_mut(), Ordering::Release); + } +} + +struct Cleanup<'a, T> { + queue: &'a Queue, + consumer: &'a mut Consumer, + complete: bool, +} + +impl Cleanup<'_, T> { + fn drain(&mut self) { + loop { + match self.consumer.pop(self.queue) { + Pop::Value(value) => { + drop(value); + } + Pop::Empty => { + self.consumer.finish(self.queue); + return; + } + } + } + } +} + +impl Drop for Cleanup<'_, T> { + fn drop(&mut self) { + if !self.complete { + // Continue reclaiming if dropping a buffered value unwinds. A second destructor panic + // follows Rust's usual double-panic behavior and aborts the process. + self.drain(); + } + } +} + +struct Backoff { + step: u32, +} + +impl Backoff { + fn new() -> Self { + Self { step: 0 } + } + + fn spin(&mut self) { + let iterations = 1 << self.step.min(6); + for _ in 0..iterations { + hint::spin_loop(); + } + self.step = self.step.saturating_add(1); + } + + fn snooze(&mut self) { + if self.step <= 6 { + self.spin(); + } else { + thread::yield_now(); + self.step = self.step.saturating_add(1); + } + } +} + +// Queue contents are never exposed by shared reference; unwinding does not abandon a reserved +// slot because allocation precedes reservation and publication invokes no user callbacks. +impl RefUnwindSafe for Queue {} +impl UnwindSafe for Queue {} +impl RefUnwindSafe for Consumer {} +impl UnwindSafe for Consumer {} + +#[cfg(test)] +mod tests; diff --git a/asyncband/src/mpsc/unbounded/queue/tests.rs b/asyncband/src/mpsc/unbounded/queue/tests.rs new file mode 100644 index 0000000..72546ad --- /dev/null +++ b/asyncband/src/mpsc/unbounded/queue/tests.rs @@ -0,0 +1,96 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::atomic::Ordering; + +use super::BLOCK_BYTES; +use super::Block; +use super::Consumer; +use super::Pop; +use super::Queue; +use super::Slot; + +fn allocated_slots(queue: &Queue, consumer: &Consumer) -> usize { + let mut block = if consumer.block.is_null() { + queue.first.load(Ordering::Relaxed) + } else { + consumer.block + }; + let mut slots = 0; + while !block.is_null() { + // SAFETY: these tests inspect the live chain with no concurrent producers or receiver. + unsafe { + let block_slots = &(*block).slots; + slots += block_slots.len(); + block = (*block).next.load(Ordering::Relaxed); + } + } + slots +} + +#[test] +fn storage_is_lazy_and_partial_draining_reclaims_blocks_in_fifo_order() { + let (queue, mut consumer) = Queue::new(); + assert_eq!(allocated_slots(&queue, &consumer), 0); + for value in 0..1024 { + queue.push([value; 128]).unwrap(); + } + let peak = allocated_slots(&queue, &consumer); + for expected in 0..512 { + assert!(matches!(consumer.pop(&queue), Pop::Value(value) if value == [expected; 128])); + } + assert!(allocated_slots(&queue, &consumer) <= peak * 3 / 4); + queue.push([1024; 128]).unwrap(); + for expected in 512..=1024 { + assert!(matches!(consumer.pop(&queue), Pop::Value(value) if value == [expected; 128])); + } + assert!(matches!(consumer.pop(&queue), Pop::Empty)); + assert!(allocated_slots(&queue, &consumer) * size_of::>() <= BLOCK_BYTES); + queue.close(&mut consumer); + assert!(queue.first.load(Ordering::Relaxed).is_null()); +} + +#[test] +fn oversized_values_use_single_slot_blocks() { + let (queue, mut consumer) = Queue::new(); + assert_eq!(allocated_slots(&queue, &consumer), 0); + queue.push([7u8; BLOCK_BYTES + 1]).unwrap(); + assert_eq!(Block::<[u8; BLOCK_BYTES + 1]>::CAPACITY, 1); + assert!(matches!(consumer.pop(&queue), Pop::Value(value) if value == [7; BLOCK_BYTES + 1])); + assert_eq!(allocated_slots(&queue, &consumer), 1); + queue.close(&mut consumer); +} + +#[test] +fn positions_wrap_at_a_block_boundary() { + let (queue, mut consumer) = Queue::new(); + queue.push(0).unwrap(); + assert!(matches!(consumer.pop(&queue), Pop::Value(0))); + // Place an empty queue immediately before the next complete lap would wrap. Low position + // bits still identify slot one, and crossing the sentinel must preserve the closed bit. + let start = usize::MAX - 61; + queue.tail.0.index.store(start, Ordering::Relaxed); + consumer.index = start; + for value in 1..100 { + queue.push(value).unwrap(); + } + for expected in 1..100 { + assert!(matches!(consumer.pop(&queue), Pop::Value(value) if value == expected)); + } + queue.close(&mut consumer); + assert_eq!(queue.push(100), Err(100)); +} diff --git a/asyncband/src/mpsc/unbounded/receiver.rs b/asyncband/src/mpsc/unbounded/receiver.rs index 4b3ea76..5d91bc1 100644 --- a/asyncband/src/mpsc/unbounded/receiver.rs +++ b/asyncband/src/mpsc/unbounded/receiver.rs @@ -15,18 +15,16 @@ // specific language governing permissions and limitations // under the License. -use std::collections::VecDeque; use std::fmt; use std::future::poll_fn; -use std::mem; use std::sync::Arc; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; use super::State; -use super::buffer::Buffer; -use super::buffer::pop_batch; -use crate::internal::mutex::Mutex; +use super::queue::Consumer; +use super::queue::Pop; use crate::mpsc::RecvError; use crate::mpsc::TryRecvError; @@ -35,9 +33,8 @@ use crate::mpsc::TryRecvError; /// Instances are created by the [`unbounded`](crate::mpsc::unbounded) function. Dropping the /// receiver discards queued values and makes subsequent sends fail. pub struct UnboundedReceiver { - shared: Arc>>, - // Only accessed through `get_mut`; the mutex preserves Sync for Send-only payloads. - batch: Mutex>, + shared: Arc>, + consumer: Consumer, } impl fmt::Debug for UnboundedReceiver { @@ -48,26 +45,17 @@ impl fmt::Debug for UnboundedReceiver { impl Drop for UnboundedReceiver { fn drop(&mut self) { - let batch = mem::take(self.batch.get_mut()); - let (buffer, waker) = { - let mut state = self.shared.lock(); - state.receiver = false; - ( - mem::replace(&mut state.buffer, Buffer::new()), - state.recv_waker.take(), - ) - }; - // Destructors may send again. A waker may also own a sender and form an ownership cycle. - drop((batch, buffer, waker)); + let waker = self.shared.recv.take(); + // Closing precedes arbitrary callbacks. Keep the waker owned locally so even a payload + // destructor panic releases a waker that owns a sender and would otherwise form a cycle. + self.shared.queue.close(&mut self.consumer); + drop(waker); } } impl UnboundedReceiver { - pub(super) fn new(shared: Arc>>) -> Self { - Self { - shared, - batch: Mutex::new(VecDeque::new()), - } + pub(super) fn new(shared: Arc>, consumer: Consumer) -> Self { + Self { shared, consumer } } /// Attempts to receive the next queued value without waiting for a new message. @@ -93,19 +81,19 @@ impl UnboundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - let batch = self.batch.get_mut(); - if batch.is_empty() { - let mut state = self.shared.lock(); - state.buffer.refill(batch); - if batch.is_empty() { - return Err(if state.senders == 0 { - TryRecvError::Disconnected - } else { - TryRecvError::Empty - }); + match self.consumer.pop(&self.shared.queue) { + Pop::Value(value) => return Ok(value), + Pop::Empty => {} + } + if self.shared.senders.0.load(Ordering::SeqCst) == 0 { + // Acquire the last sender's publication before deciding the queue is drained. + match self.consumer.pop(&self.shared.queue) { + Pop::Value(value) => Ok(value), + Pop::Empty => Err(TryRecvError::Disconnected), } + } else { + Err(TryRecvError::Empty) } - Ok(pop_batch(batch)) } /// Waits for and receives the next value. @@ -142,30 +130,25 @@ impl UnboundedReceiver { } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - let batch = self.batch.get_mut(); - if !batch.is_empty() { - return Poll::Ready(Ok(pop_batch(batch))); - } - // Waker clone/drop callbacks can send into this channel, so run them outside the lock. - let waker = cx.waker().clone(); - let mut state = self.shared.lock(); - state.buffer.refill(batch); - if !batch.is_empty() { - drop(state); - return Poll::Ready(Ok(pop_batch(batch))); + match self.try_recv() { + Ok(value) => return Poll::Ready(Ok(value)), + Err(TryRecvError::Disconnected) => return Poll::Ready(Err(RecvError::Disconnected)), + Err(TryRecvError::Empty) => {} } - if state.senders == 0 { - let old = state.recv_waker.take(); - drop(state); - drop(old); - return Poll::Ready(Err(RecvError::Disconnected)); + self.shared.recv.register(cx.waker()); + match self.try_recv() { + Ok(value) => { + drop(self.shared.recv.take()); + Poll::Ready(Ok(value)) + } + Err(TryRecvError::Disconnected) => { + drop(self.shared.recv.take()); + Poll::Ready(Err(RecvError::Disconnected)) + } + Err(TryRecvError::Empty) => Poll::Pending, } - let old = state.recv_waker.replace(waker); - drop(state); - drop(old); - Poll::Pending } } -// No operation relies on a pinned location for the receiver batch or its values. +// Moving the receiver cannot move a value in its separately allocated queue storage. impl Unpin for UnboundedReceiver {} diff --git a/asyncband/src/mpsc/unbounded/sender.rs b/asyncband/src/mpsc/unbounded/sender.rs index a90589c..57a5e85 100644 --- a/asyncband/src/mpsc/unbounded/sender.rs +++ b/asyncband/src/mpsc/unbounded/sender.rs @@ -17,21 +17,21 @@ use std::fmt; use std::sync::Arc; +use std::sync::atomic::Ordering; use super::State; -use crate::internal::mutex::Mutex; use crate::mpsc::SendError; /// The sending endpoint of an unbounded mpsc channel. /// /// Instances are created by the [`unbounded`](crate::mpsc::unbounded) function. pub struct UnboundedSender { - shared: Arc>>, + shared: Arc>, } impl Clone for UnboundedSender { fn clone(&self) -> Self { - self.shared.lock().senders += 1; + self.shared.senders.0.fetch_add(1, Ordering::Relaxed); Self { shared: self.shared.clone(), } @@ -46,23 +46,15 @@ impl fmt::Debug for UnboundedSender { impl Drop for UnboundedSender { fn drop(&mut self) { - let wake = { - let mut state = self.shared.lock(); - state.senders -= 1; - if state.senders == 0 { - state.recv_waker.take() - } else { - None - } - }; - if let Some(waker) = wake { - waker.wake(); + // Disconnection participates in the same SC order as receiver registration. + if self.shared.senders.0.fetch_sub(1, Ordering::SeqCst) == 1 { + self.shared.recv.wake(); } } } impl UnboundedSender { - pub(super) fn new(shared: Arc>>) -> Self { + pub(super) fn new(shared: Arc>) -> Self { Self { shared } } @@ -72,17 +64,8 @@ impl UnboundedSender { /// been dropped, the returned error contains `value`. Success means the message was queued; /// it does not guarantee that the receiver will consume it before being dropped. pub fn send(&self, value: T) -> Result<(), SendError> { - let waker = { - let mut state = self.shared.lock(); - if !state.receiver { - return Err(SendError::new(value)); - } - state.buffer.push(value); - state.recv_waker.take() - }; - if let Some(waker) = waker { - waker.wake(); - } + self.shared.queue.push(value).map_err(SendError::new)?; + self.shared.recv.wake(); Ok(()) } } diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 19e5ac3..60835a5 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -47,6 +47,7 @@ asyncband = { workspace = true, features = [ "waitgroup", "watch", ] } +crossbeam-channel = { workspace = true } divan = { workspace = true } flume = { workspace = true, features = ["async"] } pollster = { workspace = true } diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs index f9fe81b..4d9203c 100644 --- a/benchmarks/ecosystem/mpsc/adapters.rs +++ b/benchmarks/ecosystem/mpsc/adapters.rs @@ -49,10 +49,13 @@ pub trait UnboundedMpsc: Send + Sync + 'static { fn send(sender: &Self::Sender, value: T); fn try_recv(receiver: &mut Self::Receiver) -> T; fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T; - fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; fn recv_blocking(receiver: &mut Self::Receiver) -> T; } +pub trait AsyncUnboundedMpsc: UnboundedMpsc { + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; +} + impl BoundedMpsc for Asyncband { type Receiver = asyncband::mpsc::BoundedReceiver; type Sender = asyncband::mpsc::BoundedSender; @@ -237,15 +240,17 @@ impl UnboundedMpsc for Asyncband { poll_ready(receiver.recv(), context).unwrap() } - async fn recv_async(receiver: &mut Self::Receiver) -> T { - receiver.recv().await.unwrap() - } - fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } +impl AsyncUnboundedMpsc for Asyncband { + async fn recv_async(receiver: &mut Self::Receiver) -> T { + receiver.recv().await.unwrap() + } +} + impl UnboundedMpsc for Tokio { type Receiver = tokio::sync::mpsc::UnboundedReceiver; type Sender = tokio::sync::mpsc::UnboundedSender; @@ -266,15 +271,17 @@ impl UnboundedMpsc for Tokio { poll_ready(receiver.recv(), context).unwrap() } - async fn recv_async(receiver: &mut Self::Receiver) -> T { - receiver.recv().await.unwrap() - } - fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } +impl AsyncUnboundedMpsc for Tokio { + async fn recv_async(receiver: &mut Self::Receiver) -> T { + receiver.recv().await.unwrap() + } +} + impl UnboundedMpsc for AsyncChannel { type Receiver = async_channel::Receiver; type Sender = async_channel::Sender; @@ -295,15 +302,17 @@ impl UnboundedMpsc for AsyncChannel { poll_ready(receiver.recv(), context).unwrap() } - async fn recv_async(receiver: &mut Self::Receiver) -> T { - receiver.recv().await.unwrap() - } - fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } +impl AsyncUnboundedMpsc for AsyncChannel { + async fn recv_async(receiver: &mut Self::Receiver) -> T { + receiver.recv().await.unwrap() + } +} + impl UnboundedMpsc for Flume { type Receiver = flume::Receiver; type Sender = flume::Sender; @@ -324,11 +333,40 @@ impl UnboundedMpsc for Flume { poll_ready(receiver.recv_async(), context).unwrap() } + fn recv_blocking(receiver: &mut Self::Receiver) -> T { + pollster::block_on(receiver.recv_async()).unwrap() + } +} + +impl AsyncUnboundedMpsc for Flume { async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv_async().await.unwrap() } +} + +pub struct Crossbeam; + +impl UnboundedMpsc for Crossbeam { + type Sender = crossbeam_channel::Sender; + type Receiver = crossbeam_channel::Receiver; + + fn channel() -> (Self::Sender, Self::Receiver) { + crossbeam_channel::unbounded() + } + + fn send(sender: &Self::Sender, value: T) { + sender.send(value).unwrap(); + } + + fn try_recv(receiver: &mut Self::Receiver) -> T { + receiver.try_recv().unwrap() + } + + fn recv_ready(receiver: &mut Self::Receiver, _context: &mut Context<'_>) -> T { + receiver.try_recv().unwrap() + } fn recv_blocking(receiver: &mut Self::Receiver) -> T { - pollster::block_on(receiver.recv_async()).unwrap() + receiver.recv().unwrap() } } diff --git a/benchmarks/ecosystem/mpsc/reservation.rs b/benchmarks/ecosystem/mpsc/reservation.rs index 32defce..7bc8e28 100644 --- a/benchmarks/ecosystem/mpsc/reservation.rs +++ b/benchmarks/ecosystem/mpsc/reservation.rs @@ -25,6 +25,7 @@ use divan::counter::ItemsCount; use super::adapters::Asyncband; use super::adapters::BoundedMpsc; use super::adapters::Tokio; +use super::support::AsyncConcurrentMpsc; use super::support::BATCH_MESSAGES; use super::support::ConcurrentMpsc; use super::support::RepeatedTasks; @@ -97,9 +98,13 @@ impl ConcurrentMpsc for Reserved usize { C::recv_blocking(receiver) } +} + +impl AsyncConcurrentMpsc for Reserved { async fn send_async(sender: &Self::Sender, value: usize) { C::publish(C::reserve(sender).await, value); } + async fn recv_async(receiver: &mut Self::Receiver) -> usize { C::recv_async(receiver).await } diff --git a/benchmarks/ecosystem/mpsc/support.rs b/benchmarks/ecosystem/mpsc/support.rs index d1d64c4..4d662da 100644 --- a/benchmarks/ecosystem/mpsc/support.rs +++ b/benchmarks/ecosystem/mpsc/support.rs @@ -26,6 +26,7 @@ use std::thread::JoinHandle; use divan::black_box; +use super::adapters::AsyncUnboundedMpsc; use super::adapters::BoundedMpsc; use super::adapters::UnboundedMpsc; @@ -70,6 +71,9 @@ pub trait ConcurrentMpsc: Send + Sync + 'static { fn channel() -> (Self::Sender, Self::Receiver); fn send(sender: &Self::Sender, value: Self::Message); fn recv(receiver: &mut Self::Receiver) -> Self::Message; +} + +pub trait AsyncConcurrentMpsc: ConcurrentMpsc { fn send_async(sender: &Self::Sender, value: Self::Message) -> impl Future + Send; fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; } @@ -93,6 +97,14 @@ impl, const CAPACITY: usize, T: Message> ConcurrentMpsc C::send_blocking(sender, value); } + fn recv(receiver: &mut Self::Receiver) -> T { + C::recv_blocking(receiver) + } +} + +impl, const CAPACITY: usize, T: Message> AsyncConcurrentMpsc + for Bounded +{ async fn send_async(sender: &Self::Sender, value: T) { C::send_async(sender, value).await; } @@ -100,10 +112,6 @@ impl, const CAPACITY: usize, T: Message> ConcurrentMpsc async fn recv_async(receiver: &mut Self::Receiver) -> T { C::recv_async(receiver).await } - - fn recv(receiver: &mut Self::Receiver) -> T { - C::recv_blocking(receiver) - } } pub struct Unbounded(PhantomData); @@ -121,6 +129,12 @@ impl ConcurrentMpsc for Unbounded { C::send(sender, value); } + fn recv(receiver: &mut Self::Receiver) -> usize { + C::recv_blocking(receiver) + } +} + +impl AsyncConcurrentMpsc for Unbounded { async fn send_async(sender: &Self::Sender, value: usize) { C::send(sender, value); } @@ -128,10 +142,6 @@ impl ConcurrentMpsc for Unbounded { async fn recv_async(receiver: &mut Self::Receiver) -> usize { C::recv_async(receiver).await } - - fn recv(receiver: &mut Self::Receiver) -> usize { - C::recv_blocking(receiver) - } } // Reuse worker threads and channel storage so steady-state samples exclude thread creation. @@ -209,7 +219,7 @@ enum Receiver { // Reuse every task and the channel. The small control exchange happens once per 16,384-message // batch; it never forwards measured messages. Both payload sizes use this same start protocol. -pub struct RepeatedTasks { +pub struct RepeatedTasks { runtime: tokio::runtime::Runtime, receiver: Receiver, start: Arc, @@ -217,7 +227,7 @@ pub struct RepeatedTasks { workers: Vec>, } -impl RepeatedTasks { +impl RepeatedTasks { pub fn new(producer_count: usize, worker_threads: usize) -> Self { Self::with_receiver(producer_count, worker_threads, false) } @@ -316,7 +326,7 @@ impl RepeatedTasks { } } -impl Drop for RepeatedTasks { +impl Drop for RepeatedTasks { fn drop(&mut self) { self.stop.store(true, Ordering::Release); self.runtime.block_on(async { @@ -333,7 +343,7 @@ impl Drop for RepeatedTasks { } } -async fn receive_batch(receiver: &mut C::Receiver) -> usize { +async fn receive_batch(receiver: &mut C::Receiver) -> usize { let mut checksum = 0usize; for _ in 0..BATCH_MESSAGES { checksum = checksum.wrapping_add(C::recv_async(receiver).await.sequence()); diff --git a/benchmarks/ecosystem/mpsc/unbounded.rs b/benchmarks/ecosystem/mpsc/unbounded.rs index a01531e..4aa78c6 100644 --- a/benchmarks/ecosystem/mpsc/unbounded.rs +++ b/benchmarks/ecosystem/mpsc/unbounded.rs @@ -27,7 +27,9 @@ use divan::black_box; use divan::counter::ItemsCount; use super::adapters::AsyncChannel; +use super::adapters::AsyncUnboundedMpsc; use super::adapters::Asyncband; +use super::adapters::Crossbeam; use super::adapters::Flume; use super::adapters::Tokio; use super::adapters::UnboundedMpsc; @@ -38,7 +40,7 @@ use super::support::RepeatedTasks; use super::support::Unbounded; use crate::support::bench_context; -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume])] fn ready_round_trip(bencher: Bencher) { let mut context = bench_context(); let (sender, mut receiver) = C::channel(); @@ -49,7 +51,7 @@ fn ready_round_trip(bencher: Bencher) { }); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume])] fn try_round_trip(bencher: Bencher) { let (sender, mut receiver) = C::channel(); @@ -60,7 +62,7 @@ fn try_round_trip(bencher: Bencher) { } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume], args = [32, 1024, 65_536], sample_count = 20, sample_size = 1, @@ -70,7 +72,7 @@ fn burst_drain(bencher: Bencher, messages: usize) { } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume], consts = [64, 1024], args = [1024], sample_count = 20, @@ -84,7 +86,7 @@ fn burst_drain_inline, const SIZE: usize>( } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume], args = [1024], sample_count = 20, sample_size = 1, @@ -95,7 +97,7 @@ fn burst_drain_boxed>>(bencher: Bencher, messag } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume], args = [1024, 65_536], sample_count = 20, sample_size = 1, @@ -130,7 +132,7 @@ fn repeated_bursts, T, F: Fn() -> T>( } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume], args = PRODUCER_COUNTS, sample_count = 50, sample_size = 1, @@ -142,7 +144,7 @@ fn sustained(bencher: Bencher, producer_count: usize) { bencher.bench_local(|| batch.run()); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume])] fn clone_drop_sender(bencher: Bencher) { let (sender, _receiver) = C::channel(); bencher.bench_local(|| drop(black_box(sender.clone()))); @@ -155,7 +157,7 @@ fn clone_drop_sender(bencher: Bencher) { sample_size = 1, counter = ItemsCount::new(BATCH_MESSAGES), )] -fn scheduled(bencher: Bencher, (producers, workers): (usize, usize)) { +fn scheduled(bencher: Bencher, (producers, workers): (usize, usize)) { let mut batch = RepeatedTasks::>::new(producers, workers); batch.run(); bencher.bench_local(|| batch.run()); @@ -168,7 +170,7 @@ fn scheduled(bencher: Bencher, (producers, workers): (usize, u sample_size = 1, counter = ItemsCount::new(BATCH_MESSAGES), )] -fn scheduled_bursts_inline>( +fn scheduled_bursts_inline>( bencher: Bencher, (producers, burst_messages): (usize, usize), ) { @@ -272,3 +274,17 @@ fn check_inline_message(value: [u8; 1024], expected: &mut [u64]) -> usize { expected[producer] += 1; usize::from(value[1023]) } + +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume])] +fn create(bencher: Bencher) { + bencher.bench_local(|| black_box(C::channel())); +} + +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Crossbeam, Flume])] +fn oneshot(bencher: Bencher) { + bencher.bench_local(|| { + let (sender, mut receiver) = C::channel(); + C::send(&sender, black_box(usize::MAX)); + black_box(C::try_recv(&mut receiver)) + }); +} diff --git a/licenserc.toml b/licenserc.toml index eec57c5..a79b63a 100644 --- a/licenserc.toml +++ b/licenserc.toml @@ -22,6 +22,7 @@ builtin = "Apache-2.0-ASF" # These third-party-derived files use the applicable upstream license and copyright notices instead # of the standard ASF source header. excludes = [ + "asyncband/src/mpsc/unbounded/queue.rs", "asyncband/src/blocking/executor.rs", "asyncband/src/blocking/parker.rs", "asyncband/src/pool/bounded.rs", diff --git a/tests-integration/tests/mpsc_test/callbacks.rs b/tests-integration/tests/mpsc_test/callbacks.rs index 772c368..fa3cf4d 100644 --- a/tests-integration/tests/mpsc_test/callbacks.rs +++ b/tests-integration/tests/mpsc_test/callbacks.rs @@ -298,7 +298,7 @@ fn unbounded_replaced_and_disconnected_wakers_can_send() { } #[test] -fn unbounded_disconnect_drops_partial_and_queued_batches_outside_lock() { +fn unbounded_disconnect_drops_buffered_values_outside_lock() { struct Value(Option>); impl Drop for Value { fn drop(&mut self) { @@ -412,3 +412,31 @@ fn bounded_disconnect_finishes_cleanup_when_a_callback_panics() { } } } + +#[test] +fn unbounded_disconnect_finishes_cleanup_after_a_payload_panics() { + struct Payload { + id: usize, + drops: Arc<[AtomicUsize; 40]>, + } + impl Drop for Payload { + fn drop(&mut self) { + self.drops[self.id].fetch_add(1, Ordering::Relaxed); + assert_ne!(self.id, 2, "payload destructor panic"); + } + } + let drops = Arc::new(std::array::from_fn(|_| AtomicUsize::new(0))); + let (tx, rx) = mpsc::unbounded(); + for id in 0..40 { + assert!( + tx.send(Payload { + id, + drops: drops.clone() + }) + .is_ok() + ); + } + assert!(std::panic::catch_unwind(|| drop(rx)).is_err()); + assert!(drops.iter().all(|count| count.load(Ordering::Relaxed) == 1)); + assert!(tx.send(Payload { id: 39, drops }).is_err()); +} diff --git a/tests-integration/tests/mpsc_test/concurrency.rs b/tests-integration/tests/mpsc_test/concurrency.rs index 8937ec4..a448bb0 100644 --- a/tests-integration/tests/mpsc_test/concurrency.rs +++ b/tests-integration/tests/mpsc_test/concurrency.rs @@ -32,7 +32,6 @@ use asyncband::mpsc::TrySendError; use tests_integration::WakeCounter; use tests_integration::poll_once; use tests_integration::poll_with; -use tests_integration::test_runtime; use tokio_test::assert_ok; #[test] @@ -271,28 +270,65 @@ fn last_sender_drop_racing_with_receiver_registration_cannot_lose_wakeup() { } #[test] -#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] fn unbounded_collects_from_multiple_producers() { let (tx, mut rx) = mpsc::unbounded(); - - test_runtime().block_on(async move { - for i in 0..8 { + let messages = if cfg!(miri) { 32 } else { 128 }; + let start = Barrier::new(5); + thread::scope(|scope| { + for producer in 0..4 { let tx = tx.clone(); - tokio::spawn(async move { - tx.send(i).unwrap(); + let start = &start; + scope.spawn(move || { + start.wait(); + for sequence in 0..messages { + tx.send((producer, sequence)).unwrap(); + } }); } drop(tx); - - let mut received = Vec::new(); - while let Ok(i) = rx.recv().await { - received.push(i); + start.wait(); + let mut next = [0; 4]; + while let Ok((producer, sequence)) = pollster::block_on(rx.recv()) { + assert_eq!(sequence, next[producer]); + next[producer] += 1; } - received.sort_unstable(); - assert_eq!(received, (0..8).collect::>()); + assert_eq!(next, [messages; 4]); }); } +#[test] +fn unbounded_first_publication_racing_with_close_drops_every_payload_once() { + struct Payload(Arc); + impl Drop for Payload { + fn drop(&mut self) { + assert_eq!(self.0.fetch_add(1, Ordering::Relaxed), 0); + } + } + let messages = if cfg!(miri) { 4 } else { 128 }; + for _ in 0..if cfg!(miri) { 2 } else { 16 } { + let (tx, rx) = mpsc::unbounded(); + let drops: Vec<_> = (0..4 * messages) + .map(|_| Arc::new(AtomicUsize::new(0))) + .collect(); + let start = Barrier::new(5); + thread::scope(|scope| { + for values in drops.chunks(messages) { + let tx = tx.clone(); + let start = &start; + scope.spawn(move || { + start.wait(); + for count in values { + drop(tx.send(Payload(count.clone()))); + } + }); + } + start.wait(); + drop(rx); + }); + assert!(drops.iter().all(|count| count.load(Ordering::Relaxed) == 1)); + } +} + #[tokio::test] #[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] async fn bounded_backpressure_progresses_on_an_executor() {