Verify connections are properly configured with HealthCheckDialer

There is no way we could identify if the connection we establish
to remote is properly configured and could be used with no issues.
As a workaround on it is to issue a request to the remote when
connection is set. The simplest approach is to issue a health check
request as both gitaly and praefect expose this service outside.
As we don't want to make it after each call to Dial as it could be
missed or if we change implementation later it makes sense to define
a custom Dialer for that purpose. It should be used with the pools
we are using except of praefect connection pool to gitaly nodes,
as it would prevent praefect from start if one of gitaly nodes is not
accessible.

Closes: https://gitlab.com/gitlab-org/gitaly/-/issues/2699

2 files changed, 84 insertions, 39 deletions

diff --git a/client/dial.go b/client/dial.go
index 60be802cf..341f1cf27 100644
--- a/client/dial.go
+++ b/client/dial.go
@@ -12,6 +12,7 @@ import (
 	grpctracing "gitlab.com/gitlab-org/labkit/tracing/grpc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/keepalive"
 )
 
@@ -130,3 +131,21 @@ func FailOnNonTempDialError() []grpc.DialOption {
 		grpc.FailOnNonTempDialError(true),
 	}
 }
+
+// HealthCheckDialer uses provided dialer as an actual dialer, but issues a health check request to the remote
+// to verify the connection was set properly and could be used with no issues.
+func HealthCheckDialer(base Dialer) Dialer {
+	return func(ctx context.Context, address string, dialOptions []grpc.DialOption) (*grpc.ClientConn, error) {
+		cc, err := base(ctx, address, dialOptions)
+		if err != nil {
+			return nil, err
+		}
+
+		if _, err := healthpb.NewHealthClient(cc).Check(ctx, &healthpb.HealthCheckRequest{}); err != nil {
+			cc.Close()
+			return nil, err
+		}
+
+		return cc, nil
+	}
+}
diff --git a/client/dial_test.go b/client/dial_test.go
index fe496f9d6..9951c7da2 100644
--- a/client/dial_test.go
+++ b/client/dial_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/uber/jaeger-client-go"
+	gitalyauth "gitlab.com/gitlab-org/gitaly/auth"
 	proxytestdata "gitlab.com/gitlab-org/gitaly/internal/praefect/grpc-proxy/testdata"
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 	gitaly_x509 "gitlab.com/gitlab-org/gitaly/internal/x509"
@@ -48,64 +49,71 @@ func TestDial(t *testing.T) {
 	require.NoError(t, os.Symlink(unixSocketAbsPath, unixSocketPath))
 
 	tests := []struct {
-		name           string
-		rawAddress     string
-		envSSLCertFile string
-		dialOpts       []grpc.DialOption
-		expectFailure  bool
+		name                string
+		rawAddress          string
+		envSSLCertFile      string
+		dialOpts            []grpc.DialOption
+		expectDialFailure   bool
+		expectHealthFailure bool
 	}{
 		{
-			name:          "tcp localhost with prefix",
-			rawAddress:    "tcp://localhost:" + connectionMap["tcp"], // "tcp://localhost:1234"
-			expectFailure: false,
+			name:                "tcp localhost with prefix",
+			rawAddress:          "tcp://localhost:" + connectionMap["tcp"], // "tcp://localhost:1234"
+			expectDialFailure:   false,
+			expectHealthFailure: false,
 		},
 		{
-			name:           "tls localhost",
-			rawAddress:     "tls://localhost:" + connectionMap["tls"], // "tls://localhost:1234"
-			envSSLCertFile: "./testdata/gitalycert.pem",
-			expectFailure:  false,
+			name:                "tls localhost",
+			rawAddress:          "tls://localhost:" + connectionMap["tls"], // "tls://localhost:1234"
+			envSSLCertFile:      "./testdata/gitalycert.pem",
+			expectDialFailure:   false,
+			expectHealthFailure: false,
 		},
 		{
-			name:          "unix absolute",
-			rawAddress:    "unix:" + unixSocketAbsPath, // "unix:/tmp/temp-socket"
-			expectFailure: false,
+			name:                "unix absolute",
+			rawAddress:          "unix:" + unixSocketAbsPath, // "unix:/tmp/temp-socket"
+			expectDialFailure:   false,
+			expectHealthFailure: false,
 		},
 		{
-			name:          "unix relative",
-			rawAddress:    "unix:" + unixSocketPath, // "unix:../../tmp/temp-socket"
-			expectFailure: false,
+			name:                "unix relative",
+			rawAddress:          "unix:" + unixSocketPath, // "unix:../../tmp/temp-socket"
+			expectDialFailure:   false,
+			expectHealthFailure: false,
 		},
 		{
-			name:          "unix absolute does not exist",
-			rawAddress:    "unix:" + unixSocketAbsPath + ".does_not_exist", // "unix:/tmp/temp-socket.does_not_exist"
-			expectFailure: true,
+			name:                "unix absolute does not exist",
+			rawAddress:          "unix:" + unixSocketAbsPath + ".does_not_exist", // "unix:/tmp/temp-socket.does_not_exist"
+			expectDialFailure:   false,
+			expectHealthFailure: true,
 		},
 		{
-			name:          "unix relative does not exist",
-			rawAddress:    "unix:" + unixSocketPath + ".does_not_exist", // "unix:../../tmp/temp-socket.does_not_exist"
-			expectFailure: true,
+			name:                "unix relative does not exist",
+			rawAddress:          "unix:" + unixSocketPath + ".does_not_exist", // "unix:../../tmp/temp-socket.does_not_exist"
+			expectDialFailure:   false,
+			expectHealthFailure: true,
 		},
 		{
 			// Gitaly does not support connections that do not have a scheme.
-			name:          "tcp localhost no prefix",
-			rawAddress:    "localhost:" + connectionMap["tcp"], // "localhost:1234"
-			expectFailure: true,
+			name:              "tcp localhost no prefix",
+			rawAddress:        "localhost:" + connectionMap["tcp"], // "localhost:1234"
+			expectDialFailure: true,
 		},
 		{
-			name:          "invalid",
-			rawAddress:    ".",
-			expectFailure: true,
+			name:              "invalid",
+			rawAddress:        ".",
+			expectDialFailure: true,
 		},
 		{
-			name:          "empty",
-			rawAddress:    "",
-			expectFailure: true,
+			name:              "empty",
+			rawAddress:        "",
+			expectDialFailure: true,
 		},
 		{
-			name:          "dial fail if there is no listener on address",
-			rawAddress:    "tcp://invalid.address",
-			dialOpts:      FailOnNonTempDialError(),
-			expectFailure: true,
+			name:              "dial fail if there is no listener on address",
+			rawAddress:        "tcp://invalid.address",
+			dialOpts:          FailOnNonTempDialError(),
+			expectDialFailure: true,
 		},
 	}
 
@@ -123,15 +131,18 @@ func TestDial(t *testing.T) {
 			defer cancel()
 
 			conn, err := Dial(tt.rawAddress, tt.dialOpts)
-			if tt.expectFailure {
+			if tt.expectDialFailure {
 				require.Error(t, err)
 				return
 			}
-
 			require.NoError(t, err)
 			defer conn.Close()
 
 			_, err = healthpb.NewHealthClient(conn).Check(ctx, &healthpb.HealthCheckRequest{})
+			if tt.expectHealthFailure {
+				require.Error(t, err)
+				return
+			}
 			require.NoError(t, err)
 		})
 	}
@@ -459,3 +470,18 @@ func emitProxyWarning() bool {
 	}
 	return false
 }
+
+func TestHealthCheckDialer(t *testing.T) {
+	_, addr, cleanup := runServer(t, "token")
+	defer cleanup()
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	_, err := HealthCheckDialer(DialContext)(ctx, addr, nil)
+	require.Equal(t, status.Error(codes.Unauthenticated, "authentication required"), err, "should fail without token configured")
+
+	cc, err := HealthCheckDialer(DialContext)(ctx, addr, []grpc.DialOption{grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2("token"))})
+	require.NoError(t, err)
+	cc.Close()
+}