395 lines
12 KiB
Rust
395 lines
12 KiB
Rust
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)]
|
|
#[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 = "~/.local/bin/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)]
|
|
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<String>,
|
|
#[arg(long)]
|
|
assert_ssh_plus_ms: Option<f64>,
|
|
#[arg(long)]
|
|
assert_mosh_minus_ms: Option<f64>,
|
|
#[arg(long)]
|
|
assert_dosh_max_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 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 {
|
|
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
|
|
};
|
|
if args.no_cache && args.warm_cache {
|
|
return Err(anyhow!("--warm-cache cannot be used with --no-cache"));
|
|
}
|
|
let dosh_label = if args.warm_cache {
|
|
let _ = time_dosh_attach(&client, &args, control_path)?;
|
|
"dosh_cached_attach_ms"
|
|
} else {
|
|
"dosh_attach_ms"
|
|
};
|
|
|
|
for _ in 0..args.iterations.max(1) {
|
|
if !args.local_auth && !args.skip_ssh_baseline {
|
|
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)?);
|
|
}
|
|
|
|
dosh_times.push(time_dosh_attach(&client, &args, control_path)?);
|
|
|
|
if args.include_mosh {
|
|
mosh_times.push(time_mosh_in_pty(&args)?);
|
|
}
|
|
}
|
|
|
|
if !ssh_times.is_empty() {
|
|
println!(
|
|
"ssh_true_ms avg={:.2} samples={:?}",
|
|
avg_ms(&ssh_times),
|
|
ssh_times
|
|
);
|
|
}
|
|
println!(
|
|
"{dosh_label} avg={:.2} samples={:?}",
|
|
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!(
|
|
"--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");
|
|
}
|
|
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");
|
|
}
|
|
if let Some(max_ms) = args.assert_dosh_max_ms {
|
|
let dosh_avg = avg_ms(&dosh_times);
|
|
if dosh_avg > max_ms {
|
|
return Err(anyhow!(
|
|
"{dosh_label} avg {dosh_avg:.2}ms exceeded max {max_ms:.2}ms"
|
|
));
|
|
}
|
|
println!("gate ok: {dosh_label} avg {dosh_avg:.2}ms <= {max_ms:.2}ms");
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn time_dosh_attach(
|
|
client: &PathBuf,
|
|
args: &Args,
|
|
control_path: Option<&PathBuf>,
|
|
) -> Result<Duration> {
|
|
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);
|
|
}
|
|
time_command(&mut cmd)
|
|
}
|
|
|
|
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<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,
|
|
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"))
|
|
}
|