Add latest changes from gitlab-org/gitlab@master

2 files changed, 24 insertions, 12 deletions

diff --git a/lib/gitlab/database/background_migration/batch_optimizer.rb b/lib/gitlab/database/background_migration/batch_optimizer.rb
index 0668490dda8..58c4a214077 100644
--- a/lib/gitlab/database/background_migration/batch_optimizer.rb
+++ b/lib/gitlab/database/background_migration/batch_optimizer.rb
@@ -20,7 +20,8 @@ module Gitlab
         TARGET_EFFICIENCY = (0.9..0.95).freeze
 
         # Lower and upper bound for the batch size
-        ALLOWED_BATCH_SIZE = (1_000..2_000_000).freeze
+        MIN_BATCH_SIZE = 1_000
+        MAX_BATCH_SIZE = 2_000_000
 
         # Limit for the multiplier of the batch size
         MAX_MULTIPLIER = 1.2
@@ -43,7 +44,8 @@ module Gitlab
           return unless Feature.enabled?(:optimize_batched_migrations, type: :ops, default_enabled: :yaml)
 
           if multiplier = batch_size_multiplier
-            migration.batch_size = (migration.batch_size * multiplier).to_i.clamp(ALLOWED_BATCH_SIZE)
+            max_batch = migration.max_batch_size || MAX_BATCH_SIZE
+            migration.batch_size = (migration.batch_size * multiplier).to_i.clamp(MIN_BATCH_SIZE, max_batch)
             migration.save!
           end
         end
diff --git a/lib/gitlab/database/migrations/batched_background_migration_helpers.rb b/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
index dcaf7fad05f..a2a4a37ab87 100644
--- a/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
+++ b/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
@@ -66,6 +66,7 @@ module Gitlab
           batch_max_value: nil,
           batch_class_name: BATCH_CLASS_NAME,
           batch_size: BATCH_SIZE,
+          max_batch_size: nil,
           sub_batch_size: SUB_BATCH_SIZE
         )
 
@@ -86,7 +87,7 @@ module Gitlab
           migration_status = batch_max_value.nil? ? :finished : :active
           batch_max_value ||= batch_min_value
 
-          migration = Gitlab::Database::BackgroundMigration::BatchedMigration.create!(
+          migration = Gitlab::Database::BackgroundMigration::BatchedMigration.new(
             job_class_name: job_class_name,
             table_name: batch_table_name,
             column_name: batch_column_name,
@@ -97,19 +98,28 @@ module Gitlab
             batch_class_name: batch_class_name,
             batch_size: batch_size,
             sub_batch_size: sub_batch_size,
-            status: migration_status)
+            status: migration_status
+          )
 
-          # This guard is necessary since #total_tuple_count was only introduced schema-wise,
-          # after this migration helper had been used for the first time.
-          return migration unless migration.respond_to?(:total_tuple_count)
+          # Below `BatchedMigration` attributes were introduced after the
+          # initial `batched_background_migrations` table was created, so any
+          # migrations that ran relying on initial table schema would not know
+          # about columns introduced later on because this model is not
+          # isolated in migrations, which is why we need to check for existence
+          # of these columns first.
+          if migration.respond_to?(:max_batch_size)
+            migration.max_batch_size = max_batch_size
+          end
 
-          # We keep track of the estimated number of tuples to reason later
-          # about the overall progress of a migration.
-          migration.total_tuple_count = Gitlab::Database::SharedModel.using_connection(connection) do
-            Gitlab::Database::PgClass.for_table(batch_table_name)&.cardinality_estimate
+          if migration.respond_to?(:total_tuple_count)
+            # We keep track of the estimated number of tuples to reason later
+            # about the overall progress of a migration.
+            migration.total_tuple_count = Gitlab::Database::SharedModel.using_connection(connection) do
+              Gitlab::Database::PgClass.for_table(batch_table_name)&.cardinality_estimate
+            end
           end
-          migration.save!
 
+          migration.save!
           migration
         end
       end