From 45243c7cb886ad6bafa1537a68e758625ad788bf Mon Sep 17 00:00:00 2001 From: Orthur Date: Sat, 7 Nov 2026 22:50:14 -0500 Subject: [PATCH 1/4] fix(barrier): use a wrapping 64-bit generation counter A `Barrier` waiter kept the generation it joined. `Drop` treated the waiter as completed once the current generation differed from that number, but `poll` only once the number was less than the current one. The counter was a `usize` incremented without wrapping, so completing generation `usize::MAX` overflowed it. A debug build panicked on that increment before waking anyone and left the counter in place, so the final arrival of every later generation panicked the same way and no waiter was released again. A release build wrapped to zero. The waiters released by the wrapping generation still held `usize::MAX`, which is never less than the current number, so polling one passed its stale token to the drained waker set: it panicked if that slot was vacant, or replaced the waker of the next generation's follower in the slot and stayed pending, so that follower was never woken. On 32-bit targets this takes 2^32 generations. Complete a waiter when the generation it joined is no longer current, in one check that `poll` and `Drop` share, the same shape `Phaser` uses for its phase number. Comparing by identity makes the wrap itself harmless, but a woken waiter left unpolled for a full cycle of the counter would match the current generation again, and that cycle is only 2^32 generations for a 32-bit `usize`. Use a 64-bit counter with a wrapping increment so the cycle is out of reach on every target. Add unit tests that start the counter at `u64::MAX`, poll or drop the woken follower after the wrap, and check that the next generation's follower is still woken. The existing drop test now asserts that wake too; it passed even with the generation check in `Drop` removed. --- CHANGELOG.md | 1 + asyncband/src/barrier/mod.rs | 25 ++++++-- asyncband/src/barrier/tests.rs | 79 +++++++++++++++++++++++++ tests-integration/tests/barrier_test.rs | 1 + 4 files changed, 101 insertions(+), 5 deletions(-) create mode 100644 asyncband/src/barrier/tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 17785dc3..ab476b82 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,7 @@ 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 diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index 2a5d8bf0..98147cc8 100644 --- a/asyncband/src/barrier/mod.rs +++ b/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,10 +71,22 @@ pub struct Barrier { struct BarrierState { arrived: u32, - generation: usize, + /// Identifies the generation that is still collecting arrivals. + generation: u64, waiters: WakerSet, } +impl BarrierState { + /// Returns whether the generation a waiter joined has completed. + /// + /// The counter wraps, so this compares identity rather than order. A woken waiter would take + /// its generation for the current one again only if a full `u64` cycle of generations completed + /// before it was polled or dropped. + fn is_completed(&self, generation: u64) -> bool { + self.generation != generation + } +} + impl fmt::Debug for BarrierState { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BarrierState") @@ -186,7 +201,7 @@ impl Barrier { // the state lock, then wake the detached followers after releasing it. if state.arrived == self.n { state.arrived = 0; - state.generation += 1; + state.generation = state.generation.wrapping_add(1); let mut wakers = state.waiters.take_all(); drop(state); wakers.wake_all(); @@ -212,7 +227,7 @@ impl Barrier { #[must_use = "futures do nothing unless you `.await` or poll them"] struct BarrierWait<'a> { token: Option, - generation: usize, + generation: u64, barrier: &'a Barrier, } @@ -235,7 +250,7 @@ impl Future for BarrierWait<'_> { } = self.get_mut(); let mut state = barrier.state.lock(); - if *generation < state.generation { + if state.is_completed(*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; @@ -256,7 +271,7 @@ impl Drop for BarrierWait<'_> { } let mut state = self.barrier.state.lock(); - if self.generation != state.generation { + if state.is_completed(self.generation) { // Advancing the generation detached this registration before making the new generation // visible under the same lock. self.token = None; diff --git a/asyncband/src/barrier/tests.rs b/asyncband/src/barrier/tests.rs new file mode 100644 index 00000000..7bf9564c --- /dev/null +++ b/asyncband/src/barrier/tests.rs @@ -0,0 +1,79 @@ +// 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.0.fetch_add(1, Ordering::Relaxed); + } +} + +fn poll_wait(wait: Pin<&mut impl Future>, waker: &Waker) -> Poll { + wait.poll(&mut Context::from_waker(waker)) + .map(|result| result.is_leader()) +} + +#[test] +fn a_woken_follower_completes_across_generation_wraparound() { + 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 = 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); + + assert_eq!(poll_wait(old.as_mut(), noop), Poll::Ready(false)); + assert_eq!(poll_wait(pin!(barrier.wait()), noop), Poll::Ready(true)); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); +} + +#[test] +fn dropping_a_woken_follower_across_wraparound_keeps_the_next_generation_registered() { + 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)); + + let mut next = pin!(barrier.wait()); + assert_eq!(poll_wait(next.as_mut(), &waker), Poll::Pending); + drop(old); + + assert_eq!(poll_wait(pin!(barrier.wait()), noop), Poll::Ready(true)); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); +} diff --git a/tests-integration/tests/barrier_test.rs b/tests-integration/tests/barrier_test.rs index 6c660fe5..2b9faafe 100644 --- a/tests-integration/tests/barrier_test.rs +++ b/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()); } From 73d0654b306f41b87d6d3d022f9c5832e0bfa2e3 Mon Sep 17 00:00:00 2001 From: Orthur Date: Sat, 7 Nov 2026 22:50:14 -0500 Subject: [PATCH 2/4] perf(barrier): record the arrival in the wait future's first poll `Barrier::wait` took the state lock twice on a follower's first poll: once in the async fn body to count the arrival and once in `BarrierWait::poll` to register the waker. Move the arrival into the future's first poll so the count, the completion check and the waker registration happen under one acquisition, and let the future return the leader result itself. The leader can no longer complete a generation between a follower's arrival and its registration. An unset token marks a future that has not arrived yet, as the index does in the semaphore's `Acquire`. This relies on the `async fn` wrapper, which panics if polled after completion instead of letting the inner future record a second arrival. The wait future stays at 40 bytes on 64-bit targets and 24 bytes on i686, the same as with the separate arrival. --- CHANGELOG.md | 1 + asyncband/src/barrier/mod.rs | 72 ++++++++++++++++-------------------- 2 files changed, 32 insertions(+), 41 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ab476b82..ca7736de 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,7 @@ All notable changes to this project will be documented in this file. * 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) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index 98147cc8..c7bf7d82 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -192,43 +192,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 = state.generation.wrapping_add(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, - generation: u64, 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, } impl fmt::Debug for BarrierWait<'_> { @@ -240,24 +221,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 { - 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 state.is_completed(*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.is_completed(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 From 6f389c4d015e143e77ac69345078048e51351f64 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 8 Nov 2026 13:40:44 +0800 Subject: [PATCH 3/4] refactor(barrier): inline generation checks and consolidate wraparound tests --- asyncband/src/barrier/mod.rs | 20 ++++--------- asyncband/src/barrier/tests.rs | 55 +++++++++++++--------------------- 2 files changed, 27 insertions(+), 48 deletions(-) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index c7bf7d82..75feceb8 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -71,22 +71,13 @@ pub struct Barrier { struct BarrierState { arrived: u32, - /// Identifies the generation that is still collecting arrivals. + // 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, } -impl BarrierState { - /// Returns whether the generation a waiter joined has completed. - /// - /// The counter wraps, so this compares identity rather than order. A woken waiter would take - /// its generation for the current one again only if a full `u64` cycle of generations completed - /// before it was polled or dropped. - fn is_completed(&self, generation: u64) -> bool { - self.generation != generation - } -} - impl fmt::Debug for BarrierState { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BarrierState") @@ -163,6 +154,7 @@ impl Barrier { /// /// The barrier holds the current task until all `n` participants have arrived. The final /// participant is designated as the leader for this generation. + /// The next generation can begin before all released waiters have resumed execution. /// /// # Cancel safety /// @@ -240,7 +232,7 @@ impl Future for BarrierWait<'_> { return Poll::Ready(BarrierWaitResult(true)); } this.generation = state.generation; - } else if state.is_completed(this.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. this.token = None; @@ -261,7 +253,7 @@ impl Drop for BarrierWait<'_> { } let mut state = self.barrier.state.lock(); - if state.is_completed(self.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; diff --git a/asyncband/src/barrier/tests.rs b/asyncband/src/barrier/tests.rs index 7bf9564c..70cf0067 100644 --- a/asyncband/src/barrier/tests.rs +++ b/asyncband/src/barrier/tests.rs @@ -38,42 +38,29 @@ fn poll_wait(wait: Pin<&mut impl Future>, waker: &Wa } #[test] -fn a_woken_follower_completes_across_generation_wraparound() { - 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()); +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 = pin!(barrier.wait()); - assert_eq!(poll_wait(old.as_mut(), noop), Poll::Pending); - assert_eq!(poll_wait(pin!(barrier.wait()), noop), Poll::Ready(true)); + 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); + // 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); - assert_eq!(poll_wait(old.as_mut(), noop), Poll::Ready(false)); - assert_eq!(poll_wait(pin!(barrier.wait()), noop), Poll::Ready(true)); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); -} - -#[test] -fn dropping_a_woken_follower_across_wraparound_keeps_the_next_generation_registered() { - 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()); + if poll_before_drop { + assert_eq!(poll_wait(old.as_mut(), noop), Poll::Ready(false)); + } + drop(old); - 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)); - - let mut next = pin!(barrier.wait()); - assert_eq!(poll_wait(next.as_mut(), &waker), Poll::Pending); - 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(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)); + } } From 0166c5605d57960c352e2e6f8a528bd2f7757880 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 8 Nov 2026 13:42:34 +0800 Subject: [PATCH 4/4] docs(barrier): remove internal scheduling detail from wait docs --- asyncband/src/barrier/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index 75feceb8..c82e856b 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -154,7 +154,6 @@ impl Barrier { /// /// The barrier holds the current task until all `n` participants have arrived. The final /// participant is designated as the leader for this generation. - /// The next generation can begin before all released waiters have resumed execution. /// /// # Cancel safety ///