Prevent on_batch_complete from firing before all children are enqueued

lib/cloudtasker/batch/job.rb

@@ -206,6 +206,16 @@ def batch_completion_gid
         "#{batch_state_gid}/completed"
       end
 
+      #
+      # Return the key tracking the batch enqueuing lifecycle: 'in_progress'
+      # while the parent enqueues children, 'done' once finished.
+      #
+      # @return [String] The batch setup marker key.
+      #
+      def batch_setup_gid
+        "#{batch_state_gid}/setup_complete"
+      end
+
       #
       # Return the number of jobs in a given state
       #
@@ -371,6 +381,22 @@ def complete?
         redis.hvals(batch_state_gid).all? { |e| COMPLETION_STATUSES.include?(e) }
       end
 
+      #
+      # Return true unless the parent is still enqueuing children. The parent
+      # enqueues one at a time, so a fast child could otherwise complete off a
+      # partially-registered state and fire `on_batch_complete` prematurely
+      # (see #on_child_complete).
+      #
+      # A batch with no marker predates this feature (enqueued by an older
+      # #setup); treat it as fully enqueued so batches in flight during an
+      # upgrade complete normally instead of being stranded.
+      #
+      # @return [Boolean] True unless the batch is still enqueuing children.
+      #
+      def setup_complete?
+        redis.get(batch_setup_gid) != 'in_progress'
+      end
+
       #
       # Run worker callback. The error and dead callbacks get
       # silenced should they raise an error.
@@ -436,6 +462,11 @@ def on_child_complete(child_batch, status = :completed)
         return if status == :errored
         return unless complete?
 
+        # Do not let a child complete the batch until the parent has finished
+        # enqueuing every child. Otherwise a fast child can fire on_complete
+        # before its slower siblings have even been scheduled. See #setup_complete?.
+        return unless setup_complete?
+
         # Notify the parent batch that we are done with this batch. Use SETNX to ensure
         # only the first concurrent child to complete triggers on_complete.
         return unless redis.set(batch_completion_gid, true, nx: true, ex: BATCH_COMPLETION_TTL)
@@ -480,6 +511,7 @@ def cleanup
           m.del(batch_gid)
           m.del(batch_state_gid)
           m.del(batch_completion_gid)
+          m.del(batch_setup_gid)
           BATCH_STATUSES.each { |e| m.del(batch_state_count_gid(e)) }
         end
       end
@@ -551,8 +583,15 @@ def setup
         # Save batch
         save
 
+        # Mark as enqueuing before scheduling any child: while set, a completing
+        # child must not trigger batch completion (see #setup_complete?).
+        redis.set(batch_setup_gid, 'in_progress')
+
         # Schedule all child workers
         schedule_pending_jobs
+
+        # Mark as fully enqueued; children may now trigger completion.
+        redis.set(batch_setup_gid, 'done')
       end
 
       #

spec/cloudtasker/batch/job_spec.rb

@@ -388,6 +388,7 @@
         expect(child_worker).not_to receive(:schedule)
       end
 
+      after { expect(redis.get(batch.batch_setup_gid)).to be_nil }
       it { is_expected.to be_truthy }
     end
 
@@ -399,10 +400,95 @@
         expect(batch).to receive(:schedule_pending_jobs).and_return(true)
       end
 
+      after { expect(redis.get(batch.batch_setup_gid)).to eq('done') }
       it { is_expected.to be_truthy }
     end
   end
 
+  describe '#setup_complete?' do
+    subject { batch }
+
+    context 'when the batch is still enqueuing children' do
+      before { redis.set(batch.batch_setup_gid, 'in_progress') }
+
+      it { is_expected.not_to be_setup_complete }
+    end
+
+    context 'when the batch is fully enqueued' do
+      before { redis.set(batch.batch_setup_gid, 'done') }
+
+      it { is_expected.to be_setup_complete }
+    end
+
+    context 'without a marker (batch enqueued before this feature)' do
+      # Backward compatibility: batches in flight during an upgrade have no
+      # marker and must be treated as fully enqueued so they are not stranded.
+      it { is_expected.to be_setup_complete }
+    end
+  end
+
+  describe 'progressive batch expansion' do
+    # A running child may add more jobs to a batch that has already been set up
+    # (see docs/BATCH_JOBS.md "Expanding the parent batch"). The setup marker is
+    # written during the initial enqueue and must not interfere: the batch stays
+    # open while the expanding child runs, then completes normally.
+    let(:initial_child) { worker.new_instance }
+    let(:added_child) { worker.new_instance }
+
+    before do
+      allow_any_instance_of(Cloudtasker::Worker).to receive(:schedule)
+        .and_return(instance_double(Cloudtasker::CloudTask))
+
+      # Parent sets up with one initial child and marks the batch enqueued
+      batch.pending_jobs.push(initial_child)
+      batch.setup
+
+      # The initial child starts running, then expands the batch before completing
+      batch.update_state(initial_child.job_id, 'processing')
+      batch.pending_jobs.push(added_child)
+      batch.schedule_pending_jobs
+    end
+
+    it { expect(batch).to be_setup_complete }
+
+    context 'with the expanding child still running' do
+      before { batch.update_state(added_child.job_id, 'completed') }
+
+      # The expanding child's own 'processing' entry holds the batch open, so
+      # the freshly added child completing does not complete the batch.
+      it { expect(batch).not_to be_complete }
+    end
+
+    context 'with the expanding child and the added child both completed' do
+      before do
+        batch.update_state(added_child.job_id, 'completed')
+        batch.update_state(initial_child.job_id, 'completed')
+      end
+
+      it { expect(batch).to be_complete }
+      it { expect(batch).to be_setup_complete }
+    end
+  end
+
+  describe 'batch enqueued before the setup marker (upgrade in flight)' do
+    # A batch enqueued by a previous version has no setup marker. After an
+    # upgrade its remaining children must still be able to complete it - the
+    # parent already ran under the old code and cannot act as the fallback -
+    # otherwise the batch would be stranded and on_batch_complete never fire.
+    before do
+      batch.save
+      redis.hset(batch.batch_state_gid, child_batch.batch_id, 'processing')
+    end
+
+    it { expect(redis.get(batch.batch_setup_gid)).to be_nil }
+    it { expect(batch).to be_setup_complete }
+
+    it 'lets the final child complete the batch' do
+      expect(batch).to receive(:on_complete)
+      batch.on_child_complete(child_batch, :completed)
+    end
+  end
+
   describe '#update_state' do
     subject { batch.batch_state&.dig(child_id) }
 
@@ -592,7 +678,7 @@
     let(:complete) { true }
 
     before do
-      allow(batch).to receive_messages(complete?: complete, on_complete: true)
+      allow(batch).to receive_messages(complete?: complete, on_complete: true, setup_complete?: true)
       allow(batch).to receive(:update_state).with(child_batch.batch_id, status)
       allow(batch).to receive(:run_worker_callback)
       batch.pending_jobs.push(child_worker)
@@ -612,6 +698,15 @@
       it { is_expected.to be_truthy }
     end
 
+    context 'with batch complete but the parent still enqueuing children' do
+      before { allow(batch).to receive(:setup_complete?).and_return(false) }
+
+      after { expect(batch).to have_received(:update_state) }
+      after { expect(batch).to have_received(:run_worker_callback).with(:on_child_complete, child_batch.worker) }
+      after { expect(batch).not_to have_received(:on_complete) }
+      it { is_expected.to be_falsey }
+    end
+
     context 'with batch complete but completion already claimed by another child' do
       before do
         expect(redis).to receive(:set)
@@ -708,6 +803,7 @@
       [
         side_batch.batch_gid,
         side_batch.batch_state_gid,
+        side_batch.batch_setup_gid,
         side_batch.batch_state_count_gid('all'),
         side_batch.batch_state_count_gid('scheduled')
       ].sort