server: bound in-flight CAS writes and per-stream gRPC buffering

main.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"os/signal"
 	"runtime"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
@@ -498,6 +499,17 @@ func startGrpcServer(c *config.Config, grpcServer **grpc.Server,
 	grpcSem *semaphore.Weighted, diskCache disk.Cache) error {
 
 	opts := []grpc.ServerOption{}
+	// A static per-stream window bounds how many bytes a client can push
+	// ahead of the (disk-bound) Put goroutine, so memory per in-flight or
+	// shed write is this instead of the BDP-grown default (up to 16MB).
+	if v := os.Getenv("BAZEL_REMOTE_GRPC_INITIAL_WINDOW_SIZE"); v != "" {
+		n, err := strconv.ParseInt(v, 10, 32)
+		if err != nil || n < 65536 {
+			return fmt.Errorf("invalid BAZEL_REMOTE_GRPC_INITIAL_WINDOW_SIZE %q (bytes, >= 65536)", v)
+		}
+		log.Printf("gRPC initial stream window: %d bytes", n)
+		opts = append(opts, grpc.InitialWindowSize(int32(n)), grpc.InitialConnWindowSize(int32(n)*16))
+	}
 	streamInterceptors := []grpc.StreamServerInterceptor{}
 	unaryInterceptors := []grpc.UnaryServerInterceptor{}
 
@@ -610,6 +622,14 @@ func startGrpcServer(c *config.Config, grpcServer **grpc.Server,
 	log.Println("Starting gRPC server on address", addr)
 
 	var grpcServerOpts []server.GRPCServerOption
+	if v := os.Getenv("BAZEL_REMOTE_MAX_INFLIGHT_CAS_WRITES"); v != "" {
+		n, err := strconv.Atoi(v)
+		if err != nil || n < 0 {
+			return fmt.Errorf("invalid BAZEL_REMOTE_MAX_INFLIGHT_CAS_WRITES %q", v)
+		}
+		log.Printf("gRPC ByteStream max in-flight CAS writes: %d", n)
+		grpcServerOpts = append(grpcServerOpts, server.WithMaxInflightCASWrites(n))
+	}
 	if c.ReadChunkSizeBytes > 0 {
 		log.Printf("gRPC ByteStream read chunk size: %d bytes", c.ReadChunkSizeBytes)
 		grpcServerOpts = append(grpcServerOpts, server.WithReadChunkSizeBytes(c.ReadChunkSizeBytes))

server/grpc.go

@@ -50,7 +50,13 @@ type grpcServer struct {
 	readChunkSizeBytes     int64
 	runtimeMetrics         RuntimeMetrics
 	readLimiter            *readLimiter
-	sourceBuffers          *sourceBufferPool
+	// casWriteSlots bounds in-flight CAS ByteStream Writes past the Contains
+	// check. Each one holds a 1MB chunk buffer, the zstd output and pipe/gRPC
+	// buffers before anything reaches disk, so an unbounded burst from one
+	// tenant (8k concurrent writes observed) takes the heap past GOMEMLIMIT
+	// and stalls the process. nil = unlimited.
+	casWriteSlots chan struct{}
+	sourceBuffers *sourceBufferPool
 	// writePayloadConsumed, when set, is invoked once per fully-consumed
 	// ByteStream WriteRequest; see WithWritePayloadConsumed.
 	writePayloadConsumed func(*bytestream.WriteRequest)

server/grpc_bytestream.go

@@ -428,6 +428,14 @@ func (s *grpcServer) parseWriteResource(r string) (string, int64, casblob.Compre
 }
 
 var errWriteOffset error = errors.New("bytestream writes from non-zero offsets are unsupported")
+
+// FAILED_PRECONDITION is deliberately not a connection-class code: Bazel's
+// retrier treats UNAVAILABLE as transient and walks the full --remote_retries
+// backoff ladder per blob, which under sustained shedding becomes a
+// minutes-long upload tail after the build summary. A permanent code makes
+// the client log the upload failure once and move on; the action result is
+// never published, so nothing references the unstored blob.
+var errCASWriteShed = status.Error(codes.FailedPrecondition, "too many in-flight CAS writes")
 var errDecoderPoolFail error = errors.New("failed to get DecoderWrapper from pool")
 
 func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
@@ -549,7 +557,18 @@ func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 					rc = dec.IOReadCloser()
 				}
 
+				// Acquire last so every earlier error path returns
+				// without holding a slot; the Put goroutine owns the
+				// release.
+				if !s.tryAcquireCASWriteSlot() {
+					_ = rc.Close()
+					s.accessLogger.Printf("GRPC BYTESTREAM WRITE SHED: %s", resourceName)
+					recvResult <- errCASWriteShed
+					return
+				}
+
 				go func() {
+					defer s.releaseCASWriteSlot()
 					defer func() { _ = rc.Close() }()
 					err := s.cache.Put(srv.Context(), cache.CAS, hash, size, rc)
 					putResult <- err
@@ -625,8 +644,10 @@ func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 			}
 
 			_ = pw.CloseWithError(err)
-			s.accessLogger.Printf("GRPC BYTESTREAM WRITE FAILED: %s %s",
-				resourceName, err.Error())
+			if err != errCASWriteShed {
+				s.accessLogger.Printf("GRPC BYTESTREAM WRITE FAILED: %s %s",
+					resourceName, err.Error())
+			}
 			return err
 		}
 

server/grpc_test.go

@@ -3105,3 +3105,75 @@ func TestInsufficientStorageWhenProxyTriesToStoreAc(t *testing.T) {
 		assertStatusCodeFromError(t, err, codes.ResourceExhausted)
 	}
 }
+
+func TestGrpcByteStreamWriteShedWhenInflightCapReached(t *testing.T) {
+	fixture := grpcTestSetupInternal(t, false, WithMaxInflightCASWrites(1))
+	defer os.RemoveAll(fixture.tempdir)
+
+	writeResource := func(blob []byte, hash string) string {
+		return fmt.Sprintf("uploads/%s/blobs/%s/%d", uuid.New().String(), hash, len(blob))
+	}
+	fullWrite := func(blob []byte, hash string) error {
+		bswc, err := fixture.bsClient.Write(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = bswc.Send(&bytestream.WriteRequest{
+			ResourceName: writeResource(blob, hash),
+			FinishWrite:  true,
+			Data:         blob,
+		})
+		if err != nil && err != io.EOF {
+			t.Fatal(err)
+		}
+		_, err = bswc.CloseAndRecv()
+		return err
+	}
+
+	heldBlob, heldHash := testutils.RandomDataAndHash(64)
+	heldResource := writeResource(heldBlob, heldHash)
+	held, err := fixture.bsClient.Write(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = held.Send(&bytestream.WriteRequest{
+		ResourceName: heldResource,
+		Data:         heldBlob[:32],
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Let the server consume the first chunk and take the only slot.
+	time.Sleep(200 * time.Millisecond)
+
+	otherBlob, otherHash := testutils.RandomDataAndHash(64)
+	err = fullWrite(otherBlob, otherHash)
+	if status.Code(err) != codes.FailedPrecondition {
+		t.Fatalf("expected codes.FailedPrecondition while a write holds the slot, got %v", err)
+	}
+
+	if err := fullWrite(otherBlob, otherHash); status.Code(err) != codes.FailedPrecondition {
+		t.Fatalf("expected second shed, got %v", err)
+	}
+
+	err = held.Send(&bytestream.WriteRequest{
+		ResourceName: heldResource,
+		FinishWrite:  true,
+		Data:         heldBlob[32:],
+		WriteOffset:  32,
+	})
+	if err != nil && err != io.EOF {
+		t.Fatal(err)
+	}
+	if _, err := held.CloseAndRecv(); err != nil {
+		t.Fatalf("held write should complete: %v", err)
+	}
+
+	if err := fullWrite(otherBlob, otherHash); err != nil {
+		t.Fatalf("write after slot release should succeed: %v", err)
+	}
+	// Now present: the Contains early return answers before the gate.
+	if err := fullWrite(otherBlob, otherHash); err != nil {
+		t.Fatalf("duplicate write of present blob should succeed: %v", err)
+	}
+}

server/runtime.go

@@ -3,6 +3,8 @@ package server
 import (
 	"context"
 	"fmt"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"sync"
 	"time"
 
@@ -81,6 +83,51 @@ func WithReadChunkSizeBytes(chunkSizeBytes int64) GRPCServerOption {
 	}
 }
 
+// WithMaxInflightCASWrites caps concurrent CAS ByteStream Writes that reach
+// the Put path; excess writes are shed with FAILED_PRECONDITION. 0 = unlimited.
+func WithMaxInflightCASWrites(n int) GRPCServerOption {
+	return func(s *grpcServer) error {
+		if n < 0 {
+			return fmt.Errorf("max inflight CAS writes must not be negative: %d", n)
+		}
+		if n > 0 {
+			s.casWriteSlots = make(chan struct{}, n)
+		}
+		return nil
+	}
+}
+
+var casWriteSlotsInflight = promauto.NewGauge(prometheus.GaugeOpts{
+	Name: "bazel_remote_cas_write_slots_inflight",
+	Help: "New CAS ByteStream writes currently admitted through the in-flight write gate.",
+})
+
+var casWriteShedTotal = promauto.NewCounter(prometheus.CounterOpts{
+	Name: "bazel_remote_cas_write_shed_total",
+	Help: "New CAS ByteStream writes rejected because the in-flight write gate was full.",
+})
+
+func (s *grpcServer) tryAcquireCASWriteSlot() bool {
+	if s.casWriteSlots == nil {
+		return true
+	}
+	select {
+	case s.casWriteSlots <- struct{}{}:
+		casWriteSlotsInflight.Inc()
+		return true
+	default:
+		casWriteShedTotal.Inc()
+		return false
+	}
+}
+
+func (s *grpcServer) releaseCASWriteSlot() {
+	if s.casWriteSlots != nil {
+		<-s.casWriteSlots
+		casWriteSlotsInflight.Dec()
+	}
+}
+
 // WithWritePayloadConsumed registers a callback the ByteStream Write handler
 // invokes exactly once per received message, after that message's Data has
 // been fully consumed (the pipe write to the cache reader returned). It is