praefect: Make readiness checks injectable

It is hard to test ReadinessCheck RPC because the set of the
checks that is executed is hard or not possible to mock or
substitute. Those we made check an injectable to provide from
outside. It will help us to write better tests for the RPC.

9 files changed, 33 insertions, 13 deletions

diff --git a/cmd/praefect/main.go b/cmd/praefect/main.go
index f2389aeb7..b9edba3d1 100644
--- a/cmd/praefect/main.go
+++ b/cmd/praefect/main.go
@@ -88,6 +88,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/protoregistry"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/reconciler"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/repocleaner"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/service"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/service/transaction"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/transactions"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/sidechannel"
@@ -435,6 +436,7 @@ func run(
 			protoregistry.GitalyProtoPreregistered,
 			nodeSet.Connections(),
 			primaryGetter,
+			service.AllChecks(),
 		)
 	)
 	metricsCollectors = append(metricsCollectors, transactionManager, coordinator, repl)
diff --git a/internal/praefect/auth_test.go b/internal/praefect/auth_test.go
index ba62e0326..9f0cd9012 100644
--- a/internal/praefect/auth_test.go
+++ b/internal/praefect/auth_test.go
@@ -163,7 +163,7 @@ func runServer(t *testing.T, token string, required bool) (*grpc.Server, string,
 
 	coordinator := NewCoordinator(queue, nil, NewNodeManagerRouter(nodeMgr, nil), txMgr, conf, registry)
 
-	srv := NewGRPCServer(conf, logEntry, registry, coordinator.StreamDirector, txMgr, nil, nil, nil, nil, nil)
+	srv := NewGRPCServer(conf, logEntry, registry, coordinator.StreamDirector, txMgr, nil, nil, nil, nil, nil, nil)
 
 	serverSocketPath := testhelper.GetTemporaryGitalySocketFileName(t)
 
diff --git a/internal/praefect/remove_repository_test.go b/internal/praefect/remove_repository_test.go
index 4a1b24251..f9eaf057a 100644
--- a/internal/praefect/remove_repository_test.go
+++ b/internal/praefect/remove_repository_test.go
@@ -120,6 +120,7 @@ func TestRemoveRepositoryHandler(t *testing.T) {
 				nodeSet.Connections(),
 				nil,
 				nil,
+				nil,
 			)
 			defer srv.Stop()
 
diff --git a/internal/praefect/repository_exists_test.go b/internal/praefect/repository_exists_test.go
index 87f38d85a..3eb445e7e 100644
--- a/internal/praefect/repository_exists_test.go
+++ b/internal/praefect/repository_exists_test.go
@@ -99,6 +99,7 @@ func TestRepositoryExistsHandler(t *testing.T) {
 				nil,
 				nil,
 				nil,
+				nil,
 			)
 			defer srv.Stop()
 
diff --git a/internal/praefect/server.go b/internal/praefect/server.go
index 3192ddfdd..202eb2e6d 100644
--- a/internal/praefect/server.go
+++ b/internal/praefect/server.go
@@ -92,6 +92,7 @@ func NewGRPCServer(
 	conns Connections,
 	primaryGetter PrimaryGetter,
 	creds credentials.TransportCredentials,
+	checks []service.CheckFunc,
 	grpcOpts ...grpc.ServerOption,
 ) *grpc.Server {
 	streamInterceptors := []grpc.StreamServerInterceptor{
@@ -153,7 +154,7 @@ func NewGRPCServer(
 	warnDupeAddrs(logger, conf)
 
 	srv := grpc.NewServer(grpcOpts...)
-	registerServices(srv, txMgr, conf, rs, assignmentStore, service.Connections(conns), primaryGetter)
+	registerServices(srv, txMgr, conf, rs, assignmentStore, service.Connections(conns), primaryGetter, checks)
 
 	if conf.Failover.ElectionStrategy == config.ElectionStrategyPerRepository {
 		proxy.RegisterStreamHandlers(srv, "gitaly.RepositoryService", map[string]grpc.StreamHandler{
@@ -184,9 +185,10 @@ func registerServices(
 	assignmentStore AssignmentStore,
 	conns service.Connections,
 	primaryGetter info.PrimaryGetter,
+	checks []service.CheckFunc,
 ) {
 	// ServerServiceServer is necessary for the ServerInfo RPC
-	gitalypb.RegisterServerServiceServer(srv, server.NewServer(conf, conns))
+	gitalypb.RegisterServerServiceServer(srv, server.NewServer(conf, conns, checks))
 	gitalypb.RegisterPraefectInfoServiceServer(srv, info.NewServer(conf, rs, assignmentStore, conns, primaryGetter))
 	gitalypb.RegisterRefTransactionServer(srv, transaction.NewServer(tm))
 	healthpb.RegisterHealthServer(srv, health.NewServer())
diff --git a/internal/praefect/server_factory.go b/internal/praefect/server_factory.go
index da7fe2c7a..45424fe9d 100644
--- a/internal/praefect/server_factory.go
+++ b/internal/praefect/server_factory.go
@@ -12,6 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/grpc-proxy/proxy"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/nodes"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/protoregistry"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/service"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/transactions"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
@@ -30,6 +31,7 @@ func NewServerFactory(
 	registry *protoregistry.Registry,
 	conns Connections,
 	primaryGetter PrimaryGetter,
+	checks []service.CheckFunc,
 ) *ServerFactory {
 	return &ServerFactory{
 		conf:            conf,
@@ -43,6 +45,7 @@ func NewServerFactory(
 		registry:        registry,
 		conns:           conns,
 		primaryGetter:   primaryGetter,
+		checks:          checks,
 	}
 }
 
@@ -61,6 +64,7 @@ type ServerFactory struct {
 	secure, insecure []*grpc.Server
 	conns            Connections
 	primaryGetter    PrimaryGetter
+	checks           []service.CheckFunc
 }
 
 // Serve starts serving on the provided listener with newly created grpc.Server
@@ -131,6 +135,7 @@ func (s *ServerFactory) createGRPC(creds credentials.TransportCredentials) *grpc
 		s.conns,
 		s.primaryGetter,
 		creds,
+		s.checks,
 	)
 }
 
diff --git a/internal/praefect/server_factory_test.go b/internal/praefect/server_factory_test.go
index f9bbdd16a..7fd850e21 100644
--- a/internal/praefect/server_factory_test.go
+++ b/internal/praefect/server_factory_test.go
@@ -179,7 +179,7 @@ func TestServerFactory(t *testing.T) {
 	}
 
 	t.Run("insecure", func(t *testing.T) {
-		praefectServerFactory := NewServerFactory(conf, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil)
+		praefectServerFactory := NewServerFactory(conf, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil, nil)
 		defer praefectServerFactory.Stop()
 
 		listener, err := net.Listen(starter.TCP, "localhost:0")
@@ -212,7 +212,7 @@ func TestServerFactory(t *testing.T) {
 	})
 
 	t.Run("secure", func(t *testing.T) {
-		praefectServerFactory := NewServerFactory(conf, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil)
+		praefectServerFactory := NewServerFactory(conf, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil, nil)
 		defer praefectServerFactory.Stop()
 
 		listener, err := net.Listen(starter.TCP, "localhost:0")
@@ -254,7 +254,7 @@ func TestServerFactory(t *testing.T) {
 	t.Run("stops all listening servers", func(t *testing.T) {
 		ctx := testhelper.Context(t)
 
-		praefectServerFactory := NewServerFactory(conf, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil)
+		praefectServerFactory := NewServerFactory(conf, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil, nil)
 		defer praefectServerFactory.Stop()
 
 		// start with tcp address
@@ -322,7 +322,7 @@ func TestServerFactory(t *testing.T) {
 	t.Run("tls key path invalid", func(t *testing.T) {
 		badTLSKeyPath := conf
 		badTLSKeyPath.TLS.KeyPath = "invalid"
-		praefectServerFactory := NewServerFactory(badTLSKeyPath, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil)
+		praefectServerFactory := NewServerFactory(badTLSKeyPath, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil, nil)
 
 		err := praefectServerFactory.Serve(nil, true)
 		require.EqualError(t, err, "load certificate key pair: open invalid: no such file or directory")
@@ -331,7 +331,7 @@ func TestServerFactory(t *testing.T) {
 	t.Run("tls cert path invalid", func(t *testing.T) {
 		badTLSKeyPath := conf
 		badTLSKeyPath.TLS.CertPath = "invalid"
-		praefectServerFactory := NewServerFactory(badTLSKeyPath, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil)
+		praefectServerFactory := NewServerFactory(badTLSKeyPath, logger, coordinator.StreamDirector, nodeMgr, txMgr, queue, rs, datastore.AssignmentStore{}, registry, nil, nil, nil)
 
 		err := praefectServerFactory.Serve(nil, true)
 		require.EqualError(t, err, "load certificate key pair: open invalid: no such file or directory")
diff --git a/internal/praefect/service/server/server.go b/internal/praefect/service/server/server.go
index e15dc5bea..f49621576 100644
--- a/internal/praefect/service/server/server.go
+++ b/internal/praefect/service/server/server.go
@@ -9,15 +9,17 @@ import (
 // Server is a ServerService server
 type Server struct {
 	gitalypb.UnimplementedServerServiceServer
-	conf  config.Config
-	conns service.Connections
+	conf   config.Config
+	conns  service.Connections
+	checks []service.CheckFunc
 }
 
 // NewServer creates a new instance of a grpc ServerServiceServer
-func NewServer(conf config.Config, conns service.Connections) gitalypb.ServerServiceServer {
+func NewServer(conf config.Config, conns service.Connections, checks []service.CheckFunc) gitalypb.ServerServiceServer {
 	s := &Server{
-		conf:  conf,
-		conns: conns,
+		conf:   conf,
+		conns:  conns,
+		checks: checks,
 	}
 
 	return s
diff --git a/internal/praefect/testserver.go b/internal/praefect/testserver.go
index eec0bf4b9..142b46764 100644
--- a/internal/praefect/testserver.go
+++ b/internal/praefect/testserver.go
@@ -19,6 +19,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/mock"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/nodes"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/protoregistry"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/service"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/transactions"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper/promtest"
@@ -53,6 +54,8 @@ type BuildOptions struct {
 	WithPrimaryGetter PrimaryGetter
 	// WithRouter sets an implementation of the request router to use by praefect service.
 	WithRouter Router
+	// WithChecks sets a list of check to run when ReadinessCheck RPC is called.
+	WithChecks []service.CheckFunc
 }
 
 // WithMockBackends mocks backends with a set of passed in stubs.
@@ -215,6 +218,9 @@ func RunPraefectServer(
 	if opt.WithRouter == nil {
 		opt.WithRouter = NewNodeManagerRouter(opt.WithNodeMgr, opt.WithRepoStore)
 	}
+	if opt.WithChecks == nil {
+		opt.WithChecks = service.AllChecks()
+	}
 
 	coordinator := NewCoordinator(
 		opt.WithQueue,
@@ -246,6 +252,7 @@ func RunPraefectServer(
 		opt.WithConnections,
 		opt.WithPrimaryGetter,
 		nil,
+		opt.WithChecks,
 	)
 
 	listener, port := listenAvailPort(t)