use std::ffi::OsStr;
use std::fs::File;
use std::io::{self, Read};
use std::path::Path;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};
const WINDOW: usize = 65_536;
pub const SPAWN_BUDGET: Duration = Duration::from_secs(2);
pub fn warm_executable(path: &Path) -> io::Result<()> {
let mut file = File::open(path)?;
let metadata = file.metadata()?;
if metadata.is_dir() {
return Err(io::Error::new(
io::ErrorKind::IsADirectory,
"warmup path is a directory",
));
}
let mut remaining = metadata.len();
let mut buffer = vec![0u8; WINDOW];
while remaining > 0 {
let want = WINDOW.min(remaining as usize);
let got = file.read(&mut buffer[..want])?;
if got == 0 {
break;
}
remaining -= got as u64;
}
Ok(())
}
pub fn warm_spawn(path: &Path, args: &[impl AsRef<OsStr>]) -> io::Result<()> {
warm_spawn_until(path, args, Instant::now() + SPAWN_BUDGET)
}
pub fn warm_spawn_until(
path: &Path,
args: &[impl AsRef<OsStr>],
deadline: Instant,
) -> io::Result<()> {
let mut child = Command::new(path)
.args(args)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()?;
loop {
match child.try_wait()? {
Some(status) if status.success() => return Ok(()),
Some(status) => {
return Err(io::Error::other(format!(
"warmup {} exited {status}",
path.display()
)));
}
None if Instant::now() >= deadline => {
let _ = child.kill();
let _ = child.wait();
return Err(io::Error::new(
io::ErrorKind::TimedOut,
format!("warmup {} exceeded its budget", path.display()),
));
}
None => {
let slice = Duration::from_millis(1)
.min(deadline.saturating_duration_since(Instant::now()));
if slice.is_zero() {
continue;
}
std::thread::sleep(slice);
}
}
}
}
pub(crate) fn prefault_all(paths: impl IntoIterator<Item = impl AsRef<Path>>) {
for path in paths {
let _ = warm_executable(path.as_ref());
}
}