From ee6cd82c07b50bf08ebdad1a43f9b3f23cbbccfe Mon Sep 17 00:00:00 2001 From: joboet Date: Tue, 5 Nov 2024 19:28:06 +0100 Subject: [PATCH] std: lazily allocate the main thread handle #123550 eliminated the allocation of the main thread handle, but at the cost of greatly increased complexity. This PR proposes another approach: Instead of creating the main thread handle itself, the runtime simply remembers the thread ID of the main thread. The main thread handle is then only allocated when it is used, using the same lazy-initialization mechanism as for non-runtime use of `thread::current`, and the name method uses the thread ID to identify the main thread handle and return the correct name ("main") for it. 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. I've also moved the `Thread` code into its own module, the size of `thread/mod.rs` is just getting out of hand. --- library/std/src/lib.rs | 1 + library/std/src/panicking.rs | 9 +- library/std/src/rt.rs | 23 +- library/std/src/sys/sync/rwlock/queue.rs | 2 +- library/std/src/thread/current.rs | 2 +- library/std/src/thread/mod.rs | 315 +---------------------- library/std/src/thread/thread.rs | 311 ++++++++++++++++++++++ tests/ui/runtime/stdout-before-main.rs | 4 + 8 files changed, 340 insertions(+), 327 deletions(-) create mode 100644 library/std/src/thread/thread.rs diff --git a/library/std/src/lib.rs b/library/std/src/lib.rs index 4308b7114994f..b5450d0a20cf2 100644 --- a/library/std/src/lib.rs +++ b/library/std/src/lib.rs @@ -365,6 +365,7 @@ #![feature(std_internals)] #![feature(str_internals)] #![feature(strict_provenance_atomic_ptr)] +#![feature(sync_unsafe_cell)] #![feature(ub_checks)] // tidy-alphabetical-end // diff --git a/library/std/src/panicking.rs b/library/std/src/panicking.rs index ac1f547c9143f..aa2ac65fdba79 100644 --- a/library/std/src/panicking.rs +++ b/library/std/src/panicking.rs @@ -247,8 +247,13 @@ fn default_hook(info: &PanicHookInfo<'_>) { let location = info.location().unwrap(); let msg = payload_as_str(info.payload()); - let thread = thread::try_current(); - let name = thread.as_ref().and_then(|t| t.name()).unwrap_or(""); + let thread; + let name = if thread::is_main() { + "main" + } else { + thread = thread::try_current(); + thread.as_ref().and_then(|t| t.name()).unwrap_or("") + }; let write = #[optimize(size)] |err: &mut dyn crate::io::Write| { diff --git a/library/std/src/rt.rs b/library/std/src/rt.rs index 80e7c3c026bd7..3b528c5ea78c9 100644 --- a/library/std/src/rt.rs +++ b/library/std/src/rt.rs @@ -23,7 +23,7 @@ pub use core::panicking::{panic_display, panic_fmt}; #[rustfmt::skip] use crate::any::Any; use crate::sync::Once; -use crate::thread::{self, Thread}; +use crate::thread::{current_id, set_main_thread}; use crate::{mem, panic, sys}; // Prints to the "panic output", depending on the platform this may be: @@ -102,23 +102,10 @@ unsafe fn init(argc: isize, argv: *const *const u8, sigpipe: u8) { sys::init(argc, argv, sigpipe) }; - // 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"); + // Mark the current thread as main thread. + // SAFETY: this is the only place and time where this is called. + unsafe { + set_main_thread(current_id()); } } diff --git a/library/std/src/sys/sync/rwlock/queue.rs b/library/std/src/sys/sync/rwlock/queue.rs index 889961915f4e6..3dbeac6ec4762 100644 --- a/library/std/src/sys/sync/rwlock/queue.rs +++ b/library/std/src/sys/sync/rwlock/queue.rs @@ -201,7 +201,7 @@ impl Node { // 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(ThreadId::new())) + thread::try_current().unwrap_or_else(|| Thread::new(ThreadId::new(), None)) }); self.completed = AtomicBool::new(false); } diff --git a/library/std/src/thread/current.rs b/library/std/src/thread/current.rs index e6eb90c4c30a5..a85519cd6076d 100644 --- a/library/std/src/thread/current.rs +++ b/library/std/src/thread/current.rs @@ -204,7 +204,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 d12219b31b265..0f0d73a377960 100644 --- a/library/std/src/thread/mod.rs +++ b/library/std/src/thread/mod.rs @@ -160,18 +160,15 @@ mod tests; use crate::any::Any; use crate::cell::UnsafeCell; -use crate::ffi::CStr; use crate::marker::PhantomData; use crate::mem::{self, ManuallyDrop, forget}; use crate::num::NonZero; -use crate::pin::Pin; use crate::sync::Arc; use crate::sync::atomic::{AtomicUsize, Ordering}; -use crate::sys::sync::Parker; use crate::sys::thread as imp; use crate::sys_common::{AsInner, IntoInner}; use crate::time::{Duration, Instant}; -use crate::{env, fmt, io, panic, panicking, str}; +use crate::{env, fmt, io, panic, panicking}; #[stable(feature = "scoped_threads", since = "1.63.0")] mod scoped; @@ -185,6 +182,12 @@ mod current; pub use current::current; pub(crate) use current::{current_id, drop_current, set_current, try_current}; +mod thread; + +#[stable(feature = "rust1", since = "1.0.0")] +pub use thread::Thread; +pub(crate) use thread::{is_main, set_main_thread}; + //////////////////////////////////////////////////////////////////////////////// // Thread-local storage //////////////////////////////////////////////////////////////////////////////// @@ -477,11 +480,7 @@ impl Builder { amt }); - let id = ThreadId::new(); - let my_thread = match name { - Some(name) => Thread::new(id, name.into()), - None => Thread::new_unnamed(id), - }; + let my_thread = Thread::new(ThreadId::new(), name); let their_thread = my_thread.clone(); let my_packet: Arc> = Arc::new(Packet { @@ -1125,7 +1124,7 @@ pub fn park_timeout(dur: Duration) { let guard = PanicGuard; // SAFETY: park_timeout is called on the parker owned by this thread. unsafe { - current().inner.as_ref().parker().park_timeout(dur); + current().park_timeout(dur); } // No panic occurred, do not abort. forget(guard); @@ -1208,7 +1207,7 @@ impl ThreadId { } } - #[cfg(not(target_thread_local))] + #[allow(unused)] fn from_u64(v: u64) -> Option { NonZero::new(v).map(ThreadId) } @@ -1228,300 +1227,6 @@ 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}; - - /// 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 for ThreadNameString { - fn from(s: String) -> Self { - Self { - inner: CString::new(s).expect("thread name may not contain interior null bytes"), - } - } - } - impl ThreadName { - pub fn as_cstr(&self) -> Option<&CStr> { - match self { - ThreadName::Main => Some(c"main"), - ThreadName::Other(other) => Some(other), - ThreadName::Unnamed => None, - } - } - - 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(crate) use thread_name_string::ThreadNameString; - -/// The internal representation of a `Thread` handle -struct Inner { - name: ThreadName, // Guaranteed to be UTF-8 - id: ThreadId, - parker: Parker, -} - -impl Inner { - fn parker(self: Pin<&Self>) -> Pin<&Parker> { - unsafe { Pin::map_unchecked(self, |inner| &inner.parker) } - } -} - -#[derive(Clone)] -#[stable(feature = "rust1", since = "1.0.0")] -/// A handle to a thread. -/// -/// Threads are represented via the `Thread` type, which you can get in one of -/// two ways: -/// -/// * By spawning a new thread, e.g., using the [`thread::spawn`][`spawn`] -/// function, and calling [`thread`][`JoinHandle::thread`] on the -/// [`JoinHandle`]. -/// * By requesting the current thread, using the [`thread::current`] function. -/// -/// The [`thread::current`] function is available even for threads not spawned -/// by the APIs of this module. -/// -/// There is usually no need to create a `Thread` struct yourself, one -/// should instead use a function like `spawn` to create new threads, see the -/// docs of [`Builder`] and [`spawn`] for more details. -/// -/// [`thread::current`]: current::current -pub struct Thread { - inner: Pin>, -} - -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) - } - - 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. - // - // SAFETY: We pin the Arc immediately after creation, so its address never - // changes. - let inner = unsafe { - let mut arc = Arc::::new_uninit(); - let ptr = Arc::get_mut_unchecked(&mut arc).as_mut_ptr(); - (&raw mut (*ptr).name).write(name); - (&raw mut (*ptr).id).write(id); - Parker::new_in_place(&raw mut (*ptr).parker); - Pin::new_unchecked(arc.assume_init()) - }; - - Thread { inner } - } - - /// Like the public [`park`], but callable on any handle. This is used to - /// allow parking in TLS destructors. - /// - /// # Safety - /// May only be called from the thread to which this handle belongs. - pub(crate) unsafe fn park(&self) { - unsafe { self.inner.as_ref().parker().park() } - } - - /// Atomically makes the handle's token available if it is not already. - /// - /// Every thread is equipped with some basic low-level blocking support, via - /// the [`park`][park] function and the `unpark()` method. These can be - /// used as a more CPU-efficient implementation of a spinlock. - /// - /// See the [park documentation][park] for more details. - /// - /// # Examples - /// - /// ``` - /// use std::thread; - /// use std::time::Duration; - /// - /// let parked_thread = thread::Builder::new() - /// .spawn(|| { - /// println!("Parking thread"); - /// thread::park(); - /// println!("Thread unparked"); - /// }) - /// .unwrap(); - /// - /// // Let some time pass for the thread to be spawned. - /// thread::sleep(Duration::from_millis(10)); - /// - /// println!("Unpark the thread"); - /// parked_thread.thread().unpark(); - /// - /// parked_thread.join().unwrap(); - /// ``` - #[stable(feature = "rust1", since = "1.0.0")] - #[inline] - pub fn unpark(&self) { - self.inner.as_ref().parker().unpark(); - } - - /// Gets the thread's unique identifier. - /// - /// # Examples - /// - /// ``` - /// use std::thread; - /// - /// let other_thread = thread::spawn(|| { - /// thread::current().id() - /// }); - /// - /// let other_thread_id = other_thread.join().unwrap(); - /// assert!(thread::current().id() != other_thread_id); - /// ``` - #[stable(feature = "thread_id", since = "1.19.0")] - #[must_use] - pub fn id(&self) -> ThreadId { - self.inner.id - } - - /// Gets the thread's name. - /// - /// For more information about named threads, see - /// [this module-level documentation][naming-threads]. - /// - /// # Examples - /// - /// Threads by default have no name specified: - /// - /// ``` - /// use std::thread; - /// - /// let builder = thread::Builder::new(); - /// - /// let handler = builder.spawn(|| { - /// assert!(thread::current().name().is_none()); - /// }).unwrap(); - /// - /// handler.join().unwrap(); - /// ``` - /// - /// Thread with a specified name: - /// - /// ``` - /// use std::thread; - /// - /// let builder = thread::Builder::new() - /// .name("foo".into()); - /// - /// let handler = builder.spawn(|| { - /// assert_eq!(thread::current().name(), Some("foo")) - /// }).unwrap(); - /// - /// handler.join().unwrap(); - /// ``` - /// - /// [naming-threads]: ./index.html#naming-threads - #[stable(feature = "rust1", since = "1.0.0")] - #[must_use] - pub fn name(&self) -> Option<&str> { - self.inner.name.as_str() - } - - /// Consumes the `Thread`, returning a raw pointer. - /// - /// To avoid a memory leak the pointer must be converted - /// back into a `Thread` using [`Thread::from_raw`]. - /// - /// # Examples - /// - /// ``` - /// #![feature(thread_raw)] - /// - /// use std::thread::{self, Thread}; - /// - /// let thread = thread::current(); - /// let id = thread.id(); - /// let ptr = Thread::into_raw(thread); - /// unsafe { - /// assert_eq!(Thread::from_raw(ptr).id(), id); - /// } - /// ``` - #[unstable(feature = "thread_raw", issue = "97523")] - pub fn into_raw(self) -> *const () { - // Safety: We only expose an opaque pointer, which maintains the `Pin` invariant. - let inner = unsafe { Pin::into_inner_unchecked(self.inner) }; - Arc::into_raw(inner) as *const () - } - - /// Constructs a `Thread` from a raw pointer. - /// - /// The raw pointer must have been previously returned - /// by a call to [`Thread::into_raw`]. - /// - /// # Safety - /// - /// This function is unsafe because improper use may lead - /// to memory unsafety, even if the returned `Thread` is never - /// accessed. - /// - /// Creating a `Thread` from a pointer other than one returned - /// from [`Thread::into_raw`] is **undefined behavior**. - /// - /// Calling this function twice on the same raw pointer can lead - /// to a double-free if both `Thread` instances are dropped. - #[unstable(feature = "thread_raw", issue = "97523")] - pub unsafe fn from_raw(ptr: *const ()) -> Thread { - // Safety: Upheld by caller. - unsafe { Thread { inner: Pin::new_unchecked(Arc::from_raw(ptr as *const Inner)) } } - } - - fn cname(&self) -> Option<&CStr> { - self.inner.name.as_cstr() - } -} - -#[stable(feature = "rust1", since = "1.0.0")] -impl fmt::Debug for Thread { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Thread") - .field("id", &self.id()) - .field("name", &self.name()) - .finish_non_exhaustive() - } -} - //////////////////////////////////////////////////////////////////////////////// // JoinHandle //////////////////////////////////////////////////////////////////////////////// diff --git a/library/std/src/thread/thread.rs b/library/std/src/thread/thread.rs new file mode 100644 index 0000000000000..27fc6f67a4d77 --- /dev/null +++ b/library/std/src/thread/thread.rs @@ -0,0 +1,311 @@ +//! The [`Thread`] handle. + +use super::{ThreadId, current_id}; +use crate::ffi::{CStr, CString}; +use crate::pin::Pin; +use crate::sync::Arc; +use crate::sys::sync::Parker; +use crate::time::Duration; +use crate::{fmt, str}; + +/// The internal representation of a `Thread` handle +struct Inner { + // INVARIANT: must be valid UTF-8 + name: Option, + id: ThreadId, + parker: Parker, +} + +impl Inner { + fn parker(self: Pin<&Self>) -> Pin<&Parker> { + unsafe { Pin::map_unchecked(self, |inner| &inner.parker) } + } +} + +/// A handle to a thread. +/// +/// Threads are represented via the `Thread` type, which you can get in one of +/// two ways: +/// +/// * By spawning a new thread, e.g., using the [`thread::spawn`] function, and +/// calling [`thread`] on the [`JoinHandle`]. +/// * By requesting the current thread, using the [`thread::current`] function. +/// +/// The [`thread::current`] function is available even for threads not spawned +/// by the APIs of this module. +/// +/// There is usually no need to create a `Thread` struct yourself, one should +/// instead use a function like `spawn` to create new threads, see the docs of +/// [`Builder`] and [`spawn`] for more details. +/// +/// [`Builder`]: super::Builder +/// [`JoinHandle`]: super::JoinHandle +/// [`spawn`]: super::spawn +/// [`thread`]: super::JoinHandle::thread +/// [`thread::spawn`]: super::spawn +/// [`thread::current`]: super::current::current +#[derive(Clone)] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Thread { + inner: Pin>, +} + +impl Thread { + pub(crate) fn new(id: ThreadId, name: Option) -> Thread { + let name = name.map(|name| { + CString::new(name).expect("thread name may not contain interior null bytes") + }); + + // We have to use `unsafe` here to construct the `Parker` in-place, + // which is required for the UNIX implementation. + // + // SAFETY: We pin the Arc immediately after creation, so its address never + // changes. + let inner = unsafe { + let mut arc = Arc::::new_uninit(); + let ptr = Arc::get_mut_unchecked(&mut arc).as_mut_ptr(); + Parker::new_in_place(&raw mut (*ptr).parker); + (&raw mut (*ptr).name).write(name); + (&raw mut (*ptr).id).write(id); + Pin::new_unchecked(arc.assume_init()) + }; + + Thread { inner } + } + + /// Like the public [`park`](super::park), but callable on any handle. This is used to + /// allow parking in TLS destructors. + /// + /// # Safety + /// May only be called from the thread to which this handle belongs. + pub(crate) unsafe fn park(&self) { + unsafe { self.inner.as_ref().parker().park() } + } + + /// Like the public [`park_timeout`](super::park_timeout), but callable on any handle. + /// + /// # Safety + /// May only be called from the thread to which this handle belongs. + pub(crate) unsafe fn park_timeout(&self, dur: Duration) { + unsafe { self.inner.as_ref().parker().park_timeout(dur) } + } + + /// Atomically makes the handle's token available if it is not already. + /// + /// Every thread is equipped with some basic low-level blocking support, via + /// the [`park`](super::park) function and the `unpark()` method. These can be + /// used as a more CPU-efficient implementation of a spinlock. + /// + /// See the [park documentation](super::park) for more details. + /// + /// # Examples + /// + /// ``` + /// use std::thread; + /// use std::time::Duration; + /// + /// let parked_thread = thread::Builder::new() + /// .spawn(|| { + /// println!("Parking thread"); + /// thread::park(); + /// println!("Thread unparked"); + /// }) + /// .unwrap(); + /// + /// // Let some time pass for the thread to be spawned. + /// thread::sleep(Duration::from_millis(10)); + /// + /// println!("Unpark the thread"); + /// parked_thread.thread().unpark(); + /// + /// parked_thread.join().unwrap(); + /// ``` + #[stable(feature = "rust1", since = "1.0.0")] + #[inline] + pub fn unpark(&self) { + self.inner.as_ref().parker().unpark(); + } + + /// Gets the thread's unique identifier. + /// + /// # Examples + /// + /// ``` + /// use std::thread; + /// + /// let other_thread = thread::spawn(|| { + /// thread::current().id() + /// }); + /// + /// let other_thread_id = other_thread.join().unwrap(); + /// assert!(thread::current().id() != other_thread_id); + /// ``` + #[stable(feature = "thread_id", since = "1.19.0")] + #[must_use] + pub fn id(&self) -> ThreadId { + self.inner.id + } + + /// Gets the thread's name. + /// + /// For more information about named threads, see + /// [this module-level documentation][naming-threads]. + /// + /// # Examples + /// + /// Threads by default have no name specified: + /// + /// ``` + /// use std::thread; + /// + /// let builder = thread::Builder::new(); + /// + /// let handler = builder.spawn(|| { + /// assert!(thread::current().name().is_none()); + /// }).unwrap(); + /// + /// handler.join().unwrap(); + /// ``` + /// + /// Thread with a specified name: + /// + /// ``` + /// use std::thread; + /// + /// let builder = thread::Builder::new() + /// .name("foo".into()); + /// + /// let handler = builder.spawn(|| { + /// assert_eq!(thread::current().name(), Some("foo")) + /// }).unwrap(); + /// + /// handler.join().unwrap(); + /// ``` + /// + /// [naming-threads]: ./index.html#naming-threads + #[stable(feature = "rust1", since = "1.0.0")] + #[must_use] + pub fn name(&self) -> Option<&str> { + // SAFETY: the thread name is guaranteed to be UTF-8 by invariant. + Some(unsafe { str::from_utf8_unchecked(self.cname()?.to_bytes()) }) + } + + /// Consumes the `Thread`, returning a raw pointer. + /// + /// To avoid a memory leak the pointer must be converted + /// back into a `Thread` using [`Thread::from_raw`]. + /// + /// # Examples + /// + /// ``` + /// #![feature(thread_raw)] + /// + /// use std::thread::{self, Thread}; + /// + /// let thread = thread::current(); + /// let id = thread.id(); + /// let ptr = Thread::into_raw(thread); + /// unsafe { + /// assert_eq!(Thread::from_raw(ptr).id(), id); + /// } + /// ``` + #[unstable(feature = "thread_raw", issue = "97523")] + pub fn into_raw(self) -> *const () { + // Safety: We only expose an opaque pointer, which maintains the `Pin` invariant. + let inner = unsafe { Pin::into_inner_unchecked(self.inner) }; + Arc::into_raw(inner) as *const () + } + + /// Constructs a `Thread` from a raw pointer. + /// + /// The raw pointer must have been previously returned + /// by a call to [`Thread::into_raw`]. + /// + /// # Safety + /// + /// This function is unsafe because improper use may lead + /// to memory unsafety, even if the returned `Thread` is never + /// accessed. + /// + /// Creating a `Thread` from a pointer other than one returned + /// from [`Thread::into_raw`] is **undefined behavior**. + /// + /// Calling this function twice on the same raw pointer can lead + /// to a double-free if both `Thread` instances are dropped. + #[unstable(feature = "thread_raw", issue = "97523")] + pub unsafe fn from_raw(ptr: *const ()) -> Thread { + // Safety: Upheld by caller. + unsafe { Thread { inner: Pin::new_unchecked(Arc::from_raw(ptr as *const Inner)) } } + } + + pub(super) fn cname(&self) -> Option<&CStr> { + if let Some(name) = self.inner.name.as_deref() { + Some(name) + } else if Some(self.id()) == main_thread() { + // We don't give the main thread handle a name explicitly but detect + // it through its ID to support acquiring thread handles before main + // and to enable lazy allocation of the handle. + Some(c"main") + } else { + None + } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl fmt::Debug for Thread { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Thread") + .field("id", &self.id()) + .field("name", &self.name()) + .finish_non_exhaustive() + } +} + +cfg_if::cfg_if! { + if #[cfg(target_has_atomic = "64")] { + use crate::sync::atomic::{AtomicU64, Ordering::Relaxed}; + + static ID: AtomicU64 = AtomicU64::new(0); + + /// # Safety + /// May only be called once. + pub(crate) unsafe fn set_main_thread(id: ThreadId) { + ID.store(id.as_u64().get(), Relaxed) + } + + fn main_thread() -> Option { + ThreadId::from_u64(ID.load(Relaxed)) + } + } else { + use crate::cell::SyncUnsafeCell; + use crate::sync::atomic::{AtomicBool, Ordering::{Acquire, Release}}; + + static INIT: AtomicBool = AtomicBool::new(false); + static ID: SyncUnsafeCell> = SyncUnsafeCell::new(None); + + /// # Safety + /// May only be called once. + pub(crate) unsafe fn set_main_thread(id: ThreadId) { + unsafe { + ID.get().write(Some(id)); + INIT.store(true, Release); + } + } + + fn main_thread() -> Option { + if INIT.load(Acquire) { + unsafe { ID.get().read() } + } else { + None + } + } + } +} + +/// Returns whether this thread is the main thread. +/// +/// The result of this call is unreliable for code running outside the Rust runtime. +pub(crate) fn is_main() -> bool { + Some(current_id()) == main_thread() +} diff --git a/tests/ui/runtime/stdout-before-main.rs b/tests/ui/runtime/stdout-before-main.rs index 26c734d3e9eed..d11f44431ad77 100644 --- a/tests/ui/runtime/stdout-before-main.rs +++ b/tests/ui/runtime/stdout-before-main.rs @@ -13,6 +13,10 @@ use std::thread; #[link_section = ".init_array"] static INIT: extern "C" fn(c_int, *const *const u8, *const *const u8) = { extern "C" fn init(_argc: c_int, _argv: *const *const u8, _envp: *const *const u8) { + // Try access the thread handle before main. This shouldn't lead to abort + // later on. + drop(thread::current()); + print!("Hello from before "); }