open-standards/joule-code/joule-code-rs/crates/joule-moq-relay/examples/relay_bench.rs

//! Run the load harness (`joule_moq_relay::bench`; `draft-evens-moq-bench-00`) and write the
//! run's signed measurement record:
//!
//! ```text
//! relay_bench --profile profile.json --host-label build-1 --out record.json --key bench.key
//!     [--relays 1|2] [--relay-threads T] [--relay-cpus 0-7] [--relay-sockets S] [--per-endpoint E]
//!     [--wait-idle-s S] [--guard earlier.json --trusted <public key, hex> [tolerances]]
//!     [--url ws://127.0.0.1:4480/moq]
//! ```
//!
//! With no `--url` it starts each relay as a process of its own (this program, `relay_bench
//! relay`), on loopback over native QUIC, so that what a relay spends is read apart from its
//! clients: `--relays 1`, one relay that the publishers and subscribers all reach; `--relays
//! 2`, a cascade, the publishers on the origin and the subscribers on an edge whose upstream
//! the origin is (`[[upstream]]` over QUIC, its certificate pinned). Each relay has
//! `--relay-threads` workers (8), on `--relay-cpus` where given (Linux, by `taskset`), and its
//! listener's `limits.quic_sockets` set to `--relay-sockets` where given (else the relay's
//! default, one per core up to 16); the clients share a UDP socket by `--per-endpoint` sessions
//! (8), each socket of 4 MiB. With
//! `--url` it drives the relay there over the WebSocket binding, and reads no relay's cost.
//!
//! **The idle rule** (`bench::IDLE_LOAD_PER_CPU_PERCENT`). Before the run the host's
//! one-minute load is read; above a quarter of its processors the run waits, rereading every
//! 15 s, up to `--wait-idle-s` (0 by default), and if the host never meets the rule it writes
//! nothing and exits 3, printing the loads it read: the run is not shown. A run whose record
//! states a load above the rule is not written either. `--try` runs the tool whatever the
//! load, to try it out: it prints the figures and writes and signs nothing.
//!
//! `--key` is a file holding an Ed25519 seed as 64 hexadecimal digits; the record is signed
//! with it. `--guard` holds the run to an earlier record, read only under `--trusted` (or the
//! signing key's public key), of the same host label and profile, within the default
//! tolerances or those given (`--cpu-percent`, `--delivered-ppm`, `--lost-objects`,
//! `--variance-us`): it exits 1
//! and names each figure outside them, and 2 where the two are not comparable; a run the
//! record's host held is held to that too. Unguarded, a run the host did not hold
//! (`Record::held`: a stream object lost, a 99th percentile of lateness past
//! `bench::LATE_P99_BOUND_US`) is written all the same and exits 4. A record states what this
//! host measured on its date and nothing more.

use ed25519_dalek::SigningKey;
use joule_moq_relay::bench::{self, Host, Meter, Profile, Record, SignedRecord, Tolerances};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::config::{UpstreamConfig, UpstreamTransport};
use joule_moq_relay::transport::mux::MuxConfig;
use joule_moq_relay::transport::{quic, websocket, SharedConnection};
use joule_moq_relay::{RelayConfig, RelayServer};
use std::io::{BufRead, BufReader, Write};
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
use std::sync::{Arc, Mutex};
use std::time::Duration;

/// Bytes each client socket asks of the kernel: the sessions on it are many viewers, and a
/// small buffer would drop what no viewer's own would, measuring this program, not the relay.
const CLIENT_SOCKET_BUFFER: usize = 4 << 20;

fn arg(name: &str) -> Option<String> {
    let mut a = std::env::args().skip(1);
    while let Some(k) = a.next() {
        if k == name {
            return a.next();
        }
    }
    None
}

fn num<T: std::str::FromStr>(name: &str, default: T) -> T {
    arg(name).map_or(default, |v| {
        v.parse().unwrap_or_else(|_| {
            eprintln!("relay_bench: {name} {v} does not read");
            std::process::exit(64)
        })
    })
}

fn need(name: &str) -> String {
    arg(name).unwrap_or_else(|| {
        eprintln!("relay_bench: {name} is required");
        std::process::exit(64)
    })
}

// ── A relay of its own ──────────────────────────────────────────────────────

/// `relay_bench relay`: one relay on loopback over QUIC, as the driver asks. It prints its
/// address, its certificate (DER, hex), its leaf's SHA-256 and its relay id, answers each
/// line `cpu` on its standard input with its processor time so far in microseconds, and
/// exits when that input closes, so a killed driver leaves nothing behind. With
/// `--upstream <port> --upstream-sha256 <hex> --upstream-id <hex> --prefix <field>` it is an
/// edge whose upstream is the relay on that port, routed the namespaces under the field.
fn relay_main() {
    let threads: usize = num("--threads", 8);
    let rt = tokio::runtime::Builder::new_multi_thread().worker_threads(threads).enable_all().build().expect("a runtime");
    rt.block_on(async {
        let mut c = RelayConfig::loopback();
        c.websocket_listen = None;
        // The bench's relay holds every session a run opens; its other limits are the defaults.
        c.limits.max_sessions = 1 << 20;
        c.limits.quic_sockets = num("--sockets", c.limits.quic_sockets);
        if let Some(port) = arg("--upstream") {
            c.upstream = vec![UpstreamConfig {
                name: "origin".into(),
                url: format!("moqt://127.0.0.1:{port}{}", c.path),
                transport: UpstreamTransport::Quic,
                prefixes: vec![vec![need("--prefix")]],
                ca_path: None,
                cert_sha256: Some(need("--upstream-sha256")),
                native_roots: false,
                token_sha256: None,
                token_path: None,
                relay_key: Some(need("--upstream-id")),
            }];
        }
        let server = RelayServer::start(c).await.expect("the relay starts");
        let cert = hex::encode(server.identity.chain[0].as_ref());
        println!("{} {cert} {} {}", server.quic_addr.expect("a UDP endpoint"), hex::encode(server.identity.leaf_sha256), hex::encode(server.relay().relay_id()));
        std::io::stdout().flush().expect("stdout");
        tokio::task::spawn_blocking(|| {
            for line in std::io::stdin().lock().lines() {
                let Ok(line) = line else { break };
                if line.trim() == "cpu" {
                    println!("{}", bench::own_cpu_us());
                    let _ = std::io::stdout().flush();
                }
            }
        })
        .await
        .expect("the wait");
        std::process::exit(0);
    });
}

/// A relay process the driver started, and how to read its processor time.
struct Relay {
    role: String,
    threads: u64,
    child: Child,
    input: Mutex<ChildStdin>,
    output: Mutex<BufReader<ChildStdout>>,
    addr: std::net::SocketAddr,
    cert: Vec<u8>,
    leaf_sha256: String,
    id: String,
}

fn spawn_relay(role: &str, threads: u64, upstream: Option<(&Relay, &str)>) -> Relay {
    let exe = std::env::current_exe().expect("this binary");
    let mut cmd = match arg("--relay-cpus") {
        Some(cpus) => {
            let mut c = Command::new("taskset");
            c.args(["-c", &cpus]).arg(&exe);
            c
        }
        None => Command::new(&exe),
    };
    cmd.args(["relay", "--threads", &threads.to_string()]);
    if let Some(n) = arg("--relay-sockets") {
        cmd.args(["--sockets", &n]);
    }
    if let Some((up, prefix)) = upstream {
        cmd.args(["--upstream", &up.addr.port().to_string(), "--upstream-sha256", &up.leaf_sha256, "--upstream-id", &up.id, "--prefix", prefix]);
    }
    let mut child = cmd.stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::null()).spawn().expect("the relay process starts");
    let mut output = BufReader::new(child.stdout.take().expect("its stdout"));
    let mut line = String::new();
    output.read_line(&mut line).expect("its address");
    let f: Vec<&str> = line.split_whitespace().collect();
    let [addr, cert, leaf, id] = f[..] else { panic!("the relay said {line:?}") };
    Relay {
        role: role.into(),
        threads,
        input: Mutex::new(child.stdin.take().expect("its stdin")),
        child,
        output: Mutex::new(output),
        addr: addr.parse().expect("an address"),
        cert: hex::decode(cert).expect("hex"),
        leaf_sha256: leaf.into(),
        id: id.into(),
    }
}

impl Relay {
    fn cpu_us(&self) -> Result<u64, String> {
        writeln!(self.input.lock().unwrap(), "cpu").map_err(|e| format!("{}: {e}", self.role))?;
        let mut line = String::new();
        self.output.lock().unwrap().read_line(&mut line).map_err(|e| format!("{}: {e}", self.role))?;
        line.trim().parse().map_err(|_| format!("{} said {line:?}", self.role))
    }
}

struct Relays(Vec<Relay>);

impl Meter for Relays {
    fn relays(&self) -> Vec<(String, u64)> {
        self.0.iter().map(|r| (r.role.clone(), r.threads)).collect()
    }

    fn cpu_us(&self) -> Result<Vec<u64>, String> {
        self.0.iter().map(Relay::cpu_us).collect()
    }
}

impl Drop for Relays {
    fn drop(&mut self) {
        for r in &mut self.0 {
            let _ = r.child.kill();
            let _ = r.child.wait();
        }
    }
}

/// Client sessions over QUIC to one relay, `per` of them to a UDP socket.
struct Dialer {
    addr: std::net::SocketAddr,
    cert: Vec<u8>,
    per: usize,
    endpoints: Mutex<(usize, Vec<quinn::Endpoint>)>,
}

impl Dialer {
    async fn connect(&self) -> Result<Client, String> {
        let ep = {
            let mut e = self.endpoints.lock().unwrap();
            let n = e.0;
            e.0 += 1;
            if n / self.per >= e.1.len() {
                let bind: std::net::SocketAddr = "127.0.0.1:0".parse().expect("an address");
                let ep = quic::client_endpoint(bind, vec![rustls::pki_types::CertificateDer::from(self.cert.clone())], vec![joule_moq::draft22::ALPN.to_vec()], None).map_err(|e| e.to_string())?;
                ep.rebind(quic::bind_udp(bind, CLIENT_SOCKET_BUFFER).map_err(|e| e.to_string())?).map_err(|e| e.to_string())?;
                e.1.push(ep);
            }
            e.1[n / self.per].clone()
        };
        let conn = tokio::time::timeout(Duration::from_secs(30), quic::connect(&ep, self.addr, "localhost")).await.map_err(|_| "a QUIC handshake timed out".to_string())?.map_err(|e| e.to_string())?;
        Client::connect(Arc::new(conn), ClientOptions::default()).await
    }
}

// ── The driver ──────────────────────────────────────────────────────────────

/// Wait up to `--wait-idle-s` for the host to meet the idle rule. `Err` with the loads read
/// where it never did.
fn wait_idle(cpus: u64) -> Result<u64, String> {
    let deadline = std::time::Instant::now() + Duration::from_secs(num("--wait-idle-s", 0));
    let rule = cpus * bench::IDLE_LOAD_PER_CPU_PERCENT;
    let mut read = Vec::new();
    loop {
        let Some(l) = bench::load_x100() else { return Err("the host states no load".into()) };
        if l <= rule {
            return Ok(l);
        }
        read.push(format!("{}.{:02}", l / 100, l % 100));
        if std::time::Instant::now() >= deadline {
            return Err(format!("one-minute loads {} on {cpus} processors, the rule at most {}.{:02}", read.join(" "), rule / 100, rule % 100));
        }
        std::thread::sleep(Duration::from_secs(15));
    }
}

fn summary(r: &Record) -> String {
    let mut s = format!("{} {} on {} ({} {} {} processors)", r.date, r.profile.name, r.host.label, r.host.os, r.host.arch, r.host.cpus);
    if let Some(l) = &r.load {
        s += &format!(", load {}.{:02} before and {}.{:02} after", l.before_x100 / 100, l.before_x100 % 100, l.after_x100 / 100, l.after_x100 % 100);
    }
    if let Some(c) = &r.cost {
        let relays: Vec<String> = c.relays.iter().map(|x| format!("{} {:.2} cores", x.role, x.cpu_us as f64 / (c.wall_ms.max(1) * 1000) as f64)).collect();
        s += &format!("; {} deliveries in {:.1} s, {:.2} µs of relay CPU a delivery ({}), clients {:.2} cores", c.deliveries, c.wall_ms as f64 / 1000.0, c.cpu_ns_per_delivery as f64 / 1000.0, relays.join(", "), c.clients_cpu_us as f64 / (c.wall_ms.max(1) * 1000) as f64);
    }
    for t in &r.tracks {
        s += &format!("\n  {}: {} of {} delivered ({} ppm), lateness p50 {:.2} ms p99 {:.2} ms max {:.2} ms, {} subscribers failed", t.track, t.deliveries, t.owed, t.delivered_ppm(), t.late_p50_us as f64 / 1000.0, t.late_p99_us as f64 / 1000.0, t.late_max_us as f64 / 1000.0, t.subscribers_failed);
    }
    s
}

/// The packets the kernel dropped on loopback UDP sockets over the run, where it says (Linux).
fn udp_drops() -> Option<u64> {
    let s = std::fs::read_to_string("/proc/net/udp").ok()?;
    Some(s.lines().skip(1).filter(|l| l.split_whitespace().nth(1).is_some_and(|a| a.starts_with("0100007F:"))).filter_map(|l| l.split_whitespace().last()?.parse::<u64>().ok()).sum())
}

fn main() {
    if std::env::args().nth(1).as_deref() == Some("relay") {
        return relay_main();
    }
    let status = tokio::runtime::Builder::new_multi_thread().enable_all().build().expect("a runtime").block_on(driver());
    std::process::exit(status)
}

async fn driver() -> i32 {
    let profile: Profile = serde_json::from_str(&std::fs::read_to_string(need("--profile")).expect("the profile")).expect("the profile reads");
    let key = arg("--key").map(|k| {
        let seed: [u8; 32] = hex::decode(std::fs::read_to_string(k).expect("the key").trim()).ok().and_then(|b| b.try_into().ok()).expect("a seed of 64 hexadecimal digits");
        SigningKey::from_bytes(&seed)
    });
    let host = Host::here(&need("--host-label"));
    // `--try`: the tool itself tried out on any host, whatever its load; nothing is written or
    // signed, and a run on a host above the rule still exits 3.
    let trying = std::env::args().any(|a| a == "--try");
    if let Err(e) = wait_idle(host.cpus).or_else(|e| if trying { Ok(0) } else { Err(e) }) {
        println!("{}", serde_json::json!({ "not_shown": profile.name, "host": host, "unix_s": joule_moq_relay::relay::now_unix_ms() / 1000, "why": e }));
        eprintln!("relay_bench: not shown: {e}");
        std::process::exit(3)
    }
    let drops = udp_drops();
    let record = match arg("--url") {
        Some(url) => {
            let connect = || async {
                let conn: SharedConnection = Arc::new(websocket::connect(&url, MuxConfig::default()).await.map_err(|e| e.to_string())?);
                Client::connect(conn, ClientOptions::default()).await
            };
            bench::run(&profile, host, "websocket", connect).await
        }
        None => {
            let threads: u64 = num("--relay-threads", 8);
            let per: usize = num("--per-endpoint", 8);
            let origin = spawn_relay(if num::<u8>("--relays", 1) == 2 { "origin" } else { "relay" }, threads, None);
            let mut relays = vec![];
            let edge = match num::<u8>("--relays", 1) {
                1 => None,
                2 => {
                    let prefix = profile.tracks[0].namespace.first().cloned().unwrap_or_default();
                    if profile.tracks.iter().any(|t| t.namespace.first() != Some(&prefix)) || prefix.is_empty() || profile.meeting.is_some() {
                        eprintln!("relay_bench: a cascade runs one publisher's tracks under one first field");
                        std::process::exit(64)
                    }
                    Some(spawn_relay("edge", threads, Some((&origin, &prefix))))
                }
                n => {
                    eprintln!("relay_bench: --relays is 1 or 2, not {n}");
                    std::process::exit(64)
                }
            };
            let dial = |r: &Relay| Dialer { addr: r.addr, cert: r.cert.clone(), per, endpoints: Mutex::new((0, Vec::new())) };
            let (to_origin, to_edge) = (dial(&origin), edge.as_ref().map(dial));
            relays.push(origin);
            relays.extend(edge);
            let relays = Relays(relays);
            let subscriber = to_edge.as_ref().unwrap_or(&to_origin);
            let out = bench::run_with(&profile, host, "quic/moqt-22", || to_origin.connect(), || subscriber.connect(), Some(&relays)).await;
            drop(relays);
            out
        }
    };
    let record = record.unwrap_or_else(|e| {
        eprintln!("relay_bench: {e}");
        std::process::exit(70)
    });
    println!("{}", summary(&record));
    if let (Some(a), Some(b)) = (drops, udp_drops()) {
        println!("  the kernel dropped {} loopback UDP datagrams over the run", b.saturating_sub(a));
    }
    if !record.idle() {
        eprintln!("relay_bench: not shown: the record states a load above the idle rule, or none");
        return 3;
    }
    if trying {
        eprintln!("relay_bench: tried; nothing written");
        return 0;
    }
    if let (Some(out), Some(key)) = (arg("--out"), &key) {
        let signed = SignedRecord::sign(&record, key);
        std::fs::write(&out, serde_json::to_string_pretty(&signed).expect("JSON") + "\n").expect("the record written");
        println!("  signed by {} into {out}", signed.key);
    }
    let mut status = 0;
    if let Some(earlier) = arg("--guard") {
        let trusted: [u8; 32] = match (arg("--trusted"), &key) {
            (Some(h), _) => hex::decode(h).ok().and_then(|b| b.try_into().ok()).expect("--trusted is 64 hexadecimal digits"),
            (None, Some(k)) => k.verifying_key().to_bytes(),
            (None, None) => {
                eprintln!("relay_bench: --guard reads the earlier record under --trusted or --key");
                std::process::exit(64)
            }
        };
        let earlier: SignedRecord = serde_json::from_str(&std::fs::read_to_string(earlier).expect("the earlier record")).expect("the earlier record reads");
        let baseline = earlier.verify(&trusted).unwrap_or_else(|e| {
            eprintln!("relay_bench: the earlier record: {e}");
            std::process::exit(2)
        });
        let d = Tolerances::default();
        let tol = Tolerances { cpu_percent: num("--cpu-percent", d.cpu_percent), delivered_ppm: num("--delivered-ppm", d.delivered_ppm), lost_objects: num("--lost-objects", d.lost_objects), receive_variance_us: num("--variance-us", d.receive_variance_us), ..d };
        // A record its host held is held to that as well: a later run the host did not hold,
        // its lateness past the bound or a stream object lost, is outside it.
        let lapsed: Vec<String> = if baseline.held().is_empty() { record.held().into_iter().map(|u| format!("held by the record, not by this run: {u}")).collect() } else { Vec::new() };
        match bench::guard(&baseline, &record, &tol).map(|outside| [outside, lapsed].concat()) {
            Err(e) => {
                eprintln!("relay_bench: {e}");
                std::process::exit(2)
            }
            Ok(outside) if outside.is_empty() => eprintln!("relay_bench: within the tolerances of the record of {} (cpu {} percent, delivered {} ppm)", baseline.date, tol.cpu_percent, tol.delivered_ppm),
            Ok(outside) => {
                for o in outside {
                    eprintln!("relay_bench: {o}");
                }
                status = 1;
            }
        }
    }
    let unheld = record.held();
    for u in &unheld {
        eprintln!("relay_bench: not held: {u}");
    }
    // Guarded, the run is held to its record alone (above); else a run not held exits 4.
    if arg("--guard").is_none() && !unheld.is_empty() {
        status = 4;
    }
    status
}