xref: /drstd/src/std/sys/common/thread_local/fast_local.rs (revision 0fe3ff0054d3aec7fbf9bddecfecb10bc7d23a51)
1 use super::lazy::LazyKeyInner;
2 use crate::std::cell::Cell;
3 use crate::std::sys::thread_local_dtor::register_dtor;
4 use crate::std::{fmt, mem, panic};
5 
6 #[doc(hidden)]
7 #[allow_internal_unstable(thread_local_internals, cfg_target_thread_local, thread_local)]
8 #[allow_internal_unsafe]
9 #[rustc_macro_transparency = "semitransparent"]
10 pub macro thread_local_inner {
11     // used to generate the `LocalKey` value for const-initialized thread locals
12     (@key $t:ty, const $init:expr) => {{
13         #[inline]
14         #[deny(unsafe_op_in_unsafe_fn)]
15         unsafe fn __getit(
16             _init: $crate::std::option::Option<&mut $crate::std::option::Option<$t>>,
17         ) -> $crate::std::option::Option<&'static $t> {
18             const INIT_EXPR: $t = $init;
19             // If the platform has support for `#[thread_local]`, use it.
20             #[thread_local]
21             static mut VAL: $t = INIT_EXPR;
22 
23             // If a dtor isn't needed we can do something "very raw" and
24             // just get going.
25             if !$crate::std::mem::needs_drop::<$t>() {
26                 unsafe {
27                     return $crate::std::option::Option::Some(&VAL)
28                 }
29             }
30 
31             // 0 == dtor not registered
32             // 1 == dtor registered, dtor not run
33             // 2 == dtor registered and is running or has run
34             #[thread_local]
35             static STATE: $crate::std::cell::Cell<$crate::std::primitive::u8> = $crate::std::cell::Cell::new(0);
36 
37             // Safety: Performs `drop_in_place(ptr as *mut $t)`, and requires
38             // all that comes with it.
39             unsafe extern "C" fn destroy(ptr: *mut $crate::std::primitive::u8) {
40                 $crate::std::thread::local_impl::abort_on_dtor_unwind(|| {
41                     let old_state = STATE.replace(2);
42                     $crate::std::debug_assert_eq!(old_state, 1);
43                     // Safety: safety requirement is passed on to caller.
44                     unsafe { $crate::std::ptr::drop_in_place(ptr.cast::<$t>()); }
45                 });
46             }
47 
48             unsafe {
49                 match STATE.get() {
50                     // 0 == we haven't registered a destructor, so do
51                     //   so now.
52                     0 => {
53                         $crate::std::thread::local_impl::Key::<$t>::register_dtor(
54                             $crate::std::ptr::addr_of_mut!(VAL) as *mut $crate::std::primitive::u8,
55                             destroy,
56                         );
57                         STATE.set(1);
58                         $crate::std::option::Option::Some(&VAL)
59                     }
60                     // 1 == the destructor is registered and the value
61                     //   is valid, so return the pointer.
62                     1 => $crate::std::option::Option::Some(&VAL),
63                     // otherwise the destructor has already run, so we
64                     // can't give access.
65                     _ => $crate::std::option::Option::None,
66                 }
67             }
68         }
69 
70         unsafe {
71             $crate::std::thread::LocalKey::new(__getit)
72         }
73     }},
74 
75     // used to generate the `LocalKey` value for `thread_local!`
76     (@key $t:ty, $init:expr) => {
77         {
78             #[inline]
79             fn __init() -> $t { $init }
80 
81             #[inline]
82             unsafe fn __getit(
83                 init: $crate::std::option::Option<&mut $crate::std::option::Option<$t>>,
84             ) -> $crate::std::option::Option<&'static $t> {
85                 #[thread_local]
86                 static __KEY: $crate::std::thread::local_impl::Key<$t> =
87                     $crate::std::thread::local_impl::Key::<$t>::new();
88 
89                 unsafe {
90                     __KEY.get(move || {
91                         if let $crate::std::option::Option::Some(init) = init {
92                             if let $crate::std::option::Option::Some(value) = init.take() {
93                                 return value;
94                             } else if $crate::std::cfg!(debug_assertions) {
95                                 $crate::std::unreachable!("missing default value");
96                             }
97                         }
98                         __init()
99                     })
100                 }
101             }
102 
103             unsafe {
104                 $crate::std::thread::LocalKey::new(__getit)
105             }
106         }
107     },
108     ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {
109         $(#[$attr])* $vis const $name: $crate::std::thread::LocalKey<$t> =
110             $crate::std::thread::local_impl::thread_local_inner!(@key $t, $($init)*);
111     },
112 }
113 
114 #[derive(Copy, Clone)]
115 enum DtorState {
116     Unregistered,
117     Registered,
118     RunningOrHasRun,
119 }
120 
121 // This data structure has been carefully constructed so that the fast path
122 // only contains one branch on x86. That optimization is necessary to avoid
123 // duplicated tls lookups on OSX.
124 //
125 // LLVM issue: https://bugs.llvm.org/show_bug.cgi?id=41722
126 pub struct Key<T> {
127     // If `LazyKeyInner::get` returns `None`, that indicates either:
128     //   * The value has never been initialized
129     //   * The value is being recursively initialized
130     //   * The value has already been destroyed or is being destroyed
131     // To determine which kind of `None`, check `dtor_state`.
132     //
133     // This is very optimizer friendly for the fast path - initialized but
134     // not yet dropped.
135     inner: LazyKeyInner<T>,
136 
137     // Metadata to keep track of the state of the destructor. Remember that
138     // this variable is thread-local, not global.
139     dtor_state: Cell<DtorState>,
140 }
141 
142 impl<T> fmt::Debug for Key<T> {
143     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
144         f.debug_struct("Key").finish_non_exhaustive()
145     }
146 }
147 impl<T> Key<T> {
148     pub const fn new() -> Key<T> {
149         Key {
150             inner: LazyKeyInner::new(),
151             dtor_state: Cell::new(DtorState::Unregistered),
152         }
153     }
154 
155     // note that this is just a publicly-callable function only for the
156     // const-initialized form of thread locals, basically a way to call the
157     // free `register_dtor` function defined elsewhere in std.
158     pub unsafe fn register_dtor(a: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
159         unsafe {
160             register_dtor(a, dtor);
161         }
162     }
163 
164     pub unsafe fn get<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
165         // SAFETY: See the definitions of `LazyKeyInner::get` and
166         // `try_initialize` for more information.
167         //
168         // The caller must ensure no mutable references are ever active to
169         // the inner cell or the inner T when this is called.
170         // The `try_initialize` is dependant on the passed `init` function
171         // for this.
172         unsafe {
173             match self.inner.get() {
174                 Some(val) => Some(val),
175                 None => self.try_initialize(init),
176             }
177         }
178     }
179 
180     // `try_initialize` is only called once per fast thread local variable,
181     // except in corner cases where thread_local dtors reference other
182     // thread_local's, or it is being recursively initialized.
183     //
184     // Macos: Inlining this function can cause two `tlv_get_addr` calls to
185     // be performed for every call to `Key::get`.
186     // LLVM issue: https://bugs.llvm.org/show_bug.cgi?id=41722
187     #[inline(never)]
188     unsafe fn try_initialize<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
189         // SAFETY: See comment above (this function doc).
190         if !mem::needs_drop::<T>() || unsafe { self.try_register_dtor() } {
191             // SAFETY: See comment above (this function doc).
192             Some(unsafe { self.inner.initialize(init) })
193         } else {
194             None
195         }
196     }
197 
198     // `try_register_dtor` is only called once per fast thread local
199     // variable, except in corner cases where thread_local dtors reference
200     // other thread_local's, or it is being recursively initialized.
201     unsafe fn try_register_dtor(&self) -> bool {
202         match self.dtor_state.get() {
203             DtorState::Unregistered => {
204                 // SAFETY: dtor registration happens before initialization.
205                 // Passing `self` as a pointer while using `destroy_value<T>`
206                 // is safe because the function will build a pointer to a
207                 // Key<T>, which is the type of self and so find the correct
208                 // size.
209                 unsafe { register_dtor(self as *const _ as *mut u8, destroy_value::<T>) };
210                 self.dtor_state.set(DtorState::Registered);
211                 true
212             }
213             DtorState::Registered => {
214                 // recursively initialized
215                 true
216             }
217             DtorState::RunningOrHasRun => false,
218         }
219     }
220 }
221 
222 unsafe extern "C" fn destroy_value<T>(ptr: *mut u8) {
223     let ptr = ptr as *mut Key<T>;
224 
225     // SAFETY:
226     //
227     // The pointer `ptr` has been built just above and comes from
228     // `try_register_dtor` where it is originally a Key<T> coming from `self`,
229     // making it non-NUL and of the correct type.
230     //
231     // Right before we run the user destructor be sure to set the
232     // `Option<T>` to `None`, and `dtor_state` to `RunningOrHasRun`. This
233     // causes future calls to `get` to run `try_initialize_drop` again,
234     // which will now fail, and return `None`.
235     //
236     // Wrap the call in a catch to ensure unwinding is caught in the event
237     // a panic takes place in a destructor.
238     if let Err(_) = panic::catch_unwind(panic::AssertUnwindSafe(|| unsafe {
239         let value = (*ptr).inner.take();
240         (*ptr).dtor_state.set(DtorState::RunningOrHasRun);
241         drop(value);
242     })) {
243         rtabort!("thread local panicked on drop");
244     }
245 }
246