use std::fs::File;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use supermachine::internal::{
vmm::pool::WarmPool,
RunOptions,
VmResources,
};
#[derive(Clone, Debug)]
struct Row {
iteration: usize,
worker_restore_us: u128,
reset_vsock_us: u128,
remap_cow_us: u128,
load_meta_us: u128,
restore_snapshot_us: u128,
ram_copy_us: u128,
gic_restore_us: u128,
vcpu_restore_us: u128,
vtimer_offset_us: u128,
mmio_restore_us: u128,
listener_restore_us: u128,
library_done_us: u128,
first_200_us: i128,
port: Option<u16>,
error: Option<String>,
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut snapshot = None;
let mut iterations = 20usize;
let mut warmup = 1usize;
let mut settle_ms = 10u64;
let mut out_path = None;
let mut args = std::env::args().skip(1);
while let Some(arg) = args.next() {
match arg.as_str() {
"--snapshot" => snapshot = args.next(),
"--iterations" => {
iterations = args
.next()
.ok_or("--iterations requires a value")?
.parse()?
}
"--warmup" => warmup = args.next().ok_or("--warmup requires a value")?.parse()?,
"--settle-ms" => {
settle_ms = args.next().ok_or("--settle-ms requires a value")?.parse()?
}
"--out" => out_path = args.next(),
_ => return Err(format!("unknown argument: {arg}").into()),
}
}
let snapshot = snapshot.ok_or(
"usage: warm_pool_bench --snapshot SNAPSHOT [--iterations N] [--warmup N] [--settle-ms MS] [--out PATH]",
)?;
let out_path = out_path.unwrap_or_else(|| "/tmp/supermachine-pool-warm-library.json".to_string());
let resources = VmResources::from_snapshot(&snapshot).with_cow_restore(true);
let pool = WarmPool::start(resources, RunOptions::default())?;
for _ in 0..warmup {
let t0 = Instant::now();
let restored = pool.restore_timeout(&snapshot, Duration::from_secs(5))?;
if let Some(port) = restored.host_port {
let _ = wait_http_200(port, t0, Duration::from_secs(5));
}
if settle_ms > 0 {
std::thread::sleep(Duration::from_millis(settle_ms));
}
}
let mut rows = Vec::with_capacity(iterations);
for i in 1..=iterations {
let t0 = Instant::now();
let restored = pool.restore_timeout(&snapshot, Duration::from_secs(5))?;
let library_done_us = t0.elapsed().as_micros();
let (first_200_us, error) = match restored.host_port {
Some(port) => match wait_http_200(port, t0, Duration::from_secs(5)) {
Some(us) => (us as i128, None),
None => (-1, Some("first_200_timeout".to_string())),
},
None => (-1, Some("missing_host_port".to_string())),
};
println!(
"iter={i:02} worker={}us done={:.3}ms first_200={:.3}ms reset={}us remap={}us meta={}us state={}us ram={}us gic={}us vcpu={}us vtimer={}us mmio={}us listen={}us port={}",
restored.restore_us,
library_done_us as f64 / 1000.0,
first_200_us as f64 / 1000.0,
restored.timings.reset_vsock_us,
restored.timings.remap_cow_us,
restored.timings.load_meta_us,
restored.timings.restore_snapshot_us,
restored.timings.ram_copy_us,
restored.timings.gic_restore_us,
restored.timings.vcpu_restore_us,
restored.timings.vtimer_offset_us,
restored.timings.mmio_restore_us,
restored.timings.listener_restore_us,
restored
.host_port
.map(|p| p.to_string())
.unwrap_or_else(|| "?".to_string())
);
rows.push(Row {
iteration: i,
worker_restore_us: restored.restore_us,
reset_vsock_us: restored.timings.reset_vsock_us,
remap_cow_us: restored.timings.remap_cow_us,
load_meta_us: restored.timings.load_meta_us,
restore_snapshot_us: restored.timings.restore_snapshot_us,
ram_copy_us: restored.timings.ram_copy_us,
gic_restore_us: restored.timings.gic_restore_us,
vcpu_restore_us: restored.timings.vcpu_restore_us,
vtimer_offset_us: restored.timings.vtimer_offset_us,
mmio_restore_us: restored.timings.mmio_restore_us,
listener_restore_us: restored.timings.listener_restore_us,
library_done_us,
first_200_us,
port: restored.host_port,
error,
});
if settle_ms > 0 {
std::thread::sleep(Duration::from_millis(settle_ms));
}
}
let report = pool.shutdown()?;
write_json(
&out_path,
&snapshot,
iterations,
warmup,
settle_ms,
report.warm_restores,
&rows,
)?;
println!("wrote {out_path}");
Ok(())
}
fn wait_http_200(port: u16, t0: Instant, timeout: Duration) -> Option<u128> {
let deadline = Instant::now() + timeout;
while Instant::now() < deadline {
if http_probe(port) {
return Some(t0.elapsed().as_micros());
}
std::thread::sleep(Duration::from_millis(2));
}
None
}
fn http_probe(port: u16) -> bool {
let Ok(mut stream) = TcpStream::connect_timeout(
&std::net::SocketAddr::from(([127, 0, 0, 1], port)),
Duration::from_millis(200),
) else {
return false;
};
let _ = stream.set_read_timeout(Some(Duration::from_millis(200)));
let _ = stream.write_all(b"GET / HTTP/1.1\r\nHost: supermachine\r\nConnection: close\r\n\r\n");
let mut buf = [0u8; 256];
match stream.read(&mut buf) {
Ok(n) => buf[..n].starts_with(b"HTTP/1.1 200") || buf[..n].starts_with(b"HTTP/1.0 200"),
Err(_) => false,
}
}
fn write_json(
out_path: &str,
snapshot: &str,
iterations: usize,
warmup: usize,
settle_ms: u64,
warm_restores: u64,
rows: &[Row],
) -> std::io::Result<()> {
let mut out = File::create(out_path)?;
writeln!(out, "{{")?;
writeln!(out, " \"benchmark\": \"supermachine_pool_warm_library\",")?;
writeln!(out, " \"generated_at\": \"{}\",", generated_at())?;
writeln!(out, " \"snapshot\": \"{}\",", json_escape(snapshot))?;
writeln!(out, " \"iterations\": {iterations},")?;
writeln!(out, " \"warmup\": {warmup},")?;
writeln!(out, " \"settle_ms\": {settle_ms},")?;
writeln!(out, " \"warm_restores\": {warm_restores},")?;
write_summary(&mut out, rows)?;
writeln!(out, " \"rows\": [")?;
for (idx, row) in rows.iter().enumerate() {
if idx > 0 {
writeln!(out, ",")?;
}
write!(
out,
" {{\"iteration\":{},\"worker_restore_us\":{},\"reset_vsock_us\":{},\"remap_cow_us\":{},\"load_meta_us\":{},\"restore_snapshot_us\":{},\"ram_copy_us\":{},\"gic_restore_us\":{},\"vcpu_restore_us\":{},\"vtimer_offset_us\":{},\"mmio_restore_us\":{},\"listener_restore_us\":{},\"library_done_us\":{},\"first_200_us\":{},\"port\":{}",
row.iteration,
row.worker_restore_us,
row.reset_vsock_us,
row.remap_cow_us,
row.load_meta_us,
row.restore_snapshot_us,
row.ram_copy_us,
row.gic_restore_us,
row.vcpu_restore_us,
row.vtimer_offset_us,
row.mmio_restore_us,
row.listener_restore_us,
row.library_done_us,
row.first_200_us,
row.port
.map(|p| p.to_string())
.unwrap_or_else(|| "null".to_string())
)?;
if let Some(error) = row.error.as_ref() {
write!(out, ",\"error\":\"{}\"", json_escape(error))?;
}
write!(out, "}}")?;
}
writeln!(out, "\n ]")?;
writeln!(out, "}}")?;
Ok(())
}
fn write_summary(out: &mut File, rows: &[Row]) -> std::io::Result<()> {
let worker: Vec<u128> = rows.iter().map(|r| r.worker_restore_us).collect();
let done: Vec<u128> = rows.iter().map(|r| r.library_done_us).collect();
let first: Vec<u128> = rows
.iter()
.filter_map(|r| (r.first_200_us >= 0).then_some(r.first_200_us as u128))
.collect();
writeln!(out, " \"summary\": {{")?;
write_percentiles(out, "worker_restore_us", &worker, true)?;
write_percentiles(out, "library_done_us", &done, true)?;
write_percentiles(out, "first_200_us", &first, false)?;
writeln!(out, ",")?;
writeln!(out, " \"phase_us\": {{")?;
write_percentiles(
out,
"reset_vsock_us",
&rows.iter().map(|r| r.reset_vsock_us).collect::<Vec<_>>(),
true,
)?;
write_percentiles(
out,
"remap_cow_us",
&rows.iter().map(|r| r.remap_cow_us).collect::<Vec<_>>(),
true,
)?;
write_percentiles(
out,
"load_meta_us",
&rows.iter().map(|r| r.load_meta_us).collect::<Vec<_>>(),
true,
)?;
write_percentiles(
out,
"restore_snapshot_us",
&rows
.iter()
.map(|r| r.restore_snapshot_us)
.collect::<Vec<_>>(),
false,
)?;
writeln!(out, ",")?;
write_percentiles(
out,
"ram_copy_us",
&rows.iter().map(|r| r.ram_copy_us).collect::<Vec<_>>(),
true,
)?;
write_percentiles(
out,
"gic_restore_us",
&rows.iter().map(|r| r.gic_restore_us).collect::<Vec<_>>(),
true,
)?;
write_percentiles(
out,
"vcpu_restore_us",
&rows.iter().map(|r| r.vcpu_restore_us).collect::<Vec<_>>(),
true,
)?;
write_percentiles(
out,
"vtimer_offset_us",
&rows.iter().map(|r| r.vtimer_offset_us).collect::<Vec<_>>(),
true,
)?;
write_percentiles(
out,
"mmio_restore_us",
&rows.iter().map(|r| r.mmio_restore_us).collect::<Vec<_>>(),
true,
)?;
write_percentiles(
out,
"listener_restore_us",
&rows
.iter()
.map(|r| r.listener_restore_us)
.collect::<Vec<_>>(),
false,
)?;
writeln!(out, " }}")?;
writeln!(out, " }},")?;
Ok(())
}
fn write_percentiles(
out: &mut File,
key: &str,
values: &[u128],
trailing_comma: bool,
) -> std::io::Result<()> {
let comma = if trailing_comma { "," } else { "" };
writeln!(
out,
" \"{key}\": {{\"min\":{},\"p50\":{},\"p95\":{},\"max\":{}}}{comma}",
percentile(values, 0.0),
percentile(values, 0.5),
percentile(values, 0.95),
percentile(values, 1.0)
)
}
fn percentile(values: &[u128], q: f64) -> u128 {
if values.is_empty() {
return 0;
}
let mut sorted = values.to_vec();
sorted.sort_unstable();
let idx = ((sorted.len() - 1) as f64 * q + 0.5) as usize;
sorted[idx]
}
fn generated_at() -> String {
let secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
secs.to_string()
}
fn json_escape(s: &str) -> String {
s.replace('\\', "\\\\").replace('"', "\\\"")
}