sidekiq-buffered 0.1.0

data/lib/sidekiq/buffered/active_job_execution.rb

@@ -0,0 +1,49 @@
+module Sidekiq
+  module Buffered
+    class RetryEnqueueError < StandardError; end
+
+    module ActiveJobExecution
+      RETRY_EVENT = "enqueue_retry.active_job".freeze
+      HANDOFF_KEY = :sidekiq_buffered_active_job_handoffs
+
+      class << self
+        # Executes the batch as one Active Job. Active Job announces every
+        # retry_job through its public enqueue_retry instrumentation, and the
+        # buffered adapter records the retries Sidekiq accepted. A retry that
+        # was announced but never accepted (aborted callback, rejected push)
+        # must keep the source batch.
+        def perform(record, execution_jid)
+          data = record.items.first.active_job.merge(
+            "job_id" => record.id,
+            "provider_job_id" => execution_jid,
+            "arguments" => [record.arguments],
+          )
+          retried = false
+          observer = lambda do |_name, _start, _finish, _id, payload|
+            retried = true if payload[:job]&.job_id == record.id
+          end
+          handoffs = tracking_handoffs do
+            ::ActiveSupport::Notifications.subscribed(observer, RETRY_EVENT) { ::ActiveJob::Base.execute(data) }
+          end
+          return unless retried && !handoffs.include?(record.id)
+
+          raise RetryEnqueueError, "Buffered batch retry was not enqueued"
+        end
+
+        # Called by the queue adapter once Sidekiq accepted a retry payload.
+        def record_handoff(job_id)
+          Thread.current[HANDOFF_KEY]&.add(job_id)
+        end
+
+        private
+
+        def tracking_handoffs
+          ThreadLocal.with(HANDOFF_KEY, Set.new) do |handoffs|
+            yield
+            handoffs
+          end
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/active_job_retry.rb

@@ -0,0 +1,12 @@
+module Sidekiq
+  module Buffered
+    # Keep whole-batch Active Job retries independent of Rails engine initialization.
+    class ActiveJobRetry
+      include Sidekiq::Job
+
+      def perform(data)
+        ::ActiveJob::Base.execute(data.merge("provider_job_id" => jid))
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/batch.rb

@@ -0,0 +1,70 @@
+require "forwardable"
+
+module Sidekiq
+  module Buffered
+    # Inspection and manual intervention for one stored batch.
+    class Batch
+      extend Forwardable
+
+      class << self
+        def find(id)
+          record = Buffered.store.fetch(id)
+          new(record) if record
+        end
+
+        def failed(limit: Store::PAGE_SIZE, offset: 0, job_class: nil, group: Buffered::ANY_GROUP)
+          raise ArgumentError, "limit must be positive" unless limit.is_a?(Integer) && limit.positive?
+          raise ArgumentError, "offset must be non-negative" unless offset.is_a?(Integer) && offset >= 0
+
+          scope = Scope.new(job_class:, group:)
+          matching = each_failed.lazy.select { |batch| scope.matches?(batch.job_class, batch.group) }
+          matching.drop(offset).take(limit).force
+        end
+
+        def each_failed(&block)
+          return enum_for(__method__) unless block_given?
+
+          offset = 0
+          loop do
+            ids = Buffered.store.failed_ids(offset:, limit: Store::PAGE_SIZE)
+            ids.filter_map { |id| find(id) }.each(&block)
+            break if ids.size < Store::PAGE_SIZE
+
+            offset += ids.size
+          end
+        end
+      end
+
+      def_delegators :@record, :id, :job_class, :group, :state, :attempts, :failed_at, :last_error, :exhausted?,
+                     :input_jids
+
+      def initialize(record)
+        @record = record
+      end
+
+      def items
+        @record.arguments
+      end
+
+      # Re-reads the batch. One that has since been performed and removed keeps
+      # its last known state.
+      def reload
+        record = Buffered.store.fetch(id)
+        @record = record if record
+        self
+      end
+
+      def retry!
+        Buffered.store.replay(id)
+        reload
+      end
+
+      # The batch no longer exists in Redis after this; the state is a local marker.
+      def discard!
+        Buffered.store.discard(id)
+        @record = @record.with(state: "discarded", exhausted: false)
+        self
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/batch_context.rb

@@ -0,0 +1,37 @@
+module Sidekiq
+  module Buffered
+    # What a running batch exposes to +perform+ through +buffered_batch+.
+    class BatchContext
+      KEY = :sidekiq_buffered_batch
+
+      attr_reader :id, :job_class, :group, :input_jids, :attempts
+
+      def self.current
+        Thread.current[KEY]
+      end
+
+      def self.with(record, &)
+        ThreadLocal.with(KEY, new(record), &)
+      end
+
+      def initialize(record)
+        @id = record.id.dup.freeze
+        @job_class = record.job_class.dup.freeze
+        @group = immutable(record.group)
+        @input_jids = record.input_jids.map { |jid| jid.dup.freeze }.freeze
+        @attempts = record.attempts
+        freeze
+      end
+
+      private
+
+      def immutable(value)
+        case value
+        when Hash then value.each_value { |entry| immutable(entry) }
+        when Array then value.each { |entry| immutable(entry) }
+        end
+        value.freeze
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/batch_record.rb

@@ -0,0 +1,79 @@
+module Sidekiq
+  module Buffered
+    # One input inside a stored batch.
+    Item = Data.define(:jid, :args, :context, :active_job, :buffered_at) do
+      def self.parse(raw)
+        data = Sidekiq.load_json(raw)
+        new(
+          jid: data.fetch("jid"),
+          args: data.fetch("args"),
+          context: data.fetch("context", {}),
+          active_job: data["active_job"],
+          buffered_at: data["buffered_at"],
+        )
+      end
+
+      def active_job?
+        !active_job.nil?
+      end
+    end
+
+    # A batch as Redis holds it.
+    BatchRecord = Data.define(
+      :id, :job_class, :group, :items, :attempts, :lease_timeout, :state, :exhausted, :failed_at, :last_error,
+    ) do
+      # The fields the claim script returns for a batch that is now running.
+      def self.claimed(id, values)
+        job_class, group, raw_items, attempts, lease_timeout = values
+        new(
+          id:, job_class:, group: Input.decode_group(group), items: parse_items(raw_items),
+          attempts: attempts.to_i, lease_timeout: lease_timeout.to_i, state: "running",
+        )
+      end
+
+      # A full read of the batch hash plus its score in the failed index.
+      def self.stored(id, data, failed_at:)
+        exhausted = data["exhausted"] == "1"
+        new(
+          id:,
+          job_class: data.fetch("class"),
+          group: Input.decode_group(data["group"]),
+          items: parse_items(data.fetch("items")),
+          attempts: data.fetch("attempts", "0").to_i,
+          lease_timeout: data["lease_timeout"]&.to_i,
+          state: data["state"] || (exhausted ? "failed" : "queued"),
+          exhausted:,
+          failed_at: failed_at && Float(failed_at),
+          last_error: data["last_error"] && Sidekiq.load_json(data["last_error"]),
+        )
+      end
+
+      def self.parse_items(raw)
+        Sidekiq.load_json(raw).map { |item| Item.parse(item) }
+      end
+      private_class_method :parse_items
+
+      def initialize(id:, job_class:, items:, group: nil, attempts: 0, lease_timeout: nil, state: "queued",
+                     exhausted: false, failed_at: nil, last_error: nil)
+        super(id: id.to_s, job_class:, items:, group:, attempts:, lease_timeout:, state:, exhausted:, failed_at:,
+              last_error:)
+      end
+
+      def exhausted?
+        exhausted
+      end
+
+      def active_job?
+        items.first.active_job?
+      end
+
+      def input_jids
+        items.map(&:jid)
+      end
+
+      def arguments
+        items.map(&:args)
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/client.rb

@@ -0,0 +1,109 @@
+require "sidekiq/client"
+
+module Sidekiq
+  module Buffered
+    # Sidekiq::Client subclass used by buffered job classes. It admits every
+    # immediate input of a push_bulk chunk in one Redis script and returns
+    # input JIDs, which the plain Sidekiq::Client cannot do once the client
+    # middleware diverts a payload away from the queue.
+    class Client < Sidekiq::Client
+      def push(item)
+        return super unless Buffered.registered?(item["class"])
+
+        Collector.with(store) do |collector|
+          before = collector.size
+          super(item) || collector.jid_at(before)
+        end
+      end
+
+      def push_bulk(items)
+        return super unless Buffered.registered?(items["class"])
+
+        options = BulkOptions.new(items)
+        options.slices.flat_map { |slice| push_slice(items, slice) }
+      end
+
+      private
+
+      # Each slice is one Redis script call. A slice with +at+ holds only
+      # scheduled inputs, which Sidekiq schedules individually.
+      def push_slice(items, slice)
+        Collector.with(store) { slice.map { |payload| push(items.merge(payload)) } }
+      end
+
+      def store
+        @store ||= Store.new(redis_pool)
+      end
+
+      # Mirrors the bulk options Sidekiq::Client#push_bulk consumes and yields
+      # per-input payload overrides in batch_size slices.
+      class BulkOptions
+        DEFAULT_BATCH_SIZE = 1_000
+        MIN_SPREAD = 5
+
+        def initialize(items)
+          @args = items["args"]
+          raise ArgumentError, "Bulk arguments must be an Array of Arrays: [[1], [2]]" unless arrays?(@args)
+
+          @at = items.delete("at") || items.delete(:at)
+          @spread = items.delete(:spread_interval) || items.delete("spread_interval")
+          @jid = items.delete("jid")
+          @batch_size = items.delete(:batch_size) || items.delete("batch_size") || DEFAULT_BATCH_SIZE
+          validate_at!
+          validate_spread!
+          @at = spread_times if @spread
+        end
+
+        def slices
+          overrides = @args.map.with_index do |job_args, index|
+            override = { "args" => job_args }
+            override["jid"] = @jid if @jid
+            override["at"] = @at.is_a?(Array) ? @at[index] : @at if @at
+            override
+          end
+          overrides.each_slice(@batch_size)
+        end
+
+        private
+
+        def arrays?(value)
+          value.is_a?(Array) && value.all?(Array)
+        end
+
+        def validate_at!
+          return validate_jid! unless @at
+
+          unless Array(@at).any? && Array(@at).all?(Numeric)
+            raise ArgumentError, "Job 'at' must be a Numeric or an Array of Numeric timestamps"
+          end
+          if @at.is_a?(Array) && @at.size != @args.size
+            raise ArgumentError, "Job 'at' Array must have same size as 'args' Array"
+          end
+
+          validate_jid!
+        end
+
+        def validate_jid!
+          return unless @jid && @args.size > 1
+
+          raise ArgumentError, "Explicitly passing 'jid' when pushing more than one job is not supported"
+        end
+
+        def validate_spread!
+          return unless @spread
+
+          unless @spread.is_a?(Numeric) && @spread.positive?
+            raise ArgumentError, "Jobs 'spread_interval' must be a positive Numeric"
+          end
+          raise ArgumentError, "Only one of 'at' or 'spread_interval' can be provided" if @at
+        end
+
+        def spread_times
+          spread = [@spread, MIN_SPREAD].max
+          now = Time.now.to_f
+          Array.new(@args.size) { now + (rand * spread) }
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/client_middleware.rb

@@ -0,0 +1,38 @@
+module Sidekiq
+  module Buffered
+    # Sidekiq client middleware that diverts buffered inputs into Redis buffers.
+    # It runs after the rest of the chain so buffered inputs capture the metadata
+    # other client middleware attach (CurrentAttributes, tags, trace fields).
+    class ClientMiddleware
+      include Sidekiq::ClientMiddleware
+
+      def call(_job_class, job, _queue, redis_pool)
+        return yield if Inline.active?
+
+        buffered = Buffered.registered?(job["class"])
+        Input.normalize!(job) if buffered
+        payload = yield
+        return payload unless payload
+
+        if buffered
+          admit(payload, redis_pool)
+        else
+          Dispatch.prepare_processor_payload(payload) if payload["class"] == Processor.name
+          payload
+        end
+      end
+
+      private
+
+      def admit(payload, redis_pool)
+        # Fake mode keeps inputs in the job class's own testing array.
+        return payload if Testing.fake?
+        # Delay first, buffer second: Sidekiq schedules the admission job itself.
+        return ScheduledInput.wrap(payload) if payload.key?("at") && !Testing.inline?
+
+        Collector.admit(payload, Store.new(redis_pool))
+        false
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/collector.rb

@@ -0,0 +1,63 @@
+module Sidekiq
+  module Buffered
+    # Thread-local collection point for buffered payloads admitted through one
+    # client call. The client middleware stashes payloads here instead of writing
+    # them one by one; the outermost collector writes them atomically.
+    class Collector
+      KEY = :sidekiq_buffered_collector
+
+      class << self
+        def current
+          Thread.current[KEY]
+        end
+
+        def with(store)
+          return yield(current) if current
+
+          ThreadLocal.with(KEY, new(store)) do |collector|
+            result = yield(collector)
+            collector.write
+            result
+          end
+        end
+
+        # Middleware entry point: stash under an active collector, otherwise write now.
+        def admit(payload, store)
+          if current
+            current.stash(payload)
+          else
+            new(store).stash(payload).write
+          end
+        end
+      end
+
+      def initialize(store)
+        @store = store
+        @payloads = []
+      end
+
+      def size
+        @payloads.size
+      end
+
+      def stash(payload)
+        @payloads << payload
+        self
+      end
+
+      def jid_at(index)
+        @payloads[index]&.fetch("jid")
+      end
+
+      def write
+        return if @payloads.empty?
+
+        if Testing.inline?
+          Testing.inline(@payloads)
+        else
+          @store.append(@payloads)
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/config.rb

@@ -0,0 +1,26 @@
+module Sidekiq
+  module Buffered
+    # The +buffered+ declaration of one job class.
+    Config = Data.define(:max_size, :max_wait, :group_by, :lease_timeout) do
+      def initialize(max_size:, max_wait:, group_by: nil, lease_timeout: nil)
+        raise ArgumentError, "max_size must be a positive Integer" unless positive_integer?(max_size)
+
+        wait = max_wait.to_f
+        raise ArgumentError, "max_wait must be positive" unless wait.positive?
+
+        unless lease_timeout.nil? || positive_integer?(lease_timeout)
+          raise ArgumentError, "lease_timeout must be a positive Integer or nil"
+        end
+
+        key = group_by.to_s
+        super(max_size:, max_wait: wait, group_by: key.empty? ? nil : key, lease_timeout:)
+      end
+
+      private
+
+      def positive_integer?(value)
+        value.is_a?(Integer) && value.positive?
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/dispatch.rb

@@ -0,0 +1,95 @@
+module Sidekiq
+  module Buffered
+    # Enqueues a Processor job for each materialized batch and records the exact
+    # queued payload, so recovery can check whether it still waits in its queue.
+    class Dispatch
+      KEY = :sidekiq_buffered_dispatch
+      # Fields Sidekiq's retry machinery adds to a processor payload.
+      PROCESSOR_FIELDS = %w[display_class buffered_context retry_count retried_at failed_at error_message
+                            error_class error_backtrace discarded_at].freeze
+
+      Target = Data.define(:id, :job_class, :context)
+
+      # Called by the client middleware for every processor push, including
+      # Sidekiq's own retry and Dead set re-pushes. The stored batch context
+      # wins over whatever ambient metadata the client chain just attached.
+      # Mutates in place: Sidekiq pushes this same object.
+      def self.prepare_processor_payload(payload)
+        if payload.key?("buffered_context")
+          context = payload.fetch("buffered_context")
+          payload.select! { |key, _value| Input::PAYLOAD_FIELDS.include?(key) || PROCESSOR_FIELDS.include?(key) }
+          payload.merge!(context.except(*Input::EXECUTION_FIELDS))
+        end
+        Thread.current[KEY]&.push(payload)
+        payload
+      end
+
+      def initialize(store)
+        @store = store
+      end
+
+      def enqueue(ids)
+        targets = targets_for(ids)
+        return if targets.empty?
+
+        client = Sidekiq::Client.new(pool: @store.redis_pool)
+        queued = capture do
+          targets.group_by { |target| processor_options(target) }.each do |options, batch|
+            client.push_bulk(options.merge("class" => Processor, "args" => batch.map { |target| [target.id] }))
+          end
+        end
+        record(queued)
+      end
+
+      private
+
+      # Dispatch needs only the class and the shared middleware context, not the
+      # item payloads. Batches written before the context field existed fall
+      # back to reading their first item.
+      def targets_for(ids)
+        return [] if ids.empty?
+
+        values = @store.redis_pool.with do |conn|
+          conn.pipelined do |pipeline|
+            ids.each { |id| pipeline.call("HMGET", @store.batch_key(id), "class", "dispatch_context") }
+          end
+        end
+        ids.zip(values).filter_map do |id, (job_class, context)|
+          Target.new(id:, job_class:, context: stored_context(id, context)) if job_class
+        end
+      end
+
+      def stored_context(id, raw)
+        return Sidekiq.load_json(raw) if raw
+
+        record = @store.fetch(id)
+        record ? record.items.first.context : {}
+      end
+
+      # +wrapped+ names the job class for Sidekiq's logs, metrics, and retry
+      # callbacks; +display_class+ does the same for Sidekiq Web.
+      def processor_options(target)
+        target.context.merge(
+          "wrapped" => target.job_class, "display_class" => target.job_class, "buffered_context" => target.context,
+        )
+      end
+
+      # The middleware hands back the payload objects Sidekiq pushes. Sidekiq
+      # stamps enqueued_at on those same objects, so serializing them after the
+      # push reproduces the queue entry recovery looks for with LPOS.
+      def capture
+        ThreadLocal.with(KEY, []) do |captured|
+          yield
+          captured
+        end
+      end
+
+      def record(payloads)
+        calls = payloads.reject { |payload| payload.key?("at") }.map do |payload|
+          [[@store.batch_key(payload["args"].first)], ["queue:#{payload['queue']}", Sidekiq.dump_json(payload)]]
+        end
+        @store.redis_pool.with { |conn| Scripts::RECORD_DISPATCH.call_each(conn, calls) }
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/events.rb

@@ -0,0 +1,29 @@
+module Sidekiq
+  module Buffered
+    class << self
+      attr_reader :on_event
+
+      def on_event=(handler)
+        raise ArgumentError, "on_event must respond to call or be nil" unless handler.nil? || handler.respond_to?(:call)
+
+        @on_event = handler
+      end
+
+      def notify(event, **details)
+        @on_event&.call(event, event_snapshot(details))
+      rescue StandardError => e
+        Sidekiq.logger.warn("sidekiq-buffered event handler failed: #{e.class}: #{e.message}")
+      end
+
+      def event_snapshot(value)
+        case value
+        when Hash then value.transform_values { |item| event_snapshot(item) }.freeze
+        when Array then value.map { |item| event_snapshot(item) }.freeze
+        when String then value.dup.freeze
+        else value
+        end
+      end
+      private :event_snapshot
+    end
+  end
+end