bot-bottle/bot_bottle/resources.py

"""Locate build-time resources whether bot-bottle runs from a source
checkout or an installed wheel.

The gateway / infra / orchestrator images are built from a Docker (or Apple
`container`) build context that must contain the `bot_bottle` package,
`pyproject.toml`, and the root-level Dockerfiles as siblings. In a source
checkout that context is simply the repo root, one level above the package.
An installed wheel has no repo root: the same root-level files are shipped
inside the package under ``bot_bottle/_resources/`` (see ``setup.py``), and a
repo-root-shaped build context is staged on demand into the app-data dir.

``build_root()`` is the single source of truth — it returns a directory laid
out like a repo root (has ``bot_bottle/``, ``pyproject.toml``, the
Dockerfiles, ``nix/``, ``scripts/``). Every caller that needs a build
context, a Dockerfile path, the nix netpool module, or the netpool script
derives from it, so checkout and wheel installs share one downstream path.
"""

from __future__ import annotations

import fcntl
import hashlib
import os
import shutil
import tempfile
from pathlib import Path

from .paths import bot_bottle_root

_PKG = Path(__file__).resolve().parent          # …/bot_bottle
_CHECKOUT_ROOT = _PKG.parent                     # repo root in a checkout
_BUNDLED = _PKG / "_resources"                   # wheel-shipped copies

# Root-level files bundled into the wheel under ``_resources/`` (paths are
# relative to the checkout root, and preserved verbatim under ``_resources/``
# and in the staged build root). ``setup.py`` copies exactly this set; keep
# the two lists in sync (``test_resources`` guards that every entry exists).
BUNDLED_RESOURCES: tuple[str, ...] = (
    "pyproject.toml",
    "Dockerfile.gateway",
    "Dockerfile.orchestrator",
    "Dockerfile.orchestrator.fc",
    "nix/firecracker-netpool.nix",
    "scripts/firecracker-netpool.sh",
)

# Present at a checkout root, never in a bare installed package — the cheap
# tell for which layout we're in.
_CHECKOUT_MARKER = "Dockerfile.gateway"


class ResourceError(RuntimeError):
    """Build resources are missing from the install (corrupt/partial wheel)."""


def is_source_checkout() -> bool:
    """True when running from a source tree (the root Dockerfiles sit beside
    the package); False from an installed wheel."""
    return (_CHECKOUT_ROOT / _CHECKOUT_MARKER).is_file()


def build_root() -> Path:
    """A directory shaped like a repo root: ``bot_bottle/``, ``pyproject.toml``,
    the root Dockerfiles, ``nix/``, and ``scripts/``.

    A checkout returns the repo root itself (no copying). An installed wheel
    returns a staged copy under the app-data dir, materialized once and reused.
    The stage is keyed by a digest of the installed package + bundled resources
    (not the distribution version), so a force-reinstall of a newer commit that
    keeps ``version = 0.1.0`` still rebuilds instead of reusing a stale tree."""
    if is_source_checkout():
        return _CHECKOUT_ROOT
    return _stage_build_root()


def dockerfile(name: str) -> Path:
    """Absolute path to a root-level Dockerfile, e.g. ``Dockerfile.gateway``."""
    return build_root() / name


def nix_netpool_module() -> Path:
    """Absolute path to the firecracker netpool NixOS module."""
    return build_root() / "nix" / "firecracker-netpool.nix"


def netpool_script() -> Path:
    """Absolute path to the firecracker netpool bring-up script."""
    return build_root() / "scripts" / "firecracker-netpool.sh"


def _content_digest() -> str:
    """A 16-hex digest of the installed package + bundled resources.

    Keys the staged build root by *content*, so a force-reinstall over the same
    version string (the installer defaults to a git branch + ``pipx install
    --force``, and ``version`` stays ``0.1.0``) yields a different key and
    re-stages, rather than reusing an old commit's tree. ``_PKG`` already
    contains ``_resources``, so walking it covers both."""
    h = hashlib.sha256()
    for path in sorted(_PKG.rglob("*")):
        if not path.is_file() or "__pycache__" in path.parts or path.suffix == ".pyc":
            continue
        h.update(str(path.relative_to(_PKG)).encode())
        h.update(b"\0")
        h.update(path.read_bytes())
    return h.hexdigest()[:16]


def _stage_build_root() -> Path:
    """Materialize a repo-root-shaped build context from the installed wheel's
    bundled resources, keyed by content digest. Idempotent and concurrency-safe:
    a file lock serializes staging, a partial/stale tree is replaced, and the
    finished tree is published with an atomic rename."""
    if not _BUNDLED.is_dir():
        raise ResourceError(
            "bot-bottle build resources are missing from this install "
            f"(expected {_BUNDLED}). Reinstall the package."
        )
    base = bot_bottle_root() / "build-root"
    base.mkdir(parents=True, exist_ok=True)
    dest = base / _content_digest()
    if (dest / ".complete").is_file():
        return dest

    # Serialize staging across processes: a concurrent `start` after an install
    # must not race on the shared tree. The lock is held only around stage +
    # atomic publish; the fast path above never blocks.
    with open(base / ".stage.lock", "w", encoding="utf-8") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        if (dest / ".complete").is_file():  # another process staged while we waited
            return dest
        # Stage into a private temp dir on the same filesystem, then publish by
        # rename — never populate a shared path other processes might read.
        staging = Path(tempfile.mkdtemp(prefix=".staging-", dir=base))
        try:
            # The package itself, minus caches and the bundled-resource copies,
            # so the staged ``bot_bottle/`` matches a checkout's (keeps the
            # firecracker infra-artifact hash stable across checkout and wheel).
            shutil.copytree(
                _PKG,
                staging / "bot_bottle",
                ignore=shutil.ignore_patterns("__pycache__", "*.pyc", "_resources"),
            )
            # The bundled root files, restored to their checkout-relative layout.
            for rel in BUNDLED_RESOURCES:
                dst = staging / rel
                dst.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(_BUNDLED / rel, dst)
            (staging / ".complete").write_text("")
            # Replace any partial leftover for this digest (safe: we hold the
            # lock), then publish atomically.
            if dest.exists():
                shutil.rmtree(dest)
            os.replace(staging, dest)
            staging = None  # published; nothing to clean up
        finally:
            if staging is not None:
                shutil.rmtree(staging, ignore_errors=True)
    return dest