kernel_api/sync/once.rs
1use core::cell::UnsafeCell;
2use core::fmt;
3use core::mem::MaybeUninit;
4use core::ops::Deref;
5use core::sync::atomic::{AtomicU8, fence, Ordering};
6
7/// A low-level synchronization primitive for one-time global execution.
8///
9/// # Examples
10///
11/// ```
12/// use kernel_api::sync::Once;
13///
14/// static START: Once = Once::new();
15///
16/// START.call_once(|| {
17/// // run initialization here
18/// });
19/// ```
20pub struct Once(AtomicU8);
21
22impl fmt::Debug for Once {
23 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24 f.debug_tuple("Once")
25 .field(&State::try_from(self.0.load(Ordering::SeqCst)).unwrap())
26 .finish()
27 }
28}
29
30#[derive(Debug, Copy, Clone, Eq, PartialEq)]
31#[repr(u8)]
32enum State {
33 Uncalled = 0,
34 Running = 1,
35 Called = 2,
36 Poison = 3
37}
38
39impl State {
40 const fn const_into_u8(self) -> u8 {
41 match self {
42 State::Uncalled => 0,
43 State::Running => 1,
44 State::Called => 2,
45 State::Poison => 3
46 }
47 }
48
49 const fn const_from_u8(value: u8) -> Result<Self, ()> {
50 match value {
51 0 => Ok(State::Uncalled),
52 1 => Ok(State::Running),
53 2 => Ok(State::Called),
54 3 => Ok(State::Poison),
55 _ => Err(())
56 }
57 }
58}
59
60impl From<State> for u8 {
61 fn from(value: State) -> Self {
62 value.const_into_u8()
63 }
64}
65
66impl TryFrom<u8> for State {
67 type Error = ();
68
69 fn try_from(value: u8) -> Result<Self, Self::Error> {
70 Self::const_from_u8(value)
71 }
72}
73
74impl Once {
75 /// Creates a new `Once` value.
76 #[must_use]
77 pub const fn new() -> Self {
78 Self(AtomicU8::new(State::Uncalled.const_into_u8()))
79 }
80
81 /// Performs an initialization routine once and only once. The given closure
82 /// will be executed if this is the first time `call_once` has been called,
83 /// and otherwise the routine will *not* be invoked.
84 ///
85 /// This method will spin if another initialization
86 /// routine is currently running.
87 ///
88 /// When this function returns, it is guaranteed that some initialization
89 /// has run and completed (it might not be the closure specified). It is also
90 /// guaranteed that any memory writes performed by the executed closure can
91 /// be reliably observed by other threads at this point (there is a
92 /// happens-before relation between the closure and code executing after the
93 /// return).
94 ///
95 /// If the given closure recursively invokes `call_once` on the same [`Once`]
96 /// instance, the exact behavior is not specified: allowed outcomes are
97 /// a panic or a deadlock.
98 ///
99 /// # Examples
100 ///
101 /// ```
102 /// use kernel_api::sync::Once;
103 ///
104 /// static mut VAL: usize = 0;
105 /// static INIT: Once = Once::new();
106 ///
107 /// // Accessing a `static mut` is unsafe much of the time, but if we do so
108 /// // in a synchronized fashion (e.g., write once or read all) then we're
109 /// // good to go!
110 /// //
111 /// // This function will only call `expensive_computation` once, and will
112 /// // otherwise always return the value returned from the first invocation.
113 /// fn get_cached_val() -> usize {
114 /// unsafe {
115 /// INIT.call_once(|| {
116 /// VAL = expensive_computation();
117 /// });
118 /// VAL
119 /// }
120 /// }
121 ///
122 /// fn expensive_computation() -> usize {
123 /// // ...
124 /// # 2
125 /// }
126 /// ```
127 ///
128 /// # Panics
129 ///
130 /// The closure `f` will only be executed once even if this is called
131 /// concurrently amongst many threads. If that closure panics, however, then
132 /// it will *poison* this [`Once`] instance, causing all future invocations of
133 /// `call_once` to also panic.
134 pub fn call_once<F: FnOnce()>(&self, f: F) {
135 loop {
136 let current = self.0.compare_exchange_weak(State::Uncalled.into(), State::Running.into(), Ordering::Relaxed, Ordering::Acquire);
137 match current {
138 Ok(_) => break, // Switched from Uncalled to Running, call the function
139 Err(s) if s == State::Poison.into() => panic!("poisoned `Once`"),
140 Err(s) if s == State::Running.into() => {}, // Currently running, spin until state changes
141 Err(s) if s == State::Called.into() => return, // Already called, return immediately
142 Err(s) if s == State::Uncalled.into() => {}, // weak cas fail, try again
143 _ => unreachable!()
144 }
145 core::hint::spin_loop();
146 }
147
148 struct DropGuard<'a>(&'a Once);
149 impl Drop for DropGuard<'_> {
150 fn drop(&mut self) {
151 self.0.0.store(State::Poison.into(), Ordering::Relaxed);
152 }
153 }
154 let drop_guard = DropGuard(self);
155
156 f();
157
158 core::mem::forget(drop_guard);
159
160 self.0.store(State::Called.into(), Ordering::Release);
161 }
162
163 /// Returns `true` if some [`call_once()`] call has completed
164 /// successfully. Specifically, `is_completed` will return false in
165 /// the following situations:
166 /// * [`call_once()`] was not called at all,
167 /// * [`call_once()`] was called, but has not yet completed,
168 /// * the [`Once`] instance is poisoned
169 ///
170 /// This function returning `false` does not mean that [`Once`] has not been
171 /// executed. For example, it may have been executed in the time between
172 /// when `is_completed` starts executing and when it returns, in which case
173 /// the `false` return value would be stale (but still permissible).
174 ///
175 /// # Panics
176 ///
177 /// If this `Once` is poisoned due to the closure passed to [`call_once`] panicking.
178 ///
179 /// # Examples
180 ///
181 /// ```
182 /// use kernel_api::sync::Once;
183 ///
184 /// static INIT: Once = Once::new();
185 ///
186 /// assert_eq!(INIT.is_completed(), false);
187 /// INIT.call_once(|| {
188 /// assert_eq!(INIT.is_completed(), false);
189 /// });
190 /// assert_eq!(INIT.is_completed(), true);
191 /// ```
192 #[must_use]
193 pub fn is_complete(&self) -> bool {
194 let state = self.0.load(Ordering::Relaxed).try_into().unwrap();
195
196 matches!(state, State::Called)
197 }
198}
199
200const impl Default for Once {
201 fn default() -> Self {
202 Self::new()
203 }
204}
205
206/// A synchronization primitive which can nominally be written to only once.
207///
208/// This type is a thread-safe [`OnceCell`](`core::cell::OnceCell`), and can be used in statics.
209/// In many simple cases, you can use [`LazyLock<T, F>`] instead to get the benefits of this type
210/// with less effort: `LazyLock<T, F>` "looks like" `&T` because it initializes with `F` on deref!
211/// Where `OnceLock` shines is when `LazyLock` is too simple to support a given case, as `LazyLock`
212/// doesn't allow additional inputs to its function after you call [`LazyLock::new(|| ...)`](`LazyLock::new`).
213///
214/// A `OnceLock` can be thought of as a safe abstraction over uninitialized data that becomes
215/// initialized once written.
216/*///
217/// # Examples
218///
219/// Writing to a `OnceLock` from a separate thread:
220///
221/// ```
222/// use kernel_api::sync::OnceLock;
223///
224/// static CELL: OnceLock<usize> = OnceLock::new();
225///
226/// // `OnceLock` has not been written to yet.
227/// assert!(CELL.get().is_none());
228///
229/// // Spawn a thread and write to `OnceLock`.
230/// std::thread::spawn(|| {
231/// let value = CELL.get_or_init(|| 12345);
232/// assert_eq!(value, &12345);
233/// })
234/// .join()
235/// .unwrap();
236///
237/// // `OnceLock` now contains the value.
238/// assert_eq!(
239/// CELL.get(),
240/// Some(&12345),
241/// );
242/// ```*/
243#[derive(Debug)]
244pub struct OnceLock<T> {
245 data: UnsafeCell<MaybeUninit<T>>,
246 once: Once
247}
248
249// SAFETY: `T` is `Send` therefore exposing it to other threads is sound
250unsafe impl<T: Send> Send for OnceLock<T> {}
251// SAFETY: `T` is `Sync` therefore exposing it to other threads via ref is sound
252unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
253
254impl<T> OnceLock<T> {
255 /// Creates a new uninitialized cell.
256 #[must_use]
257 pub const fn new() -> Self {
258 Self {
259 data: UnsafeCell::new(MaybeUninit::uninit()),
260 once: Once::new()
261 }
262 }
263
264 /// Gets the reference to the underlying value.
265 ///
266 /// Returns `None` if the cell is uninitialized, or being initialized.
267 /// This method never blocks.
268 #[inline]
269 pub fn get(&self) -> Option<&T> {
270 if !self.once.is_complete() { return None; }
271 fence(Ordering::Acquire);
272
273 // SAFETY: we only expose a shared ref so mutation cannot occur
274 let ret = unsafe { &*self.data.get() };
275 // SAFETY: checked above that the OnceLock is initialised
276 unsafe { Some(ret.assume_init_ref()) }
277 }
278
279 /// Gets the mutable reference to the underlying value.
280 ///
281 /// Returns `None` if the cell is uninitialized.
282 ///
283 /// This method never blocks. Since it borrows the `OnceLock` mutably,
284 /// it is statically guaranteed that no active borrows to the `OnceLock`
285 /// exist, including from other threads.
286 #[inline]
287 pub fn get_mut(&mut self) -> Option<&mut T> {
288 if !self.once.is_complete() { return None; }
289 fence(Ordering::Acquire);
290
291 // SAFETY: checked that OnceLock has been initialized
292 unsafe {
293 Some(self.data.get_mut().assume_init_mut())
294 }
295 }
296
297 /// Gets the contents of the cell, initializing it to `f()` if the cell
298 /// was uninitialized.
299 ///
300 /// Many threads may call `get_or_init` concurrently with different
301 /// initializing functions, but it is guaranteed that only one function
302 /// will be executed if the function doesn't panic.
303 ///
304 /// # Panics
305 ///
306 /// If `f()` panics, the panic is propagated to the caller, and the cell
307 /// remains uninitialized.
308 ///
309 /// It is an error to reentrantly initialize the cell from `f`. The
310 /// exact outcome is unspecified.
311 ///
312 /// # Examples
313 ///
314 /// ```
315 /// use kernel_api::sync::OnceLock;
316 ///
317 /// let cell = OnceLock::new();
318 /// let value = cell.get_or_init(|| 92);
319 /// assert_eq!(value, &92);
320 /// let value = cell.get_or_init(|| unreachable!());
321 /// assert_eq!(value, &92);
322 /// ```
323 #[inline]
324 pub fn get_or_init(&self, f: impl FnOnce() -> T) -> &T {
325 // SAFETY: `call_once` ensures that the closure only runs once so the pointer cannot be aliased
326 self.once.call_once(|| unsafe { (*self.data.get()).write(f()); });
327 // SAFETY: checked that OnceLock has been initialized
328 unsafe { (*self.data.get()).assume_init_ref() }
329 }
330}
331
332const impl<T> Default for OnceLock<T> {
333 fn default() -> Self {
334 Self::new()
335 }
336}
337
338/// A value which is initialized on the first access.
339///
340/// This type is a thread-safe [`LazyCell`](`core::cell::LazyCell`), and can be used in statics.
341/// Since initialization may be called from multiple threads, any
342/// dereferencing call will block the calling thread if another
343/// initialization routine is currently running.
344///
345/// # Poisoning
346///
347/// If the initialization closure passed to [`LazyLock::new`] panics, the lock will be poisoned.
348/// Once the lock is poisoned, any threads that attempt to access this lock (via a dereference
349/// or via an explicit call to [`force()`](`LazyLock::force`)) will panic.
350///
351/// # Examples
352///
353/// Initialize static variables with `LazyLock`.
354/// ```
355/// use kernel_api::sync::LazyLock;
356///
357/// // Note: static items do not call [`Drop`] on program termination, so this won't be deallocated.
358/// static DEEP_THOUGHT: LazyLock<String> = LazyLock::new(|| {
359/// # mod another_crate {
360/// # pub fn great_question() -> String { "42".to_string() }
361/// # }
362/// another_crate::great_question()
363/// });
364///
365/// // The `String` is built, stored in the `LazyLock`, and returned as `&String`.
366/// let _ = &*DEEP_THOUGHT;
367/// ```
368///
369/// Initialize fields with `LazyLock`.
370/// ```
371/// use kernel_api::sync::LazyLock;
372///
373/// #[derive(Debug)]
374/// struct UseCellLock {
375/// number: LazyLock<u32>,
376/// }
377/// fn main() {
378/// let lock: LazyLock<u32> = LazyLock::new(|| 0u32);
379///
380/// let data = UseCellLock { number: lock };
381/// assert_eq!(*data.number, 0);
382/// }
383/// ```
384#[derive(Debug)]
385pub struct LazyLock<T, F = fn() -> T> {
386 once: OnceLock<T>,
387 // FIXME: actually drop this when needed
388 f: MaybeUninit<F>
389}
390
391// SAFETY: Only move out of F and never create an &F so don't need f: Sync
392// LazyLock exposes &T from &LazyLock so T must be Sync + Send for LazyLock to be Sync
393unsafe impl<T: Sync + Send, F: Send> Sync for LazyLock<T, F> {}
394
395impl<T, F: FnOnce() -> T> LazyLock<T, F> {
396 /// Creates a new lazy value with the given initializing function.
397 ///
398 /// # Examples
399 ///
400 /// ```
401 /// use kernel_api::sync::LazyLock;
402 ///
403 /// let hello = "Hello, World!".to_string();
404 ///
405 /// let lazy = LazyLock::new(|| hello.to_uppercase());
406 ///
407 /// assert_eq!(&*lazy, "HELLO, WORLD!");
408 /// ```
409 #[must_use]
410 pub const fn new(f: F) -> Self {
411 Self {
412 once: OnceLock::new(),
413 f: MaybeUninit::new(f)
414 }
415 }
416
417 /// Forces the evaluation of this lazy value and returns a reference to
418 /// result. This is equivalent to the `Deref` impl, but is explicit.
419 ///
420 /// This method will block the calling thread if another initialization
421 /// routine is currently running.
422 ///
423 /// # Panics
424 ///
425 /// If the initialization closure panics (the one that is passed to the [`new()`](`LazyLock::new`) method), the
426 /// panic is propagated to the caller, and the lock becomes poisoned. This will cause all future
427 /// accesses of the lock (via [`force()`](`LazyLock::force`) or a dereference) to panic.
428 ///
429 /// # Examples
430 ///
431 /// ```
432 /// use kernel_api::sync::LazyLock;
433 ///
434 /// let lazy = LazyLock::new(|| 92);
435 ///
436 /// assert_eq!(LazyLock::force(&lazy), &92);
437 /// assert_eq!(&*lazy, &92);
438 /// ```
439 pub fn force(this: &Self) -> &T {
440 // SAFETY: `this.once` ensures `this.f` is only moved out of once
441 this.once.get_or_init(unsafe {
442 core::ptr::read(this.f.as_ptr())
443 })
444 }
445}
446
447impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
448 type Target = T;
449
450 fn deref(&self) -> &Self::Target {
451 Self::force(self)
452 }
453}