From 4beb5ff9633229124d8386fea205eeea122f3112 Mon Sep 17 00:00:00 2001 From: tison Date: Fri, 28 Aug 2026 22:47:57 +0800 Subject: [PATCH 1/2] perf(mpsc): coalesce unbounded receiver wakeups --- CHANGELOG.md | 1 + asyncband/src/channel/mpsc/unbounded.rs | 75 ++++++++++++++++++++++--- asyncband/src/internal/atomic_waker.rs | 22 ++++++++ tests-integration/tests/mpsc_test.rs | 72 ++++++++++++++++++++++++ 4 files changed, 163 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 66b5f57..65c424a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,3 +32,4 @@ All notable changes to this project will be documented in this file. ### Improvements * Remove the `slab` dependency in favor of a focused internal waiter arena. +* Reduce unbounded MPSC contention by avoiding receiver wake operations while no receive call is parked. diff --git a/asyncband/src/channel/mpsc/unbounded.rs b/asyncband/src/channel/mpsc/unbounded.rs index 3ee892a..ddb807d 100644 --- a/asyncband/src/channel/mpsc/unbounded.rs +++ b/asyncband/src/channel/mpsc/unbounded.rs @@ -21,10 +21,12 @@ use std::fmt; use std::future::poll_fn; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use std::task::Waker; use super::RecvError; use super::SendError; @@ -43,7 +45,7 @@ use crate::internal::atomic_waker::AtomicWaker; pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { senders: AtomicUsize::new(1), - rx_waker: AtomicWaker::new(), + rx_wake: ReceiverWake::new(), }); let (sender, receiver) = std::sync::mpsc::channel(); let sender = UnboundedSender { @@ -59,7 +61,54 @@ pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { struct UnboundedState { senders: AtomicUsize, - rx_waker: AtomicWaker, + rx_wake: ReceiverWake, +} + +/// Coalesces receiver wake-ups while keeping the queue as the source of truth. +/// +/// The receiver follows `register -> arm -> recheck`; each sender follows `enqueue -> claim arm`. +/// If a sender checks before the arm is published, the recheck observes its message. Otherwise, +/// one sender changes `armed` from true to false and owns the wake. A false value after arming +/// means another sender already owns the wake or the receiver has observed readiness and disarmed. +struct ReceiverWake { + armed: AtomicBool, + waker: AtomicWaker, +} + +impl ReceiverWake { + const fn new() -> Self { + Self { + armed: AtomicBool::new(false), + waker: AtomicWaker::new(), + } + } + + fn arm(&self, waker: &Waker) { + self.waker.register(waker); + + // Together with the sender's SeqCst check, this orders arming before the receiver's queue + // recheck and enqueueing before the sender's arm check. Both sides therefore cannot miss + // each other: either the receiver observes the message or the sender observes the arm. + self.armed.store(true, Ordering::SeqCst); + } + + fn disarm(&self) { + self.armed.store(false, Ordering::SeqCst); + self.waker.clear(); + } + + fn wake(&self) { + // The load keeps the common active-receiver path read-only. The compare-exchange lets one + // sender claim each arm while concurrent senders coalesce behind it. + if self.armed.load(Ordering::SeqCst) + && self + .armed + .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + { + self.waker.wake(); + } + } } /// Send values to the associated [`UnboundedReceiver`]. @@ -95,7 +144,7 @@ impl Drop for UnboundedSender { 1 => { // If this is the last sender, we need to wake up the receiver so it can // observe the disconnected state. - self.state.rx_waker.wake(); + self.state.rx_wake.wake(); } _ => { // there are still other senders left, do nothing @@ -118,7 +167,7 @@ impl UnboundedSender { let sender = self.sender.as_ref().unwrap(); sender.send(value).map_err(|err| SendError::new(err.0))?; - self.state.rx_waker.wake(); + self.state.rx_wake.wake(); Ok(()) } @@ -136,6 +185,12 @@ pub struct UnboundedReceiver { /// That is, `UnboundedReceiver` can only be accessed by one thread at a time. unsafe impl Sync for UnboundedReceiver {} +impl Drop for UnboundedReceiver { + fn drop(&mut self) { + self.state.rx_wake.disarm(); + } +} + impl fmt::Debug for UnboundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("UnboundedReceiver").finish_non_exhaustive() @@ -241,11 +296,17 @@ impl UnboundedReceiver { Ok(v) => Poll::Ready(Ok(v)), Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), Err(TryRecvError::Empty) => { - self.state.rx_waker.register(cx.waker()); + self.state.rx_wake.arm(cx.waker()); match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), + Ok(v) => { + self.state.rx_wake.disarm(); + Poll::Ready(Ok(v)) + } + Err(TryRecvError::Disconnected) => { + self.state.rx_wake.disarm(); + Poll::Ready(Err(RecvError::Disconnected)) + } Err(TryRecvError::Empty) => Poll::Pending, } } diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs index fe6d3f2..d1617a6 100644 --- a/asyncband/src/internal/atomic_waker.rs +++ b/asyncband/src/internal/atomic_waker.rs @@ -207,6 +207,15 @@ impl AtomicWaker { } } + /// Removes the most recently registered waker without waking it. + /// + /// This method must not overlap with [`register`](Self::register). It may run concurrently with + /// [`wake`](Self::wake). + #[inline] + pub fn clear(&self) { + drop(self.take()); + } + #[inline] fn take(&self) -> Option { // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the @@ -270,6 +279,19 @@ mod tests { assert_eq!(counter.0.load(Ordering::Relaxed), 1); } + #[test] + fn clear_removes_registered_waker() { + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let atomic_waker = AtomicWaker::new(); + + atomic_waker.register(&waker); + atomic_waker.clear(); + atomic_waker.wake(); + + assert_eq!(counter.0.load(Ordering::Relaxed), 0); + } + #[test] fn reregistering_same_task_does_not_clone_waker() { let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 72256a7..7317005 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -15,7 +15,14 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; use std::task::Poll; +use std::task::Wake; +use std::task::Waker; use std::time::Instant; use asyncband::mpsc; @@ -33,6 +40,14 @@ fn expect_ready(poll: Poll) -> T { } } +struct WakeProbe(AtomicUsize); + +impl Wake for WakeProbe { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + #[test] fn test_unbounded_pressure() { let n = 1024 * 1024; @@ -209,6 +224,63 @@ fn try_recv_close_while_empty_unbounded() { assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); } +#[test] +fn unbounded_burst_wakes_parked_receiver_once() { + let (tx, mut rx) = mpsc::unbounded(); + let probe = Arc::new(WakeProbe(AtomicUsize::new(0))); + let waker = Waker::from(probe.clone()); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + for value in 0..100 { + tx.send(value).unwrap(); + } + + assert_eq!(probe.0.load(Ordering::Relaxed), 1); + assert_eq!(expect_ready(recv.as_mut().poll(&mut context)), Ok(0)); + drop(recv); + + for value in 1..100 { + assert_eq!(rx.try_recv(), Ok(value)); + } +} + +#[test] +fn unbounded_last_sender_drop_wakes_parked_receiver() { + let (tx, mut rx) = mpsc::unbounded::<()>(); + let probe = Arc::new(WakeProbe(AtomicUsize::new(0))); + let waker = Waker::from(probe.clone()); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + drop(tx); + + assert_eq!(probe.0.load(Ordering::Relaxed), 1); + assert_eq!( + expect_ready(recv.as_mut().poll(&mut context)), + Err(RecvError::Disconnected) + ); +} + +#[test] +fn dropping_unbounded_receiver_releases_registered_waker() { + let (_tx, mut rx) = mpsc::unbounded::<()>(); + let probe = Arc::new(WakeProbe(AtomicUsize::new(0))); + let waker = Waker::from(probe.clone()); + let baseline_refs = Arc::strong_count(&probe); + let mut context = Context::from_waker(&waker); + let mut recv = Box::pin(rx.recv()); + + assert!(recv.as_mut().poll(&mut context).is_pending()); + assert_eq!(Arc::strong_count(&probe), baseline_refs + 1); + + drop(recv); + drop(rx); + assert_eq!(Arc::strong_count(&probe), baseline_refs); +} + #[tokio::test] async fn send_recv_bounded() { let (tx, mut rx) = mpsc::bounded(1); From b36a9acb812c77efd8ba05f4efe46c644ffe4d70 Mon Sep 17 00:00:00 2001 From: tison Date: Fri, 28 Aug 2026 23:25:40 +0800 Subject: [PATCH 2/2] refactor(mpsc): simplify receiver wake gating --- asyncband/src/channel/mpsc/unbounded.rs | 71 +++++++------------------ asyncband/src/internal/atomic_waker.rs | 22 -------- tests-integration/tests/mpsc_test.rs | 17 ------ 3 files changed, 18 insertions(+), 92 deletions(-) diff --git a/asyncband/src/channel/mpsc/unbounded.rs b/asyncband/src/channel/mpsc/unbounded.rs index ddb807d..bb5a37d 100644 --- a/asyncband/src/channel/mpsc/unbounded.rs +++ b/asyncband/src/channel/mpsc/unbounded.rs @@ -26,7 +26,6 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; -use std::task::Waker; use super::RecvError; use super::SendError; @@ -45,7 +44,8 @@ use crate::internal::atomic_waker::AtomicWaker; pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { senders: AtomicUsize::new(1), - rx_wake: ReceiverWake::new(), + rx_waiting: AtomicBool::new(false), + rx_waker: AtomicWaker::new(), }); let (sender, receiver) = std::sync::mpsc::channel(); let sender = UnboundedSender { @@ -61,52 +61,19 @@ pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { struct UnboundedState { senders: AtomicUsize, - rx_wake: ReceiverWake, + rx_waiting: AtomicBool, + rx_waker: AtomicWaker, } -/// Coalesces receiver wake-ups while keeping the queue as the source of truth. -/// -/// The receiver follows `register -> arm -> recheck`; each sender follows `enqueue -> claim arm`. -/// If a sender checks before the arm is published, the recheck observes its message. Otherwise, -/// one sender changes `armed` from true to false and owns the wake. A false value after arming -/// means another sender already owns the wake or the receiver has observed readiness and disarmed. -struct ReceiverWake { - armed: AtomicBool, - waker: AtomicWaker, -} - -impl ReceiverWake { - const fn new() -> Self { - Self { - armed: AtomicBool::new(false), - waker: AtomicWaker::new(), - } - } - - fn arm(&self, waker: &Waker) { - self.waker.register(waker); - - // Together with the sender's SeqCst check, this orders arming before the receiver's queue - // recheck and enqueueing before the sender's arm check. Both sides therefore cannot miss - // each other: either the receiver observes the message or the sender observes the arm. - self.armed.store(true, Ordering::SeqCst); - } - - fn disarm(&self) { - self.armed.store(false, Ordering::SeqCst); - self.waker.clear(); - } - - fn wake(&self) { - // The load keeps the common active-receiver path read-only. The compare-exchange lets one - // sender claim each arm while concurrent senders coalesce behind it. - if self.armed.load(Ordering::SeqCst) +impl UnboundedState { + fn wake_receiver(&self) { + if self.rx_waiting.load(Ordering::SeqCst) && self - .armed + .rx_waiting .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst) .is_ok() { - self.waker.wake(); + self.rx_waker.wake(); } } } @@ -144,7 +111,7 @@ impl Drop for UnboundedSender { 1 => { // If this is the last sender, we need to wake up the receiver so it can // observe the disconnected state. - self.state.rx_wake.wake(); + self.state.wake_receiver(); } _ => { // there are still other senders left, do nothing @@ -167,7 +134,7 @@ impl UnboundedSender { let sender = self.sender.as_ref().unwrap(); sender.send(value).map_err(|err| SendError::new(err.0))?; - self.state.rx_wake.wake(); + self.state.wake_receiver(); Ok(()) } @@ -185,12 +152,6 @@ pub struct UnboundedReceiver { /// That is, `UnboundedReceiver` can only be accessed by one thread at a time. unsafe impl Sync for UnboundedReceiver {} -impl Drop for UnboundedReceiver { - fn drop(&mut self) { - self.state.rx_wake.disarm(); - } -} - impl fmt::Debug for UnboundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("UnboundedReceiver").finish_non_exhaustive() @@ -296,15 +257,19 @@ impl UnboundedReceiver { Ok(v) => Poll::Ready(Ok(v)), Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), Err(TryRecvError::Empty) => { - self.state.rx_wake.arm(cx.waker()); + self.state.rx_waker.register(cx.waker()); + + // Publishing the wait before rechecking the queue prevents a sender from being + // missed: it either precedes the recheck or claims this flag and wakes us. + self.state.rx_waiting.store(true, Ordering::SeqCst); match self.try_recv() { Ok(v) => { - self.state.rx_wake.disarm(); + self.state.rx_waiting.store(false, Ordering::SeqCst); Poll::Ready(Ok(v)) } Err(TryRecvError::Disconnected) => { - self.state.rx_wake.disarm(); + self.state.rx_waiting.store(false, Ordering::SeqCst); Poll::Ready(Err(RecvError::Disconnected)) } Err(TryRecvError::Empty) => Poll::Pending, diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs index d1617a6..fe6d3f2 100644 --- a/asyncband/src/internal/atomic_waker.rs +++ b/asyncband/src/internal/atomic_waker.rs @@ -207,15 +207,6 @@ impl AtomicWaker { } } - /// Removes the most recently registered waker without waking it. - /// - /// This method must not overlap with [`register`](Self::register). It may run concurrently with - /// [`wake`](Self::wake). - #[inline] - pub fn clear(&self) { - drop(self.take()); - } - #[inline] fn take(&self) -> Option { // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the @@ -279,19 +270,6 @@ mod tests { assert_eq!(counter.0.load(Ordering::Relaxed), 1); } - #[test] - fn clear_removes_registered_waker() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.register(&waker); - atomic_waker.clear(); - atomic_waker.wake(); - - assert_eq!(counter.0.load(Ordering::Relaxed), 0); - } - #[test] fn reregistering_same_task_does_not_clone_waker() { let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 7317005..6a057fc 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -264,23 +264,6 @@ fn unbounded_last_sender_drop_wakes_parked_receiver() { ); } -#[test] -fn dropping_unbounded_receiver_releases_registered_waker() { - let (_tx, mut rx) = mpsc::unbounded::<()>(); - let probe = Arc::new(WakeProbe(AtomicUsize::new(0))); - let waker = Waker::from(probe.clone()); - let baseline_refs = Arc::strong_count(&probe); - let mut context = Context::from_waker(&waker); - let mut recv = Box::pin(rx.recv()); - - assert!(recv.as_mut().poll(&mut context).is_pending()); - assert_eq!(Arc::strong_count(&probe), baseline_refs + 1); - - drop(recv); - drop(rx); - assert_eq!(Arc::strong_count(&probe), baseline_refs); -} - #[tokio::test] async fn send_recv_bounded() { let (tx, mut rx) = mpsc::bounded(1);