Revision control
Copy as Markdown
Other Tools
use core::cell::UnsafeCell;
use core::fmt;
use core::hint::spin_loop;
use core::sync::atomic::{AtomicUsize, Ordering};
/// A synchronization primitive which can be used to run a one-time global
/// initialization. Unlike its std equivalent, this is generalized so that the
/// closure returns a value and it is stored. Once therefore acts something like
/// a future, too.
pub struct Once<T> {
state: AtomicUsize,
data: UnsafeCell<Option<T>>, // TODO remove option and use mem::uninitialized
}
impl<T: fmt::Debug> fmt::Debug for Once<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.r#try() {
Some(s) => write!(f, "Once {{ data: ")
.and_then(|()| s.fmt(f))
.and_then(|()| write!(f, "}}")),
None => write!(f, "Once {{ <uninitialized> }}"),
}
}
}
// Same unsafe impls as `std::sync::RwLock`, because this also allows for
// concurrent reads.
unsafe impl<T: Send + Sync> Sync for Once<T> {}
unsafe impl<T: Send> Send for Once<T> {}
// Four states that a Once can be in, encoded into the lower bits of `state` in
// the Once structure.
const INCOMPLETE: usize = 0x0;
const RUNNING: usize = 0x1;
const COMPLETE: usize = 0x2;
const PANICKED: usize = 0x3;
use core::hint::unreachable_unchecked as unreachable;
impl<T> Once<T> {
/// Initialization constant of `Once`.
pub const INIT: Self = Once {
state: AtomicUsize::new(INCOMPLETE),
data: UnsafeCell::new(None),
};
/// Creates a new `Once` value.
pub const fn new() -> Once<T> {
Self::INIT
}
fn force_get<'a>(&'a self) -> &'a T {
match unsafe { &*self.data.get() }.as_ref() {
None => unsafe { unreachable() },
Some(p) => p,
}
}
/// Performs an initialization routine once and only once. The given closure
/// will be executed if this is the first time `call_once` has been called,
/// and otherwise the routine will *not* be invoked.
///
/// This method will block the calling thread if another initialization
/// routine is currently running.
///
/// When this function returns, it is guaranteed that some initialization
/// has run and completed (it may not be the closure specified). The
/// returned pointer will point to the result from the closure that was
/// run.
pub fn call_once<'a, F>(&'a self, builder: F) -> &'a T
where
F: FnOnce() -> T,
{
let mut status = self.state.load(Ordering::SeqCst);
if status == INCOMPLETE {
status = match self.state.compare_exchange(
INCOMPLETE,
RUNNING,
Ordering::SeqCst,
Ordering::SeqCst,
) {
Ok(status) => {
debug_assert_eq!(
status, INCOMPLETE,
"if compare_exchange succeeded, previous status must be incomplete",
);
// We init
// We use a guard (Finish) to catch panics caused by builder
let mut finish = Finish {
state: &self.state,
panicked: true,
};
unsafe { *self.data.get() = Some(builder()) };
finish.panicked = false;
self.state.store(COMPLETE, Ordering::SeqCst);
// This next line is strictly an optimization
return self.force_get();
}
Err(status) => status,
}
}
loop {
match status {
INCOMPLETE => unreachable!(),
RUNNING => {
// We spin
spin_loop();
status = self.state.load(Ordering::SeqCst)
}
PANICKED => panic!("Once has panicked"),
COMPLETE => return self.force_get(),
_ => unsafe { unreachable() },
}
}
}
/// Returns a pointer iff the `Once` was previously initialized
pub fn r#try<'a>(&'a self) -> Option<&'a T> {
match self.state.load(Ordering::SeqCst) {
COMPLETE => Some(self.force_get()),
_ => None,
}
}
/// Like try, but will spin if the `Once` is in the process of being
/// initialized
pub fn wait<'a>(&'a self) -> Option<&'a T> {
loop {
match self.state.load(Ordering::SeqCst) {
INCOMPLETE => return None,
RUNNING => {
spin_loop() // We spin
}
COMPLETE => return Some(self.force_get()),
PANICKED => panic!("Once has panicked"),
_ => unsafe { unreachable() },
}
}
}
}
struct Finish<'a> {
state: &'a AtomicUsize,
panicked: bool,
}
impl<'a> Drop for Finish<'a> {
fn drop(&mut self) {
if self.panicked {
self.state.store(PANICKED, Ordering::SeqCst);
}
}
}