diff --git a/CHANGELOG.md b/CHANGELOG.md index 14862c84..6d2cf9e3 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 `WaitGroup` completion overhead by sharing allocation ownership across participants and storing the first pending observer's waker inline. + ## v0.8.0 ### Breaking changes 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/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/internal/mod.rs b/asyncband/src/internal/mod.rs index ccaa11ac..29a3e73c 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -76,13 +76,17 @@ 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. #[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/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 b50c115e..d6f7a753 100644 --- a/asyncband/src/waitgroup/mod.rs +++ b/asyncband/src/waitgroup/mod.rs @@ -17,17 +17,18 @@ //! 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. //! -//! 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 //! @@ -39,118 +40,250 @@ //! async fn do_work() {} //! //! let group = WaitGroup::new(); -//! let mut tasks = vec![]; //! //! for _ in 0..3 { //! let participant = group.clone(); -//! tasks.push(tokio::spawn(async move { +//! tokio::spawn(async move { //! do_work().await; -//! participant.await; -//! })); +//! drop(participant); +//! }); //! } //! //! group.await; -//! for task in tasks { -//! task.await.unwrap(); -//! } //! # } //! ``` 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::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 std::task::Waker; -use crate::internal::mutex::Mutex; +use crate::internal::register_waker; +use crate::internal::rwlock::RwLock; use crate::internal::wakerset::WakerSet; use crate::internal::wakerset::WakerToken; -#[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. - handles: AtomicUsize, - waiters: Mutex, +#[derive(Default)] +enum Registration { + #[default] + Unregistered, + Inline, + Overflow(WakerToken), } -impl State { - fn new() -> Self { - Self { - handles: AtomicUsize::new(1), - waiters: Mutex::new(WakerSet::new()), +#[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), } } - 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); + #[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), + ), + } } - fn release_handle(&self) { - // 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); - debug_assert!(previous > 0, "a live handle must own one count"); - if previous != 1 { - return; + #[inline] + fn wake_all(self) { + if let Some(waker) = self.inline { + waker.wake(); + } + if let Some(overflow) = self.overflow { + overflow.wake_all(); } - - let wakers = mem::take(&mut *self.waiters.lock()); - wakers.wake_all(); } +} + +struct State { + // 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, +} - fn poll_wait(&self, token: &mut Option, cx: &mut Context<'_>) -> Poll<()> { - if self.handles.load(Ordering::Acquire) == 0 { - *token = None; +impl State { + #[inline] + fn poll_wait(&self, registration: &mut Registration, cx: &mut Context<'_>) -> Poll<()> { + if self.participants.load(Ordering::Acquire) == 0 { + *registration = Registration::Unregistered; return Poll::Ready(()); } - let mut waiters = self.waiters.lock(); - if self.handles.load(Ordering::Acquire) == 0 { - *token = None; + self.register_wait(registration, cx.waker()) + } + + fn register_wait(&self, registration: &mut Registration, waker: &Waker) -> Poll<()> { + 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 { + *registration = Registration::Unregistered; return Poll::Ready(()); } - let retired_waker = waiters.register(token, cx.waker()); + let retired_waker = waiters.register(registration, waker); drop(waiters); drop(retired_waker); Poll::Pending } - fn unregister(&self, token: &mut Option) { - if token.is_none() { + #[inline] + fn unregister(&self, registration: &mut Registration) { + if matches!(registration, Registration::Unregistered) { + return; + } + if self.participants.load(Ordering::Acquire) == 0 { + *registration = Registration::Unregistered; return; } - 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; + self.unregister_slow(registration); + } + + fn unregister_slow(&self, registration: &mut Registration) { + let mut waiters = self.waiters.write(); + if self.participants.load(Ordering::Acquire) == 0 { + *registration = Registration::Unregistered; return; } - let removed_waker = waiters.unregister(token); + let removed_waker = waiters.unregister(registration); drop(waiters); drop(removed_waker); } } +struct Shared(NonNull, PhantomData); + +// 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: 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 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, PhantomData) + } + + #[inline] + 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 inline non-final participant drop path. + #[inline(never)] + fn finish(&mut self) { + let waiters = mem::take(&mut *self.waiters.write()); + waiters.wake_all(); + self.release_ref(); + } + + #[inline] + fn clone_ref(&self) -> Self { + if self.refs.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { + std::process::abort(); + } + Self(self.0, PhantomData) + } + + #[inline] + 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())) }; + } + } +} + /// A handle representing one participant in a dynamically sized wait group. /// /// 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 live participants share one collective reference in `State::refs`. + state: Shared, } impl fmt::Debug for WaitGroup { @@ -175,9 +308,10 @@ impl WaitGroup { /// /// let wg = WaitGroup::new(); /// ``` + #[inline] pub fn new() -> Self { Self { - state: Some(Arc::new(State::new())), + state: Shared::new(), } } } @@ -187,24 +321,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: self.state.clone_participant(), + } } } impl Drop for WaitGroup { + #[inline] fn drop(&mut self) { - if let Some(state) = self.state.take() { - state.release_handle(); - } + self.state.release_participant(); } } @@ -218,10 +346,14 @@ impl IntoFuture for WaitGroup { /// /// Conversion marks this participant complete immediately. Dropping the returned future does /// not undo that completion; see [`Wait`] for observer cancellation. - 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.clone_ref(); + drop(self); + Wait { + registration: Registration::Unregistered, + state, + } } } @@ -236,8 +368,8 @@ impl IntoFuture for WaitGroup { /// and other observers are unaffected. #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct Wait { - token: Option, - state: Arc, + registration: Registration, + state: Shared, } impl Clone for Wait { @@ -248,10 +380,11 @@ impl Clone for Wait { /// # Cancellation /// /// Dropping the returned future unregisters only that observer; see [`Wait`]. + #[inline] fn clone(&self) -> Self { Wait { - token: None, - state: self.state.clone(), + registration: Registration::Unregistered, + state: self.state.clone_ref(), } } } @@ -265,14 +398,20 @@ 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) + 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.unregister(&mut self.registration); + self.state.release_ref(); } } 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..10eedd8c 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 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 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..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,9 +27,13 @@ use benchmarks::implementations; use benchmarks::support::bench_context; use benchmarks::support::poll_pending; use benchmarks::support::poll_pinned_ready; +use futures_util::FutureExt; +use tokio_util::task::TaskTracker; struct Asyncband; struct WaitgroupRs; +struct WaitgroupRsShared; +struct TokioTaskTracker; trait WaitGroup { type Group; @@ -75,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(); @@ -105,6 +158,44 @@ fn complete_worker_batch(bencher: Bencher, worker_count: usize) { }); } +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); + }); +} + +// 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 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 drain_observers(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| C::drain(&mut context)); +} + pub fn register(runner: &mut Runner) { runner.compare( "waitgroup/complete-waiter", @@ -118,4 +209,21 @@ 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/drain-4-observers", + 2_500_000, + implementations!([Asyncband, WaitgroupRsShared, TokioTaskTracker], |C, b| { + drain_observers::(b) + }), + ); } diff --git a/tests-integration/tests/waitgroup_test.rs b/tests-integration/tests/waitgroup_test.rs index 5bd36178..85bde96d 100644 --- a/tests-integration/tests/waitgroup_test.rs +++ b/tests-integration/tests/waitgroup_test.rs @@ -17,13 +17,23 @@ 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)] async fn waits_for_all_worker_handles() { let wg = WaitGroup::new(); @@ -139,3 +149,167 @@ 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 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 (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!(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() { + 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 poll_with(wait.as_mut(), waker).is_pending() { + thread::yield_now(); + } + assert!(writes.iter().all(|written| written.load(Ordering::Relaxed))); + }); + } + let mut wait = Box::pin(observer); + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + start.wait(); + drop(wait); + }); +} + +#[test] +fn racing_completion_notifies_a_poll_that_returns_pending() { + let group = WaitGroup::new(); + let worker = group.clone(); + let mut wait = Box::pin(group.into_future()); + let start = Barrier::new(2); + let (waker, wakes) = WakeCounter::new(); + + let result = thread::scope(|scope| { + scope.spawn(|| { + start.wait(); + drop(worker); + }); + start.wait(); + poll_with(wait.as_mut(), &waker) + }); + if result.is_pending() { + assert_eq!(wakes.count(), 1); + } + assert!(poll_once(wait.as_mut()).is_ready()); +} + +#[test] +fn retiring_a_registration_can_complete_the_group() { + for replace in [false, true] { + for use_overflow in [false, true] { + 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 = waker_on_drop(move || drop(worker)); + assert!(poll_with(wait.as_mut(), &retired).is_pending()); + drop(retired); + + if replace { + 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()); + }); + } + } +} + +#[test] +fn final_worker_and_cancelled_observer_release_their_registration() { + let group = WaitGroup::new(); + let worker = group.clone(); + let mut wait = Box::pin(group.into_future()); + let (waker, wakes) = WakeCounter::new(); + assert!(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!(wakes.count() <= 1); + assert_eq!(Arc::strong_count(&wakes), 1); +} + +#[test] +fn waking_can_drop_the_last_observer_during_completion() { + let group = WaitGroup::new(); + let worker = group.clone(); + let observer = Arc::new(Mutex::new(Some(Box::pin(group.into_future())))); + let waker = waker_on_wake({ + let observer = observer.clone(); + move || drop(observer.lock().unwrap().take().unwrap()) + }); + + let mut guard = observer.lock().unwrap(); + assert!(poll_with(guard.as_mut().unwrap().as_mut(), &waker).is_pending()); + drop(guard); + drop(waker); + + drop(worker); + assert!(observer.lock().unwrap().is_none()); +} 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"],