std/sys/thread_local/native/
lazy.rs

1use crate::cell::{Cell, UnsafeCell};
2use crate::mem::MaybeUninit;
3use crate::ptr;
4use crate::sys::thread_local::{abort_on_dtor_unwind, destructors};
5
6pub unsafe trait DestroyedState: Sized + Copy {
7    fn register_dtor<T>(s: &Storage<T, Self>);
8}
9
10unsafe impl DestroyedState for ! {
11    fn register_dtor<T>(_: &Storage<T, !>) {}
12}
13
14unsafe impl DestroyedState for () {
15    fn register_dtor<T>(s: &Storage<T, ()>) {
16        unsafe {
17            destructors::register(ptr::from_ref(s).cast_mut().cast(), destroy::<T>);
18        }
19    }
20}
21
22#[derive(Copy, Clone)]
23enum State<D> {
24    Uninitialized,
25    Alive,
26    Destroyed(D),
27}
28
29#[allow(missing_debug_implementations)]
30#[repr(C)]
31pub struct Storage<T, D> {
32    // This field must be first, for correctness of `#[rustc_align_static]`
33    value: UnsafeCell<MaybeUninit<T>>,
34    state: Cell<State<D>>,
35}
36
37impl<T, D> Storage<T, D>
38where
39    D: DestroyedState,
40{
41    pub const fn new() -> Storage<T, D> {
42        Storage {
43            state: Cell::new(State::Uninitialized),
44            value: UnsafeCell::new(MaybeUninit::uninit()),
45        }
46    }
47
48    /// Gets a pointer to the TLS value, potentially initializing it with the
49    /// provided parameters. If the TLS variable has been destroyed, a null
50    /// pointer is returned.
51    ///
52    /// The resulting pointer may not be used after reentrant inialialization
53    /// or thread destruction has occurred.
54    ///
55    /// # Safety
56    /// The `self` reference must remain valid until the TLS destructor is run.
57    #[inline]
58    pub unsafe fn get_or_init(&self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
59        if let State::Alive = self.state.get() {
60            self.value.get().cast()
61        } else {
62            unsafe { self.get_or_init_slow(i, f) }
63        }
64    }
65
66    /// # Safety
67    /// The `self` reference must remain valid until the TLS destructor is run.
68    #[cold]
69    unsafe fn get_or_init_slow(
70        &self,
71        i: Option<&mut Option<T>>,
72        f: impl FnOnce() -> T,
73    ) -> *const T {
74        match self.state.get() {
75            State::Uninitialized => {}
76            State::Alive => return self.value.get().cast(),
77            State::Destroyed(_) => return ptr::null(),
78        }
79
80        let v = i.and_then(Option::take).unwrap_or_else(f);
81
82        // SAFETY: we cannot be inside a `LocalKey::with` scope, as the initializer
83        // has already returned and the next scope only starts after we return
84        // the pointer. Therefore, there can be no references to the old value,
85        // even if it was initialized. Thus because we are !Sync we have exclusive
86        // access to self.value and may replace it.
87        let mut old_value = unsafe { self.value.get().replace(MaybeUninit::new(v)) };
88        match self.state.replace(State::Alive) {
89            // If the variable is not being recursively initialized, register
90            // the destructor. This might be a noop if the value does not need
91            // destruction.
92            State::Uninitialized => D::register_dtor(self),
93
94            // Recursive initialization, we only need to drop the old value
95            // as we've already registered the destructor.
96            State::Alive => unsafe { old_value.assume_init_drop() },
97
98            State::Destroyed(_) => unreachable!(),
99        }
100
101        self.value.get().cast()
102    }
103}
104
105/// Transition an `Alive` TLS variable into the `Destroyed` state, dropping its
106/// value.
107///
108/// # Safety
109/// * Must only be called at thread destruction.
110/// * `ptr` must point to an instance of `Storage<T, ()>` and be valid for
111///   accessing that instance.
112unsafe extern "C" fn destroy<T>(ptr: *mut u8) {
113    // Print a nice abort message if a panic occurs.
114    abort_on_dtor_unwind(|| {
115        let storage = unsafe { &*(ptr as *const Storage<T, ()>) };
116        if let State::Alive = storage.state.replace(State::Destroyed(())) {
117            // SAFETY: we ensured the state was Alive so the value was initialized.
118            // We also updated the state to Destroyed to prevent the destructor
119            // from accessing the thread-local variable, as this would violate
120            // the exclusive access provided by &mut T in Drop::drop.
121            unsafe { (*storage.value.get()).assume_init_drop() }
122        }
123    })
124}
std/sys/thread_local/native/lazy.rs

std/sys/thread_local/native/
lazy.rs
