use anyhow::{bail, Context, Result};
use async_trait::async_trait;
use chrono::Utc;
use std::path::Path;
use std::time::Duration;
use tokio::process::Command;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};
use tatara_core::domain::allocation::TaskRunState;
use tatara_core::domain::job::{DriverType, Task, TaskConfig};
use super::{Driver, LogEntry, TaskHandle};
#[derive(Debug)]
enum OciBackend {
AppleContainer,
Docker,
Podman,
None,
}
pub struct OciDriver {
backend: OciBackend,
}
impl OciDriver {
pub async fn detect() -> Self {
if cfg!(target_os = "macos") {
if which("container").await {
info!("OCI backend: Apple container CLI");
return Self {
backend: OciBackend::AppleContainer,
};
}
}
if which("docker").await {
info!("OCI backend: Docker");
return Self {
backend: OciBackend::Docker,
};
}
if which("podman").await {
info!("OCI backend: Podman");
return Self {
backend: OciBackend::Podman,
};
}
warn!("No OCI runtime available");
Self {
backend: OciBackend::None,
}
}
fn cli_command(&self) -> &str {
match &self.backend {
OciBackend::AppleContainer => "container",
OciBackend::Docker => "docker",
OciBackend::Podman => "podman",
OciBackend::None => unreachable!("OCI driver used without available backend"),
}
}
}
#[async_trait]
impl Driver for OciDriver {
fn name(&self) -> &str {
"oci"
}
async fn available(&self) -> bool {
!matches!(self.backend, OciBackend::None)
}
async fn start(&self, task: &Task, _alloc_dir: &Path) -> Result<TaskHandle> {
let (image, ports, volumes, entrypoint, command) = match &task.config {
TaskConfig::Oci {
image,
ports,
volumes,
entrypoint,
command,
} => (
image.clone(),
ports.clone(),
volumes.clone(),
entrypoint.clone(),
command.clone(),
),
_ => bail!("OciDriver received non-oci task config"),
};
let container_name = format!("tatara-{}", task.name);
let cli = self.cli_command();
let mut args = vec!["run".to_string(), "-d".to_string()];
args.push("--name".to_string());
args.push(container_name.clone());
for (host_port, container_port) in &ports {
args.push("-p".to_string());
args.push(format!("{}:{}", host_port, container_port));
}
for (host_path, container_path) in &volumes {
args.push("-v".to_string());
args.push(format!("{}:{}", host_path, container_path));
}
for (key, value) in &task.env {
args.push("-e".to_string());
args.push(format!("{}={}", key, value));
}
args.push(image);
if let Some(ep) = &entrypoint {
args.push("--entrypoint".to_string());
args.extend(ep.clone());
}
if let Some(cmd) = &command {
args.extend(cmd.clone());
}
let output = Command::new(cli)
.args(&args)
.output()
.await
.context("Failed to start OCI container")?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
bail!("Failed to start container: {}", stderr);
}
let container_id = String::from_utf8_lossy(&output.stdout).trim().to_string();
debug!(
task = %task.name,
container_id = %container_id,
backend = ?self.backend,
"OCI driver started container"
);
Ok(TaskHandle {
driver: DriverType::Oci,
pid: None,
container_id: Some(container_id),
started_at: Utc::now(),
})
}
async fn stop(&self, handle: &TaskHandle, timeout: Duration) -> Result<()> {
let container_id = handle
.container_id
.as_ref()
.context("No container ID in OCI task handle")?;
let cli = self.cli_command();
let timeout_secs = timeout.as_secs().to_string();
let output = Command::new(cli)
.args(["stop", "-t", &timeout_secs, container_id])
.output()
.await
.context("Failed to stop container")?;
if !output.status.success() {
warn!(
container_id = %container_id,
stderr = %String::from_utf8_lossy(&output.stderr),
"Container stop returned error"
);
}
let _ = Command::new(cli)
.args(["rm", "-f", container_id])
.output()
.await;
Ok(())
}
async fn status(&self, handle: &TaskHandle) -> Result<TaskRunState> {
let container_id = handle
.container_id
.as_ref()
.context("No container ID in OCI task handle")?;
let cli = self.cli_command();
let output = Command::new(cli)
.args(["inspect", "--format", "{{.State.Status}}", container_id])
.output()
.await
.context("Failed to inspect container")?;
if !output.status.success() {
return Ok(TaskRunState::Dead);
}
let status = String::from_utf8_lossy(&output.stdout).trim().to_string();
match status.as_str() {
"running" => Ok(TaskRunState::Running),
"created" | "restarting" => Ok(TaskRunState::Pending),
_ => Ok(TaskRunState::Dead),
}
}
async fn logs(&self, handle: &TaskHandle) -> Result<mpsc::Receiver<LogEntry>> {
let (tx, rx) = mpsc::channel(256);
let container_id = handle.container_id.clone().unwrap_or_default();
let cli = self.cli_command().to_string();
tokio::spawn(async move {
use tokio::io::{AsyncBufReadExt, BufReader};
let mut child = match Command::new(&cli)
.args(["logs", "-f", "--timestamps", &container_id])
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.spawn()
{
Ok(child) => child,
Err(e) => {
tracing::warn!(error = %e, "failed to spawn container log stream");
return;
}
};
if let Some(stdout) = child.stdout.take() {
let tx_out = tx.clone();
let cid = container_id.clone();
tokio::spawn(async move {
let mut lines = BufReader::new(stdout).lines();
while let Ok(Some(line)) = lines.next_line().await {
let entry = LogEntry {
task_name: cid.clone(),
message: line,
stream: "stdout".to_string(),
timestamp: chrono::Utc::now(),
};
if tx_out.send(entry).await.is_err() {
break;
}
}
});
}
if let Some(stderr) = child.stderr.take() {
let tx_err = tx;
let cid = container_id;
tokio::spawn(async move {
let mut lines = BufReader::new(stderr).lines();
while let Ok(Some(line)) = lines.next_line().await {
let entry = LogEntry {
task_name: cid.clone(),
message: line,
stream: "stderr".to_string(),
timestamp: chrono::Utc::now(),
};
if tx_err.send(entry).await.is_err() {
break;
}
}
});
}
let _ = child.wait().await;
});
Ok(rx)
}
}
async fn which(name: &str) -> bool {
Command::new("which")
.arg(name)
.output()
.await
.map(|o| o.status.success())
.unwrap_or(false)
}