dispatch_policy 0.1.0

data/lib/dispatch_policy/active_job_perform_all_later_patch.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  # Rails 7.1's ActiveJob.perform_all_later(*jobs) bypasses ActiveJob::Base#enqueue
+  # and calls queue_adapter.enqueue_all directly. Dispatchable hooks on #enqueue,
+  # so without this patch the batch path would skip staging.
+  module ActiveJobPerformAllLaterPatch
+    def perform_all_later(*jobs)
+      jobs.flatten!
+
+      staged, remaining = jobs.partition do |job|
+        klass = job.class
+        klass.respond_to?(:dispatch_policy?) &&
+          klass.dispatch_policy? &&
+          DispatchPolicy.enabled?
+      end
+
+      staged_count = 0
+      staged.group_by(&:class).each do |klass, group|
+        staged_count += DispatchPolicy::StagedJob.stage_many!(
+          policy: klass.resolved_dispatch_policy,
+          jobs:   group
+        )
+      end
+
+      remaining_count = remaining.empty? ? 0 : super(*remaining)
+      staged_count + remaining_count.to_i
+    end
+  end
+end
+
+ActiveJob.singleton_class.prepend(DispatchPolicy::ActiveJobPerformAllLaterPatch)

data/lib/dispatch_policy/dispatch_context.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  class DispatchContext
+    def initialize(policy:, batch:)
+      @policy     = policy
+      @cache      = {}
+      @partitions = Hash.new { |h, k| h[k] = {} }
+      batch.each { |staged| resolve_for(staged) }
+    end
+
+    def for(staged)
+      @cache[staged.id]
+    end
+
+    def record_partitions(pairs, gate:)
+      pairs.each { |staged, partition_key| @partitions[staged.id][gate.to_sym] = partition_key.to_s }
+    end
+
+    def partitions_for(staged)
+      @partitions[staged.id]
+    end
+
+    def primary_partition_for(staged)
+      @partitions[staged.id].values.first
+    end
+
+    private
+
+    def resolve_for(staged)
+      cached = staged.context
+      if cached.is_a?(Hash) && cached.present?
+        @cache[staged.id] = cached.symbolize_keys
+        return
+      end
+
+      # Fallback: recompute from the serialized args. Hit on rows staged
+      # before the context column existed, or when context_builder
+      # legitimately returned an empty hash.
+      raw  = (staged.arguments || {})["arguments"] || []
+      args = begin
+        ActiveJob::Arguments.deserialize(raw)
+      rescue StandardError => e
+        Rails.logger&.warn(
+          "[DispatchPolicy] could not deserialize args for staged=#{staged.id} " \
+          "(policy=#{staged.policy_name}): #{e.class}: #{e.message}"
+        )
+        raw
+      end
+      @cache[staged.id] = @policy.context_builder.call(args)
+    end
+  end
+end

data/lib/dispatch_policy/dispatchable.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  module Dispatchable
+    extend ActiveSupport::Concern
+
+    class_methods do
+      def dispatch_policy(&block)
+        @dispatch_policy = DispatchPolicy::Policy.new(self, &block)
+      end
+
+      def dispatch_policy?
+        !@dispatch_policy.nil?
+      end
+
+      def resolved_dispatch_policy
+        @dispatch_policy
+      end
+
+      # Walk up the ancestor chain so subclasses inherit the parent policy.
+      def inherited(subclass)
+        super
+        subclass.instance_variable_set(:@dispatch_policy, @dispatch_policy)
+      end
+    end
+
+    included do
+      attr_accessor :_dispatch_partitions, :_dispatch_admitted_at
+
+      around_perform do |job, block|
+        # queue_lag = admitted_at → perform_start. Pure signal for "is the
+        # adapter queue building up?" (high = admitting too fast) vs "are
+        # workers idle?" (near zero = ready for more). Measured BEFORE
+        # block.call so perform duration doesn't pollute it.
+        admitted_at   = job._dispatch_admitted_at
+        perform_start = Time.current
+        queue_lag_ms  = admitted_at ? ((perform_start - admitted_at) * 1000).to_i : 0
+
+        succeeded = false
+        begin
+          block.call
+          succeeded = true
+        ensure
+          policy_name = job.class.resolved_dispatch_policy&.name
+
+          if job._dispatch_partitions.present?
+            DispatchPolicy::Tick.release(
+              policy_name: policy_name,
+              partitions:  job._dispatch_partitions
+            )
+
+            # Let adaptive gates update their AIMD state first; we pick up
+            # the resulting current_max in the generic observation below
+            # so the chart surfaces the cap alongside lag + completions.
+            policy = job.class.resolved_dispatch_policy
+            job._dispatch_partitions.each do |gate_name, partition_key|
+              gate = policy&.gates&.find { |g| g.name == gate_name.to_sym }
+              next unless gate.is_a?(DispatchPolicy::Gates::AdaptiveConcurrency)
+              gate.record_observation(
+                partition_key: partition_key,
+                queue_lag_ms:  queue_lag_ms,
+                succeeded:     succeeded
+              )
+            end
+
+            # Generic observation per unique partition. Every gate with
+            # partition_by (adaptive or not) gets a sparkline this way.
+            job._dispatch_partitions.values.uniq.each do |partition_key|
+              current_max = DispatchPolicy::AdaptiveConcurrencyStats.current_max_for(
+                policy_name:   policy_name,
+                partition_key: partition_key
+              )
+              DispatchPolicy::PartitionObservation.observe!(
+                policy_name:   policy_name,
+                partition_key: partition_key,
+                queue_lag_ms:  queue_lag_ms,
+                current_max:   current_max
+              )
+            end
+          end
+          DispatchPolicy::StagedJob.mark_completed_by_active_job_id(job.job_id)
+        end
+      end
+    end
+
+    def enqueue(options = {})
+      return super unless self.class.dispatch_policy?
+      if options[:_bypass_staging]
+        return super(options.except(:_bypass_staging))
+      end
+      return super unless DispatchPolicy.enabled?
+
+      # Mirror Active Job's scheduling option handling before staging.
+      self.scheduled_at = options[:wait].seconds.from_now if options[:wait]
+      self.scheduled_at = options[:wait_until] if options[:wait_until]
+      self.queue_name   = self.class.queue_name_from_part(options[:queue]) if options[:queue]
+      self.priority     = options[:priority].to_i if options[:priority]
+
+      DispatchPolicy::StagedJob.stage!(
+        job_instance: self,
+        policy:       self.class.resolved_dispatch_policy
+      )
+      self
+    end
+
+    def serialize
+      super.merge(
+        "_dispatch_partitions"  => _dispatch_partitions || {},
+        "_dispatch_admitted_at" => _dispatch_admitted_at&.iso8601(6)
+      )
+    end
+
+    def deserialize(job_data)
+      super
+      self._dispatch_partitions  = job_data["_dispatch_partitions"]
+      ts                         = job_data["_dispatch_admitted_at"]
+      self._dispatch_admitted_at = ts ? Time.iso8601(ts) : nil
+    end
+  end
+end

data/lib/dispatch_policy/engine.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require "rails/engine"
+
+module DispatchPolicy
+  class Engine < ::Rails::Engine
+    isolate_namespace DispatchPolicy
+
+    initializer "dispatch_policy.reference_gates" do
+      config.to_prepare do
+        # Reference the built-in gates so they register in Gate.registry.
+        DispatchPolicy::Gates::Concurrency
+        DispatchPolicy::Gates::Throttle
+        DispatchPolicy::Gates::GlobalCap
+        DispatchPolicy::Gates::FairInterleave
+        DispatchPolicy::Gates::AdaptiveConcurrency
+
+        DispatchPolicy::ActiveJobPerformAllLaterPatch
+      end
+    end
+
+    initializer "dispatch_policy.boot_prune", after: :load_config_initializers do
+      config.to_prepare do
+        begin
+          DispatchPolicy::Tick.prune_orphan_gate_rows
+          DispatchPolicy::Tick.prune_idle_partitions
+          DispatchPolicy::PartitionObservation.prune!
+        rescue ActiveRecord::NoDatabaseError,
+               ActiveRecord::StatementInvalid,
+               ActiveRecord::ConnectionNotEstablished
+          # DB not ready — skip silently.
+        end
+      end
+    end
+  end
+end

data/lib/dispatch_policy/gate.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  class Gate
+    class << self
+      def registry
+        @registry ||= {}
+      end
+
+      def register(name, klass)
+        registry[name.to_sym] = klass
+      end
+    end
+
+    attr_reader :policy, :partition_by, :name
+
+    def initialize(policy:, name:, partition_by: nil, **opts)
+      @policy       = policy
+      @name         = name
+      @partition_by = partition_by
+      configure(**opts)
+    end
+
+    def configure(**_opts); end
+
+    # Resolve a partition key for a given context.
+    def partition_key_for(ctx)
+      return "default" if @partition_by.nil?
+      @partition_by.call(ctx).to_s
+    end
+
+    # Subclasses must implement.
+    def filter(_batch, _context)
+      raise NotImplementedError
+    end
+
+    # Whether this gate keeps an in-flight count that must be released
+    # when the job finishes.
+    def tracks_inflight?
+      false
+    end
+
+    protected
+
+    def resolve(value, ctx)
+      value.respond_to?(:call) ? value.call(ctx) : value
+    end
+  end
+end

data/lib/dispatch_policy/gates/adaptive_concurrency.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  module Gates
+    # Adaptive variant of :concurrency. The cap per partition (current_max)
+    # shrinks when the adapter queue backs up (recent queue_lag > target) or
+    # when performs fail; grows back when workers drain admissions quickly
+    # (queue_lag near zero). The signal is pure queue wait — admitted_at →
+    # perform_start — so it reflects "are we admitting too fast?" without
+    # getting polluted by how long the external work takes.
+    #
+    # AIMD loop on a per-partition stats row; the underlying in-flight
+    # counter is the same PartitionInflightCount used by :concurrency.
+    class AdaptiveConcurrency < Gate
+      # alpha is fast enough that a single spike is forgotten in ~3
+      # observations instead of ~15. slow_factor 0.95 halves the per-
+      # observation shrink magnitude so the cap no longer overshoots
+      # after a burst drains the adapter queue.
+      DEFAULT_EWMA_ALPHA  = 0.5
+      DEFAULT_FAIL_FACTOR = 0.5
+      DEFAULT_SLOW_FACTOR = 0.95
+
+      # target_lag_ms accepts the legacy alias `target_latency` for
+      # backwards compatibility.
+      def configure(initial_max:,
+                    target_lag_ms: nil,
+                    target_latency: nil,
+                    min: 1,
+                    ewma_alpha: DEFAULT_EWMA_ALPHA,
+                    failure_decrease_factor: DEFAULT_FAIL_FACTOR,
+                    overload_decrease_factor: DEFAULT_SLOW_FACTOR)
+        @initial_max    = initial_max
+        @min            = min
+        @target_lag_ms  = target_lag_ms || target_latency
+        @ewma_alpha     = ewma_alpha
+        @fail_factor    = failure_decrease_factor
+        @slow_factor    = overload_decrease_factor
+        raise ArgumentError, "adaptive_concurrency requires target_lag_ms" if @target_lag_ms.nil?
+      end
+
+      def tracks_inflight?
+        true
+      end
+
+      attr_reader :initial_max, :min, :target_lag_ms,
+                  :ewma_alpha, :fail_factor, :slow_factor
+
+      def filter(batch, context)
+        by_partition = batch.group_by { |staged| partition_key_for(context.for(staged)) }
+
+        # Seed any missing stats rows so the first admission has something
+        # to read. Cheap: one INSERT ... ON CONFLICT DO NOTHING per key.
+        by_partition.each_key do |key|
+          AdaptiveConcurrencyStats.seed!(
+            policy_name:   policy.name,
+            gate_name:     name,
+            partition_key: key,
+            initial_max:   resolve(@initial_max, nil).to_i
+          )
+        end
+
+        stats = AdaptiveConcurrencyStats.fetch_many(
+          policy_name:    policy.name,
+          gate_name:      name,
+          partition_keys: by_partition.keys
+        )
+
+        in_flight = PartitionInflightCount.fetch_many(
+          policy_name:    policy.name,
+          gate_name:      name,
+          partition_keys: by_partition.keys
+        )
+
+        min_v = resolve(@min, nil).to_i
+
+        admitted = []
+        by_partition.each do |partition_key, jobs|
+          effective_max = stats.dig(partition_key, :current_max) || resolve(@initial_max, nil).to_i
+          effective_max = [ effective_max, min_v ].max
+          used = in_flight.fetch(partition_key, 0)
+
+          # Safety valve: if nothing is in-flight for this partition and
+          # there's pending, the adapter queue is (or is about to be)
+          # empty and workers will idle. Ensure we hand over at least
+          # initial_max so the stream never dries up on its own.
+          if used.zero? && jobs.any?
+            effective_max = [ effective_max, resolve(@initial_max, nil).to_i ].max
+          end
+
+          jobs.each do |staged|
+            break unless used < effective_max
+            admitted << [ staged, partition_key ]
+            used += 1
+          end
+        end
+
+        context.record_partitions(admitted, gate: name)
+        admitted.map(&:first)
+      end
+
+      # Called by Dispatchable#around_perform for each adaptive gate that
+      # touched this job. Lives on the gate instance because configuration
+      # (alpha, target_latency, etc.) is per gate.
+      def record_observation(partition_key:, queue_lag_ms:, succeeded:)
+        AdaptiveConcurrencyStats.record_observation!(
+          policy_name:       policy.name,
+          gate_name:         name,
+          partition_key:     partition_key.to_s,
+          queue_lag_ms:      queue_lag_ms,
+          succeeded:         succeeded,
+          alpha:             @ewma_alpha,
+          min:               resolve(@min, nil).to_i,
+          target_lag_ms:     resolve(@target_lag_ms, nil).to_f,
+          fail_factor:       @fail_factor,
+          slow_factor:       @slow_factor,
+          initial_max:       resolve(@initial_max, nil).to_i
+        )
+      end
+    end
+
+    Gate.register(:adaptive_concurrency, AdaptiveConcurrency)
+  end
+end

data/lib/dispatch_policy/gates/concurrency.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  module Gates
+    class Concurrency < Gate
+      def configure(max:)
+        @max = max
+      end
+
+      def tracks_inflight?
+        true
+      end
+
+      def filter(batch, context)
+        by_partition = batch.group_by { |staged| partition_key_for(context.for(staged)) }
+
+        in_flight = PartitionInflightCount.fetch_many(
+          policy_name:    policy.name,
+          gate_name:      name,
+          partition_keys: by_partition.keys
+        )
+
+        admitted = []
+        by_partition.each do |partition_key, jobs|
+          jobs.each do |staged|
+            ctx   = context.for(staged)
+            limit = resolve(@max, ctx).to_i
+            used  = in_flight.fetch(partition_key, 0)
+            if used < limit
+              admitted << [ staged, partition_key ]
+              in_flight[partition_key] = used + 1
+            end
+          end
+        end
+
+        context.record_partitions(admitted, gate: name)
+        admitted.map(&:first)
+      end
+    end
+
+    Gate.register(:concurrency, Concurrency)
+  end
+end

data/lib/dispatch_policy/gates/fair_interleave.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  module Gates
+    class FairInterleave < Gate
+      def configure(**_); end
+
+      def filter(batch, context)
+        groups = batch.group_by do |staged|
+          if @partition_by
+            partition_key_for(context.for(staged))
+          else
+            context.primary_partition_for(staged) || staged.id
+          end
+        end
+        interleaved = []
+        loop do
+          taken = false
+          groups.each_value do |g|
+            next if g.empty?
+            interleaved << g.shift
+            taken = true
+          end
+          break unless taken
+        end
+        interleaved
+      end
+    end
+
+    Gate.register(:fair_interleave, FairInterleave)
+  end
+end

data/lib/dispatch_policy/gates/global_cap.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  module Gates
+    class GlobalCap < Gate
+      def configure(max:)
+        @max = max
+      end
+
+      def tracks_inflight?
+        true
+      end
+
+      def filter(batch, context)
+        limit     = resolve(@max, nil).to_i
+        in_flight = PartitionInflightCount.total_for(policy_name: policy.name, gate_name: name)
+        capacity  = [ limit - in_flight, 0 ].max
+        head      = batch.first(capacity)
+        context.record_partitions(head.map { |s| [ s, "default" ] }, gate: name)
+        head
+      end
+    end
+
+    Gate.register(:global_cap, GlobalCap)
+  end
+end

data/lib/dispatch_policy/gates/throttle.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  module Gates
+    class Throttle < Gate
+      def configure(rate:, per:, burst: nil)
+        @rate  = rate
+        @per   = per
+        @burst = burst
+      end
+
+      # Consumed tokens refill over time, no release step.
+      def tracks_inflight?
+        false
+      end
+
+      def filter(batch, context)
+        by_partition = batch.group_by { |staged| partition_key_for(context.for(staged)) }
+
+        admitted = []
+        by_partition.each do |partition_key, jobs|
+          sample_ctx = context.for(jobs.first)
+          rate       = resolve(@rate, sample_ctx).to_f
+          per        = @per.to_f
+          burst      = (resolve(@burst, sample_ctx) || rate).to_f
+
+          bucket = ThrottleBucket.lock(
+            policy_name:   policy.name,
+            gate_name:     name,
+            partition_key: partition_key,
+            burst:         burst
+          )
+          bucket.refill!(rate: rate, per: per, burst: burst)
+
+          jobs.each do |staged|
+            if bucket.consume(1)
+              admitted << [ staged, partition_key ]
+            else
+              break
+            end
+          end
+          bucket.save!
+        end
+
+        context.record_partitions(admitted, gate: name)
+        admitted.map(&:first)
+      end
+    end
+
+    Gate.register(:throttle, Throttle)
+  end
+end

data/lib/dispatch_policy/install_generator.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module DispatchPolicy
+  module Generators
+    class InstallGenerator < Rails::Generators::Base
+      include Rails::Generators::Migration
+
+      source_root File.expand_path("../../db/migrate", __dir__)
+
+      def self.next_migration_number(dirname)
+        ActiveRecord::Generators::Base.next_migration_number(dirname)
+      end
+
+      def copy_migration
+        migration_template "20260424000001_create_dispatch_policy_tables.rb",
+                           "db/migrate/create_dispatch_policy_tables.rb"
+      end
+    end
+  end
+end

data/lib/dispatch_policy/policy.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  class Policy
+    attr_reader :job_class, :gates, :snapshots, :dedupe_key_builder
+
+    def initialize(job_class, &block)
+      @job_class           = job_class
+      @name                = job_class.name.underscore.tr("/", "-")
+      @context_builder     = ->(_args) { {} }
+      @gates               = []
+      @snapshots           = {}
+      @dedupe_key_builder  = nil
+      @round_robin_builder = nil
+      instance_eval(&block) if block
+      DispatchPolicy.registry[@name] = job_class
+    end
+
+    def name(value = nil)
+      return @name if value.nil?
+      DispatchPolicy.registry.delete(@name)
+      @name = value.to_s
+      DispatchPolicy.registry[@name] = @job_class
+    end
+
+    def context(builder)
+      @context_builder = builder
+    end
+
+    def context_builder
+      @context_builder
+    end
+
+    def snapshot(key, builder)
+      @snapshots[key.to_sym] = builder
+    end
+
+    def dedupe_key(builder)
+      @dedupe_key_builder = builder
+    end
+
+    def build_dedupe_key(arguments)
+      return nil unless @dedupe_key_builder
+      key = @dedupe_key_builder.call(arguments)
+      key&.to_s
+    end
+
+    def round_robin_by(builder)
+      @round_robin_builder = builder
+    end
+
+    def round_robin?
+      !@round_robin_builder.nil?
+    end
+
+    def build_round_robin_key(arguments)
+      return nil unless @round_robin_builder
+      key = @round_robin_builder.call(arguments)
+      key.nil? || key.to_s.empty? ? nil : key.to_s
+    end
+
+    def gate(type, **opts)
+      gate_class = DispatchPolicy::Gate.registry.fetch(type.to_sym) do
+        raise ArgumentError, "Unknown gate: #{type}. Known: #{DispatchPolicy::Gate.registry.keys}"
+      end
+      @gates << gate_class.new(policy: self, name: type.to_sym, **opts)
+    end
+
+    def build_snapshot(arguments)
+      @snapshots.transform_values { |builder| builder.call(arguments) }
+    end
+  end
+end