Skip to main content

core_lib/exec/
container.rs

1use std::collections::HashMap;
2use std::time::Duration;
3
4use anyhow::Result;
5use bollard::Docker;
6use bollard::models::ContainerCreateBody;
7use bollard::query_parameters::{
8    CreateContainerOptionsBuilder, CreateImageOptionsBuilder, LogsOptionsBuilder,
9    RemoveContainerOptionsBuilder, StartContainerOptions,
10};
11use futures_util::stream::StreamExt;
12
13use tokio::fs::OpenOptions;
14use tokio::io::AsyncWriteExt;
15use tokio::time::timeout;
16
17use crate::core::id::short_id;
18
19async fn ensure_image(docker: &Docker, image: &str, log_path: &str) -> Result<()> {
20    let image_options = CreateImageOptionsBuilder::default()
21        .from_image(image)
22        .build();
23    let mut stream = docker.create_image(Some(image_options), None, None);
24
25    let mut last = String::new();
26    while let Some(status) = stream.next().await {
27        match status {
28            Ok(info) => {
29                if let Some(s) = info.status {
30                    if last != s {
31                        tracing::info!(
32                            target: "fleet",
33                            log_path = %log_path,
34                            message = %format!("Pulling {image}: {s}"),
35                        );
36                    }
37                    last = s.clone();
38                }
39            }
40            Err(e) => {
41                return Err(anyhow::anyhow!("Failed to pull image {image}: {e}"));
42            }
43        }
44    }
45    Ok(())
46}
47
48pub async fn contain_cmd(
49    image: &str,
50    cmd: Vec<String>,
51    env: Option<HashMap<String, String>>,
52    dir: &str,
53    log_path: &str,
54    timeout_secs: Option<u64>,
55) -> Result<()> {
56    tracing::info!(target: "fleet", log_path = %log_path, message = "Building image");
57    let docker = Docker::connect_with_local_defaults()?;
58    ensure_image(&docker, image, log_path).await?;
59
60    let create_option = CreateContainerOptionsBuilder::default()
61        .name(&format!("fleet-job-{}", short_id()))
62        .build();
63
64    let container_config = ContainerCreateBody {
65        image: Some(image.to_string()),
66        cmd: Some(cmd),
67        env: env.map(|m| m.iter().map(|(k, v)| format!("{k}={v}")).collect()),
68        attach_stdout: Some(true),
69        attach_stderr: Some(true),
70        tty: Some(false),
71        host_config: Some(bollard::models::HostConfig {
72            binds: Some(vec![format!("{dir}:/app",)]),
73            ..Default::default()
74        }),
75        working_dir: Some("/app".to_string()),
76        ..Default::default()
77    };
78
79    let container = docker
80        .create_container(Some(create_option), container_config)
81        .await?;
82
83    docker
84        .start_container(&container.id, None::<StartContainerOptions>)
85        .await?;
86
87    let mut log_file = OpenOptions::new()
88        .create(true)
89        .append(true)
90        .open(log_path)
91        .await?;
92
93    let logs_options = LogsOptionsBuilder::default()
94        .follow(true)
95        .stdout(true)
96        .stderr(true)
97        .build();
98
99    let mut log_stream = docker.logs(&container.id, Some(logs_options));
100
101    let logs_future = async {
102        while let Some(log) = log_stream.next().await {
103            match log? {
104                bollard::container::LogOutput::StdOut { message }
105                | bollard::container::LogOutput::StdErr { message } => {
106                    log_file.write_all(&message).await?;
107                }
108                _ => {}
109            }
110        }
111        Ok::<(), anyhow::Error>(())
112    };
113
114    let result = if let Some(secs) = timeout_secs {
115        // w Timeout
116        match timeout(Duration::from_secs(secs), logs_future).await {
117            Ok(inner) => inner,
118            Err(_) => {
119                tracing::error!(
120                    target: "fleet",
121                    log_path = %log_path,
122                    message = %format!(
123                        "Container execution timed out after {secs} seconds"
124                    ),
125                );
126                Err(anyhow::anyhow!(
127                    "Container execution timed out after {secs} seconds"
128                ))
129            }
130        }
131    } else {
132        // w no timemout
133        logs_future.await
134    };
135
136    let remove_options = RemoveContainerOptionsBuilder::default().force(true).build();
137    docker
138        .remove_container(&container.id, Some(remove_options))
139        .await?;
140    result
141}