tatara_engine/drivers/
mod.rs

1pub mod exec;
2#[cfg(all(target_os = "macos", feature = "kasou"))]
3pub mod kasou;
4pub mod kube;
5pub mod nix;
6pub mod nix_build;
7pub mod oci;
8pub mod wasi;
9
10use anyhow::Result;
11use async_trait::async_trait;
12use chrono::{DateTime, Utc};
13use serde::{Deserialize, Serialize};
14use std::path::Path;
15use std::time::Duration;
16use tokio::sync::mpsc;
17
18use tatara_core::domain::allocation::TaskRunState;
19use tatara_core::domain::job::{DriverType, Task};
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct TaskHandle {
23    pub driver: DriverType,
24    pub pid: Option<u32>,
25    pub container_id: Option<String>,
26    pub started_at: DateTime<Utc>,
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct LogEntry {
31    pub task_name: String,
32    pub message: String,
33    pub stream: String,
34    pub timestamp: DateTime<Utc>,
35}
36
37#[async_trait]
38pub trait Driver: Send + Sync {
39    fn name(&self) -> &str;
40    async fn available(&self) -> bool;
41    async fn start(&self, task: &Task, alloc_dir: &Path) -> Result<TaskHandle>;
42    async fn stop(&self, handle: &TaskHandle, timeout: Duration) -> Result<()>;
43    async fn status(&self, handle: &TaskHandle) -> Result<TaskRunState>;
44    async fn logs(&self, handle: &TaskHandle) -> Result<mpsc::Receiver<LogEntry>>;
45}
46
47pub struct DriverRegistry {
48    drivers: Vec<Box<dyn Driver>>,
49}
50
51impl DriverRegistry {
52    pub async fn new() -> Self {
53        let mut drivers: Vec<Box<dyn Driver>> = Vec::new();
54
55        let exec = exec::ExecDriver;
56        if exec.available().await {
57            drivers.push(Box::new(exec));
58        }
59
60        let oci = oci::OciDriver::detect().await;
61        if oci.available().await {
62            drivers.push(Box::new(oci));
63        }
64
65        let nix = nix::NixDriver;
66        if nix.available().await {
67            drivers.push(Box::new(nix));
68        }
69
70        let nix_build = nix_build::NixBuildDriver;
71        if nix_build.available().await {
72            drivers.push(Box::new(nix_build));
73        }
74
75        #[cfg(all(target_os = "macos", feature = "kasou"))]
76        {
77            let kasou_driver = kasou::KasouDriver::new();
78            if kasou_driver.available().await {
79                drivers.push(Box::new(kasou_driver));
80            }
81        }
82
83        // WASI driver (wasmtime)
84        let wasi_driver = wasi::WasiDriver;
85        if wasi_driver.available().await {
86            drivers.push(Box::new(wasi_driver));
87        }
88
89        // Kubernetes driver
90        let kube_driver = kube::KubeDriver::new();
91        if kube_driver.available().await {
92            drivers.push(Box::new(kube_driver));
93        }
94
95        Self { drivers }
96    }
97
98    pub fn get(&self, driver_type: &DriverType) -> Option<&dyn Driver> {
99        let name = match driver_type {
100            DriverType::Exec => "exec",
101            DriverType::Oci => "oci",
102            DriverType::Nix => "nix",
103            DriverType::NixBuild => "nix-build",
104            DriverType::Kasou => "kasou",
105            DriverType::Kube => "kube",
106            DriverType::Wasi => "wasi",
107        };
108        self.drivers
109            .iter()
110            .find(|d| d.name() == name)
111            .map(|d| d.as_ref())
112    }
113
114    pub fn available_drivers(&self) -> Vec<DriverType> {
115        self.drivers
116            .iter()
117            .filter_map(|d| match d.name() {
118                "exec" => Some(DriverType::Exec),
119                "oci" => Some(DriverType::Oci),
120                "nix" => Some(DriverType::Nix),
121                "nix-build" => Some(DriverType::NixBuild),
122                "kasou" => Some(DriverType::Kasou),
123                "kube" => Some(DriverType::Kube),
124                "wasi" => Some(DriverType::Wasi),
125                _ => None,
126            })
127            .collect()
128    }
129}
tatara_engine/drivers/mod.rs

tatara_engine/drivers/
mod.rs
