Use with_lock_retries at the migration level

For transactional migrations, apply with-lock-retries methodology at the
full transaction level. This allows us to remove usage of
with-lock-retries from helpers and we don't need to use it inside
migrations explicitly.

The overall intent here is to remove subtransactions.

See https://gitlab.com/gitlab-org/gitlab/-/issues/339115

config/initializers/active_record_migrations.rb

+# frozen_string_literal: true
+
+Gitlab::Database::Migrations::LockRetryMixin.patch!

lib/gitlab/database/migration.rb

@@ -3,6 +3,22 @@
 module Gitlab
   module Database
     class Migration
+      module LockRetriesConcern
+        extend ActiveSupport::Concern
+
+        class_methods do
+          def enable_lock_retries!
+            @enable_lock_retries = true # rubocop:disable Gitlab/ModuleWithInstanceVariables
+          end
+
+          def enable_lock_retries?
+            @enable_lock_retries
+          end
+        end
+
+        delegate :enable_lock_retries?, to: :class
+      end
+
       # This implements a simple versioning scheme for migration helpers.
       #
       # We need to be able to version helpers, so we can change their behavior without
@@ -19,6 +35,7 @@ module Gitlab
       # However, this hasn't been strictly formalized yet.
       MIGRATION_CLASSES = {
         1.0 => Class.new(ActiveRecord::Migration[6.1]) do
+          include LockRetriesConcern
           include Gitlab::Database::MigrationHelpers::V2
         end
       }.freeze

lib/gitlab/database/migrations/lock_retry_mixin.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module LockRetryMixin
+        module ActiveRecordMigrationProxyLockRetries
+          def migration_class
+            migration.class
+          end
+
+          def enable_lock_retries?
+            # regular AR migrations don't have this,
+            # only ones inheriting from Gitlab::Database::Migration have
+            return false unless migration.respond_to?(:enable_lock_retries?)
+
+            migration.enable_lock_retries?
+          end
+        end
+
+        module ActiveRecordMigratorLockRetries
+          # We patch the original method to start a transaction
+          # using the WithLockRetries methodology for the whole migration.
+          def ddl_transaction(migration, &block)
+            if use_transaction?(migration) && migration.enable_lock_retries?
+              Gitlab::Database::WithLockRetries.new(
+                klass: migration.migration_class,
+                logger: Gitlab::BackgroundMigration::Logger
+              ).run(raise_on_exhaustion: false, &block)
+            else
+              super
+            end
+          end
+        end
+
+        def self.patch!
+          ActiveRecord::MigrationProxy.prepend(ActiveRecordMigrationProxyLockRetries)
+          ActiveRecord::Migrator.prepend(ActiveRecordMigratorLockRetries)
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/migration_spec.rb

@@ -14,6 +14,10 @@ RSpec.describe Gitlab::Database::Migration do
       it 'includes migration helpers version 2' do
         expect(subject.included_modules).to include(Gitlab::Database::MigrationHelpers::V2)
       end
+
+      it 'includes LockRetriesConcern' do
+        expect(subject.included_modules).to include(Gitlab::Database::Migration::LockRetriesConcern)
+      end
     end
 
     context 'unknown version' do
@@ -31,4 +35,34 @@ RSpec.describe Gitlab::Database::Migration do
       expect(described_class[described_class.current_version].superclass).to eq(ActiveRecord::Migration::Current)
     end
   end
+
+  describe Gitlab::Database::Migration::LockRetriesConcern do
+    subject { class_def.new }
+
+    context 'when not explicitly called' do
+      let(:class_def) do
+        Class.new do
+          include Gitlab::Database::Migration::LockRetriesConcern
+        end
+      end
+
+      it 'does not disable lock retries by default' do
+        expect(subject.enable_lock_retries?).not_to be_truthy
+      end
+    end
+
+    context 'when explicitly disabled' do
+      let(:class_def) do
+        Class.new do
+          include Gitlab::Database::Migration::LockRetriesConcern
+
+          enable_lock_retries!
+        end
+      end
+
+      it 'does not disable lock retries by default' do
+        expect(subject.enable_lock_retries?).to be_truthy
+      end
+    end
+  end
 end

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

+# frozen_string_literal: true
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
+  describe Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigrationProxyLockRetries do
+    let(:migration) { double }
+    let(:return_value) { double }
+    let(:class_def) do
+      Class.new do
+        include Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigrationProxyLockRetries
+
+        attr_reader :migration
+
+        def initialize(migration)
+          @migration = migration
+        end
+      end
+    end
+
+    describe '#enable_lock_retries?' do
+      subject { class_def.new(migration).enable_lock_retries? }
+
+      it 'delegates to #migration' do
+        expect(migration).to receive(:enable_lock_retries?).and_return(return_value)
+
+        result = subject
+
+        expect(result).to eq(return_value)
+      end
+    end
+
+    describe '#migration_class' do
+      subject { class_def.new(migration).migration_class }
+
+      it 'retrieves actual migration class from #migration' do
+        expect(migration).to receive(:class).and_return(return_value)
+
+        result = subject
+
+        expect(result).to eq(return_value)
+      end
+    end
+  end
+
+  describe Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigratorLockRetries do
+    let(:class_def) do
+      Class.new do
+        attr_reader :receiver
+
+        def initialize(receiver)
+          @receiver = receiver
+        end
+
+        def ddl_transaction(migration, &block)
+          receiver.ddl_transaction(migration, &block)
+        end
+
+        def use_transaction?(migration)
+          receiver.use_transaction?(migration)
+        end
+      end.prepend(Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigratorLockRetries)
+    end
+
+    subject { class_def.new(receiver) }
+
+    before do
+      allow(migration).to receive(:migration_class).and_return('TestClass')
+      allow(receiver).to receive(:ddl_transaction)
+    end
+
+    context 'with transactions disabled' do
+      let(:migration) { double('migration', enable_lock_retries?: false) }
+      let(:receiver) { double('receiver', use_transaction?: false)}
+
+      it 'calls super method' do
+        p = proc { }
+
+        expect(receiver).to receive(:ddl_transaction).with(migration, &p)
+
+        subject.ddl_transaction(migration, &p)
+      end
+    end
+
+    context 'with transactions enabled, but lock retries disabled' do
+      let(:receiver) { double('receiver', use_transaction?: true)}
+      let(:migration) { double('migration', enable_lock_retries?: false) }
+
+      it 'calls super method' do
+        p = proc { }
+
+        expect(receiver).to receive(:ddl_transaction).with(migration, &p)
+
+        subject.ddl_transaction(migration, &p)
+      end
+    end
+
+    context 'with transactions enabled and lock retries enabled' do
+      let(:receiver) { double('receiver', use_transaction?: true)}
+      let(:migration) { double('migration', enable_lock_retries?: true) }
+
+      it 'calls super method' do
+        p = proc { }
+
+        expect(receiver).not_to receive(:ddl_transaction)
+        expect_next_instance_of(Gitlab::Database::WithLockRetries) do |retries|
+          expect(retries).to receive(:run).with(raise_on_exhaustion: false, &p)
+        end
+
+        subject.ddl_transaction(migration, &p)
+      end
+    end
+  end
+
+  describe '.patch!' do
+    subject { described_class.patch! }
+
+    it 'patches MigrationProxy' do
+      expect(ActiveRecord::MigrationProxy).to receive(:prepend).with(Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigrationProxyLockRetries)
+
+      subject
+    end
+
+    it 'patches Migrator' do
+      expect(ActiveRecord::Migrator).to receive(:prepend).with(Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigratorLockRetries)
+
+      subject
+    end
+  end
+end