win32_ecoqos/utils/
thread.rs

1use std::{
2    ffi::{OsStr, OsString},
3    os::windows::ffi::OsStringExt,
4};
5
6use windows::Win32::{
7    Foundation::{CloseHandle, HANDLE},
8    System::{
9        Diagnostics::ToolHelp::{
10            CreateToolhelp32Snapshot, Thread32First, Thread32Next, TH32CS_SNAPTHREAD, THREADENTRY32,
11        },
12        Threading::{GetThreadDescription, OpenThread, THREAD_QUERY_LIMITED_INFORMATION},
13    },
14};
15
16#[derive(Debug, PartialEq, Eq)]
17/// process information from snapshot.
18pub struct Thread {
19    /// win32 thread id
20    pub thread_id: u32,
21    /// owner win32 process id of this thread
22    pub owner_process_id: u32,
23}
24
25/// Snapshot of win32 threads.
26///
27/// ```rust
28/// use std::ffi::OsString;
29/// use win32_ecoqos::utils::Threads;
30///
31/// let _ = std::thread::Builder::new()
32///     .name("hello".into())
33///     .spawn(|| loop {
34///         std::thread::sleep(std::time::Duration::from_secs(60));
35///     });
36///
37/// let snapshot = Threads::try_new().unwrap();
38/// assert!(
39///     snapshot
40///         .find_thread_by_name(&OsString::from("hello"), true)
41///         .next()
42///         .is_some()
43/// );
44/// ```
45#[derive(Debug)]
46pub struct Threads {
47    snapshot: HANDLE,
48    last_entry: Option<THREADENTRY32>,
49}
50
51impl Drop for Threads {
52    fn drop(&mut self) {
53        let _ = unsafe { CloseHandle(self.snapshot) };
54    }
55}
56
57impl Thread {
58    pub fn get_name(&self) -> windows_result::Result<OsString> {
59        unsafe {
60            let hthread = OpenThread(THREAD_QUERY_LIMITED_INFORMATION, false, self.thread_id)?;
61            let description = GetThreadDescription(hthread)?;
62
63            Ok(OsString::from_wide(description.as_wide()))
64        }
65    }
66}
67
68impl Threads {
69    /// create a new snapshop to find threads
70    pub fn try_new() -> windows_result::Result<Self> {
71        let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, std::process::id()) }?;
72        Ok(Self {
73            snapshot,
74            last_entry: None,
75        })
76    }
77
78    /// find a thread of current process, by it's name
79    pub fn find_thread_by_name<'a>(
80        self,
81        thread_name: &'a OsStr,
82        full_match: bool,
83    ) -> impl Iterator<Item = Thread> + 'a {
84        let process_id = std::process::id();
85        self.filter(move |t| t.owner_process_id == process_id)
86            .filter(move |t| {
87                t.get_name().is_ok_and(|name| {
88                    if full_match {
89                        &name == thread_name
90                    } else {
91                        name.to_string_lossy()
92                            .contains(thread_name.to_string_lossy().as_ref())
93                    }
94                })
95            })
96    }
97}
98
99impl Iterator for Threads {
100    type Item = Thread;
101
102    fn next(&mut self) -> Option<Self::Item> {
103        let first = self.last_entry.is_none();
104        let mut entry = self.last_entry.take().unwrap_or(THREADENTRY32 {
105            dwSize: size_of::<THREADENTRY32>() as u32,
106            ..Default::default()
107        });
108
109        unsafe {
110            if first {
111                Thread32First(self.snapshot, &mut entry as *mut _)
112            } else {
113                Thread32Next(self.snapshot, &mut entry as *mut _)
114            }
115        }
116        .ok()
117        .map(|_| entry)
118        .inspect(|entry| self.last_entry = Some(*entry))
119        .map(
120            |THREADENTRY32 {
121                 th32ThreadID,
122                 th32OwnerProcessID,
123                 ..
124             }| {
125                Thread {
126                    thread_id: th32ThreadID,
127                    owner_process_id: th32OwnerProcessID,
128                }
129            },
130        )
131    }
132}
win32_ecoqos/utils/thread.rs

win32_ecoqos/utils/
thread.rs
