Batch support: fix completion races, accounting, and API polish

README.md

@@ -276,6 +276,7 @@ Here's an overview of the different options:
 It is recommended to set this value less than or equal to the queue database's connection pool size minus 2, as each worker thread uses one connection, and two additional connections are reserved for polling and heartbeat.
 - `processes`: this is the number of worker processes that will be forked by the supervisor with the settings given. By default, this is `1`, just a single process. This setting is useful if you want to dedicate more than one CPU core to a queue or queues with the same configuration. Only workers have this setting. **Note**: this option will be ignored if [running in `async` mode](#fork-vs-async-mode).
 - `concurrency_maintenance`: whether the dispatcher will perform the concurrency maintenance work. This is `true` by default, and it's useful if you don't use any [concurrency controls](#concurrency-controls) and want to disable it or if you run multiple dispatchers and want some of them to just dispatch jobs without doing anything else.
+- `batch_maintenance`: whether the dispatcher will sweep stalled [batches](#batch-jobs) as part of its maintenance work, on the same timer as concurrency maintenance (see [batch maintenance](#batch-maintenance)). This is `true` by default; disable it if you don't use batches, or if you run multiple dispatchers and want only some of them doing maintenance work.
 
 
 ### Optional scheduler configuration
@@ -632,6 +633,9 @@ and optionally trigger callbacks based on their status. It supports the followin
 - If a job is part of a batch, it can enqueue more jobs for that batch using `batch#enqueue`
 - Attaching arbitrary metadata to a batch
 
+Callback jobs are regular jobs: they don't receive any extra arguments, and they can access
+the batch they belong to through the `batch` accessor:
+
 ```rb
 class SleepyJob < ApplicationJob
   def perform(seconds_to_sleep)
@@ -641,20 +645,20 @@ class SleepyJob < ApplicationJob
 end
 
 class BatchFinishJob < ApplicationJob
-  def perform(batch) # batch is always the default first argument
-    Rails.logger.info "Good job finishing all jobs"
+  def perform
+    Rails.logger.info "Finished all #{batch.total_jobs} jobs"
   end
 end
 
 class BatchSuccessJob < ApplicationJob
-  def perform(batch) # batch is always the default first argument
-    Rails.logger.info "Good job finishing all jobs, and all of them worked!"
+  def perform
+    Rails.logger.info "All #{batch.completed_jobs} jobs worked!"
   end
 end
 
 class BatchFailureJob < ApplicationJob
-  def perform(batch) # batch is always the default first argument
-    Rails.logger.info "At least one job failed, sorry!"
+  def perform
+    Rails.logger.info "#{batch.failed_jobs} jobs failed, sorry!"
   end
 end
 
@@ -668,18 +672,117 @@ SolidQueue::Batch.enqueue(
 end
 ```
 
+A job joins the batch that's active *when it's enqueued*. Jobs instantiated inside the block
+but enqueued outside of it won't be part of the batch, and jobs instantiated elsewhere but
+enqueued inside the block will be.
+
+Callbacks can be given as a job class or as a configured job instance, e.g.
+`on_finish: BatchFinishJob.new.set(queue: :batches)`. Note that the job is serialized when
+the batch is created, so options resolved at that point (like `wait_until:` timestamps) are
+relative to batch creation, not to when the callback is eventually enqueued.
+
 ### Batch options
 
 In the case of an empty batch, a `SolidQueue::Batch::EmptyJob` is enqueued.
 
-By default, this jobs run on the `default` queue. You can specify an alternative queue for it in an initializer:
+By default, this job runs on the `default` queue. You can specify an alternative queue for it in an initializer:
 
 ```rb
 Rails.application.config.after_initialize do # or to_prepare
   SolidQueue::Batch::EmptyJob.queue_as "my_batch_queue"
 end
 ```
 
+The empty job and batch callback jobs always enqueue through Solid Queue, even when the
+job classes involved (or the application default) use a different Active Job adapter.
+
+### Batch progress and counters
+
+Batches track `total_jobs`, `completed_jobs`, `failed_jobs` and `pending_jobs`, plus a
+`progress_percentage` helper. A couple of accounting details to be aware of:
+
+- Every *attempt* counts: when a job is retried via `retry_on`, each retry is enqueued as a
+  new job in the batch, so a job that fails twice and then succeeds contributes 3 to
+  `total_jobs` (2 completed retries + 1 success).
+- Jobs discarded via `discard_on`, concurrency's `on_conflict: :discard`, or manual
+  discarding count as completed, not failed.
+- Manually retrying a failed job (via `SolidQueue::FailedExecution#retry`) doesn't re-add it
+  to its batch: if the batch already finished as failed, a successful manual retry won't
+  change the batch's status.
+
+### Batch maintenance
+
+Batch completion is normally detected as jobs finish, without ever locking the batch row
+outside a single once-per-batch moment. A few edge cases can't trigger that detection: jobs
+removed via bulk discards (which delete jobs without callbacks), a process that crashed
+after enqueueing jobs but before starting its batch, or a completion whose callback
+enqueueing failed and rolled back.
+
+The dispatcher sweeps these up automatically via `SolidQueue::Batch.sweep_stalled`, as part
+of its regular maintenance (every `concurrency_maintenance_interval` seconds, sharing a
+single maintenance timer and database connection). If you disable `batch_maintenance` (or
+don't run a dispatcher), you can run the sweep yourself, for example as a
+[recurring task](#recurring-tasks):
+
+```yml
+batch_maintenance:
+  command: "SolidQueue::Batch.sweep_stalled"
+  schedule: every 5 minutes
+```
+
+### Clearing batches
+
+Finished, non-failed batches are cleared after `config.solid_queue.clear_finished_jobs_after`,
+but only when you invoke it: like jobs, batches are cleared with
+`SolidQueue::Batch.clear_finished_in_batches`, which you'd typically run periodically
+alongside `SolidQueue::Job.clear_finished_in_batches`. Failed batches are kept, like failed
+jobs, so you can inspect them.
+
+### Upgrading existing installations
+
+If you installed Solid Queue before batches existed, add the new tables with a migration in
+`db/queue_migrate`:
+
+```ruby
+class AddSolidQueueBatches < ActiveRecord::Migration[7.1]
+  def change
+    create_table :solid_queue_batches do |t|
+      t.string :active_job_batch_id
+      t.string :description
+      t.text :on_finish
+      t.text :on_success
+      t.text :on_failure
+      t.text :metadata
+      t.integer :total_jobs, default: 0, null: false
+      t.integer :completed_jobs, default: 0, null: false
+      t.integer :failed_jobs, default: 0, null: false
+      t.datetime :enqueued_at
+      t.datetime :finished_at
+      t.datetime :failed_at
+      t.timestamps
+
+      t.index :active_job_batch_id, unique: true
+      t.index :finished_at
+    end
+
+    create_table :solid_queue_batch_executions do |t|
+      t.bigint :job_id, null: false
+      t.bigint :batch_id, null: false
+      t.datetime :created_at, null: false
+
+      t.index :job_id, unique: true
+      t.index :batch_id
+    end
+
+    add_column :solid_queue_jobs, :batch_id, :bigint
+    add_index :solid_queue_jobs, :batch_id
+
+    add_foreign_key :solid_queue_batch_executions, :solid_queue_batches, column: :batch_id, on_delete: :cascade
+    add_foreign_key :solid_queue_batch_executions, :solid_queue_jobs, column: :job_id, on_delete: :cascade
+  end
+end
+```
+
 ## Puma plugin
 
 We provide a Puma plugin if you want to run the Solid Queue's supervisor together with Puma and have Puma monitor and manage it. You just need to add

app/jobs/solid_queue/batch/empty_job.rb

@@ -3,6 +3,9 @@
 module SolidQueue
   class Batch
     class EmptyJob < (defined?(ApplicationJob) ? ApplicationJob : ActiveJob::Base)
+      # Pinned so it can't leak to another backend in mixed-adapter apps
+      self.queue_adapter = :solid_queue
+
       def perform
         # This job does nothing - it just exists to trigger batch completion
         # The batch completion will be handled by the normal job_finished! flow

app/models/solid_queue/batch.rb

@@ -2,7 +2,11 @@
 
 module SolidQueue
   class Batch < Record
-    class AlreadyFinished < StandardError; end
+    class AlreadyFinished < StandardError
+      def initialize(message = "You cannot enqueue a batch that is already finished")
+        super
+      end
+    end
 
     include Trackable, Clearable
 
@@ -18,7 +22,9 @@ class AlreadyFinished < StandardError; end
       end
     end
 
-    after_initialize :set_active_job_batch_id
+    # Reserved as a provider-agnostic batch identifier, like solid_queue_jobs.active_job_id
+    before_create :set_active_job_batch_id
+
     after_commit :start_batch, on: :create, unless: -> { ActiveRecord.respond_to?(:after_all_transactions_commit) }
 
     class << self
@@ -50,7 +56,8 @@ def wrap_in_batch_context(batch_id)
     end
 
     def enqueue(&block)
-      raise AlreadyFinished, "You cannot enqueue a batch that is already finished" if finished?
+      # Fast-fail only: the authoritative guard is in BatchExecution.create_all_from_jobs
+      raise AlreadyFinished if finished?
 
       transaction do
         save! if new_record?
@@ -73,37 +80,67 @@ def metadata
 
     def check_completion
       return if finished? || !enqueued?
-      return if batch_executions.any?
-      rows = Batch
-        .where(id: id)
-        .unfinished
-        .empty_executions
-        .update_all(finished_at: Time.current)
-
-      return if rows.zero?
-
-      with_lock do
-        failed = jobs.joins(:failed_execution).count
-        finished_attributes = {}
-        if failed > 0
-          finished_attributes[:failed_at] = Time.current
-          finished_attributes[:failed_jobs] = failed
+      return if batch_executions.exists?
+
+      transaction do
+        finished_rows = Batch.where(id: id).unfinished.enqueued.empty_executions.update_all(finished_at: Time.current)
+        finalize_completion if finished_rows.positive?
+      end
+    end
+
+    # Safety net for batches that can't finish through the normal flow
+    def self.sweep_stalled(stalled_for: 5.minutes, batch_size: 500)
+      SolidQueue.instrument(:sweep_stalled_batches, stalled_for: stalled_for, size: 0, started: 0) do |payload|
+        unfinished.empty_executions.where(enqueued_at: ...stalled_for.ago).find_each(batch_size: batch_size) do |batch|
+          payload[:size] += 1
+          batch.check_completion
         end
-        finished_attributes[:completed_jobs] = total_jobs - failed
 
-        update!(finished_attributes)
-        enqueue_callback_jobs
+        unfinished.where(enqueued_at: nil).where(created_at: ...stalled_for.ago).find_each(batch_size: batch_size) do |batch|
+          payload[:started] += 1
+          batch.start_batch
+        end
       end
     end
 
+    def start_batch
+      # Single-winner start so concurrent sweepers can't enqueue duplicate empty jobs
+      transaction do
+        if Batch.where(id: id, enqueued_at: nil).update_all(enqueued_at: Time.current).positive?
+          enqueue_empty_job if reload.total_jobs == 0
+        end
+      end
+
+      check_completion
+    end
+
     private
 
       def set_active_job_batch_id
         self.active_job_batch_id ||= SecureRandom.uuid
       end
 
-      def as_active_job(active_job_klass)
-        active_job_klass.is_a?(ActiveJob::Base) ? active_job_klass : active_job_klass.new
+      def finalize_completion
+        reload
+
+        # A blocked finisher can win with a stale NOT EXISTS on PostgreSQL; re-check under the row lock
+        raise ActiveRecord::Rollback if batch_executions.exists?
+
+        SolidQueue.instrument(:finish_batch, batch_id: id) do |payload|
+          failed = jobs.failed.count
+          finished_attributes = { completed_jobs: total_jobs - failed }
+          if failed > 0
+            finished_attributes[:failed_at] = Time.current
+            finished_attributes[:failed_jobs] = failed
+          end
+
+          update_columns(finished_attributes)
+          enqueue_callback_jobs
+
+          payload[:total_jobs] = total_jobs
+          payload[:completed_jobs] = self[:completed_jobs]
+          payload[:failed_jobs] = failed
+        end
       end
 
       def serialize_callback(value)
@@ -118,7 +155,9 @@ def serialize_callback(value)
       def enqueue_callback_job(callback_name)
         active_job = ActiveJob::Base.deserialize(send(callback_name))
         active_job.callback_batch_id = id
-        active_job.enqueue
+        # Enqueued directly so callbacks stay in Solid Queue even when the job
+        # class's adapter differs, and stay atomic with the finishing transaction
+        Job.enqueue(active_job, scheduled_at: active_job.scheduled_at || Time.current)
       end
 
       def enqueue_callback_jobs
@@ -136,10 +175,5 @@ def enqueue_empty_job
           EmptyJob.perform_later
         end
       end
-
-      def start_batch
-        enqueue_empty_job if reload.total_jobs == 0
-        update!(enqueued_at: Time.current)
-      end
   end
 end

app/models/solid_queue/batch/trackable.rb

@@ -10,24 +10,18 @@ module Trackable
         scope :succeeded, -> { finished.where(failed_at: nil) }
         scope :unfinished, -> { where(finished_at: nil) }
         scope :failed, -> { where.not(failed_at: nil) }
-        scope :empty_executions, -> {
-          where(<<~SQL)
-            NOT EXISTS (
-              SELECT 1 FROM solid_queue_batch_executions
-              WHERE solid_queue_batch_executions.batch_id = solid_queue_batches.id
-              LIMIT 1
-            )
-          SQL
-        }
+        scope :enqueued, -> { where.not(enqueued_at: nil) }
+        # Join-free so update_all keeps this condition in the completion update's own WHERE
+        scope :empty_executions, -> { where.not(id: BatchExecution.select(:batch_id)) }
       end
 
       def status
         if finished?
-          failed? ? "failed" : "completed"
+          failed? ? :failed : :completed
         elsif enqueued?
-          "enqueued"
+          :enqueued
         else
-          "pending"
+          :pending
         end
       end
 
@@ -47,12 +41,13 @@ def enqueued?
         enqueued_at.present?
       end
 
+      # Failed jobs have also lost their batch execution, so subtract them too
       def completed_jobs
-        finished? ? self[:completed_jobs] : total_jobs - batch_executions.count
+        finished? ? self[:completed_jobs] : total_jobs - pending_jobs - failed_jobs
       end
 
       def failed_jobs
-        finished? ? self[:failed_jobs] : jobs.joins(:failed_execution).count
+        finished? ? self[:failed_jobs] : jobs.failed.count
       end
 
       def pending_jobs
@@ -61,7 +56,7 @@ def pending_jobs
 
       def progress_percentage
         return 0 if total_jobs == 0
-        ((completed_jobs + failed_jobs) * 100.0 / total_jobs).round(2)
+        ((total_jobs - pending_jobs) * 100.0 / total_jobs).round(2)
       end
     end
   end