feat: Delayed::Limit (a concurrency limiter)

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    delayed (3.1.0)
+    delayed (3.2.0)
       activerecord (>= 6.0)
       concurrent-ruby
 

README.md

@@ -54,6 +54,7 @@ migration paths where possible.
   * [Priority-based Alerting Threshholds](#priority-based-alerting-threshholds)
   * [Continuous Monitoring](#continuous-monitoring)
 * [Configuration](#configuration)
+* [Rate Limiting Jobs](#rate-limiting-jobs)
 * [Migrating from other ActiveJob backends](#migrating-from-other-activejob-backends)
   * [Migrating from DelayedJob](#migrating-from-delayedjob)
 * [How to Contribute](#how-to-contribute)
@@ -531,6 +532,137 @@ Delayed.logger = Rails.logger
 Delayed.default_log_level = 'info'
 ```
 
+## Rate Limiting Jobs
+
+The `Delayed::Limit` class provides a database-backed **concurrency limiter/optimizer** for jobs
+(via a [Generic Cell Rate Algorithm](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm)
+implemented in SQL+Ruby). Use it to (e.g.) stay under a third-party API's published rate limit, or
+to keep from overwhelming a downstream datastore.
+
+The recommended interface is `with_limit`, provided via the `Delayed::Limitable` module (which is
+included in all `ActiveJob` classes by default):
+
+```ruby
+class TouchesThirdPartyApiJob < ApplicationJob
+  with_limit :third_party_api, max: 100, per: 1.minute
+
+  def perform
+    # ...
+  end
+end
+```
+
+The limiter will then attempt to maximize throughput without exceeding the limit. If the limit would
+be exceeded within a configurable timeout, the job will immediately end and enqueue a retry attempt
+with polynomial backoff. Because its state lives in the database, the limit applies across every
+worker and process at once.
+
+Plain (non-ActiveJob) classes may `include Delayed::Limitable` to use `with_limit`, but they must
+define their own rescheduling/lifecycle behavior for `Delayed::Limit::LimitExceededError` errors.
+
+#### Setup
+
+To use this feature, make sure you have a `delayed_limits` table, or run `rake
+delayed:install:migrations` to add it (see [Database Setup](#database-setup)).
+
+As of now, **only PostgreSQL and SQLite (3.35+) are supported.** The primary SQL query relies on an
+upserting `RETURNING` clause and database-native timestamp arithmetic. You can check the current
+connection at runtime with `Delayed::Limit.supported?`.
+
+#### Traffic Shaping vs Traffic Enforcement
+
+By default, the limiter will `sleep` up to 5 seconds (or a specified `wait_timeout`) before
+yielding to the limited work. (This behavior is subject to the usual GIL and OS scheduling, so
+treat the configured rate as a best-effort target rather than a hard guarantee.)
+
+Use a longer `wait_timeout` for even better throughput smoothing (at the cost of blocking threads):
+
+```ruby
+# A longer wait timeout is best for shaping outbound traffic in asynchronous contexts.
+with_limit :outbound_traffic, max: 100, per: 1.minute, wait_timeout: 30.seconds
+```
+
+Or set it to `0` to fail fast, so that the job never blocks a worker thread:
+
+```ruby
+# A zero wait timeout is best for enforcing inbound limits and shedding excess traffic.
+with_limit :inbound_traffic, max: 5, per: 1.second, wait_timeout: 0
+```
+
+#### Customizing `with_limit`
+
+The purpose defaults to the job's underscored class name, so it may be omitted entirely if the
+limit is not shared with any other class:
+
+```ruby
+with_limit max: 100, per: 1.minute
+```
+
+Use `on:` to wrap one or more other instance methods instead of `perform` (e.g. if only a portion
+of the job's work is subject to the limit):
+
+```ruby
+with_limit :third_party_api, on: :deliver!
+```
+
+**For ActiveJob classes only**, use `retry_attempts:`, `retry_wait:`, and `retry_jitter:` to
+customize the retry behavior. (By default, jobs retry indefinitely with a polynomial backoff, with
+the `wait_timeout` acting as a floor on the computed wait.) If `with_limit` is declared multiple
+times on the same class (e.g. to apply different limits to different methods), only the first
+declaration defines the job's retry behavior.
+
+#### Manually Limiting a Block of Code
+
+To rate limit code that doesn't belong to a job class, call `Delayed::Limit.within_limit`
+directly. It accepts the same `purpose`, `max:`, `per:`, and `wait_timeout:` arguments as
+`with_limit`, and wraps the limited work in a block:
+
+```ruby
+Delayed::Limit.within_limit(:widgets_api, max: 100, per: 1.minute) do
+  WidgetsApi.create_widget!(...)
+end
+```
+
+The key difference is that there is no built-in retry behavior: if the limit would be exceeded
+within the `wait_timeout`, the call raises `Delayed::Limit::LimitExceededError` immediately, and
+it is up to the caller to rescue and/or retry.
+
+#### Shared Limits
+
+To avoid repeating the same purpose's `max` and `per` across multiple call sites, register a
+limit in advance (e.g. in an initializer):
+
+```ruby
+Delayed::Limit.register!(:widgets_api, max: 100, per: 1.minute)
+```
+
+Then, reference it by just its name at each declaration:
+
+```ruby
+with_limit :widgets_api
+
+# or:
+Delayed::Limit.within_limit(:widgets_api) { ... }
+```
+
+Registered limits are cached indefinitely in memory and are not thread-safe on write, so avoid
+registering them dynamically or at runtime!
+
+#### Handling "Burst" Throughput
+
+As of now, **there is no "burst" capacity.** The limiter allows one call per "drain interval" (`per
+/ max`), so a limit of 60-per-minute behaves identically to 1-per-second (and will not allow
+more than 1 call in the first second).
+
+This is generally acceptable for background job processing (and for traffic shaping in general), but
+may be revisited in the future in order to support use cases like API traffic enforcement.
+
+#### Monitoring Limit Usage
+
+Each call emits an `ActiveSupport::Notification` (`delayed.limit.within_limit` or
+`delayed.limit.exceeded`), so you can monitor limiting activity the same way you would any other
+event (see [Monitoring Jobs & Workers](#monitoring-jobs--workers)).
+
 ## Migrating from other ActiveJob backends
 
 For the most part, standard ActiveJob APIs should be fully compatible. However, when migrating from

app/models/delayed/limit.rb

@@ -0,0 +1,184 @@
+# frozen_string_literal: true
+
+module Delayed
+  # A database-backed concurrency limiter/optimizer designed for use with
+  # `Delayed::Job`. Given a 'purpose' (stringy identifier), a limit, and a time
+  # interval, it will attempt to maximize throughput without exceeding the limit
+  # ("traffic shaping"). The `wait_timeout` parameter can also be lowered (from
+  # its default of `5.seconds`) in order to shed load more proactively ("traffic
+  # enforcement").
+  #
+  # Because the algorithm relies on database-specific timestamp arithmetic and
+  # an upserting `RETURNING` clause, only PostgreSQL and SQLite (3.35+) are
+  # supported. (See `.supported?`.)
+  #
+  # Wrap the work you want to limit in a block, and the limiter will either
+  # yield to the block (within the configured `wait_timeout`) or raise
+  # `LimitExceededError` immediately (if the `wait_timeout` would be exceeded):
+  #
+  #   Delayed::Limit.within_limit(:emails, max: 100, per: 1.minute) do
+  #     deliver_email!
+  #   end
+  #
+  # Currently, the limiter does not support "burst" capacity. A limit of 60
+  # req/minute will behave identically to a limit of 1 req/second, because the
+  # limiter only allows a single call per "drain interval" (`per / max`). When
+  # combined with the `wait_timeout`, this degree of smoothing is acceptable for
+  # background job processing (and traffic shaping in general), but may be
+  # revisited in the future to support more types of workloads.
+  #
+  # To share a limit across multiple calls, register a named 'purpose' in
+  # advance (e.g. in an initializer) and then reference it by name later:
+  #
+  #   Delayed::Limit.register!(:emails, max: 100, per: 1.minute)
+  #
+  #   Delayed::Limit.within_limit(:emails, wait_timeout: 10.seconds) do
+  #     deliver_email!
+  #   end
+  #
+  # It is **not** recommended to register limits dynamically at runtime, because
+  # registered limits are cached indefinitely in memory and are not thread-safe.
+  class Limit < ActiveRecord::Base
+    self.table_name = 'delayed_limits'
+
+    # SQLite gained support for the `RETURNING` clause in version 3.35.
+    MINIMUM_SQLITE_VERSION = Gem::Version.new('3.35.0')
+    SECONDS_PER_DAY = 86_400.0
+
+    # Raised when limit adherence would require exceeding the `wait_timeout`.
+    class LimitExceededError < StandardError; end
+
+    class << self
+      # Used only for limits registered in advance (via `.register!`):
+      def limits
+        @limits ||= {}.freeze
+      end
+
+      # The algorithm requires an upserting `RETURNING` clause and timestamp
+      # arithmetic, so only certain database adapters/versions are supported:
+      def supported?
+        case connection.adapter_name
+        when 'PostgreSQL', 'PostGIS'
+          true
+        when 'SQLite'
+          Gem::Version.new(connection.select_value('SELECT sqlite_version()')) >= MINIMUM_SQLITE_VERSION
+        else
+          false
+        end
+      end
+
+      # Register a limit policy for a given 'purpose' (stringy/symbol
+      # identifier). This is optional and should not be used for dynamic purpose
+      # names or limits, for memory and thread-safety reasons.
+      #
+      # Re-registering a purpose is a no-op if the config matches exactly, and
+      # raises ArgumentError otherwise.
+      def register!(purpose, max:, per:)
+        config = { max: max, per: per }.freeze
+
+        if limits.key?(purpose.to_sym) && limits.fetch(purpose.to_sym) != config
+          raise ArgumentError, "Limit policy '#{purpose}' is already registered and does not match #{config.inspect}"
+        end
+
+        @limits = limits.merge(purpose.to_sym => config).freeze
+      end
+
+      # This method implements a leaky bucket algorithm (or, more specifically,
+      # a Generic Cell Rate Algorithm) to enforce a per-'purpose' work limit.
+      #
+      # It will wait up to `wait_timeout` for the caller to come within the
+      # configured limit before yielding to the caller, and will raise
+      # `LimitExceededError` if the wait time would exceed that timeout (shedding
+      # the caller proactively rather than sleeping).
+      #
+      # In Generic Cell Rate Algorithm (GCRA) terms:
+      # - TAT (theoretical arrival time) -> drained_at
+      # - T (emission interval)          -> drain_interval
+      # - t0 (time of request)           -> the database's current time
+      # - τ (bucket capacity)            -> 1 call (implicitly)
+      def within_limit(purpose, max: nil, per: nil, wait_timeout: 5.seconds)
+        config = limits[purpose.to_sym]
+        if config && (max || per)
+          raise ArgumentError, "Limit policy '#{purpose}' is already registered (overriding 'max'/'per' is not supported)"
+        end
+
+        config ||= { max: max, per: per }.compact
+
+        # The drain_interval reflects the per-call rate at which the bucket
+        # empties, calculated as the overall interval (per) divided by the
+        # maximum number of calls allowed in that interval (max).
+        #
+        # e.g. for a target of 100 req/min, the drain_interval would be 0.6 sec/req.
+        drain_interval = config.fetch(:per).seconds / config.fetch(:max).to_d
+
+        # Attempt to reserve capacity via an uncached database query:
+        limit = connection.uncached do
+          find_by_sql(reserve_sql(purpose, drain_interval, wait_timeout.seconds)).first
+        end
+
+        # If 'limit' is nil, it means the WHERE clause prevented us from
+        # reserving capacity in the bucket. (This happens if the configured
+        # `wait_timeout` would be exceeded.) We assume the caller will back off
+        # and retry later, so we avoid wasting bucket capacity on a no-op.
+        if limit.nil?
+          ActiveSupport::Notifications.instrument('delayed.limit.exceeded', purpose: purpose)
+          raise LimitExceededError, "Concurrency limit exceeded for '#{purpose}'"
+        end
+
+        # If we successfully reserved capacity within the `wait_timeout`, it
+        # means that we've been told by the query how long to sleep in order to
+        # comply with the configured rate.
+        #
+        # (For best results, we MUST make a best attempt to sleep for the
+        # returned 'wait' duration before proceeding.)
+        wait = limit.wait.to_f
+        sleep(wait) if wait.positive?
+
+        ActiveSupport::Notifications.instrument('delayed.limit.within_limit', purpose: purpose)
+        yield
+      end
+
+      private
+
+      # We reserve capacity by pushing a 'drained_at' timestamp forward by the
+      # drain_interval, returning how long the caller must wait to avoid filling
+      # the bucket beyond its capacity. (The WHERE clause significantly reduces
+      # contention on the row when the bucket is full.)
+      def reserve_sql(purpose, drain_interval, max_wait)
+        case connection.adapter_name
+        when 'PostgreSQL', 'PostGIS'
+          # Postgres has native `interval` arithmetic and a statement-stable
+          # clock (`statement_timestamp()`), so we bind the intervals as ISO8601
+          # strings and let the database do the math.
+          binds = { purpose: purpose, drain_interval: drain_interval.iso8601, max_wait: max_wait.iso8601 }
+          [<<~SQL.squish, binds]
+            INSERT INTO delayed_limits (purpose, drained_at)
+            VALUES (:purpose, statement_timestamp() + :drain_interval)
+            ON CONFLICT (purpose) DO UPDATE SET
+              drained_at = GREATEST(statement_timestamp(), delayed_limits.drained_at) + :drain_interval
+            WHERE delayed_limits.drained_at - statement_timestamp() <= :max_wait
+            RETURNING EXTRACT(EPOCH FROM (drained_at - statement_timestamp() - :drain_interval)) AS wait
+          SQL
+        when 'SQLite'
+          raise NotImplementedError, "Delayed::Limit requires SQLite #{MINIMUM_SQLITE_VERSION} or newer" unless supported?
+
+          # SQLite has no interval type. Instead, `julianday` yields a
+          # sub-second-precise number of days, and we bind the intervals as a
+          # fraction of a day. (The `wait` is then scaled back to seconds.)
+          binds = { purpose: purpose, drain_interval: drain_interval.to_f / SECONDS_PER_DAY, max_wait: max_wait.to_f / SECONDS_PER_DAY }
+          [<<~SQL.squish, binds]
+            INSERT INTO delayed_limits (purpose, drained_at)
+            VALUES (:purpose, strftime('%Y-%m-%d %H:%M:%f', julianday('now') + :drain_interval))
+            ON CONFLICT (purpose) DO UPDATE SET
+              drained_at = strftime('%Y-%m-%d %H:%M:%f',
+                MAX(julianday('now'), julianday(delayed_limits.drained_at)) + :drain_interval)
+            WHERE julianday(delayed_limits.drained_at) - julianday('now') <= :max_wait
+            RETURNING (julianday(drained_at) - julianday('now') - :drain_interval) * #{SECONDS_PER_DAY} AS wait
+          SQL
+        else
+          raise NotImplementedError, "Delayed::Limit is not supported on #{connection.adapter_name.inspect}"
+        end
+      end
+    end
+  end
+end

db/migrate/10_create_delayed_limits.rb

@@ -0,0 +1,30 @@
+class CreateDelayedLimits < ActiveRecord::Migration[6.0]
+  # The `delayed_limits` table backs the `Delayed::Limit` concurrency limiter.
+  # (See `Delayed::Limit` for details.)
+  #
+  # You can delete this migration if you do not intend to use `Delayed::Limit`,
+  # but it is safe to leave the table in place.
+  def up
+    create_table :delayed_limits, primary_key: :purpose, id: :string do |t|
+      t.datetime :drained_at, null: false
+    end
+
+    return unless connection.adapter_name == 'PostgreSQL'
+
+    # As this is a small, extremely high-churn table, we make it UNLOGGED (the
+    # limiter state need not survive a crash) and tune fillfactor/autovacuum to
+    # favor in-page [HOT updates](https://www.postgresql.org/docs/current/storage-hot.html).
+    execute <<~SQL
+      ALTER TABLE delayed_limits SET UNLOGGED;
+      ALTER TABLE delayed_limits SET (
+        fillfactor = 33,
+        autovacuum_vacuum_scale_factor = 0,
+        autovacuum_vacuum_threshold = 30
+      );
+    SQL
+  end
+
+  def down
+    drop_table :delayed_limits
+  end
+end

lib/delayed.rb

@@ -16,19 +16,22 @@
 require 'delayed/helpers/migration'
 require 'delayed/worker'
 require 'delayed/job_wrapper'
+require 'delayed/limitable'
 
 if defined?(Rails::Engine)
   require 'delayed/engine'
 else
   require 'active_record'
   require_relative '../app/models/delayed/job'
+  require_relative '../app/models/delayed/limit'
 end
 
 ActiveSupport.on_load(:active_job) do
   require 'delayed/active_job_adapter'
   ActiveJob::QueueAdapters::DelayedAdapter = Class.new(Delayed::ActiveJobAdapter)
 
   include Delayed::ActiveJobAdapter::EnqueuingPatch
+  include Delayed::Limitable
 end
 
 ActiveSupport.on_load(:action_mailer) do