Allow using TCP for DB load balancing DNS lookups

Closes https://gitlab.com/gitlab-org/gitlab-ee/issues/8979

Gemfile

@@ -446,3 +446,6 @@ gem 'flipper-active_support_cache_store', '~> 0.13.0'
 # Structured logging
 gem 'lograge', '~> 0.5'
 gem 'grape_logging', '~> 1.7'
+
+# DNS Lookup
+gem 'net-dns', '~> 0.9.0'

Gemfile.lock

@@ -501,6 +501,7 @@ GEM
       mustermann (~> 1.0.0)
     mysql2 (0.4.10)
     nakayoshi_fork (0.0.4)
+    net-dns (0.9.0)
     net-ldap (0.16.0)
     net-ntp (2.1.3)
     net-ssh (5.0.1)
@@ -1093,6 +1094,7 @@ DEPENDENCIES
   minitest (~> 5.11.0)
   mysql2 (~> 0.4.10)
   nakayoshi_fork (~> 0.0.4)
+  net-dns (~> 0.9.0)
   net-ldap
   net-ntp
   net-ssh (~> 5.0)

doc/administration/database_load_balancing.md

@@ -141,6 +141,7 @@ The following options can be set:
 | `port`               | The port of the nameserver.                                                                       | 8600      |
 | `interval`           | The minimum time in seconds between checking the DNS record.                                      | 60        |
 | `disconnect_timeout` | The time in seconds after which an old connection is closed, after the list of hosts was updated. | 120       |
+| `use_tcp`            | Lookup DNS resources using TCP instead of UDP                                                     | false     |
 
 The `interval` value specifies the _minimum_ time between checks. If the A
 record has a TTL greater than this value, then service discovery will honor said
@@ -151,6 +152,10 @@ When the list of hosts is updated, it might take a while for the old connections
 to be terminated. The `disconnect_timeout` setting can be used to enforce an
 upper limit on the time it will take to terminate all old database connections.
 
+Some nameservers (like [Consul][consul-udp]) can return a truncated list of hosts when
+queried over UDP. To overcome this issue, you can use TCP for querying by setting
+`use_tcp` to `true`.
+
 ### Forking
 
 If you use an application server that forks, such as Unicorn, you _have to_
@@ -269,3 +274,4 @@ production:
 [db-req]: ../install/requirements.md#database
 [ee-3526]: https://gitlab.com/gitlab-org/gitlab-ee/merge_requests/3526
 [ee-5883]: https://gitlab.com/gitlab-org/gitlab-ee/merge_requests/5883
+[consul-udp]: https://www.consul.io/docs/agent/dns.html#udp-based-dns-queries

ee/changelogs/unreleased/allow-tcp-for-lb-service-discovery.yml

+---
+title: Allow using TCP for DB load balancing DNS lookups
+merge_request: 8961
+author:
+type: added

ee/lib/gitlab/database/load_balancing.rb

@@ -62,7 +62,8 @@ module Gitlab
           port: conf['port'] || 8600,
           record: conf['record'],
           interval: conf['interval'] || 60,
-          disconnect_timeout: conf['disconnect_timeout'] || 120
+          disconnect_timeout: conf['disconnect_timeout'] || 120,
+          use_tcp: conf['use_tcp'] || false
         }
       end
 

ee/lib/gitlab/database/load_balancing/service_discovery.rb

 # frozen_string_literal: true
 
-require 'resolv'
+require 'net/dns'
 
 module Gitlab
   module Database
@@ -12,21 +12,25 @@ module Gitlab
       # balancer with said hosts. Requests may continue to use the old hosts
       # until they complete.
       class ServiceDiscovery
-        attr_reader :resolver, :interval, :record, :disconnect_timeout
+        attr_reader :interval, :record, :disconnect_timeout
 
         MAX_SLEEP_ADJUSTMENT = 10
 
+        UnresolvableNameserverError = Class.new(StandardError)
+
         # nameserver - The nameserver to use for DNS lookups.
         # port - The port of the nameserver.
         # record - The DNS record to look up for retrieving the secondaries.
         # interval - The time to wait between lookups.
         # disconnect_timeout - The time after which an old host should be
         #                      forcefully disconnected.
-        def initialize(nameserver:, port:, record:, interval: 60, disconnect_timeout: 120)
-          @resolver = Resolv::DNS.new(nameserver_port: [[nameserver, port]])
-          @interval = interval
+        def initialize(nameserver:, port:, record:, interval: 60, disconnect_timeout: 120, use_tcp: false)
+          @nameserver = nameserver
+          @port = port
           @record = record
+          @interval = interval
           @disconnect_timeout = disconnect_timeout
+          @use_tcp = use_tcp
         end
 
         def start
@@ -96,8 +100,7 @@ module Gitlab
         # 1. The time to wait for the next check.
         # 2. An array containing the IP addresses of the DNS record.
         def addresses_from_dns
-          resources =
-            resolver.getresources(record, Resolv::DNS::Resource::IN::A)
+          resources = resolver.search(record, Net::DNS::A).answer
 
           # Addresses are sorted so we can directly compare the old and new
           # addresses, without having to use any additional data structures.
@@ -121,6 +124,28 @@ module Gitlab
         def load_balancer
           LoadBalancing.proxy.load_balancer
         end
+
+        def resolver
+          @resolver ||= Net::DNS::Resolver.new(
+            nameservers: nameserver_ip(@nameserver),
+            port: @port,
+            use_tcp: @use_tcp
+          )
+        end
+
+        private
+
+        def nameserver_ip(nameserver)
+          IPAddr.new(nameserver)
+        rescue IPAddr::InvalidAddressError
+          answer = Net::DNS::Resolver.start(nameserver, Net::DNS::A).answer
+
+          if answer.empty?
+            raise UnresolvableNameserverError, "could not resolve #{nameserver}"
+          end
+
+          answer.first.address
+        end
       end
     end
   end

ee/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb

@@ -6,6 +6,64 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   let(:service) do
     described_class.new(nameserver: 'localhost', port: 8600, record: 'foo')
   end
+  let(:localhost_packet) do
+    resource = double(:resource, address: IPAddr.new('127.0.0.1'))
+    double(:packet, answer: [resource])
+  end
+
+  before do
+    allow(Net::DNS::Resolver).to receive(:start)
+      .with('localhost', Net::DNS::A)
+      .and_return(localhost_packet)
+  end
+
+  describe '#initialize' do
+    context 'when nameserver is not an IP' do
+      it 'sets the resolver nameservers to the IP of the nameserver' do
+        allow(Net::DNS::Resolver).to receive(:start)
+          .with('foo.local', Net::DNS::A)
+          .and_return(localhost_packet)
+
+        service = described_class.new(
+          nameserver: 'foo.local',
+          port: 8600,
+          record: 'bar'
+        )
+        expect(service.resolver.nameservers).to eq(['127.0.0.1'])
+      end
+    end
+
+    context 'when nameserver is an IP' do
+      it 'sets the resolver nameservers to the IP of the nameserver' do
+        service = described_class.new(
+          nameserver: '127.0.0.2',
+          port: 8600,
+          record: 'bar'
+        )
+
+        expect(service.resolver.nameservers).to eq(['127.0.0.2'])
+      end
+    end
+
+    context 'when nameserver is unresolvable' do
+      it 'raises an exception' do
+        allow(Net::DNS::Resolver).to receive(:start)
+          .with('non-existent.localhost', Net::DNS::A)
+          .and_return(double(:packet, answer: []))
+
+        expect do
+          described_class.new(
+            nameserver: 'non-existent.localhost',
+            port: 8600,
+            record: 'bar'
+          ).resolver
+        end.to raise_exception(
+          described_class::UnresolvableNameserverError,
+          'could not resolve non-existent.localhost'
+        )
+      end
+    end
+  end
 
   describe '#start' do
     before do
@@ -130,11 +188,12 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
     it 'returns a TTL and ordered list of IP addresses' do
       res1 = double(:resource, address: '255.255.255.0', ttl: 90)
       res2 = double(:resource, address: '127.0.0.1', ttl: 90)
+      packet = double(:packet, answer: [res1, res2])
 
       allow(service.resolver)
-        .to receive(:getresources)
-        .with('foo', Resolv::DNS::Resource::IN::A)
-        .and_return([res1, res2])
+        .to receive(:search)
+        .with('foo', Net::DNS::A)
+        .and_return(packet)
 
       expect(service.addresses_from_dns)
         .to eq([90, %w[127.0.0.1 255.255.255.0]])

ee/spec/lib/gitlab/database/load_balancing_spec.rb

@@ -232,7 +232,8 @@ describe Gitlab::Database::LoadBalancing do
           port: 8600,
           record: nil,
           interval: 60,
-          disconnect_timeout: 120
+          disconnect_timeout: 120,
+          use_tcp: false
         )
       end
     end
@@ -248,7 +249,8 @@ describe Gitlab::Database::LoadBalancing do
           port: 8600,
           record: 'foo',
           interval: 60,
-          disconnect_timeout: 120
+          disconnect_timeout: 120,
+          use_tcp: false
         )
       end
     end