git.delta.rocks / jrsonnet / refs/commits / 604f09da5796

1#![allow(clippy::cast_precision_loss)]23use std::{4	ffi::OsString,5	mem,6	process::{Command, Stdio},7	time::Instant,8};910use anyhow::{Result, bail};11use nix::{libc, sys::wait::WaitStatus, unistd::Pid};1213#[derive(Debug, Clone)]14pub struct Stats {15	pub min: f64,16	pub max: f64,17	pub mean: f64,18	pub stddev: f64,19}2021impl Stats {22	fn from_samples(samples: &[f64]) -> Self {23		let n = samples.len() as f64;24		let mean = samples.iter().sum::<f64>() / n;25		let var = if samples.len() > 1 {26			samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (n - 1.0)27		} else {28			0.029		};30		Self {31			min: samples.iter().copied().fold(f64::INFINITY, f64::min),32			max: samples.iter().copied().fold(f64::NEG_INFINITY, f64::max),33			mean,34			stddev: var.sqrt(),35		}36	}37}3839#[derive(Debug, Clone)]40pub struct BenchResult {41	pub runs: u32,42	/// Wall-clock time per run, seconds.43	pub wall_secs: Stats,44	/// Peak resident set per run, KiB (Linux `ru_maxrss`).45	pub max_rss_kib: Stats,46}4748pub struct BenchOpts<'a> {49	pub program: &'a OsString,50	pub args: &'a [OsString],51	pub runs: u32,52	pub warmup: u32,53	pub output: bool,54}5556pub fn bench(opts: BenchOpts<'_>) -> Result<BenchResult> {57	if opts.runs == 0 {58		bail!("runs must be >= 1");59	}6061	for _ in 0..opts.warmup {62		run_once(opts.program, opts.args, opts.output)?;63	}6465	let mut wall = Vec::with_capacity(opts.runs as usize);66	let mut rss = Vec::with_capacity(opts.runs as usize);67	for _ in 0..opts.runs {68		let s = run_once(opts.program, opts.args, opts.output)?;69		wall.push(s.wall_secs);70		rss.push(s.max_rss_kib as f64);71	}7273	Ok(BenchResult {74		runs: opts.runs,75		wall_secs: Stats::from_samples(&wall),76		max_rss_kib: Stats::from_samples(&rss),77	})78}7980struct Sample {81	wall_secs: f64,82	max_rss_kib: i64,83}8485fn run_once(program: &OsString, args: &[OsString], silent: bool) -> Result<Sample> {86	let mut cmd = Command::new(program);87	cmd.args(args);88	if silent {89		cmd.stdout(Stdio::null()).stderr(Stdio::null());90	}9192	let start = Instant::now();93	let child = cmd.spawn()?;94	let pid = child.id() as libc::pid_t;95	// We'll reap via wait4 ourselves; don't let std touch this handle again.96	mem::forget(child);9798	let mut status: libc::c_int = 0;99	let mut ru: libc::rusage = unsafe { mem::zeroed() };100	let waited = unsafe { libc::wait4(pid, &raw mut status, 0, &raw mut ru) };101	let elapsed = start.elapsed();102	if waited < 0 {103		return Err(std::io::Error::last_os_error().into());104	}105106	match WaitStatus::from_raw(Pid::from_raw(pid), status)? {107		WaitStatus::Exited(_, 0) => {}108		WaitStatus::Exited(_, code) => bail!("child exited with code {code}"),109		WaitStatus::Signaled(_, sig, _) => bail!("child killed by signal {sig:?}"),110		other => bail!("unexpected wait status: {other:?}"),111	}112113	Ok(Sample {114		wall_secs: elapsed.as_secs_f64(),115		max_rss_kib: ru.ru_maxrss,116	})117}118119#[cfg(target_os = "linux")]120pub fn bench_cmd(args: &[String], runs: u32, warmup: u32, output: bool) -> Result<()> {121	let program = std::ffi::OsString::from(&args[0]);122	let rest: Vec<std::ffi::OsString> = args[1..].iter().map(Into::into).collect();123	let r = bench(BenchOpts {124		program: &program,125		args: &rest,126		runs,127		warmup,128		output,129	})?;130	eprintln!(131		"runs: {}  wall: {:.3}s ± {:.3}s  [{:.3}..{:.3}]",132		r.runs, r.wall_secs.mean, r.wall_secs.stddev, r.wall_secs.min, r.wall_secs.max,133	);134	eprintln!(135		"           max_rss: {} ± {} KiB  [{}..{}]",136		r.max_rss_kib.mean as i64,137		r.max_rss_kib.stddev as i64,138		r.max_rss_kib.min as i64,139		r.max_rss_kib.max as i64,140	);141	Ok(())142}