sidekiq-buffered 0.1.0

data/lib/sidekiq/buffered/execution.rb

@@ -0,0 +1,143 @@
+module Sidekiq
+  module Buffered
+    # One worker's attempt at a batch: the lease token it holds, the record it
+    # claimed, and whether perform completed. ServerMiddleware claims and
+    # acknowledges around the server chain; without it, the processor does both.
+    class Execution
+      CURRENT_KEY = :sidekiq_buffered_execution
+
+      attr_reader :batch_id, :jid, :token, :record
+
+      class << self
+        # Sidekiq's default retry delay, and the most jitter it adds on top.
+        def default_retry_delay(count)
+          (count**4) + 15
+        end
+
+        def max_retry_jitter(count)
+          10 * (count + 1)
+        end
+
+        # Sidekiq schedules a failed processor's retry on the job's thread after
+        # the server middleware has unwound, so a claimed execution stays
+        # current on its thread until the middleware or a retry callback
+        # finishes it.
+        def current
+          Thread.current[CURRENT_KEY]
+        end
+
+        def current=(execution)
+          Thread.current[CURRENT_KEY] = execution
+        end
+
+        # The current execution, if it is the one running processor payload +job+.
+        def for(job)
+          execution = current
+          execution if execution && execution.jid == job["jid"] && execution.batch_id == job["args"]&.first
+        end
+
+        # Sidekiq's retry callbacks, reached through RetryCallbacks on the job class.
+        def retry_in(job, count, exception, declared)
+          execution = self.for(job)
+          return execution.retry_in(count, exception, job, declared) if execution
+
+          Sidekiq.logger.debug { "sidekiq-buffered retry callback outside its execution: jid=#{job['jid']}" }
+          declared_callback(declared, count, exception, job)
+        end
+
+        def retries_exhausted(job, exception, declared)
+          execution = self.for(job)
+          return execution.retries_exhausted(job, exception, declared) if execution
+
+          Sidekiq.logger.debug { "sidekiq-buffered exhausted callback outside its execution: jid=#{job['jid']}" }
+          declared_callback(declared, job, exception)
+        end
+
+        # A callback the job class declared itself; its errors never change delivery.
+        def declared_callback(callback, *)
+          callback&.call(*)
+        rescue StandardError => e
+          Sidekiq.logger.warn("sidekiq-buffered retry callback: #{e.class}: #{e.message}")
+          nil
+        end
+      end
+
+      def initialize(batch_id, jid:, store: Buffered.store)
+        @batch_id = batch_id
+        @jid = jid
+        @store = store
+        @token = SecureRandom.hex(12)
+        @record = nil
+        @performed = false
+      end
+
+      # A successful claim makes this the thread's current execution.
+      def claim
+        @record = @store.claim(batch_id, token: jid, lease_token: token)
+        self.class.current = self if @record
+        @record
+      end
+
+      def performed?
+        @performed
+      end
+
+      def job_class
+        Object.const_get(record.job_class)
+      end
+
+      def run
+        BatchContext.with(record) do
+          if record.active_job?
+            ActiveJobExecution.perform(record, jid)
+          else
+            worker = job_class.new
+            worker.jid = jid if worker.respond_to?(:jid=)
+            worker.perform(record.arguments)
+          end
+        end
+        @performed = true
+      end
+
+      def ack
+        @store.ack(batch_id, token:)
+      end
+
+      def release
+        @store.release(batch_id, token:)
+      end
+
+      def exhaust(exception = nil)
+        @store.exhaust(batch_id, token:, exception:)
+      end
+
+      # Sidekiq needs nothing further from this execution.
+      def finish
+        self.class.current = nil if self.class.current.equal?(self)
+      end
+
+      # Sidekiq is about to retry the processor. The job class's own retry_in
+      # callback picks the delay, and the batch stays protected for that long.
+      def retry_in(count, exception, job, declared = nil)
+        delay = self.class.declared_callback(declared, count, exception, job)
+        case delay
+        when :discard then ack
+        when :kill then return delay # Sidekiq calls retries_exhausted next.
+        else
+          explicit = delay.respond_to?(:to_i) && delay.to_i.positive?
+          seconds = explicit ? delay.to_i : self.class.default_retry_delay(count)
+          @store.mark_retrying(batch_id, token:, exception:, delay: seconds + self.class.max_retry_jitter(count))
+        end
+        finish
+        delay
+      end
+
+      def retries_exhausted(job, exception, declared = nil)
+        exhaust(exception)
+        self.class.declared_callback(declared, job, exception)
+      ensure
+        finish
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/heartbeat.rb

@@ -0,0 +1,96 @@
+module Sidekiq
+  module Buffered
+    # One thread per process extends the leases of batches that are currently
+    # executing, so lease_timeout bounds crash recovery instead of perform time.
+    class Heartbeat < PeriodicThread
+      THREAD_NAME = "sidekiq-buffered-heartbeat".freeze
+      MIN_INTERVAL = 0.1
+      IDLE_INTERVAL = 1.0
+      Lease = Struct.new(:batch_id, :token, :timeout, :redis_pool)
+
+      @mutex = Mutex.new
+
+      class << self
+        def instance
+          @mutex.synchronize { @instance ||= new }
+        end
+
+        def track(batch_id, token:, timeout:, redis_pool:)
+          heartbeat = instance
+          heartbeat.track(batch_id, token:, timeout:, redis_pool:)
+          heartbeat.start
+        end
+
+        def untrack(batch_id)
+          @instance&.untrack(batch_id)
+        end
+
+        def tracked_ids
+          @instance&.tracked_ids || []
+        end
+
+        def start
+          instance.start
+        end
+
+        def stop
+          heartbeat = @mutex.synchronize do
+            current = @instance
+            @instance = nil
+            current
+          end
+          heartbeat&.stop
+        end
+      end
+
+      def initialize
+        super
+        @leases = {}
+        @mutex = Mutex.new
+      end
+
+      def track(batch_id, token:, timeout:, redis_pool:)
+        @mutex.synchronize { @leases[batch_id] = Lease.new(batch_id, token, Integer(timeout), redis_pool) }
+        wakeup
+      end
+
+      def untrack(batch_id)
+        @mutex.synchronize { @leases.delete(batch_id) }
+      end
+
+      def tracked_ids
+        @mutex.synchronize { @leases.keys }
+      end
+
+      def tick
+        leases = @mutex.synchronize { @leases.values }
+        return if leases.empty?
+
+        leases.group_by(&:redis_pool).each do |redis_pool, group|
+          Store.new(redis_pool).extend_leases(group).each { |lease| lost(lease) }
+        end
+      rescue StandardError => e
+        Sidekiq.logger.error("sidekiq-buffered heartbeat: #{e.class}: #{e.message}")
+      end
+
+      def interval
+        timeouts = @mutex.synchronize { @leases.values.map(&:timeout) }
+        return IDLE_INTERVAL if timeouts.empty?
+
+        [timeouts.min / 3.0, MIN_INTERVAL].max
+      end
+
+      private
+
+      # A lease that finished normally between the snapshot and the extension is
+      # already untracked; only a lease we still hold has really been lost.
+      def lost(lease)
+        still_tracked = @mutex.synchronize { @leases.delete(lease.batch_id) if @leases[lease.batch_id].equal?(lease) }
+        return unless still_tracked
+
+        Sidekiq.logger.warn("sidekiq-buffered lease lost while running: batch=#{lease.batch_id}")
+        Buffered.notify(:lease_lost, batch_id: lease.batch_id)
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/inline.rb

@@ -0,0 +1,17 @@
+module Sidekiq
+  module Buffered
+    # Marks Setter#perform_inline so the client middleware lets the payload
+    # through to synchronous execution instead of buffering it.
+    module Inline
+      KEY = :sidekiq_buffered_inline
+
+      def self.active?
+        Thread.current[KEY] == true
+      end
+
+      def self.with(&)
+        ThreadLocal.with(KEY, true, &)
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/input.rb

@@ -0,0 +1,122 @@
+module Sidekiq
+  module Buffered
+    # One buffered input as it leaves Sidekiq's client middleware chain: a job
+    # payload whose single argument is a Hash. Knows which buffer the input
+    # belongs to and how it is stored inside a batch.
+    class Input
+      # Fields Sidekiq or this gem own. Anything else on a payload is metadata
+      # other client middleware attached, and it travels with the input.
+      PAYLOAD_FIELDS = %w[class args jid queue retry retry_for retry_queue dead backtrace at enqueued_at
+                          created_at active_job wrapped pool client_class batch_size buffered_context
+                          buffered_active_job_context].freeze
+      # Sidekiq options that shape execution and therefore separate buffers.
+      EXECUTION_FIELDS = %w[queue retry retry_for retry_queue dead backtrace].freeze
+
+      # Inputs with the same identity share a buffer, and so a batch.
+      Identity = Data.define(:job_class, :group, :config, :context)
+
+      class << self
+        # Rewrites +payload+ in place so its argument list is one string-keyed
+        # Hash that carries the grouping key.
+        def normalize!(payload)
+          payload["args"] = normalize_args(payload["class"], payload["args"])
+          payload
+        end
+
+        def normalize_args(job_class, args)
+          input = new("class" => job_class, "args" => stringify_keys(args)).validate!
+          [input.arguments]
+        end
+
+        def encode_group(value)
+          Sidekiq.dump_json(value)
+        end
+
+        def decode_group(raw)
+          return if raw.nil? || raw.empty?
+
+          Sidekiq.load_json(raw)
+        end
+
+        private
+
+        def stringify_keys(value)
+          case value
+          when Hash then value.to_h { |key, nested| [key.to_s, stringify_keys(nested)] }
+          when Array then value.map { |entry| stringify_keys(entry) }
+          else value
+          end
+        end
+      end
+
+      attr_reader :payload
+
+      def initialize(payload)
+        @payload = payload
+      end
+
+      def job_class
+        @payload["class"].to_s
+      end
+
+      def jid
+        @payload["jid"]
+      end
+
+      def config
+        Buffered.config_for(job_class)
+      end
+
+      def arguments
+        @arguments ||= begin
+          args = @payload["args"]
+          unless args.is_a?(Array) && args.size == 1 && args.first.is_a?(Hash)
+            raise ArgumentError, "#{job_class} expects a single Hash argument"
+          end
+
+          args.first
+        end
+      end
+
+      # Raises unless the argument list is one Hash that carries the grouping key.
+      def validate!
+        key = config.group_by
+        return self if key.nil? || arguments.key?(key)
+
+        raise ArgumentError, "#{job_class} is missing grouping key #{key.inspect}"
+      end
+
+      # The encoded grouping value, or "" for ungrouped jobs.
+      def group
+        @group ||= begin
+          key = config.group_by
+          key ? self.class.encode_group(validate!.arguments.fetch(key)) : ""
+        end
+      end
+
+      # Middleware metadata plus the execution options, kept with the input so
+      # the batch runs under the same context its inputs were enqueued with.
+      def context
+        @context ||= begin
+          context = @payload.except(*PAYLOAD_FIELDS).merge(@payload.slice(*EXECUTION_FIELDS))
+          if (data = @payload["active_job"])
+            context["buffered_active_job_context"] = data.values_at("locale", "timezone")
+          end
+          context
+        end
+      end
+
+      def identity
+        Identity.new(job_class:, group:, config:, context:)
+      end
+
+      # The JSON stored in the buffer. Redis stamps buffered_at into it on admission.
+      def dump
+        item = { "jid" => jid, "args" => arguments }
+        item["context"] = context unless context.empty?
+        item["active_job"] = @payload["active_job"] if @payload.key?("active_job")
+        Sidekiq.dump_json(item)
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/periodic_thread.rb

@@ -0,0 +1,45 @@
+module Sidekiq
+  module Buffered
+    # A named background thread that calls +tick+, sleeps +interval+ seconds,
+    # and repeats until stopped. Subclasses define both, plus THREAD_NAME.
+    class PeriodicThread
+      def initialize
+        @stop = false
+        @thread = nil
+      end
+
+      def start
+        return if @thread&.alive?
+
+        @stop = false
+        @thread = Thread.new do
+          Thread.current.name = self.class::THREAD_NAME
+          run
+        end
+      end
+
+      def stop
+        @stop = true
+        wakeup
+      ensure
+        @thread&.join(1)
+        @thread = nil
+      end
+
+      private
+
+      def run
+        until @stop
+          tick
+          sleep(interval) unless @stop
+        end
+      end
+
+      def wakeup
+        @thread&.wakeup
+      rescue ThreadError
+        nil
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/poller.rb

@@ -0,0 +1,81 @@
+require "sidekiq/api"
+
+module Sidekiq
+  module Buffered
+    # Flushes due buffers, recovers stalled batches, and cleans up failures.
+    # Idle sleeps stretch with the number of Sidekiq processes, like Sidekiq's
+    # own scheduler, and shrink to the next known deadline.
+    class Poller < PeriodicThread
+      THREAD_NAME = "sidekiq-buffered-poller".freeze
+      MAX_ERROR_INTERVAL = 30
+      MIN_SLEEP = 0.01
+      FLUSH_LIMIT = 500
+      PROCESS_COUNT_TTL = 30
+
+      class << self
+        def start
+          @instance ||= new
+          @instance.start
+        end
+
+        def stop
+          @instance&.stop
+          @instance = nil
+        end
+      end
+
+      def initialize(store = Buffered.store)
+        super()
+        @store = store
+        @failures = 0
+        @busy = false
+        @now = nil
+        @next_at = nil
+        @process_count = nil
+        @process_counted_at = nil
+      end
+
+      def tick
+        due = @store.due(limit: FLUSH_LIMIT)
+        @store.flush_members(due.buffers, now: due.now)
+        @store.recover(due.outstanding, now: due.now)
+        @store.cleanup_failed(now: due.now)
+        @busy = due.buffers.size >= FLUSH_LIMIT || due.outstanding.size >= FLUSH_LIMIT
+        @next_at = due.next_at
+        @now = due.now
+        @failures = 0
+      rescue StandardError => e
+        @failures += 1
+        Sidekiq.logger.error("sidekiq-buffered poller: #{e.class}: #{e.message}")
+      end
+
+      # Errors back off exponentially, a full page polls again at once, and
+      # otherwise the poller sleeps until the next known deadline.
+      def interval
+        return [Buffered.poll_interval * (2**[@failures, 20].min), MAX_ERROR_INTERVAL].min if @failures.positive?
+        return MIN_SLEEP if @busy
+
+        idle = scaled_interval
+        idle = [idle, @next_at - @now + (rand * Buffered.poll_interval)].min if @next_at && @now
+        [idle, MIN_SLEEP].max
+      end
+
+      private
+
+      # Average one poll per poll_interval across the cluster, with jitter.
+      def scaled_interval
+        Buffered.poll_interval * process_count * (0.5 + rand)
+      end
+
+      def process_count
+        at = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+        return @process_count if @process_count && at - @process_counted_at < PROCESS_COUNT_TTL
+
+        @process_counted_at = at
+        @process_count = [Sidekiq::ProcessSet.new(false).size.to_i, 1].max
+      rescue StandardError
+        @process_count = 1
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/processor.rb

@@ -0,0 +1,42 @@
+module Sidekiq
+  module Buffered
+    # The Sidekiq job that runs one batch. Its only argument is the batch ID.
+    class Processor
+      include Sidekiq::Job
+
+      # Set by ServerMiddleware once it has claimed the batch.
+      attr_accessor :execution
+
+      def perform(batch_id)
+        if execution
+          execution.run # ServerMiddleware acknowledges after the rest of the chain.
+        else
+          # No middleware in front (perform_inline, tests): own the whole lifecycle.
+          self.execution = Execution.new(batch_id, jid:)
+          return unless execution.claim
+
+          execution.run
+          execution.ack
+          execution.finish
+        end
+      end
+
+      # Sidekiq 7.0 reads retry callbacks from the job instance; later versions
+      # read them from the payload's wrapped class. Both reach the job class's
+      # RetryCallbacks.
+      def sidekiq_retry_in_block
+        ->(count, exception, job) { retry_callbacks&.sidekiq_retry_in_block&.call(count, exception, job) }
+      end
+
+      def sidekiq_retries_exhausted_block
+        ->(job, exception) { retry_callbacks&.sidekiq_retries_exhausted_block&.call(job, exception) }
+      end
+
+      private
+
+      def retry_callbacks
+        execution&.job_class
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/registry.rb

@@ -0,0 +1,38 @@
+module Sidekiq
+  module Buffered
+    # The +buffered+ declarations, by job class name. Lookups accept a class or
+    # its name, since Redis and Sidekiq payloads carry names, and a subclass
+    # inherits the nearest ancestor's declaration.
+    class Registry
+      def initialize
+        @configs = {}
+        @mutex = Mutex.new
+      end
+
+      def register(klass, config)
+        @mutex.synchronize { @configs[klass.name] = config }
+      end
+
+      def registered?(klass)
+        !self[klass].nil?
+      end
+
+      def fetch(klass)
+        self[klass] || raise(KeyError, "No buffered configuration for #{klass}")
+      end
+
+      def [](klass)
+        name = klass.is_a?(Class) ? klass.name : klass.to_s
+        @configs[name] || ancestor_config(klass.is_a?(Class) ? klass : Object.const_get(name))
+      rescue NameError
+        nil
+      end
+
+      private
+
+      def ancestor_config(klass)
+        klass.superclass && self[klass.superclass]
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/retry_callbacks.rb

@@ -0,0 +1,32 @@
+module Sidekiq
+  module Buffered
+    # Sidekiq resolves a processor payload's retry callbacks through its
+    # +wrapped+ class, which is the buffered job class. Prepended to that class's
+    # singleton by +buffered+, this answers those lookups with callbacks that
+    # record the batch's retry state and then run whatever the class declared
+    # with +sidekiq_retry_in+ and +sidekiq_retries_exhausted+.
+    module RetryCallbacks
+      RetryIn = Data.define(:declared) do
+        def call(count, exception, job)
+          Execution.retry_in(job, count, exception, declared)
+        end
+      end
+
+      RetriesExhausted = Data.define(:declared) do
+        def call(job, exception)
+          Execution.retries_exhausted(job, exception, declared)
+        end
+      end
+
+      def sidekiq_retry_in_block
+        declared = defined?(super) ? super : nil
+        declared.is_a?(RetryIn) ? declared : RetryIn.new(declared)
+      end
+
+      def sidekiq_retries_exhausted_block
+        declared = defined?(super) ? super : nil
+        declared.is_a?(RetriesExhausted) ? declared : RetriesExhausted.new(declared)
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/scheduled_input.rb

@@ -0,0 +1,21 @@
+module Sidekiq
+  module Buffered
+    # Native Sidekiq scheduling and retries protect inputs waiting for admission.
+    class ScheduledInput
+      include Sidekiq::Job
+
+      # Client-only options never belong in a persisted payload.
+      CLIENT_FIELDS = %w[client_class pool].freeze
+
+      def self.wrap(payload)
+        input = Sidekiq.load_json(Sidekiq.dump_json(payload.except("at", *CLIENT_FIELDS)))
+        payload.except(*CLIENT_FIELDS).merge("class" => name, "args" => [input])
+      end
+
+      def perform(payload)
+        Object.const_get(payload.fetch("class")) # Load the buffered declaration before admission.
+        Buffered.store.append([payload])
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/scope.rb

@@ -0,0 +1,17 @@
+module Sidekiq
+  module Buffered
+    # Which job class and group an inspection or flush applies to. A class or
+    # its name narrows by class; a group value narrows by group.
+    class Scope
+      def initialize(job_class: nil, group: Buffered::ANY_GROUP)
+        @job_class = job_class.is_a?(Class) ? job_class.name : job_class&.to_s
+        @group = group
+      end
+
+      # +job_class+ is a stored class name and +group+ a decoded grouping value.
+      def matches?(job_class, group)
+        (@job_class.nil? || job_class == @job_class) && (@group.equal?(Buffered::ANY_GROUP) || group == @group)
+      end
+    end
+  end
+end