Batch Support

README.md

@@ -584,6 +584,69 @@ 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`
+- Batches can be nested within other batches, creating a hierarchy. Outer batches will not fire callbacks until all nested jobs have finished.
+- 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 MultiStepJob < ApplicationJob
+  def perform
+    batch.enqueue do
+      SleepyJob.perform_later(5)
+      # Because of this nested batch, the top-level batch won't finish until the inner,
+      # 10 second job finishes
+      # Both jobs will still run simultaneously
+      SolidQueue::Batch.enqueue do
+        SleepyJob.perform_later(10)
+      end
+    end
+  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,
+  metadata: { user_id: 123 }
+) do
+  5.times.map { |i| SleepyJob.perform_later(i) }
 end
 ```
 

app/jobs/solid_queue/batch_update_job.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class BatchUpdateJob < ActiveJob::Base
+    class UpdateFailure < RuntimeError; end
+
+    queue_as :background
+
+    discard_on ActiveRecord::RecordNotFound
+
+    def perform(batch_id, job)
+      batch = SolidQueue::BatchRecord.find_by!(batch_id: batch_id)
+
+      return if job.batch_id != batch_id
+
+      status = job.status
+      return unless status.in?([ :finished, :failed ])
+
+      batch.job_finished!(job)
+    rescue => e
+      Rails.logger.error "[SolidQueue] BatchUpdateJob failed for batch #{batch_id}, job #{job.id}: #{e.message}"
+      raise
+    end
+  end
+end

app/models/solid_queue/batch_record.rb

@@ -0,0 +1,162 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class BatchRecord < Record
+    self.table_name = "solid_queue_job_batches"
+
+    STATUSES = %w[pending processing completed failed]
+
+    belongs_to :parent_job_batch, foreign_key: :parent_job_batch_id, class_name: "SolidQueue::BatchRecord", optional: true
+    has_many :jobs, foreign_key: :batch_id, primary_key: :batch_id
+    has_many :children, foreign_key: :parent_job_batch_id, primary_key: :batch_id, class_name: "SolidQueue::BatchRecord"
+
+    serialize :on_finish, coder: JSON
+    serialize :on_success, coder: JSON
+    serialize :on_failure, coder: JSON
+    serialize :metadata, coder: JSON
+
+    validates :status, inclusion: { in: STATUSES }
+
+    scope :pending, -> { where(status: "pending") }
+    scope :processing, -> { where(status: "processing") }
+    scope :completed, -> { where(status: "completed") }
+    scope :failed, -> { where(status: "failed") }
+    scope :finished, -> { where(status: %w[completed failed]) }
+    scope :unfinished, -> { where(status: %w[pending processing]) }
+
+    after_initialize :set_batch_id
+    before_create :set_parent_job_batch_id
+
+    def on_success=(value)
+      super(serialize_callback(value))
+    end
+
+    def on_failure=(value)
+      super(serialize_callback(value))
+    end
+
+    def on_finish=(value)
+      super(serialize_callback(value))
+    end
+
+    def job_finished!(job)
+      return if finished?
+
+      transaction do
+        if job.failed_execution.present?
+          self.class.where(id: id).update_all(
+            "failed_jobs = failed_jobs + 1, pending_jobs = pending_jobs - 1"
+          )
+        else
+          self.class.where(id: id).update_all(
+            "completed_jobs = completed_jobs + 1, pending_jobs = pending_jobs - 1"
+          )
+        end
+
+        reload
+        check_completion!
+      end
+    end
+
+    def check_completion!
+      return if finished?
+
+      actual_children = children.count
+      return if actual_children < expected_children
+
+      children.find_each do |child|
+        return unless child.finished?
+      end
+
+      if pending_jobs <= 0
+        if failed_jobs > 0
+          mark_as_failed!
+        else
+          mark_as_completed!
+        end
+        clear_unpreserved_jobs
+      elsif status == "pending"
+        update!(status: "processing")
+      end
+    end
+
+    def finished?
+      status.in?(%w[completed failed])
+    end
+
+    def processing?
+      status == "processing"
+    end
+
+    def pending?
+      status == "pending"
+    end
+
+    def progress_percentage
+      return 0 if total_jobs == 0
+      ((completed_jobs + failed_jobs) * 100.0 / total_jobs).round(2)
+    end
+
+    private
+
+      def set_parent_job_batch_id
+        self.parent_job_batch_id ||= Batch.current_batch_id if Batch.current_batch_id.present?
+      end
+
+      def set_batch_id
+        self.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)
+        return value if value.blank?
+        as_active_job(value).serialize
+      end
+
+      def perform_completion_job(job_field, attrs)
+        active_job = ActiveJob::Base.deserialize(send(job_field))
+        active_job.send(:deserialize_arguments_if_needed)
+        active_job.arguments = [ Batch.new(_batch_record: self) ] + Array.wrap(active_job.arguments)
+        ActiveJob.perform_all_later([ active_job ])
+
+        active_job.provider_job_id = Job.find_by(active_job_id: active_job.job_id).id
+        attrs[job_field] = active_job.serialize
+      end
+
+      def mark_as_completed!
+        # SolidQueue does treats `discard_on` differently than failures. The job will report as being :finished,
+        #   and there is no record of the failure.
+        # GoodJob would report a discard as an error. It's possible we should do that in the future?
+        update!(status: "completed", finished_at: Time.current)
+
+        perform_completion_job(:on_success, {}) if on_success.present?
+        perform_completion_job(:on_finish, {}) if on_finish.present?
+
+        if parent_job_batch_id.present?
+          parent = BatchRecord.find_by(batch_id: parent_job_batch_id)
+          parent&.reload&.check_completion!
+        end
+      end
+
+      def mark_as_failed!
+        update!(status: "failed", finished_at: Time.current)
+        perform_completion_job(:on_failure, {}) if on_failure.present?
+        perform_completion_job(:on_finish, {}) if on_finish.present?
+
+        # Check if parent batch can now complete
+        if parent_job_batch_id.present?
+          parent = BatchRecord.find_by(batch_id: parent_job_batch_id)
+          parent&.check_completion!
+        end
+      end
+
+      def clear_unpreserved_jobs
+        SolidQueue::Batch::CleanupJob.perform_later(self) unless SolidQueue.preserve_finished_jobs?
+      end
+  end
+end
+
+require_relative "batch_record/buffer"

app/models/solid_queue/batch_record/buffer.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class BatchRecord
+    class Buffer
+      attr_reader :jobs, :child_batches
+
+      def initialize
+        @jobs = {}
+        @child_batches = []
+      end
+
+      def add(job)
+        @jobs[job.job_id] = job
+        job
+      end
+
+      def add_child_batch(batch)
+        @child_batches << batch
+        batch
+      end
+
+      def capture
+        previous_buffer = ActiveSupport::IsolatedExecutionState[:solid_queue_batch_buffer]
+        ActiveSupport::IsolatedExecutionState[:solid_queue_batch_buffer] = self
+
+        yield
+
+        @jobs
+      ensure
+        ActiveSupport::IsolatedExecutionState[:solid_queue_batch_buffer] = previous_buffer
+      end
+
+      def self.current
+        ActiveSupport::IsolatedExecutionState[:solid_queue_batch_buffer]
+      end
+
+      def self.capture_job(job)
+        current&.add(job)
+      end
+
+      def self.capture_child_batch(batch)
+        current&.add_child_batch(batch)
+      end
+    end
+  end
+end

app/models/solid_queue/execution/batchable.rb

@@ -0,0 +1,20 @@
+# 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
+          BatchUpdateJob.perform_later(job.batch_id, job)
+        rescue => e
+          Rails.logger.error "[SolidQueue] Failed to notify batch #{batch_id} about job #{id} completion: #{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
 
@@ -61,7 +61,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

app/models/solid_queue/job/batchable.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class Job
+    module Batchable
+      extend ActiveSupport::Concern
+
+      included do
+        belongs_to :job_batch, foreign_key: :batch_id, optional: true
+
+        after_update :update_batch_progress, if: :batch_id?
+      end
+
+      private
+        def update_batch_progress
+          return unless saved_change_to_finished_at? && finished_at.present?
+          return unless batch_id.present?
+
+          BatchUpdateJob.perform_later(batch_id, self)
+        rescue => e
+          Rails.logger.error "[SolidQueue] Failed to notify batch #{batch_id} about job #{id} completion: #{e.message}"
+        end
+    end
+  end
+end

app/models/solid_queue/job/executable.rb

@@ -76,8 +76,8 @@ def dispatch_bypassing_concurrency_limits
       end
 
       def finished!
-        if SolidQueue.preserve_finished_jobs?
-          touch(:finished_at)
+        if SolidQueue.preserve_finished_jobs? || batch_id.present? # We clear jobs after the batch finishes
+          update!(finished_at: Time.current)
         else
           destroy!
         end