use crate::config::Config;
use crate::daemon::log;
use crate::job::JobStatus;
use std::path::Path;
use std::time::{Duration, Instant};
const CLAIM_FILE: &str = "hook.ran";
const LOG_FILE: &str = "hook.log";
const GRACE: Duration = Duration::from_secs(2);
const OUTPUT_LIMIT: u64 = 1 << 20;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Origin {
Supervisor,
Coordinator,
}
impl Origin {
fn as_str(self) -> &'static str {
match self {
Self::Supervisor => "supervisor",
Self::Coordinator => "coordinator",
}
}
}
pub fn fire(origin: Origin, dir: &Path, status: &JobStatus) {
if !status.state.is_terminal() {
return;
}
match Config::load_short() {
Ok(cfg) => fire_with(origin, &cfg, dir, status),
Err(e) => log(&format!(
"qex did not run the stop hook of the job {}: it could not read the \
configuration ({e}). Correct the config file.",
status.id
)),
}
}
fn fire_with(origin: Origin, cfg: &Config, dir: &Path, status: &JobStatus) {
if cfg.hooks.on_stop.is_empty() || !status.state.is_terminal() {
return;
}
if !cfg.hooks.runs_on(status.state) {
return;
}
if !claim(origin, dir, status) {
return;
}
let limit = cfg.hook_timeout().unwrap_or(Duration::from_secs(30));
let verdict = match run(cfg, dir, status, limit) {
Ok(text) => text,
Err(e) => format!(
"did not start: {e}. Test the program `{}` in `[hooks] on_stop` of the config \
file. The job keeps its result.",
cfg.hooks.on_stop[0]
),
};
note(dir, &format!("qex: the stop hook {verdict}"));
log(&format!("the stop hook of the job {} {verdict}", status.id));
}
fn note(dir: &Path, text: &str) {
use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
if let Ok(mut f) = std::fs::OpenOptions::new()
.create(true)
.append(true)
.mode(0o600)
.open(dir.join(LOG_FILE))
{
write!(f, "\n{text}\n").ok();
}
}
pub fn fire_detached(dir: &Path, status: &JobStatus) {
if !status.state.is_terminal() {
return;
}
let dir = dir.to_path_buf();
let status = status.clone();
std::thread::spawn(move || fire(Origin::Coordinator, &dir, &status));
}
fn claim(origin: Origin, dir: &Path, status: &JobStatus) -> bool {
use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
match std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.mode(0o600)
.open(dir.join(CLAIM_FILE))
{
Ok(mut f) => {
writeln!(
f,
"{} {} {} {}",
status.state,
status.id,
crate::sys::now_secs(),
origin.as_str()
)
.ok();
true
}
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => false,
Err(e) => {
log(&format!(
"qex did not run the stop hook of the job {}: it could not write {} ({e})",
status.id,
dir.join(CLAIM_FILE).display()
));
false
}
}
}
fn run(cfg: &Config, dir: &Path, status: &JobStatus, limit: Duration) -> std::io::Result<String> {
use std::os::unix::fs::OpenOptionsExt;
use std::os::unix::process::CommandExt;
let out = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.mode(0o600)
.open(dir.join(LOG_FILE))?;
let (reader, writer) = pipe()?;
let writer_err = writer.try_clone()?;
let mut cmd = std::process::Command::new(&cfg.hooks.on_stop[0]);
cmd.args(&cfg.hooks.on_stop[1..])
.current_dir(match std::path::Path::new(&status.cwd) {
p if !status.cwd.is_empty() && p.is_dir() => p.to_path_buf(),
_ => std::path::PathBuf::from("/"),
})
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::from(writer))
.stderr(std::process::Stdio::from(writer_err));
for (key, value) in variables(dir, status) {
cmd.env(key, value);
}
unsafe {
cmd.pre_exec(|| {
if libc::setpgid(0, 0) == -1 {
return Err(std::io::Error::last_os_error());
}
Ok(())
});
}
let start = Instant::now();
let mut child = cmd.spawn()?;
let pid = child.id() as i32;
drop(cmd);
let cut_flag = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let cut_seen = cut_flag.clone();
let copier = std::thread::spawn(move || capture(reader, out, &cut_seen));
let deadline = start + limit;
let log_path = dir.join(LOG_FILE);
let mut too_slow = false;
let mut too_large = false;
loop {
match child.try_wait() {
Ok(Some(_)) => break,
Ok(None) => {}
Err(e) => {
unsafe {
libc::killpg(pid, libc::SIGKILL);
}
child.wait().ok();
return Ok(format!("gave an error at a test of its state: {e}"));
}
}
if Instant::now() >= deadline {
too_slow = true;
break;
}
if cut_flag.load(std::sync::atomic::Ordering::SeqCst) {
too_large = true;
break;
}
std::thread::sleep(Duration::from_millis(20));
}
if too_slow || too_large {
stop(pid, &mut child);
} else {
child.wait().ok();
}
let wrote = copier.join().unwrap_or(0);
let cut = cut_flag.load(std::sync::atomic::Ordering::SeqCst);
if cut && !too_slow {
return Ok(format!(
"wrote more than {}, so qex stopped reading it and closed the pipe. \
{} holds the first {}. Write less in the hook.",
crate::units::format_size(OUTPUT_LIMIT),
log_path.display(),
crate::units::format_size(wrote)
));
}
if too_slow {
let mut text = format!(
"used more than its time limit of {} and qex stopped it. \
Make the hook faster, or increase `[hooks] timeout`.",
crate::units::format_duration(limit)
);
if cut {
text.push_str(&format!(
" It also wrote more than {}, so qex stopped reading its output.",
crate::units::format_size(OUTPUT_LIMIT)
));
}
return Ok(text);
}
match exit_of(&mut child) {
Some(exit) if exit.success() => Ok(format!(
"ran in {}",
crate::units::format_duration(start.elapsed())
)),
Some(exit) => Ok(format!(
"stopped with {exit}. Read `qex logs {} --hook` for its output.",
status.id
)),
None => Ok("stopped, and qex could not read its result".to_string()),
}
}
fn pipe() -> std::io::Result<(std::fs::File, std::fs::File)> {
use std::os::fd::FromRawFd;
let mut fds = [0 as libc::c_int; 2];
if unsafe { libc::pipe(fds.as_mut_ptr()) } == -1 {
return Err(std::io::Error::last_os_error());
}
unsafe {
Ok((
std::fs::File::from_raw_fd(fds[0]),
std::fs::File::from_raw_fd(fds[1]),
))
}
}
fn capture(
mut reader: std::fs::File,
mut out: std::fs::File,
cut: &std::sync::atomic::AtomicBool,
) -> u64 {
use std::io::{Read, Write};
let mut buf = [0u8; 16 * 1024];
let mut written: u64 = 0;
loop {
let n = match reader.read(&mut buf) {
Ok(0) => break,
Ok(n) => n,
Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(_) => break,
};
let room = OUTPUT_LIMIT.saturating_sub(written) as usize;
let take = room.min(n);
if take > 0 && out.write_all(&buf[..take]).is_err() {
break;
}
written += take as u64;
if take < n || written >= OUTPUT_LIMIT {
cut.store(true, std::sync::atomic::Ordering::SeqCst);
break;
}
}
written
}
fn exit_of(child: &mut std::process::Child) -> Option<std::process::ExitStatus> {
child.try_wait().ok().flatten()
}
fn stop(pid: i32, child: &mut std::process::Child) {
unsafe {
libc::killpg(pid, libc::SIGTERM);
}
std::thread::sleep(GRACE);
unsafe {
libc::killpg(pid, libc::SIGKILL);
}
child.wait().ok();
}
fn variables(dir: &Path, status: &JobStatus) -> Vec<(String, String)> {
let text = |v: Option<i32>| v.map(|n| n.to_string()).unwrap_or_default();
let mut set = vec![
("QEX_JOB_ID".into(), status.id.to_string()),
("QEX_JOB_NAME".into(), crate::job::safe_name(&status.name)),
("QEX_STATE".into(), status.state.to_string()),
("QEX_EXIT_CODE".into(), text(status.exit_code)),
("QEX_SIGNAL".into(), text(status.signal)),
(
"QEX_ELAPSED_SECS".into(),
status
.elapsed()
.map(|d| d.as_secs().to_string())
.unwrap_or_default(),
),
("QEX_CWD".into(), status.cwd.clone()),
("QEX_JOB_DIR".into(), dir.display().to_string()),
("QEX_ATTEMPTS".into(), status.attempts.to_string()),
("QEX_MAX_RSS".into(), status.usage.max_rss.to_string()),
("QEX_TAGS".into(), status.tags.join(" ")),
];
for (_, value) in set.iter_mut() {
*value = crate::job::printable(value);
}
set
}
#[cfg(test)]
mod tests {
use super::*;
use crate::job::JobState;
fn status(state: JobState) -> JobStatus {
let mut s = JobStatus::new(&crate::spec::JobSpec {
id: uuid::Uuid::new_v4(),
name: "build".into(),
cwd: "/".into(),
command: vec!["true".into()],
env: Default::default(),
cpu: 1,
mem: 1 << 30,
timeout: None,
tags: vec!["ci".into()],
priority: 0,
env_capture: crate::config::EnvCapture::None,
claim_source: "explicit".into(),
group: None,
group_name: None,
locks: vec![],
retries: 0,
nice: None,
max_queue_time: None,
dedupe_key: None,
dedupe_window: 0,
learn_key: None,
needs: vec![],
after: vec![],
submitted_at: 0,
});
s.state = state;
s.exit_code = Some(3);
s.started_at = Some(100);
s.finished_at = Some(112);
s
}
struct Temp(std::path::PathBuf);
impl std::ops::Deref for Temp {
type Target = Path;
fn deref(&self) -> &Path {
&self.0
}
}
impl Drop for Temp {
fn drop(&mut self) {
std::fs::remove_dir_all(&self.0).ok();
}
}
fn temp(name: &str) -> Temp {
let dir = std::env::temp_dir().join(format!("qex-hook-{}-{name}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
Temp(dir)
}
fn cfg_with(hook: &str) -> Config {
toml::from_str(hook).unwrap()
}
#[test]
fn the_hook_of_one_job_runs_one_time_only() {
let dir = temp("once");
let mark = dir.join("count");
let cfg = cfg_with(&format!(
"[hooks]\non_stop = [\"sh\", \"-c\", \"echo x >> {}\"]\n",
mark.display()
));
let status = status(JobState::Completed);
fire_with(Origin::Supervisor, &cfg, &dir, &status);
fire_with(Origin::Supervisor, &cfg, &dir, &status);
fire_with(Origin::Supervisor, &cfg, &dir, &status);
let text = std::fs::read_to_string(&mark).unwrap();
assert_eq!(text.lines().count(), 1, "the hook ran more than one time");
assert!(dir.join(CLAIM_FILE).exists());
}
#[test]
fn the_claim_file_names_the_process_that_ran_the_hook() {
for (origin, word) in [
(Origin::Supervisor, "supervisor"),
(Origin::Coordinator, "coordinator"),
] {
let dir = temp(word);
let cfg = cfg_with("[hooks]\non_stop = [\"true\"]\n");
let s = status(JobState::Completed);
fire_with(origin, &cfg, &dir, &s);
let text = std::fs::read_to_string(dir.join(CLAIM_FILE)).unwrap();
assert_eq!(
text.split_whitespace().next_back(),
Some(word),
"the claim file must name the process: {text:?}"
);
assert!(text.starts_with("completed "), "got: {text:?}");
assert!(text.contains(&s.id.to_string()), "got: {text:?}");
}
}
#[test]
fn the_hook_receives_the_id_the_state_and_the_exit_code() {
let dir = temp("env");
let out = dir.join("env.txt");
let cfg = cfg_with(&format!(
"[hooks]\non_stop = [\"sh\", \"-c\", \"env | grep ^QEX_ > {}\"]\n",
out.display()
));
let status = status(JobState::Failed);
fire_with(Origin::Supervisor, &cfg, &dir, &status);
let text = std::fs::read_to_string(&out).unwrap();
assert!(
text.contains(&format!("QEX_JOB_ID={}", status.id)),
"{text}"
);
assert!(text.contains("QEX_JOB_NAME=build"), "{text}");
assert!(text.contains("QEX_STATE=failed"), "{text}");
assert!(text.contains("QEX_EXIT_CODE=3"), "{text}");
assert!(text.contains("QEX_ELAPSED_SECS=12"), "{text}");
assert!(text.contains("QEX_JOB_DIR="), "{text}");
assert!(text.contains("QEX_TAGS=ci"), "{text}");
assert!(text.lines().any(|l| l == "QEX_SIGNAL="), "{text}");
}
#[test]
fn a_job_name_with_shell_characters_does_not_become_a_command() {
let dir = temp("inject");
let mark = dir.join("owned");
let out = dir.join("name.txt");
let cfg = cfg_with(&format!(
"[hooks]\non_stop = [\"sh\", \"-c\", \"printf %s \\\"$QEX_JOB_NAME\\\" > {}\"]\n",
out.display()
));
let mut status = status(JobState::Completed);
status.name = format!("x; touch {}", mark.display());
fire_with(Origin::Supervisor, &cfg, &dir, &status);
assert!(
!mark.exists(),
"a job name became a command; the name must stay in the environment"
);
let got = std::fs::read_to_string(&out).unwrap();
assert!(
got.starts_with("x_touch_"),
"the safe form of the name must arrive: {got}"
);
assert!(
!got.contains(';') && !got.contains(' ') && !got.contains('/'),
"the name must carry no shell character and no path: {got}"
);
}
#[test]
fn the_files_of_the_hook_are_readable_by_the_owner_only() {
use std::os::unix::fs::PermissionsExt;
let mode_of = |path: std::path::PathBuf| {
std::fs::metadata(&path)
.unwrap_or_else(|e| panic!("{} is not there: {e}", path.display()))
.permissions()
.mode()
& 0o777
};
let dir = temp("mode");
let cfg = cfg_with("[hooks]\non_stop = [\"sh\", \"-c\", \"echo a secret\"]\n");
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
for name in [LOG_FILE, CLAIM_FILE] {
let mode = mode_of(dir.join(name));
assert_eq!(
mode, 0o600,
"{name} has the mode {mode:o}, and another user can read it"
);
}
let dir = temp("mode2");
let cfg = cfg_with("[hooks]\non_stop = [\"qex-no-such-program\"]\n");
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
let mode = mode_of(dir.join(LOG_FILE));
assert_eq!(
mode, 0o600,
"the log of a hook that did not start has the mode {mode:o}"
);
}
#[cfg(target_os = "linux")]
#[test]
fn the_hook_reads_no_standard_input() {
let dir = temp("stdin");
let out = dir.join("in.txt");
let cfg = cfg_with(&format!(
"[hooks]\non_stop = [\"sh\", \"-c\", \"readlink /proc/self/fd/0 > {}\"]\n",
out.display()
));
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
assert_eq!(
std::fs::read_to_string(&out).unwrap().trim(),
"/dev/null",
"the standard input of the hook must be /dev/null"
);
let log = std::fs::read_to_string(dir.join(LOG_FILE)).unwrap();
assert!(log.contains("ran in"), "the hook must succeed: {log}");
}
#[test]
fn the_hook_starts_in_the_directory_of_the_job_or_in_the_root() {
let dir = temp("cwd");
let out = dir.join("where.txt");
let cfg = cfg_with(&format!(
"[hooks]\non_stop = [\"sh\", \"-c\", \"pwd > {}\"]\n",
out.display()
));
let mut s = status(JobState::Completed);
s.cwd = dir.display().to_string();
fire_with(Origin::Supervisor, &cfg, &dir, &s);
let got = std::fs::read_to_string(&out).unwrap().trim().to_string();
assert_eq!(
std::fs::canonicalize(&got).unwrap(),
std::fs::canonicalize(&*dir).unwrap(),
"the hook must start in the directory of the job (it said {got})"
);
let second = temp("cwd2");
let mut s = status(JobState::Completed);
s.cwd = second.join("this-directory-is-gone").display().to_string();
fire_with(Origin::Supervisor, &cfg, &second, &s);
assert_eq!(
std::fs::read_to_string(&out).unwrap().trim(),
"/",
"a directory that is gone must not stop the hook"
);
}
#[test]
fn a_hook_that_ignores_the_first_signal_still_stops() {
let dir = temp("stubborn");
let cfg = cfg_with(
"[hooks]\non_stop = [\"sh\", \"-c\", \"trap '' TERM; sleep 60\"]\n\
timeout = \"1s\"\n",
);
let start = Instant::now();
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
let took = start.elapsed();
assert!(
took < Duration::from_secs(20),
"a hook that ignores TERM held the caller for {took:?}; only KILL ends it"
);
let log = std::fs::read_to_string(dir.join(LOG_FILE)).unwrap();
assert!(log.contains("time limit"), "got: {log}");
}
#[test]
fn a_hook_that_hangs_stops_at_its_time_limit() {
let dir = temp("hang");
let cfg = cfg_with("[hooks]\non_stop = [\"sleep\", \"60\"]\ntimeout = \"1s\"\n");
let start = Instant::now();
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
let took = start.elapsed();
assert!(
took < Duration::from_secs(10),
"the hook held the caller for {took:?}"
);
}
#[test]
fn a_hook_that_does_not_exist_is_reported_and_changes_nothing() {
let dir = temp("missing");
let cfg = cfg_with("[hooks]\non_stop = [\"qex-no-such-program\"]\n");
let status = status(JobState::Completed);
fire_with(Origin::Supervisor, &cfg, &dir, &status);
assert!(dir.join(CLAIM_FILE).exists());
}
#[test]
fn the_filter_selects_the_states_that_notify() {
let dir = temp("filter");
let mark = dir.join("ran");
let cfg = cfg_with(&format!(
"[hooks]\non_stop = [\"touch\", \"{}\"]\non_stop_states = [\"failed\"]\n",
mark.display()
));
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
assert!(!mark.exists(), "the filter must stop this state");
assert!(
!dir.join(CLAIM_FILE).exists(),
"a state that the filter stops must not take the claim"
);
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Failed));
assert!(mark.exists(), "the filter must permit this state");
}
#[test]
fn a_control_byte_in_the_data_of_a_job_still_runs_the_hook() {
let dir = temp("nul");
let out = dir.join("name.txt");
let cfg = cfg_with(&format!(
"[hooks]\non_stop = [\"sh\", \"-c\", \
\"printf '%s|%s' \\\"$QEX_JOB_NAME\\\" \\\"$QEX_TAGS\\\" > {}\"]\n",
out.display()
));
let mut status = status(JobState::Completed);
status.name = "a\0b\u{1b}[31mc\nd".to_string();
status.tags = vec!["x\0y".to_string()];
fire_with(Origin::Supervisor, &cfg, &dir, &status);
let text = std::fs::read_to_string(&out).unwrap_or_default();
let (name, tags) = text.split_once('|').unwrap_or(("", ""));
assert_eq!(name, "a_b_31mc_d", "the name must take its safe form");
assert_eq!(
tags, "x y",
"each control byte of a tag must become a space"
);
assert!(
!name.contains('\u{1b}') && !tags.contains('\u{1b}'),
"no value may carry an escape byte to a screen: {text:?}"
);
let log = std::fs::read_to_string(dir.join(LOG_FILE)).unwrap();
assert!(!log.contains("did not start"), "got: {log}");
}
#[test]
fn the_verdict_of_qex_goes_into_the_log_of_the_hook() {
let dir = temp("verdict");
let cfg = cfg_with("[hooks]\non_stop = [\"qex-no-such-program\"]\n");
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
let log = std::fs::read_to_string(dir.join(LOG_FILE)).unwrap();
assert!(log.contains("qex: the stop hook"), "got: {log}");
assert!(log.contains("qex-no-such-program"), "got: {log}");
assert!(
log.contains("on_stop"),
"the remedy must name the field: {log}"
);
}
#[test]
fn a_hook_that_writes_without_a_stop_is_stopped_at_the_cap() {
let dir = temp("flood");
let cfg = cfg_with(
"[hooks]\non_stop = [\"sh\", \"-c\", \
\"head -c 8000000 /dev/zero; sleep 60\"]\ntimeout = \"60s\"\n",
);
let start = Instant::now();
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
assert!(
start.elapsed() < Duration::from_secs(30),
"the size limit must stop the hook before the time limit"
);
let size = std::fs::metadata(dir.join(LOG_FILE)).unwrap().len();
assert!(
size <= OUTPUT_LIMIT + 4096,
"the log of the hook is {size} bytes and the limit is {OUTPUT_LIMIT}"
);
}
#[test]
fn a_hook_that_writes_a_large_file_quickly_also_meets_the_size_limit() {
let dir = temp("fastflood");
let cfg = cfg_with("[hooks]\non_stop = [\"head\", \"-c\", \"3000000\", \"/dev/zero\"]\n");
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Completed));
let size = std::fs::metadata(dir.join(LOG_FILE)).unwrap().len();
assert!(
size <= OUTPUT_LIMIT + 4096,
"the log of the hook is {size} bytes and the limit is {OUTPUT_LIMIT}"
);
let log = std::fs::read_to_string(dir.join(LOG_FILE)).unwrap_or_default();
assert!(
log.contains("wrote more than"),
"the verdict must say that qex stopped reading the hook"
);
let last = log.lines().next_back().unwrap_or("");
assert!(
last.starts_with("qex: "),
"the last line is {} bytes",
last.len()
);
}
#[test]
fn a_job_that_did_not_stop_does_not_run_the_hook() {
let dir = temp("running");
let mark = dir.join("ran");
let cfg = cfg_with(&format!(
"[hooks]\non_stop = [\"touch\", \"{}\"]\n",
mark.display()
));
fire_with(Origin::Supervisor, &cfg, &dir, &status(JobState::Running));
assert!(!mark.exists());
}
}