Move RPC byte-counting logic to separate gRPC stream wrappers

Author: iain-macdonald

enterprise/server/byte_stream_server_proxy/BUILD

@@ -15,6 +15,7 @@ go_library(
         "//server/real_environment",
         "//server/remote_cache/digest",
         "//server/util/authutil",
+        "//server/util/grpc_stream",
         "//server/util/log",
         "//server/util/status",
         "//server/util/tracing",

enterprise/server/byte_stream_server_proxy/byte_stream_server_proxy.go

@@ -12,6 +12,7 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/real_environment"
 	"github.com/buildbuddy-io/buildbuddy/server/remote_cache/digest"
 	"github.com/buildbuddy-io/buildbuddy/server/util/authutil"
+	"github.com/buildbuddy-io/buildbuddy/server/util/grpc_stream"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/status"
 	"github.com/buildbuddy-io/buildbuddy/server/util/tracing"
@@ -75,7 +76,7 @@ func (s *ByteStreamServerProxy) Read(req *bspb.ReadRequest, stream bspb.ByteStre
 		return err
 	}
 
-	localReadStream, err := s.local.Read(ctx, req)
+	localGRPCReadStream, err := s.local.Read(ctx, req)
 	if err != nil {
 		if skipRemote {
 			recordReadMetrics(metrics.MissStatusLabel, requestTypeLabel, err, 0)
@@ -96,12 +97,9 @@ func (s *ByteStreamServerProxy) Read(req *bspb.ReadRequest, stream bspb.ByteStre
 	}
 
 	responseSent := false
-	bytesRead := 0
+	localReadStream := grpc_stream.NewByteCountingServerStream(localGRPCReadStream)
 	for {
 		rsp, err := localReadStream.Recv()
-		if rsp != nil {
-			bytesRead += len(rsp.Data)
-		}
 		if err == io.EOF {
 			break
 		}
@@ -112,10 +110,10 @@ func (s *ByteStreamServerProxy) Read(req *bspb.ReadRequest, stream bspb.ByteStre
 			// the remote cache, but keep it simple for now.
 			if responseSent {
 				log.CtxInfof(ctx, "error midstream of local read: %s", err)
-				recordReadMetrics(metrics.HitStatusLabel, requestTypeLabel, err, bytesRead)
+				recordReadMetrics(metrics.HitStatusLabel, requestTypeLabel, err, int(localReadStream.GetByteCount()))
 				return err
 			} else if skipRemote {
-				recordReadMetrics(metrics.MissStatusLabel, requestTypeLabel, err, bytesRead)
+				recordReadMetrics(metrics.MissStatusLabel, requestTypeLabel, err, int(localReadStream.GetByteCount()))
 				return err
 			} else {
 				// Fall back to reading remotely if the local read fails.
@@ -126,12 +124,12 @@ func (s *ByteStreamServerProxy) Read(req *bspb.ReadRequest, stream bspb.ByteStre
 		}
 
 		if err := stream.Send(rsp); err != nil {
-			recordReadMetrics(metrics.HitStatusLabel, requestTypeLabel, err, bytesRead)
+			recordReadMetrics(metrics.HitStatusLabel, requestTypeLabel, err, int(localReadStream.GetByteCount()))
 			return err
 		}
 		responseSent = true
 	}
-	recordReadMetrics(metrics.HitStatusLabel, requestTypeLabel, nil, bytesRead)
+	recordReadMetrics(metrics.HitStatusLabel, requestTypeLabel, nil, int(localReadStream.GetByteCount()))
 	return nil
 }
 
@@ -213,11 +211,12 @@ func (s *ByteStreamServerProxy) readRemote(req *bspb.ReadRequest, stream bspb.By
 	ctx, spn := tracing.StartSpan(stream.Context())
 	defer spn.End()
 
-	remoteReadStream, err := s.remote.Read(ctx, req)
+	remoteGRPCReadStream, err := s.remote.Read(ctx, req)
 	if err != nil {
 		log.CtxInfof(ctx, "error reading from remote: %s", err)
 		return 0, err
 	}
+	remoteReadStream := grpc_stream.NewByteCountingServerStream(remoteGRPCReadStream)
 
 	var localWriteStream localWriter = &discardingLocalWriter{}
 	if req.ReadOffset == 0 && !authutil.EncryptionEnabled(ctx, s.authenticator) {
@@ -235,12 +234,8 @@ func (s *ByteStreamServerProxy) readRemote(req *bspb.ReadRequest, stream bspb.By
 		}
 	}
 
-	bytesRead := 0
 	for {
 		rsp, err := remoteReadStream.Recv()
-		if rsp != nil {
-			bytesRead += len(rsp.Data)
-		}
 		if err != nil {
 			if err == io.EOF {
 				if err := localWriteStream.commit(); err != nil {
@@ -249,18 +244,18 @@ func (s *ByteStreamServerProxy) readRemote(req *bspb.ReadRequest, stream bspb.By
 				break
 			}
 			log.CtxInfof(ctx, "error streaming from remote for read through: %s", err)
-			return bytesRead, err
+			return int(remoteReadStream.GetByteCount()), err
 		}
 
 		if err := localWriteStream.send(rsp.Data); err != nil {
 			log.CtxInfof(ctx, "Error writing locally for read through: %s", err)
 			localWriteStream = &discardingLocalWriter{}
 		}
 		if err = stream.Send(rsp); err != nil {
-			return bytesRead, err
+			return int(remoteReadStream.GetByteCount()), err
 		}
 	}
-	return bytesRead, nil
+	return int(remoteReadStream.GetByteCount()), nil
 }
 
 func (s *ByteStreamServerProxy) Write(stream bspb.ByteStream_WriteServer) error {

enterprise/server/content_addressable_storage_server_proxy/BUILD

@@ -16,6 +16,7 @@ go_library(
         "//server/real_environment",
         "//server/remote_cache/digest",
         "//server/util/authutil",
+        "//server/util/grpc_stream",
         "//server/util/log",
         "//server/util/proto",
         "//server/util/rpcutil",

enterprise/server/content_addressable_storage_server_proxy/content_addressable_storage_server_proxy.go

@@ -14,6 +14,7 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/real_environment"
 	"github.com/buildbuddy-io/buildbuddy/server/remote_cache/digest"
 	"github.com/buildbuddy-io/buildbuddy/server/util/authutil"
+	"github.com/buildbuddy-io/buildbuddy/server/util/grpc_stream"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/proto"
 	"github.com/buildbuddy-io/buildbuddy/server/util/rpcutil"
@@ -292,7 +293,6 @@ func (s *CASServerProxy) batchReadBlobsRemote(ctx context.Context, readReq *repb
 	return readResp, nil
 }
 
-// TODO(iain): record per-byte metrics here as well as above.
 func (s *CASServerProxy) GetTree(req *repb.GetTreeRequest, stream repb.ContentAddressableStorage_GetTreeServer) error {
 	if proxy_util.SkipRemote(stream.Context()) {
 		return status.UnimplementedError("Skip remote not implemented")
@@ -306,18 +306,18 @@ func (s *CASServerProxy) GetTree(req *repb.GetTreeRequest, stream repb.ContentAd
 
 func (s *CASServerProxy) getTreeWithoutCaching(req *repb.GetTreeRequest, stream repb.ContentAddressableStorage_GetTreeServer) error {
 	digests := 0
-	bytes := 0
+	remoteGRPCStream, err := s.remote.GetTree(stream.Context(), req)
+	if err != nil {
+		return err
+	}
+	remoteStream := grpc_stream.NewByteCountingServerStream(remoteGRPCStream)
 	defer func() {
 		recordMetrics(
 			"GetTree",
 			metrics.MissStatusLabel,
 			map[string]int{metrics.MissStatusLabel: digests},
-			map[string]int{metrics.MissStatusLabel: bytes})
+			map[string]int{metrics.MissStatusLabel: int(remoteStream.GetByteCount())})
 	}()
-	remoteStream, err := s.remote.GetTree(stream.Context(), req)
-	if err != nil {
-		return err
-	}
 	for {
 		rsp, err := remoteStream.Recv()
 		if err == io.EOF {
@@ -330,7 +330,6 @@ func (s *CASServerProxy) getTreeWithoutCaching(req *repb.GetTreeRequest, stream
 			digests += len(dir.GetFiles())
 			digests += len(dir.GetDirectories())
 		}
-		bytes += proto.Size(rsp)
 		if err = stream.Send(rsp); err != nil {
 			return err
 		}

server/util/grpc_stream/BUILD

@@ -0,0 +1,30 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "grpc_stream",
+    srcs = ["grpc_stream.go"],
+    importpath = "github.com/buildbuddy-io/buildbuddy/server/util/grpc_stream",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//server/util/proto",
+        "@org_golang_google_grpc//:grpc",
+        "@org_golang_google_grpc//metadata",
+        "@org_golang_google_protobuf//reflect/protoreflect",
+    ],
+)
+
+go_test(
+    name = "grpc_stream_test",
+    srcs = ["grpc_stream_test.go"],
+    deps = [
+        ":grpc_stream",
+        "//server/environment",
+        "//server/testutil/testenv",
+        "//server/testutil/testport",
+        "//server/util/grpc_client",
+        "//server/util/grpc_server",
+        "//server/util/proto",
+        "@com_github_stretchr_testify//require",
+        "@org_golang_google_genproto_googleapis_bytestream//:bytestream",
+    ],
+)

server/util/grpc_stream/grpc_stream.go

@@ -0,0 +1,133 @@
+package grpc_stream
+
+import (
+	"context"
+	"sync/atomic"
+
+	"github.com/buildbuddy-io/buildbuddy/server/util/proto"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+// gRPC has three types of streaming RPCs, as described here:
+// https://grpc.io/docs/what-is-grpc/core-concepts/#server-streaming-rpc
+// - Server streaming (client sends one message, receives many)
+// - Client treaming (client sends many messages, receives one)
+// - Bidirectional streaming (client and server send many message)
+//
+// This generic interface represents the server streaming RPC.
+type ServerStream[T protoreflect.ProtoMessage] interface {
+	Recv() (T, error)
+	grpc.ClientStream
+}
+
+// A wrapper around a ServerStream that counts the number of bytes returned by
+// the server over the lifetime of the stream.
+type ByteCountingServerStream[T protoreflect.ProtoMessage] struct {
+	bytes *atomic.Int64
+	proxy ServerStream[T]
+}
+
+func NewByteCountingServerStream[T protoreflect.ProtoMessage](proxy ServerStream[T]) ByteCountingServerStream[T] {
+	bytes := &atomic.Int64{}
+	return ByteCountingServerStream[T]{
+		bytes: bytes,
+		proxy: proxy,
+	}
+}
+
+// Returns the number of bytes received on the stream so far.
+func (s *ByteCountingServerStream[T]) GetByteCount() int64 {
+	return s.bytes.Load()
+}
+
+func (s *ByteCountingServerStream[T]) Recv() (T, error) {
+	message, err := s.proxy.Recv()
+	s.bytes.Add(int64(proto.Size(message)))
+	return message, err
+}
+
+func (s *ByteCountingServerStream[T]) Header() (metadata.MD, error) {
+	return s.proxy.Header()
+}
+
+func (s *ByteCountingServerStream[T]) Trailer() metadata.MD {
+	return s.proxy.Trailer()
+}
+
+func (s *ByteCountingServerStream[T]) CloseSend() error {
+	return s.proxy.CloseSend()
+}
+
+func (s *ByteCountingServerStream[T]) Context() context.Context {
+	return s.proxy.Context()
+}
+
+func (s *ByteCountingServerStream[T]) SendMsg(message any) error {
+	return s.proxy.SendMsg(message)
+}
+
+func (s *ByteCountingServerStream[T]) RecvMsg(message any) error {
+	return s.proxy.RecvMsg(message)
+}
+
+// This generic interface represents the client streaming RPC.
+type ClientStream[S, T protoreflect.ProtoMessage] interface {
+	Send(S) error
+	CloseAndRecv() (T, error)
+	grpc.ClientStream
+}
+
+// A wrapper around a ClientStream that counts the number of bytes sent by the
+// client over the lifetime of the stream.
+type ByteCountingClientStream[S, T protoreflect.ProtoMessage] struct {
+	bytes *atomic.Int64
+	proxy ClientStream[S, T]
+}
+
+func NewByteCountingClientStream[S, T protoreflect.ProtoMessage](proxy ClientStream[S, T]) ByteCountingClientStream[S, T] {
+	bytes := &atomic.Int64{}
+	return ByteCountingClientStream[S, T]{
+		bytes: bytes,
+		proxy: proxy,
+	}
+}
+
+// Returns the number of bytes sent to the stream so far.
+func (s *ByteCountingClientStream[S, T]) GetByteCount() int64 {
+	return s.bytes.Load()
+}
+
+func (s *ByteCountingClientStream[S, T]) Send(message S) error {
+	s.bytes.Add(int64(proto.Size(message)))
+	return s.proxy.Send(message)
+}
+
+func (s *ByteCountingClientStream[S, T]) CloseAndRecv() (T, error) {
+	return s.proxy.CloseAndRecv()
+}
+
+func (s *ByteCountingClientStream[S, T]) Header() (metadata.MD, error) {
+	return s.proxy.Header()
+}
+
+func (s *ByteCountingClientStream[S, T]) Trailer() metadata.MD {
+	return s.proxy.Trailer()
+}
+
+func (s *ByteCountingClientStream[S, T]) CloseSend() error {
+	return s.proxy.CloseSend()
+}
+
+func (s *ByteCountingClientStream[S, T]) Context() context.Context {
+	return s.proxy.Context()
+}
+
+func (s *ByteCountingClientStream[S, T]) SendMsg(message any) error {
+	return s.proxy.SendMsg(message)
+}
+
+func (s *ByteCountingClientStream[S, T]) RecvMsg(message any) error {
+	return s.proxy.RecvMsg(message)
+}