diff --git a/bot_bottle/backend/firecracker/image_builder.py b/bot_bottle/backend/firecracker/image_builder.py index c470be6..92b8d27 100644 --- a/bot_bottle/backend/firecracker/image_builder.py +++ b/bot_bottle/backend/firecracker/image_builder.py @@ -16,10 +16,14 @@ builder (booted from this same image on its own TAP) later. from __future__ import annotations +import fcntl import hashlib +import os import shutil import subprocess +from contextlib import contextmanager from pathlib import Path +from typing import Generator from ...log import die, info from . import infra_vm, util @@ -54,22 +58,46 @@ def build_agent_rootfs_dir( silent-failure image at build time rather than at first agent use.""" digest = _dockerfile_hash(dockerfile) base = util.cache_dir() / "rootfs" / f"agent-{digest}" - ready = base / ".bb-ready" - if ready.is_file(): + if (base / ".bb-ready").is_file(): info(f"using cached agent rootfs {base.name}") return base - if base.exists(): + # Serialize builds: the infra VM's buildah store + this cache dir are + # shared, so concurrent `start`s must not build into them at once. The + # lock covers the cache lookup + build + atomic publish; the ready + # fast-path above takes no lock. + with _build_lock(): + if (base / ".bb-ready").is_file(): # another build finished while we waited + info(f"using cached agent rootfs {base.name}") + return base + # Build into a temp dir and publish by atomic rename, so a partial + # build is never visible as `agent-`. + staging = util.cache_dir() / "rootfs" / f".building-{digest}" + shutil.rmtree(staging, ignore_errors=True) + staging.mkdir(parents=True) + info(f"building agent image {image_tag!r} in the infra VM") + _build_in_infra(dockerfile, staging, smoke_test, digest) + util.inject_guest_boot(staging) + (staging / ".bb-ready").write_text("ok\n") shutil.rmtree(base, ignore_errors=True) - base.mkdir(parents=True) - - info(f"building agent image {image_tag!r} in the infra VM") - _build_in_infra(dockerfile, base, smoke_test, digest) - util.inject_guest_boot(base) - ready.write_text("ok\n") + os.rename(staging, base) return base +@contextmanager +def _build_lock() -> Generator[None, None, None]: + """Host-level exclusive lock serializing agent-image builds (shared infra + buildah store + cache dir). flock auto-releases on a crash.""" + lock_path = util.cache_dir() / "rootfs" / ".build.lock" + lock_path.parent.mkdir(parents=True, exist_ok=True) + handle = open(lock_path, "w", encoding="utf-8") + try: + fcntl.flock(handle, fcntl.LOCK_EX) + yield + finally: + handle.close() + + def _build_in_infra( dockerfile: Path, base: Path, smoke_test: tuple[str, ...], digest: str, ) -> None: @@ -80,21 +108,27 @@ def _build_in_infra( key, ip = infra.private_key, infra.guest_ip tag = f"bot-bottle-agent-build-{digest}" ctx = f"/tmp/agent-build-{digest}" + smoke_ctr, export_ctr = f"{tag}-smoke", f"{tag}-export" + + def _cleanup() -> None: + # Remove only THIS build's working containers/image/context — never + # `buildah rm -a`, which would nuke a concurrent build's container. + _ssh(key, ip, + f"buildah rm {smoke_ctr} {export_ctr} >/dev/null 2>&1; " + f"buildah rmi {_STORE_FLAG} {tag} >/dev/null 2>&1; rm -rf {ctx}", + timeout=60) + + _cleanup() # clear leftovers from a crashed prior build of this digest try: - prep = _ssh(key, ip, f"rm -rf {ctx}; mkdir -p {ctx}/ctx") + prep = _ssh(key, ip, f"mkdir -p {ctx}/ctx") if prep.returncode != 0: die(f"preparing build dir in the infra VM failed: {prep.stderr.strip()}") _send_dockerfile(key, ip, dockerfile, ctx) _buildah_build(key, ip, ctx, tag) - _smoke_test(key, ip, tag, smoke_test) - _stream_rootfs(key, ip, tag, base) + _smoke_test(key, ip, tag, smoke_ctr, smoke_test) + _stream_rootfs(key, ip, tag, export_ctr, base) finally: - # Keep the ephemeral infra rootfs from accumulating build state. Builds - # are serial per launch, so removing all working containers is safe. - _ssh(key, ip, - f"buildah rm {_STORE_FLAG} -a >/dev/null 2>&1; " - f"buildah rmi {_STORE_FLAG} {tag} >/dev/null 2>&1; rm -rf {ctx}", - timeout=60) + _cleanup() def _ssh(private_key: Path, guest_ip: str, script: str, @@ -126,16 +160,18 @@ def _buildah_build(private_key: Path, guest_ip: str, ctx: str, tag: str) -> None die(f"buildah build in the infra VM failed:\n{tail}") -def _smoke_test(private_key: Path, guest_ip: str, tag: str, argv: tuple[str, ...]) -> None: +def _smoke_test(private_key: Path, guest_ip: str, tag: str, ctr: str, + argv: tuple[str, ...]) -> None: """Run the provider's smoke argv inside the freshly built image (`buildah run`, which uses the image's own PATH), failing the build - loudly if the CLI is a broken stub. No-op without a declared test.""" + loudly if the CLI is a broken stub. No-op without a declared test. Uses a + named working container (`ctr`) so cleanup is scoped to this build.""" if not argv: return cmd = ( - f"set -e; ctr=$(buildah from {_STORE_FLAG} {tag}); " - f"buildah run {_BUILD_FLAGS} \"$ctr\" -- {' '.join(argv)}; rc=$?; " - f"buildah rm \"$ctr\" >/dev/null 2>&1 || true; exit $rc" + f"set -e; buildah from {_STORE_FLAG} --name {ctr} {tag} >/dev/null; " + f"buildah run {_BUILD_FLAGS} {ctr} -- {' '.join(argv)}; rc=$?; " + f"buildah rm {ctr} >/dev/null 2>&1 || true; exit $rc" ) result = _ssh(private_key, guest_ip, cmd, timeout=120) if result.returncode != 0: @@ -144,13 +180,14 @@ def _smoke_test(private_key: Path, guest_ip: str, tag: str, argv: tuple[str, ... f"({' '.join(argv)}):\n" + "\n".join(detail)) -def _stream_rootfs(private_key: Path, guest_ip: str, tag: str, base: Path) -> None: +def _stream_rootfs(private_key: Path, guest_ip: str, tag: str, ctr: str, base: Path) -> None: """`buildah mount` the built image in the infra VM and pipe its rootfs tar straight into `base` on the host (extracted as the non-root host user, so - uid 0 isn't preserved — the guest init restores /root ownership).""" + uid 0 isn't preserved — the guest init restores /root ownership). Uses a + named working container so cleanup is scoped to this build.""" export = ( - f"set -e; ctr=$(buildah from {_STORE_FLAG} {tag}); " - f"mnt=$(buildah mount {_STORE_FLAG} \"$ctr\"); " + f"set -e; buildah from {_STORE_FLAG} --name {ctr} {tag} >/dev/null; " + f"mnt=$(buildah mount {_STORE_FLAG} {ctr}); " f"tar -C \"$mnt\" -cf - ." ) ssh_proc = subprocess.Popen( diff --git a/bot_bottle/backend/firecracker/infra_vm.py b/bot_bottle/backend/firecracker/infra_vm.py index 0fdf6f0..ee957d2 100644 --- a/bot_bottle/backend/firecracker/infra_vm.py +++ b/bot_bottle/backend/firecracker/infra_vm.py @@ -17,6 +17,7 @@ surface. from __future__ import annotations +import fcntl import hashlib import os import shlex @@ -26,8 +27,10 @@ import subprocess import time import urllib.error import urllib.request +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path +from typing import Generator from ...log import die, info from ..docker import util as docker_mod @@ -134,7 +137,11 @@ def ensure_running() -> InfraVm: """Idempotent per-host singleton. Adopt the infra VM if its control plane is already healthy (a prior launcher booted it — it outlives short-lived `start` processes); otherwise clear any stale VM and boot a fresh one. - Returns a handle usable for CA fetch / git-gate provisioning.""" + Returns a handle usable for CA fetch / git-gate provisioning. + + Concurrency-safe: the cold stop/build/boot path is serialized by a host + flock, so two simultaneous first launches don't both boot on the same + rootfs/PID. The healthy fast-path takes no lock.""" slot = netpool.orch_slot() url = f"http://{slot.guest_ip}:{CONTROL_PLANE_PORT}" key = _infra_dir() / "id_ed25519" @@ -142,11 +149,31 @@ def ensure_running() -> InfraVm: info(f"adopting running infra VM at {url}") return InfraVm(guest_ip=slot.guest_ip, private_key=key) - stop() # clear a stale/hung VM holding the link before booting fresh - ensure_built() - infra = boot() - wait_for_health(infra) - return infra + with _singleton_lock(): + # Re-check under the lock: another launcher may have booted it while + # we waited for the lock (double-checked, so we adopt not re-boot). + if key.exists() and _health_ok(url): + info(f"adopting running infra VM at {url}") + return InfraVm(guest_ip=slot.guest_ip, private_key=key) + stop() # clear a stale/hung VM holding the link before booting fresh + ensure_built() + infra = boot() + wait_for_health(infra) + return infra + + +@contextmanager +def _singleton_lock() -> Generator[None, None, None]: + """Host-level exclusive lock serializing the infra VM's cold create path + (`stop`/`ensure_built`/`boot`). flock auto-releases if the launcher + crashes, so the lock is never leaked.""" + lock_path = _infra_dir() / "singleton.lock" + handle = open(lock_path, "w", encoding="utf-8") + try: + fcntl.flock(handle, fcntl.LOCK_EX) + yield + finally: + handle.close() def stop() -> None: diff --git a/tests/unit/test_firecracker_image_builder.py b/tests/unit/test_firecracker_image_builder.py index 92c949c..40e132c 100644 --- a/tests/unit/test_firecracker_image_builder.py +++ b/tests/unit/test_firecracker_image_builder.py @@ -44,7 +44,7 @@ class TestBuildAgentRootfsDir(unittest.TestCase): build.assert_called_once() # smoke_test threads through to the VM build. self.assertEqual(("claude", "--version"), build.call_args.args[2]) - inject.assert_called_once_with(out) + inject.assert_called_once() self.assertTrue((out / ".bb-ready").is_file()) def test_content_addressed_cache_key(self): @@ -59,7 +59,7 @@ class TestBuildAgentRootfsDir(unittest.TestCase): class TestSmokeTest(unittest.TestCase): def test_empty_argv_is_noop(self): with patch.object(image_builder, "_ssh") as ssh: - image_builder._smoke_test(Path("/k"), "10.0.0.1", "tag", ()) + image_builder._smoke_test(Path("/k"), "10.0.0.1", "tag", "ctr", ()) ssh.assert_not_called() def test_failed_smoke_dies(self): @@ -67,7 +67,7 @@ class TestSmokeTest(unittest.TestCase): result = subprocess.CompletedProcess([], 1, stdout="broken", stderr="") with patch.object(image_builder, "_ssh", return_value=result), \ self.assertRaises(SystemExit): - image_builder._smoke_test(Path("/k"), "10.0.0.1", "tag", ("claude", "--version")) + image_builder._smoke_test(Path("/k"), "10.0.0.1", "tag", "ctr", ("claude", "--version")) if __name__ == "__main__":