From f07666070dcf2fb626f9f279c4f463860751ea1a Mon Sep 17 00:00:00 2001
From: Donal McBreen <donal@basecamp.com>
Date: Thu, 4 Apr 2024 11:08:00 +0100
Subject: [PATCH] Wait for all threads to finish on parallel runner

If there's an error in one thread, we wait for earlier hosts to
complete but not for later ones.

Ensure consistent behaviour across all hosts by saving the exception and
joining each thread in turn before returning.
---
 lib/sshkit/runners/parallel.rb     | 19 ++++++++++++++++++-
 test/unit/runners/test_parallel.rb | 25 +++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/lib/sshkit/runners/parallel.rb b/lib/sshkit/runners/parallel.rb
index fd5d018a..b1543029 100644
--- a/lib/sshkit/runners/parallel.rb
+++ b/lib/sshkit/runners/parallel.rb
@@ -16,7 +16,24 @@ def execute
             end
           end
         end
-        threads.each(&:join)
+
+        wait_for_threads(threads)
+      end
+
+      private
+
+      def wait_for_threads(threads)
+        exception = nil
+
+        threads.map do |t|
+          begin
+            t.join
+          rescue ExecuteError => e
+            exception ||= e
+          end
+        end
+
+        raise exception if exception
       end
     end
 
diff --git a/test/unit/runners/test_parallel.rb b/test/unit/runners/test_parallel.rb
index 4a90900b..c6a066bd 100644
--- a/test/unit/runners/test_parallel.rb
+++ b/test/unit/runners/test_parallel.rb
@@ -13,6 +13,31 @@ def test_wraps_ruby_standard_error_in_execute_error
         assert_match(/deployer@example/, error.message)
         assert_match(/oh no!/, error.message)
       end
+
+      def test_waits_for_all_threads_to_finish_on_error
+        hosts = [Host.new("deployer@example"), Host.new("deployer@example2"), Host.new("deployer@example3")]
+        completed_one, completed_three = false, false
+        runner = Parallel.new(hosts) do |host|
+          case host.hostname
+          when "example"
+            sleep 0.1
+            completed_one = true
+          when "example2"
+            raise "Boom!"
+          when "example3"
+            sleep 0.3
+            completed_three = true
+          end
+        end
+
+        error = assert_raises(SSHKit::Runner::ExecuteError) do
+          runner.execute
+        end
+        assert_match(/deployer@example2/, error.message)
+        assert_match(/Boom!/, error.message)
+        assert completed_one
+        assert completed_three
+      end
     end
   end
 end