Add latest changes from gitlab-org/gitlab@master

6 files changed, 123 insertions, 90 deletions

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index ce198061a49..f3f0d0305d3 100644
--- a/lib/gitlab/database.rb
+++ b/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
diff --git a/lib/gitlab/database/connection.rb b/lib/gitlab/database/connection.rb
index 348b7207b23..6eac115be13 100644
--- a/lib/gitlab/database/connection.rb
+++ b/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,
diff --git a/lib/gitlab/database/load_balancing.rb b/lib/gitlab/database/load_balancing.rb
index 65c35b5aaf2..d4f168ba188 100644
--- a/lib/gitlab/database/load_balancing.rb
+++ b/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,
diff --git a/lib/gitlab/database/load_balancing/configuration.rb b/lib/gitlab/database/load_balancing/configuration.rb
new file mode 100644
index 00000000000..6bd94bad5cf
--- /dev/null
+++ b/lib/gitlab/database/load_balancing/configuration.rb
@@ -0,0 +1,79 @@
+# 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
diff --git a/lib/gitlab/database/load_balancing/host.rb b/lib/gitlab/database/load_balancing/host.rb
index 4c5357ae8e3..670a373ded1 100644
--- a/lib/gitlab/database/load_balancing/host.rb
+++ b/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
diff --git a/lib/gitlab/database/load_balancing/load_balancer.rb b/lib/gitlab/database/load_balancing/load_balancer.rb
index cdaeebde4df..0fb8faae8ce 100644
--- a/lib/gitlab/database/load_balancing/load_balancer.rb
+++ b/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
           )