From 39ca0bb96363516dd000ec8db0c7dc5e7391476d Mon Sep 17 00:00:00 2001 From: Sean McArthur Date: Wed, 13 May 2026 16:08:00 -0400 Subject: [PATCH] mpsc: add array-backed internals for bounded channel --- benches/sync_mpsc.rs | 70 +++++- tokio/src/sync/mpsc/array.rs | 362 ++++++++++++++++++++++++++++++ tokio/src/sync/mpsc/bounded.rs | 342 ++++++++++++++++++++++++++-- tokio/src/sync/mpsc/chan.rs | 161 +++++++------ tokio/src/sync/mpsc/list.rs | 58 +++-- tokio/src/sync/mpsc/mod.rs | 20 ++ tokio/src/sync/mpsc/unbounded.rs | 13 +- tokio/src/sync/tests/loom_mpsc.rs | 23 ++ tokio/tests/sync_mpsc.rs | 14 ++ 9 files changed, 954 insertions(+), 109 deletions(-) create mode 100644 tokio/src/sync/mpsc/array.rs diff --git a/benches/sync_mpsc.rs b/benches/sync_mpsc.rs index 5a3f75001..e220375c9 100644 --- a/benches/sync_mpsc.rs +++ b/benches/sync_mpsc.rs @@ -19,6 +19,15 @@ impl Default for Large { } } +#[cfg(target_pointer_width = "64")] +const ARRAY_CAP: usize = 32; + +#[cfg(not(target_pointer_width = "64"))] +const ARRAY_CAP: usize = 16; + +const LIST_CAP: usize = ARRAY_CAP + 1; +const ROUNDTRIP_ITERS: usize = 1_000; + fn rt() -> tokio::runtime::Runtime { tokio::runtime::Builder::new_multi_thread() .worker_threads(6) @@ -34,6 +43,14 @@ fn create_medium(g: &mut BenchmarkGroup) { }); } +fn create_data(g: &mut BenchmarkGroup, prefix: &str) { + g.bench_function(format!("{prefix}_{SIZE}"), |b| { + b.iter(|| { + black_box(mpsc::channel::(SIZE)); + }) + }); +} + fn send_data(g: &mut BenchmarkGroup, prefix: &str) { let rt = rt(); @@ -48,6 +65,22 @@ fn send_data(g: &mut BenchmarkGroup, pr }); } +fn roundtrip_try_send_recv_data( + g: &mut BenchmarkGroup, + prefix: &str, +) { + let (tx, mut rx) = mpsc::channel::(SIZE); + + g.bench_function(format!("{prefix}_{SIZE}"), |b| { + b.iter(|| { + for _ in 0..ROUNDTRIP_ITERS { + tx.try_send(T::default()).unwrap(); + black_box(rx.try_recv().unwrap()); + } + }) + }); +} + fn contention_bounded(g: &mut BenchmarkGroup) { let rt = rt(); @@ -296,13 +329,39 @@ fn bench_create_medium(c: &mut Criterion) { group.finish(); } +fn bench_create_small(c: &mut Criterion) { + let mut group = c.benchmark_group("create_small"); + create_data::(&mut group, "medium"); + create_data::(&mut group, "medium"); + create_data::(&mut group, "medium"); + create_data::(&mut group, "medium"); + create_data::(&mut group, "large"); + create_data::(&mut group, "large"); + create_data::(&mut group, "large"); + create_data::(&mut group, "large"); + group.finish(); +} + fn bench_send(c: &mut Criterion) { let mut group = c.benchmark_group("send"); + send_data::(&mut group, "medium"); + send_data::(&mut group, "medium"); send_data::(&mut group, "medium"); + send_data::(&mut group, "large"); + send_data::(&mut group, "large"); send_data::(&mut group, "large"); group.finish(); } +fn bench_roundtrip_try_send_recv(c: &mut Criterion) { + let mut group = c.benchmark_group("roundtrip_try_send_recv"); + roundtrip_try_send_recv_data::(&mut group, "medium"); + roundtrip_try_send_recv_data::(&mut group, "medium"); + roundtrip_try_send_recv_data::(&mut group, "medium"); + roundtrip_try_send_recv_data::(&mut group, "medium"); + group.finish(); +} + fn bench_contention(c: &mut Criterion) { let mut group = c.benchmark_group("contention"); contention_bounded(&mut group); @@ -324,8 +383,17 @@ fn bench_uncontented(c: &mut Criterion) { } criterion_group!(create, bench_create_medium); +criterion_group!(create_small, bench_create_small); criterion_group!(send, bench_send); +criterion_group!(roundtrip_try_send_recv, bench_roundtrip_try_send_recv); criterion_group!(contention, bench_contention); criterion_group!(uncontented, bench_uncontented); -criterion_main!(create, send, contention, uncontented); +criterion_main!( + create, + create_small, + send, + roundtrip_try_send_recv, + contention, + uncontented +); diff --git a/tokio/src/sync/mpsc/array.rs b/tokio/src/sync/mpsc/array.rs new file mode 100644 index 000000000..301013e9f --- /dev/null +++ b/tokio/src/sync/mpsc/array.rs @@ -0,0 +1,362 @@ +//! A concurrent, lock-free, FIFO ring buffer for bounded channels. +//! +//! The semaphore in `bounded.rs` guarantees that senders only claim a slot when +//! the ring has capacity, so receiver-side slot reuse is coordinated by permit +//! release rather than by allocating new blocks. + +use crate::loom::cell::UnsafeCell; +use crate::loom::sync::atomic::AtomicUsize; +use crate::loom::sync::Arc; +use crate::sync::mpsc::block::Read; +use crate::sync::mpsc::{chan, TryPopResult}; + +use std::fmt; +use std::mem::MaybeUninit; +use std::ptr; +use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; + +const OPEN: usize = 0; +const CLOSED: usize = 1; + +pub(crate) struct Tx { + /// Shared ring buffer slots. + shared: Arc>, + + /// Position to push the next message. + tail_position: AtomicUsize, + + /// Tracks whether all sender handles have been dropped. + closed: AtomicUsize, +} + +pub(crate) struct Rx { + /// Shared ring buffer slots. + shared: Arc>, + + /// Next slot index to process. + index: usize, +} + +#[derive(Debug)] +pub(crate) struct Queue; + +struct Shared { + /// Ring buffer slots. + slots: Box<[Slot]>, + + /// Bitfield tracking slots that are ready to have their values consumed. + /// + /// The array will never be used for more than `super::BLOCK_CAP` values. + ready_slots: AtomicUsize, + + /// Receiver position, used for diagnostics. + head_position: AtomicUsize, + + /// Number of slots in the ring buffer. + capacity: usize, + + /// The position after the last valid queue position. + /// + /// Positions are kept in the range `0..position_wrap`. Using two laps of + /// the ring distinguishes an empty queue from a full one while avoiding the + /// discontinuity that integer overflow would introduce for capacities that + /// are not powers of two. + position_wrap: usize, +} + +struct Slot { + /// Storage for the slot value. + value: UnsafeCell>, +} + +unsafe impl Send for Shared {} +unsafe impl Sync for Shared {} + +impl chan::Queue for Queue { + type Tx = Tx; + type Rx = Rx; + + fn channel(bound: usize) -> (Self::Tx, Self::Rx) { + channel(bound) + } + + fn push(tx: &Self::Tx, value: T) { + tx.push(value); + } + + fn close(tx: &Self::Tx) { + tx.close(); + } + + fn is_empty(rx: &Self::Rx, tx: &Self::Tx) -> bool { + rx.is_empty(tx) + } + + fn len(rx: &Self::Rx, tx: &Self::Tx) -> usize { + rx.len(tx) + } + + fn pop(rx: &mut Self::Rx, tx: &Self::Tx) -> Option> { + rx.pop(tx) + } + + fn try_pop(rx: &mut Self::Rx, tx: &Self::Tx) -> TryPopResult { + rx.try_pop(tx) + } + + unsafe fn free(_rx: &mut Self::Rx) { + // The array queue owns its slots through `Arc>`, so there are + // no receiver-owned blocks to free. + } +} + +pub(crate) fn channel(capacity: usize) -> (Tx, Rx) { + debug_assert!(capacity > 0); + debug_assert!(capacity <= usize::BITS as usize); + + let slots = (0..capacity) + .map(|_| Slot::new()) + .collect::>() + .into_boxed_slice(); + let shared = Arc::new(Shared { + slots, + ready_slots: AtomicUsize::new(0), + head_position: AtomicUsize::new(0), + capacity, + position_wrap: 2 * capacity, + }); + + let tx = Tx { + shared: shared.clone(), + tail_position: AtomicUsize::new(0), + closed: AtomicUsize::new(OPEN), + }; + + let rx = Rx { shared, index: 0 }; + + (tx, rx) +} + +impl Tx { + /// Pushes a value into the ring buffer. + pub(crate) fn push(&self, value: T) { + // First, claim a slot for the value. The bounded channel semaphore + // ensures that this slot is empty before it is reused. + let slot_index = self.claim_position(); + let slot = self.shared.slot(slot_index); + + slot.write(value); + self.shared.set_ready(slot_index); + } + + fn claim_position(&self) -> usize { + let mut tail = self.tail_position.load(Relaxed); + + loop { + let next = self.shared.next_position(tail); + + match self + .tail_position + .compare_exchange_weak(tail, next, AcqRel, Acquire) + { + Ok(_) => return tail, + Err(actual) => tail = actual, + } + } + } + + /// Closes the send half of the queue. + /// + /// Unlike the list queue, the array queue does not need to push a fake close + /// message. The receiver can observe this flag once it has caught up to the + /// sender tail position. + pub(crate) fn close(&self) { + self.closed.store(CLOSED, Release); + } + + fn is_closed(&self) -> bool { + self.closed.load(Acquire) == CLOSED + } +} + +impl fmt::Debug for Tx { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + let tail_position = self.tail_position.load(Relaxed); + let head_position = self.shared.head_position.load(Relaxed); + + fmt.debug_struct("Tx") + .field("tail_position", &tail_position) + .field("capacity", &self.shared.capacity) + .field("len", &self.shared.distance(head_position, tail_position)) + .field("closed", &(self.closed.load(Relaxed) == CLOSED)) + .finish() + } +} + +impl Rx { + pub(crate) fn is_empty(&self, tx: &Tx) -> bool { + self.len(tx) == 0 + } + + pub(crate) fn len(&self, tx: &Tx) -> usize { + let tail = tx.tail_position.load(Acquire); + self.shared.distance(self.index, tail) + } + + /// Pops the next value off the queue. + pub(crate) fn pop(&mut self, tx: &Tx) -> Option> { + let mut tail = tx.tail_position.load(Acquire); + + if self.index == tail { + if !tx.is_closed() { + return None; + } + + // The tail claim is sequenced before closing the final sender. Load + // it again after observing `closed` so that a message claimed by + // that sender cannot be mistaken for an empty, closed queue. + tail = tx.tail_position.load(Acquire); + if self.index == tail { + return Some(Read::Closed); + } + } + + if !self.shared.is_ready(self.index) { + return None; + } + + let slot = self.shared.slot(self.index); + let value = unsafe { slot.take() }; + self.shared.clear_ready(self.index); + self.index = self.shared.next_position(self.index); + self.shared.head_position.store(self.index, Relaxed); + Some(Read::Value(value)) + } + + /// Pops the next value off the queue, detecting whether the queue is busy or + /// empty on failure. + /// + /// This function exists because `Rx::pop` can return `None` even if the + /// queue contains a message that has been completely written. This can + /// happen if the fully delivered message is behind another message that is + /// in the middle of being written, since the channel can't return messages + /// out of order. + pub(crate) fn try_pop(&mut self, tx: &Tx) -> TryPopResult { + let mut tail = tx.tail_position.load(Acquire); + + if self.index == tail { + if !tx.is_closed() { + return TryPopResult::Empty; + } + + tail = tx.tail_position.load(Acquire); + if self.index == tail { + return TryPopResult::Closed; + } + } + + if !self.shared.is_ready(self.index) { + return TryPopResult::Busy; + } + + let slot = self.shared.slot(self.index); + let value = unsafe { slot.take() }; + self.shared.clear_ready(self.index); + self.index = self.shared.next_position(self.index); + self.shared.head_position.store(self.index, Relaxed); + TryPopResult::Ok(value) + } +} + +impl fmt::Debug for Rx { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("Rx") + .field("index", &self.index) + .field("capacity", &self.shared.capacity) + .finish() + } +} + +impl Shared { + /// Returns the slot for a queue position. + fn slot(&self, position: usize) -> &Slot { + // Capacity can never be 0, guarded by the constructor. + &self.slots[position % self.capacity] + } + + /// Returns true if the slot for a queue position is ready. + fn is_ready(&self, position: usize) -> bool { + self.ready_slots.load(Acquire) & self.mask(position) != 0 + } + + /// Marks the slot for a queue position ready. + fn set_ready(&self, position: usize) { + self.ready_slots.fetch_or(self.mask(position), Release); + } + + /// Marks the slot for a queue position empty. + fn clear_ready(&self, position: usize) { + self.ready_slots.fetch_and(!self.mask(position), Release); + } + + fn next_position(&self, position: usize) -> usize { + if position + 1 == self.position_wrap { + 0 + } else { + position + 1 + } + } + + fn distance(&self, from: usize, to: usize) -> usize { + if to >= from { + to - from + } else { + self.position_wrap - from + to + } + } + + fn mask(&self, position: usize) -> usize { + 1 << (position % self.capacity) + } +} + +impl Drop for Shared { + fn drop(&mut self) { + for position in 0..self.capacity { + if self.is_ready(position) { + let slot = self.slot(position); + unsafe { slot.drop_value() }; + } + } + } +} + +impl Slot { + /// Creates an empty slot. + fn new() -> Slot { + Slot { + value: UnsafeCell::new(MaybeUninit::uninit()), + } + } + + /// Writes a value into the slot. + fn write(&self, value: T) { + self.value.with_mut(|ptr| unsafe { + ptr::write(ptr, MaybeUninit::new(value)); + }); + } + + /// Reads a value out of the slot. + unsafe fn take(&self) -> T { + let value = self.value.with(|ptr| unsafe { ptr::read(ptr) }); + + unsafe { value.assume_init() } + } + + /// Drops the value in the slot. + unsafe fn drop_value(&self) { + self.value.with_mut(|ptr| unsafe { + (*ptr).assume_init_drop(); + }); + } +} diff --git a/tokio/src/sync/mpsc/bounded.rs b/tokio/src/sync/mpsc/bounded.rs index 38a9c34f8..0816ce595 100644 --- a/tokio/src/sync/mpsc/bounded.rs +++ b/tokio/src/sync/mpsc/bounded.rs @@ -2,6 +2,9 @@ use crate::loom::sync::Arc; use crate::sync::batch_semaphore::{self as semaphore, TryAcquireError}; use crate::sync::mpsc::chan; use crate::sync::mpsc::error::{SendError, TryRecvError, TrySendError}; +use crate::sync::mpsc::list; +#[cfg(tokio_unstable)] +use crate::sync::mpsc::{array, BLOCK_CAP}; cfg_time! { use crate::sync::mpsc::error::SendTimeoutError; @@ -20,7 +23,7 @@ use std::task::{Context, Poll}; /// /// [`PollSender`]: https://docs.rs/tokio-util/latest/tokio_util/sync/struct.PollSender.html pub struct Sender { - chan: chan::Tx, + chan: TxFlavor, } /// A sender that does not prevent the channel from being closed. @@ -54,7 +57,7 @@ pub struct Sender { /// # } /// ``` pub struct WeakSender { - chan: Arc>, + chan: WeakTxFlavor, } /// Permits to send one value into the channel. @@ -65,7 +68,7 @@ pub struct WeakSender { /// [`Sender::reserve()`]: Sender::reserve /// [`Sender::try_reserve()`]: Sender::try_reserve pub struct Permit<'a, T> { - chan: &'a chan::Tx, + chan: &'a TxFlavor, } /// An [`Iterator`] of [`Permit`] that can be used to hold `n` slots in the channel. @@ -76,7 +79,7 @@ pub struct Permit<'a, T> { /// [`Sender::reserve_many()`]: Sender::reserve_many /// [`Sender::try_reserve_many()`]: Sender::try_reserve_many pub struct PermitIterator<'a, T> { - chan: &'a chan::Tx, + chan: &'a TxFlavor, n: usize, } @@ -93,7 +96,7 @@ pub struct PermitIterator<'a, T> { /// [`Sender::reserve_owned()`]: Sender::reserve_owned /// [`Sender::try_reserve_owned()`]: Sender::try_reserve_owned pub struct OwnedPermit { - chan: Option>, + chan: Option>, } /// Receives values from the associated `Sender`. @@ -105,7 +108,35 @@ pub struct OwnedPermit { /// [`ReceiverStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.ReceiverStream.html pub struct Receiver { /// The channel receiver. - chan: chan::Rx, + chan: RxFlavor, +} + +type ListTx = chan::Tx; +type ListRx = chan::Rx; +type ListChan = chan::Chan; +#[cfg(tokio_unstable)] +type ArrayTx = chan::Tx; +#[cfg(tokio_unstable)] +type ArrayRx = chan::Rx; +#[cfg(tokio_unstable)] +type ArrayChan = chan::Chan; + +enum TxFlavor { + List(ListTx), + #[cfg(tokio_unstable)] + Array(ArrayTx), +} + +enum WeakTxFlavor { + List(Arc>), + #[cfg(tokio_unstable)] + Array(Arc>), +} + +enum RxFlavor { + List(ListRx), + #[cfg(tokio_unstable)] + Array(ArrayRx), } /// Creates a bounded mpsc channel for communicating between asynchronous tasks @@ -162,10 +193,22 @@ pub fn channel(buffer: usize) -> (Sender, Receiver) { semaphore: semaphore::Semaphore::new(buffer), bound: buffer, }; - let (tx, rx) = chan::channel(semaphore); + #[cfg(tokio_unstable)] + let (tx, rx) = { + if buffer <= BLOCK_CAP { + let (tx, rx) = chan::channel::(semaphore, buffer); + (Sender::from_array(tx), Receiver::from_array(rx)) + } else { + let (tx, rx) = chan::channel::(semaphore, buffer); + (Sender::from_list(tx), Receiver::from_list(rx)) + } + }; - let tx = Sender::new(tx); - let rx = Receiver::new(rx); + #[cfg(not(tokio_unstable))] + let (tx, rx) = { + let (tx, rx) = chan::channel::(semaphore, buffer); + (Sender::from_list(tx), Receiver::from_list(rx)) + }; (tx, rx) } @@ -178,9 +221,263 @@ pub(crate) struct Semaphore { pub(crate) bound: usize, } +impl TxFlavor { + fn send(&self, value: T) { + match self { + Self::List(chan) => chan.send(value), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.send(value), + } + } + + async fn closed(&self) { + match self { + Self::List(chan) => chan.closed().await, + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.closed().await, + } + } + + fn semaphore(&self) -> &Semaphore { + match self { + Self::List(chan) => chan.semaphore(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.semaphore(), + } + } + + fn wake_rx(&self) { + match self { + Self::List(chan) => chan.wake_rx(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.wake_rx(), + } + } + + fn is_closed(&self) -> bool { + match self { + Self::List(chan) => chan.is_closed(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.is_closed(), + } + } + + fn same_channel(&self, other: &Self) -> bool { + match (self, other) { + (Self::List(a), Self::List(b)) => a.same_channel(b), + #[cfg(tokio_unstable)] + (Self::Array(a), Self::Array(b)) => a.same_channel(b), + #[cfg(tokio_unstable)] + _ => false, + } + } + + fn downgrade(&self) -> WeakTxFlavor { + match self { + Self::List(chan) => WeakTxFlavor::List(chan.downgrade()), + #[cfg(tokio_unstable)] + Self::Array(chan) => WeakTxFlavor::Array(chan.downgrade()), + } + } + + fn strong_count(&self) -> usize { + match self { + Self::List(chan) => chan.strong_count(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.strong_count(), + } + } + + fn weak_count(&self) -> usize { + match self { + Self::List(chan) => chan.weak_count(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.weak_count(), + } + } +} + +impl Clone for TxFlavor { + fn clone(&self) -> Self { + match self { + Self::List(chan) => Self::List(chan.clone()), + #[cfg(tokio_unstable)] + Self::Array(chan) => Self::Array(chan.clone()), + } + } +} + +impl fmt::Debug for TxFlavor { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::List(chan) => fmt::Debug::fmt(chan, fmt), + #[cfg(tokio_unstable)] + Self::Array(chan) => fmt::Debug::fmt(chan, fmt), + } + } +} + +impl WeakTxFlavor { + fn increment_weak_count(&self) { + match self { + Self::List(chan) => chan.increment_weak_count(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.increment_weak_count(), + } + } + + fn decrement_weak_count(&self) { + match self { + Self::List(chan) => chan.decrement_weak_count(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.decrement_weak_count(), + } + } + + fn upgrade(&self) -> Option> { + match self { + Self::List(chan) => chan::Tx::upgrade(chan.clone()).map(Sender::from_list), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan::Tx::upgrade(chan.clone()).map(Sender::from_array), + } + } + + fn strong_count(&self) -> usize { + match self { + Self::List(chan) => chan.strong_count(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.strong_count(), + } + } + + fn weak_count(&self) -> usize { + match self { + Self::List(chan) => chan.weak_count(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.weak_count(), + } + } +} + +impl fmt::Debug for WeakTxFlavor { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::List(chan) => fmt::Debug::fmt(chan, fmt), + #[cfg(tokio_unstable)] + Self::Array(chan) => fmt::Debug::fmt(chan, fmt), + } + } +} + +impl RxFlavor { + fn close(&mut self) { + match self { + Self::List(chan) => chan.close(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.close(), + } + } + + fn recv(&mut self, cx: &mut Context<'_>) -> Poll> { + match self { + Self::List(chan) => chan.recv(cx), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.recv(cx), + } + } + + fn recv_many( + &mut self, + cx: &mut Context<'_>, + buffer: &mut Vec, + limit: usize, + ) -> Poll { + match self { + Self::List(chan) => chan.recv_many(cx, buffer, limit), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.recv_many(cx, buffer, limit), + } + } + + fn try_recv(&mut self) -> Result { + match self { + Self::List(chan) => chan.try_recv(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.try_recv(), + } + } + + fn is_closed(&self) -> bool { + match self { + Self::List(chan) => chan.is_closed(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.is_closed(), + } + } + + fn is_empty(&self) -> bool { + match self { + Self::List(chan) => chan.is_empty(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.is_empty(), + } + } + + fn len(&self) -> usize { + match self { + Self::List(chan) => chan.len(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.len(), + } + } + + fn semaphore(&self) -> &Semaphore { + match self { + Self::List(chan) => chan.semaphore(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.semaphore(), + } + } + + fn sender_strong_count(&self) -> usize { + match self { + Self::List(chan) => chan.sender_strong_count(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.sender_strong_count(), + } + } + + fn sender_weak_count(&self) -> usize { + match self { + Self::List(chan) => chan.sender_weak_count(), + #[cfg(tokio_unstable)] + Self::Array(chan) => chan.sender_weak_count(), + } + } +} + +impl fmt::Debug for RxFlavor { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::List(chan) => fmt::Debug::fmt(chan, fmt), + #[cfg(tokio_unstable)] + Self::Array(chan) => fmt::Debug::fmt(chan, fmt), + } + } +} + impl Receiver { - pub(crate) fn new(chan: chan::Rx) -> Receiver { - Receiver { chan } + pub(crate) fn from_list(chan: ListRx) -> Receiver { + Receiver { + chan: RxFlavor::List(chan), + } + } + + #[cfg(tokio_unstable)] + pub(crate) fn from_array(chan: ArrayRx) -> Receiver { + Receiver { + chan: RxFlavor::Array(chan), + } } /// Receives the next value for this receiver. @@ -750,8 +1047,17 @@ impl fmt::Debug for Receiver { impl Unpin for Receiver {} impl Sender { - pub(crate) fn new(chan: chan::Tx) -> Sender { - Sender { chan } + pub(crate) fn from_list(chan: ListTx) -> Sender { + Sender { + chan: TxFlavor::List(chan), + } + } + + #[cfg(tokio_unstable)] + pub(crate) fn from_array(chan: ArrayTx) -> Sender { + Sender { + chan: TxFlavor::Array(chan), + } } /// Sends a value, waiting until there is capacity. @@ -1284,7 +1590,7 @@ impl Sender { // observe a close. On success the caller receives the permits and takes // over that job, so the guard is forgotten below. struct WakeReceiverOnDrop<'a, T> { - chan: &'a chan::Tx, + chan: &'a TxFlavor, } impl Drop for WakeReceiverOnDrop<'_, T> { @@ -1647,7 +1953,11 @@ impl Clone for WeakSender { self.chan.increment_weak_count(); WeakSender { - chan: self.chan.clone(), + chan: match &self.chan { + WeakTxFlavor::List(chan) => WeakTxFlavor::List(chan.clone()), + #[cfg(tokio_unstable)] + WeakTxFlavor::Array(chan) => WeakTxFlavor::Array(chan.clone()), + }, } } } @@ -1663,7 +1973,7 @@ impl WeakSender { /// if there are other `Sender` instances alive and the channel wasn't /// previously dropped, otherwise `None` is returned. pub fn upgrade(&self) -> Option> { - chan::Tx::upgrade(self.chan.clone()).map(Sender::new) + self.chan.upgrade() } /// Returns the number of [`Sender`] handles. diff --git a/tokio/src/sync/mpsc/chan.rs b/tokio/src/sync/mpsc/chan.rs index f22880793..d281dff0e 100644 --- a/tokio/src/sync/mpsc/chan.rs +++ b/tokio/src/sync/mpsc/chan.rs @@ -3,8 +3,9 @@ use crate::loom::future::AtomicWaker; use crate::loom::sync::atomic::AtomicUsize; use crate::loom::sync::Arc; use crate::runtime::park::CachedParkThread; +use crate::sync::mpsc::block::Read; use crate::sync::mpsc::error::TryRecvError; -use crate::sync::mpsc::{bounded, list, unbounded}; +use crate::sync::mpsc::{bounded, unbounded, TryPopResult}; use crate::sync::notify::Notify; use crate::util::cacheline::CachePadded; @@ -15,23 +16,37 @@ use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; use std::task::Poll::{Pending, Ready}; use std::task::{ready, Context, Poll}; -/// Channel sender. -pub(crate) struct Tx { - inner: Arc>, +pub(crate) trait Queue { + type Tx: fmt::Debug; + type Rx: fmt::Debug; + + fn channel(bound: usize) -> (Self::Tx, Self::Rx); + fn push(tx: &Self::Tx, value: T); + fn close(tx: &Self::Tx); + fn is_empty(rx: &Self::Rx, tx: &Self::Tx) -> bool; + fn len(rx: &Self::Rx, tx: &Self::Tx) -> usize; + fn pop(rx: &mut Self::Rx, tx: &Self::Tx) -> Option>; + fn try_pop(rx: &mut Self::Rx, tx: &Self::Tx) -> TryPopResult; + unsafe fn free(rx: &mut Self::Rx); } -impl fmt::Debug for Tx { +/// Channel sender. +pub(crate) struct Tx> { + inner: Arc>, +} + +impl> fmt::Debug for Tx { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Tx").field("inner", &self.inner).finish() } } /// Channel receiver. -pub(crate) struct Rx { - inner: Arc>, +pub(crate) struct Rx> { + inner: Arc>, } -impl fmt::Debug for Rx { +impl> fmt::Debug for Rx { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Rx").field("inner", &self.inner).finish() } @@ -49,9 +64,9 @@ pub(crate) trait Semaphore { fn is_closed(&self) -> bool; } -pub(super) struct Chan { - /// Handle to the push half of the lock-free list. - tx: CachePadded>, +pub(super) struct Chan> { + /// Handle to the push half of the queue. + tx: CachePadded, /// Receiver waker. Notified when a value is pushed into the channel. rx_waker: CachePadded, @@ -71,12 +86,13 @@ pub(super) struct Chan { tx_weak_count: AtomicUsize, /// Only accessed by `Rx` handle. - rx_fields: UnsafeCell>, + rx_fields: UnsafeCell>, } -impl fmt::Debug for Chan +impl fmt::Debug for Chan where S: fmt::Debug, + Q: Queue, { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("Chan") @@ -90,30 +106,49 @@ where } /// Fields only accessed by `Rx` handle. -struct RxFields { - /// Channel receiver. This field is only accessed by the `Receiver` type. - list: list::Rx, +struct RxFields> { + /// Queue receiver. This field is only accessed by the `Receiver` type. + queue: Q::Rx, /// `true` if `Rx::close` is called. rx_closed: bool, } -impl fmt::Debug for RxFields { +impl> fmt::Debug for RxFields { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("RxFields") - .field("list", &self.list) + .field("queue", &self.queue) .field("rx_closed", &self.rx_closed) .finish() } } -unsafe impl Send for Chan {} -unsafe impl Sync for Chan {} -impl panic::RefUnwindSafe for Chan {} -impl panic::UnwindSafe for Chan {} +unsafe impl Send for Chan +where + T: Send, + S: Send, + Q: Queue, + Q::Tx: Send, +{ +} -pub(crate) fn channel(semaphore: S) -> (Tx, Rx) { - let (tx, rx) = list::channel(); +unsafe impl Sync for Chan +where + T: Send, + S: Sync, + Q: Queue, + Q::Tx: Sync, +{ +} + +impl> panic::RefUnwindSafe for Chan {} +impl> panic::UnwindSafe for Chan {} + +pub(crate) fn channel>( + semaphore: S, + bound: usize, +) -> (Tx, Rx) { + let (tx, rx) = Q::channel(bound); let chan = Arc::new(Chan { notify_rx_closed: Notify::new(), @@ -123,7 +158,7 @@ pub(crate) fn channel(semaphore: S) -> (Tx, Rx) { tx_count: AtomicUsize::new(1), tx_weak_count: AtomicUsize::new(0), rx_fields: UnsafeCell::new(RxFields { - list: rx, + queue: rx, rx_closed: false, }), }); @@ -133,8 +168,8 @@ pub(crate) fn channel(semaphore: S) -> (Tx, Rx) { // ===== impl Tx ===== -impl Tx { - fn new(chan: Arc>) -> Tx { +impl> Tx { + fn new(chan: Arc>) -> Tx { Tx { inner: chan } } @@ -146,14 +181,14 @@ impl Tx { self.inner.tx_weak_count.load(Relaxed) } - pub(super) fn downgrade(&self) -> Arc> { + pub(super) fn downgrade(&self) -> Arc> { self.inner.increment_weak_count(); self.inner.clone() } // Returns the upgraded channel or None if the upgrade failed. - pub(super) fn upgrade(chan: Arc>) -> Option { + pub(super) fn upgrade(chan: Arc>) -> Option { let mut tx_count = chan.tx_count.load(Acquire); loop { @@ -192,7 +227,7 @@ impl Tx { } } -impl Tx { +impl> Tx { pub(crate) fn is_closed(&self) -> bool { self.inner.semaphore.is_closed() } @@ -210,8 +245,8 @@ impl Tx { } } -impl Clone for Tx { - fn clone(&self) -> Tx { +impl> Clone for Tx { + fn clone(&self) -> Tx { // Using a Relaxed ordering here is sufficient as the caller holds a // strong ref to `self`, preventing a concurrent decrement to zero. self.inner.tx_count.fetch_add(1, Relaxed); @@ -222,14 +257,13 @@ impl Clone for Tx { } } -impl Drop for Tx { +impl> Drop for Tx { fn drop(&mut self) { if self.inner.tx_count.fetch_sub(1, AcqRel) != 1 { return; } - // Close the list, which sends a `Close` message - self.inner.tx.close(); + Q::close(&self.inner.tx); // Notify the receiver self.wake_rx(); @@ -238,8 +272,8 @@ impl Drop for Tx { // ===== impl Rx ===== -impl Rx { - fn new(chan: Arc>) -> Rx { +impl> Rx { + fn new(chan: Arc>) -> Rx { Rx { inner: chan } } @@ -265,30 +299,27 @@ impl Rx { // In this case, the inner semaphore will be closed. // // 2. When all senders are dropped. - // In this case, the semaphore remains unclosed, and the `index` in the list won't - // reach the tail position. It is necessary to check the list if the last block is - // `closed`. + // In this case, the semaphore remains unclosed, and the queue won't + // report a closed state until it is drained. self.inner.semaphore.is_closed() || self.inner.tx_count.load(Acquire) == 0 } pub(crate) fn is_empty(&self) -> bool { self.inner.rx_fields.with(|rx_fields_ptr| { let rx_fields = unsafe { &*rx_fields_ptr }; - rx_fields.list.is_empty(&self.inner.tx) + Q::is_empty(&rx_fields.queue, &self.inner.tx) }) } pub(crate) fn len(&self) -> usize { self.inner.rx_fields.with(|rx_fields_ptr| { let rx_fields = unsafe { &*rx_fields_ptr }; - rx_fields.list.len(&self.inner.tx) + Q::len(&rx_fields.queue, &self.inner.tx) }) } /// Receive the next value pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll> { - use super::block::Read; - ready!(crate::trace::trace_leaf()); // Keep track of task budget @@ -299,7 +330,7 @@ impl Rx { macro_rules! try_recv { () => { - match rx_fields.list.pop(&self.inner.tx) { + match Q::pop(&mut rx_fields.queue, &self.inner.tx) { Some(Read::Value(value)) => { self.inner.semaphore.add_permit(); coop.made_progress(); @@ -347,8 +378,6 @@ impl Rx { buffer: &mut Vec, limit: usize, ) -> Poll { - use super::block::Read; - ready!(crate::trace::trace_leaf()); // Keep track of task budget @@ -367,7 +396,7 @@ impl Rx { macro_rules! try_recv { () => { while remaining > 0 { - match rx_fields.list.pop(&self.inner.tx) { + match Q::pop(&mut rx_fields.queue, &self.inner.tx) { Some(Read::Value(value)) => { remaining -= 1; buffer.push(value); @@ -422,14 +451,12 @@ impl Rx { /// Try to receive the next value. pub(crate) fn try_recv(&mut self) -> Result { - use super::list::TryPopResult; - self.inner.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; macro_rules! try_recv { () => { - match rx_fields.list.try_pop(&self.inner.tx) { + match Q::try_pop(&mut rx_fields.queue, &self.inner.tx) { TryPopResult::Ok(value) => { self.inner.semaphore.add_permit(); return Ok(value); @@ -484,37 +511,35 @@ impl Rx { } } -impl Drop for Rx { +impl> Drop for Rx { fn drop(&mut self) { - use super::block::Read::Value; - self.close(); self.inner.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; - struct Guard<'a, T, S: Semaphore> { - list: &'a mut list::Rx, - tx: &'a list::Tx, + struct Guard<'a, T, S: Semaphore, Q: Queue> { + queue: &'a mut Q::Rx, + tx: &'a Q::Tx, sem: &'a S, } - impl<'a, T, S: Semaphore> Guard<'a, T, S> { + impl<'a, T, S: Semaphore, Q: Queue> Guard<'a, T, S, Q> { fn drain(&mut self) { // call T's destructor. - while let Some(Value(_)) = self.list.pop(self.tx) { + while let Some(Read::Value(_)) = Q::pop(self.queue, self.tx) { self.sem.add_permit(); } } } - impl<'a, T, S: Semaphore> Drop for Guard<'a, T, S> { + impl<'a, T, S: Semaphore, Q: Queue> Drop for Guard<'a, T, S, Q> { fn drop(&mut self) { self.drain(); } } - let mut guard = Guard { - list: &mut rx_fields.list, + let mut guard = Guard:: { + queue: &mut rx_fields.queue, tx: &self.inner.tx, sem: &self.inner.semaphore, }; @@ -531,10 +556,10 @@ impl Drop for Rx { // ===== impl Chan ===== -impl Chan { +impl> Chan { fn send(&self, value: T) { // Push the value - self.tx.push(value); + Q::push(&self.tx, value); // Notify the rx task self.rx_waker.wake(); @@ -557,17 +582,15 @@ impl Chan { } } -impl Drop for Chan { +impl> Drop for Chan { fn drop(&mut self) { - use super::block::Read::Value; - // Safety: the only owner of the rx fields is Chan, and being // inside its own Drop means we're the last ones to touch it. self.rx_fields.with_mut(|rx_fields_ptr| { let rx_fields = unsafe { &mut *rx_fields_ptr }; - while let Some(Value(_)) = rx_fields.list.pop(&self.tx) {} - unsafe { rx_fields.list.free_blocks() }; + while let Some(Read::Value(_)) = Q::pop(&mut rx_fields.queue, &self.tx) {} + unsafe { Q::free(&mut rx_fields.queue) }; }); } } diff --git a/tokio/src/sync/mpsc/list.rs b/tokio/src/sync/mpsc/list.rs index c5f21b8b2..68d02edbe 100644 --- a/tokio/src/sync/mpsc/list.rs +++ b/tokio/src/sync/mpsc/list.rs @@ -3,6 +3,8 @@ use crate::loom::sync::atomic::{AtomicPtr, AtomicUsize}; use crate::loom::thread; use crate::sync::mpsc::block::{self, Block}; +use crate::sync::mpsc::chan; +use crate::sync::mpsc::TryPopResult; use std::fmt; use std::ptr::NonNull; @@ -30,23 +32,8 @@ pub(crate) struct Rx { free_head: NonNull>, } -/// Return value of `Rx::try_pop`. -pub(crate) enum TryPopResult { - /// Successfully popped a value. - Ok(T), - /// The channel is empty. - /// - /// Note that `list.rs` only tracks the close state set by senders. If the - /// channel is closed by `Rx::close()`, then `TryPopResult::Empty` is still - /// returned, and the close state needs to be handled by `chan.rs`. - Empty, - /// The channel is empty and closed. - /// - /// Returned when the send half is closed (all senders dropped). - Closed, - /// The channel is not empty, but the first value is being written. - Busy, -} +#[derive(Debug)] +pub(crate) struct Queue; pub(crate) fn channel() -> (Tx, Rx) { // Create the initial block shared between the tx and rx halves. @@ -69,6 +56,43 @@ pub(crate) fn channel() -> (Tx, Rx) { (tx, rx) } +impl chan::Queue for Queue { + type Tx = Tx; + type Rx = Rx; + + fn channel(_bound: usize) -> (Self::Tx, Self::Rx) { + channel() + } + + fn push(tx: &Self::Tx, value: T) { + tx.push(value); + } + + fn close(tx: &Self::Tx) { + tx.close(); + } + + fn is_empty(rx: &Self::Rx, tx: &Self::Tx) -> bool { + rx.is_empty(tx) + } + + fn len(rx: &Self::Rx, tx: &Self::Tx) -> usize { + rx.len(tx) + } + + fn pop(rx: &mut Self::Rx, tx: &Self::Tx) -> Option> { + rx.pop(tx) + } + + fn try_pop(rx: &mut Self::Rx, tx: &Self::Tx) -> TryPopResult { + rx.try_pop(tx) + } + + unsafe fn free(rx: &mut Self::Rx) { + unsafe { rx.free_blocks() }; + } +} + impl Tx { /// Pushes a value into the list. pub(crate) fn push(&self, value: T) { diff --git a/tokio/src/sync/mpsc/mod.rs b/tokio/src/sync/mpsc/mod.rs index 1947013e6..65bd639d6 100644 --- a/tokio/src/sync/mpsc/mod.rs +++ b/tokio/src/sync/mpsc/mod.rs @@ -123,6 +123,9 @@ pub use self::bounded::{ mod chan; +#[cfg(tokio_unstable)] +pub(super) mod array; + pub(super) mod list; mod unbounded; @@ -132,6 +135,23 @@ pub use self::unbounded::{ pub mod error; +pub(crate) enum TryPopResult { + /// Successfully popped a value. + Ok(T), + /// The channel is empty. + /// + /// Note that `mpsc::chan` only tracks the close state set by senders. If the + /// channel is closed by `Rx::close()`, then `TryPopResult::Empty` is still + /// returned, and the close state needs to be handled by `chan.rs`. + Empty, + /// The channel is empty and closed. + /// + /// Returned when the send half is closed (all senders dropped). + Closed, + /// The channel is not empty, but the first value is being written. + Busy, +} + /// The number of values a block can contain. /// /// This value must be a power of 2. It also must be smaller than the number of diff --git a/tokio/src/sync/mpsc/unbounded.rs b/tokio/src/sync/mpsc/unbounded.rs index dd64e498a..322e1f276 100644 --- a/tokio/src/sync/mpsc/unbounded.rs +++ b/tokio/src/sync/mpsc/unbounded.rs @@ -1,6 +1,7 @@ use crate::loom::sync::{atomic::AtomicUsize, Arc}; use crate::sync::mpsc::chan; use crate::sync::mpsc::error::{SendError, TryRecvError}; +use crate::sync::mpsc::list; use std::fmt; use std::task::{Context, Poll}; @@ -9,7 +10,7 @@ use std::task::{Context, Poll}; /// /// Instances are created by the [`unbounded_channel`] function. pub struct UnboundedSender { - chan: chan::Tx, + chan: chan::Tx, } /// An unbounded sender that does not prevent the channel from being closed. @@ -43,7 +44,7 @@ pub struct UnboundedSender { /// # } /// ``` pub struct WeakUnboundedSender { - chan: Arc>, + chan: Arc>, } impl Clone for UnboundedSender { @@ -71,7 +72,7 @@ impl fmt::Debug for UnboundedSender { /// [`UnboundedReceiverStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.UnboundedReceiverStream.html pub struct UnboundedReceiver { /// The channel receiver - chan: chan::Rx, + chan: chan::Rx, } impl fmt::Debug for UnboundedReceiver { @@ -93,7 +94,7 @@ impl fmt::Debug for UnboundedReceiver { /// the channel. Using an `unbounded` channel has the ability of causing the /// process to run out of memory. In this case, the process will be aborted. pub fn unbounded_channel() -> (UnboundedSender, UnboundedReceiver) { - let (tx, rx) = chan::channel(Semaphore(AtomicUsize::new(0))); + let (tx, rx) = chan::channel::(Semaphore(AtomicUsize::new(0)), 0); let tx = UnboundedSender::new(tx); let rx = UnboundedReceiver::new(rx); @@ -106,7 +107,7 @@ pub fn unbounded_channel() -> (UnboundedSender, UnboundedReceiver) { pub(crate) struct Semaphore(pub(crate) AtomicUsize); impl UnboundedReceiver { - pub(crate) fn new(chan: chan::Rx) -> UnboundedReceiver { + pub(crate) fn new(chan: chan::Rx) -> UnboundedReceiver { UnboundedReceiver { chan } } @@ -527,7 +528,7 @@ impl UnboundedReceiver { } impl UnboundedSender { - pub(crate) fn new(chan: chan::Tx) -> UnboundedSender { + pub(crate) fn new(chan: chan::Tx) -> UnboundedSender { UnboundedSender { chan } } diff --git a/tokio/src/sync/tests/loom_mpsc.rs b/tokio/src/sync/tests/loom_mpsc.rs index 468177143..98aa27863 100644 --- a/tokio/src/sync/tests/loom_mpsc.rs +++ b/tokio/src/sync/tests/loom_mpsc.rs @@ -223,6 +223,29 @@ fn nonempty_after_send() { }); } +#[test] +#[cfg(tokio_unstable)] +fn array_nonempty_after_send() { + loom::model(|| { + // Under loom, BLOCK_CAP is 2, so this bounded channel uses the + // array-backed queue when tokio_unstable is enabled. + let (send, recv) = mpsc::channel(2); + let send2 = send.clone(); + + let join = thread::spawn(move || { + block_on(send2.send("message2")).unwrap(); + }); + + // Loom can schedule send2 so it reserves the head slot but stalls before + // marking it ready. This send can then complete in the next slot. + // is_empty must still report that the channel contains a message. + block_on(send.send("message1")).unwrap(); + assert!(!recv.is_empty()); + + join.join().unwrap(); + }); +} + #[test] fn is_empty_during_close() { loom::model(|| { diff --git a/tokio/tests/sync_mpsc.rs b/tokio/tests/sync_mpsc.rs index 398ab6339..50229bd8d 100644 --- a/tokio/tests/sync_mpsc.rs +++ b/tokio/tests/sync_mpsc.rs @@ -70,6 +70,20 @@ async fn send_recv_with_buffer() { assert!(val.is_none()); } +#[test] +#[cfg(tokio_unstable)] +fn array_wraps_with_non_power_of_two_capacity() { + let (tx, mut rx) = mpsc::channel(3); + + // Exercise several wraps of the queue position. In particular, capacity + // three does not divide the native integer range, so relying on integer + // overflow to wrap the position would eventually select the wrong slot. + for value in 0..12 { + tx.try_send(value).unwrap(); + assert_eq!(rx.try_recv(), Ok(value)); + } +} + #[tokio::test] #[cfg(feature = "full")] async fn reserve_disarm() {