dispatch_policy 0.2.0 → 0.4.0

data/lib/dispatch_policy/tick_loop.rb

@@ -1,45 +1,78 @@
 # frozen_string_literal: true
 
 module DispatchPolicy
-  # Shared driver for DispatchTickLoopJob and any foreground tick (e.g. a
-  # rake task). Loops Tick.reap + Tick.run with an interruptible sleep and
-  # bails when stop_when returns true.
-  class TickLoop
-    def self.run(policy_name: nil, sleep_for: nil, sleep_for_busy: nil, stop_when: -> { false })
-      idle_sleep = (sleep_for      || DispatchPolicy.config.tick_sleep).to_f
-      busy_sleep = (sleep_for_busy || DispatchPolicy.config.tick_sleep_busy).to_f
+  # Drives admission until `stop_when` fires (deadline, shutdown signal, etc).
+  # Runs one Tick per policy per loop iteration; sleeps `idle_pause` when no
+  # jobs were admitted across all policies. Periodically (every
+  # `sweep_every_ticks` iterations) sweeps stale inflight rows and inactive
+  # partitions.
+  module TickLoop
+    module_function
+
+    # @param policy_name [String, nil] limit to one policy. nil = all registered.
+    # @param shard [String, nil] limit to one shard. nil = all shards.
+    def run(policy_name: nil, shard: nil, stop_when: -> { false })
+      config       = DispatchPolicy.config
+      logger       = config.logger
+      iteration    = 0
 
       loop do
         break if stop_when.call
 
+        unless DispatchPolicy.config.enabled
+          # Master switch off: stop polling. The job that drives
+          # TickLoop.run will re-schedule itself; we exit cleanly so
+          # the next iteration sees the flag and stops too.
+          logger&.info("[dispatch_policy] TickLoop exiting because config.enabled = false")
+          break
+        end
+
+        names = policy_names(policy_name)
+        if names.empty?
+          sleep(config.idle_pause)
+          next
+        end
+
         admitted = 0
-        begin
-          ActiveRecord::Base.uncached do
-            Tick.reap
-            admitted = Tick.run(policy_name: policy_name).to_i
+        names.each do |name|
+          break if stop_when.call
+
+          begin
+            result = Tick.run(policy_name: name, shard: shard)
+            admitted += result.jobs_admitted
+          rescue StandardError => e
+            logger&.error("[dispatch_policy] tick error policy=#{name} shard=#{shard.inspect} #{e.class}: #{e.message}\n#{e.backtrace.first(10).join("\n")}")
           end
-        rescue StandardError => e
-          Rails.logger&.error("[DispatchPolicy] tick error: #{e.class}: #{e.message}")
-          Rails.error.report(e, handled: true) if defined?(Rails) && Rails.respond_to?(:error)
         end
 
-        break if stop_when.call
+        iteration += 1
+        if (iteration % config.sweep_every_ticks).zero?
+          sweep!
+        end
 
-        interruptible_sleep(admitted.positive? ? busy_sleep : idle_sleep, stop_when)
+        if admitted.zero?
+          sleep(config.idle_pause)
+        elsif config.busy_pause.to_f.positive?
+          sleep(config.busy_pause)
+        end
       end
     end
 
-    def self.interruptible_sleep(total, stop_when)
-      return unless total.positive?
-
-      remaining = total
-      step = 0.1
-      while remaining.positive?
-        break if stop_when.call
-        chunk = [ remaining, step ].min
-        sleep(chunk)
-        remaining -= chunk
+    def policy_names(filter)
+      if filter
+        [filter.to_s]
+      else
+        DispatchPolicy.registry.names
       end
     end
+
+    def sweep!
+      cfg = DispatchPolicy.config
+      Repository.sweep_stale_inflight!(cutoff_seconds: cfg.inflight_stale_after)
+      Repository.sweep_inactive_partitions!(cutoff_seconds: cfg.partition_inactive_after)
+      Repository.sweep_old_tick_samples!(cutoff_seconds: cfg.metrics_retention)
+    rescue StandardError => e
+      DispatchPolicy.config.logger&.error("[dispatch_policy] sweep error: #{e.class}: #{e.message}")
+    end
   end
 end

data/lib/dispatch_policy/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module DispatchPolicy
-  VERSION = "0.2.0"
+  VERSION = "0.4.0"
 end

data/lib/dispatch_policy.rb

@@ -1,70 +1,90 @@
 # frozen_string_literal: true
 
+require "active_support"
+require "active_support/core_ext"
 require "active_job"
-require "active_record"
 
-require "dispatch_policy/version"
-require "dispatch_policy/engine" if defined?(Rails)
+require_relative "dispatch_policy/version"
+require_relative "dispatch_policy/config"
+require_relative "dispatch_policy/context"
+require_relative "dispatch_policy/policy"
+require_relative "dispatch_policy/registry"
+require_relative "dispatch_policy/serializer"
+require_relative "dispatch_policy/bypass"
+require_relative "dispatch_policy/decision"
+require_relative "dispatch_policy/gate"
+require_relative "dispatch_policy/gates/throttle"
+require_relative "dispatch_policy/gates/concurrency"
+require_relative "dispatch_policy/gates/adaptive_concurrency"
+require_relative "dispatch_policy/policy_dsl"
+require_relative "dispatch_policy/cursor_pagination"
+require_relative "dispatch_policy/pipeline"
+require_relative "dispatch_policy/repository"
+require_relative "dispatch_policy/forwarder"
+require_relative "dispatch_policy/inflight_tracker"
+require_relative "dispatch_policy/tick"
+require_relative "dispatch_policy/tick_loop"
+require_relative "dispatch_policy/job_extension"
+require_relative "dispatch_policy/operator_hints"
+require_relative "dispatch_policy/assets"
 
 module DispatchPolicy
-  Config = Struct.new(
-    :enabled,
-    :lease_duration,
-    :batch_size,
-    :round_robin_quantum,
-    :tick_max_duration,
-    :tick_sleep,
-    :tick_sleep_busy,
-    :partition_idle_ttl,
-    :admin_partition_limit,
-    keyword_init: true
-  )
+  class Error < StandardError; end
+  class PolicyAlreadyRegistered < Error; end
+  class UnknownGate < Error; end
+  class InvalidPolicy < Error; end
+  class EnqueueFailed < Error; end
 
-  def self.config
-    @config ||= Config.new(
-      enabled:               true,
-      lease_duration:        15 * 60,          # 15.minutes
-      batch_size:            500,
-      round_robin_quantum:   50,
-      tick_max_duration:     60,               # 1.minute
-      tick_sleep:            1,                # idle sleep
-      tick_sleep_busy:       0.05,             # busy sleep
-      partition_idle_ttl:    30 * 60,          # 30.minutes
-      # Hard cap on rows the admin's partition breakdown will pull per
-      # aggregation. Protects the host DB and process when a policy has
-      # tens of thousands of partitions: the admin shows the top-N most
-      # active and a truncation banner instead of dragging in everything.
-      admin_partition_limit: 5_000
-    )
-  end
+  # Adapters whose enqueue runs against ActiveRecord::Base.connection (so
+  # the adapter INSERT can join the admission TX) or whose semantics make
+  # atomicity moot (test/inline). Substring match against the adapter
+  # class name keeps the check resilient to ActiveJob's wrapper renames.
+  PG_BACKED_ADAPTER_HINTS = %w[GoodJob SolidQueue].freeze
+  EXEMPT_ADAPTER_HINTS    = %w[Test Inline Async].freeze
+
+  module_function
 
-  def self.configure
+  def configure
     yield config
   end
 
-  def self.enabled?
-    config.enabled != false
+  def config
+    @config ||= Config.new
+  end
+
+  def reset_config!
+    @config = Config.new
   end
 
-  # Registry: policy_name => job_class. Populated by Policy#initialize.
-  def self.registry
-    @registry ||= {}
+  def registry
+    @registry ||= Registry.new
   end
 
-  def self.reset_registry!
-    @registry = {}
+  def reset_registry!
+    @registry = Registry.new
+  end
+
+  # Logs a warning if the configured ActiveJob adapter is not one of the
+  # PG-backed ones the gem can guarantee atomic admission for. We do NOT
+  # raise: a host may use a custom PG-backed adapter we don't recognize,
+  # or may have accepted the trade-off knowingly. The warning is enough
+  # to surface the issue at boot.
+  def warn_unsupported_adapter
+    return unless defined?(::ActiveJob::Base)
+    adapter = ::ActiveJob::Base.queue_adapter
+    return unless adapter
+
+    klass_name = adapter.class.name.to_s
+    return if (PG_BACKED_ADAPTER_HINTS + EXEMPT_ADAPTER_HINTS).any? { |hint| klass_name.include?(hint) }
+
+    config.logger&.warn(
+      "[dispatch_policy] active_job adapter is #{klass_name}; atomic admission requires " \
+      "a PG-backed adapter that shares ActiveRecord::Base's connection (good_job, solid_queue). " \
+      "If the worker process crashes between admission COMMIT and adapter enqueue, the job is lost. " \
+      "Set DispatchPolicy.config.database_role if you use a separate DB role for queueing."
+    )
   end
 end
 
-require "dispatch_policy/policy"
-require "dispatch_policy/gate"
-require "dispatch_policy/gates/concurrency"
-require "dispatch_policy/gates/throttle"
-require "dispatch_policy/gates/global_cap"
-require "dispatch_policy/gates/fair_interleave"
-require "dispatch_policy/gates/adaptive_concurrency"
-require "dispatch_policy/dispatch_context"
-require "dispatch_policy/dispatchable"
-require "dispatch_policy/tick"
-require "dispatch_policy/tick_loop"
-require "dispatch_policy/active_job_perform_all_later_patch"
+require_relative "dispatch_policy/railtie" if defined?(Rails::Railtie)
+require_relative "dispatch_policy/engine"  if defined?(Rails::Engine)

data/lib/generators/dispatch_policy/install/install_generator.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/migration"
+
+module DispatchPolicy
+  module Generators
+    class InstallGenerator < ::Rails::Generators::Base
+      include ::Rails::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+      desc "Installs dispatch_policy: migration, initializer, and tick loop job."
+
+      def self.next_migration_number(_path)
+        Time.now.utc.strftime("%Y%m%d%H%M%S")
+      end
+
+      def copy_migration
+        migration_template "create_dispatch_policy_tables.rb.tt",
+                           "db/migrate/create_dispatch_policy_tables.rb"
+      end
+
+      def create_initializer
+        template "initializer.rb.tt", "config/initializers/dispatch_policy.rb"
+      end
+
+      def create_tick_loop_job
+        template "dispatch_tick_loop_job.rb.tt", "app/jobs/dispatch_tick_loop_job.rb"
+      end
+
+      def show_readme
+        readme_text = <<~MSG
+
+          dispatch_policy installed.
+
+          Next steps:
+            1) bin/rails db:migrate
+            2) Mount the engine in config/routes.rb:
+                 mount DispatchPolicy::Engine, at: "/dispatch_policy"
+            3) Schedule DispatchTickLoopJob (cron / good_job recurring / solid_queue recurring)
+               and start it once: DispatchTickLoopJob.perform_later
+            4) Declare a policy in any ActiveJob:
+                 dispatch_policy :name do
+                   context ->(args) { { ... } }
+                   partition_by ->(c) { c[:key] }
+                   gate :throttle, rate: 10, per: 60
+                 end
+
+        MSG
+        say readme_text, :green
+      end
+
+      private
+
+      def good_job?
+        adapter_name == "good_job"
+      end
+
+      def solid_queue?
+        adapter_name == "solid_queue"
+      end
+
+      def adapter_name
+        Rails.application.config.active_job.queue_adapter.to_s
+      rescue StandardError
+        nil
+      end
+    end
+  end
+end

data/lib/generators/dispatch_policy/install/templates/create_dispatch_policy_tables.rb.tt

@@ -0,0 +1,95 @@
+class CreateDispatchPolicyTables < ActiveRecord::Migration[<%= Rails::VERSION::STRING.to_f %>]
+  def change
+    create_table :dispatch_policy_staged_jobs do |t|
+      t.string  :policy_name,   null: false
+      t.string  :partition_key, null: false
+      t.string  :queue_name
+      t.string  :job_class,     null: false
+      t.jsonb   :job_data,      null: false
+      t.datetime :scheduled_at
+      t.integer :priority,      default: 0, null: false
+      t.datetime :enqueued_at,  null: false, default: -> { "now()" }
+      t.jsonb   :context,       null: false, default: {}
+    end
+    add_index :dispatch_policy_staged_jobs,
+              [:policy_name, :partition_key, :scheduled_at, :id],
+              name: "idx_dp_staged_admission",
+              order: { scheduled_at: "ASC NULLS FIRST", id: :asc }
+    add_index :dispatch_policy_staged_jobs, :enqueued_at,
+              name: "idx_dp_staged_enqueued_at"
+
+    create_table :dispatch_policy_partitions do |t|
+      t.string   :policy_name,        null: false
+      t.string   :partition_key,      null: false
+      t.string   :queue_name
+      t.string   :shard,              null: false, default: "default"
+      t.string   :status,             null: false, default: "active"
+      t.integer  :pending_count,      null: false, default: 0
+      t.bigint   :total_admitted,     null: false, default: 0
+      t.jsonb    :context,            null: false, default: {}
+      t.datetime :context_updated_at
+      t.datetime :last_enqueued_at
+      t.datetime :last_checked_at
+      t.datetime :last_admit_at
+      t.datetime :next_eligible_at
+      t.jsonb    :gate_state,         null: false, default: {}
+      t.float    :decayed_admits,     null: false, default: 0.0
+      t.datetime :decayed_admits_at
+      t.timestamps
+    end
+    add_index :dispatch_policy_partitions,
+              [:policy_name, :partition_key],
+              unique: true, name: "idx_dp_partitions_lookup"
+    add_index :dispatch_policy_partitions,
+              [:policy_name, :shard, :status, :next_eligible_at, :last_checked_at],
+              name:  "idx_dp_partitions_tick_order",
+              order: { next_eligible_at: "ASC NULLS FIRST", last_checked_at: "ASC NULLS FIRST" }
+
+    create_table :dispatch_policy_inflight_jobs do |t|
+      t.string   :policy_name,    null: false
+      t.string   :partition_key,  null: false
+      t.string   :active_job_id,  null: false
+      t.datetime :admitted_at,    null: false, default: -> { "now()" }
+      t.datetime :heartbeat_at,   null: false, default: -> { "now()" }
+    end
+    add_index :dispatch_policy_inflight_jobs, :active_job_id, unique: true,
+              name: "idx_dp_inflight_active_job_id"
+    add_index :dispatch_policy_inflight_jobs, [:policy_name, :partition_key],
+              name: "idx_dp_inflight_partition"
+    add_index :dispatch_policy_inflight_jobs, :heartbeat_at,
+              name: "idx_dp_inflight_heartbeat"
+
+    create_table :dispatch_policy_tick_samples do |t|
+      t.string   :policy_name,        null: false
+      t.datetime :sampled_at,         null: false, default: -> { "now()" }
+      t.integer  :duration_ms,        null: false, default: 0
+      t.integer  :partitions_seen,    null: false, default: 0
+      t.integer  :partitions_admitted, null: false, default: 0
+      t.integer  :partitions_denied,  null: false, default: 0
+      t.integer  :jobs_admitted,      null: false, default: 0
+      t.integer  :forward_failures,   null: false, default: 0
+      t.integer  :pending_total,      null: false, default: 0
+      t.integer  :inflight_total,     null: false, default: 0
+      t.jsonb    :denied_reasons,     null: false, default: {}
+    end
+    add_index :dispatch_policy_tick_samples, [:policy_name, :sampled_at],
+              name: "idx_dp_tick_samples_lookup",
+              order: { sampled_at: :desc }
+    add_index :dispatch_policy_tick_samples, :sampled_at,
+              name: "idx_dp_tick_samples_sweep"
+
+    create_table :dispatch_policy_adaptive_concurrency_stats do |t|
+      t.string   :policy_name,      null: false
+      t.string   :partition_key,    null: false
+      t.integer  :current_max,      null: false
+      t.float    :ewma_latency_ms,  null: false, default: 0.0
+      t.integer  :sample_count,     null: false, default: 0
+      t.datetime :last_observed_at
+      t.timestamps
+    end
+    add_index :dispatch_policy_adaptive_concurrency_stats,
+              [:policy_name, :partition_key],
+              unique: true,
+              name:   "idx_dp_adaptive_concurrency_lookup"
+  end
+end

data/lib/generators/dispatch_policy/install/templates/dispatch_tick_loop_job.rb.tt

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+# Run one of these per (policy, shard) pair to parallelise admission.
+# When called without arguments, processes every registered policy across
+# every shard from a single worker — fine for small deployments.
+#
+#   DispatchTickLoopJob.perform_later                # all policies, all shards
+#   DispatchTickLoopJob.perform_later("events")      # one policy, all shards
+#   DispatchTickLoopJob.perform_later("events", "shard-1")  # one policy, one shard
+#
+# When a shard argument is provided, the job is also enqueued onto a queue
+# named after the shard so the same worker pool can run both the tick loop
+# and the admitted jobs (assuming the policy's shard_by returns the queue
+# name your jobs use).
+class DispatchTickLoopJob < ApplicationJob
+  queue_as { arguments[1].presence || :dispatch_loop }
+<% if good_job? -%>
+
+  include GoodJob::ActiveJobExtensions::Concurrency
+  good_job_control_concurrency_with(
+    total_limit: 1,
+    key: -> { "dispatch_tick_loop:#{arguments[0] || 'all'}:#{arguments[1] || 'all'}" }
+  )
+<% elsif solid_queue? -%>
+
+  limits_concurrency to: 1,
+    key: -> { "dispatch_tick_loop:#{arguments[0] || 'all'}:#{arguments[1] || 'all'}" }
+<% end -%>
+
+  def perform(policy_name = nil, shard = nil)
+    deadline = Time.current + DispatchPolicy.config.tick_max_duration
+
+    DispatchPolicy::TickLoop.run(
+      policy_name: policy_name,
+      shard:       shard,
+      stop_when:   -> { adapter_shutting_down? || Time.current >= deadline }
+    )
+
+    self.class.set(wait: 1.second).perform_later(policy_name, shard)
+  end
+
+  private
+
+  def adapter_shutting_down?
+<% if good_job? -%>
+    GoodJob.current_thread_shutting_down?
+<% elsif solid_queue? -%>
+    defined?(SolidQueue::Process) && SolidQueue::Process.current_process&.shutdown?
+<% else -%>
+    false
+<% end -%>
+  end
+end

data/lib/generators/dispatch_policy/install/templates/initializer.rb.tt

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+DispatchPolicy.configure do |c|
+  c.tick_max_duration         = 25      # seconds — total time the tick job stays admitting
+  c.partition_batch_size      = 50      # partitions claimed per tick iteration
+  c.admission_batch_size      = 100     # max jobs admitted per partition per iteration
+  c.idle_pause                = 0.5     # seconds slept when no admissions happened
+  c.partition_inactive_after  = 24 * 60 * 60 # GC partitions idle this long
+  c.inflight_stale_after      = 5 * 60      # GC inflight rows whose worker stopped heartbeating
+  c.sweep_every_ticks         = 50          # how often to run the sweepers
+end