rs_zero/resil/

cpu.rs

use std::{
    sync::{Arc, Mutex},
    time::{Duration, Instant},
};

const CPU_MAX_MILLIS: u32 = 1000;

/// Provides CPU usage in millicpu-style units where `1000` means 100%.
pub trait CpuUsageProvider: Send + Sync + 'static {
    /// Returns current process or system CPU usage.
    fn cpu_usage_millis(&self) -> u32;
}

/// Default CPU provider.
///
/// It intentionally reports `0` without an additional platform dependency.
/// Applications that need CPU-aware shedding can install a provider backed by
/// their runtime or metrics system.
#[derive(Debug, Clone, Copy, Default)]
pub struct NoopCpuUsageProvider;

impl CpuUsageProvider for NoopCpuUsageProvider {
    fn cpu_usage_millis(&self) -> u32 {
        0
    }
}

/// CPU sampler configuration.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct CpuSamplerConfig {
    /// Minimum time between refreshes.
    pub refresh_interval: Duration,
    /// Exponential moving average beta. `0.95` matches go-zero's smoothing profile.
    pub beta: f64,
}

impl Default for CpuSamplerConfig {
    fn default() -> Self {
        Self {
            refresh_interval: Duration::from_millis(250),
            beta: 0.95,
        }
    }
}

/// Linux CPU provider backed by `/proc/stat`.
#[derive(Debug)]
pub struct LinuxCpuUsageProvider {
    config: CpuSamplerConfig,
    state: Mutex<LinuxCpuState>,
}

#[derive(Debug, Clone, Copy)]
struct LinuxCpuState {
    previous: Option<CpuTimes>,
    last_refresh: Option<Instant>,
    usage_millis: u32,
}

#[derive(Debug, Clone, Copy)]
struct CpuTimes {
    idle: u64,
    total: u64,
}

impl LinuxCpuUsageProvider {
    /// Creates a Linux CPU provider using go-zero style sampling defaults.
    pub fn new() -> Self {
        Self::with_config(CpuSamplerConfig::default())
    }

    /// Creates a Linux CPU provider with explicit sampling configuration.
    pub fn with_config(config: CpuSamplerConfig) -> Self {
        Self {
            config,
            state: Mutex::new(LinuxCpuState {
                previous: None,
                last_refresh: None,
                usage_millis: 0,
            }),
        }
    }

    fn refresh(&self, now: Instant) -> u32 {
        let mut state = self.state.lock().expect("linux cpu state lock");
        if let Some(last_refresh) = state.last_refresh {
            if now.duration_since(last_refresh) < self.config.refresh_interval {
                return state.usage_millis;
            }
        }

        let Some(current) = read_cpu_times() else {
            state.last_refresh = Some(now);
            return state.usage_millis;
        };

        if let Some(previous) = state.previous {
            let raw = usage_between(previous, current);
            let beta = self.config.beta.clamp(0.0, 1.0);
            state.usage_millis = ((state.usage_millis as f64 * beta) + (raw as f64 * (1.0 - beta)))
                .round()
                .clamp(0.0, CPU_MAX_MILLIS as f64) as u32;
        }
        state.previous = Some(current);
        state.last_refresh = Some(now);
        state.usage_millis
    }
}

impl Default for LinuxCpuUsageProvider {
    fn default() -> Self {
        Self::new()
    }
}

impl CpuUsageProvider for LinuxCpuUsageProvider {
    fn cpu_usage_millis(&self) -> u32 {
        self.refresh(Instant::now())
    }
}

/// Shared CPU provider handle.
pub type SharedCpuUsageProvider = Arc<dyn CpuUsageProvider>;

/// Creates the default CPU provider.
pub fn default_cpu_provider() -> SharedCpuUsageProvider {
    real_cpu_provider().unwrap_or_else(|| Arc::new(NoopCpuUsageProvider))
}

/// Creates a platform CPU provider when rs-zero can sample CPU without extra setup.
pub fn real_cpu_provider() -> Option<SharedCpuUsageProvider> {
    #[cfg(target_os = "linux")]
    {
        return Some(Arc::new(LinuxCpuUsageProvider::new()));
    }

    #[cfg(not(target_os = "linux"))]
    {
        None
    }
}

fn read_cpu_times() -> Option<CpuTimes> {
    #[cfg(target_os = "linux")]
    {
        read_linux_proc_stat()
    }

    #[cfg(not(target_os = "linux"))]
    {
        None
    }
}

#[cfg(target_os = "linux")]
fn read_linux_proc_stat() -> Option<CpuTimes> {
    let stat = std::fs::read_to_string("/proc/stat").ok()?;
    parse_proc_stat(&stat)
}

#[cfg(target_os = "linux")]
fn parse_proc_stat(input: &str) -> Option<CpuTimes> {
    let line = input.lines().find(|line| line.starts_with("cpu "))?;
    parse_proc_stat_line(line)
}

#[cfg_attr(not(target_os = "linux"), allow(dead_code))]
fn parse_proc_stat_line(line: &str) -> Option<CpuTimes> {
    let mut parts = line.split_whitespace();
    if parts.next()? != "cpu" {
        return None;
    }

    let values = parts
        .take(10)
        .map(str::parse::<u64>)
        .collect::<Result<Vec<_>, _>>()
        .ok()?;
    if values.len() < 4 {
        return None;
    }

    let idle = values[3] + values.get(4).copied().unwrap_or(0);
    let total = values.iter().copied().sum();
    Some(CpuTimes { idle, total })
}

fn usage_between(previous: CpuTimes, current: CpuTimes) -> u32 {
    let total_delta = current.total.saturating_sub(previous.total);
    if total_delta == 0 {
        return 0;
    }

    let idle_delta = current.idle.saturating_sub(previous.idle);
    let busy_delta = total_delta.saturating_sub(idle_delta);
    ((busy_delta.saturating_mul(CPU_MAX_MILLIS as u64)) / total_delta).min(CPU_MAX_MILLIS as u64)
        as u32
}

#[cfg(test)]
mod tests {
    use super::{CpuTimes, parse_proc_stat_line, usage_between};

    #[test]
    fn parses_linux_proc_stat_line() {
        let times = parse_proc_stat_line("cpu  10 20 30 40 5 6 7 8 9 10").expect("times");

        assert_eq!(times.idle, 45);
        assert_eq!(times.total, 145);
    }

    #[test]
    fn calculates_busy_cpu_millis() {
        let previous = CpuTimes {
            idle: 100,
            total: 200,
        };
        let current = CpuTimes {
            idle: 125,
            total: 300,
        };

        assert_eq!(usage_between(previous, current), 750);
    }
}