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    /// Get process stats (cpu%, memory bytes, uptime secs, disk I/O) for a given PID
96    pub fn get_stats(&self, pid: u32) -> Option<ProcessStats> {
97        let system = self.lock_system();
98        system.process(sysinfo::Pid::from_u32(pid)).map(|p| {
99            let now = std::time::SystemTime::now()
100                .duration_since(std::time::UNIX_EPOCH)
101                .map(|d| d.as_secs())
102                .unwrap_or(0);
103            let disk = p.disk_usage();
104            ProcessStats {
105                cpu_percent: p.cpu_usage(),
106                memory_bytes: p.memory(),
107                uptime_secs: now.saturating_sub(p.start_time()),
108                disk_read_bytes: disk.read_bytes,
109                disk_write_bytes: disk.written_bytes,
110            }
111        })
112    }
113}
114
115#[derive(Debug, Clone, Copy)]
116pub struct ProcessStats {
117    pub cpu_percent: f32,
118    pub memory_bytes: u64,
119    pub uptime_secs: u64,
120    pub disk_read_bytes: u64,
121    pub disk_write_bytes: u64,
122}
123
124impl ProcessStats {
125    pub fn memory_display(&self) -> String {
126        let bytes = self.memory_bytes;
127        if bytes < 1024 {
128            format!("{}B", bytes)
129        } else if bytes < 1024 * 1024 {
130            format!("{:.1}KB", bytes as f64 / 1024.0)
131        } else if bytes < 1024 * 1024 * 1024 {
132            format!("{:.1}MB", bytes as f64 / (1024.0 * 1024.0))
133        } else {
134            format!("{:.1}GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
135        }
136    }
137
138    pub fn cpu_display(&self) -> String {
139        format!("{:.1}%", self.cpu_percent)
140    }
141
142    pub fn uptime_display(&self) -> String {
143        let secs = self.uptime_secs;
144        if secs < 60 {
145            format!("{}s", secs)
146        } else if secs < 3600 {
147            format!("{}m {}s", secs / 60, secs % 60)
148        } else if secs < 86400 {
149            let hours = secs / 3600;
150            let mins = (secs % 3600) / 60;
151            format!("{}h {}m", hours, mins)
152        } else {
153            let days = secs / 86400;
154            let hours = (secs % 86400) / 3600;
155            format!("{}d {}h", days, hours)
156        }
157    }
158
159    pub fn disk_read_display(&self) -> String {
160        Self::format_bytes_per_sec(self.disk_read_bytes)
161    }
162
163    pub fn disk_write_display(&self) -> String {
164        Self::format_bytes_per_sec(self.disk_write_bytes)
165    }
166
167    fn format_bytes_per_sec(bytes: u64) -> String {
168        if bytes < 1024 {
169            format!("{}B/s", bytes)
170        } else if bytes < 1024 * 1024 {
171            format!("{:.1}KB/s", bytes as f64 / 1024.0)
172        } else if bytes < 1024 * 1024 * 1024 {
173            format!("{:.1}MB/s", bytes as f64 / (1024.0 * 1024.0))
174        } else {
175            format!("{:.1}GB/s", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
176        }
177    }
178}