244 lines
9.4 KiB
Python
244 lines
9.4 KiB
Python
"""Orchestrator: materialise + cache + verify boot files for QEMU boards.
|
|
|
|
The provider is the only object QEMU launch code needs to hold a
|
|
reference to. It owns the on-disk cache, the per-set asyncio locks
|
|
(so concurrent boot requests for the same board collapse into one
|
|
download), and the verify-then-rename ladder that gives us atomic
|
|
"a file at this path means it's correct".
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from .downloader import AssetDownloader
|
|
from .errors import BootImageError
|
|
from .integrity import decompress_zstd, sha256_file, verify_sha256
|
|
from .manifest import BootImageSpec, BootImagesManifest, ImageSetSpec
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BootImageProvider:
|
|
"""Materialise and cache boot files on demand, idempotently.
|
|
|
|
Operational notes:
|
|
|
|
* First call for a set_id pays the download cost. Subsequent calls
|
|
do an mtime/size/SHA256 cache check and return immediately if
|
|
the files are still valid.
|
|
* Concurrent ``get()`` calls for the same set_id serialise on a
|
|
per-set ``asyncio.Lock``; the second caller observes a populated
|
|
cache and skips the download.
|
|
* Different set_ids materialise in parallel.
|
|
* On SHA256 mismatch the file is left in a temp dir (never reaches
|
|
``target_path``), so a corrupt download cannot poison the cache
|
|
for the next process.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
manifest: BootImagesManifest,
|
|
downloader: AssetDownloader,
|
|
cache_dir: Path,
|
|
):
|
|
self._manifest = manifest
|
|
self._downloader = downloader
|
|
self._cache_dir = cache_dir
|
|
self._cache_dir.mkdir(parents=True, exist_ok=True)
|
|
self._locks: dict[str, asyncio.Lock] = {}
|
|
self._locks_guard = asyncio.Lock()
|
|
|
|
# ── Public API ────────────────────────────────────────────────────────
|
|
|
|
@property
|
|
def cache_dir(self) -> Path:
|
|
return self._cache_dir
|
|
|
|
@property
|
|
def manifest(self) -> BootImagesManifest:
|
|
return self._manifest
|
|
|
|
async def get(self, set_id: str) -> dict[str, Path]:
|
|
"""Return ``{image_name: absolute_path}`` for ``set_id``.
|
|
|
|
Downloads + verifies + (when applicable) decompresses on first
|
|
call. Subsequent calls are cache hits. Concurrent callers
|
|
serialise on a per-set lock.
|
|
"""
|
|
spec = self._manifest.get(set_id)
|
|
lock = await self._lock_for(set_id)
|
|
async with lock:
|
|
return await self._materialise(spec)
|
|
|
|
async def warmup(self, set_id: str) -> None:
|
|
"""Best-effort prefetch. Logs warnings on failure but never
|
|
raises — designed to be fire-and-forget from a lifespan hook
|
|
so a transient network blip doesn't break process startup.
|
|
"""
|
|
try:
|
|
await self.get(set_id)
|
|
logger.info("[boot-images] warmup complete for %r", set_id)
|
|
except BootImageError as exc:
|
|
logger.warning(
|
|
"[boot-images] warmup for %r failed: %s", set_id, exc,
|
|
)
|
|
|
|
async def warmup_all(self) -> None:
|
|
"""Concurrent warmup of every set declared in the manifest."""
|
|
await asyncio.gather(
|
|
*(self.warmup(s) for s in self._manifest.image_sets),
|
|
return_exceptions=False, # warmup() already swallows
|
|
)
|
|
|
|
def is_cached(self, set_id: str) -> bool:
|
|
"""Sync probe used by health/status endpoints.
|
|
|
|
Uses the same sidecar-SHA check ``_is_valid_cached`` does so
|
|
a manifest bump correctly reports "not cached yet" until the
|
|
next ``get()`` re-materialises the file.
|
|
"""
|
|
try:
|
|
spec = self._manifest.get(set_id)
|
|
except BootImageError:
|
|
return False
|
|
set_dir = self._cache_dir / set_id
|
|
return all(
|
|
self._is_valid_cached(set_dir / img.name, img) for img in spec.images
|
|
)
|
|
|
|
# ── Internals ─────────────────────────────────────────────────────────
|
|
|
|
async def _lock_for(self, set_id: str) -> asyncio.Lock:
|
|
async with self._locks_guard:
|
|
lock = self._locks.get(set_id)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
self._locks[set_id] = lock
|
|
return lock
|
|
|
|
async def _materialise(self, spec: ImageSetSpec) -> dict[str, Path]:
|
|
set_dir = self._cache_dir / spec.id
|
|
set_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
out: dict[str, Path] = {}
|
|
for img in spec.images:
|
|
target = set_dir / img.name
|
|
if await asyncio.to_thread(self._is_valid_cached, target, img):
|
|
logger.debug("[boot-images] cache hit %s", target)
|
|
out[img.name] = target
|
|
continue
|
|
logger.info(
|
|
"[boot-images] fetching %s/%s (asset_id=%s%s)",
|
|
spec.id,
|
|
img.name,
|
|
img.asset_id,
|
|
f", version={img.version}" if img.version else "",
|
|
)
|
|
await self._fetch_and_verify(img, target)
|
|
out[img.name] = target
|
|
return out
|
|
|
|
@staticmethod
|
|
def _sidecar(target: Path) -> Path:
|
|
"""Sidecar file recording the SHA256 of the cached payload.
|
|
|
|
Written atomically (temp + rename) after a successful
|
|
download+verify, read on every cache-validity probe. Lets the
|
|
provider detect manifest SHA bumps without re-hashing
|
|
multi-GiB files on every container start.
|
|
"""
|
|
return target.parent / f"{target.name}.sha256"
|
|
|
|
@classmethod
|
|
def _is_valid_cached(cls, path: Path, spec: BootImageSpec) -> bool:
|
|
"""O(1) cache-hit probe — presence + size + sidecar SHA match.
|
|
|
|
We deliberately do NOT re-hash the file on every probe. The
|
|
5.4 GiB Pi 3 SD image takes ~30 s to SHA256, and that cost
|
|
would be paid on every container boot pre-warm AND every user
|
|
request that triggers ``provider.get()``.
|
|
|
|
Instead, after a successful materialise we write a sidecar
|
|
``<name>.sha256`` containing the expected hash and trust it on
|
|
subsequent probes. A manifest SHA bump invalidates the sidecar
|
|
even if the size is unchanged (e.g. an in-place SD image edit
|
|
that ends up the exact same byte count), forcing a re-fetch.
|
|
|
|
If the sidecar is missing (legacy cache from before this
|
|
change, or operator tampering) the file is treated as invalid
|
|
and re-fetched. Manual operators who want to inject a file can
|
|
write the sidecar themselves: ``sha256sum file | cut -d' ' -f1
|
|
> file.sha256``.
|
|
"""
|
|
if not path.is_file():
|
|
return False
|
|
if path.stat().st_size != spec.size_bytes:
|
|
return False
|
|
sidecar = cls._sidecar(path)
|
|
if not sidecar.is_file():
|
|
return False
|
|
try:
|
|
recorded = sidecar.read_text(encoding="ascii").strip().lower()
|
|
except OSError:
|
|
return False
|
|
return recorded == spec.sha256.lower()
|
|
|
|
async def _fetch_and_verify(
|
|
self, img: BootImageSpec, target: Path,
|
|
) -> None:
|
|
if img.compressed is None:
|
|
await self._downloader.fetch(img.asset_id, target)
|
|
await asyncio.to_thread(
|
|
verify_sha256, target, img.sha256, label=img.name,
|
|
)
|
|
else:
|
|
# Compressed path: download → verify wire-format sha →
|
|
# decompress → verify decompressed sha → atomic rename to
|
|
# final cache slot.
|
|
with tempfile.TemporaryDirectory(
|
|
dir=target.parent, prefix=".staging-",
|
|
) as staging:
|
|
staging_dir = Path(staging)
|
|
compressed_path = (
|
|
staging_dir / f"{img.name}.{img.compressed.encoding}"
|
|
)
|
|
await self._downloader.fetch(img.asset_id, compressed_path)
|
|
await asyncio.to_thread(
|
|
verify_sha256,
|
|
compressed_path,
|
|
img.compressed.sha256,
|
|
label=f"{img.name} (compressed)",
|
|
)
|
|
decoded = staging_dir / img.name
|
|
if img.compressed.encoding == "zstd":
|
|
await asyncio.to_thread(decompress_zstd, compressed_path, decoded)
|
|
else:
|
|
raise BootImageError(
|
|
f"unsupported compression {img.compressed.encoding!r}"
|
|
)
|
|
await asyncio.to_thread(
|
|
verify_sha256,
|
|
decoded,
|
|
img.sha256,
|
|
label=f"{img.name} (decompressed)",
|
|
)
|
|
await asyncio.to_thread(decoded.replace, target)
|
|
# Record the expected SHA next to the file so future cache
|
|
# probes can detect manifest bumps without re-hashing the
|
|
# whole file. Sidecar write is atomic (temp + rename) so a
|
|
# process crash mid-write can't leave a half-written hash.
|
|
await asyncio.to_thread(self._write_sidecar, target, img.sha256)
|
|
|
|
@classmethod
|
|
def _write_sidecar(cls, target: Path, sha256: str) -> None:
|
|
sidecar = cls._sidecar(target)
|
|
tmp = sidecar.with_suffix(sidecar.suffix + ".tmp")
|
|
tmp.write_text(sha256.lower() + "\n", encoding="ascii")
|
|
tmp.replace(sidecar)
|