sidekiq 7.1.4 → 8.0.9

data/lib/sidekiq/job/iterable/enumerators.rb

@@ -0,0 +1,135 @@
+# frozen_string_literal: true
+
+require_relative "active_record_enumerator"
+require_relative "csv_enumerator"
+
+module Sidekiq
+  module Job
+    module Iterable
+      module Enumerators
+        # Builds Enumerator object from a given array, using +cursor+ as an offset.
+        #
+        # @param array [Array]
+        # @param cursor [Integer] offset to start iteration from
+        #
+        # @return [Enumerator]
+        #
+        # @example
+        #   array_enumerator(['build', 'enumerator', 'from', 'any', 'array'], cursor: cursor)
+        #
+        def array_enumerator(array, cursor:)
+          raise ArgumentError, "array must be an Array" unless array.is_a?(Array)
+
+          x = array.each_with_index.drop(cursor || 0)
+          x.to_enum { x.size }
+        end
+
+        # Builds Enumerator from `ActiveRecord::Relation`.
+        # Each Enumerator tick moves the cursor one row forward.
+        #
+        # @param relation [ActiveRecord::Relation] relation to iterate
+        # @param cursor [Object] offset id to start iteration from
+        # @param options [Hash] additional options that will be passed to relevant
+        #   ActiveRecord batching methods
+        #
+        # @return [ActiveRecordEnumerator]
+        #
+        # @example
+        #   def build_enumerator(cursor:)
+        #     active_record_records_enumerator(User.all, cursor: cursor)
+        #   end
+        #
+        #   def each_iteration(user)
+        #     user.notify_about_something
+        #   end
+        #
+        def active_record_records_enumerator(relation, cursor:, **options)
+          ActiveRecordEnumerator.new(relation, cursor: cursor, **options).records
+        end
+
+        # Builds Enumerator from `ActiveRecord::Relation` and enumerates on batches of records.
+        # Each Enumerator tick moves the cursor `:batch_size` rows forward.
+        # @see #active_record_records_enumerator
+        #
+        # @example
+        #   def build_enumerator(product_id, cursor:)
+        #     active_record_batches_enumerator(
+        #       Comment.where(product_id: product_id).select(:id),
+        #       cursor: cursor,
+        #       batch_size: 100
+        #     )
+        #   end
+        #
+        #   def each_iteration(batch_of_comments, product_id)
+        #     comment_ids = batch_of_comments.map(&:id)
+        #     CommentService.call(comment_ids: comment_ids)
+        #   end
+        #
+        def active_record_batches_enumerator(relation, cursor:, **options)
+          ActiveRecordEnumerator.new(relation, cursor: cursor, **options).batches
+        end
+
+        # Builds Enumerator from `ActiveRecord::Relation` and enumerates on batches,
+        # yielding `ActiveRecord::Relation`s.
+        # @see #active_record_records_enumerator
+        #
+        # @example
+        #   def build_enumerator(product_id, cursor:)
+        #     active_record_relations_enumerator(
+        #       Product.find(product_id).comments,
+        #       cursor: cursor,
+        #       batch_size: 100,
+        #     )
+        #   end
+        #
+        #   def each_iteration(batch_of_comments, product_id)
+        #     # batch_of_comments will be a Comment::ActiveRecord_Relation
+        #     batch_of_comments.update_all(deleted: true)
+        #   end
+        #
+        def active_record_relations_enumerator(relation, cursor:, **options)
+          ActiveRecordEnumerator.new(relation, cursor: cursor, **options).relations
+        end
+
+        # Builds Enumerator from a CSV file.
+        #
+        # @param csv [CSV] an instance of CSV object
+        # @param cursor [Integer] offset to start iteration from
+        #
+        # @example
+        #   def build_enumerator(import_id, cursor:)
+        #     import = Import.find(import_id)
+        #     csv_enumerator(import.csv, cursor: cursor)
+        #   end
+        #
+        #   def each_iteration(csv_row)
+        #     # insert csv_row into database
+        #   end
+        #
+        def csv_enumerator(csv, cursor:)
+          CsvEnumerator.new(csv).rows(cursor: cursor)
+        end
+
+        # Builds Enumerator from a CSV file and enumerates on batches of records.
+        #
+        # @param csv [CSV] an instance of CSV object
+        # @param cursor [Integer] offset to start iteration from
+        # @option options :batch_size [Integer] (100) size of the batch
+        #
+        # @example
+        #   def build_enumerator(import_id, cursor:)
+        #     import = Import.find(import_id)
+        #     csv_batches_enumerator(import.csv, cursor: cursor)
+        #   end
+        #
+        #   def each_iteration(batch_of_csv_rows)
+        #     # ...
+        #   end
+        #
+        def csv_batches_enumerator(csv, cursor:, **options)
+          CsvEnumerator.new(csv).batches(cursor: cursor, **options)
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/job/iterable.rb

@@ -0,0 +1,322 @@
+# frozen_string_literal: true
+
+require_relative "iterable/enumerators"
+
+module Sidekiq
+  module Job
+    class Interrupted < ::RuntimeError; end
+
+    module Iterable
+      include Enumerators
+
+      # @api private
+      def self.included(base)
+        base.extend(ClassMethods)
+      end
+
+      # @api private
+      module ClassMethods
+        def method_added(method_name)
+          raise "#{self} is an iterable job and must not define #perform" if method_name == :perform
+          super
+        end
+      end
+
+      # @api private
+      def initialize
+        super
+
+        @_executions = 0
+        @_cursor = nil
+        @_start_time = nil
+        @_runtime = 0
+        @_args = nil
+        @_cancelled = nil
+        @current_object = nil
+      end
+
+      # Access to the current object while iterating.
+      # This value is not reset so the latest element is
+      # explicitly available to cleanup/complete callbacks.
+      attr_reader :current_object
+
+      def arguments
+        @_args
+      end
+
+      # Three days is the longest period you generally need to wait for a retry to
+      # execute when using the default retry scheme. We don't want to "forget" the job
+      # is cancelled before it has a chance to execute and cancel itself.
+      CANCELLATION_PERIOD = (3 * 86_400).to_s
+
+      # Set a flag in Redis to mark this job as cancelled.
+      # Cancellation is asynchronous and is checked at the start of iteration
+      # and every 5 seconds thereafter as part of the recurring state flush.
+      def cancel!
+        return @_cancelled if cancelled?
+
+        key = iteration_key
+        _, result, _ = Sidekiq.redis do |c|
+          c.pipelined do |p|
+            p.hsetnx(key, "cancelled", Time.now.to_i)
+            p.hget(key, "cancelled")
+            p.expire(key, Sidekiq::Job::Iterable::STATE_TTL, "nx")
+          end
+        end
+        @_cancelled = result.to_i
+      end
+
+      def cancelled?
+        @_cancelled
+      end
+
+      def cursor
+        @_cursor.freeze
+      end
+
+      # A hook to override that will be called when the job starts iterating.
+      #
+      # It is called only once, for the first time.
+      #
+      def on_start
+      end
+
+      # A hook to override that will be called around each iteration.
+      #
+      # Can be useful for some metrics collection, performance tracking etc.
+      #
+      def around_iteration
+        yield
+      end
+
+      # A hook to override that will be called when the job resumes iterating.
+      #
+      def on_resume
+      end
+
+      # A hook to override that will be called each time the job is interrupted.
+      #
+      # This can be due to interruption or sidekiq stopping.
+      #
+      def on_stop
+      end
+
+      # A hook to override that will be called when the job is cancelled.
+      #
+      def on_cancel
+      end
+
+      # A hook to override that will be called when the job finished iterating.
+      #
+      def on_complete
+      end
+
+      # The enumerator to be iterated over.
+      #
+      # @return [Enumerator]
+      #
+      # @raise [NotImplementedError] with a message advising subclasses to
+      #     implement an override for this method.
+      #
+      def build_enumerator(*)
+        raise NotImplementedError, "#{self.class.name} must implement a '#build_enumerator' method"
+      end
+
+      # The action to be performed on each item from the enumerator.
+      #
+      # @return [void]
+      #
+      # @raise [NotImplementedError] with a message advising subclasses to
+      #     implement an override for this method.
+      #
+      def each_iteration(*)
+        raise NotImplementedError, "#{self.class.name} must implement an '#each_iteration' method"
+      end
+
+      def iteration_key
+        "it-#{jid}"
+      end
+
+      # @api private
+      def perform(*args)
+        @_args = args.dup.freeze
+        fetch_previous_iteration_state
+
+        @_executions += 1
+        @_start_time = mono_now
+
+        enumerator = build_enumerator(*args, cursor: @_cursor)
+        unless enumerator
+          logger.info("'#build_enumerator' returned nil, skipping the job.")
+          return
+        end
+
+        assert_enumerator!(enumerator)
+
+        if @_executions == 1
+          on_start
+        else
+          on_resume
+        end
+
+        completed = catch(:abort) do
+          iterate_with_enumerator(enumerator, args)
+        end
+
+        on_stop
+        completed = handle_completed(completed)
+
+        if completed
+          on_complete
+          cleanup
+        else
+          reenqueue_iteration_job
+        end
+      end
+
+      private
+
+      def is_cancelled?
+        @_cancelled = Sidekiq.redis { |c| c.hget(iteration_key, "cancelled") }
+      end
+
+      def fetch_previous_iteration_state
+        state = Sidekiq.redis { |conn| conn.hgetall(iteration_key) }
+
+        unless state.empty?
+          @_executions = state["ex"].to_i
+          @_cursor = Sidekiq.load_json(state["c"])
+          @_runtime = state["rt"].to_f
+        end
+      end
+
+      STATE_FLUSH_INTERVAL = 5 # seconds
+      # we need to keep the state around as long as the job
+      # might be retrying
+      STATE_TTL = 30 * 24 * 60 * 60 # one month
+
+      def iterate_with_enumerator(enumerator, arguments)
+        if is_cancelled?
+          on_cancel
+          logger.info { "Job cancelled" }
+          return true
+        end
+
+        time_limit = Sidekiq.default_configuration[:timeout]
+        found_record = false
+        state_flushed_at = mono_now
+
+        enumerator.each do |object, cursor|
+          found_record = true
+          @_cursor = cursor
+          @current_object = object
+
+          interrupt_job = interrupted? || should_interrupt?
+          if mono_now - state_flushed_at >= STATE_FLUSH_INTERVAL || interrupt_job
+            _, _, cancelled = flush_state
+            state_flushed_at = mono_now
+            if cancelled
+              @_cancelled = true
+              on_cancel
+              logger.info { "Job cancelled" }
+              return true
+            end
+          end
+
+          return false if interrupt_job
+
+          verify_iteration_time(time_limit) do
+            around_iteration do
+              each_iteration(object, *arguments)
+            rescue Exception
+              flush_state
+              raise
+            end
+          end
+        end
+
+        logger.debug("Enumerator found nothing to iterate!") unless found_record
+        true
+      ensure
+        @_runtime += (mono_now - @_start_time)
+      end
+
+      def verify_iteration_time(time_limit)
+        start = mono_now
+        yield
+        finish = mono_now
+        total = finish - start
+        if total > time_limit
+          logger.warn { "Iteration took longer (%.2f) than Sidekiq's shutdown timeout (%d). This can lead to job processing problems during deploys" % [total, time_limit] }
+        end
+      end
+
+      def reenqueue_iteration_job
+        flush_state
+        logger.debug { "Interrupting job (cursor=#{@_cursor.inspect})" }
+
+        raise Interrupted
+      end
+
+      def assert_enumerator!(enum)
+        unless enum.is_a?(Enumerator)
+          raise ArgumentError, <<~MSG
+            #build_enumerator must return an Enumerator, but returned #{enum.class}.
+            Example:
+              def build_enumerator(params, cursor:)
+                active_record_records_enumerator(
+                  Shop.find(params["shop_id"]).products,
+                  cursor: cursor
+                )
+              end
+          MSG
+        end
+      end
+
+      def should_interrupt?
+        max_iteration_runtime = Sidekiq.default_configuration[:max_iteration_runtime]
+        max_iteration_runtime && (mono_now - @_start_time > max_iteration_runtime)
+      end
+
+      def flush_state
+        key = iteration_key
+        state = {
+          "ex" => @_executions,
+          "c" => Sidekiq.dump_json(@_cursor),
+          "rt" => @_runtime
+        }
+
+        Sidekiq.redis do |conn|
+          conn.multi do |pipe|
+            pipe.hset(key, state)
+            pipe.expire(key, STATE_TTL, "nx")
+            pipe.hget(key, "cancelled")
+          end
+        end
+      end
+
+      def cleanup
+        logger.debug {
+          format("Completed iteration. executions=%d runtime=%.3f", @_executions, @_runtime)
+        }
+        Sidekiq.redis { |conn| conn.unlink(iteration_key) }
+      end
+
+      def handle_completed(completed)
+        case completed
+        when nil, # someone aborted the job but wants to call the on_complete callback
+             true
+          true
+        when false
+          false
+        else
+          raise "Unexpected thrown value: #{completed.inspect}"
+        end
+      end
+
+      def mono_now
+        ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+      end
+    end
+  end
+end

data/lib/sidekiq/job.rb

@@ -69,7 +69,11 @@ module Sidekiq
         # In practice, any option is allowed.  This is the main mechanism to configure the
         # options for a specific job.
         def sidekiq_options(opts = {})
-          opts = opts.transform_keys(&:to_s) # stringify
+          # stringify 2 levels of keys
+          opts = opts.to_h do |k, v|
+            [k.to_s, (Hash === v) ? v.transform_keys(&:to_s) : v]
+          end
+
           self.sidekiq_options_hash = get_sidekiq_options.merge(opts)
         end
 
@@ -109,7 +113,7 @@ module Sidekiq
               m = "#{name}="
               undef_method(m) if method_defined?(m) || private_method_defined?(m)
             end
-            define_singleton_method("#{name}=") do |val|
+            define_singleton_method(:"#{name}=") do |val|
               singleton_class.class_eval do
                 ACCESSOR_MUTEX.synchronize do
                   undef_method(synchronized_getter) if method_defined?(synchronized_getter) || private_method_defined?(synchronized_getter)
@@ -155,6 +159,9 @@ module Sidekiq
 
     attr_accessor :jid
 
+    # This attribute is implementation-specific and not a public API
+    attr_accessor :_context
+
     def self.included(base)
       raise ArgumentError, "Sidekiq::Job cannot be included in an ActiveJob: #{base.name}" if base.ancestors.any? { |c| c.name == "ActiveJob::Base" }
 
@@ -166,6 +173,10 @@ module Sidekiq
       Sidekiq.logger
     end
 
+    def interrupted?
+      @_context&.stopping?
+    end
+
     # This helper class encapsulates the set options for `set`, e.g.
     #
     #     SomeJob.set(queue: 'foo').perform_async(....)
@@ -237,9 +248,9 @@ module Sidekiq
       end
       alias_method :perform_sync, :perform_inline
 
-      def perform_bulk(args, batch_size: 1_000)
+      def perform_bulk(args, **options)
         client = @klass.build_client
-        client.push_bulk(@opts.merge("class" => @klass, "args" => args, :batch_size => batch_size))
+        client.push_bulk(@opts.merge({"class" => @klass, "args" => args}, options))
       end
 
       # +interval+ must be a timestamp, numeric or something that acts
@@ -366,7 +377,7 @@ module Sidekiq
 
       def build_client # :nodoc:
         pool = Thread.current[:sidekiq_redis_pool] || get_sidekiq_options["pool"] || Sidekiq.default_configuration.redis_pool
-        client_class = get_sidekiq_options["client_class"] || Sidekiq::Client
+        client_class = Thread.current[:sidekiq_client_class] || get_sidekiq_options["client_class"] || Sidekiq::Client
         client_class.new(pool: pool)
       end
     end

data/lib/sidekiq/job_logger.rb

@@ -2,22 +2,23 @@
 
 module Sidekiq
   class JobLogger
-    def initialize(logger)
-      @logger = logger
+    def initialize(config)
+      @config = config
+      @logger = @config.logger
+      @skip = !!@config[:skip_default_job_logging]
     end
 
     def call(item, queue)
       start = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
-      @logger.info("start")
+      @logger.info { "start" } unless @skip
 
       yield
 
       Sidekiq::Context.add(:elapsed, elapsed(start))
-      @logger.info("done")
+      @logger.info { "done" } unless @skip
     rescue Exception
       Sidekiq::Context.add(:elapsed, elapsed(start))
-      @logger.info("fail")
-
+      @logger.info { "fail" } unless @skip
       raise
     end
 
@@ -25,16 +26,18 @@ module Sidekiq
       # If we're using a wrapper class, like ActiveJob, use the "wrapped"
       # attribute to expose the underlying thing.
       h = {
-        class: job_hash["display_class"] || job_hash["wrapped"] || job_hash["class"],
-        jid: job_hash["jid"]
+        jid: job_hash["jid"],
+        class: job_hash["wrapped"] || job_hash["class"]
       }
-      h[:bid] = job_hash["bid"] if job_hash.has_key?("bid")
-      h[:tags] = job_hash["tags"] if job_hash.has_key?("tags")
+
+      @config[:logged_job_attributes].each do |attr|
+        h[attr.to_sym] = job_hash[attr] if job_hash.has_key?(attr)
+      end
 
       Thread.current[:sidekiq_context] = h
       level = job_hash["log_level"]
-      if level && @logger.respond_to?(:log_at)
-        @logger.log_at(level, &block)
+      if level
+        @logger.with_level(level, &block)
       else
         yield
       end