Batch Support

README.md

@@ -28,6 +28,7 @@ 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)
 - [Puma plugin](#puma-plugin)
 - [Jobs and transactional integrity](#jobs-and-transactional-integrity)
 - [Recurring tasks](#recurring-tasks)
@@ -597,6 +598,66 @@ class ApplicationMailer < ActionMailer::Base
     Rails.error.report(exception)
     raise exception
   end
+```
+
+## Batch jobs
+
+SolidQueue offers support for batching jobs. This allows you to track progress of a set of jobs,
+and optionally trigger callbacks based on their status. It supports 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. Fires even when some jobs have failed.
+  - `on_success`: Fired when all jobs have succeeded, including retries. Will not fire if any jobs have failed, but will fire if jobs have been discarded using `discard_on`
+  - `on_failure`: Fired when all jobs have finished, including retries. Will only fire if one or more jobs have failed.
+- 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
+
+```rb
+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(batch) # batch is always the default first argument
+    Rails.logger.info "Good job finishing all 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!"
+  end
+end
+
+class BatchFailureJob < ApplicationJob
+  def perform(batch) # batch is always the default first argument
+    Rails.logger.info "At least one job failed, sorry!"
+  end
+end
+
+SolidQueue::Batch.enqueue(
+  on_finish: BatchFinishJob,
+  on_success: BatchSuccessJob,
+  on_failure: BatchFailureJob,
+  user_id: 123
+) do
+  5.times.map { |i| SleepyJob.perform_later(i) }
+end
+```
+
+### 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:
+
+```rb
+Rails.application.config.after_initialize do # or to_prepare
+  SolidQueue::Batch::EmptyJob.queue_as "my_batch_queue"
 end
 ```
 

app/jobs/solid_queue/batch/empty_job.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class Batch
+    class EmptyJob < (defined?(ApplicationJob) ? ApplicationJob : ActiveJob::Base)
+      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,145 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class Batch < Record
+    class AlreadyFinished < StandardError; end
+
+    include Trackable
+
+    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
+
+    after_initialize :set_active_job_batch_id
+    after_commit :start_batch, on: :create, unless: -> { ActiveRecord.respond_to?(:after_all_transactions_commit) }
+
+    class << self
+      def enqueue(on_success: nil, on_failure: nil, on_finish: nil, **metadata, &block)
+        new.tap do |batch|
+          batch.assign_attributes(
+            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)
+      raise AlreadyFinished, "You cannot enqueue a batch that is already finished" 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.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
+        end
+        finished_attributes[:completed_jobs] = total_jobs - failed
+
+        update!(finished_attributes)
+        enqueue_callback_jobs
+      end
+    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
+      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.send(:deserialize_arguments_if_needed)
+        active_job.arguments = [ self ] + Array.wrap(active_job.arguments)
+        active_job.enqueue
+      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
+
+      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

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class Batch
+    module Trackable
+      extend ActiveSupport::Concern
+
+      included do
+        scope :finished, -> { where.not(finished_at: nil) }
+        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
+        }
+      end
+
+      def status
+        if finished?
+          failed? ? "failed" : "completed"
+        elsif enqueued?
+          "enqueued"
+        else
+          "pending"
+        end
+      end
+
+      def failed?
+        failed_at.present?
+      end
+
+      def succeeded?
+        finished? && !failed?
+      end
+
+      def finished?
+        finished_at.present?
+      end
+
+      def enqueued?
+        enqueued_at.present?
+      end
+
+      def completed_jobs
+        finished? ? self[:completed_jobs] : total_jobs - batch_executions.count
+      end
+
+      def failed_jobs
+        finished? ? self[:failed_jobs] : jobs.joins(:failed_execution).count
+      end
+
+      def pending_jobs
+        finished? ? 0 : batch_executions.count
+      end
+
+      def progress_percentage
+        return 0 if total_jobs == 0
+        ((completed_jobs + failed_jobs) * 100.0 / total_jobs).round(2)
+      end
+    end
+  end
+end

app/models/solid_queue/batch_execution.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class BatchExecution < Record
+    belongs_to :job, optional: true
+    belongs_to :batch
+
+    after_commit :check_completion, on: :destroy
+
+    private
+      def check_completion
+        batch = Batch.find_by(id: batch_id)
+        batch.check_completion if batch.present?
+      end
+
+    class << self
+      def create_all_from_jobs(jobs)
+        batch_jobs = jobs.select { |job| job.batch_id.present? }
+        return if batch_jobs.empty?
+
+        batch_jobs.group_by(&:batch_id).each do |batch_id, jobs|
+          BatchExecution.insert_all!(jobs.map { |job|
+            { batch_id:, job_id: job.respond_to?(:provider_job_id) ? job.provider_job_id : job.id }
+          })
+
+          total = jobs.size
+          SolidQueue::Batch.where(id: batch_id).update_all([ "total_jobs = total_jobs + ?", total ])
+        end
+      end
+    end
+  end
+end

app/models/solid_queue/execution/batchable.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class Execution
+    module Batchable
+      extend ActiveSupport::Concern
+
+      included do
+        after_create :update_batch_progress, if: -> { job.batch_id? }
+      end
+
+      private
+        def update_batch_progress
+          if is_a?(FailedExecution)
+            # FailedExecutions are only created when the job is done retrying
+            job.batch_execution&.destroy!
+          end
+        rescue => e
+          Rails.logger.error "[SolidQueue] Failed to notify batch #{job.batch_id} about job #{job.id} failure: #{e.message}"
+        end
+    end
+  end
+end

app/models/solid_queue/failed_execution.rb

@@ -2,7 +2,7 @@
 
 module SolidQueue
   class FailedExecution < Execution
-    include Dispatching
+    include Dispatching, Batchable
 
     serialize :error, coder: JSON
 

app/models/solid_queue/job.rb

@@ -4,7 +4,7 @@ module SolidQueue
   class Job < Record
     class EnqueueError < StandardError; end
 
-    include Executable, Clearable, Recurrable
+    include Executable, Clearable, Recurrable, Batchable
 
     serialize :arguments, coder: JSON
 
@@ -62,7 +62,8 @@ def attributes_from_active_job(active_job)
             scheduled_at: active_job.scheduled_at,
             class_name: active_job.class.name,
             arguments: active_job.serialize,
-            concurrency_key: active_job.concurrency_key
+            concurrency_key: active_job.concurrency_key,
+            batch_id: active_job.batch_id
           }
         end
     end