From c4df440c79f02b4602f0196ee79fd7b9c3cde80d Mon Sep 17 00:00:00 2001
From: Donal McBreen <donal@basecamp.com>
Date: Mon, 10 Apr 2023 15:08:48 +0100
Subject: [PATCH] Improved deploy lock acquisition

1. Don't raise lock error for non-lock issues during lock acquire
  (see https://github.com/mrsked/mrsk/pull/181)
2. If there is an error while the lock is held, don't release the lock
  and send a warning to stderr
---
 lib/mrsk/cli.rb           |  1 +
 lib/mrsk/cli/base.rb      | 12 +++++++-----
 lib/mrsk/cli/lock.rb      |  4 ++--
 test/cli/cli_test_case.rb |  6 +++++-
 test/cli/main_test.rb     | 40 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/lib/mrsk/cli.rb b/lib/mrsk/cli.rb
index 62323076..c974a814 100644
--- a/lib/mrsk/cli.rb
+++ b/lib/mrsk/cli.rb
@@ -1,4 +1,5 @@
 module Mrsk::Cli
+  class LockError < StandardError; end
 end
 
 # SSHKit uses instance eval, so we need a global const for ergonomics
diff --git a/lib/mrsk/cli/base.rb b/lib/mrsk/cli/base.rb
index ee1851bf..a0eb200a 100644
--- a/lib/mrsk/cli/base.rb
+++ b/lib/mrsk/cli/base.rb
@@ -6,8 +6,6 @@ module Mrsk::Cli
   class Base < Thor
     include SSHKit::DSL
 
-    class LockError < StandardError; end
-
     def self.exit_on_failure?() true end
 
     class_option :verbose, type: :boolean, aliases: "-v", desc: "Detailed logging"
@@ -82,8 +80,11 @@ module Mrsk::Cli
         acquire_lock
 
         yield
-      ensure
+
         release_lock
+      rescue
+        error "  \e[31mDeploy lock was not released\e[0m" if MRSK.lock_count > 0
+        raise
       end
 
       def acquire_lock
@@ -95,9 +96,10 @@ module Mrsk::Cli
       rescue SSHKit::Runner::ExecuteError => e
         if e.message =~ /cannot create directory/
           invoke "mrsk:cli:lock:status", []
+          raise LockError, "Deploy lock found"
+        else
+          raise e
         end
-
-        raise LockError, "Deploy lock found"
       end
 
       def release_lock
diff --git a/lib/mrsk/cli/lock.rb b/lib/mrsk/cli/lock.rb
index 3514e282..f3fdcf4b 100644
--- a/lib/mrsk/cli/lock.rb
+++ b/lib/mrsk/cli/lock.rb
@@ -12,7 +12,7 @@ class Mrsk::Cli::Lock < Mrsk::Cli::Base
     message = options[:message]
     handle_missing_lock do
       on(MRSK.primary_host) { execute *MRSK.lock.acquire(message, MRSK.config.version) }
-      say "Set the deploy lock"
+      say "Acquired the deploy lock"
     end
   end
 
@@ -20,7 +20,7 @@ class Mrsk::Cli::Lock < Mrsk::Cli::Base
   def release
     handle_missing_lock do
       on(MRSK.primary_host) { execute *MRSK.lock.release }
-      say "Removed the deploy lock"
+      say "Released the deploy lock"
     end
   end
 
diff --git a/test/cli/cli_test_case.rb b/test/cli/cli_test_case.rb
index 496d6a27..01cf0019 100644
--- a/test/cli/cli_test_case.rb
+++ b/test/cli/cli_test_case.rb
@@ -22,4 +22,8 @@ class CliTestCase < ActiveSupport::TestCase
     def stdouted
       capture(:stdout) { yield }.strip
     end
-end
+
+    def stderred
+      capture(:stderr) { yield }.strip
+    end
+  end
diff --git a/test/cli/main_test.rb b/test/cli/main_test.rb
index b27e5630..8d00b277 100644
--- a/test/cli/main_test.rb
+++ b/test/cli/main_test.rb
@@ -53,6 +53,46 @@ class CliMainTest < CliTestCase
     end
   end
 
+  test "deploy when locked" do
+    Thread.report_on_exception = false
+
+    SSHKit::Backend::Abstract.any_instance.stubs(:execute)
+      .with { |*arg| arg[0..1] == [:mkdir, :mrsk_lock] }
+      .raises(RuntimeError, "mkdir: cannot create directory ‘mrsk_lock’: File exists")
+
+    Mrsk::Cli::Base.any_instance.expects(:invoke).with("mrsk:cli:lock:status", [])
+
+    assert_raises(Mrsk::Cli::LockError) do
+      run_command("deploy")
+    end
+  end
+
+  test "deploy error when locking" do
+    Thread.report_on_exception = false
+
+    SSHKit::Backend::Abstract.any_instance.stubs(:execute)
+      .with { |*arg| arg[0..1] == [:mkdir, :mrsk_lock] }
+      .raises(SocketError, "getaddrinfo: nodename nor servname provided, or not known")
+
+    assert_raises(SSHKit::Runner::ExecuteError) do
+      run_command("deploy")
+    end
+  end
+
+  test "deploy errors leave lock in place" do
+    invoke_options = { "config_file" => "test/fixtures/deploy_simple.yml", "skip_broadcast" => false, "version" => "999" }
+
+    Mrsk::Cli::Main.any_instance.expects(:invoke)
+      .with("mrsk:cli:server:bootstrap", [], invoke_options)
+      .raises(RuntimeError)
+
+    assert_equal 0, MRSK.lock_count
+    assert_raises(RuntimeError) do
+      stderred { run_command("deploy") }
+    end
+    assert_equal 1, MRSK.lock_count
+  end
+
   test "redeploy" do
     invoke_options = { "config_file" => "test/fixtures/deploy_simple.yml", "skip_broadcast" => false, "version" => "999" }