Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 7ffccf1

Browse files
joboetjoboet
committedMar 12, 2024
std: rewrite native thread-local storage
1 parent 5b7343b commit 7ffccf1

File tree

3 files changed

+249
-190
lines changed

3 files changed

+249
-190
lines changed
 

‎library/std/src/sys/thread_local/fast_local.rs

Lines changed: 244 additions & 188 deletions
Original file line numberDiff line numberDiff line change
@@ -1,72 +1,93 @@
1-
use super::lazy::LazyKeyInner;
2-
use crate::cell::Cell;
1+
//! Thread local support for platforms with native TLS.
2+
//!
3+
//! To achieve the best performance, we need to differentiate between four types
4+
//! of TLS, resulting from the method of initialization used (`const` or lazy)
5+
//! and the drop requirements of the stored type, as indicated by
6+
//! [`needs_drop`](crate::mem::needs_drop). However, across these types a TLS
7+
//! variable can only be in three different different states, allowing us to
8+
//! use a common implementation specialized by type parameters. The three states
9+
//! are:
10+
//!
11+
//! 1. [`Initial`](State::Initial): the destructor has not been registered and/or
12+
//! the variable is uninitialized.
13+
//! 2. [`Alive`](State::Alive): if it exists, the destructor is registered. The
14+
//! variable is initialized.
15+
//! 3. [`Destroyed`](State::Destroyed): the TLS variable has been destroyed.
16+
//!
17+
//! Upon accessing the TLS variable through [`get_or_init`](Storage::get_or_init),
18+
//! the current state is compared:
19+
//!
20+
//! 1. If the state is `Initial`, a macro-provided closure is run with the
21+
//! parameter stored with the state, which returns the initialization value
22+
//! for the variable. The state is transitioned to `Alive` and the destructor
23+
//! is registered if the variable should be destroyed.
24+
//! 2. If the state is `Alive`, initialization was previously completed, so a
25+
//! reference to the stored value is returned.
26+
//! 3. If the state is `Destroyed`, the TLS variable was already destroyed, so
27+
//! return [`None`].
28+
//!
29+
//! The TLS destructor sets the state to `Destroyed` and drops the current value.
30+
//!
31+
//! `const`-initialized variables use the initialization value as parameter to
32+
//! the initialization closure, which is just the [`identity`](crate::convert::identity)
33+
//! function. Storing this value in the variable from the beginning means that
34+
//! no copying needs to take place (if the optimizer does a good job).
35+
//!
36+
//! Lazily initialized variables simply use `()` as parameter and perform the
37+
//! actual initialization inside of the closure.
38+
//!
39+
//! By using the `!` type (never type) as type parameter for the destroyed state,
40+
//! we can eliminate the `Destroyed` state for values that do not need a
41+
//! destructor using a generic parameter, which allows niche optimizations to
42+
//! occur for the `State` enum. Otherwise, `()` is used.
43+
44+
#![deny(unsafe_op_in_unsafe_fn)]
45+
46+
use super::abort_on_dtor_unwind;
47+
use crate::cell::UnsafeCell;
48+
use crate::hint::unreachable_unchecked;
49+
use crate::mem::forget;
50+
use crate::ptr;
351
use crate::sys::thread_local_dtor::register_dtor;
4-
use crate::{fmt, mem, panic};
552

653
#[doc(hidden)]
7-
#[allow_internal_unstable(thread_local_internals, cfg_target_thread_local, thread_local)]
54+
#[allow_internal_unstable(
55+
thread_local_internals,
56+
cfg_target_thread_local,
57+
thread_local,
58+
never_type
59+
)]
860
#[allow_internal_unsafe]
961
#[unstable(feature = "thread_local_internals", issue = "none")]
1062
#[rustc_macro_transparency = "semitransparent"]
1163
pub macro thread_local_inner {
1264
// used to generate the `LocalKey` value for const-initialized thread locals
1365
(@key $t:ty, const $init:expr) => {{
66+
const __INIT: $t = $init;
67+
1468
#[inline]
1569
#[deny(unsafe_op_in_unsafe_fn)]
16-
// FIXME: Use `SyncUnsafeCell` instead of allowing `static_mut_refs` lint
17-
#[cfg_attr(bootstrap, allow(static_mut_ref))]
18-
#[cfg_attr(not(bootstrap), allow(static_mut_refs))]
1970
unsafe fn __getit(
2071
_init: $crate::option::Option<&mut $crate::option::Option<$t>>,
2172
) -> $crate::option::Option<&'static $t> {
22-
const INIT_EXPR: $t = $init;
23-
// If the platform has support for `#[thread_local]`, use it.
24-
#[thread_local]
25-
static mut VAL: $t = INIT_EXPR;
26-
27-
// If a dtor isn't needed we can do something "very raw" and
28-
// just get going.
29-
if !$crate::mem::needs_drop::<$t>() {
73+
use $crate::thread::local_impl::Storage;
74+
use $crate::mem::needs_drop;
75+
use $crate::convert::identity;
76+
use $crate::ptr::addr_of;
77+
78+
if needs_drop::<$t>() {
79+
#[thread_local]
80+
static VAL: Storage<$t, $t, ()> = Storage::new(__INIT);
3081
unsafe {
31-
return $crate::option::Option::Some(&VAL)
82+
VAL.get_or_init(identity)
3283
}
33-
}
34-
35-
// 0 == dtor not registered
36-
// 1 == dtor registered, dtor not run
37-
// 2 == dtor registered and is running or has run
38-
#[thread_local]
39-
static STATE: $crate::cell::Cell<$crate::primitive::u8> = $crate::cell::Cell::new(0);
40-
41-
// Safety: Performs `drop_in_place(ptr as *mut $t)`, and requires
42-
// all that comes with it.
43-
unsafe extern "C" fn destroy(ptr: *mut $crate::primitive::u8) {
44-
$crate::thread::local_impl::abort_on_dtor_unwind(|| {
45-
let old_state = STATE.replace(2);
46-
$crate::debug_assert_eq!(old_state, 1);
47-
// Safety: safety requirement is passed on to caller.
48-
unsafe { $crate::ptr::drop_in_place(ptr.cast::<$t>()); }
49-
});
50-
}
51-
52-
unsafe {
53-
match STATE.get() {
54-
// 0 == we haven't registered a destructor, so do
55-
// so now.
56-
0 => {
57-
$crate::thread::local_impl::Key::<$t>::register_dtor(
58-
$crate::ptr::addr_of_mut!(VAL) as *mut $crate::primitive::u8,
59-
destroy,
60-
);
61-
STATE.set(1);
62-
$crate::option::Option::Some(&VAL)
63-
}
64-
// 1 == the destructor is registered and the value
65-
// is valid, so return the pointer.
66-
1 => $crate::option::Option::Some(&VAL),
67-
// otherwise the destructor has already run, so we
68-
// can't give access.
69-
_ => $crate::option::Option::None,
84+
} else {
85+
// Just use thread-local statics directly instead of going
86+
// through `Storage`.
87+
#[thread_local]
88+
static VAL: $t = __INIT;
89+
unsafe {
90+
$crate::option::Option::Some(&*addr_of!(VAL))
7091
}
7192
}
7293
}
@@ -77,171 +98,206 @@ pub macro thread_local_inner {
7798
}},
7899

79100
// used to generate the `LocalKey` value for `thread_local!`
80-
(@key $t:ty, $init:expr) => {
81-
{
82-
#[inline]
83-
fn __init() -> $t { $init }
84-
85-
#[inline]
86-
unsafe fn __getit(
87-
init: $crate::option::Option<&mut $crate::option::Option<$t>>,
88-
) -> $crate::option::Option<&'static $t> {
89-
#[thread_local]
90-
static __KEY: $crate::thread::local_impl::Key<$t> =
91-
$crate::thread::local_impl::Key::<$t>::new();
101+
(@key $t:ty, $init:expr) => {{
102+
#[inline]
103+
fn __init() -> $t {
104+
$init
105+
}
92106

107+
#[inline]
108+
#[deny(unsafe_op_in_unsafe_fn)]
109+
unsafe fn __getit(
110+
init: $crate::option::Option<&mut $crate::option::Option<$t>>,
111+
) -> $crate::option::Option<&'static $t> {
112+
use $crate::thread::local_impl::{Storage, take_or_call};
113+
use $crate::mem::needs_drop;
114+
115+
if needs_drop::<$t>() {
116+
#[thread_local]
117+
static VAL: Storage<(), $t, ()> = Storage::new(());
118+
unsafe {
119+
VAL.get_or_init(|()| take_or_call(init, __init))
120+
}
121+
} else {
122+
#[thread_local]
123+
static VAL: Storage<(), $t, !> = Storage::new(());
93124
unsafe {
94-
__KEY.get(move || {
95-
if let $crate::option::Option::Some(init) = init {
96-
if let $crate::option::Option::Some(value) = init.take() {
97-
return value;
98-
}
99-
if $crate::cfg!(debug_assertions) {
100-
$crate::unreachable!("missing default value");
101-
}
102-
}
103-
__init()
104-
})
125+
VAL.get_or_init(|()| take_or_call(init, __init))
105126
}
106127
}
128+
}
107129

108-
unsafe {
109-
$crate::thread::LocalKey::new(__getit)
110-
}
130+
unsafe {
131+
$crate::thread::LocalKey::new(__getit)
111132
}
112-
},
133+
}},
113134
($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {
114135
$(#[$attr])* $vis const $name: $crate::thread::LocalKey<$t> =
115136
$crate::thread::local_impl::thread_local_inner!(@key $t, $($init)*);
116137
},
117138
}
118139

119-
#[derive(Copy, Clone)]
120-
enum DtorState {
121-
Unregistered,
122-
Registered,
123-
RunningOrHasRun,
140+
#[inline]
141+
pub fn take_or_call<T, F>(i: Option<&mut Option<T>>, f: F) -> T
142+
where
143+
F: FnOnce() -> T,
144+
{
145+
i.and_then(Option::take).unwrap_or_else(f)
124146
}
125147

126-
// This data structure has been carefully constructed so that the fast path
127-
// only contains one branch on x86. That optimization is necessary to avoid
128-
// duplicated tls lookups on OSX.
129-
//
130-
// LLVM issue: https://bugs.llvm.org/show_bug.cgi?id=41722
131-
pub struct Key<T> {
132-
// If `LazyKeyInner::get` returns `None`, that indicates either:
133-
// * The value has never been initialized
134-
// * The value is being recursively initialized
135-
// * The value has already been destroyed or is being destroyed
136-
// To determine which kind of `None`, check `dtor_state`.
137-
//
138-
// This is very optimizer friendly for the fast path - initialized but
139-
// not yet dropped.
140-
inner: LazyKeyInner<T>,
141-
142-
// Metadata to keep track of the state of the destructor. Remember that
143-
// this variable is thread-local, not global.
144-
dtor_state: Cell<DtorState>,
148+
pub trait DestroyedState {
149+
const EXISTS: bool;
150+
fn create() -> Self;
145151
}
146152

147-
impl<T> fmt::Debug for Key<T> {
148-
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
149-
f.debug_struct("Key").finish_non_exhaustive()
153+
impl DestroyedState for ! {
154+
const EXISTS: bool = false;
155+
fn create() -> ! {
156+
panic!("attempted to create nonexistent state");
150157
}
151158
}
152-
impl<T> Key<T> {
153-
pub const fn new() -> Key<T> {
154-
Key { inner: LazyKeyInner::new(), dtor_state: Cell::new(DtorState::Unregistered) }
155-
}
156159

157-
// note that this is just a publicly-callable function only for the
158-
// const-initialized form of thread locals, basically a way to call the
159-
// free `register_dtor` function defined elsewhere in std.
160-
pub unsafe fn register_dtor(a: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
161-
unsafe {
162-
register_dtor(a, dtor);
163-
}
164-
}
160+
impl DestroyedState for () {
161+
const EXISTS: bool = true;
162+
#[inline]
163+
fn create() {}
164+
}
165165

166-
pub unsafe fn get<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
167-
// SAFETY: See the definitions of `LazyKeyInner::get` and
168-
// `try_initialize` for more information.
169-
//
170-
// The caller must ensure no mutable references are ever active to
171-
// the inner cell or the inner T when this is called.
172-
// The `try_initialize` is dependant on the passed `init` function
173-
// for this.
174-
unsafe {
175-
match self.inner.get() {
176-
Some(val) => Some(val),
177-
None => self.try_initialize(init),
178-
}
179-
}
166+
enum State<I, T, D> {
167+
Initial(I),
168+
Alive(T),
169+
Destroyed(D),
170+
}
171+
172+
#[allow(missing_debug_implementations)]
173+
pub struct Storage<I, T, D> {
174+
state: UnsafeCell<State<I, T, D>>,
175+
}
176+
177+
impl<I, T, D> Storage<I, T, D>
178+
where
179+
D: DestroyedState,
180+
{
181+
/// Create a new TLS storage with the provided initialization parameter.
182+
pub const fn new(param: I) -> Storage<I, T, D> {
183+
Storage { state: UnsafeCell::new(State::Initial(param)) }
180184
}
181185

182-
// `try_initialize` is only called once per fast thread local variable,
183-
// except in corner cases where thread_local dtors reference other
184-
// thread_local's, or it is being recursively initialized.
185-
//
186-
// Macos: Inlining this function can cause two `tlv_get_addr` calls to
187-
// be performed for every call to `Key::get`.
188-
// LLVM issue: https://bugs.llvm.org/show_bug.cgi?id=41722
189-
#[inline(never)]
190-
unsafe fn try_initialize<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
191-
// SAFETY: See comment above (this function doc).
192-
if !mem::needs_drop::<T>() || unsafe { self.try_register_dtor() } {
193-
// SAFETY: See comment above (this function doc).
194-
Some(unsafe { self.inner.initialize(init) })
195-
} else {
196-
None
186+
/// Get a reference to the TLS value, potentially initializing it with the
187+
/// provided closure. If the TLS variable has been destroyed already, `None`
188+
/// is returned.
189+
///
190+
/// # Safety
191+
/// * The `self` reference must remain valid until the TLS destructor is run,
192+
/// at which point the returned reference is invalidated.
193+
/// * If the closure may panic or access this variable, the initialization
194+
/// parameter must implement `Copy`.
195+
/// * The returned reference may only be used until thread destruction occurs
196+
/// and may not be used after reentrant initialization has occurred.
197+
///
198+
// FIXME(#110897): return NonNull instead of lying about the lifetime.
199+
#[inline]
200+
pub unsafe fn get_or_init<F>(&self, init: F) -> Option<&'static T>
201+
where
202+
F: FnOnce(I) -> T,
203+
{
204+
// SAFETY:
205+
// No mutable reference to the inner value exists outside the calls to
206+
// `replace`. The lifetime of the returned reference fulfills the terms
207+
// outlined above.
208+
let state = unsafe { &*self.state.get() };
209+
match state {
210+
State::Alive(v) => Some(v),
211+
State::Destroyed(_) => None,
212+
State::Initial(_) => unsafe { self.initialize(init) },
197213
}
198214
}
199215

200-
// `try_register_dtor` is only called once per fast thread local
201-
// variable, except in corner cases where thread_local dtors reference
202-
// other thread_local's, or it is being recursively initialized.
203-
unsafe fn try_register_dtor(&self) -> bool {
204-
match self.dtor_state.get() {
205-
DtorState::Unregistered => {
206-
// SAFETY: dtor registration happens before initialization.
207-
// Passing `self` as a pointer while using `destroy_value<T>`
208-
// is safe because the function will build a pointer to a
209-
// Key<T>, which is the type of self and so find the correct
210-
// size.
211-
unsafe { register_dtor(self as *const _ as *mut u8, destroy_value::<T>) };
212-
self.dtor_state.set(DtorState::Registered);
213-
true
216+
#[cold]
217+
unsafe fn initialize<F>(&self, init: F) -> Option<&'static T>
218+
where
219+
F: FnOnce(I) -> T,
220+
{
221+
// Perform initialization
222+
223+
// SAFETY:
224+
// The state is valid for reading (see above) and must be `Initial` at
225+
// this point. If the closure panics or recursively initializes the
226+
// variable, the initialization parameter is duplicated. The caller
227+
// asserts that this is valid. Otherwise the old copy is leaked below.
228+
let state = unsafe { self.state.get().read() };
229+
let State::Initial(param) = state else { unsafe { unreachable_unchecked() } };
230+
231+
let v = init(param);
232+
233+
// SAFETY:
234+
// If references to the inner value exist, they were created in `init`
235+
// and are invalidated here. The caller promises to never use them
236+
// after this.
237+
let old = unsafe { self.state.get().replace(State::Alive(v)) };
238+
239+
let recursive = match old {
240+
State::Initial(v) => {
241+
// The value was duplicated above, leak the old value to prevent
242+
// double freeing.
243+
forget(v);
244+
false
214245
}
215-
DtorState::Registered => {
216-
// recursively initialized
246+
State::Alive(v) => {
247+
// The storage was recursively initialized, so drop the previous
248+
// value. This could be changed to a panic in the future.
249+
drop(v);
217250
true
218251
}
219-
DtorState::RunningOrHasRun => false,
252+
State::Destroyed(_) => {
253+
unreachable!("thread is still alive but TLS storage was destroyed");
254+
}
255+
};
256+
257+
if D::EXISTS && !recursive {
258+
// If a `Destroyed` state exists and the variable was not initialized
259+
// before, register the destructor.
260+
unsafe {
261+
register_dtor(ptr::from_ref(self).cast_mut().cast(), destroy::<I, T, D>);
262+
}
263+
}
264+
265+
// SAFETY:
266+
// Initialization was completed and the state was set to `Alive`, so the
267+
// reference fulfills the terms outlined above.
268+
unsafe {
269+
let State::Alive(v) = &*self.state.get() else { unreachable_unchecked() };
270+
Some(v)
220271
}
221272
}
222273
}
223274

224-
unsafe extern "C" fn destroy_value<T>(ptr: *mut u8) {
225-
let ptr = ptr as *mut Key<T>;
226-
227-
// SAFETY:
228-
//
229-
// The pointer `ptr` has been built just above and comes from
230-
// `try_register_dtor` where it is originally a Key<T> coming from `self`,
231-
// making it non-NUL and of the correct type.
232-
//
233-
// Right before we run the user destructor be sure to set the
234-
// `Option<T>` to `None`, and `dtor_state` to `RunningOrHasRun`. This
235-
// causes future calls to `get` to run `try_initialize_drop` again,
236-
// which will now fail, and return `None`.
237-
//
238-
// Wrap the call in a catch to ensure unwinding is caught in the event
239-
// a panic takes place in a destructor.
240-
if let Err(_) = panic::catch_unwind(panic::AssertUnwindSafe(|| unsafe {
241-
let value = (*ptr).inner.take();
242-
(*ptr).dtor_state.set(DtorState::RunningOrHasRun);
243-
drop(value);
244-
})) {
245-
rtabort!("thread local panicked on drop");
246-
}
275+
/// Transition an `Alive` TLS variable into the `Destroyed` state, dropping its
276+
/// value. Should only be called when `D::SHOULD_DESTROY` is `true`. May abort
277+
/// otherwise.
278+
///
279+
/// # Safety
280+
/// * Must only be called at thread destruction.
281+
/// * `ptr` must point to an instance of `Storage` with matching generic parameters
282+
/// and be valid for accessing that instance.
283+
/// * Must only be called if the state is `Alive`.
284+
unsafe extern "C" fn destroy<I, T, D>(ptr: *mut u8)
285+
where
286+
D: DestroyedState,
287+
{
288+
// Print a nice abort message if a panic occurs.
289+
abort_on_dtor_unwind(|| {
290+
let storage = unsafe { &*(ptr as *const Storage<I, T, D>) };
291+
// Update the state before running the destructor as it may attempt to
292+
// access the variable.
293+
let state = unsafe { storage.state.get().replace(State::Destroyed(D::create())) };
294+
match state {
295+
State::Alive(v) => drop(v),
296+
// This should not occur but do nothing if it does.
297+
State::Destroyed(_) => {}
298+
State::Initial(_) => {
299+
unreachable!("destructor cannot run if it was not registered yet");
300+
}
301+
}
302+
})
247303
}

‎library/std/src/sys/thread_local/mod.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ cfg_if::cfg_if! {
1515
#[doc(hidden)]
1616
mod fast_local;
1717
#[doc(hidden)]
18-
pub use fast_local::{Key, thread_local_inner};
18+
pub use fast_local::{Storage, take_or_call, thread_local_inner};
1919
} else {
2020
#[doc(hidden)]
2121
mod os_local;
@@ -24,6 +24,9 @@ cfg_if::cfg_if! {
2424
}
2525
}
2626

27+
// Not used by the fast-local TLS anymore.
28+
// FIXME(#110897): remove this.
29+
#[allow(unused)]
2730
mod lazy {
2831
use crate::cell::UnsafeCell;
2932
use crate::hint;

‎library/std/src/thread/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,7 @@ cfg_if::cfg_if! {
205205
#[doc(hidden)]
206206
#[unstable(feature = "thread_local_internals", issue = "none")]
207207
pub mod local_impl {
208-
pub use crate::sys::thread_local::{thread_local_inner, Key, abort_on_dtor_unwind};
208+
pub use crate::sys::thread_local::*;
209209
}
210210
}
211211
}

0 commit comments

Comments
 (0)
Please sign in to comment.