Require a LoadBalancer for service discovery

The service discovery class now requires a load balancer to be passed
in, instead of using a default object. This decouples the class from the
global proxy/load balancer we currently use. This makes it easier to
create multiple service discovery instances, without accidentally using
the wrong load balancer.

Changelog: changed

lib/gitlab/database/load_balancing.rb

@@ -102,7 +102,9 @@ module Gitlab
       def self.start_service_discovery
         return unless service_discovery_enabled?
 
-        ServiceDiscovery.new(service_discovery_configuration).start
+        ServiceDiscovery
+          .new(proxy.load_balancer, **service_discovery_configuration)
+          .start
       end
 
       # Configures proxying of requests.
@@ -111,7 +113,9 @@ module Gitlab
 
         # Populate service discovery immediately if it is configured
         if service_discovery_enabled?
-          ServiceDiscovery.new(service_discovery_configuration).perform_service_discovery
+          ServiceDiscovery
+            .new(proxy.load_balancer, **service_discovery_configuration)
+            .perform_service_discovery
         end
       end
 

lib/gitlab/database/load_balancing/service_discovery.rb

@@ -49,14 +49,14 @@ module Gitlab
         # use_tcp - Use TCP instaed of UDP to look up resources
         # load_balancer - The load balancer instance to use
         def initialize(
+          load_balancer,
           nameserver:,
           port:,
           record:,
           record_type: 'A',
           interval: 60,
           disconnect_timeout: 120,
-          use_tcp: false,
-          load_balancer: LoadBalancing.proxy.load_balancer
+          use_tcp: false
         )
           @nameserver = nameserver
           @port = port

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

@@ -6,10 +6,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new([]) }
   let(:service) do
     described_class.new(
+      load_balancer,
       nameserver: 'localhost',
       port: 8600,
-      record: 'foo',
-      load_balancer: load_balancer
+      record: 'foo'
     )
   end
 
@@ -26,11 +26,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
     describe ':record_type' do
       subject do
         described_class.new(
+          load_balancer,
           nameserver: 'localhost',
           port: 8600,
           record: 'foo',
-          record_type: record_type,
-          load_balancer: load_balancer
+          record_type: record_type
         )
       end
 
@@ -217,11 +217,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
   describe '#addresses_from_dns' do
     let(:service) do
       described_class.new(
+        load_balancer,
         nameserver: 'localhost',
         port: 8600,
         record: 'foo',
-        record_type: record_type,
-        load_balancer: load_balancer
+        record_type: record_type
       )
     end
 

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -204,9 +204,11 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.configure_proxy' do
-    it 'configures the connection proxy' do
+    before do
       allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
+    end
 
+    it 'configures the connection proxy' do
       described_class.configure_proxy
 
       expect(ActiveRecord::Base).to have_received(:load_balancing_proxy=)
@@ -214,15 +216,22 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
 
     context 'when service discovery is enabled' do
-      let(:service_discovery) { double(Gitlab::Database::LoadBalancing::ServiceDiscovery) }
-
       it 'runs initial service discovery when configuring the connection proxy' do
+        discover = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+
         allow(described_class)
           .to receive(:configuration)
           .and_return('discover' => { 'record' => 'foo' })
 
-        expect(Gitlab::Database::LoadBalancing::ServiceDiscovery).to receive(:new).and_return(service_discovery)
-        expect(service_discovery).to receive(:perform_service_discovery)
+        expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
+          .to receive(:new)
+          .with(
+            an_instance_of(Gitlab::Database::LoadBalancing::LoadBalancer),
+            an_instance_of(Hash)
+          )
+          .and_return(discover)
+
+        expect(discover).to receive(:perform_service_discovery)
 
         described_class.configure_proxy
       end
@@ -297,10 +306,16 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         .and_return(true)
 
       instance = double(:instance)
+      lb = instance_spy(Gitlab::Database::LoadBalancing::LoadBalancer)
+      proxy = double(:proxy, load_balancer: lb)
+
+      allow(Gitlab::Database::LoadBalancing)
+        .to receive(:proxy)
+        .and_return(proxy)
 
       expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
         .to receive(:new)
-        .with(an_instance_of(Hash))
+        .with(lb, an_instance_of(Hash))
         .and_return(instance)
 
       expect(instance)