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
124#[derive(Debug, Clone, Copy)]
125pub struct ProcessStats {
126    pub cpu_percent: f32,
127    pub memory_bytes: u64,
128    pub uptime_secs: u64,
129    pub disk_read_bytes: u64,
130    pub disk_write_bytes: u64,
131}
132
133impl ProcessStats {
134    pub fn memory_display(&self) -> String {
135        let bytes = self.memory_bytes;
136        if bytes < 1024 {
137            format!("{}B", bytes)
138        } else if bytes < 1024 * 1024 {
139            format!("{:.1}KB", bytes as f64 / 1024.0)
140        } else if bytes < 1024 * 1024 * 1024 {
141            format!("{:.1}MB", bytes as f64 / (1024.0 * 1024.0))
142        } else {
143            format!("{:.1}GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
144        }
145    }
146
147    pub fn cpu_display(&self) -> String {
148        format!("{:.1}%", self.cpu_percent)
149    }
150
151    pub fn uptime_display(&self) -> String {
152        let secs = self.uptime_secs;
153        if secs < 60 {
154            format!("{}s", secs)
155        } else if secs < 3600 {
156            format!("{}m {}s", secs / 60, secs % 60)
157        } else if secs < 86400 {
158            let hours = secs / 3600;
159            let mins = (secs % 3600) / 60;
160            format!("{}h {}m", hours, mins)
161        } else {
162            let days = secs / 86400;
163            let hours = (secs % 86400) / 3600;
164            format!("{}d {}h", days, hours)
165        }
166    }
167
168    pub fn disk_read_display(&self) -> String {
169        Self::format_bytes_per_sec(self.disk_read_bytes)
170    }
171
172    pub fn disk_write_display(&self) -> String {
173        Self::format_bytes_per_sec(self.disk_write_bytes)
174    }
175
176    fn format_bytes_per_sec(bytes: u64) -> String {
177        if bytes < 1024 {
178            format!("{}B/s", bytes)
179        } else if bytes < 1024 * 1024 {
180            format!("{:.1}KB/s", bytes as f64 / 1024.0)
181        } else if bytes < 1024 * 1024 * 1024 {
182            format!("{:.1}MB/s", bytes as f64 / (1024.0 * 1024.0))
183        } else {
184            format!("{:.1}GB/s", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
185        }
186    }
187}
pitchfork_cli/procs.rs

pitchfork_cli/
procs.rs
