Skip to content

Commit ac9b2fd

Browse files
alexcrichtonalexcrichton
committed
std: Add a variant of thread locals with const init
This commit adds a variant of the `thread_local!` macro as a new `thread_local_const_init!` macro which requires that the initialization expression is constant (e.g. could be stuck into a `const` if so desired). This form of thread local allows for a more efficient implementation of `LocalKey::with` both if the value has a destructor and if it doesn't. If the value doesn't have a destructor then `with` should desugar to exactly as-if you use `#[thread_local]` given sufficient inlining. The purpose of this new form of thread locals is to precisely be equivalent to `#[thread_local]` on platforms where possible for values which fit the bill (those without destructors). This should help close the gap in performance between `thread_local!`, which is safe, relative to `#[thread_local]`, which is not easy to use in a portable fashion.
1 parent 2e6d4b8 commit ac9b2fd

File tree

4 files changed

+227
-49
lines changed

4 files changed

+227
-49
lines changed

‎std/src/lib.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,10 @@
208208
// std may use features in a platform-specific way
209209
#![allow(unused_features)]
210210
#![feature(rustc_allow_const_fn_unstable)]
211-
#![cfg_attr(test, feature(internal_output_capture, print_internals, update_panic_count))]
211+
#![cfg_attr(
212+
test,
213+
feature(internal_output_capture, print_internals, update_panic_count, thread_local_const_init)
214+
)]
212215
#![cfg_attr(
213216
all(target_vendor = "fortanix", target_env = "sgx"),
214217
feature(slice_index_methods, coerce_unsized, sgx_platform)

‎std/src/thread/local.rs‎

Lines changed: 115 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,15 @@ macro_rules! thread_local {
133133
// empty (base case for the recursion)
134134
() => {};
135135

136+
($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty = const { $init:expr }; $($rest:tt)*) => (
137+
$crate::__thread_local_inner!($(#[$attr])* $vis $name, $t, const $init);
138+
$crate::thread_local!($($rest)*);
139+
);
140+
141+
($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty = const { $init:expr }) => (
142+
$crate::__thread_local_inner!($(#[$attr])* $vis $name, $t, const $init);
143+
);
144+
136145
// process multiple declarations
137146
($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty = $init:expr; $($rest:tt)*) => (
138147
$crate::__thread_local_inner!($(#[$attr])* $vis $name, $t, $init);
@@ -151,6 +160,101 @@ macro_rules! thread_local {
151160
#[allow_internal_unstable(thread_local_internals, cfg_target_thread_local, thread_local)]
152161
#[allow_internal_unsafe]
153162
macro_rules! __thread_local_inner {
163+
// used to generate the `LocalKey` value for const-initialized thread locals
164+
(@key $t:ty, const $init:expr) => {{
165+
unsafe fn __getit() -> $crate::option::Option<&'static $t> {
166+
const _REQUIRE_UNSTABLE: () = $crate::thread::require_unstable_const_init_thread_local();
167+
168+
// wasm without atomics maps directly to `static mut`, and dtors
169+
// aren't implemented because thread dtors aren't really a thing
170+
// on wasm right now
171+
//
172+
// FIXME(#84224) this should come after the `target_thread_local`
173+
// block.
174+
#[cfg(all(target_arch = "wasm32", not(target_feature = "atomics")))]
175+
{
176+
static mut VAL: $t = $init;
177+
Some(&VAL)
178+
}
179+
180+
// If the platform has support for `#[thread_local]`, use it.
181+
#[cfg(all(
182+
target_thread_local,
183+
not(all(target_arch = "wasm32", not(target_feature = "atomics"))),
184+
))]
185+
{
186+
// If a dtor isn't needed we can do something "very raw" and
187+
// just get going.
188+
if !$crate::mem::needs_drop::<$t>() {
189+
#[thread_local]
190+
static mut VAL: $t = $init;
191+
unsafe {
192+
return Some(&VAL)
193+
}
194+
}
195+
196+
#[thread_local]
197+
static mut VAL: $t = $init;
198+
// 0 == dtor not registered
199+
// 1 == dtor registered, dtor not run
200+
// 2 == dtor registered and is running or has run
201+
#[thread_local]
202+
static mut STATE: u8 = 0;
203+
204+
unsafe extern "C" fn destroy(ptr: *mut u8) {
205+
let ptr = ptr as *mut $t;
206+
207+
unsafe {
208+
debug_assert_eq!(STATE, 1);
209+
STATE = 2;
210+
$crate::ptr::drop_in_place(ptr);
211+
}
212+
}
213+
214+
unsafe {
215+
match STATE {
216+
// 0 == we haven't registered a destructor, so do
217+
// so now.
218+
0 => {
219+
$crate::thread::__FastLocalKeyInner::<$t>::register_dtor(
220+
&VAL as *const _ as *mut u8,
221+
destroy,
222+
);
223+
STATE = 1;
224+
Some(&VAL)
225+
}
226+
// 1 == the destructor is registered and the value
227+
// is valid, so return the pointer.
228+
1 => Some(&VAL),
229+
// otherwise the destructor has already run, so we
230+
// can't give access.
231+
_ => None,
232+
}
233+
}
234+
}
235+
236+
// On platforms without `#[thread_local]` we fall back to the
237+
// same implementation as below for os thread locals.
238+
#[cfg(all(
239+
not(target_thread_local),
240+
not(all(target_arch = "wasm32", not(target_feature = "atomics"))),
241+
))]
242+
{
243+
#[inline]
244+
const fn __init() -> $t { $init }
245+
static __KEY: $crate::thread::__OsLocalKeyInner<$t> =
246+
$crate::thread::__OsLocalKeyInner::new();
247+
#[allow(unused_unsafe)]
248+
unsafe { __KEY.get(__init) }
249+
}
250+
}
251+
252+
unsafe {
253+
$crate::thread::LocalKey::new(__getit)
254+
}
255+
}};
256+
257+
// used to generate the `LocalKey` value for `thread_local!`
154258
(@key $t:ty, $init:expr) => {
155259
{
156260
#[inline]
@@ -188,9 +292,9 @@ macro_rules! __thread_local_inner {
188292
}
189293
}
190294
};
191-
($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $init:expr) => {
295+
($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {
192296
$(#[$attr])* $vis const $name: $crate::thread::LocalKey<$t> =
193-
$crate::__thread_local_inner!(@key $t, $init);
297+
$crate::__thread_local_inner!(@key $t, $($init)*);
194298
}
195299
}
196300

@@ -442,6 +546,15 @@ pub mod fast {
442546
Key { inner: LazyKeyInner::new(), dtor_state: Cell::new(DtorState::Unregistered) }
443547
}
444548

549+
// note that this is just a publically-callable function only for the
550+
// const-initialized form of thread locals, basically a way to call the
551+
// free `register_dtor` function defined elsewhere in libstd.
552+
pub unsafe fn register_dtor(a: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
553+
unsafe {
554+
register_dtor(a, dtor);
555+
}
556+
}
557+
445558
pub unsafe fn get<F: FnOnce() -> T>(&self, init: F) -> Option<&'static T> {
446559
// SAFETY: See the definitions of `LazyKeyInner::get` and
447560
// `try_initialize` for more informations.

‎std/src/thread/local/tests.rs‎

Lines changed: 101 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::cell::{Cell, UnsafeCell};
22
use crate::sync::mpsc::{channel, Sender};
3-
use crate::thread;
3+
use crate::thread::{self, LocalKey};
44
use crate::thread_local;
55

66
struct Foo(Sender<()>);
@@ -15,74 +15,90 @@ impl Drop for Foo {
1515
#[test]
1616
fn smoke_no_dtor() {
1717
thread_local!(static FOO: Cell<i32> = Cell::new(1));
18+
run(&FOO);
19+
thread_local!(static FOO2: Cell<i32> = const { Cell::new(1) });
20+
run(&FOO2);
1821

19-
FOO.with(|f| {
20-
assert_eq!(f.get(), 1);
21-
f.set(2);
22-
});
23-
let (tx, rx) = channel();
24-
let _t = thread::spawn(move || {
25-
FOO.with(|f| {
22+
fn run(key: &'static LocalKey<Cell<i32>>) {
23+
key.with(|f| {
2624
assert_eq!(f.get(), 1);
25+
f.set(2);
2726
});
28-
tx.send(()).unwrap();
29-
});
30-
rx.recv().unwrap();
27+
let t = thread::spawn(move || {
28+
key.with(|f| {
29+
assert_eq!(f.get(), 1);
30+
});
31+
});
32+
t.join().unwrap();
3133

32-
FOO.with(|f| {
33-
assert_eq!(f.get(), 2);
34-
});
34+
key.with(|f| {
35+
assert_eq!(f.get(), 2);
36+
});
37+
}
3538
}
3639

3740
#[test]
3841
fn states() {
39-
struct Foo;
42+
struct Foo(&'static LocalKey<Foo>);
4043
impl Drop for Foo {
4144
fn drop(&mut self) {
42-
assert!(FOO.try_with(|_| ()).is_err());
45+
assert!(self.0.try_with(|_| ()).is_err());
4346
}
4447
}
45-
thread_local!(static FOO: Foo = Foo);
4648

47-
thread::spawn(|| {
48-
assert!(FOO.try_with(|_| ()).is_ok());
49-
})
50-
.join()
51-
.ok()
52-
.expect("thread panicked");
49+
thread_local!(static FOO: Foo = Foo(&FOO));
50+
run(&FOO);
51+
thread_local!(static FOO2: Foo = const { Foo(&FOO2) });
52+
run(&FOO2);
53+
54+
fn run(foo: &'static LocalKey<Foo>) {
55+
thread::spawn(move || {
56+
assert!(foo.try_with(|_| ()).is_ok());
57+
})
58+
.join()
59+
.unwrap();
60+
}
5361
}
5462

5563
#[test]
5664
fn smoke_dtor() {
5765
thread_local!(static FOO: UnsafeCell<Option<Foo>> = UnsafeCell::new(None));
58-
59-
let (tx, rx) = channel();
60-
let _t = thread::spawn(move || unsafe {
61-
let mut tx = Some(tx);
62-
FOO.with(|f| {
63-
*f.get() = Some(Foo(tx.take().unwrap()));
66+
run(&FOO);
67+
thread_local!(static FOO2: UnsafeCell<Option<Foo>> = const { UnsafeCell::new(None) });
68+
run(&FOO2);
69+
70+
fn run(key: &'static LocalKey<UnsafeCell<Option<Foo>>>) {
71+
let (tx, rx) = channel();
72+
let t = thread::spawn(move || unsafe {
73+
let mut tx = Some(tx);
74+
key.with(|f| {
75+
*f.get() = Some(Foo(tx.take().unwrap()));
76+
});
6477
});
65-
});
66-
rx.recv().unwrap();
78+
rx.recv().unwrap();
79+
t.join().unwrap();
80+
}
6781
}
6882

6983
#[test]
7084
fn circular() {
71-
struct S1;
72-
struct S2;
85+
struct S1(&'static LocalKey<UnsafeCell<Option<S1>>>, &'static LocalKey<UnsafeCell<Option<S2>>>);
86+
struct S2(&'static LocalKey<UnsafeCell<Option<S1>>>, &'static LocalKey<UnsafeCell<Option<S2>>>);
7387
thread_local!(static K1: UnsafeCell<Option<S1>> = UnsafeCell::new(None));
7488
thread_local!(static K2: UnsafeCell<Option<S2>> = UnsafeCell::new(None));
75-
static mut HITS: u32 = 0;
89+
thread_local!(static K3: UnsafeCell<Option<S1>> = const { UnsafeCell::new(None) });
90+
thread_local!(static K4: UnsafeCell<Option<S2>> = const { UnsafeCell::new(None) });
91+
static mut HITS: usize = 0;
7692

7793
impl Drop for S1 {
7894
fn drop(&mut self) {
7995
unsafe {
8096
HITS += 1;
81-
if K2.try_with(|_| ()).is_err() {
97+
if self.1.try_with(|_| ()).is_err() {
8298
assert_eq!(HITS, 3);
8399
} else {
84100
if HITS == 1 {
85-
K2.with(|s| *s.get() = Some(S2));
101+
self.1.with(|s| *s.get() = Some(S2(self.0, self.1)));
86102
} else {
87103
assert_eq!(HITS, 3);
88104
}
@@ -94,38 +110,54 @@ fn circular() {
94110
fn drop(&mut self) {
95111
unsafe {
96112
HITS += 1;
97-
assert!(K1.try_with(|_| ()).is_ok());
113+
assert!(self.0.try_with(|_| ()).is_ok());
98114
assert_eq!(HITS, 2);
99-
K1.with(|s| *s.get() = Some(S1));
115+
self.0.with(|s| *s.get() = Some(S1(self.0, self.1)));
100116
}
101117
}
102118
}
103119

104120
thread::spawn(move || {
105-
drop(S1);
121+
drop(S1(&K1, &K2));
122+
})
123+
.join()
124+
.unwrap();
125+
126+
unsafe {
127+
HITS = 0;
128+
}
129+
130+
thread::spawn(move || {
131+
drop(S1(&K3, &K4));
106132
})
107133
.join()
108-
.ok()
109-
.expect("thread panicked");
134+
.unwrap();
110135
}
111136

112137
#[test]
113138
fn self_referential() {
114-
struct S1;
139+
struct S1(&'static LocalKey<UnsafeCell<Option<S1>>>);
140+
115141
thread_local!(static K1: UnsafeCell<Option<S1>> = UnsafeCell::new(None));
142+
thread_local!(static K2: UnsafeCell<Option<S1>> = const { UnsafeCell::new(None) });
116143

117144
impl Drop for S1 {
118145
fn drop(&mut self) {
119-
assert!(K1.try_with(|_| ()).is_err());
146+
assert!(self.0.try_with(|_| ()).is_err());
120147
}
121148
}
122149

123150
thread::spawn(move || unsafe {
124-
K1.with(|s| *s.get() = Some(S1));
151+
K1.with(|s| *s.get() = Some(S1(&K1)));
125152
})
126153
.join()
127-
.ok()
128-
.expect("thread panicked");
154+
.unwrap();
155+
156+
thread::spawn(move || unsafe {
157+
K2.with(|s| *s.get() = Some(S1(&K2)));
158+
})
159+
.join()
160+
.unwrap();
129161
}
130162

131163
// Note that this test will deadlock if TLS destructors aren't run (this
@@ -152,3 +184,26 @@ fn dtors_in_dtors_in_dtors() {
152184
});
153185
rx.recv().unwrap();
154186
}
187+
188+
#[test]
189+
fn dtors_in_dtors_in_dtors_const_init() {
190+
struct S1(Sender<()>);
191+
thread_local!(static K1: UnsafeCell<Option<S1>> = const { UnsafeCell::new(None) });
192+
thread_local!(static K2: UnsafeCell<Option<Foo>> = const { UnsafeCell::new(None) });
193+
194+
impl Drop for S1 {
195+
fn drop(&mut self) {
196+
let S1(ref tx) = *self;
197+
unsafe {
198+
let _ = K2.try_with(|s| *s.get() = Some(Foo(tx.clone())));
199+
}
200+
}
201+
}
202+
203+
let (tx, rx) = channel();
204+
let _t = thread::spawn(move || unsafe {
205+
let mut tx = Some(tx);
206+
K1.with(|s| *s.get() = Some(S1(tx.take().unwrap())));
207+
});
208+
rx.recv().unwrap();
209+
}

‎std/src/thread/mod.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,13 @@ pub use self::local::os::Key as __OsLocalKeyInner;
204204
#[doc(hidden)]
205205
pub use self::local::statik::Key as __StaticLocalKeyInner;
206206

207+
// This is only used to make thread locals with `const { .. }` initialization
208+
// expressions unstable. If and/or when that syntax is stabilized with thread
209+
// locals this will simply be removed.
210+
#[doc(hidden)]
211+
#[unstable(feature = "thread_local_const_init", issue = "84223")]
212+
pub const fn require_unstable_const_init_thread_local() {}
213+
207214
////////////////////////////////////////////////////////////////////////////////
208215
// Builder
209216
////////////////////////////////////////////////////////////////////////////////

0 commit comments

Comments
 (0)