pitchfork_cli/
procs.rs

1use crate::Result;
2use miette::IntoDiagnostic;
3use once_cell::sync::Lazy;
4use std::sync::Mutex;
5use sysinfo::ProcessesToUpdate;
6#[cfg(unix)]
7use sysinfo::Signal;
8
9pub struct Procs {
10    system: Mutex<sysinfo::System>,
11}
12
13pub static PROCS: Lazy<Procs> = Lazy::new(Procs::new);
14
15impl Default for Procs {
16    fn default() -> Self {
17        Self::new()
18    }
19}
20
21impl Procs {
22    pub fn new() -> Self {
23        let procs = Self {
24            system: Mutex::new(sysinfo::System::new()),
25        };
26        procs.refresh_processes();
27        procs
28    }
29
30    fn lock_system(&self) -> std::sync::MutexGuard<'_, sysinfo::System> {
31        self.system.lock().unwrap_or_else(|poisoned| {
32            warn!("System mutex was poisoned, recovering");
33            poisoned.into_inner()
34        })
35    }
36
37    pub fn title(&self, pid: u32) -> Option<String> {
38        self.lock_system()
39            .process(sysinfo::Pid::from_u32(pid))
40            .map(|p| p.name().to_string_lossy().to_string())
41    }
42
43    pub fn is_running(&self, pid: u32) -> bool {
44        self.lock_system()
45            .process(sysinfo::Pid::from_u32(pid))
46            .is_some()
47    }
48
49    pub fn all_children(&self, pid: u32) -> Vec<u32> {
50        let system = self.lock_system();
51        let all = system.processes();
52        let mut children = vec![];
53        for (child_pid, process) in all {
54            let mut process = process;
55            while let Some(parent) = process.parent() {
56                if parent == sysinfo::Pid::from_u32(pid) {
57                    children.push(child_pid.as_u32());
58                    break;
59                }
60                match system.process(parent) {
61                    Some(p) => process = p,
62                    None => break,
63                }
64            }
65        }
66        children
67    }
68
69    pub async fn kill_async(&self, pid: u32) -> Result<bool> {
70        let result = tokio::task::spawn_blocking(move || PROCS.kill(pid))
71            .await
72            .into_diagnostic()?;
73        Ok(result)
74    }
75
76    fn kill(&self, pid: u32) -> bool {
77        if let Some(process) = self.lock_system().process(sysinfo::Pid::from_u32(pid)) {
78            debug!("killing process {pid}");
79            #[cfg(windows)]
80            process.kill();
81            #[cfg(unix)]
82            process.kill_with(Signal::Term);
83            process.wait();
84            true
85        } else {
86            false
87        }
88    }
89
90    pub(crate) fn refresh_processes(&self) {
91        self.lock_system()
92            .refresh_processes(ProcessesToUpdate::All, true);
93    }
94
95    /// Refresh only specific PIDs instead of all processes.
96    /// More efficient when you only need to check a small set of known PIDs.
97    pub(crate) fn refresh_pids(&self, pids: &[u32]) {
98        let sysinfo_pids: Vec<sysinfo::Pid> =
99            pids.iter().map(|p| sysinfo::Pid::from_u32(*p)).collect();
100        self.lock_system()
101            .refresh_processes(ProcessesToUpdate::Some(&sysinfo_pids), true);
102    }
103
104    /// Get process stats (cpu%, memory bytes, uptime secs, disk I/O) for a given PID
105    pub fn get_stats(&self, pid: u32) -> Option<ProcessStats> {
106        let system = self.lock_system();
107        system.process(sysinfo::Pid::from_u32(pid)).map(|p| {
108            let now = std::time::SystemTime::now()
109                .duration_since(std::time::UNIX_EPOCH)
110                .map(|d| d.as_secs())
111                .unwrap_or(0);
112            let disk = p.disk_usage();
113            ProcessStats {
114                cpu_percent: p.cpu_usage(),
115                memory_bytes: p.memory(),
116                uptime_secs: now.saturating_sub(p.start_time()),
117                disk_read_bytes: disk.read_bytes,
118                disk_write_bytes: disk.written_bytes,
119            }
120        })
121    }
122
123    /// Get extended process information for a given PID
124    pub fn get_extended_stats(&self, pid: u32) -> Option<ExtendedProcessStats> {
125        let system = self.lock_system();
126        system.process(sysinfo::Pid::from_u32(pid)).map(|p| {
127            let now = std::time::SystemTime::now()
128                .duration_since(std::time::UNIX_EPOCH)
129                .map(|d| d.as_secs())
130                .unwrap_or(0);
131            let disk = p.disk_usage();
132
133            ExtendedProcessStats {
134                name: p.name().to_string_lossy().to_string(),
135                exe_path: p.exe().map(|e| e.to_string_lossy().to_string()),
136                cwd: p.cwd().map(|c| c.to_string_lossy().to_string()),
137                environ: p
138                    .environ()
139                    .iter()
140                    .take(20)
141                    .map(|s| s.to_string_lossy().to_string())
142                    .collect(),
143                status: format!("{:?}", p.status()),
144                cpu_percent: p.cpu_usage(),
145                memory_bytes: p.memory(),
146                virtual_memory_bytes: p.virtual_memory(),
147                uptime_secs: now.saturating_sub(p.start_time()),
148                start_time: p.start_time(),
149                disk_read_bytes: disk.read_bytes,
150                disk_write_bytes: disk.written_bytes,
151                parent_pid: p.parent().map(|pp| pp.as_u32()),
152                thread_count: p.tasks().map(|t| t.len()).unwrap_or(0),
153                user_id: p.user_id().map(|u| u.to_string()),
154            }
155        })
156    }
157}
158
159#[derive(Debug, Clone, Copy)]
160pub struct ProcessStats {
161    pub cpu_percent: f32,
162    pub memory_bytes: u64,
163    pub uptime_secs: u64,
164    pub disk_read_bytes: u64,
165    pub disk_write_bytes: u64,
166}
167
168impl ProcessStats {
169    pub fn memory_display(&self) -> String {
170        format_bytes(self.memory_bytes)
171    }
172
173    pub fn cpu_display(&self) -> String {
174        format!("{:.1}%", self.cpu_percent)
175    }
176
177    pub fn uptime_display(&self) -> String {
178        format_duration(self.uptime_secs)
179    }
180
181    pub fn disk_read_display(&self) -> String {
182        format_bytes_per_sec(self.disk_read_bytes)
183    }
184
185    pub fn disk_write_display(&self) -> String {
186        format_bytes_per_sec(self.disk_write_bytes)
187    }
188}
189
190/// Extended process stats with more detailed information
191#[derive(Debug, Clone)]
192pub struct ExtendedProcessStats {
193    pub name: String,
194    pub exe_path: Option<String>,
195    pub cwd: Option<String>,
196    pub environ: Vec<String>,
197    pub status: String,
198    pub cpu_percent: f32,
199    pub memory_bytes: u64,
200    pub virtual_memory_bytes: u64,
201    pub uptime_secs: u64,
202    pub start_time: u64,
203    pub disk_read_bytes: u64,
204    pub disk_write_bytes: u64,
205    pub parent_pid: Option<u32>,
206    pub thread_count: usize,
207    pub user_id: Option<String>,
208}
209
210impl ExtendedProcessStats {
211    pub fn memory_display(&self) -> String {
212        format_bytes(self.memory_bytes)
213    }
214
215    pub fn virtual_memory_display(&self) -> String {
216        format_bytes(self.virtual_memory_bytes)
217    }
218
219    pub fn cpu_display(&self) -> String {
220        format!("{:.1}%", self.cpu_percent)
221    }
222
223    pub fn uptime_display(&self) -> String {
224        format_duration(self.uptime_secs)
225    }
226
227    pub fn start_time_display(&self) -> String {
228        use std::time::{Duration, UNIX_EPOCH};
229        let datetime = UNIX_EPOCH + Duration::from_secs(self.start_time);
230        chrono::DateTime::<chrono::Local>::from(datetime)
231            .format("%Y-%m-%d %H:%M:%S")
232            .to_string()
233    }
234
235    pub fn disk_read_display(&self) -> String {
236        format_bytes_per_sec(self.disk_read_bytes)
237    }
238
239    pub fn disk_write_display(&self) -> String {
240        format_bytes_per_sec(self.disk_write_bytes)
241    }
242}
243
244fn format_bytes(bytes: u64) -> String {
245    if bytes < 1024 {
246        format!("{bytes}B")
247    } else if bytes < 1024 * 1024 {
248        format!("{:.1}KB", bytes as f64 / 1024.0)
249    } else if bytes < 1024 * 1024 * 1024 {
250        format!("{:.1}MB", bytes as f64 / (1024.0 * 1024.0))
251    } else {
252        format!("{:.1}GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
253    }
254}
255
256fn format_duration(secs: u64) -> String {
257    if secs < 60 {
258        format!("{secs}s")
259    } else if secs < 3600 {
260        format!("{}m {}s", secs / 60, secs % 60)
261    } else if secs < 86400 {
262        let hours = secs / 3600;
263        let mins = (secs % 3600) / 60;
264        format!("{hours}h {mins}m")
265    } else {
266        let days = secs / 86400;
267        let hours = (secs % 86400) / 3600;
268        format!("{days}d {hours}h")
269    }
270}
271
272fn format_bytes_per_sec(bytes: u64) -> String {
273    if bytes < 1024 {
274        format!("{bytes}B/s")
275    } else if bytes < 1024 * 1024 {
276        format!("{:.1}KB/s", bytes as f64 / 1024.0)
277    } else if bytes < 1024 * 1024 * 1024 {
278        format!("{:.1}MB/s", bytes as f64 / (1024.0 * 1024.0))
279    } else {
280        format!("{:.1}GB/s", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
281    }
282}
pitchfork_cli/procs.rs

pitchfork_cli/
procs.rs
