Expose our custom DNS resolver as a gRPC client dial option

gRPC depends on the target's scheme to determine which resolver to use.
Built-in DNS Resolver registers itself with "dns" scheme. We should use
a different scheme for this resolver. However, Ruby, and other
cares-based clients, don't support custom resolver. At GitLab, the gRPC
target configuration is shared between components. To ensure the
compatibility between clients, this custom resolver use same "dns"
scheme.

We don't want to mess up the default global registry. Clients are
expected to inject this resolver via exposed WithGitalyDNSResolver dial
option.

2 files changed, 118 insertions, 6 deletions

diff --git a/client/dial.go b/client/dial.go
index 4b46049a0..56125e6da 100644
--- a/client/dial.go
+++ b/client/dial.go
@@ -2,7 +2,12 @@ package client
 
 import (
 	"context"
+	"math/rand"
+	"time"
 
+	"github.com/sirupsen/logrus"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/backoff"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/dnsresolver"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/client"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/sidechannel"
 	"google.golang.org/grpc"
@@ -65,3 +70,24 @@ func HealthCheckDialer(base Dialer) Dialer {
 		return cc, nil
 	}
 }
+
+// DNSResolverBuilderConfig exposes the DNS resolver builder option. It is used to build Gitaly
+// custom DNS resolver.
+type DNSResolverBuilderConfig dnsresolver.BuilderConfig
+
+// DefaultDNSResolverBuilderConfig returns the default options for building DNS resolver.
+func DefaultDNSResolverBuilderConfig() *DNSResolverBuilderConfig {
+	return &DNSResolverBuilderConfig{
+		RefreshRate:     5 * time.Minute,
+		Logger:          logrus.StandardLogger(),
+		Backoff:         backoff.NewDefaultExponential(rand.New(rand.NewSource(time.Now().UnixNano()))),
+		DefaultGrpcPort: "443",
+	}
+}
+
+// WithGitalyDNSResolver defines a gRPC dial option for injecting Gitaly's custom DNS resolver. This
+// resolver watches for the changes of target URL periodically and update the target subchannels
+// accordingly.
+func WithGitalyDNSResolver(opts *DNSResolverBuilderConfig) grpc.DialOption {
+	return grpc.WithResolvers(dnsresolver.NewBuilder((*dnsresolver.BuilderConfig)(opts)))
+}
diff --git a/client/dial_test.go b/client/dial_test.go
index bd7b8423d..37339ade4 100644
--- a/client/dial_test.go
+++ b/client/dial_test.go
@@ -13,6 +13,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/miekg/dns"
 	"github.com/opentracing/opentracing-go"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
@@ -22,6 +23,7 @@ import (
 	internalclient "gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/client"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
 	gitalyx509 "gitlab.com/gitlab-org/gitaly/v15/internal/x509"
+	"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
 	"gitlab.com/gitlab-org/labkit/correlation"
 	grpccorrelation "gitlab.com/gitlab-org/labkit/correlation/grpc"
 	grpctracing "gitlab.com/gitlab-org/labkit/tracing/grpc"
@@ -136,7 +138,8 @@ func TestDial(t *testing.T) {
 
 			ctx := testhelper.Context(t)
 
-			conn, err := Dial(tt.rawAddress, tt.dialOpts)
+			dialOpts := append(tt.dialOpts, WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()))
+			conn, err := Dial(tt.rawAddress, dialOpts)
 			if tt.expectDialFailure {
 				require.Error(t, err)
 				return
@@ -225,7 +228,8 @@ func TestDialSidechannel(t *testing.T) {
 
 			ctx := testhelper.Context(t)
 
-			conn, err := DialSidechannel(ctx, tt.rawAddress, registry, tt.dialOpts)
+			dialOpts := append(tt.dialOpts, WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()))
+			conn, err := DialSidechannel(ctx, tt.rawAddress, registry, dialOpts)
 			require.NoError(t, err)
 			defer testhelper.MustClose(t, conn)
 
@@ -302,7 +306,9 @@ func TestDial_Correlation(t *testing.T) {
 		ctx := testhelper.Context(t)
 
 		cc, err := DialContext(ctx, "unix://"+serverSocketPath, []grpc.DialOption{
-			internalclient.UnaryInterceptor(), internalclient.StreamInterceptor(),
+			internalclient.UnaryInterceptor(),
+			internalclient.StreamInterceptor(),
+			WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()),
 		})
 		require.NoError(t, err)
 		defer testhelper.MustClose(t, cc)
@@ -337,7 +343,9 @@ func TestDial_Correlation(t *testing.T) {
 		ctx := testhelper.Context(t)
 
 		cc, err := DialContext(ctx, "unix://"+serverSocketPath, []grpc.DialOption{
-			internalclient.UnaryInterceptor(), internalclient.StreamInterceptor(),
+			internalclient.UnaryInterceptor(),
+			internalclient.StreamInterceptor(),
+			WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()),
 		})
 		require.NoError(t, err)
 		defer testhelper.MustClose(t, cc)
@@ -410,7 +418,9 @@ func TestDial_Tracing(t *testing.T) {
 		// This needs to be run after setting up the global tracer as it will cause us to
 		// create the span when executing the RPC call further down below.
 		cc, err := DialContext(ctx, "unix://"+serverSocketPath, []grpc.DialOption{
-			internalclient.UnaryInterceptor(), internalclient.StreamInterceptor(),
+			internalclient.UnaryInterceptor(),
+			internalclient.StreamInterceptor(),
+			WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()),
 		})
 		require.NoError(t, err)
 		defer testhelper.MustClose(t, cc)
@@ -469,7 +479,9 @@ func TestDial_Tracing(t *testing.T) {
 		// This needs to be run after setting up the global tracer as it will cause us to
 		// create the span when executing the RPC call further down below.
 		cc, err := DialContext(ctx, "unix://"+serverSocketPath, []grpc.DialOption{
-			internalclient.UnaryInterceptor(), internalclient.StreamInterceptor(),
+			internalclient.UnaryInterceptor(),
+			internalclient.StreamInterceptor(),
+			WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()),
 		})
 		require.NoError(t, err)
 		defer testhelper.MustClose(t, cc)
@@ -640,11 +652,85 @@ func TestHealthCheckDialer(t *testing.T) {
 		grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2("token")),
 		internalclient.UnaryInterceptor(),
 		internalclient.StreamInterceptor(),
+		WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()),
 	})
 	require.NoError(t, err)
 	require.NoError(t, cc.Close())
 }
 
+func TestWithGitalyDNSResolver_successfully(t *testing.T) {
+	t.Parallel()
+
+	serverURL := startFakeGitalyServer(t)
+	serverHost, serverPort, err := net.SplitHostPort(serverURL)
+	require.NoError(t, err)
+
+	dnsServer := testhelper.NewFakeDNSServer(t).WithHandler(dns.TypeA, func(host string) []string {
+		if host == "grpc.test." {
+			return []string{serverHost}
+		}
+		return nil
+	}).Start()
+
+	// This scheme uses our DNS resolver
+	target := fmt.Sprintf("dns://%s/grpc.test:%s", dnsServer.Addr(), serverPort)
+	conn, err := grpc.Dial(
+		target,
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()),
+	)
+	require.NoError(t, err)
+	defer testhelper.MustClose(t, conn)
+
+	client := gitalypb.NewCommitServiceClient(conn)
+	for i := 0; i < 10; i++ {
+		_, err = client.FindCommit(testhelper.Context(t), &gitalypb.FindCommitRequest{})
+		require.NoError(t, err)
+	}
+}
+
+func TestWithGitalyDNSResolver_zeroAddresses(t *testing.T) {
+	t.Parallel()
+
+	dnsServer := testhelper.NewFakeDNSServer(t).WithHandler(dns.TypeA, func(host string) []string {
+		return nil
+	}).Start()
+
+	// This scheme uses our DNS resolver
+	target := fmt.Sprintf("dns://%s/grpc.test:50051", dnsServer.Addr())
+	conn, err := grpc.Dial(
+		target,
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		WithGitalyDNSResolver(DefaultDNSResolverBuilderConfig()),
+	)
+	require.NoError(t, err)
+	defer testhelper.MustClose(t, conn)
+
+	client := gitalypb.NewCommitServiceClient(conn)
+	_, err = client.FindCommit(testhelper.Context(t), &gitalypb.FindCommitRequest{})
+	require.Equal(t, err, status.Error(codes.Unavailable, "name resolver error: produced zero addresses"))
+}
+
+type fakeCommitServer struct {
+	gitalypb.UnimplementedCommitServiceServer
+}
+
+func (s *fakeCommitServer) FindCommit(_ context.Context, _ *gitalypb.FindCommitRequest) (*gitalypb.FindCommitResponse, error) {
+	return &gitalypb.FindCommitResponse{}, nil
+}
+
+func startFakeGitalyServer(t *testing.T) string {
+	listener, err := net.Listen("tcp", "localhost:0")
+	require.NoError(t, err)
+
+	srv := grpc.NewServer(SidechannelServer(newLogger(t), insecure.NewCredentials()))
+	gitalypb.RegisterCommitServiceServer(srv, &fakeCommitServer{})
+	go testhelper.MustServe(t, srv, listener)
+	t.Cleanup(srv.Stop)
+
+	return listener.Addr().String()
+}
+
 func newLogger(tb testing.TB) *logrus.Entry {
 	return logrus.NewEntry(testhelper.NewDiscardingLogger(tb))
 }