Merge branch 'ps-refactor-NewServerWithHealth' into 'master'

Refactoring of the NewHealthServerWithListener

See merge request gitlab-org/gitaly!3428

7 files changed, 40 insertions, 70 deletions

diff --git a/internal/praefect/coordinator_pg_test.go b/internal/praefect/coordinator_pg_test.go
index dfab8af54..23edac9e2 100644
--- a/internal/praefect/coordinator_pg_test.go
+++ b/internal/praefect/coordinator_pg_test.go
@@ -121,8 +121,7 @@ func TestStreamDirectorMutator_Transaction(t *testing.T) {
 			storageNodes := make([]*config.Node, 0, len(tc.nodes))
 			for i := range tc.nodes {
 				socket := testhelper.GetTemporaryGitalySocketFileName(t)
-				server, _ := testhelper.NewServerWithHealth(t, socket)
-				defer server.Stop()
+				testhelper.NewServerWithHealth(t, socket)
 				node := &config.Node{Address: "unix://" + socket, Storage: fmt.Sprintf("node-%d", i)}
 				storageNodes = append(storageNodes, node)
 			}
diff --git a/internal/praefect/coordinator_test.go b/internal/praefect/coordinator_test.go
index 49e60ce89..d1377685c 100644
--- a/internal/praefect/coordinator_test.go
+++ b/internal/praefect/coordinator_test.go
@@ -126,10 +126,8 @@ func TestStreamDirectorReadOnlyEnforcement(t *testing.T) {
 
 func TestStreamDirectorMutator(t *testing.T) {
 	gitalySocket0, gitalySocket1 := testhelper.GetTemporaryGitalySocketFileName(t), testhelper.GetTemporaryGitalySocketFileName(t)
-	srv1, _ := testhelper.NewServerWithHealth(t, gitalySocket0)
-	defer srv1.Stop()
-	srv2, _ := testhelper.NewServerWithHealth(t, gitalySocket1)
-	defer srv2.Stop()
+	testhelper.NewServerWithHealth(t, gitalySocket0)
+	testhelper.NewServerWithHealth(t, gitalySocket1)
 
 	primaryAddress, secondaryAddress := "unix://"+gitalySocket0, "unix://"+gitalySocket1
 	primaryNode := &config.Node{Address: primaryAddress, Storage: "praefect-internal-1"}
@@ -236,8 +234,7 @@ func TestStreamDirectorMutator(t *testing.T) {
 
 func TestStreamDirectorMutator_StopTransaction(t *testing.T) {
 	socket := testhelper.GetTemporaryGitalySocketFileName(t)
-	server, _ := testhelper.NewServerWithHealth(t, socket)
-	defer server.Stop()
+	testhelper.NewServerWithHealth(t, socket)
 
 	conf := config.Config{
 		VirtualStorages: []*config.VirtualStorage{
@@ -355,8 +352,7 @@ func (m mockRouter) RouteRepositoryAccessor(ctx context.Context, virtualStorage,
 
 func TestStreamDirectorAccessor(t *testing.T) {
 	gitalySocket := testhelper.GetTemporaryGitalySocketFileName(t)
-	srv, _ := testhelper.NewServerWithHealth(t, gitalySocket)
-	defer srv.Stop()
+	testhelper.NewServerWithHealth(t, gitalySocket)
 
 	gitalyAddress := "unix://" + gitalySocket
 	conf := config.Config{
@@ -457,10 +453,8 @@ func TestStreamDirectorAccessor(t *testing.T) {
 
 func TestCoordinatorStreamDirector_distributesReads(t *testing.T) {
 	gitalySocket0, gitalySocket1 := testhelper.GetTemporaryGitalySocketFileName(t), testhelper.GetTemporaryGitalySocketFileName(t)
-	srv1, primaryHealthSrv := testhelper.NewServerWithHealth(t, gitalySocket0)
-	defer srv1.Stop()
-	srv2, healthSrv := testhelper.NewServerWithHealth(t, gitalySocket1)
-	defer srv2.Stop()
+	primaryHealthSrv := testhelper.NewServerWithHealth(t, gitalySocket0)
+	healthSrv := testhelper.NewServerWithHealth(t, gitalySocket1)
 
 	primaryNodeConf := config.Node{
 		Address: "unix://" + gitalySocket0,
@@ -835,13 +829,10 @@ func TestStreamDirector_repo_creation(t *testing.T) {
 				gitalySocket0 := testhelper.GetTemporaryGitalySocketFileName(t)
 				gitalySocket1 := testhelper.GetTemporaryGitalySocketFileName(t)
 				gitalySocket3 := testhelper.GetTemporaryGitalySocketFileName(t)
-				srv1, _ := testhelper.NewServerWithHealth(t, gitalySocket0)
-				defer srv1.Stop()
-				srv2, _ := testhelper.NewServerWithHealth(t, gitalySocket1)
-				defer srv2.Stop()
-				srv3, healthSrv3 := testhelper.NewServerWithHealth(t, gitalySocket3)
+				testhelper.NewServerWithHealth(t, gitalySocket0)
+				testhelper.NewServerWithHealth(t, gitalySocket1)
+				healthSrv3 := testhelper.NewServerWithHealth(t, gitalySocket3)
 				healthSrv3.SetServingStatus("", grpc_health_v1.HealthCheckResponse_NOT_SERVING)
-				defer srv3.Stop()
 
 				primaryNode.Address = "unix://" + gitalySocket0
 				healthySecondaryNode.Address = "unix://" + gitalySocket1
@@ -1000,8 +991,8 @@ func (m *mockPeeker) Modify(payload []byte) error {
 
 func TestAbsentCorrelationID(t *testing.T) {
 	gitalySocket0, gitalySocket1 := testhelper.GetTemporaryGitalySocketFileName(t), testhelper.GetTemporaryGitalySocketFileName(t)
-	_, healthSrv0 := testhelper.NewServerWithHealth(t, gitalySocket0)
-	_, healthSrv1 := testhelper.NewServerWithHealth(t, gitalySocket1)
+	healthSrv0 := testhelper.NewServerWithHealth(t, gitalySocket0)
+	healthSrv1 := testhelper.NewServerWithHealth(t, gitalySocket1)
 	healthSrv0.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING)
 	healthSrv1.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING)
 
@@ -1151,10 +1142,8 @@ func TestCoordinatorEnqueueFailure(t *testing.T) {
 func TestStreamDirectorStorageScope(t *testing.T) {
 	// stubs health-check requests because nodes.NewManager establishes connection on creation
 	gitalySocket0, gitalySocket1 := testhelper.GetTemporaryGitalySocketFileName(t), testhelper.GetTemporaryGitalySocketFileName(t)
-	srv1, _ := testhelper.NewServerWithHealth(t, gitalySocket0)
-	defer srv1.Stop()
-	srv2, _ := testhelper.NewServerWithHealth(t, gitalySocket1)
-	defer srv2.Stop()
+	testhelper.NewServerWithHealth(t, gitalySocket0)
+	testhelper.NewServerWithHealth(t, gitalySocket1)
 
 	primaryAddress, secondaryAddress := "unix://"+gitalySocket0, "unix://"+gitalySocket1
 	primaryGitaly := &config.Node{Address: primaryAddress, Storage: "gitaly-1"}
diff --git a/internal/praefect/nodes/local_elector_test.go b/internal/praefect/nodes/local_elector_test.go
index a8b848321..46b649d53 100644
--- a/internal/praefect/nodes/local_elector_test.go
+++ b/internal/praefect/nodes/local_elector_test.go
@@ -12,9 +12,9 @@ import (
 	"google.golang.org/grpc"
 )
 
-func setupElector(t *testing.T) (*localElector, []*nodeStatus, *grpc.ClientConn, *grpc.Server) {
+func setupElector(t *testing.T) (*localElector, []*nodeStatus, *grpc.ClientConn) {
 	socket := testhelper.GetTemporaryGitalySocketFileName(t)
-	svr, _ := testhelper.NewServerWithHealth(t, socket)
+	testhelper.NewServerWithHealth(t, socket)
 
 	cc, err := grpc.Dial(
 		"unix://"+socket,
@@ -34,12 +34,11 @@ func setupElector(t *testing.T) (*localElector, []*nodeStatus, *grpc.ClientConn,
 
 	strategy.bootstrap(time.Second)
 
-	return strategy, ns, cc, svr
+	return strategy, ns, cc
 }
 
 func TestGetShard(t *testing.T) {
-	strategy, ns, _, svr := setupElector(t)
-	defer svr.Stop()
+	strategy, ns, _ := setupElector(t)
 
 	ctx, cancel := testhelper.Context()
 	defer cancel()
@@ -52,8 +51,7 @@ func TestGetShard(t *testing.T) {
 }
 
 func TestConcurrentCheckWithPrimary(t *testing.T) {
-	strategy, ns, _, svr := setupElector(t)
-	defer svr.Stop()
+	strategy, ns, _ := setupElector(t)
 
 	iterations := 10
 	var wg sync.WaitGroup
diff --git a/internal/praefect/nodes/manager_test.go b/internal/praefect/nodes/manager_test.go
index a550049de..6aa11a59c 100644
--- a/internal/praefect/nodes/manager_test.go
+++ b/internal/praefect/nodes/manager_test.go
@@ -57,8 +57,7 @@ func assertShard(t *testing.T, exp shardAssertion, act Shard) {
 
 func TestNodeStatus(t *testing.T) {
 	socket := testhelper.GetTemporaryGitalySocketFileName(t)
-	svr, healthSvr := testhelper.NewServerWithHealth(t, socket)
-	defer svr.Stop()
+	healthSvr := testhelper.NewServerWithHealth(t, socket)
 
 	cc, err := grpc.Dial(
 		"unix://"+socket,
@@ -112,11 +111,8 @@ func TestManagerFailoverDisabledElectionStrategySQL(t *testing.T) {
 		},
 	}
 
-	srv0, healthSrv := testhelper.NewServerWithHealth(t, socket0)
-	defer srv0.Stop()
-
-	srv1, _ := testhelper.NewServerWithHealth(t, socket1)
-	defer srv1.Stop()
+	healthSrv := testhelper.NewServerWithHealth(t, socket0)
+	testhelper.NewServerWithHealth(t, socket1)
 
 	conf := config.Config{
 		Failover:        config.Failover{Enabled: false, ElectionStrategy: config.ElectionStrategySQL},
@@ -168,8 +164,7 @@ func TestDialWithUnhealthyNode(t *testing.T) {
 		},
 	}
 
-	srv, _ := testhelper.NewHealthServerWithListener(t, primaryLn)
-	defer srv.Stop()
+	testhelper.NewHealthServerWithListener(t, primaryLn)
 
 	mgr, err := NewManager(testhelper.DiscardTestEntry(t), conf, nil, nil, promtest.NewMockHistogramVec(), protoregistry.GitalyProtoPreregistered, nil, nil)
 	require.NoError(t, err)
@@ -189,11 +184,8 @@ func TestDialWithUnhealthyNode(t *testing.T) {
 
 func TestNodeManager(t *testing.T) {
 	internalSocket0, internalSocket1 := testhelper.GetTemporaryGitalySocketFileName(t), testhelper.GetTemporaryGitalySocketFileName(t)
-	srv0, healthSrv0 := testhelper.NewServerWithHealth(t, internalSocket0)
-	defer srv0.Stop()
-
-	srv1, healthSrv1 := testhelper.NewServerWithHealth(t, internalSocket1)
-	defer srv1.Stop()
+	healthSrv0 := testhelper.NewServerWithHealth(t, internalSocket0)
+	healthSrv1 := testhelper.NewServerWithHealth(t, internalSocket1)
 
 	node1 := &config.Node{
 		Storage: "praefect-internal-0",
@@ -314,13 +306,11 @@ func TestMgr_GetSyncedNode(t *testing.T) {
 	const virtualStorage = "virtual-storage-0"
 	const repoPath = "path/1"
 
-	var srvs [count]*grpc.Server
 	var healthSrvs [count]*health.Server
 	var nodes [count]*config.Node
 	for i := 0; i < count; i++ {
 		socket := testhelper.GetTemporaryGitalySocketFileName(t)
-		srvs[i], healthSrvs[i] = testhelper.NewServerWithHealth(t, socket)
-		defer srvs[i].Stop()
+		healthSrvs[i] = testhelper.NewServerWithHealth(t, socket)
 		nodes[i] = &config.Node{Storage: fmt.Sprintf("gitaly-%d", i), Address: "unix://" + socket}
 	}
 
@@ -479,8 +469,7 @@ func TestNodeStatus_IsHealthy(t *testing.T) {
 	socket := testhelper.GetTemporaryGitalySocketFileName(t)
 	address := "unix://" + socket
 
-	srv, healthSrv := testhelper.NewServerWithHealth(t, socket)
-	defer srv.Stop()
+	healthSrv := testhelper.NewServerWithHealth(t, socket)
 
 	clientConn, err := client.Dial(address, nil)
 	require.NoError(t, err)
diff --git a/internal/praefect/nodes/sql_elector_test.go b/internal/praefect/nodes/sql_elector_test.go
index 025df6588..c8b203f54 100644
--- a/internal/praefect/nodes/sql_elector_test.go
+++ b/internal/praefect/nodes/sql_elector_test.go
@@ -40,8 +40,7 @@ func TestGetPrimaryAndSecondaries(t *testing.T) {
 	}
 
 	internalSocket0 := testhelper.GetTemporaryGitalySocketFileName(t)
-	srv0, _ := testhelper.NewServerWithHealth(t, internalSocket0)
-	defer srv0.Stop()
+	testhelper.NewServerWithHealth(t, internalSocket0)
 
 	cc0, err := grpc.Dial(
 		"unix://"+internalSocket0,
@@ -77,8 +76,7 @@ func TestSqlElector_slow_execution(t *testing.T) {
 	logger := testhelper.NewTestLogger(t).WithField("test", t.Name())
 
 	gitalySocket := testhelper.GetTemporaryGitalySocketFileName(t)
-	gitalySrv, _ := testhelper.NewServerWithHealth(t, gitalySocket)
-	defer gitalySrv.Stop()
+	testhelper.NewServerWithHealth(t, gitalySocket)
 
 	gitalyConn, err := grpc.Dial(
 		"unix://"+gitalySocket,
@@ -118,11 +116,8 @@ func TestBasicFailover(t *testing.T) {
 	conf := config.Config{SocketPath: socketName}
 
 	internalSocket0, internalSocket1 := testhelper.GetTemporaryGitalySocketFileName(t), testhelper.GetTemporaryGitalySocketFileName(t)
-	srv0, healthSrv0 := testhelper.NewServerWithHealth(t, internalSocket0)
-	defer srv0.Stop()
-
-	srv1, healthSrv1 := testhelper.NewServerWithHealth(t, internalSocket1)
-	defer srv1.Stop()
+	healthSrv0 := testhelper.NewServerWithHealth(t, internalSocket0)
+	healthSrv1 := testhelper.NewServerWithHealth(t, internalSocket1)
 
 	addr0 := "unix://" + internalSocket0
 	cc0, err := grpc.Dial(
diff --git a/internal/praefect/server_test.go b/internal/praefect/server_test.go
index 991b54a76..7726fd544 100644
--- a/internal/praefect/server_test.go
+++ b/internal/praefect/server_test.go
@@ -198,7 +198,7 @@ func TestGitalyServerInfo(t *testing.T) {
 
 func TestGitalyServerInfoBadNode(t *testing.T) {
 	gitalySocket := testhelper.GetTemporaryGitalySocketFileName(t)
-	_, healthSrv := testhelper.NewServerWithHealth(t, gitalySocket)
+	healthSrv := testhelper.NewServerWithHealth(t, gitalySocket)
 	healthSrv.SetServingStatus("", grpc_health_v1.HealthCheckResponse_UNKNOWN)
 
 	conf := config.Config{
diff --git a/internal/testhelper/testserver.go b/internal/testhelper/testserver.go
index 376c36ec4..2a6552159 100644
--- a/internal/testhelper/testserver.go
+++ b/internal/testhelper/testserver.go
@@ -902,26 +902,26 @@ type HTTPSettings struct {
 	Password string `yaml:"password"`
 }
 
-// NewServerWithHealth creates a new GRPC server with the health server set up.
+// NewServerWithHealth creates a new gRPC server with the health server set up.
 // It will listen on the socket identified by `socketName`.
-func NewServerWithHealth(t testing.TB, socketName string) (*grpc.Server, *health.Server) {
+func NewServerWithHealth(t testing.TB, socketName string) *health.Server {
 	lis, err := net.Listen("unix", socketName)
 	require.NoError(t, err)
 
 	return NewHealthServerWithListener(t, lis)
 }
 
-// NewHealthServerWithListener creates a new GRPC server with the health server
+// NewHealthServerWithListener creates a new gRPC server with the health server
 // set up. It will listen on the given listener.
-func NewHealthServerWithListener(t testing.TB, listener net.Listener) (*grpc.Server, *health.Server) {
-	srv := NewTestGrpcServer(t, nil, nil)
+func NewHealthServerWithListener(t testing.TB, listener net.Listener) *health.Server {
+	srv := grpc.NewServer()
 	healthSrvr := health.NewServer()
 	healthpb.RegisterHealthServer(srv, healthSrvr)
-	healthSrvr.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
 
-	go srv.Serve(listener)
+	t.Cleanup(srv.Stop)
+	go func() { require.NoError(t, srv.Serve(listener)) }()
 
-	return srv, healthSrvr
+	return healthSrvr
 }
 
 // SetupAndStartGitlabServer creates a new GitlabTestServer, starts it and sets