test: clean-up potentially dead microVMs

We have some tests that stress the microVM in a way that _might_ cause
it to die. We do check whether that's the case, however this is
inherently racy, i.e. the microVM might die on us after the check but
before the fixture logic tries to recycle it. Such a test is
test_balloon.py::test_deflate_oom.

In order to avoid such situations, teach MicroVM.kill() that there might
be cases where the microVM process has already died and let it finish
its cleanup gracefully.

Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

tests/framework/microvm.py

@@ -301,7 +301,7 @@ def mark_killed(self):
 
         self._killed = True
 
-    def kill(self):
+    def kill(self, might_be_dead=False):
         """All clean up associated with this microVM should go here."""
         # pylint: disable=subprocess-run-check
         # if it was already killed, return
@@ -314,7 +314,7 @@ def kill(self):
 
         # Kill all background SSH connections
         for connection in self._connections:
-            connection.close()
+            connection.close(strict=not might_be_dead)
 
         # We start with vhost-user backends,
         # because if we stop Firecracker first, the backend will want
@@ -325,24 +325,27 @@ def kill(self):
 
         assert (
             "Shutting down VM after intercepting signal" not in self.log_data
+            or might_be_dead
         ), self.log_data
 
+        # pylint: disable=bare-except
         try:
             if self.firecracker_pid:
                 os.kill(self.firecracker_pid, signal.SIGKILL)
 
             if self.screen_pid:
                 os.kill(self.screen_pid, signal.SIGKILL)
         except:
-            msg = (
-                "Failed to kill Firecracker Process. Did it already die (or did the UFFD handler process die and take it down)?"
-                if self.uffd_handler
-                else "Failed to kill Firecracker Process. Did it already die?"
-            )
+            if not might_be_dead:
+                msg = (
+                    "Failed to kill Firecracker Process. Did it already die (or did the UFFD handler process die and take it down)?"
+                    if self.uffd_handler
+                    else "Failed to kill Firecracker Process. Did it already die?"
+                )
 
-            self._dump_debug_information(msg)
+                self._dump_debug_information(msg)
 
-            raise
+                raise
 
         # if microvm was spawned then check if it gets killed
         if self._spawned:

tests/host_tools/network.py

@@ -145,21 +145,33 @@ def _init_connection(self):
                 self.close()
             raise
 
-    def _check_liveness(self) -> int:
-        """Checks whether the ControlPersist connection is still alive"""
+    def _check_liveness(self, strict=True) -> int | None:
+        """Checks whether the ControlPersist connection is still alive
+
+        It will return the pid of the ControlMaster if it is still running,
+        otherwise None
+        """
         check_cmd = ["ssh", "-O", "check", *self.options, self.user_host]
 
-        _, _, stderr = self._exec(check_cmd, check=True)
+        try:
+            _, _, stderr = self._exec(check_cmd, check=True)
+        except ChildProcessError:
+            if strict:
+                raise
+
+            return None
 
         pid_match = re.match(r"Master running \(pid=(\d+)\)", stderr)
 
         assert pid_match, f"SSH ControlMaster connection not alive anymore: {stderr}"
 
         return int(pid_match.group(1))
 
-    def close(self):
+    def close(self, strict=True):
         """Closes the ControlPersist connection"""
-        master_pid = self._check_liveness()
+        master_pid = self._check_liveness(strict)
+        if master_pid is None:
+            return
 
         stop_cmd = ["ssh", "-O", "stop", *self.options, self.user_host]
 
@@ -182,7 +194,7 @@ def run(self, cmd_string, timeout=100, *, check=False, debug=False):
 
         If `debug` is set, pass `-vvv` to `ssh`. Note that this will clobber stderr.
         """
-        self._check_liveness()
+        self._check_liveness(True)
 
         command = ["ssh", *self.options, self.user_host, cmd_string]
 

tests/integration_tests/functional/test_balloon.py

@@ -228,6 +228,9 @@ def test_deflate_on_oom(uvm_plain_any, deflate_on_oom):
             assert balloon_size_after < balloon_size_before, "Balloon did not deflate"
         else:
             assert balloon_size_after >= balloon_size_before, "Balloon deflated"
+            # Kill it here, letting the infrastructure know that the process might
+            # be dead already.
+            test_microvm.kill(might_be_dead=True)
 
 
 # pylint: disable=C0103