Cleaning up praefect method signatures

9 files changed, 92 insertions, 119 deletions

diff --git a/cmd/praefect/main.go b/cmd/praefect/main.go
index 81ea610e7..ae81f1680 100644
--- a/cmd/praefect/main.go
+++ b/cmd/praefect/main.go
@@ -152,10 +152,15 @@ func run(cfgs []starter.Config, conf config.Config) error {
 		return err
 	}
 
+	registry := protoregistry.New()
+	if err = registry.RegisterFiles(protoregistry.GitalyProtoFileDescriptors...); err != nil {
+		return err
+	}
+
 	var (
 		// top level server dependencies
 		ds          = datastore.NewInMemory(conf)
-		coordinator = praefect.NewCoordinator(logger, ds, nodeManager, conf, protoregistry.GitalyProtoFileDescriptors...)
+		coordinator = praefect.NewCoordinator(logger, ds, nodeManager, conf, registry)
 		repl        = praefect.NewReplMgr(
 			conf.VirtualStorages[0].Name,
 			logger,
@@ -163,7 +168,8 @@ func run(cfgs []starter.Config, conf config.Config) error {
 			nodeManager,
 			praefect.WithLatencyMetric(latencyMetric),
 			praefect.WithQueueMetric(queueMetric))
-		srv          = praefect.NewServer(coordinator, repl, nil, logger, nodeManager, conf)
+		srv = praefect.NewServer(coordinator.StreamDirector, logger, registry, conf)
+
 		serverErrors = make(chan error, 1)
 	)
 
@@ -177,7 +183,7 @@ func run(cfgs []starter.Config, conf config.Config) error {
 		return fmt.Errorf("unable to create a bootstrap: %v", err)
 	}
 
-	srv.RegisterServices()
+	srv.RegisterServices(nodeManager, conf)
 
 	b.StopAction = srv.GracefulStop
 	for _, cfg := range cfgs {
@@ -193,8 +199,6 @@ func run(cfgs []starter.Config, conf config.Config) error {
 		serverErrors <- repl.ProcessBacklog(ctx, praefect.ExpBackoffFunc(1*time.Second, 5*time.Second))
 	}()
 
-	go coordinator.FailoverRotation()
-
 	return <-serverErrors
 }
 
diff --git a/internal/praefect/auth_test.go b/internal/praefect/auth_test.go
index 6654fee19..48f1daf7a 100644
--- a/internal/praefect/auth_test.go
+++ b/internal/praefect/auth_test.go
@@ -184,7 +184,7 @@ func runServer(t *testing.T, token string, required bool) (*Server, string, func
 	gz := proto.FileDescriptor("mock.proto")
 	fd, err := protoregistry.ExtractFileDescriptor(gz)
 	if err != nil {
-		panic(err)
+		t.Fatal(err)
 	}
 
 	logEntry := testhelper.DiscardTestEntry(t)
@@ -193,11 +193,12 @@ func runServer(t *testing.T, token string, required bool) (*Server, string, func
 	nodeMgr, err := nodes.NewManager(logEntry, conf)
 	require.NoError(t, err)
 
-	coordinator := NewCoordinator(logEntry, ds, nodeMgr, conf, fd)
+	registry := protoregistry.New()
+	require.NoError(t, registry.RegisterFiles(fd))
 
-	replMgr := NewReplMgr("praefect-internal-0", logEntry, ds, nodeMgr)
+	coordinator := NewCoordinator(logEntry, ds, nodeMgr, conf, registry)
 
-	srv := NewServer(coordinator, replMgr, nil, logEntry, nodeMgr, conf)
+	srv := NewServer(coordinator.StreamDirector, logEntry, registry, conf)
 
 	serverSocketPath := testhelper.GetTemporaryGitalySocketFileName()
 
diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index 45b0dd841..14f58ae29 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -2,13 +2,8 @@ package praefect
 
 import (
 	"context"
-	"os"
-	"os/signal"
-	"sync"
-	"syscall"
 
 	"github.com/golang/protobuf/proto"
-	"github.com/golang/protobuf/protoc-gen-go/descriptor"
 	"github.com/sirupsen/logrus"
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/config"
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/datastore"
@@ -28,35 +23,24 @@ func isDestructive(methodName string) bool {
 // downstream server. The coordinator is thread safe; concurrent calls to
 // register nodes are safe.
 type Coordinator struct {
-	nodeMgr       nodes.Manager
-	log           *logrus.Entry
-	failoverMutex sync.RWMutex
-
+	nodeMgr   nodes.Manager
+	log       logrus.FieldLogger
 	datastore datastore.Datastore
-
-	registry *protoregistry.Registry
-	conf     config.Config
+	registry  *protoregistry.Registry
+	conf      config.Config
 }
 
 // NewCoordinator returns a new Coordinator that utilizes the provided logger
-func NewCoordinator(l *logrus.Entry, ds datastore.Datastore, nodeMgr nodes.Manager, conf config.Config, fileDescriptors ...*descriptor.FileDescriptorProto) *Coordinator {
-	registry := protoregistry.New()
-	registry.RegisterFiles(fileDescriptors...)
-
+func NewCoordinator(l logrus.FieldLogger, ds datastore.Datastore, nodeMgr nodes.Manager, conf config.Config, r *protoregistry.Registry) *Coordinator {
 	return &Coordinator{
 		log:       l,
 		datastore: ds,
-		registry:  registry,
+		registry:  r,
 		nodeMgr:   nodeMgr,
 		conf:      conf,
 	}
 }
 
-// RegisterProtos allows coordinator to register new protos on the fly
-func (c *Coordinator) RegisterProtos(protos ...*descriptor.FileDescriptorProto) error {
-	return c.registry.RegisterFiles(protos...)
-}
-
 func (c *Coordinator) directRepositoryScopedMessage(ctx context.Context, mi protoregistry.MethodInfo, peeker proxy.StreamModifier, fullMethodName string, m proto.Message) (*proxy.StreamParameters, error) {
 	targetRepo, err := mi.TargetRepo(m)
 	if err != nil {
@@ -106,14 +90,11 @@ func (c *Coordinator) directRepositoryScopedMessage(ctx context.Context, mi prot
 }
 
 // streamDirector determines which downstream servers receive requests
-func (c *Coordinator) streamDirector(ctx context.Context, fullMethodName string, peeker proxy.StreamModifier) (*proxy.StreamParameters, error) {
+func (c *Coordinator) StreamDirector(ctx context.Context, fullMethodName string, peeker proxy.StreamModifier) (*proxy.StreamParameters, error) {
 	// For phase 1, we need to route messages based on the storage location
 	// to the appropriate Gitaly node.
 	c.log.Debugf("Stream director received method %s", fullMethodName)
 
-	c.failoverMutex.RLock()
-	defer c.failoverMutex.RUnlock()
-
 	mi, err := c.registry.LookupMethod(fullMethodName)
 	if err != nil {
 		return nil, err
@@ -209,22 +190,3 @@ func (c *Coordinator) createReplicaJobs(targetRepo *gitalypb.Repository, primary
 		}
 	}, nil
 }
-
-// FailoverRotation waits for the SIGUSR1 signal, then promotes the next secondary to be primary
-func (c *Coordinator) FailoverRotation() {
-	c.handleSignalAndRotate()
-}
-
-func (c *Coordinator) handleSignalAndRotate() {
-	failoverChan := make(chan os.Signal, 1)
-	signal.Notify(failoverChan, syscall.SIGUSR1)
-
-	for {
-		<-failoverChan
-
-		c.failoverMutex.Lock()
-		// TODO: update failover logic
-		c.log.Info("failover happens")
-		c.failoverMutex.Unlock()
-	}
-}
diff --git a/internal/praefect/coordinator_test.go b/internal/praefect/coordinator_test.go
index dccab79f0..9cd4d6916 100644
--- a/internal/praefect/coordinator_test.go
+++ b/internal/praefect/coordinator_test.go
@@ -60,9 +60,10 @@ func TestStreamDirector(t *testing.T) {
 
 	nodeMgr, err := nodes.NewManager(entry, conf)
 	require.NoError(t, err)
+	r := protoregistry.New()
+	require.NoError(t, r.RegisterFiles(protoregistry.GitalyProtoFileDescriptors...))
 
-	coordinator := NewCoordinator(entry, ds, nodeMgr, conf)
-	require.NoError(t, coordinator.RegisterProtos(protoregistry.GitalyProtoFileDescriptors...))
+	coordinator := NewCoordinator(entry, ds, nodeMgr, conf, r)
 
 	frame, err := proto.Marshal(&gitalypb.FetchIntoObjectPoolRequest{
 		Origin:     &targetRepo,
@@ -74,7 +75,7 @@ func TestStreamDirector(t *testing.T) {
 	fullMethod := "/gitaly.ObjectPoolService/FetchIntoObjectPool"
 
 	peeker := &mockPeeker{frame}
-	streamParams, err := coordinator.streamDirector(ctx, fullMethod, peeker)
+	streamParams, err := coordinator.StreamDirector(ctx, fullMethod, peeker)
 	require.NoError(t, err)
 	require.Equal(t, address, streamParams.Conn().Target())
 
diff --git a/internal/praefect/helper_test.go b/internal/praefect/helper_test.go
index 16094ac6c..d236cefff 100644
--- a/internal/praefect/helper_test.go
+++ b/internal/praefect/helper_test.go
@@ -72,10 +72,10 @@ func testConfig(backends int) config.Config {
 
 // setupServer wires all praefect dependencies together via dependency
 // injection
-func setupServer(t testing.TB, conf config.Config, nodeMgr nodes.Manager, l *logrus.Entry, fds []*descriptor.FileDescriptorProto) (*datastore.MemoryDatastore, *Server) {
+func setupServer(t testing.TB, conf config.Config, nodeMgr nodes.Manager, l *logrus.Entry, r *protoregistry.Registry) (*datastore.MemoryDatastore, *Server) {
 	var (
 		ds          = datastore.NewInMemory(conf)
-		coordinator = NewCoordinator(l, ds, nodeMgr, conf, fds...)
+		coordinator = NewCoordinator(l, ds, nodeMgr, conf, r)
 	)
 
 	var defaultNode *models.Node
@@ -86,20 +86,7 @@ func setupServer(t testing.TB, conf config.Config, nodeMgr nodes.Manager, l *log
 	}
 	require.NotNil(t, defaultNode)
 
-	replmgr := NewReplMgr(
-		defaultNode.Storage,
-		l,
-		ds,
-		nodeMgr,
-	)
-	server := NewServer(
-		coordinator,
-		replmgr,
-		nil,
-		l,
-		nodeMgr,
-		conf,
-	)
+	server := NewServer(coordinator.StreamDirector, l, r, conf)
 
 	return ds, server
 }
@@ -131,14 +118,17 @@ func runPraefectServerWithMock(t *testing.T, conf config.Config, backends map[st
 	require.NoError(t, err)
 	nodeMgr.Start(1*time.Millisecond, 5*time.Millisecond)
 
-	_, prf := setupServer(t, conf, nodeMgr, log.Default(), []*descriptor.FileDescriptorProto{mustLoadProtoReg(t)})
+	r := protoregistry.New()
+	require.NoError(t, r.RegisterFiles(mustLoadProtoReg(t)))
+
+	_, prf := setupServer(t, conf, nodeMgr, log.Default(), r)
 
 	listener, port := listenAvailPort(t)
 	t.Logf("praefect listening on port %d", port)
 
 	errQ := make(chan error)
 
-	prf.RegisterServices()
+	prf.RegisterServices(nodeMgr, conf)
 	go func() {
 		errQ <- prf.Serve(listener, false)
 	}()
@@ -189,7 +179,9 @@ func runPraefectServerWithGitaly(t *testing.T, conf config.Config) (*grpc.Client
 	require.NoError(t, err)
 	nodeMgr.Start(1*time.Millisecond, 5*time.Millisecond)
 
-	coordinator := NewCoordinator(logEntry, ds, nodeMgr, conf, protoregistry.GitalyProtoFileDescriptors...)
+	registry := protoregistry.New()
+	require.NoError(t, registry.RegisterFiles(protoregistry.GitalyProtoFileDescriptors...))
+	coordinator := NewCoordinator(logEntry, ds, nodeMgr, conf, registry)
 
 	replmgr := NewReplMgr(
 		conf.VirtualStorages[0].Name,
@@ -200,11 +192,9 @@ func runPraefectServerWithGitaly(t *testing.T, conf config.Config) (*grpc.Client
 		WithLatencyMetric(&promtest.MockHistogram{}),
 	)
 	prf := NewServer(
-		coordinator,
-		replmgr,
-		nil,
+		coordinator.StreamDirector,
 		logEntry,
-		nodeMgr,
+		registry,
 		conf,
 	)
 
@@ -214,7 +204,7 @@ func runPraefectServerWithGitaly(t *testing.T, conf config.Config) (*grpc.Client
 	errQ := make(chan error)
 	ctx, cancel := testhelper.Context()
 
-	prf.RegisterServices()
+	prf.RegisterServices(nodeMgr, conf)
 	go func() { errQ <- prf.Serve(listener, false) }()
 	go func() { errQ <- replmgr.ProcessBacklog(ctx, noopBackoffFunc) }()
 
diff --git a/internal/praefect/mock/mock.pb.go b/internal/praefect/mock/mock.pb.go
index edc09f622..6f66b292b 100644
--- a/internal/praefect/mock/mock.pb.go
+++ b/internal/praefect/mock/mock.pb.go
@@ -12,6 +12,8 @@ import (
 	empty "github.com/golang/protobuf/ptypes/empty"
 	gitalypb "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 	grpc "google.golang.org/grpc"
+	codes "google.golang.org/grpc/codes"
+	status "google.golang.org/grpc/status"
 )
 
 // Reference imports to suppress errors if they are not otherwise used.
@@ -151,25 +153,25 @@ func init() {
 func init() { proto.RegisterFile("mock.proto", fileDescriptor_6fa4806c90f7156d) }
 
 var fileDescriptor_6fa4806c90f7156d = []byte{
-	// 275 bytes of a gzipped FileDescriptorProto
-	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x9c, 0x50, 0xcd, 0x4a, 0xc3, 0x40,
-	0x10, 0x66, 0x4b, 0x22, 0x71, 0xaa, 0xa5, 0x5d, 0x8b, 0xc8, 0x7a, 0x91, 0x80, 0x25, 0xa7, 0x2d,
-	0xad, 0xf8, 0x00, 0x1e, 0x0a, 0x7a, 0xf0, 0x92, 0xe2, 0x03, 0xa4, 0x71, 0x8c, 0xc1, 0xa4, 0xb3,
-	0xee, 0x6e, 0x0a, 0x79, 0x92, 0x3e, 0xa4, 0x6f, 0xd0, 0x93, 0x6c, 0x96, 0x60, 0x7b, 0xf5, 0x36,
-	0x33, 0xdf, 0xcf, 0x7e, 0xfb, 0x01, 0xd4, 0x94, 0x7f, 0x49, 0xa5, 0xc9, 0x12, 0x0f, 0xdc, 0x2c,
-	0x2e, 0xcc, 0x67, 0xa6, 0xf1, 0xdd, 0xdf, 0xc4, 0x6d, 0x41, 0x54, 0x54, 0x38, 0xef, 0xb6, 0x4d,
-	0xf3, 0x31, 0xc7, 0x5a, 0xd9, 0xd6, 0x83, 0xf1, 0x3d, 0x5c, 0xae, 0xcb, 0x5a, 0x55, 0x98, 0xe2,
-	0x77, 0x83, 0xc6, 0xf2, 0x29, 0x84, 0xbb, 0xac, 0x6a, 0xf0, 0x86, 0xdd, 0xb1, 0x24, 0x4c, 0xfd,
-	0x12, 0xcf, 0x60, 0xd4, 0xd3, 0x8c, 0xa2, 0xad, 0xc1, 0x3f, 0xde, 0xe0, 0x98, 0xf7, 0x08, 0xc3,
-	0x14, 0x15, 0xf5, 0x66, 0x33, 0x08, 0x34, 0x2a, 0xea, 0xbc, 0x86, 0x4b, 0x2e, 0x8b, 0xd2, 0x66,
-	0x55, 0x2b, 0x1d, 0xc5, 0x94, 0x96, 0x74, 0x9b, 0x76, 0xf8, 0xf2, 0x87, 0xf5, 0x31, 0xd6, 0xa8,
-	0x77, 0x65, 0x8e, 0x7c, 0x05, 0x23, 0x37, 0xa2, 0x7e, 0xca, 0x73, 0x34, 0x86, 0x34, 0xbf, 0x92,
-	0xdd, 0x3f, 0x4f, 0xd2, 0x8a, 0xe9, 0xe9, 0xd1, 0x67, 0x8b, 0xa3, 0xc3, 0x3e, 0x09, 0xa2, 0xc1,
-	0x98, 0xf1, 0x17, 0x98, 0xb8, 0xc7, 0x7a, 0x93, 0xb7, 0x6d, 0xa6, 0x5b, 0x3e, 0xf1, 0xa2, 0xa3,
-	0xa0, 0xe2, 0x5a, 0xfa, 0x92, 0x64, 0x5f, 0x92, 0x5c, 0xb9, 0x92, 0xe2, 0xf3, 0xc3, 0x3e, 0x09,
-	0xa3, 0x81, 0x60, 0x0b, 0xfe, 0x0c, 0x63, 0xa7, 0x78, 0x6d, 0x6c, 0x66, 0xff, 0xed, 0xc4, 0x04,
-	0x5b, 0x6c, 0xce, 0x3a, 0xe8, 0xe1, 0x37, 0x00, 0x00, 0xff, 0xff, 0x3c, 0x0c, 0xc7, 0x57, 0xb9,
-	0x01, 0x00, 0x00,
+	// 279 bytes of a gzipped FileDescriptorProto
+	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x9c, 0x90, 0x41, 0x4a, 0xc3, 0x40,
+	0x14, 0x86, 0x99, 0x90, 0x96, 0xf0, 0xaa, 0xa5, 0x1d, 0x8b, 0x48, 0xdc, 0x48, 0x40, 0xc9, 0xa2,
+	0x4c, 0xa1, 0x2e, 0x5d, 0x29, 0xc4, 0x9d, 0x9b, 0x14, 0x0f, 0x90, 0xc6, 0x67, 0x0c, 0x26, 0x7d,
+	0xe3, 0xcc, 0xa4, 0x90, 0x93, 0xd4, 0x13, 0x79, 0x13, 0x2f, 0xd1, 0x95, 0x4c, 0x86, 0x60, 0xbb,
+	0xed, 0xee, 0xbd, 0x7f, 0xbe, 0xf9, 0xe7, 0x9f, 0x1f, 0xa0, 0xa6, 0xfc, 0x53, 0x48, 0x45, 0x86,
+	0xb8, 0x6f, 0xe7, 0xf0, 0x4c, 0x7f, 0x64, 0x0a, 0xdf, 0x9c, 0x16, 0x5e, 0x17, 0x44, 0x45, 0x85,
+	0x8b, 0x6e, 0x5b, 0x37, 0xef, 0x0b, 0xac, 0xa5, 0x69, 0xdd, 0x61, 0x74, 0x0b, 0xe7, 0xab, 0xb2,
+	0x96, 0x15, 0xa6, 0xf8, 0xd5, 0xa0, 0x36, 0x7c, 0x06, 0x83, 0x6d, 0x56, 0x35, 0x78, 0xc5, 0x6e,
+	0x58, 0x3c, 0x48, 0xdd, 0x12, 0xdd, 0xc1, 0xb8, 0xc7, 0xb4, 0xa4, 0x8d, 0xc6, 0x7f, 0xce, 0x3b,
+	0xe4, 0x1e, 0x60, 0x94, 0xa2, 0xa4, 0xde, 0x6c, 0x0e, 0xbe, 0x42, 0x49, 0x9d, 0xd7, 0x68, 0xc9,
+	0x45, 0x51, 0x9a, 0xac, 0x6a, 0x85, 0x45, 0x74, 0x69, 0x48, 0xb5, 0x4f, 0xfe, 0xf7, 0xcf, 0x9c,
+	0xa5, 0x1d, 0xb5, 0xfc, 0x65, 0x7d, 0x98, 0x15, 0xaa, 0x6d, 0x99, 0x23, 0x4f, 0x60, 0x6c, 0x47,
+	0x54, 0x8f, 0x79, 0x8e, 0x5a, 0x93, 0xe2, 0x17, 0xa2, 0xfb, 0xed, 0x51, 0xe6, 0x70, 0x76, 0x2c,
+	0xba, 0x84, 0x51, 0xb0, 0xdf, 0xc5, 0x7e, 0xe0, 0x4d, 0x18, 0x7f, 0x86, 0xa9, 0x7d, 0xb2, 0x37,
+	0x79, 0xdd, 0x64, 0xaa, 0xe5, 0x53, 0x77, 0xe9, 0x20, 0x6e, 0x78, 0x29, 0x5c, 0x55, 0xa2, 0xaf,
+	0x4a, 0x24, 0xb6, 0xaa, 0x68, 0xb8, 0xdf, 0xc5, 0x5e, 0xe0, 0xf1, 0x04, 0x26, 0x16, 0x7f, 0x69,
+	0x4c, 0x66, 0x4e, 0xb6, 0x61, 0xeb, 0x61, 0xa7, 0xdf, 0xff, 0x05, 0x00, 0x00, 0xff, 0xff, 0x6a,
+	0x35, 0x84, 0xac, 0xb9, 0x01, 0x00, 0x00,
 }
 
 // Reference imports to suppress errors if they are not otherwise used.
@@ -237,6 +239,20 @@ type SimpleServiceServer interface {
 	RepoMutatorUnary(context.Context, *RepoRequest) (*empty.Empty, error)
 }
 
+// UnimplementedSimpleServiceServer can be embedded to have forward compatible implementations.
+type UnimplementedSimpleServiceServer struct {
+}
+
+func (*UnimplementedSimpleServiceServer) ServerAccessor(ctx context.Context, req *SimpleRequest) (*SimpleResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method ServerAccessor not implemented")
+}
+func (*UnimplementedSimpleServiceServer) RepoAccessorUnary(ctx context.Context, req *RepoRequest) (*empty.Empty, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method RepoAccessorUnary not implemented")
+}
+func (*UnimplementedSimpleServiceServer) RepoMutatorUnary(ctx context.Context, req *RepoRequest) (*empty.Empty, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method RepoMutatorUnary not implemented")
+}
+
 func RegisterSimpleServiceServer(s *grpc.Server, srv SimpleServiceServer) {
 	s.RegisterService(&_SimpleService_serviceDesc, srv)
 }
diff --git a/internal/praefect/mock/mock.proto b/internal/praefect/mock/mock.proto
index a1718fdd7..8aa159efa 100644
--- a/internal/praefect/mock/mock.proto
+++ b/internal/praefect/mock/mock.proto
@@ -19,7 +19,7 @@ message SimpleResponse {
 }
 
 message RepoRequest {
-	gitaly.Repository repo = 1;
+	gitaly.Repository repo = 1 [(gitaly.target_repository)=true];
 }
 
 service SimpleService {
@@ -36,7 +36,6 @@ service SimpleService {
     option (gitaly.op_type) = {
       op: ACCESSOR
       scope_level: REPOSITORY
-      target_repository_field: "1"
     };
   }
 	
@@ -45,7 +44,6 @@ service SimpleService {
     option (gitaly.op_type) = {
       op: MUTATOR
       scope_level: REPOSITORY
-      target_repository_field: "1"
     };
   }
 }
diff --git a/internal/praefect/server.go b/internal/praefect/server.go
index 863c30a3b..76473adf5 100644
--- a/internal/praefect/server.go
+++ b/internal/praefect/server.go
@@ -18,7 +18,9 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/middleware/sentryhandler"
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/config"
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/grpc-proxy/proxy"
+	"gitlab.com/gitlab-org/gitaly/internal/praefect/middleware"
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/nodes"
+	"gitlab.com/gitlab-org/gitaly/internal/praefect/protoregistry"
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/service/server"
 	"gitlab.com/gitlab-org/gitaly/internal/server/auth"
 	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
@@ -31,11 +33,8 @@ import (
 
 // Server is a praefect server
 type Server struct {
-	nodeManager nodes.Manager
-	repl        ReplMgr
-	s           *grpc.Server
-	conf        config.Config
-	l           *logrus.Entry
+	s *grpc.Server
+	l *logrus.Entry
 }
 
 func (srv *Server) warnDupeAddrs(c config.Config) {
@@ -60,16 +59,17 @@ func (srv *Server) warnDupeAddrs(c config.Config) {
 
 // NewServer returns an initialized praefect gPRC proxy server configured
 // with the provided gRPC server options
-func NewServer(c *Coordinator, repl ReplMgr, grpcOpts []grpc.ServerOption, l *logrus.Entry, nodeManager nodes.Manager, conf config.Config) *Server {
+func NewServer(director proxy.StreamDirector, l *logrus.Entry, r *protoregistry.Registry, conf config.Config, grpcOpts ...grpc.ServerOption) *Server {
 	ctxTagOpts := []grpc_ctxtags.Option{
 		grpc_ctxtags.WithFieldExtractorForInitialReq(fieldextractors.FieldExtractor),
 	}
 
-	grpcOpts = append(grpcOpts, proxyRequiredOpts(c.streamDirector)...)
+	grpcOpts = append(grpcOpts, proxyRequiredOpts(director)...)
 	grpcOpts = append(grpcOpts, []grpc.ServerOption{
 		grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
 			grpc_ctxtags.StreamServerInterceptor(ctxTagOpts...),
 			grpccorrelation.StreamServerCorrelationInterceptor(), // Must be above the metadata handler
+			middleware.MethodTypeStreamInterceptor(r),
 			metadatahandler.StreamInterceptor,
 			grpc_prometheus.StreamServerInterceptor,
 			grpc_logrus.StreamServerInterceptor(l),
@@ -84,6 +84,7 @@ func NewServer(c *Coordinator, repl ReplMgr, grpcOpts []grpc.ServerOption, l *lo
 		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(
 			grpc_ctxtags.UnaryServerInterceptor(ctxTagOpts...),
 			grpccorrelation.UnaryServerCorrelationInterceptor(), // Must be above the metadata handler
+			middleware.MethodTypeUnaryInterceptor(r),
 			metadatahandler.UnaryInterceptor,
 			grpc_prometheus.UnaryServerInterceptor,
 			grpc_logrus.UnaryServerInterceptor(l),
@@ -98,11 +99,8 @@ func NewServer(c *Coordinator, repl ReplMgr, grpcOpts []grpc.ServerOption, l *lo
 	}...)
 
 	s := &Server{
-		s:           grpc.NewServer(grpcOpts...),
-		repl:        repl,
-		nodeManager: nodeManager,
-		conf:        conf,
-		l:           l,
+		s: grpc.NewServer(grpcOpts...),
+		l: l,
 	}
 
 	s.warnDupeAddrs(conf)
@@ -123,9 +121,9 @@ func (srv *Server) Serve(l net.Listener, secure bool) error {
 }
 
 // RegisterServices will register any services praefect needs to handle rpcs on its own
-func (srv *Server) RegisterServices() {
+func (srv *Server) RegisterServices(nm nodes.Manager, conf config.Config) {
 	// ServerServiceServer is necessary for the ServerInfo RPC
-	gitalypb.RegisterServerServiceServer(srv.s, server.NewServer(srv.conf, srv.nodeManager))
+	gitalypb.RegisterServerServiceServer(srv.s, server.NewServer(conf, nm))
 
 	healthpb.RegisterHealthServer(srv.s, health.NewServer())
 
diff --git a/internal/praefect/server_test.go b/internal/praefect/server_test.go
index 23fb06f9c..b52d08fc3 100644
--- a/internal/praefect/server_test.go
+++ b/internal/praefect/server_test.go
@@ -133,11 +133,14 @@ func TestGitalyServerInfoBadNode(t *testing.T) {
 	nodeMgr, err := nodes.NewManager(entry, conf)
 	require.NoError(t, err)
 
-	_, srv := setupServer(t, conf, nodeMgr, entry, protoregistry.GitalyProtoFileDescriptors)
+	registry := protoregistry.New()
+	require.NoError(t, registry.RegisterFiles(protoregistry.GitalyProtoFileDescriptors...))
+
+	_, srv := setupServer(t, conf, nodeMgr, entry, registry)
 
 	listener, port := listenAvailPort(t)
 	go func() {
-		srv.RegisterServices()
+		srv.RegisterServices(nodeMgr, conf)
 		srv.Serve(listener, false)
 	}()