Add public comparison benchmarks and SSH import

src/bin/dosh-bench.rs

@@ -1,7 +1,11 @@
 use anyhow::{Context, Result, anyhow};
 use clap::Parser;
+use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem};
+use std::ffi::OsStr;
+use std::io::Read;
 use std::path::PathBuf;
 use std::process::Command;
+use std::sync::mpsc;
 use std::time::{Duration, Instant};
 
 #[derive(Debug, Parser)]
@@ -36,7 +40,17 @@ struct Args {
     #[arg(long)]
     no_cache: 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<String>,
+    #[arg(long)]
     assert_ssh_plus_ms: Option<f64>,
+    #[arg(long)]
+    assert_mosh_minus_ms: Option<f64>,
 }
 
 fn main() -> Result<()> {
@@ -44,6 +58,7 @@ fn main() -> Result<()> {
     let client = args.client.clone().unwrap_or_else(default_client_path);
     let mut ssh_times = Vec::new();
     let mut dosh_times = Vec::new();
+    let mut mosh_times = Vec::new();
     let generated_control_path = if args.controlmaster {
         Some(std::env::temp_dir().join(format!("dosh-bench-control-{}", std::process::id())))
     } else {
@@ -97,6 +112,10 @@ fn main() -> Result<()> {
             cmd.arg(&args.server);
         }
         dosh_times.push(time_command(&mut cmd)?);
+
+        if args.include_mosh {
+            mosh_times.push(time_mosh_in_pty(&args)?);
+        }
     }
 
     if !ssh_times.is_empty() {
@@ -111,6 +130,13 @@ fn main() -> Result<()> {
         avg_ms(&dosh_times),
         dosh_times
     );
+    if !mosh_times.is_empty() {
+        println!(
+            "mosh_start_true_ms avg={:.2} samples={:?}",
+            avg_ms(&mosh_times),
+            mosh_times
+        );
+    }
     if let Some(margin) = args.assert_ssh_plus_ms {
         if ssh_times.is_empty() {
             return Err(anyhow!(
@@ -126,6 +152,21 @@ fn main() -> Result<()> {
         }
         println!("gate ok: dosh avg {dosh_avg:.2}ms <= ssh avg {ssh_avg:.2}ms + {margin:.2}ms");
     }
+    if let Some(margin) = args.assert_mosh_minus_ms {
+        if mosh_times.is_empty() {
+            return Err(anyhow!(
+                "--assert-mosh-minus-ms requires --include-mosh benchmark"
+            ));
+        }
+        let dosh_avg = avg_ms(&dosh_times);
+        let mosh_avg = avg_ms(&mosh_times);
+        if dosh_avg + margin > mosh_avg {
+            return Err(anyhow!(
+                "dosh attach avg {dosh_avg:.2}ms was not at least {margin:.2}ms faster than mosh avg {mosh_avg:.2}ms"
+            ));
+        }
+        println!("gate ok: dosh avg {dosh_avg:.2}ms + {margin:.2}ms <= mosh avg {mosh_avg:.2}ms");
+    }
     Ok(())
 }
 
@@ -143,6 +184,94 @@ fn add_ssh_options(cmd: &mut Command, args: &Args, control_path: Option<&PathBuf
     }
 }
 
+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<Duration> {
+    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,