Initial Dosh implementation

src/bin/dosh-bench.rs

@@ -0,0 +1,233 @@
+use anyhow::{Context, Result, anyhow};
+use clap::Parser;
+use std::path::PathBuf;
+use std::process::Command;
+use std::time::{Duration, Instant};
+
+#[derive(Debug, Parser)]
+#[command(name = "dosh-bench")]
+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<String>,
+    #[arg(long, default_value_t = 3)]
+    iterations: usize,
+    #[arg(long)]
+    local_auth: bool,
+    #[arg(long)]
+    client: Option<PathBuf>,
+    #[arg(long, default_value = "dosh-auth")]
+    ssh_auth_command: String,
+    #[arg(long)]
+    ssh_key: Option<PathBuf>,
+    #[arg(long)]
+    ssh_known_hosts: Option<PathBuf>,
+    #[arg(long)]
+    ssh_control_path: Option<PathBuf>,
+    #[arg(long)]
+    controlmaster: bool,
+    #[arg(long)]
+    no_cache: bool,
+    #[arg(long)]
+    assert_ssh_plus_ms: Option<f64>,
+}
+
+fn main() -> Result<()> {
+    let args = Args::parse();
+    let client = args.client.clone().unwrap_or_else(default_client_path);
+    let mut ssh_times = Vec::new();
+    let mut dosh_times = Vec::new();
+    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
+    };
+
+    for _ in 0..args.iterations.max(1) {
+        if !args.local_auth {
+            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)?);
+        }
+
+        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);
+        }
+        if args.local_auth {
+            cmd.arg("--local-auth").arg(&args.server);
+        } else {
+            cmd.arg("--ssh-port")
+                .arg(args.ssh_port.to_string())
+                .arg("--ssh-auth-command")
+                .arg(&args.ssh_auth_command);
+            if args.no_cache {
+                cmd.arg("--no-cache");
+            }
+            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);
+            }
+            cmd.arg(&args.server);
+        }
+        dosh_times.push(time_command(&mut cmd)?);
+    }
+
+    if !ssh_times.is_empty() {
+        println!(
+            "ssh_true_ms avg={:.2} samples={:?}",
+            avg_ms(&ssh_times),
+            ssh_times
+        );
+    }
+    println!(
+        "dosh_attach_ms avg={:.2} samples={:?}",
+        avg_ms(&dosh_times),
+        dosh_times
+    );
+    if let Some(margin) = args.assert_ssh_plus_ms {
+        if ssh_times.is_empty() {
+            return Err(anyhow!(
+                "--assert-ssh-plus-ms requires non-local SSH benchmark"
+            ));
+        }
+        let ssh_avg = avg_ms(&ssh_times);
+        let dosh_avg = avg_ms(&dosh_times);
+        if dosh_avg > ssh_avg + margin {
+            return Err(anyhow!(
+                "dosh attach avg {dosh_avg:.2}ms exceeded ssh avg {ssh_avg:.2}ms + {margin:.2}ms"
+            ));
+        }
+        println!("gate ok: dosh avg {dosh_avg:.2}ms <= ssh avg {ssh_avg:.2}ms + {margin:.2}ms");
+    }
+    Ok(())
+}
+
+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);
+    }
+}
+
+struct ControlMaster {
+    server: String,
+    ssh_port: u16,
+    ssh_key: Option<PathBuf>,
+    ssh_known_hosts: Option<PathBuf>,
+    control_path: PathBuf,
+}
+
+impl ControlMaster {
+    fn start(args: &Args, control_path: PathBuf) -> Result<Self> {
+        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<Duration> {
+    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())
+}
+
+fn avg_ms(samples: &[Duration]) -> f64 {
+    let total: f64 = samples.iter().map(|d| d.as_secs_f64() * 1000.0).sum();
+    total / samples.len() as f64
+}
+
+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"))
+}