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