From 98f06a282f3f815a06e76975879f20f28c535fc2 Mon Sep 17 00:00:00 2001
From: Adam Ruzicka <aruzicka@redhat.com>
Date: Wed, 22 Nov 2023 13:02:48 +0100
Subject: [PATCH] Fixes #36937 - Rework proxy batch triggering

Now it should be more strict about when it finishes. Previously this
action could finish too early, even before everything was dispatched to
the proxy.
---
 app/lib/actions/trigger_proxy_batch.rb        | 21 ++++++++++++-------
 test/unit/actions/trigger_proxy_batch_test.rb |  7 +++++++
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/app/lib/actions/trigger_proxy_batch.rb b/app/lib/actions/trigger_proxy_batch.rb
index 72cd5aade..d5bf45271 100644
--- a/app/lib/actions/trigger_proxy_batch.rb
+++ b/app/lib/actions/trigger_proxy_batch.rb
@@ -18,9 +18,10 @@ def run(event = nil)
           init_counts and suspend
         end
       when TriggerNextBatch
-        trigger_remote_tasks_batches(event.batches) and suspend
+        trigger_remote_tasks_batches(event.batches)
       when TriggerLastBatch
-        trigger_remote_tasks_batch and on_finish
+        output[:planning_finished] = true
+        trigger_remote_tasks_batches
       when ::Dynflow::Action::Skip
         # do nothing
       end
@@ -28,6 +29,7 @@ def run(event = nil)
 
     def trigger_remote_tasks_batches(amount = 1)
       amount.times { trigger_remote_tasks_batch }
+      done? ? on_finish : suspend
     end
 
     def trigger_remote_tasks_batch
@@ -55,15 +57,18 @@ def init_counts
     end
 
     def check_finish
-      if output[:planned_count] + output[:failed_count] + batch_size >= input[:total_count]
-        trigger_remote_tasks_batch and on_finish
-      else
-        suspend
-      end
+      return on_finish if done?
+
+      # If we're not done yet, try to trigger anything (if available)
+      # and then either finish or suspend
+      trigger_remote_tasks_batches
     end
 
     def done?
-      output[:planned_count] + output[:failed_count] >= input[:total_count]
+      # We're done when we've either:
+      # - dispatched everything
+      # - received the last message
+      output[:planned_count] + output[:failed_count] >= input[:total_count] || output[:planning_finished]
     end
 
     def remote_tasks
diff --git a/test/unit/actions/trigger_proxy_batch_test.rb b/test/unit/actions/trigger_proxy_batch_test.rb
index b85212e9a..55cc9f155 100644
--- a/test/unit/actions/trigger_proxy_batch_test.rb
+++ b/test/unit/actions/trigger_proxy_batch_test.rb
@@ -32,6 +32,13 @@ class TriggerProxyBatchTest < ActiveSupport::TestCase
           triggered.output[:planned_count] = ((total_count - 1) / batch_size) * batch_size
           run_action(triggered)
         end
+
+        it "finishes after the last batch even if the counts don't match" do
+          Actions::TriggerProxyBatch.any_instance.expects(:trigger_remote_tasks_batch).once
+          triggered.output[:planned_count] = 0
+          action = run_action(triggered, Actions::TriggerProxyBatch::TriggerLastBatch)
+          _(action.state).must_equal :success
+        end
       end
 
       describe '#trigger_remote_tasks_batch' do