kubra-sidekiq-throttled 1.5.3

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

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+require "redis_prescription"
+
+require_relative "./base"
+
+module Sidekiq
+  module Throttled
+    class Strategy
+      # Threshold throttling strategy
+      # @todo Use redis TIME command instead of sending current timestamp from
+      #   sidekiq manager. See: http://redis.io/commands/time
+      class Threshold
+        include Base
+
+        # LUA script used to limit fetch threshold.
+        # Logic behind the scene can be described in following pseudo code:
+        #
+        #     def exceeded?
+        #       limit <= LLEN(@key) && NOW - LINDEX(@key, -1) < @period
+        #     end
+        #
+        #     def increase!
+        #       LPUSH(@key, NOW)
+        #       LTRIM(@key, 0, @limit - 1)
+        #       EXPIRE(@key, @period)
+        #     end
+        #
+        #     return 1 if exceeded?
+        #
+        #     increase!
+        #     return 0
+        SCRIPT = RedisPrescription.new(File.read("#{__dir__}/threshold.lua"))
+        private_constant :SCRIPT
+
+        # @param [#to_s] strategy_key
+        # @param [#to_i, #call] limit Amount of allowed concurrent jobs
+        #   per period running for given key.
+        # @param [#to_f, #call] :period Period in seconds.
+        # @param [Proc] key_suffix Dynamic key suffix generator.
+        def initialize(strategy_key, limit:, period:, key_suffix: nil)
+          @base_key   = "#{strategy_key}:threshold"
+          @limit      = limit
+          @period     = period
+          @key_suffix = key_suffix
+        end
+
+        # @return [Float] Period in seconds
+        def period(job_args = nil)
+          return @period.to_f unless @period.respond_to? :call
+
+          @period.call(*job_args).to_f
+        end
+
+        # @return [Boolean] Whenever strategy has dynamic config
+        def dynamic?
+          @key_suffix || @limit.respond_to?(:call) || @period.respond_to?(:call)
+        end
+
+        # @return [Boolean] whenever job is throttled or not
+        def throttled?(*job_args)
+          job_limit = limit(job_args)
+          return false unless job_limit
+          return true if job_limit <= 0
+
+          keys = [key(job_args)]
+          argv = [job_limit, period(job_args), 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(*job_args)
+          job_limit = limit(job_args)
+          return 0.0 if !job_limit || count(*job_args) < job_limit
+
+          job_period = period(job_args)
+          job_key = key(job_args)
+          time_since_oldest = Time.now.to_f - Sidekiq.redis { |redis| redis.lindex(job_key, -1) }.to_f
+          if time_since_oldest > job_period
+            # The oldest job on our list is from more than the throttling period ago,
+            # which means we have not hit the limit this period.
+            0.0
+          else
+            # If we can only have X jobs every Y minutes, then wait until Y minutes have elapsed
+            # since the oldest job on our list.
+            job_period - time_since_oldest
+          end
+        end
+
+        # @return [Integer] Current count of jobs
+        def count(*job_args)
+          Sidekiq.redis { |conn| conn.llen(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
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/strategy.rb

@@ -0,0 +1,213 @@
+# frozen_string_literal: true
+
+# internal
+require_relative "./errors"
+require_relative "./strategy_collection"
+require_relative "./strategy/concurrency"
+require_relative "./strategy/threshold"
+
+module Sidekiq
+  module Throttled
+    # Meta-strategy that couples {Concurrency} and {Threshold} strategies.
+    #
+    # @private
+    class Strategy # rubocop:disable Metrics/ClassLength
+      # :enqueue means put the job back at the end of the queue immediately
+      # :schedule means schedule enqueueing the job for a later time when we expect to have capacity
+      VALID_VALUES_FOR_REQUEUE_WITH = %i[enqueue schedule].freeze
+
+      # @!attribute [r] concurrency
+      #   @return [Strategy::Concurrency, nil]
+      attr_reader :concurrency
+
+      # @!attribute [r] threshold
+      #   @return [Strategy::Threshold, nil]
+      attr_reader :threshold
+
+      # @!attribute [r] observer
+      #   @return [Proc, nil]
+      attr_reader :observer
+
+      # @!attribute [r] requeue_options
+      #   @return [Hash, nil]
+      attr_reader :requeue_options
+
+      # @param [#to_s] name
+      # @param [Hash] concurrency Concurrency options.
+      #   See keyword args of {Strategy::Concurrency#initialize} for details.
+      # @param [Hash] threshold Threshold options.
+      #   See keyword args of {Strategy::Threshold#initialize} for details.
+      # @param [#call] key_suffix Dynamic key suffix generator.
+      # @param [#call] observer Process called after throttled.
+      # @param [#call] requeue What to do with jobs that are throttled.
+      def initialize(name, concurrency: nil, threshold: nil, key_suffix: nil, observer: nil, requeue: nil) # rubocop:disable Metrics/MethodLength, Metrics/ParameterLists
+        @observer = observer
+
+        @concurrency = StrategyCollection.new(concurrency,
+          strategy:   Concurrency,
+          name:       name,
+          key_suffix: key_suffix)
+
+        @threshold = StrategyCollection.new(threshold,
+          strategy:   Threshold,
+          name:       name,
+          key_suffix: key_suffix)
+
+        @requeue_options = Throttled.config.default_requeue_options.merge(requeue || {})
+
+        validate!
+      end
+
+      # @return [Boolean] whenever strategy has dynamic config
+      def dynamic?
+        return true if @concurrency&.dynamic?
+        return true if @threshold&.dynamic?
+
+        false
+      end
+
+      # @return [Boolean] whenever job is throttled or not.
+      def throttled?(jid, *job_args)
+        if @concurrency&.throttled?(jid, *job_args)
+          @observer&.call(:concurrency, *job_args)
+          return true
+        end
+
+        if @threshold&.throttled?(*job_args)
+          @observer&.call(:threshold, *job_args)
+
+          finalize!(jid, *job_args)
+          return true
+        end
+
+        false
+      end
+
+      # @return [Proc, Symbol] How to requeue the throttled job
+      def requeue_with
+        requeue_options[:with]
+      end
+
+      # @return [String, nil] Name of the queue to re-queue the job to.
+      def requeue_to
+        requeue_options[:to]
+      end
+
+      # Return throttled job to be executed later. Implementation depends on the strategy's `requeue` options.
+      # @return [void]
+      def requeue_throttled(work) # rubocop:disable Metrics/MethodLength
+        # Resolve :with and :to options, calling them if they are Procs
+        job_args = JSON.parse(work.job)["args"]
+        with = requeue_with.respond_to?(:call) ? requeue_with.call(*job_args) : requeue_with
+        target_queue = calc_target_queue(work)
+
+        case with
+        when :enqueue
+          re_enqueue_throttled(work, target_queue)
+        when :schedule
+          # Find out when we will next be able to execute this job, and reschedule for then.
+          reschedule_throttled(work, target_queue)
+        else
+          raise "unrecognized :with option #{with}"
+        end
+      end
+
+      # Marks job as being processed.
+      # @return [void]
+      def finalize!(jid, *job_args)
+        @concurrency&.finalize!(jid, *job_args)
+      end
+
+      # Resets count of jobs of all available strategies
+      # @return [void]
+      def reset!
+        @concurrency&.reset!
+        @threshold&.reset!
+      end
+
+      private
+
+      def validate!
+        unless VALID_VALUES_FOR_REQUEUE_WITH.include?(@requeue_options[:with]) ||
+               @requeue_options[:with].respond_to?(:call)
+          raise ArgumentError, "requeue: #{@requeue_options[:with]} is not a valid value for :with"
+        end
+
+        raise ArgumentError, "Neither :concurrency nor :threshold given" unless @concurrency.any? || @threshold.any?
+      end
+
+      def calc_target_queue(work) # rubocop:disable Metrics/MethodLength
+        target = case requeue_to
+                 when Proc, Method
+                   requeue_to.call(*JSON.parse(work.job)["args"])
+                 when NilClass
+                   work.queue
+                 when String, Symbol
+                   requeue_to
+                 else
+                   raise ArgumentError, "Invalid argument for `to`"
+                 end
+
+        target = work.queue if target.nil? || target.empty?
+
+        target.to_s
+      end
+
+      # Push the job back to the head of the queue.
+      # The queue name is expected to include the "queue:" prefix, so we add it if it's missing.
+      def re_enqueue_throttled(work, target_queue)
+        target_queue = "queue:#{target_queue}" unless target_queue.start_with?("queue:")
+
+        case work.class.name
+        when "Sidekiq::Pro::SuperFetch::UnitOfWork"
+          # Calls SuperFetch UnitOfWork's requeue to remove the job from the
+          # temporary queue and push job back to the head of the target queue, so that
+          # the job won't be tried immediately after it was requeued (in most cases).
+          work.queue = target_queue
+          work.requeue
+        else
+          # This is the same operation Sidekiq performs upon `Sidekiq::Worker.perform_async` call.
+          Sidekiq.redis { |conn| conn.lpush(target_queue, work.job) }
+        end
+      end
+
+      # Reschedule the job to be executed later in the target queue.
+      # The queue name should NOT include the "queue:" prefix, so we remove it if it's present.
+      def reschedule_throttled(work, target_queue)
+        target_queue = target_queue.delete_prefix("queue:")
+        message      = JSON.parse(work.job)
+        job_class    = message.fetch("wrapped") { message.fetch("class") { return false } }
+        job_args     = message["args"]
+
+        # Re-enqueue the job to the target queue at another time as a NEW unit of work
+        # AND THEN mark this work as done, so SuperFetch doesn't think this instance is orphaned
+        # Technically, the job could processed twice if the process dies between the two lines,
+        # but your job should be idempotent anyway, right?
+        # The job running twice was already a risk with SuperFetch anyway and this doesn't really increase that risk.
+        Sidekiq::Client.enqueue_to_in(target_queue, retry_in(work), Object.const_get(job_class), *job_args)
+
+        work.acknowledge
+      end
+
+      def retry_in(work)
+        message = JSON.parse(work.job)
+        jid = message.fetch("jid") { return false }
+        job_args = message["args"]
+
+        # Ask both concurrency and threshold, if relevant, how long minimum until we can retry.
+        # If we get two answers, take the longer one.
+        intervals = [@concurrency&.retry_in(jid, *job_args), @threshold&.retry_in(*job_args)].compact
+
+        raise "Cannot compute a valid retry interval" if intervals.empty?
+
+        interval = intervals.max
+
+        # Add a random amount of jitter, proportional to the length of the minimum retry time.
+        # This helps spread out jobs more evenly and avoid clumps of jobs on the queue.
+        interval += rand(interval / 5) if interval > 10
+
+        interval
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/strategy_collection.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+# internal
+module Sidekiq
+  module Throttled
+    # Collection which transparently group several meta-strategies of one kind
+    #
+    # @private
+    class StrategyCollection
+      include Enumerable
+
+      attr_reader :strategies
+
+      # @param [Hash, Array, nil] strategies Concurrency or Threshold options
+      #   or array of options.
+      #   See keyword args of {Strategy::Concurrency#initialize} for details.
+      #   See keyword args of {Strategy::Threshold#initialize} for details.
+      # @param [Class] strategy class of strategy: Concurrency or Threshold
+      # @param [#to_s] name
+      # @param [#call] key_suffix Dynamic key suffix generator.
+      def initialize(strategies, strategy:, name:, key_suffix:)
+        @strategies = (strategies.is_a?(Hash) ? [strategies] : Array(strategies)).map do |options|
+          make_strategy(strategy, name, key_suffix, options)
+        end
+      end
+
+      # @param [#call] block
+      # Iterates each strategy in collection
+      def each(...)
+        @strategies.each(...)
+      end
+
+      # @return [Boolean] whenever any strategy in collection has dynamic config
+      def dynamic?
+        any?(&:dynamic?)
+      end
+
+      # @return [Boolean] whenever job is throttled or not
+      # by any strategy in collection.
+      def throttled?(...)
+        any? { |s| s.throttled?(...) }
+      end
+
+      # @return [Float] How long, in seconds, before we'll next be able to take on jobs
+      def retry_in(*args)
+        map { |s| s.retry_in(*args) }.max
+      end
+
+      # Marks job as being processed.
+      # @return [void]
+      def finalize!(...)
+        each { |c| c.finalize!(...) }
+      end
+
+      # Resets count of jobs of all avaliable strategies
+      # @return [void]
+      def reset!
+        each(&:reset!)
+      end
+
+      private
+
+      # @return [Base, nil]
+      def make_strategy(strategy, name, key_suffix, options)
+        return unless options
+
+        strategy.new("throttled:#{name}",
+          key_suffix: key_suffix,
+          **options)
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/version.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module Throttled
+    # Gem version
+    VERSION = "1.5.3"
+  end
+end

data/lib/sidekiq/throttled/web/stats.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module Sidekiq
+  module Throttled
+    module Web
+      # Throttle strategy stats generation helper
+      class Stats
+        TIME_CONVERSION = [
+          [60 * 60 * 24,  "day",    "days"],
+          [60 * 60,       "hour",   "hours"],
+          [60,            "minute", "minutes"],
+          [1,             "second", "seconds"]
+        ].freeze
+
+        # @param [Strategy::Concurrency, Strategy::Threshold] strategy
+        def initialize(strategy)
+          raise ArgumentError, "Can't handle dynamic strategies" if strategy&.dynamic?
+
+          @strategy = strategy
+        end
+
+        # @return [String]
+        def to_html
+          return "" unless @strategy
+
+          html = humanize_integer(@strategy.limit) << " jobs"
+
+          html << " per " << humanize_duration(@strategy.period) if @strategy.respond_to?(:period)
+
+          html << "<br />" << colorize_count(@strategy.count, @strategy.limit)
+        end
+
+        private
+
+        # @return [String]
+        def colorize_count(int, max)
+          percentile = 100.00 * int / max
+          lvl = if    80 <= percentile then "danger"
+                elsif 60 <= percentile then "warning"
+                else
+                  "success"
+                end
+
+          %(<span class="label label-#{lvl}">#{int}</span>)
+        end
+
+        # @return [String]
+        def humanize_duration(int)
+          arr = []
+
+          TIME_CONVERSION.each do |(dimension, unit, units)|
+            count = (int / dimension).to_i
+
+            next unless count.positive?
+
+            int -= count * dimension
+            arr << "#{count} #{1 == count ? unit : units}"
+          end
+
+          arr.join " "
+        end
+
+        # @return [String]
+        def humanize_integer(int)
+          digits = int.to_s.chars
+          str    = digits.shift(digits.count % 3).join
+
+          str << " " << digits.shift(3).join while digits.count.positive?
+
+          str.strip
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/throttled/web/throttled.html.erb

@@ -0,0 +1,35 @@
+<h3>Throttled</h3>
+
+<div class="table_container">
+  <table class="table table-hover table-bordered table-striped table-white">
+    <thead>
+      <tr>
+        <th>Name</th>
+        <th style="text-align:center;">Concurrency</th>
+        <th style="text-align:center;">Threshold</th>
+        <th style="text-align:center;">Actions</th>
+      </tr>
+    </thead>
+    <% Sidekiq::Throttled::Registry.each_with_static_keys do |name, strategy| %>
+      <tr>
+        <td style="vertical-align:middle;"><%= name %></td>
+        <td style="vertical-align:middle;text-align:center;">
+          <% strategy.concurrency.each do |concurrency| %>
+            <%= Sidekiq::Throttled::Web::Stats.new(concurrency).to_html %>
+          <% end %>
+        </td>
+        <td style="vertical-align:middle;text-align:center;">
+          <% strategy.threshold.each do |threshold| %>
+            <%= Sidekiq::Throttled::Web::Stats.new(threshold).to_html %>
+          <% end %>
+        </td>
+        <td style="vertical-align:middle;text-align:center;">
+          <form action="<%= root_path %>throttled/<%= CGI.escape name %>/reset" method="post">
+            <%= csrf_tag %>
+            <button class="btn btn-danger" type="submit">Reset</button>
+          </form>
+        </td>
+      </tr>
+    <% end %>
+  </table>
+</div>

data/lib/sidekiq/throttled/web.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+# stdlib
+require "pathname"
+
+# 3rd party
+require "sidekiq"
+require "sidekiq/web"
+
+# internal
+require_relative "./registry"
+require_relative "./web/stats"
+
+module Sidekiq
+  module Throttled
+    # Provides Sidekiq tab to monitor and reset throttled stats.
+    module Web
+      VIEWS         = Pathname.new(__dir__).join("web")
+      THROTTLED_TPL = VIEWS.join("throttled.html.erb").read.freeze
+
+      class << self
+        # @api private
+        def registered(app)
+          register_throttled_tab app
+        end
+
+        private
+
+        def register_throttled_tab(app)
+          app.get("/throttled") { erb THROTTLED_TPL.dup }
+
+          app.post("/throttled/:id/reset") do
+            Registry.get(params[:id], &:reset!)
+            redirect "#{root_path}throttled"
+          end
+        end
+      end
+    end
+  end
+end
+
+Sidekiq::Web.register Sidekiq::Throttled::Web
+Sidekiq::Web.tabs["Throttled"] = "throttled"

data/lib/sidekiq/throttled/worker.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require_relative "./job"
+
+module Sidekiq
+  module Throttled
+    # A new module, Sidekiq::Job, was added in Sidekiq version 6.3.0 as a
+    # simple alias for Sidekiq::Worker as the term "worker" was considered
+    # too generic and confusing. Many people call a Sidekiq process a "worker"
+    # whereas others call the thread that executes jobs a "worker".
+    Worker = Job
+  end
+end

data/lib/sidekiq/throttled.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+require "sidekiq"
+
+require_relative "./throttled/config"
+require_relative "./throttled/cooldown"
+require_relative "./throttled/job"
+require_relative "./throttled/message"
+require_relative "./throttled/middlewares/server"
+require_relative "./throttled/patches/basic_fetch"
+require_relative "./throttled/patches/super_fetch"
+require_relative "./throttled/registry"
+require_relative "./throttled/version"
+require_relative "./throttled/worker"
+
+# @see https://github.com/mperham/sidekiq/
+module Sidekiq
+  # Concurrency and threshold throttling for Sidekiq.
+  #
+  # Just add somewhere in your bootstrap:
+  #
+  #     require "sidekiq/throttled"
+  #
+  # Once you've done that you can include {Sidekiq::Throttled::Job} to your
+  # job classes and configure throttling:
+  #
+  #     class MyJob
+  #       include Sidekiq::Job
+  #       include Sidekiq::Throttled::Job
+  #
+  #       sidekiq_options :queue => :my_queue
+  #
+  #       sidekiq_throttle({
+  #         # Allow maximum 10 concurrent jobs of this class at a time.
+  #         :concurrency => { :limit => 10 },
+  #         # Allow maximum 1K jobs being processed within one hour window.
+  #         :threshold => { :limit => 1_000, :period => 1.hour }
+  #       })
+  #
+  #       def perform
+  #         # ...
+  #       end
+  #     end
+  module Throttled
+    MUTEX = Mutex.new
+    private_constant :MUTEX
+
+    @config   = Config.new.freeze
+    @cooldown = Cooldown[@config]
+
+    class << self
+      # @api internal
+      #
+      # @return [Cooldown, nil]
+      attr_reader :cooldown
+
+      # @api internal
+      #
+      # @return [Config, nil]
+      attr_reader :config
+
+      # @example
+      #   Sidekiq::Throttled.configure do |config|
+      #     config.cooldown_period = nil # Disable queues cooldown manager
+      #   end
+      #
+      # @yieldparam config [Config]
+      def configure
+        MUTEX.synchronize do
+          config = @config.dup
+
+          yield config
+
+          @config   = config.freeze
+          @cooldown = Cooldown[@config]
+        end
+      end
+
+      # Tells whenever job is throttled or not.
+      #
+      # @param [String] message Job's JSON payload
+      # @return [Boolean]
+      def throttled?(message)
+        message = Message.new(message)
+        return false unless message.job_class && message.job_id
+
+        Registry.get(message.job_class) do |strategy|
+          return strategy.throttled?(message.job_id, *message.job_args)
+        end
+
+        false
+      rescue StandardError
+        false
+      end
+
+      # Return throttled job to be executed later, delegating the details of how to do that
+      # to the Strategy for that job.
+      #
+      # @return [void]
+      def requeue_throttled(work)
+        message = JSON.parse(work.job)
+        job_class = Object.const_get(message.fetch("wrapped") { message.fetch("class") { return false } })
+
+        Registry.get job_class do |strategy|
+          strategy.requeue_throttled(work)
+        end
+      end
+    end
+  end
+
+  configure_server do |config|
+    config.server_middleware do |chain|
+      chain.add(Sidekiq::Throttled::Middlewares::Server)
+    end
+  end
+end