distro/src/container.py

import atexit
import hashlib
import os
import signal
import subprocess
import threading
from dataclasses import dataclass, field
from pathlib import Path

from src import log


@dataclass
class Mount:
    src: Path
    dst: str
    readonly: bool = False

    def as_arg(self) -> str:
        flag = ":ro" if self.readonly else ""
        return f"{self.src}:{self.dst}{flag}"


_active: set["Container"] = set()
_lock = threading.Lock()
_handlers_installed = False


def _install_handlers() -> None:
    global _handlers_installed
    if _handlers_installed:
        return
    _handlers_installed = True

    def cleanup(*_):
        _shutdown_all()

    atexit.register(_shutdown_all)
    for s in (signal.SIGINT, signal.SIGTERM):
        try:
            signal.signal(s, lambda *_: (_shutdown_all(), os._exit(130)))
        except (ValueError, OSError):
            pass


def _shutdown_all() -> None:
    with _lock:
        cs = list(_active)
        _active.clear()
    for c in cs:
        try:
            subprocess.run(
                ["podman", "rm", "-f", c.name],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                check=False,
            )
        except Exception:
            pass


@dataclass
class Container:
    name: str
    image: str
    mounts: list[Mount] = field(default_factory=list)
    tmpfs: list[str] = field(default_factory=list)
    network: bool = False
    extra_env: dict[str, str] = field(default_factory=dict)
    started: bool = False

    def __hash__(self) -> int:
        return id(self)

    def __eq__(self, other) -> bool:
        return self is other

    def start(self) -> None:
        if self.started:
            return
        _install_handlers()
        # Ensure no stale container.
        subprocess.run(
            ["podman", "rm", "-f", self.name],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        )
        argv = [
            "podman",
            "run",
            "-d",
            "--rm",
            "--name",
            self.name,
        ]
        for t in ("/tmp", "/dest", "/var/cache", *self.tmpfs):
            argv += ["--tmpfs", t]
        argv += ["--network=none"] if not self.network else []
        for m in self.mounts:
            m.src.mkdir(parents=True, exist_ok=True) if not m.src.exists() else None
            argv += ["-v", m.as_arg()]
        argv += [self.image, "sleep", "infinity"]
        log.debug(f"container start: {' '.join(argv)}")
        r = subprocess.run(argv, capture_output=True, text=True)
        if r.returncode != 0:
            raise RuntimeError(f"podman run failed: {r.stderr.strip()}")
        with _lock:
            _active.add(self)
        self.started = True

    def stop(self) -> None:
        if not self.started:
            return
        subprocess.run(
            ["podman", "rm", "-f", self.name],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        )
        with _lock:
            _active.discard(self)
        self.started = False

    def exec(
        self,
        argv: list[str],
        *,
        env: dict[str, str] | None = None,
        cwd: str | None = None,
        check: bool = True,
        user: str | None = None,
    ) -> int:
        if not self.started:
            raise RuntimeError("container not started")
        cmd = ["podman", "exec"]
        if cwd:
            cmd += ["-w", cwd]
        if user:
            cmd += ["-u", user]
        merged: dict[str, str] = {}
        merged.update(self.extra_env)
        if env:
            merged.update(env)
        for k, v in merged.items():
            cmd += ["-e", f"{k}={v}"]
        cmd += [self.name, *argv]
        log.debug(f"exec: {' '.join(argv)}")
        rc = subprocess.call(cmd)
        if check and rc != 0:
            raise RuntimeError(f"command failed (exit {rc}): {' '.join(argv)}")
        return rc

    def exec_shell(
        self,
        script: str,
        *,
        env: dict[str, str] | None = None,
        cwd: str | None = None,
        check: bool = True,
    ) -> int:
        return self.exec(["sh", "-ec", script], env=env, cwd=cwd, check=check)

    def cp_out(self, src_in: str, dst_host: Path) -> None:
        dst_host.parent.mkdir(parents=True, exist_ok=True)
        r = subprocess.run(
            ["podman", "cp", f"{self.name}:{src_in}", str(dst_host)],
            capture_output=True,
            text=True,
        )
        if r.returncode != 0:
            raise RuntimeError(f"podman cp failed: {r.stderr.strip()}")


def hash_dockerfile(p: Path) -> str:
    return hashlib.sha256(p.read_bytes()).hexdigest()


def image_exists(tag: str) -> bool:
    r = subprocess.run(["podman", "image", "exists", tag])
    return r.returncode == 0


def build_image(dockerfile: Path, tag: str) -> None:
    log.info(f"building container image {tag}")
    r = subprocess.run(
        ["podman", "build", "-t", tag, "-f", str(dockerfile), str(dockerfile.parent)]
    )
    if r.returncode != 0:
        raise RuntimeError("podman build failed")


def ensure_image(dockerfile: Path, tag: str, hash_file: Path) -> None:
    cur = hash_dockerfile(dockerfile) + "\n" + tag
    if hash_file.exists() and hash_file.read_text() == cur and image_exists(tag):
        log.debug(f"image {tag} up-to-date")
        return
    build_image(dockerfile, tag)
    hash_file.parent.mkdir(parents=True, exist_ok=True)
    hash_file.write_text(cur)