diff --git a/CHANGELOG.md b/CHANGELOG.md index 847c4a8..9ce0708 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ All notable changes to this project will be documented in this file. * Implement `broadcast::mpmc::unbounded`, an unbounded broadcast channel that retains messages until all active receivers consume them or are dropped. * Add an opt-in latest-state channel under `asyncband::watch`. +* Add an opt-in shared one-shot completion primitive under `asyncband::completion` with a single-use completer, cloneable observers, a retained borrowed result, and observable abandonment. * Add opt-in `asyncband::once::LazyCell` for values that own one asynchronous initializer and preserve its in-flight future across caller cancellation. * Add opt-in bounded and unbounded runtime-agnostic object pools under `asyncband::pool`. * Add an opt-in `asyncband::blocking::FutureExt` bridge with `block_on` and `wait_timeout` methods for waiting on runtime-agnostic futures from synchronous code. diff --git a/README.md b/README.md index 2cb481e..8e832ba 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`LazyCell`](https://docs.rs/asyncband/*/asyncband/once/struct.LazyCell.html) | `lazy-cell` | Lazily initialize a value with a stored asynchronous function. | | | [`OnceMap`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceMap.html) | `once-map` | Initialize and store one value per key. | | Task coordination | [`Barrier`](https://docs.rs/asyncband/*/asyncband/barrier/struct.Barrier.html) | `barrier` | Wait until all participants reach a synchronization point. | +| | [`Completion`](https://docs.rs/asyncband/*/asyncband/completion/struct.Completion.html) | `completion` | Publish one shared result to any number of current and future observers. | | | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a one-way countdown completes. | | | [`WaitGroup`](https://docs.rs/asyncband/*/asyncband/waitgroup/struct.WaitGroup.html) | `waitgroup` | Wait for a dynamic group of tasks to finish. | | | [`Shutdown`](https://docs.rs/asyncband/*/asyncband/shutdown/struct.Shutdown.html) | `shutdown` | Coordinate shutdown signals and completion. | diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index 9b831d7..0370aca 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -47,6 +47,7 @@ default = [] barrier = [] blocking = [] broadcast = [] +completion = [] condvar = ["mutex"] latch = [] lazy-cell = ["mutex"] diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs new file mode 100644 index 0000000..b6d805c --- /dev/null +++ b/asyncband/src/completion/mod.rs @@ -0,0 +1,298 @@ +// 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. + +//! A shared one-shot completion primitive. +//! +//! A single-use [`Completer`] publishes one value, while any number of cloned [`Completion`] +//! observers wait for that same value. Observers created after completion see it immediately. +//! If the completer is dropped without publishing a value, every observer returns [`Abandoned`]. +//! The stored value is returned by reference, so callers decide whether to borrow it, clone it, or +//! use an [`Arc`]-wrapped value when they need independently owned shared results. +//! +//! Unlike `oneshot`, which transfers one value to one receiver, completion can fan one result out +//! to many current and future observers without creating and managing one channel per observer. +//! Unlike `OnceCell`, initialization is controlled only by the distinct completer capability; +//! observers can only wait. +//! +//! # Examples +//! +//! ``` +//! use asyncband::completion; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let (completer, completion) = completion::new(); +//! let first = completion.clone(); +//! let second = completion.clone(); +//! +//! completer.complete(String::from("ready")).unwrap(); +//! +//! assert_eq!(first.wait().await.unwrap(), "ready"); +//! assert_eq!(second.wait().await.unwrap(), "ready"); +//! let late = completion.clone(); +//! assert_eq!(late.wait().await.unwrap(), "ready"); +//! # } +//! ``` + +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::Weak; +use std::task::Context; +use std::task::Poll; + +use crate::internal::mutex::Mutex; +use crate::internal::waitset::WaitSet; +use crate::internal::waitset::WakerToken; +use crate::internal::waitset::wake_all; + +/// Creates a single-use [`Completer`] and a cloneable [`Completion`] observer. +pub fn new() -> (Completer, Completion) { + let shared = Arc::new(Shared { + value: OnceLock::new(), + state: Mutex::new(State { + status: Status::Pending, + waiters: WaitSet::new(), + }), + }); + let completer = Completer { + shared: Arc::downgrade(&shared), + }; + let completion = Completion { shared }; + (completer, completion) +} + +struct Shared { + value: OnceLock, + state: Mutex, +} + +struct State { + status: Status, + waiters: WaitSet, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Status { + Pending, + Completed, + Abandoned, +} + +/// The error returned by [`Completion::wait`] when the completer was dropped without a value. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Abandoned; + +impl fmt::Display for Abandoned { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("completion was abandoned before a value was provided") + } +} + +impl std::error::Error for Abandoned {} + +/// The capability that completes a [`Completion`] with one value. +/// +/// This type deliberately does not implement [`Clone`], and [`complete`](Self::complete) consumes +/// it. Dropping it before completion abandons the primitive and wakes all pending observers. +#[must_use = "dropping the completer abandons the completion"] +pub struct Completer { + shared: Weak>, +} + +// SAFETY: The completer can only move an owned `T` into the shared `OnceLock` while holding the +// state mutex; it never exposes or accesses the stored value afterward. `Completion` retains its +// ordinary auto traits, so observers cannot cross threads unless `T` can be shared. `T: Send` also +// permits the shared allocation and its value to be destroyed by the completing thread if its +// temporary strong reference is the last one. +unsafe impl Send for Completer {} +unsafe impl Sync for Completer {} + +impl fmt::Debug for Completer { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Completer").finish_non_exhaustive() + } +} + +impl Completer { + /// Completes the primitive with `value` and wakes all pending observers. + /// + /// Returns `value` if all observers were already dropped. A successful completion does not + /// guarantee that an observer will remain alive long enough to read the value. + /// + /// # Panics + /// + /// Panics if a registered waker panics while being notified. The value is committed before + /// notification begins. Before resuming the panic, `complete` still attempts to wake every + /// remaining registered waker. + pub fn complete(mut self, value: T) -> Result<(), T> { + let Some(shared) = self.shared.upgrade() else { + return Err(value); + }; + let wakers = { + let mut state = shared.state.lock(); + assert_eq!( + state.status, + Status::Pending, + "a live completer must refer to a pending completion" + ); + + if let Err(value) = shared.value.set(value) { + drop(state); + drop(value); + panic!("pending completion value must be unset"); + } + state.status = Status::Completed; + (!state.waiters.is_empty()).then(|| state.waiters.take_wakers()) + }; + // `complete` consumes the only completer. Disarm its destructor before invoking arbitrary + // wake callbacks; the completed state no longer needs abandonment handling. + self.shared = Weak::new(); + if let Some(wakers) = wakers { + wake_all(wakers); + } + Ok(()) + } +} + +impl Drop for Completer { + fn drop(&mut self) { + let Some(shared) = self.shared.upgrade() else { + return; + }; + let wakers = { + let mut state = shared.state.lock(); + if state.status != Status::Pending { + return; + } + state.status = Status::Abandoned; + (!state.waiters.is_empty()).then(|| state.waiters.take_wakers()) + }; + if let Some(wakers) = wakers { + wake_all(wakers); + } + } +} + +/// An observer of a shared one-shot completion. +/// +/// Cloning this type creates another observer of the same eventual value. Each call to [`wait`] +/// registers independently and can be cancelled without affecting other observers. +/// +/// [`wait`]: Completion::wait +pub struct Completion { + shared: Arc>, +} + +impl Clone for Completion { + fn clone(&self) -> Self { + Self { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for Completion { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Completion").finish_non_exhaustive() + } +} + +impl Completion { + /// Waits for the shared value and returns a reference to it. + /// + /// Returns [`Abandoned`] if the completer is dropped before providing a value. Abandonment + /// remains distinct from any error stored inside `T`. + /// + /// This method is cancel safe. Dropping one pending wait unregisters only that call and does + /// not affect this observer, another wait, or the eventual result. + pub async fn wait(&self) -> Result<&T, Abandoned> { + Wait { + completion: self, + registration: None, + } + .await + } +} + +struct Wait<'a, T> { + completion: &'a Completion, + registration: Option, +} + +impl<'a, T> Future for Wait<'a, T> { + type Output = Result<&'a T, Abandoned>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let mut prepared_waker = None; + loop { + let (poll, retired_waker) = { + let mut state = this.completion.shared.state.lock(); + match state.status { + Status::Pending => { + if prepared_waker.is_none() + && state + .waiters + .registered_waker_will_wake(&this.registration, cx.waker()) + { + return Poll::Pending; + } + let Some(waker) = prepared_waker.take() else { + drop(state); + prepared_waker = Some(cx.waker().clone()); + continue; + }; + let retired = state + .waiters + .register_owned_waker(&mut this.registration, waker); + (Poll::Pending, retired) + } + Status::Completed => { + let retired = state.waiters.unregister_waker(&mut this.registration); + let completion: &'a Completion = this.completion; + let value = completion + .shared + .value + .get() + .expect("completed value must be initialized"); + (Poll::Ready(Ok(value)), retired) + } + Status::Abandoned => { + let retired = state.waiters.unregister_waker(&mut this.registration); + (Poll::Ready(Err(Abandoned)), retired) + } + } + }; + drop(retired_waker); + drop(prepared_waker); + return poll; + } + } +} + +impl Drop for Wait<'_, T> { + fn drop(&mut self) { + let waker = { + let mut state = self.completion.shared.state.lock(); + state.waiters.unregister_waker(&mut self.registration) + }; + drop(waker); + } +} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 9cabd55..6bec43f 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -21,6 +21,7 @@ pub(crate) mod atomic_waker; #[cfg(any( feature = "barrier", feature = "broadcast", + feature = "completion", feature = "latch", feature = "mpsc", feature = "mutex", @@ -49,6 +50,7 @@ pub(crate) mod value_cell; #[cfg(any( feature = "barrier", feature = "broadcast", + feature = "completion", feature = "latch", feature = "mpsc", feature = "mutex", @@ -82,12 +84,13 @@ pub(crate) mod waitlist; #[cfg(any( feature = "barrier", feature = "broadcast", + feature = "completion", feature = "latch", feature = "once", feature = "waitgroup", feature = "watch", ))] -// `barrier` constructs a wait set with `with_capacity`, while countdown-based primitives use -// `new`. One constructor is therefore unused in every single-primitive build. +// `barrier` constructs a wait set with `with_capacity`, while completion and countdown-based +// primitives use `new`. One constructor is therefore unused in every single-primitive build. #[allow(dead_code)] pub(crate) mod waitset; diff --git a/asyncband/src/internal/waitset.rs b/asyncband/src/internal/waitset.rs index 8f4ab1f..fa8282a 100644 --- a/asyncband/src/internal/waitset.rs +++ b/asyncband/src/internal/waitset.rs @@ -108,17 +108,11 @@ impl WaitSet { token: &mut Option, cx: &mut Context<'_>, ) -> Option { - if let Some(current) = token.as_ref() { - if current.epoch == self.epoch { - let waker = self - .waiters - .get_mut(current.slot) - .expect("current waker token must refer to an occupied slot"); - if !waker.will_wake(cx.waker()) { - return Some(mem::replace(waker, cx.waker().clone())); - } - return None; - } + if self.registered_waker_will_wake(token, cx.waker()) { + return None; + } + if let Some(waker) = self.current_waker(token) { + return Some(mem::replace(waker, cx.waker().clone())); } *token = Some(WakerToken { @@ -128,6 +122,45 @@ impl WaitSet { None } + /// Returns whether `token` identifies a registered waker for the same task as `waker`. + #[inline] + pub fn registered_waker_will_wake(&self, token: &Option, waker: &Waker) -> bool { + let Some(current) = token.as_ref() else { + return false; + }; + if current.epoch != self.epoch { + return false; + } + self.waiters + .get(current.slot) + .expect("current waker token must refer to an occupied slot") + .will_wake(waker) + } + + /// Registers or updates an already cloned waker in the current wake epoch. + /// + /// Any waker not retained by the wait set is returned so the caller can drop it after releasing + /// the lock that protects this wait set. + #[inline] + pub fn register_owned_waker( + &mut self, + token: &mut Option, + waker: Waker, + ) -> Option { + if let Some(current) = self.current_waker(token) { + if !current.will_wake(&waker) { + return Some(mem::replace(current, waker)); + } + return Some(waker); + } + + *token = Some(WakerToken { + epoch: self.epoch, + slot: self.waiters.insert(waker), + }); + None + } + /// Removes the waker identified by `token` if it still belongs to the current wake epoch. /// /// The returned waker must be dropped after releasing the lock that protects this wait set. @@ -140,6 +173,18 @@ impl WaitSet { None } + fn current_waker(&mut self, token: &Option) -> Option<&mut Waker> { + let current = token.as_ref()?; + if current.epoch != self.epoch { + return None; + } + Some( + self.waiters + .get_mut(current.slot) + .expect("current waker token must refer to an occupied slot"), + ) + } + #[cfg(test)] fn registered_len(&self) -> usize { self.waiters.len() diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index a29f659..349f3cc 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -61,6 +61,7 @@ //! | | [`LazyCell`](once::LazyCell) | `lazy-cell` | Lazily initialize a value with a stored asynchronous function. | //! | | [`OnceMap`](once::OnceMap) | `once-map` | Initialize and store one value per key. | //! | Task coordination | [`Barrier`](barrier::Barrier) | `barrier` | Wait until all participants reach a synchronization point. | +//! | | [`Completion`](completion::Completion) | `completion` | Publish one shared result to any number of current and future observers. | //! | | [`Latch`](latch::Latch) | `latch` | Wait until a one-way countdown completes. | //! | | [`WaitGroup`](waitgroup::WaitGroup) | `waitgroup` | Wait for a dynamic group of tasks to finish. | //! | | [`Shutdown`](shutdown::Shutdown) | `shutdown` | Coordinate shutdown signals and completion. | @@ -120,6 +121,8 @@ pub mod barrier; pub mod blocking; #[cfg(feature = "broadcast")] pub mod broadcast; +#[cfg(feature = "completion")] +pub mod completion; #[cfg(feature = "condvar")] pub mod condvar; #[cfg(feature = "latch")] diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 44accc9..ba177a2 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -29,6 +29,7 @@ asyncband = { workspace = true, features = [ "barrier", "blocking", "broadcast", + "completion", "condvar", "latch", "mpsc", diff --git a/benchmarks/asyncband/completion/mod.rs b/benchmarks/asyncband/completion/mod.rs new file mode 100644 index 0000000..2676a97 --- /dev/null +++ b/benchmarks/asyncband/completion/mod.rs @@ -0,0 +1,101 @@ +// 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::pin::pin; + +use asyncband::completion; +use divan::Bencher; +use divan::black_box; + +use crate::support::bench_context; +use crate::support::poll_pending; +use crate::support::poll_pinned_ready; +use crate::support::poll_ready; + +const OBSERVER_COUNTS: &[usize] = &[1, 2, 4, 8, 32]; + +#[divan::bench] +fn ready_wait(bencher: Bencher) { + let mut context = bench_context(); + let (completer, completion) = completion::new(); + completer.complete(1usize).unwrap(); + + bencher.bench_local(|| black_box(*poll_ready(completion.wait(), &mut context).unwrap())); +} + +#[divan::bench] +fn cancel_pending(bencher: Bencher) { + let mut context = bench_context(); + let (_completer, completion) = completion::new::(); + + bencher.bench_local(|| { + let mut wait = pin!(completion.wait()); + poll_pending(wait.as_mut(), &mut context); + }); + black_box(completion); +} + +#[divan::bench] +fn complete_then_wait(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (completer, completion) = black_box(completion::new()); + completer.complete(black_box(1usize)).unwrap(); + black_box(*poll_ready(completion.wait(), &mut context).unwrap()) + }); +} + +#[divan::bench] +fn notify_pending(bencher: Bencher) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (completer, completion) = black_box(completion::new()); + let mut wait = pin!(completion.wait()); + poll_pending(wait.as_mut(), &mut context); + + completer.complete(black_box(1usize)).unwrap(); + black_box(*poll_pinned_ready(wait.as_mut(), &mut context).unwrap()) + }); +} + +#[divan::bench(args = OBSERVER_COUNTS)] +fn notify_pending_fanout(bencher: Bencher, observer_count: usize) { + let mut context = bench_context(); + + bencher.bench_local(|| { + let (completer, first) = black_box(completion::new()); + let mut observers = Vec::with_capacity(observer_count); + observers.push(first); + for _ in 1..observer_count { + observers.push(observers[0].clone()); + } + let mut waiters = observers + .iter() + .map(|observer| Box::pin(observer.wait())) + .collect::>(); + for waiter in &mut waiters { + poll_pending(waiter.as_mut(), &mut context); + } + + completer.complete(black_box(1usize)).unwrap(); + for mut waiter in waiters { + black_box(*poll_pinned_ready(waiter.as_mut(), &mut context).unwrap()); + } + }); +} diff --git a/benchmarks/asyncband/main.rs b/benchmarks/asyncband/main.rs index 980c875..0388f7e 100644 --- a/benchmarks/asyncband/main.rs +++ b/benchmarks/asyncband/main.rs @@ -18,6 +18,7 @@ mod barrier; mod blocking; mod broadcast; +mod completion; mod condvar; mod latch; mod mpsc; diff --git a/examples/Cargo.toml b/examples/Cargo.toml index 72b1ff2..6e079a8 100644 --- a/examples/Cargo.toml +++ b/examples/Cargo.toml @@ -27,6 +27,7 @@ release = false [dependencies] asyncband = { workspace = true, features = [ + "completion", "lazy-cell", "once-cell", "shutdown", @@ -53,3 +54,7 @@ path = "src/lazy_cell_boxed_vs_inline.rs" [[example]] name = "graceful_shutdown" path = "src/graceful_shutdown.rs" + +[[example]] +name = "shared_completion" +path = "src/shared_completion.rs" diff --git a/examples/src/shared_completion.rs b/examples/src/shared_completion.rs new file mode 100644 index 0000000..4087927 --- /dev/null +++ b/examples/src/shared_completion.rs @@ -0,0 +1,50 @@ +// 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. + +//! Share one immutable job result with current and late observers. +//! +//! A oneshot channel has one receiver. Serving these independent consumers with oneshot would +//! require one channel per consumer and manual fan-out, and a late consumer would need separate +//! result storage. Completion retains the result and exposes it to every observer. + +use asyncband::completion; + +#[tokio::main(flavor = "current_thread")] +async fn main() { + let (completer, completion) = completion::new(); + let dashboard = completion.clone(); + let audit_log = completion.clone(); + + let worker = tokio::spawn(async move { + tokio::task::yield_now().await; + completer + .complete(String::from("build 42 succeeded")) + .unwrap(); + }); + + let (dashboard_result, audit_result) = tokio::join!(dashboard.wait(), audit_log.wait()); + let dashboard_result = dashboard_result.unwrap(); + let audit_result = audit_result.unwrap(); + worker.await.unwrap(); + + let late_observer = completion.clone(); + let late_result = late_observer.wait().await.unwrap(); + + assert!(std::ptr::eq(dashboard_result, audit_result)); + assert!(std::ptr::eq(dashboard_result, late_result)); + println!("{late_result}"); +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index a43060d..1762c5d 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -30,6 +30,7 @@ asyncband = { workspace = true, features = [ "barrier", "blocking", "broadcast", + "completion", "condvar", "latch", "lazy-cell", diff --git a/tests-integration/tests/completion_test.rs b/tests-integration/tests/completion_test.rs new file mode 100644 index 0000000..977487e --- /dev/null +++ b/tests-integration/tests/completion_test.rs @@ -0,0 +1,496 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cell::Cell; +use std::future::Future; +use std::mem::ManuallyDrop; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Barrier; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::RawWaker; +use std::task::RawWakerVTable; +use std::task::Wake; +use std::task::Waker; +use std::thread; +use std::time::Duration; + +use asyncband::completion; + +struct NotClone(String); + +struct TrackWake(AtomicUsize); + +impl Wake for TrackWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +struct PanicWake; + +impl Wake for PanicWake { + fn wake(self: Arc) { + panic!("wake failed"); + } +} + +struct WakeCallback(Mutex>>); + +impl Wake for WakeCallback { + fn wake(self: Arc) { + let callback = self.0.lock().unwrap().take(); + if let Some(callback) = callback { + callback(); + } + } +} + +struct DropCallbackWake(Mutex>>); + +struct CloneCallbackWake(Mutex>>); + +// This test needs a custom waker whose final `Arc` drop is observable. +#[allow(clippy::manual_noop_waker)] +impl Wake for DropCallbackWake { + fn wake(self: Arc) {} +} + +impl Drop for DropCallbackWake { + fn drop(&mut self) { + if let Some(callback) = self.0.get_mut().unwrap().take() { + callback(); + } + } +} + +unsafe fn clone_callback_waker(data: *const ()) -> RawWaker { + // SAFETY: Every raw pointer using this vtable comes from `Arc::into_raw` below. ManuallyDrop + // keeps the original waker's strong reference alive while the clone callback borrows it. + let state = ManuallyDrop::new(unsafe { Arc::::from_raw(data.cast()) }); + if let Some(callback) = state.0.lock().unwrap().take() { + callback(); + } + RawWaker::new( + Arc::into_raw(Arc::clone(&state)).cast(), + &CLONE_CALLBACK_VTABLE, + ) +} + +unsafe fn wake_clone_callback_waker(data: *const ()) { + // SAFETY: `wake` consumes the raw waker's strong reference exactly once. + drop(unsafe { Arc::::from_raw(data.cast()) }); +} + +unsafe fn wake_clone_callback_waker_by_ref(_data: *const ()) {} + +unsafe fn drop_clone_callback_waker(data: *const ()) { + // SAFETY: `drop` consumes the raw waker's strong reference exactly once. + drop(unsafe { Arc::::from_raw(data.cast()) }); +} + +static CLONE_CALLBACK_VTABLE: RawWakerVTable = RawWakerVTable::new( + clone_callback_waker, + wake_clone_callback_waker, + wake_clone_callback_waker_by_ref, + drop_clone_callback_waker, +); + +fn waker_with_clone_callback(callback: impl FnOnce() + Send + 'static) -> Waker { + let state = Arc::new(CloneCallbackWake(Mutex::new(Some(Box::new(callback))))); + let raw = RawWaker::new(Arc::into_raw(state).cast(), &CLONE_CALLBACK_VTABLE); + // SAFETY: The vtable preserves the Arc strong count and all callbacks are thread safe. + unsafe { Waker::from_raw(raw) } +} + +fn poll_with(future: Pin<&mut F>, waker: &Waker) -> Poll { + future.poll(&mut Context::from_waker(waker)) +} + +fn assert_completes_without_deadlock(message: &'static str, test: impl FnOnce() + Send + 'static) { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(test)); + finished_tx.send(()).unwrap(); + if let Err(payload) = result { + std::panic::resume_unwind(payload); + } + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect(message); + worker.join().unwrap(); +} + +#[test] +fn all_observers_borrow_the_same_non_clone_value() { + let (completer, completion) = completion::new(); + let first = completion.clone(); + let second = completion.clone(); + + assert!(completer.complete(NotClone(String::from("ready"))).is_ok()); + + let first_value = pollster::block_on(first.wait()).unwrap(); + let second_value = pollster::block_on(second.wait()).unwrap(); + let repeated = pollster::block_on(first.wait()).unwrap(); + let late = completion.clone(); + let late_value = pollster::block_on(late.wait()).unwrap(); + + assert_eq!(first_value.0.as_str(), "ready"); + assert!(std::ptr::eq(first_value, second_value)); + assert!(std::ptr::eq(first_value, repeated)); + assert!(std::ptr::eq(first_value, late_value)); +} + +#[test] +fn completer_transfers_a_send_only_value_between_threads() { + let (completer, completion) = completion::new::>(); + + thread::spawn(move || completer.complete(Cell::new(7))) + .join() + .unwrap() + .unwrap(); + + assert_eq!(pollster::block_on(completion.wait()).unwrap().get(), 7); +} + +#[test] +fn complete_returns_the_value_when_no_observers_remain() { + let (completer, completion) = completion::new(); + drop(completion); + assert_eq!( + completer.complete(String::from("unobserved")).unwrap_err(), + "unobserved" + ); +} + +#[test] +fn dropping_the_completer_abandons_every_observer() { + let (completer, first) = completion::new::(); + let second = first.clone(); + drop(completer); + + assert_eq!(pollster::block_on(first.wait()), Err(completion::Abandoned)); + assert_eq!( + pollster::block_on(second.wait()), + Err(completion::Abandoned) + ); +} + +#[test] +fn dropping_one_observer_before_or_after_completion_does_not_affect_another() { + let (completer, first) = completion::new(); + let second = first.clone(); + drop(first); + + completer.complete(5).unwrap(); + assert_eq!(pollster::block_on(second.wait()), Ok(&5)); + + let (completer, first) = completion::new(); + let second = first.clone(); + completer.complete(6).unwrap(); + drop(first); + + assert_eq!(pollster::block_on(second.wait()), Ok(&6)); +} + +#[test] +fn completed_payload_is_released_with_the_last_observer() { + let payload = Arc::new(()); + let (completer, completion) = completion::new(); + + completer.complete(payload.clone()).unwrap(); + assert_eq!(Arc::strong_count(&payload), 2); + + drop(completion); + assert_eq!(Arc::strong_count(&payload), 1); +} + +#[test] +fn abandonment_wakes_all_registered_waits() { + let (completer, first) = completion::new::(); + let second = first.clone(); + let first_tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let second_tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let first_waker = Waker::from(first_tracker.clone()); + let second_waker = Waker::from(second_tracker.clone()); + let mut first_wait = Box::pin(first.wait()); + let mut second_wait = Box::pin(second.wait()); + + assert!(poll_with(first_wait.as_mut(), &first_waker).is_pending()); + assert!(poll_with(second_wait.as_mut(), &second_waker).is_pending()); + drop(completer); + + assert_eq!(first_tracker.0.load(Ordering::Relaxed), 1); + assert_eq!(second_tracker.0.load(Ordering::Relaxed), 1); + assert_eq!( + poll_with(first_wait.as_mut(), &first_waker), + Poll::Ready(Err(completion::Abandoned)) + ); + assert_eq!( + poll_with(second_wait.as_mut(), &second_waker), + Poll::Ready(Err(completion::Abandoned)) + ); +} + +#[test] +fn payload_errors_remain_distinct_from_abandonment() { + let (completer, completion) = completion::new::>(); + completer.complete(Err("domain error")).unwrap(); + assert_eq!( + pollster::block_on(completion.wait()), + Ok(&Err("domain error")) + ); + + let (completer, completion) = completion::new::>(); + drop(completer); + assert_eq!( + pollster::block_on(completion.wait()), + Err(completion::Abandoned) + ); +} + +#[test] +fn cancelling_a_wait_releases_only_its_waker() { + let (completer, completion) = completion::new(); + let cancelled_tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waiting_tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let cancelled_waker = Waker::from(cancelled_tracker.clone()); + let waiting_waker = Waker::from(waiting_tracker.clone()); + let baseline = Arc::strong_count(&cancelled_tracker); + let mut cancelled = Box::pin(completion.wait()); + let mut waiting = Box::pin(completion.wait()); + + assert!(poll_with(cancelled.as_mut(), &cancelled_waker).is_pending()); + assert!(poll_with(waiting.as_mut(), &waiting_waker).is_pending()); + assert_eq!(Arc::strong_count(&cancelled_tracker), baseline + 1); + drop(cancelled); + assert_eq!(Arc::strong_count(&cancelled_tracker), baseline); + + completer.complete(7).unwrap(); + assert_eq!(cancelled_tracker.0.load(Ordering::Relaxed), 0); + assert_eq!(waiting_tracker.0.load(Ordering::Relaxed), 1); + assert_eq!( + poll_with(waiting.as_mut(), &waiting_waker), + Poll::Ready(Ok(&7)) + ); +} + +#[test] +fn cancelling_after_wake_does_not_consume_the_shared_result() { + let (completer, first) = completion::new(); + let second = first.clone(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let mut wait = Box::pin(first.wait()); + + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + completer.complete(9).unwrap(); + assert_eq!(tracker.0.load(Ordering::Relaxed), 1); + drop(wait); + + assert_eq!(pollster::block_on(second.wait()), Ok(&9)); +} + +#[test] +fn cancellation_and_completer_drop_have_clean_orderings() { + let (completer, completion) = completion::new::(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let baseline = Arc::strong_count(&tracker); + let mut wait = Box::pin(completion.wait()); + + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + assert_eq!(Arc::strong_count(&tracker), baseline + 1); + drop(wait); + assert_eq!(Arc::strong_count(&tracker), baseline); + drop(completer); + assert_eq!(tracker.0.load(Ordering::Relaxed), 0); + assert_eq!( + pollster::block_on(completion.wait()), + Err(completion::Abandoned) + ); + + let (completer, completion) = completion::new::(); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let waker = Waker::from(tracker.clone()); + let baseline = Arc::strong_count(&tracker); + let mut wait = Box::pin(completion.wait()); + + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + drop(completer); + assert_eq!(tracker.0.load(Ordering::Relaxed), 1); + assert_eq!(Arc::strong_count(&tracker), baseline); + drop(wait); + assert_eq!(Arc::strong_count(&tracker), baseline); + assert_eq!( + pollster::block_on(completion.wait()), + Err(completion::Abandoned) + ); +} + +#[test] +fn completion_attempts_every_waker_after_one_panics() { + let (completer, first) = completion::new(); + let second = first.clone(); + let panicking = Waker::from(Arc::new(PanicWake)); + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let tracked = Waker::from(tracker.clone()); + let mut first_wait = Box::pin(first.wait()); + let mut second_wait = Box::pin(second.wait()); + + assert!(poll_with(first_wait.as_mut(), &panicking).is_pending()); + assert!(poll_with(second_wait.as_mut(), &tracked).is_pending()); + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| completer.complete(11))); + assert!(result.is_err()); + assert_eq!(tracker.0.load(Ordering::Relaxed), 1); + assert_eq!( + poll_with(second_wait.as_mut(), &tracked), + Poll::Ready(Ok(&11)) + ); +} + +#[test] +fn wake_callbacks_run_outside_the_completion_lock() { + assert_completes_without_deadlock( + "wake callback deadlocked against the completion lock", + || { + let (completer, completion) = completion::new(); + let callback_completion = completion.clone(); + let waker = Waker::from(Arc::new(WakeCallback(Mutex::new(Some(Box::new( + move || { + assert_eq!(pollster::block_on(callback_completion.wait()), Ok(&13)); + }, + )))))); + let mut wait = Box::pin(completion.wait()); + + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + completer.complete(13).unwrap(); + assert_eq!(poll_with(wait.as_mut(), &waker), Poll::Ready(Ok(&13))); + drop(wait); + + let (completer, completion) = completion::new::(); + let callback_completion = completion.clone(); + let waker = Waker::from(Arc::new(WakeCallback(Mutex::new(Some(Box::new( + move || { + assert_eq!( + pollster::block_on(callback_completion.wait()), + Err(completion::Abandoned) + ); + }, + )))))); + let mut wait = Box::pin(completion.wait()); + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + drop(completer); + assert_eq!( + poll_with(wait.as_mut(), &waker), + Poll::Ready(Err(completion::Abandoned)) + ); + }, + ); +} + +#[test] +fn waker_clone_callbacks_run_outside_the_completion_lock() { + assert_completes_without_deadlock( + "waker clone callback deadlocked against the completion lock", + || { + let (completer, completion) = completion::new::(); + let waker = waker_with_clone_callback(move || drop(completer)); + let mut wait = Box::pin(completion.wait()); + + assert_eq!( + poll_with(wait.as_mut(), &waker), + Poll::Ready(Err(completion::Abandoned)) + ); + }, + ); +} + +#[test] +fn replaced_wakers_are_dropped_outside_the_completion_lock() { + assert_completes_without_deadlock( + "replaced waker destructor deadlocked against the completion lock", + || { + let (completer, completion) = completion::new::(); + let old_waker = Waker::from(Arc::new(DropCallbackWake(Mutex::new(Some(Box::new( + move || drop(completer), + )))))); + let mut wait = Box::pin(completion.wait()); + assert!(poll_with(wait.as_mut(), &old_waker).is_pending()); + drop(old_waker); + + let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); + let replacement = Waker::from(tracker.clone()); + assert!(poll_with(wait.as_mut(), &replacement).is_pending()); + assert_eq!(tracker.0.load(Ordering::Relaxed), 1); + assert_eq!( + poll_with(wait.as_mut(), &replacement), + Poll::Ready(Err(completion::Abandoned)) + ); + }, + ); +} + +#[test] +fn cancelled_wakers_are_dropped_outside_the_completion_lock() { + assert_completes_without_deadlock( + "cancelled waker destructor deadlocked against the completion lock", + || { + let (completer, completion) = completion::new::(); + let waker = Waker::from(Arc::new(DropCallbackWake(Mutex::new(Some(Box::new( + move || drop(completer), + )))))); + let mut wait = Box::pin(completion.wait()); + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + drop(waker); + drop(wait); + assert_eq!( + pollster::block_on(completion.wait()), + Err(completion::Abandoned) + ); + }, + ); +} + +#[test] +fn complete_and_final_observer_drop_linearize_cleanly() { + for _ in 0..100 { + let (completer, completion) = completion::new(); + let barrier = Arc::new(Barrier::new(2)); + thread::scope(|scope| { + let worker_barrier = barrier.clone(); + let worker = scope.spawn(move || { + worker_barrier.wait(); + drop(completion); + }); + + barrier.wait(); + let result = completer.complete(17); + worker.join().unwrap(); + if let Err(error) = result { + assert_eq!(error, 17); + } + }); + } +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 2a62838..34ff506 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -19,6 +19,7 @@ use std::cell::Cell; use asyncband::barrier::Barrier; use asyncband::broadcast; +use asyncband::completion; use asyncband::condvar::Condvar; use asyncband::latch::Latch; use asyncband::mpsc; @@ -70,6 +71,10 @@ 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::>(); @@ -121,6 +126,9 @@ fn movable_public_types_are_send() { let (_tx, mut rx) = watch::channel(0); assert_send_value(rx.changed()); + + let (_completer, completion) = completion::new::(); + assert_send_value(completion.wait()); } #[test] @@ -129,6 +137,9 @@ fn public_types_are_unpin() { assert_unpin::(); assert_unpin::(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::(); assert_unpin::(); assert_unpin::>>(); assert_unpin::();