From f7b139a66e34f5c8092f06e402f8f677b650d807 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 3 Nov 2026 17:25:51 +0800 Subject: [PATCH 1/8] perf(waitgroup): streamline participant completion and waiter notification --- CHANGELOG.md | 4 + asyncband/src/waitgroup/mod.rs | 159 +++++++++++-------- asyncband/src/waitgroup/waiters.rs | 178 ++++++++++++++++++++++ tests-integration/tests/waitgroup_test.rs | 159 +++++++++++++++++++ xtask/src/main.rs | 4 + 5 files changed, 440 insertions(+), 64 deletions(-) create mode 100644 asyncband/src/waitgroup/waiters.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 14862c84..0ce364c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,10 @@ All notable changes to this project will be documented in this file. ## Unreleased +### Improvements + +* Reduce the cost of creating and dropping wait group participants and notifying a single registered waiter. + ## v0.8.0 ### Breaking changes diff --git a/asyncband/src/waitgroup/mod.rs b/asyncband/src/waitgroup/mod.rs index 4fb7c403..0cc1b0a3 100644 --- a/asyncband/src/waitgroup/mod.rs +++ b/asyncband/src/waitgroup/mod.rs @@ -26,6 +26,10 @@ //! remaining handle is awaited or dropped. Existing handles may register more participants by //! cloning until the group completes; completion is one-shot. //! +//! A coordinator can hand a clone to each worker, let workers drop their handles when done, +//! and await its own handle. To share completion with independent observers, convert a handle +//! into a [`Wait`] once and clone that future. Cloning the group itself adds more work. +//! //! Completion acquires the state published before every participant completed, so work performed //! by those participants is visible after the wait returns. //! @@ -45,7 +49,7 @@ //! let participant = group.clone(); //! tasks.push(tokio::spawn(async move { //! do_work().await; -//! participant.await; +//! drop(participant); //! })); //! } //! @@ -59,88 +63,112 @@ use std::fmt; use std::future::Future; use std::future::IntoFuture; -use std::mem; +use std::ops::Deref; use std::pin::Pin; -use std::sync::Arc; +use std::ptr::NonNull; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; -use crate::internal::mutex::Mutex; -use crate::internal::wakerset::WakerSet; -use crate::internal::wakerset::WakerToken; +use crate::waitgroup::waiters::Registration; +use crate::waitgroup::waiters::Waiters; + +mod waiters; -#[derive(Debug)] struct State { - // Wait futures also own the state allocation, so Arc's strong count cannot represent handles. - // Zero is terminal: after it is published, no new handle or waiter can be registered. + // Participants hold one collective observer count until completion has detached the wakers. + // The remaining observer counts belong to individual Wait futures. handles: AtomicUsize, - waiters: Mutex, + observers: AtomicUsize, + waiters: Waiters, } impl State { fn new() -> Self { Self { handles: AtomicUsize::new(1), - waiters: Mutex::new(WakerSet::new()), + observers: AtomicUsize::new(1), + waiters: Waiters::default(), } } + #[inline] fn register_handle(&self) { // The borrowed source handle keeps the count above zero. Registration publishes no data, // so it does not need to synchronize with completion. - self.handles.fetch_add(1, Ordering::Relaxed); + if self.handles.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { + std::process::abort(); + } } - fn release_handle(&self) { + #[inline] + fn release_handle(&self) -> bool { // Every decrement is an RMW in the release sequence. A waiter that acquires zero therefore // observes work published before every preceding handle release. - let previous = self.handles.fetch_sub(1, Ordering::Release); + let previous = self.handles.fetch_sub(1, Ordering::AcqRel); debug_assert!(previous > 0, "a live handle must own one count"); if previous != 1 { - return; + return false; } - let wakers = mem::take(&mut *self.waiters.lock()); - wakers.wake_all(); + self.waiters.wake_all(); + true } - fn poll_wait(&self, token: &mut Option, cx: &mut Context<'_>) -> Poll<()> { + #[inline] + fn poll_wait(&self, token: &mut Registration, cx: &mut Context<'_>) -> Poll<()> { if self.handles.load(Ordering::Acquire) == 0 { - *token = None; + *token = Registration::Set(None); return Poll::Ready(()); } - let mut waiters = self.waiters.lock(); - if self.handles.load(Ordering::Acquire) == 0 { - *token = None; - return Poll::Ready(()); + self.waiters.register(token, cx.waker()) + } + + #[inline] + fn unregister(&self, token: &mut Registration) { + if !matches!(token, Registration::Set(None)) && self.handles.load(Ordering::Acquire) != 0 { + self.waiters.unregister(token); } + } +} - let retired_waker = waiters.register(token, cx.waker()); - drop(waiters); - drop(retired_waker); - Poll::Pending +// One observer reference anchors the allocation until the last participant has notified waiters. +// Each Wait owns an independent observer reference. The participant count keeps the anchor alive +// across registration and notification. A Shared pointer is only copied while its owner keeps a +// count alive. +#[derive(Clone, Copy)] +struct Shared(NonNull); + +// SAFETY: all shared access goes through State, which is Send + Sync. +unsafe impl Send for Shared {} +unsafe impl Sync for Shared {} + +impl Deref for Shared { + type Target = State; + #[inline] + fn deref(&self) -> &State { + // SAFETY: each caller holds a participant or an observer count. + unsafe { self.0.as_ref() } } +} - fn unregister(&self, token: &mut Option) { - if token.is_none() { - return; +impl Shared { + #[inline] + fn observe(self) -> Self { + if self.observers.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { + std::process::abort(); } + self + } - let mut waiters = self.waiters.lock(); - // Reaching zero is terminal, so the zero transition either owns this waker or has already - // taken it. Unlike reusable waker sets, no epoch is needed to disambiguate a later - // registration. - if self.handles.load(Ordering::Acquire) == 0 { - *token = None; - return; + #[inline] + fn release_observer(self) { + if self.observers.fetch_sub(1, Ordering::AcqRel) == 1 { + // SAFETY: all observers and the collective participant anchor have been released. + unsafe { drop(Box::from_raw(self.0.as_ptr())) }; } - - let removed_waker = waiters.unregister(token); - drop(waiters); - drop(removed_waker); } } @@ -148,9 +176,8 @@ impl State { /// /// See the [module level documentation](self) for more. pub struct WaitGroup { - // Keeping this optional lets `into_future` transfer the allocation to `Wait` without an - // otherwise redundant Arc increment/decrement pair. The option retains Arc's pointer niche. - state: Option>, + // All participants share one reference to the allocation. + state: Shared, } impl fmt::Debug for WaitGroup { @@ -175,9 +202,10 @@ impl WaitGroup { /// /// let wg = WaitGroup::new(); /// ``` + #[inline] pub fn new() -> Self { Self { - state: Some(Arc::new(State::new())), + state: Shared(NonNull::from(Box::leak(Box::new(State::new())))), } } } @@ -187,23 +215,18 @@ impl Clone for WaitGroup { /// /// The group completes after every participant has completed, either by awaiting or dropping /// its handle. + #[inline] fn clone(&self) -> Self { - let state = self - .state - .as_ref() - .expect("a live WaitGroup owns its state") - .clone(); - // Every handle owns one strong reference, while Wait observers may own additional ones. - // Arc's own overflow guard therefore fires before this equally wide counter can wrap. - state.register_handle(); - Self { state: Some(state) } + self.state.register_handle(); + Self { state: self.state } } } impl Drop for WaitGroup { + #[inline] fn drop(&mut self) { - if let Some(state) = self.state.take() { - state.release_handle(); + if self.state.release_handle() { + self.state.release_observer(); } } } @@ -213,10 +236,14 @@ impl IntoFuture for WaitGroup { type IntoFuture = Wait; /// Marks this participant complete and waits for every other participant to complete. - fn into_future(mut self) -> Self::IntoFuture { - let state = self.state.take().expect("a live WaitGroup owns its state"); - state.release_handle(); - Wait { token: None, state } + #[inline] + fn into_future(self) -> Self::IntoFuture { + let state = self.state.observe(); + drop(self); + Wait { + token: Registration::Set(None), + state, + } } } @@ -228,18 +255,19 @@ impl IntoFuture for WaitGroup { /// affected. #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Wait { - token: Option, - state: Arc, + token: Registration, + state: Shared, } impl Clone for Wait { /// Creates a new future that observes the same group completion. /// /// This does not register another participant. + #[inline] fn clone(&self) -> Self { Wait { - token: None, - state: self.state.clone(), + token: Registration::Set(None), + state: self.state.observe(), } } } @@ -253,6 +281,7 @@ impl fmt::Debug for Wait { impl Future for Wait { type Output = (); + #[inline] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let Self { token, state } = self.get_mut(); state.poll_wait(token, cx) @@ -260,7 +289,9 @@ impl Future for Wait { } impl Drop for Wait { + #[inline] fn drop(&mut self) { self.state.unregister(&mut self.token); + self.state.release_observer(); } } diff --git a/asyncband/src/waitgroup/waiters.rs b/asyncband/src/waitgroup/waiters.rs new file mode 100644 index 00000000..96f1ede1 --- /dev/null +++ b/asyncband/src/waitgroup/waiters.rs @@ -0,0 +1,178 @@ +// 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::cell::UnsafeCell; +use std::mem; +use std::panic::RefUnwindSafe; +use std::sync::OnceLock; +use std::sync::atomic::AtomicU8; +use std::sync::atomic::Ordering; +use std::task::Poll; +use std::task::Waker; + +use crate::internal::mutex::Mutex; +use crate::internal::wakerset::WakerSet; +use crate::internal::wakerset::WakerToken; + +const EMPTY: u8 = 0; +const REGISTERED: u8 = 1; +const BUSY: u8 = 2; +const CLOSED: u8 = 4; +const OVERFLOW: u8 = 8; + +pub enum Registration { + First, + Set(Option), +} + +impl Default for Registration { + fn default() -> Self { + Self::Set(None) + } +} + +#[derive(Default)] +pub struct Waiters { + // EMPTY and REGISTERED can be claimed as BUSY. Closing is permanent. When the closer + // sees BUSY it leaves notification to the registrar instead of waiting for it. + phase: AtomicU8, + first: UnsafeCell>, + extra: OnceLock>>, +} + +// SAFETY: the BUSY owner or the thread that closes a REGISTERED slot has exclusive access. +unsafe impl Sync for Waiters {} + +// A panic in user work never exposes the slot. Executor waker operations must not panic. +impl RefUnwindSafe for Waiters {} + +impl Waiters { + fn extra(&self) -> &Mutex { + self.extra + .get_or_init(|| Box::new(Mutex::new(WakerSet::new()))) + } + + fn claim(&self, expected: u8, mut phase: u8) -> bool { + while phase & !OVERFLOW == expected { + match self.phase.compare_exchange_weak( + phase, + BUSY | (phase & OVERFLOW), + Ordering::Acquire, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(actual) => phase = actual, + } + } + false + } + + pub fn register(&self, token: &mut Registration, waker: &Waker) -> Poll<()> { + let phase = self.phase.load(Ordering::Acquire); + match token { + Registration::First if !self.claim(REGISTERED, phase) => { + // Only this observer can change its registration. A failed claim observed CLOSED. + *token = Registration::default(); + return Poll::Ready(()); + } + Registration::First => {} + Registration::Set(None) if self.claim(EMPTY, phase) => {} + Registration::Set(id) => return self.register_extra(id, waker, phase), + } + + // SAFETY: only the BUSY owner can access the cell until it publishes REGISTERED. + let slot = unsafe { &mut *self.first.get() }; + let retired = if slot.as_ref().is_some_and(|old| old.will_wake(waker)) { + None + } else { + slot.replace(waker.clone()) + }; + *token = Registration::First; + if self.phase.fetch_sub(BUSY - REGISTERED, Ordering::AcqRel) & CLOSED != 0 { + // The closer saw BUSY and delegated the slot to us. It will never be reused. + // SAFETY: closing BUSY left exclusive access with this registrar. + let removed = unsafe { &mut *self.first.get() }.take(); + *token = Registration::Set(None); + if let Some(waker) = removed { + waker.wake(); + } + drop(retired); + Poll::Ready(()) + } else { + drop(retired); + Poll::Pending + } + } + + fn register_extra(&self, id: &mut Option, waker: &Waker, phase: u8) -> Poll<()> { + // Announce overflow before taking its lock. Closing either sees this bit and takes the + // same lock, or wins this RMW and prevents registration. + let phase = if phase & OVERFLOW == 0 { + self.phase.fetch_or(OVERFLOW, Ordering::AcqRel) + } else { + phase + }; + if phase & CLOSED != 0 { + *id = None; + return Poll::Ready(()); + } + let mut set = self.extra().lock(); + if self.phase.load(Ordering::Acquire) & CLOSED != 0 { + *id = None; + return Poll::Ready(()); + } + let retired = set.register(id, waker); + drop(set); + drop(retired); + Poll::Pending + } + + pub fn unregister(&self, token: &mut Registration) { + match mem::take(token) { + Registration::Set(None) => {} + Registration::First if self.claim(REGISTERED, self.phase.load(Ordering::Acquire)) => { + // SAFETY: we own BUSY. Release the cell before dropping its waker. + let removed = unsafe { &mut *self.first.get() }.take(); + self.phase.fetch_sub(BUSY, Ordering::Release); + drop(removed); + } + Registration::First => {} + Registration::Set(mut id) => { + let mut set = self.extra().lock(); + if self.phase.load(Ordering::Acquire) & CLOSED == 0 { + let removed = set.unregister(&mut id); + drop(set); + drop(removed); + } + } + } + } + + pub fn wake_all(&self) { + let previous = self.phase.fetch_or(CLOSED, Ordering::AcqRel); + if previous & !OVERFLOW == REGISTERED { + // SAFETY: closing REGISTERED takes the slot and permanently prevents further claims. + if let Some(waker) = unsafe { &mut *self.first.get() }.take() { + waker.wake(); + } + } + if previous & OVERFLOW != 0 { + let wakers = mem::take(&mut *self.extra().lock()); + wakers.wake_all(); + } + } +} diff --git a/tests-integration/tests/waitgroup_test.rs b/tests-integration/tests/waitgroup_test.rs index 5bd36178..448753cd 100644 --- a/tests-integration/tests/waitgroup_test.rs +++ b/tests-integration/tests/waitgroup_test.rs @@ -24,6 +24,7 @@ use asyncband::waitgroup::WaitGroup; use tests_integration::WakeCounter; use tests_integration::poll_once; +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn waits_for_all_worker_handles() { let wg = WaitGroup::new(); @@ -139,3 +140,161 @@ fn repolling_replaces_the_registered_waker() { assert_eq!(first.count(), 0); assert_eq!(second.count(), 1); } + +#[test] +fn cancelled_slots_can_be_reused_while_other_observers_are_registered() { + let wg = WaitGroup::new(); + let worker = wg.clone(); + let observer = wg.into_future(); + let mut first = Box::pin(observer.clone()); + let mut extra = Box::pin(observer.clone()); + let mut survivor = Box::pin(observer.clone()); + let (old, old_count) = WakeCounter::new(); + let (live, live_count) = WakeCounter::new(); + let (replaced, replaced_count) = WakeCounter::new(); + + assert!(tests_integration::poll_with(first.as_mut(), &old).is_pending()); + assert!(tests_integration::poll_with(extra.as_mut(), &old).is_pending()); + assert!(tests_integration::poll_with(survivor.as_mut(), &replaced).is_pending()); + drop(first); + drop(extra); + + let mut first = Box::pin(observer.clone()); + let mut extra = Box::pin(observer); + for waiter in [&mut first, &mut extra, &mut survivor] { + assert!(tests_integration::poll_with(waiter.as_mut(), &live).is_pending()); + } + drop(worker); + assert_eq!(old_count.count(), 0); + assert_eq!(replaced_count.count(), 0); + assert_eq!(live_count.count(), 3); + for waiter in [&mut first, &mut extra, &mut survivor] { + assert!(poll_once(waiter.as_mut()).is_ready()); + } +} + +#[test] +fn concurrent_registration_acquires_every_participants_work() { + use std::sync::Barrier; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + use std::thread; + + let group = WaitGroup::new(); + let start = Barrier::new(6); + let writes = [const { AtomicBool::new(false) }; 3]; + let (waker, _) = WakeCounter::new(); + thread::scope(|scope| { + for write in &writes { + let worker = group.clone(); + let start = &start; + scope.spawn(move || { + start.wait(); + write.store(true, Ordering::Relaxed); + drop(worker); + }); + } + let observer = group.into_future(); + for _ in 0..2 { + let mut wait = Box::pin(observer.clone()); + let start = &start; + let writes = &writes; + let waker = &waker; + scope.spawn(move || { + start.wait(); + while tests_integration::poll_with(wait.as_mut(), waker).is_pending() { + thread::yield_now(); + } + assert!(writes.iter().all(|v| v.load(Ordering::Relaxed))); + }); + } + let mut wait = Box::pin(observer); + assert!(tests_integration::poll_with(wait.as_mut(), &waker).is_pending()); + start.wait(); + drop(wait); + }); +} + +#[test] +fn racing_completion_notifies_a_poll_that_returns_pending() { + use std::sync::Barrier; + use std::thread; + + let group = WaitGroup::new(); + let worker = group.clone(); + let mut wait = Box::pin(group.into_future()); + let start = Barrier::new(2); + let (waker, counter) = WakeCounter::new(); + let result = thread::scope(|scope| { + scope.spawn(|| { + start.wait(); + drop(worker); + }); + start.wait(); + tests_integration::poll_with(wait.as_mut(), &waker) + }); + if result.is_pending() { + assert_eq!(counter.count(), 1); + } + assert!(poll_once(wait.as_mut()).is_ready()); +} + +#[test] +fn wait_group_and_observers_preserve_unwind_safety() { + fn assert_traits() {} + assert_traits::(); + assert_traits::(); +} + +#[test] +fn retiring_a_registration_can_complete_the_group() { + for replace in [false, true] { + for use_overflow in [false, true] { + tests_integration::assert_completes_without_deadlock(move || { + let group = WaitGroup::new(); + let worker = group.clone(); + let observer = group.into_future(); + let mut other = pin!(observer.clone()); + if use_overflow { + assert!(poll_once(other.as_mut()).is_pending()); + } + let mut wait = Box::pin(observer); + let retired = tests_integration::waker_on_drop(move || drop(worker)); + assert!(tests_integration::poll_with(wait.as_mut(), &retired).is_pending()); + drop(retired); + if replace { + let (next, counter) = WakeCounter::new(); + assert!(tests_integration::poll_with(wait.as_mut(), &next).is_pending()); + assert_eq!(counter.count(), 1); + } + drop(wait); + assert!(poll_once(other.as_mut()).is_ready()); + }); + } + } +} + +#[test] +fn final_worker_and_cancelled_observer_release_their_registration() { + use std::sync::Barrier; + use std::thread; + + let group = WaitGroup::new(); + let worker = group.clone(); + let mut wait = Box::pin(group.into_future()); + let (waker, counter) = WakeCounter::new(); + assert!(tests_integration::poll_with(wait.as_mut(), &waker).is_pending()); + drop(waker); + + let start = Barrier::new(2); + thread::scope(|scope| { + scope.spawn(|| { + start.wait(); + drop(worker); + }); + start.wait(); + drop(wait); + }); + assert!(counter.count() <= 1); + assert_eq!(Arc::strong_count(&counter), 1); +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 565d7793..67342141 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -144,6 +144,10 @@ impl CommandMiri { run_command(make_miri_cmd("tests-integration", &["--test", "mpsc_test"])); run_command(make_miri_cmd("tests-integration", &["--test", "mpmc_test"])); run_command(make_miri_cmd("tests-integration", &["--test", "spmc_test"])); + run_command(make_miri_cmd( + "tests-integration", + &["--test", "waitgroup_test"], + )); run_command(make_miri_cmd( "tests-integration", &["--test", "broadcast_test"], From cfcca1f5a763211d36c2378e963c538b49bc7df2 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 3 Nov 2026 17:25:52 +0800 Subject: [PATCH 2/8] perf(benchmarks): compare wait group fan-in and shared drain workloads --- Cargo.lock | 56 +++++++- Cargo.toml | 2 + benchmarks/Cargo.toml | 2 + benchmarks/README.md | 2 + benchmarks/benches/ecosystem/waitgroup.rs | 65 +++++++++ .../benches/ecosystem/waitgroup/drain.rs | 127 ++++++++++++++++++ 6 files changed, 252 insertions(+), 2 deletions(-) create mode 100644 benchmarks/benches/ecosystem/waitgroup/drain.rs diff --git a/Cargo.lock b/Cargo.lock index 0da2c1d4..88c0e7c6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -105,7 +105,9 @@ dependencies = [ "asyncband", "clap", "flume", + "futures-util", "tokio", + "tokio-util", "waitgroup", ] @@ -357,9 +359,20 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.33" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-macro" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] [[package]] name = "futures-sink" @@ -367,6 +380,25 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-core", + "futures-macro", + "futures-task", + "pin-project-lite", + "slab", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -806,6 +838,12 @@ dependencies = [ "libc", ] +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + [[package]] name = "smallvec" version = "1.15.2" @@ -1001,6 +1039,20 @@ dependencies = [ "tokio-stream", ] +[[package]] +name = "tokio-util" +version = "0.7.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e464cf451ba96ebfc6f9b6542f17ee8b8956e33f1e40d9690624e59d7a7f8a4b" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "futures-util", + "pin-project-lite", + "tokio", +] + [[package]] name = "unicode-ident" version = "1.0.24" diff --git a/Cargo.toml b/Cargo.toml index 229085ca..1af2a4d9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,10 +40,12 @@ async-channel = { version = "2.5.0" } cargo_metadata = { version = "0.23.1" } clap = { version = "4.6.5" } flume = { version = "0.12.0", default-features = false } +futures-util = { version = "0.3.34" } semver = { version = "1.0.28" } serde = { version = "1.0.229" } tokio = { version = "1.53.1" } tokio-test = { version = "0.4.5" } +tokio-util = { version = "0.7.20" } ureq = { version = "3.3.0", default-features = false } waitgroup = { version = "0.1.2" } which = { version = "8.0.5" } diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index fffebeee..9c1d4d7d 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -55,6 +55,8 @@ tokio = { workspace = true, features = ["rt-multi-thread", "sync"] } [dev-dependencies] async-broadcast = { workspace = true } +futures-util = { workspace = true } +tokio-util = { workspace = true, features = ["rt"] } waitgroup = { workspace = true } [lib] diff --git a/benchmarks/README.md b/benchmarks/README.md index 945e0f8b..eb8f41dd 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -44,6 +44,8 @@ The clock starts before constructing the fixture and stops after cleanup, includ MPSC task scenarios include both one and four producers, small and large buffers, and current-thread controls. Use the reported range to decide whether an apparent difference is reproducible. To compare *different* scenarios, set the same round count where their work units are equivalent; independently sized rows are not a throughput ranking. +Wait group rows cover an already registered waiter, a coordinator that arrives after work finishes, and a worker that discovers four children while the coordinator is waiting. The drain row registers four independently woken observers before the last worker exits, as in graceful shutdown. Asyncband clones its completion future, the single-waiter `waitgroup` crate uses `FutureExt::shared` to fan out completion, and Tokio uses a closed `TaskTracker` with one live token. These rows measure the primitive lifecycle on one thread. A separate scheduled row spawns four tasks on a four-thread Tokio runtime, waits for the group, and joins all four tasks each round. + ## Set and preserve a baseline 1. Choose the scenario and a fixed reference revision. For the ecosystem workloads in this revision, the reference library is Asyncband at `c37b1a4`. diff --git a/benchmarks/benches/ecosystem/waitgroup.rs b/benchmarks/benches/ecosystem/waitgroup.rs index e34bbd64..b258a8ed 100644 --- a/benchmarks/benches/ecosystem/waitgroup.rs +++ b/benchmarks/benches/ecosystem/waitgroup.rs @@ -26,6 +26,8 @@ use benchmarks::support::bench_context; use benchmarks::support::poll_pending; use benchmarks::support::poll_pinned_ready; +mod drain; + struct Asyncband; struct WaitgroupRs; @@ -105,6 +107,53 @@ fn complete_worker_batch(bencher: Bencher, worker_count: usize) { }); } +// A parent can discover children after the coordinator has started waiting. +fn nested_workers(bencher: Bencher) { + let mut context = bench_context(); + bencher.bench_local(|| { + let group = C::new(); + let parent = C::worker(&group); + let mut wait = pin!(C::wait(group)); + poll_pending(wait.as_mut(), &mut context); + let children: [_; 4] = std::array::from_fn(|_| parent.clone()); + drop(parent); + poll_pending(wait.as_mut(), &mut context); + drop(children); + poll_pinned_ready(wait.as_mut(), &mut context); + }); +} + +fn wait_after_work(bencher: Bencher) { + let mut context = bench_context(); + bencher.bench_local(|| { + let group = C::new(); + let worker = C::worker(&group); + drop(worker); + let mut wait = pin!(C::wait(group)); + poll_pinned_ready(wait.as_mut(), &mut context); + }); +} + +fn scheduled_workers(bencher: Bencher) +where + C::Worker: Send + 'static, +{ + let runtime = benchmarks::channels::runtime(4); + bencher.bench_local(|| { + runtime.block_on(async { + let group = C::new(); + let tasks: [_; 4] = std::array::from_fn(|_| { + let worker = C::worker(&group); + tokio::spawn(async move { drop(worker) }) + }); + C::wait(group).await; + for task in tasks { + task.await.unwrap(); + } + }) + }); +} + pub fn register(runner: &mut Runner) { runner.compare( "waitgroup/complete-waiter", @@ -118,4 +167,20 @@ pub fn register(runner: &mut Runner) { b, 4 )), ); + runner.compare( + "waitgroup/wait-after-work", + 5_500_000, + implementations!([Asyncband, WaitgroupRs], |C, b| wait_after_work::(b)), + ); + runner.compare( + "waitgroup/nested-4-workers", + 3_000_000, + implementations!([Asyncband, WaitgroupRs], |C, b| nested_workers::(b)), + ); + runner.compare( + "waitgroup/tasks-4-workers-4w", + 25_000, + implementations!([Asyncband, WaitgroupRs], |C, b| scheduled_workers::(b)), + ); + drain::register(runner); } diff --git a/benchmarks/benches/ecosystem/waitgroup/drain.rs b/benchmarks/benches/ecosystem/waitgroup/drain.rs new file mode 100644 index 00000000..91b564dc --- /dev/null +++ b/benchmarks/benches/ecosystem/waitgroup/drain.rs @@ -0,0 +1,127 @@ +// 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. + +//! Four independent shutdown observers wait for the same last worker. + +use std::future::Future; +use std::future::IntoFuture; +use std::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Wake; +use std::task::Waker; + +use benchmarks::harness::Bencher; +use benchmarks::harness::Runner; +use benchmarks::implementations; +use benchmarks::support::poll_pending; +use benchmarks::support::poll_pinned_ready; +use futures_util::FutureExt; +use tokio_util::task::TaskTracker; + +fn drain, W>(worker: W, observers: [F; 4], tasks: &[Task; 4]) { + let [a, b, c, d] = observers; + let mut observers = [pin!(a), pin!(b), pin!(c), pin!(d)]; + for (wait, task) in observers.iter_mut().zip(tasks) { + poll_pending(wait.as_mut(), &mut Context::from_waker(&task.waker)); + } + drop(worker); + for (wait, task) in observers.into_iter().zip(tasks) { + assert!(task.notified.0.swap(false, Ordering::Relaxed)); + poll_pinned_ready(wait, &mut Context::from_waker(&task.waker)); + } +} + +#[derive(Default)] +struct Notified(AtomicBool); + +impl Wake for Notified { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + self.0.store(true, Ordering::Relaxed); + } +} + +struct Task { + notified: Arc, + waker: Waker, +} + +impl Task { + fn new() -> Self { + let notified = Arc::::default(); + Self { + waker: Waker::from(notified.clone()), + notified, + } + } +} + +trait Drain { + fn run(tasks: &[Task; 4]); +} + +struct Asyncband; +struct WaitgroupRsShared; +struct TokioTaskTracker; + +impl Drain for Asyncband { + fn run(tasks: &[Task; 4]) { + let group = asyncband::waitgroup::WaitGroup::new(); + let worker = group.clone(); + let wait = group.into_future(); + drain(worker, std::array::from_fn(|_| wait.clone()), tasks); + } +} + +impl Drain for WaitgroupRsShared { + fn run(tasks: &[Task; 4]) { + let group = waitgroup::WaitGroup::new(); + let worker = group.worker(); + let wait = group.wait().shared(); + drain(worker, std::array::from_fn(|_| wait.clone()), tasks); + } +} + +impl Drain for TokioTaskTracker { + fn run(tasks: &[Task; 4]) { + let tracker = TaskTracker::new(); + let worker = tracker.token(); + tracker.close(); + drain(worker, std::array::from_fn(|_| tracker.wait()), tasks); + } +} + +fn drain_four(bencher: Bencher) { + let tasks = std::array::from_fn(|_| Task::new()); + bencher.bench_local(|| C::run(&tasks)); +} + +pub fn register(runner: &mut Runner) { + runner.compare( + "waitgroup/drain-4-observers", + 2_500_000, + implementations!([Asyncband, WaitgroupRsShared, TokioTaskTracker], |C, b| { + drain_four::(b) + }), + ); +} From 4a40d3d45e6799bd32848d2b1d2a613e6527e4fd Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 3 Nov 2026 17:50:25 +0800 Subject: [PATCH 3/8] perf(waitgroup): register observers under a single inline lock --- asyncband/src/waitgroup/mod.rs | 4 +- asyncband/src/waitgroup/waiters.rs | 185 ++++++++++------------------- 2 files changed, 62 insertions(+), 127 deletions(-) diff --git a/asyncband/src/waitgroup/mod.rs b/asyncband/src/waitgroup/mod.rs index 0cc1b0a3..ca99bcd0 100644 --- a/asyncband/src/waitgroup/mod.rs +++ b/asyncband/src/waitgroup/mod.rs @@ -123,13 +123,13 @@ impl State { return Poll::Ready(()); } - self.waiters.register(token, cx.waker()) + self.waiters.register(token, cx.waker(), &self.handles) } #[inline] fn unregister(&self, token: &mut Registration) { if !matches!(token, Registration::Set(None)) && self.handles.load(Ordering::Acquire) != 0 { - self.waiters.unregister(token); + self.waiters.unregister(token, &self.handles); } } } diff --git a/asyncband/src/waitgroup/waiters.rs b/asyncband/src/waitgroup/waiters.rs index 96f1ede1..b9efc1d2 100644 --- a/asyncband/src/waitgroup/waiters.rs +++ b/asyncband/src/waitgroup/waiters.rs @@ -15,164 +15,99 @@ // specific language governing permissions and limitations // under the License. -use std::cell::UnsafeCell; use std::mem; -use std::panic::RefUnwindSafe; -use std::sync::OnceLock; -use std::sync::atomic::AtomicU8; +use std::sync::PoisonError; +use std::sync::RwLock; +use std::sync::RwLockWriteGuard; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Poll; use std::task::Waker; -use crate::internal::mutex::Mutex; use crate::internal::wakerset::WakerSet; use crate::internal::wakerset::WakerToken; -const EMPTY: u8 = 0; -const REGISTERED: u8 = 1; -const BUSY: u8 = 2; -const CLOSED: u8 = 4; -const OVERFLOW: u8 = 8; - pub enum Registration { First, Set(Option), } -impl Default for Registration { - fn default() -> Self { - Self::Set(None) - } -} - #[derive(Default)] -pub struct Waiters { - // EMPTY and REGISTERED can be claimed as BUSY. Closing is permanent. When the closer - // sees BUSY it leaves notification to the registrar instead of waiting for it. - phase: AtomicU8, - first: UnsafeCell>, - extra: OnceLock>>, +struct Slots { + first: Option, + extra: Option>, } -// SAFETY: the BUSY owner or the thread that closes a REGISTERED slot has exclusive access. -unsafe impl Sync for Waiters {} - -// A panic in user work never exposes the slot. Executor waker operations must not panic. -impl RefUnwindSafe for Waiters {} +// Registry mutations need exclusive access. The standard write lock keeps the uncontended +// lock state inline on Darwin, where a standard mutex allocates a pthread mutex on first use. +#[derive(Default)] +pub struct Waiters(RwLock); impl Waiters { - fn extra(&self) -> &Mutex { - self.extra - .get_or_init(|| Box::new(Mutex::new(WakerSet::new()))) + fn slots(&self) -> RwLockWriteGuard<'_, Slots> { + self.0.write().unwrap_or_else(PoisonError::into_inner) } - fn claim(&self, expected: u8, mut phase: u8) -> bool { - while phase & !OVERFLOW == expected { - match self.phase.compare_exchange_weak( - phase, - BUSY | (phase & OVERFLOW), - Ordering::Acquire, - Ordering::Acquire, - ) { - Ok(_) => return true, - Err(actual) => phase = actual, - } + pub fn register( + &self, + token: &mut Registration, + waker: &Waker, + handles: &AtomicUsize, + ) -> Poll<()> { + let mut slots = self.slots(); + // The last participant detaches registrations under this same lock. Either it takes our + // waker, or we observe zero and need no registration. + if handles.load(Ordering::Acquire) == 0 { + *token = Registration::Set(None); + return Poll::Ready(()); } - false - } - - pub fn register(&self, token: &mut Registration, waker: &Waker) -> Poll<()> { - let phase = self.phase.load(Ordering::Acquire); - match token { - Registration::First if !self.claim(REGISTERED, phase) => { - // Only this observer can change its registration. A failed claim observed CLOSED. - *token = Registration::default(); - return Poll::Ready(()); + let retired = match token { + Registration::First => { + let current = slots.first.as_mut().expect("registered first waker"); + if current.will_wake(waker) { + None + } else { + Some(mem::replace(current, waker.clone())) + } } - Registration::First => {} - Registration::Set(None) if self.claim(EMPTY, phase) => {} - Registration::Set(id) => return self.register_extra(id, waker, phase), - } - - // SAFETY: only the BUSY owner can access the cell until it publishes REGISTERED. - let slot = unsafe { &mut *self.first.get() }; - let retired = if slot.as_ref().is_some_and(|old| old.will_wake(waker)) { - None - } else { - slot.replace(waker.clone()) - }; - *token = Registration::First; - if self.phase.fetch_sub(BUSY - REGISTERED, Ordering::AcqRel) & CLOSED != 0 { - // The closer saw BUSY and delegated the slot to us. It will never be reused. - // SAFETY: closing BUSY left exclusive access with this registrar. - let removed = unsafe { &mut *self.first.get() }.take(); - *token = Registration::Set(None); - if let Some(waker) = removed { - waker.wake(); + Registration::Set(None) if slots.first.is_none() => { + slots.first = Some(waker.clone()); + *token = Registration::First; + None } - drop(retired); - Poll::Ready(()) - } else { - drop(retired); - Poll::Pending - } - } - - fn register_extra(&self, id: &mut Option, waker: &Waker, phase: u8) -> Poll<()> { - // Announce overflow before taking its lock. Closing either sees this bit and takes the - // same lock, or wins this RMW and prevents registration. - let phase = if phase & OVERFLOW == 0 { - self.phase.fetch_or(OVERFLOW, Ordering::AcqRel) - } else { - phase + Registration::Set(id) => slots + .extra + .get_or_insert_with(Default::default) + .register(id, waker), }; - if phase & CLOSED != 0 { - *id = None; - return Poll::Ready(()); - } - let mut set = self.extra().lock(); - if self.phase.load(Ordering::Acquire) & CLOSED != 0 { - *id = None; - return Poll::Ready(()); - } - let retired = set.register(id, waker); - drop(set); + drop(slots); drop(retired); Poll::Pending } - pub fn unregister(&self, token: &mut Registration) { - match mem::take(token) { - Registration::Set(None) => {} - Registration::First if self.claim(REGISTERED, self.phase.load(Ordering::Acquire)) => { - // SAFETY: we own BUSY. Release the cell before dropping its waker. - let removed = unsafe { &mut *self.first.get() }.take(); - self.phase.fetch_sub(BUSY, Ordering::Release); - drop(removed); - } - Registration::First => {} - Registration::Set(mut id) => { - let mut set = self.extra().lock(); - if self.phase.load(Ordering::Acquire) & CLOSED == 0 { - let removed = set.unregister(&mut id); - drop(set); - drop(removed); - } - } + pub fn unregister(&self, token: &mut Registration, handles: &AtomicUsize) { + let mut slots = self.slots(); + if handles.load(Ordering::Acquire) != 0 { + let retired = match token { + Registration::First => slots.first.take(), + Registration::Set(id) => slots + .extra + .as_mut() + .expect("registered overflow waker") + .unregister(id), + }; + drop(slots); + drop(retired); } } pub fn wake_all(&self) { - let previous = self.phase.fetch_or(CLOSED, Ordering::AcqRel); - if previous & !OVERFLOW == REGISTERED { - // SAFETY: closing REGISTERED takes the slot and permanently prevents further claims. - if let Some(waker) = unsafe { &mut *self.first.get() }.take() { - waker.wake(); - } + let slots = mem::take(&mut *self.slots()); + if let Some(first) = slots.first { + first.wake(); } - if previous & OVERFLOW != 0 { - let wakers = mem::take(&mut *self.extra().lock()); - wakers.wake_all(); + if let Some(extra) = slots.extra { + extra.wake_all(); } } } From f9d4a8a3233369124ceb339dcc181bf8723da029 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 3 Nov 2026 17:58:29 +0800 Subject: [PATCH 4/8] refactor(waitgroup): tie shared pointer copies to reference acquisition --- asyncband/src/waitgroup/mod.rs | 68 +++++++++++------------ tests-integration/tests/waitgroup_test.rs | 20 +++++++ 2 files changed, 54 insertions(+), 34 deletions(-) diff --git a/asyncband/src/waitgroup/mod.rs b/asyncband/src/waitgroup/mod.rs index ca99bcd0..b2ee32d1 100644 --- a/asyncband/src/waitgroup/mod.rs +++ b/asyncband/src/waitgroup/mod.rs @@ -93,29 +93,6 @@ impl State { } } - #[inline] - fn register_handle(&self) { - // The borrowed source handle keeps the count above zero. Registration publishes no data, - // so it does not need to synchronize with completion. - if self.handles.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { - std::process::abort(); - } - } - - #[inline] - fn release_handle(&self) -> bool { - // Every decrement is an RMW in the release sequence. A waiter that acquires zero therefore - // observes work published before every preceding handle release. - let previous = self.handles.fetch_sub(1, Ordering::AcqRel); - debug_assert!(previous > 0, "a live handle must own one count"); - if previous != 1 { - return false; - } - - self.waiters.wake_all(); - true - } - #[inline] fn poll_wait(&self, token: &mut Registration, cx: &mut Context<'_>) -> Poll<()> { if self.handles.load(Ordering::Acquire) == 0 { @@ -136,9 +113,7 @@ impl State { // One observer reference anchors the allocation until the last participant has notified waiters. // Each Wait owns an independent observer reference. The participant count keeps the anchor alive -// across registration and notification. A Shared pointer is only copied while its owner keeps a -// count alive. -#[derive(Clone, Copy)] +// across registration and notification. Duplicating this pointer requires acquiring a count. struct Shared(NonNull); // SAFETY: all shared access goes through State, which is Send + Sync. @@ -156,15 +131,41 @@ impl Deref for Shared { impl Shared { #[inline] - fn observe(self) -> Self { + fn register_handle(&self) -> Self { + // Borrowing a participant keeps the count above zero. Registration publishes no data. + if self.handles.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { + std::process::abort(); + } + Self(self.0) + } + + #[inline] + fn release_handle(&mut self) { + // The last participant acquires the preceding releases before notifying observers. + let previous = self.handles.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "a live handle must own one count"); + if previous == 1 { + self.finish(); + } + } + + // Keep notification and reclamation off the nonfinal participant drop path. + #[inline(never)] + fn finish(&mut self) { + self.waiters.wake_all(); + self.release_observer(); + } + + #[inline] + fn observe(&self) -> Self { if self.observers.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { std::process::abort(); } - self + Self(self.0) } #[inline] - fn release_observer(self) { + fn release_observer(&mut self) { if self.observers.fetch_sub(1, Ordering::AcqRel) == 1 { // SAFETY: all observers and the collective participant anchor have been released. unsafe { drop(Box::from_raw(self.0.as_ptr())) }; @@ -217,17 +218,16 @@ impl Clone for WaitGroup { /// its handle. #[inline] fn clone(&self) -> Self { - self.state.register_handle(); - Self { state: self.state } + Self { + state: self.state.register_handle(), + } } } impl Drop for WaitGroup { #[inline] fn drop(&mut self) { - if self.state.release_handle() { - self.state.release_observer(); - } + self.state.release_handle(); } } diff --git a/tests-integration/tests/waitgroup_test.rs b/tests-integration/tests/waitgroup_test.rs index 448753cd..6d7b7d6f 100644 --- a/tests-integration/tests/waitgroup_test.rs +++ b/tests-integration/tests/waitgroup_test.rs @@ -298,3 +298,23 @@ fn final_worker_and_cancelled_observer_release_their_registration() { assert!(counter.count() <= 1); assert_eq!(Arc::strong_count(&counter), 1); } + +#[test] +fn waking_can_drop_the_last_observer_during_completion() { + use std::sync::Mutex; + + let group = WaitGroup::new(); + let worker = group.clone(); + let observer = Arc::new(Mutex::new(Some(Box::pin(group.into_future())))); + let waker = tests_integration::waker_on_wake({ + let observer = observer.clone(); + move || drop(observer.lock().unwrap().take().unwrap()) + }); + let mut guard = observer.lock().unwrap(); + assert!(tests_integration::poll_with(guard.as_mut().unwrap().as_mut(), &waker).is_pending()); + drop(guard); + drop(waker); + + drop(worker); + assert!(observer.lock().unwrap().is_none()); +} From 13a31ae9fca74fc7fc510e3b148f8412f4587c2c Mon Sep 17 00:00:00 2001 From: tison Date: Fri, 13 Nov 2026 21:14:00 +0800 Subject: [PATCH 5/8] refactor(waitgroup): consolidate waiter registration and align API docs Waiter registration split observer state across a submodule, overloaded `Registration::Set(None)` as the unregistered state, and duplicated single-slot waker replacement logic while leaving `internal::mutex` gated on `waitgroup` without any remaining caller. Fold waiter storage into `waitgroup::mod` with explicit `Unregistered`, `Inline`, and `Overflow(WakerToken)` states backed by `register_waker` and `WakerSet::{insert, update, remove}`. Restore `WaitGroup` and `Wait` unwind-safety auto-traits via `PhantomData`, remove the invalid unwind-safety test, drop the unused `internal::mutex` feature gate, and align module documentation and the changelog entry with current conventions. --- CHANGELOG.md | 2 +- asyncband/src/internal/mod.rs | 1 - asyncband/src/internal/wakerset.rs | 46 +++- asyncband/src/waitgroup/mod.rs | 279 +++++++++++++++------- asyncband/src/waitgroup/waiters.rs | 113 --------- tests-integration/tests/waitgroup_test.rs | 117 +++++---- 6 files changed, 290 insertions(+), 268 deletions(-) delete mode 100644 asyncband/src/waitgroup/waiters.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ce364c5..6d2cf9e3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,7 @@ All notable changes to this project will be documented in this file. ### Improvements -* Reduce the cost of creating and dropping wait group participants and notifying a single registered waiter. +* Reduce `WaitGroup` completion overhead by sharing allocation ownership across participants and storing the first pending observer's waker inline. ## v0.8.0 diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index ccaa11ac..c4108016 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -76,7 +76,6 @@ pub(crate) mod value_cell; feature = "phaser", feature = "rwlock", feature = "semaphore", - feature = "waitgroup", feature = "watch", ))] // Some primitives use only shared access, leaving `Mutex::get_mut` unused. diff --git a/asyncband/src/internal/wakerset.rs b/asyncband/src/internal/wakerset.rs index 58f19aeb..211b9c82 100644 --- a/asyncband/src/internal/wakerset.rs +++ b/asyncband/src/internal/wakerset.rs @@ -80,32 +80,54 @@ impl WakerSet { self.wakers.into_iter().for_each(Waker::wake); } + /// Inserts a new waker slot and returns its exclusive token. + #[inline] + pub fn insert(&mut self, waker: &Waker) -> WakerToken { + WakerToken(self.wakers.insert(waker.clone())) + } + + /// Replaces the waker at `token` unless it already wakes the same task. + #[inline] + #[must_use = "drop the returned waker after releasing the waker set's state lock"] + pub fn update(&mut self, token: &WakerToken, waker: &Waker) -> Option { + let current = self + .wakers + .get_mut(token.0) + .expect("waker token must refer to an occupied slot"); + if current.will_wake(waker) { + None + } else { + Some(mem::replace(current, waker.clone())) + } + } + + /// Removes and returns the waker identified by `token`. + #[inline] + #[must_use = "drop the returned waker after releasing the waker set's state lock"] + pub fn remove(&mut self, token: WakerToken) -> Waker { + self.wakers.remove(token.0) + } + /// Registers or updates a waker. /// /// Returns the previous waker only if it was replaced. #[inline] #[must_use = "drop the returned waker after releasing the waker set's state lock"] pub fn register(&mut self, token: &mut Option, waker: &Waker) -> Option { - if let Some(token) = token { - let current = self - .wakers - .get_mut(token.0) - .expect("waker token must refer to an occupied slot"); - if current.will_wake(waker) { - return None; + match token { + Some(token) => self.update(token, waker), + None => { + *token = Some(self.insert(waker)); + None } - return Some(mem::replace(current, waker.clone())); } - - *token = Some(WakerToken(self.wakers.insert(waker.clone()))); - None } /// Removes and returns the waker identified by `token`, clearing the token. #[inline] #[must_use = "drop the returned waker after releasing the waker set's state lock"] pub fn unregister(&mut self, token: &mut Option) -> Option { - token.take().map(|token| self.wakers.remove(token.0)) + token.take().map(|token| self.remove(token)) } } diff --git a/asyncband/src/waitgroup/mod.rs b/asyncband/src/waitgroup/mod.rs index a64fb4d3..3986673f 100644 --- a/asyncband/src/waitgroup/mod.rs +++ b/asyncband/src/waitgroup/mod.rs @@ -17,157 +17,273 @@ //! Coordinate completion across a dynamically sized group of participants. //! -//! A [`WaitGroup`] starts with one participant. Cloning its handle registers another participant. -//! Awaiting a handle marks that participant complete and waits until every other participant has -//! completed. Dropping a handle marks its participant complete without waiting. +//! A [`WaitGroup`] starts with one participant. Cloning a handle registers another participant, +//! dropping a handle marks its participant complete without waiting, and awaiting a handle marks +//! its participant complete before waiting for every other participant to complete. //! -//! Participants are symmetric: any number of them may wait for the same completion. A waiting -//! participant no longer keeps the group pending, and all waiters are notified when the last -//! remaining handle is awaited or dropped. Existing handles may register more participants by -//! cloning until the group completes; completion is one-shot. +//! Participants are symmetric: workers can drop their handles when done while a coordinator awaits +//! its own handle, or multiple participants can each await completion. Converting a handle into +//! [`Wait`] marks that participant complete immediately, and cloning [`Wait`] creates additional +//! observers without registering new participants. Existing handles may register more participants +//! until the last live handle completes; completion is one-shot. //! -//! A coordinator can hand a clone to each worker, let workers drop their handles when done, -//! and await its own handle. To share completion with independent observers, convert a handle -//! into a [`Wait`] once and clone that future. Cloning the group itself adds more work. -//! -//! Completion acquires the state published before every participant completed, so work performed -//! by those participants is visible after the wait returns. +//! Completion acquires the state published before every participant completed, so writes performed +//! before each handle completed are visible after [`Wait`] completes. //! //! # Examples //! //! ``` //! # #[tokio::main] //! # async fn main() { -//! use asyncband::waitgroup::WaitGroup; +//! use std::sync::Arc; +//! use std::sync::atomic::AtomicBool; +//! use std::sync::atomic::Ordering; //! -//! async fn do_work() {} +//! use asyncband::waitgroup::WaitGroup; //! //! let group = WaitGroup::new(); -//! let mut tasks = vec![]; +//! let completed = Arc::new([const { AtomicBool::new(false) }; 3]); //! -//! for _ in 0..3 { +//! for slot in 0..3 { //! let participant = group.clone(); -//! tasks.push(tokio::spawn(async move { -//! do_work().await; +//! let completed = completed.clone(); +//! tokio::spawn(async move { +//! completed[slot].store(true, Ordering::Relaxed); //! drop(participant); -//! })); +//! }); //! } //! //! group.await; -//! for task in tasks { -//! task.await.unwrap(); -//! } +//! assert!(completed.iter().all(|done| done.load(Ordering::Relaxed))); //! # } //! ``` use std::fmt; use std::future::Future; use std::future::IntoFuture; +use std::marker::PhantomData; +use std::mem; use std::ops::Deref; use std::pin::Pin; use std::ptr::NonNull; +use std::sync::PoisonError; +use std::sync::RwLock; +use std::sync::RwLockWriteGuard; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use std::task::Waker; + +use crate::internal::register_waker; +use crate::internal::wakerset::WakerSet; +use crate::internal::wakerset::WakerToken; + +#[derive(Default)] +enum Registration { + #[default] + Unregistered, + Inline, + Overflow(WakerToken), +} + +#[derive(Default)] +struct Waiters { + inline: Option, + overflow: Option>, +} + +impl Waiters { + #[inline] + #[must_use = "drop the replaced waker after releasing the waiter lock"] + fn register(&mut self, registration: &mut Registration, waker: &Waker) -> Option { + match registration { + Registration::Unregistered if self.inline.is_none() => { + self.inline = Some(waker.clone()); + *registration = Registration::Inline; + None + } + Registration::Unregistered => { + let token = self.overflow.get_or_insert_default().insert(waker); + *registration = Registration::Overflow(token); + None + } + Registration::Inline => { + debug_assert!( + self.inline.is_some(), + "inline registration must occupy the inline waker slot" + ); + register_waker(&mut self.inline, waker) + } + Registration::Overflow(token) => self + .overflow + .as_mut() + .expect("overflow registration must have an allocated waker set") + .update(token, waker), + } + } -use crate::waitgroup::waiters::Registration; -use crate::waitgroup::waiters::Waiters; + #[inline] + #[must_use = "drop the removed waker after releasing the waiter lock"] + fn unregister(&mut self, registration: &mut Registration) -> Option { + match mem::take(registration) { + Registration::Unregistered => None, + Registration::Inline => self.inline.take(), + Registration::Overflow(token) => Some( + self.overflow + .as_mut() + .expect("overflow registration must have an allocated waker set") + .remove(token), + ), + } + } -mod waiters; + #[inline] + fn wake_all(self) { + if let Some(waker) = self.inline { + waker.wake(); + } + if let Some(overflow) = self.overflow { + overflow.wake_all(); + } + } +} struct State { - // Participants hold one collective observer count until completion has detached the wakers. - // The remaining observer counts belong to individual Wait futures. - handles: AtomicUsize, - observers: AtomicUsize, - waiters: Waiters, + // Live `WaitGroup` handles share one collective reference in `refs` until the last participant + // finishes notifying waiters; each `Wait` future owns an additional reference in `refs`. + participants: AtomicUsize, + refs: AtomicUsize, + // Waiter mutations always acquire exclusive access, but `std::sync::RwLock` keeps uncontended + // lock state inline on Darwin where `std::sync::Mutex` boxes a `pthread_mutex_t` on first + // lock. + waiters: RwLock, } impl State { - fn new() -> Self { - Self { - handles: AtomicUsize::new(1), - observers: AtomicUsize::new(1), - waiters: Waiters::default(), - } + fn lock_waiters(&self) -> RwLockWriteGuard<'_, Waiters> { + self.waiters.write().unwrap_or_else(PoisonError::into_inner) } #[inline] - fn poll_wait(&self, token: &mut Registration, cx: &mut Context<'_>) -> Poll<()> { - if self.handles.load(Ordering::Acquire) == 0 { - *token = Registration::Set(None); + fn poll_wait(&self, registration: &mut Registration, cx: &mut Context<'_>) -> Poll<()> { + if self.participants.load(Ordering::Acquire) == 0 { + *registration = Registration::Unregistered; + return Poll::Ready(()); + } + + self.register_wait(registration, cx.waker()) + } + + fn register_wait(&self, registration: &mut Registration, waker: &Waker) -> Poll<()> { + let mut waiters = self.lock_waiters(); + // `finish` detaches waiters under this same lock after storing zero, so either it wakes our + // registration or we observe zero here and complete immediately. + if self.participants.load(Ordering::Acquire) == 0 { + *registration = Registration::Unregistered; return Poll::Ready(()); } - self.waiters.register(token, cx.waker(), &self.handles) + let retired_waker = waiters.register(registration, waker); + drop(waiters); + drop(retired_waker); + Poll::Pending } #[inline] - fn unregister(&self, token: &mut Registration) { - if !matches!(token, Registration::Set(None)) && self.handles.load(Ordering::Acquire) != 0 { - self.waiters.unregister(token, &self.handles); + fn unregister(&self, registration: &mut Registration) { + if matches!(registration, Registration::Unregistered) { + return; + } + if self.participants.load(Ordering::Acquire) == 0 { + *registration = Registration::Unregistered; + return; } + + self.unregister_slow(registration); + } + + fn unregister_slow(&self, registration: &mut Registration) { + let mut waiters = self.lock_waiters(); + if self.participants.load(Ordering::Acquire) == 0 { + *registration = Registration::Unregistered; + return; + } + + let removed_waker = waiters.unregister(registration); + drop(waiters); + drop(removed_waker); } } -// One observer reference anchors the allocation until the last participant has notified waiters. -// Each Wait owns an independent observer reference. The participant count keeps the anchor alive -// across registration and notification. Duplicating this pointer requires acquiring a count. -struct Shared(NonNull); +struct Shared(NonNull, PhantomData); -// SAFETY: all shared access goes through State, which is Send + Sync. +// SAFETY: Shared state is synchronized through atomics and `RwLock`, and `State: Send + +// Sync`. unsafe impl Send for Shared {} unsafe impl Sync for Shared {} impl Deref for Shared { type Target = State; + #[inline] fn deref(&self) -> &State { - // SAFETY: each caller holds a participant or an observer count. + // SAFETY: Every `Shared` owns either a participant count (anchoring the collective ref) + // or an observer ref in `State::refs`. unsafe { self.0.as_ref() } } } impl Shared { + fn new() -> Self { + let state = Box::new(State { + participants: AtomicUsize::new(1), + refs: AtomicUsize::new(1), + waiters: RwLock::new(Waiters::default()), + }); + Self(NonNull::from(Box::leak(state)), PhantomData) + } + #[inline] - fn register_handle(&self) -> Self { - // Borrowing a participant keeps the count above zero. Registration publishes no data. - if self.handles.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { + fn clone_participant(&self) -> Self { + // Borrowing a live participant keeps `participants` above zero. Registration publishes no + // caller data; abort before wrapping `isize::MAX`. + if self.participants.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { std::process::abort(); } - Self(self.0) + Self(self.0, PhantomData) } #[inline] - fn release_handle(&mut self) { - // The last participant acquires the preceding releases before notifying observers. - let previous = self.handles.fetch_sub(1, Ordering::AcqRel); - debug_assert!(previous > 0, "a live handle must own one count"); + fn release_participant(&mut self) { + // The last participant acquires all preceding releases before waking waiters. + let previous = self.participants.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "a live WaitGroup must own one participant"); if previous == 1 { self.finish(); } } - // Keep notification and reclamation off the nonfinal participant drop path. + // Keep notification and reclamation off the inline non-final participant drop path. #[inline(never)] fn finish(&mut self) { - self.waiters.wake_all(); - self.release_observer(); + let waiters = mem::take(&mut *self.lock_waiters()); + waiters.wake_all(); + self.release_ref(); } #[inline] - fn observe(&self) -> Self { - if self.observers.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { + fn clone_ref(&self) -> Self { + if self.refs.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { std::process::abort(); } - Self(self.0) + Self(self.0, PhantomData) } #[inline] - fn release_observer(&mut self) { - if self.observers.fetch_sub(1, Ordering::AcqRel) == 1 { - // SAFETY: all observers and the collective participant anchor have been released. + fn release_ref(&mut self) { + if self.refs.fetch_sub(1, Ordering::AcqRel) == 1 { + // SAFETY: All observer references and the collective participant reference have ended. unsafe { drop(Box::from_raw(self.0.as_ptr())) }; } } @@ -177,7 +293,7 @@ impl Shared { /// /// See the [module level documentation](self) for more. pub struct WaitGroup { - // All participants share one reference to the allocation. + // All live participants share one collective reference in `State::refs`. state: Shared, } @@ -206,7 +322,7 @@ impl WaitGroup { #[inline] pub fn new() -> Self { Self { - state: Shared(NonNull::from(Box::leak(Box::new(State::new())))), + state: Shared::new(), } } } @@ -219,7 +335,7 @@ impl Clone for WaitGroup { #[inline] fn clone(&self) -> Self { Self { - state: self.state.register_handle(), + state: self.state.clone_participant(), } } } @@ -227,7 +343,7 @@ impl Clone for WaitGroup { impl Drop for WaitGroup { #[inline] fn drop(&mut self) { - self.state.release_handle(); + self.state.release_participant(); } } @@ -243,10 +359,10 @@ impl IntoFuture for WaitGroup { /// not undo that completion; see [`Wait`] for observer cancellation. #[inline] fn into_future(self) -> Self::IntoFuture { - let state = self.state.observe(); + let state = self.state.clone_ref(); drop(self); Wait { - token: Registration::Set(None), + registration: Registration::Unregistered, state, } } @@ -263,7 +379,7 @@ impl IntoFuture for WaitGroup { /// and other observers are unaffected. #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Wait { - token: Registration, + registration: Registration, state: Shared, } @@ -278,8 +394,8 @@ impl Clone for Wait { #[inline] fn clone(&self) -> Self { Wait { - token: Registration::Set(None), - state: self.state.observe(), + registration: Registration::Unregistered, + state: self.state.clone_ref(), } } } @@ -295,15 +411,18 @@ impl Future for Wait { #[inline] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let Self { token, state } = self.get_mut(); - state.poll_wait(token, cx) + let Self { + registration, + state, + } = self.get_mut(); + state.poll_wait(registration, cx) } } impl Drop for Wait { #[inline] fn drop(&mut self) { - self.state.unregister(&mut self.token); - self.state.release_observer(); + self.state.unregister(&mut self.registration); + self.state.release_ref(); } } diff --git a/asyncband/src/waitgroup/waiters.rs b/asyncband/src/waitgroup/waiters.rs deleted file mode 100644 index b9efc1d2..00000000 --- a/asyncband/src/waitgroup/waiters.rs +++ /dev/null @@ -1,113 +0,0 @@ -// 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::mem; -use std::sync::PoisonError; -use std::sync::RwLock; -use std::sync::RwLockWriteGuard; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Poll; -use std::task::Waker; - -use crate::internal::wakerset::WakerSet; -use crate::internal::wakerset::WakerToken; - -pub enum Registration { - First, - Set(Option), -} - -#[derive(Default)] -struct Slots { - first: Option, - extra: Option>, -} - -// Registry mutations need exclusive access. The standard write lock keeps the uncontended -// lock state inline on Darwin, where a standard mutex allocates a pthread mutex on first use. -#[derive(Default)] -pub struct Waiters(RwLock); - -impl Waiters { - fn slots(&self) -> RwLockWriteGuard<'_, Slots> { - self.0.write().unwrap_or_else(PoisonError::into_inner) - } - - pub fn register( - &self, - token: &mut Registration, - waker: &Waker, - handles: &AtomicUsize, - ) -> Poll<()> { - let mut slots = self.slots(); - // The last participant detaches registrations under this same lock. Either it takes our - // waker, or we observe zero and need no registration. - if handles.load(Ordering::Acquire) == 0 { - *token = Registration::Set(None); - return Poll::Ready(()); - } - let retired = match token { - Registration::First => { - let current = slots.first.as_mut().expect("registered first waker"); - if current.will_wake(waker) { - None - } else { - Some(mem::replace(current, waker.clone())) - } - } - Registration::Set(None) if slots.first.is_none() => { - slots.first = Some(waker.clone()); - *token = Registration::First; - None - } - Registration::Set(id) => slots - .extra - .get_or_insert_with(Default::default) - .register(id, waker), - }; - drop(slots); - drop(retired); - Poll::Pending - } - - pub fn unregister(&self, token: &mut Registration, handles: &AtomicUsize) { - let mut slots = self.slots(); - if handles.load(Ordering::Acquire) != 0 { - let retired = match token { - Registration::First => slots.first.take(), - Registration::Set(id) => slots - .extra - .as_mut() - .expect("registered overflow waker") - .unregister(id), - }; - drop(slots); - drop(retired); - } - } - - pub fn wake_all(&self) { - let slots = mem::take(&mut *self.slots()); - if let Some(first) = slots.first { - first.wake(); - } - if let Some(extra) = slots.extra { - extra.wake_all(); - } - } -} diff --git a/tests-integration/tests/waitgroup_test.rs b/tests-integration/tests/waitgroup_test.rs index 6d7b7d6f..85bde96d 100644 --- a/tests-integration/tests/waitgroup_test.rs +++ b/tests-integration/tests/waitgroup_test.rs @@ -17,12 +17,21 @@ use std::pin::pin; use std::sync::Arc; +use std::sync::Barrier; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Waker; +use std::thread; use asyncband::waitgroup::WaitGroup; use tests_integration::WakeCounter; +use tests_integration::assert_completes_without_deadlock; use tests_integration::poll_once; +use tests_integration::poll_with; +use tests_integration::waker_on_drop; +use tests_integration::waker_on_wake; #[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] @@ -143,47 +152,44 @@ fn repolling_replaces_the_registered_waker() { #[test] fn cancelled_slots_can_be_reused_while_other_observers_are_registered() { - let wg = WaitGroup::new(); - let worker = wg.clone(); - let observer = wg.into_future(); - let mut first = Box::pin(observer.clone()); - let mut extra = Box::pin(observer.clone()); + let group = WaitGroup::new(); + let worker = group.clone(); + let observer = group.into_future(); + let mut inline = Box::pin(observer.clone()); + let mut overflow = Box::pin(observer.clone()); let mut survivor = Box::pin(observer.clone()); - let (old, old_count) = WakeCounter::new(); - let (live, live_count) = WakeCounter::new(); - let (replaced, replaced_count) = WakeCounter::new(); - - assert!(tests_integration::poll_with(first.as_mut(), &old).is_pending()); - assert!(tests_integration::poll_with(extra.as_mut(), &old).is_pending()); - assert!(tests_integration::poll_with(survivor.as_mut(), &replaced).is_pending()); - drop(first); - drop(extra); - - let mut first = Box::pin(observer.clone()); - let mut extra = Box::pin(observer); - for waiter in [&mut first, &mut extra, &mut survivor] { - assert!(tests_integration::poll_with(waiter.as_mut(), &live).is_pending()); + let (cancelled_waker, cancelled_wakes) = WakeCounter::new(); + let (replaced_waker, replaced_wakes) = WakeCounter::new(); + let (live_waker, live_wakes) = WakeCounter::new(); + + assert!(poll_with(inline.as_mut(), &cancelled_waker).is_pending()); + assert!(poll_with(overflow.as_mut(), &cancelled_waker).is_pending()); + assert!(poll_with(survivor.as_mut(), &replaced_waker).is_pending()); + drop(inline); + drop(overflow); + + let mut inline = Box::pin(observer.clone()); + let mut overflow = Box::pin(observer); + for waiter in [&mut inline, &mut overflow, &mut survivor] { + assert!(poll_with(waiter.as_mut(), &live_waker).is_pending()); } drop(worker); - assert_eq!(old_count.count(), 0); - assert_eq!(replaced_count.count(), 0); - assert_eq!(live_count.count(), 3); - for waiter in [&mut first, &mut extra, &mut survivor] { + + assert_eq!(cancelled_wakes.count(), 0); + assert_eq!(replaced_wakes.count(), 0); + assert_eq!(live_wakes.count(), 3); + for waiter in [&mut inline, &mut overflow, &mut survivor] { assert!(poll_once(waiter.as_mut()).is_ready()); } } #[test] fn concurrent_registration_acquires_every_participants_work() { - use std::sync::Barrier; - use std::sync::atomic::AtomicBool; - use std::sync::atomic::Ordering; - use std::thread; - let group = WaitGroup::new(); let start = Barrier::new(6); let writes = [const { AtomicBool::new(false) }; 3]; let (waker, _) = WakeCounter::new(); + thread::scope(|scope| { for write in &writes { let worker = group.clone(); @@ -202,14 +208,14 @@ fn concurrent_registration_acquires_every_participants_work() { let waker = &waker; scope.spawn(move || { start.wait(); - while tests_integration::poll_with(wait.as_mut(), waker).is_pending() { + while poll_with(wait.as_mut(), waker).is_pending() { thread::yield_now(); } - assert!(writes.iter().all(|v| v.load(Ordering::Relaxed))); + assert!(writes.iter().all(|written| written.load(Ordering::Relaxed))); }); } let mut wait = Box::pin(observer); - assert!(tests_integration::poll_with(wait.as_mut(), &waker).is_pending()); + assert!(poll_with(wait.as_mut(), &waker).is_pending()); start.wait(); drop(wait); }); @@ -217,40 +223,31 @@ fn concurrent_registration_acquires_every_participants_work() { #[test] fn racing_completion_notifies_a_poll_that_returns_pending() { - use std::sync::Barrier; - use std::thread; - let group = WaitGroup::new(); let worker = group.clone(); let mut wait = Box::pin(group.into_future()); let start = Barrier::new(2); - let (waker, counter) = WakeCounter::new(); + let (waker, wakes) = WakeCounter::new(); + let result = thread::scope(|scope| { scope.spawn(|| { start.wait(); drop(worker); }); start.wait(); - tests_integration::poll_with(wait.as_mut(), &waker) + poll_with(wait.as_mut(), &waker) }); if result.is_pending() { - assert_eq!(counter.count(), 1); + assert_eq!(wakes.count(), 1); } assert!(poll_once(wait.as_mut()).is_ready()); } -#[test] -fn wait_group_and_observers_preserve_unwind_safety() { - fn assert_traits() {} - assert_traits::(); - assert_traits::(); -} - #[test] fn retiring_a_registration_can_complete_the_group() { for replace in [false, true] { for use_overflow in [false, true] { - tests_integration::assert_completes_without_deadlock(move || { + assert_completes_without_deadlock(move || { let group = WaitGroup::new(); let worker = group.clone(); let observer = group.into_future(); @@ -258,14 +255,16 @@ fn retiring_a_registration_can_complete_the_group() { if use_overflow { assert!(poll_once(other.as_mut()).is_pending()); } + let mut wait = Box::pin(observer); - let retired = tests_integration::waker_on_drop(move || drop(worker)); - assert!(tests_integration::poll_with(wait.as_mut(), &retired).is_pending()); + let retired = waker_on_drop(move || drop(worker)); + assert!(poll_with(wait.as_mut(), &retired).is_pending()); drop(retired); + if replace { - let (next, counter) = WakeCounter::new(); - assert!(tests_integration::poll_with(wait.as_mut(), &next).is_pending()); - assert_eq!(counter.count(), 1); + let (next, wakes) = WakeCounter::new(); + assert!(poll_with(wait.as_mut(), &next).is_pending()); + assert_eq!(wakes.count(), 1); } drop(wait); assert!(poll_once(other.as_mut()).is_ready()); @@ -276,14 +275,11 @@ fn retiring_a_registration_can_complete_the_group() { #[test] fn final_worker_and_cancelled_observer_release_their_registration() { - use std::sync::Barrier; - use std::thread; - let group = WaitGroup::new(); let worker = group.clone(); let mut wait = Box::pin(group.into_future()); - let (waker, counter) = WakeCounter::new(); - assert!(tests_integration::poll_with(wait.as_mut(), &waker).is_pending()); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(wait.as_mut(), &waker).is_pending()); drop(waker); let start = Barrier::new(2); @@ -295,23 +291,22 @@ fn final_worker_and_cancelled_observer_release_their_registration() { start.wait(); drop(wait); }); - assert!(counter.count() <= 1); - assert_eq!(Arc::strong_count(&counter), 1); + assert!(wakes.count() <= 1); + assert_eq!(Arc::strong_count(&wakes), 1); } #[test] fn waking_can_drop_the_last_observer_during_completion() { - use std::sync::Mutex; - let group = WaitGroup::new(); let worker = group.clone(); let observer = Arc::new(Mutex::new(Some(Box::pin(group.into_future())))); - let waker = tests_integration::waker_on_wake({ + let waker = waker_on_wake({ let observer = observer.clone(); move || drop(observer.lock().unwrap().take().unwrap()) }); + let mut guard = observer.lock().unwrap(); - assert!(tests_integration::poll_with(guard.as_mut().unwrap().as_mut(), &waker).is_pending()); + assert!(poll_with(guard.as_mut().unwrap().as_mut(), &waker).is_pending()); drop(guard); drop(waker); From 7cd4393960b372f193d584c0dd42853aeee90dc9 Mon Sep 17 00:00:00 2001 From: tison Date: Fri, 13 Nov 2026 21:14:00 +0800 Subject: [PATCH 6/8] refactor(benchmarks): simplify wait group ecosystem suite The shared-drain benchmark maintained a separate submodule with custom atomic waker plumbing to assert single wakeups inside a timing loop, and `tasks-4-workers-4w` primarily measured Tokio task spawning rather than `WaitGroup` coordination. Fold `drain-4-observers` into `benches/ecosystem/waitgroup.rs` on top of the shared `bench_context` helpers and drop the noisy task-spawning scenario so the suite focuses on `WaitGroup` coordination costs. --- benchmarks/README.md | 2 +- benchmarks/benches/ecosystem/waitgroup.rs | 113 +++++++++++----- .../benches/ecosystem/waitgroup/drain.rs | 127 ------------------ 3 files changed, 79 insertions(+), 163 deletions(-) delete mode 100644 benchmarks/benches/ecosystem/waitgroup/drain.rs diff --git a/benchmarks/README.md b/benchmarks/README.md index eb8f41dd..10eedd8c 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -44,7 +44,7 @@ The clock starts before constructing the fixture and stops after cleanup, includ MPSC task scenarios include both one and four producers, small and large buffers, and current-thread controls. Use the reported range to decide whether an apparent difference is reproducible. To compare *different* scenarios, set the same round count where their work units are equivalent; independently sized rows are not a throughput ranking. -Wait group rows cover an already registered waiter, a coordinator that arrives after work finishes, and a worker that discovers four children while the coordinator is waiting. The drain row registers four independently woken observers before the last worker exits, as in graceful shutdown. Asyncband clones its completion future, the single-waiter `waitgroup` crate uses `FutureExt::shared` to fan out completion, and Tokio uses a closed `TaskTracker` with one live token. These rows measure the primitive lifecycle on one thread. A separate scheduled row spawns four tasks on a four-thread Tokio runtime, waits for the group, and joins all four tasks each round. +Wait group rows cover an already registered observer, an observer that polls after work finishes, and a worker that registers four nested workers while an observer is waiting. The drain row registers four observers before the last worker exits, as in graceful shutdown: Asyncband clones `Wait`, the single-waiter `waitgroup` crate wraps its future in `FutureExt::shared`, and Tokio uses a closed `TaskTracker` with one live token. ## Set and preserve a baseline diff --git a/benchmarks/benches/ecosystem/waitgroup.rs b/benchmarks/benches/ecosystem/waitgroup.rs index b258a8ed..df21e573 100644 --- a/benchmarks/benches/ecosystem/waitgroup.rs +++ b/benchmarks/benches/ecosystem/waitgroup.rs @@ -15,9 +15,11 @@ // specific language governing permissions and limitations // under the License. +use std::array; use std::future::Future; use std::future::IntoFuture; use std::pin::pin; +use std::task::Context; use benchmarks::harness::Bencher; use benchmarks::harness::Runner; @@ -25,11 +27,13 @@ use benchmarks::implementations; use benchmarks::support::bench_context; use benchmarks::support::poll_pending; use benchmarks::support::poll_pinned_ready; - -mod drain; +use futures_util::FutureExt; +use tokio_util::task::TaskTracker; struct Asyncband; struct WaitgroupRs; +struct WaitgroupRsShared; +struct TokioTaskTracker; trait WaitGroup { type Group; @@ -77,6 +81,53 @@ impl WaitGroup for WaitgroupRs { } } +trait SharedDrain { + fn drain(context: &mut Context<'_>); +} + +fn complete_observers, W>( + worker: W, + observers: [F; 4], + context: &mut Context<'_>, +) { + let [a, b, c, d] = observers; + let mut observers = [pin!(a), pin!(b), pin!(c), pin!(d)]; + for observer in &mut observers { + poll_pending(observer.as_mut(), context); + } + drop(worker); + for observer in observers { + poll_pinned_ready(observer, context); + } +} + +impl SharedDrain for Asyncband { + fn drain(context: &mut Context<'_>) { + let group = asyncband::waitgroup::WaitGroup::new(); + let worker = group.clone(); + let wait = group.into_future(); + complete_observers(worker, array::from_fn(|_| wait.clone()), context); + } +} + +impl SharedDrain for WaitgroupRsShared { + fn drain(context: &mut Context<'_>) { + let group = waitgroup::WaitGroup::new(); + let worker = group.worker(); + let wait = group.wait().shared(); + complete_observers(worker, array::from_fn(|_| wait.clone()), context); + } +} + +impl SharedDrain for TokioTaskTracker { + fn drain(context: &mut Context<'_>) { + let tracker = TaskTracker::new(); + let worker = tracker.token(); + tracker.close(); + complete_observers(worker, array::from_fn(|_| tracker.wait()), context); + } +} + fn complete_waiter(bencher: Bencher) { let mut context = bench_context(); @@ -107,51 +158,42 @@ fn complete_worker_batch(bencher: Bencher, worker_count: usize) { }); } -// A parent can discover children after the coordinator has started waiting. -fn nested_workers(bencher: Bencher) { +fn wait_after_work(bencher: Bencher) { let mut context = bench_context(); + bencher.bench_local(|| { let group = C::new(); - let parent = C::worker(&group); + let worker = C::worker(&group); + drop(worker); + let mut wait = pin!(C::wait(group)); - poll_pending(wait.as_mut(), &mut context); - let children: [_; 4] = std::array::from_fn(|_| parent.clone()); - drop(parent); - poll_pending(wait.as_mut(), &mut context); - drop(children); poll_pinned_ready(wait.as_mut(), &mut context); }); } -fn wait_after_work(bencher: Bencher) { +// A worker can register child work after an observer has already started waiting. +fn nested_workers(bencher: Bencher) { let mut context = bench_context(); + bencher.bench_local(|| { let group = C::new(); - let worker = C::worker(&group); - drop(worker); + let parent = C::worker(&group); let mut wait = pin!(C::wait(group)); + poll_pending(wait.as_mut(), &mut context); + + let children: [_; 4] = array::from_fn(|_| parent.clone()); + drop(parent); + poll_pending(wait.as_mut(), &mut context); + + drop(children); poll_pinned_ready(wait.as_mut(), &mut context); }); } -fn scheduled_workers(bencher: Bencher) -where - C::Worker: Send + 'static, -{ - let runtime = benchmarks::channels::runtime(4); - bencher.bench_local(|| { - runtime.block_on(async { - let group = C::new(); - let tasks: [_; 4] = std::array::from_fn(|_| { - let worker = C::worker(&group); - tokio::spawn(async move { drop(worker) }) - }); - C::wait(group).await; - for task in tasks { - task.await.unwrap(); - } - }) - }); +fn drain_observers(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| C::drain(&mut context)); } pub fn register(runner: &mut Runner) { @@ -178,9 +220,10 @@ pub fn register(runner: &mut Runner) { implementations!([Asyncband, WaitgroupRs], |C, b| nested_workers::(b)), ); runner.compare( - "waitgroup/tasks-4-workers-4w", - 25_000, - implementations!([Asyncband, WaitgroupRs], |C, b| scheduled_workers::(b)), + "waitgroup/drain-4-observers", + 2_500_000, + implementations!([Asyncband, WaitgroupRsShared, TokioTaskTracker], |C, b| { + drain_observers::(b) + }), ); - drain::register(runner); } diff --git a/benchmarks/benches/ecosystem/waitgroup/drain.rs b/benchmarks/benches/ecosystem/waitgroup/drain.rs deleted file mode 100644 index 91b564dc..00000000 --- a/benchmarks/benches/ecosystem/waitgroup/drain.rs +++ /dev/null @@ -1,127 +0,0 @@ -// 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. - -//! Four independent shutdown observers wait for the same last worker. - -use std::future::Future; -use std::future::IntoFuture; -use std::pin::pin; -use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Wake; -use std::task::Waker; - -use benchmarks::harness::Bencher; -use benchmarks::harness::Runner; -use benchmarks::implementations; -use benchmarks::support::poll_pending; -use benchmarks::support::poll_pinned_ready; -use futures_util::FutureExt; -use tokio_util::task::TaskTracker; - -fn drain, W>(worker: W, observers: [F; 4], tasks: &[Task; 4]) { - let [a, b, c, d] = observers; - let mut observers = [pin!(a), pin!(b), pin!(c), pin!(d)]; - for (wait, task) in observers.iter_mut().zip(tasks) { - poll_pending(wait.as_mut(), &mut Context::from_waker(&task.waker)); - } - drop(worker); - for (wait, task) in observers.into_iter().zip(tasks) { - assert!(task.notified.0.swap(false, Ordering::Relaxed)); - poll_pinned_ready(wait, &mut Context::from_waker(&task.waker)); - } -} - -#[derive(Default)] -struct Notified(AtomicBool); - -impl Wake for Notified { - fn wake(self: Arc) { - self.wake_by_ref(); - } - - fn wake_by_ref(self: &Arc) { - self.0.store(true, Ordering::Relaxed); - } -} - -struct Task { - notified: Arc, - waker: Waker, -} - -impl Task { - fn new() -> Self { - let notified = Arc::::default(); - Self { - waker: Waker::from(notified.clone()), - notified, - } - } -} - -trait Drain { - fn run(tasks: &[Task; 4]); -} - -struct Asyncband; -struct WaitgroupRsShared; -struct TokioTaskTracker; - -impl Drain for Asyncband { - fn run(tasks: &[Task; 4]) { - let group = asyncband::waitgroup::WaitGroup::new(); - let worker = group.clone(); - let wait = group.into_future(); - drain(worker, std::array::from_fn(|_| wait.clone()), tasks); - } -} - -impl Drain for WaitgroupRsShared { - fn run(tasks: &[Task; 4]) { - let group = waitgroup::WaitGroup::new(); - let worker = group.worker(); - let wait = group.wait().shared(); - drain(worker, std::array::from_fn(|_| wait.clone()), tasks); - } -} - -impl Drain for TokioTaskTracker { - fn run(tasks: &[Task; 4]) { - let tracker = TaskTracker::new(); - let worker = tracker.token(); - tracker.close(); - drain(worker, std::array::from_fn(|_| tracker.wait()), tasks); - } -} - -fn drain_four(bencher: Bencher) { - let tasks = std::array::from_fn(|_| Task::new()); - bencher.bench_local(|| C::run(&tasks)); -} - -pub fn register(runner: &mut Runner) { - runner.compare( - "waitgroup/drain-4-observers", - 2_500_000, - implementations!([Asyncband, WaitgroupRsShared, TokioTaskTracker], |C, b| { - drain_four::(b) - }), - ); -} From 84e5a8e4fdea0f150ad9445a7600d0faee00294b Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 14 Nov 2026 09:54:55 +0800 Subject: [PATCH 7/8] refactor(waitgroup): encapsulate waiter RwLock in internal::rwlock Add `crate::internal::rwlock::RwLock` alongside `crate::internal::mutex::Mutex` to handle lock-poisoning recovery consistently, and use it for `WaitGroup` waiter registration instead of open-coding `PoisonError::into_inner` in `waitgroup`. --- asyncband/src/internal/mod.rs | 5 ++++ asyncband/src/internal/rwlock.rs | 45 ++++++++++++++++++++++++++++++++ asyncband/src/waitgroup/mod.rs | 14 +++------- 3 files changed, 54 insertions(+), 10 deletions(-) create mode 100644 asyncband/src/internal/rwlock.rs diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index c4108016..29a3e73c 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -82,6 +82,11 @@ pub(crate) mod value_cell; #[allow(dead_code)] pub(crate) mod mutex; +#[cfg(feature = "waitgroup")] +// Some primitives use only exclusive access, leaving `read` and `get_mut` unused. +#[allow(dead_code)] +pub(crate) mod rwlock; + #[cfg(any( feature = "broadcast", feature = "mutex", diff --git a/asyncband/src/internal/rwlock.rs b/asyncband/src/internal/rwlock.rs new file mode 100644 index 00000000..a4c06c83 --- /dev/null +++ b/asyncband/src/internal/rwlock.rs @@ -0,0 +1,45 @@ +// 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::fmt; +use std::sync::PoisonError; + +pub struct RwLock(std::sync::RwLock); + +impl fmt::Debug for RwLock { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl RwLock { + pub const fn new(t: T) -> Self { + Self(std::sync::RwLock::new(t)) + } + + pub fn read(&self) -> std::sync::RwLockReadGuard<'_, T> { + self.0.read().unwrap_or_else(PoisonError::into_inner) + } + + pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, T> { + self.0.write().unwrap_or_else(PoisonError::into_inner) + } + + pub fn get_mut(&mut self) -> &mut T { + self.0.get_mut().unwrap_or_else(PoisonError::into_inner) + } +} diff --git a/asyncband/src/waitgroup/mod.rs b/asyncband/src/waitgroup/mod.rs index 3986673f..ade93b56 100644 --- a/asyncband/src/waitgroup/mod.rs +++ b/asyncband/src/waitgroup/mod.rs @@ -66,9 +66,6 @@ use std::mem; use std::ops::Deref; use std::pin::Pin; use std::ptr::NonNull; -use std::sync::PoisonError; -use std::sync::RwLock; -use std::sync::RwLockWriteGuard; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -76,6 +73,7 @@ use std::task::Poll; use std::task::Waker; use crate::internal::register_waker; +use crate::internal::rwlock::RwLock; use crate::internal::wakerset::WakerSet; use crate::internal::wakerset::WakerToken; @@ -161,10 +159,6 @@ struct State { } impl State { - fn lock_waiters(&self) -> RwLockWriteGuard<'_, Waiters> { - self.waiters.write().unwrap_or_else(PoisonError::into_inner) - } - #[inline] fn poll_wait(&self, registration: &mut Registration, cx: &mut Context<'_>) -> Poll<()> { if self.participants.load(Ordering::Acquire) == 0 { @@ -176,7 +170,7 @@ impl State { } fn register_wait(&self, registration: &mut Registration, waker: &Waker) -> Poll<()> { - let mut waiters = self.lock_waiters(); + let mut waiters = self.waiters.write(); // `finish` detaches waiters under this same lock after storing zero, so either it wakes our // registration or we observe zero here and complete immediately. if self.participants.load(Ordering::Acquire) == 0 { @@ -204,7 +198,7 @@ impl State { } fn unregister_slow(&self, registration: &mut Registration) { - let mut waiters = self.lock_waiters(); + let mut waiters = self.waiters.write(); if self.participants.load(Ordering::Acquire) == 0 { *registration = Registration::Unregistered; return; @@ -267,7 +261,7 @@ impl Shared { // Keep notification and reclamation off the inline non-final participant drop path. #[inline(never)] fn finish(&mut self) { - let waiters = mem::take(&mut *self.lock_waiters()); + let waiters = mem::take(&mut *self.waiters.write()); waiters.wake_all(); self.release_ref(); } From 63f7c76677680ed705e8955516edd38fcd1cb1f1 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 14 Nov 2026 09:54:55 +0800 Subject: [PATCH 8/8] docs: simplify WaitGroup and Barrier module examples Use `do_work().await` directly in the `WaitGroup` module example without redundant `JoinHandle` awaits or inline-const atomic arrays, and replace the inline-const atomic array in the `Barrier` module example with a single arrival counter. --- asyncband/src/barrier/mod.rs | 11 +++++------ asyncband/src/waitgroup/mod.rs | 13 ++++--------- 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index 7e33fb26..9042715a 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -33,17 +33,16 @@ //! use asyncband::barrier::Barrier; //! //! let barrier = Arc::new(Barrier::new(3)); -//! let stage = Arc::new([const { AtomicUsize::new(0) }; 3]); +//! let arrived = Arc::new(AtomicUsize::new(0)); //! let mut workers = Vec::with_capacity(3); //! -//! for slot in 0..3 { +//! for _ in 0..3 { //! let barrier = barrier.clone(); -//! let stage = stage.clone(); +//! let arrived = arrived.clone(); //! workers.push(tokio::spawn(async move { -//! stage[slot].store(slot + 1, Ordering::Release); +//! arrived.fetch_add(1, Ordering::Relaxed); //! let result = barrier.wait().await; -//! let sum: usize = stage.iter().map(|v| v.load(Ordering::Acquire)).sum(); -//! assert_eq!(sum, 6); +//! assert_eq!(arrived.load(Ordering::Relaxed), 3); //! result.is_leader() //! })); //! } diff --git a/asyncband/src/waitgroup/mod.rs b/asyncband/src/waitgroup/mod.rs index ade93b56..d6f7a753 100644 --- a/asyncband/src/waitgroup/mod.rs +++ b/asyncband/src/waitgroup/mod.rs @@ -35,26 +35,21 @@ //! ``` //! # #[tokio::main] //! # async fn main() { -//! use std::sync::Arc; -//! use std::sync::atomic::AtomicBool; -//! use std::sync::atomic::Ordering; -//! //! use asyncband::waitgroup::WaitGroup; //! +//! async fn do_work() {} +//! //! let group = WaitGroup::new(); -//! let completed = Arc::new([const { AtomicBool::new(false) }; 3]); //! -//! for slot in 0..3 { +//! for _ in 0..3 { //! let participant = group.clone(); -//! let completed = completed.clone(); //! tokio::spawn(async move { -//! completed[slot].store(true, Ordering::Relaxed); +//! do_work().await; //! drop(participant); //! }); //! } //! //! group.await; -//! assert!(completed.iter().all(|done| done.load(Ordering::Relaxed))); //! # } //! ```