use yugabyte adapter for getting the lock connection

.github/workflows/tests.yml

@@ -57,4 +57,6 @@ jobs:
           ruby-version: "${{ matrix.ruby-version }}"
       - name: Run specs
         run: |
-          bundle exec rspec 
+          bundle exec rspec
+      - name: Run Specs With YUAGBUTE_QUE_WORKER_ENABLED Enabled
+        run: YUAGBUTE_QUE_WORKER_ENABLED=true bundle exec rspec

lib/que/adapters/yugabyte.rb

@@ -12,34 +12,34 @@ def checkout_activerecord_adapter(&block)
             YugabyteRecord.connection_pool.with_connection(&block)
         end
 
-        # def establish_lock_database_connection
-        #     Thread.current["lock_database_connection_#{Thread.current.__id__}"] = LockDatabaseRecord.connection
-        # end
+        def checkout_lock_database_connection
+            # when multiple threads are running we need to make sure
+            # the acquiring and releasing of advisory locks is done by the 
+            # same connection
+            Thread.current[:db_connection] ||= LockDatabaseRecord.connection_pool.checkout
+        end
 
-        # def lock_database_connection
-        #     # connection =  @lock_database_connection[Thread.current.name]
-        #     # return connection unless connection.nil?
-        #     # @lock_database_connection[Thread.current.name] = LockDatabaseRecord.connection
-        #     @lock_database_connection ||= LockDatabaseRecord.connection
-        # end
+        def lock_database_connection
+            Thread.current[:db_connection]
+        end
 
-        def setup_lock_database_connection
-            ::LockDatabaseRecord.connection 
+        def release_lock_database_connection
+            LockDatabaseRecord.connection_pool.checkin(Thread.current[:db_connection])
         end
 
-        # def execute(command, params=[])
-        #     if command == :lock_job
-        #         queue, cursor, lock_database_connection = params
-        #         lock_job_with_lock_database(queue, cursor, lock_database_connection)
-        #     elsif command == :unlock_job
-        #         job_id, lock_database_connection = params
-        #         unlock_job(job_id, lock_database_connection)
-        #     else
-        #         super(command, params)
-        #     end
-        # end
+        def execute(command, params=[])
+            if command == :lock_job
+                queue, cursor = params
+                lock_job_with_lock_database(queue, cursor)
+            elsif command == :unlock_job
+                job_id = params[0]
+                unlock_job(job_id)
+            else
+                super(command, params)
+            end
+        end
 
-        def lock_job_with_lock_database(queue, cursor, lock_database_connection)
+        def lock_job_with_lock_database(queue, cursor)
             query = QueJob.select(:job_id, :queue, :priority, :run_at, :job_class, :retryable, :args, :error_count)
                         .select("extract(epoch from (now() - run_at)) as latency")
                         .where("queue = ? AND job_id >= ? AND run_at <= ?", queue, cursor, Time.now)
@@ -50,24 +50,24 @@ def lock_job_with_lock_database(queue, cursor, lock_database_connection)
             result = Que.execute(query)
             return result if result.empty?
 
-             if locked?(result.first['job_id'], lock_database_connection)
+             if locked?(result.first['job_id'])
                 return result
              end
 
             # continue the recursion to fetch the next available job
-            lock_job_with_lock_database(queue, result.first['job_id'], lock_database_connection)
+            lock_job_with_lock_database(queue, result.first['job_id'])
         end
 
         def cleanup!
             YugabyteRecord.connection_pool.release_connection
             LockDatabaseRecord.connection_pool.release_connection
         end
 
-        def locked?(job_id, lock_database_connection)
-          lock_database_connection.execute("SELECT pg_try_advisory_lock(#{job_id})").first["pg_try_advisory_lock"]
+        def locked?(job_id)
+            lock_database_connection.execute("SELECT pg_try_advisory_lock(#{job_id})").try(:first)&.fetch('pg_try_advisory_lock')
         end
 
-        def unlock_job(job_id, lock_database_connection)
+        def unlock_job(job_id)
            lock_database_connection.execute("SELECT pg_advisory_unlock(#{job_id})")
         end
       end

lib/que/locker.rb

@@ -60,12 +60,6 @@ def initialize(queue:, cursor_expiry:, window: nil, budget: nil, secondary_queue
       @queue_expires_at = {}
       @secondary_queues = secondary_queues
       @consolidated_queues = Array.wrap(queue).concat(secondary_queues)
-      @using_lock_database = ENV.fetch("YUGABYTE_QUE_WORKER_ENABLED", false)
-
-      if @using_lock_database
-        @lock_database_connection = LockDatabaseRecord.connection
-      end
-
       # Create a bucket that has 100% capacity, so even when we don't apply a limit we
       # have a valid bucket that we can use everywhere
       @leaky_bucket = LeakyBucket.new(window: window || 1.0, budget: budget || 1.0)
@@ -126,11 +120,7 @@ def with_locked_job
     ensure
       if job
         observe(UnlockTotal, UnlockSecondsTotal, worked_queue: job[:queue]) do
-          if @using_lock_database
-            @lock_database_connection.execute("SELECT pg_advisory_unlock(#{job["job_id"]})")
-          else
-            Que.execute("SELECT pg_advisory_unlock($1)", [job[:job_id]])
-          end
+          Que.execute(:unlock_job, [job[:job_id]])
         end
       end
     end
@@ -158,40 +148,13 @@ def exists?(job)
     end
 
     def lock_job_query(queue, cursor)
-      if @using_lock_database
-        lock_job_with_lock_database(queue, cursor)
-      else
-        Que.execute(:lock_job, [queue, cursor]).first
-      end
-    end
-
-    def lock_job_with_lock_database(queue, cursor)
-      query = QueJob.select(:job_id, :queue, :priority, :run_at, :job_class, :retryable, :args, :error_count)
-                .select("extract(epoch from (now() - run_at)) as latency")
-                .where("queue = ? AND job_id >= ? AND run_at <= ?", queue, cursor, Time.now)
-                .where(retryable: true)
-                .order(:priority, :run_at, :job_id)
-                .limit(1).to_sql
-
-      result = Que.execute(query).first
-      return result if result.nil?
-
-      return result if locked?(result['job_id'])
-
-      # continue the recursion to fetch the next available job
-      lock_job_with_lock_database(queue, result['job_id'])
-    end
-
-    def locked?(job_id)
-      @lock_database_connection.execute("SELECT pg_try_advisory_lock(#{job_id})").first["pg_try_advisory_lock"]
+      Que.execute(:lock_job, [queue, cursor]).first
     end
 
     def handle_expired_cursors!
       @consolidated_queues.each do |queue|
         queue_cursor_expires_at = @queue_expires_at.fetch(queue, monotonic_now)       
-
-
-        reset_cursor_for!(queue) if queue_cursor_expires_at <= monotonic_now
+        reset_cursor_for!(queue) if queue_cursor_expires_at < monotonic_now
       end
     end
 

lib/que/sql.rb

@@ -162,6 +162,10 @@ module Que
         WHERE locktype = 'advisory'
       ) pg USING (job_id)
     },
+
+    unlock_job: %{
+      SELECT pg_advisory_unlock($1)
+    }
   }
   # rubocop:enable Style/MutableConstant
 end