tatara_engine/drivers/
nix.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;
10
11use tatara_core::domain::allocation::TaskRunState;
12use tatara_core::domain::job::{DriverType, Task, TaskConfig};
13
14use super::exec::{is_process_alive, send_signal_and_wait};
15use super::{Driver, LogEntry, TaskHandle};
16
17pub struct NixDriver;
18
19#[async_trait]
20impl Driver for NixDriver {
21    fn name(&self) -> &str {
22        "nix"
23    }
24
25    async fn available(&self) -> bool {
26        Command::new("nix")
27            .arg("--version")
28            .output()
29            .await
30            .map(|o| o.status.success())
31            .unwrap_or(false)
32    }
33
34    async fn start(&self, task: &Task, alloc_dir: &Path) -> Result<TaskHandle> {
35        let (flake_ref, args) = match &task.config {
36            TaskConfig::Nix { flake_ref, args } => (flake_ref.clone(), args.clone()),
37            _ => bail!("NixDriver received non-nix task config"),
38        };
39
40        let log_dir = alloc_dir.join(&task.name);
41        tokio::fs::create_dir_all(&log_dir).await?;
42
43        let stdout_file = std::fs::File::create(log_dir.join("stdout.log"))
44            .context("Failed to create stdout log")?;
45        let stderr_file = std::fs::File::create(log_dir.join("stderr.log"))
46            .context("Failed to create stderr log")?;
47
48        // `nix run <flake_ref> -- <args>`
49        let mut cmd = Command::new("nix");
50        cmd.arg("run")
51            .arg(&flake_ref)
52            .arg("--")
53            .args(&args)
54            .envs(&task.env)
55            .stdout(Stdio::from(stdout_file))
56            .stderr(Stdio::from(stderr_file))
57            .kill_on_drop(false);
58
59        let child = cmd.spawn().context("Failed to spawn nix run")?;
60        let pid = child.id();
61
62        debug!(
63            task = %task.name,
64            flake_ref = %flake_ref,
65            pid = ?pid,
66            "Nix driver started process"
67        );
68
69        std::mem::forget(child);
70
71        Ok(TaskHandle {
72            driver: DriverType::Nix,
73            pid,
74            container_id: None,
75            started_at: Utc::now(),
76        })
77    }
78
79    async fn stop(&self, handle: &TaskHandle, timeout: Duration) -> Result<()> {
80        let pid = handle.pid.context("No PID in nix task handle")?;
81        send_signal_and_wait(pid, timeout).await
82    }
83
84    async fn status(&self, handle: &TaskHandle) -> Result<TaskRunState> {
85        let pid = handle.pid.context("No PID in nix task handle")?;
86        if is_process_alive(pid) {
87            Ok(TaskRunState::Running)
88        } else {
89            Ok(TaskRunState::Dead)
90        }
91    }
92
93    async fn logs(&self, _handle: &TaskHandle) -> Result<mpsc::Receiver<LogEntry>> {
94        let (tx, rx) = mpsc::channel(256);
95        tokio::spawn(async move {
96            let _ = tx;
97        });
98        Ok(rx)
99    }
100}
tatara_engine/drivers/nix.rs

tatara_engine/drivers/
nix.rs
