From 4be373074b80f5f741f7d0068176409d1607ad99 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 14:39:50 +0000 Subject: [PATCH 1/4] experimental/air: warm snapshot cache for the plain_tar path Stacked on the parallel-gzip change. Add a local warm cache for the plain_tar snapshot path, keyed by (repo, config, include_paths) under $TMPDIR/databricks/.air/: an uncompressed snapshot.tar plus a manifest of each file's size+mtime and byte range. Later runs stat the file set, copy unchanged members verbatim from the warm tar, and re-read only changed files before gzipping the upload. --no-cache bypasses it and re-packs from scratch (now also parallel-gzip, from the parent PR). The cache engages only above 64 MiB. The cache's payoff is largest when the working set does not fit the OS page cache: cold, scattered small-file reads cost seconds (6.5 s for research, 84 MB) versus ~ms to read the warm tar sequentially. When the tree is already warm in RAM, parallel gzip accounts for most of the gain and the cache adds little. Co-authored-by: Isaac --- experimental/air/cmd/run.go | 4 +- experimental/air/cmd/runsubmit.go | 4 +- experimental/air/cmd/runsubmit_test.go | 24 +- experimental/air/cmd/snapshot_cache.go | 390 ++++++++++++++++++++ experimental/air/cmd/snapshot_cache_test.go | 183 +++++++++ experimental/air/cmd/snapshot_dabs.go | 26 +- experimental/air/cmd/snapshot_package.go | 20 +- 7 files changed, 621 insertions(+), 30 deletions(-) create mode 100644 experimental/air/cmd/snapshot_cache.go create mode 100644 experimental/air/cmd/snapshot_cache_test.go diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 91b2e76065f..46f5debdb6e 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -35,6 +35,7 @@ func newRunCommand() *cobra.Command { overrides []string dryRun bool idempotencyKey string + noCache bool ) cmd := &cobra.Command{ @@ -78,6 +79,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config without submitting") cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used") + cmd.Flags().BoolVar(&noCache, "no-cache", false, "Bypass the local snapshot cache and re-pack the code tarball from scratch") _ = cmd.MarkFlagRequired("file") // --dry-run only validates the config locally, so it needs no workspace. @@ -114,7 +116,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so } w := cmdctx.WorkspaceClient(ctx) - runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey, !jsonOut) + runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey, !jsonOut, noCache) if err != nil { return err } diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 7870d2d1670..29e14a9cb3d 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -235,7 +235,7 @@ func withSpinner(ctx context.Context, show bool, msg string, fn func() error) er // upload the launch artifacts, assemble the Jobs payload, and submit it. It // returns the new run_id and its dashboard URL. showProgress enables the // stderr upload/packaging spinners (text mode only). -func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (int64, string, error) { +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress, noCache bool) (int64, string, error) { // Compute the launch dir and command_path up front — a read-only workspace lookup plus a // local path build, no writes yet — so the pre-flight validates the real command_path. The // same path is reused for the upload and submit below, so the validated path is the submitted @@ -320,7 +320,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run // Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh. err = withSpinner(ctx, showProgress, "Packaging code snapshot…", func() error { var e error - snap, e = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) + snap, e = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir, noCache) return e }) if err != nil { diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 45560157bc2..4f872d4754d 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -242,7 +242,7 @@ func TestSubmitWorkload(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) + runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false, false) require.NoError(t, err) assert.Equal(t, int64(777), runID) assert.Contains(t, dashboardURL, "/jobs/runs/777") @@ -286,7 +286,7 @@ func TestSubmitWorkloadHonorsOverride(t *testing.T) { cfg, err := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"compute.num_accelerators=4"}) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) + _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false, false) require.NoError(t, err) require.Len(t, got.Tasks, 1) @@ -329,7 +329,7 @@ code_source: // The DABs upload path logs via cmdio; the real `air run` context carries it. ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -373,7 +373,7 @@ code_source: require.NoError(t, err) ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -423,7 +423,7 @@ code_source: // The uploaded name carries a discriminator (timestamp), not the bare dir name. ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false) require.NoError(t, err) base := path.Base(snap.CodeSourcePath) assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be unique, not the bare dir name") @@ -475,9 +475,9 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false) require.NoError(t, err) - second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false) require.NoError(t, err) // Same pinned commit → identical content-addressed remote path, uploaded once @@ -519,7 +519,7 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false) require.NoError(t, err) assert.Empty(t, snap.GitStatePath) @@ -568,7 +568,7 @@ code_source: require.NoError(t, err) ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -603,7 +603,7 @@ func TestSubmitWorkloadGuards(t *testing.T) { cfg := *base cfg.UsagePolicyName = new("nope") - _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false) + _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false, false) require.ErrorContains(t, err, `no usage policy named "nope"`) for _, p := range paths { assert.NotContains(t, p, "/workspace/", "no workspace write may precede policy resolution") @@ -640,7 +640,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false) + _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false, false) require.NoError(t, err) assert.Equal(t, policyID, got.BudgetPolicyId) }) @@ -651,7 +651,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false) + _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false, false) require.NoError(t, err) assert.Equal(t, policyID, got.BudgetPolicyId) }) diff --git a/experimental/air/cmd/snapshot_cache.go b/experimental/air/cmd/snapshot_cache.go new file mode 100644 index 00000000000..82ce0c64ad7 --- /dev/null +++ b/experimental/air/cmd/snapshot_cache.go @@ -0,0 +1,390 @@ +package aircmd + +// Warm snapshot cache for the plain_tar path (working tree, no git ref). The Python +// CLI and the git_archive path re-pack the whole tree every run; for a large repo the +// file walk + read + gzip dominates submit latency. This keeps a warm, uncompressed +// tar of the tree on local disk plus a manifest of each member's identity (size+mtime) +// and byte range. On the next run we stat the file set, and rebuild the tarball by +// copying unchanged members verbatim from the warm tar — reading only changed files +// from disk — before gzipping the upload. Nothing changed is the degenerate case: we +// just recompress the warm tar. The cache is keyed by (repo path, config path, +// include_paths), so distinct repos or configs never share an entry, and it is gated +// to large trees (below the threshold a plain re-pack is cheap enough not to bother). + +import ( + "archive/tar" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "os" + "path" + "path/filepath" + "slices" + "strings" + "time" + + "github.com/databricks/cli/libs/log" + "github.com/klauspost/pgzip" +) + +const ( + // snapshotCacheVersion invalidates on-disk caches when the layout below changes. + snapshotCacheVersion = "v1" + snapshotCacheManifestName = "manifest.json" + snapshotCacheTarName = "snapshot.tar" + + // snapshotCacheMinBytes gates the cache to trees large enough that the walk+read + // cost dominates. Below it a plain re-pack is cheap and the cache bookkeeping + // isn't worth it. Rough heuristic, not exact — tune from measurement. + snapshotCacheMinBytes = 64 << 20 // 64 MiB +) + +// tarTrailer is the two zero blocks that mark end-of-archive. writeSnapshot appends +// it because members are streamed without the tar.Writer's own Close (which would +// embed a trailer between members). +var tarTrailer = make([]byte, 2*512) + +// cacheEntry records a member's identity for change detection (size+mtime) and its +// byte range within the warm snapshot.tar, so an unchanged member can be copied +// verbatim instead of re-read from disk. +type cacheEntry struct { + Size int64 `json:"size"` + ModTime int64 `json:"mtime_ns"` + Offset int64 `json:"offset"` + Length int64 `json:"length"` +} + +// snapshotManifest is the on-disk index of a warm snapshot.tar. +type snapshotManifest struct { + Version string `json:"version"` + DirName string `json:"dir_name"` + Entries map[string]cacheEntry `json:"entries"` // keyed by slash-separated relative path +} + +// snapshotCacheKey is a stable digest of the inputs that determine the tar's content +// set. Different repos, config files, or include_paths get different cache folders. +func snapshotCacheKey(absRepo, absConfig string, includePaths []string) string { + paths := slices.Clone(includePaths) + slices.Sort(paths) + material := strings.Join(append([]string{absRepo, absConfig, snapshotCacheVersion}, paths...), "\x00") + sum := sha256.Sum256([]byte(material)) + return hex.EncodeToString(sum[:]) +} + +func snapshotCacheDir(absRepo, absConfig string, includePaths []string) string { + return filepath.Join(os.TempDir(), "databricks", ".air", snapshotCacheKey(absRepo, absConfig, includePaths)) +} + +// packagePlainTarWithCache writes the working-tree tarball to outputTarball, using the +// warm cache when the tree is large enough. Small trees are packed fresh without +// touching the cache. +func packagePlainTarWithCache(ctx context.Context, repoPath, configPath string, includePaths []string, isGitRepo bool, outputTarball string) (err error) { + start := time.Now() + files, err := snapshotFiles(ctx, repoPath, includePaths, isGitRepo) + if err != nil { + return err + } + listDone := time.Now() + var total int64 + for _, f := range files { + total += f.size + } + dirName := filepath.Base(repoPath) + + // One timing line comparable to the shell path's "snapshot profile", so a + // cache-on vs --no-cache run can be compared directly. Debug-only. + mode := "rebuild-cold" + defer func() { + log.Debugf(ctx, "air snapshot cache: mode=%s files=%d uncompressed_bytes=%d list=%s pack=%s", + mode, len(files), total, listDone.Sub(start), time.Since(listDone)) + }() + + if total < snapshotCacheMinBytes { + mode = "skip-small" + return writeGzOnly(repoPath, dirName, files, outputTarball) + } + + absRepo, err := filepath.Abs(repoPath) + if err != nil { + return err + } + absConfig, err := filepath.Abs(configPath) + if err != nil { + return err + } + cacheDir := snapshotCacheDir(absRepo, absConfig, includePaths) + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + manifestPath := filepath.Join(cacheDir, snapshotCacheManifestName) + + old := loadSnapshotManifest(manifestPath) + if old != nil && old.DirName == dirName && fileExists(tarPath) && !snapshotChanged(files, old) { + mode = "hit-nochange" + return gzipFile(tarPath, outputTarball) + } + + if err := os.MkdirAll(cacheDir, 0o700); err != nil { + return fmt.Errorf("failed to create snapshot cache dir: %w", err) + } + if old != nil { + mode = "rebuild-warm" + } + return rebuildWarmSnapshot(repoPath, dirName, files, old, tarPath, outputTarball) +} + +// snapshotChanged reports whether the current file set differs from the manifest by +// any add, delete, or modification (mtime+size). Equal length plus every current file +// matching an entry means the sets are identical. +func snapshotChanged(files []snapshotFile, old *snapshotManifest) bool { + if len(files) != len(old.Entries) { + return true + } + for _, f := range files { + e, ok := old.Entries[filepath.ToSlash(f.rel)] + if !ok || e.Size != f.size || e.ModTime != f.modTime { + return true + } + } + return false +} + +// rebuildWarmSnapshot writes a fresh warm tar (copying unchanged members from the old +// one when available) and its gzipped upload copy in a single pass, then atomically +// replaces the cached tar and manifest. +func rebuildWarmSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshotManifest, tarPath, outputTarball string) (err error) { + gz, closeGz, err := newGzFile(outputTarball) + if err != nil { + return err + } + defer func() { err = firstErr(err, closeGz()) }() + + newTarPath := tarPath + ".tmp" + tarFile, err := os.Create(newTarPath) + if err != nil { + return fmt.Errorf("failed to create warm tar: %w", err) + } + defer func() { + if err != nil { + os.Remove(newTarPath) + } + }() + + var oldTar io.ReaderAt + if old != nil { + if f, e := os.Open(tarPath); e == nil { + defer f.Close() + oldTar = f + } else { + old = nil // warm tar gone; rebuild every member from disk + } + } + + manifest, err := writeSnapshot(repoPath, dirName, files, old, oldTar, tarFile, gz) + if err != nil { + tarFile.Close() + return err + } + if err = tarFile.Close(); err != nil { + return fmt.Errorf("failed to finalize warm tar: %w", err) + } + if err = closeGz(); err != nil { + return err + } + if err = os.Rename(newTarPath, tarPath); err != nil { + return fmt.Errorf("failed to install warm tar: %w", err) + } + return saveSnapshotManifest(filepath.Join(filepath.Dir(tarPath), snapshotCacheManifestName), manifest) +} + +// writeGzOnly packs files straight to a gzipped tarball without persisting a cache, +// used for trees below the cache threshold. +func writeGzOnly(repoPath, dirName string, files []snapshotFile, outputTarball string) (err error) { + gz, closeGz, err := newGzFile(outputTarball) + if err != nil { + return err + } + defer func() { err = firstErr(err, closeGz()) }() + _, err = writeSnapshot(repoPath, dirName, files, nil, nil, nil, gz) + return err +} + +// writeSnapshot streams every member (sorted for determinism) to gzDst, and to tarDst +// too when non-nil, returning the manifest that indexes each member's byte range in +// the tarDst stream. When old+oldTar are set, an unchanged member (matching size+mtime) +// is copied verbatim from oldTar rather than re-read from disk. +func writeSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshotManifest, oldTar io.ReaderAt, tarDst, gzDst io.Writer) (snapshotManifest, error) { + dst := gzDst + if tarDst != nil { + dst = io.MultiWriter(tarDst, gzDst) + } + cw := &countWriter{w: dst} + + sorted := slices.Clone(files) + slices.SortFunc(sorted, func(a, b snapshotFile) int { + return strings.Compare(filepath.ToSlash(a.rel), filepath.ToSlash(b.rel)) + }) + + entries := make(map[string]cacheEntry, len(sorted)) + for _, f := range sorted { + rel := filepath.ToSlash(f.rel) + start := cw.n + + reused := false + if old != nil && oldTar != nil { + if e, ok := old.Entries[rel]; ok && e.Size == f.size && e.ModTime == f.modTime { + if _, err := io.Copy(cw, io.NewSectionReader(oldTar, e.Offset, e.Length)); err != nil { + return snapshotManifest{}, fmt.Errorf("failed to copy cached member %q: %w", rel, err) + } + reused = true + } + } + if !reused { + if err := writeMember(cw, repoPath, path.Join(dirName, rel), f.rel); err != nil { + return snapshotManifest{}, err + } + } + entries[rel] = cacheEntry{Size: f.size, ModTime: f.modTime, Offset: start, Length: cw.n - start} + } + if _, err := cw.Write(tarTrailer); err != nil { + return snapshotManifest{}, err + } + return snapshotManifest{Version: snapshotCacheVersion, DirName: dirName, Entries: entries}, nil +} + +// writeMember streams one framed tar member (header + content + block padding) for the +// file at repoPath/rel, named name inside the archive. It deliberately Flushes rather +// than Closes the tar.Writer, so no end-of-archive trailer is written between members. +func writeMember(w io.Writer, repoPath, name, rel string) error { + full := filepath.Join(repoPath, filepath.FromSlash(rel)) + info, err := os.Lstat(full) + if err != nil { + return fmt.Errorf("failed to stat %q: %w", rel, err) + } + link := "" + if info.Mode()&os.ModeSymlink != 0 { + if link, err = os.Readlink(full); err != nil { + return fmt.Errorf("failed to read symlink %q: %w", rel, err) + } + } + hdr, err := tar.FileInfoHeader(info, link) + if err != nil { + return fmt.Errorf("failed to build tar header for %q: %w", rel, err) + } + hdr.Name = name + + tw := tar.NewWriter(w) + if err := tw.WriteHeader(hdr); err != nil { + return fmt.Errorf("failed to write tar header for %q: %w", rel, err) + } + if info.Mode().IsRegular() { + f, err := os.Open(full) + if err != nil { + return fmt.Errorf("failed to open %q: %w", rel, err) + } + defer f.Close() + if _, err := io.Copy(tw, f); err != nil { + return fmt.Errorf("failed to archive %q: %w", rel, err) + } + } + return tw.Flush() +} + +// countWriter counts the bytes written through it, to record member byte offsets. +type countWriter struct { + w io.Writer + n int64 +} + +func (c *countWriter) Write(p []byte) (int, error) { + n, err := c.w.Write(p) + c.n += int64(n) + return n, err +} + +// newGzFile creates outputTarball and returns a BestSpeed gzip writer over it plus an +// idempotent close that flushes the gzip stream and the file. Compression level is +// BestSpeed because the uploaded size does not matter for this workflow — only latency. +// +// gzip is parallel (klauspost/pgzip): compressing the whole tar is the dominant +// packaging cost on a large tree and is paid on every run — even a no-change cache hit +// re-gzips the warm tar — so it is spread across cores (measured ~18x faster than +// compress/gzip on a 470 MiB tar). pgzip buffers its own blocks, so the tar writer's +// small writes parallelize fine without extra buffering. Its output is an ordinary gzip +// stream any gunzip/tar reads, and it falls back to serial below one block — fine, since +// the cache only engages above snapshotCacheMinBytes. +func newGzFile(outputTarball string) (io.Writer, func() error, error) { + f, err := os.Create(outputTarball) + if err != nil { + return nil, nil, fmt.Errorf("failed to create tarball: %w", err) + } + gz, err := pgzip.NewWriterLevel(f, pgzip.BestSpeed) + if err != nil { + f.Close() + return nil, nil, err + } + closed := false + closeFn := func() error { + if closed { + return nil + } + closed = true + return firstErr(gz.Close(), f.Close()) + } + return gz, closeFn, nil +} + +// gzipFile writes a BestSpeed gzip of src to outputTarball. Used on a no-change cache +// hit to recompress the warm tar without re-reading the working tree. +func gzipFile(src, outputTarball string) (err error) { + in, err := os.Open(src) + if err != nil { + return fmt.Errorf("failed to open warm tar: %w", err) + } + defer in.Close() + gz, closeGz, err := newGzFile(outputTarball) + if err != nil { + return err + } + defer func() { err = firstErr(err, closeGz()) }() + _, err = io.Copy(gz, in) + return err +} + +func loadSnapshotManifest(path string) *snapshotManifest { + data, err := os.ReadFile(path) + if err != nil { + return nil + } + var m snapshotManifest + if err := json.Unmarshal(data, &m); err != nil || m.Version != snapshotCacheVersion { + return nil + } + return &m +} + +func saveSnapshotManifest(path string, m snapshotManifest) error { + data, err := json.Marshal(m) + if err != nil { + return err + } + if err := os.WriteFile(path, data, 0o600); err != nil { + return fmt.Errorf("failed to write snapshot manifest: %w", err) + } + return nil +} + +func fileExists(path string) bool { + _, err := os.Stat(path) + return err == nil +} + +func firstErr(errs ...error) error { + for _, e := range errs { + if e != nil { + return e + } + } + return nil +} diff --git a/experimental/air/cmd/snapshot_cache_test.go b/experimental/air/cmd/snapshot_cache_test.go new file mode 100644 index 00000000000..134f3418c0e --- /dev/null +++ b/experimental/air/cmd/snapshot_cache_test.go @@ -0,0 +1,183 @@ +package aircmd + +import ( + "archive/tar" + "compress/gzip" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// extractTarball returns entry name -> content for a .tar.gz. +func extractTarball(t *testing.T, path string) map[string]string { + t.Helper() + f, err := os.Open(path) + require.NoError(t, err) + defer f.Close() + gz, err := gzip.NewReader(f) + require.NoError(t, err) + defer gz.Close() + + out := map[string]string{} + tr := tar.NewReader(gz) + for { + hdr, err := tr.Next() + if err != nil { + break + } + buf := make([]byte, hdr.Size) + _, _ = tr.Read(buf) + out[hdr.Name] = string(buf) + } + return out +} + +func statFiles(t *testing.T, repo string, rels ...string) []snapshotFile { + t.Helper() + var files []snapshotFile + for _, rel := range rels { + info, err := os.Lstat(filepath.Join(repo, filepath.FromSlash(rel))) + require.NoError(t, err) + files = append(files, snapshotFile{rel: filepath.FromSlash(rel), size: info.Size(), modTime: info.ModTime().UnixNano()}) + } + return files +} + +func TestSnapshotCacheKey(t *testing.T) { + base := snapshotCacheKey("/repo", "/repo/air.yaml", nil) + assert.Equal(t, base, snapshotCacheKey("/repo", "/repo/air.yaml", nil), "stable for identical inputs") + assert.NotEqual(t, base, snapshotCacheKey("/other", "/repo/air.yaml", nil), "repo path matters") + assert.NotEqual(t, base, snapshotCacheKey("/repo", "/repo/other.yaml", nil), "config path matters") + assert.NotEqual(t, base, snapshotCacheKey("/repo", "/repo/air.yaml", []string{"src"}), "include_paths matter") + // include_paths order must not matter. + assert.Equal(t, + snapshotCacheKey("/repo", "/repo/air.yaml", []string{"a", "b"}), + snapshotCacheKey("/repo", "/repo/air.yaml", []string{"b", "a"})) +} + +func TestSnapshotChanged(t *testing.T) { + old := &snapshotManifest{Entries: map[string]cacheEntry{ + "a.txt": {Size: 1, ModTime: 100}, + "src/b.py": {Size: 2, ModTime: 200}, + }} + unchanged := []snapshotFile{{rel: "a.txt", size: 1, modTime: 100}, {rel: filepath.FromSlash("src/b.py"), size: 2, modTime: 200}} + assert.False(t, snapshotChanged(unchanged, old)) + + modified := []snapshotFile{{rel: "a.txt", size: 1, modTime: 100}, {rel: filepath.FromSlash("src/b.py"), size: 2, modTime: 999}} + assert.True(t, snapshotChanged(modified, old), "mtime change detected") + + resized := []snapshotFile{{rel: "a.txt", size: 5, modTime: 100}, {rel: filepath.FromSlash("src/b.py"), size: 2, modTime: 200}} + assert.True(t, snapshotChanged(resized, old), "size change detected") + + removed := []snapshotFile{{rel: "a.txt", size: 1, modTime: 100}} + assert.True(t, snapshotChanged(removed, old), "deletion detected") + + added := append(append([]snapshotFile(nil), unchanged...), snapshotFile{rel: "c.txt", size: 3, modTime: 300}) + assert.True(t, snapshotChanged(added, old), "addition detected") +} + +func TestWarmSnapshotColdBuild(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + writeRepoFile(t, repo, "src/model.py", "print()") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + out := filepath.Join(t.TempDir(), "snap.tar.gz") + + files := statFiles(t, repo, "a.txt", "src/model.py") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, out)) + + contents := extractTarball(t, out) + assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) + assert.Equal(t, "print()", contents[dirName+"/src/model.py"]) + + // The warm tar and manifest are persisted for the next run. + assert.FileExists(t, tarPath) + m := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, m) + assert.Equal(t, dirName, m.DirName) + assert.Len(t, m.Entries, 2) +} + +func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + writeRepoFile(t, repo, "keep.py", "keep") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + files := statFiles(t, repo, "a.txt", "keep.py") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + + old := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, old) + + // Delete keep.py from disk but keep it in the file list with its original + // size+mtime: a correct rebuild must copy its bytes from the warm tar, proving + // unchanged members are not re-read from disk. + require.NoError(t, os.Remove(filepath.Join(repo, "keep.py"))) + + out := filepath.Join(t.TempDir(), "warm.tar.gz") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, old, tarPath, out)) + + contents := extractTarball(t, out) + assert.Equal(t, "keep", contents[dirName+"/keep.py"], "unchanged member copied from warm tar") + assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) +} + +func TestWarmSnapshotRebuildEditAddDelete(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + writeRepoFile(t, repo, "b.txt", "bravo") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + files := statFiles(t, repo, "a.txt", "b.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + old := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, old) + + // Edit a.txt (changed), delete b.txt, add c.txt. + writeRepoFile(t, repo, "a.txt", "alpha-v2") + require.NoError(t, os.Remove(filepath.Join(repo, "b.txt"))) + writeRepoFile(t, repo, "c.txt", "charlie") + + out := filepath.Join(t.TempDir(), "warm.tar.gz") + newFiles := statFiles(t, repo, "a.txt", "c.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, newFiles, old, tarPath, out)) + + contents := extractTarball(t, out) + assert.Equal(t, "alpha-v2", contents[dirName+"/a.txt"], "edited file updated") + assert.Equal(t, "charlie", contents[dirName+"/c.txt"], "added file present") + _, hasB := contents[dirName+"/b.txt"] + assert.False(t, hasB, "deleted file dropped") + + // The refreshed manifest reflects the new set. + updated := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, updated) + assert.Len(t, updated.Entries, 2) +} + +func TestGzipFileRoundTrip(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + files := statFiles(t, repo, "a.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + + // gzipFile recompresses the warm tar directly (the no-change hit path). + out := filepath.Join(t.TempDir(), "reuse.tar.gz") + require.NoError(t, gzipFile(tarPath, out)) + contents := extractTarball(t, out) + assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) +} diff --git a/experimental/air/cmd/snapshot_dabs.go b/experimental/air/cmd/snapshot_dabs.go index ed809b09e01..2176bf2cce1 100644 --- a/experimental/air/cmd/snapshot_dabs.go +++ b/experimental/air/cmd/snapshot_dabs.go @@ -38,7 +38,7 @@ const uploadProvenanceSidecars = false // not reimplement workspace/volume upload. A minimal in-memory bundle carries the // local tarball path as code_source_path; ReplaceWithRemotePath rewrites it to the // artifact .internal path and Upload pushes the bytes. -func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath string, sidecarStore filer.Filer, sidecarBase string) (snapshotResult, error) { +func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath string, sidecarStore filer.Filer, sidecarBase string, noCache bool) (snapshotResult, error) { repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) if err != nil { return snapshotResult{}, err @@ -57,7 +57,7 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s if snap.RemoteVolume != nil { remoteVolume = *snap.RemoteVolume } - result, err := uploadSnapshotViaDABs(ctx, w, repoPath, plan, remoteVolume) + result, err := uploadSnapshotViaDABs(ctx, w, repoPath, configPath, plan, remoteVolume, noCache) if err != nil { return snapshotResult{}, err } @@ -141,13 +141,17 @@ func snapshotTarballName(plan snapshotPlan, dirName string) string { } // packageSnapshot writes the snapshot to tarball per the resolved plan: `git archive` -// of the pinned commit for git_archive, else a plain tar of the working tree. -func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, tarball string) error { - dirName := filepath.Base(repoPath) +// of the pinned commit for git_archive, else a plain tar of the working tree. The +// working-tree path uses the warm snapshot cache unless noCache is set, in which case +// it falls back to the shell `tar` path. +func packageSnapshot(ctx context.Context, repoPath, configPath string, plan snapshotPlan, tarball string, noCache bool) error { if plan.mode == modeGitArchive { - return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths) + return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, filepath.Base(repoPath), plan.includePaths) } - return createPlainTarball(ctx, repoPath, tarball, plan.includePaths, plan.isGitRepo) + if noCache { + return createPlainTarball(ctx, repoPath, tarball, plan.includePaths, plan.isGitRepo) + } + return packagePlainTarWithCache(ctx, repoPath, configPath, plan.includePaths, plan.isGitRepo, tarball) } // uploadSnapshotViaDABs uploads the snapshot through DABs' artifact-upload machinery @@ -159,7 +163,7 @@ func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, ta // git_archive snapshots are cacheable: the tarball name is content-addressed by // (commit, include_paths), so if the identical object is already uploaded we skip // packaging and upload entirely and just reuse the remote path. -func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath string, plan snapshotPlan, remoteVolume string) (snapshotResult, error) { +func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath, configPath string, plan snapshotPlan, remoteVolume string, noCache bool) (snapshotResult, error) { // artifactPath is where DABs uploads the tarball; GetFilerForLibraries routes to // a Workspace or Volume filer based on its prefix, then appends /.internal. artifactPath := remoteVolume @@ -236,7 +240,7 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r } // Cache miss (or plain_tar): package the tarball locally, then upload the bytes. - if err := packageSnapshot(ctx, repoPath, plan, filepath.Join(tmp, tarName)); err != nil { + if err := packageSnapshot(ctx, repoPath, configPath, plan, filepath.Join(tmp, tarName), noCache); err != nil { return snapshotResult{}, err } @@ -244,9 +248,13 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r if diags.HasError() { return snapshotResult{}, diags.Error() } + // Phase 0 profiling: isolate the upload stage from packaging so we know whether + // the cost is local (walk/gzip) or network. Debug-only; enable with -v/--debug. + uploadStart := time.Now() if diags := bundle.Apply(ctx, b, libraries.Upload(libs)); diags.HasError() { return snapshotResult{}, diags.Error() } + log.Debugf(ctx, "air snapshot profile: upload=%s", time.Since(uploadStart)) remote, err := readCodeSourcePath(b) if err != nil { diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index a2203c24970..b02101ca5e4 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -73,8 +73,8 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc args := []string{"-cf", "-", "-C", parent, "--null", "--no-recursion", "-T", "-"} cmd := exec.CommandContext(ctx, "tar", args...) var stdin bytes.Buffer - for _, file := range files { - stdin.WriteString(filepath.ToSlash(filepath.Join(dirName, file))) + for _, f := range files { + stdin.WriteString(filepath.ToSlash(filepath.Join(dirName, f.rel))) stdin.WriteByte(0) } cmd.Stdin = &stdin @@ -96,7 +96,15 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc return out.Close() } -func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, isGitRepo bool) ([]string, error) { +// snapshotFile is a file selected for the snapshot: its repo-relative path (native +// separators) with the size and mtime the warm cache uses to detect changes. +type snapshotFile struct { + rel string + size int64 + modTime int64 // Unix nanoseconds +} + +func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, isGitRepo bool) ([]snapshotFile, error) { args := []string{"-C", repoPath, "ls-files", "-z", "--cached", "--others", "--exclude-standard"} if !isGitRepo { gitDir, err := os.MkdirTemp("", "air-snapshot-git-") @@ -120,7 +128,7 @@ func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, return nil, fmt.Errorf("failed to evaluate git ignore rules: %w", err) } - var files []string + var files []snapshotFile for raw := range bytes.SplitSeq(output, []byte{0}) { if len(raw) == 0 { continue @@ -130,14 +138,14 @@ func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, if name == ".git" || strings.HasPrefix(name, ".git/") || strings.HasPrefix(base, "._") { continue } - _, err := os.Lstat(filepath.Join(repoPath, filepath.FromSlash(name))) + info, err := os.Lstat(filepath.Join(repoPath, filepath.FromSlash(name))) if errors.Is(err, os.ErrNotExist) { continue } if err != nil { return nil, fmt.Errorf("failed to inspect snapshot path %q: %w", name, err) } - files = append(files, filepath.FromSlash(name)) + files = append(files, snapshotFile{rel: filepath.FromSlash(name), size: info.Size(), modTime: info.ModTime().UnixNano()}) } return files, nil } From 85981d52cadff8a9a2480f817332cd5945642d38 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 17:09:33 +0000 Subject: [PATCH 2/4] experimental/air: use default gzip level in the warm cache too Match the parent PR: DefaultCompression rather than BestSpeed in newGzFile, so the cached tarball is re-gzipped at the same level as the --no-cache path and the upload stays small. Parallel compression makes the higher level nearly free. Co-authored-by: Isaac --- experimental/air/cmd/snapshot_cache.go | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/experimental/air/cmd/snapshot_cache.go b/experimental/air/cmd/snapshot_cache.go index 82ce0c64ad7..60a0a793896 100644 --- a/experimental/air/cmd/snapshot_cache.go +++ b/experimental/air/cmd/snapshot_cache.go @@ -303,23 +303,24 @@ func (c *countWriter) Write(p []byte) (int, error) { return n, err } -// newGzFile creates outputTarball and returns a BestSpeed gzip writer over it plus an -// idempotent close that flushes the gzip stream and the file. Compression level is -// BestSpeed because the uploaded size does not matter for this workflow — only latency. +// newGzFile creates outputTarball and returns a parallel gzip writer over it plus an +// idempotent close that flushes the gzip stream and the file. // // gzip is parallel (klauspost/pgzip): compressing the whole tar is the dominant // packaging cost on a large tree and is paid on every run — even a no-change cache hit -// re-gzips the warm tar — so it is spread across cores (measured ~18x faster than -// compress/gzip on a 470 MiB tar). pgzip buffers its own blocks, so the tar writer's -// small writes parallelize fine without extra buffering. Its output is an ordinary gzip -// stream any gunzip/tar reads, and it falls back to serial below one block — fine, since -// the cache only engages above snapshotCacheMinBytes. +// re-gzips the warm tar — so it is spread across cores (~18x faster than compress/gzip +// on a 470 MiB tar). The level is DefaultCompression, not BestSpeed: the tarball is +// re-uploaded every run so its size matters, and with parallel compression a normal +// level is nearly free (a few hundred ms for ~15-18% fewer bytes). pgzip buffers its +// own blocks, so the tar writer's small writes parallelize fine without extra buffering; +// its output is an ordinary gzip stream any gunzip/tar reads, and it falls back to serial +// below one block — fine, since the cache only engages above snapshotCacheMinBytes. func newGzFile(outputTarball string) (io.Writer, func() error, error) { f, err := os.Create(outputTarball) if err != nil { return nil, nil, fmt.Errorf("failed to create tarball: %w", err) } - gz, err := pgzip.NewWriterLevel(f, pgzip.BestSpeed) + gz, err := pgzip.NewWriterLevel(f, pgzip.DefaultCompression) if err != nil { f.Close() return nil, nil, err From 8bfd8da44a8e1fef9b5fbcc410d4fdd4eb128e0e Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 17:47:02 +0000 Subject: [PATCH 3/4] experimental/air: update config-help golden for --no-cache flag The new --no-cache flag on `air run` adds a line to its --help output, which the experimental/air/config-help acceptance test pins. Regenerate the golden. Co-authored-by: Isaac --- acceptance/experimental/air/config-help/output.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/acceptance/experimental/air/config-help/output.txt b/acceptance/experimental/air/config-help/output.txt index e9a066be055..0c00ab74539 100644 --- a/acceptance/experimental/air/config-help/output.txt +++ b/acceptance/experimental/air/config-help/output.txt @@ -22,6 +22,7 @@ Flags: -f, --file string Path to the workload YAML config -h, --help help for run --idempotency-key string Return the existing run if this key was already used + --no-cache Bypass the local snapshot cache and re-pack the code tarball from scratch --override stringArray Override a YAML field, e.g. compute.num_accelerators=8 (repeatable) --watch Stream logs until the run completes From ed6477475a9295c12b623f2886c46a3da779c431 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 19:17:31 +0000 Subject: [PATCH 4/4] experimental/air: make the warm cache crash- and concurrency-safe Isaac Review flagged two MAJOR correctness bugs in the warm cache: - Concurrent `air run` on the same cache key wrote the same snapshot.tar.tmp and raced the rename plus a non-atomic manifest write, interleaving into a corrupt tar/manifest pair. - rebuildWarmSnapshot renamed the new tar into place before saving the manifest, so a crash between the two left a manifest whose byte offsets described a different tar layout -- silently corrupting a later verbatim-reuse rebuild. Fix both by binding the manifest to a per-build, uniquely named tar (snapshot..tar) that is never overwritten, and installing the manifest atomically (unique temp + rename) only after its tar is durable. A manifest and the tar it indexes are therefore always a consistent pair: there is no window where offsets describe a mismatched tar, and concurrent rebuilds are last-writer-wins on the manifest rather than interleaving, so no lock is needed. Superseded and orphaned tars are cleaned up best-effort. Adds a rotation test. Co-authored-by: Isaac --- experimental/air/cmd/snapshot_cache.go | 110 ++++++++++++++++---- experimental/air/cmd/snapshot_cache_test.go | 65 +++++++++--- 2 files changed, 140 insertions(+), 35 deletions(-) diff --git a/experimental/air/cmd/snapshot_cache.go b/experimental/air/cmd/snapshot_cache.go index 60a0a793896..e319c84be90 100644 --- a/experimental/air/cmd/snapshot_cache.go +++ b/experimental/air/cmd/snapshot_cache.go @@ -14,6 +14,7 @@ package aircmd import ( "archive/tar" "context" + "crypto/rand" "crypto/sha256" "encoding/hex" "encoding/json" @@ -34,7 +35,10 @@ const ( // snapshotCacheVersion invalidates on-disk caches when the layout below changes. snapshotCacheVersion = "v1" snapshotCacheManifestName = "manifest.json" - snapshotCacheTarName = "snapshot.tar" + // snapshotTarPrefix begins each warm tar's filename; the rest is a per-build random + // id (snapshot..tar). A build never overwrites another build's tar, so a manifest + // and the tar it names stay a consistent pair — see rebuildWarmSnapshot. + snapshotTarPrefix = "snapshot." // snapshotCacheMinBytes gates the cache to trees large enough that the walk+read // cost dominates. Below it a plain re-pack is cheap and the cache bookkeeping @@ -57,9 +61,11 @@ type cacheEntry struct { Length int64 `json:"length"` } -// snapshotManifest is the on-disk index of a warm snapshot.tar. +// snapshotManifest is the on-disk index of a warm snapshot tar. TarName binds it to the +// exact tar its byte offsets describe, so a stale or half-written pair is never reused. type snapshotManifest struct { Version string `json:"version"` + TarName string `json:"tar_name"` // the snapshot..tar this manifest indexes DirName string `json:"dir_name"` Entries map[string]cacheEntry `json:"entries"` // keyed by slash-separated relative path } @@ -116,13 +122,16 @@ func packagePlainTarWithCache(ctx context.Context, repoPath, configPath string, return err } cacheDir := snapshotCacheDir(absRepo, absConfig, includePaths) - tarPath := filepath.Join(cacheDir, snapshotCacheTarName) manifestPath := filepath.Join(cacheDir, snapshotCacheManifestName) old := loadSnapshotManifest(manifestPath) - if old != nil && old.DirName == dirName && fileExists(tarPath) && !snapshotChanged(files, old) { + var oldTarPath string + if old != nil { + oldTarPath = filepath.Join(cacheDir, old.TarName) + } + if old != nil && old.DirName == dirName && fileExists(oldTarPath) && !snapshotChanged(files, old) { mode = "hit-nochange" - return gzipFile(tarPath, outputTarball) + return gzipFile(oldTarPath, outputTarball) } if err := os.MkdirAll(cacheDir, 0o700); err != nil { @@ -131,7 +140,7 @@ func packagePlainTarWithCache(ctx context.Context, repoPath, configPath string, if old != nil { mode = "rebuild-warm" } - return rebuildWarmSnapshot(repoPath, dirName, files, old, tarPath, outputTarball) + return rebuildWarmSnapshot(repoPath, dirName, files, old, cacheDir, oldTarPath, outputTarball) } // snapshotChanged reports whether the current file set differs from the manifest by @@ -150,52 +159,70 @@ func snapshotChanged(files []snapshotFile, old *snapshotManifest) bool { return false } -// rebuildWarmSnapshot writes a fresh warm tar (copying unchanged members from the old -// one when available) and its gzipped upload copy in a single pass, then atomically -// replaces the cached tar and manifest. -func rebuildWarmSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshotManifest, tarPath, outputTarball string) (err error) { +// rebuildWarmSnapshot writes a fresh warm tar and its gzipped upload copy in one pass +// (copying unchanged members verbatim from the previous build's tar), then installs a +// manifest pointing at the new tar. +// +// Correctness under crashes and concurrent runs rests on two properties: each build +// names its tar uniquely (snapshot..tar) and never overwrites another's, and the +// manifest is renamed into place atomically only after the tar it names is fully +// written. So a manifest and the tar it indexes are always a consistent pair — there is +// no window where the manifest's byte offsets describe a tar with a different layout +// (which would silently corrupt a later verbatim reuse), and two concurrent rebuilds are +// last-writer-wins on the manifest rather than interleaving into one file. No lock needed. +func rebuildWarmSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshotManifest, cacheDir, oldTarPath, outputTarball string) (err error) { gz, closeGz, err := newGzFile(outputTarball) if err != nil { return err } defer func() { err = firstErr(err, closeGz()) }() - newTarPath := tarPath + ".tmp" + buildID, err := randomID() + if err != nil { + return err + } + newTarName := snapshotTarPrefix + buildID + ".tar" + newTarPath := filepath.Join(cacheDir, newTarName) tarFile, err := os.Create(newTarPath) if err != nil { return fmt.Errorf("failed to create warm tar: %w", err) } defer func() { if err != nil { - os.Remove(newTarPath) + tarFile.Close() + os.Remove(newTarPath) // unreferenced on failure; don't leave an orphan } }() + // Reuse unchanged members from the previous build's tar. Its offsets come from old's + // manifest, which indexes exactly that (immutable) tar, so the ranges stay valid. var oldTar io.ReaderAt - if old != nil { - if f, e := os.Open(tarPath); e == nil { + if old != nil && oldTarPath != "" { + if f, e := os.Open(oldTarPath); e == nil { defer f.Close() oldTar = f } else { - old = nil // warm tar gone; rebuild every member from disk + old = nil // previous tar gone; rebuild every member from disk } } manifest, err := writeSnapshot(repoPath, dirName, files, old, oldTar, tarFile, gz) if err != nil { - tarFile.Close() return err } + manifest.TarName = newTarName if err = tarFile.Close(); err != nil { return fmt.Errorf("failed to finalize warm tar: %w", err) } if err = closeGz(); err != nil { return err } - if err = os.Rename(newTarPath, tarPath); err != nil { - return fmt.Errorf("failed to install warm tar: %w", err) + // Install the manifest only now that its tar is durable; the write is atomic. + if err = saveSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName), manifest); err != nil { + return err } - return saveSnapshotManifest(filepath.Join(filepath.Dir(tarPath), snapshotCacheManifestName), manifest) + cleanupOldTars(cacheDir, newTarName) + return nil } // writeGzOnly packs files straight to a gzipped tarball without persisting a cache, @@ -336,8 +363,8 @@ func newGzFile(outputTarball string) (io.Writer, func() error, error) { return gz, closeFn, nil } -// gzipFile writes a BestSpeed gzip of src to outputTarball. Used on a no-change cache -// hit to recompress the warm tar without re-reading the working tree. +// gzipFile gzips src to outputTarball via newGzFile (parallel gzip). Used on a no-change +// cache hit to recompress the warm tar without re-reading the working tree. func gzipFile(src, outputTarball string) (err error) { in, err := os.Open(src) if err != nil { @@ -365,17 +392,56 @@ func loadSnapshotManifest(path string) *snapshotManifest { return &m } +// saveSnapshotManifest writes the manifest atomically (unique temp + rename) so a reader +// or a concurrent run never sees a half-written manifest, and the pair with its tar swaps +// in as a unit. func saveSnapshotManifest(path string, m snapshotManifest) error { data, err := json.Marshal(m) if err != nil { return err } - if err := os.WriteFile(path, data, 0o600); err != nil { + tmp, err := os.CreateTemp(filepath.Dir(path), "manifest-*.json") + if err != nil { + return fmt.Errorf("failed to write snapshot manifest: %w", err) + } + tmpName := tmp.Name() + if _, err := tmp.Write(data); err != nil { + tmp.Close() + os.Remove(tmpName) + return fmt.Errorf("failed to write snapshot manifest: %w", err) + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) return fmt.Errorf("failed to write snapshot manifest: %w", err) } + if err := os.Rename(tmpName, path); err != nil { + os.Remove(tmpName) + return fmt.Errorf("failed to install snapshot manifest: %w", err) + } return nil } +// randomID returns a short random hex string used to name each warm tar uniquely. +func randomID() (string, error) { + var b [8]byte + if _, err := rand.Read(b[:]); err != nil { + return "", fmt.Errorf("failed to generate cache build id: %w", err) + } + return hex.EncodeToString(b[:]), nil +} + +// cleanupOldTars removes warm tars other than keep: superseded builds, and tars a +// concurrent rebuild orphaned. Best-effort — only keep is referenced by the manifest, +// and a wrongly removed tar costs at most a cold rebuild, never correctness. +func cleanupOldTars(cacheDir, keep string) { + matches, _ := filepath.Glob(filepath.Join(cacheDir, snapshotTarPrefix+"*.tar")) + for _, p := range matches { + if filepath.Base(p) != keep { + os.Remove(p) + } + } +} + func fileExists(path string) bool { _, err := os.Stat(path) return err == nil diff --git a/experimental/air/cmd/snapshot_cache_test.go b/experimental/air/cmd/snapshot_cache_test.go index 134f3418c0e..717eb62fdfb 100644 --- a/experimental/air/cmd/snapshot_cache_test.go +++ b/experimental/air/cmd/snapshot_cache_test.go @@ -79,6 +79,15 @@ func TestSnapshotChanged(t *testing.T) { assert.True(t, snapshotChanged(added, old), "addition detected") } +// warmTarPath returns the warm tar the cache's manifest currently points at. +func warmTarPath(t *testing.T, cacheDir string) string { + t.Helper() + m := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, m) + require.NotEmpty(t, m.TarName) + return filepath.Join(cacheDir, m.TarName) +} + func TestWarmSnapshotColdBuild(t *testing.T) { repo := t.TempDir() writeRepoFile(t, repo, "a.txt", "alpha") @@ -86,22 +95,23 @@ func TestWarmSnapshotColdBuild(t *testing.T) { dirName := filepath.Base(repo) cacheDir := t.TempDir() - tarPath := filepath.Join(cacheDir, snapshotCacheTarName) out := filepath.Join(t.TempDir(), "snap.tar.gz") files := statFiles(t, repo, "a.txt", "src/model.py") - require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, out)) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", out)) contents := extractTarball(t, out) assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) assert.Equal(t, "print()", contents[dirName+"/src/model.py"]) - // The warm tar and manifest are persisted for the next run. - assert.FileExists(t, tarPath) + // The warm tar and manifest are persisted for the next run, and the manifest names + // the tar it indexes. m := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) require.NotNil(t, m) assert.Equal(t, dirName, m.DirName) assert.Len(t, m.Entries, 2) + assert.NotEmpty(t, m.TarName) + assert.FileExists(t, filepath.Join(cacheDir, m.TarName)) } func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { @@ -111,12 +121,12 @@ func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { dirName := filepath.Base(repo) cacheDir := t.TempDir() - tarPath := filepath.Join(cacheDir, snapshotCacheTarName) files := statFiles(t, repo, "a.txt", "keep.py") - require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", filepath.Join(t.TempDir(), "cold.tar.gz"))) old := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) require.NotNil(t, old) + oldTarPath := filepath.Join(cacheDir, old.TarName) // Delete keep.py from disk but keep it in the file list with its original // size+mtime: a correct rebuild must copy its bytes from the warm tar, proving @@ -124,7 +134,7 @@ func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { require.NoError(t, os.Remove(filepath.Join(repo, "keep.py"))) out := filepath.Join(t.TempDir(), "warm.tar.gz") - require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, old, tarPath, out)) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, old, cacheDir, oldTarPath, out)) contents := extractTarball(t, out) assert.Equal(t, "keep", contents[dirName+"/keep.py"], "unchanged member copied from warm tar") @@ -138,11 +148,11 @@ func TestWarmSnapshotRebuildEditAddDelete(t *testing.T) { dirName := filepath.Base(repo) cacheDir := t.TempDir() - tarPath := filepath.Join(cacheDir, snapshotCacheTarName) files := statFiles(t, repo, "a.txt", "b.txt") - require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", filepath.Join(t.TempDir(), "cold.tar.gz"))) old := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) require.NotNil(t, old) + oldTarPath := filepath.Join(cacheDir, old.TarName) // Edit a.txt (changed), delete b.txt, add c.txt. writeRepoFile(t, repo, "a.txt", "alpha-v2") @@ -151,7 +161,7 @@ func TestWarmSnapshotRebuildEditAddDelete(t *testing.T) { out := filepath.Join(t.TempDir(), "warm.tar.gz") newFiles := statFiles(t, repo, "a.txt", "c.txt") - require.NoError(t, rebuildWarmSnapshot(repo, dirName, newFiles, old, tarPath, out)) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, newFiles, old, cacheDir, oldTarPath, out)) contents := extractTarball(t, out) assert.Equal(t, "alpha-v2", contents[dirName+"/a.txt"], "edited file updated") @@ -171,13 +181,42 @@ func TestGzipFileRoundTrip(t *testing.T) { dirName := filepath.Base(repo) cacheDir := t.TempDir() - tarPath := filepath.Join(cacheDir, snapshotCacheTarName) files := statFiles(t, repo, "a.txt") - require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", filepath.Join(t.TempDir(), "cold.tar.gz"))) // gzipFile recompresses the warm tar directly (the no-change hit path). out := filepath.Join(t.TempDir(), "reuse.tar.gz") - require.NoError(t, gzipFile(tarPath, out)) + require.NoError(t, gzipFile(warmTarPath(t, cacheDir), out)) contents := extractTarball(t, out) assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) } + +func TestWarmSnapshotRebuildRotatesTar(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + writeRepoFile(t, repo, "b.txt", "bravo") + dirName := filepath.Base(repo) + cacheDir := t.TempDir() + + files := statFiles(t, repo, "a.txt", "b.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", filepath.Join(t.TempDir(), "1.tar.gz"))) + m1 := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, m1) + tar1 := filepath.Join(cacheDir, m1.TarName) + assert.FileExists(t, tar1) + + // Change a file and rebuild: a new uniquely named tar replaces the old one, and the + // manifest points at the survivor. This is what keeps the manifest/tar pair + // consistent instead of overwriting a fixed name in place. + writeRepoFile(t, repo, "a.txt", "alpha-2") + files2 := statFiles(t, repo, "a.txt", "b.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files2, m1, cacheDir, tar1, filepath.Join(t.TempDir(), "2.tar.gz"))) + m2 := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, m2) + assert.NotEqual(t, m1.TarName, m2.TarName, "each build gets a unique tar name") + assert.FileExists(t, filepath.Join(cacheDir, m2.TarName)) + assert.NoFileExists(t, tar1, "superseded warm tar is cleaned up") + + matches, _ := filepath.Glob(filepath.Join(cacheDir, snapshotTarPrefix+"*.tar")) + assert.Len(t, matches, 1, "exactly one warm tar remains") +}