experimental/air: content-address the plain_tar upload

experimental/air/cmd/runsubmit_test.go

@@ -391,15 +391,29 @@ func testSidecarStore(t *testing.T, w *databricks.WorkspaceClient) (filer.Filer,
 	return f, base
 }
 
-// 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"})
@@ -420,14 +434,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, sidecarStore, sidecarBase)
+	first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
+	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, sidecarStore, sidecarBase)
 	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

experimental/air/cmd/snapshot_cachekey.go

@@ -7,13 +7,40 @@ package aircmd
 import (
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
+	"path/filepath"
 	"slices"
 	"strings"
 )
 
 // 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-v1"
+
+// computePlainTarKey returns a content-addressed key for a working-tree snapshot: the
+// SHA-256 over every file's path, size and mtime (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. The fingerprint is size+mtime, not content — the same
+// trade-off DABs file-sync makes — so an edit preserving both size and mtime 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\n", filepath.ToSlash(f.rel), f.size, f.modTime)
+	}
+	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).
 // Changing any input yields a different entry.

experimental/air/cmd/snapshot_cachekey_test.go

@@ -60,3 +60,26 @@ 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 path, size, or mtime.
+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]}))
+
+	// Adding or dropping a file changes the key.
+	assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey(base[:1]))
+}

experimental/air/cmd/snapshot_dabs.go

@@ -44,8 +44,9 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s
 		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
@@ -126,28 +127,33 @@ 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 — <dirName>_<cacheKey>.tar.gz keyed on (commit, include_paths) — so
-// an identical commit 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 <dirName>_<key>.tar.gz,
+// keyed on (commit, include_paths) for git_archive and on the working-tree fingerprint
+// (path+size+mtime) 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)
-		return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16])
+		return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]), nil, nil
+	}
+	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"))
+	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 {
-	dirName := filepath.Base(repoPath)
+// of the pinned commit for git_archive, else a plain tar of the pre-listed working-tree
+// files (nil for git_archive).
+func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, files []snapshotFile, tarball string) 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)
+	return createPlainTarball(ctx, repoPath, tarball, files)
 }
 
 // uploadSnapshotViaDABs uploads the snapshot through DABs' artifact-upload machinery
@@ -156,9 +162,10 @@ func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, ta
 // the remote .internal path, and uploads the bytes. When remoteVolume is set the
 // tarball goes to that UC Volume; otherwise to the user's repo_snapshots dir.
 //
-// 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.
+// The tarball name is content-addressed — by (commit, include_paths) for git_archive and
+// by the working-tree fingerprint (path+size+mtime) 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 string, plan snapshotPlan, remoteVolume string) (snapshotResult, error) {
 	// artifactPath is where DABs uploads the tarball; GetFilerForLibraries routes to
 	// a Workspace or Volume filer based on its prefix, then appends /.internal.
@@ -179,7 +186,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,
@@ -210,33 +220,31 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r
 		return snapshotResult{}, err
 	}
 
-	// git_archive is cacheable by (commit, include_paths): 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()
+	}
+	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, plan, files, filepath.Join(tmp, tarName)); err != nil {
 		return snapshotResult{}, err
 	}
 
@@ -256,8 +264,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 {

experimental/air/cmd/snapshot_package.go

@@ -34,12 +34,12 @@ 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 via `tar`. 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 {
+// top-level 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)
 	// Absolute so it resolves correctly regardless of tar's working dir (set below).
 	parent, err := filepath.Abs(filepath.Dir(repoPath))
@@ -58,18 +58,13 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc
 	}
 	outName := filepath.Base(outputTarball)
 
-	files, err := snapshotFiles(ctx, repoPath, includePaths, isGitRepo)
-	if err != nil {
-		return err
-	}
 	args := []string{"-czf", outName, "-C", parent, "--null", "--no-recursion", "-T", "-"}
-
 	cmd := exec.CommandContext(ctx, "tar", args...)
 	// Run tar in the output directory so the bare -f basename lands there.
 	cmd.Dir = outDirAbs
 	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
@@ -84,7 +79,15 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc
 	return nil
 }
 
-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 the size and mtime used to content-address the plain_tar upload.
+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-")
@@ -108,7 +111,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
@@ -118,14 +121,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
 }