From 39d84ef9fd83ea0c9264900917e77e9ff43c7897 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 3 Nov 2026 17:25:51 +0800 Subject: [PATCH 1/4] perf(waitgroup): streamline participant completion and waiter notification --- CHANGELOG.md | 1 + 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, 437 insertions(+), 64 deletions(-) create mode 100644 asyncband/src/waitgroup/waiters.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 64c75081..e8b6b469 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,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 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. ## v0.7.4 (2026-10-25) 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 19121028d400e8d603118e1159af58713b2aa862 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 3 Nov 2026 17:25:52 +0800 Subject: [PATCH 2/4] 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 39f8b9f5..8840d427 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 21509de3..260aeefc 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 e01ccffa47f70bedc18a16a57d63f151c2d9ac47 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 3 Nov 2026 17:50:25 +0800 Subject: [PATCH 3/4] 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 540858bec2fa86c7df18f0ea72195947df9a6729 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 3 Nov 2026 17:58:29 +0800 Subject: [PATCH 4/4] 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()); +}