Skip to content

Commit

Permalink
Auto merge of rust-lang#122494 - joboet:simplify_key_tls, r=m-ou-se
Browse files Browse the repository at this point in the history
Simplify key-based thread locals

This PR simplifies key-based thread-locals by:
* unifying the macro expansion of `const` and non-`const` initializers
* reducing the amount of code in the expansion
* simply reallocating on recursive initialization instead of going through `LazyKeyInner`
* replacing `catch_unwind` with the shared `abort_on_dtor_unwind`

It does not change the initialization behaviour described in rust-lang#110897.
  • Loading branch information
bors committed May 24, 2024
2 parents 213ad10 + 0e7e75e commit 9e297bf
Show file tree
Hide file tree
Showing 4 changed files with 63 additions and 189 deletions.
2 changes: 0 additions & 2 deletions library/std/src/sys/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,6 @@ pub mod cmath;
pub mod os_str;
pub mod path;
pub mod sync;
#[allow(dead_code)]
#[allow(unused_imports)]
pub mod thread_local;

// FIXME(117276): remove this, move feature implementations into individual
Expand Down
1 change: 0 additions & 1 deletion library/std/src/sys/thread_local/fast_local/lazy.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use crate::cell::UnsafeCell;
use crate::hint::unreachable_unchecked;
use crate::mem::forget;
use crate::ptr;
use crate::sys::thread_local::abort_on_dtor_unwind;
use crate::sys::thread_local_dtor::register_dtor;
Expand Down
87 changes: 3 additions & 84 deletions library/std/src/sys/thread_local/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#![unstable(feature = "thread_local_internals", reason = "should not be necessary", issue = "none")]
#![cfg_attr(test, allow(unused))]

// There are three thread-local implementations: "static", "fast", "OS".
// The "OS" thread local key type is accessed via platform-specific API calls and is slow, while the
Expand All @@ -24,94 +25,12 @@ cfg_if::cfg_if! {
}
}

// Not used by the fast-local TLS anymore.
// FIXME(#110897): remove this.
#[allow(unused)]
mod lazy {
use crate::cell::UnsafeCell;
use crate::hint;
use crate::mem;

pub struct LazyKeyInner<T> {
inner: UnsafeCell<Option<T>>,
}

impl<T> LazyKeyInner<T> {
pub const fn new() -> LazyKeyInner<T> {
LazyKeyInner { inner: UnsafeCell::new(None) }
}

pub unsafe fn get(&self) -> Option<&'static T> {
// SAFETY: The caller must ensure no reference is ever handed out to
// the inner cell nor mutable reference to the Option<T> inside said
// cell. This make it safe to hand a reference, though the lifetime
// of 'static is itself unsafe, making the get method unsafe.
unsafe { (*self.inner.get()).as_ref() }
}

/// The caller must ensure that no reference is active: this method
/// needs unique access.
pub unsafe fn initialize<F: FnOnce() -> T>(&self, init: F) -> &'static T {
// Execute the initialization up front, *then* move it into our slot,
// just in case initialization fails.
let value = init();
let ptr = self.inner.get();

// SAFETY:
//
// note that this can in theory just be `*ptr = Some(value)`, but due to
// the compiler will currently codegen that pattern with something like:
//
// ptr::drop_in_place(ptr)
// ptr::write(ptr, Some(value))
//
// Due to this pattern it's possible for the destructor of the value in
// `ptr` (e.g., if this is being recursively initialized) to re-access
// TLS, in which case there will be a `&` and `&mut` pointer to the same
// value (an aliasing violation). To avoid setting the "I'm running a
// destructor" flag we just use `mem::replace` which should sequence the
// operations a little differently and make this safe to call.
//
// The precondition also ensures that we are the only one accessing
// `self` at the moment so replacing is fine.
unsafe {
let _ = mem::replace(&mut *ptr, Some(value));
}

// SAFETY: With the call to `mem::replace` it is guaranteed there is
// a `Some` behind `ptr`, not a `None` so `unreachable_unchecked`
// will never be reached.
unsafe {
// After storing `Some` we want to get a reference to the contents of
// what we just stored. While we could use `unwrap` here and it should
// always work it empirically doesn't seem to always get optimized away,
// which means that using something like `try_with` can pull in
// panicking code and cause a large size bloat.
match *ptr {
Some(ref x) => x,
None => hint::unreachable_unchecked(),
}
}
}

/// Watch out: unsynchronized internal mutability!
///
/// # Safety
/// Causes UB if any reference to the value is used after this.
#[allow(unused)]
pub(crate) unsafe fn take(&self) -> Option<T> {
let mutable: *mut _ = UnsafeCell::get(&self.inner);
// SAFETY: That's the caller's problem.
unsafe { mutable.replace(None) }
}
}
}

/// Run a callback in a scenario which must not unwind (such as a `extern "C"
/// fn` declared in a user crate). If the callback unwinds anyway, then
/// `rtabort` with a message about thread local panicking on drop.
#[inline]
pub fn abort_on_dtor_unwind(f: impl FnOnce()) {
#[allow(dead_code)]
fn abort_on_dtor_unwind(f: impl FnOnce()) {
// Using a guard like this is lower cost.
let guard = DtorUnwindGuard;
f();
Expand Down
162 changes: 60 additions & 102 deletions library/std/src/sys/thread_local/os_local.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
use super::lazy::LazyKeyInner;
use super::abort_on_dtor_unwind;
use crate::cell::Cell;
use crate::sys_common::thread_local_key::StaticKey as OsStaticKey;
use crate::{fmt, marker, panic, ptr};
use crate::marker::PhantomData;
use crate::ptr;
use crate::sys_common::thread_local_key::StaticKey as OsKey;

#[doc(hidden)]
#[allow_internal_unstable(thread_local_internals)]
Expand All @@ -10,38 +11,9 @@ use crate::{fmt, marker, panic, ptr};
#[rustc_macro_transparency = "semitransparent"]
pub macro thread_local_inner {
// used to generate the `LocalKey` value for const-initialized thread locals
(@key $t:ty, const $init:expr) => {{
#[inline]
#[deny(unsafe_op_in_unsafe_fn)]
unsafe fn __getit(
_init: $crate::option::Option<&mut $crate::option::Option<$t>>,
) -> $crate::option::Option<&'static $t> {
const INIT_EXPR: $t = $init;

// On platforms without `#[thread_local]` we fall back to the
// same implementation as below for os thread locals.
#[inline]
const fn __init() -> $t { INIT_EXPR }
static __KEY: $crate::thread::local_impl::Key<$t> =
$crate::thread::local_impl::Key::new();
unsafe {
__KEY.get(move || {
if let $crate::option::Option::Some(init) = _init {
if let $crate::option::Option::Some(value) = init.take() {
return value;
} else if $crate::cfg!(debug_assertions) {
$crate::unreachable!("missing initial value");
}
}
__init()
})
}
}

unsafe {
$crate::thread::LocalKey::new(__getit)
}
}},
(@key $t:ty, const $init:expr) => {
$crate::thread::local_impl::thread_local_inner!(@key $t, { const INIT_EXPR: $t = $init; INIT_EXPR })
},

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

static __KEY: Key<$t> = Key::new();
unsafe {
__KEY.get(move || {
if let $crate::option::Option::Some(init) = init {
if let $crate::option::Option::Some(value) = init.take() {
return value;
} else if $crate::cfg!(debug_assertions) {
$crate::unreachable!("missing default value");
}
}
__init()
})
__KEY.get(init, __init)
}
}

Expand All @@ -85,78 +48,78 @@ pub macro thread_local_inner {

/// Use a regular global static to store this key; the state provided will then be
/// thread-local.
#[allow(missing_debug_implementations)]
pub struct Key<T> {
// OS-TLS key that we'll use to key off.
os: OsStaticKey,
marker: marker::PhantomData<Cell<T>>,
}

impl<T> fmt::Debug for Key<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Key").finish_non_exhaustive()
}
os: OsKey,
marker: PhantomData<Cell<T>>,
}

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

struct Value<T: 'static> {
inner: LazyKeyInner<T>,
value: T,
key: &'static Key<T>,
}

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

/// It is a requirement for the caller to ensure that no mutable
/// reference is active when this method is called.
pub unsafe fn get(&'static self, init: impl FnOnce() -> T) -> Option<&'static T> {
// SAFETY: See the documentation for this method.
/// Get the value associated with this key, initializating it if necessary.
///
/// # 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> {
// 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).
if let Some(ref value) = unsafe { (*ptr).inner.get() } {
return Some(value);
}
unsafe { Some(&(*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) }
}
// SAFETY: At this point we are sure we have no value and so
// initializing (or trying to) is safe.
unsafe { self.try_initialize(init) }
}

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

let ptr = if ptr.is_null() {
// If the lookup returned null, we haven't initialized our own
// local copy, so do that now.
let ptr = Box::into_raw(Box::new(Value { inner: LazyKeyInner::new(), key: self }));
// SAFETY: At this point we are sure there is no value inside
// ptr so setting it will not affect anyone else.
unsafe {
self.os.set(ptr as *mut u8);
}
ptr
} else {
// recursive initialization
ptr
};
let value = i.and_then(Option::take).unwrap_or_else(f);
let ptr = Box::into_raw(Box::new(Value { value, key: self }));
// SAFETY: (FIXME: get should actually be safe)
let old = unsafe { self.os.get() as *mut Value<T> };
// SAFETY: `ptr` is a correct pointer that can be destroyed by the key destructor.
unsafe {
self.os.set(ptr as *mut u8);
}
if !old.is_null() {
// If the variable was recursively initialized, drop the old value.
// SAFETY: We cannot be inside a `LocalKey::with` scope, as the
// initializer has already returned and the next scope only starts
// after we return the pointer. Therefore, there can be no references
// to the old value.
drop(unsafe { Box::from_raw(old) });
}

// SAFETY: ptr has been ensured as non-NUL just above an so can be
// dereferenced safely.
unsafe { Some((*ptr).inner.initialize(init)) }
// SAFETY: We just created this value above.
unsafe { Some(&(*ptr).value) }
}
}

Expand All @@ -170,16 +133,11 @@ unsafe extern "C" fn destroy_value<T: 'static>(ptr: *mut u8) {
//
// Note that to prevent an infinite loop we reset it back to null right
// before we return from the destructor ourselves.
//
// Wrap the call in a catch to ensure unwinding is caught in the event
// a panic takes place in a destructor.
if let Err(_) = panic::catch_unwind(|| unsafe {
let ptr = Box::from_raw(ptr as *mut Value<T>);
abort_on_dtor_unwind(|| {
let ptr = unsafe { Box::from_raw(ptr as *mut Value<T>) };
let key = ptr.key;
key.os.set(ptr::without_provenance_mut(1));
unsafe { key.os.set(ptr::without_provenance_mut(1)) };
drop(ptr);
key.os.set(ptr::null_mut());
}) {
rtabort!("thread local panicked on drop");
}
unsafe { key.os.set(ptr::null_mut()) };
});
}

0 comments on commit 9e297bf

Please sign in to comment.