supermachine 0.5.5

//! Full-axis benchmark across boot, cold, warm, RPS, concurrency.
//! Workload variants: trivial RPC, tiny exec, rustc compile, HTTP RPS.

use std::sync::Arc;
use std::time::{Duration, Instant};
use supermachine::{Image, VmConfig};

const RUSTC_SRC: &[u8] = b"fn main() { println!(\"hello\"); }";
const RUSTC_SH: &[&str] = &["sh", "-c", "rustc -O /tmp/main.rs -o /tmp/m && /tmp/m"];
const RUSTC_DIRECT: &[&str] = &["rustc", "-O", "/tmp/main.rs", "-o", "/tmp/m"];

fn pct(values: &mut Vec<u64>, p: f64) -> u64 {
    values.sort();
    let idx = ((values.len() - 1) as f64 * p / 100.0).round() as usize;
    values[idx.min(values.len() - 1)]
}
fn ms(us: u64) -> String {
    format!("{:.1}", us as f64 / 1000.0)
}

fn home() -> String {
    std::env::var("HOME").unwrap()
}
fn snap(name: &str) -> String {
    format!("{}/.local/supermachine-snapshots/{name}", home())
}
fn want(name: &str) -> bool {
    std::env::var("AXIS")
        .map(|v| v.split(',').any(|x| x == name))
        .unwrap_or(true)
}

// ---- Boot: time from build_pool to first acquire-ready ----
fn bench_boot(label: &str, snap_path: &str) {
    if !want("boot") {
        return;
    }
    println!("\n--- boot: {label} ---");
    let n = 5;
    let mut times = Vec::with_capacity(n);
    for i in 0..n {
        let image = match Image::from_snapshot(snap_path) {
            Ok(img) => img,
            Err(e) => {
                println!("  [{label}] iter {i}: from_snapshot failed: {e}");
                return;
            }
        };
        let t0 = Instant::now();
        let pool = match image
            .pool()
            .min(1)
            .max(1)
            .idle_timeout(Duration::MAX)
            .build()
        {
            Ok(p) => p,
            Err(e) => {
                println!("  [{label}] iter {i}: pool.build() failed: {e}");
                return;
            }
        };
        let _vm = match pool.acquire() {
            Ok(v) => v,
            Err(e) => {
                println!("  [{label}] iter {i}: pool.acquire() failed: {e}");
                return;
            }
        };
        let us = t0.elapsed().as_micros() as u64;
        times.push(us);
    }
    if times.is_empty() {
        return;
    }
    let mut t = times.clone();
    let med = pct(&mut t, 50.0);
    let p95 = pct(&mut times, 95.0);
    println!("  pool.build() + acquire: median {} ms  p95 {} ms", ms(med), ms(p95));
}

// ---- Cycle: per-iteration end-to-end with a workload ----
fn bench_cycle(
    label: &str,
    snap_path: &str,
    skip_restore: bool,
    workload: &str,
    n: usize,
) {
    if !want("cycle") {
        return;
    }
    let Ok(image) = Image::from_snapshot(snap_path) else {
        println!("\n--- cycle: {label} — snapshot missing ---");
        return;
    };
    let pool = image
        .pool()
        .min(2)
        .max(2)
        .idle_timeout(Duration::MAX)
        .restore_on_release(!skip_restore)
        .build()
        .unwrap();
    let mut times = Vec::with_capacity(n);
    for _ in 0..n {
        let t0 = Instant::now();
        let vm = pool.acquire().unwrap();
        match workload {
            "rpc" => {
                vm.write_file("/tmp/marker", b"x").unwrap();
            }
            "echo" => {
                let _ = vm
                    .exec_builder()
                    .argv(["true"].iter().copied())
                    .timeout(Duration::from_secs(10))
                    .output()
                    .unwrap();
            }
            "rustc" => {
                vm.write_file("/tmp/main.rs", RUSTC_SRC).unwrap();
                let _ = vm
                    .exec_builder()
                    .argv(RUSTC_SH.iter().copied())
                    .timeout(Duration::from_secs(60))
                    .output()
                    .unwrap();
            }
            "rustc_staged" => {
                // stage_file + chain — single RPC, no sh wrapper.
                let _ = vm
                    .exec_builder()
                    .stage_file("/tmp/main.rs", RUSTC_SRC.to_vec())
                    .argv(RUSTC_DIRECT.iter().copied())
                    .chain(["/tmp/m"].iter().copied())
                    .timeout(Duration::from_secs(60))
                    .output()
                    .unwrap();
            }
            _ => panic!("unknown workload: {workload}"),
        }
        drop(vm);
        times.push(t0.elapsed().as_micros() as u64);
    }
    let mut t = times.clone();
    let cold = times[0];
    let med = pct(&mut t, 50.0);
    let p95 = pct(&mut t, 95.0);
    let mut warm: Vec<u64> = times.iter().skip(1).copied().collect();
    let warm_med = if warm.is_empty() { 0 } else { pct(&mut warm, 50.0) };
    println!(
        "  {label} ({workload}, skip={skip_restore}): cold {} | median {} | warm-median {} | p95 {} ms",
        ms(cold),
        ms(med),
        ms(warm_med),
        ms(p95)
    );
}

// ---- Throughput: how many cycles/sec can we sustain ----
fn bench_throughput(
    label: &str,
    snap_path: &str,
    skip_restore: bool,
    workers: usize,
    duration: Duration,
) {
    if !want("rps") {
        return;
    }
    let Ok(image) = Image::from_snapshot(snap_path) else {
        println!("\n--- rps: {label} — snapshot missing ---");
        return;
    };
    let image = Arc::new(image);
    let pool = Arc::new(
        image
            .pool()
            .min(workers)
            .max(workers)
            .idle_timeout(Duration::MAX)
            .restore_on_release(!skip_restore)
            .build()
            .unwrap(),
    );
    // Warmup so timed window measures steady-state.
    for _ in 0..workers {
        let vm = pool.acquire().unwrap();
        vm.write_file("/tmp/main.rs", RUSTC_SRC).unwrap();
        let _ = vm
            .exec_builder()
            .argv(RUSTC_SH.iter().copied())
            .timeout(Duration::from_secs(60))
            .output()
            .unwrap();
        drop(vm);
    }
    let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
    let count = Arc::new(std::sync::atomic::AtomicU64::new(0));
    let t0 = Instant::now();
    let handles: Vec<_> = (0..workers)
        .map(|_| {
            let pool = Arc::clone(&pool);
            let stop = Arc::clone(&stop);
            let count = Arc::clone(&count);
            std::thread::spawn(move || {
                while !stop.load(std::sync::atomic::Ordering::Relaxed) {
                    let vm = match pool.acquire() {
                        Ok(v) => v,
                        Err(_) => continue,
                    };
                    vm.write_file("/tmp/main.rs", RUSTC_SRC).unwrap();
                    let _ = vm
                        .exec_builder()
                        .argv(RUSTC_SH.iter().copied())
                        .timeout(Duration::from_secs(60))
                        .output()
                        .unwrap();
                    drop(vm);
                    count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                }
            })
        })
        .collect();
    std::thread::sleep(duration);
    stop.store(true, std::sync::atomic::Ordering::Relaxed);
    for h in handles {
        let _ = h.join();
    }
    let elapsed = t0.elapsed().as_secs_f64();
    let total = count.load(std::sync::atomic::Ordering::Relaxed);
    println!(
        "  {label} (workers={workers}, skip={skip_restore}): {total} cycles in {:.1}s = {:.1} cycles/s",
        elapsed,
        total as f64 / elapsed
    );
}

// ---- Snapshot: capture+save latency ----
fn bench_snapshot(label: &str, snap_path: &str) {
    if !want("snapshot") {
        return;
    }
    let Ok(image) = Image::from_snapshot(snap_path) else {
        println!("\n--- snapshot: {label} — snapshot missing ---");
        return;
    };
    let dest = format!("/tmp/sm-axis-snap-{}", std::process::id());
    let _ = std::fs::remove_dir_all(&dest);
    let n = 3;
    let mut cap_us = Vec::with_capacity(n);
    let mut save_us = Vec::with_capacity(n);
    for _ in 0..n {
        let pooled = image.acquire_with(&VmConfig::new()).unwrap();
        pooled.write_file("/tmp/marker", b"x").unwrap();
        let t0 = Instant::now();
        let _ = pooled.snapshot(&dest).unwrap();
        cap_us.push(t0.elapsed().as_micros() as u64);
        // capture vs save isn't separately exposed — record total
        save_us.push(0);
        let _ = std::fs::remove_dir_all(&dest);
        drop(pooled);
    }
    cap_us.sort();
    println!(
        "  snapshot {label}: capture+save median {} ms  (n={n})",
        ms(cap_us[n / 2])
    );
}

// ---- Memory: rough RSS per worker ----
fn bench_memory(label: &str, snap_path: &str, workers: usize) {
    if !want("mem") {
        return;
    }
    let Ok(image) = Image::from_snapshot(snap_path) else {
        return;
    };
    let pool = image
        .pool()
        .min(workers)
        .max(workers)
        .idle_timeout(Duration::MAX)
        .build()
        .unwrap();
    // Hold workers idle, wait for stabilization
    let _vms: Vec<_> = (0..workers).map(|_| pool.acquire().unwrap()).collect();
    std::thread::sleep(Duration::from_millis(500));
    // ps for our supermachine-worker children
    let out = std::process::Command::new("sh")
        .arg("-c")
        .arg(format!(
            "ps -o rss= -p $(pgrep -P {} supermachine-worker | tr '\\n' ',' | sed 's/,$//') 2>/dev/null | awk '{{s+=$1}} END {{print s}}'",
            std::process::id()
        ))
        .output()
        .ok();
    let rss_kb: u64 = out
        .as_ref()
        .map(|o| String::from_utf8_lossy(&o.stdout).trim().parse().unwrap_or(0))
        .unwrap_or(0);
    println!(
        "  mem {label} (workers={workers}): total RSS {} MiB ≈ {} MiB/worker",
        rss_kb / 1024,
        rss_kb / 1024 / workers as u64
    );
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("===== axis benchmark =====");
    println!("AXIS env (csv) selects subsets; default = all");
    println!("Snapshots:");
    for s in &[
        "rust_1_slim",
        "rust_1_slim_2vcpu",
        "rust_1_slim_4vcpu",
        "rust_1_slim__warm__rustc_v1",
        "nginx_1v",
    ] {
        let p = snap(s);
        let ok = std::path::Path::new(&p).is_dir();
        println!("  {} {s}", if ok { "✓" } else { "✗" });
    }

    println!("\n========== BOOT ==========");
    bench_boot("rust_1_slim", &snap("rust_1_slim"));
    bench_boot("rust_1_slim_4vcpu", &snap("rust_1_slim_4vcpu"));
    bench_boot("warm-baked", &snap("rust_1_slim__warm__rustc_v1"));

    println!("\n========== CYCLE: trivial RPC (write_file) ==========");
    bench_cycle("rust_1_slim", &snap("rust_1_slim"), false, "rpc", 20);
    bench_cycle("rust_1_slim", &snap("rust_1_slim"), true, "rpc", 20);

    println!("\n========== CYCLE: tiny exec (`true`) ==========");
    bench_cycle("rust_1_slim", &snap("rust_1_slim"), false, "echo", 10);
    bench_cycle("rust_1_slim", &snap("rust_1_slim"), true, "echo", 10);

    println!("\n========== CYCLE: rustc hello (sh -c wrapper, separate write_file) ==========");
    bench_cycle("cold base", &snap("rust_1_slim"), false, "rustc", 10);
    bench_cycle("cold base", &snap("rust_1_slim"), true, "rustc", 10);
    bench_cycle("warm-baked", &snap("rust_1_slim__warm__rustc_v1"), false, "rustc", 10);
    bench_cycle("warm-baked", &snap("rust_1_slim__warm__rustc_v1"), true, "rustc", 10);

    println!("\n========== CYCLE: rustc hello (stage_file + chain, single RPC) ==========");
    bench_cycle("cold base", &snap("rust_1_slim"), false, "rustc_staged", 10);
    bench_cycle("cold base", &snap("rust_1_slim"), true, "rustc_staged", 10);
    bench_cycle("warm-baked", &snap("rust_1_slim__warm__rustc_v1"), false, "rustc_staged", 10);
    bench_cycle("warm-baked", &snap("rust_1_slim__warm__rustc_v1"), true, "rustc_staged", 10);

    println!("\n========== THROUGHPUT (rustc) ==========");
    bench_throughput("warm-baked, 1 worker", &snap("rust_1_slim__warm__rustc_v1"), true, 1, Duration::from_secs(5));
    bench_throughput("warm-baked, 2 workers", &snap("rust_1_slim__warm__rustc_v1"), true, 2, Duration::from_secs(5));
    bench_throughput("warm-baked, 4 workers", &snap("rust_1_slim__warm__rustc_v1"), true, 4, Duration::from_secs(5));
    bench_throughput("warm-baked, 8 workers", &snap("rust_1_slim__warm__rustc_v1"), true, 8, Duration::from_secs(5));

    println!("\n========== SNAPSHOT capture latency ==========");
    bench_snapshot("rust_1_slim", &snap("rust_1_slim"));
    bench_snapshot("warm-baked", &snap("rust_1_slim__warm__rustc_v1"));

    println!("\n========== MEMORY (rough RSS per idle worker) ==========");
    bench_memory("rust_1_slim", &snap("rust_1_slim"), 1);
    bench_memory("rust_1_slim", &snap("rust_1_slim"), 4);
    bench_memory("warm-baked", &snap("rust_1_slim__warm__rustc_v1"), 1);
    bench_memory("warm-baked", &snap("rust_1_slim__warm__rustc_v1"), 4);

    println!("\nDONE");
    Ok(())
}