Add DNS discovery support for Gitaly/Praefect

All the implementations of DNS discovery were done in this epic:
https://gitlab.com/groups/gitlab-org/-/epics/8971. Gitaly allows clients
to configure DNS discovery via dial option. This MR adds the exposed
dial options to client connection creation in Gitlab-shell.

Issue: https://gitlab.com/gitlab-org/gitaly/-/issues/4722
Changelog: added

5 files changed, 212 insertions, 163 deletions

diff --git a/client/testserver/gitalyserver.go b/client/testserver/gitalyserver.go
index d98ac01..aeaa48d 100644
--- a/client/testserver/gitalyserver.go
+++ b/client/testserver/gitalyserver.go
@@ -85,30 +85,51 @@ func (s *TestGitalyServer) SSHUploadArchive(stream pb.SSHService_SSHUploadArchiv
 	return nil
 }
 
-func StartGitalyServer(t *testing.T) (string, *TestGitalyServer) {
+func StartGitalyServer(t *testing.T, network string) (string, *TestGitalyServer) {
 	t.Helper()
 
-	tempDir, _ := os.MkdirTemp("", "gitlab-shell-test-api")
-	gitalySocketPath := path.Join(tempDir, "gitaly.sock")
-	t.Cleanup(func() { os.RemoveAll(tempDir) })
+	switch network {
+	case "unix":
+		tempDir, _ := os.MkdirTemp("", "gitlab-shell-test-api")
+		gitalySocketPath := path.Join(tempDir, "gitaly.sock")
+		t.Cleanup(func() { require.NoError(t, os.RemoveAll(tempDir)) })
 
-	err := os.MkdirAll(filepath.Dir(gitalySocketPath), 0700)
-	require.NoError(t, err)
+		err := os.MkdirAll(filepath.Dir(gitalySocketPath), 0700)
+		require.NoError(t, err)
+
+		addr, testServer := doStartTestServer(t, "unix", gitalySocketPath)
+		return fmt.Sprintf("unix:%s", addr), testServer
+
+	case "tcp":
+		addr, testServer := doStartTestServer(t, "tcp", "127.0.0.1:0")
+		return fmt.Sprintf("tcp://%s", addr), testServer
+
+	case "dns":
+		addr, testServer := doStartTestServer(t, "tcp", "127.0.0.1:0")
+		// gRPC URL with DNS scheme follows this format: https://grpc.github.io/grpc/core/md_doc_naming.html
+		// When the authority is dropped, the URL have 3 splashes.
+		return fmt.Sprintf("dns:///%s", addr), testServer
 
+	default:
+		panic(fmt.Sprintf("Unsupported network %s", network))
+	}
+}
+
+func doStartTestServer(t *testing.T, network string, path string) (string, *TestGitalyServer) {
 	server := grpc.NewServer(
 		client.SidechannelServer(log.ContextLogger(context.Background()), insecure.NewCredentials()),
 	)
 
-	listener, err := net.Listen("unix", gitalySocketPath)
+	listener, err := net.Listen(network, path)
 	require.NoError(t, err)
 
 	testServer := TestGitalyServer{}
 	pb.RegisterSSHServiceServer(server, &testServer)
 
-	go server.Serve(listener)
+	go func() {
+		require.NoError(t, server.Serve(listener))
+	}()
 	t.Cleanup(func() { server.Stop() })
 
-	gitalySocketUrl := "unix:" + gitalySocketPath
-
-	return gitalySocketUrl, &testServer
+	return listener.Addr().String(), &testServer
 }
diff --git a/internal/command/receivepack/gitalycall_test.go b/internal/command/receivepack/gitalycall_test.go
index eadc43b..9f70189 100644
--- a/internal/command/receivepack/gitalycall_test.go
+++ b/internal/command/receivepack/gitalycall_test.go
@@ -3,6 +3,7 @@ package receivepack
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -17,80 +18,85 @@ import (
 )
 
 func TestReceivePack(t *testing.T) {
-	gitalyAddress, testServer := testserver.StartGitalyServer(t)
+	for _, network := range []string{"unix", "tcp", "dns"} {
+		t.Run(fmt.Sprintf("via %s network", network), func(t *testing.T) {
+			gitalyAddress, testServer := testserver.StartGitalyServer(t, network)
+			t.Log(fmt.Sprintf("Server address: %s", gitalyAddress))
 
-	requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
-	url := testserver.StartHttpServer(t, requests)
+			requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
+			url := testserver.StartHttpServer(t, requests)
 
-	testCases := []struct {
-		username string
-		keyId    string
-	}{
-		{
-			username: "john.doe",
-		},
-		{
-			keyId: "123",
-		},
-	}
+			testCases := []struct {
+				username string
+				keyId    string
+			}{
+				{
+					username: "john.doe",
+				},
+				{
+					keyId: "123",
+				},
+			}
 
-	for _, tc := range testCases {
-		output := &bytes.Buffer{}
-		input := &bytes.Buffer{}
-		repo := "group/repo"
+			for _, tc := range testCases {
+				output := &bytes.Buffer{}
+				input := &bytes.Buffer{}
+				repo := "group/repo"
 
-		env := sshenv.Env{
-			IsSSHConnection: true,
-			OriginalCommand: "git-receive-pack " + repo,
-			RemoteAddr:      "127.0.0.1",
-		}
+				env := sshenv.Env{
+					IsSSHConnection: true,
+					OriginalCommand: "git-receive-pack " + repo,
+					RemoteAddr:      "127.0.0.1",
+				}
 
-		args := &commandargs.Shell{
-			CommandType: commandargs.ReceivePack,
-			SshArgs:     []string{"git-receive-pack", repo},
-			Env:         env,
-		}
+				args := &commandargs.Shell{
+					CommandType: commandargs.ReceivePack,
+					SshArgs:     []string{"git-receive-pack", repo},
+					Env:         env,
+				}
 
-		if tc.username != "" {
-			args.GitlabUsername = tc.username
-		} else {
-			args.GitlabKeyId = tc.keyId
-		}
+				if tc.username != "" {
+					args.GitlabUsername = tc.username
+				} else {
+					args.GitlabKeyId = tc.keyId
+				}
 
-		cfg := &config.Config{GitlabUrl: url}
-		cfg.GitalyClient.InitSidechannelRegistry(context.Background())
-		cmd := &Command{
-			Config:     cfg,
-			Args:       args,
-			ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
-		}
+				cfg := &config.Config{GitlabUrl: url}
+				cfg.GitalyClient.InitSidechannelRegistry(context.Background())
+				cmd := &Command{
+					Config:     cfg,
+					Args:       args,
+					ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
+				}
 
-		ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
-		ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
+				ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
+				ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
 
-		err := cmd.Execute(ctx)
-		require.NoError(t, err)
+				err := cmd.Execute(ctx)
+				require.NoError(t, err)
 
-		if tc.username != "" {
-			require.Equal(t, "ReceivePack: 1 "+repo, output.String())
-		} else {
-			require.Equal(t, "ReceivePack: key-123 "+repo, output.String())
-		}
+				if tc.username != "" {
+					require.Equal(t, "ReceivePack: 1 "+repo, output.String())
+				} else {
+					require.Equal(t, "ReceivePack: key-123 "+repo, output.String())
+				}
 
-		for k, v := range map[string]string{
-			"gitaly-feature-cache_invalidator":        "true",
-			"gitaly-feature-inforef_uploadpack_cache": "false",
-			"x-gitlab-client-name":                    "gitlab-shell-tests-git-receive-pack",
-			"key_id":                                  "123",
-			"user_id":                                 "1",
-			"remote_ip":                               "127.0.0.1",
-			"key_type":                                "key",
-		} {
-			actual := testServer.ReceivedMD[k]
-			require.Len(t, actual, 1)
-			require.Equal(t, v, actual[0])
-		}
-		require.Empty(t, testServer.ReceivedMD["some-other-ff"])
-		require.Equal(t, testServer.ReceivedMD["x-gitlab-correlation-id"][0], "a-correlation-id")
+				for k, v := range map[string]string{
+					"gitaly-feature-cache_invalidator":        "true",
+					"gitaly-feature-inforef_uploadpack_cache": "false",
+					"x-gitlab-client-name":                    "gitlab-shell-tests-git-receive-pack",
+					"key_id":                                  "123",
+					"user_id":                                 "1",
+					"remote_ip":                               "127.0.0.1",
+					"key_type":                                "key",
+				} {
+					actual := testServer.ReceivedMD[k]
+					require.Len(t, actual, 1)
+					require.Equal(t, v, actual[0])
+				}
+				require.Empty(t, testServer.ReceivedMD["some-other-ff"])
+				require.Equal(t, testServer.ReceivedMD["x-gitlab-correlation-id"][0], "a-correlation-id")
+			}
+		})
 	}
 }
diff --git a/internal/command/uploadarchive/gitalycall_test.go b/internal/command/uploadarchive/gitalycall_test.go
index be6032d..0479e30 100644
--- a/internal/command/uploadarchive/gitalycall_test.go
+++ b/internal/command/uploadarchive/gitalycall_test.go
@@ -3,6 +3,7 @@ package uploadarchive
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -17,59 +18,64 @@ import (
 )
 
 func TestUploadArchive(t *testing.T) {
-	gitalyAddress, testServer := testserver.StartGitalyServer(t)
+	for _, network := range []string{"unix", "tcp", "dns"} {
+		t.Run(fmt.Sprintf("via %s network", network), func(t *testing.T) {
+			gitalyAddress, testServer := testserver.StartGitalyServer(t, network)
+			t.Log(fmt.Sprintf("Server address: %s", gitalyAddress))
 
-	requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
-	url := testserver.StartHttpServer(t, requests)
+			requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
+			url := testserver.StartHttpServer(t, requests)
 
-	output := &bytes.Buffer{}
-	input := &bytes.Buffer{}
+			output := &bytes.Buffer{}
+			input := &bytes.Buffer{}
 
-	userId := "1"
-	repo := "group/repo"
+			userId := "1"
+			repo := "group/repo"
 
-	env := sshenv.Env{
-		IsSSHConnection: true,
-		OriginalCommand: "git-upload-archive " + repo,
-		RemoteAddr:      "127.0.0.1",
-	}
+			env := sshenv.Env{
+				IsSSHConnection: true,
+				OriginalCommand: "git-upload-archive " + repo,
+				RemoteAddr:      "127.0.0.1",
+			}
 
-	args := &commandargs.Shell{
-		GitlabKeyId: userId,
-		CommandType: commandargs.UploadArchive,
-		SshArgs:     []string{"git-upload-archive", repo},
-		Env:         env,
-	}
+			args := &commandargs.Shell{
+				GitlabKeyId: userId,
+				CommandType: commandargs.UploadArchive,
+				SshArgs:     []string{"git-upload-archive", repo},
+				Env:         env,
+			}
 
-	cfg := &config.Config{GitlabUrl: url}
-	cfg.GitalyClient.InitSidechannelRegistry(context.Background())
-	cmd := &Command{
-		Config:     cfg,
-		Args:       args,
-		ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
-	}
+			cfg := &config.Config{GitlabUrl: url}
+			cfg.GitalyClient.InitSidechannelRegistry(context.Background())
+			cmd := &Command{
+				Config:     cfg,
+				Args:       args,
+				ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
+			}
 
-	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
-	ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
+			ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
+			ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
 
-	err := cmd.Execute(ctx)
-	require.NoError(t, err)
+			err := cmd.Execute(ctx)
+			require.NoError(t, err)
 
-	require.Equal(t, "UploadArchive: "+repo, output.String())
+			require.Equal(t, "UploadArchive: "+repo, output.String())
 
-	for k, v := range map[string]string{
-		"gitaly-feature-cache_invalidator":        "true",
-		"gitaly-feature-inforef_uploadpack_cache": "false",
-		"x-gitlab-client-name":                    "gitlab-shell-tests-git-upload-archive",
-		"key_id":                                  "123",
-		"user_id":                                 "1",
-		"remote_ip":                               "127.0.0.1",
-		"key_type":                                "key",
-	} {
-		actual := testServer.ReceivedMD[k]
-		require.Len(t, actual, 1)
-		require.Equal(t, v, actual[0])
+			for k, v := range map[string]string{
+				"gitaly-feature-cache_invalidator":        "true",
+				"gitaly-feature-inforef_uploadpack_cache": "false",
+				"x-gitlab-client-name":                    "gitlab-shell-tests-git-upload-archive",
+				"key_id":                                  "123",
+				"user_id":                                 "1",
+				"remote_ip":                               "127.0.0.1",
+				"key_type":                                "key",
+			} {
+				actual := testServer.ReceivedMD[k]
+				require.Len(t, actual, 1)
+				require.Equal(t, v, actual[0])
+			}
+			require.Empty(t, testServer.ReceivedMD["some-other-ff"])
+			require.Equal(t, testServer.ReceivedMD["x-gitlab-correlation-id"][0], "a-correlation-id")
+		})
 	}
-	require.Empty(t, testServer.ReceivedMD["some-other-ff"])
-	require.Equal(t, testServer.ReceivedMD["x-gitlab-correlation-id"][0], "a-correlation-id")
 }
diff --git a/internal/command/uploadpack/gitalycall_test.go b/internal/command/uploadpack/gitalycall_test.go
index 6b6f009..874d12e 100644
--- a/internal/command/uploadpack/gitalycall_test.go
+++ b/internal/command/uploadpack/gitalycall_test.go
@@ -3,6 +3,7 @@ package uploadpack
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -17,60 +18,65 @@ import (
 )
 
 func TestUploadPack(t *testing.T) {
-	gitalyAddress, testServer := testserver.StartGitalyServer(t)
+	for _, network := range []string{"unix", "tcp", "dns"} {
+		t.Run(fmt.Sprintf("via %s network", network), func(t *testing.T) {
+			gitalyAddress, testServer := testserver.StartGitalyServer(t, network)
+			t.Log(fmt.Sprintf("Server address: %s", gitalyAddress))
 
-	requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
-	url := testserver.StartHttpServer(t, requests)
+			requests := requesthandlers.BuildAllowedWithGitalyHandlers(t, gitalyAddress)
+			url := testserver.StartHttpServer(t, requests)
 
-	output := &bytes.Buffer{}
-	input := &bytes.Buffer{}
+			output := &bytes.Buffer{}
+			input := &bytes.Buffer{}
 
-	userId := "1"
-	repo := "group/repo"
+			userId := "1"
+			repo := "group/repo"
 
-	env := sshenv.Env{
-		IsSSHConnection: true,
-		OriginalCommand: "git-upload-pack " + repo,
-		RemoteAddr:      "127.0.0.1",
-	}
+			env := sshenv.Env{
+				IsSSHConnection: true,
+				OriginalCommand: "git-upload-pack " + repo,
+				RemoteAddr:      "127.0.0.1",
+			}
 
-	args := &commandargs.Shell{
-		GitlabKeyId: userId,
-		CommandType: commandargs.UploadPack,
-		SshArgs:     []string{"git-upload-pack", repo},
-		Env:         env,
-	}
+			args := &commandargs.Shell{
+				GitlabKeyId: userId,
+				CommandType: commandargs.UploadPack,
+				SshArgs:     []string{"git-upload-pack", repo},
+				Env:         env,
+			}
 
-	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
-	ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
+			ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
+			ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
 
-	cfg := &config.Config{GitlabUrl: url}
-	cfg.GitalyClient.InitSidechannelRegistry(ctx)
+			cfg := &config.Config{GitlabUrl: url}
+			cfg.GitalyClient.InitSidechannelRegistry(ctx)
 
-	cmd := &Command{
-		Config:     cfg,
-		Args:       args,
-		ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
-	}
+			cmd := &Command{
+				Config:     cfg,
+				Args:       args,
+				ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
+			}
 
-	err := cmd.Execute(ctx)
-	require.NoError(t, err)
+			err := cmd.Execute(ctx)
+			require.NoError(t, err)
 
-	require.Equal(t, "SSHUploadPackWithSidechannel: "+repo, output.String())
+			require.Equal(t, "SSHUploadPackWithSidechannel: "+repo, output.String())
 
-	for k, v := range map[string]string{
-		"gitaly-feature-cache_invalidator":        "true",
-		"gitaly-feature-inforef_uploadpack_cache": "false",
-		"x-gitlab-client-name":                    "gitlab-shell-tests-git-upload-pack",
-		"key_id":                                  "123",
-		"user_id":                                 "1",
-		"remote_ip":                               "127.0.0.1",
-		"key_type":                                "key",
-	} {
-		actual := testServer.ReceivedMD[k]
-		require.Len(t, actual, 1)
-		require.Equal(t, v, actual[0])
+			for k, v := range map[string]string{
+				"gitaly-feature-cache_invalidator":        "true",
+				"gitaly-feature-inforef_uploadpack_cache": "false",
+				"x-gitlab-client-name":                    "gitlab-shell-tests-git-upload-pack",
+				"key_id":                                  "123",
+				"user_id":                                 "1",
+				"remote_ip":                               "127.0.0.1",
+				"key_type":                                "key",
+			} {
+				actual := testServer.ReceivedMD[k]
+				require.Len(t, actual, 1)
+				require.Equal(t, v, actual[0])
+			}
+			require.Empty(t, testServer.ReceivedMD["some-other-ff"])
+			require.Equal(t, testServer.ReceivedMD["x-gitlab-correlation-id"][0], "a-correlation-id")
+		})
 	}
-	require.Empty(t, testServer.ReceivedMD["some-other-ff"])
-	require.Equal(t, testServer.ReceivedMD["x-gitlab-correlation-id"][0], "a-correlation-id")
 }
diff --git a/internal/gitaly/gitaly.go b/internal/gitaly/gitaly.go
index f83c8b5..7f549cc 100644
--- a/internal/gitaly/gitaly.go
+++ b/internal/gitaly/gitaly.go
@@ -116,6 +116,16 @@ func (c *Client) newConnection(ctx context.Context, cmd Command) (conn *grpc.Cli
 				),
 			),
 		),
+
+		// In https://gitlab.com/groups/gitlab-org/-/epics/8971, we added DNS discovery support to Praefect. This was
+		// done by making two changes:
+		// - Configure client-side round-robin load-balancing in client dial options. We added that as a default option
+		// inside gitaly client in gitaly client since v15.9.0
+		// - Configure DNS resolving. Due to some technical limitations, we don't use gRPC's built-in DNS resolver.
+		// Instead, we implement our own DNS resolver. This resolver is exposed via the following configuration.
+		// Afterward, workhorse can detect and handle DNS discovery automatically. The user needs to setup and set
+		// Gitaly address to something like "dns:gitaly.service.dc1.consul"
+		gitalyclient.WithGitalyDNSResolver(gitalyclient.DefaultDNSResolverBuilderConfig()),
 	)
 
 	if cmd.Token != "" {