Add multiple DBs support to LB service discovery

This adds support for using the load balancer's service discovery
mechanism with multiple databases. The change is quite simple: instead
of depending directly on `LoadBalancing.proxy.load_balancer`, we pass
the load balancer as an argument to the `initialize` method.

lib/gitlab/database/load_balancing/service_discovery.rb

@@ -13,7 +13,8 @@ module Gitlab
       # balancer with said hosts. Requests may continue to use the old hosts
       # until they complete.
       class ServiceDiscovery
-        attr_reader :interval, :record, :record_type, :disconnect_timeout
+        attr_reader :interval, :record, :record_type, :disconnect_timeout,
+                    :load_balancer
 
         MAX_SLEEP_ADJUSTMENT = 10
 
@@ -40,7 +41,17 @@ module Gitlab
         # disconnect_timeout - The time after which an old host should be
         #                      forcefully disconnected.
         # 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)
+        # load_balancer - The load balancer instance to use
+        def initialize(
+          nameserver:,
+          port:,
+          record:,
+          record_type: 'A',
+          interval: 60,
+          disconnect_timeout: 120,
+          use_tcp: false,
+          load_balancer: LoadBalancing.proxy.load_balancer
+        )
           @nameserver = nameserver
           @port = port
           @record = record
@@ -48,6 +59,7 @@ module Gitlab
           @interval = interval
           @disconnect_timeout = disconnect_timeout
           @use_tcp = use_tcp
+          @load_balancer = load_balancer
         end
 
         def start
@@ -147,10 +159,6 @@ module Gitlab
           end.sort
         end
 
-        def load_balancer
-          LoadBalancing.proxy.load_balancer
-        end
-
         def resolver
           @resolver ||= Net::DNS::Resolver.new(
             nameservers: Resolver.new(@nameserver).resolve,

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

@@ -3,8 +3,14 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
+  let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new([]) }
   let(:service) do
-    described_class.new(nameserver: 'localhost', port: 8600, record: 'foo')
+    described_class.new(
+      nameserver: 'localhost',
+      port: 8600,
+      record: 'foo',
+      load_balancer: load_balancer
+    )
   end
 
   before do
@@ -18,7 +24,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
 
   describe '#initialize' do
     describe ':record_type' do
-      subject { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) }
+      subject do
+        described_class.new(
+          nameserver: 'localhost',
+          port: 8600,
+          record: 'foo',
+          record_type: record_type,
+          load_balancer: load_balancer
+        )
+      end
 
       context 'with a supported type' do
         let(:record_type) { 'SRV' }
@@ -162,7 +176,16 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   end
 
   describe '#addresses_from_dns' do
-    let(:service) { described_class.new(nameserver: 'localhost', port: 8600, record: 'foo', record_type: record_type) }
+    let(:service) do
+      described_class.new(
+        nameserver: 'localhost',
+        port: 8600,
+        record: 'foo',
+        record_type: record_type,
+        load_balancer: load_balancer
+      )
+    end
+
     let(:packet) { double(:packet, answer: [res1, res2]) }
 
     before do
@@ -234,13 +257,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   end
 
   describe '#addresses_from_load_balancer' do
-    it 'returns the ordered host names of the load balancer' do
-      load_balancer = Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[b a])
-
-      allow(service)
-        .to receive(:load_balancer)
-        .and_return(load_balancer)
+    let(:load_balancer) do
+      Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[b a])
+    end
 
+    it 'returns the ordered host names of the load balancer' do
       addresses = [
         described_class::Address.new('a'),
         described_class::Address.new('b')