Fix potential deadlock in nested locking

Author: rkrage

lib/pg_ha_migrations/safe_statements.rb

@@ -572,6 +572,12 @@ def safely_acquire_lock_for_table(*tables, mode: :access_exclusive, &block)
 
     until successfully_acquired_lock
       loop do
+        # If in a nested context and all of the above checks have passed,
+        # we have already acquired the lock so this check is unnecessary.
+        # In fact, it could actually cause a deadlock if a blocking query
+        # was executed shortly after the initial lock acquisition.
+        break if nested_target_tables
+
         blocking_transactions = PgHaMigrations::BlockingDatabaseTransactions.find_blocking_transactions("#{PgHaMigrations::LOCK_TIMEOUT_SECONDS} seconds")
 
         # Locking a partitioned table will also lock child tables (including sub-partitions),

spec/safe_statements_spec.rb

@@ -4478,6 +4478,45 @@ def up
               expect(locks_for_table(table_name, connection: alternate_connection)).to be_empty
             end
 
+            it "skips blocking query check for nested lock acquisition" do
+              stub_const("PgHaMigrations::LOCK_TIMEOUT_SECONDS", 1)
+
+              query_thread = nil
+
+              expect(PgHaMigrations::BlockingDatabaseTransactions).to receive(:find_blocking_transactions)
+                .once
+                .and_call_original
+
+              begin
+                migration.safely_acquire_lock_for_table(table_name) do
+                  query_thread = Thread.new { alternate_connection.execute("SELECT * FROM bogus_table") }
+
+                  sleep 2
+
+                  migration.safely_acquire_lock_for_table(table_name) do
+                    expect(locks_for_table(table_name, connection: alternate_connection_2)).to contain_exactly(
+                      having_attributes(
+                        table: "bogus_table",
+                        lock_type: "AccessExclusiveLock",
+                        granted: true,
+                        pid: kind_of(Integer),
+                      ),
+                      having_attributes(
+                        table: "bogus_table",
+                        lock_type: "AccessShareLock",
+                        granted: false,
+                        pid: kind_of(Integer),
+                      ),
+                    )
+                  end
+                end
+              ensure
+                query_thread.join if query_thread
+              end
+
+              expect(locks_for_table(table_name, connection: alternate_connection)).to be_empty
+            end
+
             it "does not allow re-entrancy when lock escalation detected" do
               expect do
                 migration.safely_acquire_lock_for_table(table_name, mode: :share) do