Schedule stages statuses bg migrations in batches

db/post_migrate/20170630111158_migrate_stages_statuses.rb

@@ -6,21 +6,22 @@ class MigrateStagesStatuses < ActiveRecord::Migration
   disable_ddl_transaction!
 
   BATCH_SIZE = 10000
+  RANGE_SIZE = 1000
   MIGRATION = 'MigrateStageStatus'.freeze
 
   class Stage < ActiveRecord::Base
     self.table_name = 'ci_stages'
+    include ::EachBatch
   end
 
   def up
-    index = 1
-
-    Stage.where(status: nil).in_batches(of: BATCH_SIZE) do |relation|
-      jobs = relation.pluck(:id).map { |id| [MIGRATION, [id]] }
+    Stage.where(status: nil).each_batch(of: BATCH_SIZE) do |relation, index|
+      relation.each_batch(of: RANGE_SIZE) do |batch|
+        range = relation.pluck('MIN(id)', 'MAX(id)').first
         schedule = index * 5.minutes
-      index += 1
 
-      BackgroundMigrationWorker.perform_bulk_in(schedule, jobs)
+        BackgroundMigrationWorker.perform_in(schedule, MIGRATION, range)
+      end
     end
   end
 

lib/gitlab/background_migration/migrate_stage_status.rb

@@ -58,7 +58,7 @@ module Gitlab
         end
       end
 
-      def perform(id)
+      def perform(start_id, stop_id)
         status_sql = Build
           .where('ci_builds.commit_id = ci_stages.pipeline_id')
           .where('ci_builds.stage = ci_stages.name')
@@ -66,7 +66,8 @@ module Gitlab
 
         sql = <<-SQL
           UPDATE ci_stages SET status = (#{status_sql})
-            WHERE ci_stages.id = #{id.to_i}
+            WHERE ci_stages.status IS NULL
+            AND ci_stages.id BETWEEN #{start_id.to_i} AND #{stop_id.to_i}
         SQL
 
         ActiveRecord::Base.connection.execute(sql)

spec/migrations/migrate_stages_statuses_spec.rb

@@ -12,6 +12,7 @@ describe MigrateStagesStatuses, :migration do
 
   before do
     stub_const("#{described_class.name}::BATCH_SIZE", 2)
+    stub_const("#{described_class.name}::RANGE_SIZE", 1)
 
     projects.create!(id: 1, name: 'gitlab1', path: 'gitlab1')
     projects.create!(id: 2, name: 'gitlab2', path: 'gitlab2')
@@ -49,9 +50,10 @@ describe MigrateStagesStatuses, :migration do
       Timecop.freeze do
         migrate!
 
-        expect(described_class::MIGRATION).to be_scheduled_migration(5.minutes, 1)
-        expect(described_class::MIGRATION).to be_scheduled_migration(5.minutes, 2)
-        expect(described_class::MIGRATION).to be_scheduled_migration(10.minutes, 3)
+        puts BackgroundMigrationWorker.jobs.inspect
+        expect(described_class::MIGRATION).to be_scheduled_migration(5.minutes, 1, 1)
+        expect(described_class::MIGRATION).to be_scheduled_migration(5.minutes, 2, 2)
+        expect(described_class::MIGRATION).to be_scheduled_migration(10.minutes, 3, 3)
         expect(BackgroundMigrationWorker.jobs.size).to eq 3
       end
     end