compute/src/governor.rs

//! The Governor — the kernel of Lumbridge Compute.
//!
//! On a unified-memory box, over-commit doesn't fail gracefully: the whole machine
//! thrashes and wedges (SSH/ping included) before the OOM killer acts. The Governor
//! makes that impossible via (1) admission control against a hard budget and
//! (2) a watchdog on `MemAvailable` that kills the newest model before thrash.
//!
//! This module currently provides the *sensing* + *admission* half. The watchdog and
//! process spawn/kill land with the `up`/`activate` commands.

use crate::config::{Model, Registry};
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream};
use std::time::{Duration, Instant};

/// Usable memory for models; the rest is reserved for the OS/desktop.
pub const DEFAULT_BUDGET_GB: f64 = 100.0;
/// Extra headroom required before admitting a new model.
pub const SAFETY_MARGIN_GB: f64 = 8.0;
/// If `MemAvailable` dips below this, the watchdog kills the newest model.
pub const WATCHDOG_FLOOR_GB: f64 = 3.0;

/// A model is considered "running" if its serving port accepts a connection.
pub fn port_open(port: u16) -> bool {
    let addr = SocketAddr::from(([127, 0, 0, 1], port));
    TcpStream::connect_timeout(&addr, Duration::from_millis(150)).is_ok()
}

pub fn is_running(m: &Model) -> bool {
    let Some(port) = m.serve.port else {
        return false;
    };
    if !port_open(port) {
        return false;
    }
    match &m.health_contains {
        Some(marker) => health_response(m, port)
            .map(|response| response.contains(marker))
            .unwrap_or(false),
        None => true,
    }
}

/// Fetch the local HTTP health endpoint without adding an HTTP client runtime.
/// Registry health URLs are deliberately localhost-only.
fn health_response(m: &Model, port: u16) -> Option<String> {
    let path = m
        .health
        .as_deref()
        .and_then(|url| url.split_once("localhost"))
        .map(|(_, tail)| tail.trim_start_matches(|c: char| c.is_ascii_digit() || c == ':'))
        .filter(|path| path.starts_with('/'))
        .unwrap_or("/");
    let addr = SocketAddr::from(([127, 0, 0, 1], port));
    let mut stream = TcpStream::connect_timeout(&addr, Duration::from_millis(300)).ok()?;
    stream
        .set_read_timeout(Some(Duration::from_millis(500)))
        .ok()?;
    stream
        .write_all(format!("GET {path} HTTP/1.0\r\nHost: localhost\r\n\r\n").as_bytes())
        .ok()?;
    let mut response = String::new();
    stream.read_to_string(&mut response).ok()?;
    Some(response)
}

pub fn running_ids(reg: &Registry) -> Vec<String> {
    reg.models
        .iter()
        .filter(|(_, m)| is_running(m))
        .map(|(id, _)| id.clone())
        .collect()
}

/// Sum of footprints of all currently-serving registered models.
pub fn committed_gb(reg: &Registry) -> f64 {
    reg.models
        .values()
        .filter(|m| is_running(m))
        .map(|m| m.footprint_gb)
        .sum()
}

/// Admission control. Two independent ceilings, and both must hold.
///
/// 1. **Declared** — `committed + add + margin <= budget`.
/// 2. **Observed** — `add + margin <= available`.
///
/// The second one is what makes the promise real. `budget_gb`, and every
/// `footprint_gb` feeding `committed_gb`, are numbers a human typed into the
/// registry. If any of them is optimistic — a model that grows past its declared
/// footprint, a KV cache larger than expected, anything started outside Compute —
/// the declared check happily passes while the box is already out of memory, and
/// on unified memory that ends in a wedged machine rather than a failed malloc.
/// `available_gb` comes from `MemAvailable`, which counts reality.
///
/// `available_gb` is `None` when sensing failed. That falls back to the declared
/// ceiling alone: refusing every start because /proc/meminfo was unreadable would
/// turn a sensing failure into a total outage.
pub fn can_admit(
    add_gb: f64,
    committed_gb: f64,
    budget_gb: f64,
    available_gb: Option<f64>,
) -> bool {
    if committed_gb + add_gb + SAFETY_MARGIN_GB > budget_gb {
        return false;
    }
    match available_gb {
        Some(available) => add_gb + SAFETY_MARGIN_GB <= available,
        None => true,
    }
}

/// What admission will actually enforce right now, for reporting. The lower of the
/// declared headroom and the observed headroom.
pub fn headroom_gb(committed_gb: f64, budget_gb: f64, available_gb: Option<f64>) -> f64 {
    let declared = budget_gb - committed_gb - SAFETY_MARGIN_GB;
    match available_gb {
        Some(available) => declared.min(available - SAFETY_MARGIN_GB),
        None => declared,
    }
}

/// Block until the exact registered model is healthy, or `timeout` elapses.
/// Models without a port are considered instantly ready.
pub fn wait_healthy(m: &Model, timeout: Duration) -> bool {
    let Some(_) = m.serve.port else {
        return true;
    };
    let start = Instant::now();
    while start.elapsed() < timeout {
        if is_running(m) {
            return true;
        }
        std::thread::sleep(Duration::from_millis(400));
    }
    false
}

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

    // A 108 GB Scene budget on a box with plenty free — the ordinary case.
    const BUDGET: f64 = 108.0;

    #[test]
    fn admits_when_both_declared_and_observed_ceilings_allow_it() {
        assert!(can_admit(66.0, 9.0, BUDGET, Some(100.0)));
    }

    #[test]
    fn refuses_when_the_declared_budget_is_exceeded() {
        // 66 + 40 + 8 margin = 114 > 108, even though the box has memory free.
        assert!(!can_admit(66.0, 40.0, BUDGET, Some(100.0)));
    }

    #[test]
    fn refuses_when_the_box_is_out_of_memory_even_though_the_paperwork_agrees() {
        // This is the case the declared check alone could never catch: nothing is
        // registered as committed, so the budget says there is 100 GB of room, but
        // MemAvailable says 20 GB. Something outside Compute is holding the pool.
        assert!(can_admit(66.0, 0.0, BUDGET, None), "declared check passes");
        assert!(!can_admit(66.0, 0.0, BUDGET, Some(20.0)));
    }

    #[test]
    fn the_safety_margin_is_enforced_against_observed_memory_too() {
        // 66 GB model with exactly 66 GB free is a refusal: the margin has to fit.
        assert!(!can_admit(66.0, 0.0, BUDGET, Some(66.0)));
        assert!(!can_admit(
            66.0,
            0.0,
            BUDGET,
            Some(66.0 + SAFETY_MARGIN_GB - 0.1)
        ));
        assert!(can_admit(66.0, 0.0, BUDGET, Some(66.0 + SAFETY_MARGIN_GB)));
    }

    #[test]
    fn failed_sensing_falls_back_to_the_declared_ceiling_rather_than_refusing_everything() {
        assert!(can_admit(66.0, 9.0, BUDGET, None));
        // ...but it must not become a way to bypass the declared budget.
        assert!(!can_admit(66.0, 40.0, BUDGET, None));
    }

    #[test]
    fn a_zero_reading_refuses_everything_rather_than_admitting_everything() {
        // parse_meminfo yields 0.0 for an unparseable /proc. That has to read as
        // "no memory", not as "no constraint".
        assert!(!can_admit(1.0, 0.0, BUDGET, Some(0.0)));
    }

    #[test]
    fn headroom_never_exceeds_what_the_machine_actually_has() {
        // The shape of a real regression: a generous declared budget on a smaller box. Reporting
        // budget-minus-committed here promises room the next admission will refuse, and both the
        // MCP tool and the HTTP API serve this number to callers that cannot check it themselves.
        let box_available = 55.8;
        let declared_budget = 100.0;
        let reported = headroom_gb(0.0, declared_budget, Some(box_available));
        assert!(
            reported <= box_available,
            "reported {reported} GB of headroom on a box with {box_available} GB free",
        );
        assert!(can_admit(
            reported,
            0.0,
            declared_budget,
            Some(box_available)
        ));
    }

    #[test]
    fn headroom_reports_the_binding_constraint_not_the_generous_one() {
        // Declared says 33 GB spare; the box says 12 GB spare minus margin.
        assert_eq!(headroom_gb(67.0, BUDGET, None), 33.0);
        assert_eq!(headroom_gb(67.0, BUDGET, Some(12.0)), 4.0);
        // And the declared ceiling still wins when it is the tighter of the two.
        assert_eq!(headroom_gb(100.0, BUDGET, Some(90.0)), 0.0);
    }
}