std: simplify key-based thread locals

Author: joboet

library/std/src/sys/thread_local/mod.rs

@@ -24,6 +24,8 @@ cfg_if::cfg_if! {
     }
 }
 
+// FIXME(#110897): remove this
+#[allow(unused)] // Not used by key-based TLS.
 mod lazy {
     use crate::cell::UnsafeCell;
     use crate::hint;

library/std/src/sys/thread_local/os_local.rs

@@ -1,7 +1,7 @@
-use super::lazy::LazyKeyInner;
 use crate::cell::Cell;
+use crate::marker::PhantomData;
 use crate::sys_common::thread_local_key::StaticKey as OsStaticKey;
-use crate::{fmt, marker, panic, ptr};
+use crate::{panic, ptr};
 
 #[doc(hidden)]
 #[allow_internal_unstable(thread_local_internals)]
@@ -10,38 +10,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) => {
@@ -85,78 +56,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()
-    }
+    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: OsStaticKey::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.
+    /// 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, init: impl FnOnce() -> T) -> Option<&'static T> {
-        // SAFETY: See the documentation for this method.
+        // 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, init) }
         }
-        // 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> };
+    /// Initialize an empty TLS key.
+    ///
+    /// # Safety
+    /// * the same requirements as for `get` apply
+    /// * `ptr` must be the current value of the key, and must be either null or the sentinel value
+    unsafe fn try_initialize(
+        &'static self,
+        ptr: *mut Value<T>,
+        init: 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 ptr = Box::into_raw(Box::new(Value { value: init(), 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) }
     }
 }