diff --git a/lib/que/adapters/active_record_with_lock.rb b/lib/que/adapters/active_record_with_lock.rb index e2556ab..711d6ad 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:) @job_connection_pool = job_connection_pool @lock_record = lock_record 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_record.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