nova-shell/stats-daemon/src/cpu.rs

246 lines
6.7 KiB
Rust

use std::fs;
use std::io::Write;
pub struct Sample {
pub idle: u64,
pub total: u64,
}
pub fn parse_stat(input: &str) -> Vec<Sample> {
input
.lines()
.filter(|l| l.starts_with("cpu"))
.map(|l| {
let vals: Vec<u64> = l
.split_whitespace()
.skip(1)
.filter_map(|s| s.parse().ok())
.collect();
let idle = vals.get(3).copied().unwrap_or(0) + vals.get(4).copied().unwrap_or(0);
let total = vals.iter().sum();
Sample { idle, total }
})
.collect()
}
pub fn pct(prev: &Sample, curr: &Sample) -> u32 {
let dt = curr.total.saturating_sub(prev.total);
let di = curr.idle.saturating_sub(prev.idle);
if dt == 0 {
return 0;
}
(dt.saturating_sub(di) * 100 / dt) as u32
}
pub fn read_stat() -> Vec<Sample> {
parse_stat(&fs::read_to_string("/proc/stat").unwrap_or_default())
}
pub fn read_core_freqs() -> Vec<f64> {
let mut freqs = Vec::new();
for i in 0.. {
let path = format!("/sys/devices/system/cpu/cpu{i}/cpufreq/scaling_cur_freq");
match fs::read_to_string(&path) {
Ok(s) => match s.trim().parse::<u64>() {
Ok(khz) => freqs.push(khz as f64 / 1_000_000.0),
Err(_) => break,
},
Err(_) => break,
}
}
freqs
}
pub fn emit_cpu(out: &mut impl Write, prev: &[Sample], curr: &[Sample], freqs: &[f64]) {
if curr.is_empty() {
return;
}
let usage = prev.first().zip(curr.first()).map_or(0, |(p, c)| pct(p, c));
let core_usage: Vec<u32> = prev
.iter()
.skip(1)
.zip(curr.iter().skip(1))
.map(|(p, c)| pct(p, c))
.collect();
let avg_freq = if freqs.is_empty() {
0.0
} else {
freqs.iter().sum::<f64>() / freqs.len() as f64
};
let n = core_usage.len().max(freqs.len());
let _ = write!(
out,
"{{\"type\":\"cpu\",\"usage\":{usage},\"freq_ghz\":{avg_freq:.3},\"cores\":["
);
for i in 0..n {
if i > 0 {
let _ = write!(out, ",");
}
let u = core_usage.get(i).copied().unwrap_or(0);
let f = freqs.get(i).copied().unwrap_or(0.0);
let _ = write!(out, "{{\"usage\":{u},\"freq_ghz\":{f:.3}}}");
}
let _ = writeln!(out, "]}}");
}
#[cfg(test)]
mod tests {
use super::*;
fn sample(idle: u64, total: u64) -> Sample {
Sample { idle, total }
}
// ── pct ──────────────────────────────────────────────────────────────
#[test]
fn pct_zero_delta_returns_zero() {
let s = Sample {
idle: 100,
total: 400,
};
assert_eq!(pct(&s, &s), 0);
}
#[test]
fn pct_all_idle() {
let prev = Sample {
idle: 0,
total: 100,
};
let curr = Sample {
idle: 100,
total: 200,
};
assert_eq!(pct(&prev, &curr), 0);
}
#[test]
fn pct_fully_busy() {
let prev = Sample {
idle: 100,
total: 200,
};
let curr = Sample {
idle: 100,
total: 300,
};
assert_eq!(pct(&prev, &curr), 100);
}
#[test]
fn pct_half_busy() {
let prev = Sample { idle: 0, total: 0 };
let curr = Sample {
idle: 50,
total: 100,
};
assert_eq!(pct(&prev, &curr), 50);
}
#[test]
fn pct_no_underflow_on_backwards_clock() {
let prev = Sample {
idle: 200,
total: 400,
};
let curr = Sample {
idle: 100,
total: 300,
};
assert_eq!(pct(&prev, &curr), 0);
}
// ── parse_stat ───────────────────────────────────────────────────────
const STAT_SAMPLE: &str = "\
cpu 100 10 50 700 40 0 0 0 0 0
cpu0 50 5 25 350 20 0 0 0 0 0
cpu1 50 5 25 350 20 0 0 0 0 0";
#[test]
fn parse_stat_count() {
assert_eq!(parse_stat(STAT_SAMPLE).len(), 3);
}
#[test]
fn parse_stat_aggregate_idle() {
assert_eq!(parse_stat(STAT_SAMPLE)[0].idle, 740);
}
#[test]
fn parse_stat_aggregate_total() {
assert_eq!(parse_stat(STAT_SAMPLE)[0].total, 900);
}
#[test]
fn parse_stat_per_core_idle() {
let s = parse_stat(STAT_SAMPLE);
assert_eq!(s[1].idle, 370);
assert_eq!(s[2].idle, 370);
}
#[test]
fn parse_stat_ignores_non_cpu_lines() {
let input = "intr 12345\ncpu 1 2 3 4 5 0 0 0 0 0\npage 0 0";
assert_eq!(parse_stat(input).len(), 1);
}
// ── emit_cpu ─────────────────────────────────────────────────────────
#[test]
fn emit_cpu_valid_json_structure() {
let prev = vec![sample(0, 0), sample(0, 0)];
let curr = vec![sample(50, 100), sample(25, 100)];
let freqs = vec![3.2, 3.1];
let mut buf = Vec::new();
emit_cpu(&mut buf, &prev, &curr, &freqs);
let s = String::from_utf8(buf).unwrap();
assert!(s.contains("\"type\":\"cpu\""));
assert!(s.contains("\"usage\":"));
assert!(s.contains("\"freq_ghz\":"));
assert!(s.contains("\"cores\":"));
assert!(s.trim().ends_with('}'));
}
#[test]
fn emit_cpu_correct_usage() {
let prev = vec![sample(0, 0), sample(0, 0)];
let curr = vec![sample(50, 100), sample(0, 0)];
let mut buf = Vec::new();
emit_cpu(&mut buf, &prev, &curr, &[]);
let s = String::from_utf8(buf).unwrap();
assert!(s.contains("\"usage\":50"), "got: {s}");
}
#[test]
fn emit_cpu_no_prev_gives_zero_usage() {
let curr = vec![sample(50, 100)];
let mut buf = Vec::new();
emit_cpu(&mut buf, &[], &curr, &[]);
let s = String::from_utf8(buf).unwrap();
assert!(s.contains("\"usage\":0"), "got: {s}");
}
#[test]
fn emit_cpu_empty_curr_produces_no_output() {
let mut buf = Vec::new();
emit_cpu(&mut buf, &[], &[], &[]);
assert!(buf.is_empty());
}
#[test]
fn emit_cpu_core_freqs_in_output() {
let curr = vec![sample(0, 100)];
let freqs = vec![3.200, 2.900];
let mut buf = Vec::new();
emit_cpu(&mut buf, &curr, &curr, &freqs);
let s = String::from_utf8(buf).unwrap();
assert!(s.contains("\"freq_ghz\":3.200"), "got: {s}");
assert!(s.contains("\"freq_ghz\":2.900"), "got: {s}");
}
}