From 6c2a94fa51181df8f0e62f97576391a0e02c5ad9 Mon Sep 17 00:00:00 2001
From: Ankitha Damodara <adamodara@gocardless.com>
Date: Wed, 7 Aug 2024 16:46:09 +0100
Subject: [PATCH] add nested with connection

---
 lib/que/adapters/active_record_with_lock.rb | 49 +++++++++++----------
 1 file changed, 26 insertions(+), 23 deletions(-)

diff --git a/lib/que/adapters/active_record_with_lock.rb b/lib/que/adapters/active_record_with_lock.rb
index e2556ab..e8e5ab8 100644
--- a/lib/que/adapters/active_record_with_lock.rb
+++ b/lib/que/adapters/active_record_with_lock.rb
@@ -4,26 +4,32 @@
 module Que
   module Adapters
     class ActiveRecordWithLock < Que::Adapters::ActiveRecord
-      LOCK_PREFIX = ENV["QUE_LOCK_PREFIX"] || 1111 # this is a random number
-      def initialize(job_connection_pool:, lock_record:)
+      def initialize(job_connection_pool:, lock_connection_pool:)
         @job_connection_pool = job_connection_pool
-        @lock_record = lock_record
+        @lock_connection_pool = lock_connection_pool
         super
       end
 
+      def checkout
+        checkout_lock_database_connection do
+          checkout_activerecord_adapter { |conn| yield conn.raw_connection }
+        end
+      rescue *AR_UNAVAILABLE_CONNECTION_ERRORS => e
+        raise UnavailableConnection, e
+      rescue ::ActiveRecord::StatementInvalid => e
+        raise e unless AR_UNAVAILABLE_CONNECTION_ERRORS.include?(e.cause.class)
+
+        # ActiveRecord::StatementInvalid is one of the most generic exceptions AR can
+        # raise, so we catch it and only handle the specific nested exceptions.
+        raise UnavailableConnection, e.cause
+      end
+
       def checkout_activerecord_adapter(&block)
         @job_connection_pool.with_connection(&block)
       end
 
-      def lock_database_connection
-        if Thread.current[:db_connection]
-          return Thread.current[:db_connection] if Thread.current[:db_connection].active?
-        end
-        # We are storing this in thread variable here to make sure
-        # same connection is used to acquire and release the advisory locks.
-        # Advisory lock will not be released if any other connection from the
-        # pool tries to release the lock
-        Thread.current[:db_connection] = @lock_record.connection
+      def checkout_lock_database_connection(&block)
+        @lock_connection_pool.with_connection(&block)
       end
 
       def execute(command, params = [])
@@ -52,26 +58,23 @@ def lock_job_with_lock_database(queue, cursor)
         result
       end
 
-      def cleanup!
-        @job_connection_pool.release_connection
-        @lock_record.remove_connection
-      end
-
       def pg_try_advisory_lock?(job_id)
-        lock_variable = "#{LOCK_PREFIX}#{job_id}".to_i
-        lock_database_connection.execute(
-          "SELECT pg_try_advisory_lock(#{lock_variable})",
-        ).try(:first)&.fetch("pg_try_advisory_lock")
+        checkout_lock_database_connection do |conn|
+          conn.execute(
+            "SELECT pg_try_advisory_lock(#{job_id})",
+          ).try(:first)&.fetch("pg_try_advisory_lock")
+        end
       end
 
       def unlock_job(job_id)
-        lock_variable = "#{LOCK_PREFIX}#{job_id}".to_i
         # If for any reason the connection that is used to get this advisory lock
         # is corrupted, the lock on this job_id would already be released when the
         # connection holding the lock goes bad.
         # Now, if a new connection tries to release the non existing lock this would just no op
         # by returning false and return a warning "WARNING:  you don't own a lock of type ExclusiveLock"
-        lock_database_connection.execute("SELECT pg_advisory_unlock(#{lock_variable})")
+        checkout_lock_database_connection do |conn|
+          conn.execute("SELECT pg_advisory_unlock(#{job_id})")
+        end
       end
     end
   end