stampeded 0.3.1

data/lib/stampeded/queue_policy.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Stampeded
+  module QueuePolicy
+    module_function
+
+    def set(queue, max_running: nil, max_wait: nil)
+      raise Error, "a queue policy needs a queue name" if queue.to_s.strip.empty?
+
+      Connection.translating_errors do
+        Connection.with do |c|
+          c.exec_params(<<~SQL, [queue.to_s, max_running&.to_i, max_wait&.to_i])
+            INSERT INTO stampeded_queue_policies (queue, max_running, max_wait_secs)
+            VALUES ($1, $2, $3)
+            ON CONFLICT (queue) DO UPDATE SET
+              max_running = EXCLUDED.max_running, max_wait_secs = EXCLUDED.max_wait_secs,
+              updated_at = now()
+          SQL
+        end
+      end
+      queue.to_s
+    end
+
+    def remove(queue)
+      Connection.with { |c| c.exec_params("DELETE FROM stampeded_queue_policies WHERE queue = $1", [queue.to_s]) }
+    end
+
+    def all
+      Connection.translating_errors do
+        Connection.with do |c|
+          c.exec(<<~SQL).to_a
+            SELECT queue, max_running, max_wait_secs, verdict, reason, decided_at
+              FROM stampeded_queue_policies ORDER BY queue
+          SQL
+        end
+      end
+    end
+  end
+end

data/lib/stampeded/railtie.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require "rails/railtie"
+
+module Stampeded
+  class Railtie < Rails::Railtie
+    rake_tasks do
+      load File.expand_path("tasks/stampeded.rake", __dir__)
+    end
+
+    initializer "stampeded.active_job", before: "active_job.set_configs" do
+      ActiveSupport.on_load(:active_job) do
+        require "active_job/queue_adapters/stampeded_adapter"
+        require "stampeded/active_job_ext"
+        require "stampeded/active_job_worker"
+        include Stampeded::ActiveJobExt
+      end
+    end
+
+    initializer "stampeded.config" do
+    end
+  end
+end

data/lib/stampeded/schedule.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Stampeded
+  module Schedule
+    module_function
+
+    def recurring(key:, job:, cron:, args: [], queue: nil, timezone: "UTC",
+                  dst_policy: nil, priority: nil)
+      raise Error, "a schedule needs a key" if key.to_s.strip.empty?
+
+      suggested = check_cron!(cron, timezone)
+      dst_policy ||= suggested
+      job_class = job.is_a?(Class) ? job : job.class
+      queue ||= resolve_queue_name(job_class) || Stampeded.config.default_queue
+
+      task, standard_args, meta = envelope(job_class, args, queue, priority)
+
+      Connection.with do |c|
+        c.exec_params(<<~SQL, [key, queue, task, JSON.generate(standard_args), JSON.generate(meta), cron, timezone, dst_policy, priority])
+          INSERT INTO stampeded_schedules
+            (key, queue, task, args, meta, cron_expr, timezone, dst_policy, priority, next_run_at)
+          VALUES ($1, $2, $3, $4, $5, $6, $7,
+                  $8, coalesce($9, 0), now())
+          ON CONFLICT (key) DO UPDATE SET
+            queue = EXCLUDED.queue, task = EXCLUDED.task, args = EXCLUDED.args, meta = EXCLUDED.meta,
+            cron_expr = EXCLUDED.cron_expr, interval_secs = NULL,
+            timezone = EXCLUDED.timezone, dst_policy = EXCLUDED.dst_policy,
+            priority = EXCLUDED.priority, updated_at = now()
+        SQL
+      end
+      key
+    end
+
+    def remove(key)
+      Connection.with { |c| c.exec_params("DELETE FROM stampeded_schedules WHERE key = $1", [key]) }
+    end
+
+    def pause(key, paused: true)
+      Connection.with do |c|
+        c.exec_params("UPDATE stampeded_schedules SET paused = $2, updated_at = now() WHERE key = $1",
+                      [key, paused])
+      end
+    end
+
+    def all
+      Connection.with do |c|
+        c.exec("SELECT key, queue, task, cron_expr, timezone, dst_policy, paused, next_run_at, last_run_at
+                  FROM stampeded_schedules ORDER BY next_run_at").to_a
+      end
+    end
+
+    def check_cron!(expr, timezone)
+      require "stampeded/native"
+      Native.cron_check(expr, timezone)
+    rescue Error => e
+      raise Error, e.message unless e.message.include?("cron")
+
+      raise Error, <<~MSG
+        #{expr.inspect}: #{e.message}
+
+        Five fields are required: minute hour day-of-month month day-of-week
+
+            "0 3 * * *"        every day at 03:00
+            "*/15 * * * *"     every fifteen minutes
+            "0 9-17 * * mon-fri"
+      MSG
+    end
+
+    def resolve_queue_name(job_class)
+      return nil unless job_class.respond_to?(:queue_name)
+
+      name = job_class.queue_name
+      name = job_class.new.queue_name if name.respond_to?(:call)
+      name&.to_s
+    rescue StandardError
+      nil
+    end
+
+    def envelope(job_class, args, queue, priority)
+      task = job_class.try(:stampeded_task)
+      context = {
+        "job_class" => job_class.name,
+        "queue_name" => queue,
+        "priority" => priority,
+        "locale" => defined?(I18n) ? I18n.locale.to_s : "en",
+        "timezone" => defined?(Time.zone) && Time.zone ? Time.zone.name : "UTC"
+      }
+
+      kwargs = args.is_a?(Hash) ? args : ({} if Array(args).empty?)
+      if kwargs && Envelope.plain?(kwargs)
+        [task || job_class.name, JSON.parse(JSON.generate(kwargs)), { "activejob" => context }]
+      elsif task
+        raise PortableArgumentsError.new(job_class.name, task)
+      else
+        context["arguments"] = ActiveJob::Arguments.serialize(Array(args))
+        [job_class.name, {}, { "activejob" => context }]
+      end
+    end
+  end
+end

data/lib/stampeded/schema.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module Stampeded
+  module Schema
+    FILES = %w[schema.sql].freeze
+
+    module_function
+
+    def root
+      candidates = [
+        File.expand_path("../../schema", __dir__),
+        File.expand_path("../../../schema", __dir__)
+      ]
+      candidates.find { |d| File.directory?(d) } ||
+        raise(Error, "cannot find schema/*.sql (looked in #{candidates.join(", ")})")
+    end
+
+    def sql
+      FILES.map { |f| File.read(File.join(root, f)) }.join("\n")
+    end
+
+    def install!(conn = nil)
+      run = lambda do |c|
+        c.exec("BEGIN")
+        begin
+          c.exec(sql)
+          c.exec("COMMIT")
+        rescue StandardError
+          c.exec("ROLLBACK")
+          raise
+        end
+      end
+
+      conn ? run.call(conn) : Connection.with(&run)
+    end
+
+    def drop!(conn = nil)
+      run = lambda do |c|
+        c.exec(<<~SQL)
+          DROP TABLE IF EXISTS stampeded_cron_occurrences, stampeded_schedules,
+                               stampeded_queue_policies, stampeded_queue_counts,
+                               stampeded_jobs, stampeded_workers CASCADE;
+          DO $$
+          DECLARE r record;
+          BEGIN
+            FOR r IN
+              SELECT p.oid::regprocedure AS sig
+                FROM pg_proc p
+                JOIN pg_namespace n ON n.oid = p.pronamespace
+               WHERE p.proname LIKE 'stampeded\_%'
+                 AND n.nspname = current_schema()
+            LOOP
+              EXECUTE 'DROP FUNCTION IF EXISTS ' || r.sig || ' CASCADE';
+            END LOOP;
+          END $$;
+        SQL
+      end
+      conn ? run.call(conn) : Connection.with(&run)
+    end
+
+    def installed?(conn = nil)
+      check = lambda do |c|
+        Stampeded.truthy?(c.exec("SELECT to_regclass('stampeded_jobs') IS NOT NULL AS ok")[0]["ok"])
+      end
+      conn ? check.call(conn) : Connection.with(&check)
+    end
+  end
+end

data/lib/stampeded/steps.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+require "json"
+require "set"
+
+module Stampeded
+  class Steps
+    class DuplicateStep < Error; end
+    class ResultTooLarge < Error; end
+
+    def initialize(context)
+      require "stampeded/native"
+      @ctx = context
+      @record = Native::Steps.new(context.checkpoint)
+      @seen = Set.new
+    end
+
+    def step(name, cursor: nil, &block)
+      name = name.to_s
+      raise DuplicateStep, "step #{name.inspect} is declared twice in this job" unless @seen.add?(name)
+      return @record.value(name) if @record.completed?(name)
+
+      result = cursor.nil? ? block.call : run_cursor_step(name, cursor, &block)
+      complete(name, result)
+      persist!
+      result
+    end
+
+    def completed?(name) = @record.completed?(name)
+
+    def completed_steps = @record.to_h["completed"]
+
+    private
+
+    def run_cursor_step(name, initial, &block)
+      saved = @record.cursor(name)
+      block.call(Cursor.new(saved.nil? ? initial : saved, self, name))
+    end
+
+    def advance(name, value)
+      @record.set_cursor(name, JSON.generate(value))
+      persist!
+    end
+
+    def complete(name, result)
+      json =
+        begin
+          result.nil? ? nil : JSON.generate(result)
+        rescue JSON::GeneratorError => e
+          raise ResultTooLarge, "step #{name.inspect} returned something not JSON-serialisable: #{e.message}"
+        end
+      @record.complete(name, json)
+    rescue ResultTooLarge
+      raise
+    rescue Error => e
+      raise ResultTooLarge, e.message
+    end
+
+    def persist! = @ctx.checkpoint!(@record.to_h)
+
+    class Cursor
+      attr_reader :value
+
+      alias cursor value
+
+      def initialize(value, steps, name)
+        @value = value
+        @steps = steps
+        @name = name
+        @since_write = 0
+      end
+
+      def advance!(value, every: 1)
+        @value = value
+        @since_write += 1
+        return value if every > 1 && @since_write < every
+
+        @since_write = 0
+        @steps.send(:advance, @name, value)
+        value
+      end
+
+      def flush!
+        return @value if @since_write.zero?
+
+        @since_write = 0
+        @steps.send(:advance, @name, @value)
+        @value
+      end
+
+      def to_s = @value.to_s
+    end
+  end
+end

data/lib/stampeded/tasks/stampeded.rake

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+namespace :stampeded do
+  desc "Install the Stampeded schema (Postgres 18+)"
+  task install: :environment do
+    require "stampeded"
+    Stampeded::Schema.install!
+    puts "schema installed"
+  end
+
+  desc "Run an ActiveJob worker (QUEUES=default,mail or QUEUES=* QUEUES_EXCEPT=dlq; CONCURRENCY=5)"
+  task work: :environment do
+    require "stampeded/active_job_worker"
+    Rails.application.eager_load! unless Rails.application.config.eager_load
+
+    worker = Stampeded::ActiveJobWorker.new(
+      queues: ENV["QUEUES"] || ENV.fetch("QUEUE", nil),
+      except: ENV.fetch("QUEUES_EXCEPT", nil),
+      concurrency: ENV.fetch("CONCURRENCY", Stampeded.config.concurrency).to_i
+    )
+    puts "stampeded worker: queues=#{worker.all_queues? ? "*" : worker.queues.join(",")} pid=#{Process.pid}"
+    worker.run
+  end
+
+  desc "Run the daemon (scheduler, reaper, pruner, dashboard). Not the worker."
+  task daemon: :environment do
+    require "stampeded/daemon"
+    args = ["--database-url", Stampeded.database_url_for_daemon, "--no-worker"]
+    args += ENV["STAMPEDED_ARGS"].to_s.split
+    Stampeded::Daemon.exec!(*args)
+  end
+end

data/lib/stampeded/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Stampeded
+  VERSION = "0.3.1" # x-release-please-version
+end

data/lib/stampeded/worker.rb

@@ -0,0 +1,224 @@
+# frozen_string_literal: true
+
+require "timeout"
+
+module Stampeded
+  class Worker
+    attr_reader :queues, :handlers
+
+    def initialize(queue: nil, queues: nil, except: nil, concurrency: nil, batch_size: nil, poll_interval: nil)
+      requested     = queues || queue || Stampeded.config.queues || Stampeded.config.default_queue
+      @all_queues   = wildcard?(requested)
+      raise ArgumentError, "except: applies only to queues: :all" if except && !@all_queues
+
+      @except       = @all_queues ? Array(except || Stampeded.config.except_queues).flat_map { it.to_s.split(",") }.map(&:strip) : []
+      @queues       = @all_queues ? [] : queue_list!(requested)
+      @concurrency  = positive!(concurrency || Stampeded.config.concurrency, "concurrency")
+      @batch_size   = positive!(batch_size || Stampeded.config.batch_size, "batch_size")
+      @poll_interval = poll_interval || Stampeded.config.poll_interval
+      @handlers     = {}
+      @stopping     = false
+      @wake_reader, @wake_writer = IO.pipe
+    end
+
+    def queue = @queues.first
+
+    def all_queues? = @all_queues
+    def register(task, &block)
+      @handlers[task.to_s] = block
+      self
+    end
+
+    def stop
+      @stopping = true
+      Stampeded.stopping = true
+      @wake_writer.write_nonblock(".", exception: false)
+    rescue IOError
+      nil
+    end
+
+    def stopping? = @stopping
+    def run
+      require "stampeded/native"
+      Native.log_init(library_log_level)
+      Stampeded.stopping = false
+
+      puller = Connection.translating_errors { start_puller }
+      log "registered #{puller.worker_id} on #{description}"
+
+      @threads = Array.new(@concurrency) { worker_thread(puller) }
+      install_signal_handlers
+      @wake_reader.wait_readable until stopping?
+
+      deadline = now_monotonic + Stampeded.config.shutdown_timeout
+      puller.stop(Stampeded.config.shutdown_timeout)
+      stopped = true
+      @threads.each { |t| t.join([deadline - now_monotonic, 0].max) }
+
+      stragglers = @threads.count(&:alive?)
+      if stragglers.positive?
+        log "#{stragglers} job(s) did not finish in #{Stampeded.config.shutdown_timeout}s; " \
+            "the reaper will reclaim them"
+      end
+      log "stopped"
+    ensure
+      @stopping = true
+      begin
+        puller.stop(0) if puller && !stopped
+      rescue Error => e
+        log "could not stop: #{e.message}"
+      end
+      @threads&.each { |t| t.join(@poll_interval * 2) }
+      puller&.close if @threads.nil? || @threads.none?(&:alive?)
+      @wake_reader.close
+      @wake_writer.close
+    end
+
+    private
+
+    def start_puller
+      Native::Puller.new(
+        url: Connection.url,
+        queues: @all_queues ? "*" : @queues.join(","),
+        except: @except.any? ? @except.join(",") : nil,
+        max_leases: @concurrency,
+        batch_size: @batch_size,
+        poll_interval: @poll_interval,
+        heartbeat_interval: Stampeded.config.heartbeat_interval
+      )
+    rescue Error => e
+      raise SchemaMissing, e if e.message.include?("stampeded_") && e.message.include?("does not exist")
+
+      raise
+    end
+
+    def description
+      return @queues.join(", ") unless @all_queues
+
+      "every queue#{" except #{@except.join(", ")}" if @except.any?}"
+    end
+
+    def worker_thread(puller)
+      Thread.new do
+        until stopping?
+          lease =
+            begin
+              puller.next(@poll_interval)
+            rescue Error => e
+              log "claim failed: #{e.message}"
+              sleep @poll_interval
+              next
+            end
+          execute(lease, puller.worker_id) if lease
+        end
+      end
+    end
+
+    def execute(lease, worker_id)
+      job = Job.new(lease.row, decoded: true)
+      started = now_monotonic
+      error = nil
+
+      begin
+        handler = @handlers[job.task]
+        raise UnknownTask, "no handler registered for task #{job.task.inspect}" if handler.nil?
+
+        with_job_timeout(job) { handler.call(JobContext.new(job, worker_id)) }
+      rescue Exception => e
+        error = e
+      end
+
+      report(lease, error)
+      log format("%s %s in %.3fs", job.task, job.id, now_monotonic - started)
+      raise error if error.is_a?(SignalException) || error.is_a?(SystemExit)
+    ensure
+      lease.release!
+    end
+
+    def report(lease, error)
+      return lease.requeue! if error.is_a?(Requeue)
+
+      message = error && error.message.to_s[0, 4000]
+      case error && failure_kind_for(error)
+      when nil then lease.complete!
+      when "unrecoverable" then lease.discard!(message)
+      when "timed_out" then lease.timed_out!(message)
+      else lease.fail!(message, retry_in_ms: error.is_a?(RetryIn) ? error.delay_ms : nil)
+      end
+    rescue Error => e
+      log "could not report #{lease.row["id"]}: #{e.message}"
+    end
+
+    def with_job_timeout(job, &block)
+      limit = Stampeded.config.job_timeout
+      return yield if limit.nil? || limit <= 0
+
+      expire = (job.row["expire_seconds"] || 3600).to_i
+      if limit >= expire
+        log "job_timeout (#{limit}s) is not shorter than #{job.task}'s expire_seconds (#{expire}s): it may run twice"
+      end
+
+      Timeout.timeout(limit, Timeout::Error, "job exceeded job_timeout of #{limit}s", &block)
+    end
+
+    def failure_kind_for(error)
+      case error
+      when UnknownTask then "unrecoverable"
+      when Timeout::Error then "timed_out"
+      else "raised"
+      end
+    end
+
+    def wildcard?(value)
+      return true if value == :all || value.to_s.strip == "*"
+
+      names = value.is_a?(Array) ? value.map(&:to_s) : value.to_s.split(",")
+      raise ArgumentError, "\"*\" cannot be mixed with queue names; use except: to leave queues out" if names.map(&:strip).include?("*")
+
+      false
+    end
+
+    def queue_list!(value)
+      names = value.is_a?(Array) ? value.map(&:to_s) : value.to_s.split(",", -1)
+      names = names.map(&:strip)
+      raise ArgumentError, "queues must name at least one queue, got #{value.inspect}" if names.empty?
+      raise ArgumentError, "queues has a blank entry: #{value.inspect}" if names.any?(&:empty?)
+
+      names.uniq
+    end
+
+    def positive!(value, name)
+      n = value.to_i
+      raise ArgumentError, "#{name} must be at least 1, got #{value.inspect}" if n < 1
+
+      n
+    end
+
+    def install_signal_handlers
+      %w[INT TERM].each do |sig|
+        trap(sig) { stop }
+      rescue ArgumentError
+        nil
+      end
+    end
+
+    def now_monotonic = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+    def log? = $stdout.tty? || ENV.key?("STAMPEDED_LOG")
+
+    LOG_LEVELS = %w[trace debug info warn error].freeze
+
+    def library_log_level
+      level = ENV.fetch("STAMPEDED_LOG", nil)&.downcase
+      return level if LOG_LEVELS.include?(level)
+
+      log? ? "info" : "warn"
+    end
+
+    def log(message)
+      return unless log?
+
+      warn "[stampeded] #{message}"
+    end
+  end
+end

data/lib/stampeded.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+require_relative "stampeded/version"
+require_relative "stampeded/errors"
+require_relative "stampeded/configuration"
+require_relative "stampeded/connection"
+require "forwardable"
+require_relative "stampeded/lifecycle"
+require_relative "stampeded/job"
+require_relative "stampeded/job_context"
+require_relative "stampeded/current"
+require_relative "stampeded/steps"
+require_relative "stampeded/client"
+require_relative "stampeded/envelope"
+require_relative "stampeded/schema"
+require_relative "stampeded/schedule"
+require_relative "stampeded/queue_policy"
+require_relative "stampeded/worker"
+require_relative "stampeded/daemon"
+
+module Stampeded
+  autoload :ActiveJobWorker, File.expand_path("stampeded/active_job_worker", __dir__)
+end
+
+require_relative "stampeded/railtie" if defined?(Rails::Railtie)
+
+module Stampeded
+  def self.truthy?(value)
+    value == true || value == "t" || value == "true"
+  end
+
+  class << self
+    def config
+      @config ||= Configuration.new
+    end
+
+    def configure
+      yield config
+    end
+
+    attr_writer :stopping
+
+    def stopping? = @stopping == true
+
+    def enqueue(queue:, task:, **opts)
+      Client.enqueue(queue: queue, task: task, **opts)
+    end
+
+    def database_url_for_daemon
+      params = Connection.connection_params
+      return params if params.is_a?(String)
+
+      user = [params[:user], params[:password]].compact.join(":")
+      host = [params[:host], params[:port]].compact.join(":")
+      "postgres://#{user}#{"@" unless user.empty?}#{host}/#{params[:dbname]}"
+    end
+
+    def create_queue(name)
+      Connection.with { |c| c.exec_params("SELECT stampeded_create_queue($1)", [name]) }
+      name
+    end
+
+    def depth(queue)
+      Connection.with do |c|
+        c.exec_params(
+          "SELECT state, n FROM stampeded_queue_counts WHERE queue = $1 AND n <> 0 ORDER BY state",
+          [queue]
+        ).each_with_object({}) { |r, h| h[r["state"]] = r["n"].to_i }
+      end
+    end
+  end
+end