Skip to main content

std\sys\thread_local/
os.rs

1use super::key::{Key, LazyKey, get, set};
2use super::{abort_on_dtor_unwind, guard};
3use crate::alloc::{self, GlobalAlloc, Layout, System};
4use crate::cell::Cell;
5use crate::marker::PhantomData;
6use crate::mem::ManuallyDrop;
7use crate::ops::Deref;
8use crate::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
9use crate::ptr::{self, NonNull};
10
11#[doc(hidden)]
12#[allow_internal_unstable(thread_local_internals)]
13#[allow_internal_unsafe]
14#[unstable(feature = "thread_local_internals", issue = "none")]
15#[rustc_macro_transparency = "semiopaque"]
16pub macro thread_local_inner {
17    // NOTE: we cannot import `Storage` or `LocalKey` with a `use` because that can shadow user
18    // provided type or type alias with a matching name. Please update the shadowing test in
19    // `tests/thread.rs` if these types are renamed.
20
21    // used to generate the `LocalKey` value for `thread_local!`.
22    (@key $t:ty, $($(#[$($align_attr:tt)*])+)?, $init:expr) => {{
23        #[inline]
24        fn __rust_std_internal_init_fn() -> $t { $init }
25
26        // NOTE: this cannot import `LocalKey` or `Storage` with a `use` because that can shadow
27        // user provided type or type alias with a matching name. Please update the shadowing test
28        // in `tests/thread.rs` if these types are renamed.
29        unsafe {
30            $crate::thread::LocalKey::new(|__rust_std_internal_init| {
31                static __RUST_STD_INTERNAL_VAL: $crate::thread::local_impl::Storage<$t, {
32                    $({
33                        // Ensure that attributes have valid syntax
34                        // and that the proper feature gate is enabled
35                        $(#[$($align_attr)*])+
36                        #[allow(unused)]
37                        static DUMMY: () = ();
38                    })?
39
40                    #[allow(unused_mut)]
41                    let mut final_align = $crate::thread::local_impl::value_align::<$t>();
42                    $($($crate::thread::local_impl::thread_local_inner!(@align final_align, $($align_attr)*);)+)?
43                    final_align
44                }>
45                    = $crate::thread::local_impl::Storage::new();
46                __RUST_STD_INTERNAL_VAL.get(__rust_std_internal_init, __rust_std_internal_init_fn)
47            })
48        }
49    }},
50
51    // process a single `rustc_align_static` attribute
52    (@align $final_align:ident, rustc_align_static($($align:tt)*) $(, $($attr_rest:tt)+)?) => {
53        let new_align: $crate::primitive::usize = $($align)*;
54        if new_align > $final_align {
55            $final_align = new_align;
56        }
57
58        $($crate::thread::local_impl::thread_local_inner!(@align $final_align, $($attr_rest)+);)?
59    },
60
61    // process a single `cfg_attr` attribute
62    // by translating it into a `cfg`ed block and recursing.
63    // https://doc.rust-lang.org/reference/conditional-compilation.html#railroad-ConfigurationPredicate
64
65    (@align $final_align:ident, cfg_attr($cfg_pred:expr, $($cfg_rhs:tt)*) $(, $($attr_rest:tt)+)?) => {
66        #[cfg($cfg_pred)]
67        {
68            $crate::thread::local_impl::thread_local_inner!(@align $final_align, $($cfg_rhs)*);
69        }
70
71        $($crate::thread::local_impl::thread_local_inner!(@align $final_align, $($attr_rest)+);)?
72    },
73}
74
75/// Use a regular global static to store this key; the state provided will then be
76/// thread-local.
77/// INVARIANT: ALIGN must be a valid alignment, and no less than `value_align::<T>`.
78#[allow(missing_debug_implementations)]
79pub struct Storage<T, const ALIGN: usize> {
80    key: LazyKey,
81    marker: PhantomData<Cell<T>>,
82}
83
84unsafe impl<T, const ALIGN: usize> Sync for Storage<T, ALIGN> {}
85
86#[repr(C)]
87struct Value<T: 'static> {
88    // This field must be first, for correctness of `#[rustc_align_static]`
89    value: T,
90    // INVARIANT: if this value is stored under a TLS key, `key` must be that `key`.
91    key: Key,
92}
93
94pub const fn value_align<T: 'static>() -> usize {
95    crate::mem::align_of::<Value<T>>()
96}
97
98/// Equivalent to `Box<Value<T>, System>`, but potentially over-aligned.
99struct AlignedSystemBox<T: 'static, const ALIGN: usize> {
100    ptr: NonNull<Value<T>>,
101}
102
103impl<T: 'static, const ALIGN: usize> AlignedSystemBox<T, ALIGN> {
104    #[inline]
105    fn new(v: Value<T>) -> Self {
106        let layout = Layout::new::<Value<T>>().align_to(ALIGN).unwrap();
107
108        // We use the System allocator here to avoid interfering with a potential
109        // Global allocator using thread-local storage.
110        let ptr: *mut Value<T> = (unsafe { System.alloc(layout) }).cast();
111        let Some(ptr) = NonNull::new(ptr) else {
112            alloc::handle_alloc_error(layout);
113        };
114        unsafe { ptr.write(v) };
115        Self { ptr }
116    }
117
118    #[inline]
119    fn into_raw(b: Self) -> *mut Value<T> {
120        let md = ManuallyDrop::new(b);
121        md.ptr.as_ptr()
122    }
123
124    #[inline]
125    unsafe fn from_raw(ptr: *mut Value<T>) -> Self {
126        Self { ptr: unsafe { NonNull::new_unchecked(ptr) } }
127    }
128}
129
130impl<T: 'static, const ALIGN: usize> Deref for AlignedSystemBox<T, ALIGN> {
131    type Target = Value<T>;
132
133    #[inline]
134    fn deref(&self) -> &Self::Target {
135        unsafe { &*(self.ptr.as_ptr()) }
136    }
137}
138
139impl<T: 'static, const ALIGN: usize> Drop for AlignedSystemBox<T, ALIGN> {
140    #[inline]
141    fn drop(&mut self) {
142        let layout = Layout::new::<Value<T>>().align_to(ALIGN).unwrap();
143
144        unsafe {
145            let unwind_result = catch_unwind(AssertUnwindSafe(|| self.ptr.drop_in_place()));
146            System.dealloc(self.ptr.as_ptr().cast(), layout);
147            if let Err(payload) = unwind_result {
148                resume_unwind(payload);
149            }
150        }
151    }
152}
153
154impl<T: 'static, const ALIGN: usize> Storage<T, ALIGN> {
155    pub const fn new() -> Storage<T, ALIGN> {
156        Storage { key: LazyKey::new(Some(destroy_value::<T, ALIGN>)), marker: PhantomData }
157    }
158
159    /// Gets a pointer to the TLS value, potentially initializing it with the
160    /// provided parameters. If the TLS variable has been destroyed, a null
161    /// pointer is returned.
162    ///
163    /// The resulting pointer may not be used after reentrant inialialization
164    /// or thread destruction has occurred.
165    #[inline]
166    pub fn get(&'static self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
167        let key = self.key.force();
168        let ptr = unsafe { get(key) as *mut Value<T> };
169        if ptr.addr() > 1 {
170            // SAFETY: the check ensured the pointer is safe (its destructor
171            // is not running) + it is coming from a trusted source (self).
172            unsafe { &(*ptr).value }
173        } else {
174            // SAFETY: trivially correct.
175            unsafe { Self::try_initialize(key, ptr, i, f) }
176        }
177    }
178
179    /// # Safety
180    /// * `key` must be the result of calling `self.key.force()`
181    /// * `ptr` must be the current value associated with `key`.
182    #[cold]
183    unsafe fn try_initialize(
184        key: Key,
185        ptr: *mut Value<T>,
186        i: Option<&mut Option<T>>,
187        f: impl FnOnce() -> T,
188    ) -> *const T {
189        if ptr.addr() == 1 {
190            // destructor is running
191            return ptr::null();
192        }
193
194        let value = AlignedSystemBox::<T, ALIGN>::new(Value {
195            value: i.and_then(Option::take).unwrap_or_else(f),
196            key,
197        });
198        let ptr = AlignedSystemBox::into_raw(value);
199
200        // SAFETY:
201        // * key came from a `LazyKey` and is thus correct.
202        // * `ptr` is a correct pointer that can be destroyed by the key destructor.
203        // * the value is stored under the key that it contains.
204        let old = unsafe {
205            let old = get(key) as *mut Value<T>;
206            set(key, ptr as *mut u8);
207            old
208        };
209
210        if !old.is_null() {
211            // If the variable was recursively initialized, drop the old value.
212            // SAFETY: We cannot be inside a `LocalKey::with` scope, as the
213            // initializer has already returned and the next scope only starts
214            // after we return the pointer. Therefore, there can be no references
215            // to the old value.
216            drop(unsafe { AlignedSystemBox::<T, ALIGN>::from_raw(old) });
217        }
218
219        // SAFETY: We just created this value above.
220        unsafe { &(*ptr).value }
221    }
222}
223
224unsafe extern "C" fn destroy_value<T: 'static, const ALIGN: usize>(ptr: *mut u8) {
225    // SAFETY:
226    //
227    // The OS TLS ensures that this key contains a null value when this
228    // destructor starts to run. We set it back to a sentinel value of 1 to
229    // ensure that any future calls to `get` for this thread will return
230    // `None`.
231    //
232    // Note that to prevent an infinite loop we reset it back to null right
233    // before we return from the destructor ourselves.
234    abort_on_dtor_unwind(|| {
235        let ptr = unsafe { AlignedSystemBox::<T, ALIGN>::from_raw(ptr as *mut Value<T>) };
236        let key = ptr.key;
237        // SAFETY: `key` is the TLS key `ptr` was stored under.
238        unsafe { set(key, ptr::without_provenance_mut(1)) };
239        drop(ptr);
240        // SAFETY: `key` is the TLS key `ptr` was stored under.
241        unsafe { set(key, ptr::null_mut()) };
242        // Make sure that the runtime cleanup will be performed
243        // after the next round of TLS destruction.
244        guard::enable();
245    });
246}
247
248#[rustc_macro_transparency = "semiopaque"]
249pub(crate) macro local_pointer {
250    () => {},
251    ($vis:vis static $name:ident; $($rest:tt)*) => {
252        $vis static $name: $crate::sys::thread_local::LocalPointer = $crate::sys::thread_local::LocalPointer::__new();
253        $crate::sys::thread_local::local_pointer! { $($rest)* }
254    },
255}
256
257pub(crate) struct LocalPointer {
258    key: LazyKey,
259}
260
261impl LocalPointer {
262    pub const fn __new() -> LocalPointer {
263        LocalPointer { key: LazyKey::new(None) }
264    }
265
266    pub fn get(&'static self) -> *mut () {
267        unsafe { get(self.key.force()) as *mut () }
268    }
269
270    pub fn set(&'static self, p: *mut ()) {
271        unsafe { set(self.key.force(), p as *mut u8) }
272    }
273}