diff --git a/CHANGELOG.md b/CHANGELOG.md index 087014c..58b46e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ All notable changes to this project will be documented in this file. ### New features * Implement `broadcast::mpmc::bounded`, a lossless bounded broadcast channel that retains at most the requested capacity and makes producers wait for the slowest active receiver. +* Add an opt-in runtime-agnostic `TaskGroup` with cloneable registration handles, explicit closing, completion-ordered output consumption, and whole-group waiting and collection. * Add opt-in bounded and unbounded `asyncband::mpmc` queues with cloneable producers and competing consumers, delivering each accepted value to exactly one receiver while a receiver remains. * Add an opt-in runtime-agnostic `Phaser` with shared observer handles, dynamic RAII participants registered individually or in batches through an owning iterator, `u64` phase numbers, split arrival/wait with cancellation-resilient retries, and a `close` operation that releases unfinished waits with `Closed`. * Add bounded MPSC `reserve` and `try_reserve` methods returning a `Permit`, allowing callers to wait for capacity before constructing a message; pending sends and reservations receive capacity in wait-queue order, and unused permits release capacity without claiming message order. diff --git a/README.md b/README.md index ed3dae3..cf3f161 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a fixed one-way countdown reaches zero. | | | [`Phaser`](https://docs.rs/asyncband/*/asyncband/phaser/struct.Phaser.html) | `phaser` | Coordinate repeated phases with a dynamic participant set. | | | [`WaitGroup`](https://docs.rs/asyncband/*/asyncband/waitgroup/struct.WaitGroup.html) | `waitgroup` | Dynamically register participants and wait until all have completed. | +| | [`TaskGroup`](https://docs.rs/asyncband/*/asyncband/task_group/struct.TaskGroup.html) | `task-group` | Track independently spawned futures and consume their outputs in completion order. | | | [`Shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/struct.Shutdown.html) | `shutdown` | Request shutdown and wait until all completion guards are dropped. | | Work coalescing | [`Once`](https://docs.rs/asyncband/*/asyncband/once/struct.Once.html) | `once` | Complete one asynchronous initialization; cancelled or panicked attempts may be retried. | | | [`OnceCell`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceCell.html) | `once-cell` | Store one value from an access-time initializer; failed, cancelled, or panicked attempts may be retried. | diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index d34a926..c2f0d95 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -65,6 +65,7 @@ rwlock = [] semaphore = [] shutdown = ["latch", "waitgroup"] singleflight = ["dep:hashbrown", "once-cell"] +task-group = [] waitgroup = [] watch = [] diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 7998e25..0c81e88 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -90,6 +90,7 @@ pub(crate) mod value_cell; feature = "phaser", feature = "rwlock", feature = "semaphore", + feature = "task-group", feature = "waitgroup", feature = "watch", ))] diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index afd6faf..caf89ce 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -65,6 +65,7 @@ //! | | [`Latch`](latch::Latch) | `latch` | Wait until a fixed one-way countdown reaches zero. | //! | | [`Phaser`](phaser::Phaser) | `phaser` | Coordinate repeated phases with a dynamic participant set. | //! | | [`WaitGroup`](waitgroup::WaitGroup) | `waitgroup` | Dynamically register participants and wait until all have completed. | +//! | | [`TaskGroup`](task_group::TaskGroup) | `task-group` | Track independently spawned futures and consume their outputs in completion order. | //! | | [`Shutdown`](shutdown::Shutdown) | `shutdown` | Request shutdown and wait until all completion guards are dropped. | //! | Work coalescing | [`Once`](once::Once) | `once` | Complete one asynchronous initialization; cancelled or panicked attempts may be retried. | //! | | [`OnceCell`](once::OnceCell) | `once-cell` | Store one value from an access-time initializer; failed, cancelled, or panicked attempts may be retried. | @@ -161,6 +162,8 @@ pub mod semaphore; pub mod shutdown; #[cfg(feature = "singleflight")] pub mod singleflight; +#[cfg(feature = "task-group")] +pub mod task_group; #[cfg(feature = "waitgroup")] pub mod waitgroup; #[cfg(feature = "watch")] @@ -168,6 +171,11 @@ pub mod watch; #[cfg(all( test, - any(feature = "once-map", feature = "phaser", feature = "singleflight") + any( + feature = "once-map", + feature = "phaser", + feature = "singleflight", + feature = "task-group" + ) ))] mod test_support; diff --git a/asyncband/src/task_group/mod.rs b/asyncband/src/task_group/mod.rs new file mode 100644 index 0000000..1bb6eca --- /dev/null +++ b/asyncband/src/task_group/mod.rs @@ -0,0 +1,703 @@ +// 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. + +//! Track independently spawned futures and consume their outputs as they complete. +//! +//! A [`TaskGroup`] owns one queue that stores outputs in the order the futures finish. Its +//! cloneable [`Registrar`] wraps futures before the caller spawns them on an executor. Asyncband +//! never chooses an executor or spawns a task itself. +//! +//! [`close`](TaskGroup::close) prevents further registration. Once the group is closed, every +//! [`Tracked`] wrapper has been dropped, and every queued output has been consumed, +//! [`join_next`](TaskGroup::join_next) returns `None`. Requiring `close` lets `join_next` tell an +//! open group with no current work from a group that will never receive more work. +//! +//! # Example +//! +//! ``` +//! # #[tokio::main] +//! # async fn main() { +//! use asyncband::task_group::TaskGroup; +//! +//! let (mut group, registrar) = TaskGroup::new(); +//! let first = tokio::spawn(registrar.track(async { 21 }).unwrap()); +//! let second = tokio::spawn(registrar.track(async { 2 }).unwrap()); +//! group.close(); +//! +//! let mut outputs = group.join().await; +//! outputs.sort_unstable(); +//! assert_eq!(outputs, [2, 21]); +//! first.await.unwrap(); +//! second.await.unwrap(); +//! # } +//! ``` +//! +//! # Task lifetime and failure +//! +//! A [`Tracked`] future remains active until its wrapper is dropped. On normal completion, its +//! output is sent before it returns [`Poll::Ready`], but the group does not finish until the +//! wrapper and its inner future are dropped. Aborted and panicking tasks produce no output and are +//! reported only by the caller's executor. +//! +//! Dropping the [`TaskGroup`] discards queued and future outputs and makes all registrars reject +//! new work. It does not cancel or abort tracked futures. Callers can wrap each future with their +//! own cancellation mechanism. +//! +//! The output queue can grow without limit until the owner consumes it or is dropped. Call +//! [`join_next`](TaskGroup::join_next) continuously if futures may finish faster than the owner can +//! consume their outputs. + +use std::any::type_name; +use std::collections::VecDeque; +use std::fmt; +use std::future::Future; +use std::mem; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Weak; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::sync::atomic::fence; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use crate::internal::mutex::Mutex; + +#[cfg(test)] +mod tests; + +/// The single owner of outputs from a dynamically registered group of futures. +/// +/// Create a group and its first [`Registrar`] with [`TaskGroup::new`]. The owner is deliberately +/// not cloneable: only one task may consume outputs in the order the futures finish. Methods that +/// consume outputs require mutable access, so only one join operation can wait at a time. +pub struct TaskGroup { + shared: Arc>, +} + +impl fmt::Debug for TaskGroup { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let state = self.shared.state.lock(); + f.debug_struct("TaskGroup") + .field("active", &self.shared.lifecycle.active_tasks()) + .field("closed", &self.shared.lifecycle.is_closed()) + .field("queued_outputs", &state.outputs.len()) + .finish_non_exhaustive() + } +} + +impl Drop for TaskGroup { + fn drop(&mut self) { + let (outputs, waiter) = self.shared.drop_owner(); + // Dropping an output or waker may use a registrar, so do it after releasing the state lock. + drop((outputs, waiter)); + } +} + +impl TaskGroup { + /// Creates an open task group and its first registration handle. + pub fn new() -> (Self, Registrar) { + let shared = Arc::new(Shared { + lifecycle: Lifecycle::new(), + state: Mutex::new(State { + owner_alive: true, + discard_outputs: false, + outputs: VecDeque::new(), + waiter: None, + }), + }); + let registrar = Registrar { + shared: Arc::downgrade(&shared), + }; + (Self { shared }, registrar) + } + + /// Returns another handle for registering futures in this group. + pub fn registrar(&self) -> Registrar { + Registrar { + shared: Arc::downgrade(&self.shared), + } + } + + /// Returns whether this group rejects new registrations. + pub fn is_closed(&self) -> bool { + self.shared.is_closed() + } + + /// Permanently prevents new futures from being registered. + /// + /// Existing tracked futures continue running. Calling this method more than once has no + /// additional effect. Joining an idle open group remains pending, so callers must close the + /// group when no more work can be registered. + /// + /// # Panics + /// + /// Panics if waking a pending join operation panics. The group remains closed. + pub fn close(&self) { + if let Some(waker) = self.shared.close() { + waker.wake(); + } + } + + /// Returns the next normally completed output, in completion order. + /// + /// Returns `None` only after the group is closed, all [`Tracked`] wrappers have been dropped, + /// and all earlier outputs have been consumed. A tracked future that is dropped, aborted, or + /// dropped after a panic produces no output. + /// + /// An output can be returned before its [`Tracked`] wrapper is dropped; final completion still + /// waits for that drop. + /// + /// Canceling this operation while it waits does not consume an output. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::task_group::TaskGroup; + /// use tokio::sync::oneshot; + /// + /// let (mut group, registrar) = TaskGroup::new(); + /// let (send, receive) = oneshot::channel(); + /// let later = tokio::spawn( + /// registrar + /// .track(async move { + /// receive.await.unwrap(); + /// "later" + /// }) + /// .unwrap(), + /// ); + /// let first = tokio::spawn(registrar.track(async { "first" }).unwrap()); + /// group.close(); + /// + /// assert_eq!(group.join_next().await, Some("first")); + /// send.send(()).unwrap(); + /// assert_eq!(group.join_next().await, Some("later")); + /// assert_eq!(group.join_next().await, None); + /// first.await.unwrap(); + /// later.await.unwrap(); + /// # } + /// ``` + pub fn join_next(&mut self) -> JoinNext<'_, T> { + JoinNext { + group: self, + registered: false, + } + } + + /// Waits for the group to finish and discards every output. + /// + /// This operation does not close the group. It remains pending while the group is open, even + /// when no futures are active. If it is canceled, outputs discarded by this call cannot be + /// recovered by a later join. + /// + /// Completed [`Tracked`] wrappers must be dropped before this returns. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::task_group::TaskGroup; + /// + /// let (mut group, registrar) = TaskGroup::new(); + /// let task = tokio::spawn(registrar.track(async { 42 }).unwrap()); + /// group.close(); + /// + /// group.wait().await; + /// assert_eq!(group.join_next().await, None); + /// task.await.unwrap(); + /// # } + /// ``` + pub async fn wait(&mut self) { + let (_discarding, outputs) = DiscardOutputs::new(self.shared.clone()); + drop(outputs); + while self.join_next().await.is_some() {} + } + + /// Waits for the group to finish and collects its outputs in completion order. + /// + /// This operation does not close the group. If it is canceled, outputs already collected by + /// this call are dropped and cannot be recovered by a later join. + /// + /// Completed [`Tracked`] wrappers must be dropped before this returns. + pub async fn join(&mut self) -> Vec { + let (queued, capacity) = self.shared.take_outputs_and_capacity_hint(); + let additional_capacity = capacity.saturating_sub(queued.len()); + let mut outputs = Vec::from(queued); + outputs.reserve(additional_capacity); + while let Some(output) = self.join_next().await { + outputs.push(output); + outputs.extend(self.shared.take_outputs()); + } + outputs + } +} + +/// A cloneable handle that registers futures in one [`TaskGroup`]. +/// +/// The handle does not keep the group alive. Registration fails after the owner is dropped or +/// [`TaskGroup::close`] is called. +pub struct Registrar { + shared: Weak>, +} + +impl Clone for Registrar { + fn clone(&self) -> Self { + Self { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for Registrar { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Registrar").finish_non_exhaustive() + } +} + +impl Registrar { + /// Returns whether this handle can no longer register new futures. + pub fn is_closed(&self) -> bool { + let Some(shared) = self.shared.upgrade() else { + return true; + }; + shared.is_closed() + } + + /// Registers `future` immediately and returns a wrapper for the caller to poll or spawn. + /// + /// Dropping the returned wrapper before it completes releases its registration without + /// producing an output. If registration fails, the error returns ownership of `future`. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::task_group::TaskGroup; + /// + /// let (group, registrar) = TaskGroup::new(); + /// group.close(); + /// + /// let future = async { 42 }; + /// let Err(error) = registrar.track(future) else { + /// panic!("a closed group accepted a future"); + /// }; + /// assert_eq!(error.into_inner().await, 42); + /// # } + /// ``` + pub fn track(&self, future: F) -> Result, TrackError> + where + F: Future, + { + let Some(shared) = self.shared.upgrade() else { + return Err(TrackError(future)); + }; + if !shared.register() { + return Err(TrackError(future)); + } + Ok(Tracked { + future, + registration: Registration { + shared, + completed: false, + }, + }) + } +} + +/// A future whose lifetime and successful output are tracked by a [`TaskGroup`]. +/// +/// On normal completion, this sends the inner future's output to the group and returns `()`. Its +/// registration remains active until the wrapper is dropped; dropping it earlier sends no output. +#[must_use = "a tracked future must be awaited, spawned, or dropped to release its registration"] +pub struct Tracked +where + F: Future, +{ + future: F, + // Fields are dropped from top to bottom, so the inner future is dropped before the active task + // count is decreased. + registration: Registration, +} + +impl fmt::Debug for Tracked +where + F: Future + fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Tracked") + .field("future", &self.future) + .finish_non_exhaustive() + } +} + +impl Future for Tracked +where + F: Future, +{ + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + assert!( + !self.as_ref().get_ref().registration.completed, + "a tracked future cannot be polled after completion" + ); + // SAFETY: This manually projects the outer pin onto `future`. Once `Tracked` is pinned, + // `future` stays at the same address until it is dropped. + let future = unsafe { self.as_mut().map_unchecked_mut(|this| &mut this.future) }; + let Poll::Ready(output) = future.poll(cx) else { + return Poll::Pending; + }; + + // SAFETY: The pin applies only to `future`; accessing `registration` does not move it. + unsafe { self.get_unchecked_mut() } + .registration + .complete(output); + Poll::Ready(()) + } +} + +/// A future returned by [`TaskGroup::join_next`]. +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct JoinNext<'a, T> { + group: &'a mut TaskGroup, + registered: bool, +} + +impl fmt::Debug for JoinNext<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("JoinNext").finish_non_exhaustive() + } +} + +impl Future for JoinNext<'_, T> { + type Output = Option; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let result = this.group.shared.poll_join_next(cx); + this.registered = result.is_pending(); + result + } +} + +impl Drop for JoinNext<'_, T> { + fn drop(&mut self) { + if self.registered { + self.group.shared.unregister_waiter(); + } + } +} + +/// A task could not be tracked because its group was closed or dropped. +/// +/// The error retains the future so the caller can recover it with [`into_inner`](Self::into_inner). +#[derive(Clone, PartialEq, Eq)] +pub struct TrackError(F); + +impl TrackError { + /// Returns a reference to the future that was not tracked. + pub fn as_inner(&self) -> &F { + &self.0 + } + + /// Consumes the error and returns the future that was not tracked. + pub fn into_inner(self) -> F { + self.0 + } +} + +impl fmt::Display for TrackError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("tracking a future in a closed task group") + } +} + +impl fmt::Debug for TrackError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "TrackError<{}>(..)", type_name::()) + } +} + +impl std::error::Error for TrackError {} + +struct Shared { + lifecycle: Lifecycle, + state: Mutex>, +} + +struct State { + owner_alive: bool, + discard_outputs: bool, + outputs: VecDeque, + waiter: Option, +} + +impl Shared { + fn register(&self) -> bool { + self.lifecycle.try_register() + } + + fn publish(&self, output: T) -> (Option, Option) { + let mut state = self.state.lock(); + let discarded = if state.owner_alive && !state.discard_outputs { + state.outputs.push_back(output); + None + } else { + Some(output) + }; + let waker = if state.owner_alive && !state.discard_outputs { + state.waiter.take() + } else { + None + }; + (discarded, waker) + } + + fn release_registration(&self) -> Option { + if !self.lifecycle.retire_task() { + return None; + } + + let mut state = self.state.lock(); + if state.owner_alive && state.outputs.is_empty() { + state.waiter.take() + } else { + None + } + } + + fn poll_join_next(&self, cx: &mut Context<'_>) -> Poll> { + let (poll, waiter) = { + let mut state = self.state.lock(); + if let Some(output) = state.outputs.pop_front() { + (Poll::Ready(Some(output)), state.waiter.take()) + } else if self.lifecycle.is_finished() { + (Poll::Ready(None), state.waiter.take()) + } else if state + .waiter + .as_ref() + .is_some_and(|waker| waker.will_wake(cx.waker())) + { + (Poll::Pending, None) + } else { + let waiter = state.waiter.replace(cx.waker().clone()); + (Poll::Pending, waiter) + } + }; + drop(waiter); + poll + } + + fn unregister_waiter(&self) { + let waiter = self.state.lock().waiter.take(); + drop(waiter); + } + + fn close(&self) -> Option { + if !self.lifecycle.close() { + return None; + } + + let mut state = self.state.lock(); + if state.outputs.is_empty() { + state.waiter.take() + } else { + None + } + } + + fn drop_owner(&self) -> (VecDeque, Option) { + let _ = self.lifecycle.close(); + let mut state = self.state.lock(); + state.owner_alive = false; + (mem::take(&mut state.outputs), state.waiter.take()) + } + + fn is_closed(&self) -> bool { + self.lifecycle.is_closed() + } + + fn take_outputs_and_capacity_hint(&self) -> (VecDeque, usize) { + let (outputs, capacity, waiter) = { + let mut state = self.state.lock(); + let active_tasks = self.lifecycle.active_tasks(); + let capacity = state.outputs.len().saturating_add(active_tasks); + let outputs = mem::take(&mut state.outputs); + let waiter = if outputs.is_empty() { + None + } else { + state.waiter.take() + }; + (outputs, capacity, waiter) + }; + drop(waiter); + (outputs, capacity) + } + + fn take_outputs(&self) -> VecDeque { + let (outputs, waiter) = { + let mut state = self.state.lock(); + let outputs = mem::take(&mut state.outputs); + let waiter = if outputs.is_empty() { + None + } else { + state.waiter.take() + }; + (outputs, waiter) + }; + drop(waiter); + outputs + } + + fn begin_discarding_outputs(&self) -> VecDeque { + let (outputs, waiter) = { + let mut state = self.state.lock(); + state.discard_outputs = true; + (mem::take(&mut state.outputs), state.waiter.take()) + }; + drop(waiter); + outputs + } + + fn end_discarding_outputs(&self) { + self.state.lock().discard_outputs = false; + } +} + +struct Lifecycle(AtomicUsize); + +impl Lifecycle { + // The high bit closes registration; the remaining bits count active tracked futures. Keeping + // both in one atomic gives registration and close a single, unambiguous order without locking + // the output queue. + const CLOSED_BIT: usize = 1 << (usize::BITS - 1); + const ACTIVE_TASKS_MASK: usize = Self::CLOSED_BIT - 1; + + const fn new() -> Self { + Self(AtomicUsize::new(0)) + } + + fn active_tasks(&self) -> usize { + // This value is used only for debug output and as an approximate vector size. + self.0.load(Ordering::Relaxed) & Self::ACTIVE_TASKS_MASK + } + + fn is_closed(&self) -> bool { + // This only reads the closed bit and does not need to make any other memory visible. + self.0.load(Ordering::Relaxed) & Self::CLOSED_BIT != 0 + } + + fn is_finished(&self) -> bool { + let current = self.0.load(Ordering::Acquire); + current & Self::CLOSED_BIT != 0 && current & Self::ACTIVE_TASKS_MASK == 0 + } + + fn try_register(&self) -> bool { + // The compare-exchange decides whether registration or close happened first. It changes + // only this count and passes no other data between threads, so Relaxed ordering is enough. + let mut current = self.0.load(Ordering::Relaxed); + loop { + if current & Self::CLOSED_BIT != 0 { + return false; + } + assert_ne!( + current, + Self::ACTIVE_TASKS_MASK, + "a task group cannot track more than isize::MAX futures" + ); + + match self.0.compare_exchange_weak( + current, + current + 1, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return true, + Err(actual) => current = actual, + } + } + } + + fn retire_task(&self) -> bool { + // A Release decrement makes this task's earlier work visible when the group finishes. The + // final task uses an Acquire fence so it also sees work released by the earlier tasks + // before waking the owner. + let previous = self.0.fetch_sub(1, Ordering::Release); + let active_tasks = previous & Self::ACTIVE_TASKS_MASK; + debug_assert!(active_tasks > 0, "a tracked future owns one active count"); + let finished = previous & Self::CLOSED_BIT != 0 && active_tasks == 1; + if finished { + fence(Ordering::Acquire); + } + finished + } + + fn close(&self) -> bool { + // If the group is already empty, Acquire makes the finished tasks' earlier work visible + // here. Otherwise, Release lets the final task see that the group was closed before it + // wakes the owner. + let previous = self.0.fetch_or(Self::CLOSED_BIT, Ordering::AcqRel); + previous & Self::CLOSED_BIT == 0 && previous & Self::ACTIVE_TASKS_MASK == 0 + } +} + +struct Registration { + shared: Arc>, + completed: bool, +} + +impl Registration { + fn complete(&mut self, output: T) { + let (discarded, waker) = self.shared.publish(output); + // Mark completion before waking or dropping output because either may panic. + self.completed = true; + if let Some(waker) = waker { + waker.wake(); + } + drop(discarded); + } +} + +impl Drop for Registration { + fn drop(&mut self) { + if let Some(waker) = self.shared.release_registration() { + waker.wake(); + } + } +} + +struct DiscardOutputs { + shared: Arc>, +} + +impl DiscardOutputs { + fn new(shared: Arc>) -> (Self, VecDeque) { + let outputs = shared.begin_discarding_outputs(); + (Self { shared }, outputs) + } +} + +impl Drop for DiscardOutputs { + fn drop(&mut self) { + self.shared.end_discarding_outputs(); + } +} diff --git a/asyncband/src/task_group/tests.rs b/asyncband/src/task_group/tests.rs new file mode 100644 index 0000000..f22e521 --- /dev/null +++ b/asyncband/src/task_group/tests.rs @@ -0,0 +1,435 @@ +// 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::future; +use std::future::Future; +use std::marker::PhantomPinned; +use std::mem; +use std::panic; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Wake; +use std::task::Waker; + +use super::Registrar; +use super::TaskGroup; +use crate::test_support::poll_once; + +struct CountWake(AtomicUsize); + +impl Wake for CountWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +struct PinnedReady { + output: usize, + _pin: PhantomPinned, +} + +struct DropFlagFuture(Arc); + +struct ReadyDropFlagFuture(Arc); + +impl Future for DropFlagFuture { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + Poll::Pending + } +} + +impl Drop for DropFlagFuture { + fn drop(&mut self) { + self.0.store(true, Ordering::Relaxed); + } +} + +impl Future for ReadyDropFlagFuture { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + Poll::Ready(()) + } +} + +impl Drop for ReadyDropFlagFuture { + fn drop(&mut self) { + self.0.store(true, Ordering::Relaxed); + } +} + +struct AssertDroppedWake { + dropped: Arc, + wakes: AtomicUsize, +} + +impl Wake for AssertDroppedWake { + fn wake(self: Arc) { + assert!(self.dropped.load(Ordering::Relaxed)); + self.wakes.fetch_add(1, Ordering::Relaxed); + } +} + +struct PanicOnDropFuture; + +struct PanicOnDropOutput; + +impl Future for PanicOnDropFuture { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + Poll::Pending + } +} + +impl Drop for PanicOnDropFuture { + fn drop(&mut self) { + panic!("future drop failed"); + } +} + +impl Drop for PanicOnDropOutput { + fn drop(&mut self) { + panic!("output drop failed"); + } +} + +impl Future for PinnedReady { + type Output = usize; + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + Poll::Ready(self.output) + } +} + +#[test] +fn task_group_and_registrar_default_to_unit_outputs() { + let (group, registrar): (TaskGroup, Registrar) = TaskGroup::new(); + + assert!(!group.is_closed()); + assert!(!registrar.is_closed()); +} + +#[test] +fn outputs_follow_completion_order() { + let (mut group, registrar) = TaskGroup::new(); + let mut first = Box::pin(registrar.track(future::ready("first")).unwrap()); + let mut second = Box::pin(registrar.track(future::ready("second")).unwrap()); + + assert_eq!(poll_once(second.as_mut()), Poll::Ready(())); + assert_eq!(poll_once(first.as_mut()), Poll::Ready(())); + group.close(); + + let mut next = Box::pin(group.join_next()); + assert_eq!(poll_once(next.as_mut()), Poll::Ready(Some("second"))); + drop(next); + let mut next = Box::pin(group.join_next()); + assert_eq!(poll_once(next.as_mut()), Poll::Ready(Some("first"))); + drop(next); + let mut next = Box::pin(group.join_next()); + assert_eq!(poll_once(next.as_mut()), Poll::Pending); + drop(second); + drop(first); + assert_eq!(poll_once(next.as_mut()), Poll::Ready(None)); +} + +#[test] +fn join_collects_queued_and_later_outputs() { + let (mut group, registrar) = TaskGroup::new(); + let mut first = Box::pin(registrar.track(future::ready(1)).unwrap()); + let mut second = Box::pin(registrar.track(future::ready(2)).unwrap()); + let mut third = Box::pin(registrar.track(future::ready(3)).unwrap()); + assert_eq!(poll_once(first.as_mut()), Poll::Ready(())); + group.close(); + + let mut join = Box::pin(group.join()); + assert!(poll_once(join.as_mut()).is_pending()); + assert_eq!(poll_once(second.as_mut()), Poll::Ready(())); + assert_eq!(poll_once(third.as_mut()), Poll::Ready(())); + drop(first); + drop(second); + drop(third); + + let Poll::Ready(outputs) = poll_once(join.as_mut()) else { + panic!("the last task completion must finish the join"); + }; + assert_eq!(outputs, [1, 2, 3]); + assert!(outputs.capacity() >= 3); +} + +#[test] +fn dropping_the_last_tracked_future_finishes_a_closed_group() { + let (mut group, registrar) = TaskGroup::<()>::new(); + let tracked = registrar.track(future::pending()).unwrap(); + group.close(); + + let counter = Arc::new(CountWake(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let mut next = Box::pin(group.join_next()); + assert!( + next.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(tracked); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + assert_eq!(poll_once(next.as_mut()), Poll::Ready(None)); +} + +#[test] +fn completed_future_is_dropped_before_wait_finishes() { + let (mut group, registrar) = TaskGroup::new(); + let dropped = Arc::new(AtomicBool::new(false)); + let mut tracked = Box::pin( + registrar + .track(ReadyDropFlagFuture(dropped.clone())) + .unwrap(), + ); + group.close(); + + assert_eq!(poll_once(tracked.as_mut()), Poll::Ready(())); + + let counter = Arc::new(AssertDroppedWake { + dropped: dropped.clone(), + wakes: AtomicUsize::new(0), + }); + let waker = Waker::from(counter.clone()); + let mut wait = Box::pin(group.wait()); + assert!( + wait.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + + drop(tracked); + assert_eq!(counter.wakes.load(Ordering::Relaxed), 1); + assert_eq!(poll_once(wait.as_mut()), Poll::Ready(())); +} + +#[test] +fn cancellation_destroys_the_inner_future_before_notifying_the_group() { + let (mut group, registrar) = TaskGroup::new(); + let dropped = Arc::new(AtomicBool::new(false)); + let tracked = registrar.track(DropFlagFuture(dropped.clone())).unwrap(); + group.close(); + + let counter = Arc::new(AssertDroppedWake { + dropped, + wakes: AtomicUsize::new(0), + }); + let waker = Waker::from(counter.clone()); + let mut join = Box::pin(group.join_next()); + assert!( + join.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + + drop(tracked); + assert_eq!(counter.wakes.load(Ordering::Relaxed), 1); + assert_eq!(poll_once(join.as_mut()), Poll::Ready(None)); +} + +#[test] +fn a_panicking_future_destructor_still_releases_the_registration() { + let (mut group, registrar) = TaskGroup::new(); + let tracked = registrar.track(PanicOnDropFuture).unwrap(); + group.close(); + + assert!(panic::catch_unwind(|| drop(tracked)).is_err()); + let mut join = Box::pin(group.join_next()); + assert_eq!(poll_once(join.as_mut()), Poll::Ready(None)); +} + +#[test] +fn cancelling_join_next_unregisters_its_waker() { + let (mut group, _registrar) = TaskGroup::<()>::new(); + let mut next = Box::pin(group.join_next()); + assert!(poll_once(next.as_mut()).is_pending()); + drop(next); + + assert!(group.shared.state.lock().waiter.is_none()); +} + +#[test] +fn wait_discards_outputs_without_intermediate_wakes() { + let (mut group, registrar) = TaskGroup::new(); + let mut first = Box::pin(registrar.track(future::ready(1)).unwrap()); + let mut second = Box::pin(registrar.track(future::ready(2)).unwrap()); + let shared = group.shared.clone(); + group.close(); + + let counter = Arc::new(CountWake(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let mut wait = Box::pin(group.wait()); + assert!( + wait.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + assert!(shared.state.lock().discard_outputs); + + assert_eq!(poll_once(first.as_mut()), Poll::Ready(())); + assert_eq!(counter.0.load(Ordering::Relaxed), 0); + assert!(shared.state.lock().outputs.is_empty()); + + assert_eq!(poll_once(second.as_mut()), Poll::Ready(())); + assert_eq!(counter.0.load(Ordering::Relaxed), 0); + drop(first); + assert_eq!(counter.0.load(Ordering::Relaxed), 0); + drop(second); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + assert_eq!(poll_once(wait.as_mut()), Poll::Ready(())); + assert!(!shared.state.lock().discard_outputs); +} + +#[test] +fn cancelling_wait_restores_result_collection() { + let (mut group, registrar) = TaskGroup::new(); + let mut discarded = Box::pin(registrar.track(future::ready(1)).unwrap()); + let mut retained = Box::pin(registrar.track(future::ready(2)).unwrap()); + let shared = group.shared.clone(); + + let mut wait = Box::pin(group.wait()); + assert!(poll_once(wait.as_mut()).is_pending()); + assert_eq!(poll_once(discarded.as_mut()), Poll::Ready(())); + drop(wait); + { + let state = shared.state.lock(); + assert!(!state.discard_outputs); + assert!(state.waiter.is_none()); + assert!(state.outputs.is_empty()); + } + + assert_eq!(poll_once(retained.as_mut()), Poll::Ready(())); + drop(discarded); + drop(retained); + group.close(); + let mut next = Box::pin(group.join_next()); + assert_eq!(poll_once(next.as_mut()), Poll::Ready(Some(2))); + drop(next); + let mut next = Box::pin(group.join_next()); + assert_eq!(poll_once(next.as_mut()), Poll::Ready(None)); +} + +#[test] +#[cfg_attr(miri, ignore = "intentionally leaks a polled future")] +fn forgetting_wait_does_not_prevent_another_wait() { + let (mut group, registrar) = TaskGroup::new(); + let mut tracked = Box::pin(registrar.track(future::ready(())).unwrap()); + let shared = group.shared.clone(); + group.close(); + + let first_counter = Arc::new(CountWake(AtomicUsize::new(0))); + let first_waker = Waker::from(first_counter.clone()); + let mut first_wait = Box::pin(group.wait()); + assert!( + first_wait + .as_mut() + .poll(&mut Context::from_waker(&first_waker)) + .is_pending() + ); + mem::forget(first_wait); + + let second_counter = Arc::new(CountWake(AtomicUsize::new(0))); + let second_waker = Waker::from(second_counter.clone()); + let mut second_wait = Box::pin(group.wait()); + assert!( + second_wait + .as_mut() + .poll(&mut Context::from_waker(&second_waker)) + .is_pending() + ); + + assert_eq!(poll_once(tracked.as_mut()), Poll::Ready(())); + assert_eq!(first_counter.0.load(Ordering::Relaxed), 0); + assert_eq!(second_counter.0.load(Ordering::Relaxed), 0); + drop(tracked); + assert_eq!(first_counter.0.load(Ordering::Relaxed), 0); + assert_eq!(second_counter.0.load(Ordering::Relaxed), 1); + assert_eq!(poll_once(second_wait.as_mut()), Poll::Ready(())); + assert!(!shared.state.lock().discard_outputs); +} + +#[test] +fn panicking_discarded_output_does_not_prevent_the_final_wake() { + let (mut group, registrar) = TaskGroup::new(); + let mut tracked = Box::pin(registrar.track(future::ready(PanicOnDropOutput)).unwrap()); + group.close(); + + let counter = Arc::new(CountWake(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let mut wait = Box::pin(group.wait()); + assert!( + wait.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + + assert!(panic::catch_unwind(panic::AssertUnwindSafe(|| poll_once(tracked.as_mut()))).is_err()); + assert_eq!(counter.0.load(Ordering::Relaxed), 0); + drop(tracked); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + assert_eq!(poll_once(wait.as_mut()), Poll::Ready(())); +} + +#[test] +fn close_rejects_registration_and_returns_the_future() { + let (group, registrar) = TaskGroup::new(); + group.close(); + + let future = future::ready(42); + let error = registrar.track(future).unwrap_err(); + assert_eq!(error.into_inner().into_inner(), 42); + assert!(group.is_closed()); + assert!(registrar.is_closed()); +} + +#[test] +fn dropping_the_owner_rejects_registration_without_stopping_tracked_futures() { + let (group, registrar) = TaskGroup::new(); + let mut tracked = Box::pin(registrar.track(future::ready(42)).unwrap()); + drop(group); + + assert!(registrar.track(future::ready(7)).is_err()); + assert_eq!(poll_once(tracked.as_mut()), Poll::Ready(())); +} + +#[test] +fn tracked_supports_non_unpin_futures() { + let (mut group, registrar) = TaskGroup::new(); + let mut tracked = Box::pin( + registrar + .track(PinnedReady { + output: 42, + _pin: PhantomPinned, + }) + .unwrap(), + ); + group.close(); + + assert_eq!(poll_once(tracked.as_mut()), Poll::Ready(())); + let mut next = Box::pin(group.join_next()); + assert_eq!(poll_once(next.as_mut()), Poll::Ready(Some(42))); +} diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index bcab8f2..ac1cbb8 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -46,6 +46,7 @@ asyncband = { workspace = true, features = [ "semaphore", "shutdown", "singleflight", + "task-group", "waitgroup", "watch", ] } diff --git a/benchmarks/asyncband/main.rs b/benchmarks/asyncband/main.rs index a389c0a..2909dc8 100644 --- a/benchmarks/asyncband/main.rs +++ b/benchmarks/asyncband/main.rs @@ -39,6 +39,7 @@ mod semaphore; mod shutdown; mod singleflight; mod support; +mod task_group; mod waitgroup; fn main() { diff --git a/benchmarks/asyncband/task_group/mod.rs b/benchmarks/asyncband/task_group/mod.rs new file mode 100644 index 0000000..eac66ca --- /dev/null +++ b/benchmarks/asyncband/task_group/mod.rs @@ -0,0 +1,250 @@ +// 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::future::Future; +use std::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Wake; +use std::task::Waker; + +use asyncband::task_group::TaskGroup; +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use crate::support::bench_context; +use crate::support::poll_pending; +use crate::support::poll_pinned_ready; +use crate::support::poll_ready; + +const BATCH_SIZES: &[usize] = &[1, 8, 32, 128]; +const THREAD_COUNTS: &[usize] = &[1, 2, 8, 32]; +const CONTENDED_SAMPLE_SIZE: u32 = 256; + +struct CountWake(AtomicUsize); + +impl Wake for CountWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +#[divan::bench(threads = THREAD_COUNTS, sample_size = CONTENDED_SAMPLE_SIZE)] +fn is_closed_contended(bencher: Bencher) { + let (group, _registrar) = TaskGroup::<()>::new(); + bencher.bench(|| black_box(group.is_closed())); +} + +#[divan::bench(threads = THREAD_COUNTS, sample_size = CONTENDED_SAMPLE_SIZE)] +fn register_then_drop_contended(bencher: Bencher) { + let (group, registrar) = TaskGroup::new(); + bencher.bench(|| drop(black_box(registrar.track(async {}).unwrap()))); + black_box(group); +} + +// Every thread completes tasks into the same group. Returning `()` avoids allocating output data, +// so the result mainly shows the cost of updating one shared count and output queue. +#[divan::bench(threads = THREAD_COUNTS, sample_size = CONTENDED_SAMPLE_SIZE)] +fn complete_contended(bencher: Bencher) { + let (group, registrar) = TaskGroup::new(); + bencher.bench(|| { + let task = registrar.track(std::future::ready(())).unwrap(); + let mut task = pin!(task); + let mut context = Context::from_waker(Waker::noop()); + black_box(task.as_mut().poll(&mut context)) + }); + black_box(group); +} + +// Keep `wait()` pending while every thread completes tasks returning `()`. In this mode the group +// discards each output instead of adding it to the queue. +#[divan::bench(threads = THREAD_COUNTS, sample_size = CONTENDED_SAMPLE_SIZE)] +fn complete_discarded_contended(bencher: Bencher) { + let (mut group, registrar) = TaskGroup::new(); + let mut context = bench_context(); + let wait = group.wait(); + let mut wait = pin!(wait); + poll_pending(wait.as_mut(), &mut context); + + bencher.bench(|| { + let task = registrar.track(std::future::ready(())).unwrap(); + let mut task = pin!(task); + let mut context = Context::from_waker(Waker::noop()); + black_box(task.as_mut().poll(&mut context)) + }); +} + +#[divan::bench(threads = THREAD_COUNTS, sample_size = CONTENDED_SAMPLE_SIZE)] +fn reject_registration_contended(bencher: Bencher) { + let (group, registrar) = TaskGroup::<()>::new(); + group.close(); + bencher.bench(|| black_box(registrar.track(async {}).is_err())); +} + +#[divan::bench] +fn complete_then_join(bencher: Bencher) { + let mut context = bench_context(); + let (mut group, registrar) = TaskGroup::new(); + + bencher.bench_local(|| { + let task = registrar.track(async { black_box(1usize) }).unwrap(); + poll_ready(task, &mut context); + black_box(poll_ready(group.join_next(), &mut context).unwrap()) + }); +} + +#[divan::bench] +fn wake_pending_join(bencher: Bencher) { + let mut context = bench_context(); + let (mut group, registrar) = TaskGroup::new(); + + bencher.bench_local(|| { + let task = registrar.track(async { black_box(1usize) }).unwrap(); + let mut join = pin!(group.join_next()); + poll_pending(join.as_mut(), &mut context); + + poll_ready(task, &mut context); + black_box(poll_pinned_ready(join.as_mut(), &mut context).unwrap()) + }); +} + +#[divan::bench(args = BATCH_SIZES)] +fn complete_batch_then_join(bencher: Bencher, task_count: usize) { + let mut context = bench_context(); + let (mut group, registrar) = TaskGroup::new(); + + bencher + .counter(ItemsCount::new(task_count)) + .bench_local(|| { + for output in 0..task_count { + let task = registrar.track(async move { black_box(output) }).unwrap(); + poll_ready(task, &mut context); + } + for _ in 0..task_count { + black_box(poll_ready(group.join_next(), &mut context).unwrap()); + } + }); +} + +// Prepare a closed group with every output already queued before timing starts. The measured work +// is `join()` moving those outputs into its result vector. +#[divan::bench(args = BATCH_SIZES)] +fn collect_completed_batch(bencher: Bencher, task_count: usize) { + let mut context = bench_context(); + + bencher + .with_inputs(|| { + let mut setup_context = bench_context(); + let (group, registrar) = TaskGroup::new(); + for output in 0..task_count { + let task = registrar.track(async move { output }).unwrap(); + poll_ready(task, &mut setup_context); + } + group.close(); + group + }) + .counter(ItemsCount::new(task_count)) + .bench_local_values(|mut group| black_box(poll_ready(group.join(), &mut context))); +} + +#[divan::bench(args = BATCH_SIZES)] +fn complete_batch_then_collect(bencher: Bencher, task_count: usize) { + let mut context = bench_context(); + + bencher + .with_inputs(|| { + let (group, registrar) = TaskGroup::new(); + let tasks = (0..task_count) + .map(|output| registrar.track(std::future::ready(output)).unwrap()) + .collect::>(); + group.close(); + (group, tasks) + }) + .counter(ItemsCount::new(task_count)) + .bench_local_values(|(mut group, tasks)| { + let mut join = pin!(group.join()); + poll_pending(join.as_mut(), &mut context); + for task in tasks { + poll_ready(task, &mut context); + } + black_box(poll_pinned_ready(join.as_mut(), &mut context)) + }); +} + +#[divan::bench(args = BATCH_SIZES)] +fn complete_batch_then_wait(bencher: Bencher, task_count: usize) { + let mut context = bench_context(); + + bencher + .with_inputs(|| { + let (group, registrar) = TaskGroup::new(); + let tasks = (0..task_count) + .map(|_| registrar.track(std::future::ready(())).unwrap()) + .collect::>(); + group.close(); + (group, tasks) + }) + .counter(ItemsCount::new(task_count)) + .bench_local_values(|(mut group, tasks)| { + let mut wait = pin!(group.wait()); + poll_pending(wait.as_mut(), &mut context); + for task in tasks { + poll_ready(task, &mut context); + } + poll_pinned_ready(wait.as_mut(), &mut context); + }); +} + +// Complete futures one at a time and poll `wait()` after every wake, as an executor normally would. +// `wait()` should wake only after the final future finishes. +#[divan::bench(args = BATCH_SIZES)] +fn complete_staggered_then_wait(bencher: Bencher, task_count: usize) { + bencher + .with_inputs(|| { + let (group, registrar) = TaskGroup::new(); + let tasks = (0..task_count) + .map(|_| registrar.track(std::future::ready(())).unwrap()) + .collect::>(); + group.close(); + + let wake_count = Arc::new(CountWake(AtomicUsize::new(0))); + let waker = Waker::from(wake_count.clone()); + (group, tasks, wake_count, waker) + }) + .counter(ItemsCount::new(task_count)) + .bench_local_values(|(mut group, tasks, wake_count, waker)| { + let mut context = Context::from_waker(&waker); + let mut wait = pin!(group.wait()); + poll_pending(wait.as_mut(), &mut context); + + let mut observed_wakes = 0; + for (index, task) in tasks.into_iter().enumerate() { + poll_ready(task, &mut context); + let current_wakes = wake_count.0.load(Ordering::Relaxed); + if current_wakes != observed_wakes && index + 1 < task_count { + observed_wakes = current_wakes; + poll_pending(wait.as_mut(), &mut context); + } + } + + poll_pinned_ready(wait.as_mut(), &mut context); + black_box(wake_count.0.load(Ordering::Relaxed)) + }); +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 4494225..02280e3 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -48,6 +48,7 @@ asyncband = { workspace = true, features = [ "semaphore", "shutdown", "singleflight", + "task-group", "waitgroup", "watch", ] } diff --git a/tests-integration/tests/task_group_test.rs b/tests-integration/tests/task_group_test.rs new file mode 100644 index 0000000..3d17c07 --- /dev/null +++ b/tests-integration/tests/task_group_test.rs @@ -0,0 +1,145 @@ +// 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::future; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use asyncband::blocking::FutureExt; +use asyncband::task_group::TaskGroup; +use tokio::sync::oneshot; + +struct PublishOnDrop<'a>(&'a AtomicUsize); + +impl Drop for PublishOnDrop<'_> { + fn drop(&mut self) { + self.0.store(1, Ordering::Relaxed); + } +} + +#[tokio::test] +async fn independently_spawned_tasks_are_joined_in_completion_order() { + let (mut group, registrar) = TaskGroup::new(); + let (first_tx, first_rx) = oneshot::channel(); + let (second_tx, second_rx) = oneshot::channel(); + + let first = tokio::spawn( + registrar + .track(async move { + first_rx.await.unwrap(); + 1 + }) + .unwrap(), + ); + let second = tokio::spawn( + registrar + .track(async move { + second_rx.await.unwrap(); + 2 + }) + .unwrap(), + ); + group.close(); + + second_tx.send(()).unwrap(); + assert_eq!(group.join_next().await, Some(2)); + first_tx.send(()).unwrap(); + assert_eq!(group.join_next().await, Some(1)); + assert_eq!(group.join_next().await, None); + + first.await.unwrap(); + second.await.unwrap(); +} + +#[tokio::test] +async fn a_child_can_register_nested_work_before_close() { + let (mut group, registrar) = TaskGroup::new(); + let child_registrar = registrar.clone(); + let (registered_tx, registered_rx) = oneshot::channel(); + + tokio::spawn( + registrar + .track(async move { + tokio::spawn(child_registrar.track(async { 2 }).unwrap()); + registered_tx.send(()).unwrap(); + 1 + }) + .unwrap(), + ); + + registered_rx.await.unwrap(); + group.close(); + let mut outputs = group.join().await; + outputs.sort_unstable(); + assert_eq!(outputs, [1, 2]); +} + +#[tokio::test] +async fn aborting_a_spawned_task_releases_its_registration() { + let (mut group, registrar) = TaskGroup::<()>::new(); + let task = tokio::spawn(registrar.track(future::pending()).unwrap()); + group.close(); + + task.abort(); + assert!(task.await.is_err()); + tokio::time::timeout(Duration::from_secs(1), group.wait()) + .await + .expect("the dropped task must not keep the group active"); +} + +#[tokio::test] +async fn wait_discards_outputs() { + let (mut group, registrar) = TaskGroup::new(); + tokio::spawn(registrar.track(async { 1 }).unwrap()); + tokio::spawn(registrar.track(async { 2 }).unwrap()); + group.close(); + + group.wait().await; + assert_eq!(group.join_next().await, None); +} + +#[test] +fn cancelled_tasks_publish_their_destructor_writes_before_wait_returns() { + let values = std::array::from_fn::<_, 8, _>(|_| AtomicUsize::new(0)); + let (mut group, registrar) = TaskGroup::new(); + let tasks = values + .iter() + .map(|value| { + let publish = PublishOnDrop(value); + registrar + .track(async move { + let _publish = publish; + future::pending::<()>().await; + }) + .unwrap() + }) + .collect::>(); + group.close(); + + std::thread::scope(|scope| { + for task in tasks { + scope.spawn(move || drop(task)); + } + FutureExt::block_on(group.wait()); + assert!( + values + .iter() + .all(|value| value.load(Ordering::Relaxed) == 1) + ); + }); +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 00d44dc..dfd8f12 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -51,6 +51,11 @@ use asyncband::shutdown::Shutdown; use asyncband::shutdown::ShutdownGuard; use asyncband::shutdown::ShutdownWatch; use asyncband::singleflight; +use asyncband::task_group::JoinNext; +use asyncband::task_group::Registrar; +use asyncband::task_group::TaskGroup; +use asyncband::task_group::TrackError; +use asyncband::task_group::Tracked; use asyncband::waitgroup::Wait; use asyncband::waitgroup::WaitGroup; use asyncband::watch; @@ -90,6 +95,9 @@ fn public_types_are_send_and_sync() { assert_send_and_sync::>(); assert_send_and_sync::>(); assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); + assert_send_and_sync::>(); assert_send_and_sync::(); assert_send_and_sync::(); assert_send_and_sync::(); @@ -157,6 +165,10 @@ fn movable_public_types_are_send() { let (_completer, completion) = completion::new::(); assert_send_value(completion.wait()); + let (mut group, registrar) = TaskGroup::new(); + assert_send_value(registrar.track(std::future::ready(42)).unwrap()); + assert_send_value(group.join_next()); + let (unbounded_sender, unbounded_receiver) = mpmc::unbounded::>(); assert_send_value(unbounded_receiver.recv()); drop(unbounded_sender); @@ -182,6 +194,11 @@ fn public_types_are_unpin() { assert_unpin::>(); assert_unpin::>(); assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>>(); + assert_unpin::>(); assert_unpin::(); assert_unpin::(); assert_unpin::();