Rollup merge of rust-lang#131233 - joboet:stdout-before-main, r=tgross35

tgross35 · web-flow · commit 7420c36fae41 · 2024-10-12T11:08:43.000-05:00
std: fix stdout-before-main Fixes rust-lang#130210. Since rust-lang#124881, `ReentrantLock` uses `ThreadId` to identify threads. This has the unfortunate consequence of breaking uses of `Stdout` before main: Locking the `ReentrantLock` that synchronizes the output will initialize the thread ID before the handle for the main thread is set in `rt::init`. But since that would overwrite the current thread ID, `thread::set_current` triggers an abort. This PR fixes the problem by using the already initialized thread ID for constructing the main thread handle and allowing `set_current` calls that do not change the thread's ID.
diff --git a/std/src/rt.rs b/std/src/rt.rs
@@ -102,9 +102,24 @@ unsafe fn init(argc: isize, argv: *const *const u8, sigpipe: u8) {
         sys::init(argc, argv, sigpipe)
     };
 
-    // Set up the current thread to give it the right name.
-    let thread = Thread::new_main();
-    thread::set_current(thread);
+    // Set up the current thread handle to give it the right name.
+    //
+    // When code running before main uses `ReentrantLock` (for example by
+    // using `println!`), the thread ID can become initialized before we
+    // create this handle. Since `set_current` fails when the ID of the
+    // handle does not match the current ID, we should attempt to use the
+    // current thread ID here instead of unconditionally creating a new
+    // one. Also see #130210.
+    let thread = Thread::new_main(thread::current_id());
+    if let Err(_thread) = thread::set_current(thread) {
+        // `thread::current` will create a new handle if none has been set yet.
+        // Thus, if someone uses it before main, this call will fail. That's a
+        // bad idea though, as we then cannot set the main thread name here.
+        //
+        // FIXME: detect the main thread in `thread::current` and use the
+        //        correct name there.
+        rtabort!("code running before main must not use thread::current");
+    }
 }
 
 /// Clean up the thread-local runtime state. This *should* be run after all other
diff --git a/std/src/thread/current.rs b/std/src/thread/current.rs
@@ -110,22 +110,24 @@ mod 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");
+/// 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> {
+    if CURRENT.get() != NONE {
+        return Err(thread);
     }
 
-    id::set(thread.id());
+    match id::get() {
+        Some(id) if id == thread.id() => {}
+        None => id::set(thread.id()),
+        _ => return Err(thread),
+    }
 
     // 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().cast_mut());
+    Ok(())
 }
 
 /// Gets the id of the thread that invokes it.
diff --git a/std/src/thread/mod.rs b/std/src/thread/mod.rs
@@ -519,9 +519,14 @@ 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 Err(_thread) = set_current(their_thread.clone()) {
+                // Both the current thread handle and the ID should not be
+                // initialized yet. Since only the C runtime and some of our
+                // platform code run before this, this point shouldn't be
+                // reachable. Use an abort to save binary size (see #123356).
+                rtabort!("something here is badly broken!");
+            }
+
             if let Some(name) = their_thread.cname() {
                 imp::Thread::set_name(name);
             }
@@ -1159,9 +1164,6 @@ 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.
     pub(crate) fn new() -> ThreadId {
         #[cold]
@@ -1173,7 +1175,7 @@ impl ThreadId {
             if #[cfg(target_has_atomic = "64")] {
                 use crate::sync::atomic::AtomicU64;
 
-                static COUNTER: AtomicU64 = AtomicU64::new(1);
+                static COUNTER: AtomicU64 = AtomicU64::new(0);
 
                 let mut last = COUNTER.load(Ordering::Relaxed);
                 loop {
@@ -1189,7 +1191,7 @@ impl ThreadId {
             } else {
                 use crate::sync::{Mutex, PoisonError};
 
-                static COUNTER: Mutex<u64> = Mutex::new(1);
+                static COUNTER: Mutex<u64> = Mutex::new(0);
 
                 let mut counter = COUNTER.lock().unwrap_or_else(PoisonError::into_inner);
                 let Some(id) = counter.checked_add(1) else {
@@ -1326,9 +1328,9 @@ impl Thread {
         Self::new_inner(id, ThreadName::Unnamed)
     }
 
-    // Used in runtime to construct main thread
-    pub(crate) fn new_main() -> Thread {
-        Self::new_inner(ThreadId::MAIN_THREAD, ThreadName::Main)
+    /// Constructs the thread handle for the main thread.
+    pub(crate) fn new_main(id: ThreadId) -> Thread {
+        Self::new_inner(id, ThreadName::Main)
     }
 
     fn new_inner(id: ThreadId, name: ThreadName) -> Thread {
