diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..78d5903 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,20 @@ +# Agent Instructions + +## Cargo XTask + +Use `cargo x` as the source of truth for repository workflows. + +- Run `cargo x --help` before choosing build, test, lint, or formatting commands. +- Run `cargo x --help` for command-specific behavior. + +## Code Style + +- Express visibility at the module boundary. Inside a restricted module, use `pub` for its module API instead of repeating `pub(crate)`. + +## Commits and Pull Requests + +Follow the semantic definition at `.github/semantic.yml`. + +- Keep title descriptions short. +- Simple PR descriptions should only include a `Summary` section. +- Complex PR descriptions may also include a `Design Notes` section. diff --git a/CHANGELOG.md b/CHANGELOG.md index a0d21b9..ef6ca49 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ All notable changes to this project will be documented in this file. * Add an opt-in `asyncband::blocking::FutureExt` bridge with `block_on` and `wait_timeout` methods for waiting on runtime-agnostic futures from synchronous code. * Add opt-in bounded and unbounded runtime-agnostic object pools under `asyncband::pool`. +* Add opt-in `asyncband::once::LazyCell` for values that own one asynchronous initializer and preserve its in-flight future across caller cancellation. ### Breaking changes diff --git a/Cargo.lock b/Cargo.lock index c72436f..4a150d3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -283,6 +283,14 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "examples" +version = "0.0.0" +dependencies = [ + "asyncband", + "tokio", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" diff --git a/Cargo.toml b/Cargo.toml index 4dde95d..2cc043c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ # under the License. [workspace] -members = ["asyncband", "benchmarks", "tests-integration", "xtask"] +members = ["asyncband", "benchmarks", "examples", "tests-integration", "xtask"] resolver = "3" [workspace.package] diff --git a/README.md b/README.md index 44ecbb6..7a8a40b 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,14 @@ async fn increment() { Public paths stay direct—such as `asyncband::mutex`, `asyncband::pool`, and `asyncband::once::OnceCell`—while Cargo features keep unused implementations out of the build. +## Examples + +Runnable examples live in the [`examples`](examples) workspace crate. For example, the [`OnceCell` versus `LazyCell` initialization guide](examples/src/once_cell_vs_lazy_cell.rs) demonstrates when a restartable fixed free function is sufficient, when initialization needs access-time context and retries, and when a lazy value needs to own and preserve a one-shot initializer. + +```shell +cargo run -p examples --example once_cell_vs_lazy_cell +``` + ## API map | Area | API | Feature | Use | @@ -65,6 +73,7 @@ Public paths stay direct—such as `asyncband::mutex`, `asyncband::pool`, and `a | | [`Condvar`](https://docs.rs/asyncband/*/asyncband/condvar/struct.Condvar.html) | `condvar` | Wait for notifications while releasing a mutex. | | 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. | +| | [`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 af8e642..3871529 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -48,6 +48,7 @@ barrier = [] blocking = [] condvar = ["mutex"] latch = [] +lazy-cell = ["mutex"] mpsc = [] mutex = [] once = ["semaphore"] diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 47c1751..cdd33e3 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -39,9 +39,17 @@ mod arena; pub(crate) mod countdown; #[cfg(any(feature = "once-map", feature = "singleflight"))] -#[cfg_attr(not(feature = "once-map"), allow(dead_code))] +// `OnceMap` and `singleflight` use different subsets of `OnceTable`, so single-feature builds +// leave some operations in the shared implementation unused. +#[allow(dead_code)] pub(crate) mod once_table; +#[cfg(any(feature = "lazy-cell", feature = "once-cell"))] +// `LazyCell` and `OnceCell` use different subsets of `ValueCell`, so single-feature builds leave +// some operations in the shared implementation unused. +#[allow(dead_code)] +pub(crate) mod value_cell; + #[cfg(any( feature = "barrier", feature = "latch", diff --git a/asyncband/src/internal/value_cell.rs b/asyncband/src/internal/value_cell.rs new file mode 100644 index 0000000..928b1b8 --- /dev/null +++ b/asyncband/src/internal/value_cell.rs @@ -0,0 +1,142 @@ +// 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::UnsafeCell; +use std::mem::MaybeUninit; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +/// Storage for a value published exactly once. +/// +/// Synchronizing initialization is deliberately left to the containing primitive. +pub struct ValueCell { + initialized: AtomicBool, + value: UnsafeCell>, +} + +// SAFETY: Shared access only exposes `&T`, and publication is synchronized by the initialized +// flag. The containing primitive is responsible for serializing writes. +unsafe impl Sync for ValueCell {} + +// SAFETY: Ownership of the cell and its value may be transferred when `T` is `Send`. +unsafe impl Send for ValueCell {} + +impl ValueCell { + pub const fn new() -> Self { + Self { + initialized: AtomicBool::new(false), + value: UnsafeCell::new(MaybeUninit::uninit()), + } + } + + pub const fn from_value(value: T) -> Self { + Self { + initialized: AtomicBool::new(true), + value: UnsafeCell::new(MaybeUninit::new(value)), + } + } + + pub fn is_initialized(&self) -> bool { + self.initialized.load(Ordering::Acquire) + } + + pub fn is_initialized_mut(&mut self) -> bool { + *self.initialized.get_mut() + } + + pub fn get(&self) -> Option<&T> { + if self.is_initialized() { + // SAFETY: The acquire load above observed publication of the value. + Some(unsafe { self.get_unchecked() }) + } else { + None + } + } + + pub fn get_mut(&mut self) -> Option<&mut T> { + if self.is_initialized_mut() { + // SAFETY: Exclusive access rules out concurrent reads or writes. + Some(unsafe { self.get_unchecked_mut() }) + } else { + None + } + } + + pub fn into_inner(mut self) -> Option { + self.take() + } + + pub fn take(&mut self) -> Option { + if self.is_initialized_mut() { + *self.initialized.get_mut() = false; + // SAFETY: The value was initialized and the flag was cleared so it will not be + // dropped a second time. + Some(unsafe { self.value.get_mut().assume_init_read() }) + } else { + None + } + } + + /// Publishes a value after the containing primitive has won initialization. + /// + /// # Safety + /// + /// The cell must be uninitialized, and no other initialization may read or write it. + pub unsafe fn set(&self, value: T) -> &T { + debug_assert!(!self.is_initialized()); + let value_ptr = self.value.get(); + unsafe { value_ptr.write(MaybeUninit::new(value)) }; + + // Publish the initialized value to readers performing an acquire load. + self.initialized.store(true, Ordering::Release); + + // SAFETY: The value was initialized and published above. + unsafe { self.get_unchecked() } + } + + pub fn set_mut(&mut self, value: T) -> &mut T { + debug_assert!(!self.is_initialized_mut()); + let value = self.value.get_mut().write(value); + *self.initialized.get_mut() = true; + value + } + + /// # Safety + /// + /// The cell must be initialized. + unsafe fn get_unchecked(&self) -> &T { + debug_assert!(self.is_initialized()); + unsafe { (&*self.value.get()).assume_init_ref() } + } + + /// # Safety + /// + /// The cell must be initialized and exclusively borrowed. + unsafe fn get_unchecked_mut(&mut self) -> &mut T { + debug_assert!(self.is_initialized_mut()); + unsafe { (&mut *self.value.get()).assume_init_mut() } + } +} + +impl Drop for ValueCell { + fn drop(&mut self) { + if self.is_initialized_mut() { + // SAFETY: The value is initialized and exclusive access rules out other users. + unsafe { self.value.get_mut().assume_init_drop() }; + } + } +} diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 1c3ffa5..41a52d8 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::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`] | `oneshot`, `mpsc` | //! | Reuse objects | [`pool::bounded`], [`pool::unbounded`] | `pool` | @@ -114,7 +114,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-cell", + feature = "once", + feature = "once-cell", + feature = "once-map" +))] pub mod once; #[cfg(feature = "oneshot")] pub mod oneshot; diff --git a/asyncband/src/once/lazy_cell/mod.rs b/asyncband/src/once/lazy_cell/mod.rs new file mode 100644 index 0000000..6d951da --- /dev/null +++ b/asyncband/src/once/lazy_cell/mod.rs @@ -0,0 +1,273 @@ +// 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::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 crate::internal::value_cell::ValueCell; +use crate::mutex::Mutex; + +type BoxFuture = Pin + Send + 'static>>; + +/// A thread-safe value initialized by a stored asynchronous function on first access. +/// +/// Initialization starts when [`force`](Self::force) is polled. Concurrent callers wait without +/// blocking their threads. If the forcing caller is cancelled, the initialization future remains +/// pinned in the cell and the next caller resumes that same future instead of starting over. +/// +/// The initialization future must be `Send + 'static`: another task may resume it after the +/// original caller is gone, potentially on another thread. The initializer itself only needs to +/// be `Send`, not `Sync`, because the cell serializes access to it. +/// +/// `LazyCell` represents one asynchronous initialization attempt. If initialization needs +/// access-time arguments or should retry after returning an error, use `OnceCell::get_or_try_init` +/// instead. A `Result` may still be the stored value when an error should be cached. +/// +/// # Poisoning +/// +/// A panic while creating or polling the initialization future permanently poisons the cell. The +/// panic is propagated to its caller, and future calls to [`force`](Self::force) or +/// [`force_mut`](Self::force_mut) panic. +/// +/// # 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 BoxFuture> { + value: ValueCell, + state: Mutex>, + poisoned: AtomicBool, +} + +struct State { + initializer: Option, + attempt: Option>, +} + +impl State { + async fn drive_attempt(&mut self, poisoned: &AtomicBool) -> T + where + F: FnOnce() -> Fut, + Fut: Future + Send + 'static, + { + if self.attempt.is_none() { + let initializer = self + .initializer + .take() + .expect("LazyCell initializer missing while uninitialized"); + let future = { + let _poison = PoisonOnPanic(poisoned); + initializer() + }; + self.attempt = Some(Box::pin(future)); + } + + let value = std::future::poll_fn(|cx| { + let _poison = PoisonOnPanic(poisoned); + self.attempt + .as_mut() + .expect("LazyCell attempt missing while initializing") + .as_mut() + .poll(cx) + }) + .await; + + // Treat panics from dropping a completed future as initializer panics as well. + let _poison = PoisonOnPanic(poisoned); + self.attempt = None; + value + } +} + +impl LazyCell { + /// Creates a new lazy value with the given asynchronous initializer. + /// + /// The initializer is not called until the first [`force`](Self::force) 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: ValueCell::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. + 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. + pub fn get_mut(this: &mut Self) -> Option<&mut T> { + this.value.get_mut() + } + + /// 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 the task + /// driving initialization 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. + pub async fn force(this: &Self) -> &T + where + F: FnOnce() -> Fut, + Fut: Future + Send + 'static, + { + if let Some(value) = this.value.get() { + return value; + } + this.assert_unpoisoned(); + + let mut state = this.state.lock().await; + if let Some(value) = this.value.get() { + return value; + } + this.assert_unpoisoned(); + + let value = state.drive_attempt(&this.poisoned).await; + + // SAFETY: The state mutex serializes initialization, and the double check above verified + // that the value was not initialized by another caller. + unsafe { this.value.set(value) } + } + + /// Initializes the value if needed and returns mutable access to it. + /// + /// # Panics + /// + /// Panics under the same conditions as [`force`](Self::force). + pub async fn force_mut(this: &mut Self) -> &mut T + where + F: FnOnce() -> Fut, + Fut: Future + Send + 'static, + { + if this.value.is_initialized_mut() { + return this + .value + .get_mut() + .expect("LazyCell value missing while initialized"); + } + if *this.poisoned.get_mut() { + panic_poisoned(); + } + + // Exclusive access makes locking and atomic value publication unnecessary. + let value = this.state.get_mut().drive_attempt(&this.poisoned).await; + this.value.set_mut(value) + } + + fn assert_unpoisoned(&self) { + if self.poisoned.load(Ordering::Acquire) { + panic_poisoned(); + } + } +} + +impl Default for LazyCell +where + T: Default, +{ + fn default() -> Self { + fn initialize() -> BoxFuture { + Box::pin(async { T::default() }) + } + + Self::new(initialize::) + } +} + +impl From for LazyCell { + fn from(value: T) -> Self { + Self { + value: ValueCell::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/mod.rs b/asyncband/src/once/mod.rs index 82c5a65..70a43d2 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-cell")] +mod lazy_cell; #[cfg(feature = "once")] mod once; #[cfg(feature = "once-cell")] @@ -24,6 +26,8 @@ mod once_cell; #[cfg(feature = "once-map")] mod once_map; +#[cfg(feature = "lazy-cell")] +pub use self::lazy_cell::LazyCell; #[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..ddb2323 100644 --- a/asyncband/src/once/once_cell/mod.rs +++ b/asyncband/src/once/once_cell/mod.rs @@ -15,18 +15,20 @@ // specific language governing permissions and limitations // under the License. -use std::cell::UnsafeCell; use std::convert::Infallible; use std::fmt; -use std::mem::MaybeUninit; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; +use crate::internal::value_cell::ValueCell; use crate::semaphore::Semaphore; use crate::semaphore::SemaphorePermit; /// A thread-safe cell which can nominally be written to only once. /// +/// Callers provide an initializer when accessing an empty cell. An initializer that returns an +/// error, panics, or is cancelled leaves the cell empty so a later caller can retry. Use +/// `LazyCell` instead when the cell should own a one-shot initializer and preserve its in-flight +/// future across caller cancellation. +/// /// # Examples /// /// ``` @@ -49,17 +51,10 @@ use crate::semaphore::SemaphorePermit; /// The outputs must be either `Results: 1, 1` or `Results: 2, 2`, i.e. once the value is set via /// an asynchronous function, the value inside the `OnceCell` will be immutable. pub struct OnceCell { - value_set: AtomicBool, - value: UnsafeCell>, + value: ValueCell, semaphore: Semaphore, } -// SAFETY: OnceCell can be shared between threads as long as T is Sync + Send. -unsafe impl Sync for OnceCell {} - -// SAFETY: OnceCell can be sent between threads as long as T is Send. -unsafe impl Send for OnceCell {} - impl Default for OnceCell { fn default() -> Self { Self::new() @@ -70,8 +65,7 @@ impl OnceCell { /// Creates a new empty `OnceCell`. pub const fn new() -> Self { Self { - value_set: AtomicBool::new(false), - value: UnsafeCell::new(MaybeUninit::uninit()), + value: ValueCell::new(), semaphore: Semaphore::new(1), } } @@ -79,20 +73,22 @@ impl OnceCell { /// Creates a new `OnceCell` initialized with the provided value. pub const fn from_value(value: T) -> Self { Self { - value_set: AtomicBool::new(true), - value: UnsafeCell::new(MaybeUninit::new(value)), + value: ValueCell::from_value(value), semaphore: Semaphore::new(1), } } /// Returns whether the internal value is set. + // `OnceMap` and `singleflight` inspect this state, while a standalone `OnceCell` build does + // not need the internal helper. + #[allow(dead_code)] pub(crate) fn initialized(&self) -> bool { - self.value_set.load(Ordering::Acquire) + self.value.is_initialized() } /// Returns whether the internal value is set. pub(crate) fn initialized_mut(&mut self) -> bool { - *self.value_set.get_mut() + self.value.is_initialized_mut() } /// Gets the reference to the underlying value. @@ -101,11 +97,7 @@ impl OnceCell { /// /// This method never blocks. pub fn get(&self) -> Option<&T> { - if self.initialized() { - Some(unsafe { self.get_unchecked() }) - } else { - None - } + self.value.get() } /// Gets the mutable reference to the underlying value. @@ -115,11 +107,7 @@ impl OnceCell { /// This method never blocks. Since it borrows the `OnceCell` mutably, it is statically /// guaranteed that no active borrows to the `OnceCell` exist, including from other threads. pub fn get_mut(&mut self) -> Option<&mut T> { - if self.initialized_mut() { - Some(unsafe { self.get_unchecked_mut() }) - } else { - None - } + self.value.get_mut() } /// Gets the reference to the internal value, initializing it with the provided asynchronous @@ -252,7 +240,10 @@ impl OnceCell { // Workaround if let Some(v) = self.get_mut() { return Ok(v); } // @see https://github.com/rust-lang/rust/issues/51545 if self.initialized_mut() { - return Ok(unsafe { self.get_unchecked_mut() }); + return Ok(self + .value + .get_mut() + .expect("OnceCell initialized value missing")); } let value = init().await?; @@ -353,14 +344,8 @@ impl OnceCell { /// assert_eq!(cell.into_inner(), Some("hello".to_string())); /// # } /// ``` - pub fn into_inner(mut self) -> Option { - if self.initialized_mut() { - // set to uninitialized for the destructor of `OnceCell` to work properly - *self.value_set.get_mut() = false; - Some(unsafe { self.value.get_mut().assume_init_read() }) - } else { - None - } + pub fn into_inner(self) -> Option { + self.value.into_inner() } /// Takes the value out of this `OnceCell`, moving it back to an uninitialized state. @@ -387,55 +372,17 @@ impl OnceCell { /// # } /// ``` pub fn take(&mut self) -> Option { - std::mem::take(self).into_inner() - } - - /// # Safety - /// - /// The cell must be initialized - #[inline] - unsafe fn get_unchecked(&self) -> &T { - debug_assert!(self.initialized()); - unsafe { (&*self.value.get()).assume_init_ref() } - } - - /// # Safety - /// - /// The cell must be initialized - #[inline] - unsafe fn get_unchecked_mut(&mut self) -> &mut T { - debug_assert!(self.initialized_mut()); - unsafe { (&mut *self.value.get()).assume_init_mut() } + self.value.take() } fn set_value(&self, value: T, permit: SemaphorePermit<'_>) -> &T { - // Hold the permit to ensure exclusive access. let _permit = permit; - - let value_ptr = self.value.get(); - unsafe { value_ptr.write(MaybeUninit::new(value)) }; - - // Use `store` with `Release` ordering to ensure that when loading it with `Acquire` - // ordering, the initialized value is visible. - self.value_set.store(true, Ordering::Release); - - // SAFETY: value initialized above - unsafe { self.get_unchecked() } + // SAFETY: Holding the only semaphore permit serializes initialization. + unsafe { self.value.set(value) } } fn set_value_mut(&mut self, value: T) -> &mut T { - let value = self.value.get_mut().write(value); - *self.value_set.get_mut() = true; - value - } -} - -impl Drop for OnceCell { - fn drop(&mut self) { - if self.initialized_mut() { - // SAFETY: The cell is initialized and being dropped, so it can't be accessed again. - unsafe { self.value.get_mut().assume_init_drop() }; - } + self.value.set_mut(value) } } diff --git a/examples/Cargo.toml b/examples/Cargo.toml new file mode 100644 index 0000000..c0bf64f --- /dev/null +++ b/examples/Cargo.toml @@ -0,0 +1,37 @@ +# 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. + +[package] +name = "examples" +publish = false + +edition.workspace = true +rust-version.workspace = true + +[package.metadata.release] +release = false + +[dependencies] +asyncband = { workspace = true, features = ["lazy-cell", "once-cell"] } +tokio = { workspace = true, features = ["macros", "rt", "sync"] } + +[lints] +workspace = true + +[[example]] +name = "once_cell_vs_lazy_cell" +path = "src/once_cell_vs_lazy_cell.rs" diff --git a/examples/src/once_cell_vs_lazy_cell.rs b/examples/src/once_cell_vs_lazy_cell.rs new file mode 100644 index 0000000..c546fdc --- /dev/null +++ b/examples/src/once_cell_vs_lazy_cell.rs @@ -0,0 +1,129 @@ +// 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::Ready; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use asyncband::once::LazyCell; +use asyncband::once::OnceCell; +use tokio::sync::Notify; + +static ONCE_ENDPOINT: OnceCell = OnceCell::new(); +static LAZY_ENDPOINT: LazyCell Ready> = + LazyCell::new(load_default_endpoint); + +fn load_default_endpoint() -> Ready { + std::future::ready("https://service.example".to_owned()) +} + +#[tokio::main(flavor = "current_thread")] +async fn main() { + fixed_static_initializer_can_use_either_cell().await; + once_cell_accepts_access_time_context().await; + lazy_cell_owns_a_local_fn_once().await; +} + +async fn fixed_static_initializer_can_use_either_cell() { + // When restarting after cancellation is acceptable, a stateless, repeatable initializer does + // not require `LazyCell`. An accessor can pass the same free function to `OnceCell`. + assert_eq!( + ONCE_ENDPOINT.get_or_init(load_default_endpoint).await, + "https://service.example" + ); + + // `LazyCell` instead stores that fixed initializer in the static itself, so callers only need + // to force the value. + assert_eq!( + LazyCell::force(&LAZY_ENDPOINT).await, + "https://service.example" + ); +} + +async fn once_cell_accepts_access_time_context() { + let endpoint = OnceCell::new(); + + // Failed attempts leave `OnceCell` empty. A later caller can retry with fresh context. + let first = endpoint + .get_or_try_init(async || Err::("service discovery unavailable")) + .await; + assert_eq!(first.unwrap_err(), "service discovery unavailable"); + + let discovered_endpoint = "https://discovered.example".to_owned(); + let endpoint = endpoint + .get_or_try_init(async move || Ok::<_, &'static str>(discovered_endpoint)) + .await + .unwrap(); + assert_eq!(endpoint, "https://discovered.example"); +} + +async fn lazy_cell_owns_a_local_fn_once() { + struct Credentials { + token: String, + } + + struct Client { + token: String, + } + + let attempts = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Notify::new()); + let resume = Arc::new(Notify::new()); + let credentials = Credentials { + token: "secret".to_owned(), + }; + + let initialize = { + let attempts = Arc::clone(&attempts); + let started = Arc::clone(&started); + let resume = Arc::clone(&resume); + + // Moving `credentials.token` out makes this local initializer `FnOnce`, not `Fn`. + move || { + let token = credentials.token; + + async move { + attempts.fetch_add(1, Ordering::SeqCst); + started.notify_one(); + resume.notified().await; + Client { token } + } + } + }; + + // A caller could pass `initialize` to `OnceCell::get_or_init` once, but the call consumes it. + // If that call is cancelled, a later caller cannot supply the same `FnOnce` again. `LazyCell` + // owns the initializer and preserves its in-flight future across callers. + let client = Arc::new(LazyCell::::new(initialize)); + + let first_caller = tokio::spawn({ + let client = Arc::clone(&client); + async move { + LazyCell::force(&client).await; + } + }); + started.notified().await; + first_caller.abort(); + assert!(first_caller.await.unwrap_err().is_cancelled()); + + // Cancellation does not consume the captured credentials or restart the initializer. The next + // caller resumes the same future. + resume.notify_one(); + assert_eq!(LazyCell::force(&client).await.token, "secret"); + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 27c0594..7ea541f 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -31,6 +31,7 @@ asyncband = { workspace = true, features = [ "blocking", "condvar", "latch", + "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..453b103 --- /dev/null +++ b/tests-integration/tests/lazy_cell_test.rs @@ -0,0 +1,309 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cell::Cell; +use std::future::Ready; +use std::rc::Rc; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use asyncband::once::LazyCell; +use tokio::sync::Notify; + +#[tokio::test] +async fn initializer_starts_when_force_is_polled() { + let attempts = Arc::new(AtomicUsize::new(0)); + let lazy = LazyCell::::new({ + let attempts = attempts.clone(); + move || { + attempts.fetch_add(1, Ordering::SeqCst); + async { 42 } + } + }); + + let force = LazyCell::force(&lazy); + assert_eq!(attempts.load(Ordering::SeqCst), 0); + drop(force); + assert_eq!(attempts.load(Ordering::SeqCst), 0); + + assert_eq!(LazyCell::force(&lazy).await, &42); + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn concurrent_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] +async fn cancellation_preserves_initialization_future() { + 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] +async fn unrelated_unwind_does_not_poison_pending_attempt() { + 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] +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] +async fn result_is_cached_as_the_value() { + let attempts = Arc::new(AtomicUsize::new(0)); + let lazy = LazyCell::, _>::new({ + let attempts = attempts.clone(); + async move || { + attempts.fetch_add(1, Ordering::SeqCst); + Err("cached") + } + }); + + assert_eq!(LazyCell::force(&lazy).await, &Err("cached")); + assert_eq!(LazyCell::force(&lazy).await, &Err("cached")); + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn initializer_poll_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 mut lazy = Arc::try_unwrap(lazy).ok().unwrap(); + let third = tokio::spawn(async move { + let _ = LazyCell::force_mut(&mut lazy).await; + }); + assert!(third.await.unwrap_err().is_panic()); +} + +#[tokio::test] +async fn initializer_creation_panic_permanently_poisons_cell() { + let lazy = Arc::new(LazyCell::::new(|| -> std::future::Ready { + panic!("initializer creation panic") + })); + + let first = { + let lazy = lazy.clone(); + tokio::spawn(async move { + let _ = LazyCell::force(&lazy).await; + }) + }; + assert!(first.await.unwrap_err().is_panic()); + + let second = tokio::spawn(async move { + let _ = LazyCell::force(&lazy).await; + }); + assert!(second.await.unwrap_err().is_panic()); +} + +#[tokio::test] +async fn force_mut_updates_value() { + let mut lazy = LazyCell::::new(async || 41); + *LazyCell::force_mut(&mut lazy).await += 1; + assert_eq!(LazyCell::get(&lazy), Some(&42)); +} + +#[tokio::test] +async fn force_mut_resumes_a_started_attempt() { + 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 { + let _ = LazyCell::force(&lazy).await; + }) + }; + started.notified().await; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + + let mut lazy = Arc::try_unwrap(lazy).ok().unwrap(); + resume.notify_one(); + assert_eq!(LazyCell::force_mut(&mut lazy).await, &mut 42); +} + +#[tokio::test] +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 local = LazyCell::>::default(); + assert_eq!(**LazyCell::force(&local).await, 0); + + let lazy: LazyCell = LazyCell::from(42); + assert_eq!(LazyCell::get(&lazy), Some(&42)); +} + +#[tokio::test] +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() -> Ready { + std::future::ready(42) +} + +static STATIC_LAZY: LazyCell Ready> = LazyCell::new(static_initializer); + +#[tokio::test] +async fn function_pointer_initializer_supports_statics() { + assert_eq!(LazyCell::force(&STATIC_LAZY).await, &42); +} diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 4419543..7872e29 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -23,6 +23,7 @@ use asyncband::latch::Latch; use asyncband::mpsc; use asyncband::mutex::Mutex; use asyncband::mutex::MutexGuard; +use asyncband::once::LazyCell; use asyncband::once::Once; use asyncband::once::OnceCell; use asyncband::once::OnceMap; @@ -67,6 +68,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::>(); @@ -114,6 +116,7 @@ fn public_types_are_unpin() { assert_unpin::(); assert_unpin::(); assert_unpin::(); + assert_unpin::>(); assert_unpin::(); assert_unpin::>(); assert_unpin::>();