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

use std::fs;
use std::io::{self, Write};
use std::thread;
use std::time::{Duration, Instant};

struct Sample {
    idle: u64,
    total: u64,
}

struct MemInfo {
    percent: u64,
    used_gb: f64,
    total_gb: f64,
    avail_gb: f64,
    cached_gb: f64,
    buffers_gb: f64,
}

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()
}

fn parse_meminfo(input: &str) -> Option<MemInfo> {
    let mut total = 0u64;
    let mut avail = 0u64;
    let mut buffers = 0u64;
    let mut cached = 0u64;
    let mut sreclaimable = 0u64;

    for line in input.lines() {
        let mut parts = line.splitn(2, ':');
        let key = parts.next().unwrap_or("").trim();
        let val: u64 = parts
            .next()
            .unwrap_or("")
            .split_whitespace()
            .next()
            .unwrap_or("")
            .parse()
            .unwrap_or(0);
        match key {
            "MemTotal" => total = val,
            "MemAvailable" => avail = val,
            "Buffers" => buffers = val,
            "Cached" => cached = val,
            "SReclaimable" => sreclaimable = val,
            _ => {}
        }
    }

    if total == 0 {
        return None;
    }

    let used = total.saturating_sub(avail);
    let cached_total = cached + sreclaimable;
    let gb = |kb: u64| kb as f64 / 1_048_576.0;

    Some(MemInfo {
        percent: used * 100 / total,
        used_gb: gb(used),
        total_gb: gb(total),
        avail_gb: gb(avail),
        cached_gb: gb(cached_total),
        buffers_gb: gb(buffers),
    })
}

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
}

fn read_stat() -> Vec<Sample> {
    parse_stat(&fs::read_to_string("/proc/stat").unwrap_or_default())
}

fn read_temp_celsius() -> Option<i32> {
    let mut max: Option<i32> = None;
    for i in 0.. {
        let path = format!("/sys/class/thermal/thermal_zone{i}/temp");
        match fs::read_to_string(&path) {
            Ok(s) => {
                if let Ok(millic) = s.trim().parse::<i32>() {
                    let c = millic / 1000;
                    max = Some(max.map_or(c, |m: i32| m.max(c)));
                }
            }
            Err(_) => break,
        }
    }
    max
}

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
}

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(|(p, c)| pct(p, c))
        .unwrap_or(0);

    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, "]}}");
}

fn emit_temp(out: &mut impl Write) {
    if let Some(c) = read_temp_celsius() {
        let _ = writeln!(out, "{{\"type\":\"temp\",\"celsius\":{c}}}");
    }
}

fn emit_mem(out: &mut impl Write) {
    let content = fs::read_to_string("/proc/meminfo").unwrap_or_default();
    if let Some(m) = parse_meminfo(&content) {
        let _ = writeln!(
            out,
            "{{\"type\":\"mem\",\"percent\":{},\"used_gb\":{:.3},\"total_gb\":{:.3},\"avail_gb\":{:.3},\"cached_gb\":{:.3},\"buffers_gb\":{:.3}}}",
            m.percent, m.used_gb, m.total_gb, m.avail_gb, m.cached_gb, m.buffers_gb,
        );
    }
}

fn parse_interval_ms() -> u64 {
    let args: Vec<String> = std::env::args().collect();
    let mut i = 1;
    while i < args.len() {
        if args[i] == "--interval" {
            if let Some(ms) = args.get(i + 1).and_then(|s| s.parse().ok()) {
                return ms;
            }
        }
        i += 1;
    }
    1000
}

fn main() {
    let interval = Duration::from_millis(parse_interval_ms());
    let stdout = io::stdout();
    let mut out = io::BufWriter::new(stdout.lock());
    let mut prev: Vec<Sample> = vec![];
    let mut freqs: Vec<f64> = vec![];
    let mut tick = 0u64;

    loop {
        let t0 = Instant::now();

        let curr = read_stat();
        if tick.is_multiple_of(2) {
            freqs = read_core_freqs();
            emit_mem(&mut out);
        }
        emit_cpu(&mut out, &prev, &curr, &freqs);
        prev = curr;

        if tick.is_multiple_of(4) {
            emit_temp(&mut out);
        }

        let _ = out.flush();
        tick += 1;

        let elapsed = t0.elapsed();
        if elapsed < interval {
            thread::sleep(interval - elapsed);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── 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,
        }; // idle went backwards
        // dt=saturating 0, di=saturating 0 → returns 0
        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() {
        // aggregate line + 2 cores = 3 samples
        let samples = parse_stat(STAT_SAMPLE);
        assert_eq!(samples.len(), 3);
    }

    #[test]
    fn parse_stat_aggregate_idle() {
        // idle=field[3], iowait=field[4] → 700+40=740
        let samples = parse_stat(STAT_SAMPLE);
        assert_eq!(samples[0].idle, 740);
    }

    #[test]
    fn parse_stat_aggregate_total() {
        // sum of all fields: 100+10+50+700+40 = 900
        let samples = parse_stat(STAT_SAMPLE);
        assert_eq!(samples[0].total, 900);
    }

    #[test]
    fn parse_stat_per_core_idle() {
        let samples = parse_stat(STAT_SAMPLE);
        assert_eq!(samples[1].idle, 370); // 350+20
        assert_eq!(samples[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";
        let samples = parse_stat(input);
        assert_eq!(samples.len(), 1);
    }

    // ── parse_meminfo ────────────────────────────────────────────────────

    const MEMINFO_SAMPLE: &str = "\
MemTotal:       16384000 kB
MemFree:         2048000 kB
MemAvailable:    4096000 kB
Buffers:          512000 kB
Cached:          3072000 kB
SReclaimable:     512000 kB
SwapTotal:       8192000 kB
SwapFree:        8192000 kB";

    #[test]
    fn parse_meminfo_percent() {
        let m = parse_meminfo(MEMINFO_SAMPLE).unwrap();
        // used = total - avail = 16384000 - 4096000 = 12288000
        // percent = 12288000 * 100 / 16384000 = 75
        assert_eq!(m.percent, 75);
    }

    #[test]
    fn parse_meminfo_total_gb() {
        let m = parse_meminfo(MEMINFO_SAMPLE).unwrap();
        // 16384000 kB / 1048576 ≈ 15.625 GB
        assert!((m.total_gb - 15.625).abs() < 0.001);
    }

    #[test]
    fn parse_meminfo_cached_includes_sreclaimable() {
        let m = parse_meminfo(MEMINFO_SAMPLE).unwrap();
        // cached = 3072000 + 512000 = 3584000 kB
        let expected = 3_584_000.0 / 1_048_576.0;
        assert!((m.cached_gb - expected).abs() < 0.001);
    }

    #[test]
    fn parse_meminfo_zero_total_returns_none() {
        assert!(parse_meminfo("MemFree: 1000 kB\n").is_none());
    }

    #[test]
    fn parse_meminfo_empty_returns_none() {
        assert!(parse_meminfo("").is_none());
    }

    // ── emit_cpu ─────────────────────────────────────────────────────────

    fn sample(idle: u64, total: u64) -> Sample {
        Sample { idle, total }
    }

    #[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() {
        // prev aggregate: idle=0, total=0 → curr: idle=50, total=100 → 50% busy
        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}");
    }

    // ── emit_mem (via parse_meminfo) ─────────────────────────────────────

    #[test]
    fn emit_mem_valid_json_structure() {
        let m = parse_meminfo(MEMINFO_SAMPLE).unwrap();
        let mut buf = Vec::new();
        let _ = writeln!(
            &mut buf,
            "{{\"type\":\"mem\",\"percent\":{},\"used_gb\":{:.3},\"total_gb\":{:.3},\"avail_gb\":{:.3},\"cached_gb\":{:.3},\"buffers_gb\":{:.3}}}",
            m.percent, m.used_gb, m.total_gb, m.avail_gb, m.cached_gb, m.buffers_gb,
        );
        let s = String::from_utf8(buf).unwrap();
        assert!(s.contains("\"type\":\"mem\""));
        assert!(s.contains("\"percent\":75"));
        assert!(s.trim().ends_with('}'));
    }
}