clawmates/crates/cm-api/src/beszel.rs

//! Bring-your-own Beszel hub client: authenticate to a workspace's Beszel hub
//! (PocketBase) and poll rich per-node metrics, mapping its `systems` to our fleet
//! nodes by hostname and upserting `node_metrics`. Server-side only.

use std::collections::HashMap;
use std::time::Duration;

use cm_db::repo::fleet_beszel::BeszelConn;
use cm_db::repo::node_metrics::NodeMetrics;
use cm_db::repo::{fleet_beszel, node_metrics, nodes};
use cm_domain::{NodeId, WorkspaceId};
use serde_json::{json, Value};
use sqlx::PgPool;

/// Authenticate to a Beszel hub → a PocketBase token. Tries the regular users
/// collection first, then superusers.
pub async fn authenticate(client: &reqwest::Client, conn: &BeszelConn) -> Result<String, String> {
    let base = conn.hub_url.trim_end_matches('/');
    for coll in ["users", "_superusers"] {
        let url = format!("{base}/api/collections/{coll}/auth-with-password");
        let resp = client
            .post(&url)
            .json(&json!({ "identity": conn.username, "password": conn.password }))
            .send()
            .await
            .map_err(|e| e.to_string())?;
        if resp.status().is_success() {
            let body: Value = resp.json().await.map_err(|e| e.to_string())?;
            if let Some(tok) = body.get("token").and_then(Value::as_str) {
                return Ok(tok.to_owned());
            }
        }
    }
    Err("auth failed (check hub URL + credentials)".into())
}

/// Fetch all `systems` records (current status + the `info` snapshot).
pub async fn fetch_systems(
    client: &reqwest::Client,
    conn: &BeszelConn,
    token: &str,
) -> Result<Vec<Value>, String> {
    let base = conn.hub_url.trim_end_matches('/');
    let resp = client
        .get(format!("{base}/api/collections/systems/records"))
        .query(&[("perPage", "500")])
        .header("Authorization", token)
        .send()
        .await
        .map_err(|e| e.to_string())?;
    if !resp.status().is_success() {
        return Err(format!("systems list HTTP {}", resp.status()));
    }
    let body: Value = resp.json().await.map_err(|e| e.to_string())?;
    Ok(body
        .get("items")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default())
}

/// Newest `1m` sample per system, in ONE request.
///
/// The alternative is a request per system per poll, which grows with the
/// fleet for data that arrives in a single sorted page. `perPage` is generous
/// rather than exact because several samples belong to the same system: sorted
/// newest-first, the FIRST row seen for a system id is its latest, so later
/// rows for that system are skipped.
///
/// A hub that cannot answer this is not an error — the caller falls back to the
/// `systems.info` snapshot, which is what it used before this existed. Losing
/// GPU and IO detail must not cost the CPU and memory that still work.
pub async fn fetch_latest_stats(
    client: &reqwest::Client,
    conn: &BeszelConn,
    token: &str,
) -> HashMap<String, Value> {
    let base = conn.hub_url.trim_end_matches('/');
    let resp = client
        .get(format!("{base}/api/collections/system_stats/records"))
        .query(&[
            ("perPage", "200"),
            ("sort", "-created"),
            ("filter", "type='1m'"),
        ])
        .header("Authorization", token)
        .send()
        .await;
    let Ok(resp) = resp else { return HashMap::new() };
    if !resp.status().is_success() {
        return HashMap::new();
    }
    let Ok(body) = resp.json::<Value>().await else {
        return HashMap::new();
    };
    let mut out: HashMap<String, Value> = HashMap::new();
    for row in body.get("items").and_then(Value::as_array).unwrap_or(&vec![]) {
        let Some(sid) = row.get("system").and_then(Value::as_str) else {
            continue;
        };
        if let Some(stats) = row.get("stats") {
            out.entry(sid.to_string()).or_insert_with(|| stats.clone());
        }
    }
    out
}

/// Proxy a system's recent 1m time-series (for the monitor-page charts).
pub async fn fetch_history(
    client: &reqwest::Client,
    conn: &BeszelConn,
    token: &str,
    system_id: &str,
) -> Result<Value, String> {
    let base = conn.hub_url.trim_end_matches('/');
    let filter = format!("system='{system_id}' && type='1m'");
    let resp = client
        .get(format!("{base}/api/collections/system_stats/records"))
        .query(&[
            ("perPage", "120"),
            ("sort", "-created"),
            ("filter", filter.as_str()),
        ])
        .header("Authorization", token)
        .send()
        .await
        .map_err(|e| e.to_string())?;
    if !resp.status().is_success() {
        return Err(format!("system_stats HTTP {}", resp.status()));
    }
    resp.json().await.map_err(|e| e.to_string())
}

fn f(v: &Value, k: &str) -> Option<f64> {
    v.get(k).and_then(Value::as_f64)
}

/// The n-th element of a numeric array field, as the integer the
/// `node_metrics` per-second columns store. Rounded rather than truncated: a
/// rate of 0.6 is traffic, and `as i64` would file it as silence.
fn pair(v: &Value, k: &str, idx: usize) -> Option<i64> {
    v.get(k)
        .and_then(Value::as_array)
        .and_then(|a| a.get(idx))
        .and_then(Value::as_f64)
        .map(|n| n.round() as i64)
}

/// Busiest GPU's utilisation percentage, from a `system_stats` sample.
///
/// `stats.g` is a MAP keyed by GPU index — `{"0":{"n":"GeForce RTX 5060 Ti",
/// "u":0,"p":4.38}}` — where `u` is utilisation and `p` is power draw. This is
/// why `gpu_pct` was null on every NVIDIA node: the old mapping read `info.g`
/// as a scalar, and `info` carries no `g` at all in Beszel 0.18. The data was
/// arriving the whole time, one collection away.
///
/// MAX rather than mean across GPUs: the question placement asks is "is there a
/// free GPU here", and averaging a saturated card with an idle one answers a
/// question nobody asked.
fn gpu_busiest(stats: &Value) -> Option<f64> {
    let gpus = stats.get("g")?.as_object()?;
    gpus.values()
        .filter_map(|g| g.get("u").and_then(Value::as_f64))
        .fold(None, |acc: Option<f64>, u| Some(acc.map_or(u, |a| a.max(u))))
}

/// Map a Beszel `systems` record into our NodeMetrics.
///
/// `stats` is the newest `system_stats` sample for this system, when there is
/// one. It carries everything the `systems.info` snapshot does not: GPU,
/// per-second network, per-second disk IO.
///
/// The array orders below were MEASURED against the hosts, not read off a
/// schema — an inverted pair here does not fail, it reports upload as download
/// forever:
///   - `b`   = [sent, recv]. `stats.ni` gives per-interface
///     `[sent_ps, recv_ps, total_sent, total_recv]`; indices 2 and 3 matched
///     `/proc/net/dev` tx_bytes and rx_bytes on all four of tank's interfaces,
///     and `b` is the sum of the per-second pair across them.
///   - `dio` = [read, write]. An 800 MB `dd` on tank moved index 1 from 7441 to
///     23688 while index 0 stayed near zero.
///
/// `info.ct` is NOT mapped to `container_count`: it reads 1 on tank (1
/// container) and also 1 on architect (4 containers), so whatever it counts, it
/// is not that.
fn metrics_from_system(system: &Value, stats: Option<&Value>) -> NodeMetrics {
    let info = system.get("info").cloned().unwrap_or_else(|| json!({}));
    let load1 = info
        .get("la")
        .and_then(Value::as_array)
        .and_then(|a| a.first())
        .and_then(Value::as_f64);
    let empty = json!({});
    let st = stats.unwrap_or(&empty);
    NodeMetrics {
        cpu_pct: f(&info, "cpu"),
        mem_pct: f(&info, "mp"),
        disk_pct: f(&info, "dp"),
        gpu_pct: gpu_busiest(st),
        temp_max: f(&info, "dt"),
        net_sent_ps: pair(st, "b", 0),
        net_recv_ps: pair(st, "b", 1),
        disk_read_ps: pair(st, "dio", 0),
        disk_write_ps: pair(st, "dio", 1),
        load1,
        container_count: None,
        data: json!({
            "beszelSystemId": system.get("id").and_then(Value::as_str),
            "status": system.get("status").and_then(Value::as_str),
            "name": system.get("name").and_then(Value::as_str),
            "host": system.get("host").and_then(Value::as_str),
            "info": info,
            // The GPU roster, so a card can name the card rather than only
            // report a percentage.
            "gpus": st.get("g").cloned().unwrap_or(Value::Null),
            "temps": st.get("t").cloned().unwrap_or(Value::Null),
        }),
    }
}

/// Poll one workspace's hub: auth, fetch systems, match to nodes by hostname,
/// upsert metrics. Returns how many nodes were updated.
pub async fn poll_workspace(
    client: &reqwest::Client,
    pool: &PgPool,
    ws: WorkspaceId,
    conn: &BeszelConn,
) -> Result<usize, String> {
    let token = authenticate(client, conn).await?;
    let systems = fetch_systems(client, conn, &token).await?;
    let stats = fetch_latest_stats(client, conn, &token).await;
    let node_rows = nodes::list(pool, ws).await.map_err(|e| e.to_string())?;
    // hostname/name (lowercased) → node id.
    let mut by_host: HashMap<String, NodeId> = HashMap::new();
    for n in &node_rows {
        if let Some(h) = n.hostname.as_deref() {
            by_host.insert(h.to_lowercase(), n.id);
        }
        by_host.entry(n.name.to_lowercase()).or_insert(n.id);
    }
    let mut updated = 0;
    for sys in &systems {
        let key = sys
            .get("name")
            .and_then(Value::as_str)
            .or_else(|| sys.get("host").and_then(Value::as_str))
            .map(str::to_lowercase);
        let Some(node_id) = key.as_deref().and_then(|k| by_host.get(k).copied()) else {
            continue;
        };
        let sample = sys
            .get("id")
            .and_then(Value::as_str)
            .and_then(|id| stats.get(id));
        if node_metrics::upsert(pool, node_id, &metrics_from_system(sys, sample))
            .await
            .is_ok()
        {
            updated += 1;
        }
    }
    Ok(updated)
}

/// Spawn the background poller: every `interval`, refresh node_metrics for each
/// workspace with a Beszel hub connected.
pub fn spawn_poller(pool: PgPool, interval: Duration) {
    tokio::spawn(async move {
        let client = reqwest::Client::new();
        let mut tick = tokio::time::interval(interval);
        loop {
            tick.tick().await;
            let workspaces = match fleet_beszel::all(&pool).await {
                Ok(w) => w,
                Err(_) => continue,
            };
            for (ws, conn) in workspaces {
                if let Err(e) = poll_workspace(&client, &pool, ws, &conn).await {
                    eprintln!("beszel poll (ws {}): {e}", ws.as_uuid());
                }
            }
        }
    });
}

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

    /// A real 0.18.7 sample, copied from tank rather than invented.
    fn sample() -> Value {
        json!({
            "b": [1830, 1811],
            "dio": [204, 23688],
            "g": { "0": { "n": "GeForce RTX 5060 Ti", "u": 37.5, "p": 4.38 } },
            "t": { "GeForce RTX 5060 Ti": 29, "k10temp_tctl": 38.38 }
        })
    }

    fn system() -> Value {
        json!({
            "id": "glo9hj260jhnlgr",
            "name": "tank",
            "host": "100.108.129.81",
            "status": "up",
            "info": { "cpu": 0.31, "mp": 7.14, "dp": 77.96, "dt": 38.85, "la": [0.03, 0.01, 0], "ct": 1 }
        })
    }

    /// GPU comes from the stats sample's MAP, not from `info`.
    ///
    /// This is the bug the whole change exists for: `info` carries no `g` in
    /// 0.18, so reading it as a scalar produced null on every NVIDIA node while
    /// the data sat one collection away. Null and "no GPU" are indistinguishable
    /// downstream, so metrics-aware placement simply never saw a GPU.
    #[test]
    fn gpu_comes_from_the_stats_sample_not_the_info_snapshot() {
        let m = metrics_from_system(&system(), Some(&sample()));
        assert_eq!(m.gpu_pct, Some(37.5));
        // No sample ⇒ no GPU claim. NOT zero: "we did not get a reading" and
        // "the card is idle" are different facts.
        assert_eq!(metrics_from_system(&system(), None).gpu_pct, None);
    }

    /// The busiest card, not the average.
    #[test]
    fn a_saturated_card_is_not_averaged_away_by_an_idle_one() {
        let two = json!({ "g": { "0": { "u": 99.0 }, "1": { "u": 1.0 } } });
        assert_eq!(gpu_busiest(&two), Some(99.0));
        assert_eq!(gpu_busiest(&json!({})), None);
        // Present but empty is still no reading.
        assert_eq!(gpu_busiest(&json!({ "g": {} })), None);
    }

    /// The measured array orders. An inverted pair does not fail — it reports
    /// upload as download, and disk reads as writes, forever.
    ///
    /// `b`   = [sent, recv]: `stats.ni` per-interface indices 2 and 3 matched
    ///         `/proc/net/dev` tx_bytes and rx_bytes on all four of tank's
    ///         interfaces, and `b` is the sum of the per-second pair.
    /// `dio` = [read, write]: an 800 MB `dd` moved index 1 from 7441 to 23688
    ///         while index 0 stayed near zero.
    #[test]
    fn the_measured_array_orders_are_not_reinverted() {
        let m = metrics_from_system(&system(), Some(&sample()));
        assert_eq!(m.net_sent_ps, Some(1830), "b[0] is SENT");
        assert_eq!(m.net_recv_ps, Some(1811), "b[1] is RECV");
        assert_eq!(m.disk_read_ps, Some(204), "dio[0] is READ");
        assert_eq!(m.disk_write_ps, Some(23688), "dio[1] is WRITE");
    }

    /// `info.ct` must not become `container_count`.
    ///
    /// It reads 1 on tank, which runs 1 container, and ALSO 1 on architect,
    /// which runs 4. It agrees with the truth exactly often enough to look
    /// right in a spot check.
    #[test]
    fn the_unidentified_ct_field_is_not_reported_as_a_container_count() {
        let m = metrics_from_system(&system(), Some(&sample()));
        assert_eq!(m.container_count, None);
        assert_eq!(system()["info"]["ct"], json!(1));
    }

    /// The snapshot fields keep working when the stats call fails.
    #[test]
    fn a_missing_stats_sample_does_not_cost_the_metrics_that_still_work() {
        let m = metrics_from_system(&system(), None);
        assert_eq!(m.cpu_pct, Some(0.31));
        assert_eq!(m.mem_pct, Some(7.14));
        assert_eq!(m.temp_max, Some(38.85));
        assert_eq!(m.load1, Some(0.03));
        assert_eq!(m.net_sent_ps, None);
    }
}