sidekiq-unique-jobs 3.0.14 → 4.0.0

data/lib/sidekiq_unique_jobs/server/middleware.rb

@@ -0,0 +1,73 @@
+require 'digest'
+require 'forwardable'
+
+module SidekiqUniqueJobs
+  module Server
+    class Middleware
+      extend Forwardable
+      def_delegators :Sidekiq, :logger
+      def_instance_delegator :@worker, :class, :worker_class
+
+      include OptionsWithFallback
+
+      def call(worker, item, queue, redis_pool = nil, &blk)
+        @worker = worker
+        @redis_pool = redis_pool
+        @queue = queue
+        @item = item
+
+        send(unique_lock, &blk)
+      end
+
+      private
+
+      attr_reader :redis_pool, :worker, :item, :worker_class
+
+      def until_executing
+        unlock
+        yield
+      end
+
+      def until_executed(&block)
+        operative = true
+        after_yield_yield(&block)
+      rescue Sidekiq::Shutdown
+        operative = false
+        raise
+      ensure
+        unlock if operative
+      end
+
+      def after_yield_yield
+        yield
+      end
+
+      def while_executing
+        lock.synchronize do
+          yield
+        end
+      rescue SidekiqUniqueJobs::RunLockFailed
+        return reschedule if reschedule_on_lock_fail
+        raise
+      end
+
+      def until_timeout
+        yield if block_given?
+      end
+
+      protected
+
+      def unlock
+        after_unlock_hook if lock.unlock(:server)
+      end
+
+      def after_unlock_hook
+        worker.after_unlock if worker.respond_to?(:after_unlock)
+      end
+
+      def reschedule
+        Sidekiq::Client.new(redis_pool).raw_push([item])
+      end
+    end
+  end
+end

data/lib/sidekiq_unique_jobs/sidekiq_unique_ext.rb

@@ -1,30 +1,81 @@
 require 'sidekiq/api'
 
 module Sidekiq
-  class Job
+  module UnlockMethod
+    def unlock(item)
+      SidekiqUniqueJobs::Unlockable.unlock(item['unique_digest'.freeze], item['jid'.freeze])
+    end
+  end
+
+  class SortedEntry
     module UniqueExtension
       def self.included(base)
         base.class_eval do
+          include UnlockMethod
           alias_method :delete_orig, :delete
           alias_method :delete, :delete_ext
+          alias_method :remove_job_orig, :remove_job
+          alias_method :remove_job, :remove_job_ext
         end
       end
 
       def delete_ext
-        unlock(payload_hash(item))
-        delete_orig
+        unlock(item) if delete_orig
+      end
+
+      private
+
+      def remove_job_ext
+        remove_job_orig do |message|
+          unlock(Sidekiq.load_json(message))
+          yield message
+        end
+      end
+    end
+
+    include UniqueExtension if Gem::Version.new(Sidekiq::VERSION) >= Gem::Version.new('3.1')
+  end
+
+  class ScheduledSet
+    module UniqueExtension
+      def self.included(base)
+        base.class_eval do
+          include UnlockMethod
+          alias_method :delete_orig, :delete
+          alias_method :delete, :delete_ext
+        end
+      end
+
+      def delete_ext
+        unlock(item) if delete_orig
       end
 
-      protected
+      def remove_job_ext
+        remove_job_orig do |message|
+          unlock(Sidekiq.load_json(message))
+          yield message
+        end
+      end
+    end
+    include UniqueExtension if Gem::Version.new(Sidekiq::VERSION) >= Gem::Version.new('3.1')
+  end
 
-      def payload_hash(item)
-        SidekiqUniqueJobs::PayloadHelper.get_payload(item['class'], item['queue'], item['args'])
+  class Job
+    module UniqueExtension
+      def self.included(base)
+        base.class_eval do
+          include UnlockMethod
+          alias_method :delete_orig, :delete
+          alias_method :delete, :delete_ext
+        end
       end
 
-      def unlock(payload_hash)
-        Sidekiq.redis { |conn| conn.del(payload_hash) }
+      def delete_ext
+        unlock(item)
+        delete_orig
       end
     end
+
     include UniqueExtension
   end
 
@@ -32,6 +83,7 @@ module Sidekiq
     module UniqueExtension
       def self.included(base)
         base.class_eval do
+          include UnlockMethod
           alias_method :clear_orig, :clear
           alias_method :clear, :clear_ext
         end
@@ -50,17 +102,27 @@ module Sidekiq
     module UniqueExtension
       def self.included(base)
         base.class_eval do
+          include UnlockMethod
           if base.method_defined?(:clear)
             alias_method :clear_orig, :clear
             alias_method :clear, :clear_ext
           end
+
+          if base.method_defined?(:delete_by_value)
+            alias_method :delete_by_value_orig, :delete_by_value
+            alias_method :delete_by_value, :delete_by_value_ext
+          end
         end
       end
 
       def clear_ext
-        each(&:delete_ext)
+        each(&:delete)
         clear_orig
       end
+
+      def delete_by_value_ext(name, value)
+        unlock(JSON.parse(value)) if delete_by_value_orig(name, value)
+      end
     end
 
     include UniqueExtension

data/lib/sidekiq_unique_jobs/testing.rb

@@ -0,0 +1,34 @@
+require 'sidekiq_unique_jobs/testing/sidekiq_overrides'
+
+module SidekiqUniqueJobs
+  module Client
+    class Middleware
+      alias_method :call_real, :call
+      def call(worker_class, item, queue, redis_pool = nil)
+        worker_class = SidekiqUniqueJobs.worker_class_constantize(worker_class)
+
+        if Sidekiq::Testing.inline?
+          _server.call(worker_class.new, item, queue, redis_pool) do
+            call_real(worker_class, item, queue, redis_pool) do
+              yield
+            end
+          end
+        else
+          call_real(worker_class, item, queue, redis_pool) do
+            yield
+          end
+        end
+      end
+
+      def _server
+        SidekiqUniqueJobs::Server::Middleware.new
+      end
+    end
+  end
+
+  class Testing
+    def mocking!
+      require 'sidekiq_unique_jobs/testing/mocking'
+    end
+  end
+end

data/lib/sidekiq_unique_jobs/testing/sidekiq_overrides.rb

@@ -0,0 +1,63 @@
+require 'sidekiq/testing'
+
+module Sidekiq
+  module Worker
+    module ClassMethods
+      include SidekiqUniqueJobs::Unlockable
+
+      # Drain and run all jobs for this worker
+      def drain
+        while (job = jobs.shift)
+          worker = new
+          worker.jid = job['jid']
+          worker.bid = job['bid'] if worker.respond_to?(:bid=)
+          execute_job(worker, job['args'])
+          unlock(job['unique_digest'], job['jid']) if Sidekiq::Testing.fake?
+        end
+      end
+
+      # Pop out a single job and perform it
+      def perform_one
+        fail(EmptyQueueError, 'perform_one called with empty job queue') if jobs.empty?
+        job = jobs.shift
+        worker = new
+        worker.jid = job['jid']
+        worker.bid = job['bid'] if worker.respond_to?(:bid=)
+        execute_job(worker, job['args'])
+        unlock(job['unique_digest'], job['jid']) if Sidekiq::Testing.fake?
+      end
+
+      # Clear all jobs for this worker
+      def clear
+        jobs.each do |job|
+          unlock(job['unique_digest'], job['jid']) if Sidekiq::Testing.fake?
+        end
+        jobs.clear
+      end
+    end
+
+    module Overrides
+      def self.included(base)
+        base.extend Sidekiq::Worker::Overrides::ClassMethods
+        base.class_eval do
+          class << self
+            alias_method :clear_all_orig, :clear_all
+            alias_method :clear_all, :clear_all_ext
+          end
+        end
+      end
+
+      module ClassMethods
+        def clear_all_ext
+          Sidekiq.redis do |c|
+            unique_keys = c.keys("#{SidekiqUniqueJobs.config.unique_prefix}:*")
+            c.del(*unique_keys) unless unique_keys.empty?
+          end
+          clear_all_orig
+        end
+      end
+    end
+
+    include Overrides
+  end
+end

data/lib/sidekiq_unique_jobs/unique_args.rb

@@ -0,0 +1,132 @@
+require 'sidekiq_unique_jobs/normalizer'
+
+module SidekiqUniqueJobs
+  # This class exists to be testable and the entire api should be considered private
+  # rubocop:disable ClassLength
+  class UniqueArgs
+    extend Forwardable
+    include Normalizer
+
+    def_delegators :SidekiqUniqueJobs, :config, :worker_class_constantize
+    def_delegators :'Sidekiq.logger', :logger, :debug, :warn, :error, :fatal
+
+    def self.digest(item)
+      new(item).unique_digest
+    end
+
+    def initialize(job)
+      Sidekiq::Logging.with_context(self.class.name) do
+        @item = job
+        @worker_class ||= worker_class_constantize(@item['class'.freeze])
+        @item['unique_prefix'.freeze] ||= unique_prefix
+        @item['unique_args'.freeze] ||= unique_args(@item['args'.freeze])
+        @item['unique_digest'.freeze] ||= unique_digest
+      end
+    end
+
+    def unique_digest
+      @unique_digest ||= begin
+        digest = Digest::MD5.hexdigest(Sidekiq.dump_json(digestable_hash))
+        digest = "#{unique_prefix}:#{digest}"
+        debug { "#{__method__} : #{digestable_hash} into #{digest}" }
+        digest
+      end
+    end
+
+    def unique_prefix
+      return config.unique_prefix unless sidekiq_worker_class?
+      @worker_class.get_sidekiq_options['unique_prefix'.freeze] || config.unique_prefix
+    end
+
+    def digestable_hash
+      hash = @item.slice('class', 'queue', 'unique_args')
+
+      if unique_on_all_queues?
+        debug { "uniqueness specified across all queues (deleting queue: #{@item['queue']} from hash)" }
+        hash.delete('queue')
+      end
+      hash
+    end
+
+    def unique_args(args)
+      if unique_args_enabled?
+        filtered_args(args)
+      else
+        debug { "#{__method__} : unique arguments disabled" }
+        args
+      end
+    rescue NameError
+      # fallback to not filtering args when class can't be instantiated
+      return args
+    end
+
+    def unique_on_all_queues?
+      return unless sidekiq_worker_class?
+      return unless unique_args_enabled?
+      @worker_class.get_sidekiq_options['unique_on_all_queues'.freeze]
+    end
+
+    def unique_args_enabled?
+      unique_args_enabled_in_worker? ||
+        config.unique_args_enabled
+    end
+
+    def unique_args_enabled_in_worker?
+      return unless sidekiq_worker_class?
+      @worker_class.get_sidekiq_options['unique_args_enabled'.freeze] ||
+        @worker_class.get_sidekiq_options['unique_args'.freeze]
+    end
+
+    def sidekiq_worker_class?
+      if @worker_class.respond_to?(:get_sidekiq_options)
+        true
+      else
+        debug { "#{@worker_class} does not respond to :get_sidekiq_options" }
+        nil
+      end
+    end
+
+    # Filters unique arguments by proc or symbol
+    # returns provided arguments for other configurations
+    def filtered_args(args)
+      return args if args.empty?
+      json_args = Normalizer.jsonify(args)
+      debug { "#filtered_args #{args} => #{json_args}" }
+
+      case unique_args_method
+      when Proc
+        filter_by_proc(json_args)
+      when Symbol
+        filter_by_symbol(json_args)
+      else
+        debug { 'arguments not filtered (the combined arguments count towards uniqueness)' }
+        json_args
+      end
+    end
+
+    def filter_by_proc(args)
+      filter_args = unique_args_method.call(args)
+      debug { "#{__method__} : #{args} -> #{filter_args}" }
+      filter_args
+    end
+
+    def filter_by_symbol(args)
+      unless @worker_class.respond_to?(unique_args_method)
+        warn do
+          "#{__method__} : #{unique_args_method}) not defined in #{@worker_class} " \
+               "returning #{args} unchanged"
+        end
+        return args
+      end
+
+      filter_args = @worker_class.send(unique_args_method, args)
+      debug { "#{__method__} : #{unique_args_method}(#{args}) => #{filter_args}" }
+      filter_args
+    end
+
+    def unique_args_method
+      @unique_args_method ||=
+        @worker_class.get_sidekiq_options['unique_args'.freeze] if sidekiq_worker_class?
+    end
+  end
+end

data/lib/sidekiq_unique_jobs/unlockable.rb

@@ -0,0 +1,26 @@
+module SidekiqUniqueJobs
+  module Unlockable
+    module_function
+
+    # rubocop:disable MethodLength
+    def unlock(unique_key, jid, redis_pool = nil)
+      result = Scripts.call(:release_lock, redis_pool,
+                            keys: [unique_key],
+                            argv: [jid])
+      case result
+      when 1
+        Sidekiq.logger.debug { "successfully unlocked #{unique_key}" }
+        true
+      when 0
+        Sidekiq.logger.debug { "expiring lock #{unique_key} is not owned by #{jid}" }
+        false
+      when -1
+        Sidekiq.logger.debug { "#{unique_key} is not a known key" }
+        false
+      else
+        fail "#{__method__} returned an unexpected value (#{result})"
+      end
+    end
+    # rubocop:enable MethodLength
+  end
+end