fix: clone wakers outside state locks

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
 

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);
         }

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<T> {
     /// 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<Arc<T>>,
     /// The version of the first message in `buffer`.
     head: u64,
@@ -396,7 +396,10 @@ impl<T> Drop for UnboundedSender<T> {
         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<T> UnboundedSender<T> {
                 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<T: Clone> UnboundedReceiver<T> {
     pub async fn recv(&mut self) -> Result<T, RecvError> {
         Recv {
             receiver: self,
-            registration: None,
+            token: None,
         }
         .await
     }
@@ -701,19 +704,19 @@ impl<T> UnboundedReceiver<T> {
 
 struct Recv<'a, T> {
     receiver: &'a mut UnboundedReceiver<T>,
-    registration: Option<WakerToken>,
+    token: Option<WakerToken>,
 }
 
 impl<T> 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<T: Clone> Future for Recv<'_, T> {
     type Output = Result<T, RecvError>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        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)))
     }
 }

asyncband/src/completion/mod.rs

@@ -158,15 +158,14 @@ impl<T> Completer<T> {
                 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<T> Drop for Completer<T> {
             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,32 +223,34 @@ impl<T> Completion<T> {
     pub async fn wait(&self) -> Result<&T, Abandoned> {
         Wait {
             completion: self,
-            registration: None,
+            token: None,
         }
         .await
     }
 }
 
 struct Wait<'a, T> {
     completion: &'a Completion<T>,
-    registration: Option<WakerToken>,
+    token: Option<WakerToken>,
 }
 
 impl<'a, T> Future for Wait<'a, T> {
     type Output = Result<&'a T, Abandoned>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         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) = {
                 let mut state = this.completion.shared.state.lock();
                 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<T> = 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<T> 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);
     }

asyncband/src/internal/arena.rs

@@ -183,15 +183,22 @@ impl<T> Arena<T> {
 
     /// 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<Item = T> + use<T> {
         let len = self.len;
         let mut values = ArenaValues {
             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() {