tawasuyu/brahman

1//! Resource accounting por workspace.
2//!
3//! Dos fuentes:
4//! - **Per-proc** (`/proc/<pid>/status` + `stat`): suma RSS y CPU ticks de
5//!   los comandos vivos del workspace. Siempre disponible. Costo: O(N pids).
6//! - **Cgroup v2** (`memory.current`, `cpu.stat`): un read por workspace si
7//!   `SomaSpec.cgroup.path` está y es leíble. Más preciso (incluye descendants).
8//!
9//! Si ambos están disponibles, devolvemos el cgroup (más preciso) y dejamos
10//! el per-proc como `sample_via_proc`.
11
12use std::path::Path;
13use std::time::Instant;
14
15#[derive(Debug, Clone, Default)]
16pub struct WorkspaceStats {
17    pub commands_alive: u32,
18    pub commands_total: u32,
19    /// RSS sumado en bytes. `None` si no se pudo medir.
20    pub rss_bytes: Option<u64>,
21    /// High-water mark de RSS (peak alguna vez observado). Cgroup v2:
22    /// `memory.peak` (≥6.5). Per-proc: suma de `VmHWM` de cada pid.
23    pub rss_peak_bytes: Option<u64>,
24    /// Tiempo CPU acumulado en microsegundos. `None` si no se pudo medir.
25    pub cpu_usec: Option<u64>,
26    /// Fuente del dato: "proc" | "cgroup" | "mixed".
27    pub source: String,
28    /// Wall-clock uptime del workspace en milisegundos.
29    pub uptime_ms: u64,
30}
31
32/// Mide stats para un set de PIDs vivos + un path de cgroup opcional.
33pub fn measure(
34    alive_pids: &[i32],
35    cgroup_path: Option<&Path>,
36    workspace_started: Instant,
37) -> WorkspaceStats {
38    let mut rss_proc: u64 = 0;
39    let mut rss_peak_proc: u64 = 0;
40    let mut cpu_proc: u64 = 0;
41    let mut proc_ok = false;
42    for &pid in alive_pids {
43        if let Some((rss, peak, cpu)) = read_proc_pid(pid) {
44            rss_proc += rss;
45            rss_peak_proc += peak;
46            cpu_proc += cpu;
47            proc_ok = true;
48        }
49    }
50
51    let cgroup = cgroup_path.and_then(read_cgroup_stats);
52
53    let (rss, rss_peak, cpu, source) = match (cgroup, proc_ok) {
54        (Some(cg), _) => (Some(cg.rss), cg.rss_peak, Some(cg.cpu_usec), "cgroup".to_string()),
55        (None, true) => (
56            Some(rss_proc),
57            Some(rss_peak_proc),
58            Some(cpu_proc),
59            "proc".to_string(),
60        ),
61        (None, false) => (None, None, None, "none".to_string()),
62    };
63
64    WorkspaceStats {
65        commands_alive: alive_pids.len() as u32,
66        commands_total: 0,
67        rss_bytes: rss,
68        rss_peak_bytes: rss_peak,
69        cpu_usec: cpu,
70        source,
71        uptime_ms: workspace_started.elapsed().as_millis() as u64,
72    }
73}
74
75struct CgroupStats {
76    rss: u64,
77    rss_peak: Option<u64>,
78    cpu_usec: u64,
79}
80
81/// Lee `(rss_bytes, rss_peak_bytes, cpu_usec)` de `/proc/<pid>/`. None si el proc desapareció.
82fn read_proc_pid(pid: i32) -> Option<(u64, u64, u64)> {
83    let (rss_kb, hwm_kb) = {
84        let status = std::fs::read_to_string(format!("/proc/{pid}/status")).ok()?;
85        let mut rss = 0u64;
86        let mut hwm = 0u64;
87        for l in status.lines() {
88            if let Some(rest) = l.strip_prefix("VmRSS:") {
89                rss = rest
90                    .trim()
91                    .split_whitespace()
92                    .next()
93                    .and_then(|s| s.parse().ok())
94                    .unwrap_or(0);
95            } else if let Some(rest) = l.strip_prefix("VmHWM:") {
96                hwm = rest
97                    .trim()
98                    .split_whitespace()
99                    .next()
100                    .and_then(|s| s.parse().ok())
101                    .unwrap_or(0);
102            }
103        }
104        (rss, hwm)
105    };
106    let cpu_usec = {
107        let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;
108        // formato: pid (comm) state ppid pgrp ... utime stime cutime cstime
109        // Cuidado: comm puede tener espacios y paréntesis. Buscamos la última `)`.
110        let end_comm = stat.rfind(')')?;
111        let after = &stat[end_comm + 1..];
112        let fields: Vec<&str> = after.split_whitespace().collect();
113        // Tras `)`, índice 0 = state, índice 11 = utime, 12 = stime.
114        let utime = fields.get(11).and_then(|s| s.parse::<u64>().ok()).unwrap_or(0);
115        let stime = fields.get(12).and_then(|s| s.parse::<u64>().ok()).unwrap_or(0);
116        let ticks = utime + stime;
117        // Convertimos ticks → microsegundos. SC_CLK_TCK típicamente 100.
118        let clk_tck = unsafe { libc::sysconf(libc::_SC_CLK_TCK) }.max(1) as u64;
119        ticks * 1_000_000 / clk_tck
120    };
121    Some((rss_kb * 1024, hwm_kb * 1024, cpu_usec))
122}
123
124/// Lee `CgroupStats` del cgroup. None si no existe o no es leíble.
125/// `memory.peak` requiere kernel ≥6.5; si falta, `rss_peak` queda None.
126fn read_cgroup_stats(cgroup_path: &Path) -> Option<CgroupStats> {
127    let mem = std::fs::read_to_string(cgroup_path.join("memory.current"))
128        .ok()
129        .and_then(|s| s.trim().parse::<u64>().ok())?;
130    let cpu_stat = std::fs::read_to_string(cgroup_path.join("cpu.stat")).ok()?;
131    let cpu_usec = cpu_stat
132        .lines()
133        .find_map(|l| l.strip_prefix("usage_usec"))
134        .and_then(|s| s.split_whitespace().next())
135        .and_then(|s| s.parse::<u64>().ok())
136        .unwrap_or(0);
137    let peak = std::fs::read_to_string(cgroup_path.join("memory.peak"))
138        .ok()
139        .and_then(|s| s.trim().parse::<u64>().ok());
140    Some(CgroupStats {
141        rss: mem,
142        rss_peak: peak,
143        cpu_usec,
144    })
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn measure_with_no_pids_returns_zero() {
153        let stats = measure(&[], None, Instant::now());
154        assert_eq!(stats.commands_alive, 0);
155        assert_eq!(stats.rss_bytes, None);
156        assert_eq!(stats.source, "none");
157    }
158
159    #[test]
160    fn measure_self_pid_returns_data() {
161        let me = std::process::id() as i32;
162        let stats = measure(&[me], None, Instant::now());
163        assert_eq!(stats.commands_alive, 1);
164        // Nuestro propio RSS debería ser > 0.
165        assert!(stats.rss_bytes.unwrap_or(0) > 0);
166        assert_eq!(stats.source, "proc");
167    }
168}