Merge branch 'db-fixes-zeitwerk' into 'master'

Fix DB-related loading of gems and classes

See merge request gitlab-org/gitlab!65317

lib/gitlab/auth/u2f_webauthn_converter.rb

 # frozen_string_literal: true
 
+require 'webauthn/u2f_migrator'
+
 module Gitlab
   module Auth
     class U2fWebauthnConverter

lib/gitlab/background_migration/migrate_u2f_webauthn.rb

 # frozen_string_literal: true
 # rubocop:disable Style/Documentation
-require "webauthn/u2f_migrator"
 
 module Gitlab
   module BackgroundMigration

lib/gitlab/database/batch_count.rb

@@ -41,159 +41,5 @@ module Gitlab
         include BatchCount
       end
     end
-
-    class BatchCounter
-      FALLBACK = -1
-      MIN_REQUIRED_BATCH_SIZE = 1_250
-      DEFAULT_SUM_BATCH_SIZE = 1_000
-      MAX_ALLOWED_LOOPS = 10_000
-      SLEEP_TIME_IN_SECONDS = 0.01 # 10 msec sleep
-      ALLOWED_MODES = [:itself, :distinct].freeze
-      FALLBACK_FINISH = 0
-      OFFSET_BY_ONE = 1
-
-      # Each query should take < 500ms https://gitlab.com/gitlab-org/gitlab/-/merge_requests/22705
-      DEFAULT_DISTINCT_BATCH_SIZE = 10_000
-      DEFAULT_BATCH_SIZE = 100_000
-
-      def initialize(relation, column: nil, operation: :count, operation_args: nil)
-        @relation = relation
-        @column = column || relation.primary_key
-        @operation = operation
-        @operation_args = operation_args
-      end
-
-      def unwanted_configuration?(finish, batch_size, start)
-        (@operation == :count && batch_size <= MIN_REQUIRED_BATCH_SIZE) ||
-          (@operation == :sum && batch_size < DEFAULT_SUM_BATCH_SIZE) ||
-          (finish - start) / batch_size >= MAX_ALLOWED_LOOPS ||
-          start >= finish
-      end
-
-      def count(batch_size: nil, mode: :itself, start: nil, finish: nil)
-        raise 'BatchCount can not be run inside a transaction' if ActiveRecord::Base.connection.transaction_open?
-
-        check_mode!(mode)
-
-        # non-distinct have better performance
-        batch_size ||= batch_size_for_mode_and_operation(mode, @operation)
-
-        start = actual_start(start)
-        finish = actual_finish(finish)
-
-        raise "Batch counting expects positive values only for #{@column}" if start < 0 || finish < 0
-        return FALLBACK if unwanted_configuration?(finish, batch_size, start)
-
-        results = nil
-        batch_start = start
-
-        while batch_start < finish
-          begin
-            batch_end = [batch_start + batch_size, finish].min
-            batch_relation = build_relation_batch(batch_start, batch_end, mode)
-
-            op_args = @operation_args
-            if @operation == :count && @operation_args.blank? && use_loose_index_scan_for_distinct_values?(mode)
-              op_args = [Gitlab::Database::LooseIndexScanDistinctCount::COLUMN_ALIAS]
-            end
-
-            results = merge_results(results, batch_relation.send(@operation, *op_args)) # rubocop:disable GitlabSecurity/PublicSend
-            batch_start = batch_end
-          rescue ActiveRecord::QueryCanceled => error
-            # retry with a safe batch size & warmer cache
-            if batch_size >= 2 * MIN_REQUIRED_BATCH_SIZE
-              batch_size /= 2
-            else
-              log_canceled_batch_fetch(batch_start, mode, batch_relation.to_sql, error)
-              return FALLBACK
-            end
-          rescue Gitlab::Database::LooseIndexScanDistinctCount::ColumnConfigurationError => error
-            Gitlab::AppJsonLogger
-              .error(
-                event: 'batch_count',
-                relation: @relation.table_name,
-                operation: @operation,
-                operation_args: @operation_args,
-                mode: mode,
-                message: "LooseIndexScanDistinctCount column error: #{error.message}"
-              )
-
-            return FALLBACK
-          end
-
-          sleep(SLEEP_TIME_IN_SECONDS)
-        end
-
-        results
-      end
-
-      def merge_results(results, object)
-        return object unless results
-
-        if object.is_a?(Hash)
-          results.merge!(object) { |_, a, b| a + b }
-        else
-          results + object
-        end
-      end
-
-      private
-
-      def build_relation_batch(start, finish, mode)
-        if use_loose_index_scan_for_distinct_values?(mode)
-          Gitlab::Database::LooseIndexScanDistinctCount.new(@relation, @column).build_query(from: start, to: finish)
-        else
-          @relation.select(@column).public_send(mode).where(between_condition(start, finish)) # rubocop:disable GitlabSecurity/PublicSend
-        end
-      end
-
-      def batch_size_for_mode_and_operation(mode, operation)
-        return DEFAULT_SUM_BATCH_SIZE if operation == :sum
-
-        mode == :distinct ? DEFAULT_DISTINCT_BATCH_SIZE : DEFAULT_BATCH_SIZE
-      end
-
-      def between_condition(start, finish)
-        return @column.between(start...finish) if @column.is_a?(Arel::Attributes::Attribute)
-
-        { @column => start...finish }
-      end
-
-      def actual_start(start)
-        start || @relation.unscope(:group, :having).minimum(@column) || 0
-      end
-
-      def actual_finish(finish)
-        (finish || @relation.unscope(:group, :having).maximum(@column) || FALLBACK_FINISH) + OFFSET_BY_ONE
-      end
-
-      def check_mode!(mode)
-        raise "The mode #{mode.inspect} is not supported" unless ALLOWED_MODES.include?(mode)
-        raise 'Use distinct count for optimized distinct counting' if @relation.limit(1).distinct_value.present? && mode != :distinct
-        raise 'Use distinct count only with non id fields' if @column == :id && mode == :distinct
-      end
-
-      def log_canceled_batch_fetch(batch_start, mode, query, error)
-        Gitlab::AppJsonLogger
-          .error(
-            event: 'batch_count',
-            relation: @relation.table_name,
-            operation: @operation,
-            operation_args: @operation_args,
-            start: batch_start,
-            mode: mode,
-            query: query,
-            message: "Query has been canceled with message: #{error.message}"
-          )
-      end
-
-      def use_loose_index_scan_for_distinct_values?(mode)
-        Feature.enabled?(:loose_index_scan_for_distinct_values) && not_group_by_query? && mode == :distinct
-      end
-
-      def not_group_by_query?
-        !@relation.is_a?(ActiveRecord::Relation) || @relation.group_values.blank?
-      end
-    end
   end
 end

lib/gitlab/database/batch_counter.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class BatchCounter
+      FALLBACK = -1
+      MIN_REQUIRED_BATCH_SIZE = 1_250
+      DEFAULT_SUM_BATCH_SIZE = 1_000
+      MAX_ALLOWED_LOOPS = 10_000
+      SLEEP_TIME_IN_SECONDS = 0.01 # 10 msec sleep
+      ALLOWED_MODES = [:itself, :distinct].freeze
+      FALLBACK_FINISH = 0
+      OFFSET_BY_ONE = 1
+
+      # Each query should take < 500ms https://gitlab.com/gitlab-org/gitlab/-/merge_requests/22705
+      DEFAULT_DISTINCT_BATCH_SIZE = 10_000
+      DEFAULT_BATCH_SIZE = 100_000
+
+      def initialize(relation, column: nil, operation: :count, operation_args: nil)
+        @relation = relation
+        @column = column || relation.primary_key
+        @operation = operation
+        @operation_args = operation_args
+      end
+
+      def unwanted_configuration?(finish, batch_size, start)
+        (@operation == :count && batch_size <= MIN_REQUIRED_BATCH_SIZE) ||
+          (@operation == :sum && batch_size < DEFAULT_SUM_BATCH_SIZE) ||
+          (finish - start) / batch_size >= MAX_ALLOWED_LOOPS ||
+          start >= finish
+      end
+
+      def count(batch_size: nil, mode: :itself, start: nil, finish: nil)
+        raise 'BatchCount can not be run inside a transaction' if ActiveRecord::Base.connection.transaction_open?
+
+        check_mode!(mode)
+
+        # non-distinct have better performance
+        batch_size ||= batch_size_for_mode_and_operation(mode, @operation)
+
+        start = actual_start(start)
+        finish = actual_finish(finish)
+
+        raise "Batch counting expects positive values only for #{@column}" if start < 0 || finish < 0
+        return FALLBACK if unwanted_configuration?(finish, batch_size, start)
+
+        results = nil
+        batch_start = start
+
+        while batch_start < finish
+          begin
+            batch_end = [batch_start + batch_size, finish].min
+            batch_relation = build_relation_batch(batch_start, batch_end, mode)
+
+            op_args = @operation_args
+            if @operation == :count && @operation_args.blank? && use_loose_index_scan_for_distinct_values?(mode)
+              op_args = [Gitlab::Database::LooseIndexScanDistinctCount::COLUMN_ALIAS]
+            end
+
+            results = merge_results(results, batch_relation.send(@operation, *op_args)) # rubocop:disable GitlabSecurity/PublicSend
+            batch_start = batch_end
+          rescue ActiveRecord::QueryCanceled => error
+            # retry with a safe batch size & warmer cache
+            if batch_size >= 2 * MIN_REQUIRED_BATCH_SIZE
+              batch_size /= 2
+            else
+              log_canceled_batch_fetch(batch_start, mode, batch_relation.to_sql, error)
+              return FALLBACK
+            end
+          rescue Gitlab::Database::LooseIndexScanDistinctCount::ColumnConfigurationError => error
+            Gitlab::AppJsonLogger
+              .error(
+                event: 'batch_count',
+                relation: @relation.table_name,
+                operation: @operation,
+                operation_args: @operation_args,
+                mode: mode,
+                message: "LooseIndexScanDistinctCount column error: #{error.message}"
+              )
+
+            return FALLBACK
+          end
+
+          sleep(SLEEP_TIME_IN_SECONDS)
+        end
+
+        results
+      end
+
+      def merge_results(results, object)
+        return object unless results
+
+        if object.is_a?(Hash)
+          results.merge!(object) { |_, a, b| a + b }
+        else
+          results + object
+        end
+      end
+
+      private
+
+      def build_relation_batch(start, finish, mode)
+        if use_loose_index_scan_for_distinct_values?(mode)
+          Gitlab::Database::LooseIndexScanDistinctCount.new(@relation, @column).build_query(from: start, to: finish)
+        else
+          @relation.select(@column).public_send(mode).where(between_condition(start, finish)) # rubocop:disable GitlabSecurity/PublicSend
+        end
+      end
+
+      def batch_size_for_mode_and_operation(mode, operation)
+        return DEFAULT_SUM_BATCH_SIZE if operation == :sum
+
+        mode == :distinct ? DEFAULT_DISTINCT_BATCH_SIZE : DEFAULT_BATCH_SIZE
+      end
+
+      def between_condition(start, finish)
+        return @column.between(start...finish) if @column.is_a?(Arel::Attributes::Attribute)
+
+        { @column => start...finish }
+      end
+
+      def actual_start(start)
+        start || @relation.unscope(:group, :having).minimum(@column) || 0
+      end
+
+      def actual_finish(finish)
+        (finish || @relation.unscope(:group, :having).maximum(@column) || FALLBACK_FINISH) + OFFSET_BY_ONE
+      end
+
+      def check_mode!(mode)
+        raise "The mode #{mode.inspect} is not supported" unless ALLOWED_MODES.include?(mode)
+        raise 'Use distinct count for optimized distinct counting' if @relation.limit(1).distinct_value.present? && mode != :distinct
+        raise 'Use distinct count only with non id fields' if @column == :id && mode == :distinct
+      end
+
+      def log_canceled_batch_fetch(batch_start, mode, query, error)
+        Gitlab::AppJsonLogger
+          .error(
+            event: 'batch_count',
+            relation: @relation.table_name,
+            operation: @operation,
+            operation_args: @operation_args,
+            start: batch_start,
+            mode: mode,
+            query: query,
+            message: "Query has been canceled with message: #{error.message}"
+          )
+      end
+
+      def use_loose_index_scan_for_distinct_values?(mode)
+        Feature.enabled?(:loose_index_scan_for_distinct_values) && not_group_by_query? && mode == :distinct
+      end
+
+      def not_group_by_query?
+        !@relation.is_a?(ActiveRecord::Relation) || @relation.group_values.blank?
+      end
+    end
+  end
+end

lib/gitlab/database/load_balancing/srv_resolver.rb

 # frozen_string_literal: true
 
+require 'net/dns'
+
 module Gitlab
   module Database
     module LoadBalancing

lib/gitlab/database/postgresql_adapter/force_disconnectable_mixin.rb

@@ -22,7 +22,7 @@ module Gitlab
         end
 
         def force_disconnect_timer
-          @force_disconnect_timer ||= ConnectionTimer.starting_now
+          @force_disconnect_timer ||= ::Gitlab::Database::ConnectionTimer.starting_now
         end
       end
     end

spec/lib/gitlab/background_migration/migrate_u2f_webauthn_spec.rb

@@ -2,6 +2,8 @@
 
 require 'spec_helper'
 
+require 'webauthn/u2f_migrator'
+
 RSpec.describe Gitlab::BackgroundMigration::MigrateU2fWebauthn, :migration, schema: 20200925125321 do
   let(:users) { table(:users) }
 

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -2230,16 +2230,14 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
   describe '#backfill_iids' do
     include MigrationsHelpers
 
-    before do
-      stub_const('Issue', Class.new(ActiveRecord::Base))
-
-      Issue.class_eval do
+    let(:issue_class) do
+      Class.new(ActiveRecord::Base) do
         include AtomicInternalId
 
         self.table_name = 'issues'
         self.inheritance_column = :_type_disabled
 
-        belongs_to :project, class_name: "::Project"
+        belongs_to :project, class_name: "::Project", inverse_of: nil
 
         has_internal_id :iid,
           scope: :project,
@@ -2262,7 +2260,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
       model.backfill_iids('issues')
 
-      issue = Issue.create!(project_id: project.id)
+      issue = issue_class.create!(project_id: project.id)
 
       expect(issue.iid).to eq(1)
     end
@@ -2273,7 +2271,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
       model.backfill_iids('issues')
 
-      issue_b = Issue.create!(project_id: project.id)
+      issue_b = issue_class.create!(project_id: project.id)
 
       expect(issue_a.reload.iid).to eq(1)
       expect(issue_b.iid).to eq(2)
@@ -2288,8 +2286,8 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
       model.backfill_iids('issues')
 
-      issue_a = Issue.create!(project_id: project_a.id)
-      issue_b = Issue.create!(project_id: project_b.id)
+      issue_a = issue_class.create!(project_id: project_a.id)
+      issue_b = issue_class.create!(project_id: project_b.id)
 
       expect(issue_a.iid).to eq(2)
       expect(issue_b.iid).to eq(3)
@@ -2303,7 +2301,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
         model.backfill_iids('issues')
 
-        issue_b = Issue.create!(project_id: project_b.id)
+        issue_b = issue_class.create!(project_id: project_b.id)
 
         expect(issue_a.reload.iid).to eq(1)
         expect(issue_b.reload.iid).to eq(1)