From 7302ca710a529291d4abceb54dec1b958e3a0367 Mon Sep 17 00:00:00 2001 From: QwQBiG Date: Sat, 29 Aug 2026 22:35:23 +0800 Subject: [PATCH 1/5] feat(completion): add a shared one-shot completion primitive Signed-off-by: QwQBiG --- CHANGELOG.md | 1 + README.md | 1 + asyncband/Cargo.toml | 1 + asyncband/src/completion/error.rs | 71 ++++ asyncband/src/completion/mod.rs | 262 ++++++++++++ asyncband/src/completion/tests.rs | 49 +++ asyncband/src/internal/mod.rs | 7 +- asyncband/src/lib.rs | 3 + benchmarks/Cargo.toml | 1 + benchmarks/asyncband/completion/mod.rs | 102 +++++ benchmarks/asyncband/main.rs | 1 + examples/Cargo.toml | 5 + examples/src/shared_completion.rs | 50 +++ tests-integration/Cargo.toml | 1 + tests-integration/tests/completion_test.rs | 453 +++++++++++++++++++++ tests-integration/tests/traits_test.rs | 12 + 16 files changed, 1018 insertions(+), 2 deletions(-) create mode 100644 asyncband/src/completion/error.rs create mode 100644 asyncband/src/completion/mod.rs create mode 100644 asyncband/src/completion/tests.rs create mode 100644 benchmarks/asyncband/completion/mod.rs create mode 100644 examples/src/shared_completion.rs create mode 100644 tests-integration/tests/completion_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 847c4a8..e791488 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` for publishing one result to current and future observers. * 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/error.rs b/asyncband/src/completion/error.rs new file mode 100644 index 0000000..0c9de57 --- /dev/null +++ b/asyncband/src/completion/error.rs @@ -0,0 +1,71 @@ +// 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::any::type_name; +use std::fmt; + +/// An error returned when a value cannot complete a [`Completion`](super::Completion). +/// +/// Completion is rejected after another value has already won or when no observers remain. The +/// rejected value can be retrieved with [`CompleteError::into_inner`]. +#[derive(Clone, PartialEq, Eq)] +pub struct CompleteError(T); + +impl CompleteError { + /// Returns a reference to the value that could not complete the primitive. + pub fn as_inner(&self) -> &T { + &self.0 + } + + /// Consumes the error and returns the value that could not complete the primitive. + pub fn into_inner(self) -> T { + self.0 + } + + pub(super) fn new(value: T) -> Self { + Self(value) + } +} + +impl fmt::Display for CompleteError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("completion rejected") + } +} + +impl fmt::Debug for CompleteError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "CompleteError<{}>(..)", type_name::()) + } +} + +impl std::error::Error for CompleteError {} + +/// An error returned when the completer is dropped before providing a value. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WaitError { + /// The completer was dropped before providing a value. + Closed, +} + +impl fmt::Display for WaitError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("completion closed before a value was provided") + } +} + +impl std::error::Error for WaitError {} diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs new file mode 100644 index 0000000..93388fc --- /dev/null +++ b/asyncband/src/completion/mod.rs @@ -0,0 +1,262 @@ +// 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 [`Completer`] publishes one value, while any number of cloned [`Completion`] observers wait +//! for that same value. Observers created after completion see it immediately. 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::channel(); +//! 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"); +//! # } +//! ``` + +mod error; + +#[cfg(test)] +mod tests; + +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::OnceLock; +use std::task::Context; +use std::task::Poll; + +pub use self::error::CompleteError; +pub use self::error::WaitError; +use crate::internal::mutex::Mutex; +use crate::internal::waitset::WaitSet; +use crate::internal::waitset::WakerToken; +use crate::internal::waitset::wake_all; + +/// Creates a shared one-shot completion primitive. +pub fn channel() -> (Completer, Completion) { + let shared = Arc::new(Shared { + value: OnceLock::new(), + state: Mutex::new(State { + status: Status::Pending, + observers: 1, + waiters: WaitSet::new(), + }), + }); + let completer = Completer { + shared: shared.clone(), + }; + let completion = Completion { shared }; + (completer, completion) +} + +struct Shared { + value: OnceLock, + state: Mutex, +} + +struct State { + status: Status, + observers: usize, + waiters: WaitSet, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Status { + Pending, + Completed, + Closed, +} + +/// The capability that completes a [`Completion`] with one value. +/// +/// This type deliberately does not implement [`Clone`]. Dropping it before completion closes the +/// primitive and wakes all pending observers. +pub struct Completer { + shared: Arc>, +} + +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. + /// + /// The value is rejected and returned if another value has already completed the primitive or + /// if no observers remain. + pub fn complete(&self, value: T) -> Result<(), CompleteError> { + let wakers = { + let mut state = self.shared.state.lock(); + if state.status != Status::Pending || state.observers == 0 { + return Err(CompleteError::new(value)); + } + + if let Err(value) = self.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()) + }; + if let Some(wakers) = wakers { + wake_all(wakers); + } + Ok(()) + } +} + +impl Drop for Completer { + fn drop(&mut self) { + let wakers = { + let mut state = self.shared.state.lock(); + if state.status != Status::Pending { + return; + } + state.status = Status::Closed; + (!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 { + let mut state = self.shared.state.lock(); + state.observers = state + .observers + .checked_add(1) + .expect("completion observer count overflowed"); + drop(state); + 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 Drop for Completion { + fn drop(&mut self) { + let mut state = self.shared.state.lock(); + state.observers -= 1; + } +} + +impl Completion { + /// Waits for the shared value and returns a reference to it. + /// + /// Returns [`WaitError::Closed`] if the completer is dropped before providing a value. This + /// transport-level closure 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, WaitError> { + 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, WaitError>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let (poll, retired_waker) = { + let mut state = this.completion.shared.state.lock(); + match state.status { + Status::Pending => { + let retired = state.waiters.register_waker(&mut this.registration, cx); + (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::Closed => { + let retired = state.waiters.unregister_waker(&mut this.registration); + (Poll::Ready(Err(WaitError::Closed)), retired) + } + } + }; + drop(retired_waker); + 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/completion/tests.rs b/asyncband/src/completion/tests.rs new file mode 100644 index 0000000..e0beb1c --- /dev/null +++ b/asyncband/src/completion/tests.rs @@ -0,0 +1,49 @@ +// 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::sync::Arc; +use std::sync::Barrier; + +use super::*; + +#[tokio::test] +async fn concurrent_complete_calls_commit_exactly_one_value() { + let (completer, completion) = channel(); + let completer = Arc::new(completer); + let barrier = Arc::new(Barrier::new(3)); + + let workers = [10, 20].map(|value| { + let completer = completer.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + barrier.wait(); + completer.complete(value) + }) + }); + + barrier.wait(); + let results = workers.map(|worker| worker.join().unwrap()); + assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 1); + + let rejected = results + .into_iter() + .find_map(Result::err) + .expect("one completion must be rejected") + .into_inner(); + let completed = *completion.wait().await.unwrap(); + assert_ne!(completed, rejected); +} 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/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..82c0e51 --- /dev/null +++ b/benchmarks/asyncband/completion/mod.rs @@ -0,0 +1,102 @@ +// 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::channel(); + completer.complete(1usize).unwrap(); + + bencher.bench_local(|| black_box(*poll_ready(completion.wait(), &mut context).unwrap())); + black_box(completer); +} + +#[divan::bench] +fn cancel_pending(bencher: Bencher) { + let mut context = bench_context(); + let (_completer, completion) = completion::channel::(); + + 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::channel()); + 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::channel()); + 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::channel()); + 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..dc49b96 --- /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::channel(); + 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..5b3983c --- /dev/null +++ b/tests-integration/tests/completion_test.rs @@ -0,0 +1,453 @@ +// 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::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::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 ReentrantRejected { + completer: Option>>, +} + +impl Drop for ReentrantRejected { + fn drop(&mut self) { + if let Some(completer) = self.completer.take() { + drop(completer.complete(Self { completer: None })); + } + } +} + +// 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(); + } + } +} + +fn poll_with(future: Pin<&mut F>, waker: &Waker) -> Poll { + future.poll(&mut Context::from_waker(waker)) +} + +#[test] +fn all_observers_borrow_the_same_non_clone_value() { + let (completer, completion) = completion::channel(); + let first = completion.clone(); + let second = completion.clone(); + + completer.complete(NotClone(String::from("ready"))).unwrap(); + + 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 rejected_completions_return_the_original_value() { + let (completer, completion) = completion::channel(); + completer.complete(String::from("winner")).unwrap(); + + let error = completer.complete(String::from("rejected")).unwrap_err(); + assert_eq!(error.as_inner(), "rejected"); + assert_eq!(error.into_inner(), "rejected"); + assert_eq!(pollster::block_on(completion.wait()).unwrap(), "winner"); + + let (completer, completion) = completion::channel(); + drop(completion); + assert_eq!( + completer + .complete(String::from("unobserved")) + .unwrap_err() + .into_inner(), + "unobserved" + ); +} + +#[test] +fn rejected_payloads_are_dropped_outside_the_completion_lock() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + let (completer, _completion) = completion::channel(); + let completer = Arc::new(completer); + completer + .complete(ReentrantRejected { completer: None }) + .unwrap(); + + drop(completer.complete(ReentrantRejected { + completer: Some(completer.clone()), + })); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("rejected payload destructor deadlocked against the completion lock"); + worker.join().unwrap(); +} + +#[test] +fn completer_drop_closes_every_observer_but_completed_drop_does_not() { + let (completer, first) = completion::channel::(); + let second = first.clone(); + drop(completer); + + assert_eq!( + pollster::block_on(first.wait()), + Err(completion::WaitError::Closed) + ); + assert_eq!( + pollster::block_on(second.wait()), + Err(completion::WaitError::Closed) + ); + + let (completer, completion) = completion::channel(); + completer.complete(42).unwrap(); + drop(completer); + assert_eq!(pollster::block_on(completion.wait()), Ok(&42)); +} + +#[test] +fn dropping_one_observer_before_or_after_completion_does_not_affect_another() { + let (completer, first) = completion::channel(); + let second = first.clone(); + drop(first); + + completer.complete(5).unwrap(); + assert_eq!(pollster::block_on(second.wait()), Ok(&5)); + + let (completer, first) = completion::channel(); + let second = first.clone(); + completer.complete(6).unwrap(); + drop(first); + + assert_eq!(pollster::block_on(second.wait()), Ok(&6)); +} + +#[test] +fn completer_drop_wakes_all_registered_waits() { + let (completer, first) = completion::channel::(); + 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::WaitError::Closed)) + ); + assert_eq!( + poll_with(second_wait.as_mut(), &second_waker), + Poll::Ready(Err(completion::WaitError::Closed)) + ); +} + +#[test] +fn payload_errors_remain_distinct_from_completion_closure() { + let (completer, completion) = completion::channel::>(); + completer.complete(Err("domain error")).unwrap(); + assert_eq!( + pollster::block_on(completion.wait()), + Ok(&Err("domain error")) + ); + + let (completer, completion) = completion::channel::>(); + drop(completer); + assert_eq!( + pollster::block_on(completion.wait()), + Err(completion::WaitError::Closed) + ); +} + +#[test] +fn cancelling_a_wait_releases_only_its_waker() { + let (completer, completion) = completion::channel(); + 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::channel(); + 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::channel::(); + 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::WaitError::Closed) + ); + + let (completer, completion) = completion::channel::(); + 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::WaitError::Closed) + ); +} + +#[test] +fn completion_attempts_every_waker_after_one_panics() { + let (completer, first) = completion::channel(); + 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() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + let (completer, completion) = completion::channel(); + let callback_completion = completion.clone(); + let waker = Waker::from(Arc::new(WakeCallback(Mutex::new(Some(Box::new( + move || { + drop(callback_completion); + }, + )))))); + 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::channel::(); + let callback_completion = completion.clone(); + let waker = Waker::from(Arc::new(WakeCallback(Mutex::new(Some(Box::new( + move || { + drop(callback_completion); + }, + )))))); + 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::WaitError::Closed)) + ); + + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("wake callback deadlocked against the completion lock"); + worker.join().unwrap(); +} + +#[test] +fn replaced_wakers_are_dropped_outside_the_completion_lock() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + let (_completer, completion) = completion::channel::(); + let callback_completion = completion.clone(); + let old_waker = Waker::from(Arc::new(DropCallbackWake(Mutex::new(Some(Box::new( + move || { + drop(callback_completion); + }, + )))))); + let mut wait = Box::pin(completion.wait()); + assert!(poll_with(wait.as_mut(), &old_waker).is_pending()); + drop(old_waker); + + let replacement = Waker::from(Arc::new(TrackWake(AtomicUsize::new(0)))); + assert!(poll_with(wait.as_mut(), &replacement).is_pending()); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("replaced waker destructor deadlocked against the completion lock"); + worker.join().unwrap(); +} + +#[test] +fn cancelled_wakers_are_dropped_outside_the_completion_lock() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + let (_completer, completion) = completion::channel::(); + let callback_completion = completion.clone(); + let waker = Waker::from(Arc::new(DropCallbackWake(Mutex::new(Some(Box::new( + move || { + drop(callback_completion); + }, + )))))); + let mut wait = Box::pin(completion.wait()); + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + drop(waker); + drop(wait); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("cancelled waker destructor deadlocked against the completion lock"); + worker.join().unwrap(); +} + +#[test] +fn complete_and_final_observer_drop_linearize_cleanly() { + for _ in 0..100 { + let (completer, completion) = completion::channel(); + 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.into_inner(), 17); + } + }); + } +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 2a62838..f1c98f4 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::channel::(); + assert_send_value(completion.wait()); } #[test] @@ -129,6 +137,10 @@ fn public_types_are_unpin() { assert_unpin::(); assert_unpin::(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::>(); + assert_unpin::(); assert_unpin::(); assert_unpin::>>(); assert_unpin::(); From 4290d5e2fee88fd3facb8ca5fbc8e5f6b35614ce Mon Sep 17 00:00:00 2001 From: QwQBiG Date: Sun, 30 Aug 2026 13:30:56 +0800 Subject: [PATCH 2/5] fix(completion): release value with final observer Signed-off-by: QwQBiG --- asyncband/src/completion/mod.rs | 23 +++++++++++++++++----- tests-integration/tests/completion_test.rs | 14 +++++++++++++ 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index 93388fc..1d5df61 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -57,6 +57,7 @@ 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; @@ -78,7 +79,7 @@ pub fn channel() -> (Completer, Completion) { }), }); let completer = Completer { - shared: shared.clone(), + shared: Arc::downgrade(&shared), }; let completion = Completion { shared }; (completer, completion) @@ -107,7 +108,7 @@ enum Status { /// This type deliberately does not implement [`Clone`]. Dropping it before completion closes the /// primitive and wakes all pending observers. pub struct Completer { - shared: Arc>, + shared: Weak>, } impl fmt::Debug for Completer { @@ -121,14 +122,23 @@ impl Completer { /// /// The value is rejected and returned if another value has already completed the primitive or /// if no observers remain. + /// + /// # Panics + /// + /// Panics if a registered waker panics while being notified. The value is committed before + /// notification begins, so subsequent completion attempts are rejected. Before resuming the + /// panic, `complete` still attempts to wake every remaining registered waker. pub fn complete(&self, value: T) -> Result<(), CompleteError> { + let Some(shared) = self.shared.upgrade() else { + return Err(CompleteError::new(value)); + }; let wakers = { - let mut state = self.shared.state.lock(); + let mut state = shared.state.lock(); if state.status != Status::Pending || state.observers == 0 { return Err(CompleteError::new(value)); } - if let Err(value) = self.shared.value.set(value) { + if let Err(value) = shared.value.set(value) { drop(state); drop(value); panic!("pending completion value must be unset"); @@ -145,8 +155,11 @@ impl Completer { impl Drop for Completer { fn drop(&mut self) { + let Some(shared) = self.shared.upgrade() else { + return; + }; let wakers = { - let mut state = self.shared.state.lock(); + let mut state = shared.state.lock(); if state.status != Status::Pending { return; } diff --git a/tests-integration/tests/completion_test.rs b/tests-integration/tests/completion_test.rs index 5b3983c..d8af017 100644 --- a/tests-integration/tests/completion_test.rs +++ b/tests-integration/tests/completion_test.rs @@ -193,6 +193,20 @@ fn dropping_one_observer_before_or_after_completion_does_not_affect_another() { 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::channel(); + + completer.complete(payload.clone()).unwrap(); + assert_eq!(Arc::strong_count(&payload), 2); + + drop(completion); + assert_eq!(Arc::strong_count(&payload), 1); + + drop(completer); +} + #[test] fn completer_drop_wakes_all_registered_waits() { let (completer, first) = completion::channel::(); From 90d4f8d750cd0002bafc113570ac5de445b93edc Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 30 Aug 2026 18:15:53 +0800 Subject: [PATCH 3/5] fix(completion): clone wakers outside state lock --- asyncband/src/completion/mod.rs | 66 +++++++++++++-------- asyncband/src/internal/waitset.rs | 67 ++++++++++++++++++---- tests-integration/tests/completion_test.rs | 63 ++++++++++++++++++++ 3 files changed, 161 insertions(+), 35 deletions(-) diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index 1d5df61..956d16b 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -236,31 +236,49 @@ impl<'a, T> Future for Wait<'a, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); - let (poll, retired_waker) = { - let mut state = this.completion.shared.state.lock(); - match state.status { - Status::Pending => { - let retired = state.waiters.register_waker(&mut this.registration, cx); - (Poll::Pending, retired) + 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::Closed => { + let retired = state.waiters.unregister_waker(&mut this.registration); + (Poll::Ready(Err(WaitError::Closed)), 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::Closed => { - let retired = state.waiters.unregister_waker(&mut this.registration); - (Poll::Ready(Err(WaitError::Closed)), retired) - } - } - }; - drop(retired_waker); - poll + }; + drop(retired_waker); + drop(prepared_waker); + return poll; + } } } 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/tests-integration/tests/completion_test.rs b/tests-integration/tests/completion_test.rs index d8af017..3b5b978 100644 --- a/tests-integration/tests/completion_test.rs +++ b/tests-integration/tests/completion_test.rs @@ -16,6 +16,7 @@ // under the License. use std::future::Future; +use std::mem::ManuallyDrop; use std::pin::Pin; use std::sync::Arc; use std::sync::Barrier; @@ -24,6 +25,8 @@ 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; @@ -62,6 +65,8 @@ impl Wake for WakeCallback { struct DropCallbackWake(Mutex>>); +struct CloneCallbackWake(Mutex>>); + struct ReentrantRejected { completer: Option>>, } @@ -88,6 +93,45 @@ impl Drop for DropCallbackWake { } } +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)) } @@ -394,6 +438,25 @@ fn wake_callbacks_run_outside_the_completion_lock() { worker.join().unwrap(); } +#[test] +fn waker_clone_callbacks_run_outside_the_completion_lock() { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + let (_completer, completion) = completion::channel::(); + let callback_completion = completion.clone(); + let waker = waker_with_clone_callback(move || drop(callback_completion)); + let mut wait = Box::pin(completion.wait()); + + assert!(poll_with(wait.as_mut(), &waker).is_pending()); + finished_tx.send(()).unwrap(); + }); + + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("waker clone callback deadlocked against the completion lock"); + worker.join().unwrap(); +} + #[test] fn replaced_wakers_are_dropped_outside_the_completion_lock() { let (finished_tx, finished_rx) = std::sync::mpsc::channel(); From 28fe9ecc3167ec45ae9052592a34c2b1cd57cf07 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 30 Aug 2026 18:17:08 +0800 Subject: [PATCH 4/5] fix(completion): preserve send-only producer values --- asyncband/src/completion/mod.rs | 8 ++++++++ tests-integration/tests/completion_test.rs | 13 +++++++++++++ tests-integration/tests/traits_test.rs | 1 + 3 files changed, 22 insertions(+) diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index 956d16b..46c8d4a 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -111,6 +111,14 @@ 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 a completing thread when 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() diff --git a/tests-integration/tests/completion_test.rs b/tests-integration/tests/completion_test.rs index 3b5b978..238af1e 100644 --- a/tests-integration/tests/completion_test.rs +++ b/tests-integration/tests/completion_test.rs @@ -15,6 +15,7 @@ // 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; @@ -156,6 +157,18 @@ fn all_observers_borrow_the_same_non_clone_value() { assert!(std::ptr::eq(first_value, late_value)); } +#[test] +fn completer_transfers_a_send_only_value_between_threads() { + let (completer, completion) = completion::channel::>(); + + thread::spawn(move || completer.complete(Cell::new(7))) + .join() + .unwrap() + .unwrap(); + + assert_eq!(pollster::block_on(completion.wait()).unwrap().get(), 7); +} + #[test] fn rejected_completions_return_the_original_value() { let (completer, completion) = completion::channel(); diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index f1c98f4..51bbb94 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -71,6 +71,7 @@ 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::>(); From f952f31a1ccc0a4e1ca3fbab379e849494c8a2f6 Mon Sep 17 00:00:00 2001 From: tison Date: Sun, 30 Aug 2026 19:36:38 +0800 Subject: [PATCH 5/5] refactor(completion): make completer single-use --- CHANGELOG.md | 2 +- asyncband/src/completion/error.rs | 71 ----- asyncband/src/completion/mod.rs | 95 +++---- asyncband/src/completion/tests.rs | 49 ---- benchmarks/asyncband/completion/mod.rs | 11 +- examples/src/shared_completion.rs | 2 +- tests-integration/tests/completion_test.rs | 313 +++++++++------------ tests-integration/tests/traits_test.rs | 8 +- 8 files changed, 189 insertions(+), 362 deletions(-) delete mode 100644 asyncband/src/completion/error.rs delete mode 100644 asyncband/src/completion/tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index e791488..9ce0708 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,7 +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` for publishing one result to current and future observers. +* 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/asyncband/src/completion/error.rs b/asyncband/src/completion/error.rs deleted file mode 100644 index 0c9de57..0000000 --- a/asyncband/src/completion/error.rs +++ /dev/null @@ -1,71 +0,0 @@ -// 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::any::type_name; -use std::fmt; - -/// An error returned when a value cannot complete a [`Completion`](super::Completion). -/// -/// Completion is rejected after another value has already won or when no observers remain. The -/// rejected value can be retrieved with [`CompleteError::into_inner`]. -#[derive(Clone, PartialEq, Eq)] -pub struct CompleteError(T); - -impl CompleteError { - /// Returns a reference to the value that could not complete the primitive. - pub fn as_inner(&self) -> &T { - &self.0 - } - - /// Consumes the error and returns the value that could not complete the primitive. - pub fn into_inner(self) -> T { - self.0 - } - - pub(super) fn new(value: T) -> Self { - Self(value) - } -} - -impl fmt::Display for CompleteError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("completion rejected") - } -} - -impl fmt::Debug for CompleteError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "CompleteError<{}>(..)", type_name::()) - } -} - -impl std::error::Error for CompleteError {} - -/// An error returned when the completer is dropped before providing a value. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum WaitError { - /// The completer was dropped before providing a value. - Closed, -} - -impl fmt::Display for WaitError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("completion closed before a value was provided") - } -} - -impl std::error::Error for WaitError {} diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index 46c8d4a..b6d805c 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -17,10 +17,11 @@ //! A shared one-shot completion primitive. //! -//! A [`Completer`] publishes one value, while any number of cloned [`Completion`] observers wait -//! for that same value. Observers created after completion see it immediately. 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. +//! 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. @@ -34,7 +35,7 @@ //! //! # #[tokio::main] //! # async fn main() { -//! let (completer, completion) = completion::channel(); +//! let (completer, completion) = completion::new(); //! let first = completion.clone(); //! let second = completion.clone(); //! @@ -47,11 +48,6 @@ //! # } //! ``` -mod error; - -#[cfg(test)] -mod tests; - use std::fmt; use std::future::Future; use std::pin::Pin; @@ -61,20 +57,17 @@ use std::sync::Weak; use std::task::Context; use std::task::Poll; -pub use self::error::CompleteError; -pub use self::error::WaitError; use crate::internal::mutex::Mutex; use crate::internal::waitset::WaitSet; use crate::internal::waitset::WakerToken; use crate::internal::waitset::wake_all; -/// Creates a shared one-shot completion primitive. -pub fn channel() -> (Completer, Completion) { +/// 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, - observers: 1, waiters: WaitSet::new(), }), }); @@ -92,7 +85,6 @@ struct Shared { struct State { status: Status, - observers: usize, waiters: WaitSet, } @@ -100,13 +92,26 @@ struct State { enum Status { Pending, Completed, - Closed, + 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`]. Dropping it before completion closes the -/// primitive and wakes all pending observers. +/// 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>, } @@ -114,7 +119,7 @@ pub struct Completer { // 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 a completing thread when its +// 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 {} @@ -128,23 +133,25 @@ impl fmt::Debug for Completer { impl Completer { /// Completes the primitive with `value` and wakes all pending observers. /// - /// The value is rejected and returned if another value has already completed the primitive or - /// if no observers remain. + /// 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, so subsequent completion attempts are rejected. Before resuming the - /// panic, `complete` still attempts to wake every remaining registered waker. - pub fn complete(&self, value: T) -> Result<(), CompleteError> { + /// 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(CompleteError::new(value)); + return Err(value); }; let wakers = { let mut state = shared.state.lock(); - if state.status != Status::Pending || state.observers == 0 { - return Err(CompleteError::new(value)); - } + 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); @@ -154,6 +161,9 @@ impl Completer { 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); } @@ -171,7 +181,7 @@ impl Drop for Completer { if state.status != Status::Pending { return; } - state.status = Status::Closed; + state.status = Status::Abandoned; (!state.waiters.is_empty()).then(|| state.waiters.take_wakers()) }; if let Some(wakers) = wakers { @@ -192,12 +202,6 @@ pub struct Completion { impl Clone for Completion { fn clone(&self) -> Self { - let mut state = self.shared.state.lock(); - state.observers = state - .observers - .checked_add(1) - .expect("completion observer count overflowed"); - drop(state); Self { shared: self.shared.clone(), } @@ -210,22 +214,15 @@ impl fmt::Debug for Completion { } } -impl Drop for Completion { - fn drop(&mut self) { - let mut state = self.shared.state.lock(); - state.observers -= 1; - } -} - impl Completion { /// Waits for the shared value and returns a reference to it. /// - /// Returns [`WaitError::Closed`] if the completer is dropped before providing a value. This - /// transport-level closure remains distinct from any error stored inside `T`. + /// 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, WaitError> { + pub async fn wait(&self) -> Result<&T, Abandoned> { Wait { completion: self, registration: None, @@ -240,7 +237,7 @@ struct Wait<'a, T> { } impl<'a, T> Future for Wait<'a, T> { - type Output = Result<&'a T, WaitError>; + type Output = Result<&'a T, Abandoned>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); @@ -277,9 +274,9 @@ impl<'a, T> Future for Wait<'a, T> { .expect("completed value must be initialized"); (Poll::Ready(Ok(value)), retired) } - Status::Closed => { + Status::Abandoned => { let retired = state.waiters.unregister_waker(&mut this.registration); - (Poll::Ready(Err(WaitError::Closed)), retired) + (Poll::Ready(Err(Abandoned)), retired) } } }; diff --git a/asyncband/src/completion/tests.rs b/asyncband/src/completion/tests.rs deleted file mode 100644 index e0beb1c..0000000 --- a/asyncband/src/completion/tests.rs +++ /dev/null @@ -1,49 +0,0 @@ -// 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::sync::Arc; -use std::sync::Barrier; - -use super::*; - -#[tokio::test] -async fn concurrent_complete_calls_commit_exactly_one_value() { - let (completer, completion) = channel(); - let completer = Arc::new(completer); - let barrier = Arc::new(Barrier::new(3)); - - let workers = [10, 20].map(|value| { - let completer = completer.clone(); - let barrier = barrier.clone(); - std::thread::spawn(move || { - barrier.wait(); - completer.complete(value) - }) - }); - - barrier.wait(); - let results = workers.map(|worker| worker.join().unwrap()); - assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 1); - - let rejected = results - .into_iter() - .find_map(Result::err) - .expect("one completion must be rejected") - .into_inner(); - let completed = *completion.wait().await.unwrap(); - assert_ne!(completed, rejected); -} diff --git a/benchmarks/asyncband/completion/mod.rs b/benchmarks/asyncband/completion/mod.rs index 82c0e51..2676a97 100644 --- a/benchmarks/asyncband/completion/mod.rs +++ b/benchmarks/asyncband/completion/mod.rs @@ -31,17 +31,16 @@ 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::channel(); + let (completer, completion) = completion::new(); completer.complete(1usize).unwrap(); bencher.bench_local(|| black_box(*poll_ready(completion.wait(), &mut context).unwrap())); - black_box(completer); } #[divan::bench] fn cancel_pending(bencher: Bencher) { let mut context = bench_context(); - let (_completer, completion) = completion::channel::(); + let (_completer, completion) = completion::new::(); bencher.bench_local(|| { let mut wait = pin!(completion.wait()); @@ -55,7 +54,7 @@ fn complete_then_wait(bencher: Bencher) { let mut context = bench_context(); bencher.bench_local(|| { - let (completer, completion) = black_box(completion::channel()); + let (completer, completion) = black_box(completion::new()); completer.complete(black_box(1usize)).unwrap(); black_box(*poll_ready(completion.wait(), &mut context).unwrap()) }); @@ -66,7 +65,7 @@ fn notify_pending(bencher: Bencher) { let mut context = bench_context(); bencher.bench_local(|| { - let (completer, completion) = black_box(completion::channel()); + let (completer, completion) = black_box(completion::new()); let mut wait = pin!(completion.wait()); poll_pending(wait.as_mut(), &mut context); @@ -80,7 +79,7 @@ fn notify_pending_fanout(bencher: Bencher, observer_count: usize) { let mut context = bench_context(); bencher.bench_local(|| { - let (completer, first) = black_box(completion::channel()); + let (completer, first) = black_box(completion::new()); let mut observers = Vec::with_capacity(observer_count); observers.push(first); for _ in 1..observer_count { diff --git a/examples/src/shared_completion.rs b/examples/src/shared_completion.rs index dc49b96..4087927 100644 --- a/examples/src/shared_completion.rs +++ b/examples/src/shared_completion.rs @@ -25,7 +25,7 @@ use asyncband::completion; #[tokio::main(flavor = "current_thread")] async fn main() { - let (completer, completion) = completion::channel(); + let (completer, completion) = completion::new(); let dashboard = completion.clone(); let audit_log = completion.clone(); diff --git a/tests-integration/tests/completion_test.rs b/tests-integration/tests/completion_test.rs index 238af1e..977487e 100644 --- a/tests-integration/tests/completion_test.rs +++ b/tests-integration/tests/completion_test.rs @@ -68,18 +68,6 @@ struct DropCallbackWake(Mutex>>); struct CloneCallbackWake(Mutex>>); -struct ReentrantRejected { - completer: Option>>, -} - -impl Drop for ReentrantRejected { - fn drop(&mut self) { - if let Some(completer) = self.completer.take() { - drop(completer.complete(Self { completer: None })); - } - } -} - // This test needs a custom waker whose final `Arc` drop is observable. #[allow(clippy::manual_noop_waker)] impl Wake for DropCallbackWake { @@ -137,13 +125,29 @@ 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::channel(); + let (completer, completion) = completion::new(); let first = completion.clone(); let second = completion.clone(); - completer.complete(NotClone(String::from("ready"))).unwrap(); + 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(); @@ -159,7 +163,7 @@ fn all_observers_borrow_the_same_non_clone_value() { #[test] fn completer_transfers_a_send_only_value_between_threads() { - let (completer, completion) = completion::channel::>(); + let (completer, completion) = completion::new::>(); thread::spawn(move || completer.complete(Cell::new(7))) .join() @@ -170,79 +174,38 @@ fn completer_transfers_a_send_only_value_between_threads() { } #[test] -fn rejected_completions_return_the_original_value() { - let (completer, completion) = completion::channel(); - completer.complete(String::from("winner")).unwrap(); - - let error = completer.complete(String::from("rejected")).unwrap_err(); - assert_eq!(error.as_inner(), "rejected"); - assert_eq!(error.into_inner(), "rejected"); - assert_eq!(pollster::block_on(completion.wait()).unwrap(), "winner"); - - let (completer, completion) = completion::channel(); +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() - .into_inner(), + completer.complete(String::from("unobserved")).unwrap_err(), "unobserved" ); } #[test] -fn rejected_payloads_are_dropped_outside_the_completion_lock() { - let (finished_tx, finished_rx) = std::sync::mpsc::channel(); - let worker = thread::spawn(move || { - let (completer, _completion) = completion::channel(); - let completer = Arc::new(completer); - completer - .complete(ReentrantRejected { completer: None }) - .unwrap(); - - drop(completer.complete(ReentrantRejected { - completer: Some(completer.clone()), - })); - finished_tx.send(()).unwrap(); - }); - - finished_rx - .recv_timeout(Duration::from_secs(10)) - .expect("rejected payload destructor deadlocked against the completion lock"); - worker.join().unwrap(); -} - -#[test] -fn completer_drop_closes_every_observer_but_completed_drop_does_not() { - let (completer, first) = completion::channel::(); +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::WaitError::Closed) - ); + assert_eq!(pollster::block_on(first.wait()), Err(completion::Abandoned)); assert_eq!( pollster::block_on(second.wait()), - Err(completion::WaitError::Closed) + Err(completion::Abandoned) ); - - let (completer, completion) = completion::channel(); - completer.complete(42).unwrap(); - drop(completer); - assert_eq!(pollster::block_on(completion.wait()), Ok(&42)); } #[test] fn dropping_one_observer_before_or_after_completion_does_not_affect_another() { - let (completer, first) = completion::channel(); + 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::channel(); + let (completer, first) = completion::new(); let second = first.clone(); completer.complete(6).unwrap(); drop(first); @@ -253,20 +216,18 @@ fn dropping_one_observer_before_or_after_completion_does_not_affect_another() { #[test] fn completed_payload_is_released_with_the_last_observer() { let payload = Arc::new(()); - let (completer, completion) = completion::channel(); + 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); - - drop(completer); } #[test] -fn completer_drop_wakes_all_registered_waits() { - let (completer, first) = completion::channel::(); +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))); @@ -283,34 +244,34 @@ fn completer_drop_wakes_all_registered_waits() { assert_eq!(second_tracker.0.load(Ordering::Relaxed), 1); assert_eq!( poll_with(first_wait.as_mut(), &first_waker), - Poll::Ready(Err(completion::WaitError::Closed)) + Poll::Ready(Err(completion::Abandoned)) ); assert_eq!( poll_with(second_wait.as_mut(), &second_waker), - Poll::Ready(Err(completion::WaitError::Closed)) + Poll::Ready(Err(completion::Abandoned)) ); } #[test] -fn payload_errors_remain_distinct_from_completion_closure() { - let (completer, completion) = completion::channel::>(); +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::channel::>(); + let (completer, completion) = completion::new::>(); drop(completer); assert_eq!( pollster::block_on(completion.wait()), - Err(completion::WaitError::Closed) + Err(completion::Abandoned) ); } #[test] fn cancelling_a_wait_releases_only_its_waker() { - let (completer, completion) = completion::channel(); + 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()); @@ -336,7 +297,7 @@ fn cancelling_a_wait_releases_only_its_waker() { #[test] fn cancelling_after_wake_does_not_consume_the_shared_result() { - let (completer, first) = completion::channel(); + let (completer, first) = completion::new(); let second = first.clone(); let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); let waker = Waker::from(tracker.clone()); @@ -352,7 +313,7 @@ fn cancelling_after_wake_does_not_consume_the_shared_result() { #[test] fn cancellation_and_completer_drop_have_clean_orderings() { - let (completer, completion) = completion::channel::(); + 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); @@ -366,10 +327,10 @@ fn cancellation_and_completer_drop_have_clean_orderings() { assert_eq!(tracker.0.load(Ordering::Relaxed), 0); assert_eq!( pollster::block_on(completion.wait()), - Err(completion::WaitError::Closed) + Err(completion::Abandoned) ); - let (completer, completion) = completion::channel::(); + 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); @@ -383,13 +344,13 @@ fn cancellation_and_completer_drop_have_clean_orderings() { assert_eq!(Arc::strong_count(&tracker), baseline); assert_eq!( pollster::block_on(completion.wait()), - Err(completion::WaitError::Closed) + Err(completion::Abandoned) ); } #[test] fn completion_attempts_every_waker_after_one_panics() { - let (completer, first) = completion::channel(); + let (completer, first) = completion::new(); let second = first.clone(); let panicking = Waker::from(Arc::new(PanicWake)); let tracker = Arc::new(TrackWake(AtomicUsize::new(0))); @@ -411,119 +372,111 @@ fn completion_attempts_every_waker_after_one_panics() { #[test] fn wake_callbacks_run_outside_the_completion_lock() { - let (finished_tx, finished_rx) = std::sync::mpsc::channel(); - let worker = thread::spawn(move || { - let (completer, completion) = completion::channel(); - let callback_completion = completion.clone(); - let waker = Waker::from(Arc::new(WakeCallback(Mutex::new(Some(Box::new( - move || { - drop(callback_completion); - }, - )))))); - 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::channel::(); - let callback_completion = completion.clone(); - let waker = Waker::from(Arc::new(WakeCallback(Mutex::new(Some(Box::new( - move || { - drop(callback_completion); - }, - )))))); - 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::WaitError::Closed)) - ); - - finished_tx.send(()).unwrap(); - }); - - finished_rx - .recv_timeout(Duration::from_secs(10)) - .expect("wake callback deadlocked against the completion lock"); - worker.join().unwrap(); + 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() { - let (finished_tx, finished_rx) = std::sync::mpsc::channel(); - let worker = thread::spawn(move || { - let (_completer, completion) = completion::channel::(); - let callback_completion = completion.clone(); - let waker = waker_with_clone_callback(move || drop(callback_completion)); - let mut wait = Box::pin(completion.wait()); - - assert!(poll_with(wait.as_mut(), &waker).is_pending()); - finished_tx.send(()).unwrap(); - }); - - finished_rx - .recv_timeout(Duration::from_secs(10)) - .expect("waker clone callback deadlocked against the completion lock"); - worker.join().unwrap(); + 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() { - let (finished_tx, finished_rx) = std::sync::mpsc::channel(); - let worker = thread::spawn(move || { - let (_completer, completion) = completion::channel::(); - let callback_completion = completion.clone(); - let old_waker = Waker::from(Arc::new(DropCallbackWake(Mutex::new(Some(Box::new( - move || { - drop(callback_completion); - }, - )))))); - let mut wait = Box::pin(completion.wait()); - assert!(poll_with(wait.as_mut(), &old_waker).is_pending()); - drop(old_waker); - - let replacement = Waker::from(Arc::new(TrackWake(AtomicUsize::new(0)))); - assert!(poll_with(wait.as_mut(), &replacement).is_pending()); - finished_tx.send(()).unwrap(); - }); - - finished_rx - .recv_timeout(Duration::from_secs(10)) - .expect("replaced waker destructor deadlocked against the completion lock"); - worker.join().unwrap(); + 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() { - let (finished_tx, finished_rx) = std::sync::mpsc::channel(); - let worker = thread::spawn(move || { - let (_completer, completion) = completion::channel::(); - let callback_completion = completion.clone(); - let waker = Waker::from(Arc::new(DropCallbackWake(Mutex::new(Some(Box::new( - move || { - drop(callback_completion); - }, - )))))); - let mut wait = Box::pin(completion.wait()); - assert!(poll_with(wait.as_mut(), &waker).is_pending()); - drop(waker); - drop(wait); - finished_tx.send(()).unwrap(); - }); - - finished_rx - .recv_timeout(Duration::from_secs(10)) - .expect("cancelled waker destructor deadlocked against the completion lock"); - worker.join().unwrap(); + 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::channel(); + let (completer, completion) = completion::new(); let barrier = Arc::new(Barrier::new(2)); thread::scope(|scope| { let worker_barrier = barrier.clone(); @@ -536,7 +489,7 @@ fn complete_and_final_observer_drop_linearize_cleanly() { let result = completer.complete(17); worker.join().unwrap(); if let Err(error) = result { - assert_eq!(error.into_inner(), 17); + assert_eq!(error, 17); } }); } diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 51bbb94..34ff506 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -74,8 +74,7 @@ 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::>(); @@ -128,7 +127,7 @@ fn movable_public_types_are_send() { let (_tx, mut rx) = watch::channel(0); assert_send_value(rx.changed()); - let (_completer, completion) = completion::channel::(); + let (_completer, completion) = completion::new::(); assert_send_value(completion.wait()); } @@ -140,8 +139,7 @@ fn public_types_are_unpin() { assert_unpin::(); assert_unpin::>(); assert_unpin::>(); - assert_unpin::>(); - assert_unpin::(); + assert_unpin::(); assert_unpin::(); assert_unpin::>>(); assert_unpin::();