Skip to main content

tatara_engine/drivers/
exec.rs

1use anyhow::{bail, Context, Result};
2use async_trait::async_trait;
3use chrono::Utc;
4use std::path::Path;
5use std::process::Stdio;
6use std::time::Duration;
7use tokio::process::Command;
8use tokio::sync::mpsc;
9use tracing::{debug, warn};
10
11use tatara_core::domain::allocation::TaskRunState;
12use tatara_core::domain::job::{DriverType, Task, TaskConfig};
13
14use super::{Driver, LogEntry, TaskHandle};
15
16pub struct ExecDriver;
17
18#[async_trait]
19impl Driver for ExecDriver {
20    fn name(&self) -> &str {
21        "exec"
22    }
23
24    async fn available(&self) -> bool {
25        true
26    }
27
28    async fn start(&self, task: &Task, alloc_dir: &Path) -> Result<TaskHandle> {
29        let (command, args, working_dir) = match &task.config {
30            TaskConfig::Exec {
31                command,
32                args,
33                working_dir,
34            } => (command.clone(), args.clone(), working_dir.clone()),
35            _ => bail!("ExecDriver received non-exec task config"),
36        };
37
38        let log_dir = alloc_dir.join(&task.name);
39        tokio::fs::create_dir_all(&log_dir).await?;
40
41        let stdout_path = log_dir.join("stdout.log");
42        let stderr_path = log_dir.join("stderr.log");
43
44        let stdout_file =
45            std::fs::File::create(&stdout_path).context("Failed to create stdout log")?;
46        let stderr_file =
47            std::fs::File::create(&stderr_path).context("Failed to create stderr log")?;
48
49        let mut cmd = Command::new(&command);
50        cmd.args(&args)
51            .envs(&task.env)
52            .stdout(Stdio::from(stdout_file))
53            .stderr(Stdio::from(stderr_file))
54            .kill_on_drop(false);
55
56        if let Some(dir) = &working_dir {
57            cmd.current_dir(dir);
58        }
59
60        let child = cmd.spawn().context("Failed to spawn process")?;
61        let pid = child.id();
62
63        debug!(
64            task = %task.name,
65            command = %command,
66            pid = ?pid,
67            "Exec driver started process"
68        );
69
70        // Detach — the process runs independently. We track by PID.
71        std::mem::forget(child);
72
73        Ok(TaskHandle {
74            driver: DriverType::Exec,
75            pid,
76            container_id: None,
77            started_at: Utc::now(),
78        })
79    }
80
81    async fn stop(&self, handle: &TaskHandle, timeout: Duration) -> Result<()> {
82        let pid = handle.pid.context("No PID in exec task handle")?;
83        send_signal_and_wait(pid, timeout).await
84    }
85
86    async fn status(&self, handle: &TaskHandle) -> Result<TaskRunState> {
87        let pid = handle.pid.context("No PID in exec task handle")?;
88
89        if is_process_alive(pid) {
90            Ok(TaskRunState::Running)
91        } else {
92            Ok(TaskRunState::Dead)
93        }
94    }
95
96    async fn logs(&self, _handle: &TaskHandle) -> Result<mpsc::Receiver<LogEntry>> {
97        let (tx, rx) = mpsc::channel(256);
98        // Logs are read via LogCollector, not streamed from the driver.
99        tokio::spawn(async move {
100            let _ = tx;
101        });
102        Ok(rx)
103    }
104}
105
106#[cfg(unix)]
107pub(crate) async fn send_signal_and_wait(pid: u32, timeout: Duration) -> Result<()> {
108    use nix::sys::signal::{kill, Signal};
109    use nix::unistd::Pid;
110
111    let nix_pid = Pid::from_raw(pid as i32);
112
113    // Send SIGTERM
114    let _ = kill(nix_pid, Signal::SIGTERM);
115
116    // Wait for graceful shutdown
117    let deadline = tokio::time::Instant::now() + timeout;
118    loop {
119        if !is_process_alive(pid) {
120            debug!(pid, "Process terminated gracefully");
121            return Ok(());
122        }
123        if tokio::time::Instant::now() >= deadline {
124            break;
125        }
126        tokio::time::sleep(Duration::from_millis(100)).await;
127    }
128
129    // Force kill
130    warn!(pid, "Process did not terminate gracefully, sending SIGKILL");
131    let _ = kill(nix_pid, Signal::SIGKILL);
132    Ok(())
133}
134
135#[cfg(not(unix))]
136pub(crate) async fn send_signal_and_wait(_pid: u32, _timeout: Duration) -> Result<()> {
137    anyhow::bail!("Process signal management not supported on this platform")
138}
139
140pub(crate) fn is_process_alive(pid: u32) -> bool {
141    #[cfg(unix)]
142    {
143        use nix::sys::signal::kill;
144        use nix::unistd::Pid;
145        kill(Pid::from_raw(pid as i32), None).is_ok()
146    }
147    #[cfg(not(unix))]
148    {
149        false
150    }
151}