std: make thread::current available in all thread_local! destructors

Author: joboet

library/std/src/rt.rs

@@ -19,9 +19,11 @@
 pub use crate::panicking::{begin_panic, panic_count};
 pub use core::panicking::{panic_display, panic_fmt};
 
+use crate::any::Any;
 use crate::sync::Once;
 use crate::sys;
 use crate::thread::{self, Thread};
+use crate::{mem, panic};
 
 // Prints to the "panic output", depending on the platform this may be:
 // - the standard error output
@@ -64,6 +66,11 @@ macro_rules! rtunwrap {
     };
 }
 
+fn handle_rt_panic(e: Box<dyn Any + Send>) {
+    mem::forget(e);
+    rtabort!("initialization or cleanup bug");
+}
+
 // One-time runtime initialization.
 // Runs before `main`.
 // SAFETY: must be called only once during runtime initialization.
@@ -99,6 +106,20 @@ unsafe fn init(argc: isize, argv: *const *const u8, sigpipe: u8) {
     thread::set_current(thread);
 }
 
+/// Clean up the thread-local runtime state. This *should* be run after all other
+/// code managed by the Rust runtime, but will not cause UB if that condition is
+/// not fulfilled. Also note that this function is not guaranteed to be run, but
+/// skipping it will cause leaks and therefore is to be avoided.
+pub(crate) fn thread_cleanup() {
+    // This function is run in situations where unwinding leads to an abort
+    // (think `extern "C"` functions). Abort here instead so that we can
+    // print a nice message.
+    panic::catch_unwind(|| {
+        crate::thread::drop_current();
+    })
+    .unwrap_or_else(handle_rt_panic);
+}
+
 // One-time runtime cleanup.
 // Runs after `main` or at program exit.
 // NOTE: this is not guaranteed to run, for example when the program aborts.
@@ -121,11 +142,6 @@ fn lang_start_internal(
     argv: *const *const u8,
     sigpipe: u8,
 ) -> Result<isize, !> {
-    use crate::{mem, panic};
-    let rt_abort = move |e| {
-        mem::forget(e);
-        rtabort!("initialization or cleanup bug");
-    };
     // Guard against the code called by this function from unwinding outside of the Rust-controlled
     // code, which is UB. This is a requirement imposed by a combination of how the
     // `#[lang="start"]` attribute is implemented as well as by the implementation of the panicking
@@ -137,16 +153,17 @@ fn lang_start_internal(
     // prevent std from accidentally introducing a panic to these functions. Another is from
     // user code from `main` or, more nefariously, as described in e.g. issue #86030.
     // SAFETY: Only called once during runtime initialization.
-    panic::catch_unwind(move || unsafe { init(argc, argv, sigpipe) }).map_err(rt_abort)?;
+    panic::catch_unwind(move || unsafe { init(argc, argv, sigpipe) })
+        .unwrap_or_else(handle_rt_panic);
     let ret_code = panic::catch_unwind(move || panic::catch_unwind(main).unwrap_or(101) as isize)
         .map_err(move |e| {
             mem::forget(e);
             rtabort!("drop of the panic payload panicked");
         });
-    panic::catch_unwind(cleanup).map_err(rt_abort)?;
+    panic::catch_unwind(cleanup).unwrap_or_else(handle_rt_panic);
     // Guard against multple threads calling `libc::exit` concurrently.
     // See the documentation for `unique_thread_exit` for more information.
-    panic::catch_unwind(|| crate::sys::exit_guard::unique_thread_exit()).map_err(rt_abort)?;
+    panic::catch_unwind(crate::sys::exit_guard::unique_thread_exit).unwrap_or_else(handle_rt_panic);
     ret_code
 }
 

library/std/src/sys/pal/hermit/mod.rs

@@ -104,6 +104,8 @@ pub unsafe extern "C" fn runtime_entry(
     let result = main(argc as isize, argv);
 
     crate::sys::thread_local::destructors::run();
+    crate::rt::thread_cleanup();
+
     hermit_abi::exit(result);
 }
 

library/std/src/sys/pal/hermit/thread.rs

@@ -50,6 +50,7 @@ impl Thread {
 
                 // run all destructors
                 crate::sys::thread_local::destructors::run();
+                crate::rt::thread_cleanup();
             }
         }
     }

library/std/src/sys/sync/rwlock/queue.rs

@@ -115,7 +115,7 @@ use crate::sync::atomic::{
     AtomicBool, AtomicPtr,
     Ordering::{AcqRel, Acquire, Relaxed, Release},
 };
-use crate::thread::{self, Thread};
+use crate::thread::{self, Thread, ThreadId};
 
 // Locking uses exponential backoff. `SPIN_COUNT` indicates how many times the
 // locking operation will be retried.
@@ -202,7 +202,9 @@ impl Node {
     fn prepare(&mut self) {
         // Fall back to creating an unnamed `Thread` handle to allow locking in
         // TLS destructors.
-        self.thread.get_or_init(|| thread::try_current().unwrap_or_else(Thread::new_unnamed));
+        self.thread.get_or_init(|| {
+            thread::try_current().unwrap_or_else(|| Thread::new_unnamed(ThreadId::new()))
+        });
         self.completed = AtomicBool::new(false);
     }
 

library/std/src/sys/thread_local/guard/apple.rs

@@ -26,6 +26,7 @@ pub fn enable() {
     unsafe extern "C" fn run_dtors(_: *mut u8) {
         unsafe {
             destructors::run();
+            crate::rt::thread_cleanup();
         }
     }
 }

library/std/src/sys/thread_local/guard/key.rs

@@ -3,10 +3,12 @@
 //! that will run all native TLS destructors in the destructor list.
 
 use crate::ptr;
-use crate::sys::thread_local::destructors;
 use crate::sys::thread_local::key::{set, LazyKey};
 
+#[cfg(target_thread_local)]
 pub fn enable() {
+    use crate::sys::thread_local::destructors;
+
     static DTORS: LazyKey = LazyKey::new(Some(run));
 
     // Setting the key value to something other than NULL will result in the
@@ -18,6 +20,41 @@ pub fn enable() {
     unsafe extern "C" fn run(_: *mut u8) {
         unsafe {
             destructors::run();
+            // On platforms with `__cxa_thread_atexit_impl`, `destructors::run`
+            // does nothing on newer systems as the TLS destructors are
+            // registered with the system. But because all of those platforms
+            // call the destructors of TLS keys after the registered ones, this
+            // function will still be run last (at the time of writing).
+            crate::rt::thread_cleanup();
+        }
+    }
+}
+
+/// On platforms with key-based TLS, the system runs the destructors for us.
+/// We still have to make sure that [`crate::rt::thread_cleanup`] is called,
+/// however. This is done by defering the execution of a TLS destructor to
+/// the next round of destruction inside the TLS destructors.
+#[cfg(not(target_thread_local))]
+pub fn enable() {
+    const DEFER: *mut u8 = ptr::without_provenance_mut(1);
+    const RUN: *mut u8 = ptr::without_provenance_mut(2);
+
+    static CLEANUP: LazyKey = LazyKey::new(Some(run));
+
+    unsafe { set(CLEANUP.force(), DEFER) }
+
+    unsafe extern "C" fn run(state: *mut u8) {
+        if state == DEFER {
+            // Make sure that this function is run again in the next round of
+            // TLS destruction. If there is no futher round, there will be leaks,
+            // but that's okay, `thread_cleanup` is not guaranteed to be called.
+            unsafe { set(CLEANUP.force(), RUN) }
+        } else {
+            debug_assert_eq!(state, RUN);
+            // If the state is still RUN in the next round of TLS destruction,
+            // it means that no other TLS destructors defined by this runtime
+            // have been run, as they would have set the state to DEFER.
+            crate::rt::thread_cleanup();
         }
     }
 }

library/std/src/sys/thread_local/guard/solid.rs

@@ -18,6 +18,9 @@ pub fn enable() {
     }
 
     unsafe extern "C" fn tls_dtor(_unused: *mut u8) {
-        unsafe { destructors::run() };
+        unsafe {
+            destructors::run();
+            crate::rt::thread_cleanup();
+        }
     }
 }

library/std/src/sys/thread_local/guard/windows.rs

@@ -92,13 +92,13 @@ unsafe extern "system" fn tls_callback(_h: *mut c_void, dw_reason: u32, _pv: *mu
     }
 
     if dw_reason == c::DLL_THREAD_DETACH || dw_reason == c::DLL_PROCESS_DETACH {
-        #[cfg(target_thread_local)]
         unsafe {
+            #[cfg(target_thread_local)]
             super::super::destructors::run();
-        }
-        #[cfg(not(target_thread_local))]
-        unsafe {
+            #[cfg(not(target_thread_local))]
             super::super::key::run_dtors();
+
+            crate::rt::thread_cleanup();
         }
     }
 }

library/std/src/sys/thread_local/key/xous.rs

@@ -213,4 +213,6 @@ unsafe fn run_dtors() {
             unsafe { cur = (*cur).next };
         }
     }
+
+    crate::rt::thread_cleanup();
 }

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

@@ -31,12 +31,15 @@ cfg_if::cfg_if! {
     ))] {
         mod statik;
         pub use statik::{EagerStorage, LazyStorage, thread_local_inner};
+        pub(crate) use statik::{LocalPointer, local_pointer};
     } else if #[cfg(target_thread_local)] {
         mod native;
         pub use native::{EagerStorage, LazyStorage, thread_local_inner};
+        pub(crate) use native::{LocalPointer, local_pointer};
     } else {
         mod os;
         pub use os::{Storage, thread_local_inner};
+        pub(crate) use os::{LocalPointer, local_pointer};
     }
 }
 
@@ -72,36 +75,47 @@ pub(crate) mod destructors {
 }
 
 /// This module provides a way to schedule the execution of the destructor list
-/// on systems without a per-variable destructor system.
-mod guard {
+/// and the [runtime cleanup](crate::rt::thread_cleanup) function. Calling `enable`
+/// should ensure that these functions are called at the right times.
+pub(crate) mod guard {
     cfg_if::cfg_if! {
         if #[cfg(all(target_thread_local, target_vendor = "apple"))] {
             mod apple;
-            pub(super) use apple::enable;
+            pub(crate) use apple::enable;
         } else if #[cfg(target_os = "windows")] {
             mod windows;
-            pub(super) use windows::enable;
+            pub(crate) use windows::enable;
         } else if #[cfg(any(
-            all(target_family = "wasm", target_feature = "atomics"),
+            target_family = "wasm",
+            target_os = "uefi",
+            target_os = "zkvm",
         ))] {
-            pub(super) fn enable() {
-                // FIXME: Right now there is no concept of "thread exit", but
-                // this is likely going to show up at some point in the form of
-                // an exported symbol that the wasm runtime is going to be
-                // expected to call. For now we just leak everything, but if
-                // such a function starts to exist it will probably need to
+            pub(crate) fn enable() {
+                // FIXME: Right now there is no concept of "thread exit" on
+                // wasm, but this is likely going to show up at some point in
+                // the form of an exported symbol that the wasm runtime is going
+                // to be expected to call. For now we just leak everything, but
+                // if such a function starts to exist it will probably need to
                 // iterate the destructor list with this function:
                 #[allow(unused)]
+                #[cfg(not(all(target_family = "wasm", not(target_feature = "atomics"))))]
                 use super::destructors::run;
+                #[allow(unused)]
+                use crate::rt::thread_cleanup;
             }
-        } else if #[cfg(target_os = "hermit")] {
-            pub(super) fn enable() {}
+        } else if #[cfg(any(
+            target_os = "hermit",
+            target_os = "xous",
+        ))] {
+            // `std` is the only runtime, so it just calls the destructor functions
+            // itself when the time comes.
+            pub(crate) fn enable() {}
         } else if #[cfg(target_os = "solid_asp3")] {
             mod solid;
-            pub(super) use solid::enable;
-        } else if #[cfg(all(target_thread_local, not(target_family = "wasm")))] {
+            pub(crate) use solid::enable;
+        } else {
             mod key;
-            pub(super) use key::enable;
+            pub(crate) use key::enable;
         }
     }
 }

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

@@ -32,6 +32,9 @@
 mod eager;
 mod lazy;
 
+use crate::cell::Cell;
+use crate::ptr;
+
 pub use eager::Storage as EagerStorage;
 pub use lazy::Storage as LazyStorage;
 
@@ -107,3 +110,31 @@ pub macro thread_local_inner {
             $crate::thread::local_impl::thread_local_inner!(@key $t, $($init)*);
     },
 }
+
+#[rustc_macro_transparency = "semitransparent"]
+pub(crate) macro local_pointer {
+    () => {},
+    ($vis:vis static $name:ident; $($rest:tt)*) => {
+        #[thread_local]
+        $vis static $name: $crate::sys::thread_local::LocalPointer = $crate::sys::thread_local::LocalPointer::__new();
+        $crate::sys::thread_local::local_pointer! { $($rest)* }
+    },
+}
+
+pub(crate) struct LocalPointer {
+    p: Cell<*mut ()>,
+}
+
+impl LocalPointer {
+    pub const fn __new() -> LocalPointer {
+        LocalPointer { p: Cell::new(ptr::null_mut()) }
+    }
+
+    pub fn get(&self) -> *mut () {
+        self.p.get()
+    }
+
+    pub fn set(&self, p: *mut ()) {
+        self.p.set(p)
+    }
+}

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

@@ -1,8 +1,9 @@
 use super::abort_on_dtor_unwind;
+use super::guard;
+use super::key::{get, set, Key, LazyKey};
 use crate::cell::Cell;
 use crate::marker::PhantomData;
 use crate::ptr;
-use crate::sys::thread_local::key::{get, set, Key, LazyKey};
 
 #[doc(hidden)]
 #[allow_internal_unstable(thread_local_internals)]
@@ -138,5 +139,35 @@ unsafe extern "C" fn destroy_value<T: 'static>(ptr: *mut u8) {
         drop(ptr);
         // SAFETY: `key` is the TLS key `ptr` was stored under.
         unsafe { set(key, ptr::null_mut()) };
+        // Make sure that the runtime cleanup will be performed
+        // after the next round of TLS destruction.
+        guard::enable();
     });
 }
+
+#[rustc_macro_transparency = "semitransparent"]
+pub(crate) macro local_pointer {
+    () => {},
+    ($vis:vis static $name:ident; $($rest:tt)*) => {
+        $vis static $name: $crate::sys::thread_local::LocalPointer = $crate::sys::thread_local::LocalPointer::__new();
+        $crate::sys::thread_local::local_pointer! { $($rest)* }
+    },
+}
+
+pub(crate) struct LocalPointer {
+    key: LazyKey,
+}
+
+impl LocalPointer {
+    pub const fn __new() -> LocalPointer {
+        LocalPointer { key: LazyKey::new(None) }
+    }
+
+    pub fn get(&'static self) -> *mut () {
+        unsafe { get(self.key.force()) as *mut () }
+    }
+
+    pub fn set(&'static self, p: *mut ()) {
+        unsafe { set(self.key.force(), p as *mut u8) }
+    }
+}

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

@@ -1,7 +1,8 @@
 //! On some targets like wasm there's no threads, so no need to generate
 //! thread locals and we can instead just use plain statics!
 
-use crate::cell::UnsafeCell;
+use crate::cell::{Cell, UnsafeCell};
+use crate::ptr;
 
 #[doc(hidden)]
 #[allow_internal_unstable(thread_local_internals)]
@@ -93,3 +94,33 @@ impl<T> LazyStorage<T> {
 
 // SAFETY: the target doesn't have threads.
 unsafe impl<T> Sync for LazyStorage<T> {}
+
+#[rustc_macro_transparency = "semitransparent"]
+pub(crate) macro local_pointer {
+    () => {},
+    ($vis:vis static $name:ident; $($rest:tt)*) => {
+        $vis static $name: $crate::sys::thread_local::LocalPointer = $crate::sys::thread_local::LocalPointer::__new();
+        $crate::sys::thread_local::local_pointer! { $($rest)* }
+    },
+}
+
+pub(crate) struct LocalPointer {
+    p: Cell<*mut ()>,
+}
+
+impl LocalPointer {
+    pub const fn __new() -> LocalPointer {
+        LocalPointer { p: Cell::new(ptr::null_mut()) }
+    }
+
+    pub fn get(&self) -> *mut () {
+        self.p.get()
+    }
+
+    pub fn set(&self, p: *mut ()) {
+        self.p.set(p)
+    }
+}
+
+// SAFETY: the target doesn't have threads.
+unsafe impl Sync for LocalPointer {}

library/std/src/thread/current.rs

@@ -0,0 +1,252 @@
+use super::{Thread, ThreadId};
+use crate::mem::ManuallyDrop;
+use crate::ptr;
+use crate::sys::thread_local::local_pointer;
+
+const NONE: *mut () = ptr::null_mut();
+const BUSY: *mut () = ptr::without_provenance_mut(1);
+const DESTROYED: *mut () = ptr::without_provenance_mut(2);
+
+local_pointer! {
+    static CURRENT;
+}
+
+/// Persistent storage for the thread ID.
+///
+/// We store the thread ID so that it never gets destroyed during the lifetime
+/// of a thread, either using `#[thread_local]` or multiple `local_pointer!`s.
+mod id {
+    use super::*;
+
+    cfg_if::cfg_if! {
+        if #[cfg(target_thread_local)] {
+            use crate::cell::Cell;
+
+            #[thread_local]
+            static ID: Cell<Option<ThreadId>> = Cell::new(None);
+
+            pub(super) const CHEAP: bool = true;
+
+            pub(super) fn get() -> Option<ThreadId> {
+                ID.get()
+            }
+
+            pub(super) fn set(id: ThreadId) {
+                ID.set(Some(id))
+            }
+        } else if #[cfg(target_pointer_width = "16")] {
+            local_pointer! {
+                static ID0;
+                static ID16;
+                static ID32;
+                static ID48;
+            }
+
+            pub(super) const CHEAP: bool = false;
+
+            pub(super) fn get() -> Option<ThreadId> {
+                let id0 = ID0.get().addr() as u64;
+                let id16 = ID16.get().addr() as u64;
+                let id32 = ID32.get().addr() as u64;
+                let id48 = ID48.get().addr() as u64;
+                ThreadId::from_u64((id48 << 48) + (id32 << 32) + (id16 << 16) + id0)
+            }
+
+            pub(super) fn set(id: ThreadId) {
+                let val = id.as_u64().get();
+                ID0.set(ptr::without_provenance_mut(val as usize));
+                ID16.set(ptr::without_provenance_mut((val >> 16) as usize));
+                ID32.set(ptr::without_provenance_mut((val >> 32) as usize));
+                ID48.set(ptr::without_provenance_mut((val >> 48) as usize));
+            }
+        } else if #[cfg(target_pointer_width = "16")] {
+            local_pointer! {
+                static ID0;
+                static ID32;
+            }
+
+            pub(super) const CHEAP: bool = false;
+
+            pub(super) fn get() -> Option<ThreadId> {
+                let id0 = ID0.get().addr() as u64;
+                let id32 = ID32.get().addr() as u64;
+                ThreadId::from_u64((id32 << 32) + id0)
+            }
+
+            pub(super) fn set(id: ThreadId) {
+                let val = id.as_u64().get();
+                ID0.set(ptr::without_provenance_mut(val as usize));
+                ID32.set(ptr::without_provenance_mut((val >> 32) as usize));
+            }
+        } else {
+            local_pointer! {
+                static ID;
+            }
+
+            pub(super) const CHEAP: bool = true;
+
+            pub(super) fn get() -> Option<ThreadId> {
+                let id = ID.get().addr() as u64;
+                ThreadId::from_u64(id)
+            }
+
+            pub(super) fn set(id: ThreadId) {
+                let val = id.as_u64().get();
+                ID.set(ptr::without_provenance_mut(val as usize));
+            }
+        }
+    }
+
+    #[inline]
+    pub(super) fn get_or_init() -> ThreadId {
+        get().unwrap_or_else(
+            #[cold]
+            || {
+                let id = ThreadId::new();
+                id::set(id);
+                id
+            },
+        )
+    }
+}
+
+/// Sets the thread handle for the current thread.
+///
+/// Aborts if the handle or the ID has been set already.
+pub(crate) fn set_current(thread: Thread) {
+    if CURRENT.get() != NONE || id::get().is_some() {
+        // Using `panic` here can add ~3kB to the binary size. We have complete
+        // control over where this is called, so just abort if there is a bug.
+        rtabort!("thread::set_current should only be called once per thread");
+    }
+
+    id::set(thread.id());
+
+    // Make sure that `crate::rt::thread_cleanup` will be run, which will
+    // call `drop_current`.
+    crate::sys::thread_local::guard::enable();
+    CURRENT.set(thread.into_raw());
+}
+
+/// Gets the id of the thread that invokes it.
+///
+/// This function will always succeed, will always return the same value for
+/// one thread and is guaranteed not to call the global allocator.
+#[inline]
+pub(crate) fn current_id() -> ThreadId {
+    // If accessing the persistant thread ID takes multiple TLS accesses, try
+    // to retrieve it from the current thread handle, which will only take one
+    // TLS access.
+    if !id::CHEAP {
+        let current = CURRENT.get();
+        if current > DESTROYED {
+            unsafe {
+                let current = ManuallyDrop::new(Thread::from_raw(current));
+                return current.id();
+            }
+        }
+    }
+
+    id::get_or_init()
+}
+
+/// Gets a handle to the thread that invokes it, if the handle has been initialized.
+pub(crate) fn try_current() -> Option<Thread> {
+    let current = CURRENT.get();
+    if current > DESTROYED {
+        unsafe {
+            let current = ManuallyDrop::new(Thread::from_raw(current));
+            Some((*current).clone())
+        }
+    } else {
+        None
+    }
+}
+
+/// Gets a handle to the thread that invokes it.
+///
+/// # Examples
+///
+/// Getting a handle to the current thread with `thread::current()`:
+///
+/// ```
+/// use std::thread;
+///
+/// let handler = thread::Builder::new()
+///     .name("named thread".into())
+///     .spawn(|| {
+///         let handle = thread::current();
+///         assert_eq!(handle.name(), Some("named thread"));
+///     })
+///     .unwrap();
+///
+/// handler.join().unwrap();
+/// ```
+#[must_use]
+#[stable(feature = "rust1", since = "1.0.0")]
+pub fn current() -> Thread {
+    let current = CURRENT.get();
+    if current > DESTROYED {
+        unsafe {
+            let current = ManuallyDrop::new(Thread::from_raw(current));
+            (*current).clone()
+        }
+    } else {
+        init_current(current)
+    }
+}
+
+#[cold]
+fn init_current(current: *mut ()) -> Thread {
+    if current == NONE {
+        CURRENT.set(BUSY);
+        // If the thread ID was initialized already, use it.
+        let id = id::get_or_init();
+        let thread = Thread::new_unnamed(id);
+
+        // Make sure that `crate::rt::thread_cleanup` will be run, which will
+        // call `drop_current`.
+        crate::sys::thread_local::guard::enable();
+        CURRENT.set(thread.clone().into_raw());
+        thread
+    } else if current == BUSY {
+        // BUSY exists solely for this check, but as it is in the slow path, the
+        // extra TLS write above shouldn't matter. The alternative is nearly always
+        // a stack overflow.
+
+        // If you came across this message, contact the author of your allocator.
+        // If you are said author: A surprising amount of functions inside the
+        // standard library (e.g. `Mutex`, `thread_local!`, `File` when using long
+        // paths, even `panic!` when using unwinding), need memory allocation, so
+        // you'll get circular dependencies all over the place when using them.
+        // I (joboet) highly recommend using only APIs from core in your allocator
+        // and implementing your own system abstractions. Still, if you feel that
+        // a particular API should be entirely allocation-free, feel free to open
+        // an issue on the Rust repository, we'll see what we can do.
+        rtabort!(
+            "\n
+            Attempted to access thread-local data while allocating said data.\n
+            Do not access functions that allocate in the global allocator!\n
+            This is a bug in the global allocator.\n
+        "
+        )
+    } else {
+        debug_assert_eq!(current, DESTROYED);
+        panic!(
+            "use of std::thread::current() is not possible after the thread's
+         local data has been destroyed"
+        )
+    }
+}
+
+/// This should be run in [`crate::rt::thread_cleanup`] to reset the thread
+/// handle.
+pub(crate) fn drop_current() {
+    let current = CURRENT.get();
+    if current > DESTROYED {
+        unsafe {
+            CURRENT.set(DESTROYED);
+            drop(Thread::from_raw(current));
+        }
+    }
+}

library/std/src/thread/local/tests.rs

@@ -1,7 +1,7 @@
 use crate::cell::{Cell, UnsafeCell};
 use crate::sync::atomic::{AtomicU8, Ordering};
 use crate::sync::{Arc, Condvar, Mutex};
-use crate::thread::{self, LocalKey};
+use crate::thread::{self, Builder, LocalKey};
 use crate::thread_local;
 
 #[derive(Clone, Default)]
@@ -340,3 +340,34 @@ fn join_orders_after_tls_destructors() {
         jh2.join().unwrap();
     }
 }
+
+// Test that thread::current is still available in TLS destructors.
+#[test]
+fn thread_current_in_dtor() {
+    // Go through one round of TLS destruction first.
+    struct Defer;
+    impl Drop for Defer {
+        fn drop(&mut self) {
+            RETRIEVE.with(|_| {});
+        }
+    }
+
+    struct RetrieveName;
+    impl Drop for RetrieveName {
+        fn drop(&mut self) {
+            *NAME.lock().unwrap() = Some(thread::current().name().unwrap().to_owned());
+        }
+    }
+
+    static NAME: Mutex<Option<String>> = Mutex::new(None);
+
+    thread_local! {
+        static DEFER: Defer = const { Defer };
+        static RETRIEVE: RetrieveName = const { RetrieveName };
+    }
+
+    Builder::new().name("test".to_owned()).spawn(|| DEFER.with(|_| {})).unwrap().join().unwrap();
+    let name = NAME.lock().unwrap();
+    let name = name.as_ref().unwrap();
+    assert_eq!(name, "test");
+}

library/std/src/thread/mod.rs

@@ -141,7 +141,7 @@
 //! [`Result`]: crate::result::Result
 //! [`Ok`]: crate::result::Result::Ok
 //! [`Err`]: crate::result::Result::Err
-//! [`thread::current`]: current
+//! [`thread::current`]: current::current
 //! [`thread::Result`]: Result
 //! [`unpark`]: Thread::unpark
 //! [`thread::park_timeout`]: park_timeout
@@ -159,7 +159,7 @@
 mod tests;
 
 use crate::any::Any;
-use crate::cell::{Cell, OnceCell, UnsafeCell};
+use crate::cell::UnsafeCell;
 use crate::env;
 use crate::ffi::CStr;
 use crate::fmt;
@@ -185,29 +185,27 @@ mod scoped;
 #[stable(feature = "scoped_threads", since = "1.63.0")]
 pub use scoped::{scope, Scope, ScopedJoinHandle};
 
+mod current;
+
+#[stable(feature = "rust1", since = "1.0.0")]
+pub use current::current;
+pub(crate) use current::{current_id, drop_current, set_current, try_current};
+
 ////////////////////////////////////////////////////////////////////////////////
 // Thread-local storage
 ////////////////////////////////////////////////////////////////////////////////
 
 #[macro_use]
 mod local;
 
-cfg_if::cfg_if! {
-    if #[cfg(test)] {
-        // Avoid duplicating the global state associated with thread-locals between this crate and
-        // realstd. Miri relies on this.
-        pub use realstd::thread::{local_impl, AccessError, LocalKey};
-    } else {
-        #[stable(feature = "rust1", since = "1.0.0")]
-        pub use self::local::{AccessError, LocalKey};
-
-        // Implementation details used by the thread_local!{} macro.
-        #[doc(hidden)]
-        #[unstable(feature = "thread_local_internals", issue = "none")]
-        pub mod local_impl {
-            pub use crate::sys::thread_local::*;
-        }
-    }
+#[stable(feature = "rust1", since = "1.0.0")]
+pub use self::local::{AccessError, LocalKey};
+
+// Implementation details used by the thread_local!{} macro.
+#[doc(hidden)]
+#[unstable(feature = "thread_local_internals", issue = "none")]
+pub mod local_impl {
+    pub use crate::sys::thread_local::*;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -487,7 +485,11 @@ impl Builder {
             amt
         });
 
-        let my_thread = name.map_or_else(Thread::new_unnamed, Thread::new);
+        let id = ThreadId::new();
+        let my_thread = match name {
+            Some(name) => Thread::new(id, name.into()),
+            None => Thread::new_unnamed(id),
+        };
         let their_thread = my_thread.clone();
 
         let my_packet: Arc<Packet<'scope, T>> = Arc::new(Packet {
@@ -526,14 +528,16 @@ impl Builder {
 
         let f = MaybeDangling::new(f);
         let main = move || {
+            // Immediately store the thread handle to avoid setting it or its ID
+            // twice, which would cause an abort.
+            set_current(their_thread.clone());
             if let Some(name) = their_thread.cname() {
                 imp::Thread::set_name(name);
             }
 
             crate::io::set_output_capture(output_capture);
 
             let f = f.into_inner();
-            set_current(their_thread);
             let try_result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
                 crate::sys::backtrace::__rust_begin_short_backtrace(f)
             }));
@@ -694,84 +698,6 @@ where
     Builder::new().spawn(f).expect("failed to spawn thread")
 }
 
-thread_local! {
-    // Invariant: `CURRENT` and `CURRENT_ID` will always be initialized together.
-    // If `CURRENT` is initialized, then `CURRENT_ID` will hold the same value
-    // as `CURRENT.id()`.
-    static CURRENT: OnceCell<Thread> = const { OnceCell::new() };
-    static CURRENT_ID: Cell<Option<ThreadId>> = const { Cell::new(None) };
-}
-
-/// Sets the thread handle for the current thread.
-///
-/// Aborts if the handle has been set already to reduce code size.
-pub(crate) fn set_current(thread: Thread) {
-    let tid = thread.id();
-    // Using `unwrap` here can add ~3kB to the binary size. We have complete
-    // control over where this is called, so just abort if there is a bug.
-    CURRENT.with(|current| match current.set(thread) {
-        Ok(()) => CURRENT_ID.set(Some(tid)),
-        Err(_) => rtabort!("thread::set_current should only be called once per thread"),
-    });
-}
-
-/// Gets a handle to the thread that invokes it.
-///
-/// In contrast to the public `current` function, this will not panic if called
-/// from inside a TLS destructor.
-pub(crate) fn try_current() -> Option<Thread> {
-    CURRENT
-        .try_with(|current| {
-            current
-                .get_or_init(|| {
-                    let thread = Thread::new_unnamed();
-                    CURRENT_ID.set(Some(thread.id()));
-                    thread
-                })
-                .clone()
-        })
-        .ok()
-}
-
-/// Gets the id of the thread that invokes it.
-#[inline]
-pub(crate) fn current_id() -> ThreadId {
-    CURRENT_ID.get().unwrap_or_else(|| {
-        // If `CURRENT_ID` isn't initialized yet, then `CURRENT` must also not be initialized.
-        // `current()` will initialize both `CURRENT` and `CURRENT_ID` so subsequent calls to
-        // `current_id()` will succeed immediately.
-        current().id()
-    })
-}
-
-/// Gets a handle to the thread that invokes it.
-///
-/// # Examples
-///
-/// Getting a handle to the current thread with `thread::current()`:
-///
-/// ```
-/// use std::thread;
-///
-/// let handler = thread::Builder::new()
-///     .name("named thread".into())
-///     .spawn(|| {
-///         let handle = thread::current();
-///         assert_eq!(handle.name(), Some("named thread"));
-///     })
-///     .unwrap();
-///
-/// handler.join().unwrap();
-/// ```
-#[must_use]
-#[stable(feature = "rust1", since = "1.0.0")]
-pub fn current() -> Thread {
-    try_current().expect(
-        "use of std::thread::current() is not possible \
-         after the thread's local data has been destroyed",
-    )
-}
-
 /// Cooperatively gives up a timeslice to the OS scheduler.
 ///
 /// This calls the underlying OS scheduler's yield primitive, signaling
@@ -1229,8 +1155,11 @@ pub fn park_timeout(dur: Duration) {
 pub struct ThreadId(NonZero<u64>);
 
 impl ThreadId {
+    // DO NOT rely on this value.
+    const MAIN_THREAD: ThreadId = ThreadId(unsafe { NonZero::new_unchecked(1) });
+
     // Generate a new unique thread ID.
-    fn new() -> ThreadId {
+    pub(crate) fn new() -> ThreadId {
         #[cold]
         fn exhausted() -> ! {
             panic!("failed to generate unique thread ID: bitspace exhausted")
@@ -1240,7 +1169,7 @@ impl ThreadId {
             if #[cfg(target_has_atomic = "64")] {
                 use crate::sync::atomic::AtomicU64;
 
-                static COUNTER: AtomicU64 = AtomicU64::new(0);
+                static COUNTER: AtomicU64 = AtomicU64::new(1);
 
                 let mut last = COUNTER.load(Ordering::Relaxed);
                 loop {
@@ -1256,7 +1185,7 @@ impl ThreadId {
             } else {
                 use crate::sync::{Mutex, PoisonError};
 
-                static COUNTER: Mutex<u64> = Mutex::new(0);
+                static COUNTER: Mutex<u64> = Mutex::new(1);
 
                 let mut counter = COUNTER.lock().unwrap_or_else(PoisonError::into_inner);
                 let Some(id) = counter.checked_add(1) else {
@@ -1273,6 +1202,11 @@ impl ThreadId {
         }
     }
 
+    #[cfg(not(target_thread_local))]
+    fn from_u64(v: u64) -> Option<ThreadId> {
+        NonZero::new(v).map(ThreadId)
+    }
+
     /// This returns a numeric identifier for the thread identified by this
     /// `ThreadId`.
     ///
@@ -1372,27 +1306,27 @@ impl Inner {
 /// should instead use a function like `spawn` to create new threads, see the
 /// docs of [`Builder`] and [`spawn`] for more details.
 ///
-/// [`thread::current`]: current
+/// [`thread::current`]: current::current
 pub struct Thread {
     inner: Pin<Arc<Inner>>,
 }
 
 impl Thread {
     /// Used only internally to construct a thread object without spawning.
-    pub(crate) fn new(name: String) -> Thread {
-        Self::new_inner(ThreadName::Other(name.into()))
+    pub(crate) fn new(id: ThreadId, name: String) -> Thread {
+        Self::new_inner(id, ThreadName::Other(name.into()))
     }
 
-    pub(crate) fn new_unnamed() -> Thread {
-        Self::new_inner(ThreadName::Unnamed)
+    pub(crate) fn new_unnamed(id: ThreadId) -> Thread {
+        Self::new_inner(id, ThreadName::Unnamed)
     }
 
     // Used in runtime to construct main thread
     pub(crate) fn new_main() -> Thread {
-        Self::new_inner(ThreadName::Main)
+        Self::new_inner(ThreadId::MAIN_THREAD, ThreadName::Main)
     }
 
-    fn new_inner(name: ThreadName) -> Thread {
+    fn new_inner(id: ThreadId, name: ThreadName) -> Thread {
         // We have to use `unsafe` here to construct the `Parker` in-place,
         // which is required for the UNIX implementation.
         //
@@ -1402,14 +1336,22 @@ impl Thread {
             let mut arc = Arc::<Inner>::new_uninit();
             let ptr = Arc::get_mut_unchecked(&mut arc).as_mut_ptr();
             addr_of_mut!((*ptr).name).write(name);
-            addr_of_mut!((*ptr).id).write(ThreadId::new());
+            addr_of_mut!((*ptr).id).write(id);
             Parker::new_in_place(addr_of_mut!((*ptr).parker));
             Pin::new_unchecked(arc.assume_init())
         };
 
         Thread { inner }
     }
 
+    fn into_raw(self) -> *mut () {
+        unsafe { Arc::into_raw(Pin::into_inner_unchecked(self.inner)) as *mut () }
+    }
+
+    unsafe fn from_raw(ptr: *mut ()) -> Thread {
+        unsafe { Thread { inner: Pin::new_unchecked(Arc::from_raw(ptr as *const Inner)) } }
+    }
+
     /// Like the public [`park`], but callable on any handle. This is used to
     /// allow parking in TLS destructors.
     ///