1use std::collections::{HashMap, VecDeque};
7
8use crate::snapshot::{ProcessInfo, Snapshot};
9use crate::source::ProcessSignal;
10
11pub const HISTORY_LEN: usize = 1024;
14
15const VIRTUAL_IFACE_PREFIXES: &[&str] = &["lo", "docker", "veth", "br-", "virbr"];
18
19#[derive(Debug, Default, Clone)]
21pub struct History {
22 samples: VecDeque<f64>,
23}
24
25impl History {
26 pub fn push(&mut self, value: f64) {
28 if self.samples.len() == HISTORY_LEN {
29 self.samples.pop_front();
30 }
31 self.samples.push_back(value);
32 }
33
34 pub fn last(&self, n: usize) -> Vec<f64> {
36 let skip = self.samples.len().saturating_sub(n);
37 self.samples.iter().skip(skip).copied().collect()
38 }
39
40 pub fn max_of_last(&self, n: usize) -> f64 {
42 self.last(n).into_iter().fold(0.0, f64::max)
43 }
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum SortKey {
49 Pid,
50 Name,
51 Threads,
52 User,
53 Memory,
54 Cpu,
55 Gpu,
56 Vram,
57}
58
59impl SortKey {
60 pub const ALL: [SortKey; 8] = [
62 SortKey::Pid,
63 SortKey::Name,
64 SortKey::Threads,
65 SortKey::User,
66 SortKey::Memory,
67 SortKey::Cpu,
68 SortKey::Gpu,
69 SortKey::Vram,
70 ];
71
72 pub fn descending_by_default(self) -> bool {
74 !matches!(self, SortKey::Pid | SortKey::Name | SortKey::User)
75 }
76}
77
78#[derive(Debug, Clone)]
80pub struct ProcessRow<'a> {
81 pub info: &'a ProcessInfo,
82 pub gpu_percent: Option<f32>,
84 pub vram_bytes: Option<u64>,
86}
87
88#[derive(Debug, Default, Clone)]
90pub struct MetricHistory {
91 pub cpu: History,
93 pub memory: History,
95 pub gpus: Vec<History>,
97 pub net_rx: History,
99 pub net_tx: History,
101 pub net_interface: Option<String>,
103}
104
105impl MetricHistory {
106 pub fn record(&mut self, snapshot: &Snapshot) {
108 self.cpu.push(snapshot.cpu.total_percent as f64);
109 let memory = &snapshot.memory;
110 if memory.total_bytes > 0 {
111 self.memory
112 .push(memory.used_bytes as f64 / memory.total_bytes as f64 * 100.0);
113 }
114 self.gpus.resize_with(snapshot.gpus.len(), History::default);
115 for (history, gpu) in self.gpus.iter_mut().zip(&snapshot.gpus) {
116 history.push(gpu.utilization_percent.unwrap_or(0.0) as f64);
117 }
118 if self.net_interface.is_none() {
119 self.net_interface = default_interface(snapshot);
120 }
121 let (rx, tx) = self
122 .net_interface
123 .as_deref()
124 .and_then(|name| snapshot.networks.iter().find(|n| n.name == name))
125 .map(|n| (n.rx_bytes_per_sec, n.tx_bytes_per_sec))
126 .unwrap_or((0.0, 0.0));
127 self.net_rx.push(rx);
128 self.net_tx.push(tx);
129 }
130
131 pub fn select_interface(&mut self, name: &str) {
133 if self.net_interface.as_deref() != Some(name) {
134 self.net_interface = Some(name.to_owned());
135 self.net_rx = History::default();
137 self.net_tx = History::default();
138 }
139 }
140
141 pub fn next_interface(&mut self, snapshot: &Snapshot) -> Option<String> {
144 let names: Vec<&str> = snapshot.networks.iter().map(|n| n.name.as_str()).collect();
145 if names.is_empty() {
146 return None;
147 }
148 let current = self
149 .net_interface
150 .as_deref()
151 .and_then(|name| names.iter().position(|n| *n == name));
152 let next = names[current.map_or(0, |i| (i + 1) % names.len())].to_owned();
153 self.select_interface(&next);
154 Some(next)
155 }
156}
157
158#[derive(Debug, Clone, Copy)]
160pub struct ProcessQuery<'q> {
161 pub sort_key: SortKey,
163 pub descending: bool,
165 pub filter: &'q str,
167 pub show_kernel_threads: bool,
169}
170
171pub fn process_rows<'a>(snapshot: &'a Snapshot, query: &ProcessQuery<'_>) -> Vec<ProcessRow<'a>> {
173 let mut gpu_usage: HashMap<u32, (Option<f32>, Option<u64>)> = HashMap::new();
175 for gpu in &snapshot.gpus {
176 for p in &gpu.processes {
177 let entry = gpu_usage.entry(p.pid).or_default();
178 if let Some(sm) = p.sm_percent {
179 entry.0 = Some(entry.0.map_or(sm, |v| v.max(sm)));
180 }
181 if let Some(mem) = p.memory_bytes {
182 entry.1 = Some(entry.1.unwrap_or(0) + mem);
183 }
184 }
185 }
186
187 let filter = query.filter.to_lowercase();
188 let mut rows: Vec<ProcessRow<'a>> = snapshot
189 .processes
190 .iter()
191 .filter(|p| query.show_kernel_threads || !p.kernel_thread)
192 .filter(|p| {
193 filter.is_empty()
194 || p.name.to_lowercase().contains(&filter)
195 || p.command.to_lowercase().contains(&filter)
196 || p.user.to_lowercase().contains(&filter)
197 || p.pid.to_string() == filter
198 })
199 .map(|info| {
200 let (gpu_percent, vram_bytes) = gpu_usage.get(&info.pid).copied().unwrap_or_default();
201 ProcessRow {
202 info,
203 gpu_percent,
204 vram_bytes,
205 }
206 })
207 .collect();
208
209 let key = query.sort_key;
210 rows.sort_by(|a, b| {
211 let ordering = match key {
212 SortKey::Pid => a.info.pid.cmp(&b.info.pid),
213 SortKey::Name => a.info.name.to_lowercase().cmp(&b.info.name.to_lowercase()),
214 SortKey::Threads => a.info.threads.cmp(&b.info.threads),
215 SortKey::User => a.info.user.cmp(&b.info.user),
216 SortKey::Memory => a.info.memory_bytes.cmp(&b.info.memory_bytes),
217 SortKey::Cpu => a.info.cpu_percent.total_cmp(&b.info.cpu_percent),
218 SortKey::Gpu => a
219 .gpu_percent
220 .unwrap_or(-1.0)
221 .total_cmp(&b.gpu_percent.unwrap_or(-1.0)),
222 SortKey::Vram => a.vram_bytes.cmp(&b.vram_bytes),
223 }
224 .then(a.info.pid.cmp(&b.info.pid));
226 if query.descending {
227 ordering.reverse()
228 } else {
229 ordering
230 }
231 });
232 rows
233}
234
235pub fn signal_name(signal: ProcessSignal) -> &'static str {
237 match signal {
238 ProcessSignal::Terminate => "SIGTERM",
239 ProcessSignal::Kill => "SIGKILL",
240 }
241}
242
243pub fn is_virtual_interface(name: &str) -> bool {
246 VIRTUAL_IFACE_PREFIXES.iter().any(|p| name.starts_with(p))
247}
248
249pub fn default_interface(snapshot: &Snapshot) -> Option<String> {
251 snapshot
252 .networks
253 .iter()
254 .filter(|n| !is_virtual_interface(&n.name))
255 .max_by_key(|n| n.rx_total_bytes + n.tx_total_bytes)
256 .or_else(|| snapshot.networks.first())
257 .map(|n| n.name.clone())
258}
259
260#[cfg(test)]
261mod tests {
262 use super::*;
263 use crate::snapshot::{GpuProcess, GpuStats, NetworkStats};
264
265 fn process(pid: u32, name: &str, cpu: f32, kernel: bool) -> ProcessInfo {
266 ProcessInfo {
267 pid,
268 name: name.to_owned(),
269 cpu_percent: cpu,
270 kernel_thread: kernel,
271 ..Default::default()
272 }
273 }
274
275 fn query(sort_key: SortKey) -> ProcessQuery<'static> {
276 ProcessQuery {
277 sort_key,
278 descending: sort_key.descending_by_default(),
279 filter: "",
280 show_kernel_threads: false,
281 }
282 }
283
284 #[test]
285 fn history_keeps_newest_samples() {
286 let mut history = History::default();
287 for i in 0..(HISTORY_LEN + 5) {
288 history.push(i as f64);
289 }
290 assert_eq!(
291 history.last(2),
292 vec![(HISTORY_LEN + 3) as f64, (HISTORY_LEN + 4) as f64]
293 );
294 assert_eq!(history.last(HISTORY_LEN * 2).len(), HISTORY_LEN);
295 assert_eq!(history.max_of_last(3), (HISTORY_LEN + 4) as f64);
296 }
297
298 #[test]
299 fn rows_filter_sort_and_join_gpu_usage() {
300 let snapshot = Snapshot {
301 processes: vec![
302 process(10, "bash", 5.0, false),
303 process(2, "kworker/0:1", 50.0, true),
304 process(30, "python", 1.0, false),
305 ],
306 gpus: vec![GpuStats {
307 processes: vec![GpuProcess {
308 pid: 30,
309 memory_bytes: Some(1 << 30),
310 sm_percent: Some(70.0),
311 }],
312 ..Default::default()
313 }],
314 ..Default::default()
315 };
316 let pids = |q: &ProcessQuery| {
317 process_rows(&snapshot, q)
318 .iter()
319 .map(|r| r.info.pid)
320 .collect::<Vec<_>>()
321 };
322 assert_eq!(pids(&query(SortKey::Cpu)), vec![10, 30]);
323 assert_eq!(pids(&query(SortKey::Vram)), vec![30, 10]);
324 assert_eq!(
325 pids(&ProcessQuery {
326 show_kernel_threads: true,
327 ..query(SortKey::Cpu)
328 }),
329 vec![2, 10, 30]
330 );
331 assert_eq!(
332 pids(&ProcessQuery {
333 filter: "PYT",
334 ..query(SortKey::Pid)
335 }),
336 vec![30]
337 );
338 let rows = process_rows(&snapshot, &query(SortKey::Pid));
339 assert_eq!(rows[1].vram_bytes, Some(1 << 30));
340 assert_eq!(rows[1].gpu_percent, Some(70.0));
341 }
342
343 #[test]
344 fn metric_history_records_and_switches_interfaces() {
345 let snapshot = Snapshot {
346 networks: vec![
347 NetworkStats {
348 name: "enp9s0".to_owned(),
349 rx_total_bytes: 10,
350 rx_bytes_per_sec: 100.0,
351 ..Default::default()
352 },
353 NetworkStats {
354 name: "wlp7s0".to_owned(),
355 tx_bytes_per_sec: 5.0,
356 ..Default::default()
357 },
358 ],
359 gpus: vec![GpuStats {
360 utilization_percent: Some(40.0),
361 ..Default::default()
362 }],
363 ..Default::default()
364 };
365 let mut history = MetricHistory::default();
366 history.record(&snapshot);
367 assert_eq!(history.net_interface.as_deref(), Some("enp9s0"));
368 assert_eq!(history.net_rx.last(1), vec![100.0]);
369 assert_eq!(history.gpus[0].last(1), vec![40.0]);
370 assert_eq!(history.next_interface(&snapshot).as_deref(), Some("wlp7s0"));
371 assert!(history.net_rx.last(10).is_empty());
373 history.record(&snapshot);
374 assert_eq!(history.net_tx.last(1), vec![5.0]);
375 }
376
377 #[test]
378 fn interface_helpers() {
379 assert!(is_virtual_interface("veth5a74356"));
380 assert!(!is_virtual_interface("enp9s0"));
381 let snapshot = Snapshot {
382 networks: vec![
383 NetworkStats {
384 name: "docker0".to_owned(),
385 rx_total_bytes: 1 << 40,
386 ..Default::default()
387 },
388 NetworkStats {
389 name: "enp9s0".to_owned(),
390 rx_total_bytes: 1 << 30,
391 ..Default::default()
392 },
393 ],
394 ..Default::default()
395 };
396 assert_eq!(default_interface(&snapshot).as_deref(), Some("enp9s0"));
397 assert_eq!(signal_name(ProcessSignal::Kill), "SIGKILL");
398 }
399}