summary refs log tree commit diff
path: root/library/std/src/thread/mod.rs
diff options
context:
space:
mode:
Diffstat (limited to 'library/std/src/thread/mod.rs')
-rw-r--r--library/std/src/thread/mod.rs32
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.