Review feedback 2

Author: moroten

pkg/grpc/reflection_relay.go

@@ -4,9 +4,7 @@ import (
 	"context"
 	"maps"
 
-	"github.com/buildbarn/bb-storage/pkg/program"
 	grpcpb "github.com/buildbarn/bb-storage/pkg/proto/configuration/grpc"
-	"github.com/buildbarn/bb-storage/pkg/util"
 	"github.com/jhump/protoreflect/v2/grpcreflect"
 	"github.com/jhump/protoreflect/v2/protoresolve"
 	"google.golang.org/grpc/reflection/grpc_reflection_v1"
@@ -25,20 +23,26 @@ var _ reflection.ServiceInfoProvider = (*combinedServiceInfoProvider)(nil)
 // GetServiceInfo returns the currently available services, which might have
 // changed since the creation of this reflection server.
 func (p *combinedServiceInfoProvider) GetServiceInfo() map[string]grpc.ServiceInfo {
-	services := make(map[string]grpc.ServiceInfo)
+	serverServiceInfo := p.server.GetServiceInfo()
+	services := make(map[string]grpc.ServiceInfo, len(p.extraServices)+len(serverServiceInfo))
 	maps.Copy(services, p.extraServices)
-	maps.Copy(services, p.server.GetServiceInfo())
+	maps.Copy(services, serverServiceInfo)
 	return services
 }
 
-// registerReflection registers the google.golang.org/grpc/reflection/ service
-// on a grpc.Server and calls remote backends in case for relayed services. The
-// connections to the backend will run with the backendCtx.
-func registerReflection(backendCtx context.Context, s *grpc.Server, serverRelayConfiguration []*grpcpb.ServerRelayConfiguration, group program.Group, grpcClientFactory ClientFactory) error {
+type serverRelayConfigWithGrpcClient struct {
+	config     *grpcpb.ServerRelayConfiguration
+	grpcClient grpc.ClientConnInterface
+}
+
+// registerReflectionServer registers the google.golang.org/grpc/reflection/
+// service on a grpc.Server and calls remote backends in case for relayed
+// services. The connections to the backend will run with the backendCtx.
+func registerReflectionServer(backendCtx context.Context, s *grpc.Server, serverRelayConfigurations []serverRelayConfigWithGrpcClient) error {
 	// Accumulate all the service names.
 	relayServices := make(map[string]grpc.ServiceInfo)
-	for _, relay := range serverRelayConfiguration {
-		for _, service := range relay.GetServices() {
+	for _, relay := range serverRelayConfigurations {
+		for _, service := range relay.config.GetServices() {
 			// According to ServiceInfoProvider docs for ServerOptions.Services,
 			// the reflection service is only interested in the service names.
 			relayServices[service] = grpc.ServiceInfo{}
@@ -47,12 +51,8 @@ func registerReflection(backendCtx context.Context, s *grpc.Server, serverRelayC
 
 	// Make a combined descriptor and extension resolver.
 	reflectionBackends := []protoresolve.Resolver{}
-	for relayIdx, relay := range serverRelayConfiguration {
-		grpcClient, err := grpcClientFactory.NewClientFromConfiguration(relay.Endpoint, group)
-		if err != nil {
-			return util.StatusWrapf(err, "Failed to create relay RPC client %d", relayIdx+1)
-		}
-		resolver := grpcreflect.NewClientAuto(backendCtx, grpcClient).AsResolver()
+	for _, relay := range serverRelayConfigurations {
+		resolver := grpcreflect.NewClientAuto(backendCtx, relay.grpcClient).AsResolver()
 		reflectionBackends = append(reflectionBackends, resolver)
 	}
 	combinedRemoteResolver := protoresolve.Combine(reflectionBackends...)

pkg/grpc/routing_stream_handler.go

@@ -17,17 +17,14 @@ import (
 func NewRoutingStreamHandler(routeTable map[string]grpc.StreamHandler) grpc.StreamHandler {
 	return func(srv any, stream grpc.ServerStream) error {
 		// All gRPC invocations has a grpc.ServerTransportStream context.
-		serviceMethod := grpc.ServerTransportStreamFromContext(stream.Context()).Method()
+		orgServiceMethod := grpc.ServerTransportStreamFromContext(stream.Context()).Method()
 		// Service and method name parsing based on grpc.Server.handleStream().
-		startIdx := 0
-		if serviceMethod != "" && serviceMethod[0] == '/' {
-			startIdx = 1
-		}
+		serviceMethod := strings.TrimPrefix(orgServiceMethod, "/")
 		endIdx := strings.LastIndex(serviceMethod, "/")
-		if endIdx <= startIdx {
-			return status.Errorf(codes.InvalidArgument, "Malformed method name %v", serviceMethod)
+		if endIdx == -1 {
+			return status.Errorf(codes.InvalidArgument, "Malformed method name %v", orgServiceMethod)
 		}
-		service := serviceMethod[startIdx:endIdx]
+		service := serviceMethod[:endIdx]
 
 		if streamHandler, ok := routeTable[service]; ok {
 			return streamHandler(srv, stream)

pkg/grpc/server.go

@@ -146,8 +146,19 @@ func NewServersFromConfigurationAndServe(configurations []*configuration.ServerC
 			}))
 		}
 
+		relayConfigWithGrpcClients := make([]serverRelayConfigWithGrpcClient, len(configuration.Relays))
+		for relayIdx, relay := range configuration.Relays {
+			grpcClient, err := grpcClientFactory.NewClientFromConfiguration(relay.Endpoint, group)
+			if err != nil {
+				return util.StatusWrapf(err, "Failed to create relay RPC client %d", relayIdx+1)
+			}
+			relayConfigWithGrpcClients[relayIdx] = serverRelayConfigWithGrpcClient{
+				config:     relay,
+				grpcClient: grpcClient,
+			}
+		}
 		if len(configuration.Relays) != 0 {
-			handler, err := newRoutingStreamHandlerFromConfiguration(configuration.Relays, grpcClientFactory, group)
+			handler, err := newRoutingStreamHandlerFromConfiguration(relayConfigWithGrpcClients)
 			if err != nil {
 				return err
 			}
@@ -169,7 +180,7 @@ func NewServersFromConfigurationAndServe(configurations []*configuration.ServerC
 
 		// Enable default services.
 		grpc_prometheus.Register(s)
-		if err := registerReflection(context.Background(), s, configuration.Relays, group, grpcClientFactory); err != nil {
+		if err := registerReflectionServer(context.Background(), s, relayConfigWithGrpcClients); err != nil {
 			return util.StatusWrap(err, "Failed to create reflection service")
 		}
 		h := health.NewServer()
@@ -218,15 +229,11 @@ func NewServersFromConfigurationAndServe(configurations []*configuration.ServerC
 	return nil
 }
 
-func newRoutingStreamHandlerFromConfiguration(serverRelayConfiguration []*grpcpb.ServerRelayConfiguration, grpcClientFactory ClientFactory, group program.Group) (grpc.StreamHandler, error) {
+func newRoutingStreamHandlerFromConfiguration(serverRelayConfigurations []serverRelayConfigWithGrpcClient) (grpc.StreamHandler, error) {
 	routeTable := make(map[string]grpc.StreamHandler)
-	for i, relay := range serverRelayConfiguration {
-		grpcClient, err := grpcClientFactory.NewClientFromConfiguration(relay.GetEndpoint(), group)
-		if err != nil {
-			return nil, util.StatusWrapf(err, "Failed to create gRPC relay RPC client at index %d", i)
-		}
-		handler := NewForwardingStreamHandler(grpcClient)
-		for _, service := range relay.GetServices() {
+	for _, relay := range serverRelayConfigurations {
+		handler := NewForwardingStreamHandler(relay.grpcClient)
+		for _, service := range relay.config.GetServices() {
 			if _, ok := routeTable[service]; ok {
 				return nil, status.Errorf(codes.InvalidArgument, "Duplicated gRPC relay for %v", service)
 			}