From 7fe24d5048f611b01076abfda26fbbd74df9b7d6 Mon Sep 17 00:00:00 2001
From: Donal McBreen <donal@basecamp.com>
Date: Mon, 1 May 2023 16:34:32 +0100
Subject: [PATCH] Check all hosts before rolling back

Hosts could end up out of sync with each other if prune commands are run
manually or when new hosts are added.

Before rolling back confirm that the required container is available on
all hosts and roles.
---
 lib/mrsk/cli/main.rb  | 23 ++++++++++++++++-------
 test/cli/main_test.rb | 19 +++++++++++++++----
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/lib/mrsk/cli/main.rb b/lib/mrsk/cli/main.rb
index 0a3d2ce0..0ead3285 100644
--- a/lib/mrsk/cli/main.rb
+++ b/lib/mrsk/cli/main.rb
@@ -236,15 +236,24 @@ class Mrsk::Cli::Main < Mrsk::Cli::Base
   subcommand "lock", Mrsk::Cli::Lock
 
   private
-    def container_available?(version, host: MRSK.primary_host)
-      available = nil
-
-      on(host) do
-        first_role = MRSK.roles_on(host).first
-        available = capture_with_info(*MRSK.app(role: first_role).container_id_for_version(version)).present?
+    def container_available?(version)
+      begin
+        on(MRSK.hosts) do
+          MRSK.roles_on(host).each do |role|
+            container_id = capture_with_info(*MRSK.app(role: role).container_id_for_version(version))
+            raise "Container not found" unless container_id.present?
+          end
+        end
+      rescue SSHKit::Runner::ExecuteError => e
+        if e.message =~ /Container not found/
+          say "Cannot rollback: #{e.message}"
+          return false
+        else
+          raise
+        end
       end
 
-      available
+      true
     end
 
     def deploy_options
diff --git a/test/cli/main_test.rb b/test/cli/main_test.rb
index 5cac1132..9ba0e54f 100644
--- a/test/cli/main_test.rb
+++ b/test/cli/main_test.rb
@@ -145,7 +145,8 @@ class CliMainTest < CliTestCase
   end
 
   test "rollback bad version" do
-    # Mrsk::Cli::Main.any_instance.stubs(:container_available?).returns(false)
+    Thread.report_on_exception = false
+
     run_command("details") # Preheat MRSK const
 
     run_command("rollback", "nonsense").tap do |output|
@@ -155,9 +156,19 @@ class CliMainTest < CliTestCase
   end
 
   test "rollback good version" do
-    Mrsk::Cli::Main.any_instance.stubs(:container_available?).returns(true)
-    SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info).with(:docker, :ps, "--filter", "label=service=app", "--filter", "label=role=web", "--filter", "status=running", "--latest", "--format", "\"{{.Names}}\"", "|", "grep -oE \"\\-[^-]+$\"", "|", "cut -c 2-").returns("version-to-rollback\n").at_least_once
-    SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info).with(:docker, :ps, "--filter", "label=service=app", "--filter", "label=role=workers", "--filter", "status=running", "--latest", "--format", "\"{{.Names}}\"", "|", "grep -oE \"\\-[^-]+$\"", "|", "cut -c 2-").returns("version-to-rollback\n").at_least_once
+    SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info)
+      .with(:docker, :container, :ls, "--all", "--filter", "name=^app-web-123$", "--quiet")
+      .returns("version-to-rollback\n").at_least_once
+    SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info)
+      .with(:docker, :container, :ls, "--all", "--filter", "name=^app-workers-123$", "--quiet")
+      .returns("version-to-rollback\n").at_least_once
+    SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info)
+      .with(:docker, :ps, "--filter", "label=service=app", "--filter", "label=role=web", "--filter", "status=running", "--latest", "--format", "\"{{.Names}}\"", "|", "grep -oE \"\\-[^-]+$\"", "|", "cut -c 2-")
+      .returns("version-to-rollback\n").at_least_once
+    SSHKit::Backend::Abstract.any_instance.expects(:capture_with_info)
+      .with(:docker, :ps, "--filter", "label=service=app", "--filter", "label=role=workers", "--filter", "status=running", "--latest", "--format", "\"{{.Names}}\"", "|", "grep -oE \"\\-[^-]+$\"", "|", "cut -c 2-")
+      .returns("version-to-rollback\n").at_least_once
+
 
     run_command("rollback", "123", config_file: "deploy_with_accessories").tap do |output|
       assert_match "Start version 123", output