[ruby] open no savepoints inside a transaction

ruby/CHANGELOG.md

@@ -16,10 +16,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- River operations called inside an open Active Record or Sequel transaction now run in that transaction and open no savepoint of their own. If an operation raises, roll back the transaction, or wrap the call in your own savepoint to recover and continue. A `driver.rollback_exception` raised inside `Driver#transaction` now reaches the enclosing transaction block, as a rollback raised in your own code would. Before, it rolled back only River's savepoint. Set `RIVER_USE_LEGACY_SUBTRANSACTIONS=1` (or `true`) to restore the savepoints. This temporary fallback will be removed in a future release. [PR #1544](https://github.com/riverqueue/river/pull/1544).
 - `work_end(job, error)` hooks now return the attempt's error, matching Go. Return the incoming error to preserve it, a replacement exception to change it, or `nil` to clear a failure. The return value feeds the next hook in registration order. Existing observing hooks must explicitly return `error`. End hooks are skipped if `work_begin` raises or the worker raises an exception outside `StandardError`. [PR #1520](https://github.com/riverqueue/river/pull/1520).
 
 ### Fixed
 
+- The scheduler no longer opens a savepoint for each due job to catch unique key conflicts. A pass that made more than 64 jobs available overflowed the Postgres subtransaction cache, which can slow queries across the server. The scheduler now finds conflicts with one statement on Postgres and checks before each write on SQLite. Conflicting jobs are still discarded with `unique_key_conflict: scheduler_discarded` metadata. On Postgres, the scheduler now waits for a due job that another transaction has locked, instead of skipping it. [PR #1544](https://github.com/riverqueue/river/pull/1544).
+
 - Default periodic jobs to their occurrence time, preserving constructor and argument-level schedules. [PR #1537](https://github.com/riverqueue/river/pull/1537).
 
 - Fail retries that omit a saved cursor-only step, while preserving their checkpoints for recovery. Jobs with an existing step marker retain their behavior, including legacy stale cursor entries. [PR #1534](https://github.com/riverqueue/river/pull/1534).

ruby/docs/README.md

@@ -184,8 +184,8 @@ For simple jobs, `River::JobArgsHash.new(:kind, hash)` avoids defining an argume
 
 ### [Transactional enqueueing](https://riverqueue.com/docs/transactional-enqueueing)
 
-Insertion hooks and middleware execute inside the insertion transaction. If
-they raise, their database writes and the enqueue roll back together. Postgres
+Insertion hooks and middleware execute inside the insertion transaction, so
+their database writes and the enqueue commit or roll back together. Postgres
 insert notifications are delivered only when the surrounding transaction commits.
 
 Inserts automatically join a transaction opened through the same Active Record connection or Sequel database object. A rollback also rolls back the job:
@@ -206,6 +206,13 @@ end
 
 The equivalent works inside `ActiveRecord::Base.transaction` with the Active Record driver.
 
+River calls inside an open transaction run in that transaction and open no
+savepoint of their own. If a River call raises, it may leave partial changes,
+so roll back the transaction. To recover from the error and continue the
+transaction, wrap the call in your own savepoint, such as
+`DB.transaction(savepoint: true)` or
+`ApplicationRecord.transaction(requires_new: true)`.
+
 ### Transactional completion
 
 Call `job.client.job_complete_tx(job)` inside a transaction to commit application

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

@@ -110,10 +110,17 @@ def migration_connection
       end
     end
 
-    # Runs the block in a new Active Record transaction or savepoint and returns
-    # the block's result.
+    # Runs the block in an Active Record transaction and returns the block's
+    # result. Joins an open, joinable transaction instead of opening a
+    # savepoint.
     def transaction(&)
-      @connection_class.transaction(requires_new: true, &)
+      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, &)
     end
 
     SQLITE_CONFLICT_WHERE = <<~SQL.chomp.freeze
@@ -278,10 +285,6 @@ def transaction(&)
       @connection_class.connection.quote(value)
     end
 
-    private def runtime_unique_violation_class
-      ::ActiveRecord::RecordNotUnique
-    end
-
     private def runtime_value(row, key)
       row[key.to_s] || row[key]
     end
@@ -310,7 +313,7 @@ def transaction(&)
     # matching migration 007 and newer River databases.
     private def sqlite_job_insert_many(insert_params_many)
       nonces = sqlite_insert_nonces(insert_params_many)
-      @connection_class.transaction(requires_new: true) do
+      transaction do
         jobs = insert_params_many.zip(nonces).map { |param, nonce| sqlite_insert_params_to_hash(param, nonce) }
         inserted_nonces = nonces.to_h { |nonce| [nonce, true] }
 

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

@@ -89,11 +89,13 @@ def migration_connection
       end
     end
 
-    # Runs the block in a Sequel transaction or savepoint and returns the
-    # block's result.
+    # Runs the block in a Sequel transaction and returns the block's result.
+    # Joins an open transaction instead of opening a savepoint, unless that
+    # transaction uses Sequel's `auto_savepoint`.
     def transaction
       block_error = nil
-      @db.transaction(savepoint: true) do
+      options = runtime_legacy_subtransactions? ? {savepoint: true} : {}
+      @db.transaction(**options) do
         yield
       rescue ArgumentError => error
         block_error = error
@@ -240,10 +242,6 @@ def transaction
       @db.literal(value)
     end
 
-    private def runtime_unique_violation_class
-      ::Sequel::UniqueConstraintViolation
-    end
-
     private def runtime_value(row, key)
       row[key] || row[key.to_s]
     end
@@ -272,7 +270,7 @@ def transaction
     # matching migration 007 and newer River databases.
     private def sqlite_job_insert_many(insert_params_array)
       nonces = sqlite_insert_nonces(insert_params_array)
-      @db.transaction(savepoint: true) do
+      transaction do
         jobs = insert_params_array.zip(nonces).map { |param, nonce| sqlite_insert_params_to_hash(param, nonce) }
         inserted_nonces = nonces.to_h { |nonce| [nonce, true] }
 

ruby/lib/driver/runtime.rb

@@ -275,35 +275,18 @@ def job_retry(id, now: Time.now.utc)
       job
     end
 
+    # Makes due scheduled and retryable jobs available. A job is discarded
+    # instead if another row already holds its unique key, or if an earlier
+    # job in the same pass shares that key. The queries find these conflicts
+    # before they write, so the pass never needs a savepoint to recover from a
+    # failed write.
     def job_schedule(now: Time.now.utc, max: 1_000)
       transaction do
-        queues = [] #: Array[String]
-        # Hold each selected row until its transition (including uniqueness
-        # conflict handling) finishes. A concurrent retry may change its due time.
-        lock = runtime_postgres? ? "FOR UPDATE SKIP LOCKED" : ""
-        rows = runtime_query_rows(<<~SQL)
-          SELECT id, queue FROM river_job
-          WHERE state IN ('retryable', 'scheduled') AND scheduled_at <= #{runtime_time(now)}
-          ORDER BY priority, scheduled_at, id
-          LIMIT #{Integer(max)} #{lock}
-        SQL
-        rows.each do |row|
-          id = runtime_value(row, :id).to_i
-          transaction do
-            runtime_execute("UPDATE river_job SET state = 'available' WHERE id = #{id} AND state IN ('retryable', 'scheduled')")
-            queues << runtime_value(row, :queue)
-          end
-        rescue runtime_unique_violation_class
-          runtime_execute(<<~SQL)
-            UPDATE river_job
-            SET state = 'discarded', finalized_at = #{runtime_time(now)},
-                metadata = #{runtime_merge_metadata("unique_key_conflict" => "scheduler_discarded")}
-            WHERE id = #{id}
-          SQL
-        end
-
+        queues = runtime_postgres? ? runtime_schedule_postgres(now, Integer(max)) : runtime_schedule_sqlite(now, Integer(max))
+        # Notify each changed queue, including queues with only discarded
+        # jobs, as Go does.
         queues.uniq.each { |queue| runtime_notify("river_insert", queue: queue) }
-        rows.length
+        queues.length
       end
     end
 
@@ -676,6 +659,12 @@ def time_precision
       runtime_postgres? ? "#{runtime_quote(encoded)}::jsonb" : "jsonb(#{runtime_quote(encoded)})"
     end
 
+    # Temporary fallback: when true, transaction helpers open a savepoint
+    # inside an open transaction instead of joining it.
+    private def runtime_legacy_subtransactions?
+      %w[1 true].include?(ENV.fetch("RIVER_USE_LEGACY_SUBTRANSACTIONS", ""))
+    end
+
     private def runtime_merge_metadata(metadata)
       raise ArgumentError, "metadata must be a Hash" unless metadata.is_a?(Hash)
 
@@ -777,6 +766,114 @@ def time_precision
       runtime_query_rows(sql).map { |row| runtime_value(row, :id).to_i }
     end
 
+    # Finds conflicts and changes the whole pass in one statement.
+    private def runtime_schedule_postgres(now, max)
+      now = runtime_time(now)
+      runtime_query_rows(<<~SQL).map { |row| runtime_value(row, :queue) }
+        WITH jobs_to_schedule AS (
+            SELECT id, unique_key, unique_states, priority, scheduled_at
+            FROM river_job
+            WHERE state IN ('retryable', 'scheduled')
+                AND priority >= 0
+                AND queue IS NOT NULL
+                AND scheduled_at <= #{now}
+            ORDER BY priority, scheduled_at, id
+            LIMIT #{max}
+            FOR UPDATE
+        ),
+        jobs_with_rownum AS (
+            SELECT
+                *,
+                CASE
+                    WHEN unique_key IS NOT NULL AND unique_states IS NOT NULL THEN
+                        ROW_NUMBER() OVER (PARTITION BY unique_key ORDER BY priority, scheduled_at, id)
+                END AS row_num
+            FROM jobs_to_schedule
+        ),
+        unique_conflicts AS (
+            SELECT river_job.unique_key
+            FROM river_job
+            JOIN jobs_with_rownum
+                ON river_job.unique_key = jobs_with_rownum.unique_key
+                AND river_job.id != jobs_with_rownum.id
+            WHERE river_job.unique_key IS NOT NULL
+                AND river_job.unique_states IS NOT NULL
+                AND river_job_state_in_bitmask(river_job.unique_states, river_job.state)
+        ),
+        job_updates AS (
+            SELECT
+                job.id,
+                job.row_num IS NOT NULL AND (uc.unique_key IS NOT NULL OR job.row_num > 1) AS discard
+            FROM jobs_with_rownum job
+            LEFT JOIN unique_conflicts uc ON job.unique_key = uc.unique_key
+        )
+        UPDATE river_job
+        SET state = CASE WHEN job_updates.discard THEN 'discarded'::river_job_state ELSE 'available'::river_job_state END,
+            finalized_at = CASE WHEN job_updates.discard THEN #{now} ELSE river_job.finalized_at END,
+            metadata = CASE WHEN job_updates.discard THEN river_job.metadata || '{"unique_key_conflict": "scheduler_discarded"}'::jsonb
+                            ELSE river_job.metadata END
+        FROM job_updates
+        WHERE river_job.id = job_updates.id
+        RETURNING river_job.queue
+      SQL
+    end
+
+    # Follows Go's SQLite driver: check each unique job in order, and make it
+    # available before checking the next one, so a later job with the same
+    # key sees it as a conflict. SQLite allows one writer at a time. Once this
+    # pass writes, other connections cannot write until it ends. If another
+    # connection writes first, this pass's next write fails instead of using
+    # stale rows.
+    private def runtime_schedule_sqlite(now, max)
+      rows = runtime_query_rows(<<~SQL)
+        SELECT id, queue, unique_key IS NOT NULL AS has_unique_key
+        FROM river_job
+        WHERE state IN ('retryable', 'scheduled') AND scheduled_at <= #{runtime_time(now)}
+        ORDER BY priority, scheduled_at, id
+        LIMIT #{max}
+      SQL
+      available = [] #: Array[Integer]
+      discarded = [] #: Array[Integer]
+      rows.each do |row|
+        id = runtime_value(row, :id).to_i
+        next available << id if runtime_value(row, :has_unique_key).to_i.zero?
+
+        conflict = runtime_query_rows(<<~SQL).any?
+          SELECT 1 FROM river_job
+          WHERE id <> #{id}
+            AND unique_key = (SELECT job.unique_key FROM river_job AS job WHERE job.id = #{id})
+            AND unique_states IS NOT NULL
+            AND CASE state
+                  WHEN 'available' THEN unique_states & (1 << 0)
+                  WHEN 'cancelled' THEN unique_states & (1 << 1)
+                  WHEN 'completed' THEN unique_states & (1 << 2)
+                  WHEN 'discarded' THEN unique_states & (1 << 3)
+                  WHEN 'pending'   THEN unique_states & (1 << 4)
+                  WHEN 'retryable' THEN unique_states & (1 << 5)
+                  WHEN 'running'   THEN unique_states & (1 << 6)
+                  WHEN 'scheduled' THEN unique_states & (1 << 7)
+                  ELSE 0
+                END >= 1
+          LIMIT 1
+        SQL
+        next discarded << id if conflict
+
+        runtime_execute("UPDATE river_job SET state = 'available' WHERE id = #{id}")
+      end
+
+      unless discarded.empty?
+        runtime_execute(<<~SQL)
+          UPDATE river_job
+          SET state = 'discarded', finalized_at = #{runtime_time(now)},
+              metadata = #{runtime_merge_metadata("unique_key_conflict" => "scheduler_discarded")}
+          WHERE id IN (#{discarded.join(",")})
+        SQL
+      end
+      runtime_execute("UPDATE river_job SET state = 'available' WHERE id IN (#{available.join(",")})") unless available.empty?
+
+      rows.map { |row| runtime_value(row, :queue) }
+    end
+
     private def runtime_state(value)
       runtime_postgres? ? "#{runtime_quote(value)}::river_job_state" : runtime_quote(value)
     end

ruby/sig/driver.rbs

@@ -177,6 +177,7 @@ module River
       def runtime_job_list: (JobListParams, ?for_delete: bool) -> Array[JobRow]
       def runtime_job_rows: (String) -> untyped
       def runtime_json: (untyped) -> untyped
+      def runtime_legacy_subtransactions?: () -> bool
       def runtime_merge_metadata: (untyped) -> untyped
       def runtime_metadata_equals: (untyped, untyped) -> untyped
       def runtime_nullable_time: (untyped) -> untyped
@@ -193,10 +194,11 @@ 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_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_unique_violation_class: () -> singleton(Exception)
       def runtime_update_value: (untyped, untyped) -> untyped
       def runtime_value: (untyped, untyped) -> untyped
       def transaction: () { () -> untyped } -> untyped

ruby/spec/client_driver_shared_examples.rb

@@ -207,7 +207,10 @@
 end
 
 RSpec.shared_examples "SQL scheduling concurrency" do
-  it "skips locked jobs without overwriting a concurrent reschedule" do
+  # Like Go, the scheduler waits for a row that another transaction has
+  # locked. After the lock is released, Postgres rechecks the row, so a job
+  # rescheduled into the future stays scheduled.
+  it "waits for locked jobs without overwriting a concurrent reschedule" do
     client = River::Client.new(@driver)
     now = Time.now.utc
     first, second = [now - 2, now - 1].map do |scheduled_at|
@@ -225,15 +228,19 @@
     end
 
     Timeout.timeout(5) { locked.pop }
-    expect(@driver.job_schedule(now: now, max: 1)).to eq(1)
-    expect(@driver.job_get_by_id(second.id).state).to eq("available")
+    scheduler = Thread.new { @driver.job_schedule(now: now) }
+    expect(scheduler.join(0.5)).to be_nil
     release << true
     writer.value
+
+    expect(scheduler.value).to eq(1)
     expect(@driver.job_get_by_id(first.id)).to have_attributes(state: "scheduled", scheduled_at: be_within(0.001).of(now + 60))
+    expect(@driver.job_get_by_id(second.id).state).to eq("available")
     expect(@driver.job_schedule(now: now)).to eq(0)
   ensure
     release << true if release
     writer&.join
+    scheduler&.join
   end
 end
 

ruby/spec/driver_runtime_feature_spec.rb

@@ -21,6 +21,14 @@
   it_behaves_like "driver timestamp precision"
 end
 
+RSpec.describe "SQLite scheduler unique conflicts", database: :sqlite do
+  around(:each) { |example| sqlite_test_transaction(&example) }
+
+  let(:driver) { River::Driver::Sequel.new(SQLITE_DB) }
+
+  it_behaves_like "scheduler unique conflicts"
+end
+
 RSpec.describe "SQLite numeric metadata counters", database: :sqlite do
   around(:each) { |example| sqlite_test_transaction(&example) }