Batch job POC

* Introduces a "batch" concept, similar to batches present in Sidekiq Pro and GoodJob

* Batches monitor a set of jobs, and when those jobs are completed can fire off a final job

* This introduces a SolidQueue::JobBatch model, as well as the ability to enqueue jobs and associate them with the batch

* There are still more ideas to figure out, but this provides a basic batch scaffolding to spark discussion

Author: jpcamara

README.md

@@ -230,6 +230,18 @@ failed_execution.discard # This will delete the job from the system
 
 We're planning to release a dashboard called _Mission Control_, where, among other things, you'll be able to examine and retry/discard failed jobs, one by one, or in bulk.
 
+## Batch jobs
+
+```rb
+SolidQueue::JobBatch.enqueue(on_finish: BatchCompletionJob) do
+  5.times.map { |i| SleepyJob.perform_later(i) }
+end
+
+SolidQueue::JobBatch.enqueue(on_success: BatchCompletionJob) do
+  5.times.map { |i| SleepyJob.perform_later(i) }
+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/models/solid_queue/claimed_execution.rb

@@ -41,6 +41,8 @@ def perform
     else
       failed_with(result.error)
     end
+
+    job.job_batch.touch(:changed_at, :last_changed_at) if job.batch_id.present?
   ensure
     job.unblock_next_blocked_job
   end

app/models/solid_queue/job.rb

@@ -6,6 +6,8 @@ class Job < Record
 
     serialize :arguments, coder: JSON
 
+    belongs_to :job_batch, foreign_key: :batch_id, optional: true
+
     class << self
       def enqueue_all(active_jobs)
         active_jobs_by_job_id = active_jobs.index_by(&:job_id)
@@ -46,14 +48,16 @@ def create_all_from_active_jobs(active_jobs)
         end
 
         def attributes_from_active_job(active_job)
+          active_job.batch_id = JobBatch.current_batch_id || active_job.batch_id
           {
             queue_name: active_job.queue_name || DEFAULT_QUEUE_NAME,
             active_job_id: active_job.job_id,
             priority: active_job.priority || DEFAULT_PRIORITY,
             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/executable.rb

@@ -78,7 +78,7 @@ def dispatch_bypassing_concurrency_limits
       end
 
       def finished!
-        if preserve_finished_jobs?
+        if preserve_finished_jobs? || batch_id.present?
           touch(:finished_at)
         else
           destroy!

app/models/solid_queue/job_batch.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class JobBatch < Record
+    belongs_to :job, foreign_key: :job_id, optional: true
+    has_many :jobs, foreign_key: :batch_id
+
+    scope :incomplete, -> {
+      where(finished_at: nil).where("changed_at IS NOT NULL OR last_changed_at < ?", 1.hour.ago)
+    }
+
+    class << self
+      def current_batch_id
+        Thread.current[:current_batch_id]
+      end
+
+      def enqueue(attributes = {})
+        previous_batch_id = current_batch_id.presence || nil
+
+        job_batch = nil
+        transaction do
+          job_batch = create!(batch_attributes(attributes))
+          Thread.current[:current_batch_id] = job_batch.id
+          yield
+        end
+
+        job_batch
+      ensure
+        Thread.current[:current_batch_id] = previous_batch_id
+      end
+
+      def dispatch_finished_batches
+        incomplete.order(:id).pluck(:id).each do |id|
+          transaction do
+            where(id:).non_blocking_lock.each(&:finish)
+          end
+        end
+      end
+
+      private
+
+        def batch_attributes(attributes)
+          attributes = case attributes
+          in { on_finish: on_finish_klass }
+            attributes.merge(
+              job_class: on_finish_klass,
+              completion_type: "success"
+            )
+          in { on_success: on_success_klass }
+            attributes.merge(
+              job_class: on_success_klass,
+              completion_type: "success"
+            )
+          end
+
+          attributes.except(:on_finish, :on_success)
+        end
+    end
+
+    def finished?
+      finished_at.present?
+    end
+
+    def finish
+      return if finished?
+      reset_changed_at
+      jobs.find_each do |next_job|
+        # FIXME: If it's failed but is going to retry, how do we know?
+        #   Because we need to know if we will determine what the failed execution means
+        # FIXME: use "success" vs "finish" vs "discard" `completion_type` to determine
+        #   how to analyze each job
+        return unless next_job.finished?
+      end
+
+      attrs = {}
+
+      if job_class.present?
+        job_klass = job_class.constantize
+        active_job = job_klass.perform_later(self)
+        attrs[:job] = Job.find_by(active_job_id: active_job.job_id)
+      end
+
+      update!({ finished_at: Time.zone.now }.merge(attrs))
+    end
+
+    private
+
+      def reset_changed_at
+        if changed_at.blank? && last_changed_at.present?
+          update_columns(last_changed_at: Time.zone.now) # wait another hour before we check again
+        else
+          update_columns(changed_at: nil) # clear out changed_at so we ignore this until the next job finishes
+        end
+      end
+  end
+end

db/migrate/20240131013203_create_solid_queue_batch_table.rb

@@ -0,0 +1,21 @@
+class CreateSolidQueueBatchTable < ActiveRecord::Migration[7.1]
+  def change
+    create_table :solid_queue_job_batches do |t|
+      t.references :job, index: { unique: true }
+      t.string :job_class
+      t.string :completion_type
+      t.datetime :finished_at
+      t.datetime :changed_at
+      t.datetime :last_changed_at
+      t.timestamps
+
+      t.index [ :finished_at ]
+      t.index [ :changed_at ]
+      t.index [ :last_changed_at ]
+    end
+
+    add_reference :solid_queue_jobs, :batch, index: true
+    add_foreign_key :solid_queue_jobs, :solid_queue_job_batches, column: :batch_id, on_delete: :cascade
+    add_foreign_key :solid_queue_job_batches, :solid_queue_jobs, column: :job_id
+  end
+end

lib/active_job/job_batch_id.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+# Inspired by active_job/core.rb docs
+# https://github.com/rails/rails/blob/1c2529b9a6ba5a1eff58be0d0373d7d9d401015b/activejob/lib/active_job/core.rb#L136
+module ActiveJob
+  module JobBatchId
+    extend ActiveSupport::Concern
+
+    included do
+      attr_accessor :batch_id
+    end
+
+    def serialize
+      super.merge('batch_id' => batch_id)
+    end
+
+    def deserialize(job_data)
+      super
+      self.batch_id = job_data['batch_id']
+    end
+  end
+end

lib/solid_queue.rb

@@ -5,6 +5,7 @@
 
 require "active_job/queue_adapters/solid_queue_adapter"
 require "active_job/concurrency_controls"
+require "active_job/job_batch_id"
 
 require "solid_queue/app_executor"
 require "solid_queue/processes/supervised"

lib/solid_queue/dispatcher.rb

@@ -30,6 +30,7 @@ def run
       def dispatch_next_batch
         with_polling_volume do
           SolidQueue::ScheduledExecution.dispatch_next_batch(batch_size)
+          SolidQueue::JobBatch.dispatch_finished_batches
         end
       end
 

lib/solid_queue/engine.rb

@@ -33,6 +33,7 @@ class Engine < ::Rails::Engine
     initializer "solid_queue.active_job.extensions" do
       ActiveSupport.on_load :active_job do
         include ActiveJob::ConcurrencyControls
+        include ActiveJob::JobBatchId
       end
     end
   end

test/dummy/app/jobs/batch_completion_job.rb

@@ -0,0 +1,7 @@
+class BatchCompletionJob < ApplicationJob
+  queue_as :background
+
+  def perform(batch)
+    Rails.logger.info "#{batch.jobs.size} jobs completed!"
+  end
+end

test/dummy/app/jobs/sleepy_job.rb

@@ -0,0 +1,10 @@
+class SleepyJob < ApplicationJob
+  queue_as :background
+
+  retry_on Exception, wait: 30.seconds, attempts: 5
+
+  def perform(seconds_to_sleep)
+    Rails.logger.info "Feeling #{seconds_to_sleep} seconds sleepy..."
+    sleep seconds_to_sleep
+  end
+end

test/dummy/db/schema.rb

@@ -46,6 +46,21 @@
     t.index ["job_id"], name: "index_solid_queue_failed_executions_on_job_id", unique: true
   end
 
+  create_table "solid_queue_job_batches", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.bigint "job_id"
+    t.string "job_class"
+    t.string "completion_type"
+    t.datetime "finished_at"
+    t.datetime "changed_at"
+    t.datetime "last_changed_at"
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+    t.index ["changed_at"], name: "index_solid_queue_job_batches_on_changed_at"
+    t.index ["finished_at"], name: "index_solid_queue_job_batches_on_finished_at"
+    t.index ["job_id"], name: "index_solid_queue_job_batches_on_job_id", unique: true
+    t.index ["last_changed_at"], name: "index_solid_queue_job_batches_on_last_changed_at"
+  end
+
   create_table "solid_queue_jobs", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
     t.string "queue_name", null: false
     t.string "class_name", null: false
@@ -57,7 +72,9 @@
     t.string "concurrency_key"
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
+    t.bigint "batch_id"
     t.index ["active_job_id"], name: "index_solid_queue_jobs_on_active_job_id"
+    t.index ["batch_id"], name: "index_solid_queue_jobs_on_batch_id"
     t.index ["class_name"], name: "index_solid_queue_jobs_on_class_name"
     t.index ["finished_at"], name: "index_solid_queue_jobs_on_finished_at"
     t.index ["queue_name", "finished_at"], name: "index_solid_queue_jobs_for_filtering"
@@ -116,6 +133,8 @@
   add_foreign_key "solid_queue_blocked_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
   add_foreign_key "solid_queue_claimed_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
   add_foreign_key "solid_queue_failed_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
+  add_foreign_key "solid_queue_job_batches", "solid_queue_jobs", column: "job_id"
+  add_foreign_key "solid_queue_jobs", "solid_queue_job_batches", column: "batch_id", on_delete: :cascade
   add_foreign_key "solid_queue_ready_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
   add_foreign_key "solid_queue_scheduled_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
 end

test/models/solid_queue/job_batch_test.rb

@@ -0,0 +1,48 @@
+require "test_helper"
+
+class SolidQueue::JobBatchTest < ActiveSupport::TestCase
+  self.use_transactional_tests = false
+
+  teardown do
+    SolidQueue::Job.destroy_all
+    SolidQueue::JobBatch.destroy_all
+  end
+
+  class NiceJob < ApplicationJob
+    retry_on Exception, wait: 1.second
+
+    def perform(arg)
+      Rails.logger.info "Hi #{arg}!"
+    end
+  end
+
+  test "batch will be completed on success" do
+    batch = SolidQueue::JobBatch.enqueue(on_finish: BatchCompletionJob) {}
+    assert_equal "success", batch.completion_type
+    assert_equal BatchCompletionJob.name, batch.job_class
+  end
+
+  test "batch will be completed on finish" do
+    batch = SolidQueue::JobBatch.enqueue(on_success: BatchCompletionJob) {}
+    assert_equal "success", batch.completion_type
+    assert_equal BatchCompletionJob.name, batch.job_class
+  end
+
+  test "sets the batch_id on jobs created inside of the enqueue block" do
+    batch = SolidQueue::JobBatch.enqueue(on_finish: BatchCompletionJob) do
+      NiceJob.perform_later("world")
+      NiceJob.perform_later("people")
+    end
+
+    assert_equal 2, SolidQueue::Job.count
+    assert_equal [batch.id] * 2, SolidQueue::Job.last(2).map(&:batch_id)
+  end
+
+  test "batch id is present inside the block" do
+    assert_nil SolidQueue::JobBatch.current_batch_id
+    SolidQueue::JobBatch.enqueue(on_finish: BatchCompletionJob) do
+      assert_not_nil SolidQueue::JobBatch.current_batch_id
+    end
+    assert_nil SolidQueue::JobBatch.current_batch_id
+  end
+end