dispatch_policy 0.1.0 → 0.3.0

data/lib/dispatch_policy/gates/throttle.rb

@@ -2,51 +2,79 @@
 
 module DispatchPolicy
   module Gates
+    # Token bucket throttle gate.
+    #
+    # Persists state in partitions.gate_state["throttle"] = {
+    #   "tokens"      => Float,   # current tokens, capped at bucket size
+    #   "refilled_at" => Float    # epoch seconds, last refill
+    # }
+    #
+    # The partition scope this gate enforces against is the policy's
+    # `partition_by` (declared in the policy DSL block, not on the gate).
+    # The bucket lives on the staged partition row — one row per
+    # `policy.partition_for(ctx)` value, one bucket per row, no dilution.
     class Throttle < Gate
-      def configure(rate:, per:, burst: nil)
-        @rate  = rate
-        @per   = per
-        @burst = burst
+      attr_reader :rate_proc, :per
+
+      def initialize(rate:, per:)
+        super()
+        @rate_proc = rate.respond_to?(:call) ? rate : ->(_ctx) { rate }
+        @per       = duration_seconds(per)
+        raise ArgumentError, "throttle :per must be > 0 (got #{@per})" unless @per.positive?
       end
 
-      # Consumed tokens refill over time, no release step.
-      def tracks_inflight?
-        false
+      def name
+        :throttle
       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!
+      def evaluate(ctx, partition, admit_budget)
+        capacity = capacity_for(ctx)
+        return Decision.deny(reason: "rate=0") if capacity <= 0
+
+        refill_rate = capacity.to_f / @per
+        state       = (partition["gate_state"] || {})["throttle"] || {}
+        tokens      = (state["tokens"] || capacity).to_f
+        refilled_at = (state["refilled_at"] || now).to_f
+
+        elapsed     = [now - refilled_at, 0.0].max
+        tokens      = [tokens + (elapsed * refill_rate), capacity.to_f].min
+
+        whole       = tokens.floor
+        if whole.zero?
+          missing      = 1.0 - tokens
+          retry_after  = missing / refill_rate
+          patch        = { "tokens" => tokens, "refilled_at" => now }
+          return Decision.new(allowed: 0,
+                              retry_after: retry_after,
+                              gate_state_patch: { "throttle" => patch },
+                              reason: "throttle_empty")
         end
 
-        context.record_partitions(admitted, gate: name)
-        admitted.map(&:first)
+        allowed = [whole, admit_budget].min
+        patch   = { "tokens" => tokens - allowed, "refilled_at" => now }
+        Decision.new(allowed: allowed, gate_state_patch: { "throttle" => patch })
       end
-    end
 
-    Gate.register(:throttle, Throttle)
+      private
+
+      def capacity_for(ctx)
+        value = @rate_proc.call(ctx)
+        value.nil? ? 0 : Integer(value)
+      end
+
+      def now
+        DispatchPolicy.config.now.to_f
+      end
+
+      def duration_seconds(value)
+        if value.is_a?(Numeric)
+          value.to_f
+        elsif value.respond_to?(:to_f) && value.respond_to?(:seconds)
+          value.to_f
+        else
+          raise ArgumentError, "throttle :per must be a numeric duration in seconds (got #{value.inspect})"
+        end
+      end
+    end
   end
 end

data/lib/dispatch_policy/inflight_tracker.rb

@@ -0,0 +1,174 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  # Around-perform that records each job execution in
+  # dispatch_policy_inflight_jobs while it runs, so the concurrency gate
+  # can count active jobs per partition.
+  #
+  # While the job runs we spawn a heartbeat thread that bumps
+  # `heartbeat_at` every `config.inflight_heartbeat_interval` seconds.
+  # Without this, jobs longer than `inflight_stale_after` (default 5 min)
+  # get their inflight row prematurely swept and the concurrency gate
+  # over-admits.
+  module InflightTracker
+    extend ActiveSupport::Concern
+
+    class_methods do
+      def dispatch_policy_inflight_tracking
+        around_perform do |job, block|
+          DispatchPolicy::InflightTracker.track(job, &block)
+        end
+      end
+    end
+
+    def self.track(job)
+      policy_name = job.class.respond_to?(:dispatch_policy_name) && job.class.dispatch_policy_name
+      return yield unless policy_name
+
+      policy = DispatchPolicy.registry.fetch(policy_name)
+      return yield unless policy
+
+      ctx           = policy.build_context(job.arguments, queue_name: job.queue_name&.to_s)
+      partition_key = policy.partition_key_for(ctx)
+
+      Repository.insert_inflight!([{
+        policy_name:    policy.name,
+        partition_key:  partition_key,
+        active_job_id:  job.job_id
+      }])
+
+      adaptive_gates = policy.gates.select { |g| g.name == :adaptive_concurrency }
+      admitted_at    = adaptive_gates.any? ? lookup_admitted_at(job.job_id) : nil
+      perform_start  = Time.current
+
+      heartbeat = start_heartbeat(job.job_id)
+
+      succeeded = false
+      begin
+        yield
+        succeeded = true
+      ensure
+        stop_heartbeat(heartbeat)
+
+        record_adaptive_observations(
+          policy:        policy,
+          gates:         adaptive_gates,
+          partition_key: partition_key,
+          admitted_at:   admitted_at,
+          perform_start: perform_start,
+          succeeded:     succeeded
+        )
+
+        begin
+          Repository.delete_inflight!(active_job_id: job.job_id)
+        rescue StandardError => e
+          DispatchPolicy.config.logger&.warn("[dispatch_policy] failed to delete inflight row #{job.job_id}: #{e.class}: #{e.message}")
+        end
+      end
+    end
+
+    # Reads the admitted_at column from the inflight row that the Tick
+    # pre-inserted. Used as the start-of-queue-wait reference for the
+    # adaptive_concurrency feedback signal (queue_lag = perform_start
+    # - admitted_at). nil if the row vanished or the lookup fails —
+    # the observation is then skipped.
+    def self.lookup_admitted_at(active_job_id)
+      result = ActiveRecord::Base.connection.exec_query(
+        "SELECT admitted_at FROM dispatch_policy_inflight_jobs WHERE active_job_id = $1 LIMIT 1",
+        "lookup_admitted_at",
+        [active_job_id]
+      )
+      row = result.first
+      return nil unless row
+      ts = row["admitted_at"]
+      ts.is_a?(Time) ? ts : Time.parse(ts.to_s)
+    rescue StandardError
+      nil
+    end
+
+    def self.record_adaptive_observations(policy:, gates:, partition_key:, admitted_at:, perform_start:, succeeded:)
+      return if gates.empty?
+
+      queue_lag_ms = if admitted_at
+        ((perform_start - admitted_at) * 1000).to_i
+      else
+        # No admitted_at means we can't measure queue wait. Treat as 0
+        # so the observation still increments sample_count and the
+        # cap can grow if everything else is healthy.
+        0
+      end
+
+      gates.each do |gate|
+        gate.record_observation(
+          policy_name:   policy.name,
+          partition_key: partition_key,
+          queue_lag_ms:  queue_lag_ms,
+          succeeded:     succeeded
+        )
+      rescue StandardError => e
+        DispatchPolicy.config.logger&.warn(
+          "[dispatch_policy] adaptive observation failed for #{policy.name}/#{partition_key}: #{e.class}: #{e.message}"
+        )
+      end
+    end
+
+    # ----- heartbeat thread -----
+
+    HEARTBEAT_KEY = :__dispatch_policy_heartbeat_token__
+
+    Heartbeat = Struct.new(:thread, :stop_flag)
+
+    def self.start_heartbeat(active_job_id)
+      interval = DispatchPolicy.config.inflight_heartbeat_interval.to_f
+      return nil if interval <= 0
+
+      stop_flag = Concurrent::AtomicBoolean.new(false) if defined?(Concurrent::AtomicBoolean)
+      stop_flag ||= ThreadSafeFlag.new
+
+      thread = Thread.new do
+        Thread.current.name = "dispatch_policy.heartbeat:#{active_job_id}"
+
+        until stop_flag.true?
+          # Sleep in small slices so stop is responsive without polling tight.
+          slept = 0.0
+          slice = [interval, 1.0].min
+          while slept < interval && !stop_flag.true?
+            sleep(slice)
+            slept += slice
+          end
+          break if stop_flag.true?
+
+          begin
+            ActiveRecord::Base.connection_pool.with_connection do
+              Repository.heartbeat_inflight!(active_job_id: active_job_id)
+            end
+          rescue StandardError => e
+            DispatchPolicy.config.logger&.warn("[dispatch_policy] heartbeat #{active_job_id} failed: #{e.class}: #{e.message}")
+          end
+        end
+      end
+
+      Heartbeat.new(thread, stop_flag)
+    end
+
+    def self.stop_heartbeat(heartbeat)
+      return if heartbeat.nil?
+
+      heartbeat.stop_flag.make_true
+      # Wake the thread out of any in-progress sleep so we don't wait the full slice.
+      heartbeat.thread.wakeup if heartbeat.thread.alive?
+      heartbeat.thread.join(1.0)
+    rescue StandardError
+      # Worst case: the thread is killed by GC; the inflight row gets a stale
+      # heartbeat_at and the sweeper will reclaim it after inflight_stale_after.
+    end
+
+    # Tiny fallback if concurrent-ruby isn't available (it's a Rails dep
+    # via active_support so it normally is).
+    class ThreadSafeFlag
+      def initialize; @mutex = Mutex.new; @value = false; end
+      def true?; @mutex.synchronize { @value }; end
+      def make_true; @mutex.synchronize { @value = true }; end
+    end
+  end
+end

data/lib/dispatch_policy/job_extension.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  # Hooks into ActiveJob::Base. Adds:
+  #   - the `dispatch_policy :name do … end` class macro
+  #   - an `around_enqueue` callback that stages jobs declaring a policy
+  #   - a `perform_all_later` patch that handles bulk enqueue
+  module JobExtension
+    extend ActiveSupport::Concern
+
+    included do
+      class_attribute :dispatch_policy_name, instance_writer: false
+    end
+
+    class_methods do
+      def dispatch_policy(name, &block)
+        policy = PolicyDSL.build(name.to_s, &block)
+        DispatchPolicy.registry.register(policy, owner: self.name)
+        self.dispatch_policy_name = policy.name
+
+        around_enqueue do |job, block|
+          DispatchPolicy::JobExtension.around_enqueue_for(job, block)
+        end
+      end
+    end
+
+    # Called by the around_enqueue lambda. Public so it can be tested directly.
+    def self.around_enqueue_for(job, block)
+      return block.call if Bypass.active?
+      return block.call unless DispatchPolicy.config.enabled
+
+      policy = DispatchPolicy.registry.fetch(job.class.dispatch_policy_name)
+      return block.call unless policy
+
+      if retry_attempt?(job) && policy.bypass_retries?
+        return block.call
+      end
+
+      # `klass.deserialize(payload)` (used elsewhere — see Forwarder, retries)
+      # only sets @serialized_arguments. ActiveJob defers the actual
+      # arguments deserialization to perform_now via the private
+      # deserialize_arguments_if_needed. If something deserializes a job
+      # and re-enqueues it without going through perform_now (e.g. a
+      # custom retry path), `job.arguments` would be []. Guard against
+      # that here so the context proc always sees the real args.
+      ensure_arguments_materialized!(job)
+
+      queue_name    = job.queue_name&.to_s || policy.queue_name
+      ctx           = policy.build_context(job.arguments, queue_name: queue_name)
+      partition_key = policy.partition_key_for(ctx)
+      shard         = policy.shard_for(ctx)
+      payload       = Serializer.serialize(job)
+
+      Repository.stage!(
+        policy_name:   policy.name,
+        partition_key: partition_key,
+        queue_name:    queue_name,
+        shard:         shard,
+        job_class:     job.class.name,
+        job_data:      payload,
+        context:       ctx.to_jsonb,
+        scheduled_at:  scheduled_time(job),
+        priority:      job.priority || 0
+      )
+
+      job.successfully_enqueued = true
+      false # halts the around_enqueue chain so the real adapter never sees the job
+    end
+
+    def self.retry_attempt?(job)
+      (job.respond_to?(:executions) ? job.executions.to_i : 0).positive?
+    end
+
+    def self.scheduled_time(job)
+      ts = job.scheduled_at
+      return nil if ts.nil?
+      return ts   if ts.is_a?(Time)
+
+      Time.at(Float(ts))
+    rescue ArgumentError, TypeError
+      nil
+    end
+
+    # ActiveJob's `arguments` getter is a plain attr_accessor that returns
+    # the in-memory @arguments. After `klass.deserialize(payload)`, that
+    # array is empty until perform_now triggers
+    # `deserialize_arguments_if_needed` (a private method). Anywhere we
+    # read `job.arguments` outside of perform we must materialize first,
+    # or the context proc receives [] and falls back to its defaults.
+    def self.ensure_arguments_materialized!(job)
+      return unless job.respond_to?(:deserialize_arguments_if_needed, true)
+      job.send(:deserialize_arguments_if_needed)
+    end
+
+    # ---- perform_all_later support -------------------------------------------
+
+    # Rails 7.1+ exposes ActiveJob.perform_all_later. We override it to route
+    # jobs declaring a dispatch_policy through a single bulk INSERT, while
+    # delegating jobs without a policy to the original enqueue_all path.
+    module BulkEnqueue
+      def perform_all_later(*jobs)
+        flat = jobs.flatten
+        return super if flat.empty?
+        # Critical: respect Bypass exactly like the per-job around_enqueue
+        # does. Forwarder.dispatch deserializes admitted jobs and calls
+        # ActiveJob.perform_all_later under Bypass.with — without this
+        # check, BulkEnqueue would re-stage them, creating an infinite
+        # admission loop with the wrong context (job.arguments is still []
+        # at that point because ActiveJob defers deserialization).
+        return super if DispatchPolicy::Bypass.active?
+        return super unless DispatchPolicy.config.enabled
+        return super unless DispatchPolicy.registry.size.positive?
+
+        with_policy, without_policy = flat.partition do |j|
+          j.class.respond_to?(:dispatch_policy_name) && j.class.dispatch_policy_name
+        end
+
+        super(without_policy) if without_policy.any?
+
+        return nil if with_policy.empty?
+
+        rows = with_policy.filter_map do |job|
+          policy = DispatchPolicy.registry.fetch(job.class.dispatch_policy_name)
+          next unless policy
+
+          # See JobExtension.ensure_arguments_materialized! — we need this
+          # for the same reason as the single-enqueue path.
+          JobExtension.ensure_arguments_materialized!(job)
+
+          queue_name    = job.queue_name&.to_s || policy.queue_name
+          ctx           = policy.build_context(job.arguments, queue_name: queue_name)
+          partition_key = policy.partition_key_for(ctx)
+          shard         = policy.shard_for(ctx)
+          payload       = Serializer.serialize(job)
+          job.successfully_enqueued = true
+
+          {
+            policy_name:   policy.name,
+            partition_key: partition_key,
+            queue_name:    queue_name,
+            shard:         shard,
+            job_class:     job.class.name,
+            job_data:      payload,
+            context:       ctx.to_jsonb,
+            scheduled_at:  JobExtension.scheduled_time(job),
+            priority:      job.priority || 0
+          }
+        end
+
+        Repository.stage_many!(rows) if rows.any?
+        nil # ActiveJob.perform_all_later contract returns nil
+      end
+    end
+  end
+end

data/lib/dispatch_policy/operator_hints.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  # Pure-Ruby hint generator. Takes a snapshot of live metrics and
+  # returns a list of {level:, message:} that the dashboard renders.
+  #
+  # Each predicate is intentionally conservative: hints fire on
+  # crossings the operator can fix from the UI or by toggling a config
+  # value, not on noise. Levels:
+  #
+  #   :info     — observation worth glancing at
+  #   :warn     — attention soon
+  #   :critical — fix now
+  module OperatorHints
+    Hint = Struct.new(:level, :message, keyword_init: true)
+
+    module_function
+
+    # `metrics` is a hash of:
+    #   tick_max_duration_ms:   int  (config tick_max_duration × 1000)
+    #   avg_tick_ms:            int
+    #   max_tick_ms:            int
+    #   pending_total:          int
+    #   admitted_per_minute:    int  (last 1m)
+    #   forward_failures:       int  (last 1m)
+    #   jobs_admitted:          int  (last 1m, denominator for fail %)
+    #   active_partitions:      int
+    #   never_checked:          int
+    #   in_backoff:             int
+    #   total_partitions:       int
+    #   adapter_target_jps:     int|nil  (config.adapter_throughput_target)
+    def for(metrics)
+      hints = []
+      m     = metrics
+
+      # ---- tick approaching deadline ---------------------------------
+      if m[:tick_max_duration_ms].to_i.positive? && m[:avg_tick_ms].to_i.positive?
+        ratio = m[:avg_tick_ms].to_f / m[:tick_max_duration_ms]
+        if ratio >= 0.6
+          hints << Hint.new(
+            level: ratio >= 0.85 ? :critical : :warn,
+            message: "Avg tick is #{format('%.0f%%', ratio * 100)} of tick_max_duration. " \
+                     "Lower admission_batch_size, set tick_admission_budget, or shard the policy."
+          )
+        end
+      end
+
+      # ---- backlog drain time ----------------------------------------
+      if m[:admitted_per_minute].to_i.positive? && m[:pending_total].to_i.positive?
+        drain_minutes = m[:pending_total].to_f / m[:admitted_per_minute]
+        if drain_minutes >= 30
+          level = drain_minutes >= 120 ? :warn : :info
+          hints << Hint.new(
+            level: level,
+            message: "At #{m[:admitted_per_minute]} admits/min, the current backlog of " \
+                     "#{m[:pending_total]} pending would take ~#{drain_minutes.round} min " \
+                     "to drain. Raise admission_batch_size, raise the gate's rate, or shard."
+          )
+        end
+      end
+
+      # ---- pending growing while admit rate is non-trivial -----------
+      # `pending_trend` compares head/tail thirds of the sparkline; a
+      # transient spike that already drained still leaves the tail
+      # average elevated. Gate on current pending > 0 so a recovered
+      # backlog does not raise a warning.
+      if m[:pending_trend] == :up &&
+         m[:admitted_per_minute].to_i.positive? &&
+         m[:pending_total].to_i.positive?
+        hints << Hint.new(
+          level: :warn,
+          message: "Pending is growing while we are admitting. Inflow > outflow — " \
+                   "either the throttle rate is below the producer rate, or the worker pool can't keep up."
+        )
+      end
+
+      # ---- forward failure rate --------------------------------------
+      if m[:jobs_admitted].to_i.positive?
+        rate = m[:forward_failures].to_f / m[:jobs_admitted]
+        if rate >= 0.01
+          hints << Hint.new(
+            level: rate >= 0.05 ? :critical : :warn,
+            message: "Forward failures at #{format('%.1f%%', rate * 100)} (#{m[:forward_failures]} / " \
+                     "#{m[:jobs_admitted]} admits). Inspect logs — adapter is rejecting enqueues."
+          )
+        end
+      end
+
+      # ---- never_checked > 0 with cardinality > batch ----------------
+      if m[:never_checked].to_i.positive?
+        hints << Hint.new(
+          level: :warn,
+          message: "#{m[:never_checked]} active partitions have never been checked. " \
+                   "The tick is not getting through them — increase partition_batch_size or shard."
+        )
+      end
+
+      # ---- partition cardinality -------------------------------------
+      if m[:total_partitions].to_i >= 50_000
+        hints << Hint.new(
+          level: :info,
+          message: "#{m[:total_partitions]} partitions in DB. claim_partitions starts to feel " \
+                   "this around 50k–100k. Consider lowering partition_inactive_after to GC " \
+                   "drained ones earlier."
+        )
+      end
+
+      # ---- adapter ceiling proximity --------------------------------
+      target_jps = m[:adapter_target_jps].to_i
+      if target_jps.positive? && m[:admitted_per_minute].to_i.positive?
+        current_jps = m[:admitted_per_minute] / 60.0
+        ratio = current_jps / target_jps
+        if ratio >= 0.7
+          hints << Hint.new(
+            level: ratio >= 0.95 ? :critical : :warn,
+            message: "Admitting #{format('%.0f', current_jps)} jobs/sec, " \
+                     "#{format('%.0f%%', ratio * 100)} of the configured adapter ceiling " \
+                     "(#{target_jps}/sec). Consider an additional shard before the next traffic spike."
+          )
+        end
+      end
+
+      hints
+    end
+  end
+end

data/lib/dispatch_policy/pipeline.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module DispatchPolicy
+  # Composes a sequence of gates into a single admission decision for one
+  # partition. Returns a value object describing how many jobs may be
+  # admitted right now and which gate-state patches to persist.
+  class Pipeline
+    Result = Struct.new(:admit_count, :retry_after, :gate_state_patch, :reasons, keyword_init: true)
+
+    def initialize(policy)
+      @policy = policy
+    end
+
+    def call(ctx, partition, max_budget)
+      budget          = max_budget
+      retry_after     = nil
+      patch           = {}
+      reasons         = []
+      decisions       = []
+
+      @policy.gates.each do |gate|
+        decision = gate.evaluate(ctx, partition, budget)
+        decisions << [gate, decision]
+        budget   = decision.allowed.finite? ? [budget, decision.allowed].min : budget
+        if decision.retry_after
+          retry_after = retry_after.nil? ? decision.retry_after : [retry_after, decision.retry_after].min
+        end
+        reasons << "#{gate.name}:#{decision.reason}" if decision.reason
+        break if budget.zero?
+      end
+
+      admit_count = budget.finite? ? budget : max_budget
+      admit_count = 0 if admit_count.negative?
+
+      decisions.each do |_, decision|
+        next unless decision.gate_state_patch
+        patch.merge!(decision.gate_state_patch)
+      end
+
+      Result.new(
+        admit_count:       admit_count,
+        retry_after:       retry_after,
+        gate_state_patch:  patch,
+        reasons:           reasons
+      )
+    end
+  end
+end