kubra-sidekiq-throttled 1.5.3

data/lib/sidekiq/throttled/job.rb

@@ -0,0 +1,143 @@
+# frozen_string_literal: true
+
+# internal
+require_relative "./registry"
+
+module Sidekiq
+  module Throttled
+    # Adds helpers to your worker classes
+    #
+    # @example Usage
+    #
+    #     class MyJob
+    #       include Sidekiq::Job
+    #       include Sidekiq::Throttled::Job
+    #
+    #       sidkiq_options :queue => :my_queue
+    #       sidekiq_throttle :threshold => { :limit => 123, :period => 1.hour },
+    #                        :requeue => { :to => :other_queue, :with => :schedule }
+    #
+    #       def perform
+    #         # ...
+    #       end
+    #     end
+    #
+    # @see ClassMethods
+    module Job
+      # Extends worker class with {ClassMethods}.
+      #
+      # @note Using `included` hook with extending worker with {ClassMethods}
+      #   in order to make API inline with `include Sidekiq::Job`.
+      #
+      # @private
+      def self.included(base)
+        base.sidekiq_class_attribute :sidekiq_throttled_requeue_options
+        base.extend(ClassMethods)
+      end
+
+      # Helper methods added to the singleton class of destination
+      module ClassMethods
+        # Registers some strategy for the worker.
+        #
+        # @example Allow max 123 MyJob jobs per hour
+        #
+        #     class MyJob
+        #       include Sidekiq::Job
+        #       include Sidekiq::Throttled::Job
+        #
+        #       sidekiq_throttle({
+        #         :threshold => { :limit => 123, :period => 1.hour }
+        #       })
+        #     end
+        #
+        # @example Allow max 10 concurrently running MyJob jobs
+        #
+        #     class MyJob
+        #       include Sidekiq::Job
+        #       include Sidekiq::Throttled::Job
+        #
+        #       sidekiq_throttle({
+        #         :concurrency => { :limit => 10 }
+        #       })
+        #     end
+        #
+        # @example Allow max 10 concurrent MyJob jobs and max 123 per hour
+        #
+        #     class MyJob
+        #       include Sidekiq::Job
+        #       include Sidekiq::Throttled::Job
+        #
+        #       sidekiq_throttle({
+        #         :threshold => { :limit => 123, :period => 1.hour },
+        #         :concurrency => { :limit => 10 }
+        #       })
+        #     end
+        #
+        # @example Allow max 123 MyJob jobs per hour; when jobs are throttled, schedule them for later in :other_queue
+        #
+        #     class MyJob
+        #       include Sidekiq::Job
+        #       include Sidekiq::Throttled::Job
+        #
+        #       sidekiq_throttle({
+        #         :threshold => { :limit => 123, :period => 1.hour },
+        #         :requeue => { :to => :other_queue, :with => :schedule }
+        #       })
+        #     end
+        #
+        # @param [Hash] requeue What to do with jobs that are throttled
+        # @see Registry.add
+        # @return [void]
+        def sidekiq_throttle(**kwargs)
+          Registry.add(self, **kwargs)
+        end
+
+        # Adds current worker to preconfigured throttling strategy. Allows
+        # sharing same pool for multiple workers.
+        #
+        # First of all we need to create shared throttling strategy:
+        #
+        #     # Create google_api throttling strategy
+        #     Sidekiq::Throttled::Registry.add(:google_api, {
+        #       :threshold => { :limit => 123, :period => 1.hour },
+        #       :concurrency => { :limit => 10 }
+        #     })
+        #
+        # Now we can assign it to our workers:
+        #
+        #     class FetchProfileJob
+        #       include Sidekiq::Job
+        #       include Sidekiq::Throttled::Job
+        #
+        #       sidekiq_throttle_as :google_api
+        #     end
+        #
+        #     class FetchCommentsJob
+        #       include Sidekiq::Job
+        #       include Sidekiq::Throttled::Job
+        #
+        #       sidekiq_throttle_as :google_api
+        #     end
+        #
+        # With the above configuration we ensure that there are maximum 10
+        # concurrently running jobs of FetchProfileJob or FetchCommentsJob
+        # allowed. And only 123 jobs of those are executed per hour.
+        #
+        # In other words, it will allow:
+        #
+        # - only `X` concurrent `FetchProfileJob`s
+        # - max `XX` `FetchProfileJob` per hour
+        # - only `Y` concurrent `FetchCommentsJob`s
+        # - max `YY` `FetchCommentsJob` per hour
+        #
+        # Where `(X + Y) == 10` and `(XX + YY) == 123`
+        #
+        # @see Registry.add_alias
+        # @return [void]
+        def sidekiq_throttle_as(name)
+          Registry.add_alias(self, name)
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/message.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module Throttled
+    class Message
+      def initialize(item)
+        @item = item.is_a?(Hash) ? item : parse(item)
+      end
+
+      def job_class
+        @item.fetch("wrapped") { @item["class"] }
+      end
+
+      def job_args
+        @item.key?("wrapped") ? @item.dig("args", 0, "arguments") : @item["args"]
+      end
+
+      def job_id
+        @item["jid"]
+      end
+
+      private
+
+      def parse(item)
+        item = Sidekiq.load_json(item)
+        item.is_a?(Hash) ? item : {}
+      rescue JSON::ParserError
+        {}
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/middlewares/server.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+# internal
+require_relative "../message"
+require_relative "../registry"
+
+module Sidekiq
+  module Throttled
+    module Middlewares
+      # Server middleware required for Sidekiq::Throttled functioning.
+      class Server
+        include Sidekiq::ServerMiddleware
+
+        def call(_worker, msg, _queue)
+          yield
+        ensure
+          message = Message.new(msg)
+
+          if message.job_class && message.job_id
+            Registry.get(message.job_class) do |strategy|
+              strategy.finalize!(message.job_id, *message.job_args)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/patches/basic_fetch.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require "sidekiq"
+require "sidekiq/fetch"
+
+require_relative "./throttled_retriever"
+
+module Sidekiq
+  module Throttled
+    module Patches
+      module BasicFetch
+        def self.prepended(base)
+          base.prepend(ThrottledRetriever)
+        end
+
+        private
+
+        # Returns list of queues to try to fetch jobs from.
+        #
+        # @note It may return an empty array.
+        # @param [Array<String>] queues
+        # @return [Array<String>]
+        def queues_cmd
+          throttled_queues = Throttled.cooldown&.queues
+          return super if throttled_queues.nil? || throttled_queues.empty?
+
+          super - throttled_queues
+        end
+      end
+    end
+  end
+end
+
+Sidekiq::BasicFetch.prepend(Sidekiq::Throttled::Patches::BasicFetch)

data/lib/sidekiq/throttled/patches/super_fetch.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require "sidekiq"
+
+require_relative "./throttled_retriever"
+
+module Sidekiq
+  module Throttled
+    module Patches
+      module SuperFetch
+        def self.prepended(base)
+          base.prepend(ThrottledRetriever)
+        end
+
+        private
+
+        # Returns list of non-paused queues to try to fetch jobs from.
+        #
+        # @note It may return an empty array.
+        # @return [Array<Array(String, String)>]
+        def active_queues
+          # Create a hash of throttled queues for fast lookup
+          throttled_queues = Throttled.cooldown&.queues&.to_h { |queue| [queue, true] }
+          return super if throttled_queues.nil? || throttled_queues.empty?
+
+          # Reject throttled queues from the list of active queues
+          super.reject { |queue, _private_queue| throttled_queues[queue] }
+        end
+      end
+    end
+  end
+end
+
+begin
+  require "sidekiq/pro/super_fetch"
+  Sidekiq::Pro::SuperFetch.prepend(Sidekiq::Throttled::Patches::SuperFetch)
+rescue LoadError
+  # Sidekiq Pro is not available
+end

data/lib/sidekiq/throttled/patches/throttled_retriever.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module Throttled
+    module Patches
+      module ThrottledRetriever
+        # Retrieves job from redis.
+        #
+        # @return [Sidekiq::BasicFetch::UnitOfWork, nil]
+        def retrieve_work
+          work = super
+
+          if work && Throttled.throttled?(work.job)
+            Throttled.cooldown&.notify_throttled(work.queue)
+            Throttled.requeue_throttled(work)
+            return nil
+          end
+
+          Throttled.cooldown&.notify_admitted(work.queue) if work
+
+          work
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/registry.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+# internal
+require_relative "./strategy"
+
+module Sidekiq
+  module Throttled
+    # Registred strategies.
+    #
+    # @private
+    module Registry
+      @strategies = {}
+      @aliases    = {}
+
+      class << self
+        # Adds strategy to the registry.
+        #
+        # @param (see Strategy#initialize)
+        # @return [Strategy]
+        def add(name, **kwargs)
+          name = name.to_s
+
+          @strategies[name] = Strategy.new(name, **kwargs)
+        end
+
+        # Adds alias for existing strategy.
+        #
+        # @param (#to_s) new_name
+        # @param (#to_s) old_name
+        # @raise [RuntimeError] if no strategy found with `old_name`
+        # @return [Strategy]
+        def add_alias(new_name, old_name)
+          new_name = new_name.to_s
+          old_name = old_name.to_s
+
+          raise "Strategy not found: #{old_name}" unless @strategies[old_name]
+
+          @aliases[new_name] = @strategies[old_name]
+        end
+
+        # @overload get(name)
+        #   @param [#to_s] name
+        #   @return [Strategy, nil] registred strategy
+        #
+        # @overload get(name, &block)
+        #   Yields control to the block if requested strategy was found.
+        #   @param [#to_s] name
+        #   @yieldparam [Strategy] strategy
+        #   @yield [strategy] Gives found strategy to the block
+        #   @return result of a block
+        def get(name)
+          strategy = find(name.to_s) || find_by_class(name)
+
+          return yield strategy if strategy && block_given?
+
+          strategy
+        end
+
+        # @overload each()
+        #   @return [Enumerator]
+        #
+        # @overload each(&block)
+        #   @yieldparam [String] name
+        #   @yieldparam [Strategy] strategy
+        #   @yield [strategy] Gives strategy to the block
+        #   @return [Registry]
+        def each
+          return to_enum(__method__) unless block_given?
+
+          @strategies.each { |*args| yield(*args) }
+          self
+        end
+
+        # @overload each_with_static_keys()
+        #   @return [Enumerator]
+        #
+        # @overload each_with_static_keys(&block)
+        #   @yieldparam [String] name
+        #   @yieldparam [Strategy] strategy
+        #   @yield [strategy] Gives strategy to the block
+        #   @return [Registry]
+        def each_with_static_keys
+          return to_enum(__method__) unless block_given?
+
+          @strategies.each do |name, strategy|
+            yield(name, strategy) unless strategy.dynamic?
+          end
+        end
+
+        private
+
+        # Find strategy by it's name.
+        #
+        # @param name [String]
+        # @return [Strategy, nil]
+        def find(name)
+          @strategies[name] || @aliases[name]
+        end
+
+        # Find strategy by class or it's parents.
+        #
+        # @param name [Class, #to_s]
+        # @return [Strategy, nil]
+        def find_by_class(name)
+          const = name.is_a?(Class) ? name : Object.const_get(name)
+          return unless const.is_a?(Class)
+
+          const.ancestors.each do |m|
+            strategy = find(m.name)
+            return strategy if strategy
+          end
+
+          nil
+        rescue NameError
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/strategy/base.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module Throttled
+    class Strategy
+      module Base
+        def limit(job_args = nil)
+          @limit.respond_to?(:call) ? @limit.call(*job_args) : @limit
+        end
+
+        private
+
+        def key(job_args)
+          key = @base_key.dup
+          return key unless @key_suffix
+
+          key << ":#{@key_suffix.call(*job_args)}"
+        rescue StandardError => e
+          Sidekiq.logger.error "Failed to get key suffix: #{e}"
+          raise e
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/strategy/concurrency.lua

@@ -0,0 +1,61 @@
+local in_progress_jobs_key = KEYS[1]
+local backlog_info_key = KEYS[2]
+local jid = ARGV[1]
+local lmt = tonumber(ARGV[2])
+local lost_job_threshold = tonumber(ARGV[3])
+local now = tonumber(ARGV[4])
+
+-- supporting functions
+local function over_limit()
+  return lmt <= redis.call("ZCARD", in_progress_jobs_key)
+end
+
+local function job_already_in_progress()
+  return redis.call("ZSCORE", in_progress_jobs_key, jid)
+end
+
+-- Estimates current backlog size. This function tends to underestimate 
+-- the actual backlog. This is intentional. Overestimates are bad as it
+-- can cause unnecessary delays in job processing. Underestimates are much
+-- safer as they only increase workload of sidekiq processors. 
+local function est_current_backlog_size()
+  local old_size = tonumber(redis.call("HGET", backlog_info_key, "size") or 0)
+  local old_timestamp = tonumber(redis.call("HGET", backlog_info_key, "timestamp") or now)
+  
+  local jobs_lost_since_old_timestamp = (now - old_timestamp) / lost_job_threshold * lmt
+
+  return math.max(old_size - jobs_lost_since_old_timestamp, 0) 
+end
+
+
+local function change_backlog_size(delta)
+  local curr_backlog_size = est_current_backlog_size()
+
+  redis.call("HSET", backlog_info_key, "size", curr_backlog_size + delta) 
+  redis.call("HSET", backlog_info_key, "timestamp", now)
+  redis.call("EXPIRE", backlog_info_key, math.ceil((lost_job_threshold * curr_backlog_size) + 1 / lmt))
+end
+
+local function register_job_in_progress()
+  redis.call("ZADD", in_progress_jobs_key, now + lost_job_threshold , jid)
+  redis.call("EXPIRE", in_progress_jobs_key, lost_job_threshold)
+end
+
+local function clear_stale_in_progress_jobs()
+  local cleared_count = redis.call("ZREMRANGEBYSCORE", in_progress_jobs_key, "-inf", "(" .. now)
+  change_backlog_size(-cleared_count)
+end
+
+-- END supporting functions
+
+clear_stale_in_progress_jobs()
+
+if over_limit() and not job_already_in_progress() then
+  change_backlog_size(1)
+  return 1
+end
+
+register_job_in_progress()
+change_backlog_size(-1)
+
+return 0

data/lib/sidekiq/throttled/strategy/concurrency.rb

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+require "redis_prescription"
+
+require_relative "./base"
+
+module Sidekiq
+  module Throttled
+    class Strategy
+      # Concurrency throttling strategy
+      class Concurrency
+        include Base
+
+        # LUA script used to limit fetch concurrency.
+        # Logic behind the scene can be described in following pseudo code:
+        #
+        #     if @limit <= LLEN(@key)
+        #       return 1
+        #     else
+        #       PUSH(@key, @jid)
+        #       return 0
+        #     end
+        SCRIPT = RedisPrescription.new(File.read("#{__dir__}/concurrency.lua"))
+        private_constant :SCRIPT
+
+        # @param [#to_s] strategy_key
+        # @param [#to_i, #call] limit Amount of allowed concurrent jobs
+        #   per processors running for given key.
+        # @param [#to_i] avg_job_duration Average number of seconds needed
+        #   to complete a job of this type. Default: 300 or 1/3 of lost_job_threshold
+        # @param [#to_i] lost_job_threshold Seconds to wait before considering
+        #   a job lost or dead. Default: 900 or 3 * avg_job_duration
+        # @param [Proc] key_suffix Dynamic key suffix generator.
+        # @deprecated @param [#to_i] ttl Obsolete alias for `lost_job_threshold`.
+        #   Default: 900 or 3 * avg_job_duration
+        def initialize(strategy_key, limit:, avg_job_duration: nil, ttl: nil, lost_job_threshold: ttl, key_suffix: nil) # rubocop:disable Metrics/ParameterLists
+          @base_key = "#{strategy_key}:concurrency.v2"
+          @limit = limit
+          @avg_job_duration, @lost_job_threshold = interp_duration_args(avg_job_duration, lost_job_threshold)
+          @key_suffix = key_suffix
+
+          raise(ArgumentError, "lost_job_threshold must be greater than avg_job_duration") if
+            @lost_job_threshold <= @avg_job_duration
+        end
+
+        # @return [Boolean] Whenever strategy has dynamic config
+        def dynamic?
+          @key_suffix || @limit.respond_to?(:call)
+        end
+
+        # @return [Boolean] whenever job is throttled or not
+        def throttled?(jid, *job_args)
+          job_limit = limit(job_args)
+          return false unless job_limit
+          return true if job_limit <= 0
+
+          keys = [key(job_args), backlog_info_key(job_args)]
+          argv = [jid.to_s, job_limit, @lost_job_threshold, Time.now.to_f]
+
+          Sidekiq.redis { |redis| (1 == SCRIPT.call(redis, keys: keys, argv: argv)) }
+        end
+
+        # @return [Float] How long, in seconds, before we'll next be able to take on jobs
+        def retry_in(_jid, *job_args)
+          job_limit = limit(job_args)
+          return 0.0 if !job_limit || count(*job_args) < job_limit
+
+          estimated_backlog_size(job_args) * @avg_job_duration / limit(job_args)
+        end
+
+        # @return [Integer] Current count of jobs
+        def count(*job_args)
+          Sidekiq.redis { |conn| conn.zcard(key(job_args)) }.to_i
+        end
+
+        # Resets count of jobs
+        # @return [void]
+        def reset!(*job_args)
+          Sidekiq.redis { |conn| conn.del(key(job_args)) }
+        end
+
+        # Remove jid from the pool of jobs in progress
+        # @return [void]
+        def finalize!(jid, *job_args)
+          Sidekiq.redis do |conn|
+            conn.zrem(key(job_args), jid.to_s)
+          end
+        end
+
+        private
+
+        def backlog_info_key(job_args)
+          "#{key(job_args)}.backlog_info"
+        end
+
+        def estimated_backlog_size(job_args)
+          old_size_str, old_timestamp_str =
+            Sidekiq.redis { |conn| conn.hmget(backlog_info_key(job_args), "size", "timestamp") }
+          old_size = (old_size_str || 0).to_f
+          old_timestamp = (old_timestamp_str || Time.now).to_f
+
+          nonneg(old_size - jobs_lost_since(old_timestamp, job_args))
+        end
+
+        def jobs_lost_since(timestamp, job_args)
+          (Time.now.to_f - timestamp) / @lost_job_threshold * limit(job_args)
+        end
+
+        def nonneg(number)
+          [number, 0].max
+        end
+
+        def interp_duration_args(avg_job_duration, lost_job_threshold)
+          if avg_job_duration && lost_job_threshold
+            [avg_job_duration.to_i, lost_job_threshold.to_i]
+          elsif avg_job_duration && lost_job_threshold.nil?
+            [avg_job_duration.to_i, avg_job_duration.to_i * 3]
+          elsif avg_job_duration.nil? && lost_job_threshold
+            [lost_job_threshold.to_i / 3, lost_job_threshold.to_i]
+          else
+            [300, 900]
+          end
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/strategy/threshold.lua

@@ -0,0 +1,14 @@
+local key = KEYS[1]
+local lmt = tonumber(ARGV[1])
+local ttl = tonumber(ARGV[2])
+local now = tonumber(ARGV[3])
+
+if lmt <= redis.call("LLEN", key) and now - redis.call("LINDEX", key, -1) < ttl then
+  return 1
+end
+
+redis.call("LPUSH", key, now)
+redis.call("LTRIM", key, 0, lmt - 1)
+redis.call("EXPIRE", key, ttl)
+
+return 0