Skip to content

Commit

Permalink
Auto merge of #125525 - joboet:tls_accessor, r=cuviper
Browse files Browse the repository at this point in the history 
Make TLS accessors closures that return pointers

The current TLS macros generate a function that returns an `Option<&'static T>`. This is both risky as we lie about lifetimes, and necessitates that those functions are `unsafe`. By returning a `*const T` instead, the accessor function do not have safety requirements any longer and can be made closures without hassle. This PR does exactly that!

For native TLS, the closure approach makes it trivial to select the right accessor function at compile-time, which could result in a slight speed-up (I have the hope that the accessors are now simple enough for the MIR-inliner to kick in).
  • Loading branch information
@bors
bors committed Jun 4, 2024
2 parents 90d6255 + 6df8d0d commit 27529d5
Show file tree
Hide file tree
Showing 8 changed files with 128 additions and 271 deletions.
34 changes: 13 additions & 21 deletions library/std/src/sys/thread_local/fast_local/eager.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,43 +21,35 @@ impl<T> Storage<T> {
Storage { state: Cell::new(State::Initial), val: UnsafeCell::new(val) }
}

/// Get a reference to the TLS value. If the TLS variable has been destroyed,
/// `None` is returned.
/// Get a pointer to the TLS value. If the TLS variable has been destroyed,
/// a null pointer is returned.
///
/// # Safety
/// * The `self` reference must remain valid until the TLS destructor has been
/// run.
/// * The returned reference may only be used until thread destruction occurs
/// and may not be used after reentrant initialization has occurred.
/// The resulting pointer may not be used after thread destruction has
/// occurred.
///
// FIXME(#110897): return NonNull instead of lying about the lifetime.
/// # Safety
/// The `self` reference must remain valid until the TLS destructor is run.
#[inline]
pub unsafe fn get(&self) -> Option<&'static T> {
pub unsafe fn get(&self) -> *const T {
match self.state.get() {
// SAFETY: as the state is not `Destroyed`, the value cannot have
// been destroyed yet. The reference fulfills the terms outlined
// above.
State::Alive => unsafe { Some(&*self.val.get()) },
State::Destroyed => None,
State::Alive => self.val.get(),
State::Destroyed => ptr::null(),
State::Initial => unsafe { self.initialize() },
}
}

#[cold]
unsafe fn initialize(&self) -> Option<&'static T> {
unsafe fn initialize(&self) -> *const T {
// Register the destructor

// SAFETY:
// * the destructor will be called at thread destruction.
// * the caller guarantees that `self` will be valid until that time.
// The caller guarantees that `self` will be valid until thread destruction.
unsafe {
register_dtor(ptr::from_ref(self).cast_mut().cast(), destroy::<T>);
}

self.state.set(State::Alive);
// SAFETY: as the state is not `Destroyed`, the value cannot have
// been destroyed yet. The reference fulfills the terms outlined
// above.
unsafe { Some(&*self.val.get()) }
self.val.get()
}
}

Expand Down
46 changes: 13 additions & 33 deletions library/std/src/sys/thread_local/fast_local/lazy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,49 +39,31 @@ where
Storage { state: UnsafeCell::new(State::Initial) }
}

/// Get a reference to the TLS value, potentially initializing it with the
/// provided parameters. If the TLS variable has been destroyed, `None` is
/// returned.
/// Get a pointer to the TLS value, potentially initializing it with the
/// provided parameters. If the TLS variable has been destroyed, a null
/// pointer is returned.
///
/// # Safety
/// * The `self` reference must remain valid until the TLS destructor is run,
/// at which point the returned reference is invalidated.
/// * The returned reference may only be used until thread destruction occurs
/// and may not be used after reentrant initialization has occurred.
/// The resulting pointer may not be used after reentrant inialialization
/// or thread destruction has occurred.
///
// FIXME(#110897): return NonNull instead of lying about the lifetime.
/// # Safety
/// The `self` reference must remain valid until the TLS destructor is run.
#[inline]
pub unsafe fn get_or_init(
&self,
i: Option<&mut Option<T>>,
f: impl FnOnce() -> T,
) -> Option<&'static T> {
// SAFETY:
// No mutable reference to the inner value exists outside the calls to
// `replace`. The lifetime of the returned reference fulfills the terms
// outlined above.
pub unsafe fn get_or_init(&self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
let state = unsafe { &*self.state.get() };
match state {
State::Alive(v) => Some(v),
State::Destroyed(_) => None,
State::Alive(v) => v,
State::Destroyed(_) => ptr::null(),
State::Initial => unsafe { self.initialize(i, f) },
}
}

#[cold]
unsafe fn initialize(
&self,
i: Option<&mut Option<T>>,
f: impl FnOnce() -> T,
) -> Option<&'static T> {
unsafe fn initialize(&self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
// Perform initialization

let v = i.and_then(Option::take).unwrap_or_else(f);

// SAFETY:
// If references to the inner value exist, they were created in `f`
// and are invalidated here. The caller promises to never use them
// after this.
let old = unsafe { self.state.get().replace(State::Alive(v)) };
match old {
// If the variable is not being recursively initialized, register
Expand All @@ -92,12 +74,10 @@ where
val => drop(val),
}

// SAFETY:
// Initialization was completed and the state was set to `Alive`, so the
// reference fulfills the terms outlined above.
// SAFETY: the state was just set to `Alive`
unsafe {
let State::Alive(v) = &*self.state.get() else { unreachable_unchecked() };
Some(v)
v
}
}
}
Expand Down
79 changes: 34 additions & 45 deletions library/std/src/sys/thread_local/fast_local/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,32 +52,26 @@ pub macro thread_local_inner {
(@key $t:ty, const $init:expr) => {{
const __INIT: $t = $init;

#[inline]
#[deny(unsafe_op_in_unsafe_fn)]
unsafe fn __getit(
_init: $crate::option::Option<&mut $crate::option::Option<$t>>,
) -> $crate::option::Option<&'static $t> {
use $crate::thread::local_impl::EagerStorage;
unsafe {
use $crate::mem::needs_drop;
use $crate::ptr::addr_of;
use $crate::thread::LocalKey;
use $crate::thread::local_impl::EagerStorage;

if needs_drop::<$t>() {
#[thread_local]
static VAL: EagerStorage<$t> = EagerStorage::new(__INIT);
unsafe {
VAL.get()
LocalKey::new(const {
if needs_drop::<$t>() {
|_| {
#[thread_local]
static VAL: EagerStorage<$t> = EagerStorage::new(__INIT);
VAL.get()
}
} else {
|_| {
#[thread_local]
static VAL: $t = __INIT;
&VAL
}
}
} else {
#[thread_local]
static VAL: $t = __INIT;
unsafe {
$crate::option::Option::Some(&*addr_of!(VAL))
}
}
}

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

Expand All @@ -88,31 +82,26 @@ pub macro thread_local_inner {
$init
}

#[inline]
#[deny(unsafe_op_in_unsafe_fn)]
unsafe fn __getit(
init: $crate::option::Option<&mut $crate::option::Option<$t>>,
) -> $crate::option::Option<&'static $t> {
use $crate::thread::local_impl::LazyStorage;
unsafe {
use $crate::mem::needs_drop;
use $crate::thread::LocalKey;
use $crate::thread::local_impl::LazyStorage;

if needs_drop::<$t>() {
#[thread_local]
static VAL: LazyStorage<$t, ()> = LazyStorage::new();
unsafe {
VAL.get_or_init(init, __init)
LocalKey::new(const {
if needs_drop::<$t>() {
|init| {
#[thread_local]
static VAL: LazyStorage<$t, ()> = LazyStorage::new();
VAL.get_or_init(init, __init)
}
} else {
|init| {
#[thread_local]
static VAL: LazyStorage<$t, !> = LazyStorage::new();
VAL.get_or_init(init, __init)
}
}
} else {
#[thread_local]
static VAL: LazyStorage<$t, !> = LazyStorage::new();
unsafe {
VAL.get_or_init(init, __init)
}
}
}

unsafe {
$crate::thread::LocalKey::new(__getit)
})
}
}},
($(#[$attr:meta])* $vis:vis $name:ident, $t:ty, $($init:tt)*) => {
Expand Down
63 changes: 25 additions & 38 deletions library/std/src/sys/thread_local/os_local.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,30 +16,22 @@ pub macro thread_local_inner {
},

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

// `#[inline] does not work on windows-gnu due to linking errors around dllimports.
// See https://github.com/rust-lang/rust/issues/109797.
#[cfg_attr(not(windows), inline)]
unsafe fn __getit(
init: $crate::option::Option<&mut $crate::option::Option<$t>>,
) -> $crate::option::Option<&'static $t> {
use $crate::thread::local_impl::Key;

static __KEY: Key<$t> = Key::new();
unsafe {
__KEY.get(init, __init)
}
}
unsafe {
use $crate::thread::LocalKey;
use $crate::thread::local_impl::Key;

unsafe {
$crate::thread::LocalKey::new(__getit)
}
// Inlining does not work on windows-gnu due to linking errors around
// dllimports. See https://github.com/rust-lang/rust/issues/109797.
LocalKey::new(#[cfg_attr(windows, inline(never))] |init| {
static VAL: Key<$t> = Key::new();
VAL.get(init, __init)
})
}
},
}},
($(#[$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)*);
Expand Down Expand Up @@ -67,38 +59,33 @@ impl<T: 'static> Key<T> {
Key { os: OsKey::new(Some(destroy_value::<T>)), marker: PhantomData }
}

/// Get the value associated with this key, initializating it if necessary.
/// Get a pointer to the TLS value, potentially initializing it with the
/// provided parameters. If the TLS variable has been destroyed, a null
/// pointer is returned.
///
/// # Safety
/// * the returned reference must not be used after recursive initialization
/// or thread destruction occurs.
pub unsafe fn get(
&'static self,
i: Option<&mut Option<T>>,
f: impl FnOnce() -> T,
) -> Option<&'static T> {
/// The resulting pointer may not be used after reentrant inialialization
/// or thread destruction has occurred.
pub fn get(&'static self, i: Option<&mut Option<T>>, f: impl FnOnce() -> T) -> *const T {
// SAFETY: (FIXME: get should actually be safe)
let ptr = unsafe { self.os.get() as *mut Value<T> };
if ptr.addr() > 1 {
// SAFETY: the check ensured the pointer is safe (its destructor
// is not running) + it is coming from a trusted source (self).
unsafe { Some(&(*ptr).value) }
unsafe { &(*ptr).value }
} else {
// SAFETY: At this point we are sure we have no value and so
// initializing (or trying to) is safe.
unsafe { self.try_initialize(ptr, i, f) }
self.try_initialize(ptr, i, f)
}
}

unsafe fn try_initialize(
fn try_initialize(
&'static self,
ptr: *mut Value<T>,
i: Option<&mut Option<T>>,
f: impl FnOnce() -> T,
) -> Option<&'static T> {
) -> *const T {
if ptr.addr() == 1 {
// destructor is running
return None;
return ptr::null();
}

let value = i.and_then(Option::take).unwrap_or_else(f);
Expand All @@ -119,7 +106,7 @@ impl<T: 'static> Key<T> {
}

// SAFETY: We just created this value above.
unsafe { Some(&(*ptr).value) }
unsafe { &(*ptr).value }
}
}

Expand Down
Loading

0 comments on commit 27529d5

Please sign in to comment.