Skip to content

Commit 9c8a58f

Browse files
committedMay 23, 2024
Auto merge of #116123 - joboet:rewrite_native_tls, r=m-ou-se
Rewrite native thread-local storage (part of #110897) The current native thread-local storage implementation has become quite messy, uses indescriptive names and unnecessarily adds code to the macro expansion. This PR tries to fix that by using a new implementation that also allows more layout optimizations and potentially increases performance by eliminating unnecessary TLS accesses. This does not change the recursive initialization behaviour I described in [this comment](#110897 (comment)), so it should be a library-only change. Changing that behaviour should be quite easy now, however. r? `@m-ou-se` `@rustbot` label +T-libs
·
1.90.01.80.0
2 parents ed172db + 60bf1ab commit 9c8a58f

File tree

10 files changed

+331
-332
lines changed

10 files changed

+331
-332
lines changed
 

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

Lines changed: 0 additions & 247 deletions
This file was deleted.
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
use crate::cell::{Cell, UnsafeCell};
2+
use crate::ptr::{self, drop_in_place};
3+
use crate::sys::thread_local::abort_on_dtor_unwind;
4+
use crate::sys::thread_local_dtor::register_dtor;
5+
6+
#[derive(Clone, Copy)]
7+
enum State {
8+
Initial,
9+
Alive,
10+
Destroyed,
11+
}
12+
13+
#[allow(missing_debug_implementations)]
14+
pub struct Storage<T> {
15+
state: Cell<State>,
16+
val: UnsafeCell<T>,
17+
}
18+
19+
impl<T> Storage<T> {
20+
pub const fn new(val: T) -> Storage<T> {
21+
Storage { state: Cell::new(State::Initial), val: UnsafeCell::new(val) }
22+
}
23+
24+
/// Get a reference to the TLS value. If the TLS variable has been destroyed,
25+
/// `None` is returned.
26+
///
27+
/// # Safety
28+
/// * The `self` reference must remain valid until the TLS destructor has been
29+
/// run.
30+
/// * The returned reference may only be used until thread destruction occurs
31+
/// and may not be used after reentrant initialization has occurred.
32+
///
33+
// FIXME(#110897): return NonNull instead of lying about the lifetime.
34+
#[inline]
35+
pub unsafe fn get(&self) -> Option<&'static T> {
36+
match self.state.get() {
37+
// SAFETY: as the state is not `Destroyed`, the value cannot have
38+
// been destroyed yet. The reference fulfills the terms outlined
39+
// above.
40+
State::Alive => unsafe { Some(&*self.val.get()) },
41+
State::Destroyed => None,
42+
State::Initial => unsafe { self.initialize() },
43+
}
44+
}
45+
46+
#[cold]
47+
unsafe fn initialize(&self) -> Option<&'static T> {
48+
// Register the destructor
49+
50+
// SAFETY:
51+
// * the destructor will be called at thread destruction.
52+
// * the caller guarantees that `self` will be valid until that time.
53+
unsafe {
54+
register_dtor(ptr::from_ref(self).cast_mut().cast(), destroy::<T>);
55+
}
56+
self.state.set(State::Alive);
57+
// SAFETY: as the state is not `Destroyed`, the value cannot have
58+
// been destroyed yet. The reference fulfills the terms outlined
59+
// above.
60+
unsafe { Some(&*self.val.get()) }
61+
}
62+
}
63+
64+
/// Transition an `Alive` TLS variable into the `Destroyed` state, dropping its
65+
/// value.
66+
///
67+
/// # Safety
68+
/// * Must only be called at thread destruction.
69+
/// * `ptr` must point to an instance of `Storage` with `Alive` state and be
70+
/// valid for accessing that instance.
71+
unsafe extern "C" fn destroy<T>(ptr: *mut u8) {
72+
// Print a nice abort message if a panic occurs.
73+
abort_on_dtor_unwind(|| {
74+
let storage = unsafe { &*(ptr as *const Storage<T>) };
75+
// Update the state before running the destructor as it may attempt to
76+
// access the variable.
77+
storage.state.set(State::Destroyed);
78+
unsafe {
79+
drop_in_place(storage.val.get());
80+
}
81+
})
82+
}
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
use crate::cell::UnsafeCell;
2+
use crate::hint::unreachable_unchecked;
3+
use crate::mem::forget;
4+
use crate::ptr;
5+
use crate::sys::thread_local::abort_on_dtor_unwind;
6+
use crate::sys::thread_local_dtor::register_dtor;
7+
8+
pub unsafe trait DestroyedState: Sized {
9+
fn register_dtor<T>(s: &Storage<T, Self>);
10+
}
11+
12+
unsafe impl DestroyedState for ! {
13+
fn register_dtor<T>(_: &Storage<T, !>) {}
14+
}
15+
16+
unsafe impl DestroyedState for () {
17+
fn register_dtor<T>(s: &Storage<T, ()>) {
18+
unsafe {
19+
register_dtor(ptr::from_ref(s).cast_mut().cast(), destroy::<T>);
20+
}
21+
}
22+
}
23+
24+
enum State<T, D> {
25+
Initial,
26+
Alive(T),
27+
Destroyed(D),
28+
}
29+
30+
#[allow(missing_debug_implementations)]
31+
pub struct Storage<T, D> {
32+
state: UnsafeCell<State<T, D>>,
33+
}
34+
35+
impl<T, D> Storage<T, D>
36+
where
37+
D: DestroyedState,
38+
{
39+
pub const fn new() -> Storage<T, D> {
40+
Storage { state: UnsafeCell::new(State::Initial) }
41+
}
42+
43+
/// Get a reference to the TLS value, potentially initializing it with the
44+
/// provided parameters. If the TLS variable has been destroyed, `None` is
45+
/// returned.
46+
///
47+
/// # Safety
48+
/// * The `self` reference must remain valid until the TLS destructor is run,
49+
/// at which point the returned reference is invalidated.
50+
/// * The returned reference may only be used until thread destruction occurs
51+
/// and may not be used after reentrant initialization has occurred.
52+
///
53+
// FIXME(#110897): return NonNull instead of lying about the lifetime.
54+
#[inline]
55+
pub unsafe fn get_or_init(
56+
&self,
57+
i: Option<&mut Option<T>>,
58+
f: impl FnOnce() -> T,
59+
) -> Option<&'static T> {
60+
// SAFETY:
61+
// No mutable reference to the inner value exists outside the calls to
62+
// `replace`. The lifetime of the returned reference fulfills the terms
63+
// outlined above.
64+
let state = unsafe { &*self.state.get() };
65+
match state {
66+
State::Alive(v) => Some(v),
67+
State::Destroyed(_) => None,
68+
State::Initial => unsafe { self.initialize(i, f) },
69+
}
70+
}
71+
72+
#[cold]
73+
unsafe fn initialize(
74+
&self,
75+
i: Option<&mut Option<T>>,
76+
f: impl FnOnce() -> T,
77+
) -> Option<&'static T> {
78+
// Perform initialization
79+
80+
let v = i.and_then(Option::take).unwrap_or_else(f);
81+
82+
// SAFETY:
83+
// If references to the inner value exist, they were created in `f`
84+
// and are invalidated here. The caller promises to never use them
85+
// after this.
86+
let old = unsafe { self.state.get().replace(State::Alive(v)) };
87+
match old {
88+
// If the variable is not being recursively initialized, register
89+
// the destructor. This might be a noop if the value does not need
90+
// destruction.
91+
State::Initial => D::register_dtor(self),
92+
// Else, drop the old value. This might be changed to a panic.
93+
val => drop(val),
94+
}
95+
96+
// SAFETY:
97+
// Initialization was completed and the state was set to `Alive`, so the
98+
// reference fulfills the terms outlined above.
99+
unsafe {
100+
let State::Alive(v) = &*self.state.get() else { unreachable_unchecked() };
101+
Some(v)
102+
}
103+
}
104+
}
105+
106+
/// Transition an `Alive` TLS variable into the `Destroyed` state, dropping its
107+
/// value.
108+
///
109+
/// # Safety
110+
/// * Must only be called at thread destruction.
111+
/// * `ptr` must point to an instance of `Storage<T, ()>` and be valid for
112+
/// accessing that instance.
113+
unsafe extern "C" fn destroy<T>(ptr: *mut u8) {
114+
// Print a nice abort message if a panic occurs.
115+
abort_on_dtor_unwind(|| {
116+
let storage = unsafe { &*(ptr as *const Storage<T, ()>) };
117+
// Update the state before running the destructor as it may attempt to
118+
// access the variable.
119+
let val = unsafe { storage.state.get().replace(State::Destroyed(())) };
120+
drop(val);
121+
})
122+
}
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
//! Thread local support for platforms with native TLS.
2+
//!
3+
//! To achieve the best performance, we choose from four different types for
4+
//! the TLS variable, depending from the method of initialization used (`const`
5+
//! or lazy) and the drop requirements of the stored type:
6+
//!
7+
//! | | `Drop` | `!Drop` |
8+
//! |--------:|:--------------------:|:-------------------:|
9+
//! | `const` | `EagerStorage<T>` | `T` |
10+
//! | lazy | `LazyStorage<T, ()>` | `LazyStorage<T, !>` |
11+
//!
12+
//! For `const` initialization and `!Drop` types, we simply use `T` directly,
13+
//! but for other situations, we implement a state machine to handle
14+
//! initialization of the variable and its destructor and destruction.
15+
//! Upon accessing the TLS variable, the current state is compared:
16+
//!
17+
//! 1. If the state is `Initial`, initialize the storage, transition the state
18+
//! to `Alive` and (if applicable) register the destructor, and return a
19+
//! reference to the value.
20+
//! 2. If the state is `Alive`, initialization was previously completed, so
21+
//! return a reference to the value.
22+
//! 3. If the state is `Destroyed`, the destructor has been run already, so
23+
//! return [`None`].
24+
//!
25+
//! The TLS destructor sets the state to `Destroyed` and drops the current value.
26+
//!
27+
//! To simplify the code, we make `LazyStorage` generic over the destroyed state
28+
//! and use the `!` type (never type) as type parameter for `!Drop` types. This
29+
//! eliminates the `Destroyed` state for these values, which can allow more niche
30+
//! optimizations to occur for the `State` enum. For `Drop` types, `()` is used.
31+
32+
#![deny(unsafe_op_in_unsafe_fn)]
33+
34+
mod eager;
35+
mod lazy;
36+
37+
pub use eager::Storage as EagerStorage;
38+
pub use lazy::Storage as LazyStorage;
39+
40+
#[doc(hidden)]
41+
#[allow_internal_unstable(
42+
thread_local_internals,
43+
cfg_target_thread_local,
44+
thread_local,
45+
never_type
46+
)]
47+
#[allow_internal_unsafe]
48+
#[unstable(feature = "thread_local_internals", issue = "none")]
49+
#[rustc_macro_transparency = "semitransparent"]
50+
pub macro thread_local_inner {
51+
// used to generate the `LocalKey` value for const-initialized thread locals
52+
(@key $t:ty, const $init:expr) => {{
53+
const __INIT: $t = $init;
54+
55+
#[inline]
56+
#[deny(unsafe_op_in_unsafe_fn)]
57+
unsafe fn __getit(
58+
_init: $crate::option::Option<&mut $crate::option::Option<$t>>,
59+
) -> $crate::option::Option<&'static $t> {
60+
use $crate::thread::local_impl::EagerStorage;
61+
use $crate::mem::needs_drop;
62+
use $crate::ptr::addr_of;
63+
64+
if needs_drop::<$t>() {
65+
#[thread_local]
66+
static VAL: EagerStorage<$t> = EagerStorage::new(__INIT);
67+
unsafe {
68+
VAL.get()
69+
}
70+
} else {
71+
#[thread_local]
72+
static VAL: $t = __INIT;
73+
unsafe {
74+
$crate::option::Option::Some(&*addr_of!(VAL))
75+
}
76+
}
77+
}
78+
79+
unsafe {
80+
$crate::thread::LocalKey::new(__getit)
81+
}
82+
}},
83+
84+
// used to generate the `LocalKey` value for `thread_local!`
85+
(@key $t:ty, $init:expr) => {{
86+
#[inline]
87+
fn __init() -> $t {
88+
$init
89+
}
90+
91+
#[inline]
92+
#[deny(unsafe_op_in_unsafe_fn)]
93+
unsafe fn __getit(
94+
init: $crate::option::Option<&mut $crate::option::Option<$t>>,
95+
) -> $crate::option::Option<&'static $t> {
96+
use $crate::thread::local_impl::LazyStorage;
97+
use $crate::mem::needs_drop;
98+
99+
if needs_drop::<$t>() {
100+
#[thread_local]
101+
static VAL: LazyStorage<$t, ()> = LazyStorage::new();
102+
unsafe {
103+
VAL.get_or_init(init, __init)
104+
}
105+
} else {
106+
#[thread_local]
107+
static VAL: LazyStorage<$t, !> = LazyStorage::new();
108+
unsafe {
109+
VAL.get_or_init(init, __init)
110+
}
111+
}
112+
}
113+
114+
unsafe {
115+
$crate::thread::LocalKey::new(__getit)
116+
}
117+
}},
118+
($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {
119+
$(#[$attr])* $vis const $name: $crate::thread::LocalKey<$t> =
120+
$crate::thread::local_impl::thread_local_inner!(@key $t, $($init)*);
121+
},
122+
}

‎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::{EagerStorage, LazyStorage, 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
}

‎src/tools/tidy/src/issues.txt‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3988,8 +3988,6 @@ ui/test-attrs/issue-52557.rs
39883988
ui/test-attrs/issue-53675-a-test-called-panic.rs
39893989
ui/threads-sendsync/issue-24313.rs
39903990
ui/threads-sendsync/issue-29488.rs
3991-
ui/threads-sendsync/issue-43733-2.rs
3992-
ui/threads-sendsync/issue-43733.rs
39933991
ui/threads-sendsync/issue-4446.rs
39943992
ui/threads-sendsync/issue-4448.rs
39953993
ui/threads-sendsync/issue-8827.rs

‎tests/ui/threads-sendsync/issue-43733-2.rs‎

Lines changed: 0 additions & 30 deletions
This file was deleted.

‎tests/ui/threads-sendsync/issue-43733.rs‎

Lines changed: 0 additions & 32 deletions
This file was deleted.

‎tests/ui/threads-sendsync/issue-43733.stderr‎

Lines changed: 0 additions & 19 deletions
This file was deleted.

0 commit comments

Comments
 (0)
Please sign in to comment.