bollard_compose/lib/src/container.rs

use crate::configuration::compose_types::DockerCompose;
use crate::helpers::dir::parent_dir_name;
use anyhow::anyhow;
use bollard::container::{Config, CreateContainerOptions, RemoveContainerOptions};
use bollard::models::{HostConfig, PortBinding, PortMap, RestartPolicy, RestartPolicyNameEnum};
use bollard::Docker;
use std::collections::HashMap;
use std::iter::Map;

pub(crate) async fn create_containers(
    compose: &DockerCompose,
    docker: &Docker,
    detach: bool,
) -> anyhow::Result<HashMap<String, String>> {
    let mut container_ids = HashMap::new();
    let parent_dir = parent_dir_name()?;
    for (name, service) in &compose.services {
        let env = create_env_vec(&service.environment);
        let conf: Config<String> = Config {
            image: service.image.clone(),
            attach_stdout: Some(!detach),
            attach_stderr: Some(!detach),
            open_stdin: Some(false),
            env,
            host_config: Some(HostConfig {
                auto_remove: service.auto_remove,
                restart_policy: parse_restart_policy(&service.restart),
                port_bindings: create_port_map(&service.ports),
                binds: parse_volume_mounts(compose, &service.volumes)?,
                ..Default::default()
            }),
            ..Default::default()
        };
        let container_name = service
            .container_name
            .clone()
            .unwrap_or_else(|| format!("{}_{}_1", parent_dir, name));
        let create_info = docker
            .create_container::<String, String>(
                Some(CreateContainerOptions {
                    name: container_name,
                    platform: None,
                }),
                conf,
            )
            .await?;
        container_ids.insert(name.clone(), create_info.id);
    }
    Ok(container_ids)
}

fn resolve_start_order(compose: &DockerCompose) -> Vec<String> {
    let mut start_order = Vec::new();
    let mut visited = HashMap::new();

    for service_name in compose.services.keys() {
        resolve_service(service_name, compose, &mut visited, &mut start_order);
    }

    start_order
}

fn resolve_service(
    service_name: &str,
    compose: &DockerCompose,
    visited: &mut HashMap<String, bool>,
    start_order: &mut Vec<String>,
) {
    if let Some(&true) = visited.get(service_name) {
        return;
    }
    visited.insert(service_name.to_string(), true);

    if let Some(service) = compose.services.get(service_name) {
        if let Some(dependencies) = &service.depends_on {
            for dependency in dependencies {
                resolve_service(dependency, compose, visited, start_order);
            }
        }
    }

    start_order.push(service_name.to_string());
}

pub(crate) async fn start_containers(
    compose: &DockerCompose,
    docker: &Docker,
    ids: HashMap<String, String>,
) -> anyhow::Result<()> {
    // resolve dependency resolution
    let start_order = resolve_start_order(compose);

    for container_name in start_order {
        docker
            .start_container::<String>(
                ids.get(&container_name).ok_or(anyhow!(
                    "no created container found with name {}",
                    container_name
                ))?,
                None,
            )
            .await?;
    }
    Ok(())
}

pub async fn stop_containers(compose: &DockerCompose, docker: &Docker) -> anyhow::Result<()> {
    let parent_dir = parent_dir_name()?;
    for (name, service) in &compose.services {
        let container_name = service
            .container_name
            .clone()
            .unwrap_or_else(|| format!("{}_{}_1", parent_dir, name));
        docker.stop_container(&container_name, None).await?;
    }
    Ok(())
}

pub async fn remove_containers(compose: &DockerCompose, docker: &Docker) -> anyhow::Result<()> {
    let parent_dir = parent_dir_name()?;
    for (name, service) in &compose.services {
        let container_name = service
            .container_name
            .clone()
            .unwrap_or_else(|| format!("{}_{}_1", parent_dir, name));
        docker
            .remove_container(
                &container_name,
                Some(RemoveContainerOptions {
                    force: true,
                    ..Default::default()
                }),
            )
            .await?;
    }
    Ok(())
}

pub(crate) async fn stop_containers_by_ids(
    docker: &Docker,
    ids: Vec<String>,
) -> anyhow::Result<()> {
    for id in ids {
        docker.stop_container(&id, None).await?;
    }
    Ok(())
}

fn create_env_vec(env: &Option<HashMap<String, String>>) -> Option<Vec<String>> {
    match env {
        Some(env) => {
            let list = env
                .iter()
                .map(|(key, value)| format!("{}={}", key, value))
                .collect();
            Some(list)
        }
        None => None,
    }
}

fn create_port_map(ports: &Option<Vec<String>>) -> Option<PortMap> {
    ports.as_ref().map(|ports| {
        ports
            .iter()
            .map(|port| {
                let parts: Vec<&str> = port.split(':').collect();
                (
                    parts[1].to_string(),
                    Some(vec![PortBinding {
                        host_ip: None,
                        host_port: Some(parts[0].to_string()),
                    }]),
                )
            })
            .collect()
    })
}

fn parse_restart_policy(restart: &Option<String>) -> Option<RestartPolicy> {
    match restart {
        None => None,
        Some(restart) => match restart.as_str() {
            "no" => Some(RestartPolicy {
                name: Some(RestartPolicyNameEnum::NO),
                ..Default::default()
            }),
            "always" => Some(RestartPolicy {
                name: Some(RestartPolicyNameEnum::ALWAYS),
                ..Default::default()
            }),
            "unless-stopped" => Some(RestartPolicy {
                name: Some(RestartPolicyNameEnum::UNLESS_STOPPED),
                ..Default::default()
            }),
            "on-failure" => Some(RestartPolicy {
                name: Some(RestartPolicyNameEnum::NO),
                ..Default::default()
            }),
            v => {
                // handle special case when retry count is specified
                if v.starts_with("on-failure:") {
                    let parts: Vec<&str> = v.split(':').collect();
                    Some(RestartPolicy {
                        name: Some(RestartPolicyNameEnum::ON_FAILURE),
                        maximum_retry_count: Some(parts[1].parse().unwrap()),
                    })
                } else {
                    None
                }
            }
        },
    }
}

fn parse_volume_mount(compose: &DockerCompose, volume: String) -> anyhow::Result<String> {
    let parts: Vec<&str> = volume.split(':').collect();
    if parts.len() != 2 {
        return Err(anyhow!("Invalid volume mount: {}", volume));
    }

    if parts[0].starts_with('/') {
        return Ok(volume);
    }

    let parent_dir = parent_dir_name()?;

    let field_missing = format!("volume field for {} missing", volume);

    // volumes map has to contain the volume defined here
    if compose
        .volumes
        .as_ref()
        .ok_or(anyhow!(field_missing.clone()))?
        .contains_key(parts[0])
    {
        Ok(format!("{}_{}:{}", parent_dir, parts[0], parts[1]))
    } else {
        Err(anyhow!(field_missing))
    }
}

fn parse_volume_mounts(
    compose: &DockerCompose,
    volumes: &Option<Vec<String>>,
) -> anyhow::Result<Option<Vec<String>>> {
    match volumes {
        None => Ok(None),
        Some(volumes) => {
            let mounts: Vec<anyhow::Result<String>> = volumes
                .iter()
                .map(|volume| parse_volume_mount(compose, volume.clone()))
                .collect();
            let mounts = mounts
                .into_iter()
                .collect::<anyhow::Result<Vec<String>>>()?;
            if mounts.is_empty() {
                Ok(None)
            } else {
                Ok(Some(mounts))
            }
        }
    }
}