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/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/asyncband/src/waitgroup/mod.rs b/asyncband/src/waitgroup/mod.rs index 4fb7c403..b2ee32d1 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,113 @@ 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(), } } - 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] + fn poll_wait(&self, token: &mut Registration, cx: &mut Context<'_>) -> Poll<()> { + if self.handles.load(Ordering::Acquire) == 0 { + *token = Registration::Set(None); + return Poll::Ready(()); + } + + self.waiters.register(token, cx.waker(), &self.handles) } - 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 unregister(&self, token: &mut Registration) { + if !matches!(token, Registration::Set(None)) && self.handles.load(Ordering::Acquire) != 0 { + self.waiters.unregister(token, &self.handles); } + } +} + +// 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); + +// SAFETY: all shared access goes through State, which is Send + Sync. +unsafe impl Send for Shared {} +unsafe impl Sync for Shared {} - let wakers = mem::take(&mut *self.waiters.lock()); - wakers.wake_all(); +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 poll_wait(&self, token: &mut Option, cx: &mut Context<'_>) -> Poll<()> { - if self.handles.load(Ordering::Acquire) == 0 { - *token = None; - return Poll::Ready(()); +impl Shared { + #[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 { + std::process::abort(); } + Self(self.0) + } - let mut waiters = self.waiters.lock(); - if self.handles.load(Ordering::Acquire) == 0 { - *token = None; - return Poll::Ready(()); + #[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(); } + } - let retired_waker = waiters.register(token, cx.waker()); - drop(waiters); - drop(retired_waker); - Poll::Pending + // Keep notification and reclamation off the nonfinal participant drop path. + #[inline(never)] + fn finish(&mut self) { + self.waiters.wake_all(); + self.release_observer(); } - fn unregister(&self, token: &mut Option) { - if token.is_none() { - return; + #[inline] + fn observe(&self) -> Self { + if self.observers.fetch_add(1, Ordering::Relaxed) >= isize::MAX as usize { + std::process::abort(); } + Self(self.0) + } - 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(&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())) }; } - - let removed_waker = waiters.unregister(token); - drop(waiters); - drop(removed_waker); } } @@ -148,9 +177,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 +203,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,24 +216,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.register_handle(), + } } } impl Drop for WaitGroup { + #[inline] fn drop(&mut self) { - if let Some(state) = self.state.take() { - state.release_handle(); - } + self.state.release_handle(); } } @@ -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..b9efc1d2 --- /dev/null +++ b/asyncband/src/waitgroup/waiters.rs @@ -0,0 +1,113 @@ +// 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/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) + }), + ); +} diff --git a/tests-integration/tests/waitgroup_test.rs b/tests-integration/tests/waitgroup_test.rs index 5bd36178..6d7b7d6f 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,181 @@ 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); +} + +#[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()); +} 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"],