Safely acquire lock on multiple tables

Author: rkrage

lib/pg_ha_migrations/safe_statements.rb

@@ -3,6 +3,14 @@ def safe_added_columns_without_default_value
     @safe_added_columns_without_default_value ||= []
   end
 
+  # This variable is used to track nested lock acquisition.
+  # Each element is an array of table objects with their lock modes.
+  # The order of the array represents the current call stack,
+  # where the most recent method call is the last element.
+  def safely_acquire_lock_for_table_history
+    @safely_acquire_lock_for_table_history ||= []
+  end
+
   def safe_create_table(table, options={}, &block)
     if options[:force]
       raise PgHaMigrations::UnsafeMigrationError.new(":force is NOT SAFE! Explicitly call unsafe_drop_table first if you want to recreate an existing table")
@@ -528,40 +536,58 @@ def exec_migration(conn, direction)
     super(conn, direction)
   end
 
-  def safely_acquire_lock_for_table(table, mode: :access_exclusive, &block)
-    nested_target_table = Thread.current[__method__]
+  def safely_acquire_lock_for_table(*tables, mode: :access_exclusive, &block)
+    # So this variable is always available in the ensure block
+    successfully_acquired_lock = false
 
     _check_postgres_adapter!
 
-    target_table = PgHaMigrations::Table.from_table_name(table, mode)
+    target_tables = tables.map do |target_table|
+      PgHaMigrations::Table.from_table_name(target_table, mode)
+    end.sort_by(&:fully_qualified_name)
+
+    # Lock mode is the same across supplied tables so just grab the first.
+    # This variable is useful for error handling and messaging.
+    target_mode = target_tables.first.mode
 
-    if nested_target_table
-      if nested_target_table != target_table
-        raise PgHaMigrations::InvalidMigrationError, "Nested lock detected! Cannot acquire lock on #{target_table.fully_qualified_name} while #{nested_target_table.fully_qualified_name} is locked."
-      elsif nested_target_table.mode < target_table.mode
-        raise PgHaMigrations::InvalidMigrationError, "Lock escalation detected! Cannot change lock level from :#{nested_target_table.mode} to :#{target_table.mode} for #{target_table.fully_qualified_name}."
+    # Grab the latest locked tables from the call stack.
+    # This will be nil if we are not in a nested context.
+    nested_target_tables = safely_acquire_lock_for_table_history.last
+
+    if nested_target_tables
+      if target_tables != nested_target_tables
+        raise PgHaMigrations::InvalidMigrationError,
+          "Nested lock detected! Cannot acquire lock on #{target_tables.map(&:fully_qualified_name).join(", ")} " \
+          "while #{nested_target_tables.map(&:fully_qualified_name).join(", ")} is locked."
       end
-    else
-      Thread.current[__method__] = target_table
-    end
 
-    # Locking a partitioned table will also lock child tables (including sub-partitions),
-    # so we need to check for blocking queries on those tables as well
-    target_tables = target_table.partitions(include_sub_partitions: true, include_self: true)
+      nested_target_mode = nested_target_tables.first.mode
 
-    successfully_acquired_lock = false
+      if nested_target_mode < target_mode
+        raise PgHaMigrations::InvalidMigrationError,
+          "Lock escalation detected! Cannot change lock level from :#{nested_target_mode} " \
+          "to :#{target_mode} for #{target_tables.map(&:fully_qualified_name).join(", ")}."
+      end
+    end
 
     until successfully_acquired_lock
-      while (
+      loop do
         blocking_transactions = PgHaMigrations::BlockingDatabaseTransactions.find_blocking_transactions("#{PgHaMigrations::LOCK_TIMEOUT_SECONDS} seconds")
-        blocking_transactions.any? do |query|
+
+        # Locking a partitioned table will also lock child tables (including sub-partitions),
+        # so we need to check for blocking queries on those tables as well
+        target_tables_for_blocking_transactions = target_tables.flat_map do |target_table|
+          target_table.partitions(include_sub_partitions: true, include_self: true)
+        end
+
+        break unless blocking_transactions.any? do |query|
           query.tables_with_locks.any? do |locked_table|
-            target_tables.any? do |target_table|
+            target_tables_for_blocking_transactions.any? do |target_table|
               target_table.conflicts_with?(locked_table)
             end
           end
         end
-      )
+
         say "Waiting on blocking transactions:"
         blocking_transactions.each do |blocking_transaction|
           say blocking_transaction.description
@@ -570,16 +596,15 @@ def safely_acquire_lock_for_table(table, mode: :access_exclusive, &block)
       end
 
       connection.transaction do
-        adjust_timeout_method = connection.postgresql_version >= 9_03_00 ? :adjust_lock_timeout : :adjust_statement_timeout
         begin
-          method(adjust_timeout_method).call(PgHaMigrations::LOCK_TIMEOUT_SECONDS) do
-            connection.execute("LOCK #{target_table.fully_qualified_name} IN #{target_table.mode.to_sql} MODE;")
+          adjust_statement_timeout(PgHaMigrations::LOCK_TIMEOUT_SECONDS) do
+            connection.execute("LOCK #{target_tables.map(&:fully_qualified_name).join(", ")} IN #{target_mode.to_sql} MODE;")
           end
           successfully_acquired_lock = true
         rescue ActiveRecord::StatementInvalid => e
-          if e.message =~ /PG::LockNotAvailable.+ lock timeout/ || e.message =~ /PG::QueryCanceled.+ statement timeout/
+          if e.message =~ /PG::QueryCanceled.+ statement timeout/
             sleep_seconds = PgHaMigrations::LOCK_FAILURE_RETRY_DELAY_MULTLIPLIER * PgHaMigrations::LOCK_TIMEOUT_SECONDS
-            say "Timed out trying to acquire #{target_table.mode.to_sql} lock on the #{target_table.fully_qualified_name} table."
+            say "Timed out trying to acquire #{target_mode.to_sql} lock on #{target_tables.map(&:fully_qualified_name).join(", ")}."
             say "Sleeping for #{sleep_seconds}s to allow potentially queued up queries to finish before continuing."
             sleep(sleep_seconds)
 
@@ -590,12 +615,13 @@ def safely_acquire_lock_for_table(table, mode: :access_exclusive, &block)
         end
 
         if successfully_acquired_lock
+          safely_acquire_lock_for_table_history.push(target_tables)
           block.call
         end
       end
     end
   ensure
-    Thread.current[__method__] = nil unless nested_target_table
+    safely_acquire_lock_for_table_history.pop if successfully_acquired_lock
   end
 
   def adjust_lock_timeout(timeout_seconds = PgHaMigrations::LOCK_TIMEOUT_SECONDS, &block)

spec/safe_statements_spec.rb

@@ -3793,14 +3793,15 @@ def up
             let(:alternate_connection_pool) do
               ActiveRecord::ConnectionAdapters::ConnectionPool.new(pool_config)
             end
-            let(:alternate_connection) do
-              alternate_connection_pool.connection
-            end
+
+            let(:alternate_connection) { alternate_connection_pool.connection }
+            let(:alternate_connection_2) { alternate_connection_pool.connection }
             let(:migration) { Class.new(migration_klass).new }
 
             before(:each) do
               ActiveRecord::Base.connection.execute(<<~SQL)
                 CREATE TABLE #{table_name}(pk SERIAL, i INTEGER);
+                CREATE TABLE #{table_name}_2(pk SERIAL, i INTEGER);
                 CREATE SCHEMA partman;
                 CREATE EXTENSION pg_partman SCHEMA partman;
               SQL
@@ -3829,6 +3830,28 @@ def up
               end
             end
 
+            it "acquires exclusive locks by default when multiple tables provided" do
+              migration.safely_acquire_lock_for_table(table_name, "bogus_table_2") do
+                expect(locks_for_table(table_name, connection: alternate_connection)).to contain_exactly(
+                  having_attributes(
+                    table: "bogus_table",
+                    lock_type: "AccessExclusiveLock",
+                    granted: true,
+                    pid: kind_of(Integer),
+                  )
+                )
+
+                expect(locks_for_table("bogus_table_2", connection: alternate_connection)).to contain_exactly(
+                  having_attributes(
+                    table: "bogus_table_2",
+                    lock_type: "AccessExclusiveLock",
+                    granted: true,
+                    pid: kind_of(Integer),
+                  )
+                )
+              end
+            end
+
             it "acquires a lock in a different mode when provided" do
               migration.safely_acquire_lock_for_table(table_name, mode: :share) do
                 expect(locks_for_table(table_name, connection: alternate_connection)).to contain_exactly(
@@ -3842,6 +3865,28 @@ def up
               end
             end
 
+            it "acquires locks in a different mode when multiple tables and mode provided" do
+              migration.safely_acquire_lock_for_table(table_name, "bogus_table_2", mode: :share_row_exclusive) do
+                expect(locks_for_table(table_name, connection: alternate_connection)).to contain_exactly(
+                  having_attributes(
+                    table: "bogus_table",
+                    lock_type: "ShareRowExclusiveLock",
+                    granted: true,
+                    pid: kind_of(Integer),
+                  )
+                )
+
+                expect(locks_for_table("bogus_table_2", connection: alternate_connection)).to contain_exactly(
+                  having_attributes(
+                    table: "bogus_table_2",
+                    lock_type: "ShareRowExclusiveLock",
+                    granted: true,
+                    pid: kind_of(Integer),
+                  )
+                )
+              end
+            end
+
             it "raises error when invalid lock mode provided" do
               expect do
                 migration.safely_acquire_lock_for_table(table_name, mode: :garbage) {}
@@ -3883,6 +3928,39 @@ def up
               end
             end
 
+            it "times out the lock query after LOCK_TIMEOUT_SECONDS when multiple tables provided" do
+              stub_const("PgHaMigrations::LOCK_TIMEOUT_SECONDS", 1)
+              stub_const("PgHaMigrations::LOCK_FAILURE_RETRY_DELAY_MULTLIPLIER", 0)
+              allow(PgHaMigrations::BlockingDatabaseTransactions).to receive(:find_blocking_transactions).and_return([])
+              allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
+
+              expect(ActiveRecord::Base.connection).to receive(:execute)
+                .with("LOCK \"public\".\"bogus_table\", \"public\".\"bogus_table_2\" IN ACCESS EXCLUSIVE MODE;")
+                .at_least(2)
+                .times
+
+              begin
+                query_thread = Thread.new do
+                  alternate_connection.execute("BEGIN; LOCK bogus_table_2;")
+                  sleep 3
+                  alternate_connection.execute("ROLLBACK")
+                end
+
+                sleep 0.5
+
+                migration.suppress_messages do
+                  migration.safely_acquire_lock_for_table(table_name, "bogus_table_2") do
+                    aggregate_failures do
+                      expect(locks_for_table(table_name, connection: alternate_connection_2)).not_to be_empty
+                      expect(locks_for_table("bogus_table_2", connection: alternate_connection_2)).not_to be_empty
+                    end
+                  end
+                end
+              ensure
+                query_thread.join
+              end
+            end
+
             it "does not wait to acquire a lock if the table has an existing but non-conflicting lock" do
               stub_const("PgHaMigrations::LOCK_TIMEOUT_SECONDS", 1)
 
@@ -4303,6 +4381,53 @@ def up
               expect(locks_for_table(table_name, connection: alternate_connection)).to be_empty
             end
 
+            it "allows re-entrancy when multiple tables provided" do
+              migration.safely_acquire_lock_for_table(table_name, "bogus_table_2") do
+                # The ordering of the args is intentional here to ensure
+                # the array sorting and equality logic works as intended
+                migration.safely_acquire_lock_for_table("bogus_table_2", table_name) do
+                  expect(locks_for_table(table_name, connection: alternate_connection)).to contain_exactly(
+                    having_attributes(
+                      table: "bogus_table",
+                      lock_type: "AccessExclusiveLock",
+                      granted: true,
+                      pid: kind_of(Integer),
+                    ),
+                  )
+
+                  expect(locks_for_table("bogus_table_2", connection: alternate_connection)).to contain_exactly(
+                    having_attributes(
+                      table: "bogus_table_2",
+                      lock_type: "AccessExclusiveLock",
+                      granted: true,
+                      pid: kind_of(Integer),
+                    ),
+                  )
+                end
+
+                expect(locks_for_table(table_name, connection: alternate_connection)).to contain_exactly(
+                  having_attributes(
+                    table: "bogus_table",
+                    lock_type: "AccessExclusiveLock",
+                    granted: true,
+                    pid: kind_of(Integer),
+                  ),
+                )
+
+                expect(locks_for_table("bogus_table_2", connection: alternate_connection)).to contain_exactly(
+                  having_attributes(
+                    table: "bogus_table_2",
+                    lock_type: "AccessExclusiveLock",
+                    granted: true,
+                    pid: kind_of(Integer),
+                  ),
+                )
+              end
+
+              expect(locks_for_table(table_name, connection: alternate_connection)).to be_empty
+              expect(locks_for_table("bogus_table_2", connection: alternate_connection)).to be_empty
+            end
+
             it "allows re-entrancy when inner lock is a lower level" do
               migration.safely_acquire_lock_for_table(table_name) do
                 migration.safely_acquire_lock_for_table(table_name, mode: :exclusive) do
@@ -4405,26 +4530,6 @@ def up
 
               expect(locks_for_table(table_name, connection: alternate_connection)).to be_empty
             end
-
-            it "uses statement_timeout instead of lock_timeout when on Postgres 9.1" do
-              allow(ActiveRecord::Base.connection).to receive(:postgresql_version).and_wrap_original do |m, *args|
-                if caller.detect { |line| line =~ /lib\/pg_ha_migrations\/blocking_database_transactions\.rb/ }
-                  # The long-running transactions check needs to know the actual
-                  # Postgres version to use the proper columns, so we don't want
-                  # to mock any calls from it.
-                  m.call(*args)
-                else
-                  9_01_12
-                end
-              end
-
-              expect do
-                migration.safely_acquire_lock_for_table(table_name) do
-                  expect(locks_for_table(table_name, connection: alternate_connection)).not_to be_empty
-                end
-                expect(locks_for_table(table_name, connection: alternate_connection)).to be_empty
-              end.not_to make_database_queries(matching: /lock_timeout/i)
-            end
           end
         end