Improve pending builds migration to use better batching

db/post_migrate/20210520105029_copy_pending_builds_to_pending_builds_table.rb

@@ -3,41 +3,48 @@
 class CopyPendingBuildsToPendingBuildsTable < ActiveRecord::Migration[6.0]
   include Gitlab::Database::MigrationHelpers
 
-  BUILDS_MAX_SIZE = 1 << (32 - 1)
-  PENDING_BUILDS_BATCH_SIZE = 1000
-  PENDING_BUILDS_MAX_BATCHES = BUILDS_MAX_SIZE / PENDING_BUILDS_BATCH_SIZE
+  BATCH_SIZE = 10000
 
   disable_ddl_transaction!
 
+  class Build < ActiveRecord::Base
+    include EachBatch
+
+    self.table_name = 'ci_builds'
+    self.inheritance_column = :_type_disabled
+
+    scope :pending, -> do
+      where(status: 'pending')
+        .where(type: 'Ci::Build')
+        .where('updated_at > ?', 24.hours.ago)
+        .order('id ASC')
+    end
+  end
+
   def up
-    1.step do |i|
-      inserts = execute <<~SQL
-        WITH pending_builds AS (
-          SELECT id,
-                 project_id
-          FROM ci_builds
-          WHERE status = 'pending'
-            AND type = 'Ci::Build'
-            AND NOT EXISTS (
-              SELECT 1 FROM ci_pending_builds
-                WHERE ci_pending_builds.build_id = ci_builds.id
-            )
-          LIMIT #{PENDING_BUILDS_BATCH_SIZE}
-        ), inserts AS (
+    Build.pending.limit(1).pluck(:id).first.try do |id|
+      Build.where('id >= ?', id).each_batch(of: BATCH_SIZE) do |batch|
+        min_id, max_id = batch.pluck('MIN(id), MAX(id)').first
+
+        execute <<~SQL
+          WITH pending_builds AS (
+            SELECT id,
+                   project_id
+            FROM ci_builds
+            WHERE status = 'pending'
+              AND type = 'Ci::Build'
+              AND NOT EXISTS (
+                SELECT 1 FROM ci_pending_builds
+                  WHERE ci_pending_builds.build_id = ci_builds.id
+              )
+            AND id BETWEEN #{min_id} AND #{max_id}
+          )
           INSERT INTO ci_pending_builds (build_id, project_id)
             SELECT id,
                    project_id
             FROM pending_builds
             ON CONFLICT DO NOTHING
-            RETURNING id
-        )
-        SELECT COUNT(*) FROM inserts;
-      SQL
-
-      break if inserts.values.flatten.first.to_i == 0
-
-      if i > PENDING_BUILDS_MAX_BATCHES
-        raise 'There are too many pending builds in your database! Aborting.'
+        SQL
       end
     end
   end

spec/migrations/copy_pending_builds_to_pending_builds_table_spec.rb

@@ -13,13 +13,14 @@ RSpec.describe CopyPendingBuildsToPendingBuildsTable do
     namespaces.create!(id: 123, name: 'sample', path: 'sample')
     projects.create!(id: 123, name: 'sample', path: 'sample', namespace_id: 123)
 
-    builds.create!(id: 1, project_id: 123, status: 'pending', type: 'Ci::Build')
-    builds.create!(id: 2, project_id: 123, status: 'pending', type: 'GenericCommitStatus')
+    builds.create!(id: 1, project_id: 123, status: 'pending', type: 'Ci::Build', created_at: 4.days.ago, updated_at: 4.days.ago)
+    builds.create!(id: 2, project_id: 123, status: 'pending', type: 'Ci::Build')
     builds.create!(id: 3, project_id: 123, status: 'pending', type: 'GenericCommitStatus')
-    builds.create!(id: 4, project_id: 123, status: 'pending', type: 'Ci::Bridge')
-    builds.create!(id: 5, project_id: 123, status: 'running', type: 'Ci::Build')
+    builds.create!(id: 4, project_id: 123, status: 'pending', type: 'GenericCommitStatus')
+    builds.create!(id: 5, project_id: 123, status: 'pending', type: 'Ci::Bridge')
     builds.create!(id: 6, project_id: 123, status: 'pending', type: 'Ci::Build')
-    builds.create!(id: 7, project_id: 123, status: 'created', type: 'Ci::Build')
+    builds.create!(id: 7, project_id: 123, status: 'running', type: 'Ci::Build')
+    builds.create!(id: 8, project_id: 123, status: 'created', type: 'Ci::Build')
   end
 
   context 'when there are new pending builds present' do
@@ -27,13 +28,13 @@ RSpec.describe CopyPendingBuildsToPendingBuildsTable do
       migrate!
 
       expect(queue.all.count).to eq 2
-      expect(queue.all.pluck(:build_id)).to match_array([1, 6])
+      expect(queue.all.pluck(:build_id)).to match_array([2, 6])
     end
   end
 
   context 'when there are pending builds already migrated present' do
     before do
-      queue.create!(id: 1, build_id: 1, project_id: 123)
+      queue.create!(id: 1, build_id: 2, project_id: 123)
     end
 
     it 'does not copy entries that have already been copied' do
@@ -42,31 +43,33 @@ RSpec.describe CopyPendingBuildsToPendingBuildsTable do
       migrate!
 
       expect(queue.all.count).to eq 2
-      expect(queue.all.pluck(:build_id)).to match_array([1, 6])
+      expect(queue.all.pluck(:build_id)).to match_array([2, 6])
     end
   end
 
   context 'when there is more than one batch of pending builds to migrate' do
     before do
-      stub_const("#{described_class}::PENDING_BUILDS_BATCH_SIZE", 1)
+      stub_const("#{described_class}::BATCH_SIZE", 2)
     end
 
     it 'correctly migrates data and exits after doing that' do
       migrate!
 
       expect(queue.all.count).to eq 2
-      expect(queue.all.pluck(:build_id)).to match_array([1, 6])
+      expect(queue.all.pluck(:build_id)).to match_array([2, 6])
     end
   end
 
-  context 'when batches limit has been exceeded' do
+  context 'when an old build has been recently updated' do
     before do
-      stub_const("#{described_class}::PENDING_BUILDS_BATCH_SIZE", 1)
-      stub_const("#{described_class}::PENDING_BUILDS_MAX_BATCHES", 1)
+      builds.find(1).touch
     end
 
-    it 'raises an error' do
-      expect { migrate! }.to raise_error(StandardError, /too many pending builds/)
+    it 'migrates it as well' do
+      migrate!
+
+      expect(queue.all.count).to eq 3
+      expect(queue.all.pluck(:build_id)).to match_array([1, 2, 6])
     end
   end