Review feedback

Author: moroten

pkg/grpc/BUILD.bazel

@@ -32,7 +32,6 @@ go_library(
         "request_metadata_tracing_interceptor.go",
         "routing_stream_handler.go",
         "server.go",
-        "server_transport_stream_context.go",
         "tls_client_certificate_authenticator.go",
     ],
     importpath = "github.com/buildbarn/bb-storage/pkg/grpc",
@@ -70,7 +69,6 @@ go_library(
         "@org_golang_google_grpc//peer",
         "@org_golang_google_grpc//reflection",
         "@org_golang_google_grpc//reflection/grpc_reflection_v1",
-        "@org_golang_google_grpc//reflection/grpc_reflection_v1alpha",
         "@org_golang_google_grpc//status",
         "@org_golang_google_grpc_security_advancedtls//:advancedtls",
         "@org_golang_google_protobuf//encoding/prototext",

pkg/grpc/forwarding_stream_handler.go

@@ -25,7 +25,8 @@ type forwardingStreamHandler struct {
 // incomingStream are forwarded to the backend stream and responses from the
 // backend stream are sent back in the incomingStream.
 func (s *forwardingStreamHandler) HandleStream(srv any, incomingStream grpc.ServerStream) error {
-	method := MustStreamMethodFromContext(incomingStream.Context())
+	// All gRPC invocations has a grpc.ServerTransportStream context.
+	method := grpc.ServerTransportStreamFromContext(incomingStream.Context()).Method()
 	desc := grpc.StreamDesc{
 		// According to grpc.StreamDesc documentation, StreamName and Handler
 		// are only used when registering handlers on a server.

pkg/grpc/forwarding_stream_handler_test.go

@@ -66,14 +66,14 @@ func TestSimpleStreamForwarder(t *testing.T) {
 	forwarder := bb_grpc.NewForwardingStreamHandler(backend)
 	serverTransportStream := mock.NewMockServerTransportStream(ctrl)
 	serverTransportStream.EXPECT().Method().Return("/serviceA/method1").AnyTimes()
-	incomingStream := mock.NewMockServerStream(ctrl)
-	outgoingStream := mock.NewMockClientStream(ctrl)
 
 	t.Run("RequestSuccess", func(t *testing.T) {
 		synctest.Test(t, func(t *testing.T) {
 			var outgoingStreamCtx context.Context
 			outgoingRecvBarrier := make(chan struct{})
 			incomingStreamCtx := grpc.NewContextWithServerTransportStream(context.Background(), serverTransportStream)
+			incomingStream := mock.NewMockServerStream(ctrl)
+			outgoingStream := mock.NewMockClientStream(ctrl)
 
 			newStreamCall := backend.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/serviceA/method1").DoAndReturn(
 				func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
@@ -114,6 +114,8 @@ func TestSimpleStreamForwarder(t *testing.T) {
 		synctest.Test(t, func(t *testing.T) {
 			var outgoingStreamCtx context.Context
 			incomingStreamCtx := grpc.NewContextWithServerTransportStream(context.Background(), serverTransportStream)
+			incomingStream := mock.NewMockServerStream(ctrl)
+			outgoingStream := mock.NewMockClientStream(ctrl)
 
 			newStreamCall := backend.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/serviceA/method1").DoAndReturn(
 				func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
@@ -145,6 +147,8 @@ func TestSimpleStreamForwarder(t *testing.T) {
 		synctest.Test(t, func(t *testing.T) {
 			var outgoingStreamCtx context.Context
 			incomingStreamCtx := grpc.NewContextWithServerTransportStream(context.Background(), serverTransportStream)
+			incomingStream := mock.NewMockServerStream(ctrl)
+			outgoingStream := mock.NewMockClientStream(ctrl)
 
 			newStreamCall := backend.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/serviceA/method1").DoAndReturn(
 				func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
@@ -178,6 +182,8 @@ func TestSimpleStreamForwarder(t *testing.T) {
 			var outgoingStreamCtx context.Context
 			incomingRecvBarrier := make(chan struct{})
 			incomingStreamCtx := grpc.NewContextWithServerTransportStream(context.Background(), serverTransportStream)
+			incomingStream := mock.NewMockServerStream(ctrl)
+			outgoingStream := mock.NewMockClientStream(ctrl)
 
 			newStreamCall := backend.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/serviceA/method1").DoAndReturn(
 				func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
@@ -216,6 +222,8 @@ func TestSimpleStreamForwarder(t *testing.T) {
 			var outgoingStreamCtx context.Context
 			incomingRecvBarrier := make(chan struct{})
 			incomingStreamCtx := grpc.NewContextWithServerTransportStream(context.Background(), serverTransportStream)
+			incomingStream := mock.NewMockServerStream(ctrl)
+			outgoingStream := mock.NewMockClientStream(ctrl)
 
 			newStreamCall := backend.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/serviceA/method1").DoAndReturn(
 				func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
@@ -249,6 +257,8 @@ func TestSimpleStreamForwarder(t *testing.T) {
 		synctest.Test(t, func(t *testing.T) {
 			var outgoingStreamCtx context.Context
 			incomingStreamCtx := grpc.NewContextWithServerTransportStream(context.Background(), serverTransportStream)
+			incomingStream := mock.NewMockServerStream(ctrl)
+			outgoingStream := mock.NewMockClientStream(ctrl)
 
 			newStreamCall := backend.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/serviceA/method1").DoAndReturn(
 				func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
@@ -280,6 +290,7 @@ func TestSimpleStreamForwarder(t *testing.T) {
 	t.Run("NewStreamError", func(t *testing.T) {
 		synctest.Test(t, func(t *testing.T) {
 			incomingStreamCtx := grpc.NewContextWithServerTransportStream(context.Background(), serverTransportStream)
+			incomingStream := mock.NewMockServerStream(ctrl)
 
 			incomingStream.EXPECT().Context().Return(incomingStreamCtx).AnyTimes()
 			backend.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/serviceA/method1").Return(nil, errors.New("no stream"))

pkg/grpc/reflection_relay.go

@@ -9,8 +9,7 @@ import (
 	"github.com/buildbarn/bb-storage/pkg/util"
 	"github.com/jhump/protoreflect/v2/grpcreflect"
 	"github.com/jhump/protoreflect/v2/protoresolve"
-	v1reflectiongrpc "google.golang.org/grpc/reflection/grpc_reflection_v1"
-	v1alphareflectiongrpc "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
+	"google.golang.org/grpc/reflection/grpc_reflection_v1"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
@@ -66,7 +65,6 @@ func registerReflection(backendCtx context.Context, s *grpc.Server, serverRelayC
 		DescriptorResolver: combinedRemoteResolver,
 		ExtensionResolver:  protoresolve.TypesFromDescriptorPool(combinedRemoteResolver),
 	}
-	v1reflectiongrpc.RegisterServerReflectionServer(s, reflection.NewServerV1(serverOptions))
-	v1alphareflectiongrpc.RegisterServerReflectionServer(s, reflection.NewServer(serverOptions))
+	grpc_reflection_v1.RegisterServerReflectionServer(s, reflection.NewServerV1(serverOptions))
 	return nil
 }

pkg/grpc/routing_stream_handler.go

@@ -16,7 +16,8 @@ import (
 // com.google.devtools.build.v1.PublishBuildEvent
 func NewRoutingStreamHandler(routeTable map[string]grpc.StreamHandler) grpc.StreamHandler {
 	return func(srv any, stream grpc.ServerStream) error {
-		serviceMethod := MustStreamMethodFromContext(stream.Context())
+		// All gRPC invocations has a grpc.ServerTransportStream context.
+		serviceMethod := grpc.ServerTransportStreamFromContext(stream.Context()).Method()
 		// Service and method name parsing based on grpc.Server.handleStream().
 		startIdx := 0
 		if serviceMethod != "" && serviceMethod[0] == '/' {