diff --git a/cache/disk/disk.go b/cache/disk/disk.go index ed809fe..c4c4156 100644 --- a/cache/disk/disk.go +++ b/cache/disk/disk.go @@ -93,6 +93,12 @@ type diskCache struct { // observation, D15). Lazily initialized; see lruCaptureBudgetBytes. lruCaptureSem *semaphore.Weighted + // Cap on total declared output-directory Tree bytes read into memory + // per ActionResult validation; see WithTreeValidationSizeLimit. + // Zero or negative disables the cap. + maxTreeValidationBytes int64 + treeValidationExceeded func(declaredBytes int64) + mu sync.Mutex lru SizedLRU @@ -1008,6 +1014,25 @@ func (c *diskCache) GetValidatedActionResult(ctx context.Context, hash string) ( pendingValidations := []*pb.Digest{} + // Pre-flight guard on validation memory (see + // WithTreeValidationSizeLimit): every referenced Tree blob is read + // wholly into memory below, and their sizes are declared in the + // ActionResult, so the total is knowable before allocating anything. + // Over-cap results are reported as a miss - semantically safe, the + // client rebuilds - and the trip is surfaced via the callback. + if c.maxTreeValidationBytes > 0 && len(result.OutputDirectories) > 0 { + var declaredTreeBytes int64 + for _, d := range result.OutputDirectories { + declaredTreeBytes += d.TreeDigest.SizeBytes + } + if declaredTreeBytes > c.maxTreeValidationBytes { + if c.treeValidationExceeded != nil { + c.treeValidationExceeded(declaredTreeBytes) + } + return nil, nil, nil // aka "not found" + } + } + // treeLeafHashes collects output_directories Tree blob hashes for the LRU // closure. The Tree blobs are recorded as closure leaves (the sweep must // keep them) but are deliberately kept OUT of pendingValidations: their diff --git a/cache/disk/options.go b/cache/disk/options.go index 0dfef63..907ef6d 100644 --- a/cache/disk/options.go +++ b/cache/disk/options.go @@ -75,6 +75,25 @@ func WithMaxEntries(n int64) Option { } } +// WithTreeValidationSizeLimit caps the total declared bytes of the +// output-directory Tree blobs that GetValidatedActionResult will read into +// memory while validating one ActionResult. Validation buffers each +// referenced Tree blob wholly (read + unmarshal), so without a cap a single +// action with a huge output directory pins an arbitrarily large allocation +// per concurrent GetActionResult. An over-cap ActionResult is reported as a +// cache miss - always semantically safe, the client rebuilds - rather than +// an error. exceeded, when non-nil, is invoked once per capped validation +// with the total declared Tree bytes; a trip means a client was just forced +// to rebuild, so it should be wired to a visible metric. maxBytes <= 0 +// disables the cap. +func WithTreeValidationSizeLimit(maxBytes int64, exceeded func(declaredBytes int64)) Option { + return func(c *CacheConfig) error { + c.diskCache.maxTreeValidationBytes = maxBytes + c.diskCache.treeValidationExceeded = exceeded + return nil + } +} + func WithMaxBlobSize(size int64) Option { return func(c *CacheConfig) error { if size <= 0 { diff --git a/cache/disk/tree_validation_guard_test.go b/cache/disk/tree_validation_guard_test.go new file mode 100644 index 0000000..824e37a --- /dev/null +++ b/cache/disk/tree_validation_guard_test.go @@ -0,0 +1,176 @@ +package disk + +// Tests for WithTreeValidationSizeLimit: GetValidatedActionResult reads every +// referenced output-directory Tree blob wholly into memory, so a cap on the +// total declared Tree bytes is its only per-request memory bound. Over-cap +// results must be reported as a miss (semantically safe: the client +// rebuilds), never as an error. + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "os" + "testing" + + pb "github.com/buchgr/bazel-remote/v2/genproto/build/bazel/remote/execution/v2" + + "github.com/buchgr/bazel-remote/v2/cache" + testutils "github.com/buchgr/bazel-remote/v2/utils" + + "google.golang.org/protobuf/proto" +) + +// putTreeBackedActionResult stores a CAS directory tree and an ActionResult +// referencing it as an output directory, returning the AC hash and the +// declared Tree blob size. +func putTreeBackedActionResult(ctx context.Context, t *testing.T, c *diskCache) (string, int64) { + t.Helper() + + fileData := []byte("tree validation guard file contents") + fileHash := sha256.Sum256(fileData) + fileHashStr := hex.EncodeToString(fileHash[:]) + if err := c.Put(ctx, cache.CAS, fileHashStr, int64(len(fileData)), + bytes.NewReader(fileData)); err != nil { + t.Fatal(err) + } + + rootDir := pb.Directory{ + Files: []*pb.FileNode{ + { + Name: "file.txt", + Digest: &pb.Digest{ + Hash: fileHashStr, + SizeBytes: int64(len(fileData)), + }, + }, + }, + } + + tree := pb.Tree{Root: &rootDir} + treeData, err := proto.Marshal(&tree) + if err != nil { + t.Fatal(err) + } + treeHash := sha256.Sum256(treeData) + treeHashStr := hex.EncodeToString(treeHash[:]) + if err := c.Put(ctx, cache.CAS, treeHashStr, int64(len(treeData)), + bytes.NewReader(treeData)); err != nil { + t.Fatal(err) + } + + ar := pb.ActionResult{ + OutputFiles: []*pb.OutputFile{ + { + Path: "file.txt", + Digest: &pb.Digest{ + Hash: fileHashStr, + SizeBytes: int64(len(fileData)), + }, + }, + }, + OutputDirectories: []*pb.OutputDirectory{ + { + Path: "out", + TreeDigest: &pb.Digest{ + Hash: treeHashStr, + SizeBytes: int64(len(treeData)), + }, + }, + }, + } + arData, err := proto.Marshal(&ar) + if err != nil { + t.Fatal(err) + } + arHash := sha256.Sum256([]byte("tree validation guard action")) + arHashStr := hex.EncodeToString(arHash[:]) + if err := c.Put(ctx, cache.AC, arHashStr, int64(len(arData)), + bytes.NewReader(arData)); err != nil { + t.Fatal(err) + } + + return arHashStr, int64(len(treeData)) +} + +func newTreeValidationCache(t *testing.T, opts ...Option) *diskCache { + t.Helper() + cacheDir := testutils.TempDir(t) + t.Cleanup(func() { _ = os.RemoveAll(cacheDir) }) + + opts = append([]Option{WithAccessLogger(testutils.NewSilentLogger())}, opts...) + c, err := New(cacheDir, 1024*32, opts...) + if err != nil { + t.Fatal(err) + } + return c.(*diskCache) +} + +func TestTreeValidationSizeLimitAllowsUnderCap(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var trips []int64 + // The limit is set after the blobs exist, via a second cache below; + // here declare a generous cap up-front. + c := newTreeValidationCache(t, + WithTreeValidationSizeLimit(1024*1024, func(declared int64) { + trips = append(trips, declared) + })) + arHash, _ := putTreeBackedActionResult(ctx, t, c) + + result, data, err := c.GetValidatedActionResult(ctx, arHash) + if err != nil { + t.Fatal(err) + } + if result == nil || data == nil { + t.Fatal("expected a validated hit under the cap") + } + if len(trips) != 0 { + t.Fatalf("expected no guard trips, got %v", trips) + } +} + +func TestTreeValidationSizeLimitReportsMissOverCap(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // Build the blobs first with no cap, then flip the cap below the + // declared Tree size on the same cache to prove the guard alone flips + // hit to miss. + c := newTreeValidationCache(t) + arHash, treeBytes := putTreeBackedActionResult(ctx, t, c) + + result, data, err := c.GetValidatedActionResult(ctx, arHash) + if err != nil || result == nil || data == nil { + t.Fatalf("expected a hit before the cap: result=%v data=%v err=%v", result, data, err) + } + + var trips []int64 + c.maxTreeValidationBytes = treeBytes - 1 + c.treeValidationExceeded = func(declared int64) { + trips = append(trips, declared) + } + + result, data, err = c.GetValidatedActionResult(ctx, arHash) + if err != nil { + t.Fatalf("over-cap validation must be a miss, not an error: %v", err) + } + if result != nil || data != nil { + t.Fatal("expected a miss over the cap") + } + if len(trips) != 1 || trips[0] != treeBytes { + t.Fatalf("expected one trip with declared bytes %d, got %v", treeBytes, trips) + } + + // An exact cap admits the result again. + c.maxTreeValidationBytes = treeBytes + result, data, err = c.GetValidatedActionResult(ctx, arHash) + if err != nil || result == nil || data == nil { + t.Fatalf("expected a hit at the exact cap: result=%v data=%v err=%v", result, data, err) + } + if len(trips) != 1 { + t.Fatalf("expected no further trips, got %v", trips) + } +} diff --git a/server/grpc.go b/server/grpc.go index ca6f4d9..23ed688 100644 --- a/server/grpc.go +++ b/server/grpc.go @@ -7,6 +7,7 @@ import ( "net" "net/http" + "golang.org/x/sync/semaphore" "google.golang.org/genproto/googleapis/bytestream" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -49,6 +50,10 @@ type grpcServer struct { runtimeMetrics RuntimeMetrics readLimiter *readLimiter sourceBuffers *sourceBufferPool + // GetTree guards; see GetTreeLimits. + getTreeSem *semaphore.Weighted + getTreeMaxResponseBytes int64 + getTreeMetrics GetTreeMetrics // writePayloadConsumed, when set, is invoked once per fully-consumed // ByteStream WriteRequest; see WithWritePayloadConsumed. writePayloadConsumed func(*bytestream.WriteRequest) diff --git a/server/grpc_cas.go b/server/grpc_cas.go index 0896b5c..f22deda 100644 --- a/server/grpc_cas.go +++ b/server/grpc_cas.go @@ -7,6 +7,7 @@ import ( "encoding/hex" "errors" "io" + "math" "google.golang.org/genproto/googleapis/rpc/code" "google.golang.org/genproto/googleapis/rpc/status" @@ -349,6 +350,21 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest, return errNilDigest } + // Fail-fast concurrency guard (see GetTreeLimits): GetTree materializes + // the whole tree in memory, so its memory bound is this slot count times + // the response byte cap. No queueing - a saturated caller retries. + if s.getTreeSem != nil { + if !s.getTreeSem.TryAcquire(1) { + if s.getTreeMetrics != nil { + s.getTreeMetrics.GetTreeDenied(GetTreeDeniedSaturated) + } + s.accessLogger.Printf("%s %s DENIED: concurrency limit", errorPrefix, in.RootDigest.Hash) + return grpc_status.Error(codes.ResourceExhausted, + "too many concurrent GetTree requests, please retry") + } + defer s.getTreeSem.Release(1) + } + err := s.validateHash(in.RootDigest.Hash, in.RootDigest.SizeBytes, errorPrefix) if err != nil { return err @@ -365,6 +381,19 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest, return grpc_status.Error(codes.Unknown, err.Error()) } + // Running response byte budget (see GetTreeLimits). The response size is + // discovered directory-by-directory during the walk, so this is a + // mid-traversal check on accumulated serialized bytes, not an up-front + // reservation. Zero (disabled) means an effectively unlimited budget. + budget := s.getTreeMaxResponseBytes + if budget <= 0 { + budget = math.MaxInt64 + } + budget -= int64(len(data)) + if budget < 0 { + return s.getTreeOverBudget(errorPrefix, in.RootDigest.Hash) + } + dir := pb.Directory{} err = proto.Unmarshal(data, &dir) if err != nil { @@ -372,7 +401,10 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest, return grpc_status.Error(codes.DataLoss, err.Error()) } - err = s.fillDirectories(stream.Context(), &resp, &dir, errorPrefix) + err = s.fillDirectories(stream.Context(), &resp, &dir, &budget, errorPrefix) + if err == errGetTreeOverBudget { + return s.getTreeOverBudget(errorPrefix, in.RootDigest.Hash) + } if err != nil { return err } @@ -388,9 +420,27 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest, return nil } +// errGetTreeOverBudget aborts the traversal when the accumulated response +// exceeds the configured byte cap. Mapped to ResourceExhausted in GetTree. +var errGetTreeOverBudget = errors.New("GetTree response byte budget exceeded") + +func (s *grpcServer) getTreeOverBudget(errorPrefix, rootHash string) error { + if s.getTreeMetrics != nil { + s.getTreeMetrics.GetTreeDenied(GetTreeDeniedResponseBytes) + } + s.accessLogger.Printf("%s %s DENIED: response over %d byte limit", + errorPrefix, rootHash, s.getTreeMaxResponseBytes) + return grpc_status.Errorf(codes.ResourceExhausted, + "tree exceeds the server's %d byte GetTree response limit", + s.getTreeMaxResponseBytes) +} + // Attempt to populate `resp`. Return errors for invalid requests, but -// otherwise attempt to return as many blobs as possible. -func (s *grpcServer) fillDirectories(ctx context.Context, resp *pb.GetTreeResponse, dir *pb.Directory, errorPrefix string) error { +// otherwise attempt to return as many blobs as possible. budget is the +// remaining serialized-byte allowance for the accumulated response; it is +// decremented as the walk discovers directories, and crossing it aborts +// with errGetTreeOverBudget. +func (s *grpcServer) fillDirectories(ctx context.Context, resp *pb.GetTreeResponse, dir *pb.Directory, budget *int64, errorPrefix string) error { // Add this dir. resp.Directories = append(resp.Directories, dir) @@ -403,14 +453,23 @@ func (s *grpcServer) fillDirectories(ctx context.Context, resp *pb.GetTreeRespon return err } + // Check the declared size before fetching, so the guard also + // bounds the transient getBlobData allocation. + *budget -= dirNode.Digest.SizeBytes + if *budget < 0 { + return errGetTreeOverBudget + } + data, err := s.getBlobData(ctx, dirNode.Digest.Hash, dirNode.Digest.SizeBytes) if err == errBlobNotFound { s.accessLogger.Printf("GRPC GETTREEREQUEST BLOB %s NOT FOUND", dirNode.Digest.Hash) + *budget += dirNode.Digest.SizeBytes continue } if err != nil { s.accessLogger.Printf("GRPC GETTREEREQUEST BLOB %s ERR: %v", err) + *budget += dirNode.Digest.SizeBytes continue } @@ -418,13 +477,14 @@ func (s *grpcServer) fillDirectories(ctx context.Context, resp *pb.GetTreeRespon err = proto.Unmarshal(data, &dirMsg) if err != nil { s.accessLogger.Printf("GRPC GETTREEREQUEST BAD BLOB: %v", err) + *budget += dirNode.Digest.SizeBytes continue } s.accessLogger.Printf("GRPC GETTREEREQUEST BLOB %s ADDED OK", dirNode.Digest.Hash) - err = s.fillDirectories(ctx, resp, &dirMsg, errorPrefix) + err = s.fillDirectories(ctx, resp, &dirMsg, budget, errorPrefix) if err != nil { return err } diff --git a/server/grpc_gettree_guard_test.go b/server/grpc_gettree_guard_test.go new file mode 100644 index 0000000..e7f4f79 --- /dev/null +++ b/server/grpc_gettree_guard_test.go @@ -0,0 +1,277 @@ +package server + +// Tests for the GetTree guards (WithGetTreeLimits): the fail-fast concurrency +// slot and the running response byte cap. GetTree materializes the whole +// directory tree into one in-memory response, so these guards are its only +// memory bound. + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "os" + "sync" + "testing" + + "golang.org/x/sync/semaphore" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + + pb "github.com/buchgr/bazel-remote/v2/genproto/build/bazel/remote/execution/v2" + + "github.com/buchgr/bazel-remote/v2/cache" + "github.com/buchgr/bazel-remote/v2/cache/disk" + testutils "github.com/buchgr/bazel-remote/v2/utils" +) + +// recordingGetTreeMetrics counts guard trips by reason. +type recordingGetTreeMetrics struct { + mu sync.Mutex + reasons []string +} + +func (m *recordingGetTreeMetrics) GetTreeDenied(reason string) { + m.mu.Lock() + defer m.mu.Unlock() + m.reasons = append(m.reasons, reason) +} + +func (m *recordingGetTreeMetrics) denied() []string { + m.mu.Lock() + defer m.mu.Unlock() + return append([]string(nil), m.reasons...) +} + +// fakeGetTreeStream satisfies pb.ContentAddressableStorage_GetTreeServer for +// direct handler invocation. Only Context and Send are exercised. +type fakeGetTreeStream struct { + grpc.ServerStream + ctx context.Context + responses []*pb.GetTreeResponse +} + +func (s *fakeGetTreeStream) Context() context.Context { return s.ctx } + +func (s *fakeGetTreeStream) Send(r *pb.GetTreeResponse) error { + s.responses = append(s.responses, r) + return nil +} + +// getTreeGuardHarness holds a CAS with a two-level directory tree and a +// directly-constructed grpcServer, so guard behavior can be tested without +// goroutines or wire plumbing. +type getTreeGuardHarness struct { + server *grpcServer + metrics *recordingGetTreeMetrics + rootDigest *pb.Digest + // Serialized bytes of the root directory blob and of all directory + // blobs together, for sizing byte caps precisely. + rootBytes int64 + totalBytes int64 + numDirs int +} + +func newGetTreeGuardHarness(t *testing.T) *getTreeGuardHarness { + t.Helper() + + ctx := context.Background() + cacheDir, err := os.MkdirTemp("", "bazel-remote-gettree-guard-"+t.Name()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.RemoveAll(cacheDir) }) + + diskCache, err := disk.New(cacheDir, 1024*1024, + disk.WithAccessLogger(testutils.NewSilentLogger())) + if err != nil { + t.Fatal(err) + } + + putDir := func(dir *pb.Directory) *pb.Digest { + data, err := proto.Marshal(dir) + if err != nil { + t.Fatal(err) + } + hash := sha256.Sum256(data) + hashStr := hex.EncodeToString(hash[:]) + err = diskCache.Put(ctx, cache.CAS, hashStr, int64(len(data)), + bytes.NewReader(data)) + if err != nil { + t.Fatal(err) + } + return &pb.Digest{Hash: hashStr, SizeBytes: int64(len(data))} + } + + fileBlob := []byte("gettree guard test file contents") + fileHash := sha256.Sum256(fileBlob) + fileHashStr := hex.EncodeToString(fileHash[:]) + err = diskCache.Put(ctx, cache.CAS, fileHashStr, int64(len(fileBlob)), + bytes.NewReader(fileBlob)) + if err != nil { + t.Fatal(err) + } + fileNode := &pb.FileNode{ + Name: "file.txt", + Digest: &pb.Digest{Hash: fileHashStr, SizeBytes: int64(len(fileBlob))}, + } + + subDigest := putDir(&pb.Directory{Files: []*pb.FileNode{fileNode}}) + rootDigest := putDir(&pb.Directory{ + Files: []*pb.FileNode{fileNode}, + Directories: []*pb.DirectoryNode{ + {Name: "subdir", Digest: subDigest}, + }, + }) + + metrics := &recordingGetTreeMetrics{} + return &getTreeGuardHarness{ + server: &grpcServer{ + cache: diskCache, + accessLogger: testutils.NewSilentLogger(), + errorLogger: testutils.NewSilentLogger(), + getTreeMetrics: metrics, + }, + metrics: metrics, + rootDigest: rootDigest, + rootBytes: rootDigest.SizeBytes, + totalBytes: rootDigest.SizeBytes + subDigest.SizeBytes, + numDirs: 2, + } +} + +func (h *getTreeGuardHarness) getTree(t *testing.T) ([]*pb.GetTreeResponse, error) { + t.Helper() + stream := &fakeGetTreeStream{ctx: context.Background()} + err := h.server.GetTree(&pb.GetTreeRequest{RootDigest: h.rootDigest}, stream) + return stream.responses, err +} + +func requireResourceExhausted(t *testing.T, err error) { + t.Helper() + if err == nil { + t.Fatal("expected ResourceExhausted, got success") + } + st, ok := status.FromError(err) + if !ok || st.Code() != codes.ResourceExhausted { + t.Fatalf("expected ResourceExhausted, got %v", err) + } +} + +func TestGetTreeConcurrencyGuardFailsFast(t *testing.T) { + h := newGetTreeGuardHarness(t) + h.server.getTreeSem = semaphore.NewWeighted(1) + + // With the only slot held, GetTree is denied immediately. + if !h.server.getTreeSem.TryAcquire(1) { + t.Fatal("failed to occupy the GetTree slot") + } + _, err := h.getTree(t) + requireResourceExhausted(t, err) + if got := h.metrics.denied(); len(got) != 1 || got[0] != GetTreeDeniedSaturated { + t.Fatalf("expected one %q denial, got %v", GetTreeDeniedSaturated, got) + } + + // Releasing the slot makes GetTree succeed, and the handler releases + // the slot on completion so a subsequent call also succeeds. + h.server.getTreeSem.Release(1) + for i := 0; i < 2; i++ { + responses, err := h.getTree(t) + if err != nil { + t.Fatalf("GetTree attempt %d failed after slot release: %v", i, err) + } + if len(responses) != 1 || len(responses[0].Directories) != h.numDirs { + t.Fatalf("expected one response with %d directories, got %v", h.numDirs, responses) + } + } + if got := h.metrics.denied(); len(got) != 1 { + t.Fatalf("expected no further denials, got %v", got) + } +} + +func TestGetTreeResponseByteCap(t *testing.T) { + h := newGetTreeGuardHarness(t) + + // A budget that covers the root but not the child directory aborts + // mid-traversal. + h.server.getTreeMaxResponseBytes = h.totalBytes - 1 + _, err := h.getTree(t) + requireResourceExhausted(t, err) + if got := h.metrics.denied(); len(got) != 1 || got[0] != GetTreeDeniedResponseBytes { + t.Fatalf("expected one %q denial, got %v", GetTreeDeniedResponseBytes, got) + } + + // A budget smaller than the root blob is denied before unmarshaling + // anything. + h.server.getTreeMaxResponseBytes = h.rootBytes - 1 + _, err = h.getTree(t) + requireResourceExhausted(t, err) + + // An exact budget serves the full tree. + h.server.getTreeMaxResponseBytes = h.totalBytes + responses, err := h.getTree(t) + if err != nil { + t.Fatalf("GetTree failed with an exact byte budget: %v", err) + } + if len(responses) != 1 || len(responses[0].Directories) != h.numDirs { + t.Fatalf("expected one response with %d directories, got %v", h.numDirs, responses) + } +} + +// TestGetTreeGuardsEndToEnd proves the option plumbing: a server configured +// through WithGetTreeLimits denies an over-budget tree over the wire. +func TestGetTreeGuardsEndToEnd(t *testing.T) { + metrics := &recordingGetTreeMetrics{} + fixture := grpcTestSetupInternal(t, false, WithGetTreeLimits(GetTreeLimits{ + MaxConcurrent: 1, + MaxResponseBytes: 1, + Metrics: metrics, + })) + defer func() { _ = os.RemoveAll(fixture.tempdir) }() + + dir := pb.Directory{} + dirData, err := proto.Marshal(&dir) + if err != nil { + t.Fatal(err) + } + // An empty Directory marshals to zero bytes, which would pass any + // budget; add a file to make it non-empty. + fileBlob := []byte("x") + fileHash := sha256.Sum256(fileBlob) + dir.Files = []*pb.FileNode{{ + Name: "f", + Digest: &pb.Digest{Hash: hex.EncodeToString(fileHash[:]), SizeBytes: 1}, + }} + dirData, err = proto.Marshal(&dir) + if err != nil { + t.Fatal(err) + } + dirHash := sha256.Sum256(dirData) + dirDigest := pb.Digest{ + Hash: hex.EncodeToString(dirHash[:]), + SizeBytes: int64(len(dirData)), + } + + upReq := pb.BatchUpdateBlobsRequest{ + Requests: []*pb.BatchUpdateBlobsRequest_Request{ + {Digest: &dirDigest, Data: dirData}, + }, + } + _, err = fixture.casClient.BatchUpdateBlobs(ctx, &upReq) + if err != nil { + t.Fatal(err) + } + + stream, err := fixture.casClient.GetTree(ctx, + &pb.GetTreeRequest{RootDigest: &dirDigest}) + if err != nil { + t.Fatal(err) + } + _, err = stream.Recv() + requireResourceExhausted(t, err) + if got := metrics.denied(); len(got) != 1 || got[0] != GetTreeDeniedResponseBytes { + t.Fatalf("expected one %q denial, got %v", GetTreeDeniedResponseBytes, got) + } +} diff --git a/server/runtime.go b/server/runtime.go index 6934402..3f7539f 100644 --- a/server/runtime.go +++ b/server/runtime.go @@ -106,6 +106,61 @@ func WithWritePayloadConsumed(fn func(*bytestream.WriteRequest)) GRPCServerOptio } } +// GetTree denial reasons reported to GetTreeMetrics. +const ( + // GetTreeDeniedSaturated means no concurrency slot was free. + GetTreeDeniedSaturated = "saturated" + // GetTreeDeniedResponseBytes means the materialized response exceeded + // the configured byte cap mid-traversal. + GetTreeDeniedResponseBytes = "response_bytes" +) + +// GetTreeMetrics receives GetTree guard trips. Implementations must be safe +// for concurrent use. +type GetTreeMetrics interface { + GetTreeDenied(reason string) +} + +// GetTreeLimits guards the GetTree handler, which materializes the entire +// directory tree into one in-memory response before sending (there is no +// pagination). Its response size is unknowable before the traversal runs, so +// unlike the batch paths it cannot take an up-front reservation; the guard is +// a fail-fast concurrency slot plus a running byte cap checked as the walk +// discovers directories. Both trips return ResourceExhausted (retryable) and +// are reported to Metrics; they are expected to stay silent, and a trip is +// the evidence bar for building pagination. +type GetTreeLimits struct { + // MaxConcurrent bounds simultaneous GetTree handlers (fail-fast, no + // queueing). Zero disables the concurrency guard. + MaxConcurrent int64 + // MaxResponseBytes bounds the serialized size of the directory data + // accumulated into one response. Zero disables the byte cap. + MaxResponseBytes int64 + // Metrics, when non-nil, receives guard trips. + Metrics GetTreeMetrics +} + +// WithGetTreeLimits installs the GetTree guards. Zero limits preserve +// bazel-remote's existing unbounded behavior. +func WithGetTreeLimits(limits GetTreeLimits) GRPCServerOption { + return func(s *grpcServer) error { + if limits.MaxConcurrent < 0 { + return fmt.Errorf("max concurrent GetTree must not be negative: %d", limits.MaxConcurrent) + } + if limits.MaxResponseBytes < 0 { + return fmt.Errorf("max GetTree response bytes must not be negative: %d", limits.MaxResponseBytes) + } + if limits.MaxConcurrent > 0 { + s.getTreeSem = semaphore.NewWeighted(limits.MaxConcurrent) + } else { + s.getTreeSem = nil + } + s.getTreeMaxResponseBytes = limits.MaxResponseBytes + s.getTreeMetrics = limits.Metrics + return nil + } +} + // WithMaxBatchTotalSizeBytes limits the total declared blob bytes in one // batch CAS request and advertises that limit through // CacheCapabilities.MaxBatchTotalSizeBytes. Per REAPI that capability bounds