1 files changed, 39 insertions, 2 deletions
diff --git a/lib/gitlab/database/background_migration/batched_migration.rb b/lib/gitlab/database/background_migration/batched_migration.rb
index 4aa33ed7946..e85162f355e 100644
--- a/lib/gitlab/database/background_migration/batched_migration.rb
+++ b/lib/gitlab/database/background_migration/batched_migration.rb
@@ -20,9 +20,12 @@ module Gitlab
           paused: 0,
           active: 1,
           aborted: 2,
-          finished: 3
+          finished: 3,
+          failed: 4
         }
 
+        attribute :pause_ms, :integer, default: 100
+
         def self.active_migration
           active.queue_order.first
         end
@@ -35,7 +38,13 @@ module Gitlab
         end
 
         def create_batched_job!(min, max)
-          batched_jobs.create!(min_value: min, max_value: max, batch_size: batch_size, sub_batch_size: sub_batch_size)
+          batched_jobs.create!(
+            min_value: min,
+            max_value: max,
+            batch_size: batch_size,
+            sub_batch_size: sub_batch_size,
+            pause_ms: pause_ms
+          )
         end
 
         def next_min_value
@@ -58,12 +67,40 @@ module Gitlab
           write_attribute(:batch_class_name, class_name.demodulize)
         end
 
+        def migrated_tuple_count
+          batched_jobs.succeeded.sum(:batch_size)
+        end
+
         def prometheus_labels
           @prometheus_labels ||= {
             migration_id: id,
             migration_identifier: "%s/%s.%s" % [job_class_name, table_name, column_name]
           }
         end
+
+        def smoothed_time_efficiency(number_of_jobs: 10, alpha: 0.2)
+          jobs = batched_jobs.successful_in_execution_order.reverse_order.limit(number_of_jobs)
+
+          return if jobs.size < number_of_jobs
+
+          efficiencies = jobs.map(&:time_efficiency).reject(&:nil?).each_with_index
+
+          dividend = efficiencies.reduce(0) do |total, (job_eff, i)|
+            total + job_eff * (1 - alpha)**i
+          end
+
+          divisor = efficiencies.reduce(0) do |total, (job_eff, i)|
+            total + (1 - alpha)**i
+          end
+
+          return if divisor == 0
+
+          (dividend / divisor).round(2)
+        end
+
+        def optimize!
+          BatchOptimizer.new(self).optimize!
+        end
       end
     end
   end