tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-09-07 00:00:08 +02:00
Code Issues Packages Projects Releases Wiki Activity

sync: clean up OnceCell (#3945)

Browse Source
This commit is contained in:
Alice Ryhl
2021-07-12 11:19:14 +02:00
committed by GitHub
parent 3b38ebd7f5
commit 80d8d40a34
1 changed files with 208 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/sync/once_cell.rs
+208 -157
View File
@@ -1,4 +1,4 @@
use super::Semaphore; use super::{Semaphore, SemaphorePermit, TryAcquireError};
use crate::loom::cell::UnsafeCell; use crate::loom::cell::UnsafeCell;
use std::error::Error; use std::error::Error;
use std::fmt; use std::fmt;
@@ -8,15 +8,30 @@ use std::ops::Drop;
use std::ptr; use std::ptr;
use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::atomic::{AtomicBool, Ordering};
/// A thread-safe cell which can be written to only once. // This file contains an implementation of an OnceCell. The principle
// behind the safety the of the cell is that any thread with an `&OnceCell` may
// access the `value` field according the following rules:
//
// 1. When `value_set` is false, the `value` field may be modified by the
// thread holding the permit on the semaphore.
// 2. When `value_set` is true, the `value` field may be accessed immutably by
// any thread.
//
// It is an invariant that if the semaphore is closed, then `value_set` is true.
// The reverse does not necessarily hold — but if not, the semaphore may not
// have any available permits.
//
// A thread with a `&mut OnceCell` may modify the value in any way it wants as
// long as the invariants are upheld.
/// A thread-safe cell that can be written to only once.
/// ///
/// Provides the functionality to either set the value, in case `OnceCell` /// A `OnceCell` is typically used for global variables that need to be
/// is uninitialized, or get the already initialized value by using an async /// initialized once on first use, but need no further changes. The `OnceCell`
/// function via [`OnceCell::get_or_init`]. /// in Tokio allows the initialization procedure to be asynchronous.
///
/// [`OnceCell::get_or_init`]: crate::sync::OnceCell::get_or_init
/// ///
/// # Examples /// # Examples
///
/// ``` /// ```
/// use tokio::sync::OnceCell; /// use tokio::sync::OnceCell;
/// ///
@@ -28,8 +43,28 @@ use std::sync::atomic::{AtomicBool, Ordering};
/// ///
/// #[tokio::main] /// #[tokio::main]
/// async fn main() { /// async fn main() {
/// let result1 = ONCE.get_or_init(some_computation).await; /// let result = ONCE.get_or_init(some_computation).await;
/// assert_eq!(*result1, 2); /// assert_eq!(*result, 2);
/// }
/// ```
///
/// It is often useful to write a wrapper method for accessing the value.
///
/// ```
/// use tokio::sync::OnceCell;
///
/// static ONCE: OnceCell<u32> = OnceCell::const_new();
///
/// async fn get_global_integer() -> &'static u32 {
/// ONCE.get_or_init(|| async {
/// 1 + 1
/// }).await
/// }
///
/// #[tokio::main]
/// async fn main() {
/// let result = get_global_integer().await;
/// assert_eq!(*result, 2);
/// } /// }
/// ``` /// ```
pub struct OnceCell<T> { pub struct OnceCell<T> {
@@ -68,7 +103,7 @@ impl<T: Eq> Eq for OnceCell<T> {}
impl<T> Drop for OnceCell<T> { impl<T> Drop for OnceCell<T> {
fn drop(&mut self) { fn drop(&mut self) {
if self.initialized() { if self.initialized_mut() {
unsafe { unsafe {
self.value self.value
.with_mut(|ptr| ptr::drop_in_place((&mut *ptr).as_mut_ptr())); .with_mut(|ptr| ptr::drop_in_place((&mut *ptr).as_mut_ptr()));
@@ -90,7 +125,7 @@ impl<T> From<T> for OnceCell<T> {
} }
impl<T> OnceCell<T> { impl<T> OnceCell<T> {
/// Creates a new uninitialized OnceCell instance. /// Creates a new empty `OnceCell` instance.
pub fn new() -> Self { pub fn new() -> Self {
OnceCell { OnceCell {
value_set: AtomicBool::new(false), value_set: AtomicBool::new(false),
@@ -99,8 +134,9 @@ impl<T> OnceCell<T> {
} }
} }
/// Creates a new initialized OnceCell instance if `value` is `Some`, otherwise /// Creates a new `OnceCell` that contains the provided value, if any.
/// has the same functionality as [`OnceCell::new`]. ///
/// If the `Option` is `None`, this is equivalent to `OnceCell::new`.
/// ///
/// [`OnceCell::new`]: crate::sync::OnceCell::new /// [`OnceCell::new`]: crate::sync::OnceCell::new
pub fn new_with(value: Option<T>) -> Self { pub fn new_with(value: Option<T>) -> Self {
@@ -111,8 +147,31 @@ impl<T> OnceCell<T> {
} }
} }
/// Creates a new uninitialized OnceCell instance. /// Creates a new empty `OnceCell` instance.
#[cfg(all(feature = "parking_lot", not(all(loom, test)),))] ///
/// Equivalent to `OnceCell::new`, except that it can be used in static
/// variables.
///
/// # Example
///
/// ```
/// use tokio::sync::OnceCell;
///
/// static ONCE: OnceCell<u32> = OnceCell::const_new();
///
/// async fn get_global_integer() -> &'static u32 {
/// ONCE.get_or_init(|| async {
/// 1 + 1
/// }).await
/// }
///
/// #[tokio::main]
/// async fn main() {
/// let result = get_global_integer().await;
/// assert_eq!(*result, 2);
/// }
/// ```
#[cfg(all(feature = "parking_lot", not(all(loom, test))))]
#[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))]
pub const fn const_new() -> Self { pub const fn const_new() -> Self {
OnceCell { OnceCell {
@@ -122,33 +181,48 @@ impl<T> OnceCell<T> {
} }
} }
/// Whether the value of the OnceCell is set or not. /// Returns `true` if the `OnceCell` currently contains a value, and `false`
/// otherwise.
pub fn initialized(&self) -> bool { pub fn initialized(&self) -> bool {
// Using acquire ordering so any threads that read a true from this
// atomic is able to read the value.
self.value_set.load(Ordering::Acquire) self.value_set.load(Ordering::Acquire)
} }
// SAFETY: safe to call only once self.initialized() is true /// Returns `true` if the `OnceCell` currently contains a value, and `false`
/// otherwise.
fn initialized_mut(&mut self) -> bool {
*self.value_set.get_mut()
}
// SAFETY: The OnceCell must not be empty.
unsafe fn get_unchecked(&self) -> &T { unsafe fn get_unchecked(&self) -> &T {
&*self.value.with(|ptr| (*ptr).as_ptr()) &*self.value.with(|ptr| (*ptr).as_ptr())
} }
// SAFETY: safe to call only once self.initialized() is true. Safe because // SAFETY: The OnceCell must not be empty.
// because of the mutable reference.
unsafe fn get_unchecked_mut(&mut self) -> &mut T { unsafe fn get_unchecked_mut(&mut self) -> &mut T {
&mut *self.value.with_mut(|ptr| (*ptr).as_mut_ptr()) &mut *self.value.with_mut(|ptr| (*ptr).as_mut_ptr())
} }
// SAFETY: safe to call only once a permit on the semaphore has been fn set_value(&self, value: T, permit: SemaphorePermit<'_>) -> &T {
// acquired // SAFETY: We are holding the only permit on the semaphore.
unsafe fn set_value(&self, value: T) { unsafe {
self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value)); self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value));
}
// Using release ordering so any threads that read a true from this
// atomic is able to read the value we just stored.
self.value_set.store(true, Ordering::Release); self.value_set.store(true, Ordering::Release);
self.semaphore.close(); self.semaphore.close();
permit.forget();
// SAFETY: We just initialized the cell.
unsafe { self.get_unchecked() }
} }
/// Tries to get a reference to the value of the OnceCell. /// Returns a reference to the value currently stored in the `OnceCell`, or
/// /// `None` if the `OnceCell` is empty.
/// Returns None if the value of the OnceCell hasn't previously been initialized.
pub fn get(&self) -> Option<&T> { pub fn get(&self) -> Option<&T> {
if self.initialized() { if self.initialized() {
Some(unsafe { self.get_unchecked() }) Some(unsafe { self.get_unchecked() })
@@ -157,179 +231,161 @@ impl<T> OnceCell<T> {
} }
} }
/// Tries to return a mutable reference to the value of the cell. /// Returns a mutable reference to the value currently stored in the
/// `OnceCell`, or `None` if the `OnceCell` is empty.
/// ///
/// Returns None if the cell hasn't previously been initialized. /// Since this call borrows the `OnceCell` mutably, it is safe to mutate the
/// value inside the `OnceCell` — the mutable borrow statically guarantees
/// no other references exist.
pub fn get_mut(&mut self) -> Option<&mut T> { pub fn get_mut(&mut self) -> Option<&mut T> {
if self.initialized() { if self.initialized_mut() {
Some(unsafe { self.get_unchecked_mut() }) Some(unsafe { self.get_unchecked_mut() })
} else { } else {
None None
} }
} }
/// Sets the value of the OnceCell to the argument value. /// Set the value of the `OnceCell` to the given value if the `OnceCell` is
/// empty.
/// ///
/// If the value of the OnceCell was already set prior to this call /// If the `OnceCell` already has a value, this call will fail with an
/// then [`SetError::AlreadyInitializedError`] is returned. If another thread /// [`SetError::AlreadyInitializedError`].
/// is initializing the cell while this method is called, ///
/// [`SetError::InitializingError`] is returned. In order to wait /// If the `OnceCell` is empty, but some other task is currently trying to
/// for an ongoing initialization to finish, call /// set the value, this call will fail with [`SetError::InitializingError`].
/// [`OnceCell::get_or_init`] instead.
/// ///
/// [`SetError::AlreadyInitializedError`]: crate::sync::SetError::AlreadyInitializedError /// [`SetError::AlreadyInitializedError`]: crate::sync::SetError::AlreadyInitializedError
/// [`SetError::InitializingError`]: crate::sync::SetError::InitializingError /// [`SetError::InitializingError`]: crate::sync::SetError::InitializingError
/// ['OnceCell::get_or_init`]: crate::sync::OnceCell::get_or_init
pub fn set(&self, value: T) -> Result<(), SetError<T>> { pub fn set(&self, value: T) -> Result<(), SetError<T>> {
if !self.initialized() { if self.initialized() {
// Another thread might be initializing the cell, in which case `try_acquire` will return Err(SetError::AlreadyInitializedError(value));
// return an error
match self.semaphore.try_acquire() {
Ok(_permit) => {
if !self.initialized() {
// SAFETY: There is only one permit on the semaphore, hence only one
// mutable reference is created
unsafe { self.set_value(value) };
return Ok(());
} else {
unreachable!(
"acquired the permit after OnceCell value was already initialized."
);
}
}
_ => {
// Couldn't acquire the permit, look if initializing process is already completed
if !self.initialized() {
return Err(SetError::InitializingError(value));
}
}
}
} }
Err(SetError::AlreadyInitializedError(value)) // Another task might be initializing the cell, in which case
// `try_acquire` will return an error. If we succeed to acquire the
// permit, then we can set the value.
match self.semaphore.try_acquire() {
Ok(permit) => {
debug_assert!(!self.initialized());
self.set_value(value, permit);
Ok(())
}
Err(TryAcquireError::NoPermits) => {
// Some other task is holding the permit. That task is
// currently trying to initialize the value.
Err(SetError::InitializingError(value))
}
Err(TryAcquireError::Closed) => {
// The semaphore was closed. Some other task has initialized
// the value.
Err(SetError::AlreadyInitializedError(value))
}
}
} }
/// Tries to initialize the value of the OnceCell using the async function `f`. /// Get the value currently in the `OnceCell`, or initialize it with the
/// If the value of the OnceCell was already initialized prior to this call, /// given asynchronous operation.
/// a reference to that initialized value is returned. If some other thread
/// initiated the initialization prior to this call and the initialization
/// hasn't completed, this call waits until the initialization is finished.
/// ///
/// This will deadlock if `f` tries to initialize the cell itself. /// If some other task is currently working on initializing the `OnceCell`,
/// this call will wait for that other task to finish, then return the value
/// that the other task produced.
///
/// If the provided operation is cancelled or panics, the initialization
/// attempt is cancelled. If there are other tasks waiting for the value to
/// be initialized, one of them will start another attempt at initializing
/// the value.
///
/// This will deadlock if `f` tries to initialize the cell recursively.
pub async fn get_or_init<F, Fut>(&self, f: F) -> &T pub async fn get_or_init<F, Fut>(&self, f: F) -> &T
where where
F: FnOnce() -> Fut, F: FnOnce() -> Fut,
Fut: Future<Output = T>, Fut: Future<Output = T>,
{ {
if self.initialized() { if self.initialized() {
// SAFETY: once the value is initialized, no mutable references are given out, so // SAFETY: The OnceCell has been fully initialized.
// we can give out arbitrarily many immutable references
unsafe { self.get_unchecked() } unsafe { self.get_unchecked() }
} else { } else {
// After acquire().await we have either acquired a permit while self.value // Here we try to acquire the semaphore permit. Holding the permit
// is still uninitialized, or the current thread is awoken after another thread // will allow us to set the value of the OnceCell, and prevents
// has initialized the value and closed the semaphore, in which case self.initialized // other tasks from initializing the OnceCell while we are holding
// is true and we don't set the value here // it.
match self.semaphore.acquire().await { match self.semaphore.acquire().await {
Ok(_permit) => { Ok(permit) => {
if !self.initialized() { debug_assert!(!self.initialized());
// If `f()` panics or `select!` is called, this `get_or_init` call
// is aborted and the semaphore permit is dropped.
let value = f().await;
// SAFETY: There is only one permit on the semaphore, hence only one // If `f()` panics or `select!` is called, this
// mutable reference is created // `get_or_init` call is aborted and the semaphore permit is
unsafe { self.set_value(value) }; // dropped.
let value = f().await;
// SAFETY: once the value is initialized, no mutable references are given out, so self.set_value(value, permit)
// we can give out arbitrarily many immutable references
unsafe { self.get_unchecked() }
} else {
unreachable!("acquired semaphore after value was already initialized.");
}
} }
Err(_) => { Err(_) => {
if self.initialized() { debug_assert!(self.initialized());
// SAFETY: once the value is initialized, no mutable references are given out, so
// we can give out arbitrarily many immutable references // SAFETY: The semaphore has been closed. This only happens
unsafe { self.get_unchecked() } // when the OnceCell is fully initialized.
} else { unsafe { self.get_unchecked() }
unreachable!(
"Semaphore closed, but the OnceCell has not been initialized."
);
}
} }
} }
} }
} }
/// Tries to initialize the value of the OnceCell using the async function `f`. /// Get the value currently in the `OnceCell`, or initialize it with the
/// If the value of the OnceCell was already initialized prior to this call, /// given asynchronous operation.
/// a reference to that initialized value is returned. If some other thread
/// initiated the initialization prior to this call and the initialization
/// hasn't completed, this call waits until the initialization is finished.
/// If the function argument `f` returns an error, `get_or_try_init`
/// returns that error, otherwise the result of `f` will be stored in the cell.
/// ///
/// This will deadlock if `f` tries to initialize the cell itself. /// If some other task is currently working on initializing the `OnceCell`,
/// this call will wait for that other task to finish, then return the value
/// that the other task produced.
///
/// If the provided operation returns an error, is cancelled or panics, the
/// initialization attempt is cancelled. If there are other tasks waiting
/// for the value to be initialized, one of them will start another attempt
/// at initializing the value.
///
/// This will deadlock if `f` tries to initialize the cell recursively.
pub async fn get_or_try_init<E, F, Fut>(&self, f: F) -> Result<&T, E> pub async fn get_or_try_init<E, F, Fut>(&self, f: F) -> Result<&T, E>
where where
F: FnOnce() -> Fut, F: FnOnce() -> Fut,
Fut: Future<Output = Result<T, E>>, Fut: Future<Output = Result<T, E>>,
{ {
if self.initialized() { if self.initialized() {
// SAFETY: once the value is initialized, no mutable references are given out, so // SAFETY: The OnceCell has been fully initialized.
// we can give out arbitrarily many immutable references
unsafe { Ok(self.get_unchecked()) } unsafe { Ok(self.get_unchecked()) }
} else { } else {
// After acquire().await we have either acquired a permit while self.value // Here we try to acquire the semaphore permit. Holding the permit
// is still uninitialized, or the current thread is awoken after another thread // will allow us to set the value of the OnceCell, and prevents
// has initialized the value and closed the semaphore, in which case self.initialized // other tasks from initializing the OnceCell while we are holding
// is true and we don't set the value here // it.
match self.semaphore.acquire().await { match self.semaphore.acquire().await {
Ok(_permit) => { Ok(permit) => {
if !self.initialized() { debug_assert!(!self.initialized());
// If `f()` panics or `select!` is called, this `get_or_try_init` call
// is aborted and the semaphore permit is dropped.
let value = f().await;
match value { // If `f()` panics or `select!` is called, this
Ok(value) => { // `get_or_try_init` call is aborted and the semaphore
// SAFETY: There is only one permit on the semaphore, hence only one // permit is dropped.
// mutable reference is created let value = f().await;
unsafe { self.set_value(value) };
// SAFETY: once the value is initialized, no mutable references are given out, so match value {
// we can give out arbitrarily many immutable references Ok(value) => Ok(self.set_value(value, permit)),
unsafe { Ok(self.get_unchecked()) } Err(e) => Err(e),
}
Err(e) => Err(e),
}
} else {
unreachable!("acquired semaphore after value was already initialized.");
} }
} }
Err(_) => { Err(_) => {
if self.initialized() { debug_assert!(self.initialized());
// SAFETY: once the value is initialized, no mutable references are given out, so
// we can give out arbitrarily many immutable references // SAFETY: The semaphore has been closed. This only happens
unsafe { Ok(self.get_unchecked()) } // when the OnceCell is fully initialized.
} else { unsafe { Ok(self.get_unchecked()) }
unreachable!(
"Semaphore closed, but the OnceCell has not been initialized."
);
}
} }
} }
} }
} }
/// Moves the value out of the cell, destroying the cell in the process. /// Take the value from the cell, destroying the cell in the process.
/// /// Returns `None` if the cell is empty.
/// Returns `None` if the cell is uninitialized.
pub fn into_inner(mut self) -> Option<T> { pub fn into_inner(mut self) -> Option<T> {
if self.initialized() { if self.initialized_mut() {
// Set to uninitialized for the destructor of `OnceCell` to work properly // Set to uninitialized for the destructor of `OnceCell` to work properly
*self.value_set.get_mut() = false; *self.value_set.get_mut() = false;
Some(unsafe { self.value.with(|ptr| ptr::read(ptr).assume_init()) }) Some(unsafe { self.value.with(|ptr| ptr::read(ptr).assume_init()) })
@@ -338,20 +394,18 @@ impl<T> OnceCell<T> {
} }
} }
/// Takes ownership of the current value, leaving the cell uninitialized. /// Takes ownership of the current value, leaving the cell empty. Returns
/// /// `None` if the cell is empty.
/// Returns `None` if the cell is uninitialized.
pub fn take(&mut self) -> Option<T> { pub fn take(&mut self) -> Option<T> {
std::mem::take(self).into_inner() std::mem::take(self).into_inner()
} }
} }
// Since `get` gives us access to immutable references of the // Since `get` gives us access to immutable references of the OnceCell, OnceCell
// OnceCell, OnceCell can only be Sync if T is Sync, otherwise // can only be Sync if T is Sync, otherwise OnceCell would allow sharing
// OnceCell would allow sharing references of !Sync values across // references of !Sync values across threads. We need T to be Send in order for
// threads. We need T to be Send in order for OnceCell to by Sync // OnceCell to by Sync because we can use `set` on `&OnceCell<T>` to send values
// because we can use `set` on `&OnceCell<T>` to send // (of type T) across threads.
// values (of type T) across threads.
unsafe impl<T: Sync + Send> Sync for OnceCell<T> {} unsafe impl<T: Sync + Send> Sync for OnceCell<T> {}
// Access to OnceCell's value is guarded by the semaphore permit // Access to OnceCell's value is guarded by the semaphore permit
@@ -359,20 +413,17 @@ unsafe impl<T: Sync + Send> Sync for OnceCell<T> {}
// it's safe to send it to another thread // it's safe to send it to another thread
unsafe impl<T: Send> Send for OnceCell<T> {} unsafe impl<T: Send> Send for OnceCell<T> {}
/// Errors that can be returned from [`OnceCell::set`] /// Errors that can be returned from [`OnceCell::set`].
/// ///
/// [`OnceCell::set`]: crate::sync::OnceCell::set /// [`OnceCell::set`]: crate::sync::OnceCell::set
#[derive(Debug, PartialEq)] #[derive(Debug, PartialEq)]
pub enum SetError<T> { pub enum SetError<T> {
/// Error resulting from [`OnceCell::set`] calls if the cell was previously initialized. /// The cell was already initialized when [`OnceCell::set`] was called.
/// ///
/// [`OnceCell::set`]: crate::sync::OnceCell::set /// [`OnceCell::set`]: crate::sync::OnceCell::set
AlreadyInitializedError(T), AlreadyInitializedError(T),
/// Error resulting from [`OnceCell::set`] calls when the cell is currently being /// The cell is currently being initialized.
/// initialized during the calls to that method.
///
/// [`OnceCell::set`]: crate::sync::OnceCell::set
InitializingError(T), InitializingError(T),
} }