Skip to main content

rustpython_host_env/
thread.rs

1#[cfg(any(target_os = "linux", target_os = "macos"))]
2use alloc::ffi::CString;
3
4/// OS thread-name cap, matching CPython `_thread._NAME_MAXLEN`.
5pub const NAME_MAXLEN: usize = {
6    if cfg!(windows) {
7        100
8    } else if cfg!(target_os = "linux") {
9        15
10    } else if cfg!(target_os = "macos") {
11        63
12    } else {
13        16
14    }
15};
16
17#[cfg(any(target_os = "linux", target_os = "macos"))]
18fn truncate_thread_name_bytes(name: &[u8]) -> &[u8] {
19    let name = name.split(|&b| b == 0).next().unwrap_or(b"");
20    name.get(..NAME_MAXLEN.min(name.len())).unwrap_or(b"")
21}
22
23#[cfg(unix)]
24pub fn current_thread_id() -> u64 {
25    unsafe { libc::pthread_self() as u64 }
26}
27
28#[cfg(windows)]
29pub fn current_thread_id() -> u64 {
30    unsafe { windows_sys::Win32::System::Threading::GetCurrentThreadId() as u64 }
31}
32
33/// Kernel-assigned thread id, unique system-wide until the thread exits.
34#[cfg(windows)]
35pub fn native_id() -> u64 {
36    unsafe { windows_sys::Win32::System::Threading::GetCurrentThreadId() as u64 }
37}
38
39#[cfg(any(target_os = "macos", target_os = "ios"))]
40pub fn native_id() -> u64 {
41    let mut tid = 0u64;
42    unsafe {
43        libc::pthread_threadid_np(0 as libc::pthread_t, &mut tid);
44    }
45    tid
46}
47
48#[cfg(any(target_os = "linux", target_os = "android"))]
49pub fn native_id() -> u64 {
50    unsafe { libc::syscall(libc::SYS_gettid) as u64 }
51}
52
53#[cfg(target_os = "freebsd")]
54pub fn native_id() -> u64 {
55    unsafe { libc::pthread_getthreadid_np() as u64 }
56}
57
58#[cfg(target_os = "openbsd")]
59pub fn native_id() -> u64 {
60    unsafe { libc::getthrid() as u64 }
61}
62
63#[cfg(target_os = "netbsd")]
64pub fn native_id() -> u64 {
65    unsafe { libc::_lwp_self() as u64 }
66}
67
68#[cfg(windows)]
69#[allow(clippy::not_unsafe_ptr_arg_deref)]
70pub fn thread_id_from_handle(handle: *mut core::ffi::c_void) -> u64 {
71    unsafe { windows_sys::Win32::System::Threading::GetThreadId(handle) as u64 }
72}
73
74#[cfg(any(target_os = "linux", target_os = "macos"))]
75pub fn set_current_thread_name(name: &str) {
76    set_current_thread_name_bytes(name.as_bytes());
77}
78
79/// Set the OS thread name from filesystem-encoded bytes (NUL-truncated).
80#[cfg(any(target_os = "linux", target_os = "macos"))]
81pub fn set_current_thread_name_bytes(name: &[u8]) {
82    let truncated = truncate_thread_name_bytes(name);
83    let Ok(c_name) = CString::new(truncated) else {
84        return;
85    };
86    unsafe {
87        #[cfg(target_os = "linux")]
88        libc::pthread_setname_np(libc::pthread_self(), c_name.as_ptr());
89        #[cfg(target_os = "macos")]
90        libc::pthread_setname_np(c_name.as_ptr());
91    }
92}
93
94#[cfg(windows)]
95pub fn set_current_thread_name(name: &str) {
96    let wide: Vec<u16> = name.encode_utf16().chain(core::iter::once(0)).collect();
97    let _ = set_current_thread_name_wide(&wide);
98}
99
100#[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
101pub fn set_current_thread_name(_name: &str) {}
102
103/// The current thread's description as UTF-16 code units, without the terminator.
104#[cfg(windows)]
105pub fn current_thread_name_wide() -> std::io::Result<Vec<u16>> {
106    use windows_sys::Win32::{
107        Foundation::LocalFree,
108        System::Threading::{GetCurrentThread, GetThreadDescription},
109    };
110
111    let mut raw = core::ptr::null_mut();
112    let status = unsafe { GetThreadDescription(GetCurrentThread(), &mut raw) };
113    if status < 0 {
114        return Err(std::io::Error::from_raw_os_error(status));
115    }
116    if raw.is_null() {
117        return Ok(Vec::new());
118    }
119    let mut len = 0usize;
120    unsafe {
121        while *raw.add(len) != 0 {
122            len += 1;
123        }
124    }
125    let name = unsafe { core::slice::from_raw_parts(raw, len) }.to_vec();
126    unsafe { LocalFree(raw.cast()) };
127    Ok(name)
128}
129
130/// `SetThreadDescription` on the calling thread. `name` must be NUL-terminated.
131#[cfg(windows)]
132pub fn set_current_thread_name_wide(name: &[u16]) -> std::io::Result<()> {
133    use windows_sys::Win32::System::Threading::{GetCurrentThread, SetThreadDescription};
134
135    if name.last() != Some(&0) {
136        return Err(std::io::Error::new(
137            std::io::ErrorKind::InvalidInput,
138            "thread name must be NUL-terminated",
139        ));
140    }
141    let status = unsafe { SetThreadDescription(GetCurrentThread(), name.as_ptr()) };
142    if status < 0 {
143        Err(std::io::Error::from_raw_os_error(status))
144    } else {
145        Ok(())
146    }
147}
148
149/// The current thread name as bytes, without the terminator.
150#[cfg(any(target_os = "linux", target_os = "macos"))]
151pub fn current_thread_name(buf_len: usize) -> std::io::Result<Vec<u8>> {
152    let mut buffer = vec![0u8; buf_len];
153    let status = unsafe {
154        libc::pthread_getname_np(
155            libc::pthread_self(),
156            buffer.as_mut_ptr().cast(),
157            buffer.len(),
158        )
159    };
160    if status != 0 {
161        return Err(std::io::Error::from_raw_os_error(status));
162    }
163    let len = buffer.iter().position(|&b| b == 0).unwrap_or(buffer.len());
164    buffer.truncate(len);
165    Ok(buffer)
166}