use std::collections::HashMap;
use std::io::Read;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use portable_pty::{native_pty_system, CommandBuilder, PtySize};
use thiserror::Error;
use tokio::sync::{mpsc, oneshot};
use tokio::task;
use crate::beholders::{registry_with_user_beholders, BeholderSelect};
use crate::store::{RunFilter, StoreError, TaskStore};
use crate::types::{BeholderStatus, Initiator, OutputChunk, RunStatus, Stream, TaskRunId, TaskRunMeta};
const DEFAULT_GRACE: Duration = Duration::from_secs(5);
const READ_BUF_SIZE: usize = 4096;
const SIGTERM: i32 = 15;
const SIGKILL: i32 = 9;
#[derive(Debug, Error)]
pub enum DriverError {
#[error("store: {0}")]
Store(#[from] StoreError),
#[error("pty: {0}")]
Pty(String),
#[error("run not found: {0}")]
NotFound(String),
#[error("io: {0}")]
Io(#[from] std::io::Error),
}
#[derive(Debug, Clone)]
pub struct SpawnOpts {
pub cwd: PathBuf,
pub env: Vec<(String, String)>,
pub label: Option<String>,
pub initiator: Initiator,
pub pty_cols: u16,
pub pty_rows: u16,
pub stdin_enabled: bool,
pub pin: bool,
pub beholder_select: BeholderSelect,
pub verbatim_output: bool,
pub log_fd_enabled: bool,
pub origin: Option<String>,
pub argv: Option<Vec<String>>,
pub pipe: bool,
pub pipefail: bool,
}
const PIPEFAIL_PRELUDE: &str = "if (set -o pipefail) 2>/dev/null; then set -o pipefail; fi\n";
impl Default for SpawnOpts {
fn default() -> Self {
Self {
cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/")),
env: vec![],
label: None,
initiator: Initiator::Human { camp: "local".to_string() },
pty_cols: 80,
pty_rows: 24,
stdin_enabled: false,
pin: false,
beholder_select: BeholderSelect::Auto,
verbatim_output: false,
log_fd_enabled: true,
origin: None,
argv: None,
pipe: false,
pipefail: false,
}
}
}
#[derive(Default)]
pub struct DriverChannels {
pub completion: Option<mpsc::UnboundedSender<(TaskRunId, RunStatus)>>,
pub output: Option<mpsc::UnboundedSender<OutputChunk>>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum StaleRunPolicy {
#[default]
LostOnDisappear,
AdoptLiveHosts { origins: Vec<String> },
}
impl StaleRunPolicy {
fn tombstones(&self, meta: &TaskRunMeta) -> bool {
match self {
StaleRunPolicy::LostOnDisappear => true,
StaleRunPolicy::AdoptLiveHosts { origins } => {
let exempt_origin = origins.is_empty()
|| meta
.origin
.as_deref()
.is_some_and(|o| origins.iter().any(|want| want == o));
if !exempt_origin {
return true;
}
match meta.host_pid {
Some(pid) => !host_process_alive(pid),
None => true,
}
}
}
}
}
#[cfg(unix)]
fn host_process_alive(pid: u32) -> bool {
if pid == 0 {
return false;
}
if pid == std::process::id() {
return true;
}
let rc = unsafe { libc::kill(pid as libc::pid_t, 0) };
rc == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
}
#[cfg(not(unix))]
fn host_process_alive(_pid: u32) -> bool {
false
}
struct RunControl {
kill_tx: mpsc::Sender<KillRequest>,
stdin_tx: Option<mpsc::Sender<Vec<u8>>>,
master: Option<Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>>>,
origin: Option<String>,
last_attached_at: Instant,
}
struct ReaderDone(Option<oneshot::Sender<()>>);
impl Drop for ReaderDone {
fn drop(&mut self) {
if let Some(tx) = self.0.take() {
let _ = tx.send(());
}
}
}
#[derive(Debug)]
struct KillRequest {
signal: i32,
}
#[cfg(unix)]
#[derive(serde::Deserialize)]
struct ShimRecord {
level: String,
target: String,
msg: String,
#[serde(default)]
fields: serde_json::Value,
#[serde(rename = "_lib", default)]
lib: Option<String>,
#[serde(rename = "_lib_ver", default)]
lib_version: Option<String>,
}
#[cfg(unix)]
struct FdCloser(libc::c_int);
#[cfg(unix)]
impl Drop for FdCloser {
fn drop(&mut self) {
unsafe { libc::close(self.0) };
}
}
#[cfg(unix)]
unsafe impl Send for FdCloser {}
pub struct TaskDriver {
store: Arc<TaskStore>,
active: Arc<Mutex<HashMap<String, RunControl>>>,
channels: DriverChannels,
}
impl TaskDriver {
pub async fn new(store: Arc<TaskStore>) -> Result<Self, DriverError> {
Self::with_channels(store, DriverChannels::default()).await
}
pub async fn with_channels(
store: Arc<TaskStore>,
channels: DriverChannels,
) -> Result<Self, DriverError> {
Self::with_config(store, channels, StaleRunPolicy::default()).await
}
pub async fn with_config(
store: Arc<TaskStore>,
channels: DriverChannels,
stale_policy: StaleRunPolicy,
) -> Result<Self, DriverError> {
let stale = store
.list_runs(&RunFilter {
status: Some("running".to_string()),
..Default::default()
})
.await?;
for meta in stale {
if !stale_policy.tombstones(&meta) {
continue;
}
let status = RunStatus::Lost {
reason: "daemon restarted while run was in-flight".to_string(),
};
store.update_status(&meta.id, &status).await?;
if let Some(ref tx) = channels.completion {
let _ = tx.send((meta.id.clone(), status));
}
}
Ok(Self {
store,
active: Arc::new(Mutex::new(HashMap::new())),
channels,
})
}
pub async fn spawn_run(&self, cmd: &str, opts: SpawnOpts) -> Result<TaskRunId, DriverError> {
let id = TaskRunId::new();
let started_at = unix_now_secs();
let started_at_ms: u64 = started_at.saturating_mul(1000);
let user_dir = std::env::var_os("YAH_BEHOLDERS_DIR")
.map(std::path::PathBuf::from)
.or_else(|| {
std::env::var_os("HOME")
.map(|h| std::path::PathBuf::from(h).join(".yah/beholders"))
});
let registry = registry_with_user_beholders(user_dir.as_deref());
let select = if opts.argv.is_some() {
&BeholderSelect::None
} else {
&opts.beholder_select
};
let attach = registry.attach(cmd, select, opts.verbatim_output);
let effective_cmd = match &attach.status.rewrite_added {
Some(added) if !added.is_empty() && !attach.argv.is_empty() => attach.argv.join(" "),
_ => cmd.to_string(),
};
self.store.insert_run(&TaskRunMeta {
id: id.clone(),
command: cmd.to_string(),
cwd: opts.cwd.clone(),
env: opts.env.clone(),
started_at,
status: RunStatus::Running,
label: opts.label.clone(),
initiator: opts.initiator.clone(),
beholder_status: Some(attach.status),
pinned: opts.pin,
origin: opts.origin.clone(),
host_pid: Some(std::process::id()),
}).await?;
let (program, args): (String, Vec<String>) = match opts.argv.as_deref() {
Some([p, rest @ ..]) => (p.clone(), rest.to_vec()),
_ => {
let line = if opts.pipefail {
format!("{PIPEFAIL_PRELUDE}{effective_cmd}")
} else {
effective_cmd.clone()
};
("sh".to_string(), vec!["-c".to_string(), line])
}
};
#[cfg(unix)]
let log_fifo: Option<(libc::c_int, FdCloser, std::path::PathBuf)> = if opts.log_fd_enabled {
let fifo_path = std::env::temp_dir().join(format!("yah-log-{}.fifo", id));
let path_cstr = match std::ffi::CString::new(fifo_path.to_string_lossy().as_bytes()) {
Ok(s) => s,
Err(_) => {
return Err(DriverError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"log FIFO path contained nul byte",
)));
}
};
let mkfifo_ret = unsafe { libc::mkfifo(path_cstr.as_ptr(), 0o600) };
if mkfifo_ret != 0 {
None } else {
let rfd = unsafe {
libc::open(path_cstr.as_ptr(), libc::O_RDONLY | libc::O_NONBLOCK)
};
if rfd < 0 {
let _ = unsafe { libc::unlink(path_cstr.as_ptr()) };
None
} else {
unsafe { libc::fcntl(rfd, libc::F_SETFL, 0) };
let wfd = unsafe {
libc::open(path_cstr.as_ptr(), libc::O_WRONLY)
};
if wfd < 0 {
unsafe { libc::close(rfd) };
let _ = unsafe { libc::unlink(path_cstr.as_ptr()) };
None
} else {
Some((rfd, FdCloser(wfd), fifo_path))
}
}
}
} else {
None
};
#[cfg(unix)]
let fifo_env: Option<(String, String)> = log_fifo
.as_ref()
.map(|(_, _, path)| (id.to_string(), path.to_string_lossy().into_owned()));
#[cfg(not(unix))]
let fifo_env: Option<(String, String)> = None;
let pid: u32;
let reap: Box<dyn FnOnce() -> Option<u32> + Send>;
let stdin_tx: Option<mpsc::Sender<Vec<u8>>>;
let master: Option<Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>>>;
let mut sources: Vec<(Box<dyn Read + Send>, Stream)> = Vec::new();
if opts.pipe {
use std::process::{Command, Stdio};
let mut cmd = Command::new(&program);
cmd.args(&args);
cmd.current_dir(&opts.cwd);
for (k, v) in &opts.env {
cmd.env(k, v);
}
if let Some((run_id_env, fifo_path)) = &fifo_env {
cmd.env("YAH_TASK_RUN", run_id_env);
cmd.env("YAH_LOG_PIPE", fifo_path);
}
cmd.stdout(Stdio::piped());
cmd.stderr(Stdio::piped());
cmd.stdin(if opts.stdin_enabled { Stdio::piped() } else { Stdio::null() });
let mut child = cmd.spawn().map_err(DriverError::Io)?;
pid = child.id();
if let Some(out) = child.stdout.take() {
sources.push((Box::new(out), Stream::Stdout));
}
if let Some(err) = child.stderr.take() {
sources.push((Box::new(err), Stream::Stderr));
}
stdin_tx = child.stdin.take().map(|mut writer| {
let (tx, mut rx) = mpsc::channel::<Vec<u8>>(64);
task::spawn(async move {
use std::io::Write;
while let Some(bytes) = rx.recv().await {
let _ = writer.write_all(&bytes);
let _ = writer.flush();
}
});
tx
});
master = None;
reap = Box::new(move || child.wait().ok().and_then(|s| s.code()).map(|c| c as u32));
} else {
let pty_sys = native_pty_system();
let pair = pty_sys
.openpty(PtySize {
rows: opts.pty_rows,
cols: opts.pty_cols,
pixel_width: 0,
pixel_height: 0,
})
.map_err(|e| DriverError::Pty(e.to_string()))?;
let pty_reader = pair
.master
.try_clone_reader()
.map_err(|e| DriverError::Pty(e.to_string()))?;
sources.push((Box::new(pty_reader), Stream::Stdout));
stdin_tx = if opts.stdin_enabled {
let mut writer = pair
.master
.take_writer()
.map_err(|e| DriverError::Pty(e.to_string()))?;
let (tx, mut rx) = mpsc::channel::<Vec<u8>>(64);
task::spawn(async move {
use std::io::Write;
while let Some(bytes) = rx.recv().await {
let _ = writer.write_all(&bytes);
let _ = writer.flush();
}
});
Some(tx)
} else {
None
};
let mut cb = CommandBuilder::new(&program);
cb.args(&args);
cb.cwd(&opts.cwd);
for (k, v) in &opts.env {
cb.env(k, v);
}
cb.env("TERM", "xterm-256color");
if let Some((run_id_env, fifo_path)) = &fifo_env {
cb.env("YAH_TASK_RUN", run_id_env);
cb.env("YAH_LOG_PIPE", fifo_path);
}
let child = pair
.slave
.spawn_command(cb)
.map_err(|e| DriverError::Pty(e.to_string()))?;
drop(pair.slave);
pid = child.process_id().unwrap_or(0);
let m: Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>> =
Arc::new(Mutex::new(pair.master));
master = Some(Arc::clone(&m));
reap = Box::new(move || {
let mut c = child;
let _m = m; c.wait().ok().map(|s| s.exit_code())
});
}
#[cfg(unix)]
let log_wfd_holder: Option<FdCloser> = if let Some((rfd, wfd, fifo_path)) = log_fifo {
let store_log = Arc::clone(&self.store);
let id_log = id.clone();
let rt = tokio::runtime::Handle::current();
tokio::task::spawn_blocking(move || {
run_log_receiver(rt, store_log, id_log, rfd, fifo_path, started_at_ms);
});
Some(wfd)
} else {
None
};
let (kill_tx, kill_rx) = mpsc::channel::<KillRequest>(4);
let (reader_done_tx, reader_done_rx) = oneshot::channel::<()>();
{
let done = Arc::new(ReaderDone(Some(reader_done_tx)));
let mut beholder = attach.beholder;
for (reader, stream) in sources {
spawn_output_pump(
reader,
stream,
Arc::clone(&self.store),
id.clone(),
started_at_ms,
self.channels.output.clone(),
if stream == Stream::Stdout { beholder.take() } else { None },
Arc::clone(&done),
);
}
}
{
let store_l = Arc::clone(&self.store);
let active_l = Arc::clone(&self.active);
let id_l = id.clone();
let completion_tx_l = self.channels.completion.clone();
#[cfg(unix)]
let wfd_l = log_wfd_holder;
task::spawn(async move {
run_lifecycle(
store_l,
active_l,
id_l,
pid,
reap,
kill_rx,
reader_done_rx,
completion_tx_l,
#[cfg(unix)]
wfd_l,
)
.await;
});
}
self.active
.lock()
.unwrap()
.insert(
id.to_string(),
RunControl {
kill_tx,
stdin_tx,
master,
origin: opts.origin.clone(),
last_attached_at: Instant::now(),
},
);
Ok(id)
}
pub async fn resize_run(
&self,
id: &TaskRunId,
cols: u16,
rows: u16,
) -> Result<(), DriverError> {
let master = self
.active
.lock()
.unwrap()
.get(&id.to_string())
.and_then(|c| c.master.as_ref().map(Arc::clone));
match master {
Some(m) => {
let size = PtySize { rows, cols, pixel_width: 0, pixel_height: 0 };
m.lock()
.unwrap()
.resize(size)
.map_err(|e| DriverError::Pty(e.to_string()))
}
None => Err(DriverError::NotFound(id.to_string())),
}
}
pub fn foreground_pid(&self, id: &TaskRunId) -> Option<u32> {
let master = self
.active
.lock()
.unwrap()
.get(&id.to_string())
.and_then(|c| c.master.as_ref().map(Arc::clone))?;
#[cfg(unix)]
{
let pid = master.lock().unwrap().process_group_leader()?;
u32::try_from(pid).ok()
}
#[cfg(not(unix))]
{
let _ = master;
None
}
}
pub async fn kill_run(&self, id: &TaskRunId, signal: Option<i32>) -> Result<(), DriverError> {
let kill_tx = self
.active
.lock()
.unwrap()
.get(&id.to_string())
.map(|c| c.kill_tx.clone());
match kill_tx {
Some(tx) => tx
.send(KillRequest { signal: signal.unwrap_or(SIGTERM) })
.await
.map_err(|_| DriverError::NotFound(id.to_string())),
None => Err(DriverError::NotFound(id.to_string())),
}
}
pub async fn send_stdin(&self, id: &TaskRunId, bytes: Vec<u8>) -> Result<(), DriverError> {
let stdin_tx = self
.active
.lock()
.unwrap()
.get(&id.to_string())
.and_then(|c| c.stdin_tx.clone());
match stdin_tx {
Some(tx) => tx
.send(bytes)
.await
.map_err(|_| DriverError::NotFound(id.to_string())),
None => Err(DriverError::NotFound(id.to_string())),
}
}
pub fn note_attached(&self, id: &TaskRunId) {
if let Some(control) = self.active.lock().unwrap().get_mut(&id.to_string()) {
control.last_attached_at = Instant::now();
}
}
pub fn attached_age(&self, id: &TaskRunId) -> Option<Duration> {
self.active
.lock()
.unwrap()
.get(&id.to_string())
.map(|c| c.last_attached_at.elapsed())
}
pub async fn reap_unattached(&self, idle: Duration, origins: &[String]) -> Vec<TaskRunId> {
if origins.is_empty() {
return Vec::new();
}
let candidates: Vec<TaskRunId> = {
let active = self.active.lock().unwrap();
active
.iter()
.filter(|(_, c)| {
c.origin
.as_deref()
.is_some_and(|o| origins.iter().any(|want| want == o))
&& c.last_attached_at.elapsed() >= idle
})
.filter_map(|(id, _)| id.parse::<TaskRunId>().ok())
.collect()
};
let mut reaped = Vec::new();
for id in candidates {
if self.kill_run(&id, None).await.is_ok() {
reaped.push(id);
}
}
reaped
}
}
#[cfg(unix)]
fn run_log_receiver(
rt: tokio::runtime::Handle,
store: Arc<TaskStore>,
run_id: TaskRunId,
read_fd: libc::c_int,
fifo_path: std::path::PathBuf,
started_at_ms: u64,
) {
use std::io::BufRead;
use std::os::unix::io::FromRawFd;
let file = unsafe { std::fs::File::from_raw_fd(read_fd) };
let reader = std::io::BufReader::new(file);
for line in reader.lines() {
let line = match line {
Ok(l) => l,
Err(_) => break,
};
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let rec: ShimRecord = match serde_json::from_str(trimmed) {
Ok(r) => r,
Err(_) => continue, };
let level = rec.level.parse::<crate::types::Level>().unwrap_or(crate::types::Level::Info);
let source = crate::types::EventSource::Shim {
lib: rec.lib.unwrap_or_else(|| "unknown".to_string()),
version: rec.lib_version.unwrap_or_else(|| "0.0.0".to_string()),
};
let fields = if rec.fields.is_object() {
rec.fields
} else {
serde_json::Value::Object(Default::default())
};
let offset = elapsed_ms(started_at_ms);
let _ = rt.block_on(store.append_event(
&run_id,
offset,
level,
&rec.target,
&rec.msg,
&fields,
None,
&source,
));
}
let _ = std::fs::remove_file(&fifo_path);
}
#[allow(clippy::too_many_arguments)]
fn spawn_output_pump(
reader: Box<dyn Read + Send>,
stream: Stream,
store: Arc<TaskStore>,
id: TaskRunId,
started_at_ms: u64,
output_tx: Option<mpsc::UnboundedSender<OutputChunk>>,
beholder: Option<Box<dyn crate::beholders::Beholder>>,
done: Arc<ReaderDone>,
) {
let rt = tokio::runtime::Handle::current();
tokio::task::spawn_blocking(move || {
let _done = done;
let mut beholder = beholder;
let mut buf = [0u8; READ_BUF_SIZE];
let mut reader = reader;
loop {
match reader.read(&mut buf) {
Ok(0) | Err(_) => break,
Ok(n) => {
let offset = elapsed_ms(started_at_ms);
let append_res =
rt.block_on(store.append_chunk(&id, offset, stream, &buf[..n]));
if let Ok(seq) = append_res {
let chunk = (output_tx.is_some() || beholder.is_some()).then(|| {
OutputChunk {
run_id: id.clone(),
seq,
offset_ms: offset,
stream,
bytes: buf[..n].to_vec(),
}
});
if let (Some(tx), Some(c)) = (&output_tx, &chunk) {
let _ = tx.send(c.clone());
}
let mut detach_beholder = false;
if let (Some(b), Some(chunk)) = (beholder.as_mut(), &chunk) {
for ev in b.parse_chunk(chunk) {
let _ = rt.block_on(store.append_event(
&ev.run_id,
ev.offset_ms,
ev.level,
&ev.target,
&ev.msg,
&ev.fields,
ev.anchor.as_ref().map(|a| a.seq),
&ev.source,
));
}
if let Some(reason) = b.unknown_format_reason() {
let new_status =
BeholderStatus::unknown_format_with_reason(b.name(), reason);
let _ =
rt.block_on(store.update_beholder_status(&id, &new_status));
detach_beholder = true;
}
}
if detach_beholder {
beholder = None;
}
}
}
}
}
if let Some(ref mut b) = beholder {
let final_offset = elapsed_ms(started_at_ms);
for ev in b.on_done(&id, final_offset) {
let _ = rt.block_on(store.append_event(
&ev.run_id,
ev.offset_ms,
ev.level,
&ev.target,
&ev.msg,
&ev.fields,
ev.anchor.as_ref().map(|a| a.seq),
&ev.source,
));
}
if let Some(reason) = b.unknown_format_reason() {
let new_status = BeholderStatus::unknown_format_with_reason(b.name(), reason);
let _ = rt.block_on(store.update_beholder_status(&id, &new_status));
}
}
});
}
#[allow(clippy::too_many_arguments)]
async fn run_lifecycle(
store: Arc<TaskStore>,
active: Arc<Mutex<HashMap<String, RunControl>>>,
id: TaskRunId,
pid: u32,
reap: Box<dyn FnOnce() -> Option<u32> + Send>,
mut kill_rx: mpsc::Receiver<KillRequest>,
reader_done_rx: oneshot::Receiver<()>,
completion_tx: Option<tokio::sync::mpsc::UnboundedSender<(TaskRunId, RunStatus)>>,
#[cfg(unix)]
_log_wfd: Option<FdCloser>,
) {
let reader_done = async { reader_done_rx.await.ok(); };
tokio::pin!(reader_done);
let sent_signal: Option<i32>;
tokio::select! {
req = kill_rx.recv() => {
match req {
Some(KillRequest { signal }) => {
send_unix_signal(pid, signal);
if signal == SIGKILL {
sent_signal = Some(SIGKILL);
} else {
tokio::select! {
_ = &mut reader_done => {
sent_signal = Some(signal);
}
_ = tokio::time::sleep(DEFAULT_GRACE) => {
send_unix_signal(pid, SIGKILL);
sent_signal = Some(SIGKILL);
}
}
}
}
None => {
send_unix_signal(pid, SIGKILL);
sent_signal = Some(SIGKILL);
}
}
}
_ = &mut reader_done => {
sent_signal = None;
}
}
let exit_code = task::spawn_blocking(reap).await.ok().flatten();
let ended_at = unix_now_secs();
let status = match sent_signal {
Some(sig) => RunStatus::Killed { signal: sig, ended_at },
None => match exit_code {
Some(code) => RunStatus::Done { exit_code: code as i32, ended_at },
None => RunStatus::Lost {
reason: "process exited without an exit code".to_string(),
},
},
};
if let Err(e) = store.update_status(&id, &status).await {
eprintln!("[yah task-runs] failed to record terminal status for run {id}: {e}");
}
if let Some(ref tx) = completion_tx {
let _ = tx.send((id.clone(), status));
}
active.lock().unwrap().remove(&id.to_string());
}
fn send_unix_signal(pid: u32, signal: i32) {
#[cfg(unix)]
unsafe {
libc::kill(pid as libc::pid_t, signal);
}
}
fn unix_now_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
fn elapsed_ms(started_at_ms: u64) -> u32 {
let now_ms = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
now_ms.saturating_sub(started_at_ms).min(u32::MAX as u64) as u32
}
#[cfg(test)]
mod tests {
use super::*;
use crate::store::ChunkFilter;
async fn open_store(dir: &tempfile::TempDir) -> Arc<TaskStore> {
Arc::new(TaskStore::open(&dir.path().join("tr.turso")).await.unwrap())
}
#[tokio::test]
async fn lost_on_disappear_marks_stale_running_runs() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let stale_id = TaskRunId::new();
store
.insert_run(&TaskRunMeta {
id: stale_id.clone(),
command: "sleep 9999".to_string(),
cwd: "/tmp".into(),
env: vec![],
started_at: unix_now_secs() - 60,
status: RunStatus::Running,
label: None,
initiator: Initiator::Human { camp: "test".to_string() },
beholder_status: None,
pinned: false,
origin: None,
host_pid: None,
})
.await
.unwrap();
let _driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let meta = store.get_run(&stale_id).await.unwrap().unwrap();
assert!(
matches!(meta.status, RunStatus::Lost { .. }),
"stale run should be Lost, got {:?}",
meta.status
);
}
#[tokio::test]
async fn a_sweep_tombstone_fires_the_completion_channel() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let stale_id = TaskRunId::new();
store
.insert_run(&TaskRunMeta {
id: stale_id.clone(),
command: "sleep 9999".to_string(),
cwd: "/tmp".into(),
env: vec![],
started_at: unix_now_secs() - 60,
status: RunStatus::Running,
label: None,
initiator: Initiator::Human { camp: "test".to_string() },
beholder_status: None,
pinned: false,
origin: None,
host_pid: None,
})
.await
.unwrap();
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let _driver = TaskDriver::with_channels(
Arc::clone(&store),
DriverChannels { completion: Some(tx), output: None },
)
.await
.unwrap();
let (id, status) = rx.try_recv().expect("the sweep must announce what it tombstoned");
assert_eq!(id, stale_id);
assert!(
matches!(status, RunStatus::Lost { .. }),
"expected Lost, got {status:?}"
);
}
async fn plant_running(
store: &Arc<TaskStore>,
origin: Option<&str>,
host_pid: Option<u32>,
) -> TaskRunId {
let id = TaskRunId::new();
store
.insert_run(&TaskRunMeta {
id: id.clone(),
command: "sleep 9999".to_string(),
cwd: "/tmp".into(),
env: vec![],
started_at: unix_now_secs() - 60,
status: RunStatus::Running,
label: None,
initiator: Initiator::Human {
camp: "test".to_string(),
},
beholder_status: None,
pinned: false,
origin: origin.map(str::to_string),
host_pid,
})
.await
.unwrap();
id
}
async fn is_lost(store: &Arc<TaskStore>, id: &TaskRunId) -> bool {
matches!(
store.get_run(id).await.unwrap().unwrap().status,
RunStatus::Lost { .. }
)
}
fn adopt_terminal() -> StaleRunPolicy {
StaleRunPolicy::AdoptLiveHosts {
origins: vec!["terminal".to_string()],
}
}
#[tokio::test]
async fn a_run_owned_by_a_live_host_survives_a_new_driver() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let id = plant_running(&store, Some("terminal"), Some(std::process::id())).await;
let _driver = TaskDriver::with_config(
Arc::clone(&store),
DriverChannels::default(),
adopt_terminal(),
)
.await
.unwrap();
assert!(
!is_lost(&store, &id).await,
"a terminal run whose owner is alive must stay Running — \
tombstoning it is what made a surviving shell read as dead"
);
}
#[tokio::test]
async fn a_run_whose_host_is_gone_is_still_tombstoned() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let dead_pid = {
let child = std::process::Command::new("true").spawn().unwrap();
let pid = child.id();
let mut child = child;
let _ = child.wait();
pid
};
let id = plant_running(&store, Some("terminal"), Some(dead_pid)).await;
let _driver = TaskDriver::with_config(
Arc::clone(&store),
DriverChannels::default(),
adopt_terminal(),
)
.await
.unwrap();
assert!(
is_lost(&store, &id).await,
"pid {dead_pid} was reaped; its run has no owner left and must be Lost"
);
}
#[tokio::test]
async fn a_non_matching_origin_is_tombstoned_even_with_a_live_host() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let job = plant_running(&store, None, Some(std::process::id())).await;
let other = plant_running(&store, Some("gnome"), Some(std::process::id())).await;
let _driver = TaskDriver::with_config(
Arc::clone(&store),
DriverChannels::default(),
adopt_terminal(),
)
.await
.unwrap();
assert!(is_lost(&store, &job).await, "an origin-less job is not exempt");
assert!(
is_lost(&store, &other).await,
"an origin outside the list is not exempt"
);
}
#[tokio::test]
async fn an_unattributed_run_is_tombstoned() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let id = plant_running(&store, Some("terminal"), None).await;
let _driver = TaskDriver::with_config(
Arc::clone(&store),
DriverChannels::default(),
adopt_terminal(),
)
.await
.unwrap();
assert!(is_lost(&store, &id).await);
}
#[tokio::test]
async fn the_default_policy_is_still_lost_on_disappear() {
assert_eq!(StaleRunPolicy::default(), StaleRunPolicy::LostOnDisappear);
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let id = plant_running(&store, Some("terminal"), Some(std::process::id())).await;
let _driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
assert!(
is_lost(&store, &id).await,
"the default must tombstone regardless of origin or owner liveness"
);
}
#[tokio::test]
async fn spawn_run_stamps_this_process_as_the_owner() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"true",
SpawnOpts {
cwd: "/tmp".into(),
origin: Some("terminal".to_string()),
..Default::default()
},
)
.await
.unwrap();
let meta = store.get_run(&id).await.unwrap().unwrap();
assert_eq!(meta.host_pid, Some(std::process::id()));
}
#[tokio::test]
async fn new_driver_does_not_touch_completed_runs() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let done_id = TaskRunId::new();
store
.insert_run(&TaskRunMeta {
id: done_id.clone(),
command: "true".to_string(),
cwd: "/tmp".into(),
env: vec![],
started_at: unix_now_secs() - 10,
status: RunStatus::Running,
label: None,
initiator: Initiator::Human { camp: "test".to_string() },
beholder_status: None,
pinned: false,
origin: None,
host_pid: None,
})
.await
.unwrap();
store
.update_status(&done_id, &RunStatus::Done { exit_code: 0, ended_at: unix_now_secs() })
.await
.unwrap();
let _driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let meta = store.get_run(&done_id).await.unwrap().unwrap();
assert!(
matches!(meta.status, RunStatus::Done { .. }),
"completed run must not be touched"
);
}
#[tokio::test]
async fn spawn_echo_and_read_chunks() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"echo hello_world",
SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
)
.await
.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(5);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete in time, status={:?}", meta.status);
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
let chunks = store
.get_chunks(&id, &ChunkFilter::default())
.await
.unwrap();
let output: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
let text = String::from_utf8_lossy(&output);
assert!(
text.contains("hello_world"),
"expected 'hello_world' in output, got: {text:?}"
);
let meta = store.get_run(&id).await.unwrap().unwrap();
assert!(
matches!(meta.status, RunStatus::Done { exit_code: 0, .. }),
"expected Done(0), got {:?}",
meta.status
);
}
async fn run_to_completion(
store: &Arc<TaskStore>,
driver: &TaskDriver,
cmd: &str,
opts: SpawnOpts,
) -> Vec<OutputChunk> {
let id = driver.spawn_run(cmd, opts).await.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(10);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete in time, status={:?}", meta.status);
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
store.get_chunks(&id, &ChunkFilter::default()).await.unwrap()
}
fn joined(chunks: &[OutputChunk]) -> Vec<u8> {
chunks.iter().flat_map(|c| c.bytes.clone()).collect()
}
fn joined_stream(chunks: &[OutputChunk], stream: Stream) -> Vec<u8> {
chunks
.iter()
.filter(|c| c.stream == stream)
.flat_map(|c| c.bytes.clone())
.collect()
}
#[tokio::test]
async fn pipe_mode_child_sees_no_tty_on_stdout() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let cmd = "if [ -t 1 ]; then echo TTY; else echo PIPE; fi";
let piped = run_to_completion(
&store,
&driver,
cmd,
SpawnOpts { cwd: "/tmp".into(), pipe: true, ..Default::default() },
)
.await;
assert_eq!(joined(&piped), b"PIPE\n");
let ptied = run_to_completion(
&store,
&driver,
cmd,
SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
)
.await;
assert_eq!(joined(&ptied), b"TTY\r\n");
}
#[tokio::test]
async fn pipe_mode_does_not_translate_newlines() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let piped = run_to_completion(
&store,
&driver,
r"printf 'a\nb\n'",
SpawnOpts { cwd: "/tmp".into(), pipe: true, ..Default::default() },
)
.await;
assert_eq!(joined(&piped), b"a\nb\n");
}
#[tokio::test]
async fn pipe_mode_keeps_stderr_separate_from_stdout() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let cmd = "printf 'to-out\n'; printf 'to-err\n' >&2";
let piped = run_to_completion(
&store,
&driver,
cmd,
SpawnOpts { cwd: "/tmp".into(), pipe: true, ..Default::default() },
)
.await;
assert_eq!(joined_stream(&piped, Stream::Stdout), b"to-out\n");
assert_eq!(joined_stream(&piped, Stream::Stderr), b"to-err\n");
let ptied = run_to_completion(
&store,
&driver,
cmd,
SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
)
.await;
assert!(
joined_stream(&ptied, Stream::Stderr).is_empty(),
"PTY runs have no stderr chunks; that is the behaviour pipe mode exists to fix",
);
}
#[tokio::test]
async fn pipe_mode_drains_both_streams_before_the_run_is_terminal() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let piped = run_to_completion(
&store,
&driver,
"head -c 4096 /dev/zero | tr '\\0' 'x'; head -c 4096 /dev/zero | tr '\\0' 'y' >&2",
SpawnOpts { cwd: "/tmp".into(), pipe: true, ..Default::default() },
)
.await;
assert_eq!(joined_stream(&piped, Stream::Stdout).len(), 4096);
assert_eq!(joined_stream(&piped, Stream::Stderr).len(), 4096);
}
#[tokio::test]
async fn pipe_mode_records_the_childs_exit_code() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"exit 101",
SpawnOpts { cwd: "/tmp".into(), pipe: true, ..Default::default() },
)
.await
.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(10);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
match meta.status {
RunStatus::Done { exit_code, .. } => {
assert_eq!(exit_code, 101);
return;
}
RunStatus::Lost { .. } | RunStatus::Killed { .. } => {
panic!("unexpected terminal status {:?}", meta.status)
}
_ => {}
}
if std::time::Instant::now() > deadline {
panic!("run did not complete in time");
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
}
#[tokio::test]
async fn pipe_mode_has_no_terminal_to_resize_or_read_a_foreground_pid_from() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"sleep 2",
SpawnOpts { cwd: "/tmp".into(), pipe: true, ..Default::default() },
)
.await
.unwrap();
assert!(matches!(
driver.resize_run(&id, 100, 40).await,
Err(DriverError::NotFound(_))
));
assert_eq!(driver.foreground_pid(&id), None);
let _ = driver.kill_run(&id, Some(SIGKILL)).await;
}
async fn await_done(store: &TaskStore, id: &TaskRunId) -> TaskRunMeta {
let deadline = std::time::Instant::now() + Duration::from_secs(5);
loop {
let meta = store.get_run(id).await.unwrap().unwrap();
if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
return meta;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete in time, status={:?}", meta.status);
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
async fn output_of(store: &TaskStore, id: &TaskRunId) -> String {
let chunks = store.get_chunks(id, &ChunkFilter::default()).await.unwrap();
let bytes: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
String::from_utf8_lossy(&bytes).into_owned()
}
#[cfg(unix)]
#[tokio::test]
async fn a_multi_line_command_is_not_flattened_into_one_line() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let id = driver
.spawn_run(
"echo one\necho two",
SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
)
.await
.unwrap();
await_done(&store, &id).await;
let out = output_of(&store, &id).await;
assert!(out.contains("one"), "got: {out:?}");
assert!(
out.contains("two"),
"the second line must have run as its own command; got: {out:?}"
);
assert!(
!out.contains("one echo two"),
"the newline was flattened into a space; got: {out:?}"
);
}
#[cfg(unix)]
#[tokio::test]
async fn a_wrapper_the_caller_wrote_is_not_stripped_from_the_spawned_command() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let id = driver
.spawn_run(
"npx r739s2-nonexistent-tool --version",
SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
)
.await
.unwrap();
let meta = await_done(&store, &id).await;
let out = output_of(&store, &id).await;
assert!(
!matches!(meta.status, RunStatus::Done { exit_code: 0, .. }),
"expected a failure, got {:?} with output {out:?}",
meta.status
);
assert!(
!out.contains("--version: "),
"the wrapper was stripped and the shell tried to run the flag; got: {out:?}"
);
}
#[tokio::test]
async fn explicit_argv_execs_the_program_directly() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"unused-because-argv-wins",
SpawnOpts {
cwd: "/tmp".into(),
argv: Some(vec![
"/bin/sh".into(),
"-c".into(),
"printf 'argv0=%s\\n' \"$0\"".into(),
"direct-exec-marker".into(),
]),
..Default::default()
},
)
.await
.unwrap();
await_done(&store, &id).await;
let text = output_of(&store, &id).await;
assert!(
text.contains("argv0=direct-exec-marker"),
"argv should have been exec'd verbatim, got: {text:?}"
);
}
#[tokio::test]
async fn explicit_argv_still_records_the_requested_command() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"$SHELL",
SpawnOpts {
cwd: "/tmp".into(),
argv: Some(vec!["/bin/sh".into(), "-c".into(), "true".into()]),
..Default::default()
},
)
.await
.unwrap();
let meta = await_done(&store, &id).await;
assert_eq!(meta.command, "$SHELL");
assert!(
matches!(meta.status, RunStatus::Done { exit_code: 0, .. }),
"expected Done(0), got {:?}",
meta.status
);
}
#[tokio::test]
async fn pipefail_reports_the_failing_stage_and_posix_reports_the_last_one() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let line = "(exit 101) | tail -1";
let posix = driver
.spawn_run(
line,
SpawnOpts { cwd: "/tmp".into(), pipefail: false, ..Default::default() },
)
.await
.unwrap();
let meta = await_done(&store, &posix).await;
assert!(
matches!(meta.status, RunStatus::Done { exit_code: 0, .. }),
"POSIX pipeline status is the LAST stage's — expected Done(0), got {:?}",
meta.status
);
let failing = driver
.spawn_run(
line,
SpawnOpts { cwd: "/tmp".into(), pipefail: true, ..Default::default() },
)
.await
.unwrap();
let meta = await_done(&store, &failing).await;
assert!(
matches!(meta.status, RunStatus::Done { exit_code: 101, .. }),
"pipefail must surface the producer's 101, got {:?}",
meta.status
);
}
#[tokio::test]
async fn pipefail_does_not_leak_into_the_recorded_command() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"echo recorded-verbatim | cat",
SpawnOpts { cwd: "/tmp".into(), pipefail: true, ..Default::default() },
)
.await
.unwrap();
let meta = await_done(&store, &id).await;
assert_eq!(meta.command, "echo recorded-verbatim | cat");
assert!(
!meta.command.contains("pipefail"),
"the prelude leaked into the recorded command: {:?}",
meta.command
);
}
#[tokio::test]
async fn the_pipefail_probe_never_costs_the_command_that_follows_it() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"echo probe-survived",
SpawnOpts { cwd: "/tmp".into(), pipefail: true, ..Default::default() },
)
.await
.unwrap();
let meta = await_done(&store, &id).await;
let text = output_of(&store, &id).await;
assert!(
matches!(meta.status, RunStatus::Done { exit_code: 0, .. }),
"expected Done(0), got {:?}",
meta.status
);
assert!(text.contains("probe-survived"), "command did not run, got: {text:?}");
assert!(
!text.contains("pipefail"),
"the probe printed a diagnostic into the build's own output: {text:?}"
);
}
#[tokio::test]
async fn empty_argv_falls_back_to_the_shell_path() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"echo empty_argv_fallback",
SpawnOpts { cwd: "/tmp".into(), argv: Some(vec![]), ..Default::default() },
)
.await
.unwrap();
await_done(&store, &id).await;
let text = output_of(&store, &id).await;
assert!(
text.contains("empty_argv_fallback"),
"empty argv must not spawn nothing, got: {text:?}"
);
}
#[tokio::test]
async fn spawn_failing_command_records_nonzero_exit() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = driver
.spawn_run(
"exit 42",
SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
)
.await
.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(5);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
match meta.status {
RunStatus::Done { exit_code, .. } => {
assert_ne!(exit_code, 0, "exit 42 should produce a non-zero exit code");
}
other => panic!("unexpected status: {other:?}"),
}
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete in time");
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
#[cfg(unix)]
#[tokio::test]
async fn kill_with_sigterm_transitions_to_killed() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let id = driver
.spawn_run(
"sleep 60",
SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(100)).await;
driver.kill_run(&id, Some(SIGTERM)).await.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(10);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if matches!(meta.status, RunStatus::Killed { .. } | RunStatus::Lost { .. }) {
assert!(
matches!(meta.status, RunStatus::Killed { .. }),
"expected Killed, got {:?}",
meta.status
);
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not become Killed in time, status={:?}", meta.status);
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
#[cfg(unix)]
#[tokio::test]
async fn kill_run_returns_not_found_after_exit() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let id = driver
.spawn_run(
"echo done",
SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
)
.await
.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(5);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete");
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
let result = driver.kill_run(&id, None).await;
assert!(
matches!(result, Err(DriverError::NotFound(_))),
"expected NotFound, got {result:?}"
);
}
#[cfg(unix)]
#[tokio::test]
async fn stdin_send_reaches_child() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let id = driver
.spawn_run(
"read line && echo got_$line",
SpawnOpts {
cwd: "/tmp".into(),
stdin_enabled: true,
..Default::default()
},
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(150)).await;
driver.send_stdin(&id, b"hello\n".to_vec()).await.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(5);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete after stdin input");
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
let chunks = store.get_chunks(&id, &ChunkFilter::default()).await.unwrap();
let raw: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
let text = String::from_utf8_lossy(&raw);
assert!(
text.contains("got_hello"),
"expected 'got_hello' in output, got: {text:?}"
);
}
#[tokio::test]
async fn resize_run_changes_geometry_the_child_sees() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let id = driver
.spawn_run(
"read line && stty size",
SpawnOpts {
cwd: "/tmp".into(),
stdin_enabled: true,
..Default::default()
},
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(150)).await;
driver.resize_run(&id, 120, 40).await.unwrap();
driver.send_stdin(&id, b"go\n".to_vec()).await.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(5);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete after stdin input");
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
let chunks = store.get_chunks(&id, &ChunkFilter::default()).await.unwrap();
let raw: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
let text = String::from_utf8_lossy(&raw);
assert!(
text.contains("40 120"),
"expected resized geometry '40 120' in output, got: {text:?}"
);
}
#[tokio::test]
async fn resize_run_returns_not_found_after_exit() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let id = driver
.spawn_run("true", SpawnOpts { cwd: "/tmp".into(), ..Default::default() })
.await
.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(5);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not exit");
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
assert!(matches!(
driver.resize_run(&id, 100, 30).await,
Err(DriverError::NotFound(_))
));
}
#[cfg(unix)]
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn log_pipe_events_land_in_store() {
use crate::store::EventFilter;
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let cmd = r#"printf '{"level":"warn","target":"test.shim","msg":"hello-from-pipe","fields":{"x":42},"_lib":"test-shim","_lib_ver":"0.1.0"}\n' >> "$YAH_LOG_PIPE""#;
let id = driver
.spawn_run(cmd, SpawnOpts { cwd: "/tmp".into(), ..Default::default() })
.await
.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(20);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete in time");
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
tokio::time::sleep(Duration::from_millis(500)).await;
let events = store.query_events(&id, &EventFilter::default()).await.unwrap();
assert!(
!events.is_empty(),
"expected at least one shim event, got none"
);
let ev = events.iter().find(|e| e.target == "test.shim");
let ev = ev.expect("event with target 'test.shim' not found");
assert_eq!(ev.msg, "hello-from-pipe");
assert_eq!(ev.level, crate::types::Level::Warn);
assert!(
matches!(&ev.source, crate::types::EventSource::Shim { lib, .. } if lib == "test-shim"),
"unexpected source: {:?}",
ev.source
);
assert_eq!(ev.fields.get("x"), Some(&serde_json::json!(42)));
}
#[cfg(unix)]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn log_pipe_disabled_produces_no_events() {
use crate::store::EventFilter;
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
let cmd = r#"[ -n "$YAH_LOG_PIPE" ] && printf '{"level":"info","target":"t","msg":"m","fields":{}}\n' >> "$YAH_LOG_PIPE" || true"#;
let id = driver
.spawn_run(
cmd,
SpawnOpts { cwd: "/tmp".into(), log_fd_enabled: false, ..Default::default() },
)
.await
.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(5);
loop {
let meta = store.get_run(&id).await.unwrap().unwrap();
if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
break;
}
if std::time::Instant::now() > deadline {
panic!("run did not complete");
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
tokio::time::sleep(Duration::from_millis(100)).await;
let events = store.query_events(&id, &EventFilter::default()).await.unwrap();
assert!(
events.is_empty(),
"expected no shim events when log_fd_enabled=false, got {}",
events.len()
);
}
const BUILD_RUN: &str = "build-run";
fn opted_in() -> Vec<String> {
vec![BUILD_RUN.to_string()]
}
async fn spawn_long_run(driver: &TaskDriver, origin: &str) -> TaskRunId {
driver
.spawn_run(
"sleep 30",
SpawnOpts {
cwd: "/tmp".into(),
origin: Some(origin.to_string()),
..Default::default()
},
)
.await
.unwrap()
}
async fn await_status(
store: &Arc<TaskStore>,
id: &TaskRunId,
want: fn(&RunStatus) -> bool,
) -> RunStatus {
let deadline = std::time::Instant::now() + Duration::from_secs(10);
loop {
let status = store.get_run(id).await.unwrap().unwrap().status;
if want(&status) {
return status;
}
if std::time::Instant::now() > deadline {
panic!("run never reached the expected status, last={status:?}");
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
}
#[tokio::test]
async fn an_unpolled_run_of_an_opted_in_origin_is_reaped() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = spawn_long_run(&driver, BUILD_RUN).await;
tokio::time::sleep(Duration::from_millis(300)).await;
let reaped = driver
.reap_unattached(Duration::from_millis(200), &opted_in())
.await;
assert_eq!(reaped, vec![id.clone()], "the unattached run should be reaped");
let status = await_status(&store, &id, |s| {
matches!(s, RunStatus::Killed { .. } | RunStatus::Done { .. })
})
.await;
assert!(
matches!(status, RunStatus::Killed { .. }),
"a reaped run ends Killed, got {status:?}"
);
}
#[tokio::test]
async fn a_run_a_client_is_still_polling_is_never_reaped() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = spawn_long_run(&driver, BUILD_RUN).await;
for _ in 0..6 {
tokio::time::sleep(Duration::from_millis(100)).await;
driver.note_attached(&id);
let reaped = driver
.reap_unattached(Duration::from_millis(200), &opted_in())
.await;
assert!(reaped.is_empty(), "a polled run must survive, reaped {reaped:?}");
}
assert!(
matches!(
store.get_run(&id).await.unwrap().unwrap().status,
RunStatus::Running
),
"the polled run should still be running"
);
tokio::time::sleep(Duration::from_millis(300)).await;
let reaped = driver
.reap_unattached(Duration::from_millis(200), &opted_in())
.await;
assert_eq!(reaped, vec![id], "a run that stopped being polled is reapable");
}
#[tokio::test]
async fn a_terminal_tile_is_never_reaped_however_long_it_idles() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = spawn_long_run(&driver, "terminal").await;
tokio::time::sleep(Duration::from_millis(300)).await;
for _ in 0..3 {
let reaped = driver.reap_unattached(Duration::ZERO, &opted_in()).await;
assert!(
reaped.is_empty(),
"a terminal tile is outside the opted-in origins, reaped {reaped:?}"
);
tokio::time::sleep(Duration::from_millis(50)).await;
}
assert!(
matches!(
store.get_run(&id).await.unwrap().unwrap().status,
RunStatus::Running
),
"the terminal run must still be running"
);
driver.kill_run(&id, Some(SIGKILL)).await.unwrap();
}
#[tokio::test]
async fn an_origin_less_run_and_an_empty_opt_in_list_reap_nothing() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let plain = driver
.spawn_run("sleep 30", SpawnOpts { cwd: "/tmp".into(), ..Default::default() })
.await
.unwrap();
let build = spawn_long_run(&driver, BUILD_RUN).await;
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(
driver.reap_unattached(Duration::ZERO, &[]).await.is_empty(),
"an empty opt-in list must reap nothing, not everything"
);
assert_eq!(
driver.reap_unattached(Duration::ZERO, &opted_in()).await,
vec![build],
"only the opted-in origin is reapable"
);
assert!(
matches!(
store.get_run(&plain).await.unwrap().unwrap().status,
RunStatus::Running
),
"the origin-less run must be untouched"
);
driver.kill_run(&plain, Some(SIGKILL)).await.unwrap();
}
#[tokio::test]
async fn attached_age_resets_on_a_poll_and_is_none_for_a_foreign_run() {
let dir = tempfile::tempdir().unwrap();
let store = open_store(&dir).await;
let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
let id = spawn_long_run(&driver, BUILD_RUN).await;
tokio::time::sleep(Duration::from_millis(150)).await;
let aged = driver.attached_age(&id).expect("driver owns this run");
assert!(aged >= Duration::from_millis(100), "age should have grown, got {aged:?}");
driver.note_attached(&id);
let fresh = driver.attached_age(&id).unwrap();
assert!(fresh < aged, "a poll resets the age: {fresh:?} vs {aged:?}");
assert!(
driver.attached_age(&TaskRunId::new()).is_none(),
"a run this driver does not own has no attachment age"
);
driver.kill_run(&id, Some(SIGKILL)).await.unwrap();
}
}