Instrument traditional background migrations

Add a rake task,
gitlab:db:migration_testing:sample_background_migrations,
to sample queued background migrations currently in redis.

Sampling output is stored in tmp/migration-testing/background_migrations
so that the migration testing pipeline can access it.

lib/gitlab/database/migrations/runner.rb

@@ -17,6 +17,10 @@ module Gitlab
             Runner.new(direction: :down, migrations: migrations_for_down, result_dir: BASE_RESULT_DIR.join('down'))
           end
 
+          def background_migrations
+            TestBackgroundRunner.new(result_dir: BASE_RESULT_DIR.join('background_migrations'))
+          end
+
           def migration_context
             @migration_context ||= ApplicationRecord.connection.migration_context
           end

lib/gitlab/database/migrations/test_background_runner.rb

@@ -4,12 +4,10 @@ module Gitlab
   module Database
     module Migrations
       class TestBackgroundRunner
-        # TODO - build a rake task to call this method, and support it in the gitlab-com-database-testing project.
-        # Until then, we will inject a migration with a very high timestamp during database testing
-        # that calls this class to run jobs
-        # See https://gitlab.com/gitlab-org/database-team/gitlab-com-database-testing/-/issues/41 for details
+        attr_reader :result_dir
 
-        def initialize
+        def initialize(result_dir:)
+          @result_dir = result_dir
           @job_coordinator = Gitlab::BackgroundMigration.coordinator_for_database(Gitlab::Database::MAIN_DATABASE_NAME)
         end
 
@@ -24,18 +22,30 @@ module Gitlab
           # without .to_f, we do integer division
           # For example, 3.minutes / 2 == 1.minute whereas 3.minutes / 2.to_f == (1.minute + 30.seconds)
           duration_per_migration_type = for_duration / jobs_to_run.count.to_f
-          jobs_to_run.each do |_migration_name, jobs|
+          jobs_to_run.each do |migration_name, jobs|
             run_until = duration_per_migration_type.from_now
-            jobs.shuffle.each do |j|
-              break if run_until <= Time.current
 
-              run_job(j)
-            end
+            run_jobs_for_migration(migration_name: migration_name, jobs: jobs, run_until: run_until)
           end
         end
 
         private
 
+        def run_jobs_for_migration(migration_name:, jobs:, run_until:)
+          per_background_migration_result_dir = File.join(@result_dir, migration_name)
+
+          instrumentation = Instrumentation.new(result_dir: per_background_migration_result_dir)
+          batch_names = (1..).each.lazy.map { |i| "batch_#{i}"}
+
+          jobs.shuffle.each do |j|
+            break if run_until <= Time.current
+
+            instrumentation.observe(version: nil, name: batch_names.next, connection: ActiveRecord::Migration.connection) do
+              run_job(j)
+            end
+          end
+        end
+
         def run_job(job)
           Gitlab::BackgroundMigration.perform(job.args[0], job.args[1])
         end

lib/tasks/gitlab/db.rake

@@ -307,6 +307,13 @@ namespace :gitlab do
       task down: :environment do
         Gitlab::Database::Migrations::Runner.down.run
       end
+
+      desc 'Sample traditional background migrations with instrumentation'
+      task :sample_background_migrations, [:duration_s] => [:environment] do |_t, args|
+        duration = args[:duration_s]&.to_i&.seconds || 30.minutes # Default of 30 minutes
+
+        Gitlab::Database::Migrations::Runner.background_migrations.run_jobs(for_duration: duration)
+      end
     end
 
     desc 'Run all pending batched migrations'

spec/lib/gitlab/database/migrations/runner_spec.rb

@@ -124,4 +124,16 @@ RSpec.describe Gitlab::Database::Migrations::Runner do
       expect(metadata).to match('version' => described_class::SCHEMA_VERSION)
     end
   end
+
+  describe '.background_migrations' do
+    it 'is a TestBackgroundRunner' do
+      expect(described_class.background_migrations).to be_a(Gitlab::Database::Migrations::TestBackgroundRunner)
+    end
+
+    it 'is configured with a result dir of /background_migrations' do
+      runner = described_class.background_migrations
+
+      expect(runner.result_dir).to eq(described_class::BASE_RESULT_DIR.join( 'background_migrations'))
+    end
+  end
 end

spec/lib/gitlab/database/migrations/test_background_runner_spec.rb

@@ -11,11 +11,17 @@ RSpec.describe Gitlab::Database::Migrations::TestBackgroundRunner, :redis do
     Sidekiq::Testing.disable! { ex.run }
   end
 
+  let(:result_dir) { Dir.mktmpdir }
+
+  after do
+    FileUtils.rm_rf(result_dir)
+  end
+
   context 'without jobs to run' do
     it 'returns immediately' do
-      runner = described_class.new
+      runner = described_class.new(result_dir: result_dir)
       expect(runner).not_to receive(:run_job)
-      described_class.new.run_jobs(for_duration: 1.second)
+      described_class.new(result_dir: result_dir).run_jobs(for_duration: 1.second)
     end
   end
 
@@ -30,7 +36,7 @@ RSpec.describe Gitlab::Database::Migrations::TestBackgroundRunner, :redis do
 
     context 'finding pending background jobs' do
       it 'finds all the migrations' do
-        expect(described_class.new.traditional_background_migrations.to_a.size).to eq(5)
+        expect(described_class.new(result_dir: result_dir).traditional_background_migrations.to_a.size).to eq(5)
       end
     end
 
@@ -53,12 +59,28 @@ RSpec.describe Gitlab::Database::Migrations::TestBackgroundRunner, :redis do
         end
       end
 
+      def expect_recorded_migration_runs(migrations_to_runs)
+        migrations_to_runs.each do |migration, runs|
+          path = File.join(result_dir, migration.name.demodulize)
+          num_subdirs = Pathname(path).children.count(&:directory?)
+          expect(num_subdirs).to eq(runs)
+        end
+      end
+
+      def expect_migration_runs(migrations_to_run_counts)
+        expect_migration_call_counts(migrations_to_run_counts)
+
+        yield
+
+        expect_recorded_migration_runs(migrations_to_run_counts)
+      end
+
       it 'runs the migration class correctly' do
         calls = []
         define_background_migration(migration_name) do |i|
           calls << i
         end
-        described_class.new.run_jobs(for_duration: 1.second) # Any time would work here as we do not advance time
+        described_class.new(result_dir: result_dir).run_jobs(for_duration: 1.second) # Any time would work here as we do not advance time
         expect(calls).to contain_exactly(1, 2, 3, 4, 5)
       end
 
@@ -67,9 +89,9 @@ RSpec.describe Gitlab::Database::Migrations::TestBackgroundRunner, :redis do
           travel(1.minute)
         end
 
-        expect_migration_call_counts(migration => 3)
-
-        described_class.new.run_jobs(for_duration: 3.minutes)
+        expect_migration_runs(migration => 3) do
+          described_class.new(result_dir: result_dir).run_jobs(for_duration: 3.minutes)
+        end
       end
 
       context 'with multiple migrations to run' do
@@ -90,12 +112,12 @@ RSpec.describe Gitlab::Database::Migrations::TestBackgroundRunner, :redis do
             travel(2.minutes)
           end
 
-          expect_migration_call_counts(
+          expect_migration_runs(
             migration => 2, # 1 minute jobs for 90 seconds, can finish the first and start the second
             other_migration => 1 # 2 minute jobs for 90 seconds, past deadline after a single job
-          )
-
-          described_class.new.run_jobs(for_duration: 3.minutes)
+          ) do
+            described_class.new(result_dir: result_dir).run_jobs(for_duration: 3.minutes)
+          end
         end
 
         it 'does not give leftover time to extra migrations' do
@@ -107,12 +129,13 @@ RSpec.describe Gitlab::Database::Migrations::TestBackgroundRunner, :redis do
           other_migration = define_background_migration(other_migration_name) do
             travel(1.minute)
           end
-          expect_migration_call_counts(
+
+          expect_migration_runs(
             migration => 5,
             other_migration => 2
-          )
-
-          described_class.new.run_jobs(for_duration: 3.minutes)
+          ) do
+            described_class.new(result_dir: result_dir).run_jobs(for_duration: 3.minutes)
+          end
         end
       end
     end

spec/tasks/gitlab/db_rake_spec.rb

@@ -538,6 +538,20 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
         subject
       end
     end
+
+    describe '#sample_background_migrations' do
+      it 'delegates to the migration runner with a default sample duration' do
+        expect(::Gitlab::Database::Migrations::Runner).to receive_message_chain(:background_migrations, :run_jobs).with(for_duration: 30.minutes)
+
+        run_rake_task('gitlab:db:migration_testing:sample_background_migrations')
+      end
+
+      it 'delegates to the migration runner with a configured sample duration' do
+        expect(::Gitlab::Database::Migrations::Runner).to receive_message_chain(:background_migrations, :run_jobs).with(for_duration: 100.seconds)
+
+        run_rake_task('gitlab:db:migration_testing:sample_background_migrations', '[100]')
+      end
+    end
   end
 
   describe '#execute_batched_migrations' do