[ruby] Optimizations for massive improvement in SQL statements and CPU use

ruby/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Reduce job processing overhead in both drivers by returning claimed and finished jobs directly from updates and locking Postgres completion targets by primary key. Avoid per-job connection resets in Active Record, and keep job completions and insertions from interrupting maintenance error backoff. Thanks to [@jpcamara](https://github.com/jpcamara) for the improvements! [PR #1566](https://github.com/riverqueue/river/pull/1566).
+
 ## [0.15.0] - 2026-10-30
 
 ### Added

ruby/bench/drain.rb

@@ -0,0 +1,226 @@
+# frozen_string_literal: true
+
+# Run from ruby/rails/riverqueue-rails with its bundle, against a disposable
+# Postgres server started with shared_preload_libraries=pg_stat_statements:
+#   bundle exec ruby ../../bench/drain.rb --database-url postgres://localhost:55432/river_bench
+# To compare checkouts, use each checkout's Rails Gemfile as BUNDLE_GEMFILE
+# while running this same script. Alternate baseline and candidate runs.
+# Insertion, migration, and worker boot are outside the measured drain.
+# SQL statistics are reset server-wide, so use a dedicated benchmark server.
+
+require "bundler/setup"
+require "json"
+require "logger"
+require "open3"
+require "optparse"
+require "rbconfig"
+require "securerandom"
+require "timeout"
+require "pg"
+require "riverqueue-rails"
+
+class DrainJob < ActiveJob::Base
+  class << self
+    attr_accessor :executions, :mutex
+  end
+
+  # The only work is an in-memory audit of actual Active Job deliveries.
+  def perform(number)
+    self.class.mutex.synchronize { self.class.executions << number }
+  end
+end
+
+def connect(options, schema)
+  ActiveRecord::Base.establish_connection(
+    url: options.fetch(:database_url), pool: options.fetch(:threads) + 6,
+    schema_search_path: "#{schema},public"
+  )
+  River::Driver::ActiveRecord.new
+end
+
+def monotonic_now
+  Process.clock_gettime(Process::CLOCK_MONOTONIC)
+end
+
+def postgres_cpu(postmaster)
+  output, status = Open3.capture2("ps", "-axo", "pid=,ppid=,time=")
+  raise "ps failed" unless status.success?
+
+  output.lines.sum do |line|
+    pid, parent, elapsed = line.split
+    next 0.0 unless pid.to_i == postmaster || parent.to_i == postmaster
+
+    days, clock = elapsed.include?("-") ? elapsed.split("-", 2) : ["0", elapsed]
+    days.to_i * 86_400 + clock.split(":").map(&:to_f).reduce { |total, part| total * 60 + part }
+  end
+end
+
+def run_worker(options, schema, command_read, result_write)
+  result_write.sync = true
+  status = 0
+  begin
+    driver = connect(options, schema)
+    DrainJob.executions = []
+    DrainJob.mutex = Mutex.new
+    client = River::Client.new(driver, config: River::Config.new(
+      fetch_cooldown: options.fetch(:cooldown), fetch_poll_interval: 1.0,
+      logger: Logger.new(File::NULL), queues: {"default" => options.fetch(:threads)},
+      workers: River::Workers.new.add(River::Rails::Worker)
+    ))
+    result_write.puts("ready")
+    raise "missing start command" unless command_read.gets&.strip == "start"
+
+    before = Process.times
+    client.start
+    raise "missing stop command" unless command_read.gets&.strip == "stop"
+
+    after = Process.times
+    client.stop
+    result_write.puts(JSON.generate(cpu_seconds: after.utime + after.stime - before.utime - before.stime,
+      executions: DrainJob.executions))
+  rescue => error
+    result_write.puts(JSON.generate(error: error.full_message))
+    status = 1
+  ensure
+    client&.stop_and_cancel
+    ActiveRecord::Base.connection_pool.disconnect!
+    result_write.close
+    command_read.close
+  end
+  status
+end
+
+def start_worker(options, schema)
+  command_read, command_write = IO.pipe
+  command_write.sync = true
+  result_read, result_write = IO.pipe
+  # Boot fresh interpreters: forking a process that has used libpq can crash
+  # in native connection setup on macOS, even after its connections close.
+  pid = Process.spawn({"RIVER_DRAIN_WORKER" => JSON.generate(options: options, schema: schema)},
+    RbConfig.ruby, File.expand_path(__FILE__), in: command_read, out: result_write)
+  command_read.close
+  result_write.close
+  {command: command_write, pid: pid, result: result_read}
+end
+
+if ENV["RIVER_DRAIN_WORKER"]
+  worker = JSON.parse(ENV.fetch("RIVER_DRAIN_WORKER"))
+  ActiveJob::Base.logger = Logger.new(File::NULL)
+  exit run_worker(worker.fetch("options").transform_keys(&:to_sym), worker.fetch("schema"), $stdin, $stdout)
+end
+
+options = {cooldown: 0.001, jobs: 50_000, label: "candidate", processes: 8, threads: 20, timeout: 600.0}
+OptionParser.new do |parser|
+  parser.on("--cooldown SECONDS", Float) { |value| options[:cooldown] = value }
+  parser.on("--database-url URL") { |value| options[:database_url] = value }
+  parser.on("--jobs COUNT", Integer) { |value| options[:jobs] = value }
+  parser.on("--label LABEL") { |value| options[:label] = value }
+  parser.on("--processes COUNT", Integer) { |value| options[:processes] = value }
+  parser.on("--threads COUNT", Integer) { |value| options[:threads] = value }
+  parser.on("--timeout SECONDS", Float) { |value| options[:timeout] = value }
+end.parse!
+raise "--database-url is required (use a dedicated benchmark server)" unless options[:database_url]
+raise "jobs, processes, threads, and timeout must be positive" unless options.values_at(:jobs, :processes, :threads, :timeout).all?(&:positive?)
+
+ActiveJob::Base.logger = Logger.new(File::NULL)
+schema = "river_drain_#{SecureRandom.hex(8)}"
+workers = []
+begin
+  admin = PG.connect(options.fetch(:database_url))
+  admin.exec("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
+  admin.exec("SET pg_stat_statements.track = 'none'")
+  admin.exec("CREATE SCHEMA #{schema}")
+  admin.exec("SET search_path = #{schema},public")
+  driver = connect(options, schema)
+  River::Migrator.new(driver).migrate
+  client = River::Client.new(driver)
+  adapter = ActiveJob::QueueAdapters::RiverAdapter.new(client: client)
+  options.fetch(:jobs).times.each_slice(1_000) do |numbers|
+    adapter.enqueue_all(numbers.map { |number| DrainJob.new(number) })
+  end
+  admin.exec("VACUUM ANALYZE river_job")
+  ActiveRecord::Base.connection_pool.disconnect!
+  admin.close
+
+  options.fetch(:processes).times { workers << start_worker(options, schema) }
+  Timeout.timeout(60) do
+    workers.each do |worker|
+      ready = worker.fetch(:result).gets&.strip
+      raise "worker failed to start: #{ready}" unless ready == "ready"
+    end
+  end
+  admin = PG.connect(options.fetch(:database_url))
+  admin.exec("SET pg_stat_statements.track = 'none'")
+  admin.exec("SET search_path = #{schema},public")
+  backend = admin.exec("SELECT pg_backend_pid()").getvalue(0, 0)
+  postmaster = Open3.capture2("ps", "-o", "ppid=", "-p", backend).first.to_i
+  raise "could not find Postgres postmaster" unless postmaster.positive?
+
+  admin.exec("SELECT pg_stat_statements_reset()")
+  postgres_before = postgres_cpu(postmaster)
+  started = monotonic_now
+  workers.each { |worker| worker.fetch(:command).puts("start") }
+  loop do
+    if (ready = IO.select(workers.map { |worker| worker.fetch(:result) }, nil, nil, 0))
+      raise "worker exited during drain: #{ready.first.first.gets}"
+    end
+    counts = admin.exec("SELECT state, count(*) FROM river_job GROUP BY state").to_h { |row| [row.fetch("state"), row.fetch("count").to_i] }
+    break if counts["completed"] == options.fetch(:jobs)
+    raise "jobs failed: #{counts}" if counts.keys.intersect?(["cancelled", "discarded", "retryable"])
+    raise "drain timed out: #{counts}" if monotonic_now - started > options.fetch(:timeout)
+
+    sleep 0.05
+  end
+  elapsed = monotonic_now - started
+  postgres_seconds = postgres_cpu(postmaster) - postgres_before
+  statements = admin.exec(<<~SQL).map do |row|
+    SELECT query, calls, total_exec_time, shared_blks_hit, shared_blks_read
+    FROM pg_stat_statements WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database())
+    ORDER BY calls DESC
+  SQL
+    row.transform_values.with_index { |value, index| index.zero? ? value : value.to_f }
+  end
+  workers.each { |worker| worker.fetch(:command).puts("stop") }
+  results = Timeout.timeout(60) do
+    workers.map do |worker|
+      result = JSON.parse(worker.fetch(:result).readline)
+      raise result.fetch("error") if result.key?("error")
+
+      Process.wait(worker.fetch(:pid))
+      worker[:reaped] = true
+      raise "worker exited unsuccessfully" unless $?.success?
+
+      result
+    end
+  end
+  executions = results.flat_map { |result| result.fetch("executions") }
+  raise "missing or duplicate deliveries" unless executions.sort == (0...options.fetch(:jobs)).to_a
+  raise "retried jobs" unless admin.exec("SELECT count(*) FROM river_job WHERE attempt <> 1").getvalue(0, 0) == "0"
+
+  ruby_seconds = results.sum { |result| result.fetch("cpu_seconds") }
+  calls = statements.sum { |row| row.fetch("calls") }
+  puts JSON.generate(options.except(:database_url).merge(
+    active_record_version: ActiveRecord.gem_version.to_s, ruby_version: RUBY_VERSION, postgres_version: admin.server_version,
+    elapsed_seconds: elapsed, jobs_per_second: options.fetch(:jobs) / elapsed, exactly_once: true,
+    ruby_cpu_ms_per_job: ruby_seconds * 1_000 / options.fetch(:jobs),
+    postgres_cpu_ms_per_job: postgres_seconds * 1_000 / options.fetch(:jobs),
+    statements_per_job: calls / options.fetch(:jobs), statements: statements
+  ))
+ensure
+  workers.each do |worker|
+    unless worker[:reaped]
+      begin
+        Process.kill("TERM", worker.fetch(:pid))
+        Process.wait(worker.fetch(:pid))
+      rescue Errno::ESRCH, Errno::ECHILD
+        # The worker may have exited before cleanup.
+      end
+    end
+    worker.fetch(:command).close
+    worker.fetch(:result).close
+  end
+  ActiveRecord::Base.connection_handler.clear_all_connections!
+  admin = PG.connect(options.fetch(:database_url)) if !admin || admin.finished?
+  admin.exec("DROP SCHEMA IF EXISTS #{schema} CASCADE")
+  admin.close
+end

ruby/driver/riverqueue-activerecord/lib/driver.rb

@@ -31,7 +31,7 @@ def initialize(connection_class: ::ActiveRecord::Base)
 
       @connection_class = connection_class
       @postgres_capabilities_cache = PostgresCapabilities::Cache.new
-      @is_sqlite = connection_class.connection.adapter_name.downcase.include?("sqlite")
+      @is_sqlite = with_connection { |connection| connection.adapter_name.downcase.include?("sqlite") }
       # Do not inherit application scopes, callbacks, or schema caches. Each
       # driver owns its model, but all operations use the selected class's pool.
       @job_model = Class.new(::ActiveRecord::Base) do
@@ -115,12 +115,16 @@ def migration_connection
     # savepoint.
     def transaction(&)
       legacy = runtime_legacy_subtransactions?
-      # A joined Active Record block swallows ActiveRecord::Rollback without
-      # rolling anything back. Yield directly so that the transaction's owner
-      # receives it and rolls back, as Sequel does with Sequel::Rollback.
-      return yield if !legacy && @connection_class.connection.current_transaction.joinable?
-
-      @connection_class.transaction(requires_new: legacy, &)
+      with_connection do |connection|
+        # A joined Active Record block swallows ActiveRecord::Rollback without
+        # rolling anything back. Yield directly so that the transaction's owner
+        # receives it and rolls back, as Sequel does with Sequel::Rollback.
+        if !legacy && connection.current_transaction.joinable?
+          yield
+        else
+          connection.transaction(requires_new: legacy, &)
+        end
+      end
     end
 
     SQLITE_CONFLICT_WHERE = <<~SQL.chomp.freeze
@@ -245,7 +249,7 @@ def transaction(&)
     end
 
     private def runtime_execute(sql)
-      @connection_class.connection.execute(sql)
+      with_connection { |connection| connection.execute(sql) }
     end
 
     private def runtime_job_list_without_params
@@ -275,14 +279,27 @@ def transaction(&)
 
     private def runtime_query_rows(sql)
       if @is_sqlite
-        @connection_class.connection.raw_connection.execute(sql).to_a
+        with_connection { |connection| connection.raw_connection.execute(sql).to_a }
       else
-        @connection_class.connection.select_all(sql).to_a
+        with_connection { |connection| connection.select_all(sql).to_a }
       end
     end
 
     private def runtime_quote(value)
-      @connection_class.connection.quote(value)
+      with_connection { |connection| connection.quote(value) }
+    end
+
+    # Runs an UPDATE on river_job and decodes the rows it changed.
+    private def runtime_returning_jobs(sql)
+      if @is_sqlite
+        rows = with_connection { |connection| connection.raw_connection.execute("#{sql}\nRETURNING #{SQLITE_JOB_COLUMNS}").to_a }
+        rows.map { |row| runtime_decoded_job { sqlite_to_job_row_from_raw(row) } }
+      else
+        result = with_connection { |connection| connection.select_all("#{sql}\nRETURNING river_job.*") }
+        result.rows.map do |row|
+          runtime_decoded_job { postgres_to_job_row_from_raw(row, result.columns, result.column_types).first }
+        end
+      end
     end
 
     private def runtime_value(row, key)
@@ -355,7 +372,7 @@ def transaction(&)
           RETURNING #{SQLITE_JOB_COLUMNS}
         SQL
 
-        rows = @connection_class.connection.raw_connection.execute(sql, [JSON.generate(jobs)])
+        rows = with_connection { |connection| connection.raw_connection.execute(sql, [JSON.generate(jobs)]) }
         sqlite_notify_insert(insert_params_many)
 
         rows.map do |row|
@@ -367,7 +384,7 @@ def transaction(&)
 
     private def sqlite_job_rows(suffix, binds = [])
       sql = "SELECT #{SQLITE_JOB_COLUMNS} FROM river_job #{suffix}"
-      @connection_class.connection.raw_connection.execute(sql, binds)
+      with_connection { |connection| connection.raw_connection.execute(sql, binds) }
     end
 
     private def sqlite_notify_insert(insert_params_many)
@@ -381,13 +398,15 @@ def transaction(&)
         {payload: JSON.generate({queue: queue}), topic: "river_insert"}
       end
 
-      @connection_class.connection.raw_connection.execute(<<~SQL, [JSON.generate(notifications)])
-        INSERT INTO river_notification (payload, topic)
-        SELECT
-          json_extract(value, '$.payload'),
-          json_extract(value, '$.topic')
-        FROM json_each(cast(? AS blob))
-      SQL
+      with_connection do |connection|
+        connection.raw_connection.execute(<<~SQL, [JSON.generate(notifications)])
+          INSERT INTO river_notification (payload, topic)
+          SELECT
+            json_extract(value, '$.payload'),
+            json_extract(value, '$.topic')
+          FROM json_each(cast(? AS blob))
+        SQL
+      end
     end
 
     private def sqlite_to_job_row_from_model(river_job)
@@ -402,5 +421,12 @@ def transaction(&)
     private def to_job_row_from_model(river_job)
       @is_sqlite ? sqlite_to_job_row_from_model(river_job) : postgres_to_job_row_from_model(river_job)
     end
+
+    # Borrows the thread's current connection, or checks one out for the call
+    # only. A permanent lease would outlive the short-lived job threads, and the
+    # pool would reset each one when it reaps it.
+    private def with_connection(&)
+      @connection_class.connection_pool.with_connection(&)
+    end
   end
 end