fix(barrier): handle generation counter wraparound on 32-bit targets

CHANGELOG.md

@@ -38,12 +38,14 @@ All notable changes to this project will be documented in this file.
 
 * Prevent `OnceMap` from keeping an empty entry when a computation fails, panics, or is cancelled while a concurrent `get` looks up the same key.
 * Prevent `OnceMap` and `singleflight::Group` from leaking entries or values when a duplicate key's destructor panics, including when the original computation is cancelled during that destructor.
+* Fix `Barrier` on 32-bit targets after 2^32 completed generations, where debug builds panicked and stopped releasing waiters, and release builds could panic or make a waiter miss its wakeup.
 
 ### Improvements
 
 * Reduce the cost of initializing an uncontended `OnceCell` by avoiding its internal wait-queue lock.
 * Reduce the cost of polling a pending oneshot receive again from the same task, which now keeps its registered waker instead of cloning and dropping it.
 * Reduce the state kept by futures waiting for a `Mutex`, `RwLock`, `Semaphore`, `Once`, `OnceCell`, or bounded broadcast capacity, making lock, permit acquisition, and bounded broadcast send futures smaller.
+* Reduce the cost of `Barrier::wait` for participants that are not the last to arrive, which now record their arrival and register for wakeup under one lock acquisition instead of two.
 
 ## v0.7.4 (2026-10-25)
 

asyncband/src/barrier/mod.rs

@@ -57,6 +57,9 @@ use crate::internal::mutex::Mutex;
 use crate::internal::wakerset::WakerSet;
 use crate::internal::wakerset::WakerToken;
 
+#[cfg(test)]
+mod tests;
+
 /// A synchronization primitive for multiple tasks that need to wait for each other.
 ///
 /// See the [module level documentation](self) for more.
@@ -68,7 +71,10 @@ pub struct Barrier {
 
 struct BarrierState {
     arrived: u32,
-    generation: usize,
+    // Identifies the generation collecting arrivals. Compare identity, not order: the counter
+    // wraps, and woken waiters may remain unpolled while later generations complete. Using u64
+    // keeps an old waiter from matching again unless it survives a full 2^64-generation cycle.
+    generation: u64,
     waiters: WakerSet,
 }
 
@@ -177,43 +183,24 @@ impl Barrier {
     /// # }
     /// ```
     pub async fn wait(&self) -> BarrierWaitResult {
-        let generation = {
-            let mut state = self.state.lock();
-            let generation = state.generation;
-            state.arrived += 1;
-
-            // The final arrival completes this generation. Advance the generation while holding
-            // the state lock, then wake the detached followers after releasing it.
-            if state.arrived == self.n {
-                state.arrived = 0;
-                state.generation += 1;
-                let mut wakers = state.waiters.take_all();
-                drop(state);
-                wakers.wake_all();
-                return BarrierWaitResult(true);
-            }
-
-            generation
-        };
-
-        let fut = BarrierWait {
-            token: None,
-            generation,
+        BarrierWait {
             barrier: self,
-        };
-        fut.await;
-        BarrierWaitResult(false)
+            generation: 0,
+            token: None,
+        }
+        .await
     }
 }
 
-/// A future returned by [`Barrier::wait()`].
-///
-/// This future will complete when all tasks have reached the barrier point.
+/// The future awaited by [`Barrier::wait()`]. Its first poll records the arrival.
 #[must_use = "futures do nothing unless you `.await` or poll them"]
 struct BarrierWait<'a> {
-    token: Option<WakerToken>,
-    generation: usize,
     barrier: &'a Barrier,
+    /// The generation joined by the first poll, meaningful while `token` is set.
+    generation: u64,
+    /// The follower's registration. `None` also means not yet arrived, so this must not be polled
+    /// after `Ready`, which the `async fn` in [`Barrier::wait`] guarantees.
+    token: Option<WakerToken>,
 }
 
 impl fmt::Debug for BarrierWait<'_> {
@@ -225,24 +212,33 @@ impl fmt::Debug for BarrierWait<'_> {
 }
 
 impl Future for BarrierWait<'_> {
-    type Output = ();
+    type Output = BarrierWaitResult;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let Self {
-            token,
-            generation,
-            barrier,
-        } = self.get_mut();
+        let this = self.get_mut();
+        let mut state = this.barrier.state.lock();
 
-        let mut state = barrier.state.lock();
-        if *generation < state.generation {
+        if this.token.is_none() {
+            state.arrived += 1;
+            // The final arrival completes this generation. Advance the generation while holding
+            // the state lock, then wake the detached followers after releasing it.
+            if state.arrived == this.barrier.n {
+                state.arrived = 0;
+                state.generation = state.generation.wrapping_add(1);
+                let mut wakers = state.waiters.take_all();
+                drop(state);
+                wakers.wake_all();
+                return Poll::Ready(BarrierWaitResult(true));
+            }
+            this.generation = state.generation;
+        } else if state.generation != this.generation {
             // Completion advances the generation and detaches its old waiters under this same lock,
             // so no registration represented by this token remains in the waker set.
-            *token = None;
-            return Poll::Ready(());
+            this.token = None;
+            return Poll::Ready(BarrierWaitResult(false));
         }
 
-        let retired_waker = state.waiters.register(token, cx.waker());
+        let retired_waker = state.waiters.register(&mut this.token, cx.waker());
         drop(state);
         drop(retired_waker);
         Poll::Pending
@@ -256,7 +252,7 @@ impl Drop for BarrierWait<'_> {
         }
 
         let mut state = self.barrier.state.lock();
-        if self.generation != state.generation {
+        if state.generation != self.generation {
             // Advancing the generation detached this registration before making the new generation
             // visible under the same lock.
             self.token = None;

asyncband/src/barrier/tests.rs

@@ -0,0 +1,66 @@
+// 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::pin::pin;
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::task::Wake;
+use std::task::Waker;
+
+use super::*;
+
+struct WakeCounter(AtomicUsize);
+
+impl Wake for WakeCounter {
+    fn wake(self: Arc<Self>) {
+        self.0.fetch_add(1, Ordering::Relaxed);
+    }
+}
+
+fn poll_wait(wait: Pin<&mut impl Future<Output = BarrierWaitResult>>, waker: &Waker) -> Poll<bool> {
+    wait.poll(&mut Context::from_waker(waker))
+        .map(|result| result.is_leader())
+}
+
+#[test]
+fn a_woken_follower_across_wraparound_preserves_the_next_generation() {
+    for poll_before_drop in [false, true] {
+        let barrier = Barrier::new(2);
+        barrier.state.lock().generation = u64::MAX;
+        let noop = Waker::noop();
+        let counter = Arc::new(WakeCounter(AtomicUsize::new(0)));
+        let waker = Waker::from(counter.clone());
+
+        let mut old = Box::pin(barrier.wait());
+        assert_eq!(poll_wait(old.as_mut(), noop), Poll::Pending);
+        assert_eq!(poll_wait(pin!(barrier.wait()), noop), Poll::Ready(true));
+
+        // The next generation reuses the slot the old follower registered in.
+        let mut next = pin!(barrier.wait());
+        assert_eq!(poll_wait(next.as_mut(), &waker), Poll::Pending);
+
+        if poll_before_drop {
+            assert_eq!(poll_wait(old.as_mut(), noop), Poll::Ready(false));
+        }
+        drop(old);
+
+        assert_eq!(poll_wait(pin!(barrier.wait()), noop), Poll::Ready(true));
+        assert_eq!(counter.0.load(Ordering::Relaxed), 1);
+        assert_eq!(poll_wait(next.as_mut(), &waker), Poll::Ready(false));
+    }
+}

tests-integration/tests/barrier_test.rs

@@ -73,6 +73,7 @@ fn dropping_a_woken_wait_keeps_the_next_generation_registered() {
 
     let mut next_leader = spawn(b.wait());
     assert!(assert_ready!(next_leader.poll()).is_leader());
+    assert!(next.is_woken());
     assert!(!assert_ready!(next.poll()).is_leader());
 }