use std::collections::VecDeque;
use std::fs::File;
use std::io::{BufRead, BufReader, Read};
use std::path::Path;
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TreeRun {
Passed,
Failed(String),
TimedOut,
Cancelled,
Spawn(String),
}
const TAIL: usize = 64;
const POLL: Duration = Duration::from_millis(20);
const GRACE: Duration = Duration::from_millis(500);
#[cfg(unix)]
const NICENESS: i32 = 10;
#[cfg(unix)]
extern "C" {
#[link_name = "setpriority"]
fn libc_setpriority(which: i32, who: u32, prio: i32) -> i32;
#[link_name = "kill"]
fn libc_kill(pid: i32, sig: i32) -> i32;
#[link_name = "pipe"]
fn libc_pipe(fds: *mut i32) -> i32;
}
#[cfg(unix)]
fn merged_pipe() -> Option<(File, Stdio, Stdio)> {
use std::os::unix::io::FromRawFd;
let mut fds = [0i32; 2];
if unsafe { libc_pipe(fds.as_mut_ptr()) } != 0 {
return None;
}
let (read, write) = unsafe { (File::from_raw_fd(fds[0]), File::from_raw_fd(fds[1])) };
let write2 = write.try_clone().ok()?;
Some((read, Stdio::from(write), Stdio::from(write2)))
}
#[cfg(unix)]
const SIGTERM: i32 = 15;
#[cfg(unix)]
const SIGKILL: i32 = 9;
#[cfg(unix)]
fn signal_group(pgid: u32, sig: i32) -> bool {
unsafe { libc_kill(-(pgid as i32), sig) == 0 }
}
#[cfg(unix)]
pub fn group_alive(pgid: u32) -> bool {
signal_group(pgid, 0)
}
fn kill_tree(child: &mut Child) {
#[cfg(unix)]
{
let pgid = child.id();
signal_group(pgid, SIGTERM);
let until = Instant::now() + GRACE;
while Instant::now() < until {
let _ = child.try_wait();
if !group_alive(pgid) {
break;
}
std::thread::sleep(POLL);
}
if group_alive(pgid) {
signal_group(pgid, SIGKILL);
}
}
#[cfg(not(unix))]
{
let _ = Command::new("taskkill")
.args(["/PID", &child.id().to_string(), "/T", "/F"])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status();
}
let _ = child.kill();
let _ = child.wait();
}
fn drain(
stream: impl Read + Send + 'static,
tail: Arc<Mutex<VecDeque<String>>>,
) -> std::thread::JoinHandle<()> {
std::thread::spawn(move || {
for line in BufReader::new(stream).lines().map_while(Result::ok) {
if let Ok(mut t) = tail.lock() {
if t.len() == TAIL {
t.pop_front();
}
t.push_back(line);
}
}
})
}
fn summary(tail: &Mutex<VecDeque<String>>) -> String {
tail.lock()
.ok()
.and_then(|t| {
t.iter()
.rev()
.map(|l| l.trim())
.find(|l| !l.is_empty())
.map(str::to_string)
})
.unwrap_or_else(|| "no output".to_string())
}
pub fn run(argv: &[String], cwd: &Path, deadline: Instant, cancel: &AtomicBool) -> TreeRun {
run_inner(argv, cwd, deadline, cancel, false).0
}
pub fn run_output(
argv: &[String],
cwd: &Path,
deadline: Instant,
cancel: &AtomicBool,
) -> Result<String, TreeRun> {
match run_inner(argv, cwd, deadline, cancel, true) {
(TreeRun::Passed, out) => Ok(out),
(other, _) => Err(other),
}
}
fn run_inner(
argv: &[String],
cwd: &Path,
deadline: Instant,
cancel: &AtomicBool,
need_output: bool,
) -> (TreeRun, String) {
let Some((program, rest)) = argv.split_first() else {
return (TreeRun::Spawn("empty command".into()), String::new());
};
let mut cmd = Command::new(program);
cmd.args(rest).current_dir(cwd).stdin(Stdio::null());
#[cfg(unix)]
let merged = merged_pipe();
#[cfg(not(unix))]
let merged: Option<(File, Stdio, Stdio)> = None;
let reader = match merged {
Some((read, out, err)) => {
cmd.stdout(out).stderr(err);
Some(read)
}
None => {
cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
None
}
};
crate::hooks::common::strip_git_env(&mut cmd);
#[cfg(unix)]
{
use std::os::unix::process::CommandExt;
cmd.process_group(0);
unsafe {
cmd.pre_exec(|| {
libc_setpriority(0, 0, NICENESS);
Ok(())
});
}
}
let mut child = match cmd.spawn() {
Ok(c) => c,
Err(e) => return (TreeRun::Spawn(format!("{program}: {e}")), String::new()),
};
drop(cmd);
let tail = Arc::new(Mutex::new(VecDeque::with_capacity(TAIL)));
let mut drains = Vec::new();
if let Some(read) = reader {
drains.push(drain(read, Arc::clone(&tail)));
}
if let Some(out) = child.stdout.take() {
drains.push(drain(out, Arc::clone(&tail)));
}
if let Some(err) = child.stderr.take() {
drains.push(drain(err, Arc::clone(&tail)));
}
loop {
match child.try_wait() {
Ok(Some(status)) => {
let until = if need_output {
deadline
} else {
Instant::now() + GRACE
};
while Instant::now() < until
&& !(need_output && cancel.load(Ordering::SeqCst))
&& drains.iter().any(|d| !d.is_finished())
{
std::thread::sleep(Duration::from_millis(5));
}
if need_output && drains.iter().any(|d| !d.is_finished()) {
return (TreeRun::TimedOut, String::new());
}
let out = tail
.lock()
.map(|t| t.iter().cloned().collect::<Vec<_>>().join("\n"))
.unwrap_or_default();
return if status.success() {
(TreeRun::Passed, out)
} else {
(TreeRun::Failed(summary(&tail)), out)
};
}
Ok(None) => {}
Err(e) => {
kill_tree(&mut child);
return (TreeRun::Spawn(format!("{program}: {e}")), String::new());
}
}
if cancel.load(Ordering::SeqCst) {
kill_tree(&mut child);
return (TreeRun::Cancelled, String::new());
}
if Instant::now() >= deadline {
kill_tree(&mut child);
return (TreeRun::TimedOut, String::new());
}
std::thread::sleep(POLL);
}
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
fn sh(script: &str) -> Vec<String> {
vec!["sh".into(), "-c".into(), script.into()]
}
fn soon(ms: u64) -> Instant {
Instant::now() + Duration::from_millis(ms)
}
#[test]
fn tree_run_passed() {
let r = run(
&sh("exit 0"),
Path::new("."),
soon(5_000),
&AtomicBool::new(false),
);
assert_eq!(r, TreeRun::Passed);
}
#[test]
fn tree_run_failed_summarises_the_last_line() {
let r = run(
&sh("echo 'a.ts: bad'; echo '3 problems' >&2; exit 1"),
Path::new("."),
soon(5_000),
&AtomicBool::new(false),
);
assert_eq!(r, TreeRun::Failed("3 problems".into()));
}
#[test]
fn tree_run_spawn_error() {
let r = run(
&["amont-no-such-program-xyz".to_string()],
Path::new("."),
soon(5_000),
&AtomicBool::new(false),
);
assert!(matches!(r, TreeRun::Spawn(_)), "{r:?}");
}
#[test]
fn tree_run_timed_out_leaves_no_child_alive() {
let dir = std::env::temp_dir().join(format!("amont-tree-run-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let pidfile = dir.join("grandchild.pid");
let script = format!("sleep 30 & echo $! > {}; wait", pidfile.display());
let r = run(&sh(&script), &dir, soon(400), &AtomicBool::new(false));
assert_eq!(r, TreeRun::TimedOut);
let pid: i32 = std::fs::read_to_string(&pidfile)
.unwrap()
.trim()
.parse()
.unwrap();
std::thread::sleep(Duration::from_millis(100));
let alive = unsafe { libc_kill(pid, 0) == 0 };
assert!(!alive, "grandchild {pid} survived the deadline");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn tree_run_cancelled() {
let flag = Arc::new(AtomicBool::new(false));
let setter = Arc::clone(&flag);
std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(150));
setter.store(true, Ordering::SeqCst);
});
let r = run(&sh("sleep 30"), Path::new("."), soon(10_000), &flag);
assert_eq!(r, TreeRun::Cancelled);
}
}