server: add --max_inflight_bytes to bound concurrent request memory

cache/disk/disk_test.go

@@ -1854,3 +1854,87 @@ func TestResultFromProxyTooLargeToReserve(t *testing.T) {
 
 	close(fakeProxy.getEvents)
 }
+
+// --- Read/decompress hot-path benchmarks ---
+//
+// CAS blobs are stored zstd-compressed (the default storageMode), so serving an
+// uncompressed Get decompresses on the fly. Unlike the write path (which
+// allocated a fresh compressed-output buffer per 1 MiB chunk), the read path
+// gets its zstd decoder from a sync.Pool and reuses it, so steady-state
+// per-request churn is expected to be modest. The memory that drives OOM under
+// download bursts is the *live* per-stream state (pooled decoder window buffers
+// plus gRPC send buffers) multiplied by unbounded concurrency, not per-request
+// garbage. These benchmarks quantify the per-request read allocations; the
+// serial-vs-parallel gap hints at how much extra state each concurrent reader
+// pins (e.g. an additional pooled decoder).
+const readBenchBlobSize = 16 * 1024 * 1024 // 16 MiB
+const readBenchChunk = 2 * 1024 * 1024     // mirror the server's maxChunkSize reads
+
+// benchReadSetup creates a cache holding one large incompressible CAS blob and
+// returns the cache and the blob's hash.
+func benchReadSetup(b *testing.B) (Cache, string) {
+	b.Helper()
+	dir := b.TempDir()
+	dc, err := New(dir, readBenchBlobSize*4, WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		b.Fatal(err)
+	}
+	data, hash := testutils.RandomDataAndHash(readBenchBlobSize)
+	if err := dc.Put(context.Background(), cache.CAS, hash, readBenchBlobSize, bytes.NewReader(data)); err != nil {
+		b.Fatal(err)
+	}
+	return dc, hash
+}
+
+// drainGet performs one Get + full decompressing read into a reused buffer.
+func drainGet(b *testing.B, dc Cache, hash string, buf []byte) {
+	rc, _, err := dc.Get(context.Background(), cache.CAS, hash, readBenchBlobSize, 0)
+	if err != nil {
+		b.Fatal(err)
+	}
+	for {
+		_, rerr := rc.Read(buf)
+		if rerr == io.EOF {
+			break
+		}
+		if rerr != nil {
+			_ = rc.Close()
+			b.Fatal(rerr)
+		}
+	}
+	if err := rc.Close(); err != nil {
+		b.Fatal(err)
+	}
+}
+
+// BenchmarkDiskCacheGetDecompress measures per-request allocations for serving
+// (decompressing) a single blob. Run with -benchmem.
+func BenchmarkDiskCacheGetDecompress(b *testing.B) {
+	dc, hash := benchReadSetup(b)
+	buf := make([]byte, readBenchChunk)
+
+	b.SetBytes(readBenchBlobSize)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		drainGet(b, dc, hash, buf)
+	}
+}
+
+// BenchmarkDiskCacheGetDecompressParallel serves the blob from many goroutines
+// at once, approximating a download burst. Each goroutine uses its own read
+// buffer; any growth in B/op vs the serial benchmark reflects extra live state
+// pinned per concurrent reader (notably additional pooled zstd decoders).
+func BenchmarkDiskCacheGetDecompressParallel(b *testing.B) {
+	dc, hash := benchReadSetup(b)
+
+	b.SetBytes(readBenchBlobSize)
+	b.ReportAllocs()
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		buf := make([]byte, readBenchChunk)
+		for pb.Next() {
+			drainGet(b, dc, hash, buf)
+		}
+	})
+}

cache/grpcproxy/grpcproxy_test.go

@@ -234,7 +234,7 @@ func newFixture(t *testing.T, proxy cache.Proxy, storageMode string) *fixture {
 	grpcServer := grpc.NewServer()
 
 	go func() {
-		err := server.ServeGRPC(listener, grpcServer, false, false, true, unlimitedMaxCasBlobSize, diskCache, logger, logger)
+		err := server.ServeGRPC(listener, grpcServer, false, false, true, unlimitedMaxCasBlobSize, 0, diskCache, logger, logger)
 		if err != nil {
 			logger.Printf("%s", err.Error())
 		}

config/config.go

@@ -127,6 +127,7 @@ type Config struct {
 	LogTimezone                 string                    `yaml:"log_timezone"`
 	MaxBlobSize                 int64                     `yaml:"max_blob_size"`
 	MaxProxyBlobSize            int64                     `yaml:"max_proxy_blob_size"`
+	MaxInflightBytes            int64                     `yaml:"max_inflight_bytes"`
 
 	// Fields that are created by combinations of the flags above.
 	ProxyBackend cache.Proxy
@@ -185,7 +186,8 @@ func newFromArgs(dir string, maxSize int, storageMode string, zstdImplementation
 	logTimezone string,
 	maxSizeHardLimit int,
 	maxBlobSize int64,
-	maxProxyBlobSize int64) (*Config, error) {
+	maxProxyBlobSize int64,
+	maxInflightBytes int64) (*Config, error) {
 
 	c := Config{
 		HTTPAddress:                 httpAddress,
@@ -224,6 +226,7 @@ func newFromArgs(dir string, maxSize int, storageMode string, zstdImplementation
 		LogTimezone:                 logTimezone,
 		MaxBlobSize:                 maxBlobSize,
 		MaxProxyBlobSize:            maxProxyBlobSize,
+		MaxInflightBytes:            maxInflightBytes,
 	}
 
 	err := validateConfig(&c)
@@ -399,6 +402,10 @@ func validateConfig(c *Config) error {
 		return errors.New("the 'max_proxy_blob_size' flag/key must be a positive integer")
 	}
 
+	if c.MaxInflightBytes < 0 {
+		return errors.New("the 'max_inflight_bytes' flag/key must not be negative (0 disables the limit)")
+	}
+
 	if c.GoogleCloudStorage != nil && c.HTTPBackend != nil && c.S3CloudStorage != nil {
 		return errors.New("one can specify at most one proxying backend")
 	}
@@ -679,5 +686,6 @@ func get(ctx *cli.Context) (*Config, error) {
 		ctx.Int("max_size_hard_limit"),
 		ctx.Int64("max_blob_size"),
 		ctx.Int64("max_proxy_blob_size"),
+		ctx.Int64("max_inflight_bytes"),
 	)
 }

main.go

@@ -462,6 +462,7 @@ func startGrpcServer(c *config.Config, grpcServer **grpc.Server,
 		c.EnableACKeyInstanceMangling,
 		enableRemoteAssetAPI,
 		c.MaxBlobSize,
+		c.MaxInflightBytes,
 		diskCache, c.AccessLogger, c.ErrorLogger)
 }
 

server/BUILD.bazel

@@ -43,6 +43,7 @@ go_library(
         "@org_golang_google_grpc//status:go_default_library",
         "@org_golang_google_protobuf//encoding/protojson:go_default_library",
         "@org_golang_google_protobuf//proto:go_default_library",
+        "@org_golang_x_sync//semaphore:go_default_library",
     ],
 )
 

server/grpc.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"sync"
 
 	"google.golang.org/genproto/googleapis/bytestream"
 	"google.golang.org/grpc"
@@ -17,6 +18,8 @@ import (
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
 
+	"golang.org/x/sync/semaphore"
+
 	asset "github.com/buchgr/bazel-remote/v2/genproto/build/bazel/remote/asset/v1"
 	pb "github.com/buchgr/bazel-remote/v2/genproto/build/bazel/remote/execution/v2"
 	"github.com/buchgr/bazel-remote/v2/genproto/build/bazel/semver"
@@ -43,6 +46,39 @@ type grpcServer struct {
 	depsCheck           bool
 	mangleACKeys        bool
 	maxCasBlobSizeBytes int64
+
+	// maxInflightBytes bounds the total size of blobs being served/received
+	// concurrently, providing memory backpressure. Zero disables the limit.
+	// inflightSem is nil when the limit is disabled.
+	maxInflightBytes int64
+	inflightSem      *semaphore.Weighted
+}
+
+// acquireInflight reserves budget for a request handling a blob of the given
+// size, blocking (with backpressure) until the budget is available. It returns
+// a release function that must be called when the request completes. The weight
+// is clamped to [1, maxInflightBytes] so that a single blob larger than the
+// whole budget is still served (serialized) rather than deadlocking. When the
+// limit is disabled the returned release is a no-op.
+func (s *grpcServer) acquireInflight(ctx context.Context, size int64) (func(), error) {
+	if s.inflightSem == nil {
+		return func() {}, nil
+	}
+
+	w := size
+	if w < 1 {
+		w = 1
+	}
+	if w > s.maxInflightBytes {
+		w = s.maxInflightBytes
+	}
+
+	if err := s.inflightSem.Acquire(ctx, w); err != nil {
+		return func() {}, err
+	}
+
+	var once sync.Once
+	return func() { once.Do(func() { s.inflightSem.Release(w) }) }, nil
 }
 
 var readOnlyMethods = map[string]struct{}{
@@ -64,30 +100,39 @@ func ListenAndServeGRPC(
 	mangleACKeys bool,
 	enableRemoteAssetAPI bool,
 	maxCasBlobSizeBytes int64,
+	maxInflightBytes int64,
 	c disk.Cache, a cache.Logger, e cache.Logger) error {
 
 	listener, err := net.Listen(network, addr)
 	if err != nil {
 		return err
 	}
 
-	return ServeGRPC(listener, srv, validateACDeps, mangleACKeys, enableRemoteAssetAPI, maxCasBlobSizeBytes, c, a, e)
+	return ServeGRPC(listener, srv, validateACDeps, mangleACKeys, enableRemoteAssetAPI, maxCasBlobSizeBytes, maxInflightBytes, c, a, e)
 }
 
 func ServeGRPC(l net.Listener, srv *grpc.Server,
 	validateACDepsCheck bool,
 	mangleACKeys bool,
 	enableRemoteAssetAPI bool,
 	maxCasBlobSizeBytes int64,
+	maxInflightBytes int64,
 	c disk.Cache, a cache.Logger, e cache.Logger) error {
 
+	var inflightSem *semaphore.Weighted
+	if maxInflightBytes > 0 {
+		inflightSem = semaphore.NewWeighted(maxInflightBytes)
+	}
+
 	s := &grpcServer{
 		cache:               c,
 		accessLogger:        a,
 		errorLogger:         e,
 		depsCheck:           validateACDepsCheck,
 		mangleACKeys:        mangleACKeys,
 		maxCasBlobSizeBytes: maxCasBlobSizeBytes,
+		maxInflightBytes:    maxInflightBytes,
+		inflightSem:         inflightSem,
 	}
 	pb.RegisterActionCacheServer(srv, s)
 	pb.RegisterCapabilitiesServer(srv, s)

server/grpc_bytestream.go

@@ -101,6 +101,15 @@ func (s *grpcServer) Read(req *bytestream.ReadRequest,
 		return status.Error(codes.OutOfRange, msg)
 	}
 
+	// Bound the memory used to serve concurrent reads: hold budget
+	// proportional to the blob size for the whole duration of the stream
+	// (the gRPC send buffers accumulate while streaming, after Get returns).
+	releaseInflight, err := s.acquireInflight(resp.Context(), size)
+	if err != nil {
+		return status.Error(codes.Canceled, err.Error())
+	}
+	defer releaseInflight()
+
 	var rc io.ReadCloser
 	var foundSize int64
 
@@ -366,6 +375,18 @@ func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 
 	cmp := casblob.Identity
 
+	// releaseInflight is assigned once (in the receive goroutine, after the
+	// blob size is known) and released when Write returns. Reading it here is
+	// safe: Write only returns after receiving from recvResult/putResult, which
+	// happens-after the assignment. It stays nil (no-op) on paths that never
+	// acquire (e.g. bad resource, blob already present, or limit disabled).
+	var releaseInflight func()
+	defer func() {
+		if releaseInflight != nil {
+			releaseInflight()
+		}
+	}()
+
 	go func() {
 		firstIteration := true
 		var resourceName string
@@ -435,6 +456,15 @@ func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 					return
 				}
 
+				// Bound the memory used to receive/compress concurrent uploads.
+				// Held for the whole Write (released by the outer defer).
+				release, aerr := s.acquireInflight(srv.Context(), size)
+				if aerr != nil {
+					recvResult <- status.Error(codes.Canceled, aerr.Error())
+					return
+				}
+				releaseInflight = release
+
 				var rc io.ReadCloser = pr
 				if cmp == casblob.Zstandard {
 					dec, ok := decoderPool.Get().(*syncpool.DecoderWrapper)