OpenPXE/crates/core/src/metrics.rs

//! Tiny lock-free Prometheus-compatible metrics.
//!
//! We don't pull in `prometheus` or `metrics-rs` for this — they bring
//! their own runtime, registry, and complexity. OpenPXE has a fixed,
//! tiny set of counters/gauges and the exposition format is plain text.
//! A handful of `AtomicU64`s and a `Display` impl gets us everything
//! Prometheus / Grafana / VictoriaMetrics needs to scrape:
//!
//!   openpxe_dhcp_replies_total       counter (per arch label)
//!   openpxe_tftp_transfers_total     counter (per status label)
//!   openpxe_tftp_bytes_total         counter
//!   openpxe_http_requests_total      counter (per route label)
//!   openpxe_iso_count                gauge
//!   openpxe_client_count             gauge
//!   openpxe_queue_count               gauge
//!   openpxe_queue_imaging             gauge
//!   openpxe_uptime_seconds           gauge
//!   openpxe_build_info{version}      gauge (always 1)
//!
//! Cheap to clone — internal state is a couple of arcs. Counters use
//! `Relaxed` ordering: we don't synchronise across counters, just need
//! per-counter monotonicity.

use std::fmt::Write as _;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

#[derive(Debug, Default)]
#[allow(clippy::struct_field_names)]
struct Inner {
    // DHCP proxy
    dhcp_replies_legacy: AtomicU64,
    dhcp_replies_uefi: AtomicU64,
    dhcp_replies_arm64: AtomicU64,
    dhcp_replies_unknown: AtomicU64,
    dhcp_declined: AtomicU64,
    // TFTP
    tftp_transfers_ok: AtomicU64,
    tftp_transfers_err: AtomicU64,
    tftp_bytes: AtomicU64,
    // HTTP
    http_boot_script: AtomicU64,
    http_iso_range: AtomicU64,
    http_iso_inner: AtomicU64,
    http_ipxe_binary: AtomicU64,
    http_api: AtomicU64,
    // Gauges (set explicitly; not cumulative)
    iso_count: AtomicU64,
    client_count: AtomicU64,
    queue_count: AtomicU64,
    queue_imaging: AtomicU64,
    nfs_mounts_active: AtomicU64,
}

#[derive(Debug, Clone, Default)]
pub struct Metrics {
    inner: Arc<Inner>,
}

impl Metrics {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    // ── DHCP ───────────────────────────────────────────────────────────

    pub fn record_dhcp_reply(&self, arch: &str) {
        let counter = match arch {
            "bios" => &self.inner.dhcp_replies_legacy,
            "uefi-x64" | "uefi-ia32" => &self.inner.dhcp_replies_uefi,
            "uefi-arm64" => &self.inner.dhcp_replies_arm64,
            _ => &self.inner.dhcp_replies_unknown,
        };
        counter.fetch_add(1, Ordering::Relaxed);
    }

    pub fn record_dhcp_decline(&self) {
        self.inner.dhcp_declined.fetch_add(1, Ordering::Relaxed);
    }

    // ── TFTP ───────────────────────────────────────────────────────────

    pub fn record_tftp_ok(&self, bytes: u64) {
        self.inner.tftp_transfers_ok.fetch_add(1, Ordering::Relaxed);
        self.inner.tftp_bytes.fetch_add(bytes, Ordering::Relaxed);
    }

    pub fn record_tftp_err(&self) {
        self.inner.tftp_transfers_err.fetch_add(1, Ordering::Relaxed);
    }

    // ── HTTP ───────────────────────────────────────────────────────────

    pub fn record_http(&self, route: HttpRoute) {
        let counter = match route {
            HttpRoute::BootScript => &self.inner.http_boot_script,
            HttpRoute::IsoRange => &self.inner.http_iso_range,
            HttpRoute::IsoInner => &self.inner.http_iso_inner,
            HttpRoute::IpxeBinary => &self.inner.http_ipxe_binary,
            HttpRoute::Api => &self.inner.http_api,
        };
        counter.fetch_add(1, Ordering::Relaxed);
    }

    // ── Gauges ─────────────────────────────────────────────────────────

    pub fn set_iso_count(&self, n: u64) {
        self.inner.iso_count.store(n, Ordering::Relaxed);
    }

    pub fn set_client_count(&self, n: u64) {
        self.inner.client_count.store(n, Ordering::Relaxed);
    }

    pub fn set_queue_counts(&self, total: u64, imaging: u64) {
        self.inner.queue_count.store(total, Ordering::Relaxed);
        self.inner.queue_imaging.store(imaging, Ordering::Relaxed);
    }

    pub fn set_nfs_active(&self, n: u64) {
        self.inner.nfs_mounts_active.store(n, Ordering::Relaxed);
    }

    /// Render in the Prometheus text exposition format.
    /// Uptime is supplied by the caller because `Metrics` doesn't own
    /// the start instant; the HTTP layer does.
    #[must_use]
    pub fn render(&self, version: &str, uptime_secs: u64) -> String {
        let mut out = String::with_capacity(2048);
        let i = &self.inner;

        // Helper closures.
        let write_counter = |o: &mut String, name: &str, help: &str, val: u64, lbl: &str| {
            let _ = writeln!(o, "# HELP {name} {help}");
            let _ = writeln!(o, "# TYPE {name} counter");
            if lbl.is_empty() {
                let _ = writeln!(o, "{name} {val}");
            } else {
                let _ = writeln!(o, "{name}{{{lbl}}} {val}");
            }
        };
        let write_gauge = |o: &mut String, name: &str, help: &str, val: u64, lbl: &str| {
            let _ = writeln!(o, "# HELP {name} {help}");
            let _ = writeln!(o, "# TYPE {name} gauge");
            if lbl.is_empty() {
                let _ = writeln!(o, "{name} {val}");
            } else {
                let _ = writeln!(o, "{name}{{{lbl}}} {val}");
            }
        };

        // Counters with one HELP/TYPE per metric name and per-label rows.
        let _ = writeln!(out, "# HELP openpxe_dhcp_replies_total Number of proxyDHCP replies sent, by client architecture.");
        let _ = writeln!(out, "# TYPE openpxe_dhcp_replies_total counter");
        let _ = writeln!(
            out,
            "openpxe_dhcp_replies_total{{arch=\"bios\"}} {}",
            i.dhcp_replies_legacy.load(Ordering::Relaxed)
        );
        let _ = writeln!(
            out,
            "openpxe_dhcp_replies_total{{arch=\"uefi\"}} {}",
            i.dhcp_replies_uefi.load(Ordering::Relaxed)
        );
        let _ = writeln!(
            out,
            "openpxe_dhcp_replies_total{{arch=\"arm64\"}} {}",
            i.dhcp_replies_arm64.load(Ordering::Relaxed)
        );
        let _ = writeln!(
            out,
            "openpxe_dhcp_replies_total{{arch=\"unknown\"}} {}",
            i.dhcp_replies_unknown.load(Ordering::Relaxed)
        );
        write_counter(
            &mut out,
            "openpxe_dhcp_declined_total",
            "DHCP requests we saw but did not reply to (mac filter, arch unsupported, etc).",
            i.dhcp_declined.load(Ordering::Relaxed),
            "",
        );

        let _ = writeln!(out, "# HELP openpxe_tftp_transfers_total TFTP transfers, by status.");
        let _ = writeln!(out, "# TYPE openpxe_tftp_transfers_total counter");
        let _ = writeln!(
            out,
            "openpxe_tftp_transfers_total{{status=\"ok\"}} {}",
            i.tftp_transfers_ok.load(Ordering::Relaxed)
        );
        let _ = writeln!(
            out,
            "openpxe_tftp_transfers_total{{status=\"err\"}} {}",
            i.tftp_transfers_err.load(Ordering::Relaxed)
        );
        write_counter(
            &mut out,
            "openpxe_tftp_bytes_total",
            "Total bytes successfully delivered over TFTP.",
            i.tftp_bytes.load(Ordering::Relaxed),
            "",
        );

        let _ = writeln!(out, "# HELP openpxe_http_requests_total HTTP requests served, by route family.");
        let _ = writeln!(out, "# TYPE openpxe_http_requests_total counter");
        for (label, counter) in [
            ("boot_script", &i.http_boot_script),
            ("iso_range", &i.http_iso_range),
            ("iso_inner", &i.http_iso_inner),
            ("ipxe_binary", &i.http_ipxe_binary),
            ("api", &i.http_api),
        ] {
            let _ = writeln!(
                out,
                "openpxe_http_requests_total{{route=\"{label}\"}} {}",
                counter.load(Ordering::Relaxed)
            );
        }

        // Gauges.
        write_gauge(&mut out, "openpxe_iso_count", "ISOs currently registered (local + NFS).", i.iso_count.load(Ordering::Relaxed), "");
        write_gauge(&mut out, "openpxe_client_count", "PXE clients seen this process lifetime.", i.client_count.load(Ordering::Relaxed), "");
        write_gauge(&mut out, "openpxe_queue_count", "Clients currently waiting at the deployment queue.", i.queue_count.load(Ordering::Relaxed), "");
        write_gauge(&mut out, "openpxe_queue_imaging", "Clients currently imaging (queue + assigned target).", i.queue_imaging.load(Ordering::Relaxed), "");
        write_gauge(&mut out, "openpxe_nfs_mounts_active", "NFS shares currently mounted.", i.nfs_mounts_active.load(Ordering::Relaxed), "");
        write_gauge(&mut out, "openpxe_uptime_seconds", "Seconds since this OpenPXE instance started.", uptime_secs, "");

        let _ = writeln!(out, "# HELP openpxe_build_info Build metadata. Always 1; the version is in the label.");
        let _ = writeln!(out, "# TYPE openpxe_build_info gauge");
        let _ = writeln!(out, "openpxe_build_info{{version=\"{version}\"}} 1");

        out
    }
}

/// Stable label values for the HTTP route counter. Adding a new route
/// here without updating `record_http` will break compilation, which is
/// exactly the safety we want — Prometheus alerts on cardinality drift,
/// so accidental new label values matter.
#[derive(Debug, Clone, Copy)]
pub enum HttpRoute {
    BootScript,
    IsoRange,
    IsoInner,
    IpxeBinary,
    Api,
}

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

    #[test]
    fn render_emits_each_metric_family_once() {
        let m = Metrics::new();
        m.record_dhcp_reply("uefi-x64");
        m.record_dhcp_reply("bios");
        m.record_tftp_ok(1024);
        m.record_http(HttpRoute::Api);
        m.set_iso_count(3);
        let out = m.render("0.2.0", 42);
        assert_eq!(out.matches("# TYPE openpxe_dhcp_replies_total counter").count(), 1);
        assert_eq!(out.matches("# TYPE openpxe_iso_count gauge").count(), 1);
        assert!(out.contains("openpxe_dhcp_replies_total{arch=\"uefi\"} 1"));
        assert!(out.contains("openpxe_dhcp_replies_total{arch=\"bios\"} 1"));
        assert!(out.contains("openpxe_tftp_transfers_total{status=\"ok\"} 1"));
        assert!(out.contains("openpxe_tftp_bytes_total 1024"));
        assert!(out.contains("openpxe_http_requests_total{route=\"api\"} 1"));
        assert!(out.contains("openpxe_iso_count 3"));
        assert!(out.contains("openpxe_uptime_seconds 42"));
        assert!(out.contains("openpxe_build_info{version=\"0.2.0\"} 1"));
    }

    #[test]
    fn cloned_metrics_share_state() {
        let a = Metrics::new();
        let b = a.clone();
        a.record_dhcp_reply("bios");
        b.record_dhcp_reply("bios");
        let out = a.render("test", 0);
        assert!(out.contains("openpxe_dhcp_replies_total{arch=\"bios\"} 2"));
    }
}