tokio/util/
once_cell.rs

1#![allow(dead_code)]
2use std::cell::UnsafeCell;
3use std::mem::MaybeUninit;
4use std::sync::Once;
5
6pub(crate) struct OnceCell<T> {
7    once: Once,
8    value: UnsafeCell<MaybeUninit<T>>,
9}
10
11unsafe impl<T: Send + Sync> Send for OnceCell<T> {}
12unsafe impl<T: Send + Sync> Sync for OnceCell<T> {}
13
14impl<T> OnceCell<T> {
15    pub(crate) const fn new() -> Self {
16        Self {
17            once: Once::new(),
18            value: UnsafeCell::new(MaybeUninit::uninit()),
19        }
20    }
21
22    /// Get the value inside this cell, initializing it using the provided
23    /// function if necessary.
24    ///
25    /// If the `init` closure panics, then the `OnceCell` is poisoned and all
26    /// future calls to `get` will panic.
27    #[inline]
28    pub(crate) fn get(&self, init: impl FnOnce() -> T) -> &T {
29        if !self.once.is_completed() {
30            self.do_init(init);
31        }
32
33        // Safety: The `std::sync::Once` guarantees that we can only reach this
34        // line if a `call_once` closure has been run exactly once and without
35        // panicking. Thus, the value is not uninitialized.
36        //
37        // There is also no race because the only `&self` method that modifies
38        // `value` is `do_init`, but if the `call_once` closure is still
39        // running, then no thread has gotten past the `call_once`.
40        unsafe { &*(self.value.get() as *const T) }
41    }
42
43    #[cold]
44    fn do_init(&self, init: impl FnOnce() -> T) {
45        let value_ptr = self.value.get() as *mut T;
46
47        self.once.call_once(|| {
48            let set_to = init();
49
50            // Safety: The `std::sync::Once` guarantees that this initialization
51            // will run at most once, and that no thread can get past the
52            // `call_once` until it has run exactly once. Thus, we have
53            // exclusive access to `value`.
54            unsafe {
55                std::ptr::write(value_ptr, set_to);
56            }
57        });
58    }
59}
60
61impl<T> Drop for OnceCell<T> {
62    fn drop(&mut self) {
63        if self.once.is_completed() {
64            let value_ptr = self.value.get() as *mut T;
65            unsafe {
66                std::ptr::drop_in_place(value_ptr);
67            }
68        }
69    }
70}