unshell/examples/protocol_bench.rs

use std::hint::black_box;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::Instant;

use unshell::protocol::tree::{
    ChildRoute, Endpoint, Ingress, LeafSpec, LocalEvent, ProtocolEndpoint,
};
use unshell::protocol::{CallMessage, PacketHeader, PacketType, decode_frame, encode_packet};

const SAMPLES: usize = 500;
const ITERS: usize = 1_000;
const TOOL_ITERS: usize = 10_000;

fn main() {
    if std::env::args().nth(1).as_deref() == Some("tools") {
        run_external_tools();
        return;
    }

    println!("protocol benchmark");
    println!("samples: {SAMPLES}");
    println!("iterations/sample: {ITERS}");
    println!();

    let benches = [
        bench_encode_call(),
        bench_decode_call(),
        bench_forward_call_receive(),
        bench_local_call_receive(),
        bench_hook_data_receive(),
    ];

    println!(
        "{:32} {:>14} {:>14} {:>14}",
        "benchmark", "mean ns/op", "stddev", "samples"
    );
    for bench in benches {
        println!(
            "{:32} {:>14.2} {:>14.2} {:>14}",
            bench.name, bench.mean_ns, bench.stddev_ns, bench.samples
        );
    }

    println!();
    println!("Run `cargo run --example protocol_bench -- tools` to build and execute");
    println!("the standalone operation binaries under strace, perf, and heaptrack.");
}

struct BenchResult {
    name: &'static str,
    mean_ns: f64,
    stddev_ns: f64,
    samples: usize,
}

fn bench_encode_call() -> BenchResult {
    let header = PacketHeader {
        packet_type: PacketType::Call,
        src_path: path(&["root"]),
        dst_path: path(&["root", "worker"]),
        dst_leaf: Some(String::from("service")),
        hook_id: None,
    };
    let message = CallMessage {
        procedure_id: String::from("example.service.v1.invoke"),
        data: vec![7; 64],
        response_hook: None,
    };

    run_bench("encode_call", || {
        let frame =
            encode_packet(black_box(&header), black_box(&message)).expect("encode should work");
        black_box(frame.len());
    })
}

fn bench_decode_call() -> BenchResult {
    let header = PacketHeader {
        packet_type: PacketType::Call,
        src_path: path(&["root"]),
        dst_path: path(&["root", "worker"]),
        dst_leaf: Some(String::from("service")),
        hook_id: None,
    };
    let message = CallMessage {
        procedure_id: String::from("example.service.v1.invoke"),
        data: vec![9; 64],
        response_hook: None,
    };
    let frame = encode_packet(&header, &message).expect("seed frame should encode");

    run_bench("decode_call", || {
        let parsed = decode_frame(black_box(frame.as_slice())).expect("decode should work");
        let call = parsed.deserialize_call().expect("call should deserialize");
        black_box(call.data.len());
    })
}

fn bench_forward_call_receive() -> BenchResult {
    run_prebuilt_bench(
        "forward_call_receive",
        build_forward_call_cases,
        |(mut root, frame)| {
            let outcome = root
                .receive(&Ingress::Local, frame)
                .expect("forward receive should work");
            black_box(outcome.forward.is_some());
        },
    )
}

fn bench_local_call_receive() -> BenchResult {
    run_prebuilt_bench(
        "local_call_receive",
        build_local_call_cases,
        |(mut endpoint, frame)| {
            let outcome = endpoint
                .receive(&Ingress::Parent, frame)
                .expect("local call should work");
            match black_box(outcome.event) {
                Some(LocalEvent::Call { .. }) => {}
                other => panic!("expected local call event, got {other:?}"),
            }
        },
    )
}

fn bench_hook_data_receive() -> BenchResult {
    run_prebuilt_bench(
        "hook_data_receive",
        build_hook_data_cases,
        |(mut host, frame)| {
            let outcome = host
                .receive(&Ingress::Child(path(&["worker"])), frame)
                .expect("hook data should work");
            match black_box(outcome.event) {
                Some(LocalEvent::Data { .. }) => {}
                other => panic!("expected local data event, got {other:?}"),
            }
        },
    )
}

fn run_bench(name: &'static str, mut op: impl FnMut()) -> BenchResult {
    let mut samples = Vec::with_capacity(SAMPLES);
    for _ in 0..SAMPLES {
        let start = Instant::now();
        for _ in 0..ITERS {
            op();
        }
        let elapsed = start.elapsed().as_nanos() as f64 / ITERS as f64;
        samples.push(elapsed);
    }
    summarize(name, &samples)
}

fn run_prebuilt_bench<T, F>(
    name: &'static str,
    mut build_cases: F,
    mut op: impl FnMut(T),
) -> BenchResult
where
    F: FnMut() -> Vec<T>,
{
    let mut repeated = Vec::with_capacity(SAMPLES);
    for _ in 0..SAMPLES {
        let mut cases = build_cases();
        assert_eq!(cases.len(), ITERS);
        let start = Instant::now();
        for case in cases.drain(..) {
            op(case);
        }
        let elapsed = start.elapsed().as_nanos() as f64 / ITERS as f64;
        repeated.push(elapsed);
    }
    summarize(name, &repeated)
}

fn build_forward_call_cases() -> Vec<(ProtocolEndpoint, unshell::protocol::FrameBytes)> {
    (0..ITERS)
        .map(|_| {
            let mut root = ProtocolEndpoint::new(
                Vec::new(),
                None,
                vec![ChildRoute::registered(path(&["edge"]))],
                Vec::new(),
            );
            let hook_id = root.allocate_hook_id();
            let frame = root
                .make_call(
                    path(&["edge", "worker"]),
                    Some(String::from("service")),
                    String::from("example.service.v1.invoke"),
                    Some(hook_id),
                    vec![1; 32],
                )
                .expect("seed call should encode");
            (root, frame)
        })
        .collect()
}

fn build_local_call_cases() -> Vec<(ProtocolEndpoint, unshell::protocol::FrameBytes)> {
    (0..ITERS)
        .map(|_| {
            let endpoint = ProtocolEndpoint::new(
                path(&["worker"]),
                Some(Vec::new()),
                Vec::new(),
                vec![LeafSpec {
                    name: String::from("service"),
                    procedures: vec![String::from("example.service.v1.invoke")],
                }],
            );
            let frame = encode_packet(
                &PacketHeader {
                    packet_type: PacketType::Call,
                    src_path: Vec::new(),
                    dst_path: path(&["worker"]),
                    dst_leaf: Some(String::from("service")),
                    hook_id: None,
                },
                &CallMessage {
                    procedure_id: String::from("example.service.v1.invoke"),
                    data: vec![2; 32],
                    response_hook: Some(unshell::protocol::HookTarget {
                        hook_id: 42,
                        return_path: Vec::new(),
                    }),
                },
            )
            .expect("seed local call should encode");
            (endpoint, frame)
        })
        .collect()
}

fn build_hook_data_cases() -> Vec<(ProtocolEndpoint, unshell::protocol::FrameBytes)> {
    (0..ITERS)
        .map(|_| {
            let mut host = ProtocolEndpoint::new(
                Vec::new(),
                None,
                vec![ChildRoute::registered(path(&["worker"]))],
                Vec::new(),
            );
            let hook_id = host.allocate_hook_id();
            host.make_call(
                path(&["worker"]),
                None,
                String::from("example.service.v1.invoke"),
                Some(hook_id),
                vec![3; 8],
            )
            .expect("seed active hook should encode");
            let frame = encode_packet(
                &PacketHeader {
                    packet_type: PacketType::Data,
                    src_path: path(&["worker"]),
                    dst_path: Vec::new(),
                    dst_leaf: None,
                    hook_id: Some(hook_id),
                },
                &unshell::protocol::DataMessage {
                    procedure_id: String::from("example.service.v1.invoke"),
                    data: vec![4; 16],
                    end_hook: false,
                },
            )
            .expect("seed data should encode");
            (host, frame)
        })
        .collect()
}

fn summarize(name: &'static str, samples: &[f64]) -> BenchResult {
    let mean = samples.iter().sum::<f64>() / samples.len() as f64;
    let variance = samples
        .iter()
        .map(|sample| {
            let delta = sample - mean;
            delta * delta
        })
        .sum::<f64>()
        / samples.len() as f64;

    BenchResult {
        name,
        mean_ns: mean,
        stddev_ns: variance.sqrt(),
        samples: samples.len(),
    }
}

fn path(parts: &[&str]) -> Vec<String> {
    parts.iter().map(|part| String::from(*part)).collect()
}

fn run_external_tools() {
    let root = Path::new(env!("CARGO_MANIFEST_DIR"));
    build_examples(root);

    let ops = [
        ("encode_call", "protocol_op_encode_call"),
        ("decode_call", "protocol_op_decode_call"),
        ("forward_call_receive", "protocol_op_forward_call_receive"),
        ("local_call_receive", "protocol_op_local_call_receive"),
        ("hook_data_receive", "protocol_op_hook_data_receive"),
    ];

    let heap_dir = root.join("heaptrack-cli");
    std::fs::create_dir_all(&heap_dir).expect("heaptrack-cli directory should be creatable");

    for (name, binary) in ops {
        let binary_path = root.join("target/debug/examples").join(binary);
        println!();
        println!("=== {name} ===");
        run_binary(&binary_path, TOOL_ITERS, "direct run");
        run_strace(&binary_path, TOOL_ITERS);
        run_perf(&binary_path, TOOL_ITERS);
        run_heaptrack(root, &heap_dir, name, &binary_path, TOOL_ITERS);
    }
}

fn build_examples(root: &Path) {
    run_command(
        "cargo build --examples",
        Command::new("cargo")
            .arg("build")
            .arg("--examples")
            .current_dir(root),
    );
}

fn run_binary(binary: &Path, iterations: usize, label: &str) {
    run_command(
        label,
        Command::new(binary).arg(iterations.to_string()),
    );
}

fn run_strace(binary: &Path, iterations: usize) {
    run_command(
        "strace -c memory syscalls",
        Command::new("strace")
            .arg("-qq")
            .arg("-c")
            .arg("-e")
            .arg("trace=brk,mmap,mremap,munmap,mprotect,madvise")
            .arg(binary)
            .arg(iterations.to_string()),
    );
}

fn run_perf(binary: &Path, iterations: usize) {
    run_command(
        "perf stat",
        Command::new("perf")
            .arg("stat")
            .arg("-e")
            .arg("task-clock,cycles,instructions,branches,branch-misses,cache-references,cache-misses")
            .arg(binary)
            .arg(iterations.to_string()),
    );
}

fn run_heaptrack(root: &Path, heap_dir: &Path, name: &str, binary: &Path, iterations: usize) {
    let prefix = heap_dir.join(format!("{name}.zst"));
    run_command(
        "heaptrack --record-only",
        Command::new("heaptrack")
            .arg("--record-only")
            .arg("-o")
            .arg(&prefix)
            .arg(binary)
            .arg(iterations.to_string())
            .current_dir(root),
    );

    let recorded = PathBuf::from(format!("{}.zst", prefix.display()));
    run_command(
        "heaptrack_print summary",
        Command::new("heaptrack_print")
            .arg("-f")
            .arg(recorded)
            .arg("-n")
            .arg("4")
            .arg("-s")
            .arg("2")
            .current_dir(root),
    );
}

fn run_command(label: &str, command: &mut Command) {
    println!("--- {label} ---");
    let output = command.output().unwrap_or_else(|error| panic!("{label} failed to launch: {error}"));
    if !output.stdout.is_empty() {
        print!("{}", String::from_utf8_lossy(&output.stdout));
    }
    if !output.stderr.is_empty() {
        print!("{}", String::from_utf8_lossy(&output.stderr));
    }
    assert!(output.status.success(), "{label} failed with status {}", output.status);
}