refactor protoregistry interface

3 files changed, 18 insertions, 18 deletions

diff --git a/internal/praefect/protoregistry/protoregistry.go b/internal/praefect/protoregistry/protoregistry.go
index 83f8a16bd..eed1b38d1 100644
--- a/internal/praefect/protoregistry/protoregistry.go
+++ b/internal/praefect/protoregistry/protoregistry.go
@@ -66,33 +66,35 @@ func (mi MethodInfo) TargetRepo(msg proto.Message) (*gitalypb.Repository, error)
 // Registry contains info about RPC methods
 type Registry struct {
 	sync.RWMutex
-	protos map[string]map[string]MethodInfo
+	protos map[string]MethodInfo
 }
 
 // New creates a new ProtoRegistry
 func New() *Registry {
 	return &Registry{
-		protos: make(map[string]map[string]MethodInfo),
+		protos: make(map[string]MethodInfo),
 	}
 }
 
-// RegisterFiles takes one or more descriptor.FileDescriptorProto and populates the registry with its info
+// RegisterFiles takes one or more descriptor.FileDescriptorProto and populates
+// the registry with its info
 func (pr *Registry) RegisterFiles(protos ...*descriptor.FileDescriptorProto) error {
 	pr.Lock()
 	defer pr.Unlock()
 
 	for _, p := range protos {
-		for _, serviceDescriptorProto := range p.GetService() {
-			for _, methodDescriptorProto := range serviceDescriptorProto.GetMethod() {
-				mi, err := parseMethodInfo(methodDescriptorProto)
+		for _, svc := range p.GetService() {
+			for _, method := range svc.GetMethod() {
+				mi, err := parseMethodInfo(method)
 				if err != nil {
 					return err
 				}
 
-				if _, ok := pr.protos[serviceDescriptorProto.GetName()]; !ok {
-					pr.protos[serviceDescriptorProto.GetName()] = make(map[string]MethodInfo)
-				}
-				pr.protos[serviceDescriptorProto.GetName()][methodDescriptorProto.GetName()] = mi
+				fullMethodName := fmt.Sprintf(
+					"/%s.%s/%s",
+					p.GetPackage(), svc.GetName(), method.GetName(),
+				)
+				pr.protos[fullMethodName] = mi
 			}
 		}
 	}
@@ -189,16 +191,13 @@ func parseOID(rawFieldOID string) ([]int, error) {
 }
 
 // LookupMethod looks up an MethodInfo by service and method name
-func (pr *Registry) LookupMethod(service, method string) (MethodInfo, error) {
+func (pr *Registry) LookupMethod(fullMethodName string) (MethodInfo, error) {
 	pr.RLock()
 	defer pr.RUnlock()
 
-	if _, ok := pr.protos[service]; !ok {
-		return MethodInfo{}, fmt.Errorf("service not found: %v", service)
-	}
-	methodInfo, ok := pr.protos[service][method]
+	methodInfo, ok := pr.protos[fullMethodName]
 	if !ok {
-		return MethodInfo{}, fmt.Errorf("method not found: %v", method)
+		return MethodInfo{}, fmt.Errorf("Full method name not found: %v", fullMethodName)
 	}
 	return methodInfo, nil
 }
diff --git a/internal/praefect/protoregistry/protoregistry_test.go b/internal/praefect/protoregistry/protoregistry_test.go
index 56cffe255..3d896d6d7 100644
--- a/internal/praefect/protoregistry/protoregistry_test.go
+++ b/internal/praefect/protoregistry/protoregistry_test.go
@@ -1,6 +1,7 @@
 package protoregistry_test
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -187,7 +188,7 @@ func TestPopulatesProtoRegistry(t *testing.T) {
 
 	for serviceName, methods := range expectedResults {
 		for methodName, opType := range methods {
-			methodInfo, err := r.LookupMethod(serviceName, methodName)
+			methodInfo, err := r.LookupMethod(fmt.Sprintf("/gitaly.%s/%s", serviceName, methodName))
 			require.NoError(t, err)
 			assert.Equalf(t, opType, methodInfo.Operation, "expect %s:%s to have the correct op type", serviceName, methodName)
 		}
diff --git a/internal/praefect/protoregistry/targetrepo_test.go b/internal/praefect/protoregistry/targetrepo_test.go
index fc016807d..f2c1f394e 100644
--- a/internal/praefect/protoregistry/targetrepo_test.go
+++ b/internal/praefect/protoregistry/targetrepo_test.go
@@ -63,7 +63,7 @@ func TestProtoRegistryTargetRepo(t *testing.T) {
 	for _, tc := range testcases {
 		desc := fmt.Sprintf("%s:%s %s", tc.svc, tc.method, tc.desc)
 		t.Run(desc, func(t *testing.T) {
-			info, err := r.LookupMethod(tc.svc, tc.method)
+			info, err := r.LookupMethod(fmt.Sprintf("/gitaly.%s/%s", tc.svc, tc.method))
 			require.NoError(t, err)
 
 			actualTarget, actualErr := info.TargetRepo(tc.pbMsg)