navian-memcheck-cli 0.1.0

//! `navian-memcheck` — command-line soak runner.
//!
//! Runs a target command under sustained conditions, polls its resident memory
//! on a fixed interval, and applies the same plateau assertion the library uses
//! in `cargo test`: after warmup, RSS must level off. Exits non-zero when memory
//! keeps climbing or blows past a cap, so it gates CI directly.
//!
//! ```text
//! navian-memcheck soak --cmd "./target/release/my-server --load" \
//!   --duration 20m --interval 5s --slope-budget 2mb --max 6gb --format json
//! ```

use std::io::Write as _;
use std::process::{Command, Stdio};
use std::thread::sleep;
use std::time::{Duration, Instant};

use navian_memcheck::{report_from_samples, SoakConfig, Verdict};

const USAGE: &str = "\
navian-memcheck — soak a command and prove its memory plateaus

USAGE:
    navian-memcheck soak --cmd <COMMAND> [OPTIONS]

REQUIRED:
    --cmd <COMMAND>        Command to run under load (executed via `sh -c`)

STOP CONDITION (default --duration 10m; if both are given, whichever is hit first):
    --duration <DUR>       Wall-clock budget, e.g. 90s, 20m, 1h
    --samples <N>          Stop after N samples

OPTIONS:
    --interval <DUR>       Time between RSS samples (default 5s)
    --warmup <FRAC>        Fraction discarded before the plateau fit (default 0.5)
    --slope-budget <SIZE>  Max back-half growth per sample, e.g. 2mb (default 1mb)
    --max <SIZE>           Hard ceiling on peak RSS, e.g. 6gb (default: none)
    --min-movement <SIZE>  Require RSS to span at least SIZE (peak-trough), else
                           report INCONCLUSIVE — proves the run actually exercised
                           memory (default: off; a flat run passes but reports moved)
    --format <text|json>   Output format (default text)
    -h, --help             Print this help

The target runs in its own process group; the whole tree's RSS is summed, and the
whole group is stopped at the budget. If the target exits on its own before the
budget, the run is INCONCLUSIVE (a truncated soak can't certify a plateau) — for a
finite workload, size --samples so the tool stops it first.

LIMITATION: RSS is summed over this one process GROUP. A target that re-groups —
setsid(2), or a double-fork daemonize — leaves the group and its memory stops being
counted, so a leak there can read as a flat (under-counted) plateau. Don't soak a
process that daemonizes; run its foreground/no-detach mode instead.

EXIT CODES:
    0  memory plateaued (and stayed under --max), full soak completed
    1  still growing, or exceeded --max
    2  target crashed or exited early, too few samples, or a usage error
";

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    match run(&args) {
        Ok(code) => std::process::exit(code),
        Err(e) => {
            eprintln!("navian-memcheck: {e}");
            std::process::exit(2);
        }
    }
}

fn run(args: &[String]) -> Result<i32, String> {
    if args.is_empty() || args.iter().any(|a| a == "-h" || a == "--help") {
        print!("{USAGE}");
        return Ok(if args.is_empty() { 2 } else { 0 });
    }
    if args[0] != "soak" {
        return Err(format!(
            "unknown subcommand `{}` (expected `soak`)",
            args[0]
        ));
    }

    let mut opts = Opts::default();
    let mut it = args[1..].iter();
    while let Some(flag) = it.next() {
        let mut val = || {
            it.next()
                .cloned()
                .ok_or_else(|| format!("{flag} needs a value"))
        };
        match flag.as_str() {
            "--cmd" => opts.cmd = Some(val()?),
            "--duration" => opts.duration = Some(parse_duration(&val()?)?),
            "--samples" => opts.max_samples = Some(parse_u64(&val()?)?),
            "--interval" => {
                opts.interval = parse_duration(&val()?)?;
                if opts.interval.is_zero() {
                    // A zero interval either certifies a "soak" over ~no elapsed time
                    // (with --samples) or busy-spins spawning `ps` (without). Reject it.
                    return Err("--interval must be greater than zero".into());
                }
            }
            "--warmup" => {
                let w = parse_f64(&val()?)?;
                if !w.is_finite() || !(0.0..=0.95).contains(&w) {
                    return Err(format!(
                        "--warmup must be a number in 0.0..=0.95, got `{w}`"
                    ));
                }
                opts.warmup = w;
            }
            "--slope-budget" => opts.slope_budget = parse_size(&val()?)? as f64,
            "--max" => opts.max_bytes = Some(parse_size(&val()?)?),
            "--min-movement" => opts.min_movement = parse_size(&val()?)?,
            "--format" => {
                opts.json = match val()?.as_str() {
                    "json" => true,
                    "text" => false,
                    other => {
                        return Err(format!("--format must be `json` or `text`, got `{other}`"))
                    }
                };
            }
            other => return Err(format!("unknown flag `{other}` (try --help)")),
        }
    }

    let cmd = opts.cmd.clone().ok_or("--cmd is required")?;
    if opts.duration.is_none() && opts.max_samples.is_none() {
        opts.duration = Some(Duration::from_secs(600)); // default 10m
    }

    // Spawn the target under a shell so pipes/flags work as written. On Unix we put
    // it in its own process group so we can sample and signal the whole tree — not
    // just the shell, which typically forks the real workload as a child.
    let mut command = Command::new("sh");
    command
        .arg("-c")
        .arg(&cmd)
        .stdout(Stdio::null())
        .stderr(Stdio::null());
    #[cfg(unix)]
    {
        use std::os::unix::process::CommandExt;
        command.process_group(0);
    }
    let mut child = command
        .spawn()
        .map_err(|e| format!("failed to spawn command: {e}"))?;
    // The child is a new group leader, so its process-group id equals its pid, and
    // every descendant inherits it. We sample and kill by group.
    let pgid = child.id();

    // Kill the whole soaked group on EVERY exit path — including the error returns in
    // the loop below (RSS-read failure, wait error). Without this, an early error left
    // the workload running. Signalling an already-dead group is a harmless no-op.
    let _group_guard = GroupGuard(pgid);

    if !opts.json {
        eprintln!("navian-memcheck: soaking process group {pgid} — {cmd}");
    }

    let mut samples: Vec<(u64, u64)> = Vec::new();
    let start = Instant::now();
    let mut tick: u64 = 0;
    let mut child_exited: Option<i32> = None;
    let mut consecutive_misses = 0u32;

    loop {
        sleep(opts.interval);

        // Did the target exit on its own? Any self-exit makes the run INCONCLUSIVE:
        // the tool only certifies a plateau over a soak window IT bounded (via
        // --samples/--duration, stopping a still-running target). A target that exits
        // first — even one tick before the budget — may have crashed or simply
        // finished, and a truncated lifetime can't prove memory levels off. Erring
        // toward inconclusive (never a false pass) is the whole point. For a finite
        // workload, size --samples so the tool stops it first.
        match child.try_wait() {
            Ok(Some(status)) => {
                child_exited = Some(status.code().unwrap_or(-1));
                break;
            }
            Ok(None) => {}
            Err(e) => return Err(format!("waiting on child failed: {e}")),
        }

        if let Some(bytes) = read_group_rss_bytes(pgid) {
            consecutive_misses = 0;
            samples.push((tick, bytes));
            tick += 1;
            if !opts.json {
                eprint!("\r  sample {tick}: {:.1} MB   ", bytes as f64 / 1_048_576.0);
                let _ = std::io::stderr().flush();
            }
        } else {
            // A failed read is ambiguous: the child may have exited, or `ps` may
            // have hiccupped transiently. Re-check the child before concluding.
            if let Ok(Some(status)) = child.try_wait() {
                child_exited = Some(status.code().unwrap_or(-1));
                break;
            }
            // Child is still alive — treat as a transient miss and keep going,
            // bounded so a permanently-broken `ps` can't loop forever.
            consecutive_misses += 1;
            if consecutive_misses > 5 {
                return Err(
                    "could not read the target's RSS after repeated attempts (is `ps` available?)"
                        .into(),
                );
            }
        }

        if let Some(max) = opts.max_samples {
            if tick >= max {
                break;
            }
        }
        if let Some(dur) = opts.duration {
            if start.elapsed() >= dur {
                break;
            }
        }
    }
    if !opts.json {
        eprintln!();
    }

    // Stop the whole process group if it's still running — not just the shell.
    kill_group(pgid);
    let _ = child.wait();

    if samples.len() < 3 {
        return Err(format!(
            "collected only {} sample(s) — the command exited too fast or --interval is too large",
            samples.len()
        ));
    }

    let cfg = SoakConfig {
        iterations: tick,
        sample_every: 1, // one tick per sample: slope is already per-sample
        warmup_frac: opts.warmup,
        max_bytes: opts.max_bytes,
        slope_bytes_per_sample: opts.slope_budget,
        min_r2_for_growth: 0.0, // r² floor off by default (see the library docs)
        min_movement_bytes: opts.min_movement, // movement enforcement (off unless set)
    };
    let report = report_from_samples(&cfg, samples, "rss");

    // The tool only certifies a plateau when IT ran the full soak. If the target
    // ended the run itself — crashed, or exited before the budget — the result is
    // inconclusive (exit 2): a truncated slice of its lifetime can't prove memory
    // levels off. (`child_exited` is set only when the target stopped on its own;
    // when we reach the budget we kill it, leaving this `None`.)
    if opts.json {
        print_json(&report, child_exited);
    } else {
        println!("navian-memcheck: {}", report.summary());
        match child_exited {
            Some(0) => {
                println!(
                    "  INCONCLUSIVE — target exited before the soak budget; the run was truncated"
                );
            }
            Some(code) => {
                println!("  FAIL — target exited with code {code} before the soak budget elapsed");
            }
            None => {}
        }
    }

    // 2 = target ended early / crashed, or inconclusive; 1 = memory failed; 0 = pass.
    let code = if child_exited.is_some()
        || matches!(report.verdict, Verdict::InsufficientSamples { .. })
    {
        2
    } else {
        i32::from(!report.passed())
    };
    Ok(code)
}

/// Sum the resident memory (bytes) of an entire process group via `ps` (portable
/// across macOS and Linux; `ps` reports RSS in kilobytes). This captures the real
/// workload even when the shell forks it as a child. Returns `None` if the group
/// is gone or `ps` is unavailable, so the caller can tell that from a real reading.
fn read_group_rss_bytes(pgid: u32) -> Option<u64> {
    let out = Command::new("ps")
        .args(["-A", "-o", "pgid=,rss="])
        .output()
        .ok()?;
    if !out.status.success() {
        return None;
    }
    let text = String::from_utf8_lossy(&out.stdout);
    let mut total = 0u64;
    let mut found = false;
    for line in text.lines() {
        let mut it = line.split_whitespace();
        if let (Some(g), Some(r)) = (it.next(), it.next()) {
            if let (Ok(g), Ok(kb)) = (g.parse::<u32>(), r.parse::<u64>()) {
                if g == pgid {
                    total += kb.saturating_mul(1024);
                    found = true;
                }
            }
        }
    }
    found.then_some(total)
}

/// Kill an entire process group (the target and every descendant). Uses the `kill`
/// utility with a negative pid so the crate stays free of `unsafe` libc calls.
#[cfg(unix)]
fn kill_group(pgid: u32) {
    let _ = Command::new("kill")
        .arg("-KILL")
        .arg(format!("-{pgid}"))
        .output();
}
#[cfg(not(unix))]
fn kill_group(_pgid: u32) {}

/// RAII guard: kills the soaked process group when it drops, so no exit path from
/// the sampling loop (normal, error, or panic) can leave the workload running.
struct GroupGuard(u32);
impl Drop for GroupGuard {
    fn drop(&mut self) {
        kill_group(self.0);
    }
}

fn print_json(report: &navian_memcheck::SoakReport, child_exited: Option<i32>) {
    let (verdict, detail) = match report.verdict {
        Verdict::Pass => ("pass", String::from("null")),
        Verdict::ExceededCap { peak, cap } => (
            "exceeded_cap",
            format!("{{\"peak_bytes\":{peak},\"cap_bytes\":{cap}}}"),
        ),
        Verdict::StillGrowing { slope, budget } => (
            "still_growing",
            format!(
                "{{\"slope_bytes_per_sample\":{slope:.1},\"budget_bytes_per_sample\":{budget:.1}}}"
            ),
        ),
        Verdict::InsufficientSamples { reason } => (
            "insufficient_samples",
            format!("{{\"reason\":\"{reason}\"}}"),
        ),
    };
    let exited = child_exited.map_or_else(|| "null".into(), |c| c.to_string());
    // JSON array of the advisory trust-warnings (each string escaped).
    let warnings = {
        let items: Vec<String> = report
            .trust_warnings
            .iter()
            .map(|w| format!("\"{}\"", json_escape(w)))
            .collect();
        format!("[{}]", items.join(","))
    };
    // `ok` is the overall gate (matches exit 0): memory passed AND the target ran
    // the full soak. `passed` is just the memory verdict, kept for detail.
    let ok = child_exited.is_none() && report.passed();
    println!(
        "{{\"ok\":{ok},\"passed\":{},\"verdict\":\"{verdict}\",\"detail\":{detail},\
\"sampler\":\"{}\",\"samples\":{},\"baseline_bytes\":{},\"peak_bytes\":{},\"moved_bytes\":{},\
\"back_half_slope_bytes_per_sample\":{:.1},\"back_half_r2\":{:.4},\"trust_warnings\":{warnings},\"child_exit_code\":{exited}}}",
        report.passed(),
        report.sampler_kind,
        report.samples.len(),
        report.baseline,
        report.peak,
        report.moved_bytes,
        report.back_half_slope,
        report.back_half_r2,
    );
}

/// Minimal JSON string escaping for the (tool-authored) warning strings: the two
/// characters that would break a JSON string literal, plus control chars.
fn json_escape(s: &str) -> String {
    use std::fmt::Write as _;
    let mut out = String::with_capacity(s.len());
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => {
                let _ = write!(out, "\\u{:04x}", c as u32);
            }
            c => out.push(c),
        }
    }
    out
}

struct Opts {
    cmd: Option<String>,
    duration: Option<Duration>,
    max_samples: Option<u64>,
    interval: Duration,
    warmup: f64,
    slope_budget: f64,
    max_bytes: Option<u64>,
    min_movement: u64,
    json: bool,
}

impl Default for Opts {
    fn default() -> Self {
        Opts {
            cmd: None,
            duration: None,
            max_samples: None,
            interval: Duration::from_secs(5),
            warmup: 0.5,
            slope_budget: 1024.0 * 1024.0, // 1 MB/sample
            max_bytes: None,
            min_movement: 0, // movement enforcement off unless --min-movement is set
            json: false,
        }
    }
}

// ── small parsers (no external deps) ──

fn parse_u64(s: &str) -> Result<u64, String> {
    s.trim()
        .parse()
        .map_err(|_| format!("`{s}` is not an integer"))
}

fn parse_f64(s: &str) -> Result<f64, String> {
    s.trim()
        .parse()
        .map_err(|_| format!("`{s}` is not a number"))
}

/// Parse a duration like `90s`, `20m`, `1h`, or a bare number (seconds).
fn parse_duration(s: &str) -> Result<Duration, String> {
    let s = s.trim();
    let (num, mult) = match s.chars().last() {
        Some('s') => (&s[..s.len() - 1], 1),
        Some('m') => (&s[..s.len() - 1], 60),
        Some('h') => (&s[..s.len() - 1], 3600),
        Some(c) if c.is_ascii_digit() => (s, 1),
        _ => return Err(format!("bad duration `{s}` (use 90s, 20m, 1h)")),
    };
    let n: u64 = num
        .trim()
        .parse()
        .map_err(|_| format!("bad duration `{s}`"))?;
    let secs = n
        .checked_mul(mult)
        .ok_or_else(|| format!("duration too large: `{s}`"))?;
    Ok(Duration::from_secs(secs))
}

/// Parse a byte size like `4096`, `500kb`, `2mb`, `6gb` (case-insensitive).
fn parse_size(s: &str) -> Result<u64, String> {
    let s = s.trim().to_ascii_lowercase();
    let (num, mult): (&str, u64) = if let Some(p) = s.strip_suffix("gb") {
        (p, 1024 * 1024 * 1024)
    } else if let Some(p) = s.strip_suffix("mb") {
        (p, 1024 * 1024)
    } else if let Some(p) = s.strip_suffix("kb") {
        (p, 1024)
    } else if let Some(p) = s.strip_suffix('b') {
        (p, 1)
    } else {
        (s.as_str(), 1)
    };
    let n: f64 = num.trim().parse().map_err(|_| format!("bad size `{s}`"))?;
    if !n.is_finite() || n < 0.0 {
        return Err(format!("size must be a finite, non-negative number: `{s}`"));
    }
    // Check the PRODUCT, not just `n`: a large finite `n` (e.g. 1e300) times the unit
    // overflows to +inf, and `inf as u64` saturates to u64::MAX — which would silently
    // DISABLE a budget/cap instead of rejecting the input. Reject non-representable sizes.
    let bytes = n * mult as f64;
    // `>=`, not `>`: `u64::MAX as f64` rounds UP to 2^64, so exactly 2^64 would slip
    // past `>` and then `as u64` saturates to u64::MAX — silently disabling the
    // budget/cap. Reject anything that isn't strictly below that boundary.
    if !bytes.is_finite() || bytes >= u64::MAX as f64 {
        return Err(format!("size out of range: `{s}`"));
    }
    Ok(bytes as u64)
}

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

    #[test]
    fn durations() {
        assert_eq!(parse_duration("90s").unwrap(), Duration::from_secs(90));
        assert_eq!(parse_duration("20m").unwrap(), Duration::from_secs(1200));
        assert_eq!(parse_duration("1h").unwrap(), Duration::from_secs(3600));
        assert_eq!(parse_duration("45").unwrap(), Duration::from_secs(45));
        assert!(parse_duration("abc").is_err());
    }

    #[test]
    fn sizes() {
        assert_eq!(parse_size("4096").unwrap(), 4096);
        assert_eq!(parse_size("2mb").unwrap(), 2 * 1024 * 1024);
        assert_eq!(parse_size("6GB").unwrap(), 6 * 1024 * 1024 * 1024);
        assert_eq!(parse_size("500kb").unwrap(), 500 * 1024);
        assert!(parse_size("-1mb").is_err());
        // Overflow: a huge finite size must be REJECTED, not saturate to u64::MAX
        // (which would silently disable a budget/cap).
        assert!(parse_size("1e300gb").is_err());
        assert!(parse_size("inf").is_err());
        // 2^64 exactly must be rejected (u64::MAX as f64 rounds up to 2^64).
        assert!(parse_size("18446744073709551616b").is_err());
    }

    #[test]
    fn zero_interval_is_rejected() {
        assert!(run(&["soak".into(), "--cmd".into(), "true".into(), "--interval".into(), "0s".into()]).is_err());
    }
}