compose.rs

//! Docker compose command wrapper.
use std::io;
use std::path::{Path, PathBuf};
use std::process::{Command, Output};
use std::time::{Duration, Instant};

use tokio::time::sleep;

#[derive(Clone, Debug)]
pub struct DockerCompose {
    file: PathBuf,
    project: String,
    env_vars: Vec<(String, String)>,
}

#[derive(Debug)]
pub struct RunningCompose {
    compose: DockerCompose,
    is_active: bool,
}

impl Drop for RunningCompose {
    fn drop(&mut self) {
        if !self.is_active {
            return;
        }

        if let Err(error) = self.compose.down() {
            tracing::error!(
                "Failed to stop compose project '{}' from '{}': {error}",
                self.compose.project,
                self.compose.file.display()
            );
        }
    }
}

impl RunningCompose {
    /// Returns the compose project name for this running stack.
    #[must_use]
    pub fn project(&self) -> &str {
        &self.compose.project
    }

    /// Disables the automatic teardown so containers are left running after this
    /// guard is dropped.  Useful for post-run debugging.
    pub fn keep(&mut self) {
        self.is_active = false;
    }
}

impl DockerCompose {
    #[must_use]
    pub fn new(file: &Path, project: &str) -> Self {
        Self {
            file: file.to_path_buf(),
            project: project.to_string(),
            env_vars: vec![],
        }
    }

    #[must_use]
    pub fn with_env(mut self, key: &str, value: &str) -> Self {
        self.env_vars.push((key.to_string(), value.to_string()));
        self
    }

    /// Runs docker compose up and returns a guard that will always run `down --volumes` on drop.
    ///
    /// # Errors
    ///
    /// Returns an error when docker compose fails to start all services.
    pub fn up(&self, no_build: bool) -> io::Result<RunningCompose> {
        let mut args = vec!["up", "--wait", "--detach"];
        if no_build {
            args.push("--no-build");
        }

        let output = self.run_compose(&args)?;

        if output.status.success() {
            Ok(RunningCompose {
                compose: self.clone(),
                is_active: true,
            })
        } else {
            Err(io::Error::other(format!(
                "docker compose up failed for file '{}' and project '{}': {}",
                self.file.display(),
                self.project,
                String::from_utf8_lossy(&output.stderr)
            )))
        }
    }

    /// Builds images defined in the compose file.
    ///
    /// Build output is streamed live to stdout/stderr so progress is visible.
    ///
    /// # Errors
    ///
    /// Returns an error when docker compose build fails.
    pub fn build(&self) -> io::Result<()> {
        let mut command = Command::new("docker");
        command.envs(self.env_vars.iter().map(|(key, value)| (key, value)));
        command.arg("compose");
        command.arg("-f").arg(&self.file);
        command.arg("-p").arg(&self.project);
        command.arg("build");

        tracing::info!("Running docker compose command: {:?}", command);

        let status = command.status()?;
        if status.success() {
            Ok(())
        } else {
            Err(io::Error::other(format!(
                "docker compose build failed for file '{}' and project '{}'",
                self.file.display(),
                self.project,
            )))
        }
    }

    /// Runs docker compose down --volumes.
    ///
    /// # Errors
    ///
    /// Returns an error when docker compose cannot stop and remove resources.
    pub fn down(&self) -> io::Result<()> {
        let output = self.run_compose(&["down", "--volumes"])?;

        if output.status.success() {
            Ok(())
        } else {
            Err(io::Error::other(format!(
                "docker compose down failed for file '{}' and project '{}': {}",
                self.file.display(),
                self.project,
                String::from_utf8_lossy(&output.stderr)
            )))
        }
    }

    /// Resolves an ephemeral host port from a service published container port.
    ///
    /// # Errors
    ///
    /// Returns an error when the compose command fails or port parsing fails.
    pub fn port(&self, service: &str, container_port: u16) -> io::Result<u16> {
        let output = self.run_compose(&["port", service, &container_port.to_string()])?;

        if !output.status.success() {
            return Err(io::Error::other(format!(
                "docker compose port failed for file '{}' and project '{}', service '{}' and port '{}': stderr: {} stdout: {}",
                self.file.display(),
                self.project,
                service,
                container_port,
                String::from_utf8_lossy(&output.stderr),
                String::from_utf8_lossy(&output.stdout)
            )));
        }

        let stdout = String::from_utf8_lossy(&output.stdout);
        let first_line = stdout
            .lines()
            .next()
            .ok_or_else(|| io::Error::other("docker compose port returned no output"))?;

        let host_port = first_line
            .rsplit(':')
            .next()
            .ok_or_else(|| io::Error::other("docker compose port output has no ':' separator"))?
            .parse::<u16>()
            .map_err(|_| io::Error::other(format!("invalid host port in output: '{first_line}'")))?;

        Ok(host_port)
    }

    /// Waits until a service has a resolved host port mapping.
    ///
    /// This helper retries `docker compose port` until it succeeds, the timeout
    /// expires, or the target service exits.
    ///
    /// # Errors
    ///
    /// Returns an error when the service exits, port mapping cannot be resolved
    /// before timeout, or compose commands fail while gathering diagnostics.
    pub async fn wait_for_port_mapping(
        &self,
        service: &str,
        container_port: u16,
        timeout: Duration,
        poll_interval: Duration,
        extra_log_services: &[&str],
    ) -> io::Result<u16> {
        let deadline = Instant::now() + timeout;

        loop {
            if let Ok(ps_output) = self.ps()
                && compose_service_has_exited(&ps_output, service)
            {
                let logs_output = self
                    .logs(&[service])
                    .unwrap_or_else(|error| format!("failed to collect compose logs output: {error}"));

                return Err(io::Error::other(format!(
                    "compose service '{service}' exited while waiting for port mapping '{container_port}'.\nCompose ps:\n{ps_output}\nCompose logs:\n{logs_output}"
                )));
            }

            match self.port(service, container_port) {
                Ok(host_port) => return Ok(host_port),
                Err(_) => {
                    tracing::info!("Waiting for compose port mapping for service '{service}'");
                }
            }

            if Instant::now() >= deadline {
                let ps_output = self
                    .ps()
                    .unwrap_or_else(|error| format!("failed to collect compose ps output: {error}"));

                let mut log_services = Vec::with_capacity(1 + extra_log_services.len());
                log_services.push(service);
                for extra_service in extra_log_services {
                    if *extra_service != service {
                        log_services.push(*extra_service);
                    }
                }

                let logs_output = self
                    .logs(&log_services)
                    .unwrap_or_else(|error| format!("failed to collect compose logs output: {error}"));

                return Err(io::Error::new(
                    io::ErrorKind::TimedOut,
                    format!(
                        "timed out waiting for compose port mapping for service '{service}' and port '{container_port}'.\nCompose ps:\n{ps_output}\nCompose logs:\n{logs_output}"
                    ),
                ));
            }

            sleep(poll_interval).await;
        }
    }

    /// Runs `docker compose exec` in non-interactive mode for scripted commands.
    ///
    /// # Errors
    ///
    /// Returns an error when command execution fails.
    pub fn exec(&self, service: &str, cmd: &[&str]) -> io::Result<Output> {
        let mut args = vec!["exec".to_string(), "-T".to_string(), service.to_string()];
        args.extend(cmd.iter().map(|value| (*value).to_string()));

        self.run_compose_strings(&args)
    }

    /// Runs `docker compose ps -a` and returns stdout.
    ///
    /// # Errors
    ///
    /// Returns an error when the compose command fails.
    pub fn ps(&self) -> io::Result<String> {
        let output = self.run_compose(&["ps", "-a"])?;

        if output.status.success() {
            Ok(String::from_utf8_lossy(&output.stdout).to_string())
        } else {
            Err(io::Error::other(format!(
                "docker compose ps failed for file '{}' and project '{}': {}",
                self.file.display(),
                self.project,
                String::from_utf8_lossy(&output.stderr)
            )))
        }
    }

    /// Runs `docker compose logs --no-color <services...>` and returns stdout.
    ///
    /// # Errors
    ///
    /// Returns an error when the compose command fails.
    pub fn logs(&self, services: &[&str]) -> io::Result<String> {
        let mut args = vec!["logs".to_string(), "--no-color".to_string()];
        args.extend(services.iter().map(|service| (*service).to_string()));

        let output = self.run_compose_strings(&args)?;

        if output.status.success() {
            Ok(String::from_utf8_lossy(&output.stdout).to_string())
        } else {
            Err(io::Error::other(format!(
                "docker compose logs failed for file '{}' and project '{}': {}",
                self.file.display(),
                self.project,
                String::from_utf8_lossy(&output.stderr)
            )))
        }
    }

    fn run_compose(&self, args: &[&str]) -> io::Result<Output> {
        let args_as_strings: Vec<String> = args.iter().map(|value| (*value).to_string()).collect();
        self.run_compose_strings(&args_as_strings)
    }

    fn run_compose_strings(&self, args: &[String]) -> io::Result<Output> {
        let mut command = Command::new("docker");
        command.envs(self.env_vars.iter().map(|(key, value)| (key, value)));
        command.arg("compose");
        command.arg("-f").arg(&self.file);
        command.arg("-p").arg(&self.project);
        command.args(args);

        tracing::info!("Running docker compose command: {:?}", command);

        command.output()
    }
}

fn compose_service_has_exited(ps_output: &str, service_name: &str) -> bool {
    ps_output.lines().any(|line| {
        line.contains(service_name)
            && (line.contains("exited") || line.contains("dead") || line.contains("created") || line.contains("removing"))
    })
}