#![allow(unsafe_code)]
use core::convert::Infallible;
use core::error::Error;
use core::fmt;
use core::future::Future;
use core::mem::{self, MaybeUninit};
use core::ptr;
use super::contend::Closed;
use super::internal::{AtomicPtr, Ordering, UnsafeCell, take_ptr};
use super::raw_semaphore::{RawPermit, RawSemaphore, observed_closed};
const fn inline_marker<T>() -> *mut T {
ptr::without_provenance_mut(usize::MAX)
}
#[derive(Debug, PartialEq, Eq)]
pub enum SetError<T> {
AlreadyInitializedError(T),
InitializingError(T),
}
impl<T> fmt::Display for SetError<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::AlreadyInitializedError(_) => f.write_str("AlreadyInitializedError"),
Self::InitializingError(_) => f.write_str("InitializingError"),
}
}
}
impl<T: fmt::Debug> Error for SetError<T> {}
impl<T> SetError<T> {
pub fn is_already_init_err(&self) -> bool {
matches!(self, Self::AlreadyInitializedError(_))
}
pub fn is_initializing_err(&self) -> bool {
matches!(self, Self::InitializingError(_))
}
}
pub struct OnceCell<T> {
value: AtomicPtr<T>,
slot: UnsafeCell<MaybeUninit<T>>,
init: RawSemaphore,
}
unsafe impl<T: Send + Sync> Sync for OnceCell<T> {}
impl<T> OnceCell<T> {
loom_const_fn! {
pub fn new() -> Self {
Self {
value: AtomicPtr::new(ptr::null_mut()),
slot: UnsafeCell::new(MaybeUninit::uninit()),
init: RawSemaphore::without_spin(1),
}
}
}
loom_const_fn! {
pub fn const_new() -> Self {
Self::new()
}
}
pub fn new_with(value: Option<T>) -> Self {
match value {
Some(value) => Self::from(value),
None => Self::new(),
}
}
loom_const_fn! {
pub fn const_new_with(value: T) -> Self {
Self {
value: AtomicPtr::new(inline_marker()),
slot: UnsafeCell::new(MaybeUninit::new(value)),
init: RawSemaphore::new_closed(),
}
}
}
pub fn initialized(&self) -> bool {
!self.value.load(Ordering::Acquire).is_null()
}
fn resolve(&self, stored: *mut T) -> &T {
if stored == inline_marker() {
unsafe { &*self.slot.with(|slot| slot.cast::<T>()) }
} else {
unsafe { &*stored }
}
}
pub fn get(&self) -> Option<&T> {
let stored = self.value.load(Ordering::Acquire);
(!stored.is_null()).then(|| self.resolve(stored))
}
pub fn get_mut(&mut self) -> Option<&mut T> {
let stored = self.value.load(Ordering::Acquire);
if stored.is_null() {
None
} else if stored == inline_marker() {
Some(unsafe { &mut *self.slot.with_mut(|slot| slot.cast::<T>()) })
} else {
Some(unsafe { &mut *stored })
}
}
fn take_value(&mut self) -> Option<T> {
let stored = take_ptr(&mut self.value);
if stored.is_null() {
None
} else if stored == inline_marker() {
Some(unsafe { self.slot.with_mut(|slot| (*slot).assume_init_read()) })
} else {
Some(*unsafe { Box::from_raw(stored) })
}
}
pub fn into_inner(mut self) -> Option<T> {
self.take_value()
}
pub fn take(&mut self) -> Option<T> {
mem::take(self).into_inner()
}
pub fn set(&self, value: T) -> Result<(), SetError<T>> {
if self.initialized() {
return Err(SetError::AlreadyInitializedError(value));
}
match self.init.try_acquire_permit(1) {
Ok(turn) => self
.publish_inline(value, turn)
.map(drop)
.map_err(|(value, _)| SetError::AlreadyInitializedError(value)),
Err(seen) if observed_closed(seen) => Err(SetError::AlreadyInitializedError(value)),
Err(_) => Err(SetError::InitializingError(value)),
}
}
pub fn get_or_init_racy(&self, init: impl FnOnce() -> T) -> &T {
match self.get() {
Some(value) => value,
None => self.publish_boxed(init()),
}
}
pub async fn get_or_init<F, Fut>(&self, f: F) -> &T
where
F: FnOnce() -> Fut,
Fut: Future<Output = T>,
{
match self
.get_or_try_init(|| async move { Ok::<T, Infallible>(f().await) })
.await
{
Ok(value) => value,
Err(never) => match never {},
}
}
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 let Some(value) = self.get() {
return Ok(value);
}
match self.init.acquire(1).await {
Ok(turn) => {
if let Some(value) = self.get() {
return Ok(value);
}
let value = f().await?;
Ok(match self.publish_inline(value, turn) {
Ok(value) | Err((_, value)) => value,
})
}
Err(Closed) => match self.get() {
Some(value) => Ok(value),
None => Ok(self.publish_boxed(f().await?)),
},
}
}
fn publish_inline(&self, value: T, turn: RawPermit<'_>) -> Result<&T, (T, &T)> {
self.slot.with_mut(|slot| unsafe {
(*slot).write(value);
});
match self.value.compare_exchange(
ptr::null_mut(),
inline_marker(),
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
turn.forget();
self.init.close();
Ok(self.resolve(inline_marker()))
}
Err(winner) => Err((
unsafe { self.slot.with_mut(|slot| (*slot).assume_init_read()) },
self.resolve(winner),
)),
}
}
fn publish_boxed(&self, value: T) -> &T {
let fresh = Box::into_raw(Box::new(value));
match self.value.compare_exchange(
ptr::null_mut(),
fresh,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
self.init.close();
unsafe { &*fresh }
}
Err(winner) => {
drop(unsafe { Box::from_raw(fresh) });
self.resolve(winner)
}
}
}
}
impl<T> Drop for OnceCell<T> {
fn drop(&mut self) {
drop(self.take_value());
}
}
impl<T> Default for OnceCell<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: fmt::Debug> fmt::Debug for OnceCell<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("OnceCell")
.field("value", &self.get())
.finish()
}
}
impl<T: Clone> Clone for OnceCell<T> {
fn clone(&self) -> Self {
Self::new_with(self.get().cloned())
}
}
impl<T: PartialEq> PartialEq for OnceCell<T> {
fn eq(&self, other: &Self) -> bool {
self.get() == other.get()
}
}
impl<T: Eq> Eq for OnceCell<T> {}
impl<T> From<T> for OnceCell<T> {
fn from(value: T) -> Self {
Self::const_new_with(value)
}
}
#[cfg(all(test, not(loom)))]
mod tests {
use super::super::queue::Policy;
use super::*;
use crate::sync::Event;
use crate::sync::test_support::{Polled, block_on};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering as StdOrdering};
#[test]
fn a_waiter_for_the_initializer_turn_never_spins() {
let cell = OnceCell::<u8>::new();
assert_eq!(cell.init.queue().spin.rounds(), 0);
let opened = Event::new();
let mut first = Polled::new(cell.get_or_init(|| async {
opened.wait().await;
1
}));
first.pending();
let mut second = Polled::new(cell.get_or_init(|| async { 2 }));
second.pending();
opened.set();
assert_eq!(*first.ready(), 1);
assert_eq!(*second.ready(), 1);
assert_eq!(cell.init.queue().spin.rounds(), 0, "the turn never spins");
}
#[test]
fn set_publishes_once() {
let cell = OnceCell::new();
assert_eq!(cell.get(), None);
assert!(!cell.initialized());
assert_eq!(cell.set(1), Ok(()));
assert!(cell.initialized());
assert_eq!(cell.set(2), Err(SetError::AlreadyInitializedError(2)));
assert_eq!(cell.get(), Some(&1));
assert_eq!(cell.into_inner(), Some(1));
}
#[test]
fn set_fails_while_an_initializer_holds_the_turn() {
let cell = OnceCell::new();
let mut running = Polled::new(cell.get_or_init(core::future::pending::<u8>));
running.pending();
let err = cell.set(3).unwrap_err();
assert!(err.is_initializing_err() && !err.is_already_init_err());
assert_eq!(err, SetError::InitializingError(3));
assert_eq!(err.to_string(), "InitializingError");
assert_eq!(cell.get(), None);
drop(running);
assert_eq!(cell.set(3), Ok(()), "the turn was given up");
let err = cell.set(4).unwrap_err();
assert!(err.is_already_init_err());
assert_eq!(err.to_string(), "AlreadyInitializedError");
}
#[test]
fn a_racing_builder_publishes_beside_a_running_initializer() {
let cell = OnceCell::new();
let mut stuck = Polled::new(cell.get_or_init(core::future::pending));
stuck.pending();
let mut waiting = Polled::new(cell.get_or_init(|| async { 2 }));
waiting.pending();
assert_eq!(*cell.get_or_init_racy(|| 1), 1);
assert_eq!(waiting.wakes(), 1, "the publisher wakes the queue");
assert_eq!(*waiting.ready(), 1);
let mut late = Polled::new(cell.get_or_init(|| async { 3 }));
assert_eq!(*late.ready(), 1);
assert_eq!(
cell.set(9),
Err(SetError::AlreadyInitializedError(9)),
"the cell is full, not being initialized"
);
drop(stuck);
}
#[test]
fn an_initializer_that_lost_to_a_racing_builder_drops_its_value() {
let cell = OnceCell::new();
let opened = Event::new();
let mut slow = Polled::new(cell.get_or_init(|| async {
opened.wait().await;
Arc::new(5)
}));
slow.pending();
let racy = Arc::new(7);
assert_eq!(**cell.get_or_init_racy(|| Arc::clone(&racy)), 7);
opened.set();
assert_eq!(**slow.ready(), 7, "the racer's value won");
assert_eq!(Arc::strong_count(&racy), 2, "the cell holds the winner");
}
#[test]
fn racing_builders_publish_exactly_one_value() {
let cell = Arc::new(OnceCell::new());
let built = Arc::new(AtomicUsize::new(0));
let seen: Vec<usize> = (0..4)
.map(|i| {
let (cell, built) = (Arc::clone(&cell), Arc::clone(&built));
std::thread::spawn(move || {
*cell.get_or_init_racy(|| {
built.fetch_add(1, StdOrdering::SeqCst);
i
})
})
})
.map(|thread| thread.join().expect("racer"))
.collect();
assert!(seen.iter().all(|v| *v == seen[0]), "{seen:?}");
assert!(built.load(StdOrdering::SeqCst) >= 1);
}
#[test]
fn one_initializer_runs_while_the_others_wait() {
let cell = OnceCell::new();
let opened = Event::new();
let mut first = Polled::new(cell.get_or_init(|| async {
opened.wait().await;
1
}));
first.pending();
let mut second = Polled::new(cell.get_or_init(|| async { 2 }));
second.pending();
opened.set();
assert_eq!(*first.ready(), 1);
assert_eq!(second.wakes(), 1);
assert_eq!(*second.ready(), 1);
}
#[test]
fn a_closed_turn_queue_over_an_empty_cell_still_answers() {
let cell = OnceCell::<u8>::new();
cell.init.close();
assert_eq!(*block_on(cell.get_or_init(|| async { 5 })), 5);
assert_eq!(cell.get(), Some(&5));
let failed = OnceCell::<u8>::new();
failed.init.close();
assert_eq!(
block_on(failed.get_or_try_init(|| async { Err::<u8, &str>("no") })),
Err("no")
);
assert_eq!(failed.get(), None);
}
#[test]
fn the_initializer_is_not_called_when_the_cell_is_full() {
let cell = OnceCell::from(1);
let mut called = false;
let value = block_on(cell.get_or_init(|| {
called = true;
async { 2 }
}));
assert_eq!(*value, 1);
assert!(!called);
let tried = block_on(cell.get_or_try_init(|| async { Err::<u8, ()>(()) }));
assert_eq!(tried, Ok(&1));
}
#[test]
fn a_dropped_initializer_passes_the_turn_on() {
let cell = OnceCell::new();
let mut first = Polled::new(cell.get_or_init(core::future::pending));
first.pending();
let mut second = Polled::new(cell.get_or_init(|| async { 2 }));
second.pending();
drop(first);
assert_eq!(second.wakes(), 1);
assert_eq!(*second.ready(), 2);
}
#[test]
fn a_failed_initializer_lets_the_next_one_try() {
let cell = OnceCell::new();
let opened = Event::new();
let mut failing = Polled::new(cell.get_or_try_init(|| async {
opened.wait().await;
Err::<u32, &str>("no")
}));
failing.pending();
let mut next = Polled::new(cell.get_or_try_init(|| async { Ok::<u32, &str>(3) }));
next.pending();
opened.set();
assert_eq!(failing.ready(), Err("no"));
assert_eq!(next.ready(), Ok(&3));
}
#[test]
fn constructors_conversions_and_traits() {
static FULL: OnceCell<u32> = OnceCell::const_new_with(4);
static EMPTY: OnceCell<u32> = OnceCell::const_new();
assert_eq!(FULL.get(), Some(&4));
assert_eq!(EMPTY.get(), None);
assert_eq!(FULL.set(5), Err(SetError::AlreadyInitializedError(5)));
assert_eq!(*block_on(FULL.get_or_init(|| async { 6 })), 4);
assert_eq!(OnceCell::new_with(Some(1)).get(), Some(&1));
assert_eq!(OnceCell::<u8>::new_with(None).get(), None);
assert_eq!(OnceCell::<u8>::default().get(), None);
assert_eq!(OnceCell::from(2).into_inner(), Some(2));
assert_eq!(
format!("{:?}", OnceCell::from(2)),
"OnceCell { value: Some(2) }"
);
assert_eq!(
format!("{:?}", OnceCell::<u8>::new()),
"OnceCell { value: None }"
);
let cell = OnceCell::from(String::from("a"));
let copy = cell.clone();
assert_eq!(copy, cell);
assert_eq!(copy.get().map(String::as_str), Some("a"));
assert_ne!(cell, OnceCell::new());
assert_eq!(OnceCell::<u8>::new(), OnceCell::new());
assert_eq!(OnceCell::<String>::new().clone().get(), None);
}
#[test]
fn get_mut_and_take_work_for_either_home_of_the_value() {
let mut inline = OnceCell::new();
assert_eq!(inline.get_mut(), None);
assert_eq!(inline.set(vec![1]), Ok(()));
inline.get_mut().expect("full").push(2);
assert_eq!(inline.take(), Some(vec![1, 2]));
assert_eq!(inline.get(), None);
assert_eq!(inline.set(vec![3]), Ok(()), "a taken cell is set again");
let mut boxed = OnceCell::new();
assert_eq!(*boxed.get_or_init_racy(|| vec![1]), vec![1]);
boxed.get_mut().expect("full").push(2);
assert_eq!(boxed.take(), Some(vec![1, 2]));
assert_eq!(boxed.take(), None);
assert_eq!(*boxed.get_or_init_racy(|| vec![9]), vec![9]);
let mut const_full = OnceCell::const_new_with(String::from("c"));
const_full.get_mut().expect("full").push('d');
assert_eq!(const_full.into_inner().as_deref(), Some("cd"));
}
#[test]
fn every_value_is_dropped_exactly_once() {
let marker = Arc::new(());
let racy = OnceCell::new();
racy.get_or_init_racy(|| Arc::clone(&marker));
let inline = OnceCell::new();
assert!(inline.set(Arc::clone(&marker)).is_ok());
let lost = OnceCell::new();
assert!(lost.set(Arc::clone(&marker)).is_ok());
assert!(lost.set(Arc::clone(&marker)).is_err());
let from = OnceCell::from(Arc::clone(&marker));
assert_eq!(Arc::strong_count(&marker), 5);
drop((racy, inline, lost, from));
assert_eq!(Arc::strong_count(&marker), 1);
let zst = OnceCell::new();
assert!(zst.set(()).is_ok());
assert_eq!(zst.get(), Some(&()));
let zst_racy = OnceCell::<()>::new();
assert_eq!(zst_racy.get_or_init_racy(|| ()), &());
}
#[test]
fn threads_setting_and_initializing_agree_on_one_value() {
let rounds = if cfg!(miri) { 3 } else { 200 };
for _ in 0..rounds {
let cell = Arc::new(OnceCell::new());
let setter = {
let cell = Arc::clone(&cell);
std::thread::spawn(move || cell.set(1).is_ok())
};
let init = {
let cell = Arc::clone(&cell);
std::thread::spawn(move || *block_on(cell.get_or_init(|| async { 2 })))
};
let set_won = setter.join().expect("setter");
let got = init.join().expect("init");
assert_eq!(got, *cell.get().expect("full"));
if set_won {
assert_eq!(got, 1);
}
}
}
}