Merge branch 'yorick/load-balancing-config-class' into 'master'

Move load balancing configuration into a dedicated class

See merge request gitlab-org/gitlab!69233

config/initializers/database_config.rb

@@ -24,7 +24,7 @@ ActiveRecord::Base.establish_connection(Gitlab::Database.main.db_config_with_def
 
 Gitlab.ee do
   if Gitlab::Runtime.sidekiq? && Gitlab::Geo.geo_database_configured?
-    Rails.configuration.geo_database['pool'] = Gitlab::Database.main.default_pool_size
+    Rails.configuration.geo_database['pool'] = Gitlab::Database.default_pool_size
     Geo::TrackingBase.establish_connection(Rails.configuration.geo_database)
   end
 end

lib/gitlab/database.rb

@@ -6,6 +6,7 @@ module Gitlab
 
     MAIN_DATABASE_NAME = 'main'
     CI_DATABASE_NAME = 'ci'
+    DEFAULT_POOL_HEADROOM = 10
 
     # This constant is used when renaming tables concurrently.
     # If you plan to rename a table using the `rename_table_safely` method, add your table here one milestone before the rename.
@@ -62,6 +63,20 @@ module Gitlab
       DATABASES[PRIMARY_DATABASE_NAME]
     end
 
+    # We configure the database connection pool size automatically based on the
+    # configured concurrency. We also add some headroom, to make sure we don't
+    # run out of connections when more threads besides the 'user-facing' ones
+    # are running.
+    #
+    # Read more about this in
+    # doc/development/database/client_side_connection_pool.md
+    def self.default_pool_size
+      headroom =
+        (ENV["DB_POOL_HEADROOM"].presence || DEFAULT_POOL_HEADROOM).to_i
+
+      Gitlab::Runtime.max_threads + headroom
+    end
+
     def self.has_config?(database_name)
       Gitlab::Application.config.database_configuration[Rails.env].include?(database_name.to_s)
     end

lib/gitlab/database/connection.rb

@@ -5,8 +5,6 @@ module Gitlab
     # Configuration settings and methods for interacting with a PostgreSQL
     # database, with support for multiple databases.
     class Connection
-      DEFAULT_POOL_HEADROOM = 10
-
       attr_reader :scope
 
       # Initializes a new `Database`.
@@ -20,20 +18,6 @@ module Gitlab
         @open_transactions_baseline = 0
       end
 
-      # We configure the database connection pool size automatically based on
-      # the configured concurrency. We also add some headroom, to make sure we
-      # don't run out of connections when more threads besides the 'user-facing'
-      # ones are running.
-      #
-      # Read more about this in
-      # doc/development/database/client_side_connection_pool.md
-      def default_pool_size
-        headroom =
-          (ENV["DB_POOL_HEADROOM"].presence || DEFAULT_POOL_HEADROOM).to_i
-
-        Gitlab::Runtime.max_threads + headroom
-      end
-
       def config
         # The result of this method must not be cached, as other methods may use
         # it after making configuration changes and expect those changes to be
@@ -48,7 +32,7 @@ module Gitlab
       end
 
       def pool_size
-        config[:pool] || default_pool_size
+        config[:pool] || Database.default_pool_size
       end
 
       def username
@@ -77,7 +61,9 @@ module Gitlab
 
       def db_config_with_default_pool_size
         db_config_object = scope.connection_db_config
-        config = db_config_object.configuration_hash.merge(pool: default_pool_size)
+        config = db_config_object
+          .configuration_hash
+          .merge(pool: Database.default_pool_size)
 
         ActiveRecord::DatabaseConfigurations::HashConfig.new(
           db_config_object.env_name,

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/connection_spec.rb

@@ -5,29 +5,14 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::Connection do
   let(:connection) { described_class.new }
 
-  describe '#default_pool_size' do
-    before do
-      allow(Gitlab::Runtime).to receive(:max_threads).and_return(7)
-    end
-
-    it 'returns the max thread size plus a fixed headroom of 10' do
-      expect(connection.default_pool_size).to eq(17)
-    end
-
-    it 'returns the max thread size plus a DB_POOL_HEADROOM if this env var is present' do
-      stub_env('DB_POOL_HEADROOM', '7')
-
-      expect(connection.default_pool_size).to eq(14)
-    end
-  end
-
   describe '#config' do
     it 'returns a HashWithIndifferentAccess' do
       expect(connection.config).to be_an_instance_of(HashWithIndifferentAccess)
     end
 
     it 'returns a default pool size' do
-      expect(connection.config).to include(pool: connection.default_pool_size)
+      expect(connection.config)
+        .to include(pool: Gitlab::Database.default_pool_size)
     end
 
     it 'does not cache its results' do
@@ -43,7 +28,7 @@ RSpec.describe Gitlab::Database::Connection do
       it 'returns the default pool size' do
         expect(connection).to receive(:config).and_return({ pool: nil })
 
-        expect(connection.pool_size).to eq(connection.default_pool_size)
+        expect(connection.pool_size).to eq(Gitlab::Database.default_pool_size)
       end
     end
 
@@ -129,7 +114,7 @@ RSpec.describe Gitlab::Database::Connection do
 
   describe '#db_config_with_default_pool_size' do
     it 'returns db_config with our default pool size' do
-      allow(connection).to receive(:default_pool_size).and_return(9)
+      allow(Gitlab::Database).to receive(:default_pool_size).and_return(9)
 
       expect(connection.db_config_with_default_pool_size.pool).to eq(9)
     end

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/lib/gitlab/database_spec.rb

@@ -15,6 +15,22 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.default_pool_size' do
+    before do
+      allow(Gitlab::Runtime).to receive(:max_threads).and_return(7)
+    end
+
+    it 'returns the max thread size plus a fixed headroom of 10' do
+      expect(described_class.default_pool_size).to eq(17)
+    end
+
+    it 'returns the max thread size plus a DB_POOL_HEADROOM if this env var is present' do
+      stub_env('DB_POOL_HEADROOM', '7')
+
+      expect(described_class.default_pool_size).to eq(14)
+    end
+  end
+
   describe '.has_config?' do
     context 'two tier database config' do
       before 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)