proc_ctl/
proc_query.rs

1use crate::common::{resolve_pid, MaybeHasPid};
2use crate::{Pid, ProcCtlError, ProcCtlResult};
3use std::path::PathBuf;
4use std::process::Child;
5use std::sync::Mutex;
6use std::sync::OnceLock;
7use sysinfo::{Process, ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System};
8
9/// Information about a process
10#[derive(Debug, Clone)]
11pub struct ProcInfo {
12    /// The name
13    pub name: String,
14    /// The command used to launch the process
15    pub cmd: Vec<String>,
16    /// The path to the executable the process is running
17    pub exe: Option<PathBuf>,
18    /// The process ID
19    pub pid: Pid,
20    /// Parent process ID if relevant
21    pub parent: Option<Pid>,
22    /// Environment variables available to the process
23    pub env: Vec<String>,
24    /// The current working directory of the process
25    pub cwd: Option<PathBuf>,
26}
27
28/// Get information about a process
29#[derive(Debug)]
30pub struct ProcQuery {
31    process_id: Option<Pid>,
32    name: Option<String>,
33    min_num_children: Option<usize>,
34}
35
36impl ProcQuery {
37    /// Create a new process query
38    pub fn new() -> Self {
39        ProcQuery {
40            process_id: None,
41            name: None,
42            min_num_children: None,
43        }
44    }
45
46    /// Set the process ID to match
47    ///
48    /// One of this, [ProcQuery::process_name] or [ProcQuery::process_id_from_child] must be called before the query is usable.
49    pub fn process_id(mut self, pid: Pid) -> Self {
50        self.process_id = Some(pid);
51        self
52    }
53
54    /// Set the process name to match
55    ///
56    /// One of this, [ProcQuery::process_id] or [ProcQuery::process_id_from_child] must be called before the query is usable.
57    pub fn process_name(mut self, name: impl AsRef<str>) -> Self {
58        let name = name.as_ref().to_string();
59        #[cfg(target_os = "windows")]
60        let name = {
61            let mut name = name;
62            if !name.ends_with(".exe") {
63                name.push_str(".exe");
64            }
65            name
66        };
67        self.name = Some(name);
68        self
69    }
70
71    /// Get the process ID of a child process
72    ///
73    /// Either this function or `process_id` are required to be called before the query is usable.
74    pub fn process_id_from_child(self, child: &Child) -> Self {
75        self.process_id(child.id())
76    }
77
78    /// Require at least `num_children` children to have been started by the matched process for the query to succeed.
79    pub fn expect_min_num_children(mut self, num_children: usize) -> Self {
80        self.min_num_children = Some(num_children);
81        self
82    }
83
84    /// List all processes matching the current filters.
85    pub fn list_processes(&self) -> ProcCtlResult<Vec<ProcInfo>> {
86        let mut sys_handle = sys_handle().lock().unwrap();
87        sys_handle.refresh_processes_specifics(
88            ProcessesToUpdate::All,
89            true,
90            ProcessRefreshKind::everything(),
91        );
92        let processes = sys_handle.processes();
93
94        let infos: Vec<ProcInfo> = processes
95            .values()
96            .filter(|p| {
97                // Should ignore threads, we're only looking for top-level processes.
98                if p.thread_kind().is_some() {
99                    return false;
100                }
101
102                if let Some(pid) = self.process_id {
103                    if p.pid().as_u32() != pid {
104                        return false;
105                    }
106                }
107
108                if let Some(name) = &self.name {
109                    if p.name().to_string_lossy().as_ref() != name {
110                        return false;
111                    }
112                }
113
114                true
115            })
116            .map(|p| p.into())
117            .collect();
118
119        Ok(infos)
120    }
121
122    /// Find the children of the selected process
123    pub fn children(&self) -> ProcCtlResult<Vec<ProcInfo>> {
124        let pid = resolve_pid(self)?;
125
126        let mut sys_handle = sys_handle().lock().unwrap();
127        sys_handle.refresh_processes(ProcessesToUpdate::All, true);
128        let processes = sys_handle.processes();
129        let children: Vec<ProcInfo> = processes
130            .values()
131            .filter(|p| p.parent() == Some(sysinfo::Pid::from(pid as usize)))
132            .map(|p| p.into())
133            .collect();
134
135        if let Some(num) = &self.min_num_children {
136            if children.len() < *num {
137                return Err(ProcCtlError::TooFewChildren(children.len(), *num));
138            }
139        }
140
141        Ok(children)
142    }
143
144    /// Execute the query and retry until it succeeds or exhausts the configured retries
145    #[cfg(feature = "resilience")]
146    pub fn children_with_retry_sync(
147        &self,
148        delay: std::time::Duration,
149        count: usize,
150    ) -> ProcCtlResult<Vec<ProcInfo>> {
151        retry::retry(retry::delay::Fixed::from(delay).take(count), || {
152            self.children()
153        })
154        .map_err(|e| e.error)
155    }
156
157    /// Async equivalent of `children_with_retry_sync`
158    #[cfg(feature = "async")]
159    #[async_recursion::async_recursion]
160    pub async fn children_with_retry(
161        &self,
162        delay: std::time::Duration,
163        count: usize,
164    ) -> ProcCtlResult<Vec<ProcInfo>> {
165        match self.children() {
166            Ok(infos) => Ok(infos),
167            Err(e) => {
168                if count == 0 {
169                    Err(e)
170                } else {
171                    tokio::time::sleep(delay).await;
172                    self.children_with_retry(delay, count - 1).await
173                }
174            }
175        }
176    }
177}
178
179fn sys_handle() -> &'static Mutex<System> {
180    static SYS_HANDLE: OnceLock<Mutex<System>> = OnceLock::new();
181    SYS_HANDLE.get_or_init(|| {
182        let mut sys = System::new_with_specifics(
183            RefreshKind::new().with_processes(ProcessRefreshKind::new()),
184        );
185        sys.refresh_processes(ProcessesToUpdate::All, true);
186
187        Mutex::new(sys)
188    })
189}
190
191impl From<&Process> for ProcInfo {
192    fn from(value: &Process) -> Self {
193        ProcInfo {
194            name: value.name().to_string_lossy().to_string(),
195            cmd: value
196                .cmd()
197                .iter()
198                .map(|p| p.to_string_lossy().to_string())
199                .collect(),
200            exe: value.exe().map(|p| p.to_owned()),
201            pid: value.pid().as_u32() as Pid,
202            parent: value.parent().map(|p| p.as_u32() as Pid),
203            env: value
204                .environ()
205                .iter()
206                .map(|p| p.to_string_lossy().to_string())
207                .collect(),
208            cwd: value.cwd().map(|p| p.to_owned()),
209        }
210    }
211}
212
213impl MaybeHasPid for ProcQuery {
214    fn get_pid(&self) -> Option<Pid> {
215        self.process_id
216    }
217}
218
219impl Default for ProcQuery {
220    fn default() -> Self {
221        ProcQuery::new()
222    }
223}
proc_ctl/proc_query.rs

proc_ctl/
proc_query.rs
