std: lazily allocate the main thread handle

Thereby, we also allow accessing thread::current before main: as the runtime no longer tries to install its own handle, this will no longer trigger an abort. Rather, the name returned from name will only be "main" after the runtime initialization code has run, but I think that is acceptable.

This new approach also requires some changes to the signal handling code, as calling `thread::current` would now allocate when called on the main thread, which is not acceptable. I fixed this by adding a new function (`with_current_name`) that performs all the naming logic without allocation or without initializing the thread ID (which could allocate on some platforms).

2 files changed, 117 insertions, 69 deletions

diff --git a/library/std/src/thread/current.rs b/library/std/src/thread/current.rs
index 3d2c288b360..fa88059fe64 100644
--- a/library/std/src/thread/current.rs
+++ b/library/std/src/thread/current.rs
@@ -15,7 +15,7 @@ local_pointer! {
 ///
 /// 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 {
+pub(super) mod id {
     use super::*;
 
     cfg_if::cfg_if! {
@@ -27,7 +27,7 @@ mod id {
 
             pub(super) const CHEAP: bool = true;
 
-            pub(super) fn get() -> Option<ThreadId> {
+            pub(crate) fn get() -> Option<ThreadId> {
                 ID.get()
             }
 
@@ -44,7 +44,7 @@ mod id {
 
             pub(super) const CHEAP: bool = false;
 
-            pub(super) fn get() -> Option<ThreadId> {
+            pub(crate) 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;
@@ -67,7 +67,7 @@ mod id {
 
             pub(super) const CHEAP: bool = false;
 
-            pub(super) fn get() -> Option<ThreadId> {
+            pub(crate) fn get() -> Option<ThreadId> {
                 let id0 = ID0.get().addr() as u64;
                 let id32 = ID32.get().addr() as u64;
                 ThreadId::from_u64((id32 << 32) + id0)
@@ -85,7 +85,7 @@ mod id {
 
             pub(super) const CHEAP: bool = true;
 
-            pub(super) fn get() -> Option<ThreadId> {
+            pub(crate) fn get() -> Option<ThreadId> {
                 let id = ID.get().addr() as u64;
                 ThreadId::from_u64(id)
             }
@@ -112,7 +112,7 @@ mod id {
 
 /// Tries to set the thread handle for the current thread. Fails if a handle was
 /// already set or if the thread ID of `thread` would change an already-set ID.
-pub(crate) fn set_current(thread: Thread) -> Result<(), Thread> {
+pub(super) fn set_current(thread: Thread) -> Result<(), Thread> {
     if CURRENT.get() != NONE {
         return Err(thread);
     }
@@ -140,28 +140,28 @@ pub(crate) fn current_id() -> ThreadId {
     // 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();
-            }
+        if let Some(id) = try_with_current(|t| t.map(|t| t.id())) {
+            return 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> {
+/// Gets a reference to the handle of the thread that invokes it, if the handle
+/// has been initialized.
+pub(super) fn try_with_current<F, R>(f: F) -> R
+where
+    F: FnOnce(Option<&Thread>) -> R,
+{
     let current = CURRENT.get();
     if current > DESTROYED {
         unsafe {
             let current = ManuallyDrop::new(Thread::from_raw(current));
-            Some((*current).clone())
+            f(Some(&current))
         }
     } else {
-        None
+        f(None)
     }
 }
 
@@ -176,7 +176,7 @@ pub(crate) fn current_or_unnamed() -> Thread {
             (*current).clone()
         }
     } else if current == DESTROYED {
-        Thread::new_unnamed(id::get_or_init())
+        Thread::new(id::get_or_init(), None)
     } else {
         init_current(current)
     }
@@ -221,7 +221,7 @@ fn init_current(current: *mut ()) -> Thread {
         CURRENT.set(BUSY);
         // If the thread ID was initialized already, use it.
         let id = id::get_or_init();
-        let thread = Thread::new_unnamed(id);
+        let thread = Thread::new(id, None);
 
         // Make sure that `crate::rt::thread_cleanup` will be run, which will
         // call `drop_current`.
diff --git a/library/std/src/thread/mod.rs b/library/std/src/thread/mod.rs
index 81d75641ab5..cf7a951d100 100644
--- a/library/std/src/thread/mod.rs
+++ b/library/std/src/thread/mod.rs
@@ -183,7 +183,8 @@ mod current;
 
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use current::current;
-pub(crate) use current::{current_id, current_or_unnamed, drop_current, set_current, try_current};
+pub(crate) use current::{current_id, current_or_unnamed, drop_current};
+use current::{set_current, try_with_current};
 
 mod spawnhook;
 
@@ -498,10 +499,7 @@ impl Builder {
         });
 
         let id = ThreadId::new();
-        let my_thread = match name {
-            Some(name) => Thread::new(id, name),
-            None => Thread::new_unnamed(id),
-        };
+        let my_thread = Thread::new(id, name);
 
         let hooks = if no_hooks {
             spawnhook::ChildSpawnHooks::default()
@@ -1232,7 +1230,7 @@ impl ThreadId {
         }
     }
 
-    #[cfg(not(target_thread_local))]
+    #[cfg(any(not(target_thread_local), target_has_atomic = "64"))]
     fn from_u64(v: u64) -> Option<ThreadId> {
         NonZero::new(v).map(ThreadId)
     }
@@ -1256,30 +1254,16 @@ impl ThreadId {
 // Thread
 ////////////////////////////////////////////////////////////////////////////////
 
-/// The internal representation of a `Thread`'s name.
-enum ThreadName {
-    Main,
-    Other(ThreadNameString),
-    Unnamed,
-}
-
 // This module ensures private fields are kept private, which is necessary to enforce the safety requirements.
 mod thread_name_string {
-    use core::str;
-
-    use super::ThreadName;
     use crate::ffi::{CStr, CString};
+    use crate::str;
 
     /// Like a `String` it's guaranteed UTF-8 and like a `CString` it's null terminated.
     pub(crate) struct ThreadNameString {
         inner: CString,
     }
-    impl core::ops::Deref for ThreadNameString {
-        type Target = CStr;
-        fn deref(&self) -> &CStr {
-            &self.inner
-        }
-    }
+
     impl From<String> for ThreadNameString {
         fn from(s: String) -> Self {
             Self {
@@ -1287,27 +1271,91 @@ mod thread_name_string {
             }
         }
     }
-    impl ThreadName {
-        pub fn as_cstr(&self) -> Option<&CStr> {
-            match self {
-                ThreadName::Main => Some(c"main"),
-                ThreadName::Other(other) => Some(other),
-                ThreadName::Unnamed => None,
-            }
+
+    impl ThreadNameString {
+        pub fn as_cstr(&self) -> &CStr {
+            &self.inner
         }
 
-        pub fn as_str(&self) -> Option<&str> {
-            // SAFETY: `as_cstr` can only return `Some` for a fixed CStr or a `ThreadNameString`,
-            // which is guaranteed to be UTF-8.
-            self.as_cstr().map(|s| unsafe { str::from_utf8_unchecked(s.to_bytes()) })
+        pub fn as_str(&self) -> &str {
+            // SAFETY: `ThreadNameString` is guaranteed to be UTF-8.
+            unsafe { str::from_utf8_unchecked(self.inner.to_bytes()) }
+        }
+    }
+}
+
+use thread_name_string::ThreadNameString;
+
+pub(crate) mod main_thread {
+    cfg_if::cfg_if! {
+        if #[cfg(target_has_atomic = "64")] {
+            use super::ThreadId;
+            use crate::sync::atomic::AtomicU64;
+            use crate::sync::atomic::Ordering::Relaxed;
+
+            static MAIN: AtomicU64 = AtomicU64::new(0);
+
+            pub(super) fn get() -> Option<ThreadId> {
+                ThreadId::from_u64(MAIN.load(Relaxed))
+            }
+
+            /// # Safety
+            /// May only be called once.
+            pub(crate) unsafe fn set(id: ThreadId) {
+                MAIN.store(id.as_u64().get(), Relaxed)
+            }
+        } else {
+            use super::ThreadId;
+            use crate::mem::MaybeUninit;
+            use crate::sync::atomic::AtomicBool;
+            use crate::sync::atomic::Ordering::{Acquire, Release};
+
+            static INIT: AtomicBool = AtomicBool::new(false);
+            static mut MAIN: MaybeUninit<ThreadId> = MaybeUninit::uninit();
+
+            pub(super) fn get() -> Option<ThreadId> {
+                if INIT.load(Acquire) {
+                    Some(unsafe { MAIN.assume_init() })
+                } else {
+                    None
+                }
+            }
+
+            /// # Safety
+            /// May only be called once.
+            pub(crate) unsafe fn set(id: ThreadId) {
+                unsafe { MAIN = MaybeUninit::new(id) };
+                INIT.store(true, Release);
+            }
         }
     }
 }
-pub(crate) use thread_name_string::ThreadNameString;
+
+pub(crate) fn with_current_name<F, R>(f: F) -> R
+where
+    F: FnOnce(Option<&str>) -> R,
+{
+    try_with_current(|thread| {
+        if let Some(thread) = thread {
+            if let Some(name) = &thread.inner.name {
+                return f(Some(name.as_str()));
+            } else if Some(thread.inner.id) == main_thread::get() {
+                return f(Some("main"));
+            }
+        } else if let Some(main) = main_thread::get()
+            && let Some(id) = current::id::get()
+            && id == main
+        {
+            return f(Some("main"));
+        }
+
+        f(None)
+    })
+}
 
 /// The internal representation of a `Thread` handle
 struct Inner {
-    name: ThreadName, // Guaranteed to be UTF-8
+    name: Option<ThreadNameString>,
     id: ThreadId,
     parker: Parker,
 }
@@ -1343,21 +1391,9 @@ pub struct Thread {
 }
 
 impl Thread {
-    /// Used only internally to construct a thread object without spawning.
-    pub(crate) fn new(id: ThreadId, name: String) -> Thread {
-        Self::new_inner(id, ThreadName::Other(name.into()))
-    }
-
-    pub(crate) fn new_unnamed(id: ThreadId) -> Thread {
-        Self::new_inner(id, ThreadName::Unnamed)
-    }
-
-    /// Constructs the thread handle for the main thread.
-    pub(crate) fn new_main(id: ThreadId) -> Thread {
-        Self::new_inner(id, ThreadName::Main)
-    }
+    pub(crate) fn new(id: ThreadId, name: Option<String>) -> Thread {
+        let name = name.map(ThreadNameString::from);
 
-    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.
         //
@@ -1489,7 +1525,13 @@ impl Thread {
     #[stable(feature = "rust1", since = "1.0.0")]
     #[must_use]
     pub fn name(&self) -> Option<&str> {
-        self.inner.name.as_str()
+        if let Some(name) = &self.inner.name {
+            Some(name.as_str())
+        } else if main_thread::get() == Some(self.inner.id) {
+            Some("main")
+        } else {
+            None
+        }
     }
 
     /// Consumes the `Thread`, returning a raw pointer.
@@ -1541,7 +1583,13 @@ impl Thread {
     }
 
     fn cname(&self) -> Option<&CStr> {
-        self.inner.name.as_cstr()
+        if let Some(name) = &self.inner.name {
+            Some(name.as_cstr())
+        } else if main_thread::get() == Some(self.inner.id) {
+            Some(c"main")
+        } else {
+            None
+        }
     }
 }