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    pub fn get(&'static self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
166        let key = self.key.force();
167        let ptr = unsafe { get(key) as *mut Value<T> };
168        if ptr.addr() > 1 {
169            // SAFETY: the check ensured the pointer is safe (its destructor
170            // is not running) + it is coming from a trusted source (self).
171            unsafe { &(*ptr).value }
172        } else {
173            // SAFETY: trivially correct.
174            unsafe { Self::try_initialize(key, ptr, i, f) }
175        }
176    }
177
178    /// # Safety
179    /// * `key` must be the result of calling `self.key.force()`
180    /// * `ptr` must be the current value associated with `key`.
181    unsafe fn try_initialize(
182        key: Key,
183        ptr: *mut Value<T>,
184        i: Option<&mut Option<T>>,
185        f: impl FnOnce() -> T,
186    ) -> *const T {
187        if ptr.addr() == 1 {
188            // destructor is running
189            return ptr::null();
190        }
191
192        let value = AlignedSystemBox::<T, ALIGN>::new(Value {
193            value: i.and_then(Option::take).unwrap_or_else(f),
194            key,
195        });
196        let ptr = AlignedSystemBox::into_raw(value);
197
198        // SAFETY:
199        // * key came from a `LazyKey` and is thus correct.
200        // * `ptr` is a correct pointer that can be destroyed by the key destructor.
201        // * the value is stored under the key that it contains.
202        let old = unsafe {
203            let old = get(key) as *mut Value<T>;
204            set(key, ptr as *mut u8);
205            old
206        };
207
208        if !old.is_null() {
209            // If the variable was recursively initialized, drop the old value.
210            // SAFETY: We cannot be inside a `LocalKey::with` scope, as the
211            // initializer has already returned and the next scope only starts
212            // after we return the pointer. Therefore, there can be no references
213            // to the old value.
214            drop(unsafe { AlignedSystemBox::<T, ALIGN>::from_raw(old) });
215        }
216
217        // SAFETY: We just created this value above.
218        unsafe { &(*ptr).value }
219    }
220}
221
222unsafe extern "C" fn destroy_value<T: 'static, const ALIGN: usize>(ptr: *mut u8) {
223    // SAFETY:
224    //
225    // The OS TLS ensures that this key contains a null value when this
226    // destructor starts to run. We set it back to a sentinel value of 1 to
227    // ensure that any future calls to `get` for this thread will return
228    // `None`.
229    //
230    // Note that to prevent an infinite loop we reset it back to null right
231    // before we return from the destructor ourselves.
232    abort_on_dtor_unwind(|| {
233        let ptr = unsafe { AlignedSystemBox::<T, ALIGN>::from_raw(ptr as *mut Value<T>) };
234        let key = ptr.key;
235        // SAFETY: `key` is the TLS key `ptr` was stored under.
236        unsafe { set(key, ptr::without_provenance_mut(1)) };
237        drop(ptr);
238        // SAFETY: `key` is the TLS key `ptr` was stored under.
239        unsafe { set(key, ptr::null_mut()) };
240        // Make sure that the runtime cleanup will be performed
241        // after the next round of TLS destruction.
242        guard::enable();
243    });
244}
245
246#[rustc_macro_transparency = "semiopaque"]
247pub(crate) macro local_pointer {
248    () => {},
249    ($vis:vis static $name:ident; $($rest:tt)*) => {
250        $vis static $name: $crate::sys::thread_local::LocalPointer = $crate::sys::thread_local::LocalPointer::__new();
251        $crate::sys::thread_local::local_pointer! { $($rest)* }
252    },
253}
254
255pub(crate) struct LocalPointer {
256    key: LazyKey,
257}
258
259impl LocalPointer {
260    pub const fn __new() -> LocalPointer {
261        LocalPointer { key: LazyKey::new(None) }
262    }
263
264    pub fn get(&'static self) -> *mut () {
265        unsafe { get(self.key.force()) as *mut () }
266    }
267
268    pub fn set(&'static self, p: *mut ()) {
269        unsafe { set(self.key.force(), p as *mut u8) }
270    }
271}