use anyhow::{Context, Result, anyhow}; use clap::Parser; use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem}; use std::ffi::OsStr; use std::fmt::Write as _; use std::fs; use std::io::Read; use std::path::PathBuf; use std::process::Command; use std::sync::mpsc; use std::time::{Duration, Instant}; #[derive(Debug, Parser)] #[command( name = "dosh-bench", version = dosh::build_info::VERSION, long_version = dosh::build_info::LONG_VERSION )] struct Args { #[arg(long, default_value = "local")] server: String, #[arg(long, default_value = "default")] session: String, #[arg(long, default_value_t = 22)] ssh_port: u16, #[arg(long, default_value_t = 50000)] dosh_port: u16, #[arg(long)] dosh_host: Option, #[arg(long, default_value_t = 3)] iterations: usize, #[arg(long)] local_auth: bool, /// Benchmark native cold auth (no cache, native handshake) terminal-ready time. #[arg(long)] cold_native: bool, /// Benchmark cached attach-ticket terminal-ready time (warms the cache first). #[arg(long)] cached_ticket: bool, /// Benchmark UDP resume terminal-ready time (warms the cache first). #[arg(long)] resume: bool, #[arg(long)] client: Option, #[arg(long, default_value = "~/.local/bin/dosh-auth")] ssh_auth_command: String, #[arg(long)] ssh_key: Option, #[arg(long)] ssh_known_hosts: Option, #[arg(long)] ssh_control_path: Option, #[arg(long)] controlmaster: bool, #[arg(long)] no_cache: bool, #[arg(long)] warm_cache: bool, #[arg(long)] skip_ssh_baseline: bool, #[arg(long)] include_mosh: bool, #[arg(long, default_value = "mosh")] mosh: PathBuf, #[arg(long, default_value = "mosh-server")] mosh_server_command: String, #[arg(long)] mosh_port: Option, /// Emit machine-readable JSON (one object per metric, with raw samples). #[arg(long)] json: bool, /// Emit a publishable Markdown benchmark report. #[arg(long)] markdown: bool, /// Write the report to a file instead of stdout. #[arg(long)] output: Option, /// Optional label printed in summary/JSON output (e.g. machine/OS identifier). #[arg(long)] label: Option, #[arg(long)] assert_ssh_plus_ms: Option, #[arg(long)] assert_mosh_minus_ms: Option, #[arg(long)] assert_dosh_max_ms: Option, } /// One Dosh attach path to benchmark. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum DoshPath { /// `--auth native --no-cache`: full native handshake, no cache. ColdNative, /// Cached attach-ticket fast path (requires a warmed cache). CachedTicket, /// UDP resume fast path (requires a warmed cache + `cache_attach_tickets = false`). Resume, /// `--local-auth`: self-contained local bootstrap, no SSH. LocalAuth, /// SSH-bootstrap cold attach (cache controlled by `--no-cache` / `--warm-cache`). SshBootstrap, } impl DoshPath { fn metric(self) -> &'static str { match self { DoshPath::ColdNative => "dosh_cold_native_ms", DoshPath::CachedTicket => "dosh_cached_attach_ms", DoshPath::Resume => "dosh_resume_ms", DoshPath::LocalAuth => "dosh_local_attach_ms", DoshPath::SshBootstrap => "dosh_attach_ms", } } /// Whether this path consumes a warmed credential cache. fn needs_warm(self) -> bool { matches!(self, DoshPath::CachedTicket | DoshPath::Resume) } } fn main() -> Result<()> { let args = Args::parse(); let client = args.client.clone().unwrap_or_else(default_client_path); if args.no_cache && args.warm_cache { return Err(anyhow!("--warm-cache cannot be used with --no-cache")); } let generated_control_path = if args.controlmaster { Some(std::env::temp_dir().join(format!("dosh-bench-control-{}", std::process::id()))) } else { None }; let control_path = generated_control_path .as_ref() .or(args.ssh_control_path.as_ref()); let _controlmaster = if let Some(path) = generated_control_path.as_ref() { Some(ControlMaster::start(&args, path.clone())?) } else { None }; // Resolve which Dosh paths to benchmark. Explicit path flags select an // explicit matrix; otherwise fall back to the legacy single-path behavior // so existing callers (CI docker scripts) keep working unchanged. let explicit_paths = explicit_dosh_paths(&args); let dosh_paths = if explicit_paths.is_empty() { vec![legacy_dosh_path(&args)] } else { explicit_paths }; let mut results: Vec = Vec::new(); // SSH baseline (shared across the matrix; the comparison is per-iteration // ssh-true vs the dosh path startup that replaces it). let run_ssh = !args.local_auth && !args.skip_ssh_baseline && !dosh_only(&dosh_paths); if run_ssh { let mut ssh_times = Vec::new(); for _ in 0..args.iterations.max(1) { let mut ssh = Command::new("ssh"); add_ssh_options(&mut ssh, &args, control_path); ssh.arg(&args.server).arg("true"); ssh_times.push(time_command(&mut ssh)?); } results.push(MetricSamples::new("ssh_true_ms", ssh_times)); } for path in &dosh_paths { if path.needs_warm() { // Prime the cache with one attach so the fast path has credentials. let _ = time_dosh_attach(&client, &args, *path, control_path, true)?; } let mut samples = Vec::new(); for _ in 0..args.iterations.max(1) { samples.push(time_dosh_attach( &client, &args, *path, control_path, false, )?); } results.push(MetricSamples::new(path.metric(), samples)); } if args.include_mosh { let mut mosh_times = Vec::new(); for _ in 0..args.iterations.max(1) { mosh_times.push(time_mosh_in_pty(&args)?); } results.push(MetricSamples::new("mosh_start_true_ms", mosh_times)); } let report = if args.json { render_json(&args, &results) } else if args.markdown { render_markdown(&args, &results) } else { render_table(&args, &results) }; write_report(&args, &report)?; run_assertions(&args, &results)?; Ok(()) } /// Paths explicitly requested via flags. fn explicit_dosh_paths(args: &Args) -> Vec { let mut paths = Vec::new(); if args.cold_native { paths.push(DoshPath::ColdNative); } if args.cached_ticket { paths.push(DoshPath::CachedTicket); } if args.resume { paths.push(DoshPath::Resume); } paths } /// The single path implied by legacy flags when no explicit path is requested. fn legacy_dosh_path(args: &Args) -> DoshPath { if args.local_auth { DoshPath::LocalAuth } else if args.warm_cache { DoshPath::CachedTicket } else { DoshPath::SshBootstrap } } /// True when none of the selected paths needs an SSH baseline comparison /// (i.e. all are local-auth or cache fast paths driven without SSH). fn dosh_only(paths: &[DoshPath]) -> bool { paths.iter().all(|p| { matches!( p, DoshPath::LocalAuth | DoshPath::CachedTicket | DoshPath::Resume ) }) } fn time_dosh_attach( client: &PathBuf, args: &Args, path: DoshPath, control_path: Option<&PathBuf>, warm: bool, ) -> Result { let mut cmd = Command::new(client); cmd.arg("--attach-only") .arg("--session") .arg(&args.session) .arg("--dosh-port") .arg(args.dosh_port.to_string()); if let Some(host) = &args.dosh_host { cmd.arg("--dosh-host").arg(host); } match path { DoshPath::LocalAuth => { cmd.arg("--local-auth"); // While warming, write the cache; measured runs honor --no-cache. if args.no_cache && !warm { cmd.arg("--no-cache"); } cmd.arg(&args.server); } DoshPath::CachedTicket | DoshPath::Resume => { // Fast paths must read the warmed cache, so never pass --no-cache here. // Whether the cache uses tickets vs resume is decided by the client // config (`cache_attach_tickets`) the harness wrote into HOME. if args.local_auth { cmd.arg("--local-auth"); } add_bootstrap_args(&mut cmd, args, control_path); cmd.arg(&args.server); } DoshPath::ColdNative => { cmd.arg("--auth").arg("native"); // Cold path: never use a warmed cache. cmd.arg("--no-cache"); add_bootstrap_args(&mut cmd, args, control_path); cmd.arg(&args.server); } DoshPath::SshBootstrap => { add_bootstrap_args(&mut cmd, args, control_path); // Cold SSH bootstrap honors --no-cache on measured runs. if args.no_cache && !warm { cmd.arg("--no-cache"); } cmd.arg(&args.server); } } time_command(&mut cmd) } /// Append the SSH bootstrap arguments shared by the non-local paths. Native /// auth reuses the same SSH key/known-hosts/port plumbing. fn add_bootstrap_args(cmd: &mut Command, args: &Args, control_path: Option<&PathBuf>) { if args.local_auth { return; } cmd.arg("--ssh-port") .arg(args.ssh_port.to_string()) .arg("--ssh-auth-command") .arg(&args.ssh_auth_command); if let Some(key) = &args.ssh_key { cmd.arg("--ssh-key").arg(key); } if let Some(known_hosts) = &args.ssh_known_hosts { cmd.arg("--ssh-known-hosts").arg(known_hosts); } if let Some(control_path) = control_path { cmd.arg("--ssh-control-path").arg(control_path); } } fn add_ssh_options(cmd: &mut Command, args: &Args, control_path: Option<&PathBuf>) { cmd.arg("-p").arg(args.ssh_port.to_string()).arg("-T"); if let Some(key) = &args.ssh_key { cmd.arg("-i").arg(key); } if let Some(known_hosts) = &args.ssh_known_hosts { cmd.arg("-o") .arg(format!("UserKnownHostsFile={}", known_hosts.display())); } if let Some(control_path) = control_path { cmd.arg("-S").arg(control_path); } } fn mosh_ssh_command(args: &Args) -> String { let mut parts = vec![ "ssh".to_string(), "-p".to_string(), args.ssh_port.to_string(), ]; if let Some(key) = &args.ssh_key { parts.push("-i".to_string()); parts.push(shell_word(key.as_os_str())); } if let Some(known_hosts) = &args.ssh_known_hosts { parts.push("-o".to_string()); parts.push(shell_word(OsStr::new(&format!( "UserKnownHostsFile={}", known_hosts.display() )))); } parts.join(" ") } fn time_mosh_in_pty(args: &Args) -> Result { let pty_system = NativePtySystem::default(); let pair = pty_system .openpty(PtySize { rows: 24, cols: 80, pixel_width: 0, pixel_height: 0, }) .context("open pty for mosh benchmark")?; let mut command_line = vec![ shell_word(args.mosh.as_os_str()), shell_word(OsStr::new(&format!("--ssh={}", mosh_ssh_command(args)))), shell_word(OsStr::new(&format!( "--server={}", args.mosh_server_command ))), "--no-init".to_string(), ]; if let Some(port) = &args.mosh_port { command_line.push("-p".to_string()); command_line.push(shell_word(OsStr::new(port))); } command_line.push(shell_word(OsStr::new(&args.server))); command_line.push("--".to_string()); command_line.push("true".to_string()); let mut cmd = CommandBuilder::new("/bin/sh"); cmd.arg("-c"); cmd.arg(command_line.join(" ")); let start = Instant::now(); let mut child = pair .slave .spawn_command(cmd) .context("spawn mosh benchmark in pty")?; drop(pair.slave); let mut reader = pair.master.try_clone_reader().context("clone pty reader")?; let (tx, rx) = mpsc::channel(); std::thread::spawn(move || { let mut output = Vec::new(); let _ = reader.read_to_end(&mut output); let _ = tx.send(output); }); { let _writer = pair.master.take_writer().context("take pty writer")?; } let status = child.wait().context("wait for mosh benchmark")?; let elapsed = start.elapsed(); drop(pair.master); let output = rx.recv().unwrap_or_default(); if !status.success() { return Err(anyhow!( "mosh command failed with status {status:?}\noutput:\n{}", String::from_utf8_lossy(&output) )); } Ok(elapsed) } fn shell_word(value: &OsStr) -> String { let value = value.to_string_lossy(); format!("'{}'", value.replace('\'', "'\\''")) } struct ControlMaster { server: String, ssh_port: u16, ssh_key: Option, ssh_known_hosts: Option, control_path: PathBuf, } impl ControlMaster { fn start(args: &Args, control_path: PathBuf) -> Result { let mut cmd = Command::new("ssh"); cmd.arg("-p") .arg(args.ssh_port.to_string()) .arg("-T") .arg("-M") .arg("-S") .arg(&control_path) .arg("-f") .arg("-N") .arg("-o") .arg("ControlPersist=60") .arg("-o") .arg("ExitOnForwardFailure=yes"); if let Some(key) = &args.ssh_key { cmd.arg("-i").arg(key); } if let Some(known_hosts) = &args.ssh_known_hosts { cmd.arg("-o") .arg(format!("UserKnownHostsFile={}", known_hosts.display())); } cmd.arg(&args.server); time_command(&mut cmd).context("start SSH ControlMaster")?; Ok(Self { server: args.server.clone(), ssh_port: args.ssh_port, ssh_key: args.ssh_key.clone(), ssh_known_hosts: args.ssh_known_hosts.clone(), control_path, }) } } impl Drop for ControlMaster { fn drop(&mut self) { let mut cmd = Command::new("ssh"); cmd.arg("-p") .arg(self.ssh_port.to_string()) .arg("-T") .arg("-S") .arg(&self.control_path) .arg("-O") .arg("exit"); if let Some(key) = &self.ssh_key { cmd.arg("-i").arg(key); } if let Some(known_hosts) = &self.ssh_known_hosts { cmd.arg("-o") .arg(format!("UserKnownHostsFile={}", known_hosts.display())); } let _ = cmd.arg(&self.server).output(); } } fn time_command(cmd: &mut Command) -> Result { let start = Instant::now(); let output = cmd.output().with_context(|| format!("run {:?}", cmd))?; if !output.status.success() { return Err(anyhow!( "command failed {:?}\nstdout:\n{}\nstderr:\n{}", cmd, String::from_utf8_lossy(&output.stdout), String::from_utf8_lossy(&output.stderr) )); } Ok(start.elapsed()) } /// Per-metric samples with summary statistics. struct MetricSamples { name: &'static str, samples: Vec, } impl MetricSamples { fn new(name: &'static str, samples: Vec) -> Self { Self { name, samples } } fn ms(&self) -> Vec { self.samples .iter() .map(|d| d.as_secs_f64() * 1000.0) .collect() } fn stats(&self) -> Stats { Stats::from_ms(&self.ms()) } } /// Summary statistics for a set of millisecond samples. struct Stats { count: usize, min: f64, median: f64, p95: f64, mean: f64, max: f64, } impl Stats { fn from_ms(values: &[f64]) -> Self { if values.is_empty() { return Self { count: 0, min: 0.0, median: 0.0, p95: 0.0, mean: 0.0, max: 0.0, }; } let mut sorted = values.to_vec(); sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); let count = sorted.len(); let mean = sorted.iter().sum::() / count as f64; Self { count, min: sorted[0], median: percentile(&sorted, 50.0), p95: percentile(&sorted, 95.0), mean, max: sorted[count - 1], } } } /// Linear-interpolated percentile over a pre-sorted slice. fn percentile(sorted: &[f64], pct: f64) -> f64 { if sorted.is_empty() { return 0.0; } if sorted.len() == 1 { return sorted[0]; } let rank = (pct / 100.0) * (sorted.len() - 1) as f64; let lo = rank.floor() as usize; let hi = rank.ceil() as usize; if lo == hi { sorted[lo] } else { let frac = rank - lo as f64; sorted[lo] + (sorted[hi] - sorted[lo]) * frac } } fn render_table(args: &Args, results: &[MetricSamples]) -> String { let mut out = String::new(); if let Some(label) = &args.label { let _ = writeln!(out, "# label: {label}"); } let _ = writeln!( out, "{:<24} {:>6} {:>9} {:>9} {:>9} {:>9} {:>9}", "metric", "n", "min", "median", "p95", "mean", "max" ); for metric in results { let s = metric.stats(); let _ = writeln!( out, "{:<24} {:>6} {:>9.2} {:>9.2} {:>9.2} {:>9.2} {:>9.2}", metric.name, s.count, s.min, s.median, s.p95, s.mean, s.max ); } // Raw per-iteration samples, so published numbers can include raw data. for metric in results { let ms: Vec = metric.ms().iter().map(|v| format!("{v:.2}")).collect(); let _ = writeln!(out, "{} samples_ms=[{}]", metric.name, ms.join(", ")); } out } fn render_json(args: &Args, results: &[MetricSamples]) -> String { let mut entries = Vec::new(); for metric in results { let s = metric.stats(); let samples: Vec = metric.ms().iter().map(|v| format!("{v:.4}")).collect(); entries.push(format!( "{{\"metric\":\"{}\",\"count\":{},\"min\":{:.4},\"median\":{:.4},\"p95\":{:.4},\"mean\":{:.4},\"max\":{:.4},\"samples_ms\":[{}]}}", metric.name, s.count, s.min, s.median, s.p95, s.mean, s.max, samples.join(",") )); } let label = match &args.label { Some(label) => format!("\"{}\"", label.replace('"', "\\\"")), None => "null".to_string(), }; format!( "{{\"label\":{label},\"iterations\":{},\"metrics\":[{}]}}", args.iterations.max(1), entries.join(",") ) } fn render_markdown(args: &Args, results: &[MetricSamples]) -> String { let mut out = String::new(); let label = args.label.as_deref().unwrap_or("Dosh benchmark"); let _ = writeln!(out, "# {label}\n"); let _ = writeln!(out, "- server: `{}`", args.server); let _ = writeln!(out, "- iterations: `{}`", args.iterations.max(1)); let _ = writeln!( out, "- generated_by: `dosh-bench {}`\n", env!("CARGO_PKG_VERSION") ); let _ = writeln!( out, "| metric | n | min ms | median ms | p95 ms | mean ms | max ms |" ); let _ = writeln!(out, "|---|---:|---:|---:|---:|---:|---:|"); for metric in results { let s = metric.stats(); let _ = writeln!( out, "| `{}` | {} | {:.2} | {:.2} | {:.2} | {:.2} | {:.2} |", metric.name, s.count, s.min, s.median, s.p95, s.mean, s.max ); } let _ = writeln!(out, "\n## Raw Samples\n"); for metric in results { let ms: Vec = metric.ms().iter().map(|v| format!("{v:.2}")).collect(); let _ = writeln!(out, "- `{}`: [{}]", metric.name, ms.join(", ")); } out } fn write_report(args: &Args, report: &str) -> Result<()> { if let Some(path) = &args.output { if let Some(parent) = path.parent() && !parent.as_os_str().is_empty() { fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?; } fs::write(path, report).with_context(|| format!("write {}", path.display()))?; println!("wrote {}", path.display()); } else { println!("{report}"); } Ok(()) } fn metric_mean(results: &[MetricSamples], name: &str) -> Option { results .iter() .find(|m| m.name == name) .map(|m| m.stats().mean) } /// The headline dosh metric for assertions: prefer cached, then resume, then /// cold native, then ssh-bootstrap, then local-auth. fn primary_dosh(results: &[MetricSamples]) -> Option<(&'static str, f64)> { const ORDER: [&str; 5] = [ "dosh_cached_attach_ms", "dosh_resume_ms", "dosh_cold_native_ms", "dosh_attach_ms", "dosh_local_attach_ms", ]; for name in ORDER { if let Some(mean) = metric_mean(results, name) { return Some((name, mean)); } } None } fn run_assertions(args: &Args, results: &[MetricSamples]) -> Result<()> { if let Some(margin) = args.assert_ssh_plus_ms { let ssh = metric_mean(results, "ssh_true_ms") .ok_or_else(|| anyhow!("--assert-ssh-plus-ms requires non-local SSH benchmark"))?; let (name, dosh) = primary_dosh(results).ok_or_else(|| anyhow!("no dosh metric to assert against"))?; if dosh > ssh + margin { return Err(anyhow!( "{name} avg {dosh:.2}ms exceeded ssh avg {ssh:.2}ms + {margin:.2}ms" )); } println!("gate ok: {name} avg {dosh:.2}ms <= ssh avg {ssh:.2}ms + {margin:.2}ms"); } if let Some(margin) = args.assert_mosh_minus_ms { let mosh = metric_mean(results, "mosh_start_true_ms") .ok_or_else(|| anyhow!("--assert-mosh-minus-ms requires --include-mosh benchmark"))?; let (name, dosh) = primary_dosh(results).ok_or_else(|| anyhow!("no dosh metric to assert against"))?; if dosh + margin > mosh { return Err(anyhow!( "{name} avg {dosh:.2}ms was not at least {margin:.2}ms faster than mosh avg {mosh:.2}ms" )); } println!("gate ok: {name} avg {dosh:.2}ms + {margin:.2}ms <= mosh avg {mosh:.2}ms"); } if let Some(max_ms) = args.assert_dosh_max_ms { let (name, dosh) = primary_dosh(results).ok_or_else(|| anyhow!("no dosh metric to assert against"))?; if dosh > max_ms { return Err(anyhow!("{name} avg {dosh:.2}ms exceeded max {max_ms:.2}ms")); } println!("gate ok: {name} avg {dosh:.2}ms <= {max_ms:.2}ms"); } Ok(()) } fn default_client_path() -> PathBuf { std::env::current_exe() .ok() .and_then(|path| path.parent().map(|parent| parent.join("dosh-client"))) .unwrap_or_else(|| PathBuf::from("dosh-client")) } #[cfg(test)] mod tests { use super::*; fn args_for_render() -> Args { Args { server: "local".to_string(), session: "default".to_string(), ssh_port: 22, dosh_port: 50000, dosh_host: None, iterations: 2, local_auth: true, cold_native: false, cached_ticket: false, resume: false, client: None, ssh_auth_command: "~/.local/bin/dosh-auth".to_string(), ssh_key: None, ssh_known_hosts: None, ssh_control_path: None, controlmaster: false, no_cache: false, warm_cache: false, skip_ssh_baseline: true, include_mosh: false, mosh: PathBuf::from("mosh"), mosh_server_command: "mosh-server".to_string(), mosh_port: None, json: false, markdown: true, output: None, label: Some("test host".to_string()), assert_ssh_plus_ms: None, assert_mosh_minus_ms: None, assert_dosh_max_ms: None, } } #[test] fn markdown_report_contains_summary_and_samples() { let args = args_for_render(); let metrics = vec![MetricSamples::new( "dosh_cached_attach_ms", vec![Duration::from_millis(3), Duration::from_millis(5)], )]; let report = render_markdown(&args, &metrics); assert!(report.contains("# test host")); assert!(report.contains("| `dosh_cached_attach_ms` | 2 | 3.00")); assert!(report.contains("- `dosh_cached_attach_ms`: [3.00, 5.00]")); } #[test] fn json_report_contains_metrics() { let args = args_for_render(); let metrics = vec![MetricSamples::new( "dosh_local_attach_ms", vec![Duration::from_millis(7)], )]; let report = render_json(&args, &metrics); assert!(report.contains("\"label\":\"test host\"")); assert!(report.contains("\"metric\":\"dosh_local_attach_ms\"")); assert!(report.contains("\"samples_ms\":[7.0000]")); } }