Auto merge of #122494 - joboet:simplify_key_tls, r=m-ou-se

bors · bors · commit 9e297bf54d31 · 2024-05-24T15:34:07.000Z
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 #110897.
diff --git a/library/std/src/sys/mod.rs b/library/std/src/sys/mod.rs
@@ -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
diff --git a/library/std/src/sys/thread_local/fast_local/lazy.rs b/library/std/src/sys/thread_local/fast_local/lazy.rs
@@ -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;
diff --git a/library/std/src/sys/thread_local/mod.rs b/library/std/src/sys/thread_local/mod.rs
@@ -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
@@ -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();
diff --git a/library/std/src/sys/thread_local/os_local.rs b/library/std/src/sys/thread_local/os_local.rs
@@ -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)]
@@ -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) => {
@@ -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)
                 }
             }
 
@@ -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) }
     }
 }
 
@@ -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()) };
+    });
 }
