Rollup merge of #121666 - ChrisDenton:thread-name, r=cuviper

matthiaskrgr · web-flow · commit 0f544f280a52 · 2024-03-02T16:53:14.000+01:00
Use the OS thread name by default if `THREAD_INFO` has not been initialized

Currently if `THREAD_INFO` hasn't been initialized then the name will be set to `None`.  This PR changes it to use the OS thread name by default. This mostly affects foreign threads at the moment but we could expand this to make more use of the OS thread name in the future.

Note: I've only implemented `Thread::get_name` for windows, linux and macos (and macos adjacent) targets. The rest just return `None`.
diff --git a/library/std/src/sys/pal/hermit/thread.rs b/library/std/src/sys/pal/hermit/thread.rs
@@ -2,7 +2,7 @@
 
 use super::abi;
 use super::thread_local_dtor::run_dtors;
-use crate::ffi::CStr;
+use crate::ffi::{CStr, CString};
 use crate::io;
 use crate::mem;
 use crate::num::NonZero;
@@ -71,6 +71,10 @@ impl Thread {
         // nope
     }
 
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     #[inline]
     pub fn sleep(dur: Duration) {
         unsafe {
diff --git a/library/std/src/sys/pal/itron/thread.rs b/library/std/src/sys/pal/itron/thread.rs
@@ -8,7 +8,7 @@ use super::{
 };
 use crate::{
     cell::UnsafeCell,
-    ffi::CStr,
+    ffi::{CStr, CString},
     hint, io,
     mem::ManuallyDrop,
     num::NonZero,
@@ -204,6 +204,10 @@ impl Thread {
         // nope
     }
 
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     pub fn sleep(dur: Duration) {
         for timeout in dur2reltims(dur) {
             expect_success(unsafe { abi::dly_tsk(timeout) }, &"dly_tsk");
diff --git a/library/std/src/sys/pal/sgx/thread.rs b/library/std/src/sys/pal/sgx/thread.rs
@@ -1,6 +1,6 @@
 #![cfg_attr(test, allow(dead_code))] // why is this necessary?
 use super::unsupported;
-use crate::ffi::CStr;
+use crate::ffi::{CStr, CString};
 use crate::io;
 use crate::num::NonZero;
 use crate::time::Duration;
@@ -133,6 +133,10 @@ impl Thread {
         // which succeeds as-is with the SGX target.
     }
 
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     pub fn sleep(dur: Duration) {
         usercalls::wait_timeout(0, dur, || true);
     }
diff --git a/library/std/src/sys/pal/teeos/thread.rs b/library/std/src/sys/pal/teeos/thread.rs
@@ -1,7 +1,7 @@
 use core::convert::TryInto;
 
 use crate::cmp;
-use crate::ffi::CStr;
+use crate::ffi::{CStr, CString};
 use crate::io;
 use crate::mem;
 use crate::num::NonZero;
@@ -101,6 +101,10 @@ impl Thread {
         // contact the teeos rustzone team.
     }
 
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     /// only main thread could wait for sometime in teeos
     pub fn sleep(dur: Duration) {
         let sleep_millis = dur.as_millis();
diff --git a/library/std/src/sys/pal/uefi/thread.rs b/library/std/src/sys/pal/uefi/thread.rs
@@ -1,5 +1,5 @@
 use super::unsupported;
-use crate::ffi::CStr;
+use crate::ffi::{CStr, CString};
 use crate::io;
 use crate::num::NonZero;
 use crate::ptr::NonNull;
@@ -23,6 +23,10 @@ impl Thread {
         // nope
     }
 
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     pub fn sleep(dur: Duration) {
         let boot_services: NonNull<r_efi::efi::BootServices> =
             crate::os::uefi::env::boot_services().expect("can't sleep").cast();
diff --git a/library/std/src/sys/pal/unix/thread.rs b/library/std/src/sys/pal/unix/thread.rs
@@ -1,5 +1,5 @@
 use crate::cmp;
-use crate::ffi::CStr;
+use crate::ffi::{CStr, CString};
 use crate::io;
 use crate::mem;
 use crate::num::NonZero;
@@ -225,6 +225,44 @@ impl Thread {
         // Newlib, Emscripten, and VxWorks have no way to set a thread name.
     }
 
+    #[cfg(target_os = "linux")]
+    pub fn get_name() -> Option<CString> {
+        const TASK_COMM_LEN: usize = 16;
+        let mut name = vec![0u8; TASK_COMM_LEN];
+        let res = unsafe {
+            libc::pthread_getname_np(libc::pthread_self(), name.as_mut_ptr().cast(), name.len())
+        };
+        if res != 0 {
+            return None;
+        }
+        name.truncate(name.iter().position(|&c| c == 0)?);
+        CString::new(name).ok()
+    }
+
+    #[cfg(any(target_os = "macos", target_os = "ios", target_os = "tvos", target_os = "watchos"))]
+    pub fn get_name() -> Option<CString> {
+        let mut name = vec![0u8; libc::MAXTHREADNAMESIZE];
+        let res = unsafe {
+            libc::pthread_getname_np(libc::pthread_self(), name.as_mut_ptr().cast(), name.len())
+        };
+        if res != 0 {
+            return None;
+        }
+        name.truncate(name.iter().position(|&c| c == 0)?);
+        CString::new(name).ok()
+    }
+
+    #[cfg(not(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "tvos",
+        target_os = "watchos"
+    )))]
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     #[cfg(not(target_os = "espidf"))]
     pub fn sleep(dur: Duration) {
         let mut secs = dur.as_secs();
diff --git a/library/std/src/sys/pal/unsupported/thread.rs b/library/std/src/sys/pal/unsupported/thread.rs
@@ -1,5 +1,5 @@
 use super::unsupported;
-use crate::ffi::CStr;
+use crate::ffi::{CStr, CString};
 use crate::io;
 use crate::num::NonZero;
 use crate::time::Duration;
@@ -22,6 +22,10 @@ impl Thread {
         // nope
     }
 
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     pub fn sleep(_dur: Duration) {
         panic!("can't sleep");
     }
diff --git a/library/std/src/sys/pal/wasi/thread.rs b/library/std/src/sys/pal/wasi/thread.rs
@@ -1,4 +1,4 @@
-use crate::ffi::CStr;
+use crate::ffi::{CStr, CString};
 use crate::io;
 use crate::mem;
 use crate::num::NonZero;
@@ -134,6 +134,10 @@ impl Thread {
         // nope
     }
 
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     pub fn sleep(dur: Duration) {
         let nanos = dur.as_nanos();
         assert!(nanos <= u64::MAX as u128);
diff --git a/library/std/src/sys/pal/windows/c.rs b/library/std/src/sys/pal/windows/c.rs
@@ -344,6 +344,12 @@ compat_fn_with_fallback! {
         SetLastError(ERROR_CALL_NOT_IMPLEMENTED as DWORD); E_NOTIMPL
     }
 
+    // >= Win10 1607
+    // https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-getthreaddescription
+    pub fn GetThreadDescription(hthread: HANDLE, lpthreaddescription: *mut PWSTR) -> HRESULT {
+        SetLastError(ERROR_CALL_NOT_IMPLEMENTED as DWORD); E_NOTIMPL
+    }
+
     // >= Win8 / Server 2012
     // https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-getsystemtimepreciseasfiletime
     pub fn GetSystemTimePreciseAsFileTime(lpsystemtimeasfiletime: *mut FILETIME) -> () {
diff --git a/library/std/src/sys/pal/windows/c/bindings.txt b/library/std/src/sys/pal/windows/c/bindings.txt
@@ -1923,6 +1923,7 @@ Windows.Win32.Foundation.HANDLE_FLAG_INHERIT
 Windows.Win32.Foundation.HANDLE_FLAG_PROTECT_FROM_CLOSE
 Windows.Win32.Foundation.HANDLE_FLAGS
 Windows.Win32.Foundation.HMODULE
+Windows.Win32.Foundation.LocalFree
 Windows.Win32.Foundation.MAX_PATH
 Windows.Win32.Foundation.NO_ERROR
 Windows.Win32.Foundation.NTSTATUS
diff --git a/library/std/src/sys/pal/windows/c/windows_sys.rs b/library/std/src/sys/pal/windows/c/windows_sys.rs
@@ -379,6 +379,10 @@ extern "system" {
     ) -> BOOL;
 }
 #[link(name = "kernel32")]
+extern "system" {
+    pub fn LocalFree(hmem: HLOCAL) -> HLOCAL;
+}
+#[link(name = "kernel32")]
 extern "system" {
     pub fn MoveFileExW(
         lpexistingfilename: PCWSTR,
@@ -3441,6 +3445,7 @@ pub type HANDLE_FLAGS = u32;
 pub const HANDLE_FLAG_INHERIT: HANDLE_FLAGS = 1u32;
 pub const HANDLE_FLAG_PROTECT_FROM_CLOSE: HANDLE_FLAGS = 2u32;
 pub const HIGH_PRIORITY_CLASS: PROCESS_CREATION_FLAGS = 128u32;
+pub type HLOCAL = *mut ::core::ffi::c_void;
 pub type HMODULE = *mut ::core::ffi::c_void;
 pub type HRESULT = i32;
 pub const IDLE_PRIORITY_CLASS: PROCESS_CREATION_FLAGS = 64u32;
diff --git a/library/std/src/sys/pal/windows/thread.rs b/library/std/src/sys/pal/windows/thread.rs
@@ -9,7 +9,7 @@ use crate::sys::handle::Handle;
 use crate::sys::stack_overflow;
 use crate::sys_common::FromInner;
 use crate::time::Duration;
-
+use alloc::ffi::CString;
 use core::ffi::c_void;
 
 use super::time::WaitableTimer;
@@ -71,6 +71,29 @@ impl Thread {
         };
     }
 
+    pub fn get_name() -> Option<CString> {
+        unsafe {
+            let mut ptr = core::ptr::null_mut();
+            let result = c::GetThreadDescription(c::GetCurrentThread(), &mut ptr);
+            if result < 0 {
+                return None;
+            }
+            let name = String::from_utf16_lossy({
+                let mut len = 0;
+                while *ptr.add(len) != 0 {
+                    len += 1;
+                }
+                core::slice::from_raw_parts(ptr, len)
+            })
+            .into_bytes();
+            // Attempt to free the memory.
+            // This should never fail but if it does then there's not much we can do about it.
+            let result = c::LocalFree(ptr.cast::<c_void>());
+            debug_assert!(result.is_null());
+            if name.is_empty() { None } else { Some(CString::from_vec_unchecked(name)) }
+        }
+    }
+
     pub fn join(self) {
         let rc = unsafe { c::WaitForSingleObject(self.handle.as_raw_handle(), c::INFINITE) };
         if rc == c::WAIT_FAILED {
diff --git a/library/std/src/sys/pal/xous/thread.rs b/library/std/src/sys/pal/xous/thread.rs
@@ -1,4 +1,4 @@
-use crate::ffi::CStr;
+use crate::ffi::{CStr, CString};
 use crate::io;
 use crate::num::NonZero;
 use crate::os::xous::ffi::{
@@ -113,6 +113,10 @@ impl Thread {
         // nope
     }
 
+    pub fn get_name() -> Option<CString> {
+        None
+    }
+
     pub fn sleep(dur: Duration) {
         // Because the sleep server works on units of `usized milliseconds`, split
         // the messages up into these chunks. This means we may run into issues
diff --git a/library/std/src/sys_common/thread_info.rs b/library/std/src/sys_common/thread_info.rs
@@ -1,6 +1,7 @@
 #![allow(dead_code)] // stack_guard isn't used right now on all platforms
 
 use crate::cell::OnceCell;
+use crate::sys;
 use crate::sys::thread::guard::Guard;
 use crate::thread::Thread;
 
@@ -23,7 +24,8 @@ impl ThreadInfo {
     {
         THREAD_INFO
             .try_with(move |thread_info| {
-                let thread = thread_info.thread.get_or_init(|| Thread::new(None));
+                let thread =
+                    thread_info.thread.get_or_init(|| Thread::new(sys::thread::Thread::get_name()));
                 f(thread, &thread_info.stack_guard)
             })
             .ok()
diff --git a/library/std/src/thread/tests.rs b/library/std/src/thread/tests.rs
@@ -69,6 +69,25 @@ fn test_named_thread_truncation() {
     result.unwrap().join().unwrap();
 }
 
+#[cfg(any(
+    target_os = "windows",
+    target_os = "linux",
+    target_os = "macos",
+    target_os = "ios",
+    target_os = "tvos",
+    target_os = "watchos"
+))]
+#[test]
+fn test_get_os_named_thread() {
+    use crate::sys::thread::Thread;
+    let handler = thread::spawn(|| {
+        let name = c"test me please";
+        Thread::set_name(name);
+        assert_eq!(name, Thread::get_name().unwrap().as_c_str());
+    });
+    handler.join().unwrap();
+}
+
 #[test]
 #[should_panic]
 fn test_invalid_named_thread() {
