Batch Support

README.md

@@ -30,6 +30,12 @@ Solid Queue can be used with SQL databases such as MySQL, PostgreSQL, or SQLite,
   - [Performance considerations](#performance-considerations)
 - [Failed jobs and retries](#failed-jobs-and-retries)
   - [Error reporting on jobs](#error-reporting-on-jobs)
+- [Batch jobs](#batch-jobs)
+  - [Empty batches](#empty-batches)
+  - [Batch progress and counters](#batch-progress-and-counters)
+  - [Batch maintenance](#batch-maintenance)
+  - [Clearing batches](#clearing-batches)
+  - [Upgrading existing installations](#upgrading-existing-installations)
 - [Puma plugin](#puma-plugin)
 - [Jobs and transactional integrity](#jobs-and-transactional-integrity)
 - [Recurring tasks](#recurring-tasks)
@@ -278,6 +284,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
@@ -637,6 +644,151 @@ class ApplicationMailer < ActionMailer::Base
     Rails.error.report(exception)
     raise exception
   end
+```
+
+## Batch jobs
+
+Solid Queue supports grouping jobs into batches, so you can track the progress of the set as a whole and optionally fire callbacks based on its status. Batches support the following:
+
+- Relating jobs to a batch, to track their status
+- Three available callbacks to fire:
+  - `on_finish`: fired when all jobs have finished, including retries, even when some jobs have failed.
+  - `on_success`: fired when all jobs have succeeded, including retries. It won't fire if any jobs have failed, but it will fire if jobs have been discarded using `discard_on`.
+  - `on_failure`: fired when all jobs have finished, including retries, and one or more of them have failed.
+- Enqueuing more jobs for a batch from inside one of its jobs, with `batch.enqueue`
+- Attaching a description and arbitrary metadata to a batch
+
+Callback jobs are regular jobs: the batch doesn't pass them any arguments (although you can configure your own), and they can access the batch they belong to through the `batch` accessor:
+
+```ruby
+class SleepyJob < ApplicationJob
+  def perform(seconds_to_sleep)
+    Rails.logger.info "Feeling #{seconds_to_sleep} seconds sleepy..."
+    sleep seconds_to_sleep
+  end
+end
+
+class BatchFinishJob < ApplicationJob
+  def perform
+    Rails.logger.info "Finished all #{batch.total_jobs} jobs"
+  end
+end
+
+class BatchSuccessJob < ApplicationJob
+  def perform
+    Rails.logger.info "All #{batch.completed_jobs} jobs worked!"
+  end
+end
+
+class BatchFailureJob < ApplicationJob
+  def perform
+    Rails.logger.info "#{batch.failed_jobs} jobs failed, sorry!"
+  end
+end
+
+SolidQueue::Batch.enqueue(
+  on_finish: BatchFinishJob,
+  on_success: BatchSuccessJob,
+  on_failure: BatchFailureJob,
+  user_id: 123
+) do
+  5.times { |i| SleepyJob.perform_later(i) }
+end
+```
+
+A job joins the batch that's active *when its enqueue is requested*—this also works when Rails defers the actual enqueue until after the surrounding transaction commits. In particular:
+
+- A job created outside a batch and enqueued inside one joins that batch.
+- Creating a job inside a batch without enqueueing it doesn't keep the batch open.
+- If a job already carries a batch ID but is enqueued inside another active batch, the active batch takes precedence.
+
+Besides the callbacks, `SolidQueue::Batch.enqueue` accepts a `description:`, to label the batch, and stores any other keyword arguments (like `user_id: 123` above) as the batch's `metadata`.
+
+Callbacks can be given as a job class or as a configured job instance—for example, `on_finish: BatchFinishJob.new.set(queue: :batches)` or `on_success: BatchSuccessJob.new("some argument")`. 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.
+
+### Empty batches
+
+In the case of an empty batch, a `SolidQueue::Batch::EmptyJob` is enqueued, so the batch can still finish and fire its callbacks. By default, this job runs on the `default` queue, and you can specify an alternative queue for it in an initializer:
+
+```ruby
+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`—the two retried attempts count as completed, plus the final 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 with `SolidQueue::Batch.clear_finished_in_batches` after `config.solid_queue.clear_finished_jobs_after`, but only when you invoke it. Failed batches are kept, like failed jobs, so you can inspect them. Installing Solid Queue configures [a recurring task](#recurring-tasks) that clears finished jobs every hour; you can add a matching entry for batches to your `recurring.yml`:
+
+```yml
+clear_solid_queue_finished_batches:
+  command: "SolidQueue::Batch.clear_finished_in_batches(sleep_between_batches: 0.3)"
+  schedule: every hour at minute 12
+```
+
+### 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
 ```
 

app/jobs/solid_queue/batch/empty_job.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class Batch
+    class EmptyJob < (defined?(ApplicationJob) ? ApplicationJob : ActiveJob::Base)
+      # Always use Solid Queue, even when ApplicationJob uses another adapter.
+      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
+      end
+    end
+  end
+end

app/models/solid_queue/batch.rb

@@ -0,0 +1,196 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class Batch < Record
+    class AlreadyFinished < StandardError
+      def initialize(message = "You cannot enqueue a batch that is already finished")
+        super
+      end
+    end
+
+    include Trackable, Clearable
+
+    has_many :jobs
+    has_many :batch_executions, class_name: "SolidQueue::BatchExecution", dependent: :destroy
+
+    serialize :metadata, coder: JSON
+    %w[ finish success failure ].each do |callback_type|
+      serialize "on_#{callback_type}", coder: JSON
+
+      define_method("on_#{callback_type}=") do |callback|
+        super serialize_callback(callback)
+      end
+    end
+
+    # Provider-agnostic batch identifier, analogous to 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
+      def enqueue(description: nil, on_success: nil, on_failure: nil, on_finish: nil, **metadata, &block)
+        new.tap do |batch|
+          batch.assign_attributes(
+            description: description,
+            on_success: on_success,
+            on_failure: on_failure,
+            on_finish: on_finish,
+            metadata: metadata
+          )
+
+          batch.enqueue(&block)
+        end
+      end
+
+      def current_batch_id
+        ActiveSupport::IsolatedExecutionState[:current_batch_id]
+      end
+
+      def wrap_in_batch_context(batch_id)
+        previous_batch_id = current_batch_id.presence
+        ActiveSupport::IsolatedExecutionState[:current_batch_id] = batch_id
+        yield
+      ensure
+        ActiveSupport::IsolatedExecutionState[:current_batch_id] = previous_batch_id
+      end
+    end
+
+    def enqueue(&block)
+      # Fast-fail for the common case. create_all_from_jobs atomically guards
+      # concurrent additions when it creates their tracking rows.
+      raise AlreadyFinished if finished?
+
+      transaction do
+        save! if new_record?
+
+        Batch.wrap_in_batch_context(id) do
+          block&.call(self)
+        end
+
+        if ActiveRecord.respond_to?(:after_all_transactions_commit)
+          ActiveRecord.after_all_transactions_commit do
+            start_batch
+          end
+        end
+      end
+    end
+
+    def metadata
+      (super || {}).with_indifferent_access
+    end
+
+    def check_completion
+      return if finished? || !enqueued?
+      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
+
+    COMPLETION_GRACE = 3.seconds
+
+    def self.sweep_stalled(stalled_for: 5.minutes, batch_size: 500)
+      SolidQueue.instrument(:sweep_stalled_batches, stalled_for: stalled_for, size: 0, started: 0, repaired: 0) do |payload|
+        # BatchExecution rows represent outstanding work. A row for a resolved
+        # job violates that invariant, so remove it immediately; destroy's
+        # after_commit callback retries the batch completion check.
+        [ BatchExecution.for_finished_jobs, BatchExecution.for_failed_jobs ].each do |leaked|
+          leaked.find_each(batch_size: batch_size) do |batch_execution|
+            payload[:repaired] += 1
+            batch_execution.destroy
+          end
+        end
+
+        # A started batch with no tracking rows can finish, but allow time for a
+        # transaction-deferred EmptyJob enqueue to become visible.
+        unfinished.empty_executions.where(enqueued_at: ...COMPLETION_GRACE.ago).find_each(batch_size: batch_size) do |batch|
+          payload[:size] += 1
+          batch.check_completion
+        end
+
+        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 finalize_completion
+        reload
+
+        # PostgreSQL can let a blocked CAS win from a stale NOT EXISTS snapshot.
+        # Re-check in a new statement while this transaction holds 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)
+        if value.present?
+          active_job = value.is_a?(ActiveJob::Base) ? value : value.new
+          # We can pick up batch ids from context, but callbacks should never be considered a part of the batch
+          active_job.batch_id = nil
+          active_job.serialize
+        end
+      end
+
+      def enqueue_callback_job(callback_name)
+        active_job = ActiveJob::Base.deserialize(send(callback_name))
+        active_job.callback_batch_id = id
+        # Bypass the job class's adapter so callbacks stay in Solid Queue and
+        # their enqueue stays in this transaction, while honoring enqueue callbacks.
+        active_job.run_callbacks(:enqueue) do
+          Job.enqueue(active_job, scheduled_at: active_job.scheduled_at || Time.current)
+        end
+      end
+
+      def enqueue_callback_jobs
+        if failed_at?
+          enqueue_callback_job(:on_failure) if on_failure.present?
+        else
+          enqueue_callback_job(:on_success) if on_success.present?
+        end
+
+        enqueue_callback_job(:on_finish) if on_finish.present?
+      end
+
+      def enqueue_empty_job
+        Batch.wrap_in_batch_context(id) do
+          EmptyJob.perform_later
+        end
+      end
+  end
+end