Skip to main content

kaizen/shell/
sampler_cmd.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! Detached process sampler (`__sampler-run`).
3
4use crate::core::config;
5use crate::store::Store;
6use anyhow::Result;
7use std::path::{Path, PathBuf};
8use std::time::{Duration, SystemTime, UNIX_EPOCH};
9use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System};
10
11fn stop_path(workspace: &Path, session_id: &str) -> PathBuf {
12    workspace.join(".kaizen/sampler-stop").join(session_id)
13}
14
15/// Sample `pid` until stop file, cap, or process exit.
16pub fn cmd_sampler_run(workspace: &Path, session_id: &str, pid: u32) -> Result<()> {
17    let cfg = config::load(workspace)?;
18    let s = &cfg.collect.system_sampler;
19    if !s.enabled {
20        return Ok(());
21    }
22    let db = workspace.join(".kaizen/kaizen.db");
23    let store = Store::open(&db)?;
24    let target = Pid::from_u32(pid);
25    let mut sys = System::new();
26    let kind = ProcessRefreshKind::nothing().with_cpu().with_memory();
27    let mut n: u32 = 0;
28    while n < s.max_samples_per_session {
29        if stop_path(workspace, session_id).exists() {
30            break;
31        }
32        std::thread::sleep(Duration::from_millis(s.sample_ms.max(100)));
33        sys.refresh_processes_specifics(ProcessesToUpdate::Some(&[target]), false, kind);
34        let ts = now_ms();
35        let row = sys.process(target);
36        let (cpu, rss) = match row {
37            Some(p) => (Some(p.cpu_usage() as f64), Some(p.memory())),
38            None => break,
39        };
40        store.append_session_sample(session_id, ts, pid, cpu, rss)?;
41        n += 1;
42    }
43    Ok(())
44}
45
46fn now_ms() -> u64 {
47    SystemTime::now()
48        .duration_since(UNIX_EPOCH)
49        .map(|d| d.as_millis() as u64)
50        .unwrap_or(0)
51}