fix(targethasher): make directory hashing cancellation-aware (audit #14.3)

core/itg/graph/update.go

@@ -85,7 +85,7 @@ func (g *OptimizedGraph) UpdateGraph(
 		}
 	}
 	// compute hashes for source file, package group, and rule common targets
-	if err := computeAvailableHashes(sourceHasher, targets); err != nil {
+	if err := computeAvailableHashes(ctx, sourceHasher, targets); err != nil {
 		return err
 	}
 
@@ -147,6 +147,7 @@ func (g *OptimizedGraph) upsertExternalRuleTarget(target *targethasher.Target, i
 
 // computeAvailableHashes computes hashes that are available without dep traversal.
 func computeAvailableHashes(
+	ctx context.Context,
 	hasher targethasher.SourceHasher,
 	targets map[string]*targethasher.Target,
 ) error {
@@ -161,7 +162,7 @@ func computeAvailableHashes(
 			h.Write([]byte(name))
 			hash = h.Sum(nil)
 		case targethasher.SourceFileType:
-			h, err := hasher.HashSourceFile(target.SourceFile)
+			h, err := hasher.HashSourceFile(ctx, target.SourceFile)
 			if err != nil {
 				return err
 			}

core/itg/graph/update_test.go

@@ -29,9 +29,11 @@ import (
 type fakeSourceHasher struct {
 	result []byte
 	err    error
+	ctx    context.Context
 }
 
-func (f *fakeSourceHasher) HashSourceFile(_ *buildpb.SourceFile) ([]byte, error) {
+func (f *fakeSourceHasher) HashSourceFile(ctx context.Context, _ *buildpb.SourceFile) ([]byte, error) {
+	f.ctx = ctx
 	return f.result, f.err
 }
 
@@ -44,6 +46,7 @@ func TestComputeAvailableHashes(t *testing.T) {
 		t.Parallel()
 		expected := []byte{0xAB, 0xCD}
 		hasher := &fakeSourceHasher{result: expected}
+		ctx := context.WithValue(context.Background(), struct{}{}, "source-hash")
 		name := "//pkg:file.go"
 		targets := map[string]*targethasher.Target{
 			name: {
@@ -53,8 +56,9 @@ func TestComputeAvailableHashes(t *testing.T) {
 			},
 		}
 
-		require.NoError(t, computeAvailableHashes(hasher, targets))
+		require.NoError(t, computeAvailableHashes(ctx, hasher, targets))
 		assert.Equal(t, expected, targets[name].Hash)
+		assert.Equal(t, ctx, hasher.ctx)
 	})
 
 	t.Run("package group hashed by name", func(t *testing.T) {
@@ -64,7 +68,7 @@ func TestComputeAvailableHashes(t *testing.T) {
 			name: {Name: name, RuleType: targethasher.PackageGroup},
 		}
 
-		require.NoError(t, computeAvailableHashes(&fakeSourceHasher{}, targets))
+		require.NoError(t, computeAvailableHashes(context.Background(), &fakeSourceHasher{}, targets))
 
 		h := sha1.New()
 		h.Write([]byte(name))
@@ -78,7 +82,7 @@ func TestComputeAvailableHashes(t *testing.T) {
 			name: {Name: name, RuleType: targethasher.ExternalRuleType},
 		}
 
-		require.NoError(t, computeAvailableHashes(&fakeSourceHasher{}, targets))
+		require.NoError(t, computeAvailableHashes(context.Background(), &fakeSourceHasher{}, targets))
 		assert.Nil(t, targets[name].Hash, "external rule targets should not get a hash here")
 	})
 
@@ -89,7 +93,7 @@ func TestComputeAvailableHashes(t *testing.T) {
 			name: {Name: name, RuleType: targethasher.GeneratedFileType},
 		}
 
-		require.NoError(t, computeAvailableHashes(&fakeSourceHasher{}, targets))
+		require.NoError(t, computeAvailableHashes(context.Background(), &fakeSourceHasher{}, targets))
 		assert.Nil(t, targets[name].Hash, "generated file hash is resolved later")
 	})
 
@@ -106,7 +110,7 @@ func TestComputeAvailableHashes(t *testing.T) {
 			},
 		}
 
-		require.NoError(t, computeAvailableHashes(&fakeSourceHasher{}, targets))
+		require.NoError(t, computeAvailableHashes(context.Background(), &fakeSourceHasher{}, targets))
 		assert.NotNil(t, targets[name].HashWithoutDeps, "rule should have HashWithoutDeps after hashing")
 		assert.Nil(t, targets[name].Hash, "full hash is not computed here — deps are needed")
 	})
@@ -118,7 +122,7 @@ func TestComputeAvailableHashes(t *testing.T) {
 			"//pkg:f": {Name: "//pkg:f", RuleType: targethasher.SourceFileType, SourceFile: &buildpb.SourceFile{}},
 		}
 
-		err := computeAvailableHashes(hasher, targets)
+		err := computeAvailableHashes(context.Background(), hasher, targets)
 		assert.Error(t, err)
 	})
 }

core/targethasher/graph.go

@@ -643,7 +643,7 @@ func HashRecursively(ctx context.Context, p HashParam) ([]byte, error) {
 
 	switch target.RuleType {
 	case SourceFileType:
-		h, err := p.Hasher.HashSourceFile(target.SourceFile)
+		h, err := p.Hasher.HashSourceFile(ctx, target.SourceFile)
 		if err != nil {
 			return nil, err
 		}

core/targethasher/graph_test.go

@@ -297,11 +297,12 @@ func assertEqualTargetHash(t *testing.T, expected, actual Target) {
 
 func Test_fromProto(t *testing.T) {
 	ctrl := gomock.NewController(t)
+	ctx := context.WithValue(context.Background(), struct{}{}, "source-hash")
 
 	mockHasher := NewMockSourceHasher(ctrl)
 	mockHasher.EXPECT().
-		HashSourceFile(gomock.Any()).
-		DoAndReturn(func(s *buildpb.SourceFile) ([]byte, error) {
+		HashSourceFile(gomock.Eq(ctx), gomock.Any()).
+		DoAndReturn(func(_ context.Context, s *buildpb.SourceFile) ([]byte, error) {
 			h := newHash()
 			io.WriteString(h, s.GetName())
 			return h.Sum(nil), nil
@@ -311,7 +312,7 @@ func Test_fromProto(t *testing.T) {
 	q, err := bazel.FromFile("testdata/test.proto.bin")
 	require.NoError(t, err)
 
-	a, err := fromProto(context.Background(), q, mockHasher, "", set.NewSet[string](), set.NewSet[string](), nil, true)
+	a, err := fromProto(ctx, q, mockHasher, "", set.NewSet[string](), set.NewSet[string](), nil, true)
 	require.NoError(t, err)
 
 	assert.Empty(t, a.Warnings)

core/targethasher/sourcehasher.go

@@ -15,6 +15,7 @@
 package targethasher
 
 import (
+	"context"
 	"crypto/sha1"
 	"fmt"
 	"hash"
@@ -37,11 +38,15 @@ const (
 	_defaultSourceFileVisibility = "//visibility:private"
 )
 
+// cancelCheckInterval is the number of files hashed between cancellation
+// checks during a directory walk.
+const cancelCheckInterval = 1024
+
 // SourceHasher provides hashes for source nodes in the target graph. These
 // can be calculated based on disk contents or form other sources such as a
 // vcs system.
 type SourceHasher interface {
-	HashSourceFile(s *buildpb.SourceFile) ([]byte, error)
+	HashSourceFile(ctx context.Context, s *buildpb.SourceFile) ([]byte, error)
 }
 
 // diskHashHelper is a SourceHasher that provides hashes based on disk
@@ -73,11 +78,11 @@ func NewSourceHasher(p Params) SourceHasher {
 }
 
 // HashSourceFile does a no-op hash for the noOpHasher.
-func (hh *noOpHasher) HashSourceFile(sourceFile *buildpb.SourceFile) ([]byte, error) {
+func (hh *noOpHasher) HashSourceFile(_ context.Context, sourceFile *buildpb.SourceFile) ([]byte, error) {
 	return nil, nil
 }
 
-func (hh *diskHashHelper) HashSourceFile(sourceFile *buildpb.SourceFile) ([]byte, error) {
+func (hh *diskHashHelper) HashSourceFile(ctx context.Context, sourceFile *buildpb.SourceFile) ([]byte, error) {
 	nonDefaultVisibilities := filterVisibilityLabels(sourceFile.GetVisibilityLabel())
 	// The location may look like /foo/decl.go:1:1
 	location, _, _ := strings.Cut(sourceFile.GetLocation(), ":")
@@ -94,7 +99,7 @@ func (hh *diskHashHelper) HashSourceFile(sourceFile *buildpb.SourceFile) ([]byte
 
 	var hash hash.Hash
 	if fi.IsDir() {
-		hash, err = hashDir(location)
+		hash, err = hashDir(ctx, location)
 	} else {
 		hash, err = hashFile(location)
 	}
@@ -131,14 +136,25 @@ func hashFile(path string) (hash.Hash, error) {
 	return hash, nil
 }
 
-func hashDir(root string) (hash.Hash, error) {
+func hashDir(ctx context.Context, root string) (hash.Hash, error) {
+	if err := context.Cause(ctx); err != nil {
+		return nil, err
+	}
+
 	dirHash := newHash()
+	var fileCount int
 	walkDirFunc := func(path string, d fs.DirEntry, err error) error {
 		if err != nil {
 			return err
 		}
 
 		if d.Type().IsRegular() {
+			if fileCount%cancelCheckInterval == 0 {
+				if err := context.Cause(ctx); err != nil {
+					return err
+				}
+			}
+			fileCount++
 			fileHash, err := hashFile(path)
 			if err != nil {
 				return err

core/targethasher/sourcehasher_test.go

@@ -15,6 +15,8 @@
 package targethasher
 
 import (
+	"context"
+	"errors"
 	"os"
 	"path/filepath"
 	"testing"
@@ -41,7 +43,7 @@ func TestNewSourceHasher_buildsMaps(t *testing.T) {
 func TestNoOpHasher_returnsNil(t *testing.T) {
 	h := &noOpHasher{}
 	sf := &buildpb.SourceFile{Name: strPtr("//:dummy")}
-	got, err := h.HashSourceFile(sf)
+	got, err := h.HashSourceFile(context.Background(), sf)
 	assert.NoError(t, err, "unexpected error: %v", err)
 	assert.Nil(t, got, "expected nil hash, got %v", got)
 }
@@ -68,7 +70,7 @@ func TestDiskHashHelper_KnownFileHashUsed(t *testing.T) {
 		Location:        strPtr(abs + ":1:1"),
 		VisibilityLabel: []string{"//visibility:private"},
 	}
-	got, err := h.HashSourceFile(sf)
+	got, err := h.HashSourceFile(context.Background(), sf)
 	assert.NoError(t, err, "unexpected err: %v", err)
 	assert.Equal(t, string(known), string(got), "expected known hash %q, got %q", known, got)
 }
@@ -94,7 +96,7 @@ func TestDiskHashHelper_NonDefaultVisibilityForcesDiskHash(t *testing.T) {
 		Location:        strPtr(abs + ":1:1"),
 		VisibilityLabel: []string{"//visibility:public"}, // non-default
 	}
-	got, err := h.HashSourceFile(sf)
+	got, err := h.HashSourceFile(context.Background(), sf)
 	assert.NoError(t, err, "unexpected err: %v", err)
 	assert.NotEqual(t, string(got), "KNOWN", "expected disk hash, but got known hash")
 	assert.NotEqual(t, []byte{}, got, "expected non-empty disk hash")
@@ -121,7 +123,7 @@ func TestDiskHashHelper_HashesFileFromDisk(t *testing.T) {
 		Location:        strPtr(abs),
 		VisibilityLabel: []string{"//visibility:private"},
 	}
-	got, err := h.HashSourceFile(sf)
+	got, err := h.HashSourceFile(context.Background(), sf)
 	assert.NoError(t, err, "unexpected err: %v", err)
 	assert.NotEmpty(t, got, "expected non-empty hash from disk")
 }
@@ -150,9 +152,32 @@ func TestDiskHashHelper_HashesDirectory(t *testing.T) {
 		Location:        strPtr(dirAbs),
 		VisibilityLabel: []string{"//visibility:private"},
 	}
-	got, err := h.HashSourceFile(sf)
+	got, err := h.HashSourceFile(context.Background(), sf)
 	assert.NoError(t, err, "unexpected err: %v", err)
 	assert.NotEmpty(t, got, "expected non-empty hash for directory")
 }
 
 func strPtr(s string) *string { return &s }
+
+func TestDiskHashHelper_RespectsContextCancellation(t *testing.T) {
+	tmp := t.TempDir()
+	require.NoError(t, os.WriteFile(filepath.Join(tmp, "file.txt"), []byte("x"), 0o644))
+
+	h := &diskHashHelper{
+		workspaceroot:   tmp,
+		knownFileHashes: map[string][]byte{},
+	}
+	sf := &buildpb.SourceFile{
+		Name:            strPtr("//:directory"),
+		Location:        strPtr(tmp),
+		VisibilityLabel: []string{"//visibility:private"},
+	}
+
+	cause := errors.New("stop hashing")
+	ctx, cancel := context.WithCancelCause(context.Background())
+	cancel(cause)
+
+	_, err := h.HashSourceFile(ctx, sf)
+	require.Error(t, err)
+	assert.ErrorIs(t, err, cause)
+}