Merge branch 'morefice/add-migration-lock-retries-helper' into 'master'

Add LockRetriesHelpers

See merge request !101333



Merged-by: Diogo Frazão <dfrazao@gitlab.com>
Approved-by: Omar Qunsul <oqunsul@gitlab.com>
Approved-by: Diogo Frazão <dfrazao@gitlab.com>
Co-authored-by: Maxime Orefice <morefice@gitlab.com>

lib/gitlab/database/migration_helpers.rb

@@ -6,6 +6,7 @@ module MigrationHelpers
       include Migrations::ReestablishedConnectionStack
       include Migrations::BackgroundMigrationHelpers
       include Migrations::BatchedBackgroundMigrationHelpers
+      include Migrations::LockRetriesHelpers
       include DynamicModelHelpers
       include RenameTableHelpers
       include AsyncIndexes::MigrationHelpers
@@ -405,52 +406,6 @@ def disable_statement_timeout
         end
       end
 
-      # 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+.
-      #
-      # Note this helper uses subtransactions when run inside an already open transaction.
-      #
-      # ==== 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)
-        raise_on_exhaustion = !!kwargs.delete(:raise_on_exhaustion)
-        merged_args = {
-          connection: connection,
-          klass: self.class,
-          logger: Gitlab::BackgroundMigration::Logger,
-          allow_savepoints: true
-        }.merge(kwargs)
-
-        Gitlab::Database::WithLockRetries.new(**merged_args)
-          .run(raise_on_exhaustion: raise_on_exhaustion, &block)
-      end
-
       def true_value
         Database.true_value
       end

lib/gitlab/database/migrations/lock_retries_helpers.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module LockRetriesHelpers
+        # 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+.
+        #
+        # Note this helper uses subtransactions when run inside an already open transaction.
+        #
+        # ==== 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)
+          raise_on_exhaustion = !!kwargs.delete(:raise_on_exhaustion)
+          merged_args = {
+            connection: connection,
+            klass: self.class,
+            logger: Gitlab::BackgroundMigration::Logger,
+            allow_savepoints: true
+          }.merge(kwargs)
+
+          Gitlab::Database::WithLockRetries.new(**merged_args)
+            .run(raise_on_exhaustion: raise_on_exhaustion, &block)
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -2526,48 +2526,6 @@ def self.name
     end
   end
 
-  describe '#with_lock_retries' do
-    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
-
-    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 allowing subtransactions' do
-      with_lock_retries = double
-
-      expect(Gitlab::Database::WithLockRetries).to receive(:new).with(hash_including(allow_savepoints: true)).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
-
   describe '#backfill_iids' do
     include MigrationsHelpers
 

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::LockRetriesHelpers do
+  let(:model) do
+    ActiveRecord::Migration.new.extend(described_class)
+  end
+
+  describe '#with_lock_retries' do
+    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
+
+    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 allowing subtransactions' do
+      with_lock_retries = double
+
+      expect(Gitlab::Database::WithLockRetries)
+        .to receive(:new).with(hash_including(allow_savepoints: true)).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