Update docs to have record_type in discovery

This will allow us to have SRV lookups vs A lookups.

A lookups will remain the default

doc/administration/database_load_balancing.md

@@ -122,6 +122,7 @@ production:
     discover:
       nameserver: localhost
       record: secondary.postgresql.service.consul
+      record_type: A
       port: 8600
       interval: 60
       disconnect_timeout: 120
@@ -137,12 +138,16 @@ The following options can be set:
 | Option               | Description                                                                                       | Default   |
 |----------------------|---------------------------------------------------------------------------------------------------|-----------|
 | `nameserver`         | The nameserver to use for looking up the DNS record.                                              | localhost |
-| `record`             | The A record to look up. This option is required for service discovery to work.                   |           |
+| `record`             | The record to look up. This option is required for service discovery to work.                     |           |
+| `record_type`        | Optional record type to look up, this can be either A or SRV (since GitLab 12.3)                  | A         |
 | `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     |
 
+If `record_type` is set to `SRV`, GitLab will continue to use a round-robin algorithm
+and will ignore the `weight` and `priority` in the record.
+
 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
 TTL. For example, if the TTL of the A record is 90 seconds, then service

ee/changelogs/unreleased/db_load_balancing_service_discovery_srv.yml

+---
+title: 'DB Load Balancing: Support SRV lookups'
+merge_request: 15558
+author:
+type: changed

ee/lib/gitlab/database/load_balancing.rb

@@ -63,6 +63,7 @@ module Gitlab
           nameserver: conf['nameserver'] || 'localhost',
           port: conf['port'] || 8600,
           record: conf['record'],
+          record_type: conf['record_type'] || 'A',
           interval: conf['interval'] || 60,
           disconnect_timeout: conf['disconnect_timeout'] || 120,
           use_tcp: conf['use_tcp'] || false

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

@@ -5,7 +5,7 @@ module Gitlab
     module LoadBalancing
       # A single database host used for load balancing.
       class Host
-        attr_reader :pool, :last_checked_at, :intervals, :load_balancer, :host
+        attr_reader :pool, :last_checked_at, :intervals, :load_balancer, :host, :port
 
         delegate :connection, :release_connection, to: :pool
 
@@ -25,10 +25,11 @@ module Gitlab
 
         # host - The address of the database.
         # load_balancer - The LoadBalancer that manages this Host.
-        def initialize(host, load_balancer)
+        def initialize(host, load_balancer, port: nil)
           @host = host
+          @port = port
           @load_balancer = load_balancer
-          @pool = Database.create_connection_pool(LoadBalancing.pool_size, host)
+          @pool = Database.create_connection_pool(LoadBalancing.pool_size, host, port)
           @online = true
           @last_checked_at = Time.zone.now
 

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

@@ -20,8 +20,8 @@ module Gitlab
           @mutex.synchronize { @hosts.length }
         end
 
-        def host_names
-          @mutex.synchronize { @hosts.map(&:host) }
+        def host_names_and_ports
+          @mutex.synchronize { @hosts.map { |host| [host.host, host.port] } }
         end
 
         def hosts=(hosts)

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

@@ -13,20 +13,38 @@ module Gitlab
       # balancer with said hosts. Requests may continue to use the old hosts
       # until they complete.
       class ServiceDiscovery
-        attr_reader :interval, :record, :disconnect_timeout
+        attr_reader :interval, :record, :record_type, :disconnect_timeout
 
         MAX_SLEEP_ADJUSTMENT = 10
 
+        RECORD_TYPES = {
+          'A' => Net::DNS::A,
+          'SRV' => Net::DNS::SRV
+        }.freeze
+
+        Address = Struct.new(:hostname, :port) do
+          def to_s
+            port ? "#{hostname}:#{port}" : hostname
+          end
+
+          def <=>(other)
+            self.to_s <=> other.to_s
+          end
+        end
+
         # 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.
+        # record_type - The type of DNS record to look up
         # 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, use_tcp: false)
+        # use_tcp - Use TCP instaed of UDP to look up resources
+        def initialize(nameserver:, port:, record:, record_type: 'A', interval: 60, disconnect_timeout: 120, use_tcp: false)
           @nameserver = nameserver
           @port = port
           @record = record
+          @record_type = record_type_for(record_type)
           @interval = interval
           @disconnect_timeout = disconnect_timeout
           @use_tcp = use_tcp
@@ -76,12 +94,12 @@ module Gitlab
         # Replaces all the hosts in the load balancer with the new ones,
         # disconnecting the old connections.
         #
-        # addresses - An Array of IP addresses to use for the new hosts.
+        # addresses - An Array of Address structs to use for the new hosts.
         def replace_hosts(addresses)
           old_hosts = load_balancer.host_list.hosts
 
           load_balancer.host_list.hosts = addresses.map do |addr|
-            Host.new(addr, load_balancer)
+            Host.new(addr.hostname, load_balancer, port: addr.port)
           end
 
           # We must explicitly disconnect the old connections, otherwise we may
@@ -97,15 +115,21 @@ module Gitlab
         # Returns an Array containing:
         #
         # 1. The time to wait for the next check.
-        # 2. An array containing the IP addresses of the DNS record.
+        # 2. An array containing the hostnames of the DNS record.
         def addresses_from_dns
-          resources = resolver.search(record, Net::DNS::A).answer
+          resources = resolver.search(record, record_type).answer
+
+          addresses =
+            case record_type
+            when Net::DNS::A
+              addresses_from_a_record(resources)
+            when Net::DNS::SRV
+              addresses_from_srv_record(resources)
+            end
 
           # Addresses are sorted so we can directly compare the old and new
           # addresses, without having to use any additional data structures.
-          addresses = resources.map { |r| r.address.to_s }.sort
-
-          [new_wait_time_for(resources), addresses]
+          [new_wait_time_for(resources), addresses.sort]
         end
 
         def new_wait_time_for(resources)
@@ -117,7 +141,9 @@ module Gitlab
         end
 
         def addresses_from_load_balancer
-          load_balancer.host_list.host_names.sort
+          load_balancer.host_list.host_names_and_ports.map do |hostname, port|
+            Address.new(hostname, port)
+          end.sort
         end
 
         def load_balancer
@@ -131,6 +157,22 @@ module Gitlab
             use_tcp: @use_tcp
           )
         end
+
+        private
+
+        def record_type_for(type)
+          RECORD_TYPES.fetch(type) do
+            raise(ArgumentError, "Unsupported record type: #{type}")
+          end
+        end
+
+        def addresses_from_srv_record(resources)
+          resources.map { |r| Address.new(r.host.to_s, r.port) }
+        end
+
+        def addresses_from_a_record(resources)
+          resources.map { |r| Address.new(r.address.to_s) }
+        end
       end
     end
   end

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

@@ -11,7 +11,7 @@ describe Gitlab::Database::LoadBalancing::HostList do
 
   let(:host_list) do
     hosts = Array.new(2) do
-      Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer)
+      Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer, port: 5432)
     end
 
     described_class.new(hosts)
@@ -23,9 +23,35 @@ describe Gitlab::Database::LoadBalancing::HostList do
     end
   end
 
-  describe '#host_names' do
-    it 'returns the host names of all hosts' do
-      expect(host_list.host_names).to eq(%w[localhost localhost])
+  describe '#host_names_and_ports' do
+    context 'with ports' do
+      it 'returns the host names of all hosts' do
+        hosts = [
+          ['localhost', 5432],
+          ['localhost', 5432]
+        ]
+
+        expect(host_list.host_names_and_ports).to eq(hosts)
+      end
+    end
+
+    context 'without ports' do
+      let(:host_list) do
+        hosts = Array.new(2) do
+          Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer)
+        end
+
+        described_class.new(hosts)
+      end
+
+      it 'returns the host names of all hosts' do
+        hosts = [
+          ['localhost', nil],
+          ['localhost', nil]
+        ]
+
+        expect(host_list.host_names_and_ports).to eq(hosts)
+      end
     end
   end
 

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

@@ -16,6 +16,26 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
       .and_return(packet)
   end
 
+  describe '#initialize' do
+    describe ':record_type' do
+      subject { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) }
+
+      context 'with a supported type' do
+        let(:record_type) { 'SRV' }
+
+        it { expect(subject.record_type).to eq Net::DNS::SRV }
+      end
+
+      context 'with an unsupported type' do
+        let(:record_type) { 'AAAA' }
+
+        it 'raises an argument error' do
+          expect { subject }.to raise_error(ArgumentError, 'Unsupported record type: AAAA')
+        end
+      end
+    end
+  end
+
   describe '#start' do
     before do
       allow(service)
@@ -63,6 +83,9 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   end
 
   describe '#refresh_if_necessary' do
+    let(:address_foo) { described_class::Address.new('foo') }
+    let(:address_bar) { described_class::Address.new('bar') }
+
     context 'when a refresh is necessary' do
       before do
         allow(service)
@@ -71,13 +94,13 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
 
         allow(service)
           .to receive(:addresses_from_dns)
-          .and_return([10, %w[foo bar]])
+          .and_return([10, [address_foo, address_bar]])
       end
 
       it 'refreshes the load balancer hosts' do
         expect(service)
           .to receive(:replace_hosts)
-          .with(%w[foo bar])
+          .with([address_foo, address_bar])
 
         expect(service.refresh_if_necessary).to eq(10)
       end
@@ -104,8 +127,11 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   end
 
   describe '#replace_hosts' do
+    let(:address_foo) { described_class::Address.new('foo') }
+    let(:address_bar) { described_class::Address.new('bar') }
+
     let(:load_balancer) do
-      Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[foo])
+      Gitlab::Database::LoadBalancing::LoadBalancer.new([address_foo])
     end
 
     before do
@@ -115,9 +141,9 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
     end
 
     it 'replaces the hosts of the load balancer' do
-      service.replace_hosts(%w[bar])
+      service.replace_hosts([address_bar])
 
-      expect(load_balancer.host_list.host_names).to eq(%w[bar])
+      expect(load_balancer.host_list.host_names_and_ports).to eq([['bar', nil]])
     end
 
     it 'disconnects the old connections' do
@@ -131,23 +157,51 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
         .to receive(:disconnect!)
         .with(2)
 
-      service.replace_hosts(%w[bar])
+      service.replace_hosts([address_bar])
     end
   end
 
   describe '#addresses_from_dns' 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])
+    let(:service) { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) }
+    let(:packet) { double(:packet, answer: [res1, res2]) }
 
+    before do
       allow(service.resolver)
         .to receive(:search)
-        .with('foo', Net::DNS::A)
+        .with('foo', described_class::RECORD_TYPES[record_type])
         .and_return(packet)
+    end
+
+    context 'with an A record' do
+      let(:record_type) { 'A' }
+
+      let(:res1) { double(:resource, address: '255.255.255.0', ttl: 90) }
+      let(:res2) { double(:resource, address: '127.0.0.1', ttl: 90) }
+
+      it 'returns a TTL and ordered list of IP addresses' do
+        addresses = [
+          described_class::Address.new('127.0.0.1'),
+          described_class::Address.new('255.255.255.0')
+        ]
+
+        expect(service.addresses_from_dns).to eq([90, addresses])
+      end
+    end
+
+    context 'with an SRV record' do
+      let(:record_type) { 'SRV' }
+
+      let(:res1) { double(:resource, host: '255.255.255.0', port: 5432, weight: 1, priority: 1, ttl: 90) }
+      let(:res2) { double(:resource, host: '127.0.0.1', port: 5433, weight: 1, priority: 1, ttl: 90) }
 
-      expect(service.addresses_from_dns)
-        .to eq([90, %w[127.0.0.1 255.255.255.0]])
+      it 'returns a TTL and ordered list of hosts' do
+        addresses = [
+          described_class::Address.new('127.0.0.1', 5433),
+          described_class::Address.new('255.255.255.0', 5432)
+        ]
+
+        expect(service.addresses_from_dns).to eq([90, addresses])
+      end
     end
   end
 
@@ -177,7 +231,12 @@ describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
         .to receive(:load_balancer)
         .and_return(load_balancer)
 
-      expect(service.addresses_from_load_balancer).to eq(%w[a b])
+      addresses = [
+        described_class::Address.new('a'),
+        described_class::Address.new('b')
+      ]
+
+      expect(service.addresses_from_load_balancer).to eq(addresses)
     end
   end
 end

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

@@ -220,6 +220,7 @@ describe Gitlab::Database::LoadBalancing do
           nameserver: 'localhost',
           port: 8600,
           record: nil,
+          record_type: 'A',
           interval: 60,
           disconnect_timeout: 120,
           use_tcp: false
@@ -231,12 +232,13 @@ describe Gitlab::Database::LoadBalancing do
       it 'returns a Hash including the custom configuration' do
         allow(described_class)
           .to receive(:configuration)
-          .and_return('discover' => { 'record' => 'foo' })
+          .and_return('discover' => { 'record' => 'foo', 'record_type' => 'SRV' })
 
         expect(described_class.service_discovery_configuration).to eq(
           nameserver: 'localhost',
           port: 8600,
           record: 'foo',
+          record_type: 'SRV',
           interval: 60,
           disconnect_timeout: 120,
           use_tcp: false

lib/gitlab/database.rb

@@ -195,13 +195,14 @@ module Gitlab
 
     # pool_size - The size of the DB pool.
     # host - An optional host name to use instead of the default one.
-    def self.create_connection_pool(pool_size, host = nil)
+    def self.create_connection_pool(pool_size, host = nil, port = nil)
       # See activerecord-4.2.7.1/lib/active_record/connection_adapters/connection_specification.rb
       env = Rails.env
       original_config = ActiveRecord::Base.configurations
 
       env_config = original_config[env].merge('pool' => pool_size)
       env_config['host'] = host if host
+      env_config['port'] = port if port
 
       config = original_config.merge(env => env_config)
 

spec/lib/gitlab/database_spec.rb

@@ -347,6 +347,17 @@ describe Gitlab::Database do
         pool.disconnect!
       end
     end
+
+    it 'allows setting of a custom hostname and port' do
+      pool = described_class.create_connection_pool(5, '127.0.0.1', 5432)
+
+      begin
+        expect(pool.spec.config[:host]).to eq('127.0.0.1')
+        expect(pool.spec.config[:port]).to eq(5432)
+      ensure
+        pool.disconnect!
+      end
+    end
   end
 
   describe '.cached_column_exists?' do