From 6d9b1c659a9df341289a9141c479ded948a011e9 Mon Sep 17 00:00:00 2001 From: Matthew Hambrecht Date: Thu, 20 Aug 2026 23:42:10 -0400 Subject: [PATCH 1/4] feat: Add async LazyLock implementation --- HISTORY.md | 1 + README.md | 1 + asyncband/Cargo.toml | 1 + asyncband/src/lib.rs | 9 +- asyncband/src/once/lazy_lock/mod.rs | 405 ++++++++++++++++++++++ asyncband/src/once/mod.rs | 6 + asyncband/src/once/once_cell/mod.rs | 8 +- tests-integration/Cargo.toml | 1 + tests-integration/tests/lazy_lock_test.rs | 240 +++++++++++++ tests-integration/tests/traits_test.rs | 3 + 10 files changed, 672 insertions(+), 3 deletions(-) create mode 100644 asyncband/src/once/lazy_lock/mod.rs create mode 100644 tests-integration/tests/lazy_lock_test.rs diff --git a/HISTORY.md b/HISTORY.md index f570577..3e559d9 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -10,6 +10,7 @@ Asyncband collects runtime-agnostic synchronization primitives informed by sever - `condvar::Condvar` is inspired by [`std::sync::Condvar`](https://doc.rust-lang.org/std/sync/struct.Condvar.html) and [`async_std::sync::Condvar`](https://docs.rs/async-std/latest/async_std/sync/struct.Condvar.html), with a fair FIFO waiter queue and standard non-buffered notification semantics. - `latch::Latch` is inspired by [`latches`](https://github.com/mirromutth/latches), with a different implementation based on the internal `CountdownState` primitive. - `mutex::Mutex` is derived from [`tokio::sync::Mutex`](https://docs.rs/tokio/latest/tokio/sync/struct.Mutex.html). +- `once::LazyLock` is inspired by [`std::sync::LazyLock`](https://doc.rust-lang.org/std/sync/struct.LazyLock.html) and [`async-lazy`](https://github.com/Jules-Bertholet/async-lazy), with restart-on-cancellation semantics built from Asyncband primitives. - `once::OnceCell` is derived from [`tokio::sync::OnceCell`](https://docs.rs/tokio/latest/tokio/sync/struct.OnceCell.html), but uses Asyncband's semaphore implementation. - `once::OnceMap` is inspired by [`uv-once-map`](https://github.com/astral-sh/uv/tree/main/crates/uv-once-map), with a redesigned interface and implementation. - `oneshot::channel` is derived from the [`oneshot`](https://github.com/faern/oneshot) crate, with significant simplifications because Asyncband does not provide synchronized receive operations. diff --git a/README.md b/README.md index 8f59d4d..1c8d9b3 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,7 @@ The crate enables no primitives by default. Categories describe each primitive's | | [`Condvar`](https://docs.rs/asyncband/*/asyncband/condvar/struct.Condvar.html) | `condvar` | Wait for notifications while releasing a mutex. | | One-time initialization | [`Once`](https://docs.rs/asyncband/*/asyncband/once/struct.Once.html) | `once` | Run asynchronous initialization exactly once. | | | [`OnceCell`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceCell.html) | `once-cell` | Initialize and store one asynchronous value. | +| | [`LazyLock`](https://docs.rs/asyncband/*/asyncband/once/struct.LazyLock.html) | `lazy-lock` | 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. | | | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a one-way countdown completes. | diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index c36f13f..af6a4ee 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -43,6 +43,7 @@ blocking = [] broadcast = [] condvar = ["mutex"] latch = [] +lazy-lock = ["mutex", "once-cell"] mpsc = [] mutex = [] once = ["semaphore"] diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 3d810e9..040d403 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -54,7 +54,7 @@ //! | Use case | APIs | Cargo features | //! | -------------------------- | ------------------------------------------------------------------------------------- | ------------------------------------------- | //! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | -//! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::OnceMap`] | `once`, `once-cell`, `once-map` | +//! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::LazyLock`], [`once::OnceMap`] | `once`, `once-cell`, `lazy-lock`, `once-map` | //! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | //! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::overflow`] | `oneshot`, `mpsc`, `broadcast` | //! | Control workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | @@ -103,7 +103,12 @@ pub mod latch; pub mod mpsc; #[cfg(feature = "mutex")] pub mod mutex; -#[cfg(any(feature = "once", feature = "once-cell", feature = "once-map"))] +#[cfg(any( + feature = "lazy-lock", + feature = "once", + feature = "once-cell", + feature = "once-map" +))] pub mod once; #[cfg(feature = "oneshot")] pub mod oneshot; diff --git a/asyncband/src/once/lazy_lock/mod.rs b/asyncband/src/once/lazy_lock/mod.rs new file mode 100644 index 0000000..e37bb80 --- /dev/null +++ b/asyncband/src/once/lazy_lock/mod.rs @@ -0,0 +1,405 @@ +// 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::convert::Infallible; +use std::fmt; +use std::future::Future; +use std::panic::RefUnwindSafe; +use std::panic::UnwindSafe; +use std::pin::Pin; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use super::OnceCell; +use crate::mutex::Mutex; + +/// A boxed future suitable for the default [`LazyLock`] initializer type. +pub type LazyLockFuture = Pin + Send + 'static>>; + +/// A value initialized by an asynchronous function on first access. +/// +/// Initialization starts when [`force`](Self::force) or +/// [`try_force`](Self::try_force) is polled. Concurrent callers wait without +/// blocking their threads. +/// +/// If an attempt is cancelled, its future is dropped and the initializer is +/// retained. The next caller starts a new attempt. Initializers must therefore +/// be safe to invoke again after cancellation, therefore it is up to the user +/// to ensure idempotency. +/// +/// # Poisoning +/// +/// A panic from the initializer permanently poisons the lock. The panic is +/// propagated to its caller, and future calls to `force`, `try_force`, +/// `force_mut`, or `try_force_mut` panic. Errors returned through `Result` do +/// not poison the lock and are there to indicate initialization is possible. +/// +/// # Examples +/// +/// ``` +/// # #[tokio::main] +/// # async fn main() { +/// use asyncband::once::LazyLock; +/// +/// let lazy = LazyLock::::new(async || "ready".to_owned()); +/// +/// assert_eq!(LazyLock::get(&lazy), None); +/// assert_eq!(LazyLock::force(&lazy).await, "ready"); +/// assert_eq!(LazyLock::get(&lazy).map(String::as_str), Some("ready")); +/// # } +/// ``` +pub struct LazyLock LazyLockFuture> { + value: OnceCell, + initializer: Mutex>, + poisoned: AtomicBool, +} + +impl LazyLock { + /// Creates a new lazy value with the given asynchronous initializer. + /// + /// The initializer is not called until the first initialization future is + /// polled. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyLock; + /// + /// let lazy = LazyLock::::new(async || 92); + /// assert_eq!(*LazyLock::force(&lazy).await, 92); + /// # } + /// ``` + pub const fn new(initializer: F) -> Self { + Self { + value: OnceCell::new(), + initializer: Mutex::new(Some(initializer)), + poisoned: AtomicBool::new(false), + } + } + + /// Returns a reference to the value if initialized. + /// + /// This method never starts initialization or waits for an active attempt. + /// It returns `None` when the lock is uninitialized, initializing, or + /// poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyLock; + /// + /// let lazy = LazyLock::::new(async || 92); + /// assert_eq!(LazyLock::get(&lazy), None); + /// LazyLock::force(&lazy).await; + /// assert_eq!(LazyLock::get(&lazy), Some(&92)); + /// # } + /// ``` + pub fn get(this: &Self) -> Option<&T> { + this.value.get() + } + + /// Returns a mutable reference to the value if initialized. + /// + /// This method never starts initialization. It returns `None` when the lock + /// is uninitialized or poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyLock; + /// + /// let mut lazy = LazyLock::::new(async || 92); + /// assert_eq!(LazyLock::get_mut(&mut lazy), None); + /// LazyLock::force(&lazy).await; + /// *LazyLock::get_mut(&mut lazy).unwrap() = 44; + /// assert_eq!(LazyLock::get(&lazy), Some(&44)); + /// # } + /// ``` + pub fn get_mut(this: &mut Self) -> Option<&mut T> { + this.value.get_mut() + } + + /// Consumes the lock and returns its value or initializer. + /// + /// Returns `Ok(value)` when initialized and `Err(initializer)` otherwise. + /// + /// # Panics + /// + /// Panics if the lock is poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyLock; + /// + /// let lazy = LazyLock::::new(async || 92); + /// LazyLock::force(&lazy).await; + /// assert_eq!(LazyLock::into_inner(lazy).ok(), Some(92)); + /// # } + /// ``` + pub fn into_inner(this: Self) -> Result { + let Self { + value, + initializer, + poisoned, + } = this; + + if poisoned.into_inner() { + panic_poisoned(); + } + + match value.into_inner() { + Some(value) => Ok(value), + None => Err(initializer + .into_inner() + .expect("LazyLock initializer missing while uninitialized")), + } + } + + /// Internal helper to initialize the value with a fallible initializer. + /// Waiters are queued to restart on cancellation or error, winner removes + /// the initializer and ensures those waiters see the value. + async fn initialize(&self, run: G) -> Result<&T, E> + where + G: AsyncFnOnce(&mut F) -> Result, + { + if let Some(value) = self.value.get() { + return Ok(value); + } + self.assert_unpoisoned(); + + let mut initializer = self.initializer.lock().await; + if let Some(value) = self.value.get() { + return Ok(value); + } + self.assert_unpoisoned(); + + // Running initializer under a lock ensures we queue waiters to restart + // on cancellation or Err. + let _poison = PoisonOnPanic(&self.poisoned); + let value = run(initializer + .as_mut() + .expect("LazyLock initializer missing while uninitialized")) + .await?; + + drop(initializer.take()); // Avoid reinitialization possibility + let value = unsafe { self.value.set_value_unchecked(value) }; + + Ok(value) + } + + /// Panics if the lock is poisoned. + fn assert_unpoisoned(&self) { + if self.poisoned.load(Ordering::Acquire) { + panic_poisoned(); + } + } +} + +impl LazyLock +where + F: AsyncFnMut() -> T, +{ + /// Initializes the value if needed and returns a reference to it. + /// + /// If another task is initializing the lock, this call waits for that + /// attempt. If the active attempt is cancelled, this or another waiting + /// caller starts a new attempt. + /// + /// # Panics + /// + /// Panics if the initializer panics or the lock was previously poisoned. + /// Recursive initialization of the same lock deadlocks. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyLock; + /// + /// let lazy = LazyLock::::new(async || 92); + /// assert_eq!(LazyLock::force(&lazy).await, &92); + /// # } + /// ``` + pub async fn force(this: &Self) -> &T { + // Templated to Infalliable for non-monadic initializers + match this + .initialize(async |initializer| Ok::(initializer().await)) + .await + { + Ok(value) => value, + Err(error) => match error {}, + } + } + + /// Initializes the value if needed and returns a mutable reference to it. + /// + /// # Panics + /// + /// Panics if the initializer panics or the lock was previously poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyLock; + /// + /// let mut lazy = LazyLock::::new(async || 92); + /// *LazyLock::force_mut(&mut lazy).await = 44; + /// assert_eq!(LazyLock::get(&lazy), Some(&44)); + /// # } + /// ``` + pub async fn force_mut(this: &mut Self) -> &mut T { + let _ = Self::force(this).await; + this.value + .get_mut() + .expect("LazyLock value missing after success") + } +} + +impl LazyLock { + /// Initializes the value with a fallible initializer. + /// + /// An error is returned only to the caller whose attempt produced it. The + /// lock remains uninitialized, and the next waiting caller starts a new + /// serialized attempt. Initializers must therefore be idempotent and safe + /// to invoke again after cancellation or error. + /// + /// # Panics + /// + /// Panics if the initializer panics or the lock was previously poisoned. + /// Recursive initialization of the same lock deadlocks. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyLock; + /// + /// let lazy = LazyLock::::new(async || Ok::<_, std::io::Error>(92)); + /// assert_eq!(LazyLock::try_force(&lazy).await.unwrap(), &92); + /// # } + /// ``` + pub async fn try_force(this: &Self) -> Result<&T, E> + where + F: AsyncFnMut() -> Result, + { + this.initialize(async |initializer| initializer().await) + .await + } + + /// Initializes the value with a fallible initializer and returns mutable access. + /// + /// An error leaves the lock uninitialized so a later caller can retry. + /// + /// # Panics + /// + /// Panics if the initializer panics or the lock was previously poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyLock; + /// + /// let mut lazy = LazyLock::::new(async || Ok::<_, ()>(92)); + /// *LazyLock::try_force_mut(&mut lazy).await.unwrap() = 44; + /// assert_eq!(LazyLock::get(&lazy), Some(&44)); + /// # } + /// ``` + pub async fn try_force_mut(this: &mut Self) -> Result<&mut T, E> + where + F: AsyncFnMut() -> Result, + { + let _ = Self::try_force(this).await?; + Ok(this + .value + .get_mut() + .expect("LazyLock value missing after success")) + } +} + +impl Default for LazyLock +where + T: Default + Send + 'static, +{ + /// Creates a lazy value initialized with [`Default::default`]. + fn default() -> Self { + fn initialize() -> LazyLockFuture + where + T: Default + Send + 'static, + { + Box::pin(async { T::default() }) + } + + Self::new(initialize::) + } +} + +impl From for LazyLock { + /// Creates an already initialized lazy value. + fn from(value: T) -> Self { + Self { + value: OnceCell::from_value(value), + initializer: Mutex::new(None), + poisoned: AtomicBool::new(false), + } + } +} + +impl fmt::Debug for LazyLock { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut tuple = f.debug_tuple("LazyLock"); + match Self::get(self) { + Some(value) => tuple.field(value), + None => tuple.field(&format_args!("")), + }; + tuple.finish() + } +} + +impl UnwindSafe for LazyLock {} + +impl RefUnwindSafe for LazyLock {} + +struct PoisonOnPanic<'a>(&'a AtomicBool); + +impl Drop for PoisonOnPanic<'_> { + fn drop(&mut self) { + if std::thread::panicking() { + self.0.store(true, Ordering::Release); + } + } +} + +#[cold] +#[inline(never)] +fn panic_poisoned() -> ! { + panic!("LazyLock instance has previously been poisoned") +} diff --git a/asyncband/src/once/mod.rs b/asyncband/src/once/mod.rs index e2f7e86..2e2bbbf 100644 --- a/asyncband/src/once/mod.rs +++ b/asyncband/src/once/mod.rs @@ -17,6 +17,8 @@ //! Asynchronous primitives for one-time coordination. +#[cfg(feature = "lazy-lock")] +mod lazy_lock; #[cfg(feature = "once")] #[allow(clippy::module_inception)] mod once; @@ -25,6 +27,10 @@ mod once_cell; #[cfg(feature = "once-map")] mod once_map; +#[cfg(feature = "lazy-lock")] +pub use self::lazy_lock::LazyLock; +#[cfg(feature = "lazy-lock")] +pub use self::lazy_lock::LazyLockFuture; #[cfg(feature = "once")] pub use self::once::Once; #[cfg(feature = "once-cell")] diff --git a/asyncband/src/once/once_cell/mod.rs b/asyncband/src/once/once_cell/mod.rs index 3fdd906..a07c7b2 100644 --- a/asyncband/src/once/once_cell/mod.rs +++ b/asyncband/src/once/once_cell/mod.rs @@ -409,9 +409,15 @@ impl OnceCell { } fn set_value(&self, value: T, permit: SemaphorePermit<'_>) -> &T { - // Hold the permit to ensure exclusive access. let _permit = permit; + unsafe { self.set_value_unchecked(value) } + } + /// # Safety + /// + /// No other initialization may access the cell. + pub(crate) unsafe fn set_value_unchecked(&self, value: T) -> &T { + debug_assert!(!self.initialized()); let value_ptr = self.value.get(); unsafe { value_ptr.write(MaybeUninit::new(value)) }; diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 104171c..16c7221 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -32,6 +32,7 @@ asyncband = { workspace = true, features = [ "broadcast", "condvar", "latch", + "lazy-lock", "mpsc", "mutex", "once", diff --git a/tests-integration/tests/lazy_lock_test.rs b/tests-integration/tests/lazy_lock_test.rs new file mode 100644 index 0000000..e46ad34 --- /dev/null +++ b/tests-integration/tests/lazy_lock_test.rs @@ -0,0 +1,240 @@ +// 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::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use asyncband::once::LazyLock; +use asyncband::once::LazyLockFuture; +use tokio::sync::Notify; + +#[tokio::test] +/// Ensure that multiple concurrent calls to a successful `force` only run the +/// initializer once. +async fn force_runs_initializer_once() { + let attempts = Arc::new(AtomicUsize::new(0)); + let lazy = Arc::new(LazyLock::::new({ + let attempts = attempts.clone(); + async move || { + attempts.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + 42 + } + })); + + let mut tasks = Vec::new(); + for _ in 0..32 { + let lazy = lazy.clone(); + tasks.push(tokio::spawn(async move { *LazyLock::force(&lazy).await })); + } + + for task in tasks { + assert_eq!(task.await.unwrap(), 42); + } + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +/// Ensure that cancellation of a `force` call does not prevent future calls from +/// rerunning the initializer. +async fn cancellation_restarts_initialization() { + let attempts = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Notify::new()); + let lazy = Arc::new(LazyLock::::new({ + let attempts = attempts.clone(); + let started = started.clone(); + async move || { + let attempt = attempts.fetch_add(1, Ordering::SeqCst); + if attempt == 0 { + started.notify_one(); + std::future::pending::<()>().await; + } + 42 + } + })); + + let task = { + let lazy = lazy.clone(); + tokio::spawn(async move { *LazyLock::force(&lazy).await }) + }; + started.notified().await; + assert_eq!(LazyLock::get(&lazy), None); + + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + + assert_eq!(*LazyLock::force(&lazy).await, 42); + assert_eq!(attempts.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +/// Validates falliabile initialization. If the initializer returns an error, +/// the value is not stored and future calls may retry it. +async fn queued_callers_retry_after_errors() { + let attempts = Arc::new(AtomicUsize::new(0)); + let lazy = Arc::new(LazyLock::::new({ + let attempts = attempts.clone(); + async move || { + let attempt = attempts.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + if attempt < 2 { Err("retry") } else { Ok(42) } + } + })); + + let mut tasks = Vec::new(); + for _ in 0..3 { + let lazy = lazy.clone(); + tasks.push(tokio::spawn(async move { + LazyLock::try_force(&lazy).await.copied() + })); + } + + let mut errors = 0; + let mut successes = 0; + for task in tasks { + match task.await.unwrap() { + Ok(42) => successes += 1, + Err("retry") => errors += 1, + result => panic!("unexpected result: {result:?}"), + } + } + + assert_eq!(errors, 2); + assert_eq!(successes, 1); + assert_eq!(attempts.load(Ordering::SeqCst), 3); + assert_eq!(LazyLock::try_force(&lazy).await, Ok(&42)); + assert_eq!(attempts.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +/// Ensure that a panic in the initializer permanently poisons the lock, preventing future calls +/// from succeeding. +async fn panic_permanently_poisons_lock() { + let attempts = Arc::new(AtomicUsize::new(0)); + let lazy = Arc::new(LazyLock::::new({ + let attempts = attempts.clone(); + async move || { + attempts.fetch_add(1, Ordering::SeqCst); + panic!("initializer panic"); + } + })); + + let first = { + let lazy = lazy.clone(); + tokio::spawn(async move { + let _ = LazyLock::force(&lazy).await; + }) + }; + assert!(first.await.unwrap_err().is_panic()); + assert_eq!(LazyLock::get(&lazy), None); + + let second = { + let lazy = lazy.clone(); + tokio::spawn(async move { + let _ = LazyLock::force(&lazy).await; + }) + }; + assert!(second.await.unwrap_err().is_panic()); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + + let lazy = Arc::try_unwrap(lazy).ok().unwrap(); + let result = std::panic::catch_unwind(|| LazyLock::into_inner(lazy)); + assert!(result.is_err()); +} + +#[tokio::test] +/// Ensure that `force_mut` and `try_force_mut` can be used to mutate the value +/// after it has been initialized. +async fn mutable_force_updates_value() { + let mut lazy = LazyLock::::new(async || 41); + *LazyLock::force_mut(&mut lazy).await += 1; + assert_eq!(LazyLock::get(&lazy), Some(&42)); + + let mut fallible = LazyLock::::new(async || Ok::<_, ()>(41)); + *LazyLock::try_force_mut(&mut fallible).await.unwrap() += 1; + assert_eq!(LazyLock::get(&fallible), Some(&42)); +} + +#[tokio::test] +/// Ensure that `into_inner` returns the value if it has been initialized, or +/// returns the initializer if it has not been initialized. +async fn into_inner_returns_value_or_initializer() { + let lazy = LazyLock::::new(async || 42); + let initializer = LazyLock::into_inner(lazy).unwrap_err(); + assert_eq!(initializer().await, 42); + + let lazy = LazyLock::::new(async || 42); + LazyLock::force(&lazy).await; + assert!(matches!(LazyLock::into_inner(lazy), Ok(42))); +} + +#[tokio::test] +/// Validates that `Debug` and `Default` trait implementations work as +/// expected. +async fn default_from_and_debug_match_lazy_lock() { + let lazy = LazyLock::::default(); + assert_eq!(format!("{lazy:?}"), "LazyLock()"); + assert_eq!(LazyLock::force(&lazy).await, &0); + assert_eq!(format!("{lazy:?}"), "LazyLock(0)"); + + let lazy: LazyLock = LazyLock::from(42); + assert_eq!(LazyLock::get(&lazy), Some(&42)); +} + +#[tokio::test] +/// Ensure that the initializer does not need to be `Sync` in order for the `LazyLock` to be +/// `Sync`. This is important for cases where the initializer captures non-`Sync` state, such as a +/// `Cell`. This is guaranteed by the internal mutex. +async fn initializer_need_not_be_sync() { + fn assert_sync(_: &T) {} + + let count = Cell::new(0); + let lazy = LazyLock::::new(async move || { + count.set(count.get() + 1); + count.get() + }); + + assert_sync(&lazy); + assert_eq!(LazyLock::force(&lazy).await, &1); +} + +fn static_initializer() -> LazyLockFuture { + Box::pin(async { 42 }) +} + +static STATIC_LAZY: LazyLock = LazyLock::new(static_initializer); + +#[tokio::test] +async fn default_initializer_type_supports_statics() { + assert_eq!(LazyLock::force(&STATIC_LAZY).await, &42); +} + +fn fallible_static_initializer() -> LazyLockFuture> { + Box::pin(async { Ok(42) }) +} + +type FallibleStaticInitializer = fn() -> LazyLockFuture>; + +static FALLIBLE_STATIC_LAZY: LazyLock = + LazyLock::new(fallible_static_initializer); + +#[tokio::test] +async fn one_type_supports_fallible_statics() { + assert_eq!(LazyLock::try_force(&FALLIBLE_STATIC_LAZY).await, Ok(&42)); +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index e894246..d074830 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -22,6 +22,7 @@ use asyncband::latch::Latch; use asyncband::mpsc; use asyncband::mutex::Mutex; use asyncband::mutex::MutexGuard; +use asyncband::once::LazyLock; use asyncband::once::Once; use asyncband::once::OnceCell; use asyncband::once::OnceMap; @@ -44,6 +45,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::>(); @@ -89,6 +91,7 @@ fn public_types_are_unpin() { assert_unpin::(); assert_unpin::(); assert_unpin::(); + assert_unpin::>(); assert_unpin::(); assert_unpin::>(); assert_unpin::>(); From 683cc8e3f01ed20ab1d5c64f7f96186b2888371a Mon Sep 17 00:00:00 2001 From: Matthew Hambrecht Date: Sat, 22 Aug 2026 21:29:37 -0400 Subject: [PATCH 2/4] add resumability + initialization args --- HISTORY.md | 2 +- README.md | 2 +- asyncband/Cargo.toml | 2 +- asyncband/src/lib.rs | 4 +- asyncband/src/once/lazy_cell/mod.rs | 668 ++++++++++++++++++++++ asyncband/src/once/lazy_lock/mod.rs | 405 ------------- asyncband/src/once/mod.rs | 12 +- tests-integration/Cargo.toml | 2 +- tests-integration/tests/lazy_cell_test.rs | 344 +++++++++++ tests-integration/tests/lazy_lock_test.rs | 240 -------- tests-integration/tests/traits_test.rs | 6 +- 11 files changed, 1027 insertions(+), 660 deletions(-) create mode 100644 asyncband/src/once/lazy_cell/mod.rs delete mode 100644 asyncband/src/once/lazy_lock/mod.rs create mode 100644 tests-integration/tests/lazy_cell_test.rs delete mode 100644 tests-integration/tests/lazy_lock_test.rs diff --git a/HISTORY.md b/HISTORY.md index 3e559d9..ad0b883 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -10,7 +10,7 @@ Asyncband collects runtime-agnostic synchronization primitives informed by sever - `condvar::Condvar` is inspired by [`std::sync::Condvar`](https://doc.rust-lang.org/std/sync/struct.Condvar.html) and [`async_std::sync::Condvar`](https://docs.rs/async-std/latest/async_std/sync/struct.Condvar.html), with a fair FIFO waiter queue and standard non-buffered notification semantics. - `latch::Latch` is inspired by [`latches`](https://github.com/mirromutth/latches), with a different implementation based on the internal `CountdownState` primitive. - `mutex::Mutex` is derived from [`tokio::sync::Mutex`](https://docs.rs/tokio/latest/tokio/sync/struct.Mutex.html). -- `once::LazyLock` is inspired by [`std::sync::LazyLock`](https://doc.rust-lang.org/std/sync/struct.LazyLock.html) and [`async-lazy`](https://github.com/Jules-Bertholet/async-lazy), with restart-on-cancellation semantics built from Asyncband primitives. +- `once::LazyCell` is inspired by [`std::sync::LazyLock`](https://doc.rust-lang.org/std/sync/struct.LazyLock.html) and [`async-lazy`](https://github.com/Jules-Bertholet/async-lazy), with resumable cancellation semantics built from Asyncband primitives. - `once::OnceCell` is derived from [`tokio::sync::OnceCell`](https://docs.rs/tokio/latest/tokio/sync/struct.OnceCell.html), but uses Asyncband's semaphore implementation. - `once::OnceMap` is inspired by [`uv-once-map`](https://github.com/astral-sh/uv/tree/main/crates/uv-once-map), with a redesigned interface and implementation. - `oneshot::channel` is derived from the [`oneshot`](https://github.com/faern/oneshot) crate, with significant simplifications because Asyncband does not provide synchronized receive operations. diff --git a/README.md b/README.md index 1c8d9b3..b18d892 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ The crate enables no primitives by default. Categories describe each primitive's | | [`Condvar`](https://docs.rs/asyncband/*/asyncband/condvar/struct.Condvar.html) | `condvar` | Wait for notifications while releasing a mutex. | | One-time initialization | [`Once`](https://docs.rs/asyncband/*/asyncband/once/struct.Once.html) | `once` | Run asynchronous initialization exactly once. | | | [`OnceCell`](https://docs.rs/asyncband/*/asyncband/once/struct.OnceCell.html) | `once-cell` | Initialize and store one asynchronous value. | -| | [`LazyLock`](https://docs.rs/asyncband/*/asyncband/once/struct.LazyLock.html) | `lazy-lock` | Lazily initialize a value with a stored asynchronous function. | +| | [`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. | | | [`Latch`](https://docs.rs/asyncband/*/asyncband/latch/struct.Latch.html) | `latch` | Wait until a one-way countdown completes. | diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index af6a4ee..c3dcdd8 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -43,7 +43,7 @@ blocking = [] broadcast = [] condvar = ["mutex"] latch = [] -lazy-lock = ["mutex", "once-cell"] +lazy-cell = ["mutex", "once-cell"] mpsc = [] mutex = [] once = ["semaphore"] diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 040d403..56d9305 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -54,7 +54,7 @@ //! | Use case | APIs | Cargo features | //! | -------------------------- | ------------------------------------------------------------------------------------- | ------------------------------------------- | //! | Protect shared state | [`mutex::Mutex`], [`rwlock::RwLock`], [`condvar::Condvar`] | `mutex`, `rwlock`, `condvar` | -//! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::LazyLock`], [`once::OnceMap`] | `once`, `once-cell`, `lazy-lock`, `once-map` | +//! | Initialize values once | [`once::Once`], [`once::OnceCell`], [`once::LazyCell`], [`once::OnceMap`] | `once`, `once-cell`, `lazy-cell`, `once-map` | //! | Coordinate tasks | [`barrier::Barrier`], [`latch::Latch`], [`waitgroup::WaitGroup`], [`shutdown`] | `barrier`, `latch`, `waitgroup`, `shutdown` | //! | Send values | [`oneshot::channel`], [`mpsc::bounded`], [`mpsc::unbounded`], [`broadcast::overflow`] | `oneshot`, `mpsc`, `broadcast` | //! | Control workloads | [`semaphore::Semaphore`], [`singleflight::Group`] | `semaphore`, `singleflight` | @@ -104,7 +104,7 @@ pub mod mpsc; #[cfg(feature = "mutex")] pub mod mutex; #[cfg(any( - feature = "lazy-lock", + feature = "lazy-cell", feature = "once", feature = "once-cell", feature = "once-map" diff --git a/asyncband/src/once/lazy_cell/mod.rs b/asyncband/src/once/lazy_cell/mod.rs new file mode 100644 index 0000000..04fd4d1 --- /dev/null +++ b/asyncband/src/once/lazy_cell/mod.rs @@ -0,0 +1,668 @@ +// 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::Any; +use std::any::TypeId; +use std::fmt; +use std::future::Future; +use std::panic::RefUnwindSafe; +use std::panic::UnwindSafe; +use std::pin::Pin; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use super::OnceCell; +use crate::mutex::Mutex; + +/// A boxed future suitable for the default [`LazyCell`] initializer type. +pub type LazyCellFuture = Pin + Send + 'static>>; + +/// A thread-safe value initialized by an asynchronous function on first access. +/// +/// Initialization starts when [`force`](Self::force) or +/// [`try_force`](Self::try_force) is polled. Concurrent callers wait without +/// blocking their threads. +/// +/// If a caller is cancelled, the initialization future remains pinned in the +/// cell. The next caller resumes that same future. For fallible initialization, +/// an error ends the current attempt and the next caller starts a new attempt. +/// The stored future must be `Send + 'static` because a different task may +/// resume it. Call-time arguments may contain references only when the returned +/// future does not retain them. +/// +/// Infallible initializers are called once and may move captured values directly +/// into their future. Fallible initializers must remain callable after an error. +/// A captured value can be cloned into each owned attempt: +/// +/// ``` +/// # #[tokio::main] +/// # async fn main() { +/// use std::sync::Arc; +/// +/// use asyncband::once::LazyCell; +/// +/// let client = Arc::new("client".to_owned()); +/// let lazy = LazyCell::::new(move || { +/// let client = client.clone(); +/// async move { Ok::<_, ()>(client.len()) } +/// }); +/// +/// assert_eq!(LazyCell::try_force(&lazy).await, Ok(&6)); +/// # } +/// ``` +/// +/// # Poisoning +/// +/// A panic from the initializer permanently poisons the cell. The panic is +/// propagated to its caller, and future calls to `force`, `try_force`, +/// `force_mut`, or `try_force_mut` panic. Errors returned through `Result` do +/// not poison the cell and allow a later caller to retry initialization. +/// +/// # Examples +/// +/// ``` +/// # #[tokio::main] +/// # async fn main() { +/// use asyncband::once::LazyCell; +/// +/// let lazy = LazyCell::::new(async || "ready".to_owned()); +/// +/// assert_eq!(LazyCell::get(&lazy), None); +/// assert_eq!(LazyCell::force(&lazy).await, "ready"); +/// assert_eq!(LazyCell::get(&lazy).map(String::as_str), Some("ready")); +/// # } +/// ``` +pub struct LazyCell LazyCellFuture> { + value: OnceCell, + state: Mutex>, + poisoned: AtomicBool, +} + +type Attempt = Pin> + Send + 'static>>; + +enum AttemptOutput { + Value(T), + Error(Box), +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum AttemptKind { + Infallible, + Fallible(TypeId), +} + +struct State { + initializer: Option, + attempt: Option<(AttemptKind, Attempt)>, +} + +impl LazyCell { + /// Creates a new lazy value with the given asynchronous initializer. + /// + /// The initializer is not called until the first initialization future is + /// polled. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let lazy = LazyCell::::new(async || 92); + /// assert_eq!(*LazyCell::force(&lazy).await, 92); + /// # } + /// ``` + pub const fn new(initializer: F) -> Self { + Self { + value: OnceCell::new(), + state: Mutex::new(State { + initializer: Some(initializer), + attempt: None, + }), + poisoned: AtomicBool::new(false), + } + } + + /// Returns a reference to the value if initialized. + /// + /// This method never starts initialization or waits for an active attempt. + /// It returns `None` when the cell is uninitialized, initializing, or + /// poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let lazy = LazyCell::::new(async || 92); + /// assert_eq!(LazyCell::get(&lazy), None); + /// LazyCell::force(&lazy).await; + /// assert_eq!(LazyCell::get(&lazy), Some(&92)); + /// # } + /// ``` + pub fn get(this: &Self) -> Option<&T> { + this.value.get() + } + + /// Returns a mutable reference to the value if initialized. + /// + /// This method never starts initialization. It returns `None` when the cell + /// is uninitialized or poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let mut lazy = LazyCell::::new(async || 92); + /// assert_eq!(LazyCell::get_mut(&mut lazy), None); + /// LazyCell::force(&lazy).await; + /// *LazyCell::get_mut(&mut lazy).unwrap() = 44; + /// assert_eq!(LazyCell::get(&lazy), Some(&44)); + /// # } + /// ``` + pub fn get_mut(this: &mut Self) -> Option<&mut T> { + this.value.get_mut() + } + + /// Consumes the cell and returns its value or initializer. + /// + /// Returns `Ok(value)` when initialized and `Err(initializer)` otherwise. + /// + /// # Panics + /// + /// Panics if the cell is poisoned or a one-shot initializer was started but + /// did not complete. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let lazy = LazyCell::::new(async || 92); + /// LazyCell::force(&lazy).await; + /// assert_eq!(LazyCell::into_inner(lazy).ok(), Some(92)); + /// # } + /// ``` + pub fn into_inner(this: Self) -> Result { + let Self { + value, + state, + poisoned, + } = this; + + if poisoned.into_inner() { + panic_poisoned(); + } + + let State { + initializer, + attempt, + } = state.into_inner(); + drop(attempt); + + match value.into_inner() { + Some(value) => Ok(value), + None => Err(initializer.expect("LazyCell one-shot initializer has already started")), + } + } + + /// Initializes the value if needed and returns a reference to it. + async fn initialize_once(&self, start: G) -> &T + where + G: FnOnce(F) -> Fut, + Fut: Future + Send + 'static, + { + if let Some(value) = self.value.get() { + return value; + } + self.assert_unpoisoned(); + + let mut state = self.state.lock().await; + if let Some(value) = self.value.get() { + return value; + } + self.assert_unpoisoned(); + + let _poison = PoisonOnPanic(&self.poisoned); + + if state.attempt.is_none() { + let initializer = state + .initializer + .take() + .expect("LazyCell initializer missing while uninitialized"); + let future = start(initializer); + let attempt = Box::pin(async move { AttemptOutput::Value(future.await) }); + state.attempt = Some((AttemptKind::Infallible, attempt)); + } + + let (kind, attempt) = state + .attempt + .as_mut() + .expect("LazyCell attempt missing while initializing"); + assert!( + *kind == AttemptKind::Infallible, + "LazyCell force method does not match the active attempt" + ); + + let output = attempt.as_mut().await; + state.attempt = None; + let AttemptOutput::Value(value) = output else { + unreachable!("infallible LazyCell attempt returned an error") + }; + unsafe { self.value.set_value_unchecked(value) } + } + + /// Initializes the value with a fallible initializer and returns a reference to it. + async fn initialize_retry(&self, start: G) -> Result<&T, E> + where + G: FnOnce(&mut F) -> Fut, + Fut: Future> + Send + 'static, + E: 'static, + { + if let Some(value) = self.value.get() { + return Ok(value); + } + self.assert_unpoisoned(); + + let mut state = self.state.lock().await; + if let Some(value) = self.value.get() { + return Ok(value); + } + self.assert_unpoisoned(); + + let _poison = PoisonOnPanic(&self.poisoned); + let kind = AttemptKind::Fallible(TypeId::of::()); + + if state.attempt.is_none() { + let future = state + .initializer + .as_mut() + .map(start) + .expect("LazyCell initializer missing while uninitialized"); + let attempt = Box::pin(async move { + match future.await { + Ok(value) => AttemptOutput::Value(value), + Err(error) => AttemptOutput::Error(Box::new(error)), + } + }); + state.attempt = Some((kind, attempt)); + } + + let (active_kind, attempt) = state + .attempt + .as_mut() + .expect("LazyCell attempt missing while initializing"); + assert!( + *active_kind == kind, + "LazyCell force method does not match the active attempt" + ); + + let output = attempt.as_mut().await; + state.attempt = None; + match output { + AttemptOutput::Value(value) => { + state.initializer = None; + Ok(unsafe { self.value.set_value_unchecked(value) }) + } + AttemptOutput::Error(error) => { + let error = error + .downcast::() + .expect("LazyCell attempt error type changed"); + Err(*error) + } + } + } + + /// Panics if the cell is poisoned. + fn assert_unpoisoned(&self) { + if self.poisoned.load(Ordering::Acquire) { + panic_poisoned(); + } + } +} + +impl LazyCell { + /// Initializes the value if needed and returns a reference to it. + /// + /// If another task is initializing the cell, this call waits for that + /// attempt. If its caller is cancelled, a later caller resumes the same + /// pinned future. + /// + /// # Panics + /// + /// Panics if the initializer panics or the cell was previously poisoned. + /// Recursive initialization of the same cell deadlocks. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let lazy = LazyCell::::new(async || 92); + /// assert_eq!(LazyCell::force(&lazy).await, &92); + /// # } + /// ``` + pub async fn force(this: &Self) -> &T + where + F: FnOnce() -> Fut, + Fut: Future + Send + 'static, + { + this.initialize_once(|initializer| initializer()).await + } + + /// Initializes the value using call-time arguments. + /// + /// The arguments are passed to the initializer only when this call starts a + /// new attempt. If an attempt is already active, this call resumes it and + /// its arguments are unused. + /// + /// # Panics + /// + /// Panics if the initializer panics or the cell was previously poisoned. + /// Recursive initialization of the same cell deadlocks. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let lazy = LazyCell::::new(async |name: String| name.to_uppercase()); + /// assert_eq!( + /// LazyCell::force_with(&lazy, "asyncband".to_owned()).await, + /// "ASYNCBAND" + /// ); + /// # } + /// ``` + pub async fn force_with(this: &Self, args: A) -> &T + where + F: FnOnce(A) -> Fut, + Fut: Future + Send + 'static, + { + this.initialize_once(|initializer| initializer(args)).await + } + + /// Initializes the value if needed and returns a mutable reference to it. + /// + /// # Panics + /// + /// Panics if the initializer panics or the cell was previously poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let mut lazy = LazyCell::::new(async || 92); + /// *LazyCell::force_mut(&mut lazy).await = 44; + /// assert_eq!(LazyCell::get(&lazy), Some(&44)); + /// # } + /// ``` + pub async fn force_mut(this: &mut Self) -> &mut T + where + F: FnOnce() -> Fut, + Fut: Future + Send + 'static, + { + let _ = Self::force(this).await; + this.value + .get_mut() + .expect("LazyCell value missing after success") + } + + /// Initializes the value with call-time arguments and returns mutable access. + /// + /// # Panics + /// + /// Panics if the initializer panics or the cell was previously poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let mut lazy = LazyCell::::new(async |value| value); + /// *LazyCell::force_mut_with(&mut lazy, 42).await += 1; + /// assert_eq!(LazyCell::get(&lazy), Some(&43)); + /// # } + /// ``` + pub async fn force_mut_with(this: &mut Self, args: A) -> &mut T + where + F: FnOnce(A) -> Fut, + Fut: Future + Send + 'static, + { + let _ = Self::force_with(this, args).await; + this.value + .get_mut() + .expect("LazyCell value missing after success") + } +} + +impl LazyCell { + /// Initializes the value with a fallible initializer. + /// + /// An error is returned only to the caller whose attempt produced it. The + /// cell remains uninitialized, and the next waiting caller starts a new + /// serialized attempt. Cancellation preserves the active future and does + /// not invoke the initializer again. + /// + /// # Panics + /// + /// Panics if the initializer panics or the cell was previously poisoned. + /// Recursive initialization of the same cell deadlocks. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let lazy = LazyCell::::new(async || Ok::<_, std::io::Error>(92)); + /// assert_eq!(LazyCell::try_force(&lazy).await.unwrap(), &92); + /// # } + /// ``` + pub async fn try_force(this: &Self) -> Result<&T, E> + where + F: FnMut() -> Fut, + Fut: Future> + Send + 'static, + E: 'static, + { + this.initialize_retry(|initializer| initializer()).await + } + + /// Initializes the value fallibly using call-time arguments. + /// + /// The arguments are passed to the initializer only when this call starts a + /// new attempt. If an attempt is already active, this call resumes it and + /// its arguments are unused. An error leaves the cell uninitialized so the + /// next caller can start a new attempt with its own arguments. + /// + /// # Panics + /// + /// Panics if the initializer panics or the cell was previously poisoned. + /// Recursive initialization of the same cell deadlocks. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let lazy = LazyCell::::new(async |(value, valid): (u32, bool)| { + /// valid.then_some(value).ok_or("invalid") + /// }); + /// + /// assert_eq!( + /// LazyCell::try_force_with(&lazy, (1, false)).await, + /// Err("invalid") + /// ); + /// assert_eq!(LazyCell::try_force_with(&lazy, (42, true)).await, Ok(&42)); + /// # } + /// ``` + pub async fn try_force_with(this: &Self, args: A) -> Result<&T, E> + where + F: FnMut(A) -> Fut, + Fut: Future> + Send + 'static, + E: 'static, + { + this.initialize_retry(|initializer| initializer(args)).await + } + + /// Initializes the value with a fallible initializer and returns mutable access. + /// + /// An error leaves the cell uninitialized so a later caller can retry. + /// + /// # Panics + /// + /// Panics if the initializer panics or the cell was previously poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let mut lazy = LazyCell::::new(async || Ok::<_, ()>(92)); + /// *LazyCell::try_force_mut(&mut lazy).await.unwrap() = 44; + /// assert_eq!(LazyCell::get(&lazy), Some(&44)); + /// # } + /// ``` + pub async fn try_force_mut(this: &mut Self) -> Result<&mut T, E> + where + F: FnMut() -> Fut, + Fut: Future> + Send + 'static, + E: 'static, + { + let _ = Self::try_force(this).await?; + Ok(this + .value + .get_mut() + .expect("LazyCell value missing after success")) + } + + /// Initializes the value fallibly with call-time arguments and returns mutable access. + /// + /// An error leaves the cell uninitialized so a later caller can retry. + /// + /// # Panics + /// + /// Panics if the initializer panics or the cell was previously poisoned. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::once::LazyCell; + /// + /// let mut lazy = LazyCell::::new(async |value| Ok::<_, ()>(value)); + /// *LazyCell::try_force_mut_with(&mut lazy, 42).await.unwrap() += 1; + /// assert_eq!(LazyCell::get(&lazy), Some(&43)); + /// # } + /// ``` + pub async fn try_force_mut_with(this: &mut Self, args: A) -> Result<&mut T, E> + where + F: FnMut(A) -> Fut, + Fut: Future> + Send + 'static, + E: 'static, + { + let _ = Self::try_force_with(this, args).await?; + Ok(this + .value + .get_mut() + .expect("LazyCell value missing after success")) + } +} + +impl Default for LazyCell +where + T: Default + Send + 'static, +{ + /// Creates a lazy value initialized with [`Default::default`]. + fn default() -> Self { + fn initialize() -> LazyCellFuture + where + T: Default + Send + 'static, + { + Box::pin(async { T::default() }) + } + + Self::new(initialize::) + } +} + +impl From for LazyCell { + /// Creates an already initialized lazy value. + fn from(value: T) -> Self { + Self { + value: OnceCell::from_value(value), + state: Mutex::new(State { + initializer: None, + attempt: None, + }), + poisoned: AtomicBool::new(false), + } + } +} + +impl fmt::Debug for LazyCell { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut tuple = f.debug_tuple("LazyCell"); + match Self::get(self) { + Some(value) => tuple.field(value), + None => tuple.field(&format_args!("")), + }; + tuple.finish() + } +} + +impl UnwindSafe for LazyCell {} + +impl RefUnwindSafe for LazyCell {} + +struct PoisonOnPanic<'a>(&'a AtomicBool); + +impl Drop for PoisonOnPanic<'_> { + fn drop(&mut self) { + if std::thread::panicking() { + self.0.store(true, Ordering::Release); + } + } +} + +#[cold] +#[inline(never)] +fn panic_poisoned() -> ! { + panic!("LazyCell instance has previously been poisoned") +} diff --git a/asyncband/src/once/lazy_lock/mod.rs b/asyncband/src/once/lazy_lock/mod.rs deleted file mode 100644 index e37bb80..0000000 --- a/asyncband/src/once/lazy_lock/mod.rs +++ /dev/null @@ -1,405 +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::convert::Infallible; -use std::fmt; -use std::future::Future; -use std::panic::RefUnwindSafe; -use std::panic::UnwindSafe; -use std::pin::Pin; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; - -use super::OnceCell; -use crate::mutex::Mutex; - -/// A boxed future suitable for the default [`LazyLock`] initializer type. -pub type LazyLockFuture = Pin + Send + 'static>>; - -/// A value initialized by an asynchronous function on first access. -/// -/// Initialization starts when [`force`](Self::force) or -/// [`try_force`](Self::try_force) is polled. Concurrent callers wait without -/// blocking their threads. -/// -/// If an attempt is cancelled, its future is dropped and the initializer is -/// retained. The next caller starts a new attempt. Initializers must therefore -/// be safe to invoke again after cancellation, therefore it is up to the user -/// to ensure idempotency. -/// -/// # Poisoning -/// -/// A panic from the initializer permanently poisons the lock. The panic is -/// propagated to its caller, and future calls to `force`, `try_force`, -/// `force_mut`, or `try_force_mut` panic. Errors returned through `Result` do -/// not poison the lock and are there to indicate initialization is possible. -/// -/// # Examples -/// -/// ``` -/// # #[tokio::main] -/// # async fn main() { -/// use asyncband::once::LazyLock; -/// -/// let lazy = LazyLock::::new(async || "ready".to_owned()); -/// -/// assert_eq!(LazyLock::get(&lazy), None); -/// assert_eq!(LazyLock::force(&lazy).await, "ready"); -/// assert_eq!(LazyLock::get(&lazy).map(String::as_str), Some("ready")); -/// # } -/// ``` -pub struct LazyLock LazyLockFuture> { - value: OnceCell, - initializer: Mutex>, - poisoned: AtomicBool, -} - -impl LazyLock { - /// Creates a new lazy value with the given asynchronous initializer. - /// - /// The initializer is not called until the first initialization future is - /// polled. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::once::LazyLock; - /// - /// let lazy = LazyLock::::new(async || 92); - /// assert_eq!(*LazyLock::force(&lazy).await, 92); - /// # } - /// ``` - pub const fn new(initializer: F) -> Self { - Self { - value: OnceCell::new(), - initializer: Mutex::new(Some(initializer)), - poisoned: AtomicBool::new(false), - } - } - - /// Returns a reference to the value if initialized. - /// - /// This method never starts initialization or waits for an active attempt. - /// It returns `None` when the lock is uninitialized, initializing, or - /// poisoned. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::once::LazyLock; - /// - /// let lazy = LazyLock::::new(async || 92); - /// assert_eq!(LazyLock::get(&lazy), None); - /// LazyLock::force(&lazy).await; - /// assert_eq!(LazyLock::get(&lazy), Some(&92)); - /// # } - /// ``` - pub fn get(this: &Self) -> Option<&T> { - this.value.get() - } - - /// Returns a mutable reference to the value if initialized. - /// - /// This method never starts initialization. It returns `None` when the lock - /// is uninitialized or poisoned. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::once::LazyLock; - /// - /// let mut lazy = LazyLock::::new(async || 92); - /// assert_eq!(LazyLock::get_mut(&mut lazy), None); - /// LazyLock::force(&lazy).await; - /// *LazyLock::get_mut(&mut lazy).unwrap() = 44; - /// assert_eq!(LazyLock::get(&lazy), Some(&44)); - /// # } - /// ``` - pub fn get_mut(this: &mut Self) -> Option<&mut T> { - this.value.get_mut() - } - - /// Consumes the lock and returns its value or initializer. - /// - /// Returns `Ok(value)` when initialized and `Err(initializer)` otherwise. - /// - /// # Panics - /// - /// Panics if the lock is poisoned. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::once::LazyLock; - /// - /// let lazy = LazyLock::::new(async || 92); - /// LazyLock::force(&lazy).await; - /// assert_eq!(LazyLock::into_inner(lazy).ok(), Some(92)); - /// # } - /// ``` - pub fn into_inner(this: Self) -> Result { - let Self { - value, - initializer, - poisoned, - } = this; - - if poisoned.into_inner() { - panic_poisoned(); - } - - match value.into_inner() { - Some(value) => Ok(value), - None => Err(initializer - .into_inner() - .expect("LazyLock initializer missing while uninitialized")), - } - } - - /// Internal helper to initialize the value with a fallible initializer. - /// Waiters are queued to restart on cancellation or error, winner removes - /// the initializer and ensures those waiters see the value. - async fn initialize(&self, run: G) -> Result<&T, E> - where - G: AsyncFnOnce(&mut F) -> Result, - { - if let Some(value) = self.value.get() { - return Ok(value); - } - self.assert_unpoisoned(); - - let mut initializer = self.initializer.lock().await; - if let Some(value) = self.value.get() { - return Ok(value); - } - self.assert_unpoisoned(); - - // Running initializer under a lock ensures we queue waiters to restart - // on cancellation or Err. - let _poison = PoisonOnPanic(&self.poisoned); - let value = run(initializer - .as_mut() - .expect("LazyLock initializer missing while uninitialized")) - .await?; - - drop(initializer.take()); // Avoid reinitialization possibility - let value = unsafe { self.value.set_value_unchecked(value) }; - - Ok(value) - } - - /// Panics if the lock is poisoned. - fn assert_unpoisoned(&self) { - if self.poisoned.load(Ordering::Acquire) { - panic_poisoned(); - } - } -} - -impl LazyLock -where - F: AsyncFnMut() -> T, -{ - /// Initializes the value if needed and returns a reference to it. - /// - /// If another task is initializing the lock, this call waits for that - /// attempt. If the active attempt is cancelled, this or another waiting - /// caller starts a new attempt. - /// - /// # Panics - /// - /// Panics if the initializer panics or the lock was previously poisoned. - /// Recursive initialization of the same lock deadlocks. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::once::LazyLock; - /// - /// let lazy = LazyLock::::new(async || 92); - /// assert_eq!(LazyLock::force(&lazy).await, &92); - /// # } - /// ``` - pub async fn force(this: &Self) -> &T { - // Templated to Infalliable for non-monadic initializers - match this - .initialize(async |initializer| Ok::(initializer().await)) - .await - { - Ok(value) => value, - Err(error) => match error {}, - } - } - - /// Initializes the value if needed and returns a mutable reference to it. - /// - /// # Panics - /// - /// Panics if the initializer panics or the lock was previously poisoned. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::once::LazyLock; - /// - /// let mut lazy = LazyLock::::new(async || 92); - /// *LazyLock::force_mut(&mut lazy).await = 44; - /// assert_eq!(LazyLock::get(&lazy), Some(&44)); - /// # } - /// ``` - pub async fn force_mut(this: &mut Self) -> &mut T { - let _ = Self::force(this).await; - this.value - .get_mut() - .expect("LazyLock value missing after success") - } -} - -impl LazyLock { - /// Initializes the value with a fallible initializer. - /// - /// An error is returned only to the caller whose attempt produced it. The - /// lock remains uninitialized, and the next waiting caller starts a new - /// serialized attempt. Initializers must therefore be idempotent and safe - /// to invoke again after cancellation or error. - /// - /// # Panics - /// - /// Panics if the initializer panics or the lock was previously poisoned. - /// Recursive initialization of the same lock deadlocks. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::once::LazyLock; - /// - /// let lazy = LazyLock::::new(async || Ok::<_, std::io::Error>(92)); - /// assert_eq!(LazyLock::try_force(&lazy).await.unwrap(), &92); - /// # } - /// ``` - pub async fn try_force(this: &Self) -> Result<&T, E> - where - F: AsyncFnMut() -> Result, - { - this.initialize(async |initializer| initializer().await) - .await - } - - /// Initializes the value with a fallible initializer and returns mutable access. - /// - /// An error leaves the lock uninitialized so a later caller can retry. - /// - /// # Panics - /// - /// Panics if the initializer panics or the lock was previously poisoned. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::once::LazyLock; - /// - /// let mut lazy = LazyLock::::new(async || Ok::<_, ()>(92)); - /// *LazyLock::try_force_mut(&mut lazy).await.unwrap() = 44; - /// assert_eq!(LazyLock::get(&lazy), Some(&44)); - /// # } - /// ``` - pub async fn try_force_mut(this: &mut Self) -> Result<&mut T, E> - where - F: AsyncFnMut() -> Result, - { - let _ = Self::try_force(this).await?; - Ok(this - .value - .get_mut() - .expect("LazyLock value missing after success")) - } -} - -impl Default for LazyLock -where - T: Default + Send + 'static, -{ - /// Creates a lazy value initialized with [`Default::default`]. - fn default() -> Self { - fn initialize() -> LazyLockFuture - where - T: Default + Send + 'static, - { - Box::pin(async { T::default() }) - } - - Self::new(initialize::) - } -} - -impl From for LazyLock { - /// Creates an already initialized lazy value. - fn from(value: T) -> Self { - Self { - value: OnceCell::from_value(value), - initializer: Mutex::new(None), - poisoned: AtomicBool::new(false), - } - } -} - -impl fmt::Debug for LazyLock { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let mut tuple = f.debug_tuple("LazyLock"); - match Self::get(self) { - Some(value) => tuple.field(value), - None => tuple.field(&format_args!("")), - }; - tuple.finish() - } -} - -impl UnwindSafe for LazyLock {} - -impl RefUnwindSafe for LazyLock {} - -struct PoisonOnPanic<'a>(&'a AtomicBool); - -impl Drop for PoisonOnPanic<'_> { - fn drop(&mut self) { - if std::thread::panicking() { - self.0.store(true, Ordering::Release); - } - } -} - -#[cold] -#[inline(never)] -fn panic_poisoned() -> ! { - panic!("LazyLock instance has previously been poisoned") -} diff --git a/asyncband/src/once/mod.rs b/asyncband/src/once/mod.rs index 2e2bbbf..b4b0ca7 100644 --- a/asyncband/src/once/mod.rs +++ b/asyncband/src/once/mod.rs @@ -17,8 +17,8 @@ //! Asynchronous primitives for one-time coordination. -#[cfg(feature = "lazy-lock")] -mod lazy_lock; +#[cfg(feature = "lazy-cell")] +mod lazy_cell; #[cfg(feature = "once")] #[allow(clippy::module_inception)] mod once; @@ -27,10 +27,10 @@ mod once_cell; #[cfg(feature = "once-map")] mod once_map; -#[cfg(feature = "lazy-lock")] -pub use self::lazy_lock::LazyLock; -#[cfg(feature = "lazy-lock")] -pub use self::lazy_lock::LazyLockFuture; +#[cfg(feature = "lazy-cell")] +pub use self::lazy_cell::LazyCell; +#[cfg(feature = "lazy-cell")] +pub use self::lazy_cell::LazyCellFuture; #[cfg(feature = "once")] pub use self::once::Once; #[cfg(feature = "once-cell")] diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 16c7221..40c0a2f 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -32,7 +32,7 @@ asyncband = { workspace = true, features = [ "broadcast", "condvar", "latch", - "lazy-lock", + "lazy-cell", "mpsc", "mutex", "once", diff --git a/tests-integration/tests/lazy_cell_test.rs b/tests-integration/tests/lazy_cell_test.rs new file mode 100644 index 0000000..d670a45 --- /dev/null +++ b/tests-integration/tests/lazy_cell_test.rs @@ -0,0 +1,344 @@ +// 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::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use asyncband::once::LazyCell; +use asyncband::once::LazyCellFuture; +use tokio::sync::Notify; + +#[tokio::test] +/// Ensure that multiple concurrent calls to a successful `force` only run the +/// initializer once. +async fn force_runs_initializer_once() { + let attempts = Arc::new(AtomicUsize::new(0)); + let lazy = Arc::new(LazyCell::::new({ + let attempts = attempts.clone(); + async move || { + attempts.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + 42 + } + })); + + let mut tasks = Vec::new(); + for _ in 0..32 { + let lazy = lazy.clone(); + tasks.push(tokio::spawn(async move { *LazyCell::force(&lazy).await })); + } + + for task in tasks { + assert_eq!(task.await.unwrap(), 42); + } + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +/// Ensure that cancellation preserves the active initialization future. +async fn cancellation_resumes_initialization() { + let attempts = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Notify::new()); + let resume = Arc::new(Notify::new()); + let lazy = Arc::new(LazyCell::::new({ + let attempts = attempts.clone(); + let started = started.clone(); + let resume = resume.clone(); + async move || { + attempts.fetch_add(1, Ordering::SeqCst); + started.notify_one(); + resume.notified().await; + 42 + } + })); + + let task = { + let lazy = lazy.clone(); + tokio::spawn(async move { *LazyCell::force(&lazy).await }) + }; + started.notified().await; + assert_eq!(LazyCell::get(&lazy), None); + + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + + resume.notify_one(); + assert_eq!(*LazyCell::force(&lazy).await, 42); + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +/// Ensure that dropping the cell drops a suspended initialization future. +async fn dropping_cell_drops_suspended_attempt() { + let held = Arc::new(()); + let weak = Arc::downgrade(&held); + let started = Arc::new(Notify::new()); + let lazy = Arc::new(LazyCell::::new({ + let started = started.clone(); + async move || { + started.notify_one(); + std::future::pending::<()>().await; + drop(held); + 42 + } + })); + + let task = { + let lazy = lazy.clone(); + tokio::spawn(async move { *LazyCell::force(&lazy).await }) + }; + started.notified().await; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + + drop(Arc::try_unwrap(lazy).ok().unwrap()); + assert!(weak.upgrade().is_none()); +} + +#[tokio::test] +/// Validates fallible initialization. If the initializer returns an error, +/// the value is not stored and future calls may retry it. +async fn queued_callers_retry_after_errors() { + let attempts = Arc::new(AtomicUsize::new(0)); + let lazy = Arc::new(LazyCell::::new({ + let attempts = attempts.clone(); + move || { + let attempts = attempts.clone(); + async move { + let attempt = attempts.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + if attempt < 2 { Err("retry") } else { Ok(42) } + } + } + })); + + let mut tasks = Vec::new(); + for _ in 0..3 { + let lazy = lazy.clone(); + tasks.push(tokio::spawn(async move { + LazyCell::try_force(&lazy).await.copied() + })); + } + + let mut errors = 0; + let mut successes = 0; + for task in tasks { + match task.await.unwrap() { + Ok(42) => successes += 1, + Err("retry") => errors += 1, + result => panic!("unexpected result: {result:?}"), + } + } + + assert_eq!(errors, 2); + assert_eq!(successes, 1); + assert_eq!(attempts.load(Ordering::SeqCst), 3); + assert_eq!(LazyCell::try_force(&lazy).await, Ok(&42)); + assert_eq!(attempts.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +/// Ensure that fallible attempts can receive different call-time arguments. +async fn call_time_arguments_are_used_for_retries() { + let lazy = LazyCell::::new( + async |(value, succeed): (u32, bool)| { + if succeed { Ok(value) } else { Err("retry") } + }, + ); + + assert_eq!( + LazyCell::try_force_with(&lazy, (1, false)).await, + Err("retry") + ); + assert_eq!(LazyCell::try_force_with(&lazy, (42, true)).await, Ok(&42)); +} + +#[tokio::test] +/// Ensure that a resumed attempt keeps its original call-time arguments. +async fn cancellation_preserves_active_arguments() { + let attempts = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Notify::new()); + let resume = Arc::new(Notify::new()); + let lazy = Arc::new(LazyCell::::new( + async |(value, attempts, started, resume): ( + u32, + Arc, + Arc, + Arc, + )| { + attempts.fetch_add(1, Ordering::SeqCst); + started.notify_one(); + resume.notified().await; + Ok::<_, &'static str>(value) + }, + )); + + let task = { + let lazy = lazy.clone(); + let attempts = attempts.clone(); + let started = started.clone(); + let resume = resume.clone(); + tokio::spawn(async move { + LazyCell::try_force_with(&lazy, (41, attempts, started, resume)) + .await + .copied() + }) + }; + started.notified().await; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + + resume.notify_one(); + assert_eq!( + LazyCell::try_force_with( + &lazy, + (99, attempts.clone(), started.clone(), resume.clone()) + ) + .await, + Ok(&41) + ); + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +/// Ensure that a panic in the initializer permanently poisons the cell, preventing future calls +/// from succeeding. +async fn panic_permanently_poisons_cell() { + let attempts = Arc::new(AtomicUsize::new(0)); + let lazy = Arc::new(LazyCell::::new({ + let attempts = attempts.clone(); + async move || { + attempts.fetch_add(1, Ordering::SeqCst); + panic!("initializer panic"); + } + })); + + let first = { + let lazy = lazy.clone(); + tokio::spawn(async move { + let _ = LazyCell::force(&lazy).await; + }) + }; + assert!(first.await.unwrap_err().is_panic()); + assert_eq!(LazyCell::get(&lazy), None); + + let second = { + let lazy = lazy.clone(); + tokio::spawn(async move { + let _ = LazyCell::force(&lazy).await; + }) + }; + assert!(second.await.unwrap_err().is_panic()); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + + let lazy = Arc::try_unwrap(lazy).ok().unwrap(); + let result = std::panic::catch_unwind(|| LazyCell::into_inner(lazy)); + assert!(result.is_err()); +} + +#[tokio::test] +/// Ensure that `force_mut` and `try_force_mut` can be used to mutate the value +/// after it has been initialized. +async fn mutable_force_updates_value() { + let mut lazy = LazyCell::::new(async || 41); + *LazyCell::force_mut(&mut lazy).await += 1; + assert_eq!(LazyCell::get(&lazy), Some(&42)); + + let mut fallible = LazyCell::::new(async || Ok::<_, ()>(41)); + *LazyCell::try_force_mut(&mut fallible).await.unwrap() += 1; + assert_eq!(LazyCell::get(&fallible), Some(&42)); + + let mut with_args = LazyCell::::new(async |value| Ok::<_, ()>(value)); + *LazyCell::try_force_mut_with(&mut with_args, 41) + .await + .unwrap() += 1; + assert_eq!(LazyCell::get(&with_args), Some(&42)); + + let mut infallible_with_args = LazyCell::::new(async |value| value); + *LazyCell::force_mut_with(&mut infallible_with_args, 41).await += 1; + assert_eq!(LazyCell::get(&infallible_with_args), Some(&42)); +} + +#[tokio::test] +/// Ensure that `into_inner` returns the value if it has been initialized, or +/// returns the initializer if it has not been initialized. +async fn into_inner_returns_value_or_initializer() { + let lazy = LazyCell::::new(async || 42); + let initializer = LazyCell::into_inner(lazy).unwrap_err(); + assert_eq!(initializer().await, 42); + + let lazy = LazyCell::::new(async || 42); + LazyCell::force(&lazy).await; + assert!(matches!(LazyCell::into_inner(lazy), Ok(42))); +} + +#[tokio::test] +/// Validates that `Debug` and `Default` trait implementations work as +/// expected. +async fn default_from_and_debug_match_lazy_cell() { + let lazy = LazyCell::::default(); + assert_eq!(format!("{lazy:?}"), "LazyCell()"); + assert_eq!(LazyCell::force(&lazy).await, &0); + assert_eq!(format!("{lazy:?}"), "LazyCell(0)"); + + let lazy: LazyCell = LazyCell::from(42); + assert_eq!(LazyCell::get(&lazy), Some(&42)); +} + +#[tokio::test] +/// Ensure that the initializer does not need to be `Sync` in order for the `LazyCell` to be +/// `Sync`. This is important for cases where the initializer captures non-`Sync` state, such as a +/// `Cell`. This is guaranteed by the internal mutex. +async fn initializer_need_not_be_sync() { + fn assert_sync(_: &T) {} + + let count = Cell::new(0); + let lazy = LazyCell::::new(async move || { + count.set(count.get() + 1); + count.get() + }); + + assert_sync(&lazy); + assert_eq!(LazyCell::force(&lazy).await, &1); +} + +fn static_initializer() -> LazyCellFuture { + Box::pin(async { 42 }) +} + +static STATIC_LAZY: LazyCell = LazyCell::new(static_initializer); + +#[tokio::test] +async fn default_initializer_type_supports_statics() { + assert_eq!(LazyCell::force(&STATIC_LAZY).await, &42); +} + +fn fallible_static_initializer() -> LazyCellFuture> { + Box::pin(async { Ok(42) }) +} + +type FallibleStaticInitializer = fn() -> LazyCellFuture>; + +static FALLIBLE_STATIC_LAZY: LazyCell = + LazyCell::new(fallible_static_initializer); + +#[tokio::test] +async fn one_type_supports_fallible_statics() { + assert_eq!(LazyCell::try_force(&FALLIBLE_STATIC_LAZY).await, Ok(&42)); +} diff --git a/tests-integration/tests/lazy_lock_test.rs b/tests-integration/tests/lazy_lock_test.rs deleted file mode 100644 index e46ad34..0000000 --- a/tests-integration/tests/lazy_lock_test.rs +++ /dev/null @@ -1,240 +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::cell::Cell; -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; - -use asyncband::once::LazyLock; -use asyncband::once::LazyLockFuture; -use tokio::sync::Notify; - -#[tokio::test] -/// Ensure that multiple concurrent calls to a successful `force` only run the -/// initializer once. -async fn force_runs_initializer_once() { - let attempts = Arc::new(AtomicUsize::new(0)); - let lazy = Arc::new(LazyLock::::new({ - let attempts = attempts.clone(); - async move || { - attempts.fetch_add(1, Ordering::SeqCst); - tokio::task::yield_now().await; - 42 - } - })); - - let mut tasks = Vec::new(); - for _ in 0..32 { - let lazy = lazy.clone(); - tasks.push(tokio::spawn(async move { *LazyLock::force(&lazy).await })); - } - - for task in tasks { - assert_eq!(task.await.unwrap(), 42); - } - assert_eq!(attempts.load(Ordering::SeqCst), 1); -} - -#[tokio::test] -/// Ensure that cancellation of a `force` call does not prevent future calls from -/// rerunning the initializer. -async fn cancellation_restarts_initialization() { - let attempts = Arc::new(AtomicUsize::new(0)); - let started = Arc::new(Notify::new()); - let lazy = Arc::new(LazyLock::::new({ - let attempts = attempts.clone(); - let started = started.clone(); - async move || { - let attempt = attempts.fetch_add(1, Ordering::SeqCst); - if attempt == 0 { - started.notify_one(); - std::future::pending::<()>().await; - } - 42 - } - })); - - let task = { - let lazy = lazy.clone(); - tokio::spawn(async move { *LazyLock::force(&lazy).await }) - }; - started.notified().await; - assert_eq!(LazyLock::get(&lazy), None); - - task.abort(); - assert!(task.await.unwrap_err().is_cancelled()); - - assert_eq!(*LazyLock::force(&lazy).await, 42); - assert_eq!(attempts.load(Ordering::SeqCst), 2); -} - -#[tokio::test] -/// Validates falliabile initialization. If the initializer returns an error, -/// the value is not stored and future calls may retry it. -async fn queued_callers_retry_after_errors() { - let attempts = Arc::new(AtomicUsize::new(0)); - let lazy = Arc::new(LazyLock::::new({ - let attempts = attempts.clone(); - async move || { - let attempt = attempts.fetch_add(1, Ordering::SeqCst); - tokio::task::yield_now().await; - if attempt < 2 { Err("retry") } else { Ok(42) } - } - })); - - let mut tasks = Vec::new(); - for _ in 0..3 { - let lazy = lazy.clone(); - tasks.push(tokio::spawn(async move { - LazyLock::try_force(&lazy).await.copied() - })); - } - - let mut errors = 0; - let mut successes = 0; - for task in tasks { - match task.await.unwrap() { - Ok(42) => successes += 1, - Err("retry") => errors += 1, - result => panic!("unexpected result: {result:?}"), - } - } - - assert_eq!(errors, 2); - assert_eq!(successes, 1); - assert_eq!(attempts.load(Ordering::SeqCst), 3); - assert_eq!(LazyLock::try_force(&lazy).await, Ok(&42)); - assert_eq!(attempts.load(Ordering::SeqCst), 3); -} - -#[tokio::test] -/// Ensure that a panic in the initializer permanently poisons the lock, preventing future calls -/// from succeeding. -async fn panic_permanently_poisons_lock() { - let attempts = Arc::new(AtomicUsize::new(0)); - let lazy = Arc::new(LazyLock::::new({ - let attempts = attempts.clone(); - async move || { - attempts.fetch_add(1, Ordering::SeqCst); - panic!("initializer panic"); - } - })); - - let first = { - let lazy = lazy.clone(); - tokio::spawn(async move { - let _ = LazyLock::force(&lazy).await; - }) - }; - assert!(first.await.unwrap_err().is_panic()); - assert_eq!(LazyLock::get(&lazy), None); - - let second = { - let lazy = lazy.clone(); - tokio::spawn(async move { - let _ = LazyLock::force(&lazy).await; - }) - }; - assert!(second.await.unwrap_err().is_panic()); - assert_eq!(attempts.load(Ordering::SeqCst), 1); - - let lazy = Arc::try_unwrap(lazy).ok().unwrap(); - let result = std::panic::catch_unwind(|| LazyLock::into_inner(lazy)); - assert!(result.is_err()); -} - -#[tokio::test] -/// Ensure that `force_mut` and `try_force_mut` can be used to mutate the value -/// after it has been initialized. -async fn mutable_force_updates_value() { - let mut lazy = LazyLock::::new(async || 41); - *LazyLock::force_mut(&mut lazy).await += 1; - assert_eq!(LazyLock::get(&lazy), Some(&42)); - - let mut fallible = LazyLock::::new(async || Ok::<_, ()>(41)); - *LazyLock::try_force_mut(&mut fallible).await.unwrap() += 1; - assert_eq!(LazyLock::get(&fallible), Some(&42)); -} - -#[tokio::test] -/// Ensure that `into_inner` returns the value if it has been initialized, or -/// returns the initializer if it has not been initialized. -async fn into_inner_returns_value_or_initializer() { - let lazy = LazyLock::::new(async || 42); - let initializer = LazyLock::into_inner(lazy).unwrap_err(); - assert_eq!(initializer().await, 42); - - let lazy = LazyLock::::new(async || 42); - LazyLock::force(&lazy).await; - assert!(matches!(LazyLock::into_inner(lazy), Ok(42))); -} - -#[tokio::test] -/// Validates that `Debug` and `Default` trait implementations work as -/// expected. -async fn default_from_and_debug_match_lazy_lock() { - let lazy = LazyLock::::default(); - assert_eq!(format!("{lazy:?}"), "LazyLock()"); - assert_eq!(LazyLock::force(&lazy).await, &0); - assert_eq!(format!("{lazy:?}"), "LazyLock(0)"); - - let lazy: LazyLock = LazyLock::from(42); - assert_eq!(LazyLock::get(&lazy), Some(&42)); -} - -#[tokio::test] -/// Ensure that the initializer does not need to be `Sync` in order for the `LazyLock` to be -/// `Sync`. This is important for cases where the initializer captures non-`Sync` state, such as a -/// `Cell`. This is guaranteed by the internal mutex. -async fn initializer_need_not_be_sync() { - fn assert_sync(_: &T) {} - - let count = Cell::new(0); - let lazy = LazyLock::::new(async move || { - count.set(count.get() + 1); - count.get() - }); - - assert_sync(&lazy); - assert_eq!(LazyLock::force(&lazy).await, &1); -} - -fn static_initializer() -> LazyLockFuture { - Box::pin(async { 42 }) -} - -static STATIC_LAZY: LazyLock = LazyLock::new(static_initializer); - -#[tokio::test] -async fn default_initializer_type_supports_statics() { - assert_eq!(LazyLock::force(&STATIC_LAZY).await, &42); -} - -fn fallible_static_initializer() -> LazyLockFuture> { - Box::pin(async { Ok(42) }) -} - -type FallibleStaticInitializer = fn() -> LazyLockFuture>; - -static FALLIBLE_STATIC_LAZY: LazyLock = - LazyLock::new(fallible_static_initializer); - -#[tokio::test] -async fn one_type_supports_fallible_statics() { - assert_eq!(LazyLock::try_force(&FALLIBLE_STATIC_LAZY).await, Ok(&42)); -} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index d074830..c096c56 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -22,7 +22,7 @@ use asyncband::latch::Latch; use asyncband::mpsc; use asyncband::mutex::Mutex; use asyncband::mutex::MutexGuard; -use asyncband::once::LazyLock; +use asyncband::once::LazyCell; use asyncband::once::Once; use asyncband::once::OnceCell; use asyncband::once::OnceMap; @@ -45,7 +45,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::>(); @@ -91,7 +91,7 @@ fn public_types_are_unpin() { assert_unpin::(); assert_unpin::(); assert_unpin::(); - assert_unpin::>(); + assert_unpin::>(); assert_unpin::(); assert_unpin::>(); assert_unpin::>(); From 8bed51d050b43b6b1ed5155683d76ed78d943d60 Mon Sep 17 00:00:00 2001 From: Matthew Hambrecht Date: Sat, 22 Aug 2026 22:56:43 -0400 Subject: [PATCH 3/4] fix potential deadlock footgun + cleaner api description + prevent unrelated panic poisoning --- asyncband/src/once/lazy_cell/mod.rs | 45 ++++++++--- tests-integration/tests/lazy_cell_test.rs | 91 +++++++++++++++++++++++ 2 files changed, 127 insertions(+), 9 deletions(-) diff --git a/asyncband/src/once/lazy_cell/mod.rs b/asyncband/src/once/lazy_cell/mod.rs index 04fd4d1..606c305 100644 --- a/asyncband/src/once/lazy_cell/mod.rs +++ b/asyncband/src/once/lazy_cell/mod.rs @@ -45,8 +45,11 @@ pub type LazyCellFuture = Pin + Send + 'static>>; /// future does not retain them. /// /// Infallible initializers are called once and may move captured values directly -/// into their future. Fallible initializers must remain callable after an error. -/// A captured value can be cloned into each owned attempt: +/// into their future. Fallible initializers are `FnMut` factories that must +/// remain callable after an error and return an owned future. Lending +/// `AsyncFnMut` closures that borrow captured state into the returned future are +/// not supported. State can instead be updated before creating the future or +/// cloned into each owned attempt: /// /// ``` /// # #[tokio::main] @@ -245,17 +248,24 @@ impl LazyCell { } self.assert_unpoisoned(); - let _poison = PoisonOnPanic(&self.poisoned); + let mut start = Some(start); + // Initialize the value if no other task has started an attempt. if state.attempt.is_none() { let initializer = state .initializer .take() .expect("LazyCell initializer missing while uninitialized"); - let future = start(initializer); + let future = { + let _poison = PoisonOnPanic(&self.poisoned); + start.take().expect("LazyCell initializer start missing")(initializer) + }; let attempt = Box::pin(async move { AttemptOutput::Value(future.await) }); state.attempt = Some((AttemptKind::Infallible, attempt)); } + // Drop unused caller arguments before resuming the active attempt. Retaining a + // guard could deadlock the attempt on a resource it needs. + drop(start); let (kind, attempt) = state .attempt @@ -266,7 +276,12 @@ impl LazyCell { "LazyCell force method does not match the active attempt" ); - let output = attempt.as_mut().await; + // Avoid unrelated panics from the initializer poisoning the cell for all future callers. + let output = std::future::poll_fn(|cx| { + let _poison = PoisonOnPanic(&self.poisoned); + attempt.as_mut().poll(cx) + }) + .await; state.attempt = None; let AttemptOutput::Value(value) = output else { unreachable!("infallible LazyCell attempt returned an error") @@ -292,15 +307,19 @@ impl LazyCell { } self.assert_unpoisoned(); - let _poison = PoisonOnPanic(&self.poisoned); let kind = AttemptKind::Fallible(TypeId::of::()); + let mut start = Some(start); + // Initialize the value if no other task has started an attempt. if state.attempt.is_none() { - let future = state + let initializer = state .initializer .as_mut() - .map(start) .expect("LazyCell initializer missing while uninitialized"); + let future = { + let _poison = PoisonOnPanic(&self.poisoned); + start.take().expect("LazyCell initializer start missing")(initializer) + }; let attempt = Box::pin(async move { match future.await { Ok(value) => AttemptOutput::Value(value), @@ -309,6 +328,9 @@ impl LazyCell { }); state.attempt = Some((kind, attempt)); } + // Drop unused caller arguments before resuming the active attempt. Retaining a + // guard could deadlock the attempt on a resource it needs. + drop(start); let (active_kind, attempt) = state .attempt @@ -319,7 +341,12 @@ impl LazyCell { "LazyCell force method does not match the active attempt" ); - let output = attempt.as_mut().await; + // Avoid unrelated panics from the initializer poisoning the cell for all future callers. + let output = std::future::poll_fn(|cx| { + let _poison = PoisonOnPanic(&self.poisoned); + attempt.as_mut().poll(cx) + }) + .await; state.attempt = None; match output { AttemptOutput::Value(value) => { diff --git a/tests-integration/tests/lazy_cell_test.rs b/tests-integration/tests/lazy_cell_test.rs index d670a45..fbf6130 100644 --- a/tests-integration/tests/lazy_cell_test.rs +++ b/tests-integration/tests/lazy_cell_test.rs @@ -83,6 +83,40 @@ async fn cancellation_resumes_initialization() { assert_eq!(attempts.load(Ordering::SeqCst), 1); } +#[tokio::test] +/// Ensure unrelated task unwinding does not poison a pending attempt. +async fn unrelated_unwind_does_not_poison_cell() { + let started = Arc::new(Notify::new()); + let resume = Arc::new(Notify::new()); + let lazy = Arc::new(LazyCell::::new({ + let started = started.clone(); + let resume = resume.clone(); + async move || { + started.notify_one(); + resume.notified().await; + 42 + } + })); + + let task = { + let lazy = lazy.clone(); + tokio::spawn(async move { + tokio::select! { + biased; + _ = LazyCell::force(&lazy) => {} + _ = async { + started.notified().await; + panic!("unrelated panic"); + } => {} + } + }) + }; + assert!(task.await.unwrap_err().is_panic()); + + resume.notify_one(); + assert_eq!(LazyCell::force(&lazy).await, &42); +} + #[tokio::test] /// Ensure that dropping the cell drops a suspended initialization future. async fn dropping_cell_drops_suspended_attempt() { @@ -216,6 +250,63 @@ async fn cancellation_preserves_active_arguments() { assert_eq!(attempts.load(Ordering::SeqCst), 1); } +struct NotifyOnDrop(Option>); + +impl Drop for NotifyOnDrop { + fn drop(&mut self) { + if let Some(notify) = self.0.take() { + notify.notify_one(); + } + } +} + +#[tokio::test] +/// Ensure unused caller arguments are dropped before an attempt resumes. +async fn unused_arguments_are_dropped_before_resume() { + let started = Arc::new(Notify::new()); + let resume = Arc::new(Notify::new()); + let dropped = Arc::new(Notify::new()); + let lazy = Arc::new(LazyCell::::new({ + let started = started.clone(); + let resume = resume.clone(); + move |(value, _): (u32, NotifyOnDrop)| { + let started = started.clone(); + let resume = resume.clone(); + async move { + started.notify_one(); + resume.notified().await; + Ok::<_, ()>(value) + } + } + })); + + let first = { + let lazy = lazy.clone(); + tokio::spawn(async move { + LazyCell::try_force_with(&lazy, (41, NotifyOnDrop(None))) + .await + .copied() + }) + }; + started.notified().await; + first.abort(); + assert!(first.await.unwrap_err().is_cancelled()); + + let second = { + let lazy = lazy.clone(); + let dropped = dropped.clone(); + tokio::spawn(async move { + LazyCell::try_force_with(&lazy, (99, NotifyOnDrop(Some(dropped)))) + .await + .copied() + }) + }; + + dropped.notified().await; + resume.notify_one(); + assert_eq!(second.await.unwrap(), Ok(41)); +} + #[tokio::test] /// Ensure that a panic in the initializer permanently poisons the cell, preventing future calls /// from succeeding. From b29d3f377271274e27cb20f406b7d38e8985835a Mon Sep 17 00:00:00 2001 From: Matthew Hambrecht Date: Sat, 22 Aug 2026 23:07:02 -0400 Subject: [PATCH 4/4] clearer comments --- asyncband/src/once/lazy_cell/mod.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/asyncband/src/once/lazy_cell/mod.rs b/asyncband/src/once/lazy_cell/mod.rs index 606c305..cfbcd4f 100644 --- a/asyncband/src/once/lazy_cell/mod.rs +++ b/asyncband/src/once/lazy_cell/mod.rs @@ -250,7 +250,7 @@ impl LazyCell { let mut start = Some(start); - // Initialize the value if no other task has started an attempt. + // Start an attempt when none is active. if state.attempt.is_none() { let initializer = state .initializer @@ -276,7 +276,7 @@ impl LazyCell { "LazyCell force method does not match the active attempt" ); - // Avoid unrelated panics from the initializer poisoning the cell for all future callers. + // Scope poisoning to initializer polls so cancellation does not poison. let output = std::future::poll_fn(|cx| { let _poison = PoisonOnPanic(&self.poisoned); attempt.as_mut().poll(cx) @@ -286,6 +286,7 @@ impl LazyCell { let AttemptOutput::Value(value) = output else { unreachable!("infallible LazyCell attempt returned an error") }; + // SAFETY: The state mutex serializes initialization. unsafe { self.value.set_value_unchecked(value) } } @@ -310,7 +311,7 @@ impl LazyCell { let kind = AttemptKind::Fallible(TypeId::of::()); let mut start = Some(start); - // Initialize the value if no other task has started an attempt. + // Start an attempt when none is active. if state.attempt.is_none() { let initializer = state .initializer @@ -341,7 +342,7 @@ impl LazyCell { "LazyCell force method does not match the active attempt" ); - // Avoid unrelated panics from the initializer poisoning the cell for all future callers. + // Scope poisoning to initializer polls so cancellation does not poison. let output = std::future::poll_fn(|cx| { let _poison = PoisonOnPanic(&self.poisoned); attempt.as_mut().poll(cx) @@ -351,6 +352,7 @@ impl LazyCell { match output { AttemptOutput::Value(value) => { state.initializer = None; + // SAFETY: The state mutex serializes initialization. Ok(unsafe { self.value.set_value_unchecked(value) }) } AttemptOutput::Error(error) => {