Move load balancing configuration into a class

This moves load balancing configuration settings, such as the
replication lag settings, into a dedicated class. This allows us to
scope load balancer settings to load balancer instances, instead of
being forced to use global settings. This in turn will make it easier to
always enable the load balancing code, and to eventually support load
balancing of multiple databases.

lib/gitlab/database/load_balancing.rb

@@ -36,94 +36,42 @@ module Gitlab
 
       # Returns a Hash containing the load balancing configuration.
       def self.configuration
-        Gitlab::Database.main.config[:load_balancing] || {}
-      end
-
-      # Returns the maximum replica lag size in bytes.
-      def self.max_replication_difference
-        (configuration['max_replication_difference'] || 8.megabytes).to_i
-      end
-
-      # Returns the maximum lag time for a replica.
-      def self.max_replication_lag_time
-        (configuration['max_replication_lag_time'] || 60.0).to_f
-      end
-
-      # Returns the interval (in seconds) to use for checking the status of a
-      # replica.
-      def self.replica_check_interval
-        (configuration['replica_check_interval'] || 60).to_f
-      end
-
-      # Returns the additional hosts to use for load balancing.
-      def self.hosts
-        configuration['hosts'] || []
-      end
-
-      def self.service_discovery_enabled?
-        configuration.dig('discover', 'record').present?
-      end
-
-      def self.service_discovery_configuration
-        conf = configuration['discover'] || {}
-
-        {
-          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
-        }
-      end
-
-      def self.pool_size
-        Gitlab::Database.main.pool_size
+        @configuration ||= Configuration.for_model(ActiveRecord::Base)
       end
 
       # Returns true if load balancing is to be enabled.
       def self.enable?
         return false if Gitlab::Runtime.rake?
-        return false unless self.configured?
 
-        true
+        configured?
       end
 
-      # Returns true if load balancing has been configured. Since
-      # Sidekiq does not currently use load balancing, we
-      # may want Web application servers to detect replication lag by
-      # posting the write location of the database if load balancing is
-      # configured.
       def self.configured?
-        hosts.any? || service_discovery_enabled?
+        configuration.load_balancing_enabled? ||
+          configuration.service_discovery_enabled?
       end
 
       def self.start_service_discovery
-        return unless service_discovery_enabled?
+        return unless configuration.service_discovery_enabled?
 
         ServiceDiscovery
-          .new(proxy.load_balancer, **service_discovery_configuration)
+          .new(proxy.load_balancer, **configuration.service_discovery)
           .start
       end
 
       # Configures proxying of requests.
       def self.configure_proxy
-        lb = LoadBalancer.new(hosts, primary_only: !enable?)
+        lb = LoadBalancer.new(configuration, primary_only: !enable?)
         ActiveRecord::Base.load_balancing_proxy = ConnectionProxy.new(lb)
 
         # Populate service discovery immediately if it is configured
-        if service_discovery_enabled?
+        if configuration.service_discovery_enabled?
           ServiceDiscovery
-            .new(lb, **service_discovery_configuration)
+            .new(lb, **configuration.service_discovery)
             .perform_service_discovery
         end
       end
 
-      def self.active_record_models
-        ActiveRecord::Base.descendants
-      end
-
       DB_ROLES = [
         ROLE_PRIMARY = :primary,
         ROLE_REPLICA = :replica,

lib/gitlab/database/load_balancing/configuration.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LoadBalancing
+      # Configuration settings for a single LoadBalancer instance.
+      class Configuration
+        attr_accessor :hosts, :max_replication_difference,
+                      :max_replication_lag_time, :replica_check_interval,
+                      :service_discovery, :pool_size, :model
+
+        # Creates a configuration object for the given ActiveRecord model.
+        def self.for_model(model)
+          cfg = model.connection_db_config.configuration_hash
+          lb_cfg = cfg[:load_balancing] || {}
+          config = new(model)
+
+          if (size = cfg[:pool])
+            config.pool_size = size
+          end
+
+          if (diff = lb_cfg[:max_replication_difference])
+            config.max_replication_difference = diff
+          end
+
+          if (lag = lb_cfg[:max_replication_lag_time])
+            config.max_replication_lag_time = lag.to_f
+          end
+
+          if (interval = lb_cfg[:replica_check_interval])
+            config.replica_check_interval = interval.to_f
+          end
+
+          if (hosts = lb_cfg[:hosts])
+            config.hosts = hosts
+          end
+
+          discover = (lb_cfg[:discover] || {}).symbolize_keys
+
+          # We iterate over the known/default keys so we don't end up with
+          # random keys in our configuration hash.
+          config.service_discovery.each do |key, _|
+            if (value = discover[key])
+              config.service_discovery[key] = value
+            end
+          end
+
+          config
+        end
+
+        def initialize(model, hosts = [])
+          @max_replication_difference = 8.megabytes
+          @max_replication_lag_time = 60.0
+          @replica_check_interval = 60.0
+          @model = model
+          @hosts = hosts
+          @pool_size = Database.default_pool_size
+          @service_discovery = {
+            nameserver: 'localhost',
+            port: 8600,
+            record: nil,
+            record_type: 'A',
+            interval: 60,
+            disconnect_timeout: 120,
+            use_tcp: false
+          }
+        end
+
+        def load_balancing_enabled?
+          hosts.any? || service_discovery_enabled?
+        end
+
+        def service_discovery_enabled?
+          service_discovery[:record].present?
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/load_balancing/host.rb

@@ -29,11 +29,15 @@ module Gitlab
           @host = host
           @port = port
           @load_balancer = load_balancer
-          @pool = load_balancer.create_replica_connection_pool(::Gitlab::Database::LoadBalancing.pool_size, host, port)
+          @pool = load_balancer.create_replica_connection_pool(
+            load_balancer.configuration.pool_size,
+            host,
+            port
+          )
           @online = true
           @last_checked_at = Time.zone.now
 
-          interval = ::Gitlab::Database::LoadBalancing.replica_check_interval
+          interval = load_balancer.configuration.replica_check_interval
           @intervals = (interval..(interval * 2)).step(0.5).to_a
         end
 
@@ -108,7 +112,7 @@ module Gitlab
 
         def replication_lag_below_threshold?
           if (lag_time = replication_lag_time)
-            lag_time <= ::Gitlab::Database::LoadBalancing.max_replication_lag_time
+            lag_time <= load_balancer.configuration.max_replication_lag_time
           else
             false
           end
@@ -125,7 +129,7 @@ module Gitlab
           # only do this if we haven't replicated in a while so we only need
           # to connect to the primary when truly necessary.
           if (lag_size = replication_lag_size)
-            lag_size <= ::Gitlab::Database::LoadBalancing.max_replication_difference
+            lag_size <= load_balancer.configuration.max_replication_difference
           else
             false
           end

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -12,20 +12,21 @@ module Gitlab
 
         REPLICA_SUFFIX = '_replica'
 
-        attr_reader :host_list
+        attr_reader :host_list, :configuration
 
-        # hosts - The hostnames/addresses of the additional databases.
-        # model - The ActiveRecord base model the load balancer is enabled for.
+        # configuration - An instance of `LoadBalancing::Configuration` that
+        #                 contains the configuration details (such as the hosts)
+        #                 for this load balancer.
         # primary_only - If set, the replicas are ignored and the primary is
         #                always used.
-        def initialize(hosts = [], model = ActiveRecord::Base, primary_only: false)
+        def initialize(configuration, primary_only: false)
+          @configuration = configuration
           @primary_only = primary_only
-          @model = model
           @host_list =
             if primary_only
               HostList.new([PrimaryHost.new(self)])
             else
-              HostList.new(hosts.map { |addr| Host.new(addr, self) })
+              HostList.new(configuration.hosts.map { |addr| Host.new(addr, self) })
             end
         end
 
@@ -231,7 +232,7 @@ module Gitlab
         # leverage that.
         def pool
           ActiveRecord::Base.connection_handler.retrieve_connection_pool(
-            @model.connection_specification_name,
+            @configuration.model.connection_specification_name,
             role: ActiveRecord::Base.writing_role,
             shard: ActiveRecord::Base.default_shard
           )

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
+  describe '.for_model' do
+    let(:model) do
+      config = ActiveRecord::DatabaseConfigurations::HashConfig
+        .new('main', 'test', configuration_hash)
+
+      double(:model, connection_db_config: config)
+    end
+
+    context 'when load balancing is not configured' do
+      let(:configuration_hash) { {} }
+
+      it 'uses the default settings' do
+        config = described_class.for_model(model)
+
+        expect(config.hosts).to eq([])
+        expect(config.max_replication_difference).to eq(8.megabytes)
+        expect(config.max_replication_lag_time).to eq(60.0)
+        expect(config.replica_check_interval).to eq(60.0)
+        expect(config.service_discovery).to eq(
+          nameserver: 'localhost',
+          port: 8600,
+          record: nil,
+          record_type: 'A',
+          interval: 60,
+          disconnect_timeout: 120,
+          use_tcp: false
+        )
+        expect(config.pool_size).to eq(Gitlab::Database.default_pool_size)
+      end
+    end
+
+    context 'when load balancing is configured' do
+      let(:configuration_hash) do
+        {
+          pool: 4,
+          load_balancing: {
+            max_replication_difference: 1,
+            max_replication_lag_time: 2,
+            replica_check_interval: 3,
+            hosts: %w[foo bar],
+            discover: {
+              'record' => 'foo.example.com'
+            }
+          }
+        }
+      end
+
+      it 'uses the custom configuration settings' do
+        config = described_class.for_model(model)
+
+        expect(config.hosts).to eq(%w[foo bar])
+        expect(config.max_replication_difference).to eq(1)
+        expect(config.max_replication_lag_time).to eq(2.0)
+        expect(config.replica_check_interval).to eq(3.0)
+        expect(config.service_discovery).to eq(
+          nameserver: 'localhost',
+          port: 8600,
+          record: 'foo.example.com',
+          record_type: 'A',
+          interval: 60,
+          disconnect_timeout: 120,
+          use_tcp: false
+        )
+        expect(config.pool_size).to eq(4)
+      end
+    end
+  end
+
+  describe '#load_balancing_enabled?' do
+    it 'returns true when hosts are configured' do
+      config = described_class.new(ActiveRecord::Base, %w[foo bar])
+
+      expect(config.load_balancing_enabled?).to eq(true)
+    end
+
+    it 'returns true when a service discovery record is configured' do
+      config = described_class.new(ActiveRecord::Base)
+      config.service_discovery[:record] = 'foo'
+
+      expect(config.load_balancing_enabled?).to eq(true)
+    end
+
+    it 'returns false when no hosts are configured and service discovery is disabled' do
+      config = described_class.new(ActiveRecord::Base)
+
+      expect(config.load_balancing_enabled?).to eq(false)
+    end
+  end
+
+  describe '#service_discovery_enabled?' do
+    it 'returns true when a record is configured' do
+      config = described_class.new(ActiveRecord::Base)
+      config.service_discovery[:record] = 'foo'
+
+      expect(config.service_discovery_enabled?).to eq(true)
+    end
+
+    it 'returns false when no record is configured' do
+      config = described_class.new(ActiveRecord::Base)
+
+      expect(config.service_discovery_enabled?).to eq(false)
+    end
+  end
+end

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

@@ -4,7 +4,10 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
   let(:proxy) do
-    described_class.new(Gitlab::Database::LoadBalancing::LoadBalancer.new([]))
+    config = Gitlab::Database::LoadBalancing::Configuration
+      .new(ActiveRecord::Base)
+
+    described_class.new(Gitlab::Database::LoadBalancing::LoadBalancer.new(config))
   end
 
   describe '#select' do

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

@@ -4,7 +4,12 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::HostList do
   let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
-  let(:load_balancer) { double(:load_balancer) }
+  let(:load_balancer) do
+    Gitlab::Database::LoadBalancing::LoadBalancer.new(
+      Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
+    )
+  end
+
   let(:host_count) { 2 }
   let(:hosts) { Array.new(host_count) { Gitlab::Database::LoadBalancing::Host.new(db_host, load_balancer, port: 5432) } }
   let(:host_list) { described_class.new(hosts) }

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

@@ -3,7 +3,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::Host do
-  let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new }
+  let(:load_balancer) do
+    Gitlab::Database::LoadBalancing::LoadBalancer
+      .new(Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base))
+  end
 
   let(:host) do
     Gitlab::Database::LoadBalancing::Host.new('localhost', load_balancer)
@@ -274,7 +277,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Host do
     end
 
     it 'returns false when the data is not recent enough' do
-      diff = Gitlab::Database::LoadBalancing.max_replication_difference * 2
+      diff = load_balancer.configuration.max_replication_difference * 2
 
       expect(host)
         .to receive(:query_and_release)

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

@@ -5,7 +5,12 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   let(:conflict_error) { Class.new(RuntimeError) }
   let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
-  let(:lb) { described_class.new([db_host, db_host]) }
+  let(:config) do
+    Gitlab::Database::LoadBalancing::Configuration
+      .new(ActiveRecord::Base, [db_host, db_host])
+  end
+
+  let(:lb) { described_class.new(config) }
   let(:request_cache) { lb.send(:request_cache) }
 
   before do
@@ -43,7 +48,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
 
   describe '#initialize' do
     it 'ignores the hosts when the primary_only option is enabled' do
-      lb = described_class.new([db_host], primary_only: true)
+      config = Gitlab::Database::LoadBalancing::Configuration
+        .new(ActiveRecord::Base, [db_host])
+      lb = described_class.new(config, primary_only: true)
       hosts = lb.host_list.hosts
 
       expect(hosts.length).to eq(1)
@@ -134,7 +141,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
 
     it 'uses the primary when the primary_only option is enabled' do
-      lb = described_class.new(primary_only: true)
+      config = Gitlab::Database::LoadBalancing::Configuration
+        .new(ActiveRecord::Base)
+      lb = described_class.new(config, primary_only: true)
 
       # When no hosts are configured, we don't want to produce any warnings, as
       # they aren't useful/too noisy.
@@ -174,8 +183,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
 
     it 'does not create conflicts with other load balancers when caching hosts' do
-      lb1 = described_class.new([db_host, db_host], ActiveRecord::Base)
-      lb2 = described_class.new([db_host, db_host], Ci::CiDatabaseRecord)
+      ci_config = Gitlab::Database::LoadBalancing::Configuration
+        .new(Ci::CiDatabaseRecord, [db_host, db_host])
+
+      lb1 = described_class.new(config)
+      lb2 = described_class.new(ci_config)
 
       host1 = lb1.host
       host2 = lb2.host

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

@@ -3,7 +3,12 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::PrimaryHost do
-  let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new }
+  let(:load_balancer) do
+    Gitlab::Database::LoadBalancing::LoadBalancer.new(
+      Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
+    )
+  end
+
   let(:host) { Gitlab::Database::LoadBalancing::PrimaryHost.new(load_balancer) }
 
   describe '#connection' do

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

@@ -3,7 +3,12 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
-  let(:load_balancer) { Gitlab::Database::LoadBalancing::LoadBalancer.new([]) }
+  let(:load_balancer) do
+    Gitlab::Database::LoadBalancing::LoadBalancer.new(
+      Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
+    )
+  end
+
   let(:service) do
     described_class.new(
       load_balancer,
@@ -184,7 +189,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
     let(:address_bar) { described_class::Address.new('bar') }
 
     let(:load_balancer) do
-      Gitlab::Database::LoadBalancing::LoadBalancer.new([address_foo])
+      Gitlab::Database::LoadBalancing::LoadBalancer.new(
+        Gitlab::Database::LoadBalancing::Configuration
+          .new(ActiveRecord::Base, [address_foo])
+      )
     end
 
     before do
@@ -307,7 +315,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
 
   describe '#addresses_from_load_balancer' do
     let(:load_balancer) do
-      Gitlab::Database::LoadBalancing::LoadBalancer.new(%w[b a])
+      Gitlab::Database::LoadBalancing::LoadBalancer.new(
+        Gitlab::Database::LoadBalancing::Configuration
+          .new(ActiveRecord::Base, %w[b a])
+      )
     end
 
     it 'returns the ordered host names of the load balancer' do

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -40,106 +40,25 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.configuration' do
-    it 'returns a Hash' do
-      lb_config = { 'hosts' => %w(foo) }
+    it 'returns the configuration for the load balancer' do
+      raw = ActiveRecord::Base.connection_db_config.configuration_hash
+      cfg = described_class.configuration
 
-      original_db_config = Gitlab::Database.main.config
-      modified_db_config = original_db_config.merge(load_balancing: lb_config)
-      expect(Gitlab::Database.main).to receive(:config).and_return(modified_db_config)
-
-      expect(described_class.configuration).to eq(lb_config)
-    end
-  end
-
-  describe '.max_replication_difference' do
-    context 'without an explicitly configured value' do
-      it 'returns the default value' do
-        allow(described_class)
-          .to receive(:configuration)
-          .and_return({})
-
-        expect(described_class.max_replication_difference).to eq(8.megabytes)
-      end
-    end
-
-    context 'with an explicitly configured value' do
-      it 'returns the configured value' do
-        allow(described_class)
-          .to receive(:configuration)
-          .and_return({ 'max_replication_difference' => 4 })
-
-        expect(described_class.max_replication_difference).to eq(4)
-      end
-    end
-  end
-
-  describe '.max_replication_lag_time' do
-    context 'without an explicitly configured value' do
-      it 'returns the default value' do
-        allow(described_class)
-          .to receive(:configuration)
-          .and_return({})
-
-        expect(described_class.max_replication_lag_time).to eq(60)
-      end
-    end
-
-    context 'with an explicitly configured value' do
-      it 'returns the configured value' do
-        allow(described_class)
-          .to receive(:configuration)
-          .and_return({ 'max_replication_lag_time' => 4 })
-
-        expect(described_class.max_replication_lag_time).to eq(4)
-      end
-    end
-  end
-
-  describe '.replica_check_interval' do
-    context 'without an explicitly configured value' do
-      it 'returns the default value' do
-        allow(described_class)
-          .to receive(:configuration)
-          .and_return({})
-
-        expect(described_class.replica_check_interval).to eq(60)
-      end
-    end
-
-    context 'with an explicitly configured value' do
-      it 'returns the configured value' do
-        allow(described_class)
-          .to receive(:configuration)
-          .and_return({ 'replica_check_interval' => 4 })
-
-        expect(described_class.replica_check_interval).to eq(4)
-      end
-    end
-  end
-
-  describe '.hosts' do
-    it 'returns a list of hosts' do
-      allow(described_class)
-        .to receive(:configuration)
-        .and_return({ 'hosts' => %w(foo bar baz) })
-
-      expect(described_class.hosts).to eq(%w(foo bar baz))
-    end
-  end
-
-  describe '.pool_size' do
-    it 'returns a Fixnum' do
-      expect(described_class.pool_size).to be_a_kind_of(Integer)
+      # There isn't much to test here as the load balancing settings might not
+      # (and likely aren't) set when running tests.
+      expect(cfg.pool_size).to eq(raw[:pool])
     end
   end
 
   describe '.enable?' do
     before do
-      allow(described_class).to receive(:hosts).and_return(%w(foo))
+      allow(described_class.configuration)
+        .to receive(:hosts)
+        .and_return(%w(foo))
     end
 
     it 'returns false when no hosts are specified' do
-      allow(described_class).to receive(:hosts).and_return([])
+      allow(described_class.configuration).to receive(:hosts).and_return([])
 
       expect(described_class.enable?).to eq(false)
     end
@@ -163,10 +82,10 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
 
     it 'returns true when service discovery is enabled' do
-      allow(described_class).to receive(:hosts).and_return([])
+      allow(described_class.configuration).to receive(:hosts).and_return([])
       allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(false)
 
-      allow(described_class)
+      allow(described_class.configuration)
         .to receive(:service_discovery_enabled?)
         .and_return(true)
 
@@ -175,17 +94,17 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.configured?' do
-    it 'returns true when Sidekiq is being used' do
-      allow(described_class).to receive(:hosts).and_return(%w(foo))
-      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+    it 'returns true when hosts are configured' do
+      allow(described_class.configuration)
+        .to receive(:hosts)
+        .and_return(%w[foo])
+
       expect(described_class.configured?).to eq(true)
     end
 
-    it 'returns true when service discovery is enabled in Sidekiq' do
-      allow(described_class).to receive(:hosts).and_return([])
-      allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
-
-      allow(described_class)
+    it 'returns true when service discovery is enabled' do
+      allow(described_class.configuration).to receive(:hosts).and_return([])
+      allow(described_class.configuration)
         .to receive(:service_discovery_enabled?)
         .and_return(true)
 
@@ -193,9 +112,8 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
 
     it 'returns false when neither service discovery nor hosts are configured' do
-      allow(described_class).to receive(:hosts).and_return([])
-
-      allow(described_class)
+      allow(described_class.configuration).to receive(:hosts).and_return([])
+      allow(described_class.configuration)
         .to receive(:service_discovery_enabled?)
         .and_return(false)
 
@@ -219,9 +137,9 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       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' })
+        allow(described_class.configuration)
+          .to receive(:service_discovery)
+          .and_return({ record: 'foo' })
 
         expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
           .to receive(:new)
@@ -238,60 +156,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
   end
 
-  describe '.active_record_models' do
-    it 'returns an Array' do
-      expect(described_class.active_record_models).to be_an_instance_of(Array)
-    end
-  end
-
-  describe '.service_discovery_enabled?' do
-    it 'returns true if service discovery is enabled' do
-      allow(described_class)
-        .to receive(:configuration)
-        .and_return('discover' => { 'record' => 'foo' })
-
-      expect(described_class.service_discovery_enabled?).to eq(true)
-    end
-
-    it 'returns false if service discovery is disabled' do
-      expect(described_class.service_discovery_enabled?).to eq(false)
-    end
-  end
-
-  describe '.service_discovery_configuration' do
-    context 'when no configuration is provided' do
-      it 'returns a default configuration Hash' do
-        expect(described_class.service_discovery_configuration).to eq(
-          nameserver: 'localhost',
-          port: 8600,
-          record: nil,
-          record_type: 'A',
-          interval: 60,
-          disconnect_timeout: 120,
-          use_tcp: false
-        )
-      end
-    end
-
-    context 'when configuration is provided' do
-      it 'returns a Hash including the custom configuration' do
-        allow(described_class)
-          .to receive(:configuration)
-          .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
-        )
-      end
-    end
-  end
-
   describe '.start_service_discovery' do
     it 'does not start if service discovery is disabled' do
       expect(Gitlab::Database::LoadBalancing::ServiceDiscovery)
@@ -301,12 +165,14 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
 
     it 'starts service discovery if enabled' do
-      allow(described_class)
+      allow(described_class.configuration)
         .to receive(:service_discovery_enabled?)
         .and_return(true)
 
       instance = double(:instance)
-      lb = Gitlab::Database::LoadBalancing::LoadBalancer.new([])
+      config = Gitlab::Database::LoadBalancing::Configuration
+        .new(ActiveRecord::Base)
+      lb = Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
       proxy = Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
 
       allow(described_class)
@@ -345,7 +211,12 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
     context 'when the load balancing is configured' do
       let(:db_host) { ActiveRecord::Base.connection_pool.db_config.host }
-      let(:load_balancer) { described_class::LoadBalancer.new([db_host]) }
+      let(:config) do
+        Gitlab::Database::LoadBalancing::Configuration
+          .new(ActiveRecord::Base, [db_host])
+      end
+
+      let(:load_balancer) { described_class::LoadBalancer.new(config) }
       let(:proxy) { described_class::ConnectionProxy.new(load_balancer) }
 
       context 'when a proxy connection is used' do

spec/support/database_load_balancing.rb

@@ -4,8 +4,9 @@ RSpec.configure do |config|
   config.before(:each, :db_load_balancing) do
     allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
 
-    lb = ::Gitlab::Database::LoadBalancing::LoadBalancer
-      .new([Gitlab::Database.main.config['host']])
+    config = Gitlab::Database::LoadBalancing::Configuration
+      .new(ActiveRecord::Base, [Gitlab::Database.main.config['host']])
+    lb = ::Gitlab::Database::LoadBalancing::LoadBalancer.new(config)
     proxy = ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
 
     allow(ActiveRecord::Base).to receive(:load_balancing_proxy).and_return(proxy)