urus/examples/causal_bench.rs

//! Causal-profiling bench (RFC 007): a real urus webserver as the workload,
//! with a planted, known-answer bottleneck.
//!
//! Where smarm's `causal_pipeline` demo is synthetic, this exercises the
//! exact machinery the causal commits touched, for real:
//!
//!   - epoll parks (`wait_readable_timeout` / `wait_writable_timeout`)
//!     → park-gated resume credit
//!   - keep-alive / request / write deadlines → timer-heap virtual time
//!   - the experiment controller's fixed windows → wall-anchored timers
//!   - mixed CPU (parse/serialize) and IO (socket) → non-trivial ranking
//!
//! Shape: `LOAD_CONNS` plain OS threads hammer `GET /order/:id` over
//! blocking loopback TCP (keep-alive). OS threads on the client side are
//! deliberate: they can't absorb injected virtual delay, so the only
//! delayable code is the server's — the same trick the smarm-side timer
//! tests use. The handler calls a single store actor (crud's pattern:
//! one actor owns the data, serialization is structural) which burns
//! `STORE_US` of calibrated *work* — work-shaped, not timed, because a
//! timed wait absorbs injected delay and flattens every experiment.
//!
//! Known answer: the store is serialized and saturated, so it is the
//! throughput ceiling (1e6/STORE_US rps). A virtual speedup of `store`
//! by p% predicts throughput ×1/(1-p): +33% @25%, +100% @50%. The other
//! five sites (parse, router, pipeline, serialize, socket-write — the
//! urus lib instrumentation) are tens of µs and parallel across
//! connections: predicted impact ≈0. Per the smarm-side validation runs,
//! absolute magnitudes are trustworthy to roughly ±15% while the
//! *ranking* is solid — the verdict thresholds reflect that.
//!
//! Run (needs cores; the verdict is SKIPPED below 4):
//!   cargo run --release --example causal_bench --features smarm-causal
//!
//! Prints the summary, writes `profile.coz` next to the CWD, and exits
//! nonzero if the expected separation doesn't hold.

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};

use smarm::{channel, Sender};
use urus::{serve_with_shutdown, shutdown_handle, Config, Conn, Next, Pipeline, Router};

// ---------------------------------------------------------------------------
// Tunables
// ---------------------------------------------------------------------------

/// Planted bottleneck cost: calibrated work per store request, µs.
const STORE_US: u64 = 400;
/// Handler-side render work per request, µs. Runs under the `pipeline`
/// site (no inner guard), parallel across connection actors.
const RENDER_US: u64 = 50;
/// Client connections = OS load threads. Plenty to keep a 400µs store
/// saturated even at a 50% virtual speedup (see saturation note in main).
const LOAD_CONNS: usize = 16;

// ---------------------------------------------------------------------------
// Calibrated work (lifted from smarm's causal_pipeline demo)
// ---------------------------------------------------------------------------

/// LCG-mix `iters` times in dependent sequence (unvectorizable, un-elidable),
/// staying preemptible — and causal-sampleable/delayable — via `check!()`.
fn work_iters(iters: u64) {
    let mut acc = 0x2545_f491_4f6c_dd1du64;
    let mut i = 0u64;
    while i < iters {
        let chunk_end = (i + 256).min(iters);
        while i < chunk_end {
            acc = acc.wrapping_mul(6364136223846793005).wrapping_add(i);
            i += 1;
        }
        std::hint::black_box(acc);
        smarm::check!();
    }
}

/// Iterations per microsecond on this machine, so stage costs are
/// meaningful in time while staying work-shaped.
fn calibrate_iters_per_us() -> u64 {
    let n = 8_000_000u64;
    let t = Instant::now();
    work_iters(n);
    (n / (t.elapsed().as_micros().max(1) as u64)).max(1)
}

static PER_US: OnceLock<u64> = OnceLock::new();

fn work_us(us: u64) {
    work_iters(us * PER_US.get().copied().unwrap_or(1));
}

// ---------------------------------------------------------------------------
// Store actor — the planted bottleneck (crud's once-cell pattern)
// ---------------------------------------------------------------------------

enum Req {
    Get { id: u64, reply: Sender<u64> },
}

/// Set at teardown; the store polls it (a cross-thread `Sender::send`
/// can't wake a parked actor, so a bare `recv` would hang AllDone —
/// same limitation crud's store works around).
static STORE_STOP: AtomicBool = AtomicBool::new(false);

static STORE_TX: OnceLock<Sender<Req>> = OnceLock::new();

fn store_loop(rx: smarm::Receiver<Req>) {
    loop {
        let req = match rx.recv_timeout(Duration::from_millis(250)) {
            Ok(r) => r,
            Err(smarm::channel::RecvTimeoutError::Timeout) => {
                if STORE_STOP.load(Ordering::Relaxed) {
                    return;
                }
                continue;
            }
            Err(smarm::channel::RecvTimeoutError::Disconnected) => return,
        };
        let Req::Get { id, reply } = req;
        {
            // The known-answer site: serialized by construction (one
            // actor, one request at a time), saturated by the load.
            let _g = smarm::causal_site!("store");
            work_us(STORE_US);
        }
        let _ = reply.send(id);
    }
}

/// Spawned from the first handler that runs — connection actors are
/// inside the runtime, which `spawn` requires; the static then shares
/// the Sender with every later handler.
fn store() -> &'static Sender<Req> {
    STORE_TX.get_or_init(|| {
        let (tx, rx) = channel::<Req>();
        smarm::spawn(move || store_loop(rx));
        tx
    })
}

// ---------------------------------------------------------------------------
// Handler
// ---------------------------------------------------------------------------

fn order(conn: Conn, _next: Next) -> Conn {
    let id: u64 = conn.params.get("id").and_then(|s| s.parse().ok()).unwrap_or(0);
    let (tx, rx) = channel::<u64>();
    store().send(Req::Get { id, reply: tx }).ok();
    let got = rx.recv().expect("store dropped");
    // Render work: attributed to the enclosing `pipeline` site.
    work_us(RENDER_US);
    conn.put_status(200).put_body(got.to_string())
}

// ---------------------------------------------------------------------------
// Load generation — plain OS threads, blocking loopback TCP, keep-alive
// ---------------------------------------------------------------------------

/// Read one HTTP/1.1 response (headers + content-length body) into `buf`,
/// draining exactly what was consumed. Errors mean: reconnect.
fn read_response(s: &mut TcpStream, buf: &mut Vec<u8>) -> std::io::Result<()> {
    let mut tmp = [0u8; 4096];
    let header_end = loop {
        if let Some(i) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
            break i + 4;
        }
        let n = s.read(&mut tmp)?;
        if n == 0 {
            return Err(std::io::ErrorKind::UnexpectedEof.into());
        }
        buf.extend_from_slice(&tmp[..n]);
    };
    let head = String::from_utf8_lossy(&buf[..header_end]).to_ascii_lowercase();
    let len: usize = head
        .lines()
        .find_map(|l| l.strip_prefix("content-length:"))
        .and_then(|v| v.trim().parse().ok())
        .unwrap_or(0);
    while buf.len() < header_end + len {
        let n = s.read(&mut tmp)?;
        if n == 0 {
            return Err(std::io::ErrorKind::UnexpectedEof.into());
        }
        buf.extend_from_slice(&tmp[..n]);
    }
    buf.drain(..header_end + len);
    Ok(())
}

fn load_loop(port: u16, stop: Arc<AtomicBool>) {
    let mut n: u64 = 1;
    'outer: while !stop.load(Ordering::Relaxed) {
        let mut s = match TcpStream::connect(("127.0.0.1", port)) {
            Ok(s) => s,
            Err(_) => {
                std::thread::sleep(Duration::from_millis(10));
                continue;
            }
        };
        s.set_read_timeout(Some(Duration::from_secs(5))).ok();
        s.set_nodelay(true).ok();
        let mut buf = Vec::with_capacity(4096);
        while !stop.load(Ordering::Relaxed) {
            let req = format!("GET /order/{n} HTTP/1.1\r\nhost: bench\r\n\r\n");
            if s.write_all(req.as_bytes()).is_err() {
                continue 'outer;
            }
            if read_response(&mut s, &mut buf).is_err() {
                continue 'outer;
            }
            n += 1;
        }
    }
}

// ---------------------------------------------------------------------------
// main — boot, load, sweep, verdict
// ---------------------------------------------------------------------------

fn main() {
    let per_us = calibrate_iters_per_us();
    PER_US.set(per_us).expect("PER_US set twice");
    println!("calibration: {per_us} work iters/µs");
    println!(
        "planted: store {STORE_US}µs serialized (ceiling ≈{} rps), render {RENDER_US}µs, {LOAD_CONNS} conns",
        1_000_000 / STORE_US
    );
    println!("theory: store +33% @25%, +100% @50%; every other site ≈0");

    // Ephemeral port via bind-and-release (integration-test pattern; the
    // race window before the server rebinds is far below flake threshold).
    let port = {
        let l = TcpListener::bind("127.0.0.1:0").expect("bind probe");
        l.local_addr().expect("local_addr").port()
    };
    let addr: SocketAddr = format!("127.0.0.1:{port}").parse().expect("addr");

    let (handle, signal) = shutdown_handle();
    let server = std::thread::spawn(move || {
        let pipe = Pipeline::new().plug(Router::new().get("/order/:id", order));
        serve_with_shutdown(Config::new(addr), smarm::Config::default(), pipe, signal).expect("serve");
    });

    // Wait until it's accepting.
    let up = (0..100).any(|_| {
        std::thread::sleep(Duration::from_millis(50));
        TcpStream::connect(addr).is_ok()
    });
    assert!(up, "server didn't come up on {addr}");

    let stop = Arc::new(AtomicBool::new(false));
    let loaders: Vec<_> = (0..LOAD_CONNS)
        .map(|_| {
            let stop = stop.clone();
            std::thread::spawn(move || load_loop(port, stop))
        })
        .collect();

    // Warm up: queues to steady state, every site + the progress point
    // registered by real traffic before the sweep enumerates them.
    std::thread::sleep(Duration::from_millis(500));

    // The controller runs on this plain OS thread: its windows are wall
    // time by construction here; inside a runtime they'd be wall-anchored
    // via the efbc254 timer path (the thing the jobrunner sweep confirms).
    let results = smarm::causal::run_experiments(&smarm::causal::ExperimentPlan {
        speedups_pct: vec![0, 25, 50],
        experiment: Duration::from_millis(700),
        cooldown: Duration::from_millis(150),
    });

    stop.store(true, Ordering::Relaxed);
    for l in loaders {
        l.join().expect("loader panicked");
    }
    STORE_STOP.store(true, Ordering::Relaxed);
    handle.shutdown();
    server.join().expect("server panicked");

    print!("{}", smarm::causal::render_summary(&results));
    match std::fs::write("profile.coz", smarm::causal::render_coz(&results)) {
        Ok(()) => println!("\nwrote profile.coz"),
        Err(e) => eprintln!("\nfailed to write profile.coz: {e}"),
    }

    // Verdict — same discipline as the smarm demo: skip where the
    // separation physically can't exist.
    let cores = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1);
    if cores < 4 {
        println!("verdict: SKIPPED ({cores} cores; separation needs real parallelism)");
        return;
    }
    let mut failures: Vec<String> = Vec::new();
    let impact =
        |site: &str| smarm::causal::impact_pct(&results, site, 25, "responses");
    let mut expect = |site: &str, ok: &dyn Fn(f64) -> bool, want: &str| match impact(site) {
        Some(p) => {
            let verdict = if ok(p) { "ok" } else { "FAIL" };
            println!("verdict: {site} @25% -> {p:+.1}% (want {want}) {verdict}");
            if !ok(p) {
                failures.push(format!("{site}: {p:+.1}% (want {want})"));
            }
        }
        None => {
            println!("verdict: {site} @25% -> missing cell FAIL");
            failures.push(format!("{site}: missing cell"));
        }
    };
    expect("store", &|p| p > 15.0, "> +15%");
    for site in ["parse", "router", "pipeline", "serialize", "socket-write"] {
        expect(site, &|p| p < 10.0, "< +10%");
    }

    if failures.is_empty() {
        println!("verdict: PASS — planted bottleneck found, cold sites quiet");
    } else {
        println!("verdict: FAIL — {}", failures.join("; "));
        std::process::exit(1);
    }
}