sidekiq-buffered 0.1.0

data/lib/sidekiq/buffered/store.rb

@@ -0,0 +1,363 @@
+require "digest"
+require_relative "scripts"
+
+module Sidekiq
+  module Buffered
+    # Redis persistence for buffers and batches. Every mutation is one Lua
+    # script, so concurrent producers, pollers, and workers never see a batch
+    # half-written.
+    class Store
+      PREFIX = "sidekiq-buffered:v1".freeze
+      DEADLINES = "#{PREFIX}:deadlines".freeze
+      BATCH_SEQ = "#{PREFIX}:batch-seq".freeze
+      BATCH_PREFIX = "#{PREFIX}:batch:".freeze
+      LEASE_PREFIX = "#{PREFIX}:lease:".freeze
+      META = "#{PREFIX}:meta".freeze
+      OUTSTANDING = "#{PREFIX}:outstanding".freeze
+      FAILED = "#{PREFIX}:failed".freeze
+      DEFAULT_CONTEXT = { "queue" => "default", "retry" => true }.freeze
+      PAGE_SIZE = 100
+      MAX_ERROR_MESSAGE_BYTES = 2048
+
+      Due = Struct.new(:now, :buffers, :outstanding, :next_at)
+
+      attr_reader :redis_pool
+
+      # Redis is the clock unless +clock+ (a callable returning Unix seconds) is
+      # given, so skewed producer and worker clocks cannot disagree about deadlines.
+      def initialize(redis_pool, clock: nil)
+        @redis_pool = redis_pool
+        @clock = clock
+      end
+
+      def now
+        return @clock.call if @clock
+
+        seconds, microseconds = redis { |conn| conn.call("TIME") }
+        seconds.to_i + (microseconds.to_i / 1_000_000.0)
+      end
+
+      def batch_key(batch_id)
+        "#{BATCH_PREFIX}#{batch_id}"
+      end
+
+      def lease_key(batch_id)
+        "#{LEASE_PREFIX}#{batch_id}"
+      end
+
+      # Admission is one script call. Redis stamps the time itself unless a clock is pinned.
+      def append(payloads, now: nil)
+        groups = payloads.map { |payload| Input.new(payload) }.group_by(&:identity).map do |identity, inputs|
+          keys, args = append_call(identity, inputs)
+          { keys:, args: }
+        end
+        keys = groups.flat_map { |group| group[:keys] }.uniq
+        ids = redis { |conn| Scripts::APPEND.call(conn, keys:, args: [Sidekiq.dump_json(groups), pinned(now)]) }
+        dispatch(ids)
+        Buffered.notify(:appended, input_count: payloads.size, batch_ids: ids)
+        ids
+      end
+
+      def flush(scope = Scope.new)
+        at = now
+        ids = []
+        each_buffer do |member, meta|
+          next unless scope.matches?(meta["class"], Input.decode_group(meta["group"]))
+
+          keys, args = flush_call(member, now: at, force: true)
+          ids.concat(redis { |conn| Scripts::FLUSH_DUE.call(conn, keys:, args:) } || [])
+        end
+        dispatch(ids)
+        ids
+      end
+
+      def each_buffer
+        cursor = "0"
+        loop do
+          cursor, entries = redis { |conn| conn.call("ZSCAN", DEADLINES, cursor, "COUNT", PAGE_SIZE) }
+          entries.each_slice(2) do |member, _score|
+            raw = redis { |conn| conn.call("HGET", META, member) }
+            yield member, Sidekiq.load_json(raw) if raw
+          end
+          break if cursor == "0"
+        end
+      end
+
+      # Snapshot of everything the poller acts on, in one round trip.
+      def due(now: nil, limit: 500)
+        values = redis { |conn| Scripts::TICK.call(conn, keys: [DEADLINES, OUTSTANDING], args: [pinned(now), limit]) }
+        upcoming = values.values_at(3, 4).filter_map { |raw| Float(raw) unless raw.nil? || raw.empty? }
+        Due.new(Float(values.fetch(0)), values.fetch(1), values.fetch(2), upcoming.min)
+      end
+
+      def flush_due(now: nil, limit: 500)
+        snapshot = due(now:, limit:)
+        flush_members(snapshot.buffers, now: snapshot.now)
+      end
+
+      def flush_members(members, now:)
+        return [] if members.empty?
+
+        results = redis { |conn| Scripts::FLUSH_DUE.call_each(conn, members.map { |member| flush_call(member, now:) }) }
+        ids = results.flat_map { |flushed| Array(flushed) }
+        dispatch(ids)
+        ids
+      end
+
+      def recover_outstanding(now: nil, limit: 500)
+        snapshot = due(now:, limit:)
+        recover(snapshot.outstanding, now: snapshot.now)
+      end
+
+      def recover(ids, now:)
+        return [] if ids.empty?
+
+        recover_at = now + Buffered.recovery_interval
+        results = redis do |conn|
+          calls = ids.map do |id|
+            [[OUTSTANDING, lease_key(id), batch_key(id)], [id, now, recover_at, Buffered.recovery_max_interval]]
+          end
+          Scripts::RECOVER.call_each(conn, calls)
+        end
+        reserved = ids.zip(results).filter_map { |id, result| id if result == 1 }
+        return reserved if reserved.empty?
+
+        dispatch(reserved)
+        Buffered.notify(:recovered, batch_ids: reserved)
+        Sidekiq.logger.info("sidekiq-buffered recovered batches: #{reserved.join(', ')}")
+        reserved
+      end
+
+      def claim(batch_id, token:, lease_token: token, now: nil)
+        now ||= self.now
+        values = redis do |conn|
+          Scripts::CLAIM.call(
+            conn,
+            keys: batch_keys(batch_id),
+            args: [token.to_s, Buffered.lease_timeout, now, batch_id, lease_token.to_s],
+          )
+        end
+        unless values
+          Buffered.notify(:claim_skipped, batch_id:)
+          Sidekiq.logger.debug { "sidekiq-buffered claim skipped: batch=#{batch_id}" }
+          return
+        end
+
+        record = BatchRecord.claimed(batch_id, values)
+        Buffered.notify(:claimed, batch_id:, attempts: record.attempts)
+        record
+      end
+
+      # Extends every lease still owned by its token and returns the leases that
+      # could not be extended.
+      def extend_leases(leases, now: nil)
+        now ||= self.now
+        results = redis do |conn|
+          calls = leases.map do |lease|
+            [
+              [lease_key(lease.batch_id), OUTSTANDING],
+              [lease.token, lease.timeout, now + lease.timeout, lease.batch_id],
+            ]
+          end
+          Scripts::EXTEND.call_each(conn, calls)
+        end
+        leases.zip(results).filter_map { |lease, extended| lease if extended.to_i.zero? }
+      end
+
+      def release(batch_id, token:, now: nil)
+        now ||= self.now
+        released = redis do |conn|
+          Scripts::RELEASE.call(
+            conn,
+            keys: [batch_key(batch_id), lease_key(batch_id), OUTSTANDING],
+            args: [token.to_s, now + Buffered.recovery_interval, batch_id],
+          )
+        end
+        Buffered.notify(:released, batch_id:) if released == 1
+        released
+      end
+
+      def ack(batch_id, token:)
+        acknowledged = redis { |conn| Scripts::ACK.call(conn, keys: batch_keys(batch_id), args: [token.to_s, batch_id]) }
+        Sidekiq.logger.warn("sidekiq-buffered acknowledgement lost: batch=#{batch_id}") if acknowledged.zero?
+        Buffered.notify(acknowledged == 1 ? :acknowledged : :acknowledgement_lost, batch_id:)
+        acknowledged
+      end
+
+      # Protects recovery through the latest possible retry time plus the
+      # batch's lease timeout. The original Sidekiq JID can reclaim regardless.
+      def mark_retrying(batch_id, token:, exception:, delay:, now: nil)
+        now ||= self.now
+        updated = redis do |conn|
+          Scripts::RETRYING.call(
+            conn,
+            keys: [batch_key(batch_id), lease_key(batch_id), OUTSTANDING],
+            args: [batch_id, token, delay, now, Buffered.lease_timeout, error_details(exception, now)],
+          )
+        end
+        Buffered.notify(:retrying, batch_id:) if updated == 1
+        updated
+      end
+
+      def exhaust(batch_id, token: nil, exception: nil, now: nil)
+        now ||= self.now
+        updated = redis do |conn|
+          Scripts::EXHAUST.call(
+            conn,
+            keys: batch_keys(batch_id),
+            args: [batch_id, now, token.to_s, exception ? error_details(exception, now) : ""],
+          )
+        end
+        Buffered.notify(:failed, batch_id:) if updated == 1
+        updated
+      end
+
+      def fetch(batch_id)
+        redis do |conn|
+          data = conn.hgetall(batch_key(batch_id))
+          next if data.nil? || data.empty?
+
+          BatchRecord.stored(batch_id, data, failed_at: conn.zscore(FAILED, batch_id))
+        end
+      end
+
+      def failed_ids(limit: PAGE_SIZE, offset: 0)
+        redis { |conn| conn.call("ZRANGE", FAILED, offset, offset + limit - 1) }
+      end
+
+      def replay(batch_id, now: nil)
+        now ||= self.now
+        result = redis do |conn|
+          Scripts::REPLAY.call(conn, keys: batch_keys(batch_id), args: [batch_id, now + Buffered.recovery_interval])
+        end
+        remove_dead_entry(exhausted_result!(batch_id, result))
+        dispatch([batch_id])
+        batch_id
+      end
+
+      def discard(batch_id)
+        result = redis { |conn| Scripts::DISCARD.call(conn, keys: batch_keys(batch_id), args: [batch_id]) }
+        remove_dead_entry(exhausted_result!(batch_id, result))
+        Buffered.notify(:discarded, batch_id:)
+        batch_id
+      end
+
+      def cleanup_failed(now: nil, limit: PAGE_SIZE)
+        retention = Buffered.failed_retention
+        return [] if retention.nil?
+
+        cutoff = (now || self.now) - retention
+        redis do |conn|
+          Scripts::CLEANUP_FAILED.call(
+            conn,
+            keys: [FAILED, OUTSTANDING],
+            args: [BATCH_PREFIX, LEASE_PREFIX, cutoff, limit],
+          )
+        end || []
+      end
+
+      private
+
+      def redis(&)
+        @redis_pool.with(&)
+      end
+
+      def dispatch(ids)
+        Dispatch.new(self).enqueue(ids)
+      end
+
+      # The clock value handed to scripts that can read Redis TIME themselves:
+      # empty unless the caller or an injected clock pins it.
+      def pinned(now)
+        now ||= @clock&.call
+        now.nil? ? "" : format("%.6f", now)
+      end
+
+      # The KEYS most lifecycle scripts operate on.
+      def batch_keys(batch_id)
+        [batch_key(batch_id), lease_key(batch_id), OUTSTANDING, FAILED]
+      end
+
+      # Replay and discard return {1, previous_jid} on success, 0 when the batch
+      # is not exhausted, and -1 when it does not exist.
+      def exhausted_result!(batch_id, result)
+        status, previous_jid = Array(result)
+        case status.to_i
+        when 1 then previous_jid
+        when 0 then raise ArgumentError, "Buffered batch #{batch_id} is not exhausted"
+        else raise ArgumentError, "Buffered batch #{batch_id} was not found"
+        end
+      end
+
+      def remove_dead_entry(jid)
+        return if jid.nil? || jid.empty?
+
+        redis do |conn|
+          cursor = "0"
+          loop do
+            cursor, entries = conn.call("ZSCAN", "dead", cursor, "MATCH", "*\"jid\":\"#{jid}\"*", "COUNT", PAGE_SIZE)
+            values = entries.each_slice(2).map(&:first)
+            conn.call("ZREM", "dead", *values) unless values.empty?
+            break if cursor == "0"
+          end
+        end
+      end
+
+      def append_call(identity, inputs)
+        config = identity.config
+        member = buffer_key(identity.job_class, identity.group, identity.context)
+        [
+          [member, DEADLINES, BATCH_SEQ, META, OUTSTANDING],
+          [
+            identity.job_class, identity.group, config.max_size, config.max_wait, member, BATCH_PREFIX,
+            buffer_meta(identity), Buffered.recovery_interval, *inputs.map(&:dump),
+          ],
+        ]
+      end
+
+      # +context+ is stored as a JSON string so Lua copies it verbatim into each
+      # batch; every input in a buffer shares the same middleware context.
+      def buffer_meta(identity)
+        Sidekiq.dump_json(
+          "class" => identity.job_class,
+          "group" => identity.group,
+          "id_prefix" => SecureRandom.uuid,
+          "capacity_key" => buffer_key(identity.job_class, identity.group),
+          "max_size" => identity.config.max_size,
+          "lease_timeout" => identity.config.lease_timeout || 0,
+          "context" => Sidekiq.dump_json(identity.context),
+        )
+      end
+
+      def flush_call(member, now:, force: false)
+        [
+          [member, DEADLINES, BATCH_SEQ, META, OUTSTANDING],
+          [member, now, BATCH_PREFIX, now + Buffered.recovery_interval, SecureRandom.uuid, force ? "1" : "0"],
+        ]
+      end
+
+      def error_details(exception, now)
+        message = exception.message.to_s.encode("UTF-8", invalid: :replace, undef: :replace)
+        Sidekiq.dump_json(
+          "class" => exception.class.name,
+          "message" => message.byteslice(0, MAX_ERROR_MESSAGE_BYTES).scrub(""),
+          "at" => now,
+        )
+      end
+
+      def canonical(value)
+        case value
+        when Hash then value.sort.to_h.transform_values { |entry| canonical(entry) }
+        when Array then value.map { |entry| canonical(entry) }
+        else value
+        end
+      end
+
+      def buffer_key(job_class, group, context = {})
+        context = context.reject { |key, value| DEFAULT_CONTEXT.key?(key) && DEFAULT_CONTEXT[key] == value }
+        group = Sidekiq.dump_json([group, canonical(context)]) unless context.empty?
+        "#{PREFIX}:buffer:#{job_class}:#{Digest::SHA256.hexdigest(group)}"
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/testing.rb

@@ -0,0 +1,65 @@
+module Sidekiq
+  module Buffered
+    # Sidekiq::Testing support: fake mode keeps inputs in the job class's jobs
+    # array and drains them in batches; inline mode performs partial batches
+    # immediately. Neither touches Redis.
+    module Testing
+      class << self
+        def enabled?
+          defined?(Sidekiq::Testing) && Sidekiq::Testing.enabled?
+        end
+
+        def fake?
+          enabled? && Sidekiq::Testing.fake?
+        end
+
+        def inline?
+          enabled? && Sidekiq::Testing.inline?
+        end
+
+        def inline(payloads)
+          payloads.map { |payload| Input.new(payload) }.group_by(&:identity).each do |identity, inputs|
+            inputs.each_slice(identity.config.max_size) { |batch| execute(batch) }
+          end
+        end
+
+        def drain(klass)
+          perform_one(klass) until Sidekiq::Queues.jobs_by_class[klass.name].empty?
+        end
+
+        def perform_one(klass)
+          payloads = Sidekiq::Queues.jobs_by_class[klass.name]
+          raise Sidekiq::EmptyQueueError, "perform_one called with no jobs available" if payloads.empty?
+
+          inputs = payloads.map { |payload| Input.new(payload) }
+          identity = inputs.first.identity
+          batch = inputs.select { |input| input.identity == identity }.first(identity.config.max_size)
+          batch.each { |input| Sidekiq::Queues.delete_for(input.jid, input.payload["queue"], klass.name) }
+          execute(batch)
+        end
+
+        private
+
+        def execute(inputs)
+          first = inputs.first
+          record = BatchRecord.new(
+            id: SecureRandom.uuid,
+            job_class: first.job_class,
+            group: Input.decode_group(first.group),
+            items: inputs.map { |input| Item.parse(input.dump) },
+            attempts: 1,
+          )
+          BatchContext.with(record) { perform(record, first) }
+        end
+
+        def perform(record, first)
+          if record.active_job?
+            ActiveJobExecution.perform(record, first.jid)
+          else
+            Object.const_get(record.job_class).process_job(first.payload.merge("args" => [record.arguments]))
+          end
+        end
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/thread_local.rb

@@ -0,0 +1,14 @@
+module Sidekiq
+  module Buffered
+    # Sets a thread-local for the duration of a block, then restores what was there.
+    module ThreadLocal
+      def self.with(key, value)
+        previous = Thread.current[key]
+        Thread.current[key] = value
+        yield value
+      ensure
+        Thread.current[key] = previous
+      end
+    end
+  end
+end

data/lib/sidekiq/buffered/version.rb

@@ -0,0 +1,5 @@
+module Sidekiq
+  module Buffered
+    VERSION = "0.1.0".freeze
+  end
+end

data/lib/sidekiq/buffered.rb

@@ -0,0 +1,172 @@
+require "sidekiq"
+require_relative "buffered/version"
+require_relative "buffered/config"
+require_relative "buffered/registry"
+require_relative "buffered/scope"
+require_relative "buffered/thread_local"
+require_relative "buffered/input"
+require_relative "buffered/batch_record"
+require_relative "buffered/store"
+require_relative "buffered/dispatch"
+require_relative "buffered/collector"
+require_relative "buffered/inline"
+require_relative "buffered/client"
+require_relative "buffered/client_middleware"
+require_relative "buffered/setter"
+require_relative "buffered/testing"
+require_relative "buffered/scheduled_input"
+require_relative "buffered/periodic_thread"
+require_relative "buffered/poller"
+require_relative "buffered/heartbeat"
+require_relative "buffered/execution"
+require_relative "buffered/retry_callbacks"
+require_relative "buffered/processor"
+require_relative "buffered/server_middleware"
+require_relative "buffered/batch"
+require_relative "buffered/batch_context"
+require_relative "buffered/statistics"
+require_relative "buffered/events"
+
+module Sidekiq
+  module Buffered
+    ANY_GROUP = Object.new.freeze
+    DEFAULT_POLL_INTERVAL = 0.25
+    DEFAULT_RECOVERY_INTERVAL = 30
+    DEFAULT_RECOVERY_MAX_INTERVAL = 300
+    DEFAULT_LEASE_TIMEOUT = 60
+    DEFAULT_FAILED_RETENTION = nil
+
+    @registry = Registry.new
+    @poll_interval = DEFAULT_POLL_INTERVAL
+    @recovery_interval = DEFAULT_RECOVERY_INTERVAL
+    @recovery_max_interval = DEFAULT_RECOVERY_MAX_INTERVAL
+    @lease_timeout = DEFAULT_LEASE_TIMEOUT
+    @failed_retention = DEFAULT_FAILED_RETENTION
+
+    class << self
+      attr_reader :registry, :poll_interval, :lease_timeout, :failed_retention, :recovery_interval,
+                  :recovery_max_interval
+
+      def current_batch
+        BatchContext.current
+      end
+
+      # A store on Sidekiq's current default Redis pool.
+      def store
+        Store.new(Sidekiq.redis_pool)
+      end
+
+      def flush(job_class: nil, group: ANY_GROUP)
+        store.flush(Scope.new(job_class:, group:))
+      end
+
+      def stats(**)
+        Statistics.new(**).call
+      end
+
+      def included(base)
+        base.extend(ClassMethods)
+      end
+
+      def poll_interval=(value)
+        interval = value.to_f
+        raise ArgumentError, "poll_interval must be positive" unless interval.positive?
+
+        @poll_interval = interval
+      end
+
+      def recovery_interval=(value)
+        @recovery_interval = positive_integer(value, "recovery_interval")
+      end
+
+      def recovery_max_interval=(value)
+        @recovery_max_interval = positive_integer(value, "recovery_max_interval")
+      end
+
+      def lease_timeout=(value)
+        @lease_timeout = positive_integer(value, "lease_timeout")
+      end
+
+      def failed_retention=(value)
+        @failed_retention = positive_integer(value, "failed_retention", nullable: true)
+      end
+
+      def registered?(klass)
+        registry.registered?(klass)
+      end
+
+      def config_for(klass)
+        registry.fetch(klass)
+      end
+
+      private
+
+      def positive_integer(value, name, nullable: false)
+        return value if value.is_a?(Integer) && value.positive?
+        return if nullable && value.nil?
+
+        raise ArgumentError, "#{name} must be a positive Integer#{' or nil' if nullable}"
+      end
+    end
+
+    def buffered_batch
+      Buffered.current_batch
+    end
+
+    module ClassMethods
+      def buffered(max_size:, max_wait:, group_by: nil, lease_timeout: nil)
+        Buffered.registry.register(self, Config.new(max_size:, max_wait:, group_by:, lease_timeout:))
+        singleton_class.prepend(RetryCallbacks)
+        install_native_job_support if include?(Sidekiq::Job)
+      end
+
+      def flush(group: ANY_GROUP)
+        Buffered.flush(job_class: self, group:)
+      end
+
+      def drain
+        return super unless Buffered.registered?(self) && Testing.enabled?
+
+        Testing.drain(self)
+      end
+
+      def perform_one
+        return super unless Buffered.registered?(self) && Testing.enabled?
+
+        Testing.perform_one(self)
+      end
+
+      def sidekiq_buffered_config
+        return unless Buffered.registered?(self)
+
+        Buffered.config_for(self)
+      end
+
+      private
+
+      # Buffered pushes go through Sidekiq::Buffered::Client so callers keep
+      # receiving input JIDs and bulk chunks are admitted atomically.
+      def install_native_job_support
+        sidekiq_options("client_class" => Client)
+        extend(NativeJobMethods) unless singleton_class.include?(NativeJobMethods)
+      end
+    end
+  end
+end
+
+Sidekiq.default_configuration.client_middleware do |chain|
+  chain.prepend(Sidekiq::Buffered::ClientMiddleware)
+end
+
+Sidekiq.configure_server do |config|
+  config.server_middleware { |chain| chain.prepend(Sidekiq::Buffered::ServerMiddleware) }
+  config.on(:startup) do
+    Sidekiq::Buffered::Poller.start
+    Sidekiq::Buffered::Heartbeat.start
+  end
+  config.on(:quiet) { Sidekiq::Buffered::Poller.stop }
+  config.on(:shutdown) do
+    Sidekiq::Buffered::Poller.stop
+    Sidekiq::Buffered::Heartbeat.stop
+  end
+end

data/lib/sidekiq-buffered.rb

@@ -0,0 +1 @@
+require_relative "sidekiq/buffered"

data/sidekiq-buffered.gemspec

@@ -0,0 +1,37 @@
+require_relative "lib/sidekiq/buffered/version"
+
+Gem::Specification.new do |spec|
+  spec.name = "sidekiq-buffered"
+  spec.version = Sidekiq::Buffered::VERSION
+  spec.authors = ["Sarunas Kujalis"]
+  spec.email = ["sarunas@kujalis.lt"]
+
+  spec.summary = "Buffer Sidekiq job inputs and perform them in batches by size or wait time."
+  spec.description = <<~DESC
+    Declares a Sidekiq job as buffered, keeps ordinary enqueue APIs, and later
+    invokes perform with an array of the original argument hashes.
+  DESC
+  spec.homepage = "https://github.com/ksarunas/sidekiq-buffered"
+  spec.license = "MIT"
+  spec.required_ruby_version = ">= 3.2.0"
+
+  spec.metadata["homepage_uri"] = spec.homepage
+  spec.metadata["source_code_uri"] = spec.homepage
+  spec.metadata["changelog_uri"] = "#{spec.homepage}/blob/main/CHANGELOG.md"
+  spec.metadata["rubygems_mfa_required"] = "true"
+
+  spec.files = Dir.chdir(__dir__) do
+    `git ls-files -z`.split("\x0").reject do |f|
+      f.start_with?("bin/", "spec/", "features/", ".git")
+    end
+  end
+  spec.require_paths = ["lib"]
+
+  spec.add_dependency "sidekiq", ">= 7.0", "< 9"
+  spec.add_development_dependency "activejob", ">= 7.0", "< 9"
+  spec.add_development_dependency "rake", "~> 13.0"
+  spec.add_development_dependency "rspec", "~> 3.13"
+  spec.add_development_dependency "rubocop", "~> 1.81"
+  spec.add_development_dependency "rubocop-performance", "~> 1.25"
+  spec.add_development_dependency "rubocop-rspec", "~> 3.6"
+end