launcher/src/docker/build.rs

use super::ServeOpts;
use super::shq;
use crate::config::{Cluster, Config, Node, Recipe};
use crate::ssh::run_ssh;
use anyhow::{Context, Result};
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, HashMap};
use std::io::Read;
use std::path::Path;
use tracing::info;

/// Default upper bound on a single remote image build before the launcher
/// gives up (90 min). A hung build (stuck tmux, network stall) would
/// otherwise loop the 30s poll forever, hanging the launcher.
pub const DEFAULT_MAX_BUILD_SECS: u64 = 90 * 60;

/// Compute a SHA-256 checksum over all files in `dir` (recursive).
/// Sorted by relative path to ensure reproducibility.
pub fn compute_dir_checksum(dir: &Path) -> Result<String> {
    let mut hasher = Sha256::new();
    let mut entries: Vec<_> = Vec::new();

    for entry in walkdir::WalkDir::new(dir).sort_by_file_name() {
        let entry = entry?;
        if entry.file_type().is_dir() {
            continue;
        }
        let rel = entry
            .path()
            .strip_prefix(dir)
            .unwrap_or(entry.path())
            .to_string_lossy()
            .to_string();
        entries.push((rel, entry.path().to_path_buf()));
    }
    entries.sort_by(|a, b| a.0.cmp(&b.0));

    for (rel, path) in &entries {
        let mut file = std::fs::File::open(path)?;
        let mut buf = Vec::new();
        file.read_to_end(&mut buf)?;
        hasher.update(rel.as_bytes());
        hasher.update(b"\0");
        hasher.update(&buf);
    }

    let hash = hasher.finalize();
    Ok(hash.iter().map(|b| format!("{:02x}", b)).collect())
}

/// Parse a `versions.env` file into a sorted map. Only `KEY=VAL` lines are
/// captured; comments (`#`) and blank lines are dropped. The value is
/// everything after the first `=` (so values may contain `=`).
fn parse_versions_env(raw: &str) -> BTreeMap<String, String> {
    let mut pins = BTreeMap::new();
    for line in raw.lines() {
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        if let Some((k, v)) = line.split_once('=') {
            let k = k.trim();
            if !k.is_empty() {
                pins.insert(k.to_string(), v.trim().to_string());
            }
        }
    }
    pins
}

/// Serialize a pin map back to `versions.env` text (sorted by key; one
/// `KEY=VAL` per line). Comments are not preserved — this is a generated
/// artifact written to the remote build dir, not the source of truth.
fn serialize_versions_env(pins: &BTreeMap<String, String>) -> String {
    let mut out = String::from(
        "# Merged versions.env \u{2014} generated by sm121 (recipe [build.versions] overrides applied)\n",
    );
    for (k, v) in pins {
        out.push_str(k);
        out.push('=');
        out.push_str(v);
        out.push('\n');
    }
    out
}

/// Compute the build checksum, folding in recipe `[build.versions]` overrides.
///
/// Returns `(merged_versions_env, checksum)`:
/// - No overrides -> `(None, dir_checksum)`: the image's `versions.env` is
///   uploaded as-is and the checksum is unchanged (existing images skip
///   rebuild, backward-compatible).
/// - Overrides present -> `(Some(merged), combined)`: `merged` is the image's
///   `versions.env` with recipe pins applied (recipe wins) for the launcher
///   to write to the remote build dir; `combined` hashes the dir checksum
///   together with the merged env so any override change forces a rebuild.
fn build_checksum(
    local_dir: &Path,
    overrides: &HashMap<String, String>,
) -> Result<(Option<String>, String)> {
    let dir_checksum = compute_dir_checksum(local_dir)?;
    if overrides.is_empty() {
        return Ok((None, dir_checksum));
    }
    let env_path = local_dir.join("versions.env");
    let raw = std::fs::read_to_string(&env_path).context(format!(
        "recipe carries [build.versions] overrides but '{}' has no versions.env",
        env_path.display()
    ))?;
    let mut pins = parse_versions_env(&raw);
    for (k, v) in overrides {
        pins.insert(k.clone(), v.clone());
    }
    let merged = serialize_versions_env(&pins);
    let mut h = Sha256::new();
    h.update(dir_checksum.as_bytes());
    h.update(b"\0");
    h.update(merged.as_bytes());
    let combined: String = h.finalize().iter().map(|b| format!("{:02x}", b)).collect();
    Ok((Some(merged), combined))
}

/// Start a Docker image build on a remote node in a tmux session.
/// After the build completes, stamps the checksum label on the image.
pub async fn build_image_on_node(
    config: &Config,
    node: &Node,
    local_dir: &Path,
    recipe_name: &str,
    image_tag: &str,
    checksum: &str,
    merged_versions_env: Option<&str>,
) -> Result<()> {
    let session_name = format!("sm121-build-{}", recipe_name);
    let logfile = format!("/tmp/{}-build.log", recipe_name);
    let remote_dir = format!("/tmp/{}-build", recipe_name);

    info!(
        "Building image '{}' on '{}' (tmux: {})...",
        image_tag, node.name, session_name
    );
    crate::ssh::upload_dir(config, node, local_dir, &remote_dir, false).await?;

    // If the recipe overrode any `versions.env` pins, overwrite the uploaded
    // copy on the node with the merged content (recipe pins win over the
    // image's defaults). `build.sh` sources `versions.env`, so this is how
    // recipe overrides reach the Dockerfile `--build-arg`s.
    if let Some(merged) = merged_versions_env {
        let b64 = STANDARD.encode(merged.as_bytes());
        run_ssh(
            config,
            node,
            format!(
                "printf '%s' '{}' | base64 -d > {}/versions.env",
                b64,
                shq(&remote_dir)
            ),
            false,
        )
        .await?;
        info!(
            "Applied recipe [build.versions] overrides to versions.env on '{}'",
            node.name
        );
    }

    // Kill any stale session from a previous (interrupted or falsely reported
    // dead) run so `tmux new-session` does not collide with "duplicate session".
    run_ssh(
        config,
        node,
        format!(
            "tmux kill-session -t {} 2>/dev/null || true",
            shq(&session_name)
        ),
        false,
    )
    .await?;

    let tmux_cmd = format!(
        "tmux new-session -d -s {} -c {} \"TAG={} bash build.sh 2>&1 | tee {}\"",
        shq(&session_name),
        shq(&remote_dir),
        shq(image_tag),
        shq(&logfile)
    );
    run_ssh(config, node, tmux_cmd, false).await?;

    let start = std::time::Instant::now();
    loop {
        tokio::time::sleep(std::time::Duration::from_secs(30)).await;
        let elapsed = start.elapsed();

        if elapsed.as_secs() >= DEFAULT_MAX_BUILD_SECS {
            anyhow::bail!(
                "Build on '{}' of '{}' timed out after {} — no completion marker. \
                 Inspect tmux session '{}' / log '{}' on the node.",
                node.name,
                image_tag,
                format_elapsed(elapsed),
                session_name,
                logfile
            );
        }

        // `tmux has-session` exits 0 and prints nothing when the session is
        // alive; the `|| echo DEAD` makes it print DEAD when the session is
        // gone (still exit 0, so run_ssh stays Ok). An empty `alive` therefore
        // means ALIVE — it must NOT be treated as dead (that was the original
        // bug: a still-running build was reported "failed after 0m" at the
        // first 30s poll, orphaning the tmux session).
        let alive = run_ssh(
            config,
            node,
            format!(
                "tmux has-session -t {} 2>&1 || echo DEAD",
                shq(&session_name)
            ),
            false,
        )
        .await?;
        let dead = alive.trim().contains("DEAD");

        if let Ok(logs) = run_ssh(
            config,
            node,
            format!(
                "tail -10 {} 2>/dev/null || echo '(no log yet)'",
                shq(&logfile)
            ),
            false,
        )
        .await
        {
            for line in logs.lines().filter(|l| !l.is_empty()) {
                info!("  [build] {}", line.trim());
            }
        }

        if dead {
            // Detect the success marker anywhere in the full log (a `tail -5`
            // window can miss it if trailing output pushes it past the end).
            let marker = format!("Done: {}", image_tag);
            let ok = run_ssh(
                config,
                node,
                format!(
                    "grep -F -- {} {} >/dev/null 2>&1 && echo OK || echo NO",
                    shq(&marker),
                    shq(&logfile)
                ),
                false,
            )
            .await
            .map(|s| s.trim() == "OK")
            .unwrap_or(false);
            if ok {
                info!(
                    "Image '{}' built on '{}' ({})",
                    image_tag,
                    node.name,
                    format_elapsed(elapsed)
                );
            } else {
                let exit_log = run_ssh(config, node, format!("tail -20 {}", shq(&logfile)), false)
                    .await
                    .unwrap_or_default();
                anyhow::bail!(
                    "Build on '{}' failed after {}.\n{}",
                    node.name,
                    format_elapsed(elapsed),
                    exit_log
                );
            }
            break;
        }
        info!("  Build running for {}...", format_elapsed(elapsed));
    }

    // Stamp checksum label on the built image.
    let stamp_cmd = format!(
        "printf 'FROM %s\n' {} | docker build --label sm121.checksum={} -t {} -",
        shq(image_tag),
        checksum,
        shq(image_tag)
    );
    let _stamp = run_ssh(config, node, stamp_cmd, false).await?;
    info!("Stamped checksum label '{}' on '{}'", checksum, image_tag);
    Ok(())
}

/// Decide whether to build and dispatch the build for the serve flow.
pub async fn handle_build(
    config: &Config,
    cluster: &Cluster,
    recipe_name: &str,
    recipe: &Recipe,
    opts: &ServeOpts,
    build_node_name: &str,
) -> Result<()> {
    if opts.dry_run {
        let dir = recipe.docker_image_dir.as_deref().unwrap_or(recipe_name);
        let overrides = if recipe.build.versions.is_empty() {
            String::new()
        } else {
            format!(
                "\n[dry-run] recipe [build.versions] overrides {} pin(s) \u{2014} merged into versions.env on the node",
                recipe.build.versions.len()
            )
        };
        println!(
            "[dry-run] would build '{}' on '{}' from docker-images/{}/build.sh{}",
            recipe.docker_image, build_node_name, dir, overrides
        );
        return Ok(());
    }
    let image_dir = recipe
        .docker_image_dir
        .clone()
        .unwrap_or_else(|| recipe_name.to_string());
    let local_dir = std::path::Path::new("docker-images").join(&image_dir);
    if !local_dir.join("build.sh").exists() {
        anyhow::bail!(
            "No build system at '{}' (missing build.sh)",
            local_dir.display()
        );
    }
    let build_node = cluster
        .nodes
        .iter()
        .find(|n| n.name == build_node_name)
        .ok_or_else(|| anyhow::anyhow!("build-node '{}' not in cluster", build_node_name))?;

    let (merged_env, checksum) = build_checksum(local_dir.as_path(), &recipe.build.versions)?;

    let should_build = if opts.force_build {
        true
    } else {
        let img_id = run_ssh(
            config,
            build_node,
            format!(
                "docker image ls --format '{{{{.ID}}}}' {} | head -1",
                shq(&recipe.docker_image)
            ),
            false,
        )
        .await?;
        if img_id.trim().is_empty() {
            true
        } else {
            let label = run_ssh(
                config,
                build_node,
                format!(
                    "docker inspect --format '{{{{index .Config.Labels \"sm121.checksum\"}}}}' {} 2>/dev/null || echo MISSING",
                    shq(&recipe.docker_image)
                ),
                false,
            )
            .await?;
            let label = label.trim();
            if label == checksum {
                info!(
                    "Image '{}' checksum matches on '{}' — skipping build",
                    recipe.docker_image, build_node.name
                );
                false
            } else {
                info!(
                    "Image '{}' checksum mismatch on '{}' (got {}, want {}) — rebuilding",
                    recipe.docker_image, build_node.name, label, checksum
                );
                true
            }
        }
    };

    if should_build {
        build_image_on_node(
            config,
            build_node,
            local_dir.as_path(),
            recipe_name,
            &recipe.docker_image,
            &checksum,
            merged_env.as_deref(),
        )
        .await?;
    }
    Ok(())
}

/// Format a `Duration` as a compact human string: `45s`, `1m 30s`, `12m 04s`.
fn format_elapsed(d: std::time::Duration) -> String {
    let total_secs = d.as_secs();
    let mins = total_secs / 60;
    let secs = total_secs % 60;
    if mins == 0 {
        format!("{}s", secs)
    } else {
        format!("{}m {:02}s", mins, secs)
    }
}

#[cfg(test)]
mod tests {
    use super::{build_checksum, format_elapsed, parse_versions_env, serialize_versions_env};
    use std::collections::HashMap;
    use std::fs;
    use std::path::PathBuf;
    use std::time::Duration;

    fn temp_test_dir(name: &str) -> PathBuf {
        let p =
            std::env::temp_dir().join(format!("sm121-build-test-{}-{}", std::process::id(), name));
        fs::remove_dir_all(&p).ok();
        fs::create_dir_all(&p).unwrap();
        p
    }

    #[test]
    fn format_elapsed_seconds_and_minutes() {
        assert_eq!(format_elapsed(Duration::from_secs(0)), "0s");
        assert_eq!(format_elapsed(Duration::from_secs(30)), "30s");
        assert_eq!(format_elapsed(Duration::from_secs(59)), "59s");
        assert_eq!(format_elapsed(Duration::from_secs(60)), "1m 00s");
        assert_eq!(format_elapsed(Duration::from_secs(90)), "1m 30s");
        assert_eq!(format_elapsed(Duration::from_secs(125)), "2m 05s");
    }

    #[test]
    fn parse_versions_env_skips_comments_and_blanks() {
        let raw = "# comment\n\nKEY=val\nOTHER=a=b\n  SPACED  =  x  \nnoeq\n";
        let pins = parse_versions_env(raw);
        assert_eq!(pins.get("KEY"), Some(&"val".to_string()));
        assert_eq!(pins.get("OTHER"), Some(&"a=b".to_string()));
        assert_eq!(pins.get("SPACED"), Some(&"x".to_string()));
        assert!(!pins.contains_key("noeq"));
        assert_eq!(pins.len(), 3);
    }

    #[test]
    fn serialize_versions_env_sorted_one_per_line() {
        let mut pins = std::collections::BTreeMap::new();
        pins.insert("B".into(), "2".into());
        pins.insert("A".into(), "1".into());
        let s = serialize_versions_env(&pins);
        assert_eq!(
            s,
            "# Merged versions.env \u{2014} generated by sm121 (recipe [build.versions] overrides applied)\nA=1\nB=2\n"
        );
    }

    #[test]
    fn build_checksum_no_overrides_returns_none_and_dir_checksum() {
        let dir = temp_test_dir("no_overrides");
        fs::write(dir.join("versions.env"), "A=1\nB=2\n").unwrap();
        let (merged, ck) = build_checksum(&dir, &HashMap::new()).unwrap();
        assert!(merged.is_none());
        let (_, ck2) = build_checksum(&dir, &HashMap::new()).unwrap();
        assert_eq!(ck, ck2);
        fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn build_checksum_overrides_change_checksum_and_produce_merged() {
        let dir = temp_test_dir("overrides");
        fs::write(dir.join("versions.env"), "A=1\nB=2\n").unwrap();
        let no_over = build_checksum(&dir, &HashMap::new()).unwrap().1;
        let mut over = HashMap::new();
        over.insert("A".to_string(), "99".to_string());
        over.insert("C".to_string(), "3".to_string());
        let (merged, with_over) = build_checksum(&dir, &over).unwrap();
        assert_ne!(no_over, with_over);
        let m = merged.unwrap();
        assert!(m.contains("A=99"));
        assert!(m.contains("C=3"));
        assert!(m.contains("B=2"));
        assert!(!m.contains("\nA=1\n"));
        over.insert("A".to_string(), "100".to_string());
        let (_, ck_changed) = build_checksum(&dir, &over).unwrap();
        assert_ne!(with_over, ck_changed);
        fs::remove_dir_all(&dir).ok();
    }
}