Performance improvements

* By querying batch executions remaining, the query times remain very fast.

* When we are constantly updating the single batch row counts, it becomes a hotspot. Fast executing jobs quickly accumulate and slow down overall job processing (processing a few thousand jobs goes for 10ish seconds to 40ish seconds). This still adds a bit of overhead, but significantly less (10ish seconds to 15ish seconds)

* Handle batch completion in an after_commit to make sure the transaction is visible before checking executions. This may mean we need to introduce some monitoring in the cases an after_commit fails to fire due network issues or a database issue

Author: jpcamara

app/models/solid_queue/batch.rb

@@ -50,16 +50,28 @@ def on_finish=(value)
 
     def check_completion!
       return if finished? || !ready?
+      return if batch_executions.limit(1).exists?
 
-      with_lock do
-        return if finished_at? || !ready?
+      rows = Batch
+        .by_batch_id(batch_id)
+        .unfinished
+        .empty_executions
+        .update_all(finished_at: Time.current)
+
+      return if rows.zero?
 
-        if pending_jobs == 0
-          finished_attributes = { finished_at: Time.current }
-          finished_attributes[:failed_at] = Time.current if failed_jobs > 0
-          update!(finished_attributes)
-          execute_callbacks
+      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
+        finished_attributes[:pending_jobs] = 0
+
+        update!(finished_attributes)
+        execute_callbacks
       end
     end
 

app/models/solid_queue/batch/trackable.rb

@@ -10,6 +10,16 @@ module Trackable
         scope :succeeded, -> { finished.where(failed_at: nil) }
         scope :unfinished, -> { where(finished_at: nil) }
         scope :failed, -> { where.not(failed_at: nil) }
+        scope :by_batch_id, ->(batch_id) { where(batch_id:) }
+        scope :empty_executions, -> {
+          where(<<~SQL)
+            NOT EXISTS (
+              SELECT 1 FROM solid_queue_batch_executions
+              WHERE solid_queue_batch_executions.batch_id = solid_queue_batches.batch_id
+              LIMIT 1
+            )
+          SQL
+        }
       end
 
       def status
@@ -38,6 +48,18 @@ def ready?
         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? ? self[:pending_jobs] : batch_executions.count
+      end
+
       def progress_percentage
         return 0 if total_jobs == 0
         ((completed_jobs + failed_jobs) * 100.0 / total_jobs).round(2)

app/models/solid_queue/batch_execution.rb

@@ -5,6 +5,14 @@ class BatchExecution < Record
     belongs_to :job, optional: true
     belongs_to :batch, foreign_key: :batch_id, primary_key: :batch_id
 
+    after_commit :check_completion, on: :destroy
+
+    private
+      def check_completion
+        batch = Batch.find_by(batch_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? }
@@ -23,30 +31,6 @@ def create_all_from_jobs(jobs)
         end
       end
 
-      def process_job_completion(job, status)
-        batch_id = job.batch_id
-        batch_execution = job.batch_execution
-
-        return if batch_execution.blank?
-
-        transaction do
-          batch_execution.destroy!
-
-          if status == "failed"
-            Batch.where(batch_id: batch_id).update_all(
-              "pending_jobs = pending_jobs - 1, failed_jobs = failed_jobs + 1"
-            )
-          else
-            Batch.where(batch_id: batch_id).update_all(
-              "pending_jobs = pending_jobs - 1, completed_jobs = completed_jobs + 1"
-            )
-          end
-        end
-
-        batch = Batch.find_by(batch_id: batch_id)
-        batch.check_completion! if batch.present?
-      end
-
       private
 
         def provider_upsert_options

app/models/solid_queue/execution/batchable.rb

@@ -11,9 +11,9 @@ module Batchable
 
       private
         def update_batch_progress
-          # FailedExecutions are only created when the job is done retrying
           if is_a?(FailedExecution)
-            BatchExecution.process_job_completion(job, "failed")
+            # 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}"

app/models/solid_queue/job/batchable.rb

@@ -28,9 +28,7 @@ def update_batch_progress
           return unless saved_change_to_finished_at? && finished_at.present?
           return unless batch_id.present?
 
-          # Jobs marked as finished are considered completed
-          # (even if they failed and are being retried - we don't know that here)
-          BatchExecution.process_job_completion(self, "completed")
+          batch_execution&.destroy!
         rescue => e
           Rails.logger.error "[SolidQueue] Failed to update batch #{batch_id} progress for job #{id}: #{e.message}"
         end