diff --git a/CHANGELOG.md b/CHANGELOG.md index 214b359..d6d6f5e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,7 @@ All notable changes to this project will be documented in this file. * Release cancelled wait registrations promptly and reclaim fulfilled `Semaphore::reduce_permits` debt nodes. * Preserve fan-out notifications when one registered waker panics. +* Clone task wakers outside primitive state locks so reentrant clone callbacks cannot deadlock wait registration. ### Improvements diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index a3c2af7..69cef81 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -180,7 +180,8 @@ impl Barrier { state: Mutex::new(BarrierState { arrived: 0, generation: 0, - waiters: WaitSet::with_capacity(n as usize), + // The final participant completes the generation without parking. + waiters: WaitSet::with_capacity((n - 1) as usize), }), } } @@ -239,12 +240,12 @@ impl Barrier { let generation = state.generation; state.arrived += 1; - // the last arriver is the leader; - // wake up other waiters, increment the generation, and return + // The final arrival completes this generation. Advance the generation while holding + // the state lock, then wake the drained followers after releasing it. if state.arrived == self.n { state.arrived = 0; state.generation += 1; - let wakers = state.waiters.take_wakers(); + let wakers = state.waiters.drain(); drop(state); wake_all(wakers); return BarrierWaitResult(true); @@ -291,16 +292,24 @@ impl Future for BarrierWait<'_> { barrier, } = self.get_mut(); - let replaced_waker = { - let mut state = barrier.state.lock(); - if *generation < state.generation { - // Advancing the generation drains its registrations under this same lock. - *token = None; - return Poll::Ready(()); - } - state.waiters.register_waker(token, cx) - }; - drop(replaced_waker); + // A follower normally parks once, so cloning first keeps its common pending path to one + // state-lock acquisition. Cloning may reenter and complete the barrier; checking the + // generation afterward closes that race. The completion poll may clone an unused waker, + // which is the deliberate cost of avoiding a second lock-and-recheck phase here. + let waker = cx.waker().clone(); + let mut state = barrier.state.lock(); + if *generation < state.generation { + // Completion advances the generation and drains its old waiters under this same lock, + // so no registration represented by this token remains in the wait set. + *token = None; + drop(state); + drop(waker); + return Poll::Ready(()); + } + + let retired_waker = state.waiters.register(token, waker); + drop(state); + drop(retired_waker); Poll::Pending } } @@ -310,7 +319,7 @@ impl Drop for BarrierWait<'_> { if self.token.is_some() { let removed_waker = { let mut state = self.barrier.state.lock(); - state.waiters.unregister_waker(&mut self.token) + state.waiters.unregister(&mut self.token) }; drop(removed_waker); } diff --git a/asyncband/src/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/broadcast/mpmc/unbounded/mod.rs index 9e17089..2207d59 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -39,9 +39,9 @@ //! [`UnboundedReceiver::resubscribe`] to create a receiver that starts at the current tail. //! //! Messages are reclaimed once the slowest receiver moves past them, which scans one slot per -//! receiver. Only the receiver that advances the slowest cursor pays for that scan, and the -//! channel keeps a slot for every receiver it hands out, so the cost follows the largest number -//! of receivers that were ever active at once rather than the number active now. +//! receiver. Only the receiver that advances the slowest cursor pays for that scan. The channel +//! keeps a slot for every receiver it hands out, so the cost follows the largest number of +//! receivers that were ever active at once rather than the number active now. //! //! # Examples //! @@ -188,10 +188,10 @@ const MIN_RETAINED_CAPACITY: usize = 64; struct Inner { /// Messages whose versions are in the range `[head, tail)`. /// - /// Each message is held behind an `Arc` so a receive can hand the payload out of the critical - /// section. Cloning the `Arc` under the lock keeps `T::clone` — and, for reclaimed messages, - /// `T::drop` — outside it, which matters because both are arbitrary user code that may call - /// back into this channel. + /// Each message is held behind an `Arc` so the receive path can move the payload out of the + /// critical section. Cloning the `Arc` under the lock keeps `T::clone` — and, for reclaimed + /// messages, `T::drop` — outside it, which matters because both are arbitrary user code that + /// may call back into this channel. buffer: VecDeque>, /// The version of the first message in `buffer`. head: u64, @@ -396,7 +396,10 @@ impl Drop for UnboundedSender { match self.shared.senders.fetch_sub(1, Ordering::AcqRel) { 1 => { // Wake every parked receiver so it can observe the channel's disconnected state. - let wakers = self.shared.inner.lock().waiters.take_wakers(); + let wakers = { + let mut inner = self.shared.inner.lock(); + inner.waiters.drain() + }; wake_all(wakers); } _ => { @@ -454,7 +457,7 @@ impl UnboundedSender { inner.peak_len = inner.peak_len.max(inner.buffer.len()); } - inner.waiters.take_wakers() + inner.waiters.drain() }; // Notify all waiting receivers. An unsent message is dropped here too, once the lock is @@ -568,7 +571,7 @@ impl UnboundedReceiver { pub async fn recv(&mut self) -> Result { Recv { receiver: self, - registration: None, + token: None, } .await } @@ -701,19 +704,19 @@ impl UnboundedReceiver { struct Recv<'a, T> { receiver: &'a mut UnboundedReceiver, - registration: Option, + token: Option, } impl Drop for Recv<'_, T> { fn drop(&mut self) { - // Ready paths clear the registration, so only a cancelled pending receive takes this lock. - if self.registration.is_none() { + // Ready paths clear the token, so only a cancelled pending receive takes this lock. + if self.token.is_none() { return; } let waker = { let mut inner = self.receiver.shared.inner.lock(); - inner.waiters.unregister_waker(&mut self.registration) + inner.waiters.unregister(&mut self.token) }; drop(waker); } @@ -723,35 +726,46 @@ impl Future for Recv<'_, T> { type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let Self { - receiver, - registration, - } = self.get_mut(); - - // One critical section decides between all three outcomes. Senders append messages and - // drain the wait set under this same lock, so registering here cannot miss a wake-up and - // cannot report disconnection while a message remains for this receiver. - let received = { + let Self { receiver, token } = self.get_mut(); + + // Buffered messages and repeated polls with the same task waker require no clone. If the + // pending path needs a new waker, release the lock, clone, and repeat the full state check + // before registration. Senders publish messages and drain waiters under the same lock, so + // the recheck cannot miss a send, disconnection, or state change made by a reentrant clone + // callback. The loop executes at most twice. + let mut prepared_waker = None; + let received = loop { let mut inner = receiver.shared.inner.lock(); match inner.receive(receiver.key) { - Some(received) => received, + Some(received) => break received, None => { if receiver.shared.senders.load(Ordering::Acquire) == 0 { - *registration = None; + *token = None; + drop(inner); + drop(prepared_waker); return Poll::Ready(Err(RecvError::Disconnected)); } - let waker = inner.waiters.register_waker(registration, cx); + if prepared_waker.is_none() && inner.waiters.will_wake(token, cx.waker()) { + return Poll::Pending; + } + let Some(waker) = prepared_waker.take() else { + drop(inner); + prepared_waker = Some(cx.waker().clone()); + continue; + }; + let retired_waker = inner.waiters.register(token, waker); drop(inner); - drop(waker); + drop(retired_waker); return Poll::Pending; } } }; + drop(prepared_waker); let (msg, reclaimed) = received; - *registration = None; + *token = None; Poll::Ready(Ok(take_msg(msg, reclaimed))) } } diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index b6d805c..81a8bff 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -158,15 +158,14 @@ impl Completer { drop(value); panic!("pending completion value must be unset"); } + // Publish the value before making completion observable and detaching its waiters. state.status = Status::Completed; - (!state.waiters.is_empty()).then(|| state.waiters.take_wakers()) + state.waiters.drain() }; // `complete` consumes the only completer. Disarm its destructor before invoking arbitrary // wake callbacks; the completed state no longer needs abandonment handling. self.shared = Weak::new(); - if let Some(wakers) = wakers { - wake_all(wakers); - } + wake_all(wakers); Ok(()) } } @@ -181,12 +180,11 @@ impl Drop for Completer { if state.status != Status::Pending { return; } + // Publish abandonment and detach its waiters atomically with respect to registration. state.status = Status::Abandoned; - (!state.waiters.is_empty()).then(|| state.waiters.take_wakers()) + state.waiters.drain() }; - if let Some(wakers) = wakers { - wake_all(wakers); - } + wake_all(wakers); } } @@ -225,7 +223,7 @@ impl Completion { pub async fn wait(&self) -> Result<&T, Abandoned> { Wait { completion: self, - registration: None, + token: None, } .await } @@ -233,7 +231,7 @@ impl Completion { struct Wait<'a, T> { completion: &'a Completion, - registration: Option, + token: Option, } impl<'a, T> Future for Wait<'a, T> { @@ -241,6 +239,10 @@ impl<'a, T> Future for Wait<'a, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); + // Terminal waits require no waker, so inspect the state before cloning. If a pending wait + // needs a new waker, release the lock, clone, and repeat the full state check before + // registration. The loop executes at most twice: once to prepare a waker and once to + // register it or observe the intervening terminal transition. let mut prepared_waker = None; loop { let (poll, retired_waker) = { @@ -248,9 +250,7 @@ impl<'a, T> Future for Wait<'a, T> { match state.status { Status::Pending => { if prepared_waker.is_none() - && state - .waiters - .registered_waker_will_wake(&this.registration, cx.waker()) + && state.waiters.will_wake(&this.token, cx.waker()) { return Poll::Pending; } @@ -259,13 +259,11 @@ impl<'a, T> Future for Wait<'a, T> { prepared_waker = Some(cx.waker().clone()); continue; }; - let retired = state - .waiters - .register_owned_waker(&mut this.registration, waker); + let retired = state.waiters.register(&mut this.token, waker); (Poll::Pending, retired) } Status::Completed => { - let retired = state.waiters.unregister_waker(&mut this.registration); + let retired = state.waiters.unregister(&mut this.token); let completion: &'a Completion = this.completion; let value = completion .shared @@ -275,7 +273,7 @@ impl<'a, T> Future for Wait<'a, T> { (Poll::Ready(Ok(value)), retired) } Status::Abandoned => { - let retired = state.waiters.unregister_waker(&mut this.registration); + let retired = state.waiters.unregister(&mut this.token); (Poll::Ready(Err(Abandoned)), retired) } } @@ -289,9 +287,13 @@ impl<'a, T> Future for Wait<'a, T> { impl Drop for Wait<'_, T> { fn drop(&mut self) { + if self.token.is_none() { + return; + } + let waker = { let mut state = self.completion.shared.state.lock(); - state.waiters.unregister_waker(&mut self.registration) + state.waiters.unregister(&mut self.token) }; drop(waker); } diff --git a/asyncband/src/internal/arena.rs b/asyncband/src/internal/arena.rs index 3158d31..2cb762c 100644 --- a/asyncband/src/internal/arena.rs +++ b/asyncband/src/internal/arena.rs @@ -183,8 +183,9 @@ impl Arena { /// Takes every occupied value in slot order while retaining the allocation for reuse. /// - /// Every previously issued slot ID becomes invalid, including IDs for slots that were already - /// vacant. Consumers that retain IDs across this operation must supply their own epoch check. + /// After a non-empty take, every previously issued slot ID becomes invalid, including IDs for + /// slots that were already vacant. Consumers that retain IDs across this operation must supply + /// their own epoch check. #[inline] pub fn take_all(&mut self) -> impl Iterator + use { let len = self.len; @@ -192,6 +193,12 @@ impl Arena { first: None, rest: Vec::new(), }; + if len == 0 { + // Individually removed values leave vacant slots behind. Keep their free list intact + // instead of scanning the arena's historical high-water mark to drain no values. + return values.into_iter(); + } + for slot in self.slots.drain(..) { if let Slot::Occupied(value) = slot { if values.first.is_none() { diff --git a/asyncband/src/internal/countdown.rs b/asyncband/src/internal/countdown.rs index 524a792..10ec6d6 100644 --- a/asyncband/src/internal/countdown.rs +++ b/asyncband/src/internal/countdown.rs @@ -39,32 +39,25 @@ impl CountdownState { } } - /// Performs volatile read on `state`. - /// - /// All other writes to `state` should be at least [`Ordering::Release`]. + /// Loads the current count, acquiring state published before a transition to zero. pub fn state(&self) -> u32 { self.state.load(Ordering::Acquire) } - /// Performs volatile CAS on `state`. - /// - /// If the comparison succeeds, performs read-modify-write operation with [`Ordering::Relaxed`] - /// for read, and [`Ordering::Release`] for write; if the comparison fails, performs load - /// operation with [`Ordering::Relaxed`]. + /// Attempts to replace `current` with `new`, publishing the new count on success. /// - /// @see https://doc.rust-lang.org/std/sync/atomic/struct.AtomicU32.html#method.compare_exchange_weak - /// @see https://en.cppreference.com/w/cpp/atomic/atomic_compare_exchange + /// A spurious or contended failure returns the observed count so the caller can retry. fn cas_state(&self, current: u32, new: u32) -> Result<(), u32> { self.state .compare_exchange_weak(current, new, Ordering::Release, Ordering::Relaxed) .map(|_| ()) } - /// Drain and wake up all waiters. + /// Drains the waiter set under its lock, then wakes every waiter after releasing the lock. pub fn wake_all(&self) { let wakers = { let mut waiters = self.waiters.lock(); - waiters.take_wakers() + waiters.drain() }; wake_all(wakers); @@ -79,25 +72,32 @@ impl CountdownState { return Poll::Ready(()); } - let replaced_waker = { - let mut waiters = self.waiters.lock(); - if self.state() == 0 { - // A concurrent zero transition will drain after this lock is released. - *token = None; - return Poll::Ready(()); - } - waiters.register_waker(token, cx) - }; - drop(replaced_waker); + // The atomic probe keeps an already-ready poll free of both waker cloning and waiter + // locking. An active countdown normally parks, so eagerly clone before the waiter lock to + // keep that path to one acquisition. Rechecking afterward observes a zero transition made + // concurrently or by a reentrant clone callback. + let waker = cx.waker().clone(); + let mut waiters = self.waiters.lock(); + if self.state() == 0 { + // A concurrent zero transition will drain after this lock is released. + *token = None; + drop(waiters); + drop(waker); + return Poll::Ready(()); + } + + let retired_waker = waiters.register(token, waker); + drop(waiters); + drop(retired_waker); Poll::Pending } #[inline] - pub fn unregister_waker(&self, token: &mut Option) { + pub fn unregister(&self, token: &mut Option) { if token.is_some() { let removed_waker = { let mut waiters = self.waiters.lock(); - waiters.unregister_waker(token) + waiters.unregister(token) }; drop(removed_waker); } @@ -140,8 +140,7 @@ impl CountdownState { let mut cnt = self.state(); loop { if cnt == 0 { - // the one who decrements the counter to zero should wake up all waiters, not this - // one + // Only the operation that performs the transition to zero owns waiter notification. return false; } diff --git a/asyncband/src/internal/waitset.rs b/asyncband/src/internal/waitset.rs index ded38a0..1e7b983 100644 --- a/asyncband/src/internal/waitset.rs +++ b/asyncband/src/internal/waitset.rs @@ -15,10 +15,17 @@ // specific language governing permissions and limitations // under the License. +//! Cancellable storage for task wakers. +//! +//! A `WaitSet` is protected by the state lock of its owning primitive, but it never invokes a +//! [`Waker`] callback itself. It accepts owned wakers and returns every waker it stops owning, so +//! callers can clone before locking—or unlock, clone, and recheck—and can drop or wake returned +//! wakers after unlocking. This keeps arbitrary clone, drop, and wake callbacks outside primitive +//! critical sections. + use std::mem; use std::panic; use std::panic::AssertUnwindSafe; -use std::task::Context; use std::task::Waker; use crate::internal::arena::Arena; @@ -52,10 +59,11 @@ pub fn wake_all(mut wakers: impl Iterator) { } } -/// A single-owner token for a waker registered in a [`WaitSet`]. +/// An exclusive handle to one waiter slot in a [`WaitSet`]. /// -/// This deliberately does not implement `Clone` or `Copy`: duplicating a token could let a stale -/// token refer to a slot reused by another waiter in the same epoch. +/// The wait set owns the registered waker; this token only lets its future update or cancel that +/// registration. It deliberately does not implement `Clone` or `Copy`, because duplicating the +/// handle could let a stale token refer to a slot reused by another waiter in the same epoch. #[derive(Debug)] pub struct WakerToken { epoch: u64, @@ -85,46 +93,25 @@ impl WaitSet { } } - /// Returns whether no wakers are currently registered. - #[inline] - pub fn is_empty(&self) -> bool { - self.waiters.is_empty() - } - - /// Takes all registered wakers as an owning iterator without waking them. - #[inline] - pub fn take_wakers(&mut self) -> impl Iterator + 'static { - self.epoch = self.epoch.checked_add(1).expect("wait set epoch overflow"); - self.waiters.take_all() - } - - /// Registers or updates a waker in the current wake epoch. + /// Drains all registered wakers as an owning iterator without waking them. /// - /// If an existing waker is replaced, it is returned so the caller can drop it after releasing - /// the lock that protects this wait set. + /// A non-empty drain starts a new epoch so tokens retained by the drained futures cannot alias + /// slots reused by later registrations. The caller must consume or drop the iterator after + /// releasing the lock that protects this wait set. #[inline] - pub fn register_waker( - &mut self, - token: &mut Option, - cx: &mut Context<'_>, - ) -> Option { - if self.registered_waker_will_wake(token, cx.waker()) { - return None; + pub fn drain(&mut self) -> impl Iterator + 'static { + if !self.waiters.is_empty() { + self.epoch = self.epoch.checked_add(1).expect("wait set epoch overflow"); } - if let Some(waker) = self.current_waker(token) { - return Some(mem::replace(waker, cx.waker().clone())); - } - - *token = Some(WakerToken { - epoch: self.epoch, - slot: self.waiters.insert(cx.waker().clone()), - }); - None + self.waiters.take_all() } /// Returns whether `token` identifies a registered waker for the same task as `waker`. + /// + /// This comparison neither clones a waker nor invokes its `RawWaker` callbacks, so callers may + /// use it while holding the lock that protects this wait set. #[inline] - pub fn registered_waker_will_wake(&self, token: &Option, waker: &Waker) -> bool { + pub fn will_wake(&self, token: &Option, waker: &Waker) -> bool { let Some(current) = token.as_ref() else { return false; }; @@ -139,14 +126,15 @@ impl WaitSet { /// Registers or updates an already cloned waker in the current wake epoch. /// + /// The caller must obtain the owned waker without holding the lock that protects this wait set, + /// because cloning can invoke arbitrary user code. If it released that lock to clone, it must + /// recheck the primitive's state after reacquiring the lock and before calling this method. + /// /// Any waker not retained by the wait set is returned so the caller can drop it after releasing /// the lock that protects this wait set. #[inline] - pub fn register_owned_waker( - &mut self, - token: &mut Option, - waker: Waker, - ) -> Option { + #[must_use = "drop the returned waker after releasing the wait set's state lock"] + pub fn register(&mut self, token: &mut Option, waker: Waker) -> Option { if let Some(current) = self.current_waker(token) { if !current.will_wake(&waker) { return Some(mem::replace(current, waker)); @@ -165,11 +153,13 @@ impl WaitSet { /// /// The returned waker must be dropped after releasing the lock that protects this wait set. #[inline] - pub fn unregister_waker(&mut self, token: &mut Option) -> Option { + #[must_use = "drop the returned waker after releasing the wait set's state lock"] + pub fn unregister(&mut self, token: &mut Option) -> Option { let token = token.take()?; if token.epoch == self.epoch { return Some(self.waiters.remove(token.slot)); } + // A drain advanced the epoch and already took ownership of this token's waker. None } @@ -253,7 +243,7 @@ mod tests { token: &mut Option, waker: &Waker, ) -> Option { - waiters.register_waker(token, &mut Context::from_waker(waker)) + waiters.register(token, waker.clone()) } #[test] @@ -267,12 +257,12 @@ mod tests { let mut second_token = None; register(&mut waiters, &mut first_token, &first_waker); - assert_eq!(waiters.take_wakers().count(), 1); + assert_eq!(waiters.drain().count(), 1); register(&mut waiters, &mut second_token, &second_waker); register(&mut waiters, &mut first_token, &first_waker); - let registered = waiters.take_wakers().collect::>(); + let registered = waiters.drain().collect::>(); assert_eq!(registered.len(), 2); wake_all(registered.into_iter()); assert_eq!(first_task.0.load(Ordering::Relaxed), 1); @@ -294,7 +284,7 @@ mod tests { register(&mut waiters, &mut third, &tracked); let result = panic::catch_unwind(AssertUnwindSafe(|| { - wake_all(waiters.take_wakers()); + wake_all(waiters.drain()); })); assert!(result.is_err()); assert_eq!(tracker.0.load(Ordering::Relaxed), 1); @@ -309,7 +299,7 @@ mod tests { register(&mut waiters, &mut token, &waker); drop(waker); - let removed = waiters.unregister_waker(&mut token); + let removed = waiters.unregister(&mut token); assert_eq!(waiters.registered_len(), 0); assert!(!dropped.load(Ordering::Relaxed)); diff --git a/asyncband/src/latch/mod.rs b/asyncband/src/latch/mod.rs index 19cfa02..cfb4e1f 100644 --- a/asyncband/src/latch/mod.rs +++ b/asyncband/src/latch/mod.rs @@ -284,9 +284,7 @@ impl Future for LatchWait<'_> { impl Drop for LatchWait<'_> { fn drop(&mut self) { - if self.token.is_some() { - self.latch.state.unregister_waker(&mut self.token); - } + self.latch.state.unregister(&mut self.token); } } @@ -316,8 +314,6 @@ impl Future for OwnedLatchWait { impl Drop for OwnedLatchWait { fn drop(&mut self) { - if self.token.is_some() { - self.latch.state.unregister_waker(&mut self.token); - } + self.latch.state.unregister(&mut self.token); } } diff --git a/asyncband/src/once/once/mod.rs b/asyncband/src/once/once/mod.rs index 6267331..6db2bb3 100644 --- a/asyncband/src/once/once/mod.rs +++ b/asyncband/src/once/once/mod.rs @@ -238,8 +238,6 @@ impl Future for OnceWait<'_> { impl Drop for OnceWait<'_> { fn drop(&mut self) { - if self.token.is_some() { - self.once.done.unregister_waker(&mut self.token); - } + self.once.done.unregister(&mut self.token); } } diff --git a/asyncband/src/waitgroup/mod.rs b/asyncband/src/waitgroup/mod.rs index ee21d96..93567f3 100644 --- a/asyncband/src/waitgroup/mod.rs +++ b/asyncband/src/waitgroup/mod.rs @@ -182,8 +182,6 @@ impl Future for Wait { impl Drop for Wait { fn drop(&mut self) { - if self.token.is_some() { - self.state.unregister_waker(&mut self.token); - } + self.state.unregister(&mut self.token); } } diff --git a/asyncband/src/watch/mod.rs b/asyncband/src/watch/mod.rs index a541904..9ddae32 100644 --- a/asyncband/src/watch/mod.rs +++ b/asyncband/src/watch/mod.rs @@ -135,14 +135,16 @@ impl fmt::Debug for Sender { impl Drop for Sender { fn drop(&mut self) { + // Only the final sender detaches the parked receivers; their wake callbacks run unlocked. let wakers = { let mut state = self.shared.state.lock(); state.senders -= 1; - (state.senders == 0 && !state.waiters.is_empty()).then(|| state.waiters.take_wakers()) + if state.senders != 0 { + return; + } + state.waiters.drain() }; - if let Some(wakers) = wakers { - wake_all(wakers); - } + wake_all(wakers); } } @@ -166,12 +168,11 @@ impl Sender { .expect("watch channel version counter overflowed"); let replaced = mem::replace(&mut state.value, Arc::new(value)); state.version = version; - let wakers = (!state.waiters.is_empty()).then(|| state.waiters.take_wakers()); + let wakers = state.waiters.drain(); (wakers, replaced) }; - if let Some(wakers) = wakers { - wake_all(wakers); - } + // Waker callbacks and the replaced value's destructor may reenter this channel. + wake_all(wakers); drop(replaced); Ok(()) } @@ -263,7 +264,7 @@ impl Receiver { pub async fn changed(&mut self) -> Result, RecvError> { Changed { receiver: self, - registration: None, + token: None, } .await } @@ -279,7 +280,7 @@ impl Receiver { struct Changed<'a, T> { receiver: &'a mut Receiver, - registration: Option, + token: Option, } impl Future for Changed<'_, T> { @@ -287,30 +288,45 @@ impl Future for Changed<'_, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); - let (poll, retired_waker) = { - let mut state = this.receiver.shared.state.lock(); - if state.version != this.receiver.seen { - let retired = state.waiters.unregister_waker(&mut this.registration); - this.receiver.seen = state.version; - (Poll::Ready(Ok(state.value.clone())), retired) - } else if state.senders == 0 { - let retired = state.waiters.unregister_waker(&mut this.registration); - (Poll::Ready(Err(RecvError::Disconnected)), retired) - } else { - let retired = state.waiters.register_waker(&mut this.registration, cx); - (Poll::Pending, retired) - } - }; + // A changed call normally parks once and then consumes one update, so cloning first keeps + // the common pending path to one state-lock acquisition. Checking the version and sender + // count afterward closes races caused by concurrent work or a reentrant clone callback. + // Ready and disconnected polls may consequently clone an unused waker. + let waker = cx.waker().clone(); + let mut state = this.receiver.shared.state.lock(); + if state.version != this.receiver.seen { + let retired_waker = state.waiters.unregister(&mut this.token); + this.receiver.seen = state.version; + let value = state.value.clone(); + drop(state); + drop(retired_waker); + drop(waker); + return Poll::Ready(Ok(value)); + } + if state.senders == 0 { + let retired_waker = state.waiters.unregister(&mut this.token); + drop(state); + drop(retired_waker); + drop(waker); + return Poll::Ready(Err(RecvError::Disconnected)); + } + + let retired_waker = state.waiters.register(&mut this.token, waker); + drop(state); drop(retired_waker); - poll + Poll::Pending } } impl Drop for Changed<'_, T> { fn drop(&mut self) { + if self.token.is_none() { + return; + } + let waker = { let mut state = self.receiver.shared.state.lock(); - state.waiters.unregister_waker(&mut self.registration) + state.waiters.unregister(&mut self.token) }; drop(waker); } diff --git a/tests-integration/tests/completion_test.rs b/tests-integration/tests/completion_test.rs index 977487e..59006d3 100644 --- a/tests-integration/tests/completion_test.rs +++ b/tests-integration/tests/completion_test.rs @@ -17,7 +17,6 @@ use std::cell::Cell; use std::future::Future; -use std::mem::ManuallyDrop; use std::pin::Pin; use std::sync::Arc; use std::sync::Barrier; @@ -26,8 +25,6 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; -use std::task::RawWaker; -use std::task::RawWakerVTable; use std::task::Wake; use std::task::Waker; use std::thread; @@ -66,8 +63,6 @@ impl Wake for WakeCallback { struct DropCallbackWake(Mutex>>); -struct CloneCallbackWake(Mutex>>); - // This test needs a custom waker whose final `Arc` drop is observable. #[allow(clippy::manual_noop_waker)] impl Wake for DropCallbackWake { @@ -82,45 +77,6 @@ impl Drop for DropCallbackWake { } } -unsafe fn clone_callback_waker(data: *const ()) -> RawWaker { - // SAFETY: Every raw pointer using this vtable comes from `Arc::into_raw` below. ManuallyDrop - // keeps the original waker's strong reference alive while the clone callback borrows it. - let state = ManuallyDrop::new(unsafe { Arc::::from_raw(data.cast()) }); - if let Some(callback) = state.0.lock().unwrap().take() { - callback(); - } - RawWaker::new( - Arc::into_raw(Arc::clone(&state)).cast(), - &CLONE_CALLBACK_VTABLE, - ) -} - -unsafe fn wake_clone_callback_waker(data: *const ()) { - // SAFETY: `wake` consumes the raw waker's strong reference exactly once. - drop(unsafe { Arc::::from_raw(data.cast()) }); -} - -unsafe fn wake_clone_callback_waker_by_ref(_data: *const ()) {} - -unsafe fn drop_clone_callback_waker(data: *const ()) { - // SAFETY: `drop` consumes the raw waker's strong reference exactly once. - drop(unsafe { Arc::::from_raw(data.cast()) }); -} - -static CLONE_CALLBACK_VTABLE: RawWakerVTable = RawWakerVTable::new( - clone_callback_waker, - wake_clone_callback_waker, - wake_clone_callback_waker_by_ref, - drop_clone_callback_waker, -); - -fn waker_with_clone_callback(callback: impl FnOnce() + Send + 'static) -> Waker { - let state = Arc::new(CloneCallbackWake(Mutex::new(Some(Box::new(callback))))); - let raw = RawWaker::new(Arc::into_raw(state).cast(), &CLONE_CALLBACK_VTABLE); - // SAFETY: The vtable preserves the Arc strong count and all callbacks are thread safe. - unsafe { Waker::from_raw(raw) } -} - fn poll_with(future: Pin<&mut F>, waker: &Waker) -> Poll { future.poll(&mut Context::from_waker(waker)) } @@ -410,23 +366,6 @@ fn wake_callbacks_run_outside_the_completion_lock() { ); } -#[test] -fn waker_clone_callbacks_run_outside_the_completion_lock() { - assert_completes_without_deadlock( - "waker clone callback deadlocked against the completion lock", - || { - let (completer, completion) = completion::new::(); - let waker = waker_with_clone_callback(move || drop(completer)); - let mut wait = Box::pin(completion.wait()); - - assert_eq!( - poll_with(wait.as_mut(), &waker), - Poll::Ready(Err(completion::Abandoned)) - ); - }, - ); -} - #[test] fn replaced_wakers_are_dropped_outside_the_completion_lock() { assert_completes_without_deadlock( diff --git a/tests-integration/tests/waitset_reentrancy_test.rs b/tests-integration/tests/waitset_reentrancy_test.rs new file mode 100644 index 0000000..2df903a --- /dev/null +++ b/tests-integration/tests/waitset_reentrancy_test.rs @@ -0,0 +1,185 @@ +// 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::future::Future; +use std::mem::ManuallyDrop; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Mutex; +use std::task::Context; +use std::task::Poll; +use std::task::RawWaker; +use std::task::RawWakerVTable; +use std::task::Waker; +use std::thread; +use std::time::Duration; + +use asyncband::barrier::Barrier; +use asyncband::broadcast::mpmc; +use asyncband::completion; +use asyncband::latch::Latch; +use asyncband::watch; + +struct CloneCallback(Mutex>>); + +unsafe fn clone_callback_waker(data: *const ()) -> RawWaker { + // SAFETY: Every raw pointer using this vtable comes from `Arc::into_raw` below. `ManuallyDrop` + // keeps the original waker's strong reference alive while the clone callback borrows it. + let callback = ManuallyDrop::new(unsafe { Arc::::from_raw(data.cast()) }); + if let Some(callback) = callback.0.lock().unwrap().take() { + callback(); + } + RawWaker::new( + Arc::into_raw(Arc::clone(&callback)).cast(), + &CLONE_CALLBACK_VTABLE, + ) +} + +unsafe fn wake_clone_callback_waker(data: *const ()) { + // SAFETY: `wake` consumes the raw waker's strong reference exactly once. + drop(unsafe { Arc::::from_raw(data.cast()) }); +} + +unsafe fn wake_clone_callback_waker_by_ref(_data: *const ()) {} + +unsafe fn drop_clone_callback_waker(data: *const ()) { + // SAFETY: `drop` consumes the raw waker's strong reference exactly once. + drop(unsafe { Arc::::from_raw(data.cast()) }); +} + +static CLONE_CALLBACK_VTABLE: RawWakerVTable = RawWakerVTable::new( + clone_callback_waker, + wake_clone_callback_waker, + wake_clone_callback_waker_by_ref, + drop_clone_callback_waker, +); + +fn waker_with_clone_callback(callback: impl FnOnce() + Send + 'static) -> Waker { + let callback = Arc::new(CloneCallback(Mutex::new(Some(Box::new(callback))))); + let raw = RawWaker::new(Arc::into_raw(callback).cast(), &CLONE_CALLBACK_VTABLE); + // SAFETY: The vtable preserves the Arc strong count and all callbacks are thread safe. + unsafe { Waker::from_raw(raw) } +} + +fn poll_with(future: Pin<&mut F>, waker: &Waker) -> Poll { + future.poll(&mut Context::from_waker(waker)) +} + +fn assert_completes_without_deadlock(message: &'static str, test: impl FnOnce() + Send + 'static) { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(test)); + finished_tx.send(()).unwrap(); + if let Err(payload) = result { + std::panic::resume_unwind(payload); + } + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect(message); + worker.join().unwrap(); +} + +#[test] +fn completion_clones_wakers_outside_its_state_lock() { + assert_completes_without_deadlock( + "waker clone callback deadlocked against the completion lock", + || { + let (completer, completion) = completion::new::(); + let waker = waker_with_clone_callback(move || drop(completer)); + let mut wait = Box::pin(completion.wait()); + + assert_eq!( + poll_with(wait.as_mut(), &waker), + Poll::Ready(Err(completion::Abandoned)) + ); + }, + ); +} + +#[test] +fn latch_clones_wakers_outside_its_waiter_lock() { + assert_completes_without_deadlock( + "waker clone callback deadlocked against the latch lock", + || { + let latch = Arc::new(Latch::new(1)); + let callback_latch = latch.clone(); + let waker = waker_with_clone_callback(move || callback_latch.count_down()); + let mut wait = Box::pin(latch.wait()); + + assert_eq!(poll_with(wait.as_mut(), &waker), Poll::Ready(())); + }, + ); +} + +#[test] +fn barrier_clones_wakers_outside_its_state_lock() { + assert_completes_without_deadlock( + "waker clone callback deadlocked against the barrier lock", + || { + let barrier = Arc::new(Barrier::new(2)); + let callback_barrier = barrier.clone(); + let waker = waker_with_clone_callback(move || { + let mut wait = Box::pin(callback_barrier.wait()); + let Poll::Ready(result) = poll_with(wait.as_mut(), Waker::noop()) else { + panic!("second barrier participant must complete the generation"); + }; + assert!(result.is_leader()); + }); + let mut wait = Box::pin(barrier.wait()); + + let Poll::Ready(result) = poll_with(wait.as_mut(), &waker) else { + panic!("first barrier participant must observe the completed generation"); + }; + assert!(!result.is_leader()); + }, + ); +} + +#[test] +fn watch_clones_wakers_outside_its_state_lock() { + assert_completes_without_deadlock( + "waker clone callback deadlocked against the watch lock", + || { + let (sender, mut receiver) = watch::channel(0); + let callback_sender = sender.clone(); + let waker = waker_with_clone_callback(move || callback_sender.send(1).unwrap()); + let mut changed = Box::pin(receiver.changed()); + + assert_eq!( + poll_with(changed.as_mut(), &waker), + Poll::Ready(Ok(Arc::new(1))) + ); + }, + ); +} + +#[test] +fn broadcast_clones_wakers_outside_its_state_lock() { + assert_completes_without_deadlock( + "waker clone callback deadlocked against the broadcast lock", + || { + let (sender, mut receiver) = mpmc::unbounded(); + let callback_sender = sender.clone(); + let waker = waker_with_clone_callback(move || callback_sender.send(1)); + let mut recv = Box::pin(receiver.recv()); + + assert_eq!(poll_with(recv.as_mut(), &waker), Poll::Ready(Ok(1))); + }, + ); +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs index f84e1b6..5e0ba84 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -117,6 +117,10 @@ impl CommandMiri { "tests-integration", &["--test", "unsafe_paths_test"], )); + run_command(make_miri_cmd( + "tests-integration", + &["--test", "waitset_reentrancy_test"], + )); } }