Merge pull request #156 from doctolib/Fix_lock_timeout_increase

fix the increase lock timeout on retry

Author: CharlesDelannoy

.github/workflows/continuous-integration-workflow.yml

@@ -32,7 +32,7 @@ jobs:
     strategy:
       matrix:
         postgres: [ 11.7, 12.14, 15.2 ]
-        ruby: [ 3.0, 3.1, 3.2 ]
+        ruby: [ 3.1, 3.2 ]
     services:
       postgres:
         image: postgres:${{ matrix.postgres }}

Gemfile.lock

@@ -99,4 +99,4 @@ DEPENDENCIES
   strong_migrations
 
 BUNDLED WITH
-   2.2.16
+   2.6.6

gemfiles/activerecord61.gemfile

@@ -6,6 +6,7 @@ gemspec path: '..'
 
 gem 'activerecord', '~> 6.1.0'
 gem 'bundler'
+gem 'concurrent-ruby', '1.3.4'
 gem 'debug'
 gem 'minitest', '>= 5'
 gem 'mocha'

gemfiles/activerecord70.gemfile

@@ -6,6 +6,7 @@ gemspec path: '..'
 
 gem 'activerecord', '~> 7.0.0'
 gem 'bundler'
+gem 'concurrent-ruby', '1.3.4'
 gem 'debug'
 gem 'minitest', '>= 5'
 gem 'mocha'

lib/safe-pg-migrations/configuration.rb

@@ -1,9 +1,12 @@
 # frozen_string_literal: true
 
 require 'active_support/core_ext/numeric/time'
+require 'safe-pg-migrations/helpers/pg_helper'
 
 module SafePgMigrations
   class Configuration
+    include Helpers::PgHelper
+
     attr_accessor(*%i[
                     backfill_batch_size
                     backfill_pause
@@ -77,12 +80,5 @@ def pg_lock_timeout
       # By reducing the lock timeout by a very small margin, we ensure that the lock timeout is raised in priority
       pg_duration safe_timeout * 0.99
     end
-
-    private
-
-    def pg_duration(duration)
-      value, unit = duration.integer? ? [duration, 's'] : [(duration * 1000).to_i, 'ms']
-      "#{value}#{unit}"
-    end
   end
 end

lib/safe-pg-migrations/helpers/pg_helper.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module SafePgMigrations
+  module Helpers
+    module PgHelper
+      def pg_duration(duration)
+        value, unit = duration.integer? ? [duration, 's'] : [(duration * 1000).to_i, 'ms']
+        "#{value}#{unit}"
+      end
+    end
+  end
+end

lib/safe-pg-migrations/plugins/statement_retrier.rb

@@ -3,6 +3,7 @@
 module SafePgMigrations
   module StatementRetrier
     include Helpers::StatementsHelper
+    include Helpers::PgHelper
 
     RETRIABLE_SCHEMA_STATEMENTS.each do |method|
       define_method method do |*args, **options, &block|
@@ -13,17 +14,14 @@ module StatementRetrier
     private
 
     def retry_if_lock_timeout
-      initial_lock_timeout = SafePgMigrations.config.lock_timeout
+      @lock_timeout = SafePgMigrations.config.lock_timeout
       number_of_retries = 0
       begin
         number_of_retries += 1
         yield
-      rescue ActiveRecord::LockWaitTimeout
+      rescue ActiveRecord::LockWaitTimeout => e
         # Retrying is useless if we're inside a transaction.
-        if transaction_open? || number_of_retries >= SafePgMigrations.config.max_tries
-          SafePgMigrations.config.lock_timeout = initial_lock_timeout
-          raise
-        end
+        raise e if transaction_open? || number_of_retries >= SafePgMigrations.config.max_tries
 
         retry_delay = SafePgMigrations.config.retry_delay
         Helpers::Logger.say "Retrying in #{retry_delay} seconds...", sub_item: true
@@ -39,20 +37,23 @@ def retry_if_lock_timeout
     end
 
     def increase_lock_timeout
-      Helpers::Logger.say "  Increasing the lock timeout... Currently set to #{SafePgMigrations.config.lock_timeout}",
+      Helpers::Logger.say "  Increasing the lock timeout... Currently set to #{pg_duration(@lock_timeout)}",
                           sub_item: true
-      SafePgMigrations.config.lock_timeout = (SafePgMigrations.config.lock_timeout + lock_timeout_step)
-      unless SafePgMigrations.config.lock_timeout < SafePgMigrations.config.max_lock_timeout_for_retry
-        SafePgMigrations.config.lock_timeout = SafePgMigrations.config.max_lock_timeout_for_retry
+      @lock_timeout += lock_timeout_step
+      unless @lock_timeout < SafePgMigrations.config.max_lock_timeout_for_retry
+        @lock_timeout = SafePgMigrations.config.max_lock_timeout_for_retry
       end
-      Helpers::Logger.say "  Lock timeout is now set to #{SafePgMigrations.config.lock_timeout}", sub_item: true
+      execute("SET lock_timeout TO '#{pg_duration(@lock_timeout)}'")
+      Helpers::Logger.say "  Lock timeout is now set to #{pg_duration(@lock_timeout)}", sub_item: true
     end
 
     def lock_timeout_step
+      return @lock_timeout_step if defined?(@lock_timeout_step)
+
       max_lock_timeout_for_retry = SafePgMigrations.config.max_lock_timeout_for_retry
       lock_timeout = SafePgMigrations.config.lock_timeout
       max_tries = SafePgMigrations.config.max_tries
-      @lock_timeout_step ||= (max_lock_timeout_for_retry - lock_timeout) / (max_tries - 1)
+      @lock_timeout_step = (max_lock_timeout_for_retry - lock_timeout) / (max_tries - 1)
     end
   end
 end

test/statement_retrier_test.rb

@@ -7,26 +7,28 @@ def test_lock_timeout_increase_on_retry
     SafePgMigrations.config.lock_timeout = 0.1.seconds
     SafePgMigrations.config.increase_lock_timeout_on_retry = true
 
-    calls = calls_for_lock_timeout_migration
+    2.times do
+      calls = calls_for_lock_timeout_migration
 
-    assert_equal [
-      '   -> Retrying in 60 seconds...',
-      '   ->   Increasing the lock timeout... Currently set to 0.1',
-      '   ->   Lock timeout is now set to 0.325',
-      '   -> Retrying now.',
-      '   -> Retrying in 60 seconds...',
-      '   ->   Increasing the lock timeout... Currently set to 0.325',
-      '   ->   Lock timeout is now set to 0.55',
-      '   -> Retrying now.',
-      '   -> Retrying in 60 seconds...',
-      '   ->   Increasing the lock timeout... Currently set to 0.55',
-      '   ->   Lock timeout is now set to 0.775',
-      '   -> Retrying now.',
-      '   -> Retrying in 60 seconds...',
-      '   ->   Increasing the lock timeout... Currently set to 0.775',
-      '   ->   Lock timeout is now set to 1',
-      '   -> Retrying now.',
-    ], calls[1..].map(&:first)
+      assert_equal [
+        '   -> Retrying in 60 seconds...',
+        '   ->   Increasing the lock timeout... Currently set to 100ms',
+        '   ->   Lock timeout is now set to 325ms',
+        '   -> Retrying now.',
+        '   -> Retrying in 60 seconds...',
+        '   ->   Increasing the lock timeout... Currently set to 325ms',
+        '   ->   Lock timeout is now set to 550ms',
+        '   -> Retrying now.',
+        '   -> Retrying in 60 seconds...',
+        '   ->   Increasing the lock timeout... Currently set to 550ms',
+        '   ->   Lock timeout is now set to 775ms',
+        '   -> Retrying now.',
+        '   -> Retrying in 60 seconds...',
+        '   ->   Increasing the lock timeout... Currently set to 775ms',
+        '   ->   Lock timeout is now set to 1s',
+        '   -> Retrying now.',
+      ], calls[1..].map(&:first)
+    end
   end
 
   def test_no_lock_timeout_increase_on_retry_if_disabled