dispatch_policy 0.2.0 → 0.3.0

data/lib/dispatch_policy/policy.rb

@@ -2,85 +2,87 @@
 
 module DispatchPolicy
   class Policy
-    attr_reader :job_class, :gates, :snapshots, :dedupe_key_builder
+    DEFAULT_SHARD = "default"
 
-    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
-      @round_robin_weight  = :equal
-      @round_robin_window  = 60
-      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
+    attr_reader :name, :context_proc, :gates, :retry_strategy, :queue_name,
+                :admission_batch_size, :shard_by_proc, :partition_by_proc,
+                :fairness_half_life_seconds, :tick_admission_budget
 
-    def context(builder)
-      @context_builder = builder
-    end
+    def initialize(name:, context_proc:, gates:, retry_strategy: :restage,
+                   queue_name: nil, admission_batch_size: nil, shard_by_proc: nil,
+                   partition_by_proc: nil,
+                   fairness_half_life_seconds: nil, tick_admission_budget: nil)
+      @name                 = name.to_s
+      @context_proc         = context_proc
+      @gates                = gates.freeze
+      @retry_strategy       = retry_strategy
+      @queue_name           = queue_name
+      @admission_batch_size = admission_batch_size
+      @shard_by_proc        = shard_by_proc
+      @partition_by_proc    = partition_by_proc
+      @fairness_half_life_seconds = fairness_half_life_seconds
+      @tick_admission_budget = tick_admission_budget
 
-    def context_builder
-      @context_builder
+      validate!
     end
 
-    def snapshot(key, builder)
-      @snapshots[key.to_sym] = builder
+    # Builds the Context the gates and shard_by will see at admission time.
+    # The user's context_proc receives the job's arguments. The gem then
+    # enriches the resulting hash with `queue_name` (the ActiveJob queue)
+    # so shard_by/partition_by can route by queue without the user having
+    # to thread it through their proc.
+    def build_context(arguments, queue_name: nil)
+      base = context_proc ? context_proc.call(arguments) : {}
+      base = (base || {}).to_h
+      base = base.merge(queue_name: queue_name) if queue_name
+      Context.wrap(base)
     end
 
-    def dedupe_key(builder)
-      @dedupe_key_builder = builder
+    def partition_key_for(ctx)
+      partition_for(ctx)
     end
 
-    def build_dedupe_key(arguments)
-      return nil unless @dedupe_key_builder
-      key = @dedupe_key_builder.call(arguments)
-      key&.to_s
+    # Policy-level partition scope. Both the staged_jobs row and the
+    # concurrency gate's inflight_jobs row use this single canonical
+    # value as their partition_key, so all gates enforce their state
+    # at exactly the same scope. Required: validate! raises if the
+    # policy is built without one.
+    def partition_for(ctx)
+      value = @partition_by_proc.call(ctx)
+      value.nil? ? "" : value.to_s
     end
 
-    def round_robin_by(builder, weight: :equal, window: 60)
-      raise ArgumentError, "weight must be :equal or :time" unless %i[equal time].include?(weight)
-      @round_robin_builder = builder
-      @round_robin_weight  = weight
-      @round_robin_window  = window
-    end
+    # The shard a partition belongs to. Stable per (policy, partition_key)
+    # via first-writer-wins in Repository.upsert_partition!. If no shard_by
+    # is declared the partition lives on the "default" shard.
+    def shard_for(ctx)
+      return DEFAULT_SHARD unless @shard_by_proc
 
-    def round_robin?
-      !@round_robin_builder.nil?
+      value = @shard_by_proc.call(ctx)
+      value.nil? ? DEFAULT_SHARD : value.to_s
     end
 
-    def round_robin_weight
-      @round_robin_weight
+    def restage_retries?
+      retry_strategy == :restage
     end
 
-    def round_robin_window
-      @round_robin_window
+    def bypass_retries?
+      retry_strategy == :bypass
     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
+    private
 
-    def gate(type, **opts)
-      gate_class = DispatchPolicy::Gate.registry.fetch(type.to_sym) do
-        raise ArgumentError, "Unknown gate: #{type}. Known: #{DispatchPolicy::Gate.registry.keys}"
+    def validate!
+      raise InvalidPolicy, "policy name required" if @name.empty?
+      raise InvalidPolicy, "partition_by required" unless @partition_by_proc
+      unless %i[restage bypass].include?(@retry_strategy)
+        raise InvalidPolicy, "retry_strategy must be :restage or :bypass"
       end
-      @gates << gate_class.new(policy: self, name: type.to_sym, **opts)
-    end
-
-    def build_snapshot(arguments)
-      @snapshots.transform_values { |builder| builder.call(arguments) }
+      # Note: gates are NOT required. A policy with no gates uses the
+      # admission_batch_size (or tick_admission_budget when set) as its
+      # only ceiling, with the in-tick fairness reorder distributing
+      # admissions across partitions. Useful for "balance N tenants
+      # fairly without rate-limiting any of them" workloads.
     end
   end
 end

data/lib/dispatch_policy/policy_dsl.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  class PolicyDSL
+    GATE_TYPES = {
+      throttle:             Gates::Throttle,
+      concurrency:          Gates::Concurrency,
+      adaptive_concurrency: Gates::AdaptiveConcurrency
+    }.freeze
+
+    def self.build(name, &block)
+      dsl = new(name)
+      dsl.instance_eval(&block) if block
+      dsl.to_policy
+    end
+
+    def initialize(name)
+      @name                 = name
+      @context_proc         = nil
+      @gates                = []
+      @retry_strategy       = :restage
+      @queue_name           = nil
+      @admission_batch_size = nil
+      @shard_by_proc        = nil
+      @partition_by_proc    = nil
+      @fairness_half_life_seconds = nil
+      @tick_admission_budget = nil
+    end
+
+    def context(callable = nil, &block)
+      @context_proc = callable || block
+    end
+
+    def gate(type, **options)
+      klass = GATE_TYPES[type] || raise(UnknownGate, "unknown gate type: #{type.inspect}")
+      @gates << klass.new(**options)
+    end
+
+    def retry_strategy(strategy)
+      @retry_strategy = strategy
+    end
+
+    def queue_name(name)
+      @queue_name = name.to_s
+    end
+
+    def admission_batch_size(size)
+      @admission_batch_size = Integer(size)
+    end
+
+    # Per-policy override for the EWMA half-life used to weigh recent
+    # admissions when reordering claimed partitions inside the tick.
+    # Accepts a Numeric (seconds) or any object responding to `to_f`
+    # (so ActiveSupport durations like `30.seconds` work too).
+    #   fairness half_life: 30.seconds
+    def fairness(half_life: nil)
+      @fairness_half_life_seconds = Float(half_life) if half_life
+    end
+
+    # Per-policy override for the global tick admission cap. nil
+    # (default) means use config.tick_admission_budget; if that's also
+    # nil, no global cap is enforced and per-partition admission_batch_size
+    # is the only ceiling.
+    def tick_admission_budget(value)
+      @tick_admission_budget = Integer(value)
+    end
+
+    # Defines the partition scope. Required — every policy declares
+    # exactly one. Every gate in the policy uses this proc to compute
+    # the scope it enforces against (the staged_jobs row, the throttle
+    # bucket on that row, and the concurrency gate's inflight rows all
+    # share the same canonical key).
+    #
+    #   dispatch_policy :endpoints do
+    #     partition_by ->(ctx) { ctx[:endpoint_id] }
+    #     gate :throttle,    rate: 60, per: 60
+    #     gate :concurrency, max: 5
+    #   end
+    #
+    # If you need different scopes per gate (e.g. throttle by endpoint
+    # AND concurrency by account), use two policies and let one chain
+    # into the other.
+    def partition_by(callable = nil, &block)
+      @partition_by_proc = callable || block
+    end
+
+    # Routes a partition to a specific shard. The proc receives the
+    # enriched Context (which includes :queue_name from the job) and
+    # returns a string. Tick loops can be scoped per-shard so multiple
+    # workers can process a single policy in parallel.
+    #
+    #   shard_by ->(ctx) { ctx[:queue_name] }                   # shard = job's queue
+    #   shard_by ->(ctx) { "shard-#{ctx[:account_id].hash % 4}" } # explicit hash
+    #
+    # IMPORTANT: shard_by must be CONSISTENT with the gate's
+    # `partition_by` of any rate/concurrency budget you want to enforce
+    # globally. A throttle gate's bucket lives on the partition row, so
+    # if two staged_partitions sharing the same throttle key end up on
+    # different shards, each shard runs its own bucket and the effective
+    # rate becomes rate × N_shards.
+    def shard_by(callable = nil, &block)
+      @shard_by_proc = callable || block
+    end
+
+    def to_policy
+      Policy.new(
+        name:                 @name,
+        context_proc:         @context_proc,
+        gates:                @gates,
+        retry_strategy:       @retry_strategy,
+        queue_name:           @queue_name,
+        admission_batch_size: @admission_batch_size,
+        shard_by_proc:        @shard_by_proc,
+        partition_by_proc:    @partition_by_proc,
+        fairness_half_life_seconds: @fairness_half_life_seconds,
+        tick_admission_budget: @tick_admission_budget
+      )
+    end
+  end
+end

data/lib/dispatch_policy/railtie.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require "rails/railtie"
+
+module DispatchPolicy
+  class Railtie < ::Rails::Railtie
+    initializer "dispatch_policy.active_job" do
+      ActiveSupport.on_load(:active_job) do
+        include DispatchPolicy::JobExtension
+      end
+
+      ActiveSupport.on_load(:active_job) do
+        if defined?(ActiveJob) && ActiveJob.respond_to?(:perform_all_later)
+          singleton = ActiveJob.singleton_class
+          unless singleton.include?(DispatchPolicy::JobExtension::BulkEnqueue)
+            singleton.prepend(DispatchPolicy::JobExtension::BulkEnqueue)
+          end
+        end
+      end
+    end
+
+    # Hosts copy the gem's migration into their own db/migrate via
+    # `rails railties:install:migrations` (or hand-write a cutover
+    # migration like opstasks did). We deliberately do NOT auto-merge
+    # the gem's db/migrate into the host's lookup paths — that
+    # surfaces an `ActiveRecord::DuplicateMigrationNameError` for
+    # any host already carrying a migration named
+    # `CreateDispatchPolicyTables` (e.g. one copied from the
+    # upstream tick-hardening branch during a cutover).
+
+    config.after_initialize do
+      DispatchPolicy.warn_unsupported_adapter
+    end
+  end
+end

data/lib/dispatch_policy/registry.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  class Registry
+    def initialize
+      @mutex    = Mutex.new
+      @policies = {}
+    end
+
+    def register(policy, owner: nil)
+      @mutex.synchronize do
+        existing = @policies[policy.name]
+        if existing && existing[:owner] != owner
+          raise PolicyAlreadyRegistered, "policy #{policy.name.inspect} already registered for #{existing[:owner]}"
+        end
+        @policies[policy.name] = { policy: policy, owner: owner }
+      end
+      policy
+    end
+
+    def fetch(name)
+      entry = @policies[name.to_s]
+      entry && entry[:policy]
+    end
+
+    def [](name)
+      fetch(name)
+    end
+
+    def names
+      @policies.keys
+    end
+
+    def each(&block)
+      @policies.values.map { |e| e[:policy] }.each(&block)
+    end
+
+    def size
+      @policies.size
+    end
+
+    def clear
+      @mutex.synchronize { @policies.clear }
+    end
+  end
+end