Add method to split batched migration jobs

Adds a method to split the jobs to prevent failures due to long running
jobs not being able to complete successfully

lib/gitlab/database/background_migration/batched_job.rb

@@ -44,6 +44,52 @@ module Gitlab
           # TODO: Switch to individual job interval (prereq: https://gitlab.com/gitlab-org/gitlab/-/issues/328801)
           duration.to_f / batched_migration.interval
         end
+
+        def split_and_retry!
+          with_lock do
+            raise 'Only failed jobs can be split' unless failed?
+
+            new_batch_size = batch_size / 2
+
+            raise 'Job cannot be split further' if new_batch_size < 1
+
+            batching_strategy = batched_migration.batch_class.new
+            next_batch_bounds = batching_strategy.next_batch(
+              batched_migration.table_name,
+              batched_migration.column_name,
+              batch_min_value: min_value,
+              batch_size: new_batch_size
+            )
+            midpoint = next_batch_bounds.last
+
+            # We don't want the midpoint to go over the existing max_value because
+            # those IDs would already be in the next batched migration job.
+            # This could happen when a lot of records in the current batch are deleted.
+            #
+            # In this case, we just lower the batch size so that future calls to this
+            # method could eventually split the job if it continues to fail.
+            if midpoint >= max_value
+              update!(batch_size: new_batch_size, status: :pending)
+            else
+              old_max_value = max_value
+
+              update!(
+                batch_size: new_batch_size,
+                max_value: midpoint,
+                attempts: 0,
+                status: :pending,
+                started_at: nil,
+                finished_at: nil,
+                metrics: {}
+              )
+
+              new_record = dup
+              new_record.min_value = midpoint.next
+              new_record.max_value = old_max_value
+              new_record.save!
+            end
+          end
+        end
       end
     end
   end

spec/lib/gitlab/database/background_migration/batched_job_spec.rb

@@ -124,4 +124,73 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedJob, type: :model d
       end
     end
   end
+
+  describe '#split_and_retry!' do
+    let!(:job) { create(:batched_background_migration_job, batch_size: 10, min_value: 6, max_value: 15, status: :failed) }
+
+    it 'splits the job into two and marks them as pending' do
+      allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
+        allow(batch_class).to receive(:next_batch).with(anything, anything, batch_min_value: 6, batch_size: 5).and_return([6, 10])
+      end
+
+      expect { job.split_and_retry! }.to change { described_class.count }.by(1)
+
+      expect(job).to have_attributes(
+        min_value: 6,
+        max_value: 10,
+        batch_size: 5,
+        status: 'pending',
+        attempts: 0,
+        started_at: nil,
+        finished_at: nil,
+        metrics: {}
+      )
+
+      new_job = described_class.last
+
+      expect(new_job).to have_attributes(
+        batched_background_migration_id: job.batched_background_migration_id,
+        min_value: 11,
+        max_value: 15,
+        batch_size: 5,
+        status: 'pending',
+        attempts: 0,
+        started_at: nil,
+        finished_at: nil,
+        metrics: {}
+      )
+      expect(new_job.created_at).not_to eq(job.created_at)
+    end
+
+    context 'when job is not failed' do
+      let!(:job) { create(:batched_background_migration_job, status: :succeeded) }
+
+      it 'raises an exception' do
+        expect { job.split_and_retry! }.to raise_error 'Only failed jobs can be split'
+      end
+    end
+
+    context 'when batch size is already 1' do
+      let!(:job) { create(:batched_background_migration_job, batch_size: 1, status: :failed) }
+
+      it 'raises an exception' do
+        expect { job.split_and_retry! }.to raise_error 'Job cannot be split further'
+      end
+    end
+
+    context 'when computed midpoint is larger than the max value of the batch' do
+      before do
+        allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
+          allow(batch_class).to receive(:next_batch).with(anything, anything, batch_min_value: 6, batch_size: 5).and_return([6, 16])
+        end
+      end
+
+      it 'lowers the batch size and marks the job as pending' do
+        expect { job.split_and_retry! }.not_to change { described_class.count }
+
+        expect(job.batch_size).to eq(5)
+        expect(job.status).to eq('pending')
+      end
+    end
+  end
 end