diff --git a/acceptance/experimental/air/config-help/output.txt b/acceptance/experimental/air/config-help/output.txt index a17388562a..ee0ce7f668 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 diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 2e78986912..ae50cfb19b 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 72e128b8f3..fb2a7d1323 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -273,7 +273,7 @@ func stageRunArtifacts(ctx context.Context, launchWriter fileWriter, items []upl // 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 // staging spinner (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 @@ -346,7 +346,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run snapshotArtifactPath := path.Join(base, ".air", "repo_snapshots") uploadSnapshot = func(ctx context.Context) (snapshotResult, error) { // Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh. - return snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, snapshotArtifactPath, fc, funcDir) + return snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, snapshotArtifactPath, fc, funcDir, noCache) } } diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index cf781c42f5..9613d7db2d 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -337,7 +337,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") @@ -378,7 +378,7 @@ func TestSubmitWorkloadStagingErrorPreventsSubmit(t *testing.T) { cfg, err := loadRunConfig(cfgPath) 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.Error(t, err) assert.ErrorContains(t, err, "failed to create launch directory") assert.Zero(t, submitCalls.Load()) @@ -408,7 +408,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) @@ -439,7 +439,7 @@ environment: cfg, err := loadRunConfig(cfgPath) 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) tasks, ok := got["tasks"].([]any) @@ -490,7 +490,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 @@ -535,7 +535,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 @@ -555,15 +555,29 @@ func testSidecarStore(t *testing.T, w *databricks.WorkspaceClient) (filer.Filer, const testSnapshotArtifactPath = "/Workspace/Users/tester@databricks.com/.air/repo_snapshots" -// A plain-tar (working-tree) snapshot is uploaded under a unique, timestamped name so -// two concurrent submissions of the same root_path don't clobber each other's upload. -func TestSubmitWorkloadPlainTarNameIsUnique(t *testing.T) { +// A plain-tar (working-tree) snapshot is content-addressed by its file fingerprint +// (path+size+mtime): submitting the same unchanged tree twice reuses the already-uploaded +// tarball and skips the second upload, resolving to the identical remote path. +func TestSubmitWorkloadPlainTarContentAddressed(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { return jobs.SubmitRunResponse{RunId: 555} }) + // Count snapshot import-file calls, preserving fake-workspace persistence so the + // second submit's existence Stat sees the first upload. A count (not a set keyed by + // path) is what proves the skip: both submits resolve to the same content-addressed + // name, so a set could not tell a skipped second submit from one that re-uploaded to + // that same path. + snapshotUploads := 0 + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { + p := req.Vars["path"] + if strings.Contains(p, "/.air/repo_snapshots/") { + snapshotUploads++ + } + return req.Workspace.WorkspaceFilesImportFile(p, req.Body, req.URL.Query().Get("overwrite") == "true") + }) stubValidateConfig(server) testserver.AddDefaultHandlers(server) w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) @@ -584,14 +598,24 @@ code_source: loaded, err := loadRunConfig(cfgPath) require.NoError(t, err) - // 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, testSnapshotArtifactPath, sidecarStore, sidecarBase) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) + require.NoError(t, err) + require.NotZero(t, snapshotUploads, "first submit should upload the tarball") + afterFirst := snapshotUploads + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, 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") - assert.Regexp(t, `^src_\d{8}_\d{6}\.tar\.gz$`, base) + + // Content-addressed name: not the bare dir, but a 16-hex-char fingerprint. + base := path.Base(first.CodeSourcePath) + assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be content-addressed, not the bare dir name") + assert.Regexp(t, `^src_[0-9a-f]{16}\.tar\.gz$`, base) + + // Same unchanged tree → identical remote path, and the second submit moved no bytes: + // zero new import-file calls (a real skip), not a re-upload to the same name. + assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath) + assert.Equal(t, afterFirst, snapshotUploads, "unchanged plain_tar should skip the second upload") } // A git_archive snapshot is content-addressed by (commit, include_paths): submitting @@ -639,9 +663,9 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) require.NoError(t, err) - second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) require.NoError(t, err) // Same pinned commit → identical content-addressed remote path, uploaded once @@ -683,7 +707,7 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) require.NoError(t, err) assert.Empty(t, snap.GitStatePath) @@ -732,7 +756,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 @@ -767,7 +791,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") @@ -804,7 +828,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) }) @@ -815,7 +839,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 0000000000..5f1f6fdfe9 --- /dev/null +++ b/experimental/air/cmd/snapshot_cache.go @@ -0,0 +1,466 @@ +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 metadata identity +// 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/rand" + "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 = "v2" + snapshotCacheManifestName = "manifest.json" + // 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 + // 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 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"` + Mode uint32 `json:"mode"` + LinkTarget string `json:"link_target,omitempty"` + Offset int64 `json:"offset"` + Length int64 `json:"length"` +} + +// 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 +} + +// 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) + manifestPath := filepath.Join(cacheDir, snapshotCacheManifestName) + + old := loadSnapshotManifest(manifestPath) + 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(oldTarPath, 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, cacheDir, oldTarPath, outputTarball) +} + +// snapshotChanged reports whether the current file set differs from the manifest by +// any add, delete, or metadata modification. 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 || e.Mode != f.mode || e.LinkTarget != f.linkTarget { + return true + } + } + return false +} + +// 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()) }() + + 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 { + 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 && oldTarPath != "" { + if f, e := os.Open(oldTarPath); e == nil { + defer f.Close() + oldTar = f + } else { + old = nil // previous tar gone; rebuild every member from disk + } + } + + manifest, err := writeSnapshot(repoPath, dirName, files, old, oldTar, tarFile, gz) + if err != nil { + 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 + } + // 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 + } + cleanupOldTars(cacheDir, newTarName) + return nil +} + +// 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 metadata) +// 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 && e.Mode == f.mode && e.LinkTarget == f.linkTarget { + 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, + Mode: f.mode, + LinkTarget: f.linkTarget, + 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 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 (~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.DefaultCompression) + 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 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 { + 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 +} + +// 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 + } + 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 +} + +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 0000000000..ae0686708a --- /dev/null +++ b/experimental/air/cmd/snapshot_cache_test.go @@ -0,0 +1,255 @@ +package aircmd + +import ( + "archive/tar" + "compress/gzip" + "os" + "path/filepath" + "slices" + "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 { + fullPath := filepath.Join(repo, filepath.FromSlash(rel)) + info, err := os.Lstat(fullPath) + require.NoError(t, err) + linkTarget := "" + if info.Mode()&os.ModeSymlink != 0 { + linkTarget, err = os.Readlink(fullPath) + require.NoError(t, err) + } + files = append(files, snapshotFile{ + rel: filepath.FromSlash(rel), + size: info.Size(), + modTime: info.ModTime().UnixNano(), + mode: uint32(info.Mode()), + linkTarget: linkTarget, + }) + } + 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 TestSnapshotChangedDetectsModeAndSymlinkTarget(t *testing.T) { + old := &snapshotManifest{Entries: map[string]cacheEntry{ + "run.sh": {Size: 1, ModTime: 100, Mode: 0o644}, + "model": {Size: 6, ModTime: 200, Mode: uint32(os.ModeSymlink | 0o777), LinkTarget: "v1.bin"}, + }} + unchanged := []snapshotFile{ + {rel: "run.sh", size: 1, modTime: 100, mode: 0o644}, + {rel: "model", size: 6, modTime: 200, mode: uint32(os.ModeSymlink | 0o777), linkTarget: "v1.bin"}, + } + assert.False(t, snapshotChanged(unchanged, old)) + + modeChanged := slices.Clone(unchanged) + modeChanged[0].mode = 0o755 + assert.True(t, snapshotChanged(modeChanged, old)) + + linkChanged := slices.Clone(unchanged) + linkChanged[1].linkTarget = "v2.bin" + assert.True(t, snapshotChanged(linkChanged, old)) +} + +// 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") + writeRepoFile(t, repo, "src/model.py", "print()") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + 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, 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, 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) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + writeRepoFile(t, repo, "keep.py", "keep") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + files := statFiles(t, repo, "a.txt", "keep.py") + 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 + // cached metadata: 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, cacheDir, oldTarPath, 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() + files := statFiles(t, repo, "a.txt", "b.txt") + 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") + 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, cacheDir, oldTarPath, 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() + files := statFiles(t, repo, "a.txt") + 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(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") +} diff --git a/experimental/air/cmd/snapshot_cachekey.go b/experimental/air/cmd/snapshot_cachekey.go index 44c5cb8f35..e3b782303f 100644 --- a/experimental/air/cmd/snapshot_cachekey.go +++ b/experimental/air/cmd/snapshot_cachekey.go @@ -7,6 +7,8 @@ package aircmd import ( "crypto/sha256" "encoding/hex" + "fmt" + "path/filepath" "slices" "strings" ) @@ -14,6 +16,32 @@ import ( // snapshotPackagingVersion is bumped when packaging logic changes in a way that invalidates existing caches const snapshotPackagingVersion = "v1" +// plainTarKeyVersion namespaces the plain_tar working-tree key (so it can never collide +// with a git_archive key) and lets us invalidate it if the fingerprint scheme changes. +// computePlainTarKey also folds in the shared snapshotPackagingVersion, so a +// packaging-logic bump invalidates both modes' keys. +const plainTarKeyVersion = "plaintar-v2" + +// computePlainTarKey returns a content-addressed key for a working-tree snapshot: the +// SHA-256 over every file's path, size, mtime, mode, and symlink target (sorted for stability). An unchanged +// tree yields the same key, so an already-uploaded tarball can be reused instead of +// re-packaged and re-uploaded. Regular-file content detection is size+mtime, not content — +// the same trade-off DABs file-sync makes — so an edit preserving both is not seen. +func computePlainTarKey(files []snapshotFile) string { + sorted := slices.Clone(files) + slices.SortFunc(sorted, func(a, b snapshotFile) int { + return strings.Compare(a.rel, b.rel) + }) + + h := sha256.New() + for _, f := range sorted { + fmt.Fprintf(h, "%s\x00%d\x00%d\x00%d\x00%s\n", + filepath.ToSlash(f.rel), f.size, f.modTime, f.mode, f.linkTarget) + } + fmt.Fprintf(h, "%s\x00%s", plainTarKeyVersion, snapshotPackagingVersion) + return hex.EncodeToString(h.Sum(nil)) +} + // computeSnapshotCacheKey returns a stable cache key for a snapshot tarball: the // SHA-256 digest of (commitSHA, normalized includePaths, snapshotPackagingVersion, // subtreePrefix). Changing any input yields a different entry. An empty subtree diff --git a/experimental/air/cmd/snapshot_cachekey_test.go b/experimental/air/cmd/snapshot_cachekey_test.go index 2f933aa40b..755e78030c 100644 --- a/experimental/air/cmd/snapshot_cachekey_test.go +++ b/experimental/air/cmd/snapshot_cachekey_test.go @@ -70,3 +70,28 @@ func TestComputeSnapshotCacheKeyProperties(t *testing.T) { // The version constant participates: a different version is a different key. assert.NotEqual(t, snapshotPackagingVersion, "") } + +// TestComputePlainTarKeyProperties pins the working-tree fingerprint behavior: it is +// order-independent and reacts to any change in a file's metadata identity. +func TestComputePlainTarKeyProperties(t *testing.T) { + base := []snapshotFile{ + {rel: "a.txt", size: 10, modTime: 100}, + {rel: "src/b.py", size: 20, modTime: 200}, + } + + // Order-independent: the files are sorted by path before hashing. + assert.Equal(t, + computePlainTarKey(base), + computePlainTarKey([]snapshotFile{base[1], base[0]}), + ) + + // A changed size, mtime, or path each yields a different key. + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 11, modTime: 100}, base[1]})) + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 10, modTime: 101}, base[1]})) + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "renamed.txt", size: 10, modTime: 100}, base[1]})) + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 10, modTime: 100, mode: 0o755}, base[1]})) + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 10, modTime: 100, linkTarget: "target"}, base[1]})) + + // Adding or dropping a file changes the key. + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey(base[:1])) +} diff --git a/experimental/air/cmd/snapshot_dabs.go b/experimental/air/cmd/snapshot_dabs.go index fd753d050d..7763526f56 100644 --- a/experimental/air/cmd/snapshot_dabs.go +++ b/experimental/air/cmd/snapshot_dabs.go @@ -38,14 +38,15 @@ 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, snapshotArtifactPath string, sidecarStore filer.Filer, sidecarBase string) (snapshotResult, error) { +func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath, snapshotArtifactPath 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 } - // Resolve how to package before touching the tarball: git_archive (pinned commit, - // cacheable) vs plain_tar (working tree, not cacheable). + // Resolve how to package before touching the tarball: git_archive (pinned commit) vs + // plain_tar (working tree). Both are content-addressed, so an unchanged input reuses + // the already-uploaded tarball instead of re-packaging and re-uploading. plan, err := resolveSnapshotPlan(ctx, newGitRepo(repoPath), snap.Git, snap.IncludePaths) if err != nil { return snapshotResult{}, err @@ -56,7 +57,7 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s if snap.RemoteVolume != nil { snapshotArtifactPath = *snap.RemoteVolume } - result, err := uploadSnapshotViaDABs(ctx, w, repoPath, plan, snapshotArtifactPath) + result, err := uploadSnapshotViaDABs(ctx, w, repoPath, configPath, plan, snapshotArtifactPath, noCache) if err != nil { return snapshotResult{}, err } @@ -126,28 +127,42 @@ func uploadSnapshotSidecars(ctx context.Context, sidecarStore filer.Filer, sidec return path.Join(sidecarBase, gitStateName), diffPath } -// snapshotTarballName is the uploaded filename for the snapshot. It is deterministic -// for git_archive — _.tar.gz keyed on (commit, include_paths, -// root_path subtree) — so an identical snapshot reuses the same remote object (see -// the cache check below). For plain_tar it is timestamped so concurrent submissions -// of the same directory don't clobber each other's upload (working-tree content -// isn't pinned to a SHA, so it can't be content-addressed). -func snapshotTarballName(plan snapshotPlan, dirName string) string { +// snapshotTarName resolves the content-addressed upload filename for the snapshot and, +// for plain_tar, the working-tree file listing used to build both the key and the tarball +// (nil for git_archive, which lists nothing locally). The name is _.tar.gz, +// keyed on (commit, include_paths, root_path subtree) for git_archive and on the +// working-tree metadata fingerprint for plain_tar, so an identical input +// reuses the same remote object (see the skip in uploadSnapshotViaDABs). +func snapshotTarName(ctx context.Context, repoPath string, plan snapshotPlan) (string, []snapshotFile, error) { + dirName := filepath.Base(repoPath) if plan.mode == modeGitArchive { key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths, plan.subtreePrefix) - return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]) + return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]), nil, nil + } + start := time.Now() + files, err := snapshotFiles(ctx, repoPath, plan.includePaths, plan.isGitRepo) + if err != nil { + return "", nil, err } - return fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) + var totalBytes int64 + for _, file := range files { + totalBytes += file.size + } + log.Debugf(ctx, "air snapshot listing: files=%d uncompressed_bytes=%d duration=%s", len(files), totalBytes, time.Since(start)) + return fmt.Sprintf("%s_%s.tar.gz", dirName, computePlainTarKey(files)[:16]), files, nil } -// 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 { +// packageSnapshot writes the snapshot to tarball per the resolved plan. The working-tree +// path uses the warm snapshot cache unless noCache is set. +func packageSnapshot(ctx context.Context, repoPath, configPath string, plan snapshotPlan, files []snapshotFile, tarball string, noCache bool) error { dirName := filepath.Base(repoPath) if plan.mode == modeGitArchive { return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths, plan.subtreePrefix) } - return createPlainTarball(ctx, repoPath, tarball, plan.includePaths, plan.isGitRepo) + if noCache { + return createPlainTarball(ctx, repoPath, tarball, files) + } + return packagePlainTarWithCache(ctx, repoPath, configPath, plan.includePaths, plan.isGitRepo, tarball) } // uploadSnapshotViaDABs uploads the snapshot through DABs' artifact-upload machinery @@ -156,10 +171,11 @@ func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, ta // the remote .internal path, and uploads the bytes. artifactPath is either the // already-resolved user's repo_snapshots directory or a configured UC Volume. // -// git_archive snapshots are cacheable: the tarball name is content-addressed by -// (commit, include_paths, root_path subtree), 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, artifactPath string) (snapshotResult, error) { +// The tarball name is content-addressed — by (commit, include_paths, root_path subtree) +// for git_archive and by the working-tree metadata fingerprint for plain_tar — +// so if the identical object is already uploaded we skip packaging and upload entirely +// and reuse the remote path. +func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath, configPath string, plan snapshotPlan, artifactPath 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. tmp, err := os.MkdirTemp("", "air-snapshot-*") @@ -168,7 +184,10 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r } defer os.RemoveAll(tmp) - tarName := snapshotTarballName(plan, filepath.Base(repoPath)) + tarName, files, err := snapshotTarName(ctx, repoPath, plan) + if err != nil { + return snapshotResult{}, err + } b := &bundle.Bundle{ BundleRootPath: tmp, @@ -199,33 +218,36 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r return snapshotResult{}, err } - // git_archive is cacheable by (commit, include_paths, root_path subtree): if the - // identical tarball is already uploaded, skip packaging + upload and reuse it. - // Only the config-path rewrite (ReplaceWithRemotePath) runs — no bytes move. - if plan.mode == modeGitArchive { - f, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b) - if diags.HasError() { + // Both modes are content-addressed by tarName: if the identical tarball is already + // uploaded, skip packaging + upload and reuse it. Only the config-path rewrite + // (ReplaceWithRemotePath) runs — no bytes move. + f, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b) + if diags.HasError() { + return snapshotResult{}, diags.Error() + } + if result, handled, err := maybeUploadSnapshotOverlay( + ctx, f, uploadPath, repoPath, configPath, plan, files, tmp, artifactPath, noCache, + ); handled { + return result, err + } + exists, err := snapshotExists(ctx, f, tarName) + if err != nil { + return snapshotResult{}, err + } + if exists { + if _, diags := libraries.ReplaceWithRemotePath(ctx, b); diags.HasError() { return snapshotResult{}, diags.Error() } - exists, err := snapshotExists(ctx, f, tarName) + remote, err := readCodeSourcePath(b) if err != nil { return snapshotResult{}, err } - if exists { - if _, diags := libraries.ReplaceWithRemotePath(ctx, b); diags.HasError() { - return snapshotResult{}, diags.Error() - } - remote, err := readCodeSourcePath(b) - if err != nil { - return snapshotResult{}, err - } - log.Debugf(ctx, "snapshot cache hit for %s at %s", shortSHA(plan.commitSHA), path.Join(uploadPath, tarName)) - return snapshotResult{CodeSourcePath: remote}, nil - } + log.Debugf(ctx, "snapshot upload skipped; reusing %s", path.Join(uploadPath, tarName)) + return snapshotResult{CodeSourcePath: remote}, nil } - // 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 { + // Miss: package the tarball locally, then upload the bytes. + if err := packageSnapshot(ctx, repoPath, configPath, plan, files, filepath.Join(tmp, tarName), noCache); err != nil { return snapshotResult{}, err } @@ -245,8 +267,8 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r } // snapshotExists reports whether name already exists in the artifact store, used to -// short-circuit a cacheable git_archive upload. A not-found is a clean miss (false, -// nil); any other error is surfaced. +// short-circuit a content-addressed upload (either mode). A not-found is a clean miss +// (false, nil); any other error is surfaced. func snapshotExists(ctx context.Context, store filer.Filer, name string) (bool, error) { _, err := store.Stat(ctx, name) if err == nil { diff --git a/experimental/air/cmd/snapshot_overlay.go b/experimental/air/cmd/snapshot_overlay.go new file mode 100644 index 0000000000..873aa43cea --- /dev/null +++ b/experimental/air/cmd/snapshot_overlay.go @@ -0,0 +1,560 @@ +package aircmd + +import ( + "archive/tar" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path" + "path/filepath" + "slices" + "strings" + "time" + + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/log" + "github.com/klauspost/pgzip" +) + +const ( + snapshotOverlayVersion = "v1" + snapshotOverlayDescriptorName = ".databricks/air-source-overlay-v1/descriptor" + snapshotOverlayDeletionsRoot = ".databricks/air-source-overlay-v1/deletions" + + snapshotAnchorMaxAge = 30 * 24 * time.Hour + snapshotOverlayMaxChanged = 1024 + snapshotOverlayMaxBytes int64 = 128 << 20 + snapshotOverlayMaxRatio = 0.25 +) + +type snapshotOverlayEntry struct { + Size int64 `json:"size"` + ModTime int64 `json:"mtime_ns"` + Mode uint32 `json:"mode"` + LinkTarget string `json:"link_target,omitempty"` +} + +type snapshotAnchorState struct { + Version string `json:"version"` + Branch string `json:"branch"` + CreatedAtUnix int64 `json:"created_at_unix"` + AnchorName string `json:"anchor_name"` + AnchorFingerprint string `json:"anchor_fingerprint"` + Component string `json:"component"` + Entries map[string]snapshotOverlayEntry `json:"entries"` + Results map[string]string `json:"results"` +} + +type snapshotOverlayDescriptor struct { + FormatVersion int `json:"format_version"` + Anchor string `json:"anchor"` + AnchorFingerprint string `json:"anchor_fingerprint"` + ResultFingerprint string `json:"result_fingerprint"` + Component string `json:"component"` +} + +type snapshotOverlayDiff struct { + Changed []snapshotFile + Deleted []string + ChangedBytes int64 + TotalBytes int64 +} + +func snapshotOverlayManifest(files []snapshotFile) map[string]snapshotOverlayEntry { + entries := make(map[string]snapshotOverlayEntry, len(files)) + for _, file := range files { + entries[filepath.ToSlash(file.rel)] = snapshotOverlayEntry{ + Size: file.size, + ModTime: file.modTime, + Mode: file.mode, + LinkTarget: file.linkTarget, + } + } + return entries +} + +func snapshotOverlayFingerprint(files []snapshotFile) string { + return computePlainTarKey(files) +} + +func computeSnapshotOverlayDiff(files []snapshotFile, anchor map[string]snapshotOverlayEntry) snapshotOverlayDiff { + current := make(map[string]struct{}, len(files)) + diff := snapshotOverlayDiff{} + for _, file := range files { + rel := filepath.ToSlash(file.rel) + current[rel] = struct{}{} + diff.TotalBytes += file.size + entry, ok := anchor[rel] + if !ok || entry.Size != file.size || entry.ModTime != file.modTime || + entry.Mode != file.mode || entry.LinkTarget != file.linkTarget { + diff.Changed = append(diff.Changed, file) + diff.ChangedBytes += file.size + } + } + for rel := range anchor { + if _, ok := current[rel]; !ok { + diff.Deleted = append(diff.Deleted, rel) + } + } + slices.SortFunc(diff.Changed, func(a, b snapshotFile) int { + return strings.Compare(filepath.ToSlash(a.rel), filepath.ToSlash(b.rel)) + }) + slices.Sort(diff.Deleted) + return diff +} + +func snapshotOverlayAnchorReason(diff snapshotOverlayDiff) string { + changedPaths := len(diff.Changed) + len(diff.Deleted) + if changedPaths > snapshotOverlayMaxChanged { + return "changed-path-count" + } + if diff.ChangedBytes > snapshotOverlayMaxBytes { + return "changed-bytes" + } + if diff.TotalBytes > 0 && float64(diff.ChangedBytes)/float64(diff.TotalBytes) > snapshotOverlayMaxRatio { + return "changed-byte-ratio" + } + return "" +} + +func snapshotOverlayStatePath(repoPath, configPath string, includePaths []string, branch string) (string, error) { + absRepo, err := filepath.Abs(repoPath) + if err != nil { + return "", err + } + absConfig, err := filepath.Abs(configPath) + if err != nil { + return "", err + } + branchKey := sha256.Sum256([]byte(snapshotOverlayVersion + "\x00" + branch)) + return filepath.Join( + snapshotCacheDir(absRepo, absConfig, includePaths), + "overlay-anchor-"+hex.EncodeToString(branchKey[:8])+".json", + ), nil +} + +func loadSnapshotAnchorState(statePath, branch, component string) *snapshotAnchorState { + file, err := os.Open(statePath) + if err != nil { + return nil + } + defer file.Close() + decoder := json.NewDecoder(io.LimitReader(file, 64<<20)) + decoder.DisallowUnknownFields() + var state snapshotAnchorState + if err := decoder.Decode(&state); err != nil { + return nil + } + var extra any + if err := decoder.Decode(&extra); err != io.EOF || state.Version != snapshotOverlayVersion || + state.Branch != branch || state.Component != component || state.CreatedAtUnix <= 0 || + state.CreatedAtUnix > time.Now().Add(time.Hour).Unix() || + state.AnchorName == "" || state.AnchorFingerprint == "" || state.Entries == nil || state.Results == nil { + return nil + } + if !validSnapshotObjectNameForKind(state.AnchorName, "anchor") || !validSnapshotFingerprint(state.AnchorFingerprint) { + return nil + } + anchorFiles := make([]snapshotFile, 0, len(state.Entries)) + for rel, entry := range state.Entries { + if validateSnapshotRelativePath(rel) != nil { + return nil + } + mode := os.FileMode(entry.Mode) + if entry.Size < 0 || (!mode.IsRegular() && mode&os.ModeSymlink == 0) || + (mode&os.ModeSymlink == 0 && entry.LinkTarget != "") { + return nil + } + anchorFiles = append(anchorFiles, snapshotFile{ + rel: filepath.FromSlash(rel), + size: entry.Size, + modTime: entry.ModTime, + mode: entry.Mode, + linkTarget: entry.LinkTarget, + }) + } + if snapshotOverlayFingerprint(anchorFiles) != state.AnchorFingerprint || state.Results[state.AnchorFingerprint] != state.AnchorName { + return nil + } + for fingerprint, name := range state.Results { + if !validSnapshotFingerprint(fingerprint) || !validSnapshotObjectName(name) { + return nil + } + } + return &state +} + +func validSnapshotObjectName(name string) bool { + return validSnapshotObjectNameForKind(name, "anchor") || validSnapshotObjectNameForKind(name, "overlay") +} + +func validSnapshotObjectNameForKind(name, kind string) bool { + if len(name) > 255 || path.Base(name) != name || strings.ContainsAny(name, "/\\") || !strings.HasSuffix(name, ".tar.gz") { + return false + } + prefix := "air_" + kind + "_" + snapshotOverlayVersion + "_" + stem := strings.TrimSuffix(name, ".tar.gz") + return strings.HasPrefix(stem, prefix) && validSnapshotFingerprint(strings.TrimPrefix(stem, prefix)) +} + +func validSnapshotFingerprint(fingerprint string) bool { + if len(fingerprint) != sha256.Size*2 { + return false + } + _, err := hex.DecodeString(fingerprint) + return err == nil +} + +func saveSnapshotAnchorState(statePath string, state snapshotAnchorState) error { + data, err := json.Marshal(state) + if err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(statePath), 0o700); err != nil { + return fmt.Errorf("failed to create overlay state directory: %w", err) + } + tmp, err := os.CreateTemp(filepath.Dir(statePath), "overlay-anchor-*.json") + if err != nil { + return fmt.Errorf("failed to write overlay state: %w", err) + } + tmpName := tmp.Name() + if _, err := tmp.Write(data); err != nil { + tmp.Close() + os.Remove(tmpName) + return fmt.Errorf("failed to write overlay state: %w", err) + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return fmt.Errorf("failed to write overlay state: %w", err) + } + if err := os.Rename(tmpName, statePath); err != nil { + os.Remove(tmpName) + return fmt.Errorf("failed to install overlay state: %w", err) + } + return nil +} + +// maybeUploadSnapshotOverlay handles the WSFS incremental path. A false +// handled result leaves the caller on the existing full-snapshot path. +func maybeUploadSnapshotOverlay( + ctx context.Context, + store filer.Filer, + uploadPath string, + repoPath string, + configPath string, + plan snapshotPlan, + files []snapshotFile, + tmp string, + artifactPath string, + noCache bool, +) (result snapshotResult, handled bool, err error) { + isVolumePath := artifactPath == "/Volumes" || strings.HasPrefix(artifactPath, "/Volumes/") + if noCache || isVolumePath || plan.mode != modePlainTar { + return snapshotResult{}, false, nil + } + + var totalBytes int64 + for _, file := range files { + totalBytes += file.size + } + if totalBytes < snapshotCacheMinBytes || filepath.Base(repoPath) == ".databricks" { + log.Debugf(ctx, "air snapshot overlay: decision=full-existing reason=small-or-reserved files=%d uncompressed_bytes=%d", len(files), totalBytes) + return snapshotResult{}, false, nil + } + + branch := "" + if plan.isGitRepo { + branch = newGitRepo(repoPath).currentBranch(ctx) + } + component := filepath.Base(repoPath) + statePath, err := snapshotOverlayStatePath(repoPath, configPath, plan.includePaths, branch) + if err != nil { + return snapshotResult{}, true, err + } + fingerprint := snapshotOverlayFingerprint(files) + state := loadSnapshotAnchorState(statePath, branch, component) + anchorReason := "" + if state == nil { + anchorReason = "missing-or-invalid-state" + } else if time.Since(time.Unix(state.CreatedAtUnix, 0)) >= snapshotAnchorMaxAge { + anchorReason = "anchor-age" + } else { + exists, statErr := snapshotExists(ctx, store, state.AnchorName) + if statErr != nil { + return snapshotResult{}, true, statErr + } + if !exists { + anchorReason = "missing-remote-anchor" + } + } + + if anchorReason == "" { + if name, ok := state.Results[fingerprint]; ok { + exists, statErr := snapshotExists(ctx, store, name) + if statErr != nil { + return snapshotResult{}, true, statErr + } + if exists { + log.Debugf(ctx, "air snapshot overlay: decision=reuse branch=%q object=%s", branch, name) + return snapshotResult{CodeSourcePath: path.Join(uploadPath, name)}, true, nil + } + delete(state.Results, fingerprint) + } + + diff := computeSnapshotOverlayDiff(files, state.Entries) + anchorReason = snapshotOverlayAnchorReason(diff) + changedPaths := len(diff.Changed) + len(diff.Deleted) + ratio := 0.0 + if diff.TotalBytes > 0 { + ratio = float64(diff.ChangedBytes) / float64(diff.TotalBytes) + } + log.Debugf(ctx, "air snapshot overlay: decision-input branch=%q changed_paths=%d changed_bytes=%d total_bytes=%d changed_ratio=%.4f reanchor_reason=%q", + branch, changedPaths, diff.ChangedBytes, diff.TotalBytes, ratio, anchorReason) + if anchorReason == "" { + archivePath := filepath.Join(tmp, "overlay.tar.gz") + packageStart := time.Now() + if err := createSnapshotOverlay(repoPath, component, state, fingerprint, diff, archivePath); err != nil { + return snapshotResult{}, true, err + } + afterFiles, err := snapshotFiles(ctx, repoPath, plan.includePaths, plan.isGitRepo) + if err != nil { + return snapshotResult{}, true, err + } + if snapshotOverlayFingerprint(afterFiles) != fingerprint { + return snapshotResult{}, true, errors.New("code source changed while creating its overlay snapshot; retry submission") + } + packageDuration := time.Since(packageStart) + archiveInfo, err := os.Stat(archivePath) + if err != nil { + return snapshotResult{}, true, err + } + publishStart := time.Now() + name, err := publishImmutableSnapshot(ctx, store, "air_overlay_"+snapshotOverlayVersion, archivePath) + if err != nil { + return snapshotResult{}, true, err + } + log.Debugf(ctx, "air snapshot overlay: decision=overlay package=%s publish=%s compressed_bytes=%d object=%s", + packageDuration, time.Since(publishStart), archiveInfo.Size(), name) + state.Results[fingerprint] = name + if err := saveSnapshotAnchorState(statePath, *state); err != nil { + return snapshotResult{}, true, err + } + return snapshotResult{CodeSourcePath: path.Join(uploadPath, name)}, true, nil + } + } + + log.Debugf(ctx, "air snapshot overlay: decision=new-anchor branch=%q reason=%s files=%d uncompressed_bytes=%d", branch, anchorReason, len(files), totalBytes) + archivePath := filepath.Join(tmp, "anchor.tar.gz") + packageStart := time.Now() + if err := packageSnapshot(ctx, repoPath, configPath, plan, files, archivePath, false); err != nil { + return snapshotResult{}, true, err + } + afterFiles, err := snapshotFiles(ctx, repoPath, plan.includePaths, plan.isGitRepo) + if err != nil { + return snapshotResult{}, true, err + } + if snapshotOverlayFingerprint(afterFiles) != fingerprint { + return snapshotResult{}, true, errors.New("code source changed while creating its anchor snapshot; retry submission") + } + packageDuration := time.Since(packageStart) + archiveInfo, err := os.Stat(archivePath) + if err != nil { + return snapshotResult{}, true, err + } + publishStart := time.Now() + name, err := publishImmutableSnapshot(ctx, store, "air_anchor_"+snapshotOverlayVersion, archivePath) + if err != nil { + return snapshotResult{}, true, err + } + log.Debugf(ctx, "air snapshot overlay: decision=anchor package=%s publish=%s compressed_bytes=%d object=%s", + packageDuration, time.Since(publishStart), archiveInfo.Size(), name) + newState := snapshotAnchorState{ + Version: snapshotOverlayVersion, + Branch: branch, + CreatedAtUnix: time.Now().Unix(), + AnchorName: name, + AnchorFingerprint: fingerprint, + Component: component, + Entries: snapshotOverlayManifest(files), + Results: map[string]string{fingerprint: name}, + } + if err := saveSnapshotAnchorState(statePath, newState); err != nil { + return snapshotResult{}, true, err + } + return snapshotResult{CodeSourcePath: path.Join(uploadPath, name)}, true, nil +} + +func createSnapshotOverlay( + repoPath string, + component string, + state *snapshotAnchorState, + resultFingerprint string, + diff snapshotOverlayDiff, + outputPath string, +) (err error) { + output, err := os.Create(outputPath) + if err != nil { + return fmt.Errorf("failed to create overlay: %w", err) + } + gz, err := pgzip.NewWriterLevel(output, pgzip.DefaultCompression) + if err != nil { + output.Close() + return err + } + tw := tar.NewWriter(gz) + defer func() { + err = firstErr(err, tw.Close(), gz.Close(), output.Close()) + }() + + descriptor, err := json.Marshal(snapshotOverlayDescriptor{ + FormatVersion: 1, + Anchor: state.AnchorName, + AnchorFingerprint: state.AnchorFingerprint, + ResultFingerprint: resultFingerprint, + Component: component, + }) + if err != nil { + return err + } + if err := tw.WriteHeader(&tar.Header{ + Name: snapshotOverlayDescriptorName, + Mode: 0o600, + Size: int64(len(descriptor)), + }); err != nil { + return fmt.Errorf("failed to write overlay descriptor: %w", err) + } + if _, err := tw.Write(descriptor); err != nil { + return fmt.Errorf("failed to write overlay descriptor: %w", err) + } + + for _, rel := range diff.Deleted { + if err := validateSnapshotRelativePath(rel); err != nil { + return err + } + if err := tw.WriteHeader(&tar.Header{ + Name: path.Join(snapshotOverlayDeletionsRoot, rel), + Mode: 0o600, + }); err != nil { + return fmt.Errorf("failed to write deletion marker for %q: %w", rel, err) + } + } + for _, file := range diff.Changed { + if err := writeSnapshotOverlayMember(tw, repoPath, component, file); err != nil { + return err + } + } + return nil +} + +func validateSnapshotRelativePath(rel string) error { + clean := path.Clean(filepath.ToSlash(rel)) + if clean == "." || clean == ".." || strings.HasPrefix(clean, "../") || path.IsAbs(clean) || clean != filepath.ToSlash(rel) { + return fmt.Errorf("invalid snapshot path %q", rel) + } + return nil +} + +func writeSnapshotOverlayMember(tw *tar.Writer, repoPath, component string, file snapshotFile) error { + rel := filepath.ToSlash(file.rel) + if err := validateSnapshotRelativePath(rel); err != nil { + return err + } + fullPath := filepath.Join(repoPath, file.rel) + info, linkTarget, err := inspectSnapshotOverlayFile(fullPath) + if err != nil { + return fmt.Errorf("failed to inspect changed snapshot file %q: %w", rel, err) + } + if !snapshotFileMatches(file, info, linkTarget) { + return fmt.Errorf("snapshot file changed while packaging: %q", rel) + } + if !info.Mode().IsRegular() && info.Mode()&os.ModeSymlink == 0 { + return nil + } + header, err := tar.FileInfoHeader(info, linkTarget) + if err != nil { + return fmt.Errorf("failed to build overlay header for %q: %w", rel, err) + } + header.Name = path.Join(component, rel) + if err := tw.WriteHeader(header); err != nil { + return fmt.Errorf("failed to write overlay header for %q: %w", rel, err) + } + if info.Mode().IsRegular() { + input, err := os.Open(fullPath) + if err != nil { + return fmt.Errorf("failed to open changed snapshot file %q: %w", rel, err) + } + _, copyErr := io.Copy(tw, input) + closeErr := input.Close() + if copyErr != nil { + return fmt.Errorf("failed to write changed snapshot file %q: %w", rel, copyErr) + } + if closeErr != nil { + return fmt.Errorf("failed to close changed snapshot file %q: %w", rel, closeErr) + } + } + after, afterLink, err := inspectSnapshotOverlayFile(fullPath) + if err != nil || !snapshotFileMatches(file, after, afterLink) { + return fmt.Errorf("snapshot file changed while packaging: %q", rel) + } + return nil +} + +func inspectSnapshotOverlayFile(name string) (os.FileInfo, string, error) { + info, err := os.Lstat(name) + if err != nil { + return nil, "", err + } + linkTarget := "" + if info.Mode()&os.ModeSymlink != 0 { + linkTarget, err = os.Readlink(name) + } + return info, linkTarget, err +} + +func snapshotFileMatches(file snapshotFile, info os.FileInfo, linkTarget string) bool { + return info != nil && file.size == info.Size() && file.modTime == info.ModTime().UnixNano() && + file.mode == uint32(info.Mode()) && file.linkTarget == linkTarget +} + +func publishImmutableSnapshot(ctx context.Context, store filer.Filer, kind, archivePath string) (string, error) { + input, err := os.Open(archivePath) + if err != nil { + return "", err + } + hash := sha256.New() + if _, err := io.Copy(hash, input); err != nil { + input.Close() + return "", fmt.Errorf("failed to hash snapshot archive: %w", err) + } + if _, err := input.Seek(0, io.SeekStart); err != nil { + input.Close() + return "", err + } + name := kind + "_" + hex.EncodeToString(hash.Sum(nil)) + ".tar.gz" + err = store.Write(ctx, name, input, filer.CreateParentDirectories) + closeErr := input.Close() + if err != nil && !errors.Is(err, fs.ErrExist) { + return "", fmt.Errorf("failed to publish immutable snapshot %s: %w", name, err) + } + if closeErr != nil { + return "", closeErr + } + localInfo, err := os.Stat(archivePath) + if err != nil { + return "", err + } + remoteInfo, err := store.Stat(ctx, name) + if err != nil { + return "", fmt.Errorf("failed to verify immutable snapshot %s: %w", name, err) + } + if remoteInfo.Size() != localInfo.Size() { + return "", fmt.Errorf("immutable snapshot %s has unexpected size: got %d, want %d", name, remoteInfo.Size(), localInfo.Size()) + } + return name, nil +} diff --git a/experimental/air/cmd/snapshot_overlay_test.go b/experimental/air/cmd/snapshot_overlay_test.go new file mode 100644 index 0000000000..3118e83655 --- /dev/null +++ b/experimental/air/cmd/snapshot_overlay_test.go @@ -0,0 +1,217 @@ +package aircmd + +import ( + "archive/tar" + "compress/gzip" + "encoding/json" + "io" + "os" + "path/filepath" + "testing" + "time" + + "github.com/databricks/cli/libs/filer" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func overlayTestFile(t *testing.T, repo, rel string) snapshotFile { + t.Helper() + info, err := os.Lstat(filepath.Join(repo, filepath.FromSlash(rel))) + require.NoError(t, err) + linkTarget := "" + if info.Mode()&os.ModeSymlink != 0 { + linkTarget, err = os.Readlink(filepath.Join(repo, filepath.FromSlash(rel))) + require.NoError(t, err) + } + return snapshotFile{ + rel: filepath.FromSlash(rel), + size: info.Size(), + modTime: info.ModTime().UnixNano(), + mode: uint32(info.Mode()), + linkTarget: linkTarget, + } +} + +func TestComputeSnapshotOverlayDiff(t *testing.T) { + anchor := map[string]snapshotOverlayEntry{ + "same": {Size: 1, ModTime: 1, Mode: 0o644}, + "mode": {Size: 2, ModTime: 2, Mode: 0o644}, + "link": {Size: 3, ModTime: 3, Mode: uint32(os.ModeSymlink | 0o777), LinkTarget: "old"}, + "deleted": {Size: 10, ModTime: 4, Mode: 0o644}, + } + files := []snapshotFile{ + {rel: "same", size: 1, modTime: 1, mode: 0o644}, + {rel: "mode", size: 2, modTime: 2, mode: 0o755}, + {rel: "link", size: 3, modTime: 3, mode: uint32(os.ModeSymlink | 0o777), linkTarget: "new"}, + {rel: "added", size: 4, modTime: 5, mode: 0o600}, + } + + diff := computeSnapshotOverlayDiff(files, anchor) + + require.Len(t, diff.Changed, 3) + assert.Equal(t, []string{"added", "link", "mode"}, []string{ + filepath.ToSlash(diff.Changed[0].rel), + filepath.ToSlash(diff.Changed[1].rel), + filepath.ToSlash(diff.Changed[2].rel), + }) + assert.Equal(t, []string{"deleted"}, diff.Deleted) + assert.Equal(t, int64(9), diff.ChangedBytes) + assert.Equal(t, int64(10), diff.TotalBytes) +} + +func TestSnapshotOverlayAnchorReason(t *testing.T) { + assert.Empty(t, snapshotOverlayAnchorReason(snapshotOverlayDiff{ChangedBytes: 1, TotalBytes: 100})) + assert.Equal(t, "changed-path-count", snapshotOverlayAnchorReason(snapshotOverlayDiff{ + Changed: make([]snapshotFile, snapshotOverlayMaxChanged+1), + })) + assert.Equal(t, "changed-bytes", snapshotOverlayAnchorReason(snapshotOverlayDiff{ + ChangedBytes: snapshotOverlayMaxBytes + 1, + TotalBytes: snapshotOverlayMaxBytes * 10, + })) + assert.Equal(t, "changed-byte-ratio", snapshotOverlayAnchorReason(snapshotOverlayDiff{ + ChangedBytes: 26, + TotalBytes: 100, + })) +} + +func TestCreateSnapshotOverlay(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "changed.sh", "#!/bin/sh\necho changed\n") + require.NoError(t, os.Chmod(filepath.Join(repo, "changed.sh"), 0o755)) + writeRepoFile(t, repo, "added.txt", "added") + require.NoError(t, os.Symlink("added.txt", filepath.Join(repo, "current"))) + files := []snapshotFile{ + overlayTestFile(t, repo, "changed.sh"), + overlayTestFile(t, repo, "added.txt"), + overlayTestFile(t, repo, "current"), + } + state := &snapshotAnchorState{ + AnchorName: "air_anchor_v1_aabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccdd.tar.gz", + AnchorFingerprint: "aabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccdd", + } + out := filepath.Join(t.TempDir(), "overlay.tar.gz") + require.NoError(t, createSnapshotOverlay(repo, filepath.Base(repo), state, "b"+state.AnchorFingerprint[1:], snapshotOverlayDiff{ + Changed: files, + Deleted: []string{"deleted.txt"}, + }, out)) + + archive, err := os.Open(out) + require.NoError(t, err) + defer archive.Close() + gz, err := gzip.NewReader(archive) + require.NoError(t, err) + defer gz.Close() + reader := tar.NewReader(gz) + + first, err := reader.Next() + require.NoError(t, err) + assert.Equal(t, snapshotOverlayDescriptorName, first.Name) + descriptorBytes, err := io.ReadAll(reader) + require.NoError(t, err) + var descriptor snapshotOverlayDescriptor + require.NoError(t, json.Unmarshal(descriptorBytes, &descriptor)) + assert.Equal(t, state.AnchorName, descriptor.Anchor) + assert.Equal(t, filepath.Base(repo), descriptor.Component) + + entries := map[string]*tar.Header{} + for { + header, err := reader.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + copy := *header + entries[header.Name] = © + } + assert.Contains(t, entries, snapshotOverlayDeletionsRoot+"/deleted.txt") + assert.Equal(t, int64(0), entries[snapshotOverlayDeletionsRoot+"/deleted.txt"].Size) + assert.Equal(t, int64(0o755), entries[filepath.Base(repo)+"/changed.sh"].Mode) + assert.Equal(t, byte(tar.TypeSymlink), entries[filepath.Base(repo)+"/current"].Typeflag) + assert.Equal(t, "added.txt", entries[filepath.Base(repo)+"/current"].Linkname) +} + +func TestSnapshotAnchorStateIsBranchLocalAndStrict(t *testing.T) { + repo := t.TempDir() + config := filepath.Join(repo, "run.yaml") + mainPath, err := snapshotOverlayStatePath(repo, config, []string{"src"}, "main") + require.NoError(t, err) + featurePath, err := snapshotOverlayStatePath(repo, config, []string{"src"}, "feature") + require.NoError(t, err) + assert.NotEqual(t, mainPath, featurePath) + + entries := map[string]snapshotOverlayEntry{"src/train.py": {Size: 1}} + anchorFingerprint := snapshotOverlayFingerprint([]snapshotFile{{rel: filepath.FromSlash("src/train.py"), size: 1}}) + anchorName := "air_anchor_v1_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.tar.gz" + state := snapshotAnchorState{ + Version: snapshotOverlayVersion, + Branch: "main", + CreatedAtUnix: time.Now().Unix(), + AnchorName: anchorName, + AnchorFingerprint: anchorFingerprint, + Component: filepath.Base(repo), + Entries: entries, + Results: map[string]string{anchorFingerprint: anchorName}, + } + require.NoError(t, saveSnapshotAnchorState(mainPath, state)) + assert.NotNil(t, loadSnapshotAnchorState(mainPath, "main", filepath.Base(repo))) + assert.Nil(t, loadSnapshotAnchorState(mainPath, "other", filepath.Base(repo))) + + require.NoError(t, os.WriteFile(mainPath, append([]byte(`{"unknown":true}`), '\n'), 0o600)) + assert.Nil(t, loadSnapshotAnchorState(mainPath, "main", filepath.Base(repo))) +} + +func TestMaybeUploadSnapshotOverlayLifecycle(t *testing.T) { + repo := t.TempDir() + largePath := filepath.Join(repo, "large.bin") + large, err := os.Create(largePath) + require.NoError(t, err) + require.NoError(t, large.Truncate(snapshotCacheMinBytes+1)) + require.NoError(t, large.Close()) + writeRepoFile(t, repo, "train.py", "print('one')") + configPath := filepath.Join(repo, "run.yaml") + require.NoError(t, os.WriteFile(configPath, []byte("test"), 0o600)) + + remoteDir := t.TempDir() + store, err := filer.NewLocalClient(remoteDir) + require.NoError(t, err) + plan := snapshotPlan{mode: modePlainTar, isGitRepo: false} + files, err := snapshotFiles(t.Context(), repo, nil, false) + require.NoError(t, err) + statePath, err := snapshotOverlayStatePath(repo, configPath, nil, "") + require.NoError(t, err) + t.Cleanup(func() { _ = os.RemoveAll(filepath.Dir(statePath)) }) + + anchor, handled, err := maybeUploadSnapshotOverlay( + t.Context(), store, "/Workspace/test/.internal", repo, configPath, plan, files, + t.TempDir(), "", false, + ) + require.NoError(t, err) + require.True(t, handled) + assert.Contains(t, filepath.Base(anchor.CodeSourcePath), "air_anchor_v1_") + + oldInfo, err := os.Stat(filepath.Join(repo, "train.py")) + require.NoError(t, err) + writeRepoFile(t, repo, "train.py", "print('two')") + require.NoError(t, os.Chtimes( + filepath.Join(repo, "train.py"), oldInfo.ModTime().Add(time.Second), oldInfo.ModTime().Add(time.Second), + )) + files, err = snapshotFiles(t.Context(), repo, nil, false) + require.NoError(t, err) + overlay, handled, err := maybeUploadSnapshotOverlay( + t.Context(), store, "/Workspace/test/.internal", repo, configPath, plan, files, + t.TempDir(), "", false, + ) + require.NoError(t, err) + require.True(t, handled) + assert.Contains(t, filepath.Base(overlay.CodeSourcePath), "air_overlay_v1_") + assert.NotEqual(t, anchor.CodeSourcePath, overlay.CodeSourcePath) + + reused, handled, err := maybeUploadSnapshotOverlay( + t.Context(), store, "/Workspace/test/.internal", repo, configPath, plan, files, + t.TempDir(), "", false, + ) + require.NoError(t, err) + require.True(t, handled) + assert.Equal(t, overlay.CodeSourcePath, reused.CodeSourcePath) +} diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index 131580c0aa..e13417011a 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -49,23 +49,18 @@ func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outpu return nil } -// createPlainTarball writes a gzipped tar of repoPath's working tree to +// createPlainTarball writes a gzipped tar of the given working-tree files to // outputTarball. The archive preserves repoPath's directory name as the top-level -// entry. When includePaths is set, only those paths (nested under the directory -// name) are archived. .git and macOS AppleDouble files are always excluded; a -// .gitignore at repoPath is honored. -func createPlainTarball(ctx context.Context, repoPath, outputTarball string, includePaths []string, isGitRepo bool) error { +// entry. files come pre-resolved from snapshotFiles (.gitignore honored, .git and +// macOS AppleDouble excluded), so the caller can reuse the same listing to +// content-address the upload. +func createPlainTarball(ctx context.Context, repoPath, outputTarball string, files []snapshotFile) error { dirName := filepath.Base(repoPath) - - files, err := snapshotFiles(ctx, repoPath, includePaths, isGitRepo) - if err != nil { - return err - } entries := make([]tarpack.Entry, len(files)) - for i, rel := range files { + for i, f := range files { entries[i] = tarpack.Entry{ - Name: filepath.ToSlash(filepath.Join(dirName, rel)), - Path: filepath.Join(repoPath, rel), + Name: filepath.ToSlash(filepath.Join(dirName, f.rel)), + Path: filepath.Join(repoPath, f.rel), } } @@ -95,7 +90,17 @@ 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) plus metadata used by content addressing, overlays, and the warm cache. +type snapshotFile struct { + rel string + size int64 + modTime int64 // Unix nanoseconds + mode uint32 + linkTarget string +} + +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-") @@ -119,7 +124,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 @@ -129,14 +134,32 @@ 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)) + // tarpack omits directories and other special files. Filter them from the + // manifest too, so an overlay describes exactly the full archive's contents. + if !info.Mode().IsRegular() && info.Mode()&os.ModeSymlink == 0 { + continue + } + linkTarget := "" + if info.Mode()&os.ModeSymlink != 0 { + linkTarget, err = os.Readlink(filepath.Join(repoPath, filepath.FromSlash(name))) + if err != nil { + return nil, fmt.Errorf("failed to inspect snapshot symlink %q: %w", name, err) + } + } + files = append(files, snapshotFile{ + rel: filepath.FromSlash(name), + size: info.Size(), + modTime: info.ModTime().UnixNano(), + mode: uint32(info.Mode()), + linkTarget: linkTarget, + }) } return files, nil } diff --git a/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go index 9c4f07c879..9bda10dad4 100644 --- a/experimental/air/cmd/snapshot_package_test.go +++ b/experimental/air/cmd/snapshot_package_test.go @@ -132,7 +132,9 @@ func TestCreatePlainTarball(t *testing.T) { writeRepoFile(t, repo, ".git/config", "x") out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, nil, false)) + files, err := snapshotFiles(ctx, repo, nil, false) + require.NoError(t, err) + require.NoError(t, createPlainTarball(ctx, repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) @@ -152,7 +154,9 @@ func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { writeRepoFile(t, repo, ".gitignore", "*.log\n") out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, nil, false)) + files, err := snapshotFiles(ctx, repo, nil, false) + require.NoError(t, err) + require.NoError(t, createPlainTarball(ctx, repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) @@ -167,7 +171,9 @@ func TestCreatePlainTarball_IncludePaths(t *testing.T) { writeRepoFile(t, repo, "src/model.py", "print()") out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, []string{"src"}, false)) + files, err := snapshotFiles(ctx, repo, []string{"src"}, false) + require.NoError(t, err) + require.NoError(t, createPlainTarball(ctx, repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) @@ -187,7 +193,9 @@ func TestCreatePlainTarball_HonorsNestedGitignoreAndNegation(t *testing.T) { writeRepoFile(t, repo, "nested/keep.tmp", "keep") out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(t.Context(), repo, out, nil, false)) + files, err := snapshotFiles(t.Context(), repo, nil, false) + require.NoError(t, err) + require.NoError(t, createPlainTarball(t.Context(), repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) @@ -207,7 +215,9 @@ func TestCreatePlainTarball_SkipsDeletedTrackedFiles(t *testing.T) { require.NoError(t, os.Remove(filepath.Join(repo, "deleted.txt"))) out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(t.Context(), repo, out, nil, true)) + files, err := snapshotFiles(t.Context(), repo, nil, true) + require.NoError(t, err) + require.NoError(t, createPlainTarball(t.Context(), repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go index d799dd96a7..df465fdb50 100644 --- a/experimental/air/cmd/snapshot_resolve.go +++ b/experimental/air/cmd/snapshot_resolve.go @@ -20,7 +20,8 @@ const ( // root_path subtree). modeGitArchive snapshotMode = iota // modePlainTar packages the working tree (including uncommitted changes) via - // `tar`. Not cacheable — working-tree content isn't pinned to a SHA. + // `tar`. Content-addressed by the working-tree metadata fingerprint, so an + // unchanged tree reuses the uploaded tarball; a same-size, same-mtime regular-file edit is missed. modePlainTar )