From 58bdde8e68ac7682b13bc3c41c446e0dc4ed3268 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 30 Aug 2026 20:02:01 +0800 Subject: [PATCH 1/4] fix: clone wakers outside state locks --- CHANGELOG.md | 1 + asyncband/src/barrier/mod.rs | 15 +- asyncband/src/broadcast/mpmc/unbounded/mod.rs | 30 ++- asyncband/src/completion/mod.rs | 4 +- asyncband/src/internal/countdown.rs | 15 +- asyncband/src/internal/waitset.rs | 32 +-- asyncband/src/watch/mod.rs | 18 +- tests-integration/tests/completion_test.rs | 61 ------ .../tests/waitset_reentrancy_test.rs | 185 ++++++++++++++++++ xtask/src/main.rs | 4 + 10 files changed, 252 insertions(+), 113 deletions(-) create mode 100644 tests-integration/tests/waitset_reentrancy_test.rs 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..8971a64 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -291,17 +291,22 @@ impl Future for BarrierWait<'_> { barrier, } = self.get_mut(); - let replaced_waker = { + // Waker cloning may call back into the barrier, so it must happen before taking the state + // lock. The generation is checked afterward to close the resulting race. + let waker = cx.waker().clone(); + let (poll, retired_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(()); + (Poll::Ready(()), Some(waker)) + } else { + let retired = state.waiters.register_waker(token, waker); + (Poll::Pending, retired) } - state.waiters.register_waker(token, cx) }; - drop(replaced_waker); - Poll::Pending + drop(retired_waker); + poll } } diff --git a/asyncband/src/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/broadcast/mpmc/unbounded/mod.rs index 9e17089..ce00154 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -728,27 +728,43 @@ impl Future for Recv<'_, T> { 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 mut prepared_waker = None; + // Senders append messages and drain the wait set under this same lock. If the receive is + // pending, clone its waker without the lock and retry the complete decision before + // registering it, so neither a wake-up nor a reentrant clone callback can be missed. + 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; + 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 + .registered_waker_will_wake(registration, 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_waker(registration, waker); drop(inner); - drop(waker); + drop(retired_waker); return Poll::Pending; } } }; + drop(prepared_waker); let (msg, reclaimed) = received; *registration = None; diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index b6d805c..ae237f7 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -259,9 +259,7 @@ 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_waker(&mut this.registration, waker); (Poll::Pending, retired) } Status::Completed => { diff --git a/asyncband/src/internal/countdown.rs b/asyncband/src/internal/countdown.rs index 524a792..b6a87e2 100644 --- a/asyncband/src/internal/countdown.rs +++ b/asyncband/src/internal/countdown.rs @@ -79,17 +79,22 @@ impl CountdownState { return Poll::Ready(()); } - let replaced_waker = { + // Cloning a waker can invoke arbitrary user code. Do it before taking the waiter lock, then + // recheck the countdown so a transition made by the clone callback cannot be missed. + let waker = cx.waker().clone(); + let (poll, retired_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(()); + (Poll::Ready(()), Some(waker)) + } else { + let retired = waiters.register_waker(token, waker); + (Poll::Pending, retired) } - waiters.register_waker(token, cx) }; - drop(replaced_waker); - Poll::Pending + drop(retired_waker); + poll } #[inline] diff --git a/asyncband/src/internal/waitset.rs b/asyncband/src/internal/waitset.rs index ded38a0..a1adf81 100644 --- a/asyncband/src/internal/waitset.rs +++ b/asyncband/src/internal/waitset.rs @@ -18,7 +18,6 @@ use std::mem; use std::panic; use std::panic::AssertUnwindSafe; -use std::task::Context; use std::task::Waker; use crate::internal::arena::Arena; @@ -98,30 +97,6 @@ impl WaitSet { self.waiters.take_all() } - /// Registers or updates a waker in the current wake epoch. - /// - /// If an existing waker is replaced, it is returned so the caller can drop it 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; - } - 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 - } - /// Returns whether `token` identifies a registered waker for the same task as `waker`. #[inline] pub fn registered_waker_will_wake(&self, token: &Option, waker: &Waker) -> bool { @@ -139,10 +114,13 @@ impl WaitSet { /// Registers or updates an already cloned waker in the current wake epoch. /// + /// The caller must clone the waker before acquiring the lock that protects this wait set, + /// because cloning can invoke arbitrary user code. + /// /// 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( + pub fn register_waker( &mut self, token: &mut Option, waker: Waker, @@ -253,7 +231,7 @@ mod tests { token: &mut Option, waker: &Waker, ) -> Option { - waiters.register_waker(token, &mut Context::from_waker(waker)) + waiters.register_waker(token, waker.clone()) } #[test] diff --git a/asyncband/src/watch/mod.rs b/asyncband/src/watch/mod.rs index a541904..ca1291d 100644 --- a/asyncband/src/watch/mod.rs +++ b/asyncband/src/watch/mod.rs @@ -287,21 +287,29 @@ impl Future for Changed<'_, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); - let (poll, retired_waker) = { + // A changed call normally parks once and then consumes one update. Preparing its waker + // before the lock avoids a second lock acquisition on that common pending path. + let waker = cx.waker().clone(); + let (poll, retired_waker, unused_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) + (Poll::Ready(Ok(state.value.clone())), retired, Some(waker)) } else if state.senders == 0 { let retired = state.waiters.unregister_waker(&mut this.registration); - (Poll::Ready(Err(RecvError::Disconnected)), retired) + ( + Poll::Ready(Err(RecvError::Disconnected)), + retired, + Some(waker), + ) } else { - let retired = state.waiters.register_waker(&mut this.registration, cx); - (Poll::Pending, retired) + let retired = state.waiters.register_waker(&mut this.registration, waker); + (Poll::Pending, retired, None) } }; drop(retired_waker); + drop(unused_waker); poll } } 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"], + )); } } From 19475433b1e9eedc17ec19965037280297944be8 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 30 Aug 2026 20:52:57 +0800 Subject: [PATCH 2/4] refactor: streamline wait set lifecycle --- asyncband/src/barrier/mod.rs | 30 ++++----- asyncband/src/broadcast/mpmc/unbounded/mod.rs | 34 +++++----- asyncband/src/completion/mod.rs | 32 +++++---- asyncband/src/internal/arena.rs | 4 ++ asyncband/src/internal/countdown.rs | 30 ++++----- asyncband/src/internal/waitset.rs | 46 +++++++------ asyncband/src/latch/mod.rs | 8 +-- asyncband/src/once/once/mod.rs | 4 +- asyncband/src/waitgroup/mod.rs | 4 +- asyncband/src/watch/mod.rs | 66 ++++++++++--------- 10 files changed, 128 insertions(+), 130 deletions(-) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index 8971a64..36bf119 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -180,7 +180,7 @@ impl Barrier { state: Mutex::new(BarrierState { arrived: 0, generation: 0, - waiters: WaitSet::with_capacity(n as usize), + waiters: WaitSet::with_capacity(n.saturating_sub(1) as usize), }), } } @@ -244,7 +244,7 @@ impl Barrier { 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); @@ -294,19 +294,19 @@ impl Future for BarrierWait<'_> { // Waker cloning may call back into the barrier, so it must happen before taking the state // lock. The generation is checked afterward to close the resulting race. let waker = cx.waker().clone(); - let (poll, retired_waker) = { - let mut state = barrier.state.lock(); - if *generation < state.generation { - // Advancing the generation drains its registrations under this same lock. - *token = None; - (Poll::Ready(()), Some(waker)) - } else { - let retired = state.waiters.register_waker(token, waker); - (Poll::Pending, retired) - } - }; + let mut state = barrier.state.lock(); + if *generation < state.generation { + // Advancing the generation drains its registrations under this same lock. + *token = None; + drop(state); + drop(waker); + return Poll::Ready(()); + } + + let retired_waker = state.waiters.register(token, waker); + drop(state); drop(retired_waker); - poll + Poll::Pending } } @@ -315,7 +315,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 ce00154..db7f721 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -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,10 +726,7 @@ 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(); + let Self { receiver, token } = self.get_mut(); let mut prepared_waker = None; // Senders append messages and drain the wait set under this same lock. If the receive is @@ -739,17 +739,13 @@ impl Future for Recv<'_, T> { 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)); } - if prepared_waker.is_none() - && inner - .waiters - .registered_waker_will_wake(registration, cx.waker()) - { + if prepared_waker.is_none() && inner.waiters.will_wake(token, cx.waker()) { return Poll::Pending; } let Some(waker) = prepared_waker.take() else { @@ -757,7 +753,7 @@ impl Future for Recv<'_, T> { prepared_waker = Some(cx.waker().clone()); continue; }; - let retired_waker = inner.waiters.register_waker(registration, waker); + let retired_waker = inner.waiters.register(token, waker); drop(inner); drop(retired_waker); return Poll::Pending; @@ -767,7 +763,7 @@ impl Future for Recv<'_, T> { 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 ae237f7..eb65f86 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -159,14 +159,12 @@ impl Completer { panic!("pending completion value must be unset"); } 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(()) } } @@ -182,11 +180,9 @@ impl Drop for Completer { return; } 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 +221,7 @@ impl Completion { pub async fn wait(&self) -> Result<&T, Abandoned> { Wait { completion: self, - registration: None, + token: None, } .await } @@ -233,7 +229,7 @@ impl Completion { struct Wait<'a, T> { completion: &'a Completion, - registration: Option, + token: Option, } impl<'a, T> Future for Wait<'a, T> { @@ -248,9 +244,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,11 +253,11 @@ impl<'a, T> Future for Wait<'a, T> { prepared_waker = Some(cx.waker().clone()); continue; }; - let retired = state.waiters.register_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 @@ -273,7 +267,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) } } @@ -287,9 +281,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..c2e61d7 100644 --- a/asyncband/src/internal/arena.rs +++ b/asyncband/src/internal/arena.rs @@ -192,6 +192,10 @@ impl Arena { first: None, rest: Vec::new(), }; + if len == 0 { + 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 b6a87e2..e7e937e 100644 --- a/asyncband/src/internal/countdown.rs +++ b/asyncband/src/internal/countdown.rs @@ -64,7 +64,7 @@ impl CountdownState { pub fn wake_all(&self) { let wakers = { let mut waiters = self.waiters.lock(); - waiters.take_wakers() + waiters.drain() }; wake_all(wakers); @@ -82,27 +82,27 @@ impl CountdownState { // Cloning a waker can invoke arbitrary user code. Do it before taking the waiter lock, then // recheck the countdown so a transition made by the clone callback cannot be missed. let waker = cx.waker().clone(); - let (poll, retired_waker) = { - let mut waiters = self.waiters.lock(); - if self.state() == 0 { - // A concurrent zero transition will drain after this lock is released. - *token = None; - (Poll::Ready(()), Some(waker)) - } else { - let retired = waiters.register_waker(token, waker); - (Poll::Pending, retired) - } - }; + 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 + 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); } diff --git a/asyncband/src/internal/waitset.rs b/asyncband/src/internal/waitset.rs index a1adf81..b259634 100644 --- a/asyncband/src/internal/waitset.rs +++ b/asyncband/src/internal/waitset.rs @@ -15,6 +15,13 @@ // 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. It only moves already owned +//! wakers while that lock is held: cloning must happen before taking the lock, and unused, removed, +//! or drained wakers must be dropped or woken after releasing it. This keeps arbitrary waker +//! callbacks outside primitive critical sections. + use std::mem; use std::panic; use std::panic::AssertUnwindSafe; @@ -84,22 +91,21 @@ 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. + /// Drains all registered wakers as an owning iterator without waking them. + /// + /// A non-empty drain starts a new epoch so tokens retained by the drained futures cannot alias + /// slots reused by later registrations. #[inline] - pub fn take_wakers(&mut self) -> impl Iterator + 'static { - self.epoch = self.epoch.checked_add(1).expect("wait set epoch overflow"); + 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"); + } self.waiters.take_all() } /// Returns whether `token` identifies a registered waker for the same task as `waker`. #[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; }; @@ -120,11 +126,8 @@ impl WaitSet { /// 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_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)); @@ -143,7 +146,8 @@ 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)); @@ -231,7 +235,7 @@ mod tests { token: &mut Option, waker: &Waker, ) -> Option { - waiters.register_waker(token, waker.clone()) + waiters.register(token, waker.clone()) } #[test] @@ -245,12 +249,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); @@ -272,7 +276,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); @@ -287,7 +291,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 ca1291d..030aca4 100644 --- a/asyncband/src/watch/mod.rs +++ b/asyncband/src/watch/mod.rs @@ -138,11 +138,12 @@ impl Drop for Sender { 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 +167,10 @@ 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); - } + wake_all(wakers); drop(replaced); Ok(()) } @@ -263,7 +262,7 @@ impl Receiver { pub async fn changed(&mut self) -> Result, RecvError> { Changed { receiver: self, - registration: None, + token: None, } .await } @@ -279,7 +278,7 @@ impl Receiver { struct Changed<'a, T> { receiver: &'a mut Receiver, - registration: Option, + token: Option, } impl Future for Changed<'_, T> { @@ -290,35 +289,40 @@ impl Future for Changed<'_, T> { // A changed call normally parks once and then consumes one update. Preparing its waker // before the lock avoids a second lock acquisition on that common pending path. let waker = cx.waker().clone(); - let (poll, retired_waker, unused_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, Some(waker)) - } else if state.senders == 0 { - let retired = state.waiters.unregister_waker(&mut this.registration); - ( - Poll::Ready(Err(RecvError::Disconnected)), - retired, - Some(waker), - ) - } else { - let retired = state.waiters.register_waker(&mut this.registration, waker); - (Poll::Pending, retired, None) - } - }; + 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); - drop(unused_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); } From a7843b7ea12b42499f9793bae5fc6a714e70a7d4 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 30 Aug 2026 21:31:31 +0800 Subject: [PATCH 3/4] refactor: simplify barrier waiter capacity --- asyncband/src/barrier/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index 36bf119..75d5170 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -180,7 +180,7 @@ impl Barrier { state: Mutex::new(BarrierState { arrived: 0, generation: 0, - waiters: WaitSet::with_capacity(n.saturating_sub(1) as usize), + waiters: WaitSet::with_capacity((n - 1) as usize), }), } } From b0ee3f2b4e8b0153e50bbf087c9bf8ce1c9bda12 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 30 Aug 2026 21:43:05 +0800 Subject: [PATCH 4/4] docs: clarify waiter polling tradeoffs --- asyncband/src/barrier/mod.rs | 14 ++++++---- asyncband/src/broadcast/mpmc/unbounded/mod.rs | 22 ++++++++------- asyncband/src/completion/mod.rs | 6 ++++ asyncband/src/internal/arena.rs | 7 +++-- asyncband/src/internal/countdown.rs | 24 ++++++---------- asyncband/src/internal/waitset.rs | 28 ++++++++++++------- asyncband/src/watch/mod.rs | 8 ++++-- 7 files changed, 65 insertions(+), 44 deletions(-) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index 75d5170..69cef81 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -180,6 +180,7 @@ impl Barrier { state: Mutex::new(BarrierState { arrived: 0, generation: 0, + // The final participant completes the generation without parking. waiters: WaitSet::with_capacity((n - 1) as usize), }), } @@ -239,8 +240,8 @@ 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; @@ -291,12 +292,15 @@ impl Future for BarrierWait<'_> { barrier, } = self.get_mut(); - // Waker cloning may call back into the barrier, so it must happen before taking the state - // lock. The generation is checked afterward to close the resulting race. + // 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 { - // Advancing the generation drains its registrations under this same lock. + // 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); diff --git a/asyncband/src/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/broadcast/mpmc/unbounded/mod.rs index db7f721..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, @@ -728,10 +728,12 @@ impl Future for Recv<'_, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 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; - // Senders append messages and drain the wait set under this same lock. If the receive is - // pending, clone its waker without the lock and retry the complete decision before - // registering it, so neither a wake-up nor a reentrant clone callback can be missed. let received = loop { let mut inner = receiver.shared.inner.lock(); diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index eb65f86..81a8bff 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -158,6 +158,7 @@ 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.drain() }; @@ -179,6 +180,7 @@ 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.drain() }; @@ -237,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) = { diff --git a/asyncband/src/internal/arena.rs b/asyncband/src/internal/arena.rs index c2e61d7..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; @@ -193,6 +194,8 @@ impl Arena { 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(); } diff --git a/asyncband/src/internal/countdown.rs b/asyncband/src/internal/countdown.rs index e7e937e..10ec6d6 100644 --- a/asyncband/src/internal/countdown.rs +++ b/asyncband/src/internal/countdown.rs @@ -39,28 +39,21 @@ 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(); @@ -79,8 +72,10 @@ impl CountdownState { return Poll::Ready(()); } - // Cloning a waker can invoke arbitrary user code. Do it before taking the waiter lock, then - // recheck the countdown so a transition made by the clone callback cannot be missed. + // 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 { @@ -145,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 b259634..1e7b983 100644 --- a/asyncband/src/internal/waitset.rs +++ b/asyncband/src/internal/waitset.rs @@ -17,10 +17,11 @@ //! Cancellable storage for task wakers. //! -//! A `WaitSet` is protected by the state lock of its owning primitive. It only moves already owned -//! wakers while that lock is held: cloning must happen before taking the lock, and unused, removed, -//! or drained wakers must be dropped or woken after releasing it. This keeps arbitrary waker -//! callbacks outside primitive critical sections. +//! 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; @@ -58,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, @@ -94,7 +96,8 @@ impl WaitSet { /// Drains all registered wakers as an owning iterator without waking them. /// /// A non-empty drain starts a new epoch so tokens retained by the drained futures cannot alias - /// slots reused by later registrations. + /// 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 drain(&mut self) -> impl Iterator + 'static { if !self.waiters.is_empty() { @@ -104,6 +107,9 @@ impl WaitSet { } /// 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 will_wake(&self, token: &Option, waker: &Waker) -> bool { let Some(current) = token.as_ref() else { @@ -120,8 +126,9 @@ impl WaitSet { /// Registers or updates an already cloned waker in the current wake epoch. /// - /// The caller must clone the waker before acquiring the lock that protects this wait set, - /// because cloning can invoke arbitrary user code. + /// 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. @@ -152,6 +159,7 @@ impl WaitSet { 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 } diff --git a/asyncband/src/watch/mod.rs b/asyncband/src/watch/mod.rs index 030aca4..9ddae32 100644 --- a/asyncband/src/watch/mod.rs +++ b/asyncband/src/watch/mod.rs @@ -135,6 +135,7 @@ 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; @@ -170,6 +171,7 @@ impl Sender { let wakers = state.waiters.drain(); (wakers, replaced) }; + // Waker callbacks and the replaced value's destructor may reenter this channel. wake_all(wakers); drop(replaced); Ok(()) @@ -286,8 +288,10 @@ impl Future for Changed<'_, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); - // A changed call normally parks once and then consumes one update. Preparing its waker - // before the lock avoids a second lock acquisition on that common pending path. + // 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 {