diff --git a/main.go b/main.go index 3c6b8f6..5940ffa 100644 --- a/main.go +++ b/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)) diff --git a/server/grpc.go b/server/grpc.go index 1d427e8..bf036e6 100644 --- a/server/grpc.go +++ b/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) diff --git a/server/grpc_bytestream.go b/server/grpc_bytestream.go index e9e8ad4..04e42a4 100644 --- a/server/grpc_bytestream.go +++ b/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 } diff --git a/server/grpc_test.go b/server/grpc_test.go index d770cd6..68385cd 100644 --- a/server/grpc_test.go +++ b/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) + } +} diff --git a/server/runtime.go b/server/runtime.go index fa85ac1..607fc5e 100644 --- a/server/runtime.go +++ b/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