Skip to content

Commit

Permalink
std: make TLS accessors closures that return pointers
Browse files Browse the repository at this point in the history
  • Loading branch information
joboet committed May 24, 2024
1 parent 5af5f81 commit 078095a
Show file tree
Hide file tree
Showing 6 changed files with 107 additions and 180 deletions.
34 changes: 13 additions & 21 deletions std/src/sys/thread_local/fast_local/eager.rs
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 std/src/sys/thread_local/fast_local/lazy.rs
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 std/src/sys/thread_local/fast_local/mod.rs
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 std/src/sys/thread_local/os_local.rs
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 078095a

Please sign in to comment.