dtact_util/sync/
once_cell.rs

1//! Async once-initialized cell.
2
3use super::wait_queue::WaitQueue;
4use std::cell::UnsafeCell;
5use std::future::Future;
6use std::sync::atomic::{AtomicU8, Ordering};
7use std::task::{Context, Poll};
8
9const UNINIT: u8 = 0;
10const INITIALIZING: u8 = 1;
11const INIT: u8 = 2;
12
13/// A cell that's initialized at most once, asynchronously.
14///
15/// Concurrent callers of [`OnceCell::get_or_init`] that race to be first
16/// all wait for the same single initialization to complete rather than
17/// each running their own initializer.
18#[repr(align(64))]
19pub struct OnceCell<T> {
20    state: AtomicU8,
21    data: UnsafeCell<Option<T>>,
22    wait: WaitQueue,
23}
24
25// SAFETY: `state`'s CAS transitions are the sole gate on `data` access —
26// exactly one caller ever holds `INITIALIZING` and writes `data`, every
27// other reader only reads `data` after observing `INIT`.
28unsafe impl<T: Send> Send for OnceCell<T> {}
29unsafe impl<T: Send + Sync> Sync for OnceCell<T> {}
30
31impl<T> Default for OnceCell<T> {
32    fn default() -> Self {
33        Self::new()
34    }
35}
36
37impl<T> OnceCell<T> {
38    /// Create an empty, uninitialized cell.
39    #[must_use]
40    pub const fn new() -> Self {
41        Self {
42            state: AtomicU8::new(UNINIT),
43            data: UnsafeCell::new(None),
44            wait: WaitQueue::new(),
45        }
46    }
47
48    /// Create a cell already initialized with `value`.
49    #[must_use]
50    pub const fn new_with(value: T) -> Self {
51        Self {
52            state: AtomicU8::new(INIT),
53            data: UnsafeCell::new(Some(value)),
54            wait: WaitQueue::new(),
55        }
56    }
57
58    /// The cell's value, if it's already been initialized.
59    #[must_use]
60    #[inline(always)]
61    pub fn get(&self) -> Option<&T> {
62        (self.state.load(Ordering::Acquire) == INIT)
63            // SAFETY: `INIT` observed under Acquire pairs with the
64            // Release store in `get_or_init`/`set` that wrote `data`.
65            .then(|| unsafe { (*self.data.get()).as_ref() })
66            .flatten()
67    }
68
69    /// Set the cell's value if it isn't already initialized.
70    ///
71    /// # Errors
72    /// Returns `value` back if the cell was already initialized (or is
73    /// concurrently being initialized by [`Self::get_or_init`]).
74    #[inline(always)]
75    pub fn set(&self, value: T) -> Result<(), T> {
76        if self
77            .state
78            .compare_exchange(UNINIT, INITIALIZING, Ordering::AcqRel, Ordering::Acquire)
79            .is_err()
80        {
81            return Err(value);
82        }
83        // SAFETY: we hold the sole `INITIALIZING` token for this cell.
84        unsafe {
85            *self.data.get() = Some(value);
86        }
87        self.state.store(INIT, Ordering::Release);
88        self.wait.wake_all();
89        Ok(())
90    }
91
92    /// Return the cell's value, initializing it with (the result of
93    /// awaiting) `init` first if necessary. Concurrent callers racing to
94    /// initialize the same cell all observe the *same* initialization —
95    /// only the winner's `init` future actually runs.
96    pub async fn get_or_init<F, Fut>(&self, init: F) -> &T
97    where
98        F: FnOnce() -> Fut,
99        Fut: Future<Output = T>,
100        T: Send + Sync,
101    {
102        loop {
103            if let Some(v) = self.get() {
104                return v;
105            }
106            if self
107                .state
108                .compare_exchange(UNINIT, INITIALIZING, Ordering::AcqRel, Ordering::Acquire)
109                .is_ok()
110            {
111                let value = init().await;
112                // SAFETY: we hold the sole `INITIALIZING` token, so
113                // writing `data` then publishing `INIT` (and reading the
114                // just-written reference back out) is exclusive — no
115                // `.expect()`/panic path needed the way going back
116                // through `self.get()` would require.
117                let value_ref = unsafe {
118                    let slot = &mut *self.data.get();
119                    *slot = Some(value);
120                    slot.as_ref().unwrap_unchecked()
121                };
122                self.state.store(INIT, Ordering::Release);
123                self.wait.wake_all();
124                return value_ref;
125            }
126            // Someone else is initializing (or just finished, in which
127            // case `self.get()` at the top of the next iteration returns
128            // immediately) — wait for them.
129            std::future::poll_fn(|cx| self.poll_wait_for_init(cx)).await;
130        }
131    }
132
133    #[inline]
134    fn poll_wait_for_init(&self, cx: &Context<'_>) -> Poll<()> {
135        if self.state.load(Ordering::Acquire) == INIT {
136            return Poll::Ready(());
137        }
138        let token = self.wait.register(cx.waker());
139        if self.state.load(Ordering::Acquire) == INIT {
140            self.wait.cancel(token);
141            return Poll::Ready(());
142        }
143        Poll::Pending
144    }
145
146    /// Consume the cell, returning its value if initialized.
147    #[inline(always)]
148    pub fn into_inner(self) -> Option<T> {
149        self.data.into_inner()
150    }
151}
dtact_util/sync/once_cell.rs

dtact_util/sync/
once_cell.rs
