diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-17 12:02:58 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-17 12:02:58 +0000 |
commit | 698f8c2f01ea549d77d7dc3338a12e04c11057b9 (patch) | |
tree | 173a775858bd501c378080a10dca74132f05bc50 /vendor/tokio/src/sync/once_cell.rs | |
parent | Initial commit. (diff) | |
download | rustc-698f8c2f01ea549d77d7dc3338a12e04c11057b9.tar.xz rustc-698f8c2f01ea549d77d7dc3338a12e04c11057b9.zip |
Adding upstream version 1.64.0+dfsg1.upstream/1.64.0+dfsg1
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'vendor/tokio/src/sync/once_cell.rs')
-rw-r--r-- | vendor/tokio/src/sync/once_cell.rs | 406 |
1 files changed, 406 insertions, 0 deletions
diff --git a/vendor/tokio/src/sync/once_cell.rs b/vendor/tokio/src/sync/once_cell.rs new file mode 100644 index 000000000..ce55d9e35 --- /dev/null +++ b/vendor/tokio/src/sync/once_cell.rs @@ -0,0 +1,406 @@ +use super::Semaphore; +use crate::loom::cell::UnsafeCell; +use std::error::Error; +use std::fmt; +use std::future::Future; +use std::mem::MaybeUninit; +use std::ops::Drop; +use std::ptr; +use std::sync::atomic::{AtomicBool, Ordering}; + +/// A thread-safe cell which can be written to only once. +/// +/// Provides the functionality to either set the value, in case `OnceCell` +/// is uninitialized, or get the already initialized value by using an async +/// function via [`OnceCell::get_or_init`]. +/// +/// [`OnceCell::get_or_init`]: crate::sync::OnceCell::get_or_init +/// +/// # Examples +/// ``` +/// use tokio::sync::OnceCell; +/// +/// async fn some_computation() -> u32 { +/// 1 + 1 +/// } +/// +/// static ONCE: OnceCell<u32> = OnceCell::const_new(); +/// +/// #[tokio::main] +/// async fn main() { +/// let result1 = ONCE.get_or_init(some_computation).await; +/// assert_eq!(*result1, 2); +/// } +/// ``` +pub struct OnceCell<T> { + value_set: AtomicBool, + value: UnsafeCell<MaybeUninit<T>>, + semaphore: Semaphore, +} + +impl<T> Default for OnceCell<T> { + fn default() -> OnceCell<T> { + OnceCell::new() + } +} + +impl<T: fmt::Debug> fmt::Debug for OnceCell<T> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("OnceCell") + .field("value", &self.get()) + .finish() + } +} + +impl<T: Clone> Clone for OnceCell<T> { + fn clone(&self) -> OnceCell<T> { + OnceCell::new_with(self.get().cloned()) + } +} + +impl<T: PartialEq> PartialEq for OnceCell<T> { + fn eq(&self, other: &OnceCell<T>) -> bool { + self.get() == other.get() + } +} + +impl<T: Eq> Eq for OnceCell<T> {} + +impl<T> Drop for OnceCell<T> { + fn drop(&mut self) { + if self.initialized() { + unsafe { + self.value + .with_mut(|ptr| ptr::drop_in_place((&mut *ptr).as_mut_ptr())); + }; + } + } +} + +impl<T> From<T> for OnceCell<T> { + fn from(value: T) -> Self { + let semaphore = Semaphore::new(0); + semaphore.close(); + OnceCell { + value_set: AtomicBool::new(true), + value: UnsafeCell::new(MaybeUninit::new(value)), + semaphore, + } + } +} + +impl<T> OnceCell<T> { + /// Creates a new uninitialized OnceCell instance. + pub fn new() -> Self { + OnceCell { + value_set: AtomicBool::new(false), + value: UnsafeCell::new(MaybeUninit::uninit()), + semaphore: Semaphore::new(1), + } + } + + /// Creates a new initialized OnceCell instance if `value` is `Some`, otherwise + /// has the same functionality as [`OnceCell::new`]. + /// + /// [`OnceCell::new`]: crate::sync::OnceCell::new + pub fn new_with(value: Option<T>) -> Self { + if let Some(v) = value { + OnceCell::from(v) + } else { + OnceCell::new() + } + } + + /// Creates a new uninitialized OnceCell instance. + #[cfg(all(feature = "parking_lot", not(all(loom, test)),))] + #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] + pub const fn const_new() -> Self { + OnceCell { + value_set: AtomicBool::new(false), + value: UnsafeCell::new(MaybeUninit::uninit()), + semaphore: Semaphore::const_new(1), + } + } + + /// Whether the value of the OnceCell is set or not. + pub fn initialized(&self) -> bool { + self.value_set.load(Ordering::Acquire) + } + + // SAFETY: safe to call only once self.initialized() is true + unsafe fn get_unchecked(&self) -> &T { + &*self.value.with(|ptr| (*ptr).as_ptr()) + } + + // SAFETY: safe to call only once self.initialized() is true. Safe because + // because of the mutable reference. + unsafe fn get_unchecked_mut(&mut self) -> &mut T { + &mut *self.value.with_mut(|ptr| (*ptr).as_mut_ptr()) + } + + // SAFETY: safe to call only once a permit on the semaphore has been + // acquired + unsafe fn set_value(&self, value: T) { + self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value)); + self.value_set.store(true, Ordering::Release); + self.semaphore.close(); + } + + /// Tries to get a reference to the value of the OnceCell. + /// + /// Returns None if the value of the OnceCell hasn't previously been initialized. + pub fn get(&self) -> Option<&T> { + if self.initialized() { + Some(unsafe { self.get_unchecked() }) + } else { + None + } + } + + /// Tries to return a mutable reference to the value of the cell. + /// + /// Returns None if the cell hasn't previously been initialized. + pub fn get_mut(&mut self) -> Option<&mut T> { + if self.initialized() { + Some(unsafe { self.get_unchecked_mut() }) + } else { + None + } + } + + /// Sets the value of the OnceCell to the argument value. + /// + /// If the value of the OnceCell was already set prior to this call + /// then [`SetError::AlreadyInitializedError`] is returned. If another thread + /// is initializing the cell while this method is called, + /// [`SetError::InitializingError`] is returned. In order to wait + /// for an ongoing initialization to finish, call + /// [`OnceCell::get_or_init`] instead. + /// + /// [`SetError::AlreadyInitializedError`]: crate::sync::SetError::AlreadyInitializedError + /// [`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>> { + if !self.initialized() { + // Another thread might be initializing the cell, in which case `try_acquire` will + // 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)) + } + + /// Tries to initialize the value of the OnceCell using the async function `f`. + /// If the value of the OnceCell was already initialized prior to this call, + /// 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. + pub async fn get_or_init<F, Fut>(&self, f: F) -> &T + where + F: FnOnce() -> Fut, + Fut: Future<Output = T>, + { + if self.initialized() { + // SAFETY: once the value is initialized, no mutable references are given out, so + // we can give out arbitrarily many immutable references + unsafe { self.get_unchecked() } + } else { + // After acquire().await we have either acquired a permit while self.value + // is still uninitialized, or the current thread is awoken after another thread + // has initialized the value and closed the semaphore, in which case self.initialized + // is true and we don't set the value here + match self.semaphore.acquire().await { + Ok(_permit) => { + if !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 + // mutable reference is created + unsafe { self.set_value(value) }; + + // SAFETY: once the value is initialized, no mutable references are given out, so + // we can give out arbitrarily many immutable references + unsafe { self.get_unchecked() } + } else { + unreachable!("acquired semaphore after value was already initialized."); + } + } + Err(_) => { + if self.initialized() { + // SAFETY: once the value is initialized, no mutable references are given out, so + // we can give out arbitrarily many immutable references + unsafe { self.get_unchecked() } + } else { + unreachable!( + "Semaphore closed, but the OnceCell has not been initialized." + ); + } + } + } + } + } + + /// Tries to initialize the value of the OnceCell using the async function `f`. + /// If the value of the OnceCell was already initialized prior to this call, + /// 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. + pub async fn get_or_try_init<E, F, Fut>(&self, f: F) -> Result<&T, E> + where + F: FnOnce() -> Fut, + Fut: Future<Output = Result<T, E>>, + { + if self.initialized() { + // SAFETY: once the value is initialized, no mutable references are given out, so + // we can give out arbitrarily many immutable references + unsafe { Ok(self.get_unchecked()) } + } else { + // After acquire().await we have either acquired a permit while self.value + // is still uninitialized, or the current thread is awoken after another thread + // has initialized the value and closed the semaphore, in which case self.initialized + // is true and we don't set the value here + match self.semaphore.acquire().await { + Ok(_permit) => { + if !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 { + Ok(value) => { + // SAFETY: There is only one permit on the semaphore, hence only one + // mutable reference is created + unsafe { self.set_value(value) }; + + // SAFETY: once the value is initialized, no mutable references are given out, so + // we can give out arbitrarily many immutable references + unsafe { Ok(self.get_unchecked()) } + } + Err(e) => Err(e), + } + } else { + unreachable!("acquired semaphore after value was already initialized."); + } + } + Err(_) => { + if self.initialized() { + // SAFETY: once the value is initialized, no mutable references are given out, so + // we can give out arbitrarily many immutable references + unsafe { Ok(self.get_unchecked()) } + } else { + unreachable!( + "Semaphore closed, but the OnceCell has not been initialized." + ); + } + } + } + } + } + + /// Moves the value out of the cell, destroying the cell in the process. + /// + /// Returns `None` if the cell is uninitialized. + pub fn into_inner(mut self) -> Option<T> { + if self.initialized() { + // Set to uninitialized for the destructor of `OnceCell` to work properly + *self.value_set.get_mut() = false; + Some(unsafe { self.value.with(|ptr| ptr::read(ptr).assume_init()) }) + } else { + None + } + } + + /// Takes ownership of the current value, leaving the cell uninitialized. + /// + /// Returns `None` if the cell is uninitialized. + pub fn take(&mut self) -> Option<T> { + std::mem::take(self).into_inner() + } +} + +// Since `get` gives us access to immutable references of the +// OnceCell, OnceCell can only be Sync if T is Sync, otherwise +// OnceCell would allow sharing references of !Sync values across +// threads. We need T to be Send in order for OnceCell to by Sync +// because we can use `set` on `&OnceCell<T>` to send +// values (of type T) across threads. +unsafe impl<T: Sync + Send> Sync for OnceCell<T> {} + +// Access to OnceCell's value is guarded by the semaphore permit +// and atomic operations on `value_set`, so as long as T itself is Send +// it's safe to send it to another thread +unsafe impl<T: Send> Send for OnceCell<T> {} + +/// Errors that can be returned from [`OnceCell::set`] +/// +/// [`OnceCell::set`]: crate::sync::OnceCell::set +#[derive(Debug, PartialEq)] +pub enum SetError<T> { + /// Error resulting from [`OnceCell::set`] calls if the cell was previously initialized. + /// + /// [`OnceCell::set`]: crate::sync::OnceCell::set + AlreadyInitializedError(T), + + /// Error resulting from [`OnceCell::set`] calls when the cell is currently being + /// initialized during the calls to that method. + /// + /// [`OnceCell::set`]: crate::sync::OnceCell::set + InitializingError(T), +} + +impl<T> fmt::Display for SetError<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + SetError::AlreadyInitializedError(_) => write!(f, "AlreadyInitializedError"), + SetError::InitializingError(_) => write!(f, "InitializingError"), + } + } +} + +impl<T: fmt::Debug> Error for SetError<T> {} + +impl<T> SetError<T> { + /// Whether `SetError` is `SetError::AlreadyInitializedError`. + pub fn is_already_init_err(&self) -> bool { + match self { + SetError::AlreadyInitializedError(_) => true, + SetError::InitializingError(_) => false, + } + } + + /// Whether `SetError` is `SetError::InitializingError` + pub fn is_initializing_err(&self) -> bool { + match self { + SetError::AlreadyInitializedError(_) => false, + SetError::InitializingError(_) => true, + } + } +} |