status-monitor/agent/src/cpu.rs

use anyhow::anyhow;
use common::metrics::CpuBreakdown;
use std::collections::HashMap;

#[derive(Default, Debug, Clone, Copy)]
struct CpuValues {
    pub user: u64,
    pub nice: u64,
    pub system: u64,
    pub idle: u64,
    pub iowait: u64,
    pub irq: u64,
    pub softirq: u64,
    pub steal: u64,
    pub guest: u64,
    pub guest_nice: u64,
}

impl CpuValues {
    pub fn total(&self) -> u64 {
        self.user
            + self.nice
            + self.system
            + self.idle
            + self.iowait
            + self.irq
            + self.softirq
            + self.steal
            + self.guest
            + self.guest_nice
    }
}

#[cfg_attr(not(target_os = "linux"), derive(Default))]
pub struct CpuStatReader {
    previous_stats: HashMap<String, CpuValues>,
}

const CPU_LINE_PATH: &str = "/proc/stat";
const CPU_TEMP_PATH: &str = "/sys/class/thermal/thermal_zone0/temp";

impl CpuStatReader {
    pub fn new() -> Self {
        Self {
            previous_stats: HashMap::new(),
        }
    }

    #[cfg(target_os = "linux")]
    pub fn cpu_breakdown(&mut self) -> anyhow::Result<CpuBreakdown> {
        let content = std::fs::read_to_string(CPU_LINE_PATH)?;

        let cpu_line = content
            .lines()
            .find(|line| line.starts_with("cpu "))
            .unwrap();

        self.parse_cpu_line(cpu_line)
    }

    #[cfg(target_os = "linux")]
    pub fn cpu_temperature(&mut self) -> anyhow::Result<f32> {
        let content = std::fs::read_to_string(CPU_TEMP_PATH)?;
        let temp_millicelsius = content.trim().parse::<f32>()?;

        let temp_celsius = temp_millicelsius / 1000.0;

        Ok(temp_celsius)
    }

    fn parse_cpu_line(&mut self, line: &str) -> anyhow::Result<CpuBreakdown> {
        let mut parts = line.split_whitespace();

        let cpu_name = parts
            .next()
            .ok_or_else(|| anyhow!("missing cpu name"))?
            .to_string();

        let mut values = CpuValues::default();
        let mut fields = [
            &mut values.user,
            &mut values.nice,
            &mut values.system,
            &mut values.idle,
            &mut values.iowait,
            &mut values.irq,
            &mut values.softirq,
            &mut values.steal,
            &mut values.guest,
            &mut values.guest_nice,
        ];

        for (field, part) in fields.iter_mut().zip(&mut parts) {
            **field = part.parse()?;
        }

        if fields.iter().any(|f| **f == 0 && parts.next().is_some()) {
            return Err(anyhow!("invalid cpu line"));
        }

        let previous = self.previous_stats.get(&cpu_name).copied();
        let percentages = self.calculate_percentages(&values, previous);

        self.previous_stats.insert(cpu_name, values);

        Ok(percentages)
    }

    fn calculate_percentages(
        &self,
        current: &CpuValues,
        previous: Option<CpuValues>,
    ) -> CpuBreakdown {
        let Some(prev) = previous else {
            return CpuBreakdown::default();
        };

        let total_delta = current.total().saturating_sub(prev.total()) as f32;
        if total_delta <= 0.0 {
            return CpuBreakdown::default();
        }

        let idle = current.idle.saturating_sub(prev.idle) as f32 / total_delta;
        let total = 1.0 - idle as f32;

        let calculate_pct =
            |curr: u64, prev: u64| total * (curr.saturating_sub(prev) as f32 / total_delta) * 100.0;

        CpuBreakdown {
            system: calculate_pct(current.system, prev.system),
            user: calculate_pct(current.user, prev.user),
            idle: calculate_pct(current.idle, prev.idle),
            steal: calculate_pct(current.steal, prev.steal),
            iowait: calculate_pct(current.iowait, prev.iowait),
        }
    }
}