From 6b0f2c87933bed61739208fc0c370747aa63defd Mon Sep 17 00:00:00 2001
From: MouadHilali <mouad.hilali@doctolib.com>
Date: Tue, 26 Mar 2024 15:28:32 +0100
Subject: [PATCH 01/12] increment gradually the lock_timeout value

---
 lib/safe-pg-migrations/configuration.rb             | 3 ++-
 lib/safe-pg-migrations/plugins/statement_retrier.rb | 9 +++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index 03c4876e..77d9d192 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -13,8 +13,9 @@ class Configuration
                     retry_delay
                     max_tries
                     sensitive_logger
+                    lock_timeout
+                    safe_timeout
                   ])
-    attr_reader :lock_timeout, :safe_timeout
 
     def initialize
       self.default_value_backfill_threshold = nil
diff --git a/lib/safe-pg-migrations/plugins/statement_retrier.rb b/lib/safe-pg-migrations/plugins/statement_retrier.rb
index ff872d9a..bbf96904 100644
--- a/lib/safe-pg-migrations/plugins/statement_retrier.rb
+++ b/lib/safe-pg-migrations/plugins/statement_retrier.rb
@@ -21,6 +21,15 @@ def retry_if_lock_timeout
         raise if transaction_open? # Retrying is useless if we're inside a transaction.
         raise unless remaining_tries > 0
 
+        number_of_retries = SafePgMigrations.config.max_tries - remaining_tries
+        if SafePgMigrations.config.max_tries - remaining_tries <= 20
+          SafePgMigrations.config.lock_timeout = SafePgMigrations.config.lock_timeout * number_of_retries
+          SafePgMigrations.config.safe_timeout = SafePgMigrations.config.safe_timeout * number_of_retries
+        else
+          SafePgMigrations.config.lock_timeout = SafePgMigrations.config.lock_timeout * 20
+          SafePgMigrations.config.safe_timeout = SafePgMigrations.config.safe_timeout * 20
+        end
+
         retry_delay = SafePgMigrations.config.retry_delay
         Helpers::Logger.say "Retrying in #{retry_delay} seconds...", sub_item: true
         sleep retry_delay

From 2f59f8065a496d358b2e8f33c7d7178b51ca784e Mon Sep 17 00:00:00 2001
From: MouadHilali <mouad.hilali@doctolib.com>
Date: Wed, 27 Mar 2024 14:51:44 +0100
Subject: [PATCH 02/12] introduce max_lock_timeout

---
 lib/safe-pg-migrations/configuration.rb             |  6 ++++--
 lib/safe-pg-migrations/plugins/statement_retrier.rb | 13 +++++--------
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index 77d9d192..d4c5da6d 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -12,10 +12,11 @@ class Configuration
                     backfill_pause
                     retry_delay
                     max_tries
+                    max_lock_timeout
                     sensitive_logger
                     lock_timeout
-                    safe_timeout
                   ])
+    attr_reader :safe_timeout
 
     def initialize
       self.default_value_backfill_threshold = nil
@@ -27,6 +28,7 @@ def initialize
       self.backfill_pause = 0.5.second
       self.retry_delay = 1.minute
       self.max_tries = 5
+      self.max_lock_timeout = 1.second
       self.sensitive_logger = nil
     end
 
@@ -43,7 +45,7 @@ def lock_timeout=(value)
     def safe_timeout=(value)
       raise 'Setting safe timeout to 0 disables the safe timeout and is dangerous' unless value
 
-      unless lock_timeout.nil? || value > lock_timeout
+      unless lock_timeout.nil? || value > lock_timeout || value > max_lock_timeout
         raise ArgumentError, "Safe timeout (#{value}) cannot be less than lock timeout (#{lock_timeout})"
       end
 
diff --git a/lib/safe-pg-migrations/plugins/statement_retrier.rb b/lib/safe-pg-migrations/plugins/statement_retrier.rb
index bbf96904..0037ee6f 100644
--- a/lib/safe-pg-migrations/plugins/statement_retrier.rb
+++ b/lib/safe-pg-migrations/plugins/statement_retrier.rb
@@ -13,21 +13,18 @@ module StatementRetrier
     private
 
     def retry_if_lock_timeout
-      remaining_tries = SafePgMigrations.config.max_tries
+      number_of_retries = 0
       begin
-        remaining_tries -= 1
+        number_of_retries += 1
         yield
       rescue ActiveRecord::LockWaitTimeout
         raise if transaction_open? # Retrying is useless if we're inside a transaction.
-        raise unless remaining_tries > 0
+        raise if number_of_retries >= SafePgMigrations.config.max_tries
 
-        number_of_retries = SafePgMigrations.config.max_tries - remaining_tries
-        if SafePgMigrations.config.max_tries - remaining_tries <= 20
+        if SafePgMigrations.config.lock_timeout < SafePgMigrations.config.max_lock_timeout
           SafePgMigrations.config.lock_timeout = SafePgMigrations.config.lock_timeout * number_of_retries
-          SafePgMigrations.config.safe_timeout = SafePgMigrations.config.safe_timeout * number_of_retries
         else
-          SafePgMigrations.config.lock_timeout = SafePgMigrations.config.lock_timeout * 20
-          SafePgMigrations.config.safe_timeout = SafePgMigrations.config.safe_timeout * 20
+          SafePgMigrations.config.lock_timeout = SafePgMigrations.config.max_lock_timeout
         end
 
         retry_delay = SafePgMigrations.config.retry_delay

From df907d5a34bd8426850c37924e3208cfe21a3ec8 Mon Sep 17 00:00:00 2001
From: MouadHilali <mouad.hilali@doctolib.com>
Date: Fri, 29 Mar 2024 12:53:07 +0100
Subject: [PATCH 03/12] ignore if lock timeout is nil

---
 lib/safe-pg-migrations/configuration.rb             |  8 ++++----
 lib/safe-pg-migrations/plugins/statement_retrier.rb | 10 ++++++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index d4c5da6d..847dc24b 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -5,16 +5,16 @@
 module SafePgMigrations
   class Configuration
     attr_accessor(*%i[
+                    backfill_batch_size
+                    backfill_pause
                     blocking_activity_logger_margin
                     blocking_activity_logger_verbose
                     default_value_backfill_threshold
-                    backfill_batch_size
-                    backfill_pause
-                    retry_delay
+                    lock_timeout
                     max_tries
                     max_lock_timeout
+                    retry_delay
                     sensitive_logger
-                    lock_timeout
                   ])
     attr_reader :safe_timeout
 
diff --git a/lib/safe-pg-migrations/plugins/statement_retrier.rb b/lib/safe-pg-migrations/plugins/statement_retrier.rb
index 0037ee6f..d96f58b7 100644
--- a/lib/safe-pg-migrations/plugins/statement_retrier.rb
+++ b/lib/safe-pg-migrations/plugins/statement_retrier.rb
@@ -21,10 +21,12 @@ def retry_if_lock_timeout
         raise if transaction_open? # Retrying is useless if we're inside a transaction.
         raise if number_of_retries >= SafePgMigrations.config.max_tries
 
-        if SafePgMigrations.config.lock_timeout < SafePgMigrations.config.max_lock_timeout
-          SafePgMigrations.config.lock_timeout = SafePgMigrations.config.lock_timeout * number_of_retries
-        else
-          SafePgMigrations.config.lock_timeout = SafePgMigrations.config.max_lock_timeout
+        unless SafePgMigrations.config.lock_timeout.nil?
+          if SafePgMigrations.config.lock_timeout < SafePgMigrations.config.max_lock_timeout
+            SafePgMigrations.config.lock_timeout = SafePgMigrations.config.lock_timeout * number_of_retries
+          else
+            SafePgMigrations.config.lock_timeout = SafePgMigrations.config.max_lock_timeout
+          end
         end
 
         retry_delay = SafePgMigrations.config.retry_delay

From bb6df4473b73332bb34a4adc4d6c6e108f3301a8 Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Wed, 17 Apr 2024 16:46:00 +0200
Subject: [PATCH 04/12] Setup the config setter for lock timeout options +
 update README.md

---
 README.md                               |  2 ++
 lib/safe-pg-migrations/configuration.rb | 24 +++++++++++++++---------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 9ce30272..5161aa52 100644
--- a/README.md
+++ b/README.md
@@ -329,6 +329,8 @@ SafePgMigrations.config.safe_timeout = 5.seconds # Statement timeout used for al
 
 SafePgMigrations.config.lock_timeout = nil # Lock timeout used for all DDL operations except from CREATE / DROP INDEX. If not set, safe_timeout will be used with a deduction of 1% to ensure that the lock timeout is raised in priority
 
+SafePgMigrations.config.max_lock_timeout_for_retry = 1.second # Max lock timeout for the retries for all DDL operations except from CREATE / DROP INDEX. Each retry will increase the lock_timeout by (max_lock_timeout_for_retry - lock_timeout) / max_tries
+
 SafePgMigrations.config.blocking_activity_logger_verbose = true # Outputs the raw blocking queries on timeout. When false, outputs information about the lock instead
 
 SafePgMigrations.config.sensitive_logger = nil # When given, sensitive data will be sent to this logger instead of the standard output. Must implement method `info`.
diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index 847dc24b..dcc7f7f5 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -10,13 +10,11 @@ class Configuration
                     blocking_activity_logger_margin
                     blocking_activity_logger_verbose
                     default_value_backfill_threshold
-                    lock_timeout
                     max_tries
-                    max_lock_timeout
                     retry_delay
                     sensitive_logger
                   ])
-    attr_reader :safe_timeout
+    attr_reader :lock_timeout, :safe_timeout, :max_lock_timeout_for_retry
 
     def initialize
       self.default_value_backfill_threshold = nil
@@ -28,30 +26,38 @@ def initialize
       self.backfill_pause = 0.5.second
       self.retry_delay = 1.minute
       self.max_tries = 5
-      self.max_lock_timeout = 1.second
+      self.max_lock_timeout_for_retry = 1.second
       self.sensitive_logger = nil
     end
 
     def lock_timeout=(value)
       raise 'Setting lock timeout to 0 disables the lock timeout and is dangerous' if value == 0.seconds
 
-      unless value.nil? || value < safe_timeout
-        raise ArgumentError, "Lock timeout (#{value}) cannot be greater than safe timeout (#{safe_timeout})"
+      unless value.nil? || (value < safe_timeout && value < max_lock_timeout_on_retry)
+        raise ArgumentError, "Lock timeout (#{value}) cannot be greater than the safe timeout (#{safe_timeout}) or the max lock timeout for retry (#{max_lock_timeout_for_retry})"
       end
 
       @lock_timeout = value
     end
 
     def safe_timeout=(value)
-      raise 'Setting safe timeout to 0 disables the safe timeout and is dangerous' unless value
+      raise 'Setting safe timeout to 0 or nil disables the safe timeout and is dangerous' unless value && value > 0.seconds
 
-      unless lock_timeout.nil? || value > lock_timeout || value > max_lock_timeout
-        raise ArgumentError, "Safe timeout (#{value}) cannot be less than lock timeout (#{lock_timeout})"
+      unless lock_timeout.nil? || (value > lock_timeout && value > max_lock_timeout_for_retry)
+        raise ArgumentError, "Safe timeout (#{value}) cannot be lower than the lock timeout (#{lock_timeout}) or the max lock timeout for retry (#{max_lock_timeout_for_retry})"
       end
 
       @safe_timeout = value
     end
 
+    def max_lock_timeout_for_retry(value)
+      unless lock_timeout.nil? || (value > lock_timeout && value < safe_timeout)
+        raise ArgumentError, "Max lock timeout for retry (#{value}) cannot be lower than the lock timeout (#{lock_timeout}) and greater than the safe timeout (#{safe_timeout})"
+      end
+
+      @max_lock_timeout_for_retry = value
+    end
+
     def pg_statement_timeout
       pg_duration safe_timeout
     end

From fe242fa72b1caf7472945423fc9333006b83760c Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Wed, 17 Apr 2024 17:01:00 +0200
Subject: [PATCH 05/12] update statement_retrier.rb to setup step to increment

---
 lib/safe-pg-migrations/configuration.rb       |  2 +-
 .../plugins/statement_retrier.rb              | 22 ++++++++++++-------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index dcc7f7f5..cd79aa5a 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -33,7 +33,7 @@ def initialize
     def lock_timeout=(value)
       raise 'Setting lock timeout to 0 disables the lock timeout and is dangerous' if value == 0.seconds
 
-      unless value.nil? || (value < safe_timeout && value < max_lock_timeout_on_retry)
+      unless value.nil? || (value < safe_timeout && value < max_lock_timeout_for_retry)
         raise ArgumentError, "Lock timeout (#{value}) cannot be greater than the safe timeout (#{safe_timeout}) or the max lock timeout for retry (#{max_lock_timeout_for_retry})"
       end
 
diff --git a/lib/safe-pg-migrations/plugins/statement_retrier.rb b/lib/safe-pg-migrations/plugins/statement_retrier.rb
index d96f58b7..5f5811f8 100644
--- a/lib/safe-pg-migrations/plugins/statement_retrier.rb
+++ b/lib/safe-pg-migrations/plugins/statement_retrier.rb
@@ -13,21 +13,16 @@ module StatementRetrier
     private
 
     def retry_if_lock_timeout
+      lock_timeout_step = get_lock_timeout_step(SafePgMigrations.config)
       number_of_retries = 0
       begin
         number_of_retries += 1
         yield
       rescue ActiveRecord::LockWaitTimeout
         raise if transaction_open? # Retrying is useless if we're inside a transaction.
-        raise if number_of_retries >= SafePgMigrations.config.max_tries
+        raise if number_of_retries >= max_tries
 
-        unless SafePgMigrations.config.lock_timeout.nil?
-          if SafePgMigrations.config.lock_timeout < SafePgMigrations.config.max_lock_timeout
-            SafePgMigrations.config.lock_timeout = SafePgMigrations.config.lock_timeout * number_of_retries
-          else
-            SafePgMigrations.config.lock_timeout = SafePgMigrations.config.max_lock_timeout
-          end
-        end
+        increase_lock_timeout(lock_timeout_step) unless SafePgMigrations.config.lock_timeout.nil?
 
         retry_delay = SafePgMigrations.config.retry_delay
         Helpers::Logger.say "Retrying in #{retry_delay} seconds...", sub_item: true
@@ -36,5 +31,16 @@ def retry_if_lock_timeout
         retry
       end
     end
+
+    def increase_lock_timeout(lock_timeout_step)
+      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
+      end
+    end
+
+    def get_lock_timeout_step(config)
+      (config.max_lock_timeout_for_retry - config.lock_timeout) / config.max_tries
+    end
   end
 end

From adab2ba0a65c542a953a9e85379212385a987fcf Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Wed, 17 Apr 2024 17:26:56 +0200
Subject: [PATCH 06/12] Fix broken tests

---
 lib/safe-pg-migrations/configuration.rb         |  8 ++++----
 .../plugins/statement_retrier.rb                | 17 ++++++++++-------
 test/statement_retrier_test.rb                  |  4 ++++
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index cd79aa5a..4e2e0a2d 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -33,7 +33,7 @@ def initialize
     def lock_timeout=(value)
       raise 'Setting lock timeout to 0 disables the lock timeout and is dangerous' if value == 0.seconds
 
-      unless value.nil? || (value < safe_timeout && value < max_lock_timeout_for_retry)
+      unless value.nil? || (value < safe_timeout && value <= max_lock_timeout_for_retry)
         raise ArgumentError, "Lock timeout (#{value}) cannot be greater than the safe timeout (#{safe_timeout}) or the max lock timeout for retry (#{max_lock_timeout_for_retry})"
       end
 
@@ -43,15 +43,15 @@ def lock_timeout=(value)
     def safe_timeout=(value)
       raise 'Setting safe timeout to 0 or nil disables the safe timeout and is dangerous' unless value && value > 0.seconds
 
-      unless lock_timeout.nil? || (value > lock_timeout && value > max_lock_timeout_for_retry)
+      unless lock_timeout.nil? || (value > lock_timeout && value >= max_lock_timeout_for_retry)
         raise ArgumentError, "Safe timeout (#{value}) cannot be lower than the lock timeout (#{lock_timeout}) or the max lock timeout for retry (#{max_lock_timeout_for_retry})"
       end
 
       @safe_timeout = value
     end
 
-    def max_lock_timeout_for_retry(value)
-      unless lock_timeout.nil? || (value > lock_timeout && value < safe_timeout)
+    def max_lock_timeout_for_retry=(value)
+      unless lock_timeout.nil? || (value >= lock_timeout && value <= safe_timeout)
         raise ArgumentError, "Max lock timeout for retry (#{value}) cannot be lower than the lock timeout (#{lock_timeout}) and greater than the safe timeout (#{safe_timeout})"
       end
 
diff --git a/lib/safe-pg-migrations/plugins/statement_retrier.rb b/lib/safe-pg-migrations/plugins/statement_retrier.rb
index 5f5811f8..e3dc93b8 100644
--- a/lib/safe-pg-migrations/plugins/statement_retrier.rb
+++ b/lib/safe-pg-migrations/plugins/statement_retrier.rb
@@ -13,16 +13,19 @@ module StatementRetrier
     private
 
     def retry_if_lock_timeout
-      lock_timeout_step = get_lock_timeout_step(SafePgMigrations.config)
+      initial_lock_timeout = SafePgMigrations.config.lock_timeout
       number_of_retries = 0
       begin
         number_of_retries += 1
         yield
       rescue ActiveRecord::LockWaitTimeout
-        raise if transaction_open? # Retrying is useless if we're inside a transaction.
-        raise if number_of_retries >= max_tries
+        # 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
 
-        increase_lock_timeout(lock_timeout_step) unless SafePgMigrations.config.lock_timeout.nil?
+        increase_lock_timeout unless SafePgMigrations.config.lock_timeout.nil?
 
         retry_delay = SafePgMigrations.config.retry_delay
         Helpers::Logger.say "Retrying in #{retry_delay} seconds...", sub_item: true
@@ -32,15 +35,15 @@ def retry_if_lock_timeout
       end
     end
 
-    def increase_lock_timeout(lock_timeout_step)
+    def increase_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
       end
     end
 
-    def get_lock_timeout_step(config)
-      (config.max_lock_timeout_for_retry - config.lock_timeout) / config.max_tries
+    def lock_timeout_step
+      (SafePgMigrations.config.max_lock_timeout_for_retry - SafePgMigrations.config.lock_timeout) / SafePgMigrations.config.max_tries
     end
   end
 end
diff --git a/test/statement_retrier_test.rb b/test/statement_retrier_test.rb
index 68a08e43..34581628 100644
--- a/test/statement_retrier_test.rb
+++ b/test/statement_retrier_test.rb
@@ -48,4 +48,8 @@ def test_statement_retry
       '   -> Retrying now.',
     ], calls[7..9]
   end
+
+  def lock_timout_increase_on_retry
+
+  end
 end

From a3e7ff4f83d483a5b2adbbe991ae7cbca2ba7333 Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Mon, 22 Apr 2024 10:39:41 +0200
Subject: [PATCH 07/12] Add tests for lock timeout increase

---
 .../plugins/statement_retrier.rb              | 11 +--
 test/statement_retrier_test.rb                | 67 ++++++++++++++-----
 2 files changed, 56 insertions(+), 22 deletions(-)

diff --git a/lib/safe-pg-migrations/plugins/statement_retrier.rb b/lib/safe-pg-migrations/plugins/statement_retrier.rb
index e3dc93b8..d1e97bad 100644
--- a/lib/safe-pg-migrations/plugins/statement_retrier.rb
+++ b/lib/safe-pg-migrations/plugins/statement_retrier.rb
@@ -25,10 +25,11 @@ def retry_if_lock_timeout
           raise
         end
 
-        increase_lock_timeout unless SafePgMigrations.config.lock_timeout.nil?
-
         retry_delay = SafePgMigrations.config.retry_delay
         Helpers::Logger.say "Retrying in #{retry_delay} seconds...", sub_item: true
+
+        increase_lock_timeout unless SafePgMigrations.config.lock_timeout.nil?
+
         sleep retry_delay
         Helpers::Logger.say 'Retrying now.', sub_item: true
         retry
@@ -36,14 +37,16 @@ def retry_if_lock_timeout
     end
 
     def increase_lock_timeout
-      SafePgMigrations.config.lock_timeout += lock_timeout_step
+      Helpers::Logger.say "  Increasing the lock timeout... Currently set to #{SafePgMigrations.config.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
       end
+      Helpers::Logger.say "  Lock timeout is now set to #{SafePgMigrations.config.lock_timeout}", sub_item: true
     end
 
     def lock_timeout_step
-      (SafePgMigrations.config.max_lock_timeout_for_retry - SafePgMigrations.config.lock_timeout) / SafePgMigrations.config.max_tries
+      @lock_timeout_step ||= (SafePgMigrations.config.max_lock_timeout_for_retry - SafePgMigrations.config.lock_timeout) / (SafePgMigrations.config.max_tries - 1)
     end
   end
 end
diff --git a/test/statement_retrier_test.rb b/test/statement_retrier_test.rb
index 34581628..c827fdcd 100644
--- a/test/statement_retrier_test.rb
+++ b/test/statement_retrier_test.rb
@@ -4,24 +4,8 @@
 
 class StatementRetrierTest < Minitest::Test
   def test_retry_if_lock_timeout
-    @migration =
-      Class.new(ActiveRecord::Migration::Current) do
-        def up
-          connection.send(:retry_if_lock_timeout) do
-            raise ActiveRecord::LockWaitTimeout, 'PG::LockNotAvailable: ERROR:  canceling statement due to lock timeout'
-          end
-        end
-      end.new
+    calls = calls_for_lock_timeout_migration
 
-    @connection.expects(:sleep).times(4)
-    calls =
-      record_calls(@migration, :write) do
-        run_migration
-        flunk 'run_migration should raise'
-      rescue StandardError => e
-        assert_instance_of ActiveRecord::LockWaitTimeout, e.cause
-        assert_includes e.cause.message, 'canceling statement due to lock timeout'
-      end
     assert_equal [
       '   -> Retrying in 60 seconds...',
       '   -> Retrying now.',
@@ -49,7 +33,54 @@ def test_statement_retry
     ], calls[7..9]
   end
 
-  def lock_timout_increase_on_retry
+  def test_lock_timout_increase_on_retry
+    SafePgMigrations.config.lock_timeout = 0.1.seconds
+
+    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..-1].map(&:first)
+
+  end
+
+  private
+
+  def calls_for_lock_timeout_migration
+    @migration = Class.new(ActiveRecord::Migration::Current) do
+      def up
+        connection.send(:retry_if_lock_timeout) do
+          raise ActiveRecord::LockWaitTimeout, 'PG::LockNotAvailable: ERROR:  canceling statement due to lock timeout'
+        end
+      end
+    end.new
+
+    @connection.expects(:sleep).times(4)
+
+    calls =
+      record_calls(@migration, :write) do
+        run_migration
+        flunk 'run_migration should raise'
+      rescue StandardError => e
+        assert_instance_of ActiveRecord::LockWaitTimeout, e.cause
+        assert_includes e.cause.message, 'canceling statement due to lock timeout'
+      end
 
+    calls
   end
 end

From aeab5ae95521643e1cc9c5b2c6a613205b7ec2c1 Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Mon, 22 Apr 2024 10:52:46 +0200
Subject: [PATCH 08/12] add the feature as an option + update README

---
 README.md                                           | 4 +++-
 lib/safe-pg-migrations/configuration.rb             | 2 ++
 lib/safe-pg-migrations/plugins/statement_retrier.rb | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 5161aa52..bf120efb 100644
--- a/README.md
+++ b/README.md
@@ -329,7 +329,9 @@ SafePgMigrations.config.safe_timeout = 5.seconds # Statement timeout used for al
 
 SafePgMigrations.config.lock_timeout = nil # Lock timeout used for all DDL operations except from CREATE / DROP INDEX. If not set, safe_timeout will be used with a deduction of 1% to ensure that the lock timeout is raised in priority
 
-SafePgMigrations.config.max_lock_timeout_for_retry = 1.second # Max lock timeout for the retries for all DDL operations except from CREATE / DROP INDEX. Each retry will increase the lock_timeout by (max_lock_timeout_for_retry - lock_timeout) / max_tries
+SafePgMigrations.config.increase_lock_timeout_on_retry # Activate the lock timeout increase feature on retry if set to true. See max_lock_timeout_for_retry for more information.
+
+SafePgMigrations.config.max_lock_timeout_for_retry = 1.second # Max lock timeout for the retries for all DDL operations except from CREATE / DROP INDEX. Each retry will increase the lock_timeout (if increase_lock_timeout_on_retry option is set to true) by (max_lock_timeout_for_retry - lock_timeout) / max_tries
 
 SafePgMigrations.config.blocking_activity_logger_verbose = true # Outputs the raw blocking queries on timeout. When false, outputs information about the lock instead
 
diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index 4e2e0a2d..c62d7af1 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -13,6 +13,7 @@ class Configuration
                     max_tries
                     retry_delay
                     sensitive_logger
+                    increase_lock_timeout_on_retry
                   ])
     attr_reader :lock_timeout, :safe_timeout, :max_lock_timeout_for_retry
 
@@ -28,6 +29,7 @@ def initialize
       self.max_tries = 5
       self.max_lock_timeout_for_retry = 1.second
       self.sensitive_logger = nil
+      self.increase_lock_timeout_on_retry = false
     end
 
     def lock_timeout=(value)
diff --git a/lib/safe-pg-migrations/plugins/statement_retrier.rb b/lib/safe-pg-migrations/plugins/statement_retrier.rb
index d1e97bad..feb90b52 100644
--- a/lib/safe-pg-migrations/plugins/statement_retrier.rb
+++ b/lib/safe-pg-migrations/plugins/statement_retrier.rb
@@ -28,7 +28,7 @@ def retry_if_lock_timeout
         retry_delay = SafePgMigrations.config.retry_delay
         Helpers::Logger.say "Retrying in #{retry_delay} seconds...", sub_item: true
 
-        increase_lock_timeout unless SafePgMigrations.config.lock_timeout.nil?
+        increase_lock_timeout if SafePgMigrations.config.increase_lock_timeout_on_retry && !SafePgMigrations.config.lock_timeout.nil?
 
         sleep retry_delay
         Helpers::Logger.say 'Retrying now.', sub_item: true

From 7a7f854703af442a30f6e666503bb522f4e36156 Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Mon, 22 Apr 2024 10:53:04 +0200
Subject: [PATCH 09/12] Add tests

---
 test/statement_retrier_test.rb | 71 +++++++++++++++++++++-------------
 1 file changed, 45 insertions(+), 26 deletions(-)

diff --git a/test/statement_retrier_test.rb b/test/statement_retrier_test.rb
index c827fdcd..7f31d7eb 100644
--- a/test/statement_retrier_test.rb
+++ b/test/statement_retrier_test.rb
@@ -3,6 +3,51 @@
 require 'test_helper'
 
 class StatementRetrierTest < Minitest::Test
+
+  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
+
+    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..-1].map(&:first)
+  end
+
+  def test_no_lock_timeout_increase_on_retry_if_disabled
+    SafePgMigrations.config.lock_timeout = 0.1.seconds
+    SafePgMigrations.config.increase_lock_timeout_on_retry = false
+
+    calls = calls_for_lock_timeout_migration
+
+    assert_equal [
+                   '   -> Retrying in 60 seconds...',
+                   '   -> Retrying now.',
+                   '   -> Retrying in 60 seconds...',
+                   '   -> Retrying now.',
+                   '   -> Retrying in 60 seconds...',
+                   '   -> Retrying now.',
+                   '   -> Retrying in 60 seconds...',
+                   '   -> Retrying now.',
+                 ], calls[1..-1].map(&:first)
+  end
+
   def test_retry_if_lock_timeout
     calls = calls_for_lock_timeout_migration
 
@@ -33,32 +78,6 @@ def test_statement_retry
     ], calls[7..9]
   end
 
-  def test_lock_timout_increase_on_retry
-    SafePgMigrations.config.lock_timeout = 0.1.seconds
-
-    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..-1].map(&:first)
-
-  end
-
   private
 
   def calls_for_lock_timeout_migration

From 94fa5951257331cdc1f78d99bf65fc0bdd8b4c13 Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Mon, 22 Apr 2024 10:57:22 +0200
Subject: [PATCH 10/12] fix rubocop rule

---
 test/statement_retrier_test.rb | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/test/statement_retrier_test.rb b/test/statement_retrier_test.rb
index 7f31d7eb..09a1dd30 100644
--- a/test/statement_retrier_test.rb
+++ b/test/statement_retrier_test.rb
@@ -91,15 +91,13 @@ def up
 
     @connection.expects(:sleep).times(4)
 
-    calls =
-      record_calls(@migration, :write) do
-        run_migration
-        flunk 'run_migration should raise'
-      rescue StandardError => e
-        assert_instance_of ActiveRecord::LockWaitTimeout, e.cause
-        assert_includes e.cause.message, 'canceling statement due to lock timeout'
-      end
 
-    calls
+    record_calls(@migration, :write) do
+      run_migration
+      flunk 'run_migration should raise'
+    rescue StandardError => e
+      assert_instance_of ActiveRecord::LockWaitTimeout, e.cause
+      assert_includes e.cause.message, 'canceling statement due to lock timeout'
+    end
   end
 end

From 57261faaf94cb6bb6a9091365db8495f8907a53b Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Mon, 22 Apr 2024 11:46:44 +0200
Subject: [PATCH 11/12] fix rubocop rules

---
 lib/safe-pg-migrations/configuration.rb       | 13 +++--
 .../plugins/statement_retrier.rb              | 12 +++--
 test/statement_retrier_test.rb                | 54 +++++++++----------
 3 files changed, 44 insertions(+), 35 deletions(-)

diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index c62d7af1..05feedd0 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -36,17 +36,21 @@ def lock_timeout=(value)
       raise 'Setting lock timeout to 0 disables the lock timeout and is dangerous' if value == 0.seconds
 
       unless value.nil? || (value < safe_timeout && value <= max_lock_timeout_for_retry)
-        raise ArgumentError, "Lock timeout (#{value}) cannot be greater than the safe timeout (#{safe_timeout}) or the max lock timeout for retry (#{max_lock_timeout_for_retry})"
+        raise ArgumentError, "Lock timeout (#{value}) cannot be greater than the safe timeout (#{safe_timeout}) or the
+                              max lock timeout for retry (#{max_lock_timeout_for_retry})"
       end
 
       @lock_timeout = value
     end
 
     def safe_timeout=(value)
-      raise 'Setting safe timeout to 0 or nil disables the safe timeout and is dangerous' unless value && value > 0.seconds
+      unless value && value > 0.seconds
+        raise 'Setting safe timeout to 0 or nil disables the safe timeout and is dangerous'
+      end
 
       unless lock_timeout.nil? || (value > lock_timeout && value >= max_lock_timeout_for_retry)
-        raise ArgumentError, "Safe timeout (#{value}) cannot be lower than the lock timeout (#{lock_timeout}) or the max lock timeout for retry (#{max_lock_timeout_for_retry})"
+        raise ArgumentError, "Safe timeout (#{value}) cannot be lower than the lock timeout (#{lock_timeout}) or the
+                              max lock timeout for retry (#{max_lock_timeout_for_retry})"
       end
 
       @safe_timeout = value
@@ -54,7 +58,8 @@ def safe_timeout=(value)
 
     def max_lock_timeout_for_retry=(value)
       unless lock_timeout.nil? || (value >= lock_timeout && value <= safe_timeout)
-        raise ArgumentError, "Max lock timeout for retry (#{value}) cannot be lower than the lock timeout (#{lock_timeout}) and greater than the safe timeout (#{safe_timeout})"
+        raise ArgumentError, "Max lock timeout for retry (#{value}) cannot be lower than the lock timeout
+                              (#{lock_timeout}) and greater than the safe timeout (#{safe_timeout})"
       end
 
       @max_lock_timeout_for_retry = value
diff --git a/lib/safe-pg-migrations/plugins/statement_retrier.rb b/lib/safe-pg-migrations/plugins/statement_retrier.rb
index feb90b52..6c91f4fe 100644
--- a/lib/safe-pg-migrations/plugins/statement_retrier.rb
+++ b/lib/safe-pg-migrations/plugins/statement_retrier.rb
@@ -28,7 +28,9 @@ def retry_if_lock_timeout
         retry_delay = SafePgMigrations.config.retry_delay
         Helpers::Logger.say "Retrying in #{retry_delay} seconds...", sub_item: true
 
-        increase_lock_timeout if SafePgMigrations.config.increase_lock_timeout_on_retry && !SafePgMigrations.config.lock_timeout.nil?
+        if SafePgMigrations.config.increase_lock_timeout_on_retry && !SafePgMigrations.config.lock_timeout.nil?
+          increase_lock_timeout
+        end
 
         sleep retry_delay
         Helpers::Logger.say 'Retrying now.', sub_item: true
@@ -37,7 +39,8 @@ def retry_if_lock_timeout
     end
 
     def increase_lock_timeout
-      Helpers::Logger.say "  Increasing the lock timeout... Currently set to #{SafePgMigrations.config.lock_timeout}", sub_item: true
+      Helpers::Logger.say "  Increasing the lock timeout... Currently set to #{SafePgMigrations.config.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
@@ -46,7 +49,10 @@ def increase_lock_timeout
     end
 
     def lock_timeout_step
-      @lock_timeout_step ||= (SafePgMigrations.config.max_lock_timeout_for_retry - SafePgMigrations.config.lock_timeout) / (SafePgMigrations.config.max_tries - 1)
+      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)
     end
   end
 end
diff --git a/test/statement_retrier_test.rb b/test/statement_retrier_test.rb
index 09a1dd30..094f89d8 100644
--- a/test/statement_retrier_test.rb
+++ b/test/statement_retrier_test.rb
@@ -3,7 +3,6 @@
 require 'test_helper'
 
 class StatementRetrierTest < Minitest::Test
-
   def test_lock_timeout_increase_on_retry
     SafePgMigrations.config.lock_timeout = 0.1.seconds
     SafePgMigrations.config.increase_lock_timeout_on_retry = true
@@ -11,23 +10,23 @@ def test_lock_timeout_increase_on_retry
     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..-1].map(&:first)
+      '   -> 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)
   end
 
   def test_no_lock_timeout_increase_on_retry_if_disabled
@@ -37,15 +36,15 @@ def test_no_lock_timeout_increase_on_retry_if_disabled
     calls = calls_for_lock_timeout_migration
 
     assert_equal [
-                   '   -> Retrying in 60 seconds...',
-                   '   -> Retrying now.',
-                   '   -> Retrying in 60 seconds...',
-                   '   -> Retrying now.',
-                   '   -> Retrying in 60 seconds...',
-                   '   -> Retrying now.',
-                   '   -> Retrying in 60 seconds...',
-                   '   -> Retrying now.',
-                 ], calls[1..-1].map(&:first)
+      '   -> Retrying in 60 seconds...',
+      '   -> Retrying now.',
+      '   -> Retrying in 60 seconds...',
+      '   -> Retrying now.',
+      '   -> Retrying in 60 seconds...',
+      '   -> Retrying now.',
+      '   -> Retrying in 60 seconds...',
+      '   -> Retrying now.',
+    ], calls[1..].map(&:first)
   end
 
   def test_retry_if_lock_timeout
@@ -91,7 +90,6 @@ def up
 
     @connection.expects(:sleep).times(4)
 
-
     record_calls(@migration, :write) do
       run_migration
       flunk 'run_migration should raise'

From 43f730af8487e5d3e31397fe9acf7699c423e40e Mon Sep 17 00:00:00 2001
From: CharlesDelannoy <charles.delannoy@doctolib.com>
Date: Tue, 23 Apr 2024 11:55:07 +0200
Subject: [PATCH 12/12] Sorting

---
 lib/safe-pg-migrations/configuration.rb | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/lib/safe-pg-migrations/configuration.rb b/lib/safe-pg-migrations/configuration.rb
index 05feedd0..2fbe2fae 100644
--- a/lib/safe-pg-migrations/configuration.rb
+++ b/lib/safe-pg-migrations/configuration.rb
@@ -10,26 +10,26 @@ class Configuration
                     blocking_activity_logger_margin
                     blocking_activity_logger_verbose
                     default_value_backfill_threshold
+                    increase_lock_timeout_on_retry
                     max_tries
                     retry_delay
                     sensitive_logger
-                    increase_lock_timeout_on_retry
                   ])
     attr_reader :lock_timeout, :safe_timeout, :max_lock_timeout_for_retry
 
     def initialize
-      self.default_value_backfill_threshold = nil
-      self.safe_timeout = 5.seconds
-      self.lock_timeout = nil
-      self.blocking_activity_logger_margin = 1.second
-      self.blocking_activity_logger_verbose = true
       self.backfill_batch_size = 100_000
       self.backfill_pause = 0.5.second
-      self.retry_delay = 1.minute
-      self.max_tries = 5
+      self.blocking_activity_logger_margin = 1.second
+      self.blocking_activity_logger_verbose = true
+      self.default_value_backfill_threshold = nil
+      self.increase_lock_timeout_on_retry = false
+      self.lock_timeout = nil
       self.max_lock_timeout_for_retry = 1.second
+      self.max_tries = 5
+      self.retry_delay = 1.minute
+      self.safe_timeout = 5.seconds
       self.sensitive_logger = nil
-      self.increase_lock_timeout_on_retry = false
     end
 
     def lock_timeout=(value)