1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
use super::lazy::LazyKeyInner;
use crate::cell::Cell;
use crate::sys::thread_local_dtor::register_dtor;
use crate::{fmt, mem, panic};

#[doc(hidden)]
#[allow_internal_unstable(thread_local_internals, cfg_target_thread_local, thread_local)]
#[allow_internal_unsafe]
#[unstable(feature = "thread_local_internals", issue = "none")]
#[rustc_macro_transparency = "semitransparent"]
pub macro thread_local_inner {
    // used to generate the `LocalKey` value for const-initialized thread locals
    (@key $t:ty, const $init:expr) => {{
        #[inline]
        #[deny(unsafe_op_in_unsafe_fn)]
        unsafe fn __getit(
            _init: $crate::option::Option<&mut $crate::option::Option<$t>>,
        ) -> $crate::option::Option<&'static $t> {
            const INIT_EXPR: $t = $init;
            // If the platform has support for `#[thread_local]`, use it.
            #[thread_local]
            static mut VAL: $t = INIT_EXPR;

            // If a dtor isn't needed we can do something "very raw" and
            // just get going.
            if !$crate::mem::needs_drop::<$t>() {
                unsafe {
                    return $crate::option::Option::Some(&VAL)
                }
            }

            // 0 == dtor not registered
            // 1 == dtor registered, dtor not run
            // 2 == dtor registered and is running or has run
            #[thread_local]
            static STATE: $crate::cell::Cell<$crate::primitive::u8> = $crate::cell::Cell::new(0);

            // Safety: Performs `drop_in_place(ptr as *mut $t)`, and requires
            // all that comes with it.
            unsafe extern "C" fn destroy(ptr: *mut $crate::primitive::u8) {
                $crate::thread::local_impl::abort_on_dtor_unwind(|| {
                    let old_state = STATE.replace(2);
                    $crate::debug_assert_eq!(old_state, 1);
                    // Safety: safety requirement is passed on to caller.
                    unsafe { $crate::ptr::drop_in_place(ptr.cast::<$t>()); }
                });
            }

            unsafe {
                match STATE.get() {
                    // 0 == we haven't registered a destructor, so do
                    //   so now.
                    0 => {
                        $crate::thread::local_impl::Key::<$t>::register_dtor(
                            $crate::ptr::addr_of_mut!(VAL) as *mut $crate::primitive::u8,
                            destroy,
                        );
                        STATE.set(1);
                        $crate::option::Option::Some(&VAL)
                    }
                    // 1 == the destructor is registered and the value
                    //   is valid, so return the pointer.
                    1 => $crate::option::Option::Some(&VAL),
                    // otherwise the destructor has already run, so we
                    // can't give access.
                    _ => $crate::option::Option::None,
                }
            }
        }

        unsafe {
            $crate::thread::LocalKey::new(__getit)
        }
    }},

    // used to generate the `LocalKey` value for `thread_local!`
    (@key $t:ty, $init:expr) => {
        {
            #[inline]
            fn __init() -> $t { $init }

            #[inline]
            unsafe fn __getit(
                init: $crate::option::Option<&mut $crate::option::Option<$t>>,
            ) -> $crate::option::Option<&'static $t> {
                #[thread_local]
                static __KEY: $crate::thread::local_impl::Key<$t> =
                    $crate::thread::local_impl::Key::<$t>::new();

                // FIXME: remove the #[allow(...)] marker when macros don't
                // raise warning for missing/extraneous unsafe blocks anymore.
                // See https://github.com/rust-lang/rust/issues/74838.
                #[allow(unused_unsafe)]
                unsafe {
                    __KEY.get(move || {
                        if let $crate::option::Option::Some(init) = init {
                            if let $crate::option::Option::Some(value) = init.take() {
                                return value;
                            } else if $crate::cfg!(debug_assertions) {
                                $crate::unreachable!("missing default value");
                            }
                        }
                        __init()
                    })
                }
            }

            unsafe {
                $crate::thread::LocalKey::new(__getit)
            }
        }
    },
    ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {
        $(#[$attr])* $vis const $name: $crate::thread::LocalKey<$t> =
            $crate::thread::local_impl::thread_local_inner!(@key $t, $($init)*);
    },
}

#[derive(Copy, Clone)]
enum DtorState {
    Unregistered,
    Registered,
    RunningOrHasRun,
}

// This data structure has been carefully constructed so that the fast path
// only contains one branch on x86. That optimization is necessary to avoid
// duplicated tls lookups on OSX.
//
// LLVM issue: https://bugs.llvm.org/show_bug.cgi?id=41722
pub struct Key<T> {
    // If `LazyKeyInner::get` returns `None`, that indicates either:
    //   * The value has never been initialized
    //   * The value is being recursively initialized
    //   * The value has already been destroyed or is being destroyed
    // To determine which kind of `None`, check `dtor_state`.
    //
    // This is very optimizer friendly for the fast path - initialized but
    // not yet dropped.
    inner: LazyKeyInner<T>,

    // Metadata to keep track of the state of the destructor. Remember that
    // this variable is thread-local, not global.
    dtor_state: Cell<DtorState>,
}

impl<T> fmt::Debug for Key<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Key").finish_non_exhaustive()
    }
}
impl<T> Key<T> {
    pub const fn new() -> Key<T> {
        Key { inner: LazyKeyInner::new(), dtor_state: Cell::new(DtorState::Unregistered) }
    }

    // note that this is just a publicly-callable function only for the
    // const-initialized form of thread locals, basically a way to call the
    // free `register_dtor` function defined elsewhere in std.
    pub unsafe fn register_dtor(a: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
        unsafe {
            register_dtor(a, dtor);
        }
    }

    pub unsafe fn get<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
        // SAFETY: See the definitions of `LazyKeyInner::get` and
        // `try_initialize` for more information.
        //
        // The caller must ensure no mutable references are ever active to
        // the inner cell or the inner T when this is called.
        // The `try_initialize` is dependant on the passed `init` function
        // for this.
        unsafe {
            match self.inner.get() {
                Some(val) => Some(val),
                None => self.try_initialize(init),
            }
        }
    }

    // `try_initialize` is only called once per fast thread local variable,
    // except in corner cases where thread_local dtors reference other
    // thread_local's, or it is being recursively initialized.
    //
    // Macos: Inlining this function can cause two `tlv_get_addr` calls to
    // be performed for every call to `Key::get`.
    // LLVM issue: https://bugs.llvm.org/show_bug.cgi?id=41722
    #[inline(never)]
    unsafe fn try_initialize<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
        // SAFETY: See comment above (this function doc).
        if !mem::needs_drop::<T>() || unsafe { self.try_register_dtor() } {
            // SAFETY: See comment above (this function doc).
            Some(unsafe { self.inner.initialize(init) })
        } else {
            None
        }
    }

    // `try_register_dtor` is only called once per fast thread local
    // variable, except in corner cases where thread_local dtors reference
    // other thread_local's, or it is being recursively initialized.
    unsafe fn try_register_dtor(&self) -> bool {
        match self.dtor_state.get() {
            DtorState::Unregistered => {
                // SAFETY: dtor registration happens before initialization.
                // Passing `self` as a pointer while using `destroy_value<T>`
                // is safe because the function will build a pointer to a
                // Key<T>, which is the type of self and so find the correct
                // size.
                unsafe { register_dtor(self as *const _ as *mut u8, destroy_value::<T>) };
                self.dtor_state.set(DtorState::Registered);
                true
            }
            DtorState::Registered => {
                // recursively initialized
                true
            }
            DtorState::RunningOrHasRun => false,
        }
    }
}

unsafe extern "C" fn destroy_value<T>(ptr: *mut u8) {
    let ptr = ptr as *mut Key<T>;

    // SAFETY:
    //
    // The pointer `ptr` has been built just above and comes from
    // `try_register_dtor` where it is originally a Key<T> coming from `self`,
    // making it non-NUL and of the correct type.
    //
    // Right before we run the user destructor be sure to set the
    // `Option<T>` to `None`, and `dtor_state` to `RunningOrHasRun`. This
    // causes future calls to `get` to run `try_initialize_drop` again,
    // which will now fail, and return `None`.
    //
    // Wrap the call in a catch to ensure unwinding is caught in the event
    // a panic takes place in a destructor.
    if let Err(_) = panic::catch_unwind(panic::AssertUnwindSafe(|| unsafe {
        let value = (*ptr).inner.take();
        (*ptr).dtor_state.set(DtorState::RunningOrHasRun);
        drop(value);
    })) {
        rtabort!("thread local panicked on drop");
    }
}