Skip to content

Commit 5fd4a84

Browse files
joboetjoboet
committed
std: simplify key-based thread locals
1 parent e69f14b commit 5fd4a84

File tree

2 files changed

+62
-102
lines changed

2 files changed

+62
-102
lines changed

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

+2
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ cfg_if::cfg_if! {
2424
}
2525
}
2626

27+
// FIXME(#110897): remove this
28+
#[allow(unused)] // Not used by key-based TLS.
2729
mod lazy {
2830
use crate::cell::UnsafeCell;
2931
use crate::hint;

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

+60-102
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
use super::lazy::LazyKeyInner;
1+
use super::abort_on_dtor_unwind;
22
use crate::cell::Cell;
3-
use crate::sys_common::thread_local_key::StaticKey as OsStaticKey;
4-
use crate::{fmt, marker, panic, ptr};
3+
use crate::marker::PhantomData;
4+
use crate::sys_common::thread_local_key::StaticKey as OsKey;
5+
use crate::{panic, ptr};
56

67
#[doc(hidden)]
78
#[allow_internal_unstable(thread_local_internals)]
@@ -10,38 +11,9 @@ use crate::{fmt, marker, panic, ptr};
1011
#[rustc_macro_transparency = "semitransparent"]
1112
pub macro thread_local_inner {
1213
// used to generate the `LocalKey` value for const-initialized thread locals
13-
(@key $t:ty, const $init:expr) => {{
14-
#[inline]
15-
#[deny(unsafe_op_in_unsafe_fn)]
16-
unsafe fn __getit(
17-
_init: $crate::option::Option<&mut $crate::option::Option<$t>>,
18-
) -> $crate::option::Option<&'static $t> {
19-
const INIT_EXPR: $t = $init;
20-
21-
// On platforms without `#[thread_local]` we fall back to the
22-
// same implementation as below for os thread locals.
23-
#[inline]
24-
const fn __init() -> $t { INIT_EXPR }
25-
static __KEY: $crate::thread::local_impl::Key<$t> =
26-
$crate::thread::local_impl::Key::new();
27-
unsafe {
28-
__KEY.get(move || {
29-
if let $crate::option::Option::Some(init) = _init {
30-
if let $crate::option::Option::Some(value) = init.take() {
31-
return value;
32-
} else if $crate::cfg!(debug_assertions) {
33-
$crate::unreachable!("missing initial value");
34-
}
35-
}
36-
__init()
37-
})
38-
}
39-
}
40-
41-
unsafe {
42-
$crate::thread::LocalKey::new(__getit)
43-
}
44-
}},
14+
(@key $t:ty, const $init:expr) => {
15+
$crate::thread::local_impl::thread_local_inner!(@key $t, { const INIT_EXPR: $t = $init; INIT_EXPR })
16+
},
4517

4618
// used to generate the `LocalKey` value for `thread_local!`
4719
(@key $t:ty, $init:expr) => {
@@ -55,20 +27,11 @@ pub macro thread_local_inner {
5527
unsafe fn __getit(
5628
init: $crate::option::Option<&mut $crate::option::Option<$t>>,
5729
) -> $crate::option::Option<&'static $t> {
58-
static __KEY: $crate::thread::local_impl::Key<$t> =
59-
$crate::thread::local_impl::Key::new();
30+
use $crate::thread::local_impl::Key;
6031

32+
static __KEY: Key<$t> = Key::new();
6133
unsafe {
62-
__KEY.get(move || {
63-
if let $crate::option::Option::Some(init) = init {
64-
if let $crate::option::Option::Some(value) = init.take() {
65-
return value;
66-
} else if $crate::cfg!(debug_assertions) {
67-
$crate::unreachable!("missing default value");
68-
}
69-
}
70-
__init()
71-
})
34+
__KEY.get(init, __init)
7235
}
7336
}
7437

@@ -85,78 +48,78 @@ pub macro thread_local_inner {
8548

8649
/// Use a regular global static to store this key; the state provided will then be
8750
/// thread-local.
51+
#[allow(missing_debug_implementations)]
8852
pub struct Key<T> {
89-
// OS-TLS key that we'll use to key off.
90-
os: OsStaticKey,
91-
marker: marker::PhantomData<Cell<T>>,
92-
}
93-
94-
impl<T> fmt::Debug for Key<T> {
95-
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96-
f.debug_struct("Key").finish_non_exhaustive()
97-
}
53+
os: OsKey,
54+
marker: PhantomData<Cell<T>>,
9855
}
9956

10057
unsafe impl<T> Sync for Key<T> {}
10158

10259
struct Value<T: 'static> {
103-
inner: LazyKeyInner<T>,
60+
value: T,
10461
key: &'static Key<T>,
10562
}
10663

10764
impl<T: 'static> Key<T> {
10865
#[rustc_const_unstable(feature = "thread_local_internals", issue = "none")]
10966
pub const fn new() -> Key<T> {
110-
Key { os: OsStaticKey::new(Some(destroy_value::<T>)), marker: marker::PhantomData }
67+
Key { os: OsKey::new(Some(destroy_value::<T>)), marker: PhantomData }
11168
}
11269

113-
/// It is a requirement for the caller to ensure that no mutable
114-
/// reference is active when this method is called.
115-
pub unsafe fn get(&'static self, init: impl FnOnce() -> T) -> Option<&'static T> {
116-
// SAFETY: See the documentation for this method.
70+
/// Get the value associated with this key, initializating it if necessary.
71+
///
72+
/// # Safety
73+
/// * the returned reference must not be used after recursive initialization
74+
/// or thread destruction occurs.
75+
pub unsafe fn get(
76+
&'static self,
77+
i: Option<&mut Option<T>>,
78+
f: impl FnOnce() -> T,
79+
) -> Option<&'static T> {
80+
// SAFETY: (FIXME: get should actually be safe)
11781
let ptr = unsafe { self.os.get() as *mut Value<T> };
11882
if ptr.addr() > 1 {
11983
// SAFETY: the check ensured the pointer is safe (its destructor
12084
// is not running) + it is coming from a trusted source (self).
121-
if let Some(ref value) = unsafe { (*ptr).inner.get() } {
122-
return Some(value);
123-
}
85+
unsafe { Some(&(*ptr).value) }
86+
} else {
87+
// SAFETY: At this point we are sure we have no value and so
88+
// initializing (or trying to) is safe.
89+
unsafe { self.try_initialize(ptr, i, f) }
12490
}
125-
// SAFETY: At this point we are sure we have no value and so
126-
// initializing (or trying to) is safe.
127-
unsafe { self.try_initialize(init) }
12891
}
12992

130-
// `try_initialize` is only called once per os thread local variable,
131-
// except in corner cases where thread_local dtors reference other
132-
// thread_local's, or it is being recursively initialized.
133-
unsafe fn try_initialize(&'static self, init: impl FnOnce() -> T) -> Option<&'static T> {
134-
// SAFETY: No mutable references are ever handed out meaning getting
135-
// the value is ok.
136-
let ptr = unsafe { self.os.get() as *mut Value<T> };
93+
unsafe fn try_initialize(
94+
&'static self,
95+
ptr: *mut Value<T>,
96+
i: Option<&mut Option<T>>,
97+
f: impl FnOnce() -> T,
98+
) -> Option<&'static T> {
13799
if ptr.addr() == 1 {
138100
// destructor is running
139101
return None;
140102
}
141103

142-
let ptr = if ptr.is_null() {
143-
// If the lookup returned null, we haven't initialized our own
144-
// local copy, so do that now.
145-
let ptr = Box::into_raw(Box::new(Value { inner: LazyKeyInner::new(), key: self }));
146-
// SAFETY: At this point we are sure there is no value inside
147-
// ptr so setting it will not affect anyone else.
148-
unsafe {
149-
self.os.set(ptr as *mut u8);
150-
}
151-
ptr
152-
} else {
153-
// recursive initialization
154-
ptr
155-
};
104+
let value = i.and_then(Option::take).unwrap_or_else(f);
105+
let ptr = Box::into_raw(Box::new(Value { value, key: self }));
106+
// SAFETY: (FIXME: get should actually be safe)
107+
let old = unsafe { self.os.get() as *mut Value<T> };
108+
// SAFETY: `ptr` is a correct pointer that can be destroyed by the key destructor.
109+
unsafe {
110+
self.os.set(ptr as *mut u8);
111+
}
112+
if !old.is_null() {
113+
// If the variable was recursively initialized, drop the old value.
114+
// SAFETY: We cannot be inside a `LocalKey::with` scope, as the
115+
// initializer has already returned and the next scope only starts
116+
// after we return the pointer. Therefore, there can be no references
117+
// to the old value.
118+
drop(unsafe { Box::from_raw(old) });
119+
}
156120

157-
// SAFETY: ptr has been ensured as non-NUL just above an so can be
158-
// dereferenced safely.
159-
unsafe { Some((*ptr).inner.initialize(init)) }
121+
// SAFETY: We just created this value above.
122+
unsafe { Some(&(*ptr).value) }
160123
}
161124
}
162125

@@ -170,16 +133,11 @@ unsafe extern "C" fn destroy_value<T: 'static>(ptr: *mut u8) {
170133
//
171134
// Note that to prevent an infinite loop we reset it back to null right
172135
// before we return from the destructor ourselves.
173-
//
174-
// Wrap the call in a catch to ensure unwinding is caught in the event
175-
// a panic takes place in a destructor.
176-
if let Err(_) = panic::catch_unwind(|| unsafe {
177-
let ptr = Box::from_raw(ptr as *mut Value<T>);
136+
abort_on_dtor_unwind(|| {
137+
let ptr = unsafe { Box::from_raw(ptr as *mut Value<T>) };
178138
let key = ptr.key;
179-
key.os.set(ptr::without_provenance_mut(1));
139+
unsafe { key.os.set(ptr::without_provenance_mut(1)) };
180140
drop(ptr);
181-
key.os.set(ptr::null_mut());
182-
}) {
183-
rtabort!("thread local panicked on drop");
184-
}
141+
unsafe { key.os.set(ptr::null_mut()) };
142+
});
185143
}

0 commit comments

Comments
 (0)