dotfiles/agent/src/health.rs

//! System health information and checking

use std::sync::Arc;

use serde::Serialize;
use tokio::sync::watch;

const MEMORY_USAGE_CRITICAL_THRESHOLD: f64 = 90.0;
const CPU_USAGE_CRITICAL_THRESHOLD: f32 = 90.0;
const DISK_USAGE_CRITICAL_THRESHOLD: f32 = 90.0;

pub struct System {
    sys: sysinfo::System,
    disks: sysinfo::Disks,
}

impl System {
    pub fn new() -> Self {
        Self {
            sys: sysinfo::System::new(),
            disks: sysinfo::Disks::new(),
        }
    }

    pub fn refresh_resources(&mut self) {
        use sysinfo::{CpuRefreshKind, MemoryRefreshKind};

        self.sys.refresh_specifics(
            sysinfo::RefreshKind::new()
                .with_memory(MemoryRefreshKind::new().with_ram().with_swap())
                .with_cpu(CpuRefreshKind::new().with_cpu_usage()),
        );

        self.disks.refresh_list();
    }

    pub fn system(&self) -> &sysinfo::System {
        &self.sys
    }

    pub fn disks(&self) -> &sysinfo::Disks {
        &self.disks
    }
}

#[derive(Debug, Default, PartialEq, Serialize)]
pub enum Status {
    #[default]
    Normal,
    MemoryWarning,
    MemoryCritical,
    CpuWarning,
    CpuCritical,
    DiskWarning,
    DiskCritical,
}

#[derive(Debug, Serialize)]
pub struct Health {
    status: Vec<Status>,
}

#[derive(Clone)]
pub struct HealthMonitor(Arc<watch::Sender<Health>>);

impl HealthMonitor {
    pub fn new() -> Self {
        let (sender, _) = watch::channel(Health {
            status: [Status::default()].into(),
        });

        Self(Arc::new(sender))
    }

    pub fn check_system(&self, system: &System) {
        let s = system.system();

        let memory_usage = if s.total_memory() > 0 {
            s.used_memory() as f64 * 100.0 / s.total_memory() as f64
        } else {
            0.0
        };

        let cpu_usage = s.global_cpu_info().cpu_usage();

        // for d in system.disks().list() {
        //     let avail = if d.total_space() > 0 {
        //         (d.available_space() * 100 / d.total_space()) as u8
        //     } else {
        //         0 as u8
        //     };
        // }

        self.0.send_if_modified(|health| {
            let cpu_changed = if cpu_usage > CPU_USAGE_CRITICAL_THRESHOLD {
                health
                    .status
                    .iter()
                    .find(|&it| *it == Status::CpuCritical)
                    .and(None)
                    .unwrap_or_else(|| {
                        health.status.push(Status::CpuCritical);
                        true
                    })
            } else {
                false
            };

            let mem_changed = if memory_usage > MEMORY_USAGE_CRITICAL_THRESHOLD {
                health
                    .status
                    .iter()
                    .find(|&it| *it == Status::MemoryCritical)
                    .and(None)
                    .unwrap_or_else(|| {
                        health.status.push(Status::MemoryCritical);
                        true
                    })
            } else {
                false
            };

            dbg!(cpu_changed || mem_changed);
            cpu_changed || mem_changed
        });
    }

    pub fn monitor(&self) -> watch::Receiver<Health> {
        self.0.subscribe()
    }
}