Skip to main content

rivet/sync/
once.rs

1//! A minimal no_std `Once` cell: written exactly once (at boot), read-only
2//! afterwards. Used for the `Channel::split`-once pattern (plan.md [B8]):
3//! the sender/receiver halves of a static channel are stored here once at
4//! boot and borrowed by tasks for the lifetime of the program.
5
6use core::cell::UnsafeCell;
7use core::mem::MaybeUninit;
8
9/// Single-writer, multi-reader `Once` cell. `set` may be called at most
10/// once (a second call returns `Err` with the value back); `get` returns
11/// `None` until then and `Some(&T)` forever after.
12pub struct Once<T> {
13    set: crate::sync::atomic::AtomicBool,
14    cell: UnsafeCell<MaybeUninit<T>>,
15}
16
17// Safety: `set` is single-threaded (boot); after publication (Release
18// store of the flag) the value is immutable, so shared `&T` reads are
19// sound. `T` must be `Sync` for `&T` to be shareable.
20unsafe impl<T: Sync> Sync for Once<T> {}
21
22impl<T> Once<T> {
23    /// Create an empty `Once`.
24    #[cfg(not(loom))]
25    pub const fn new() -> Self {
26        Self::new_impl()
27    }
28
29    #[cfg(loom)]
30    pub fn new() -> Self {
31        Self::new_impl()
32    }
33
34    #[cfg(not(loom))]
35    const fn new_impl() -> Self {
36        Self {
37            set: crate::sync::atomic::AtomicBool::new(false),
38            cell: UnsafeCell::new(MaybeUninit::uninit()),
39        }
40    }
41
42    #[cfg(loom)]
43    fn new_impl() -> Self {
44        Self {
45            set: crate::sync::atomic::AtomicBool::new(false),
46            cell: UnsafeCell::new(MaybeUninit::uninit()),
47        }
48    }
49
50    /// Store the value (boot time, single writer). Returns `Err(value)`
51    /// if already set.
52    pub fn set(&self, value: T) -> Result<(), T> {
53        if self.set.load(crate::sync::atomic::Ordering::Acquire) {
54            return Err(value);
55        }
56        // SAFETY: single writer (caller contract); the value is written
57        // before the flag is published with Release, so readers that see
58        // the flag also see the value (Release→Acquire).
59        unsafe {
60            (*self.cell.get()).write(value);
61        }
62        self.set.store(true, crate::sync::atomic::Ordering::Release);
63        Ok(())
64    }
65
66    /// Borrow the stored value, or `None` if not yet set.
67    pub fn get(&self) -> Option<&T> {
68        if !self.set.load(crate::sync::atomic::Ordering::Acquire) {
69            return None;
70        }
71        // SAFETY: the value was written before the flag was published; a
72        // reader that sees the flag (Acquire) observes the fully written
73        // value, and it is never mutated again.
74        unsafe { Some(&*(*self.cell.get()).as_ptr()) }
75    }
76}
77
78impl<T> Default for Once<T> {
79    fn default() -> Self {
80        Self::new()
81    }
82}