diff --git a/asyncband/src/internal/atomic_option_box.rs b/asyncband/src/internal/atomic_option_box.rs deleted file mode 100644 index 26d9227..0000000 --- a/asyncband/src/internal/atomic_option_box.rs +++ /dev/null @@ -1,167 +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. - -// This is derived from https://github.com/jorendorff/atomicbox/blob/07756444/src/atomic_option_box.rs. - -use std::marker::PhantomData; -use std::ptr; -use std::sync::atomic::AtomicPtr; -use std::sync::atomic::Ordering; - -/// Internal owning atomic pointer used to store an optional receiver waker. -pub(crate) struct AtomicOptionBox { - /// Pointer to a `T` value in the heap, representing `Some(t)`; - /// or a null pointer for `None`. - ptr: AtomicPtr, - - /// This effectively makes `AtomicOptionBox` non-`Send` and non-`Sync` - /// if `T` is non-`Send`. - phantom: PhantomData>, -} - -/// Mark `AtomicOptionBox` as safe to share across threads. -/// -/// This is safe because shared access to an `AtomicOptionBox` does not -/// provide shared access to any `T` value. However, it does provide the -/// ability to get a `Box` from another thread, so `T: Send` is required. -unsafe impl Sync for AtomicOptionBox where T: Send {} - -fn into_ptr(value: Option>) -> *mut T { - match value { - Some(box_value) => Box::into_raw(box_value), - None => ptr::null_mut(), - } -} - -// SAFETY: The caller must ensure that `ptr` was obtained from `Box::into_raw` or is null. -unsafe fn from_ptr(ptr: *mut T) -> Option> { - if ptr.is_null() { - None - } else { - Some(unsafe { Box::from_raw(ptr) }) - } -} - -impl AtomicOptionBox { - /// Creates a new `AtomicOptionBox` with no value. - pub(crate) const fn none() -> Self { - Self { - ptr: AtomicPtr::new(ptr::null_mut()), - phantom: PhantomData, - } - } - - fn swap(&self, other: Option>) -> Option> { - let order = match other { - Some(_) => Ordering::AcqRel, - None => Ordering::Acquire, - }; - let new_ptr = into_ptr(other); - let old_ptr = self.ptr.swap(new_ptr, order); - unsafe { from_ptr(old_ptr) } - } - - /// Stores `other` and drops the previous value. - pub(crate) fn store(&self, other: Option>) { - drop(self.swap(other)); - } - - /// Replaces the value with `None` and returns the previous value. - pub(crate) fn take(&self) -> Option> { - self.swap(None) - } -} - -impl Drop for AtomicOptionBox { - /// Dropping an `AtomicOptionBox` drops the final `Box` value (if any) stored in it. - fn drop(&mut self) { - let ptr = *self.ptr.get_mut(); - unsafe { drop(from_ptr(ptr)) } - } -} - -#[cfg(test)] -mod tests { - use core::sync::atomic::Ordering; - use std::sync::Arc; - use std::sync::atomic::AtomicUsize; - - use super::*; - - #[test] - fn atomic_option_box_swap_works() { - let b = AtomicOptionBox::none(); - let bis = Box::new("bis"); - assert_eq!(b.swap(Some(bis)), None); - assert_eq!(b.swap(None), Some(Box::new("bis"))); - } - - #[test] - fn atomic_option_box_store_works() { - let b = AtomicOptionBox::none(); - let bis = Box::new("bis"); - b.store(Some(bis)); - assert_eq!(b.take(), Some(Box::new("bis"))); - assert_eq!(b.take(), None); - } - - #[test] - fn atomic_option_box_pointer_identity() { - let box1 = Box::new(1); - let p1 = &*box1 as *const i32; - let atom = AtomicOptionBox::none(); - atom.store(Some(box1)); - - let box2 = Box::new(2); - let p2 = &*box2 as *const i32; - assert_ne!(p2, p1); - - let box3 = atom.swap(Some(box2)).unwrap(); // box1 out, box2 in - let p3 = &*box3 as *const i32; - assert_eq!(p3, p1); // box3 is box1 - - let box4 = atom.swap(None).unwrap(); // box2 out, None in - let p4 = &*box4 as *const i32; - assert_eq!(p4, p2); // box4 is box2 - } - - #[test] - fn stored_values_are_dropped() { - struct K(Arc, usize); - - impl Drop for K { - fn drop(&mut self) { - self.0.fetch_add(self.1, Ordering::Relaxed); - } - } - - let n = Arc::new(AtomicUsize::new(0)); - { - let ab = AtomicOptionBox::none(); - ab.store(Some(Box::new(K(n.clone(), 5)))); - assert_eq!(n.load(Ordering::Relaxed), 0); - let first = ab.swap(None); - assert_eq!(n.load(Ordering::Relaxed), 0); - drop(first); - assert_eq!(n.load(Ordering::Relaxed), 5); - let second = ab.swap(Some(Box::new(K(n.clone(), 13)))); - assert!(second.is_none()); - assert_eq!(n.load(Ordering::Relaxed), 5); - } - assert_eq!(n.load(Ordering::Relaxed), 5 + 13); - } -} diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs new file mode 100644 index 0000000..fe6d3f2 --- /dev/null +++ b/asyncband/src/internal/atomic_waker.rs @@ -0,0 +1,463 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// This state machine is derived from the futures-rs `AtomicWaker`, licensed under +// Apache-2.0 OR MIT: https://github.com/rust-lang/futures-rs/blob/0.3.34/futures-core/src/task/__internal/atomic_waker.rs. +// Its panic recovery is informed by Tokio's `AtomicWaker`, licensed under MIT: +// https://github.com/tokio-rs/tokio/blob/tokio-1.53.1/tokio/src/sync/task/atomic_waker.rs. + +use std::cell::UnsafeCell; +use std::panic::AssertUnwindSafe; +use std::panic::RefUnwindSafe; +use std::panic::UnwindSafe; +use std::panic::catch_unwind; +use std::panic::resume_unwind; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Waker; + +const WAITING: usize = 0; +const REGISTERING: usize = 0b01; +const WAKING: usize = 0b10; + +/// A single-registerer, multi-notifier cell for task wake-up. +/// +/// The atomic state both grants exclusive access to `waker` and records one coalesced wake request. +/// The operation that moves the state out of `WAITING` remains the only slot owner until it returns +/// the state to `WAITING`. +/// +/// * `WAITING`: the slot is unlocked and may contain a registered waker. +/// * `REGISTERING`: `register` exclusively owns the slot and no concurrent wake is pending. +/// * `WAKING`: `wake` exclusively owns the slot. A racing `register` self-wakes without touching +/// the slot. +/// * `REGISTERING | WAKING`: `register` still owns the slot and must complete a concurrent wake +/// before returning to `WAITING`. +/// +/// Valid state transitions are: +/// +/// ```text +/// register: WAITING ----------------Acquire CAS---------------> REGISTERING +/// REGISTERING ------------AcqRel CAS----------------> WAITING +/// +/// wake: WAITING ----------------AcqRel fetch_or-----------> WAKING +/// WAKING -----------------Release swap--------------> WAITING +/// +/// race: REGISTERING ------------AcqRel fetch_or-----------> REGISTERING | WAKING +/// REGISTERING | WAKING ---AcqRel swap---------------> WAITING +/// ``` +/// +/// Additional calls to `wake` while `WAKING` is set are coalesced. A wake completed before a +/// registration starts is not remembered, so callers must register before rechecking the condition +/// that determines whether to return `Pending`. +/// +/// Every transition that acquires slot ownership has an Acquire operation paired with the previous +/// owner's Release transition to `WAITING`. The Release half of `wake` also publishes the caller's +/// preceding condition update; a racing `register` acquires that publication before it returns. +pub struct AtomicWaker { + state: AtomicUsize, + waker: UnsafeCell>, +} + +// SAFETY: `state` grants exclusive access to `waker`, and losing concurrent registrations do not +// touch the slot. `Waker` itself is `Send + Sync`. +unsafe impl Sync for AtomicWaker {} + +// `Waker` callbacks may unwind, but no panic leaves a state bit owned by the unwinding operation. A +// failed clone leaves the old slot intact and completes any raced wake, while wake and drop +// callbacks run after that operation's critical section has been released. +impl RefUnwindSafe for AtomicWaker {} +impl UnwindSafe for AtomicWaker {} + +impl AtomicWaker { + #[inline] + pub const fn new() -> Self { + Self { + state: AtomicUsize::new(WAITING), + waker: UnsafeCell::new(None), + } + } + + /// Registers `waker`, replacing a previously registered task if it differs. + /// + /// Calls to this method must not overlap. It may run concurrently with any number of calls to + /// [`wake`](Self::wake). + #[inline] + pub fn register(&self, waker: &Waker) { + // ORDERING: On success, Acquire pairs with the Release operation that last returned the + // state to WAITING and transfers exclusive ownership of the waker slot to this thread. On + // failure, Acquire matters when this reads WAKING from a notifier's AcqRel fetch_or: it + // receives the condition update that preceded that wake before this method returns. + match self + .state + .compare_exchange(WAITING, REGISTERING, Ordering::Acquire, Ordering::Acquire) + .unwrap_or_else(|state| state) + { + WAITING => { + // SAFETY: changing WAITING to REGISTERING grants this thread exclusive access to + // the waker slot until the state is returned to WAITING. + unsafe { self.register_locked(waker) } + } + WAKING => { + // A concurrent wake owns the slot. Self-waking ensures that this registration is + // not lost even though it cannot replace the slot right now. + waker.wake_by_ref(); + } + state => { + // Concurrent registration violates this type's contract. Ignoring the losing + // registration preserves memory safety and lets the winner provide notification. + debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); + } + } + } + + /// Registers a waker after this thread has acquired the REGISTERING state. + /// + /// # Safety + /// + /// The caller must have changed `state` from WAITING to REGISTERING and must be the only + /// thread accessing `waker`. + #[inline] + unsafe fn register_locked(&self, waker: &Waker) { + // Avoid both cloning and dropping the common case where an executor polls the receiver + // repeatedly with the same task waker. + let needs_replacement = match unsafe { &*self.waker.get() } { + Some(current) => !current.will_wake(waker), + None => true, + }; + + let mut clone_panic = None; + let old_waker = if needs_replacement { + match catch_unwind(AssertUnwindSafe(|| waker.clone())) { + Ok(new_waker) => unsafe { (*self.waker.get()).replace(new_waker) }, + Err(payload) => { + clone_panic = Some(payload); + None + } + } + } else { + None + }; + + // ORDERING: Release publishes a newly registered waker when the CAS succeeds. If it fails, + // Acquire receives the concurrent notifier's Release publication before the wake is + // completed below. AcqRel is the weakest success ordering that permits an Acquire failure + // ordering, although its Acquire half is not otherwise relied upon on the success path. + let concurrent_wake = match self.state.compare_exchange( + REGISTERING, + WAITING, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => None, + Err(state) => { + debug_assert_eq!(state, REGISTERING | WAKING); + + // SAFETY: REGISTERING remains set, so this thread still owns the waker slot. + let registered = unsafe { (*self.waker.get()).take() }; + + // ORDERING: Acquire receives all coalesced wake publications. Release publishes + // the empty slot and makes it available to the next register or wake operation. + self.state.swap(WAITING, Ordering::AcqRel); + registered + } + }; + + if let Some(payload) = clone_panic { + // Preserve the original clone panic while still completing a wake that raced with it. + if let Some(waker) = concurrent_wake { + let _ = catch_unwind(AssertUnwindSafe(|| waker.wake())); + } + resume_unwind(payload); + } + + // User waker code runs only after the state machine is back in WAITING, so a panic cannot + // leave the cell locked. If the wake raced with a replacement, notify both tasks: the + // concurrent call may have targeted the old registration, while future progress relies on + // the new one. A panic from the superseded waker must not prevent the new task from waking. + if let Some(waker) = concurrent_wake { + if let Some(old_waker) = old_waker { + let _ = catch_unwind(AssertUnwindSafe(|| old_waker.wake())); + } + waker.wake(); + } else { + // Drop a replaced waker only after releasing the state lock. + drop(old_waker); + } + } + + /// Wakes and removes the most recently registered waker, if any. + #[inline] + pub fn wake(&self) { + if let Some(waker) = self.take() { + waker.wake(); + } + } + + #[inline] + fn take(&self) -> Option { + // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the + // previous owner. Release publishes the condition update that the caller performed before + // calling wake, including when a registering thread already owns the slot. + match self.state.fetch_or(WAKING, Ordering::AcqRel) { + WAITING => { + // SAFETY: changing WAITING to WAKING grants this thread exclusive access to the + // waker slot until the state is returned to WAITING. + let waker = unsafe { (*self.waker.get()).take() }; + + // ORDERING: Release publishes the emptied slot before another operation acquires + // it. The fetch_or above already performed the required Acquire operation. + let old_state = self.state.swap(WAITING, Ordering::Release); + debug_assert_eq!(old_state, WAKING); + waker + } + state => { + // The thread registering a waker observes WAKING and completes this notification, + // or another waking thread has already taken responsibility for it. + debug_assert!( + state == REGISTERING || state == REGISTERING | WAKING || state == WAKING + ); + None + } + } + } +} + +#[cfg(test)] +mod tests { + use std::ptr; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::task::RawWaker; + use std::task::RawWakerVTable; + use std::task::Wake; + + use super::*; + + struct WakeCounter(AtomicUsize); + + impl Wake for WakeCounter { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[test] + fn wake_notifies_once() { + 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.wake(); + atomic_waker.wake(); + + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + } + + #[test] + fn reregistering_same_task_does_not_clone_waker() { + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let atomic_waker = AtomicWaker::new(); + + atomic_waker.register(&waker); + let registered_refs = Arc::strong_count(&counter); + atomic_waker.register(&waker); + + assert_eq!(Arc::strong_count(&counter), registered_refs); + } + + #[test] + fn wake_before_register_is_not_remembered() { + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let atomic_waker = AtomicWaker::new(); + + atomic_waker.wake(); + atomic_waker.register(&waker); + + assert_eq!(counter.0.load(Ordering::Relaxed), 0); + atomic_waker.wake(); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + } + + #[test] + fn wake_during_replacement_notifies_old_and_new_tasks() { + let old_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let old_waker = Waker::from(old_counter.clone()); + let new_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let new_waker = Waker::from(new_counter.clone()); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(&old_waker); + + assert_eq!( + atomic_waker.state.compare_exchange( + WAITING, + REGISTERING, + Ordering::AcqRel, + Ordering::Acquire, + ), + Ok(WAITING) + ); + std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); + + // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching + // the slot. Calling the helper completes the interrupted registration. + unsafe { atomic_waker.register_locked(&new_waker) }; + + assert_eq!(old_counter.0.load(Ordering::Relaxed), 1); + assert_eq!(new_counter.0.load(Ordering::Relaxed), 1); + } + + #[test] + fn failed_wake_synchronizes_with_next_registration() { + for _ in 0..1_000 { + let did_publish = AtomicBool::new(false); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(Waker::noop()); + + std::thread::scope(|scope| { + let wake = scope.spawn(|| { + did_publish.store(true, Ordering::Relaxed); + atomic_waker.take() + }); + + let local_waker = atomic_waker.take(); + atomic_waker.register(Waker::noop()); + + let publication_is_visible = did_publish.load(Ordering::Relaxed); + let concurrent_thread_took_waker = wake.join().unwrap().is_some(); + assert!(publication_is_visible || concurrent_thread_took_waker); + drop(local_waker); + }); + } + } + + #[cfg(panic = "unwind")] + #[test] + fn clone_panic_does_not_poison_state() { + static PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( + |_| panic!("clone failed"), + |_| unreachable!(), + |_| unreachable!(), + |_| {}, + ); + + let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) }; + let atomic_waker = AtomicWaker::new(); + + assert!( + catch_unwind(|| { + atomic_waker.register(&panicking); + }) + .is_err() + ); + + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + atomic_waker.register(&Waker::from(counter.clone())); + atomic_waker.wake(); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + } + + #[cfg(panic = "unwind")] + #[test] + fn clone_panic_completes_concurrent_wake() { + static PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( + |_| panic!("clone failed"), + |_| unreachable!(), + |_| unreachable!(), + |_| {}, + ); + + let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) }; + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(&Waker::from(counter.clone())); + + assert_eq!( + atomic_waker.state.compare_exchange( + WAITING, + REGISTERING, + Ordering::Acquire, + Ordering::Acquire, + ), + Ok(WAITING) + ); + std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); + + // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching + // the state. Calling the helper completes the interrupted registration. + assert!( + catch_unwind(|| unsafe { + atomic_waker.register_locked(&panicking); + }) + .is_err() + ); + + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + + let next_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + atomic_waker.register(&Waker::from(next_counter.clone())); + atomic_waker.wake(); + assert_eq!(next_counter.0.load(Ordering::Relaxed), 1); + } + + #[cfg(panic = "unwind")] + #[test] + fn drop_panic_does_not_poison_state() { + unsafe fn clone_drop_panicker(data: *const ()) -> RawWaker { + RawWaker::new(data, &DROP_PANICKING_VTABLE) + } + + unsafe fn wake_drop_panicker(_: *const ()) {} + + unsafe fn drop_drop_panicker(data: *const ()) { + // SAFETY: the test keeps the pointed-to AtomicBool alive until every derived waker has + // been dropped. + let should_panic = unsafe { &*data.cast::() }; + if should_panic.swap(false, Ordering::Relaxed) { + panic!("drop failed"); + } + } + + static DROP_PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( + clone_drop_panicker, + wake_drop_panicker, + wake_drop_panicker, + drop_drop_panicker, + ); + + let should_panic = AtomicBool::new(true); + let old_waker = unsafe { + Waker::from_raw(RawWaker::new( + ptr::from_ref(&should_panic).cast(), + &DROP_PANICKING_VTABLE, + )) + }; + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let new_waker = Waker::from(counter.clone()); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(&old_waker); + + assert!(catch_unwind(|| atomic_waker.register(&new_waker)).is_err()); + + atomic_waker.wake(); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + } +} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index caac4a2..248f9ee 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -16,7 +16,7 @@ // under the License. #[cfg(feature = "mpsc")] -pub(crate) mod atomic_option_box; +pub(crate) mod atomic_waker; #[cfg(any( feature = "barrier", diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index 0e90bbf..60f0da0 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -27,9 +27,8 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; -use std::task::Waker; -use crate::internal::atomic_option_box::AtomicOptionBox; +use crate::internal::atomic_waker::AtomicWaker; use crate::internal::semaphore::Acquire; use crate::internal::semaphore::Semaphore; use crate::mpsc::RecvError; @@ -49,7 +48,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { let state = Arc::new(BoundedState { senders: AtomicUsize::new(1), tx_permits: Semaphore::new(0), - rx_task: AtomicOptionBox::none(), + rx_waker: AtomicWaker::new(), }); let (sender, receiver) = std::sync::mpsc::sync_channel(buffer); let sender = BoundedSender { @@ -66,7 +65,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { struct BoundedState { senders: AtomicUsize, tx_permits: Semaphore, - rx_task: AtomicOptionBox, + rx_waker: AtomicWaker, } /// Send values to the associated [`BoundedReceiver`]. @@ -102,9 +101,7 @@ impl Drop for BoundedSender { 1 => { // If this is the last sender, we need to wake up the receiver so it can // observe the disconnected state. - if let Some(waker) = self.state.rx_task.take() { - waker.wake(); - } + self.state.rx_waker.wake(); } _ => { // there are still other senders left, do nothing @@ -203,9 +200,7 @@ impl BoundedSender { let sender = self.sender.as_ref().unwrap(); match sender.try_send(value) { Ok(()) => { - if let Some(waker) = self.state.rx_task.take() { - waker.wake(); - } + self.state.rx_waker.wake(); Ok(()) } @@ -347,8 +342,7 @@ impl BoundedReceiver { Ok(v) => Poll::Ready(Ok(v)), Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), Err(TryRecvError::Empty) => { - let waker = Some(Box::new(cx.waker().clone())); - self.state.rx_task.store(waker); + self.state.rx_waker.register(cx.waker()); match self.try_recv() { Ok(v) => Poll::Ready(Ok(v)), diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index 06be655..eb446be 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -25,9 +25,8 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; -use std::task::Waker; -use crate::internal::atomic_option_box::AtomicOptionBox; +use crate::internal::atomic_waker::AtomicWaker; use crate::mpsc::RecvError; use crate::mpsc::SendError; use crate::mpsc::TryRecvError; @@ -44,7 +43,7 @@ use crate::mpsc::TryRecvError; pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { senders: AtomicUsize::new(1), - rx_task: AtomicOptionBox::none(), + rx_waker: AtomicWaker::new(), }); let (sender, receiver) = std::sync::mpsc::channel(); let sender = UnboundedSender { @@ -60,7 +59,7 @@ pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { struct UnboundedState { senders: AtomicUsize, - rx_task: AtomicOptionBox, + rx_waker: AtomicWaker, } /// Send values to the associated [`UnboundedReceiver`]. @@ -96,9 +95,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. - if let Some(waker) = self.state.rx_task.take() { - waker.wake(); - } + self.state.rx_waker.wake(); } _ => { // there are still other senders left, do nothing @@ -121,9 +118,7 @@ impl UnboundedSender { let sender = self.sender.as_ref().unwrap(); sender.send(value).map_err(|err| SendError::new(err.0))?; - if let Some(waker) = self.state.rx_task.take() { - waker.wake(); - } + self.state.rx_waker.wake(); Ok(()) } @@ -246,8 +241,7 @@ impl UnboundedReceiver { Ok(v) => Poll::Ready(Ok(v)), Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), Err(TryRecvError::Empty) => { - let waker = Some(Box::new(cx.waker().clone())); - self.state.rx_task.store(waker); + self.state.rx_waker.register(cx.waker()); match self.try_recv() { Ok(v) => Poll::Ready(Ok(v)), diff --git a/benchmarks/mpsc.rs b/benchmarks/mpsc.rs index 8764734..f86a7ca 100644 --- a/benchmarks/mpsc.rs +++ b/benchmarks/mpsc.rs @@ -25,6 +25,29 @@ use super::support::poll_pinned_ready; const SENDER_COUNTS: &[usize] = &[1, 8, 32]; +#[divan::bench] +fn reregister_pending_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (_sender, mut receiver) = mpsc::unbounded::(); + let mut recv = Box::pin(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + + bencher.bench_local(|| poll_pending(recv.as_mut(), &mut context)); +} + +#[divan::bench] +fn wake_pending_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (sender, mut receiver) = mpsc::unbounded(); + + bencher.bench_local(|| { + let mut recv = Box::pin(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + sender.send(black_box(usize::MAX)).unwrap(); + black_box(poll_pinned_ready(recv.as_mut(), &mut context).unwrap()) + }); +} + #[divan::bench(args = SENDER_COUNTS)] fn cancel_backpressured_senders(bencher: Bencher, sender_count: usize) { let mut context = bench_context();