AkurAI-Monitor / src/collectors.rs

Latest change 8bade644cef8ddaa358cc88eb6017c8e64d5a6e4 - AKURAI-MONITOR-7 monitor processes, docker, systemd; fix #errors alert delivery by Ólafur Búi Ólafsson

use crate::config;
use std::collections::HashMap;

#[derive(Debug, Clone, serde::Serialize)]
pub struct Metric {
    pub name: String,
    pub value: f64,
}

pub fn collect_all() -> Vec<Metric> {
    let mut metrics = Vec::new();
    metrics.extend(collect_cpu());
    metrics.extend(collect_memory());
    metrics.extend(collect_disk());
    metrics.extend(collect_load());
    metrics.extend(collect_uptime());
    metrics.extend(collect_network());
    metrics
}

fn collect_cpu() -> Vec<Metric> {
    let read_cpu = || -> Option<(u64, u64)> {
        let stat = std::fs::read_to_string("/proc/stat").ok()?;
        let line = stat.lines().next()?;
        let parts: Vec<u64> = line
            .split_whitespace()
            .skip(1)
            .filter_map(|s| s.parse().ok())
            .collect();
        if parts.len() < 4 {
            return None;
        }
        let idle = parts[3];
        let total: u64 = parts.iter().sum();
        Some((idle, total))
    };

    let Some((idle1, total1)) = read_cpu() else {
        return vec![];
    };
    std::thread::sleep(std::time::Duration::from_secs(1));
    let Some((idle2, total2)) = read_cpu() else {
        return vec![];
    };

    let idle_delta = idle2.saturating_sub(idle1) as f64;
    let total_delta = total2.saturating_sub(total1) as f64;
    if total_delta == 0.0 {
        return vec![];
    }

    let usage = (1.0 - idle_delta / total_delta) * 100.0;
    vec![Metric {
        name: "cpu.usage".to_string(),
        value: (usage * 10.0).round() / 10.0,
    }]
}

fn collect_memory() -> Vec<Metric> {
    let Ok(meminfo) = std::fs::read_to_string("/proc/meminfo") else {
        return vec![];
    };

    let mut values: HashMap<&str, u64> = HashMap::new();
    for line in meminfo.lines() {
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() >= 2 {
            let key = parts[0].trim_end_matches(':');
            if let Ok(val) = parts[1].parse::<u64>() {
                values.insert(key, val);
            }
        }
    }

    let total = *values.get("MemTotal").unwrap_or(&0);
    let available = *values.get("MemAvailable").unwrap_or(&0);
    if total == 0 {
        return vec![];
    }

    let used = total.saturating_sub(available);
    let used_pct = (used as f64 / total as f64) * 100.0;

    vec![
        Metric {
            name: "mem.total_kb".to_string(),
            value: total as f64,
        },
        Metric {
            name: "mem.used_kb".to_string(),
            value: used as f64,
        },
        Metric {
            name: "mem.available_kb".to_string(),
            value: available as f64,
        },
        Metric {
            name: "mem.used_pct".to_string(),
            value: (used_pct * 10.0).round() / 10.0,
        },
    ]
}

fn collect_disk() -> Vec<Metric> {
    let Ok(mounts) = std::fs::read_to_string("/proc/mounts") else {
        return vec![];
    };

    let mut metrics = Vec::new();
    for line in mounts.lines() {
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() < 3 {
            continue;
        }
        let mount = parts[1];
        // The remote Titan wrapper and this in-process collector share a
        // root-disk-only contract; other mounts are intentionally excluded.
        if mount != "/" {
            continue;
        }
        let fstype = parts[2];

        // Only real filesystems
        if !matches!(fstype, "ext4" | "ext3" | "xfs" | "btrfs" | "zfs") {
            continue;
        }

        unsafe {
            let c_path = std::ffi::CString::new(mount).unwrap();
            let mut stat: libc::statvfs = std::mem::zeroed();
            if libc::statvfs(c_path.as_ptr(), &mut stat) != 0 {
                continue;
            }

            let total = stat.f_blocks * stat.f_frsize as u64;
            let free = stat.f_bfree * stat.f_frsize as u64;
            if total == 0 {
                continue;
            }
            let used = total - free;
            let used_pct = (used as f64 / total as f64) * 100.0;

            metrics.push(Metric {
                name: "disk.root.used_pct".to_string(),
                value: (used_pct * 10.0).round() / 10.0,
            });
            metrics.push(Metric {
                name: "disk.root.total_gb".to_string(),
                value: (total as f64 / 1_073_741_824.0 * 10.0).round() / 10.0,
            });
            metrics.push(Metric {
                name: "disk.root.used_gb".to_string(),
                value: (used as f64 / 1_073_741_824.0 * 10.0).round() / 10.0,
            });
        }
    }

    metrics
}

fn collect_load() -> Vec<Metric> {
    let Ok(loadavg) = std::fs::read_to_string("/proc/loadavg") else {
        return vec![];
    };

    let parts: Vec<&str> = loadavg.split_whitespace().collect();
    if parts.len() < 3 {
        return vec![];
    }

    let mut metrics = Vec::new();
    if let Ok(v) = parts[0].parse::<f64>() {
        metrics.push(Metric {
            name: "load.1m".to_string(),
            value: v,
        });
    }
    if let Ok(v) = parts[1].parse::<f64>() {
        metrics.push(Metric {
            name: "load.5m".to_string(),
            value: v,
        });
    }
    if let Ok(v) = parts[2].parse::<f64>() {
        metrics.push(Metric {
            name: "load.15m".to_string(),
            value: v,
        });
    }

    metrics
}

fn collect_uptime() -> Vec<Metric> {
    let Ok(uptime) = std::fs::read_to_string("/proc/uptime") else {
        return vec![];
    };

    let parts: Vec<&str> = uptime.split_whitespace().collect();
    if let Some(Ok(secs)) = parts.first().map(|s| s.parse::<f64>()) {
        vec![Metric {
            name: "uptime.seconds".to_string(),
            value: secs.round(),
        }]
    } else {
        vec![]
    }
}

fn collect_network() -> Vec<Metric> {
    let read_net = || -> Option<HashMap<String, (u64, u64)>> {
        let dev = std::fs::read_to_string("/proc/net/dev").ok()?;
        let mut map = HashMap::new();
        for line in dev.lines().skip(2) {
            let line = line.trim();
            let (iface, rest) = line.split_once(':')?;
            let iface = iface.trim();
            if iface == "lo" {
                continue;
            }
            let vals: Vec<u64> = rest
                .split_whitespace()
                .filter_map(|s| s.parse().ok())
                .collect();
            if vals.len() >= 9 {
                map.insert(iface.to_string(), (vals[0], vals[8])); // rx_bytes, tx_bytes
            }
        }
        Some(map)
    };

    let Some(snap1) = read_net() else {
        return vec![];
    };
    // Sample over a 1s window so rx/tx deltas reflect actual throughput.
    std::thread::sleep(std::time::Duration::from_secs(1));
    let Some(snap2) = read_net() else {
        return vec![];
    };

    let mut metrics = Vec::new();
    for (iface, (rx2, tx2)) in &snap2 {
        if let Some((rx1, tx1)) = snap1.get(iface) {
            metrics.push(Metric {
                name: format!("net.{iface}.rx_bytes"),
                value: rx2.saturating_sub(*rx1) as f64,
            });
            metrics.push(Metric {
                name: format!("net.{iface}.tx_bytes"),
                value: tx2.saturating_sub(*tx1) as f64,
            });
        }
    }

    metrics
}

/// Per-process, per-service-unit, and per-container entity metrics. Sampled
/// only on the persist cadence (see `main.rs`), never the live tick — these
/// are shell/proc scans an order of magnitude heavier than the host
/// counters, and archiving them every 1-5s would multiply metrics-table
/// growth for no benefit at sub-minute resolution.
pub fn collect_entities() -> Vec<Metric> {
    let mut metrics = Vec::new();
    metrics.extend(collect_processes());
    metrics.extend(collect_systemd());
    metrics.extend(collect_docker());
    metrics
}

/// Sanitize a process/unit/container name into a metric-name-safe segment:
/// alphanumeric, `-`, and `_` only, everything else collapsed to `_`. Metric
/// names are dot-delimited (`proc.<name>.mem_kb`), so the name itself must
/// never contain a `.`.
fn sanitize_segment(raw: &str) -> String {
    let cleaned: String = raw
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
                c
            } else {
                '_'
            }
        })
        .collect();
    let trimmed = cleaned.trim_matches('_');
    if trimmed.is_empty() {
        "unknown".to_string()
    } else {
        trimmed.to_string()
    }
}

/// Aggregate resident memory by process command name across every PID in
/// /proc, then emit the top-N (by RSS) as `proc.<comm>.mem_kb` /
/// `proc.<comm>.count`. Aggregating by comm (not PID) keeps the metric
/// namespace stable across restarts/respawns of the same program; kernel
/// threads report zero resident pages and naturally sort out of the top-N.
fn collect_processes() -> Vec<Metric> {
    let Ok(entries) = std::fs::read_dir("/proc") else {
        return vec![];
    };

    let page_kb = (unsafe { libc::sysconf(libc::_SC_PAGESIZE) }).max(0) as u64 / 1024;
    let mut by_name: HashMap<String, (u64, u64)> = HashMap::new(); // name -> (rss_kb, count)

    for entry in entries.flatten() {
        if entry.file_name().to_string_lossy().parse::<u32>().is_err() {
            continue; // not a PID directory (self, net, sys, ...)
        }
        let base = entry.path();
        let Ok(comm) = std::fs::read_to_string(base.join("comm")) else {
            continue; // process exited between readdir and read
        };
        let Ok(statm) = std::fs::read_to_string(base.join("statm")) else {
            continue;
        };
        let Some(resident_pages) = statm
            .split_whitespace()
            .nth(1)
            .and_then(|s| s.parse::<u64>().ok())
        else {
            continue;
        };
        let name = sanitize_segment(comm.trim());
        let slot = by_name.entry(name).or_insert((0, 0));
        slot.0 += resident_pages * page_kb;
        slot.1 += 1;
    }

    let top_n = config::get().proc_top_n;
    let mut ranked: Vec<(String, u64, u64)> = by_name
        .into_iter()
        .map(|(name, (rss_kb, count))| (name, rss_kb, count))
        .collect();
    ranked.sort_by_key(|entry| std::cmp::Reverse(entry.1));
    ranked.truncate(top_n);

    let mut metrics = Vec::with_capacity(ranked.len() * 2);
    for (name, rss_kb, count) in ranked {
        metrics.push(Metric {
            name: format!("proc.{name}.mem_kb"),
            value: rss_kb as f64,
        });
        metrics.push(Metric {
            name: format!("proc.{name}.count"),
            value: count as f64,
        });
    }
    metrics
}

/// Memory, active state, and restart count for allowlisted systemd service
/// units (`MONITOR_SYSTEMD_MATCH` prefixes, default = journal match list),
/// plus a host-wide failed-unit count that is always emitted regardless of
/// the allowlist so "any service failed" alerting needs no configuration.
fn collect_systemd() -> Vec<Metric> {
    let mut metrics = Vec::new();
    let Ok(list_output) = std::process::Command::new("systemctl")
        .args([
            "list-units",
            "--type=service",
            "--all",
            "--no-legend",
            "--plain",
            "--no-pager",
        ])
        .output()
    else {
        return metrics; // systemctl unavailable (e.g. non-systemd host)
    };
    if !list_output.status.success() {
        return metrics;
    }

    let prefixes = &config::get().systemd_match;
    let mut failed_count = 0u64;
    let mut matched_units = Vec::new();
    for line in String::from_utf8_lossy(&list_output.stdout).lines() {
        let mut columns = line.split_whitespace();
        let Some(unit) = columns.next().filter(|u| u.ends_with(".service")) else {
            continue;
        };
        // Columns: UNIT LOAD ACTIVE SUB DESCRIPTION...
        if columns.nth(1) == Some("failed") {
            failed_count += 1;
        }
        let base = unit.trim_end_matches(".service");
        if prefixes
            .iter()
            .any(|prefix| base.starts_with(prefix.as_str()))
        {
            matched_units.push(unit.to_string());
        }
    }
    metrics.push(Metric {
        name: "systemd.units.failed".to_string(),
        value: failed_count as f64,
    });

    for unit in matched_units {
        let Ok(output) = std::process::Command::new("systemctl")
            .args(["show", &unit, "-p", "MemoryCurrent,ActiveState,NRestarts"])
            .output()
        else {
            continue;
        };
        if !output.status.success() {
            continue;
        }
        let text = String::from_utf8_lossy(&output.stdout);
        let mut mem_current: Option<u64> = None;
        let mut active = false;
        let mut restarts: Option<u64> = None;
        for line in text.lines() {
            let Some((key, value)) = line.split_once('=') else {
                continue;
            };
            match key {
                "MemoryCurrent" => mem_current = value.parse().ok(),
                "ActiveState" => active = value == "active",
                "NRestarts" => restarts = value.parse().ok(),
                _ => {}
            }
        }
        let name = sanitize_segment(unit.trim_end_matches(".service"));
        metrics.push(Metric {
            name: format!("systemd.{name}.active"),
            value: if active { 1.0 } else { 0.0 },
        });
        // MemoryCurrent reads back as u64::MAX ("[not set]") when memory
        // accounting is disabled for the unit; skip rather than emit a
        // bogus multi-exabyte sample.
        if let Some(bytes) = mem_current
            && bytes != u64::MAX
        {
            metrics.push(Metric {
                name: format!("systemd.{name}.mem_kb"),
                value: (bytes / 1024) as f64,
            });
        }
        if let Some(restarts) = restarts {
            metrics.push(Metric {
                name: format!("systemd.{name}.restarts"),
                value: restarts as f64,
            });
        }
    }

    metrics
}

/// Memory usage of every running container via the Docker CLI. Returns
/// empty when `docker` is absent, unreachable, or the caller lacks socket
/// access — Docker is an optional per-host dependency, not a requirement.
fn collect_docker() -> Vec<Metric> {
    let Ok(output) = std::process::Command::new("docker")
        .args([
            "stats",
            "--no-stream",
            "--format",
            "{{.Name}}\t{{.MemUsage}}",
        ])
        .output()
    else {
        return vec![];
    };
    if !output.status.success() {
        return vec![];
    }

    let mut metrics = Vec::new();
    let mut running = 0u64;
    for line in String::from_utf8_lossy(&output.stdout).lines() {
        let mut columns = line.splitn(2, '\t');
        let (Some(raw_name), Some(mem_usage)) = (columns.next(), columns.next()) else {
            continue;
        };
        // MemUsage is "<used> / <limit>", e.g. "128.3MiB / 1.9GiB".
        let Some(used_kb) = mem_usage
            .split('/')
            .next()
            .and_then(|used| parse_docker_memory_kb(used.trim()))
        else {
            continue;
        };
        running += 1;
        metrics.push(Metric {
            name: format!("docker.{}.mem_kb", sanitize_segment(raw_name)),
            value: used_kb,
        });
    }
    metrics.push(Metric {
        name: "docker.containers.running".to_string(),
        value: running as f64,
    });
    metrics
}

/// Parse Docker's human-readable memory string ("128.3MiB", "1.9GiB",
/// "512B") into kilobytes.
fn parse_docker_memory_kb(raw: &str) -> Option<f64> {
    let split_at = raw.find(|c: char| c.is_ascii_alphabetic())?;
    let (number, unit) = raw.split_at(split_at);
    let value: f64 = number.parse().ok()?;
    let kb_per_unit = match unit {
        "B" => 1.0 / 1024.0,
        "KiB" | "kB" => 1.0,
        "MiB" | "MB" => 1024.0,
        "GiB" | "GB" => 1024.0 * 1024.0,
        "TiB" | "TB" => 1024.0 * 1024.0 * 1024.0,
        _ => return None,
    };
    Some(value * kb_per_unit)
}

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

    #[test]
    fn sanitize_segment_collapses_unsafe_characters() {
        assert_eq!(sanitize_segment("nginx"), "nginx");
        assert_eq!(sanitize_segment("[kworker/0:1]"), "kworker_0_1");
        assert_eq!(
            sanitize_segment("akurai-monitor.service"),
            "akurai-monitor_service"
        );
        assert_eq!(sanitize_segment("..."), "unknown");
    }

    #[test]
    fn parses_docker_memory_units() {
        assert_eq!(parse_docker_memory_kb("512B"), Some(0.5));
        assert_eq!(parse_docker_memory_kb("128.3MiB"), Some(128.3 * 1024.0));
        assert_eq!(
            parse_docker_memory_kb("1.9GiB"),
            Some(1.9 * 1024.0 * 1024.0)
        );
        assert_eq!(parse_docker_memory_kb("4KiB"), Some(4.0));
        assert_eq!(parse_docker_memory_kb("bogus"), None);
    }
}