diff --git a/CHANGELOG.md b/CHANGELOG.md index 9543c02..eb8fc85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,3 +39,4 @@ All notable changes to this project will be documented in this file. * Remove the `slab` dependency in favor of a focused internal waiter arena. * Describe disconnected channel states consistently in channel error messages. +* Replace the legacy standard-library MPSC backend with an owned queue core and remove the receiver types' manual `Sync` implementations. diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index 4f888f3..5b49328 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -32,6 +32,8 @@ use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; +use super::queue::BoundedQueue; +use super::queue::PushError; use crate::internal::atomic_waker::AtomicWaker; use crate::internal::semaphore::Acquire; use crate::internal::semaphore::Semaphore; @@ -50,23 +52,20 @@ use crate::internal::semaphore::Semaphore; pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); let state = Arc::new(BoundedState { + queue: BoundedQueue::new(buffer), senders: AtomicUsize::new(1), tx_permits: Semaphore::new(0), rx_waker: AtomicWaker::new(), }); - let (sender, receiver) = std::sync::mpsc::sync_channel(buffer); let sender = BoundedSender { state: state.clone(), - sender: Some(sender), - }; - let receiver = BoundedReceiver { - state: state.clone(), - receiver: Some(receiver), }; + let receiver = BoundedReceiver { state }; (sender, receiver) } -struct BoundedState { +struct BoundedState { + queue: BoundedQueue, senders: AtomicUsize, tx_permits: Semaphore, rx_waker: AtomicWaker, @@ -76,8 +75,7 @@ struct BoundedState { /// /// Instances are created by the [`bounded`] function. pub struct BoundedSender { - state: Arc, - sender: Option>, + state: Arc>, } impl Clone for BoundedSender { @@ -85,7 +83,6 @@ impl Clone for BoundedSender { self.state.senders.fetch_add(1, Ordering::Release); BoundedSender { state: self.state.clone(), - sender: self.sender.clone(), } } } @@ -98,9 +95,6 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { - // Dropping the final underlying sender disconnects the channel. - drop(self.sender.take()); - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { 1 => { // Wake the receiver so it can observe the channel's disconnected state. @@ -197,18 +191,14 @@ impl BoundedSender { /// # } /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - // INVARIANT: A shared borrow of the endpoint cannot overlap its destructor. - let sender = self.sender.as_ref().unwrap(); - match sender.try_send(value) { + match self.state.queue.try_push(value) { Ok(()) => { self.state.rx_waker.wake(); Ok(()) } - Err(std::sync::mpsc::TrySendError::Full(value)) => Err(TrySendError::Full(value)), - Err(std::sync::mpsc::TrySendError::Disconnected(value)) => { - Err(TrySendError::Disconnected(value)) - } + Err(PushError::Full(value)) => Err(TrySendError::Full(value)), + Err(PushError::Disconnected(value)) => Err(TrySendError::Disconnected(value)), } } } @@ -217,14 +207,9 @@ impl BoundedSender { /// /// Instances are created by the [`bounded`] function. pub struct BoundedReceiver { - state: Arc, - receiver: Option>, + state: Arc>, } -/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`. -/// That is, `BoundedReceiver` can only be accessed by one thread at a time. -unsafe impl Sync for BoundedReceiver {} - impl fmt::Debug for BoundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BoundedReceiver").finish_non_exhaustive() @@ -233,7 +218,7 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - drop(self.receiver.take()); + self.state.queue.disconnect_receiver(); self.state.tx_permits.notify_all(); } } @@ -274,15 +259,20 @@ impl BoundedReceiver { /// # } /// ``` pub fn try_recv(&mut self) -> Result { - // INVARIANT: A mutable borrow of the endpoint cannot overlap its destructor. - let receiver = self.receiver.as_ref().unwrap(); - match receiver.try_recv() { - Ok(v) => { + if let Some(value) = self.state.queue.pop() { + self.state.tx_permits.release_if_nonempty(1); + Ok(value) + } else if self.state.senders.load(Ordering::Acquire) == 0 { + // The final sender can enqueue between the first empty observation and decrementing + // the sender count, so check the queue again before reporting disconnection. + if let Some(value) = self.state.queue.pop() { self.state.tx_permits.release_if_nonempty(1); - Ok(v) + Ok(value) + } else { + Err(TryRecvError::Disconnected) } - Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected), - Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty), + } else { + Err(TryRecvError::Empty) } } diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs index 040c98b..30ef65e 100644 --- a/asyncband/src/mpsc/mod.rs +++ b/asyncband/src/mpsc/mod.rs @@ -23,6 +23,7 @@ mod bounded; mod error; +mod queue; mod unbounded; pub use self::bounded::BoundedReceiver; diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs new file mode 100644 index 0000000..b0838cd --- /dev/null +++ b/asyncband/src/mpsc/queue.rs @@ -0,0 +1,352 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cell::UnsafeCell; +use std::collections::VecDeque; +use std::hint::spin_loop; +use std::mem; +use std::mem::MaybeUninit; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::sync::atomic::fence; + +use crate::internal::mutex::Mutex; + +pub(super) struct UnboundedQueue { + inner: Mutex>, +} + +struct UnboundedInner { + // Storage and receiver liveness share one lock so a send is linearized either before receiver + // disconnection, with its value in the queue, or after it, with the value returned to sender. + messages: VecDeque, + receiver_alive: bool, +} + +pub(super) struct UnboundedConsumer { + local: Mutex>, +} + +pub(super) enum PushError { + Full(T), + Disconnected(T), +} + +impl UnboundedQueue { + pub(super) const fn new() -> Self { + Self { + inner: Mutex::new(UnboundedInner { + messages: VecDeque::new(), + receiver_alive: true, + }), + } + } + + pub(super) fn push(&self, value: T) -> Result<(), PushError> { + let mut inner = self.inner.lock(); + if !inner.receiver_alive { + return Err(PushError::Disconnected(value)); + } + inner.messages.push_back(value); + Ok(()) + } + + pub(super) fn pop(&self, consumer: &UnboundedConsumer) -> Option { + let mut local = consumer.local.lock(); + if local.is_empty() { + let mut inner = self.inner.lock(); + mem::swap(&mut *local, &mut inner.messages); + } + local.pop_front() + } + + pub(super) fn disconnect_receiver(&self, consumer: &UnboundedConsumer) { + let (local, shared) = { + let mut local = consumer.local.lock(); + let mut inner = self.inner.lock(); + inner.receiver_alive = false; + (mem::take(&mut *local), mem::take(&mut inner.messages)) + }; + drop((local, shared)); + } +} + +impl UnboundedConsumer { + pub(super) const fn new() -> Self { + Self { + local: Mutex::new(VecDeque::new()), + } + } +} + +pub(super) struct BoundedQueue { + slots: Box<[Slot]>, + head: CachePadded, + tail: CachePadded, + capacity: usize, + one_lap: usize, + mark_bit: usize, +} + +#[repr(align(64))] +struct CachePadded(T); + +impl std::ops::Deref for CachePadded { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +struct Slot { + stamp: AtomicUsize, + value: UnsafeCell>, +} + +// SAFETY: A successful tail CAS gives one producer exclusive access to a slot. That producer +// initializes the value before publishing the next stamp with Release ordering. The single +// consumer reads only after acquiring that stamp and publishes the following lap before reuse. +unsafe impl Sync for Slot {} + +impl BoundedQueue { + pub(super) fn new(capacity: usize) -> Self { + assert!(capacity <= usize::MAX / 4, "mpsc capacity is too large"); + let mark_bit = (capacity + 1).next_power_of_two(); + let one_lap = mark_bit * 2; + let slots = (0..capacity) + .map(|index| Slot { + stamp: AtomicUsize::new(index), + value: UnsafeCell::new(MaybeUninit::uninit()), + }) + .collect(); + Self { + slots, + head: CachePadded(AtomicUsize::new(0)), + tail: CachePadded(AtomicUsize::new(0)), + capacity, + one_lap, + mark_bit, + } + } + + pub(super) fn try_push(&self, value: T) -> Result<(), PushError> { + let mut tail = self.tail.load(Ordering::Relaxed); + let mut backoff = 0; + loop { + if tail & self.mark_bit != 0 { + return Err(PushError::Disconnected(value)); + } + + let index = tail & (self.mark_bit - 1); + let slot = &self.slots[index]; + let stamp = slot.stamp.load(Ordering::Acquire); + if stamp == tail { + let next_tail = self.advance(tail); + match self.tail.compare_exchange_weak( + tail, + next_tail, + Ordering::SeqCst, + Ordering::Relaxed, + ) { + Ok(_) => { + // SAFETY: The successful CAS reserved this slot exclusively, and its + // matching stamp proves the consumer completed its previous lap. + unsafe { (*slot.value.get()).write(value) }; + slot.stamp.store(tail.wrapping_add(1), Ordering::Release); + return Ok(()); + } + Err(actual) => tail = actual, + } + } else if stamp.wrapping_add(self.one_lap) == tail.wrapping_add(1) { + fence(Ordering::SeqCst); + if self.head.load(Ordering::Relaxed).wrapping_add(self.one_lap) == tail { + return Err(PushError::Full(value)); + } + tail = self.tail.load(Ordering::Relaxed); + } else { + tail = self.tail.load(Ordering::Relaxed); + } + Self::spin(&mut backoff); + } + } + + pub(super) fn pop(&self) -> Option { + let mut head = self.head.load(Ordering::Relaxed); + let mut backoff = 0; + loop { + let index = head & (self.mark_bit - 1); + let slot = &self.slots[index]; + let stamp = slot.stamp.load(Ordering::Acquire); + if stamp == head.wrapping_add(1) { + let next_head = self.advance(head); + // SAFETY: Acquiring the matching stamp observes initialization by the producer. + // There is one consumer, so the value is read exactly once. + let value = unsafe { (*slot.value.get()).assume_init_read() }; + slot.stamp + .store(head.wrapping_add(self.one_lap), Ordering::Release); + self.head.store(next_head, Ordering::SeqCst); + return Some(value); + } + + if stamp == head { + fence(Ordering::SeqCst); + if self.tail.load(Ordering::Relaxed) & !self.mark_bit == head { + return None; + } + } + if backoff == 8 { + return None; + } + Self::spin(&mut backoff); + head = self.head.load(Ordering::Relaxed); + } + } + + pub(super) fn disconnect_receiver(&self) { + let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; + self.discard_until(tail); + } + + fn advance(&self, position: usize) -> usize { + let index = position & (self.mark_bit - 1); + if index + 1 < self.capacity { + position + 1 + } else { + let lap = position & !(self.one_lap - 1); + lap.wrapping_add(self.one_lap) + } + } + + fn discard_until(&self, tail: usize) { + let mut head = self.head.load(Ordering::Relaxed); + let mut backoff = 0; + while head != tail { + let index = head & (self.mark_bit - 1); + let slot = &self.slots[index]; + if slot.stamp.load(Ordering::Acquire) == head.wrapping_add(1) { + let next_head = self.advance(head); + // Move the head before dropping the value so unwinding cannot drop it twice. + slot.stamp + .store(head.wrapping_add(self.one_lap), Ordering::Release); + self.head.store(next_head, Ordering::SeqCst); + // SAFETY: The acquired matching stamp proves the slot contains an initialized + // value, and advancing the single-consumer head claims it exactly once. + unsafe { (*slot.value.get()).assume_init_drop() }; + head = next_head; + backoff = 0; + } else { + Self::spin(&mut backoff); + } + } + } + + fn spin(step: &mut u32) { + for _ in 0..(*step).min(6).pow(2) { + spin_loop(); + } + *step = (*step).saturating_add(1); + } +} + +impl Drop for BoundedQueue { + fn drop(&mut self) { + let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; + self.discard_until(tail); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::thread; + + use super::BoundedQueue; + use super::PushError; + use super::UnboundedConsumer; + use super::UnboundedQueue; + + #[test] + fn bounded_queue_preserves_capacity_and_fifo_order() { + let queue = BoundedQueue::new(3); + for value in 0..3 { + assert!(queue.try_push(value).is_ok()); + } + assert!(matches!(queue.try_push(3), Err(PushError::Full(3)))); + for value in 0..3 { + assert_eq!(queue.pop(), Some(value)); + } + assert_eq!(queue.pop(), None); + + for value in 3..12 { + assert!(queue.try_push(value).is_ok()); + assert_eq!(queue.pop(), Some(value)); + } + } + + #[test] + fn bounded_queue_coordinates_multiple_producers() { + let queue = Arc::new(BoundedQueue::new(4)); + let producers: Vec<_> = (0..2) + .map(|producer| { + let queue = queue.clone(); + thread::spawn(move || { + for offset in 0..32 { + let mut value = producer * 32 + offset; + loop { + match queue.try_push(value) { + Ok(()) => break, + Err(PushError::Full(returned)) => { + value = returned; + thread::yield_now(); + } + Err(PushError::Disconnected(_)) => panic!("queue disconnected"), + } + } + } + }) + }) + .collect(); + + let mut values = Vec::new(); + while values.len() < 64 { + if let Some(value) = queue.pop() { + values.push(value); + } else { + thread::yield_now(); + } + } + for producer in producers { + producer.join().unwrap(); + } + values.sort_unstable(); + assert_eq!(values, (0..64).collect::>()); + } + + #[test] + fn unbounded_queue_batches_without_reordering() { + let queue = UnboundedQueue::new(); + let consumer = UnboundedConsumer::new(); + assert!(queue.push(1).is_ok()); + assert!(queue.push(2).is_ok()); + assert_eq!(queue.pop(&consumer), Some(1)); + assert!(queue.push(3).is_ok()); + assert_eq!(queue.pop(&consumer), Some(2)); + assert_eq!(queue.pop(&consumer), Some(3)); + assert_eq!(queue.pop(&consumer), None); + } +} diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index ed6eb56..8533a4d 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -29,6 +29,9 @@ use std::task::Poll; use super::RecvError; use super::SendError; use super::TryRecvError; +use super::queue::PushError; +use super::queue::UnboundedConsumer; +use super::queue::UnboundedQueue; use crate::internal::atomic_waker::AtomicWaker; /// Creates an unbounded mpsc channel for communicating between asynchronous @@ -42,22 +45,22 @@ use crate::internal::atomic_waker::AtomicWaker; /// process to run out of memory. In this case, the process will be aborted. pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { + queue: UnboundedQueue::new(), senders: AtomicUsize::new(1), rx_waker: AtomicWaker::new(), }); - let (sender, receiver) = std::sync::mpsc::channel(); let sender = UnboundedSender { state: state.clone(), - sender: Some(sender), }; let receiver = UnboundedReceiver { - state: state.clone(), - receiver, + state, + consumer: UnboundedConsumer::new(), }; (sender, receiver) } -struct UnboundedState { +struct UnboundedState { + queue: UnboundedQueue, senders: AtomicUsize, rx_waker: AtomicWaker, } @@ -66,8 +69,7 @@ struct UnboundedState { /// /// Instances are created by the [`unbounded`] function. pub struct UnboundedSender { - state: Arc, - sender: Option>, + state: Arc>, } impl Clone for UnboundedSender { @@ -75,7 +77,6 @@ impl Clone for UnboundedSender { self.state.senders.fetch_add(1, Ordering::Release); UnboundedSender { state: self.state.clone(), - sender: self.sender.clone(), } } } @@ -88,9 +89,6 @@ impl fmt::Debug for UnboundedSender { impl Drop for UnboundedSender { fn drop(&mut self) { - // Dropping the final underlying sender disconnects the channel. - drop(self.sender.take()); - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { 1 => { // Wake the receiver so it can observe the channel's disconnected state. @@ -113,9 +111,11 @@ impl UnboundedSender { /// If the receiver has been dropped, this function returns an error. The error includes /// the value passed to `send`. pub fn send(&self, value: T) -> Result<(), SendError> { - // INVARIANT: A shared borrow of the endpoint cannot overlap its destructor. - let sender = self.sender.as_ref().unwrap(); - sender.send(value).map_err(|err| SendError::new(err.0))?; + match self.state.queue.push(value) { + Ok(()) => {} + Err(PushError::Disconnected(value)) => return Err(SendError::new(value)), + Err(PushError::Full(_)) => unreachable!("unbounded queue cannot be full"), + } self.state.rx_waker.wake(); @@ -127,20 +127,22 @@ impl UnboundedSender { /// /// Instances are created by the [`unbounded`] function. pub struct UnboundedReceiver { - state: Arc, - receiver: std::sync::mpsc::Receiver, + state: Arc>, + consumer: UnboundedConsumer, } -/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`. -/// That is, `UnboundedReceiver` can only be accessed by one thread at a time. -unsafe impl Sync for UnboundedReceiver {} - impl fmt::Debug for UnboundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("UnboundedReceiver").finish_non_exhaustive() } } +impl Drop for UnboundedReceiver { + fn drop(&mut self) { + self.state.queue.disconnect_receiver(&self.consumer); + } +} + impl UnboundedReceiver { /// Tries to receive the next value for this receiver. /// @@ -176,10 +178,17 @@ impl UnboundedReceiver { /// # } /// ``` pub fn try_recv(&mut self) -> Result { - match self.receiver.try_recv() { - Ok(v) => Ok(v), - Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected), - Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty), + if let Some(value) = self.state.queue.pop(&self.consumer) { + Ok(value) + } else if self.state.senders.load(Ordering::Acquire) == 0 { + // The final sender can enqueue between the first empty observation and decrementing + // the sender count, so check the queue again before reporting disconnection. + self.state + .queue + .pop(&self.consumer) + .ok_or(TryRecvError::Disconnected) + } else { + Err(TryRecvError::Empty) } } diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 50d5a10..5c696f2 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -176,6 +176,44 @@ fn unbounded_try_recv_preserves_order_and_reports_state() { assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); } +#[test] +fn cancelled_receive_does_not_consume_a_later_message() { + let (unbounded_tx, mut unbounded_rx) = mpsc::unbounded(); + { + let mut receive = Box::pin(unbounded_rx.recv()); + assert!(poll_once(receive.as_mut()).is_pending()); + } + unbounded_tx.send(1).unwrap(); + assert_eq!(unbounded_rx.try_recv(), Ok(1)); + + let (bounded_tx, mut bounded_rx) = mpsc::bounded(1); + { + let mut receive = Box::pin(bounded_rx.recv()); + assert!(poll_once(receive.as_mut()).is_pending()); + } + bounded_tx.try_send(2).unwrap(); + assert_eq!(bounded_rx.try_recv(), Ok(2)); +} + +#[test] +fn buffered_messages_are_drained_before_disconnection() { + let (unbounded_tx, mut unbounded_rx) = mpsc::unbounded(); + unbounded_tx.send(1).unwrap(); + unbounded_tx.send(2).unwrap(); + drop(unbounded_tx); + assert_eq!(unbounded_rx.try_recv(), Ok(1)); + assert_eq!(unbounded_rx.try_recv(), Ok(2)); + assert_eq!(unbounded_rx.try_recv(), Err(TryRecvError::Disconnected)); + + let (bounded_tx, mut bounded_rx) = mpsc::bounded(2); + bounded_tx.try_send(3).unwrap(); + bounded_tx.try_send(4).unwrap(); + drop(bounded_tx); + assert_eq!(bounded_rx.try_recv(), Ok(3)); + assert_eq!(bounded_rx.try_recv(), Ok(4)); + assert_eq!(bounded_rx.try_recv(), Err(TryRecvError::Disconnected)); +} + #[tokio::test] async fn send_recv_bounded() { let (tx, mut rx) = mpsc::bounded(1);