From a088fbba399e4d6644c700d81dc18fd138f9bae2 Mon Sep 17 00:00:00 2001
From: Chris Roberts <croberts@hashicorp.com>
Date: Tue, 13 Apr 2021 15:07:09 -0700
Subject: [PATCH] Properly wait for reboot process to start

    This adds a check to the `#wait_for_reboot` method on the linux
    guest reboot capability to determine if the a reboot is still
    in process. This prevents the reboot process from being initiated
    and the `#ready?` check on the guest being called before the
    system shutdown process has shutdown the communicator process.
---
 plugins/guests/linux/cap/reboot.rb            | 24 +++++++-
 .../plugins/guests/linux/cap/reboot_test.rb   | 60 ++++++++++++++++++-
 2 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/plugins/guests/linux/cap/reboot.rb b/plugins/guests/linux/cap/reboot.rb
index 0cf29ff06e4..3b3eb19a8a7 100644
--- a/plugins/guests/linux/cap/reboot.rb
+++ b/plugins/guests/linux/cap/reboot.rb
@@ -12,13 +12,16 @@ class Reboot
 
         def self.reboot(machine)
           @logger = Log4r::Logger.new("vagrant::linux::reboot")
+          reboot_script = "ps -q 1 -o comm=,start= > /tmp/.vagrant-reboot"
+
           if systemd?(machine.communicate)
-            reboot_script = "systemctl reboot"
+            reboot_cmd = "systemctl reboot"
           else
-            reboot_script = "reboot"
+            reboot_cmd = "reboot"
           end
 
           comm = machine.communicate
+          reboot_script += "; #{reboot_cmd}"
 
           @logger.debug("Issuing reboot command for guest")
           comm.sudo(reboot_script)
@@ -43,8 +46,23 @@ def self.reboot(machine)
         end
 
         def self.wait_for_reboot(machine)
-          while !machine.guest.ready?
+          caught = false
+          begin
+            check_script = 'grep "$(ps -q 1 -o comm=,start=)" /tmp/.vagrant-reboot'
+            while machine.guest.ready? && machine.communicate.execute(check_script, error_check: false) == 0
+              sleep 10
+            end
+          rescue
+            # The check script execution may result in an exception
+            # getting raised depending on the state of the communicator
+            # when executing. We'll allow for it to happen once, and then
+            # raise if we get an exception again
+            if caught
+              raise
+            end
+            caught = true
             sleep 10
+            retry
           end
         end
       end
diff --git a/test/unit/plugins/guests/linux/cap/reboot_test.rb b/test/unit/plugins/guests/linux/cap/reboot_test.rb
index 66c38a5c4c2..858f506544d 100644
--- a/test/unit/plugins/guests/linux/cap/reboot_test.rb
+++ b/test/unit/plugins/guests/linux/cap/reboot_test.rb
@@ -7,7 +7,7 @@
     VagrantPlugins::GuestLinux::Plugin.components.guest_capabilities[:linux].get(:wait_for_reboot)
   end
 
-  let(:machine) { double("machine") }
+  let(:machine) { double("machine", guest: guest) }
   let(:guest) { double("guest") }
   let(:communicator) { VagrantTests::DummyCommunicator::Communicator.new(machine) }
   let(:ui) { double("ui") }
@@ -106,4 +106,62 @@
       end
     end
   end
+
+  describe ".wait_for_reboot" do
+    before do
+      allow(machine).to receive(:communicate).and_return(communicator)
+      allow(communicator).to receive(:execute).and_return(0)
+      allow(guest).to receive(:ready?).and_return(false)
+    end
+
+    context "when guest is ready" do
+      before { expect(guest).to receive(:ready?).and_return(true) }
+
+      it "should sleep" do
+        expect(described_class).to receive(:sleep).with(10)
+        described_class.wait_for_reboot(machine)
+      end
+
+      context "when check script fails" do
+        before { expect(communicator).to receive(:execute).with(/grep/, any_args).and_return(1) }
+
+        it "should not sleep" do
+          expect(described_class).not_to receive(:sleep)
+          described_class.wait_for_reboot(machine)
+        end
+      end
+
+      context "when communicator raises an error" do
+        let(:error) { Class.new(StandardError) }
+
+        before do
+          expect(communicator).to receive(:execute).with(/grep/, any_args).and_raise(error)
+          expect(guest).to receive(:ready?).and_return(true)
+        end
+
+        it "should sleep once for exception and once for the guest being ready" do
+          expect(described_class).to receive(:sleep).with(10).twice
+          described_class.wait_for_reboot(machine)
+        end
+
+        context "when communicator raises error more than once" do
+          before { expect(communicator).to receive(:execute).with(/grep/, any_args).and_raise(error) }
+
+          it "should sleep once and raise error" do
+            expect(described_class).to receive(:sleep).with(10)
+            expect { described_class.wait_for_reboot(machine) }.to raise_error(error)
+          end
+        end
+      end
+    end
+
+    context "when guest is not ready" do
+      before { expect(guest).to receive(:ready?).and_return(false) }
+
+      it "should not sleep" do
+        expect(described_class).not_to receive(:sleep)
+        described_class.wait_for_reboot(machine)
+      end
+    end
+  end
 end