crankshaft-docker 0.6.0

//! Builders for containers.

use std::path::PathBuf;

use bollard::Docker;
use bollard::models::Mount;
use bollard::models::ServiceSpec;
use bollard::models::ServiceSpecMode;
use bollard::models::ServiceSpecModeReplicated;
use bollard::models::TaskSpec;
use bollard::models::TaskSpecContainerSpec;
use bollard::models::TaskSpecResources;
use bollard::models::TaskSpecRestartPolicy;
use bollard::models::TaskSpecRestartPolicyConditionEnum;
use indexmap::IndexMap;
use tracing::warn;

use super::Service;
use crate::Error;
use crate::Result;

/// A builder for a [`Service`].
pub struct Builder {
    /// A reference to the [`Docker`] client that will be used to create this
    /// container.
    client: Docker,

    /// The name of the service.
    name: Option<String>,

    /// The image (e.g., `ubuntu:latest`).
    image: Option<String>,

    /// The program to run.
    program: Option<String>,

    /// The arguments to the command.
    args: Vec<String>,

    /// The file path to write the container's stdout stream to.
    stdout: Option<PathBuf>,

    /// The file path to write the container's stderr stream to.
    stderr: Option<PathBuf>,

    /// Environment variables.
    env: IndexMap<String, String>,

    /// The working directory.
    work_dir: Option<String>,

    /// The mounts for the service's task template.
    mounts: Vec<Mount>,

    /// The task resources for the service.
    resources: Option<TaskSpecResources>,
}

impl Builder {
    /// Creates a new [`Builder`].
    pub fn new(client: Docker) -> Self {
        Self {
            client,
            name: None,
            image: Default::default(),
            program: Default::default(),
            args: Default::default(),
            stdout: None,
            stderr: None,
            env: Default::default(),
            work_dir: Default::default(),
            mounts: Default::default(),
            resources: Default::default(),
        }
    }

    /// Sets the name of the service.
    pub fn name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }

    /// Adds an image name.
    pub fn image(mut self, image: impl Into<String>) -> Self {
        self.image = Some(image.into());
        self
    }

    /// Sets the program to run.
    pub fn program(mut self, program: impl Into<String>) -> Self {
        self.program = Some(program.into());
        self
    }

    /// Sets an argument.
    pub fn arg(mut self, arg: impl Into<String>) -> Self {
        self.args.push(arg.into());
        self
    }

    /// Sets multiple arguments.
    pub fn args(mut self, args: impl IntoIterator<Item = impl Into<String>>) -> Self {
        self.args.extend(args.into_iter().map(Into::into));
        self
    }

    /// Sets an environment variable.
    pub fn env(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.env.insert(name.into(), value.into());
        self
    }

    /// Sets multiple environment variables.
    pub fn envs(
        mut self,
        variables: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
    ) -> Self {
        self.env
            .extend(variables.into_iter().map(|(k, v)| (k.into(), v.into())));
        self
    }

    /// Sets the file to write the container's stdout stream to.
    pub fn stdout(mut self, path: impl Into<PathBuf>) -> Self {
        self.stdout = Some(path.into());
        self
    }

    /// Sets the file to write the container's stderr stream to.
    pub fn stderr(mut self, path: impl Into<PathBuf>) -> Self {
        self.stderr = Some(path.into());
        self
    }

    /// Sets the working directory.
    pub fn work_dir(mut self, work_dir: impl Into<String>) -> Self {
        self.work_dir = Some(work_dir.into());
        self
    }

    /// Sets a mount for the service.
    pub fn mount(mut self, mount: impl Into<Mount>) -> Self {
        self.mounts.push(mount.into());
        self
    }

    /// Sets multiple mounts for the service.
    pub fn mounts(mut self, mounts: impl IntoIterator<Item = impl Into<Mount>>) -> Self {
        self.mounts.extend(mounts.into_iter().map(Into::into));
        self
    }

    /// Sets the task resources.
    pub fn resources(mut self, resources: TaskSpecResources) -> Self {
        self.resources = Some(resources);
        self
    }

    /// Consumes `self` and attempts to create a Docker service.
    pub async fn try_build(self) -> Result<Service> {
        let image = self
            .image
            .ok_or_else(|| Error::MissingBuilderField("image"))?;
        let program = self
            .program
            .ok_or_else(|| Error::MissingBuilderField("program"))?;

        let response = self
            .client
            .create_service(
                ServiceSpec {
                    name: self.name,
                    mode: Some(ServiceSpecMode {
                        replicated: Some(ServiceSpecModeReplicated { replicas: Some(1) }),
                        ..Default::default()
                    }),
                    task_template: Some(TaskSpec {
                        container_spec: Some(TaskSpecContainerSpec {
                            image: Some(image),
                            command: Some(vec![program]),
                            args: Some(self.args),
                            dir: self.work_dir,
                            env: Some(self.env.iter().map(|(k, v)| format!("{k}={v}")).collect()),
                            mounts: Some(self.mounts),
                            // Ensure the caller's group id is added so that the container can
                            // access the mounts and working directory
                            #[cfg(unix)]
                            groups: Some(vec![nix::unistd::Gid::effective().to_string()]),
                            ..Default::default()
                        }),
                        resources: self.resources,
                        restart_policy: Some(TaskSpecRestartPolicy {
                            condition: Some(TaskSpecRestartPolicyConditionEnum::NONE),
                            ..Default::default()
                        }),
                        ..Default::default()
                    }),
                    ..Default::default()
                },
                None,
            )
            .await
            .map_err(Error::Docker)?;

        for warning in response.warnings.unwrap_or_default() {
            warn!("Docker daemon: {warning}");
        }

        Ok(Service {
            client: self.client,
            id: response.id.ok_or_else(|| {
                Error::Message("Docker daemon response did not contain a service identifier".into())
            })?,
            stdout: self.stdout,
            stderr: self.stderr,
        })
    }
}