diff options
Diffstat (limited to 'library/std/src/thread/mod.rs')
| -rw-r--r-- | library/std/src/thread/mod.rs | 32 |
1 files changed, 29 insertions, 3 deletions
diff --git a/library/std/src/thread/mod.rs b/library/std/src/thread/mod.rs index c8ee365392f..7b98f2ae763 100644 --- a/library/std/src/thread/mod.rs +++ b/library/std/src/thread/mod.rs @@ -159,7 +159,7 @@ mod tests; use crate::any::Any; -use crate::cell::{OnceCell, UnsafeCell}; +use crate::cell::{Cell, OnceCell, UnsafeCell}; use crate::env; use crate::ffi::{CStr, CString}; use crate::fmt; @@ -699,17 +699,22 @@ where } 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(()) => {} + Ok(()) => CURRENT_ID.set(Some(tid)), Err(_) => rtabort!("thread::set_current should only be called once per thread"), }); } @@ -719,7 +724,28 @@ pub(crate) fn set_current(thread: Thread) { /// 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(|| Thread::new_unnamed()).clone()).ok() + 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. |
