diff --git a/CHANGELOG.md b/CHANGELOG.md index 17785dc3..ca7736de 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index 2a5d8bf0..c82e856b 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,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, - 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, } 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 { - 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; diff --git a/asyncband/src/barrier/tests.rs b/asyncband/src/barrier/tests.rs new file mode 100644 index 00000000..70cf0067 --- /dev/null +++ b/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.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_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)); + } +} 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()); }