//! children run from a fresh copy of the bench binary: on linux a child's
//! ru_maxrss starts at its parent's rss, which can be tens of MiB into seeding.
use std::io::{Read, Write as _};
use std::os::unix::process::ExitStatusExt as _;
use std::path::{Path, PathBuf};
use std::process::{Command, ExitStatus, Stdio};
use std::sync::atomic::{AtomicI32, AtomicU64, Ordering};
use std::thread::JoinHandle;
use std::time::Instant;
const MAXRSS_UNIT: u64 = if cfg!(target_os = "macos") { 1 } else { 1024 };
/// cpu and peak rss include reaped descendants (git, nginx), not detached ones.
#[derive(Debug, Clone, Copy, Default, serde::Serialize, serde::Deserialize)]
pub struct Usage {
pub wall_s: f64,
pub answered_s: Option<f64>,
pub user_s: f64,
pub sys_s: f64,
pub max_rss_bytes: u64,
/// sampled every 50 ms on linux, which can miss a short peak. none elsewhere.
#[serde(default)]
pub peak_anon_bytes: Option<u64>,
#[serde(default)]
pub peak_file_bytes: Option<u64>,
pub out_blocks: u64,
}
pub struct Ran {
pub status: ExitStatus,
pub stdout: Vec<u8>,
pub stderr: Vec<u8>,
pub usage: Usage,
}
impl Ran {
pub fn stdout(&self) -> String {
String::from_utf8_lossy(&self.stdout).into_owned()
}
pub fn stderr(&self) -> String {
String::from_utf8_lossy(&self.stderr).into_owned()
}
pub fn succeeded(&self) -> bool {
self.status.success()
|| (self.usage.answered_s.is_some() && self.status.signal() == Some(libc::SIGTERM))
}
pub fn ok(self, what: &str) -> Self {
assert!(
self.succeeded(),
"{what} failed ({}): stdout={} stderr={}",
self.status,
self.stdout(),
self.stderr()
);
self
}
}
pub fn run(command: &Command, stdin: Option<&[u8]>) -> Ran {
run_with(command, stdin, false)
}
/// `run`, but the child gets SIGTERM once a line reached stdout. a write answers before
/// it waits for its hostname to serve the rev, which a bench knot without tls never does.
pub fn run_until_answered(command: &Command, stdin: Option<&[u8]>) -> Ran {
run_with(command, stdin, true)
}
fn run_with(command: &Command, stdin: Option<&[u8]>, stop_on_answer: bool) -> Ran {
let usage_file = usage_file();
let mut wrapped = wrap(command, &usage_file);
wrapped
.stdin(if stdin.is_some() {
Stdio::piped()
} else {
Stdio::null()
})
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let started = Instant::now();
let mut child = wrapped.spawn().expect("spawn measured child");
let pipe = child.stdin.take();
let input = stdin.map(<[u8]>::to_vec);
let writer = std::thread::spawn(move || {
if let (Some(mut pipe), Some(input)) = (pipe, input) {
let _ = pipe.write_all(&input);
}
});
let stop = stop_on_answer.then_some(child.id());
let stdout = drain(child.stdout.take().unwrap(), started, stop);
let stderr = drain(child.stderr.take().unwrap(), started, None);
let (status, wrapper) = reap(child.id(), started);
writer.join().unwrap();
let (stdout, answered_s) = stdout.join().unwrap();
let mut usage = read_usage(&usage_file).unwrap_or(wrapper);
usage.wall_s = wrapper.wall_s;
usage.answered_s = answered_s;
Ran {
status,
stdout,
stderr: stderr.join().unwrap().0,
usage,
}
}
fn usage_file() -> PathBuf {
static COUNTER: AtomicU64 = AtomicU64::new(0);
std::env::temp_dir().join(format!(
"ks-bench-usage-{}-{}.json",
std::process::id(),
COUNTER.fetch_add(1, Ordering::Relaxed)
))
}
pub fn read_usage(file: &Path) -> Option<Usage> {
let usage = serde_json::from_slice(&std::fs::read(file).ok()?).ok();
let _ = std::fs::remove_file(file);
usage
}
fn wrap(command: &Command, usage_file: &Path) -> Command {
let mut wrapped = Command::new(std::env::current_exe().unwrap());
wrapped
.arg(command.get_program())
.args(command.get_args())
.env(USAGE_ENV, usage_file);
for (name, value) in command.get_envs() {
match value {
Some(value) => wrapped.env(name, value),
None => wrapped.env_remove(name),
};
}
if let Some(dir) = command.get_current_dir() {
wrapped.current_dir(dir);
}
wrapped
}
pub const USAGE_ENV: &str = "KS_BENCH_USAGE";
static WRAPPED: AtomicI32 = AtomicI32::new(0);
extern "C" fn forward_term(_: libc::c_int) {
// SAFETY: kill is async-signal-safe
unsafe { libc::kill(WRAPPED.load(Ordering::Relaxed), libc::SIGTERM) };
}
fn run_measured(mut command: Command, usage_file: &Path) -> ! {
let started = Instant::now();
let child = command
.env_remove(USAGE_ENV)
.spawn()
.expect("spawn measured child");
WRAPPED.store(child.id() as i32, Ordering::Relaxed);
// SAFETY: the handler only forwards the signal with kill
unsafe {
libc::signal(
libc::SIGTERM,
forward_term as *const () as libc::sighandler_t,
)
};
let done = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let sampler = {
let (root, done) = (child.id(), done.clone());
std::thread::spawn(move || resident::peak(root, &done))
};
let (status, mut usage) = reap(child.id(), started);
done.store(true, Ordering::Relaxed);
if let Some((anon, file)) = sampler.join().unwrap() {
usage.peak_anon_bytes = Some(anon);
usage.peak_file_bytes = Some(file);
}
std::fs::write(usage_file, serde_json::to_vec(&usage).unwrap()).unwrap();
if let Some(signal) = status.signal() {
// SAFETY: back to the default action, then die of the child's signal
unsafe {
libc::signal(signal, libc::SIG_DFL);
libc::raise(signal);
}
}
std::process::exit(status.code().unwrap_or(1));
}
pub fn wrapper(usage_file: &Path) -> ! {
let mut args = std::env::args_os().skip(1);
let mut command = Command::new(args.next().expect("a program to measure"));
command.args(args);
run_measured(command, usage_file)
}
fn drain(
mut pipe: impl Read + Send + 'static,
started: Instant,
stop: Option<u32>,
) -> JoinHandle<(Vec<u8>, Option<f64>)> {
std::thread::spawn(move || {
let mut bytes = Vec::new();
let mut answered = None;
let mut chunk = [0; 1 << 16];
while let Ok(n @ 1..) = pipe.read(&mut chunk) {
bytes.extend_from_slice(&chunk[..n]);
if answered.is_none() && bytes.contains(&b'\n') {
answered = Some(started.elapsed().as_secs_f64());
if let Some(pid) = stop {
// SAFETY: a plain signal to our own child, which is unreaped until we return
unsafe { libc::kill(pid as libc::pid_t, libc::SIGTERM) };
}
}
}
(bytes, answered)
})
}
/// `wait4` rather than `Child::wait`, which drops the rusage.
fn reap(pid: u32, started: Instant) -> (ExitStatus, Usage) {
let mut status = 0;
// SAFETY: rusage is plain old data that wait4 fills in
let mut usage: libc::rusage = unsafe { std::mem::zeroed() };
loop {
// SAFETY: pid is our unreaped child, both pointers outlive the call
let reaped = unsafe { libc::wait4(pid as libc::pid_t, &mut status, 0, &mut usage) };
if reaped == pid as libc::pid_t {
break;
}
let error = std::io::Error::last_os_error();
assert_eq!(
error.kind(),
std::io::ErrorKind::Interrupted,
"wait4 {pid}: {error}"
);
}
let seconds = |time: libc::timeval| time.tv_sec as f64 + time.tv_usec as f64 / 1e6;
(
ExitStatus::from_raw(status),
Usage {
wall_s: started.elapsed().as_secs_f64(),
answered_s: None,
user_s: seconds(usage.ru_utime),
sys_s: seconds(usage.ru_stime),
max_rss_bytes: usage.ru_maxrss as u64 * MAXRSS_UNIT,
peak_anon_bytes: None,
peak_file_bytes: None,
out_blocks: usage.ru_oublock as u64,
},
)
}
pub const SSH_USAGE_ENV: &str = "KS_BENCH_SSH_USAGE";
pub fn ssh_stand_in(usage_file: &Path) -> ! {
let original = std::env::args().last().expect("a remote command");
let binary = std::env::var_os("KS_BENCH_BIN").expect("KS_BENCH_BIN");
let config = std::env::var_os("KS_BENCH_CONFIG").expect("KS_BENCH_CONFIG");
let mut command = Command::new(binary);
command
.arg("--config")
.arg(config)
.args(["git-ssh", "--operator"])
.env("SSH_ORIGINAL_COMMAND", original)
.env_remove(SSH_USAGE_ENV);
run_measured(command, usage_file)
}
#[cfg(target_os = "linux")]
mod resident {
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
const SAMPLE: std::time::Duration = std::time::Duration::from_millis(50);
pub fn peak(root: u32, done: &AtomicBool) -> Option<(u64, u64)> {
let (mut anon, mut file) = (0, 0);
while !done.load(Ordering::Relaxed) {
let (now_anon, now_file) = tree(root);
anon = anon.max(now_anon);
file = file.max(now_file);
std::thread::sleep(SAMPLE);
}
Some((anon, file))
}
fn tree(root: u32) -> (u64, u64) {
let mut parents = HashMap::new();
for entry in std::fs::read_dir("/proc").into_iter().flatten().flatten() {
let Some(pid) = entry
.file_name()
.to_str()
.and_then(|n| n.parse::<u32>().ok())
else {
continue;
};
// the command name can hold spaces and parens, which is why the split starts at the last `)`.
let ppid = std::fs::read_to_string(entry.path().join("stat"))
.ok()
.and_then(|stat| {
stat.rsplit_once(')')?
.1
.split_whitespace()
.nth(1)?
.parse()
.ok()
});
if let Some(ppid) = ppid {
parents.insert(pid, ppid);
}
}
let below_root = |mut pid: u32| loop {
if pid == root {
return true;
}
match parents.get(&pid) {
Some(&parent) if parent != 0 => pid = parent,
_ => return false,
}
};
let (mut anon, mut file) = (0, 0);
for &pid in parents.keys().filter(|&&pid| below_root(pid)) {
let Ok(status) = std::fs::read_to_string(format!("/proc/{pid}/status")) else {
continue;
};
for line in status.lines() {
let kib = |line: &str| {
line.split_whitespace()
.nth(1)
.and_then(|n| n.parse::<u64>().ok())
.unwrap_or(0)
* 1024
};
if line.starts_with("RssAnon:") {
anon += kib(line);
} else if line.starts_with("RssFile:") {
file += kib(line);
}
}
}
(anon, file)
}
}
#[cfg(not(target_os = "linux"))]
mod resident {
pub fn peak(_: u32, _: &std::sync::atomic::AtomicBool) -> Option<(u64, u64)> {
None
}
}