sidekiq-tasks 1.1.0 → 1.2.0

data/lib/sidekiq/tasks/magic_comments.rb

@@ -0,0 +1,20 @@
+module Sidekiq
+  module Tasks
+    module MagicComments
+      autoload :Comment, "sidekiq/tasks/magic_comments/comment"
+      autoload :Comments, "sidekiq/tasks/magic_comments/comments"
+      autoload :Registry, "sidekiq/tasks/magic_comments/registry"
+      autoload :Scanner, "sidekiq/tasks/magic_comments/scanner"
+
+      module Handlers
+        autoload :Base, "sidekiq/tasks/magic_comments/handlers/base"
+        autoload :Enable, "sidekiq/tasks/magic_comments/handlers/enable"
+        autoload :Disable, "sidekiq/tasks/magic_comments/handlers/disable"
+        autoload :DryRun, "sidekiq/tasks/magic_comments/handlers/dry_run"
+        autoload :SidekiqOptions, "sidekiq/tasks/magic_comments/handlers/sidekiq_options"
+      end
+    end
+  end
+end
+
+require_relative "magic_comments/registration"

data/lib/sidekiq/tasks/run.rb

@@ -0,0 +1,161 @@
+module Sidekiq
+  module Tasks
+    # The context of a task execution, available from the task code through `Sidekiq::Tasks.current_run`.
+    class Run
+      THREAD_KEY = :sidekiq_tasks_run
+
+      attr_reader :report
+
+      # Outside of a Sidekiq::Tasks job (e.g. `rake` in a terminal), dry run is disabled
+      # and the report is printed on exit.
+      def self.current
+        Thread.current[THREAD_KEY] || standalone
+      end
+
+      def self.with(run)
+        previous_run = Thread.current[THREAD_KEY]
+        Thread.current[THREAD_KEY] = run
+        yield
+      ensure
+        Thread.current[THREAD_KEY] = previous_run
+      end
+
+      def self.standalone
+        @_standalone ||= new.tap do |run|
+          at_exit { puts JSON.pretty_generate(run.report.to_h) unless run.report.empty? }
+        end
+      end
+      private_class_method :standalone
+
+      def initialize(dry_run: false, report_max_bytes: Sidekiq::Tasks.config.report_max_bytes)
+        @dry_run = dry_run
+        @report = Report.new(max_bytes: report_max_bytes)
+      end
+
+      def dry_run?
+        @dry_run
+      end
+    end
+
+    # A structured execution report, stored in the task history and displayed in the web interface.
+    #
+    # The report size is bounded by `max_bytes`, see Sidekiq::Tasks::Storage::Base#report_max_bytes:
+    # - the serialized report never exceeds `max_bytes`, the last items of the longest lists are dropped
+    #   and the counts remain exact,
+    # - while the task runs, items stop being kept beyond MEMORY_FACTOR times that size, protecting worker memory,
+    # - strings are truncated to MAX_STRING_LENGTH, so that a single long message does not fill the whole report,
+    # - at most MAX_KEYS values, lists and item keys, guarding against keys built from the data.
+    class Report
+      MAX_KEYS = 50
+      MAX_STRING_LENGTH = 255
+      MEMORY_FACTOR = 4
+
+      def initialize(max_bytes: Sidekiq::Tasks.config.report_max_bytes)
+        @max_bytes = max_bytes
+        @kept_bytes = 0
+        @values = {}
+        @lists = {}
+        @truncated = false
+      end
+
+      def []=(key, value)
+        key = sanitize_key(key)
+        @values[key] = sanitize_scalar(value) if accept_key?(@values, key)
+      end
+
+      def increment(key, by = 1)
+        key = sanitize_key(key)
+        @values[key] = sanitize_scalar((@values[key] || 0) + by) if accept_key?(@values, key)
+      end
+
+      # Adds an item (a scalar or a flat Hash) to a list.
+      def add(key, item)
+        key = sanitize_key(key)
+        return unless accept_key?(@lists, key)
+
+        list = @lists[key] ||= {"items" => [], "count" => 0}
+        list["count"] += 1
+        keep_item(list, sanitize_item(item))
+      end
+
+      def empty?
+        @values.empty? && @lists.empty?
+      end
+
+      # @return [Hash] The report, serializable as JSON and at most `max_bytes` long.
+      def to_h
+        lists = @lists.transform_values { |list| list.merge("items" => list["items"].dup) }
+        report = {"values" => @values, "lists" => lists}
+        report["truncated"] = true if @truncated
+        return report if Sidekiq.dump_json(report).bytesize <= @max_bytes
+
+        fit(report.merge("truncated" => true))
+      end
+
+      private
+
+      # Falls back to an empty report when values alone exceed max_bytes.
+      def fit(report)
+        drop_items(report, Sidekiq.dump_json(report).bytesize)
+        return report if Sidekiq.dump_json(report).bytesize <= @max_bytes
+
+        {"values" => {}, "lists" => {}, "truncated" => true}
+      end
+
+      # Sizes are tracked incrementally so that the report is not serialized again for each dropped item.
+      def drop_items(report, size)
+        lists = report["lists"].values
+
+        while size > @max_bytes
+          longest_list = lists.max_by { |list| list["items"].size }
+          break if longest_list.nil? || longest_list["items"].empty?
+
+          size -= Sidekiq.dump_json(longest_list["items"].pop).bytesize + 1
+        end
+      end
+
+      # Items are kept with a margin over max_bytes so that to_h can still balance the lists,
+      # but never beyond it, so that the worker memory stays bounded.
+      def keep_item(list, item)
+        if @kept_bytes > @max_bytes * MEMORY_FACTOR
+          @truncated = true
+          return
+        end
+
+        @kept_bytes += Sidekiq.dump_json(item).bytesize
+        list["items"] << item
+      end
+
+      def accept_key?(hash, key)
+        return true if hash.key?(key) || hash.size < MAX_KEYS
+
+        @truncated = true
+        false
+      end
+
+      def sanitize_key(key)
+        truncate(key.to_s)
+      end
+
+      def sanitize_item(item)
+        return sanitize_scalar(item) unless item.is_a?(Hash)
+
+        item.first(MAX_KEYS).to_h { |key, value| [sanitize_key(key), sanitize_scalar(value)] }
+      end
+
+      def sanitize_scalar(value)
+        case value
+        when Float then value.finite? ? value : value.to_s
+        when Numeric, TrueClass, FalseClass, NilClass then value
+        else truncate(value.to_s)
+        end
+      end
+
+      def truncate(string)
+        return string if string.length <= MAX_STRING_LENGTH
+
+        "#{string[0...(MAX_STRING_LENGTH - 3)]}..."
+      end
+    end
+  end
+end

data/lib/sidekiq/tasks/sidekiq_options_validator.rb

@@ -0,0 +1,25 @@
+module Sidekiq
+  module Tasks
+    module SidekiqOptionsValidator
+      KEYS = {
+        queue: [String],
+        retry: [NilClass, TrueClass, FalseClass, Integer],
+        retry_for: [NilClass, Integer, Float],
+        dead: [NilClass, TrueClass, FalseClass],
+        backtrace: [NilClass, TrueClass, FalseClass, Integer],
+        pool: [NilClass, String],
+        tags: [NilClass, Array],
+      }.freeze
+
+      def self.validate!(options, allow_partial: false)
+        Sidekiq::Tasks::Validations.validate_class!(options, [Hash], "sidekiq_options")
+
+        KEYS.each do |key, classes|
+          next if allow_partial && !options.key?(key)
+
+          Sidekiq::Tasks::Validations.validate_class!(options[key], classes, key)
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/tasks/storage/base.rb

@@ -60,6 +60,28 @@ module Sidekiq
           raise NotImplementedError, "Storage must implement #store_execution_error"
         end
 
+        # Marks a history entry as a dry run.
+        # Optional: storages that don't implement it simply don't display the dry run badge.
+        #
+        # @param jid [String] The Sidekiq job ID.
+        def store_dry_run(_jid)
+        end
+
+        # Stores the execution report for a specific history entry.
+        # Optional: storages that don't implement it simply don't display reports.
+        #
+        # @param jid [String] The Sidekiq job ID.
+        # @param report [Hash] The report, see Sidekiq::Tasks::Report#to_h.
+        def store_execution_report(_jid, _report)
+        end
+
+        # The maximum size of a stored report, lowered by storages with their own memory constraints.
+        #
+        # @return [Integer]
+        def report_max_bytes
+          Sidekiq::Tasks.config.report_max_bytes
+        end
+
         protected
 
         def truncate_message(message, max_length)

data/lib/sidekiq/tasks/storage/redis.rb

@@ -3,6 +3,10 @@ module Sidekiq
     module Storage
       class Redis < Base
         JID_PREFIX = "task".freeze
+
+        # Maximum size of the reports stored in the history of a task, whatever the history_limit.
+        TASK_REPORTS_BUDGET = 10 * 1024 * 1024
+
         def jid_key
           "#{JID_PREFIX}:#{task_name}"
         end
@@ -11,22 +15,20 @@ module Sidekiq
           "#{jid_key}:history"
         end
 
+        def execution_key(jid)
+          "#{jid_key}:execution:#{jid}"
+        end
+
         def last_enqueue_at
           stored_time("last_enqueue_at")
         end
 
         def history
-          raw_entries = Sidekiq.redis { |conn| conn.lrange(history_key, 0, -1) }
+          entries = Sidekiq.redis { |conn| conn.lrange(history_key, 0, -1) }.map { |raw| Sidekiq.load_json(raw) }
 
-          return [] unless raw_entries
+          return [] if entries.empty?
 
-          raw_entries.map do |raw|
-            entry = Sidekiq.load_json(raw)
-            ["enqueued_at", "executed_at", "finished_at"].each do |key|
-              entry[key] = Time.at(entry[key]) if entry[key]
-            end
-            entry
-          end
+          entries.zip(executions(entries)).map { |entry, execution| parse_entry(entry.merge(execution)) }
         end
 
         def store_history(jid, task_args, time, user: nil)
@@ -34,6 +36,7 @@ module Sidekiq
             task_trace = {jid: jid, task_name: task_name, args: task_args, enqueued_at: time.to_f}
             task_trace[:user] = user if user
             conn.lpush(history_key, Sidekiq.dump_json(task_trace))
+            delete_executions(conn, conn.lrange(history_key, history_limit, -1))
             conn.ltrim(history_key, 0, history_limit - 1)
           end
         end
@@ -45,16 +48,24 @@ module Sidekiq
         end
 
         def store_execution(jid, time_key)
-          update_history_entry(jid) do |entry|
-            entry.merge(time_key => Time.now.to_f)
-          end
+          Sidekiq.redis { |conn| conn.hset(execution_key(jid), time_key, Time.now.to_f) }
         end
 
         def store_execution_error(jid, error)
-          update_history_entry(jid) do |entry|
-            error_message = truncate_message("#{error.class}: #{error.message}", ERROR_MESSAGE_MAX_LENGTH)
-            entry.merge("error" => error_message)
-          end
+          error_message = truncate_message("#{error.class}: #{error.message}", ERROR_MESSAGE_MAX_LENGTH)
+          Sidekiq.redis { |conn| conn.hset(execution_key(jid), "error", error_message) }
+        end
+
+        def store_dry_run(jid)
+          Sidekiq.redis { |conn| conn.hset(execution_key(jid), "dry_run", Sidekiq.dump_json(true)) }
+        end
+
+        def store_execution_report(jid, report)
+          Sidekiq.redis { |conn| conn.hset(execution_key(jid), "report", Sidekiq.dump_json(report)) }
+        end
+
+        def report_max_bytes
+          [super, TASK_REPORTS_BUDGET / history_limit].min
         end
 
         private
@@ -69,20 +80,30 @@ module Sidekiq
           [nil, ""].include?(timestamp) ? nil : Time.at(timestamp.to_f)
         end
 
-        def update_history_entry(jid)
+        def executions(entries)
           Sidekiq.redis do |conn|
-            entries = conn.lrange(history_key, 0, -1)
-
-            entries.each_with_index do |raw, index|
-              entry = Sidekiq.load_json(raw)
-              next unless entry["jid"] == jid
-
-              updated_entry = yield(entry)
-              conn.lset(history_key, index, Sidekiq.dump_json(updated_entry))
-              break
+            conn.pipelined do |pipeline|
+              entries.each { |entry| pipeline.hgetall(execution_key(entry["jid"])) }
             end
           end
         end
+
+        def parse_entry(entry)
+          ["enqueued_at", "executed_at", "finished_at"].each do |key|
+            entry[key] = Time.at(entry[key].to_f) if entry[key]
+          end
+          ["dry_run", "report"].each do |key|
+            entry[key] = Sidekiq.load_json(entry[key]) if entry[key]
+          end
+
+          entry
+        end
+
+        def delete_executions(conn, raw_entries)
+          return if raw_entries.empty?
+
+          conn.del(*raw_entries.map { |raw| execution_key(Sidekiq.load_json(raw)["jid"]) })
+        end
       end
     end
   end

data/lib/sidekiq/tasks/strategies/base.rb

@@ -51,9 +51,12 @@ module Sidekiq
         #
         # @param name [String] The name of the task to enqueue.
         # @param params [Hash] Parameters to pass to the task.
+        # @param sidekiq_options [Hash] Per-task Sidekiq options that override class-level defaults.
+        # @param dry_run [Boolean] Whether the task must be executed in dry run mode.
         # @return [String] The JID of the sidekiq job that will execute the task.
-        def enqueue_task(name, params = {})
-          Sidekiq::Tasks::Job.perform_async(name, params.to_json)
+        def enqueue_task(name, params = {}, sidekiq_options: {}, dry_run: false)
+          job = sidekiq_options.empty? ? Sidekiq::Tasks::Job : Sidekiq::Tasks::Job.set(sidekiq_options)
+          job.perform_async(name, params.to_json, dry_run)
         end
 
         # Returns all the tasks that should be executed.

data/lib/sidekiq/tasks/strategies/rake_task.rb

@@ -2,6 +2,11 @@ module Sidekiq
   module Tasks
     module Strategies
       class RakeTask < Base
+        def initialize(rules: [], scanner: Sidekiq::Tasks::MagicComments::Scanner.new)
+          super(rules: rules)
+          @scanner = scanner
+        end
+
         def load_tasks
           Rake::TaskManager.record_task_metadata = true
           Rake.application.load_rakefile
@@ -12,14 +17,33 @@ module Sidekiq
           Sidekiq::Tasks::TaskMetadata.new(
             name: task.name,
             desc: task.full_comment,
-            file: task.locations.first.split(":").first,
-            args: task.arg_names
+            file: task_file(task),
+            args: task.arg_names,
+            **magic_comment_options(task)
           )
+        rescue Sidekiq::Tasks::ArgumentError => e
+          Sidekiq::Tasks::TaskMetadata.new(name: task.name, file: task_file(task), error: e.message)
         end
 
         def execute_task(name, args = nil)
           Rake::Task[name].execute(args)
         end
+
+        private
+
+        def task_file(task)
+          task.locations.first&.split(":")&.first
+        end
+
+        def magic_comment_options(task)
+          comments = @scanner.scan(task)
+          handlers = Sidekiq::Tasks::MagicComments::Handlers
+
+          {
+            sidekiq_options: comments.fetch(handlers::SidekiqOptions.name_token, default: {}),
+            dry_run: comments.any?(handlers::DryRun.name_token),
+          }
+        end
       end
     end
   end

data/lib/sidekiq/tasks/strategies/rules/disable_with_comment.rb

@@ -9,8 +9,8 @@ module Sidekiq
 
           protected
 
-          def magic_comment_regex
-            /sidekiq-tasks:disable/
+          def magic_comment_token
+            Sidekiq::Tasks::MagicComments::Handlers::Disable.name_token
           end
         end
       end

data/lib/sidekiq/tasks/strategies/rules/enable_with_comment.rb

@@ -3,67 +3,19 @@ module Sidekiq
     module Strategies
       module Rules
         class EnableWithComment < Base
-          def respected?(task)
-            file, line_number = task.locations.first.split(":")
-            line_number = line_number.to_i
-
-            lines = read_file_lines(file)
-
-            return false if lines.nil?
-
-            return true if task_has_magic_comment?(lines, line_number)
-            return true if desc_has_magic_comment?(task.full_comment, lines, line_number)
-
-            namespace_line_index = find_namespace_line_index(lines, task)
-            return false unless namespace_line_index
-
-            namespace_has_magic_comment?(lines, namespace_line_index)
+          def initialize(scanner: Sidekiq::Tasks::MagicComments::Scanner.new)
+            super()
+            @scanner = scanner
           end
 
-          protected
-
-          def magic_comment_regex
-            /sidekiq-tasks:enable/
-          end
-
-          private
-
-          def read_file_lines(file)
-            File.read(file).split("\n")
-          rescue Errno::ENOENT
-            raise ArgumentError, "File '#{file}' not found"
-          end
-
-          def desc_has_magic_comment?(desc, lines, task_line)
-            return false unless desc
-
-            desc_line_index = lines[0...task_line].rindex { |line| line.strip.match?(/^desc\b/) }
-
-            return false unless desc_line_index
-
-            comment_line = lines[desc_line_index - 1]&.strip
-            !!(comment_line =~ magic_comment_regex)
-          end
-
-          def task_has_magic_comment?(lines, task_line)
-            context_range = (task_line - 3..task_line).to_a.select { |i| i >= 0 }
-            context_range.reverse.any? do |i|
-              lines[i]&.strip&.match?(magic_comment_regex)
-            end
-          end
-
-          def find_namespace_line_index(lines, task)
-            namespace = namespace_name(task)
-            lines.find_index { |line| line.strip.match?(/^namespace\s+:#{Regexp.escape(namespace)}/) }
+          def respected?(task)
+            @scanner.scan(task).any?(magic_comment_token)
           end
 
-          def namespace_has_magic_comment?(lines, namespace_line_index)
-            comment_line = lines[namespace_line_index - 1]&.strip
-            comment_line&.match?(magic_comment_regex)
-          end
+          protected
 
-          def namespace_name(task)
-            task.name.split(":").first
+          def magic_comment_token
+            Sidekiq::Tasks::MagicComments::Handlers::Enable.name_token
           end
         end
       end

data/lib/sidekiq/tasks/task.rb

@@ -6,8 +6,9 @@ module Sidekiq
       extend Forwardable
       include Sidekiq::Tasks::Validations
 
-      def_delegators :metadata, :name, :desc, :file, :args
+      def_delegators :metadata, :name, :desc, :file, :args, :sidekiq_options, :error, :error?
       def_delegators :storage, :last_enqueue_at, :history
+      def_delegator :metadata, :dry_run?, :dry_run_supported?
 
       attr_reader :metadata, :strategy
 
@@ -22,29 +23,45 @@ module Sidekiq
         validate_class!(strategy, [Sidekiq::Tasks::Strategies::Base], "strategy")
       end
 
-      def enqueue(params = {}, user: nil)
-        jid = strategy.enqueue_task(name, params)
+      def enqueue(params = {}, user: nil, dry_run: false)
+        raise Sidekiq::Tasks::ArgumentError, "cannot enqueue broken task '#{name}': #{error}" if error?
+
+        enqueue_options = {}
+        enqueue_options[:sidekiq_options] = sidekiq_options if sidekiq_options.any?
+        enqueue_options[:dry_run] = true if dry_run
+
+        jid = strategy.enqueue_task(name, params, **enqueue_options)
 
         storage.store_enqueue(jid, params, user: user)
+        storage.store_dry_run(jid) if dry_run
       end
 
-      def execute(params = {}, jid: nil)
+      def execute(params = {}, jid: nil, dry_run: false)
+        run = Sidekiq::Tasks::Run.new(dry_run: dry_run, report_max_bytes: storage.report_max_bytes)
         storage.store_execution(jid, "executed_at")
 
         begin
-          strategy.execute_task(name, params)
+          Sidekiq::Tasks::Run.with(run) { strategy.execute_task(name, params) }
         rescue StandardError, SystemExit => e
-          storage.store_execution(jid, "finished_at")
+          finish_execution(jid, run)
           storage.store_execution_error(jid, e)
           raise
         end
 
-        storage.store_execution(jid, "finished_at")
+        finish_execution(jid, run)
       end
 
       def storage
         @_storage ||= Sidekiq::Tasks.config.storage.new(name, history_limit: Sidekiq::Tasks.config.history_limit)
       end
+
+      private
+
+      # The report is stored before finished_at so that it is displayed as soon as the execution appears finished.
+      def finish_execution(jid, run)
+        storage.store_execution_report(jid, run.report.to_h) unless run.report.empty?
+        storage.store_execution(jid, "finished_at")
+      end
     end
   end
 end

data/lib/sidekiq/tasks/task_metadata.rb

@@ -3,24 +3,40 @@ module Sidekiq
     class TaskMetadata
       include Sidekiq::Tasks::Validations
 
-      attr_reader :name, :desc, :file, :args
+      attr_reader :name, :desc, :file, :args, :sidekiq_options, :error, :dry_run
 
-      def initialize(name:, file:, desc: "", args: [])
+      def initialize(name:, file:, desc: "", args: [], sidekiq_options: {}, error: nil, dry_run: false)
         @name = name
         @file = file
         @desc = desc
         @args = args
+        @sidekiq_options = sidekiq_options
+        @error = error
+        @dry_run = dry_run
 
         validate_params!
       end
 
+      def error?
+        !@error.nil?
+      end
+
+      def dry_run?
+        @dry_run
+      end
+
       private
 
       def validate_params!
         validate_class!(name, [String, Symbol], "name")
         validate_class!(file, [String, NilClass], "file")
+        validate_class!(error, [String, NilClass], "error")
+        return if error?
+
         validate_class!(desc, [String, NilClass], "desc")
         validate_array_classes!(args, [String, Symbol], "args")
+        validate_class!(sidekiq_options, [Hash], "sidekiq_options")
+        validate_class!(dry_run, [TrueClass, FalseClass], "dry_run")
       end
     end
   end

data/lib/sidekiq/tasks/version.rb

@@ -2,6 +2,6 @@
 
 module Sidekiq
   module Tasks
-    VERSION = "1.1.0"
+    VERSION = "1.2.0"
   end
 end

data/lib/sidekiq/tasks/web/extension.rb

@@ -62,7 +62,8 @@ module Sidekiq
             args = Sidekiq::Tasks::Web::Params.new(task, fetch_param("args")).permit!
 
             current_user = Sidekiq::Tasks.config.current_user&.call(env)
-            task.enqueue(args, user: current_user)
+            dry_run = task.dry_run_supported? && fetch_param("dry_run") == "1"
+            task.enqueue(args, user: current_user, dry_run: dry_run)
 
             redirect(task_url(root_path, task))
           rescue Sidekiq::Tasks::ArgumentError => e

data/lib/sidekiq/tasks/web/helpers/tag_helper.rb

@@ -6,7 +6,7 @@ module Sidekiq
           extend self
 
           def build_tag(tag, content = nil, **attributes, &block)
-            attr_string = attributes.map { |key, value| "#{key}=\"#{value}\"" }.join(" ")
+            attr_string = attributes.map { |key, value| "#{key}=\"#{ERB::Util.html_escape(value)}\"" }.join(" ")
             attr_string = " #{attr_string}" unless attr_string.empty?
 
             "<#{tag}#{attr_string}>#{block&.call || content}</#{tag}>"

data/lib/sidekiq/tasks/web/helpers/task_helper.rb

@@ -36,6 +36,13 @@ module Sidekiq
             end
           end
 
+          # Returns the table columns for a report list: the keys of Hash items, or nil for scalar items.
+          def report_list_columns(items)
+            return nil unless items.any? && items.all?(Hash)
+
+            items.flat_map(&:keys).uniq
+          end
+
           def format_task_duration(start_time, end_time)
             return "-" unless start_time && end_time
 
@@ -48,6 +55,18 @@ module Sidekiq
               "#{[duration_in_milliseconds, 1].max}ms"
             end
           end
+
+          def format_task_sidekiq_options(task)
+            task.sidekiq_options.map { |key, value| "#{key}: #{format_sidekiq_option_value(value)}" }.join(", ")
+          end
+
+          def format_sidekiq_option_value(value)
+            case value
+            when Array then "[#{value.join(", ")}]"
+            when nil then "-"
+            else value.to_s
+            end
+          end
         end
       end
     end