Don't delete batches until all their callback jobs complete

Connects to #1387

Add logic to delete batches only after their callback jobs have completed.

app/models/good_job/batch_record.rb

@@ -18,7 +18,7 @@ class BatchRecord < BaseRecord
     scope :not_discarded, -> { where(discarded_at: nil) }
     scope :succeeded, -> { finished.not_discarded }
 
-    scope :finished_before, ->(timestamp) { where(arel_table['finished_at'].lteq(bind_value('finished_at', timestamp, ActiveRecord::Type::DateTime))) }
+    scope :finished_before, ->(timestamp) { where(arel_table['callback_jobs_finished_at'].lteq(bind_value('callback_jobs_finished_at', timestamp, ActiveRecord::Type::DateTime))) }
 
     alias_attribute :enqueued?, :enqueued_at
     alias_attribute :discarded?, :discarded_at
@@ -68,6 +68,10 @@ def _continue_discard_or_finish(execution = nil, lock: true)
               on_success.constantize.set(priority: callback_priority, queue: callback_queue_name).perform_later(to_batch, { event: :success }) if !discarded_at && on_success.present?
               on_finish.constantize.set(priority: callback_priority, queue: callback_queue_name).perform_later(to_batch, { event: :finish }) if on_finish.present?
             end
+
+            if finished_at && callback_jobs.where(finished_at: nil).count.zero?
+              update(callback_jobs_finished_at: Time.current)
+            end
           end
         end
       end

spec/app/models/good_job/batch_record_spec.rb

@@ -42,4 +42,35 @@
       expect(result.last).to eq last_job
     end
   end
+
+  describe 'callback_jobs_finished_at' do
+    it 'is set when all callback jobs are finished' do
+      batch = described_class.create!
+      batch.update(enqueued_at: Time.current, finished_at: Time.current)
+      batch.callback_jobs.create!(finished_at: nil)
+      batch.callback_jobs.create!(finished_at: Time.current)
+
+      batch._continue_discard_or_finish
+
+      expect(batch.reload.callback_jobs_finished_at).to be_nil
+
+      batch.callback_jobs.update_all(finished_at: Time.current)
+      batch._continue_discard_or_finish
+
+      expect(batch.reload.callback_jobs_finished_at).to be_within(1.second).of(Time.current)
+    end
+  end
+
+  describe 'deletion logic' do
+    it 'checks callback_jobs_finished_at instead of finished_at' do
+      batch = described_class.create!
+      batch.update(enqueued_at: Time.current, finished_at: Time.current, callback_jobs_finished_at: nil)
+
+      expect { described_class.finished_before(Time.current).delete_all }.not_to change { described_class.count }
+
+      batch.update(callback_jobs_finished_at: Time.current)
+
+      expect { described_class.finished_before(Time.current).delete_all }.to change { described_class.count }.by(-1)
+    end
+  end
 end

spec/integration/batch_spec.rb

@@ -336,4 +336,31 @@ def perform(_batch, _params)
       expect(batch.active_jobs.count).to eq 2
     end
   end
+
+  describe 'batch deletion' do
+    it 'deletes batches only after their callback jobs have completed' do
+      batch = GoodJob::Batch.new
+      batch.on_finish = "BatchCallbackJob"
+      batch.enqueue do
+        TestJob.perform_later
+      end
+
+      GoodJob.perform_inline
+
+      batch.reload
+      expect(batch).to be_finished
+      expect(batch.callback_jobs_finished_at).to be_nil
+
+      batch.callback_active_jobs.each do |job|
+        job.update(finished_at: Time.current)
+      end
+
+      batch._record._continue_discard_or_finish
+
+      batch.reload
+      expect(batch.callback_jobs_finished_at).to be_within(1.second).of(Time.current)
+
+      expect { GoodJob::BatchRecord.finished_before(Time.current).delete_all }.to change { GoodJob::BatchRecord.count }.by(-1)
+    end
+  end
 end