Disallow subtrans for v2/with_lock_retries

lib/gitlab/database/migration_helpers.rb

@@ -380,6 +380,8 @@ module Gitlab
       # The timings can be controlled via the +timing_configuration+ parameter.
       # If the lock was not acquired within the retry period, a last attempt is made without using +lock_timeout+.
       #
+      # Note this helper uses subtransactions when run inside an already open transaction.
+      #
       # ==== Examples
       #   # Invoking without parameters
       #   with_lock_retries do

lib/gitlab/database/migration_helpers/v2.rb

@@ -6,6 +6,44 @@ module Gitlab
       module V2
         include Gitlab::Database::MigrationHelpers
 
+        # Executes the block with a retry mechanism that alters the +lock_timeout+ and +sleep_time+ between attempts.
+        # The timings can be controlled via the +timing_configuration+ parameter.
+        # If the lock was not acquired within the retry period, a last attempt is made without using +lock_timeout+.
+        #
+        # In order to retry the block, the method wraps the block into a transaction.
+        # Note it cannot be used inside an already open transaction and will raise an error in that case.
+        #
+        # ==== Examples
+        #   # Invoking without parameters
+        #   with_lock_retries do
+        #     drop_table :my_table
+        #   end
+        #
+        #   # Invoking with custom +timing_configuration+
+        #   t = [
+        #     [1.second, 1.second],
+        #     [2.seconds, 2.seconds]
+        #   ]
+        #
+        #   with_lock_retries(timing_configuration: t) do
+        #     drop_table :my_table # this will be retried twice
+        #   end
+        #
+        #   # Disabling the retries using an environment variable
+        #   > export DISABLE_LOCK_RETRIES=true
+        #
+        #   with_lock_retries do
+        #     drop_table :my_table # one invocation, it will not retry at all
+        #   end
+        #
+        # ==== Parameters
+        # * +timing_configuration+ - [[ActiveSupport::Duration, ActiveSupport::Duration], ...] lock timeout for the block, sleep time before the next iteration, defaults to `Gitlab::Database::WithLockRetries::DEFAULT_TIMING_CONFIGURATION`
+        # * +logger+ - [Gitlab::JsonLogger]
+        # * +env+ - [Hash] custom environment hash, see the example with `DISABLE_LOCK_RETRIES`
+        def with_lock_retries(*args, **kwargs, &block)
+          super(*args, **kwargs.merge(allow_subtrans: false), &block)
+        end
+
         # Renames a column without requiring downtime.
         #
         # Concurrent renames work by using database triggers to ensure both the

spec/lib/gitlab/database/migration_helpers/v2_spec.rb

@@ -11,6 +11,8 @@ RSpec.describe Gitlab::Database::MigrationHelpers::V2 do
 
   before do
     allow(migration).to receive(:puts)
+
+    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
   end
 
   shared_examples_for 'Setting up to rename a column' do
@@ -218,4 +220,52 @@ RSpec.describe Gitlab::Database::MigrationHelpers::V2 do
       let(:added_column) { :original }
     end
   end
+
+  describe '#with_lock_retries' do
+    let(:model) do
+      ActiveRecord::Migration.new.extend(described_class)
+    end
+
+    let(:buffer) { StringIO.new }
+    let(:in_memory_logger) { Gitlab::JsonLogger.new(buffer) }
+    let(:env) { { 'DISABLE_LOCK_RETRIES' => 'true' } }
+
+    it 'sets the migration class name in the logs' do
+      model.with_lock_retries(env: env, logger: in_memory_logger) { }
+
+      buffer.rewind
+      expect(buffer.read).to include("\"class\":\"#{model.class}\"")
+    end
+
+    using RSpec::Parameterized::TableSyntax
+
+    where(raise_on_exhaustion: [true, false])
+
+    with_them do
+      it 'sets raise_on_exhaustion as requested' do
+        with_lock_retries = double
+        expect(Gitlab::Database::WithLockRetries).to receive(:new).and_return(with_lock_retries)
+        expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: raise_on_exhaustion)
+
+        model.with_lock_retries(env: env, logger: in_memory_logger, raise_on_exhaustion: raise_on_exhaustion) { }
+      end
+    end
+
+    it 'does not raise on exhaustion by default' do
+      with_lock_retries = double
+      expect(Gitlab::Database::WithLockRetries).to receive(:new).and_return(with_lock_retries)
+      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
+
+      model.with_lock_retries(env: env, logger: in_memory_logger) { }
+    end
+
+    it 'defaults to disallowing subtransactions' do
+      with_lock_retries = double
+
+      expect(Gitlab::Database::WithLockRetries).to receive(:new).with(hash_including(allow_subtrans: false)).and_return(with_lock_retries)
+      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
+
+      model.with_lock_retries(env: env, logger: in_memory_logger) { }
+    end
+  end
 end