Last active
April 2, 2021 10:03
-
-
Save thomcc/beb17f576951f8ddba35cd3c904f6632 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
use core::sync::atomic::{*, Ordering::*}; | |
use core::mem::MaybeUninit; | |
use core::cell::UnsafeCell; | |
/// A macro similar to `lazy_static::lazy_static!`, but that will attempt to | |
/// eagerly initialize the static value using [`on_startup!`](crate::on_startup). | |
#[macro_export] | |
macro_rules! eager_static { | |
($( | |
$(#[$attr:meta])* | |
$v:vis static ref $N:ident : $T:ty = $init:expr; | |
)*) => { | |
$( | |
$(#[$attr])* | |
$v static $N: $crate::_private::EagerStatic<$T> = | |
$crate::_private::EagerStatic::<$T>::__new(|| $init); | |
)* | |
$crate::on_startup!{ | |
$(let _: &$T = &*$N;)* | |
} | |
}; | |
} | |
/// `EagerStatic` is a wrapper type that `Deref`s to the target type, and can be | |
/// used as such. It's mostly private, but is generated by the `eager_static!` | |
/// macro. | |
/// | |
/// When you do: | |
/// | |
/// ``` | |
/// startup::eager_static! { | |
/// static ref FOO: Vec<i32> = vec![1, 2, 3, 4]; | |
/// } | |
/// ``` | |
/// | |
/// The actual item that is produced is something like: | |
/// | |
/// ``` | |
/// # stringify! { | |
/// static FOO: startup::_private::EagerStatic<T> = ...; | |
/// # } | |
/// ``` | |
pub struct EagerStatic<T, F = fn() -> T> { | |
state: AtomicUsize, | |
inner: UnsafeCell<MaybeUninit<T>>, | |
init_fn: F, | |
} | |
const UNINIT: usize = 0; | |
const INITIALIZING: usize = 1; | |
const INITIALIZED: usize = 2; | |
// FIXME: poisoning sucks. | |
const PANICED: usize = 3; | |
impl<T, F> EagerStatic<T, F> { | |
#[inline] | |
#[doc(hidden)] | |
pub const fn __new(init_fn: F) -> Self { | |
EagerStatic { | |
state: AtomicUsize::new(0), | |
inner: UnsafeCell::new(MaybeUninit::uninit()), | |
init_fn | |
} | |
} | |
} | |
impl<T> core::ops::Deref for EagerStatic<T> { | |
type Target = T; | |
#[inline] | |
fn deref(&self) -> &T { | |
// use a double-check sorta scheme to avoid needing to RMW in the common case. | |
let v = self.state.load(Acquire); | |
if v != INITIALIZED { | |
self.try_init(); | |
debug_assert_eq!(self.state.load(Acquire), INITIALIZED); | |
} | |
unsafe { &*self.inner.get().cast::<T>() } | |
} | |
} | |
struct StateResetOnDrop<'a>(&'a AtomicUsize); | |
impl<'a> Drop for StateResetOnDrop<'a> { | |
fn drop(&mut self) { | |
self.0.store(PANICED, Release); | |
} | |
} | |
#[inline(never)] | |
#[cold] | |
fn poisoned() -> ! { | |
panic!("I don't really want to commit to supporting this behavior, but for now we poison on panic"); | |
} | |
impl<T> EagerStatic<T> { | |
// separated to discourage inlining | |
fn try_init(&self) { | |
match self.state.compare_exchange(UNINIT, INITIALIZING, Acquire, Relaxed) { | |
Err(PANICED) => poisoned(), | |
Err(INITIALIZED) => {}, | |
Err(INITIALIZING) => self.wait_for_initialized(), | |
Ok(_) => unsafe { | |
let panic_guard = StateResetOnDrop(&self.state); | |
self.inner.get().write(MaybeUninit::new((self.init_fn)())); | |
core::mem::forget(panic_guard); | |
self.state.store(INITIALIZED, Release); | |
} | |
// could use unchecked unreachable in release code... | |
Err(state) if cfg!(debug_assertions) => unreachable!("bug: {}", state), | |
_ => {} | |
} | |
} | |
// Realistically this shouldn't get called outside of init cycles or when | |
// running on a platform with no support for static init. We just spin | |
#[cold] | |
#[inline(never)] | |
fn wait_for_initialized(&self) { | |
loop { | |
// if we cared this could be load(Relaxed) w/ a fence(Acquire) before | |
// returning. In practice this function should never be called | |
// though. | |
match self.state.load(Acquire) { | |
INITIALIZED => break, | |
INITIALIZING => core::hint::spin_loop(), | |
PANICED => poisoned(), | |
// could use unchecked unreachable in release code... | |
state if cfg!(debug_assertions) => unreachable!("bug: {}", state), | |
_ => {} | |
} | |
} | |
} | |
} | |
unsafe impl<T: Send> Send for EagerStatic<T> {} | |
unsafe impl<T: Send + Sync> Sync for EagerStatic<T> {} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment