internal/agent/imagecache/imagecache.go
Ref: Size: 18.2 KiB History
// Package imagecache downloads and verifies content-addressed base images
// (decoded to raw in-process, LRU-evicted beyond MaxBytes). Layout:
// <dir>/<sha256>.raw — keyed by checksum (spec). LRU: a hit refreshes the
// file's mtime, and after every successful Ensure the oldest .raw images
// beyond MaxBytes are evicted (never the one just ensured).
//
// Eviction is safe for running VMs: PrepareRootDisk copies (reflink) the base, so
// nothing references it after create. With concurrent per-VM creates there is a
// narrow window where one VM's evict can remove a base another VM just resolved
// but has not yet copied — it requires the two in-flight images to exceed
// MaxBytes between them, since evict removes least-recently-used first and a
// just-ensured image carries the newest mtime. The cost is one failed create
// attempt, retried on the next tick. Close it with a recency floor in evict if
// it ever bites in practice.
package imagecache
import (
"bytes"
"compress/gzip"
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"os"
"path/filepath"
"sort"
"time"
"github.com/a73x/eitri/internal/agent/permanent"
"github.com/a73x/eitri/internal/names"
"github.com/lima-vm/go-qcow2reader"
"github.com/lima-vm/go-qcow2reader/convert"
"github.com/lima-vm/go-qcow2reader/image"
"github.com/lima-vm/go-qcow2reader/image/qcow2"
"github.com/lima-vm/go-qcow2reader/image/raw"
"golang.org/x/sync/singleflight"
)
// httpDoer is the slice of *http.Client the cache needs to fetch images. It is
// a field on Cache, not a package global, so a test can stub the transport —
// exercising slow, hung, and error responses without reaching the network.
type httpDoer interface {
Do(*http.Request) (*http.Response, error)
}
type Cache struct {
dir string
http httpDoer
// MaxBytes caps the summed size of cached images; 0 disables eviction.
// Best-effort: the just-ensured image is never evicted, even if it alone
// exceeds the cap. Accounting uses APPARENT size (fi.Size), which
// overstates sparse raw images — the error direction is safe (evicts too
// eagerly, never too late), but a huge-virtual-size image can pin the
// cache over cap; use block-based accounting if that ever bites.
MaxBytes int64
// progressStep is the package constant, held per-Cache so a test can watch
// a download advance over a handful of kilobytes instead of gigabytes.
// Written once at construction and only read after.
progressStep int64
// fetching collapses concurrent Ensure calls for the same image into one
// download+decode. Per-VM reconcile workers made creates concurrent, so a
// fleet rolling out one image now fetches it from every VM's worker at once;
// without this that is N identical multi-GB downloads and N decodes, all
// racing to rename onto the same final path.
fetching singleflight.Group
}
// New returns a cache rooted at dir. Its HTTP client carries a generous timeout
// so large images on slow links still complete, but a hung connection cannot
// stall a reconcile worker forever.
func New(dir string) *Cache {
return &Cache{dir: dir, http: &http.Client{Timeout: 10 * time.Minute}, progressStep: defaultProgressStep}
}
// defaultProgressStep is how much has to move before a download says so again.
// A callback per Read would fire thousands of times for one image; every 32 MiB
// is often enough that a watcher sees the number climb and rare enough that
// what it feeds — a report row published under a lock — costs nothing.
const defaultProgressStep = 32 << 20
// progressWriter counts bytes past it and calls report as they pass, at most
// once per step. It writes nothing: it rides the download's existing
// MultiWriter, beside the file and the hash, so nothing in the copy path has to
// know it is there.
//
// total is what the server said the body is, and it is <= 0 whenever the server
// declined to say — a chunked response has no length. Consumers render that as
// "how far", not "how far of what".
type progressWriter struct {
total int64
step int64
done int64
marked int64
report func(done, total int64)
}
func (w *progressWriter) Write(p []byte) (int, error) {
w.done += int64(len(p))
if w.done-w.marked >= w.step {
w.marked = w.done
w.report(w.done, w.total)
}
return len(p), nil
}
// fetch downloads url into a temp file in the cache dir and verifies its
// sha256. Returns the temp path; the caller owns renaming or removing it.
//
// progress, when non-nil, is told how far the body has come as it comes. It is
// called from the copy, so it must not block: a slow callback is a slow
// download.
func (c *Cache) fetch(ctx context.Context, url, sha string, progress func(done, total int64)) (string, error) {
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
if err != nil {
return "", err
}
resp, err := c.http.Do(req)
if err != nil {
return "", err
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
return "", fmt.Errorf("download %s: HTTP %d", url, resp.StatusCode)
}
tmp, err := os.CreateTemp(c.dir, "download-*")
if err != nil {
return "", err
}
h := sha256.New()
sink := []io.Writer{tmp, h}
if progress != nil {
sink = append(sink, &progressWriter{total: resp.ContentLength, step: c.progressStep, report: progress})
}
if _, err := io.Copy(io.MultiWriter(sink...), resp.Body); err != nil {
tmp.Close()
os.Remove(tmp.Name())
return "", err
}
if err := tmp.Close(); err != nil {
os.Remove(tmp.Name())
return "", err
}
if got := hex.EncodeToString(h.Sum(nil)); got != sha {
os.Remove(tmp.Name())
return "", fmt.Errorf("checksum mismatch for %s: got %s want %s", url, got, sha)
}
return tmp.Name(), nil
}
// Ensure returns the local path to the raw base image for sha, fetching and
// decoding it if the cache does not already hold it.
//
// Concurrent calls for the SAME sha are collapsed into one fetch: the first
// caller does the work and the rest wait for its result. Images are content-
// addressed, so a shared result is by definition the right one.
//
// Two consequences of sharing, both tolerable because the reconcile loop is
// level-triggered and simply retries on the next tick: the shared fetch is
// bounded by the FIRST caller's context, so if that caller's pass times out the
// waiters inherit its error; and a waiter whose own context expires first stops
// waiting without cancelling the fetch, which continues for the others.
//
// progress, when non-nil, is told how far the download has come. It belongs to
// the caller that DOES the work: a caller collapsed onto another's fetch is not
// the one reading the body, so its callback never fires and it sees only that
// an image is being fetched. That is the honest reading — the bytes are not
// its download — and it is why the caller's own wording has to stand on its
// own without any numbers in it.
func (c *Cache) Ensure(ctx context.Context, url, sha string, progress func(done, total int64)) (string, error) {
// Guard path traversal: sha becomes part of the cache file path. Checked
// before the singleflight so a bad digest can never key an entry.
if !names.IsSHA256Hex(sha) {
return "", fmt.Errorf("invalid sha256: %q", sha)
}
ch := c.fetching.DoChan(sha, func() (any, error) { return c.ensureOnce(ctx, url, sha, progress) })
select {
case r := <-ch:
if r.Err != nil {
return "", r.Err
}
return r.Val.(string), nil
case <-ctx.Done():
return "", ctx.Err()
}
}
// ensureOnce is the un-deduplicated body of Ensure: at most one runs per sha at
// a time.
func (c *Cache) ensureOnce(ctx context.Context, url, sha string, progress func(done, total int64)) (string, error) {
final := filepath.Join(c.dir, sha+".raw")
if _, err := os.Stat(final); err == nil {
// Hit: refresh recency so frequently-used images sort as recent.
now := time.Now()
_ = os.Chtimes(final, now, now)
c.evict(final)
return final, nil
}
tmp, err := c.fetch(ctx, url, sha, progress)
if err != nil {
return "", err
}
// A no-op once the raw path has renamed tmp away; every other path needs it.
defer os.Remove(tmp)
// Two authorities, and which answers is decided by what the question is.
// sniffFile answers a TRANSPORT question — is this gzipped? — because gzip
// wraps a disk image rather than being one. Everything else is a disk format,
// and go-qcow2reader is the authority there: it parses qcow2, vmdk, vhdx,
// vdi, parallels, vpc and asif headers, and falls back to raw for what none
// of them claim. That fallback is the same default-to-raw rule this package
// started with, now with six fewer formats misfiling under it.
format, err := sniffFile(tmp)
if err != nil {
return "", fmt.Errorf("imagecache: identify %s: %w", url, err)
}
if format != formatGzip {
passthrough, err := isRawImage(url, tmp)
if err != nil {
return "", err
}
if passthrough {
// The verified download IS the base image. Renaming it onto final is
// atomic (same directory) and costs neither a decode pass nor a second
// copy of a multi-gigabyte file.
if err := os.Chmod(tmp, 0o644); err != nil {
return "", fmt.Errorf("imagecache: chmod image: %w", err)
}
if err := os.Rename(tmp, final); err != nil {
return "", fmt.Errorf("imagecache: rename to final: %w", err)
}
c.evict(final)
return final, nil
}
}
// Materialize into a UNIQUE temp file first, then rename onto final atomically
// so a crash mid-write cannot leave a corrupt file at the final path. The temp
// name is unique (not a fixed "<sha>.converting") so two concurrent Ensure
// calls for the same sha can never converge on one file and corrupt it. The
// suffix deliberately does NOT end in ".raw", so evict()'s "*.raw" glob never
// sees a half-built temp.
cf, err := os.CreateTemp(c.dir, sha+".raw.converting-*")
if err != nil {
return "", fmt.Errorf("imagecache: create temp: %w", err)
}
converting := cf.Name()
_ = cf.Close()
_ = os.Remove(converting) // reserve the unique name; let the writer create it fresh
defer os.Remove(converting)
// Gzip returned raw above only if it was not gzip, so this is the whole
// space: unwrap the transport, or decode the container.
if format == formatGzip {
err = decompress(url, tmp, converting)
} else {
err = decode(url, tmp, converting)
}
if err != nil {
return "", err
}
if err := os.Rename(converting, final); err != nil {
return "", fmt.Errorf("imagecache: rename to final: %w", err)
}
c.evict(final)
return final, nil
}
// decompress writes the gzipped raw image at src to dst, punching the long runs
// of zeros a cloud image is mostly made of into holes.
func decompress(url, src, dst string) error {
in, err := os.Open(src)
if err != nil {
return fmt.Errorf("imagecache: open download: %w", err)
}
defer in.Close()
zr, err := gzip.NewReader(in)
if err != nil {
// The download's sha already matched, so a broken gzip header is a
// property of the published artifact: re-fetching returns the same bytes
// and fails the same way.
return permanent.Errorf("imagecache: gzip header: %v", err)
}
defer zr.Close()
// Gzip wraps a RAW image. Sniff what is inside before writing a byte: a
// .qcow2.gz would otherwise decompress into a cache entry named <sha>.raw
// holding qcow2 bytes, and the guest would fail to boot for a reason nothing
// in the error names — exactly the illegible failure this package exists to
// prevent. The peeked bytes go back in front of the stream so the writer
// still sees the whole image. Failing before the output file is created is
// what leaves no temp behind.
head := make([]byte, magicLen)
n, err := io.ReadFull(zr, head)
if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) {
return fmt.Errorf("imagecache: read decompressed header: %w", err)
}
head = head[:n]
if sniff(head) == formatQcow2 {
// Permanent for the same reason a broken header is: the sha matched, so
// these bytes are what the publisher published.
return permanent.Errorf("imagecache: %s decompresses to a qcow2 image; gzip must wrap a RAW image — publish the qcow2 ungzipped (eitri decodes it) or gzip the raw", url)
}
out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o644)
if err != nil {
return fmt.Errorf("imagecache: create image: %w", err)
}
if err := writeSparse(out, io.MultiReader(bytes.NewReader(head), zr)); err != nil {
out.Close()
return fmt.Errorf("imagecache: decompress: %w", err)
}
if err := out.Close(); err != nil {
return fmt.Errorf("imagecache: close image: %w", err)
}
return nil
}
// isRawImage reports whether the download at path is already a raw image and
// can be used as-is. Anything go-qcow2reader recognizes as a container — qcow2,
// vmdk, vhdx, vdi, parallels, vpc, asif — needs decoding first.
func isRawImage(url, path string) (bool, error) {
img, err := openImage(url, path)
if err != nil {
return false, err
}
defer img.Close()
return img.Type() == raw.Type, nil
}
// openImage parses the download's header. A header this cannot parse is a
// property of the artifact, not of the attempt: the checksum already matched,
// so a re-fetch returns the same bytes and fails identically.
func openImage(url, path string) (*imageFile, error) {
f, err := os.Open(path)
if err != nil {
return nil, fmt.Errorf("imagecache: open download: %w", err)
}
img, err := qcow2reader.Open(readerAtOnly{f})
if err != nil {
f.Close()
return nil, permanent.Errorf("imagecache: %s: %v", url, err)
}
// A base image stands alone: everything the guest disk holds comes from the
// bytes whose checksum was verified, and nothing is read in from the host
// filesystem beside them. qcow2 is the only format go-qcow2reader follows an
// external reference for — the rest are stubs that decode nothing — so this
// is the whole of that boundary. Checked before Readable() so the refusal
// speaks about the image rather than repeating the path it asked for.
if q, ok := img.(*qcow2.Qcow2); ok && (q.BackingFileOffset != 0 || q.BackingFile != "") {
img.Close()
f.Close()
return nil, permanent.Errorf("imagecache: %s declares a backing file; eitri boots only self-contained images — flatten it (qemu-img convert) and publish the result", url)
}
if err := img.Readable(); err != nil {
img.Close()
f.Close()
return nil, permanent.Errorf("imagecache: %s is a %s image eitri cannot read: %v", url, img.Type(), err)
}
return &imageFile{Image: img, f: f}, nil
}
// readerAtOnly hands the decoder the download's bytes and nothing else. The
// qcow2 decoder resolves a relative backing path against the name of the file
// it was given, and reaches that name through an interface an *os.File
// satisfies; a value that offers only ReaderAt cannot answer it. Closing stays
// with imageFile, which owns the file.
type readerAtOnly struct{ io.ReaderAt }
// imageFile ties the decoded image to the file underneath it, so one Close
// releases both.
type imageFile struct {
image.Image
f *os.File
}
func (i *imageFile) Close() error {
err := i.Image.Close()
if ferr := i.f.Close(); err == nil {
err = ferr
}
return err
}
// decode writes the disk image at src out as raw at dst. The decoder skips
// unallocated and all-zero extents rather than materialising them, so the
// result is sparse for the same reason the gzip path's writer produces one —
// but exactly, from the image's own extent map, instead of by scanning for
// zeros.
//
// The closing Truncate is load-bearing in the same way it is there: Convert
// never writes a trailing zero extent, so without it an image ending in zeros
// lands shorter than it claims, and PrepareRootDisk sizes a guest's disk
// against that number.
func decode(url, src, dst string) error {
img, err := openImage(url, src)
if err != nil {
return err
}
defer img.Close()
out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o644)
if err != nil {
return fmt.Errorf("imagecache: create image: %w", err)
}
if err := convert.Convert(out, img, convert.Options{}); err != nil {
out.Close()
return fmt.Errorf("imagecache: decode %s: %w", url, err)
}
if err := out.Truncate(img.Size()); err != nil {
out.Close()
return fmt.Errorf("imagecache: size image: %w", err)
}
if err := out.Close(); err != nil {
return fmt.Errorf("imagecache: close image: %w", err)
}
return nil
}
// SweepTemps removes abandoned download and decode temporaries — the two temp
// forms Ensure writes before its atomic rename. Neither matches evict's "*.raw"
// glob, so nothing else ever reclaims them, and an agent killed mid-fetch (the
// deploy path SIGTERMs it) leaves one behind per in-flight image.
//
// MUST be called at agent start, before any Ensure: it cannot distinguish an
// abandoned temp from one a concurrent fetch is still writing. At process start
// this agent has no fetch in flight, and the deploy script confirms the previous
// agent is gone before starting the new one, so every temp present is abandoned.
// Best-effort: a file that cannot be removed is left for the next start.
func (c *Cache) SweepTemps() {
for _, pat := range []string{"download-*", "*.raw.converting-*"} {
matches, err := filepath.Glob(filepath.Join(c.dir, pat))
if err != nil {
continue
}
for _, p := range matches {
if err := os.Remove(p); err == nil {
slog.Info("imagecache: swept abandoned temp", "path", p)
}
}
}
}
// evict removes least-recently-used .raw images until the cache fits
// MaxBytes, never touching keep (the image the current create is about to
// use). Failures are logged-by-omission best-effort: eviction must never
// fail an Ensure that already succeeded.
func (c *Cache) evict(keep string) {
if c.MaxBytes <= 0 {
return
}
entries, err := filepath.Glob(filepath.Join(c.dir, "*.raw"))
if err != nil {
return
}
type img struct {
path string
size int64
mtime time.Time
}
var imgs []img
var total int64
for _, p := range entries {
fi, err := os.Stat(p)
if err != nil {
continue
}
total += fi.Size()
if p != keep {
imgs = append(imgs, img{p, fi.Size(), fi.ModTime()})
}
}
sort.Slice(imgs, func(i, j int) bool { return imgs[i].mtime.Before(imgs[j].mtime) })
for _, im := range imgs {
if total <= c.MaxBytes {
return
}
if os.Remove(im.path) == nil {
total -= im.size
// The one observable trace of eviction: cap-thrash (an oversized
// image forcing constant re-downloads) is invisible without it.
slog.Info("imagecache: evicted", "path", im.path, "size", im.size)
}
}
}