diff --git a/ruby/CHANGELOG.md b/ruby/CHANGELOG.md index 7964a27e9..61eb332c0 100644 --- a/ruby/CHANGELOG.md +++ b/ruby/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Enforce the 32 MiB encoded UTF-8 JSON limit for `Job#output=`, preserving previous output when validation fails. [PR #1569](https://github.com/riverqueue/river/pull/1569). +- 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 diff --git a/ruby/bench/drain.rb b/ruby/bench/drain.rb new file mode 100644 index 000000000..69ca19cf5 --- /dev/null +++ b/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 diff --git a/ruby/driver/riverqueue-activerecord/lib/driver.rb b/ruby/driver/riverqueue-activerecord/lib/driver.rb index ba4e44f10..59ea1122e 100644 --- a/ruby/driver/riverqueue-activerecord/lib/driver.rb +++ b/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 diff --git a/ruby/driver/riverqueue-activerecord/spec/client_spec.rb b/ruby/driver/riverqueue-activerecord/spec/client_spec.rb index 6ed38ce4a..ff4da6537 100644 --- a/ruby/driver/riverqueue-activerecord/spec/client_spec.rb +++ b/ruby/driver/riverqueue-activerecord/spec/client_spec.rb @@ -37,6 +37,58 @@ it_behaves_like "Postgres finalized job list plans" if adapter == :postgres it_behaves_like "Postgres rescue concurrency" if adapter == :postgres it_behaves_like "dedicated worker process" + + it "returns borrowed connections after each operation on a job thread" do + leases = Thread.new do + pool = ActiveRecord::Base.connection_pool + observed = {} + record = lambda do |operation, &block| + result = block.call + observed[operation] = !!pool.active_connection? + result + end + + driver = record.call(:initialize) { River::Driver::ActiveRecord.new } + client = record.call(:client) { River::Client.new(driver) } + row = record.call(:insert) { client.insert(River::JobArgsHash.new("borrowed", {})).job } + claimed = record.call(:claim) { driver.job_claim(id: row.id, attempted_by: "borrower") } + record.call(:complete) { driver.job_complete(id: row.id, expected_job: claimed, finalized_at: Time.now.utc) } + record.call(:read) { driver.job_get_by_id(row.id) } + record.call(:cancelled) { driver.job_get_cancelled_ids([row.id]) } + record.call(:transaction) { driver.transaction { driver.job_get_by_id(row.id) } } + record.call(:rollback) do + driver.transaction do + client.insert(River::JobArgsHash.new("rolled_back", {})) + raise driver.rollback_exception + end + end + record.call(:error) do + driver.transaction { raise "transaction failed" } + rescue RuntimeError + # Check the lease after the exception leaves the transaction. + end + observed + ensure + pool&.release_connection + end.value + + expect(leases).to eq(%i[initialize client insert claim complete read cancelled transaction rollback error].to_h { |operation| [operation, false] }) + expect(@driver.job_list.map(&:kind)).to eq(["borrowed"]) + end + + it "keeps the caller's leased connection through nested driver operations" do + Thread.new do + pool = ActiveRecord::Base.connection_pool + connection = pool.lease_connection + @driver.transaction do + @driver.transaction { @driver.job_get_cancelled_ids([1]) } + expect(pool.active_connection?).to equal(connection) + end + expect(pool.active_connection?).to equal(connection) + ensure + pool&.release_connection + end.value + end end end end diff --git a/ruby/driver/riverqueue-sequel/lib/driver.rb b/ruby/driver/riverqueue-sequel/lib/driver.rb index f550a85a5..6a2351ea8 100644 --- a/ruby/driver/riverqueue-sequel/lib/driver.rb +++ b/ruby/driver/riverqueue-sequel/lib/driver.rb @@ -242,6 +242,15 @@ def transaction @db.literal(value) end + # Runs an UPDATE on river_job and decodes the rows it changed. + private def runtime_returning_jobs(sql) + if @is_sqlite + @db.fetch("#{sql}\nRETURNING #{SQLITE_JOB_COLUMNS}").all.map { |row| runtime_decoded_job { sqlite_to_job_row_from_raw(row) } } + else + @db.fetch("#{sql}\nRETURNING river_job.*").all.map { |row| runtime_decoded_job { postgres_to_job_row(row) } } + end + end + private def runtime_value(row, key) row[key] || row[key.to_s] end diff --git a/ruby/lib/client_runtime.rb b/ruby/lib/client_runtime.rb index 47cb2acaf..7b3f49a5e 100644 --- a/ruby/lib/client_runtime.rb +++ b/ruby/lib/client_runtime.rb @@ -589,7 +589,7 @@ def wake end rescue => error @config.logger.error("River maintenance stopped: #{error.full_message}") - wait(5) + wait(5, @maintenance_condition) retry unless stopping? end @@ -906,8 +906,8 @@ def wake @mutex.synchronize { @stop_requested } end - private def wait(duration) - @mutex.synchronize { @condition.wait(@mutex, duration) unless @stop_requested } + private def wait(duration, condition = @condition) + @mutex.synchronize { condition.wait(@mutex, duration) unless @stop_requested } end private def wait_for_maintenance(duration, generation) diff --git a/ruby/lib/driver/runtime.rb b/ruby/lib/driver/runtime.rb index 39a46c0f0..2bfdd0475 100644 --- a/ruby/lib/driver/runtime.rb +++ b/ruby/lib/driver/runtime.rb @@ -82,16 +82,10 @@ def job_complete(id:, finalized_at:, expected_job: nil, metadata: nil, now: Time assignments = ["state = 'completed'", "finalized_at = #{runtime_time(finalized_at)}"] assignments << "metadata = #{runtime_merge_metadata(metadata)}" unless metadata.nil? || metadata.empty? - transaction do - updated_id = runtime_returning_ids(<<~SQL).first - UPDATE river_job SET #{assignments.join(", ")} - WHERE #{predicate} AND state = 'running' AND NOT (#{runtime_cancel_attempted}) - RETURNING id - SQL - next runtime_read_job(updated_id) if updated_id + completed = runtime_returning_jobs(runtime_update_running_job(id, assignments, expected_job: expected_job, uncancelled: true)).first + return completed if completed - :cancelled if runtime_returning_ids("SELECT id FROM river_job WHERE #{predicate} AND (#{runtime_cancel_attempted})").any? - end + :cancelled if runtime_returning_ids("SELECT id FROM river_job WHERE #{predicate} AND (#{runtime_cancel_attempted})").any? end def job_delete(id) @@ -128,11 +122,7 @@ def job_delete_finalized(retention:, now: Time.now.utc, max: 1_000) # completion. Return :cancelled so the runtime can record the cancellation. def job_delete_if_running(id) id = Integer(id) - deleted = runtime_returning_ids(<<~SQL).any? - DELETE FROM river_job - WHERE id = #{id} AND state = 'running' AND NOT (#{runtime_cancel_attempted}) - RETURNING id - SQL + deleted = runtime_returning_ids("#{runtime_delete_running_job(id, uncancelled: true)} RETURNING river_job.id").any? return true if deleted return :cancelled if job_get_cancelled_ids([id]).include?(id) @@ -307,13 +297,7 @@ def job_set_state_if_running(id:, state:, now: Time.now.utc, attempt: nil, assignments << "state = CASE WHEN #{cancel_path} THEN 'cancelled' ELSE #{runtime_state(state)} END" - id = runtime_returning_ids(<<~SQL).first - UPDATE river_job - SET #{assignments.join(",\n ")} - WHERE id = #{Integer(id)} AND state = 'running' - RETURNING id - SQL - id ? runtime_read_job(id) : nil + runtime_returning_jobs(runtime_update_running_job(id, assignments)).first end def job_update(id, params) @@ -487,55 +471,60 @@ def time_precision runtime_postgres? ? "metadata ? 'cancel_attempted_at'" : "(CASE WHEN json_valid(metadata, 10) THEN metadata -> 'cancel_attempted_at' END) IS NOT NULL" end + # One statement claims the jobs and returns their rows, like Go's + # JobGetAvailable. The fetch subquery locks rows it then rejects, including + # jobs other producers already claimed, until the statement commits. private def runtime_claim_jobs(predicate, max, attempted_by, now) - transaction do - attempted_by_sql = if runtime_postgres? - "array_append(CASE WHEN cardinality(attempted_by) >= 100 THEN attempted_by[(cardinality(attempted_by) - 98):] ELSE attempted_by END, #{runtime_quote(attempted_by)})" - else - # Preserve corrupt history so decoding can fail this attempt instead - # of silently repairing it or aborting the entire claim batch. Trim - # to the newest 99 entries before appending, preserving JSON types. - <<~SQL - CASE WHEN NOT json_valid(attempted_by, 10) AND attempted_by IS NOT NULL THEN attempted_by - ELSE jsonb(json_insert(json( - CASE WHEN json_array_length(attempted_by) >= 100 THEN ( - SELECT json_group_array(json(value)) FROM ( - SELECT attempted_by -> ('$[' || key || ']') AS value FROM json_each(attempted_by) - WHERE key >= json_array_length(attempted_by) - 99 ORDER BY key - ) - ) ELSE coalesce(attempted_by, jsonb('[]')) END - ), '$[#]', #{runtime_quote(attempted_by)})) END - SQL - end - - lock_clause = runtime_postgres? ? "FOR UPDATE SKIP LOCKED" : "" - # An administratively requeued job may already be at the counter limit. - # Saturate instead of letting that row abort the entire claim batch. - ids = runtime_returning_ids(<<~SQL) - UPDATE river_job - SET attempt = CASE WHEN attempt < #{River::MAX_ATTEMPTS_LIMIT} THEN attempt + 1 ELSE attempt END, - attempted_at = #{runtime_time(now)}, - attempted_by = #{attempted_by_sql}, - state = 'running' - WHERE id IN ( - SELECT id FROM river_job - WHERE #{predicate} - ORDER BY priority ASC, scheduled_at ASC, id ASC - LIMIT #{Integer(max)} - #{lock_clause} - ) - RETURNING id + attempted_by_sql = if runtime_postgres? + "array_append(CASE WHEN cardinality(attempted_by) >= 100 THEN attempted_by[(cardinality(attempted_by) - 98):] ELSE attempted_by END, #{runtime_quote(attempted_by)})" + else + # Preserve corrupt history so decoding can fail this attempt instead + # of silently repairing it or aborting the entire claim batch. Trim + # to the newest 99 entries before appending, preserving JSON types. + <<~SQL + CASE WHEN NOT json_valid(attempted_by, 10) AND attempted_by IS NOT NULL THEN attempted_by + ELSE jsonb(json_insert(json( + CASE WHEN json_array_length(attempted_by) >= 100 THEN ( + SELECT json_group_array(json(value)) FROM ( + SELECT attempted_by -> ('$[' || key || ']') AS value FROM json_each(attempted_by) + WHERE key >= json_array_length(attempted_by) - 99 ORDER BY key + ) + ) ELSE coalesce(attempted_by, jsonb('[]')) END + ), '$[#]', #{runtime_quote(attempted_by)})) END SQL - ids.map { |id| runtime_read_job(id) } end + + lock_clause = runtime_postgres? ? "FOR UPDATE SKIP LOCKED" : "" + # An administratively requeued job may already be at the counter limit. + # Saturate instead of letting that row abort the entire claim batch. + runtime_returning_jobs(<<~SQL) + UPDATE river_job + SET attempt = CASE WHEN attempt < #{River::MAX_ATTEMPTS_LIMIT} THEN attempt + 1 ELSE attempt END, + attempted_at = #{runtime_time(now)}, + attempted_by = #{attempted_by_sql}, + state = 'running' + WHERE id IN ( + SELECT id FROM river_job + WHERE #{predicate} + ORDER BY priority ASC, scheduled_at ASC, id ASC + LIMIT #{Integer(max)} + #{lock_clause} + ) + SQL end - private def runtime_read_job(id) - job_get_by_id(id) + # A row that can't be decoded comes back with its decode error attached, + # so one corrupt job fails only its own attempt. + private def runtime_decoded_job + yield rescue River::JobRowDecodeError => error error.job end + private def runtime_read_job(id) + runtime_decoded_job { job_get_by_id(id) } + end + private def runtime_decode_postgres_job(row) row = row.transform_keys(&:to_sym) decoder = JobRowDecoder.new @@ -615,6 +604,11 @@ def time_precision "(#{column} IS NULL OR #{column} #{comparison} #{value} OR (#{column} = #{value} AND #{id_clause}))" end + private def runtime_delete_running_job(id, uncancelled: false) + lock, source, condition = runtime_running_job(id, uncancelled) + "#{lock}DELETE FROM river_job#{" USING #{source}" if source} WHERE #{condition}" + end + private def runtime_job_attempt_predicate(job) "attempt = #{Integer(job.attempt)} AND attempted_at IS NOT DISTINCT FROM #{runtime_nullable_time(job.attempted_at)}" end @@ -766,6 +760,34 @@ def time_precision runtime_query_rows(sql).map { |row| runtime_value(row, :id).to_i } end + # Targets one job only while it's running and, when uncancelled, only while + # no cancellation was requested. Postgres locks the row by primary key + # first and tests the locked row, the way Go's JobRetry and JobUpdate lock + # their target. With the state test in the statement's own WHERE, Postgres + # can plan it on the fetching index (state, queue, ...) when few jobs are + # running, and that scan reads every running entry, including the dead + # ones claims and completions leave behind. Returns the CTE, the extra + # source to join, and the condition. + private def runtime_running_job(id, uncancelled, expected_job: nil) + predicate = "id = #{Integer(id)}" + predicate += " AND #{runtime_job_attempt_predicate(expected_job)}" if expected_job + unless runtime_postgres? + condition = "#{predicate} AND state = 'running'" + condition += " AND NOT (#{runtime_cancel_attempted})" if uncancelled + return ["", nil, condition] + end + + lock = <<~SQL + WITH running_job AS ( + SELECT id, state AS locked_state, #{runtime_cancel_attempted} AS cancel_attempted + FROM river_job WHERE #{predicate} FOR UPDATE + ) + SQL + condition = "river_job.id = running_job.id AND running_job.locked_state = 'running'" + condition += " AND NOT running_job.cancel_attempted" if uncancelled + [lock, "running_job", condition] + end + # Finds conflicts and changes the whole pass in one statement. private def runtime_schedule_postgres(now, max) now = runtime_time(now) @@ -898,6 +920,11 @@ def time_precision "#{runtime_quote(encoded)}#{cast}" end + private def runtime_update_running_job(id, assignments, expected_job: nil, uncancelled: false) + lock, source, condition = runtime_running_job(id, uncancelled, expected_job: expected_job) + "#{lock}UPDATE river_job SET #{assignments.join(", ")}#{" FROM #{source}" if source} WHERE #{condition}" + end + private def runtime_update_value(field, value) case field when :attempt diff --git a/ruby/sig/driver.rbs b/ruby/sig/driver.rbs index 7e0601398..9141d03cc 100644 --- a/ruby/sig/driver.rbs +++ b/ruby/sig/driver.rbs @@ -124,6 +124,7 @@ module River module Runtime def runtime_read_job: (untyped) -> untyped + def runtime_decoded_job: () { () -> JobRow } -> JobRow def runtime_decode_sqlite_job: (untyped) -> JobRow def runtime_decode_postgres_job: (untyped) -> JobRow @postgres_capabilities_cache: PostgresCapabilities::Cache @@ -172,6 +173,7 @@ module River def runtime_cursor_clause: (untyped) -> String def runtime_execute: (String) -> untyped def runtime_in_clause: (untyped, untyped) -> untyped + def runtime_delete_running_job: (untyped, ?uncancelled: bool) -> String def runtime_job_attempt_predicate: (untyped) -> String def runtime_job_list_without_params: () -> untyped def runtime_job_list: (JobListParams, ?for_delete: bool) -> Array[JobRow] @@ -194,11 +196,14 @@ module River def runtime_quote: (untyped) -> String def runtime_rescue_retry: (untyped, untyped, untyped, untyped, untyped) -> Time def runtime_returning_ids: (untyped) -> untyped + def runtime_returning_jobs: (String) -> Array[JobRow] + def runtime_running_job: (untyped, bool, ?expected_job: JobRow?) -> [String, String?, String] def runtime_schedule_postgres: (Time, Integer) -> Array[String] def runtime_schedule_sqlite: (Time, Integer) -> Array[String] def runtime_state: (untyped) -> untyped def runtime_tag_contains: (untyped) -> untyped def runtime_time: (untyped) -> untyped + def runtime_update_running_job: (untyped, Array[String], ?expected_job: JobRow?, ?uncancelled: bool) -> String def runtime_update_value: (untyped, untyped) -> untyped def runtime_value: (untyped, untyped) -> untyped def transaction: () { () -> untyped } -> untyped diff --git a/ruby/sig/runtime.rbs b/ruby/sig/runtime.rbs index 2782a3465..c186cb5ae 100644 --- a/ruby/sig/runtime.rbs +++ b/ruby/sig/runtime.rbs @@ -439,7 +439,7 @@ module River def start_notifications: () -> untyped def start_producer: (String, QueueConfig) -> untyped def stopping?: () -> bool - def wait: (Float | Integer) -> untyped + def wait: (Float | Integer, ?ConditionVariable) -> untyped def wait_for_maintenance: (Float | Integer, Integer) -> untyped def wait_for_queue: (String, Float | Integer, Integer) -> untyped def wait_for_running_jobs: (String, Float | Integer) -> untyped diff --git a/ruby/spec/client_runtime_branch_spec.rb b/ruby/spec/client_runtime_branch_spec.rb index 5510cad55..c95a522f7 100644 --- a/ruby/spec/client_runtime_branch_spec.rb +++ b/ruby/spec/client_runtime_branch_spec.rb @@ -567,10 +567,12 @@ def execute(runtime, value) checks >= 3 end - value.define_singleton_method(:wait) { |_duration| } + waits = [] + value.define_singleton_method(:wait) { |duration, condition| waits << [duration, condition] } expect { value.send(:maintenance_loop) }.not_to raise_error expect(checks).to eq(3) + expect(waits).to eq([[5, value.instance_variable_get(:@maintenance_condition)]]) end it "does not retry failed maintenance after stop begins" do @@ -583,10 +585,12 @@ def execute(runtime, value) checks >= 2 end - value.define_singleton_method(:wait) { |_duration| } + waits = [] + value.define_singleton_method(:wait) { |duration, condition| waits << [duration, condition] } expect { value.send(:maintenance_loop) }.not_to raise_error expect(checks).to eq(2) + expect(waits).to eq([[5, value.instance_variable_get(:@maintenance_condition)]]) end it "isolates a failing service while continuing maintenance and renewing leadership" do diff --git a/ruby/spec/driver_runtime_feature_spec.rb b/ruby/spec/driver_runtime_feature_spec.rb index 8cfef94f7..6d3ca6742 100644 --- a/ruby/spec/driver_runtime_feature_spec.rb +++ b/ruby/spec/driver_runtime_feature_spec.rb @@ -86,8 +86,7 @@ driver = Object.new.extend(River::Driver::Runtime) driver.define_singleton_method(:runtime_postgres?) { true } driver.define_singleton_method(:runtime_quote) { |value| "'#{value}'" } - driver.define_singleton_method(:transaction) { |&block| block.call } - driver.define_singleton_method(:runtime_returning_ids) do |statement| + driver.define_singleton_method(:runtime_returning_jobs) do |statement| sql = statement [] end @@ -96,6 +95,22 @@ expect(sql).to include("array_append", "FOR UPDATE SKIP LOCKED") end + it "locks a running Postgres job by primary key before completing it" do + sql = nil + driver = Object.new.extend(River::Driver::Runtime) + driver.define_singleton_method(:runtime_postgres?) { true } + driver.define_singleton_method(:runtime_quote) { |value| "'#{value}'" } + driver.define_singleton_method(:runtime_returning_jobs) do |statement| + sql = statement + [:completed] + end + + expect(driver.job_complete(id: 1, finalized_at: Time.now.utc)).to eq(:completed) + expect(sql).to include("FROM river_job WHERE id = 1 FOR UPDATE", "FROM running_job", + "running_job.locked_state = 'running'", "NOT running_job.cancel_attempted") + expect(sql).not_to match(/WHERE[^()]*\bstate = 'running'/) + end + it "rejects nil runtime timestamps" do expect { postgres_driver.send(:runtime_time, nil) }.to raise_error(ArgumentError, "time cannot be nil") end diff --git a/ruby/spec/driver_runtime_shared_examples.rb b/ruby/spec/driver_runtime_shared_examples.rb index 785b8a906..fd83a264a 100644 --- a/ruby/spec/driver_runtime_shared_examples.rb +++ b/ruby/spec/driver_runtime_shared_examples.rb @@ -289,6 +289,22 @@ def insert_scheduler_job(value, **options) expect(driver.job_complete(id: -1, finalized_at: now)).to be_nil end + it "returns claimed and finished rows without reading jobs again" do + first = insert_job(1, scheduled_at: now - 1) + second = insert_job(2, scheduled_at: now - 1) + driver.define_singleton_method(:job_get_by_id) { |_id| raise "claim and completion must return their own rows" } + + claimed = driver.job_get_available(queue: first.queue, max: 2, attempted_by: "test", now: now) + expect(claimed.map(&:id)).to contain_exactly(first.id, second.id) + expect(claimed).to all(have_attributes(attempt: 1, attempted_by: ["test"], state: "running")) + expect(driver.job_complete(id: first.id, finalized_at: now, metadata: {"output" => "done"})).to have_attributes( + id: first.id, metadata: include("output" => "done"), state: "completed" + ) + expect(driver.job_set_state_if_running(id: second.id, state: "retryable", scheduled_at: now + 60)).to have_attributes( + id: second.id, scheduled_at: driver.normalize_time(now + 60), state: "retryable" + ) + end + it "leaves cancellation-marked jobs for normal cancellation handling" do row = insert_job(1, scheduled_at: now - 1) driver.job_claim(id: row.id, attempted_by: "test", now: now) @@ -305,7 +321,7 @@ def insert_scheduler_job(value, **options) row = insert_job(1, scheduled_at: now - 1) driver.job_claim(id: row.id, attempted_by: "test", now: now) cancel_before_write = true - driver.define_singleton_method(:runtime_returning_ids) do |sql| + driver.define_singleton_method(:runtime_returning_jobs) do |sql| if cancel_before_write cancel_before_write = false job_cancel(row.id)