Skip to main content

tatara_engine/drivers/
nix_build.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, info, warn};
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
17/// Driver for `nix build` — produces store paths and optionally pushes to Attic cache.
18pub struct NixBuildDriver;
19
20#[async_trait]
21impl Driver for NixBuildDriver {
22    fn name(&self) -> &str {
23        "nix-build"
24    }
25
26    async fn available(&self) -> bool {
27        Command::new("nix")
28            .arg("--version")
29            .output()
30            .await
31            .map(|o| o.status.success())
32            .unwrap_or(false)
33    }
34
35    async fn start(&self, task: &Task, alloc_dir: &Path) -> Result<TaskHandle> {
36        let (flake_ref, system, extra_args, attic_cache) = match &task.config {
37            TaskConfig::NixBuild {
38                flake_ref,
39                system,
40                extra_args,
41                attic_cache,
42            } => (
43                flake_ref.clone(),
44                system.clone(),
45                extra_args.clone(),
46                attic_cache.clone(),
47            ),
48            _ => bail!("NixBuildDriver received non-nix-build task config"),
49        };
50
51        let log_dir = alloc_dir.join(&task.name);
52        tokio::fs::create_dir_all(&log_dir).await?;
53
54        let stdout_file = std::fs::File::create(log_dir.join("stdout.log"))
55            .context("Failed to create stdout log")?;
56        let stderr_file = std::fs::File::create(log_dir.join("stderr.log"))
57            .context("Failed to create stderr log")?;
58
59        // Build the nix command
60        let mut cmd = Command::new("nix");
61        cmd.arg("build")
62            .arg(&flake_ref)
63            .arg("--print-out-paths")
64            .arg("--no-link");
65
66        if let Some(sys) = &system {
67            cmd.arg("--system").arg(sys);
68        }
69
70        for arg in &extra_args {
71            cmd.arg(arg);
72        }
73
74        cmd.envs(&task.env)
75            .stdout(Stdio::from(stdout_file))
76            .stderr(Stdio::from(stderr_file))
77            .kill_on_drop(false);
78
79        info!(
80            flake_ref = %flake_ref,
81            system = ?system,
82            attic_cache = ?attic_cache,
83            "Starting nix build"
84        );
85
86        let child = cmd.spawn().context("Failed to spawn nix build")?;
87        let pid = child.id();
88
89        // Spawn background task to wait for completion and optionally push to Attic
90        let log_dir_clone = log_dir.clone();
91        tokio::spawn(async move {
92            let output = child.wait_with_output().await;
93            match output {
94                Ok(out) if out.status.success() => {
95                    // Read store path from stdout log
96                    if let Ok(stdout) =
97                        tokio::fs::read_to_string(log_dir_clone.join("stdout.log")).await
98                    {
99                        let store_path = stdout.trim();
100                        info!(store_path = %store_path, "nix build complete");
101
102                        // Push to Attic if configured
103                        if let Some(cache_name) = attic_cache {
104                            info!(cache = %cache_name, "Pushing to Attic cache");
105                            let push_result = Command::new("attic")
106                                .arg("push")
107                                .arg(&cache_name)
108                                .arg(store_path)
109                                .output()
110                                .await;
111
112                            match push_result {
113                                Ok(r) if r.status.success() => {
114                                    info!(cache = %cache_name, "Attic push complete");
115                                }
116                                Ok(r) => {
117                                    warn!(
118                                        cache = %cache_name,
119                                        stderr = %String::from_utf8_lossy(&r.stderr),
120                                        "Attic push failed (non-fatal)"
121                                    );
122                                }
123                                Err(e) => {
124                                    warn!(cache = %cache_name, error = %e, "Attic push command failed");
125                                }
126                            }
127                        }
128                    }
129                }
130                Ok(out) => {
131                    warn!(exit_code = ?out.status.code(), "nix build failed");
132                }
133                Err(e) => {
134                    warn!(error = %e, "Failed to wait for nix build");
135                }
136            }
137        });
138
139        Ok(TaskHandle {
140            driver: DriverType::NixBuild,
141            pid,
142            container_id: None,
143            started_at: Utc::now(),
144        })
145    }
146
147    async fn stop(&self, handle: &TaskHandle, timeout: Duration) -> Result<()> {
148        if let Some(pid) = handle.pid {
149            send_signal_and_wait(pid, timeout).await?;
150        }
151        Ok(())
152    }
153
154    async fn status(&self, handle: &TaskHandle) -> Result<TaskRunState> {
155        match handle.pid {
156            Some(pid) if is_process_alive(pid) => Ok(TaskRunState::Running),
157            Some(_) => Ok(TaskRunState::Dead),
158            None => Ok(TaskRunState::Dead),
159        }
160    }
161
162    async fn logs(&self, handle: &TaskHandle) -> Result<mpsc::Receiver<LogEntry>> {
163        // TODO: Stream from build log file
164        let (_tx, rx) = mpsc::channel(1);
165        Ok(rx)
166    }
167}