1use std::collections::HashMap;
62use std::io::Read;
63use std::path::PathBuf;
64use std::sync::{Arc, Mutex};
65use std::time::{Duration, SystemTime, UNIX_EPOCH};
66
67use portable_pty::{native_pty_system, CommandBuilder, PtySize};
68use thiserror::Error;
69use tokio::sync::{mpsc, oneshot};
70use tokio::task;
71
72use crate::beholders::{registry_with_user_beholders, BeholderSelect};
73use crate::store::{RunFilter, StoreError, TaskStore};
74use crate::types::{BeholderStatus, Initiator, OutputChunk, RunStatus, Stream, TaskRunId, TaskRunMeta};
75
76const DEFAULT_GRACE: Duration = Duration::from_secs(5);
77const READ_BUF_SIZE: usize = 4096;
78const SIGTERM: i32 = 15;
79const SIGKILL: i32 = 9;
80
81#[derive(Debug, Error)]
84pub enum DriverError {
85 #[error("store: {0}")]
86 Store(#[from] StoreError),
87 #[error("pty: {0}")]
88 Pty(String),
89 #[error("run not found: {0}")]
90 NotFound(String),
91 #[error("io: {0}")]
92 Io(#[from] std::io::Error),
93}
94
95#[derive(Debug, Clone)]
99pub struct SpawnOpts {
100 pub cwd: PathBuf,
101 pub env: Vec<(String, String)>,
103 pub label: Option<String>,
104 pub initiator: Initiator,
105 pub pty_cols: u16,
107 pub pty_rows: u16,
109 pub stdin_enabled: bool,
111 pub pin: bool,
113 pub beholder_select: BeholderSelect,
115 pub tty_attached: bool,
118 pub log_fd_enabled: bool,
122 pub origin: Option<String>,
125}
126
127impl Default for SpawnOpts {
128 fn default() -> Self {
129 Self {
130 cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/")),
131 env: vec![],
132 label: None,
133 initiator: Initiator::Human { camp: "local".to_string() },
134 pty_cols: 80,
135 pty_rows: 24,
136 stdin_enabled: false,
137 pin: false,
138 beholder_select: BeholderSelect::Auto,
139 tty_attached: false,
140 log_fd_enabled: true,
141 origin: None,
142 }
143 }
144}
145
146#[derive(Default)]
151pub struct DriverChannels {
152 pub completion: Option<mpsc::UnboundedSender<(TaskRunId, RunStatus)>>,
155 pub output: Option<mpsc::UnboundedSender<OutputChunk>>,
162}
163
164struct RunControl {
167 kill_tx: mpsc::Sender<KillRequest>,
168 stdin_tx: Option<mpsc::Sender<Vec<u8>>>,
169 master: Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>>,
173}
174
175#[derive(Debug)]
176struct KillRequest {
177 signal: i32,
178}
179
180#[cfg(unix)]
187#[derive(serde::Deserialize)]
188struct ShimRecord {
189 level: String,
190 target: String,
191 msg: String,
192 #[serde(default)]
193 fields: serde_json::Value,
194 #[serde(rename = "_lib", default)]
197 lib: Option<String>,
198 #[serde(rename = "_lib_ver", default)]
200 lib_version: Option<String>,
201}
202
203#[cfg(unix)]
210struct FdCloser(libc::c_int);
211
212#[cfg(unix)]
213impl Drop for FdCloser {
214 fn drop(&mut self) {
215 unsafe { libc::close(self.0) };
216 }
217}
218
219#[cfg(unix)]
222unsafe impl Send for FdCloser {}
223
224pub struct TaskDriver {
230 store: Arc<TaskStore>,
231 active: Arc<Mutex<HashMap<String, RunControl>>>,
232 channels: DriverChannels,
234}
235
236impl TaskDriver {
237 pub async fn new(store: Arc<TaskStore>) -> Result<Self, DriverError> {
242 Self::with_channels(store, DriverChannels::default()).await
243 }
244
245 pub async fn with_channels(
248 store: Arc<TaskStore>,
249 channels: DriverChannels,
250 ) -> Result<Self, DriverError> {
251 let stale = store.list_runs(&RunFilter {
252 status: Some("running".to_string()),
253 ..Default::default()
254 }).await?;
255 for meta in stale {
256 store.update_status(
257 &meta.id,
258 &RunStatus::Lost {
259 reason: "daemon restarted while run was in-flight".to_string(),
260 },
261 ).await?;
262 }
263 Ok(Self {
264 store,
265 active: Arc::new(Mutex::new(HashMap::new())),
266 channels,
267 })
268 }
269
270 pub async fn spawn_run(&self, cmd: &str, opts: SpawnOpts) -> Result<TaskRunId, DriverError> {
283 let id = TaskRunId::new();
284 let started_at = unix_now_secs();
285 let started_at_ms: u64 = started_at.saturating_mul(1000);
286
287 let user_dir = std::env::var_os("YAH_BEHOLDERS_DIR")
290 .map(std::path::PathBuf::from)
291 .or_else(|| {
292 std::env::var_os("HOME")
293 .map(|h| std::path::PathBuf::from(h).join(".yah/beholders"))
294 });
295 let registry = registry_with_user_beholders(user_dir.as_deref());
296 let attach = registry.attach(cmd, &opts.beholder_select, opts.tty_attached);
297 let effective_cmd = if attach.argv.is_empty() {
299 cmd.to_string()
300 } else {
301 attach.argv.join(" ")
302 };
303
304 self.store.insert_run(&TaskRunMeta {
305 id: id.clone(),
306 command: cmd.to_string(),
307 cwd: opts.cwd.clone(),
308 env: opts.env.clone(),
309 started_at,
310 status: RunStatus::Running,
311 label: opts.label.clone(),
312 initiator: opts.initiator.clone(),
313 beholder_status: Some(attach.status),
314 pinned: opts.pin,
315 origin: opts.origin.clone(),
316 }).await?;
317
318 let pty_sys = native_pty_system();
320 let pair = pty_sys
321 .openpty(PtySize {
322 rows: opts.pty_rows,
323 cols: opts.pty_cols,
324 pixel_width: 0,
325 pixel_height: 0,
326 })
327 .map_err(|e| DriverError::Pty(e.to_string()))?;
328
329 let pty_reader = pair
331 .master
332 .try_clone_reader()
333 .map_err(|e| DriverError::Pty(e.to_string()))?;
334
335 let stdin_tx: Option<mpsc::Sender<Vec<u8>>> = if opts.stdin_enabled {
337 let mut writer = pair
338 .master
339 .take_writer()
340 .map_err(|e| DriverError::Pty(e.to_string()))?;
341 let (tx, mut rx) = mpsc::channel::<Vec<u8>>(64);
342 task::spawn(async move {
343 use std::io::Write;
344 while let Some(bytes) = rx.recv().await {
345 let _ = writer.write_all(&bytes);
346 let _ = writer.flush();
347 }
348 });
349 Some(tx)
350 } else {
351 None
352 };
353
354 #[cfg(unix)]
367 let log_fifo: Option<(libc::c_int, FdCloser, std::path::PathBuf)> = if opts.log_fd_enabled {
368 let fifo_path = std::env::temp_dir().join(format!("yah-log-{}.fifo", id));
369 let path_cstr = match std::ffi::CString::new(fifo_path.to_string_lossy().as_bytes()) {
370 Ok(s) => s,
371 Err(_) => {
372 return Err(DriverError::Io(std::io::Error::new(
374 std::io::ErrorKind::InvalidInput,
375 "log FIFO path contained nul byte",
376 )));
377 }
378 };
379 let mkfifo_ret = unsafe { libc::mkfifo(path_cstr.as_ptr(), 0o600) };
380 if mkfifo_ret != 0 {
381 None } else {
383 let rfd = unsafe {
385 libc::open(path_cstr.as_ptr(), libc::O_RDONLY | libc::O_NONBLOCK)
386 };
387 if rfd < 0 {
388 let _ = unsafe { libc::unlink(path_cstr.as_ptr()) };
389 None
390 } else {
391 unsafe { libc::fcntl(rfd, libc::F_SETFL, 0) };
393 let wfd = unsafe {
395 libc::open(path_cstr.as_ptr(), libc::O_WRONLY)
396 };
397 if wfd < 0 {
398 unsafe { libc::close(rfd) };
399 let _ = unsafe { libc::unlink(path_cstr.as_ptr()) };
400 None
401 } else {
402 Some((rfd, FdCloser(wfd), fifo_path))
403 }
404 }
405 }
406 } else {
407 None
408 };
409
410 let mut cb = CommandBuilder::new("sh");
412 cb.args(["-c", &effective_cmd]);
413 cb.cwd(&opts.cwd);
414 for (k, v) in &opts.env {
415 cb.env(k, v);
416 }
417 cb.env("TERM", "xterm-256color");
418
419 #[cfg(unix)]
421 if let Some((_, _, ref fifo_path)) = log_fifo {
422 cb.env("YAH_TASK_RUN", id.to_string());
423 cb.env("YAH_LOG_PIPE", fifo_path.to_string_lossy().as_ref());
424 }
425
426 let child = pair
427 .slave
428 .spawn_command(cb)
429 .map_err(|e| DriverError::Pty(e.to_string()))?;
430 drop(pair.slave);
432
433 let master: Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>> =
436 Arc::new(Mutex::new(pair.master));
437
438 let pid = child.process_id().unwrap_or(0);
439
440 #[cfg(unix)]
448 let log_wfd_holder: Option<FdCloser> = if let Some((rfd, wfd, fifo_path)) = log_fifo {
449 let store_log = Arc::clone(&self.store);
450 let id_log = id.clone();
451 let rt = tokio::runtime::Handle::current();
452 tokio::task::spawn_blocking(move || {
455 run_log_receiver(rt, store_log, id_log, rfd, fifo_path, started_at_ms);
456 });
457 Some(wfd)
458 } else {
459 None
460 };
461
462 let (kill_tx, kill_rx) = mpsc::channel::<KillRequest>(4);
464 let (reader_done_tx, reader_done_rx) = oneshot::channel::<()>();
465
466 {
469 let store_r = Arc::clone(&self.store);
470 let id_r = id.clone();
471 let mut beholder = attach.beholder;
472 let output_tx = self.channels.output.clone();
473 let rt = tokio::runtime::Handle::current();
474 tokio::task::spawn_blocking(move || {
475 let mut buf = [0u8; READ_BUF_SIZE];
476 let mut reader = pty_reader;
477 loop {
478 match reader.read(&mut buf) {
479 Ok(0) | Err(_) => break,
480 Ok(n) => {
481 let offset = elapsed_ms(started_at_ms);
482 let append_res = rt.block_on(store_r.append_chunk(
483 &id_r,
484 offset,
485 Stream::Stdout,
486 &buf[..n],
487 ));
488 if let Ok(seq) = append_res {
489 let chunk = (output_tx.is_some() || beholder.is_some()).then(|| {
493 OutputChunk {
494 run_id: id_r.clone(),
495 seq,
496 offset_ms: offset,
497 stream: Stream::Stdout,
498 bytes: buf[..n].to_vec(),
499 }
500 });
501 if let (Some(tx), Some(c)) = (&output_tx, &chunk) {
505 let _ = tx.send(c.clone());
506 }
507 let mut detach_beholder = false;
508 if let (Some(b), Some(chunk)) = (beholder.as_mut(), &chunk) {
509 for ev in b.parse_chunk(chunk) {
510 let _ = rt.block_on(store_r.append_event(
511 &ev.run_id,
512 ev.offset_ms,
513 ev.level,
514 &ev.target,
515 &ev.msg,
516 &ev.fields,
517 ev.anchor.as_ref().map(|a| a.seq),
518 &ev.source,
519 ));
520 }
521 if let Some(reason) = b.unknown_format_reason() {
522 let new_status = BeholderStatus::unknown_format_with_reason(
523 b.name(),
524 reason,
525 );
526 let _ = rt.block_on(
527 store_r.update_beholder_status(&id_r, &new_status),
528 );
529 detach_beholder = true;
530 }
531 }
532 if detach_beholder {
533 beholder = None;
534 }
535 }
536 }
537 }
538 }
539 if let Some(ref mut b) = beholder {
540 let final_offset = elapsed_ms(started_at_ms);
541 for ev in b.on_done(&id_r, final_offset) {
542 let _ = rt.block_on(store_r.append_event(
543 &ev.run_id,
544 ev.offset_ms,
545 ev.level,
546 &ev.target,
547 &ev.msg,
548 &ev.fields,
549 ev.anchor.as_ref().map(|a| a.seq),
550 &ev.source,
551 ));
552 }
553 if let Some(reason) = b.unknown_format_reason() {
554 let new_status = BeholderStatus::unknown_format_with_reason(b.name(), reason);
555 let _ = rt.block_on(store_r.update_beholder_status(&id_r, &new_status));
556 }
557 }
558 let _ = reader_done_tx.send(());
559 });
560 }
561
562 {
566 let store_l = Arc::clone(&self.store);
567 let active_l = Arc::clone(&self.active);
568 let id_l = id.clone();
569 let master_l = Arc::clone(&master);
570 let completion_tx_l = self.channels.completion.clone();
571 #[cfg(unix)]
572 let wfd_l = log_wfd_holder;
573 task::spawn(async move {
574 run_lifecycle(
575 store_l,
576 active_l,
577 id_l,
578 pid,
579 child,
580 master_l,
581 kill_rx,
582 reader_done_rx,
583 completion_tx_l,
584 #[cfg(unix)]
585 wfd_l,
586 )
587 .await;
588 });
589 }
590
591 self.active
592 .lock()
593 .unwrap()
594 .insert(id.to_string(), RunControl { kill_tx, stdin_tx, master });
595
596 Ok(id)
597 }
598
599 pub async fn resize_run(
606 &self,
607 id: &TaskRunId,
608 cols: u16,
609 rows: u16,
610 ) -> Result<(), DriverError> {
611 let master = self
612 .active
613 .lock()
614 .unwrap()
615 .get(&id.to_string())
616 .map(|c| Arc::clone(&c.master));
617
618 match master {
619 Some(m) => {
620 let size = PtySize { rows, cols, pixel_width: 0, pixel_height: 0 };
621 m.lock()
622 .unwrap()
623 .resize(size)
624 .map_err(|e| DriverError::Pty(e.to_string()))
625 }
626 None => Err(DriverError::NotFound(id.to_string())),
627 }
628 }
629
630 pub async fn kill_run(&self, id: &TaskRunId, signal: Option<i32>) -> Result<(), DriverError> {
637 let kill_tx = self
638 .active
639 .lock()
640 .unwrap()
641 .get(&id.to_string())
642 .map(|c| c.kill_tx.clone());
643
644 match kill_tx {
645 Some(tx) => tx
646 .send(KillRequest { signal: signal.unwrap_or(SIGTERM) })
647 .await
648 .map_err(|_| DriverError::NotFound(id.to_string())),
649 None => Err(DriverError::NotFound(id.to_string())),
650 }
651 }
652
653 pub async fn send_stdin(&self, id: &TaskRunId, bytes: Vec<u8>) -> Result<(), DriverError> {
655 let stdin_tx = self
656 .active
657 .lock()
658 .unwrap()
659 .get(&id.to_string())
660 .and_then(|c| c.stdin_tx.clone());
661
662 match stdin_tx {
663 Some(tx) => tx
664 .send(bytes)
665 .await
666 .map_err(|_| DriverError::NotFound(id.to_string())),
667 None => Err(DriverError::NotFound(id.to_string())),
668 }
669 }
670}
671
672#[cfg(unix)]
681fn run_log_receiver(
682 rt: tokio::runtime::Handle,
683 store: Arc<TaskStore>,
684 run_id: TaskRunId,
685 read_fd: libc::c_int,
686 fifo_path: std::path::PathBuf,
687 started_at_ms: u64,
688) {
689 use std::io::BufRead;
690 use std::os::unix::io::FromRawFd;
691
692 let file = unsafe { std::fs::File::from_raw_fd(read_fd) };
695 let reader = std::io::BufReader::new(file);
696
697 for line in reader.lines() {
698 let line = match line {
699 Ok(l) => l,
700 Err(_) => break,
701 };
702 let trimmed = line.trim();
703 if trimmed.is_empty() {
704 continue;
705 }
706 let rec: ShimRecord = match serde_json::from_str(trimmed) {
707 Ok(r) => r,
708 Err(_) => continue, };
710 let level = rec.level.parse::<crate::types::Level>().unwrap_or(crate::types::Level::Info);
711 let source = crate::types::EventSource::Shim {
712 lib: rec.lib.unwrap_or_else(|| "unknown".to_string()),
713 version: rec.lib_version.unwrap_or_else(|| "0.0.0".to_string()),
714 };
715 let fields = if rec.fields.is_object() {
716 rec.fields
717 } else {
718 serde_json::Value::Object(Default::default())
719 };
720 let offset = elapsed_ms(started_at_ms);
721 let _ = rt.block_on(store.append_event(
722 &run_id,
723 offset,
724 level,
725 &rec.target,
726 &rec.msg,
727 &fields,
728 None,
729 &source,
730 ));
731 }
732
733 let _ = std::fs::remove_file(&fifo_path);
735}
736
737async fn run_lifecycle(
740 store: Arc<TaskStore>,
741 active: Arc<Mutex<HashMap<String, RunControl>>>,
742 id: TaskRunId,
743 pid: u32,
744 child: Box<dyn portable_pty::Child + Send>,
745 master: Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>>,
746 mut kill_rx: mpsc::Receiver<KillRequest>,
747 reader_done_rx: oneshot::Receiver<()>,
748 completion_tx: Option<tokio::sync::mpsc::UnboundedSender<(TaskRunId, RunStatus)>>,
749 #[cfg(unix)]
753 _log_wfd: Option<FdCloser>,
754) {
755 let reader_done = async { reader_done_rx.await.ok(); };
758 tokio::pin!(reader_done);
759
760 let sent_signal: Option<i32>;
761
762 tokio::select! {
763 req = kill_rx.recv() => {
764 match req {
765 Some(KillRequest { signal }) => {
766 send_unix_signal(pid, signal);
767 if signal == SIGKILL {
768 sent_signal = Some(SIGKILL);
769 } else {
770 tokio::select! {
772 _ = &mut reader_done => {
773 sent_signal = Some(signal);
775 }
776 _ = tokio::time::sleep(DEFAULT_GRACE) => {
777 send_unix_signal(pid, SIGKILL);
779 sent_signal = Some(SIGKILL);
780 }
781 }
782 }
783 }
784 None => {
786 send_unix_signal(pid, SIGKILL);
787 sent_signal = Some(SIGKILL);
788 }
789 }
790 }
791 _ = &mut reader_done => {
792 sent_signal = None;
793 }
794 }
795
796 let exit_code = task::spawn_blocking(move || {
801 let mut c = child;
802 let _m = master; c.wait().ok().map(|s| s.exit_code())
804 })
805 .await
806 .ok()
807 .flatten();
808
809 let ended_at = unix_now_secs();
810 let status = match sent_signal {
811 Some(sig) => RunStatus::Killed { signal: sig, ended_at },
812 None => match exit_code {
813 Some(code) => RunStatus::Done { exit_code: code as i32, ended_at },
814 None => RunStatus::Lost {
815 reason: "process exited without an exit code".to_string(),
816 },
817 },
818 };
819
820 let _ = store.update_status(&id, &status).await;
821 if let Some(ref tx) = completion_tx {
822 let _ = tx.send((id.clone(), status));
823 }
824 active.lock().unwrap().remove(&id.to_string());
825}
826
827fn send_unix_signal(pid: u32, signal: i32) {
830 #[cfg(unix)]
831 unsafe {
832 libc::kill(pid as libc::pid_t, signal);
833 }
834 }
836
837fn unix_now_secs() -> u64 {
838 SystemTime::now()
839 .duration_since(UNIX_EPOCH)
840 .unwrap_or_default()
841 .as_secs()
842}
843
844fn elapsed_ms(started_at_ms: u64) -> u32 {
845 let now_ms = SystemTime::now()
846 .duration_since(UNIX_EPOCH)
847 .unwrap_or_default()
848 .as_millis() as u64;
849 now_ms.saturating_sub(started_at_ms).min(u32::MAX as u64) as u32
850}
851
852#[cfg(test)]
855mod tests {
856 use super::*;
857 use crate::store::ChunkFilter;
858
859 async fn open_store(dir: &tempfile::TempDir) -> Arc<TaskStore> {
860 Arc::new(TaskStore::open(&dir.path().join("tr.turso")).await.unwrap())
861 }
862
863 #[tokio::test]
866 async fn lost_on_disappear_marks_stale_running_runs() {
867 let dir = tempfile::tempdir().unwrap();
868 let store = open_store(&dir).await;
869
870 let stale_id = TaskRunId::new();
872 store
873 .insert_run(&TaskRunMeta {
874 id: stale_id.clone(),
875 command: "sleep 9999".to_string(),
876 cwd: "/tmp".into(),
877 env: vec![],
878 started_at: unix_now_secs() - 60,
879 status: RunStatus::Running,
880 label: None,
881 initiator: Initiator::Human { camp: "test".to_string() },
882 beholder_status: None,
883 pinned: false,
884 origin: None,
885 })
886 .await
887 .unwrap();
888
889 let _driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
891
892 let meta = store.get_run(&stale_id).await.unwrap().unwrap();
893 assert!(
894 matches!(meta.status, RunStatus::Lost { .. }),
895 "stale run should be Lost, got {:?}",
896 meta.status
897 );
898 }
899
900 #[tokio::test]
901 async fn new_driver_does_not_touch_completed_runs() {
902 let dir = tempfile::tempdir().unwrap();
903 let store = open_store(&dir).await;
904
905 let done_id = TaskRunId::new();
906 store
907 .insert_run(&TaskRunMeta {
908 id: done_id.clone(),
909 command: "true".to_string(),
910 cwd: "/tmp".into(),
911 env: vec![],
912 started_at: unix_now_secs() - 10,
913 status: RunStatus::Running,
914 label: None,
915 initiator: Initiator::Human { camp: "test".to_string() },
916 beholder_status: None,
917 pinned: false,
918 origin: None,
919 })
920 .await
921 .unwrap();
922 store
923 .update_status(&done_id, &RunStatus::Done { exit_code: 0, ended_at: unix_now_secs() })
924 .await
925 .unwrap();
926
927 let _driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
928
929 let meta = store.get_run(&done_id).await.unwrap().unwrap();
930 assert!(
931 matches!(meta.status, RunStatus::Done { .. }),
932 "completed run must not be touched"
933 );
934 }
935
936 #[tokio::test]
939 async fn spawn_echo_and_read_chunks() {
940 let dir = tempfile::tempdir().unwrap();
941 let store = open_store(&dir).await;
942 let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
943
944 let id = driver
945 .spawn_run(
946 "echo hello_world",
947 SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
948 )
949 .await
950 .unwrap();
951
952 let deadline = std::time::Instant::now() + Duration::from_secs(5);
954 loop {
955 let meta = store.get_run(&id).await.unwrap().unwrap();
956 if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
957 break;
958 }
959 if std::time::Instant::now() > deadline {
960 panic!("run did not complete in time, status={:?}", meta.status);
961 }
962 tokio::time::sleep(Duration::from_millis(50)).await;
963 }
964
965 let chunks = store
967 .get_chunks(&id, &ChunkFilter::default())
968 .await
969 .unwrap();
970 let output: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
971 let text = String::from_utf8_lossy(&output);
972 assert!(
973 text.contains("hello_world"),
974 "expected 'hello_world' in output, got: {text:?}"
975 );
976
977 let meta = store.get_run(&id).await.unwrap().unwrap();
978 assert!(
979 matches!(meta.status, RunStatus::Done { exit_code: 0, .. }),
980 "expected Done(0), got {:?}",
981 meta.status
982 );
983 }
984
985 #[tokio::test]
986 async fn spawn_failing_command_records_nonzero_exit() {
987 let dir = tempfile::tempdir().unwrap();
988 let store = open_store(&dir).await;
989 let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
990
991 let id = driver
992 .spawn_run(
993 "exit 42",
994 SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
995 )
996 .await
997 .unwrap();
998
999 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1000 loop {
1001 let meta = store.get_run(&id).await.unwrap().unwrap();
1002 if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
1003 match meta.status {
1004 RunStatus::Done { exit_code, .. } => {
1005 assert_ne!(exit_code, 0, "exit 42 should produce a non-zero exit code");
1006 }
1007 other => panic!("unexpected status: {other:?}"),
1008 }
1009 break;
1010 }
1011 if std::time::Instant::now() > deadline {
1012 panic!("run did not complete in time");
1013 }
1014 tokio::time::sleep(Duration::from_millis(50)).await;
1015 }
1016 }
1017
1018 #[cfg(unix)]
1021 #[tokio::test]
1022 async fn kill_with_sigterm_transitions_to_killed() {
1023 let dir = tempfile::tempdir().unwrap();
1024 let store = open_store(&dir).await;
1025 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1026
1027 let id = driver
1028 .spawn_run(
1029 "sleep 60",
1030 SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
1031 )
1032 .await
1033 .unwrap();
1034
1035 tokio::time::sleep(Duration::from_millis(100)).await;
1037
1038 driver.kill_run(&id, Some(SIGTERM)).await.unwrap();
1039
1040 let deadline = std::time::Instant::now() + Duration::from_secs(10);
1041 loop {
1042 let meta = store.get_run(&id).await.unwrap().unwrap();
1043 if matches!(meta.status, RunStatus::Killed { .. } | RunStatus::Lost { .. }) {
1044 assert!(
1045 matches!(meta.status, RunStatus::Killed { .. }),
1046 "expected Killed, got {:?}",
1047 meta.status
1048 );
1049 break;
1050 }
1051 if std::time::Instant::now() > deadline {
1052 panic!("run did not become Killed in time, status={:?}", meta.status);
1053 }
1054 tokio::time::sleep(Duration::from_millis(50)).await;
1055 }
1056 }
1057
1058 #[cfg(unix)]
1059 #[tokio::test]
1060 async fn kill_run_returns_not_found_after_exit() {
1061 let dir = tempfile::tempdir().unwrap();
1062 let store = open_store(&dir).await;
1063 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1064
1065 let id = driver
1066 .spawn_run(
1067 "echo done",
1068 SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
1069 )
1070 .await
1071 .unwrap();
1072
1073 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1075 loop {
1076 let meta = store.get_run(&id).await.unwrap().unwrap();
1077 if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
1078 break;
1079 }
1080 if std::time::Instant::now() > deadline {
1081 panic!("run did not complete");
1082 }
1083 tokio::time::sleep(Duration::from_millis(50)).await;
1084 }
1085
1086 let result = driver.kill_run(&id, None).await;
1088 assert!(
1089 matches!(result, Err(DriverError::NotFound(_))),
1090 "expected NotFound, got {result:?}"
1091 );
1092 }
1093
1094 #[cfg(unix)]
1097 #[tokio::test]
1098 async fn stdin_send_reaches_child() {
1099 let dir = tempfile::tempdir().unwrap();
1100 let store = open_store(&dir).await;
1101 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1102
1103 let id = driver
1105 .spawn_run(
1106 "read line && echo got_$line",
1107 SpawnOpts {
1108 cwd: "/tmp".into(),
1109 stdin_enabled: true,
1110 ..Default::default()
1111 },
1112 )
1113 .await
1114 .unwrap();
1115
1116 tokio::time::sleep(Duration::from_millis(150)).await;
1117 driver.send_stdin(&id, b"hello\n".to_vec()).await.unwrap();
1118
1119 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1120 loop {
1121 let meta = store.get_run(&id).await.unwrap().unwrap();
1122 if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
1123 break;
1124 }
1125 if std::time::Instant::now() > deadline {
1126 panic!("run did not complete after stdin input");
1127 }
1128 tokio::time::sleep(Duration::from_millis(50)).await;
1129 }
1130
1131 let chunks = store.get_chunks(&id, &ChunkFilter::default()).await.unwrap();
1132 let raw: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
1133 let text = String::from_utf8_lossy(&raw);
1134 assert!(
1135 text.contains("got_hello"),
1136 "expected 'got_hello' in output, got: {text:?}"
1137 );
1138 }
1139
1140 #[cfg(unix)]
1148 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1149 async fn log_pipe_events_land_in_store() {
1150 use crate::store::EventFilter;
1151
1152 let dir = tempfile::tempdir().unwrap();
1153 let store = open_store(&dir).await;
1154 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1155
1156 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""#;
1159
1160 let id = driver
1161 .spawn_run(cmd, SpawnOpts { cwd: "/tmp".into(), ..Default::default() })
1162 .await
1163 .unwrap();
1164
1165 let deadline = std::time::Instant::now() + Duration::from_secs(20);
1169 loop {
1170 let meta = store.get_run(&id).await.unwrap().unwrap();
1171 if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
1172 break;
1173 }
1174 if std::time::Instant::now() > deadline {
1175 panic!("run did not complete in time");
1176 }
1177 tokio::time::sleep(Duration::from_millis(50)).await;
1178 }
1179
1180 tokio::time::sleep(Duration::from_millis(500)).await;
1183
1184 let events = store.query_events(&id, &EventFilter::default()).await.unwrap();
1185 assert!(
1186 !events.is_empty(),
1187 "expected at least one shim event, got none"
1188 );
1189 let ev = events.iter().find(|e| e.target == "test.shim");
1190 let ev = ev.expect("event with target 'test.shim' not found");
1191 assert_eq!(ev.msg, "hello-from-pipe");
1192 assert_eq!(ev.level, crate::types::Level::Warn);
1193 assert!(
1194 matches!(&ev.source, crate::types::EventSource::Shim { lib, .. } if lib == "test-shim"),
1195 "unexpected source: {:?}",
1196 ev.source
1197 );
1198 assert_eq!(ev.fields.get("x"), Some(&serde_json::json!(42)));
1199 }
1200
1201 #[cfg(unix)]
1204 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1205 async fn log_pipe_disabled_produces_no_events() {
1206 use crate::store::EventFilter;
1207
1208 let dir = tempfile::tempdir().unwrap();
1209 let store = open_store(&dir).await;
1210 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1211
1212 let cmd = r#"[ -n "$YAH_LOG_PIPE" ] && printf '{"level":"info","target":"t","msg":"m","fields":{}}\n' >> "$YAH_LOG_PIPE" || true"#;
1215
1216 let id = driver
1217 .spawn_run(
1218 cmd,
1219 SpawnOpts { cwd: "/tmp".into(), log_fd_enabled: false, ..Default::default() },
1220 )
1221 .await
1222 .unwrap();
1223
1224 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1225 loop {
1226 let meta = store.get_run(&id).await.unwrap().unwrap();
1227 if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
1228 break;
1229 }
1230 if std::time::Instant::now() > deadline {
1231 panic!("run did not complete");
1232 }
1233 tokio::time::sleep(Duration::from_millis(50)).await;
1234 }
1235
1236 tokio::time::sleep(Duration::from_millis(100)).await;
1237
1238 let events = store.query_events(&id, &EventFilter::default()).await.unwrap();
1239 assert!(
1240 events.is_empty(),
1241 "expected no shim events when log_fd_enabled=false, got {}",
1242 events.len()
1243 );
1244 }
1245}