dispatch_policy 0.1.0 → 0.3.0

data/app/models/dispatch_policy/partition_inflight_count.rb

@@ -1,42 +0,0 @@
-# frozen_string_literal: true
-
-module DispatchPolicy
-  class PartitionInflightCount < ApplicationRecord
-    self.table_name = "dispatch_policy_partition_counts"
-
-    def self.fetch_many(policy_name:, gate_name:, partition_keys:)
-      return {} if partition_keys.empty?
-
-      where(policy_name: policy_name, gate_name: gate_name.to_s, partition_key: partition_keys)
-        .pluck(:partition_key, :in_flight).to_h
-        .tap { |h| partition_keys.each { |k| h[k] ||= 0 } }
-    end
-
-    def self.total_for(policy_name:, gate_name:)
-      where(policy_name: policy_name, gate_name: gate_name.to_s).sum(:in_flight)
-    end
-
-    def self.increment(policy_name:, gate_name:, partition_key:, by: 1)
-      now = Time.current
-      sql = <<~SQL.squish
-        INSERT INTO #{quoted_table_name}
-          (policy_name, gate_name, partition_key, in_flight, created_at, updated_at)
-        VALUES (?, ?, ?, ?, ?, ?)
-        ON CONFLICT (policy_name, gate_name, partition_key)
-        DO UPDATE SET
-          in_flight  = #{quoted_table_name}.in_flight + EXCLUDED.in_flight,
-          updated_at = EXCLUDED.updated_at
-      SQL
-      connection.exec_update(
-        sanitize_sql_array([ sql, policy_name, gate_name.to_s, partition_key.to_s, by, now, now ])
-      )
-    end
-
-    def self.decrement(policy_name:, gate_name:, partition_key:, by: 1)
-      where(policy_name: policy_name, gate_name: gate_name.to_s, partition_key: partition_key.to_s)
-        .update_all([
-          "in_flight = GREATEST(in_flight - ?, 0), updated_at = ?", by, Time.current
-        ])
-    end
-  end
-end

data/app/models/dispatch_policy/partition_observation.rb

@@ -1,49 +0,0 @@
-# frozen_string_literal: true
-
-module DispatchPolicy
-  # Minute-bucketed observability per (policy, partition). Any gate with
-  # partition_by gets an observation row here — adaptive, throttle,
-  # concurrency, whatever — so the admin chart shows queue lag / throughput
-  # for all partitioned policies, not just the adaptive ones.
-  #
-  # One row per (policy, partition, minute): total_lag_ms accumulates the
-  # sum of queue_lag_ms observations in that minute, observation_count
-  # increments, max_lag_ms tracks the worst spike. Average lag for the
-  # bucket is derived on read as total / count.
-  class PartitionObservation < ApplicationRecord
-    self.table_name = "dispatch_policy_partition_observations"
-
-    OBSERVATION_TTL = 2 * 60 * 60  # 2 hours
-
-    def self.observe!(policy_name:, partition_key:, queue_lag_ms:, current_max: nil)
-      return if partition_key.nil? || partition_key.to_s.empty?
-
-      now = Time.current
-      lag = queue_lag_ms.to_i
-      sql = <<~SQL.squish
-        INSERT INTO #{quoted_table_name}
-          (policy_name, partition_key, minute_bucket,
-           total_lag_ms, observation_count, max_lag_ms, current_max,
-           created_at, updated_at)
-        VALUES (?, ?, date_trunc('minute', ?::timestamp), ?, 1, ?, ?, ?, ?)
-        ON CONFLICT (policy_name, partition_key, minute_bucket)
-        DO UPDATE SET
-          total_lag_ms      = #{quoted_table_name}.total_lag_ms + EXCLUDED.total_lag_ms,
-          observation_count = #{quoted_table_name}.observation_count + 1,
-          max_lag_ms        = GREATEST(#{quoted_table_name}.max_lag_ms, EXCLUDED.max_lag_ms),
-          current_max       = COALESCE(EXCLUDED.current_max, #{quoted_table_name}.current_max),
-          updated_at        = EXCLUDED.updated_at
-      SQL
-      connection.exec_update(
-        sanitize_sql_array([
-          sql, policy_name, partition_key.to_s, now,
-          lag, lag, current_max, now, now
-        ])
-      )
-    end
-
-    def self.prune!
-      where("minute_bucket < ?", Time.current - OBSERVATION_TTL).delete_all
-    end
-  end
-end

data/app/models/dispatch_policy/throttle_bucket.rb

@@ -1,41 +0,0 @@
-# frozen_string_literal: true
-
-module DispatchPolicy
-  class ThrottleBucket < ApplicationRecord
-    self.table_name = "dispatch_policy_throttle_buckets"
-
-    def self.lock(policy_name:, gate_name:, partition_key:, burst:)
-      now = Time.current
-      seed_sql = <<~SQL.squish
-        INSERT INTO #{quoted_table_name}
-          (policy_name, gate_name, partition_key, tokens, refilled_at, created_at, updated_at)
-        VALUES (?, ?, ?, ?, ?, ?, ?)
-        ON CONFLICT (policy_name, gate_name, partition_key) DO NOTHING
-      SQL
-      connection.exec_update(
-        sanitize_sql_array([
-          seed_sql, policy_name, gate_name.to_s, partition_key.to_s,
-          burst.to_f, now, now, now
-        ])
-      )
-
-      where(policy_name: policy_name, gate_name: gate_name.to_s, partition_key: partition_key.to_s)
-        .lock("FOR UPDATE")
-        .first!
-    end
-
-    def refill!(rate:, per:, burst:)
-      now = Time.current
-      elapsed = (now - refilled_at).to_f
-      new_tokens = tokens + (rate * elapsed / per)
-      self.tokens = [ new_tokens, burst.to_f ].min
-      self.refilled_at = now
-    end
-
-    def consume(n = 1)
-      return false if tokens < n
-      self.tokens -= n
-      true
-    end
-  end
-end

data/db/migrate/20260424000001_create_dispatch_policy_tables.rb

@@ -1,80 +0,0 @@
-# frozen_string_literal: true
-
-class CreateDispatchPolicyTables < ActiveRecord::Migration[7.1]
-  def change
-    create_table :dispatch_policy_staged_jobs do |t|
-      t.string   :job_class,       null: false
-      t.string   :policy_name,     null: false
-      t.jsonb    :arguments,       null: false
-      t.jsonb    :snapshot,        null: false, default: {}
-      t.jsonb    :context,         null: false, default: {}
-      t.integer  :priority,        null: false, default: 100
-      t.datetime :not_before_at
-      t.datetime :staged_at,       null: false
-      t.datetime :admitted_at
-      t.datetime :completed_at
-      t.datetime :lease_expires_at
-      t.string   :active_job_id
-      t.string   :dedupe_key
-      t.string   :round_robin_key
-      t.jsonb    :partitions,      null: false, default: {}
-
-      t.timestamps
-    end
-
-    add_index :dispatch_policy_staged_jobs,
-      %i[policy_name priority staged_at],
-      where: "admitted_at IS NULL",
-      name: "idx_dp_staged_dispatch_order"
-
-    add_index :dispatch_policy_staged_jobs,
-      %i[policy_name dedupe_key],
-      unique: true,
-      where: "dedupe_key IS NOT NULL AND completed_at IS NULL",
-      name: "idx_dp_staged_dedupe_active"
-
-    add_index :dispatch_policy_staged_jobs,
-      %i[lease_expires_at],
-      where: "admitted_at IS NOT NULL",
-      name: "idx_dp_staged_lease_expires"
-
-    add_index :dispatch_policy_staged_jobs,
-      %i[completed_at],
-      where: "completed_at IS NOT NULL",
-      name: "idx_dp_staged_completed_at"
-
-    add_index :dispatch_policy_staged_jobs,
-      %i[policy_name round_robin_key priority staged_at],
-      where: "admitted_at IS NULL AND round_robin_key IS NOT NULL",
-      name: "idx_dp_staged_round_robin"
-
-    create_table :dispatch_policy_partition_counts do |t|
-      t.string  :policy_name,   null: false
-      t.string  :gate_name,     null: false
-      t.string  :partition_key, null: false, default: "default"
-      t.integer :in_flight,     null: false, default: 0
-
-      t.timestamps
-    end
-
-    add_index :dispatch_policy_partition_counts,
-      %i[policy_name gate_name partition_key],
-      unique: true,
-      name: "idx_dp_partition_counts_unique"
-
-    create_table :dispatch_policy_throttle_buckets do |t|
-      t.string   :policy_name,   null: false
-      t.string   :gate_name,     null: false
-      t.string   :partition_key, null: false, default: "default"
-      t.float    :tokens,        null: false
-      t.datetime :refilled_at,   null: false
-
-      t.timestamps
-    end
-
-    add_index :dispatch_policy_throttle_buckets,
-      %i[policy_name gate_name partition_key],
-      unique: true,
-      name: "idx_dp_throttle_buckets_unique"
-  end
-end

data/db/migrate/20260424000002_create_adaptive_concurrency_stats.rb

@@ -1,22 +0,0 @@
-# frozen_string_literal: true
-
-class CreateAdaptiveConcurrencyStats < ActiveRecord::Migration[7.1]
-  def change
-    create_table :dispatch_policy_adaptive_concurrency_stats do |t|
-      t.string   :policy_name,     null: false
-      t.string   :gate_name,       null: false
-      t.string   :partition_key,   null: false, default: "default"
-      t.integer  :current_max,     null: false
-      t.float    :ewma_latency_ms, null: false, default: 0
-      t.integer  :sample_count,    null: false, default: 0
-      t.datetime :last_observed_at
-
-      t.timestamps
-    end
-
-    add_index :dispatch_policy_adaptive_concurrency_stats,
-      %i[policy_name gate_name partition_key],
-      unique: true,
-      name: "idx_dp_adaptive_concurrency_stats_unique"
-  end
-end

data/db/migrate/20260424000003_create_adaptive_concurrency_samples.rb

@@ -1,25 +0,0 @@
-# frozen_string_literal: true
-
-class CreateAdaptiveConcurrencySamples < ActiveRecord::Migration[7.1]
-  def change
-    create_table :dispatch_policy_adaptive_concurrency_samples do |t|
-      t.string   :policy_name,     null: false
-      t.string   :gate_name,       null: false
-      t.string   :partition_key,   null: false
-      t.datetime :minute_bucket,   null: false
-      t.float    :ewma_latency_ms, null: false, default: 0
-      t.integer  :current_max,     null: false
-
-      t.timestamps
-    end
-
-    add_index :dispatch_policy_adaptive_concurrency_samples,
-      %i[policy_name gate_name partition_key minute_bucket],
-      unique: true,
-      name: "idx_dp_adaptive_concurrency_samples_unique"
-
-    add_index :dispatch_policy_adaptive_concurrency_samples,
-      :minute_bucket,
-      name: "idx_dp_adaptive_concurrency_samples_time"
-  end
-end

data/db/migrate/20260424000004_rename_samples_to_partition_observations.rb

@@ -1,32 +0,0 @@
-# frozen_string_literal: true
-
-class RenameSamplesToPartitionObservations < ActiveRecord::Migration[7.1]
-  def up
-    drop_table :dispatch_policy_adaptive_concurrency_samples, if_exists: true
-
-    create_table :dispatch_policy_partition_observations do |t|
-      t.string   :policy_name,       null: false
-      t.string   :partition_key,     null: false
-      t.datetime :minute_bucket,     null: false
-      t.bigint   :total_lag_ms,      null: false, default: 0
-      t.integer  :observation_count, null: false, default: 0
-      t.integer  :max_lag_ms,        null: false, default: 0
-      t.integer  :current_max
-
-      t.timestamps
-    end
-
-    add_index :dispatch_policy_partition_observations,
-      %i[policy_name partition_key minute_bucket],
-      unique: true,
-      name: "idx_dp_partition_observations_unique"
-
-    add_index :dispatch_policy_partition_observations,
-      :minute_bucket,
-      name: "idx_dp_partition_observations_time"
-  end
-
-  def down
-    drop_table :dispatch_policy_partition_observations
-  end
-end

data/lib/dispatch_policy/active_job_perform_all_later_patch.rb

@@ -1,32 +0,0 @@
-# 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

@@ -1,53 +0,0 @@
-# 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

@@ -1,120 +0,0 @@
-# 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/gates/fair_interleave.rb

@@ -1,32 +0,0 @@
-# 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

@@ -1,26 +0,0 @@
-# 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/install_generator.rb

@@ -1,23 +0,0 @@
-# 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