diff --git a/cache/disk/disk_test.go b/cache/disk/disk_test.go index d3eebb90c..4eabf9fba 100644 --- a/cache/disk/disk_test.go +++ b/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) + } + }) +} diff --git a/cache/grpcproxy/grpcproxy_test.go b/cache/grpcproxy/grpcproxy_test.go index 79d7546a4..76588de89 100644 --- a/cache/grpcproxy/grpcproxy_test.go +++ b/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()) } diff --git a/config/config.go b/config/config.go index 05cb40172..aa0126a8c 100644 --- a/config/config.go +++ b/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"), ) } diff --git a/main.go b/main.go index 895816225..88a26a700 100644 --- a/main.go +++ b/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) } diff --git a/server/BUILD.bazel b/server/BUILD.bazel index 9411df57c..1a333b7cd 100644 --- a/server/BUILD.bazel +++ b/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", ], ) diff --git a/server/grpc.go b/server/grpc.go index c4300683e..4f2ee1ba7 100644 --- a/server/grpc.go +++ b/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,6 +100,7 @@ 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) @@ -71,7 +108,7 @@ func ListenAndServeGRPC( 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, @@ -79,8 +116,14 @@ func ServeGRPC(l net.Listener, srv *grpc.Server, 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, @@ -88,6 +131,8 @@ func ServeGRPC(l net.Listener, srv *grpc.Server, depsCheck: validateACDepsCheck, mangleACKeys: mangleACKeys, maxCasBlobSizeBytes: maxCasBlobSizeBytes, + maxInflightBytes: maxInflightBytes, + inflightSem: inflightSem, } pb.RegisterActionCacheServer(srv, s) pb.RegisterCapabilitiesServer(srv, s) diff --git a/server/grpc_bytestream.go b/server/grpc_bytestream.go index 1c5f738ea..a81aa068b 100644 --- a/server/grpc_bytestream.go +++ b/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) diff --git a/server/grpc_test.go b/server/grpc_test.go index ea73b6d8c..e600b8bca 100644 --- a/server/grpc_test.go +++ b/server/grpc_test.go @@ -11,6 +11,8 @@ import ( "net" "net/http" "os" + "sync" + "sync/atomic" "testing" "time" @@ -77,6 +79,12 @@ func grpcTestSetup(t *testing.T) (tc grpcTestFixtureWithTmpDirCache) { var testMaxCasBlobSizeBytes int64 = 123456789 +// testMaxInflightBytes is the in-flight memory budget passed to ServeGRPC by the +// shared test setup. It defaults to 0 (unlimited) so existing tests are +// unaffected; tests that exercise the limit set it (non-parallel) and restore +// it via defer. +var testMaxInflightBytes int64 = 0 + func grpcTestSetupInternal(t *testing.T, mangleACKeys bool) (tc grpcTestFixtureWithTmpDirCache) { dir, err := os.MkdirTemp("", "bazel-remote-grpc-tests-"+t.Name()) @@ -124,6 +132,7 @@ func grpcTestSetupWithCustomCache(t *testing.T, mangleACKeys bool, validateAC bo mangleACKeys, enableRemoteAssetAPI, testMaxCasBlobSizeBytes, + testMaxInflightBytes, diskCache, accessLogger, errorLogger) if err2 != nil { fmt.Println(err2) @@ -2838,3 +2847,104 @@ func TestInsufficientStorageWhenProxyTriesToStoreAc(t *testing.T) { assertStatusCodeFromError(t, err, codes.ResourceExhausted) } } + +// blockingCountingCache is a disk.Cache whose Get blocks until released, while +// recording the number of concurrent Get calls. It lets a test observe how many +// reads the in-flight-bytes limit allows to proceed at once. +type blockingCountingCache struct { + *StubCache + blobSize int64 + entered chan struct{} // one send per Get entry + release chan struct{} // closed to unblock all Gets + concurrent int32 + maxConcurrent int32 +} + +func (c *blockingCountingCache) Get(ctx context.Context, kind cache.EntryKind, hash string, size int64, offset int64) (io.ReadCloser, int64, error) { + n := atomic.AddInt32(&c.concurrent, 1) + for { + m := atomic.LoadInt32(&c.maxConcurrent) + if n <= m || atomic.CompareAndSwapInt32(&c.maxConcurrent, m, n) { + break + } + } + c.entered <- struct{}{} + <-c.release + atomic.AddInt32(&c.concurrent, -1) + return io.NopCloser(bytes.NewReader(make([]byte, c.blobSize))), c.blobSize, nil +} + +// TestGrpcInflightLimitBoundsConcurrentReads verifies that --max_inflight_bytes +// caps the number of concurrent reads: with a budget of K*blobSize, at most K +// reads may be in flight at once, and the rest block (backpressure) until +// budget frees. Without the limit, all N would proceed and OOM under a burst. +func TestGrpcInflightLimitBoundsConcurrentReads(t *testing.T) { + // Not parallel: mutates the package-global test budget. + const blobSize = int64(1 << 20) // 1 MiB + const budgetK = int64(4) + const numReaders = 16 + + prev := testMaxInflightBytes + testMaxInflightBytes = budgetK * blobSize + defer func() { testMaxInflightBytes = prev }() + + bc := &blockingCountingCache{ + StubCache: &StubCache{}, + blobSize: blobSize, + entered: make(chan struct{}, numReaders), + release: make(chan struct{}), + } + + fixture := grpcTestSetupWithCustomCache(t, false, true, bc) + + _, hash := testutils.RandomDataAndHash(blobSize) + resource := fmt.Sprintf("blobs/%s/%d", hash, blobSize) + + var wg sync.WaitGroup + for i := 0; i < numReaders; i++ { + wg.Add(1) + go func() { + defer wg.Done() + rc, err := fixture.bsClient.Read(ctx, &bytestream.ReadRequest{ResourceName: resource}) + if err != nil { + return + } + for { + if _, e := rc.Recv(); e != nil { + return + } + } + }() + } + + // Exactly budgetK reads should get through the limiter and reach Get. + for i := int64(0); i < budgetK; i++ { + select { + case <-bc.entered: + case <-time.After(5 * time.Second): + t.Fatalf("only %d/%d reads reached Get; expected %d", i, numReaders, budgetK) + } + } + + // A (budgetK+1)th read must NOT reach Get while the first K hold the budget. + select { + case <-bc.entered: + t.Fatalf("more than budget/blobSize=%d concurrent reads reached Get: limit not enforced", budgetK) + case <-time.After(300 * time.Millisecond): + } + + // Release everything; all reads should now drain through in waves. + close(bc.release) + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(15 * time.Second): + t.Fatal("timed out waiting for reads to complete after release") + } + + if got := atomic.LoadInt32(&bc.maxConcurrent); int64(got) > budgetK { + t.Fatalf("max concurrent reads in Get = %d, want <= %d (budget/blobSize)", got, budgetK) + } +} diff --git a/utils/flags/flags.go b/utils/flags/flags.go index 38018c576..79f22d068 100644 --- a/utils/flags/flags.go +++ b/utils/flags/flags.go @@ -206,6 +206,12 @@ func GetCliFlags() []cli.Flag { Usage: "When using proxy backends, sets the number of Goroutines to process parallel uploads to backend.", EnvVars: []string{"BAZEL_REMOTE_NUM_UPLOADERS"}, }, + &cli.Int64Flag{ + Name: "max_inflight_bytes", + Value: 0, + Usage: "Limits the total logical/uncompressed size of gRPC blobs being uploaded or downloaded concurrently, providing memory backpressure under bursts. Requests block until budget is available rather than exhausting memory. 0 (the default) disables the limit. A single blob larger than the limit is still served, serialized against others.", + EnvVars: []string{"BAZEL_REMOTE_MAX_INFLIGHT_BYTES"}, + }, &cli.StringFlag{ Name: "grpc_proxy.url", Value: "",