static_generics/

once.rs

//! Lazy, one-time initialization on top of
//! [`generic_static`](crate::namespace::Namespace::generic_static).
//!
//! [`generic_static`](crate::namespace::Namespace::generic_static). only supports
//! [`bytemuck::Zeroable`]
//! types, so the storage is always zeroed. That allows
//! `Atomic*`, [`UnsafeCell`],
//! [`Cell`](core::cell::Cell), and other similar types to be stored — but not types that need a
//! non-zero initial value.
//!
//! [`OnceSlot`] allows to store non Zeroable types as a generic static:
//! It is itself zeroable, but it can be lazily initialized provided by initializer closure
//! into [`OnceSlot::get_or_init`](OnceSlot::get)
//!
//! ```rust
//! use static_generics::{define_namespace, namespace::NamespaceExt};
//!
//! define_namespace!(MyNs);
//!
//! struct Config {
//!     retries: u32,
//! }
//!
//! let cfg = MyNs::once::<Config>(|| Config { retries: 3 });
//! assert_eq!(cfg.retries, 3);
//! // Second call ignores the closure and returns the same address.
//! let again = MyNs::once::<Config>(|| Config { retries: 99 });
//! assert!(core::ptr::eq(cfg, again));
//! ```

use core::cell::UnsafeCell;
use core::mem::MaybeUninit;
use core::sync::atomic::{AtomicU8, Ordering};

#[cfg(feature = "std")]
use parking_lot::{Condvar, Mutex};

const UNINIT: u8 = 0;
const BUSY: u8 = 1;
const INIT: u8 = 2;
const POISONED: u8 = 3;

struct PoisonGuard<'a, T> {
    slot: &'a OnceSlot<T>,
    done: bool,
}

impl<T> Drop for PoisonGuard<'_, T> {
    fn drop(&mut self) {
        if !self.done {
            self.slot.finish(POISONED);
        }
    }
}

/// Zero-initialized slot for one-time initialization of a `T`.
///
/// The slot itself starts zeroed (state `UNINIT`), which is what makes it
/// usable with `generic_static`.
/// The `T` value is written exactly once by the thread that wins the
/// claim race; other threads block until `INIT` (or `POISONED`) becomes
/// visible.
///
/// With the `std` feature, waiters block on a [`Mutex`]/[`Condvar`] pair
/// (no spinning); without it they spin with [`core::hint::spin_loop`].
///
/// # Panics
///
/// [`get_or_init`](OnceSlot::get_or_init) must not be called reentrantly
/// with the same slot (the initializer calling back into the same slot).
/// That deadlocks: the thread already holds the `BUSY` claim, so the inner
/// call blocks forever waiting for itself.
///
/// If the initializer panics, the slot is poisoned.
pub struct OnceSlot<T> {
    state: AtomicU8,
    #[cfg(feature = "std")]
    lock: Mutex<()>,
    #[cfg(feature = "std")]
    cvar: Condvar,
    value: UnsafeCell<MaybeUninit<T>>,
}
// SAFETY: `state` is 0 for UNINIT, value is zeroed via `MaybeUninit`, and lock/cvar in parking_lot are zeroable.
unsafe impl<T> bytemuck::Zeroable for OnceSlot<T> {}

unsafe impl<T: Send + Sync> Sync for OnceSlot<T> {}
unsafe impl<T: Send> Send for OnceSlot<T> {}

impl<T> OnceSlot<T> {
    /// Create a new uninitialized slot.
    #[must_use]
    pub const fn new() -> Self {
        Self {
            state: AtomicU8::new(UNINIT),
            #[cfg(feature = "std")]
            lock: Mutex::new(()),
            #[cfg(feature = "std")]
            cvar: Condvar::new(),
            value: UnsafeCell::new(MaybeUninit::uninit()),
        }
    }

    fn finish(&self, next: u8) {
        debug_assert!(next == INIT || next == POISONED);
        #[cfg(feature = "std")]
        {
            let _guard = self.lock.lock();
            self.state.store(next, Ordering::Release);
            self.cvar.notify_all();
        }
        #[cfg(not(feature = "std"))]
        {
            self.state.store(next, Ordering::Release);
        }
    }

    /// Returns the value if already initialized, otherwise `None`.
    #[must_use]
    pub fn get(&'static self) -> Option<&'static T> {
        if self.state.load(Ordering::Acquire) == INIT {
            // SAFETY: `INIT` is only stored after the payload was fully
            // written, with `Release` ordering. `Acquire` here synchronizes
            // with that store, so the payload is initialized.
            Some(unsafe { (*self.value.get()).assume_init_ref() })
        } else {
            None
        }
    }

    /// Returns `true` if initialization completed successfully.
    #[must_use]
    pub fn is_completed(&'static self) -> bool {
        self.state.load(Ordering::Acquire) == INIT
    }

    /// Returns `true` if the initializer panicked and poisoned the slot.
    #[must_use]
    pub fn is_poisoned(&'static self) -> bool {
        self.state.load(Ordering::Acquire) == POISONED
    }

    /// Block until the winner finishes, then return or panic on poison.
    fn wait(&'static self) -> &'static T {
        #[cfg(feature = "std")]
        {
            let mut guard = self.lock.lock();
            loop {
                match self.state.load(Ordering::Acquire) {
                    INIT => {
                        // SAFETY: same publication argument as in `get`.
                        return unsafe { (*self.value.get()).assume_init_ref() };
                    }
                    POISONED => panic!("OnceSlot has previously been poisoned"),
                    _ => self.cvar.wait(&mut guard),
                }
            }
        }
        #[cfg(not(feature = "std"))]
        {
            loop {
                match self.state.load(Ordering::Acquire) {
                    INIT => {
                        // SAFETY: same publication argument as in `get`.
                        return unsafe { (*self.value.get()).assume_init_ref() };
                    }
                    POISONED => panic!("OnceSlot has previously been poisoned"),
                    _ => core::hint::spin_loop(),
                }
            }
        }
    }

    /// Returns the value, running `init` exactly once to produce it.
    ///
    /// Concurrent callers block (on a condvar with `std`, by spinning
    /// without it) until the winner finishes. If `init` panics, the slot
    /// is poisoned and every caller — past, present, and future — panics.
    ///
    /// # Panics
    ///
    /// Panics if the slot is poisoned (including when a concurrent
    /// initializer panics while this call is blocked). Must not be called
    /// reentrantly with the same slot.
    pub fn get_or_init(&'static self, init: impl FnOnce() -> T) -> &'static T {
        if let Some(value) = self.get() {
            return value;
        }

        self.initialize(init)
    }

    #[cold]
    #[inline(never)]
    fn initialize(&'static self, init: impl FnOnce() -> T) -> &'static T {
        match self
            .state
            .compare_exchange(UNINIT, BUSY, Ordering::AcqRel, Ordering::Acquire)
        {
            Ok(_) => {
                let mut guard = PoisonGuard {
                    slot: self,
                    done: false,
                };
                let value = init();
                // SAFETY: we hold the unique `BUSY` claim, so no other
                // thread reads or writes the payload concurrently.
                unsafe {
                    (*self.value.get()).write(value);
                }
                self.finish(INIT);
                guard.done = true;
                // SAFETY: just stored `INIT` after writing the payload.
                unsafe { (*self.value.get()).assume_init_ref() }
            }
            Err(INIT) => {
                // Winner finished between our fast-path `get()` and the
                // `CAS`; the payload is initialized.
                // SAFETY: `INIT` observed, same as in `get`.
                unsafe { (*self.value.get()).assume_init_ref() }
            }
            Err(POISONED) => panic!("OnceSlot has previously been poisoned"),
            Err(_) => self.wait(),
        }
    }
}

impl<T> Default for OnceSlot<T> {
    fn default() -> Self {
        Self::new()
    }
}