From 26915dcb72c36318c89461362487973f64e4a28f Mon Sep 17 00:00:00 2001
From: James Falcon <therealfalcon@gmail.com>
Date: Wed, 13 Jan 2021 10:34:12 -0600
Subject: [PATCH] Track boot id on each boot.

When we signal reboot and attempt to reconnect, we can wind up
reconnecting before the reboot ever took place. This change updates the boot
process to ensure we've obtained a new boot id before proceeding.
---
 pycloudlib/instance.py            | 32 +++++++++++++++----------------
 pycloudlib/tests/test_instance.py |  4 ++--
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/pycloudlib/instance.py b/pycloudlib/instance.py
index b13ceaa8..d55471b1 100644
--- a/pycloudlib/instance.py
+++ b/pycloudlib/instance.py
@@ -35,6 +35,7 @@ def __init__(self, key_pair):
         self.port = '22'
         self.username = 'ubuntu'
         self.connect_timeout = 60
+        self.boot_id = ''
 
     @property
     @abstractmethod
@@ -99,7 +100,7 @@ def wait(self, *, raise_on_cloudinit_failure: bool = True):
             complete successfully.
         """
         self._wait_for_instance_start()
-        self._wait_for_execute()
+        self._wait_for_boot_id()
         self._wait_for_cloudinit(raise_on_failure=raise_on_cloudinit_failure)
 
     @abstractmethod
@@ -370,29 +371,28 @@ def _tmpfile(self):
         self._tmp_count += 1
         return path
 
-    def _wait_for_execute(self):
+    def _wait_for_boot_id(self):
         """Wait until we can execute a command in the instance."""
-        self._log.debug('_wait_for_execute to complete')
+        self._log.debug('_wait_for_boot_id to complete')
 
-        test_instance_command = "whoami"
-        result = self.execute(test_instance_command)
-        if result.failed:
-            retries = 10
-            while retries:
-                result = self.execute(test_instance_command)
-
-                if result.ok:
+        boot_id_command = "cat /proc/sys/kernel/random/boot_id"
+        retries = 300
+        while retries:
+            result = self.execute(boot_id_command)
+            if result.ok:
+                boot_id = result.stdout.strip()
+                if boot_id != self.boot_id:
+                    self.boot_id = boot_id
                     break
 
-                retries -= 1
-                time.sleep(10)
-
-        if result.failed:
+            retries -= 1
+            time.sleep(1)
+        else:
             raise OSError(
                 "{}\n{}".format(
                     "Instance can't be reached",
                     "Failed to execute {} command".format(
-                        test_instance_command)
+                        boot_id_command)
                 )
             )
 
diff --git a/pycloudlib/tests/test_instance.py b/pycloudlib/tests/test_instance.py
index 3eab0901..4b4a01ff 100644
--- a/pycloudlib/tests/test_instance.py
+++ b/pycloudlib/tests/test_instance.py
@@ -65,7 +65,7 @@ def test_wait(self, raise_on_cloudinit_failure, concrete_instance_cls):
         with mock.patch.multiple(
             instance,
             _wait_for_instance_start=mock.DEFAULT,
-            _wait_for_execute=mock.DEFAULT,
+            _wait_for_boot_id=mock.DEFAULT,
             _wait_for_cloudinit=mock.DEFAULT,
         ) as mocks:
             if raise_on_cloudinit_failure is not None:
@@ -76,7 +76,7 @@ def test_wait(self, raise_on_cloudinit_failure, concrete_instance_cls):
                 instance.wait()
 
         assert 1 == mocks["_wait_for_instance_start"].call_count
-        assert 1 == mocks["_wait_for_execute"].call_count
+        assert 1 == mocks["_wait_for_boot_id"].call_count
         assert 1 == mocks["_wait_for_cloudinit"].call_count
         kwargs = mocks["_wait_for_cloudinit"].call_args[1]
         if raise_on_cloudinit_failure is None: