Repository: firecracker-microvm/firecracker
File: tests/host_tools/network.py (Python)

Code before:
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "ConnectTimeout=2",
            # Set up a persistent background connection
            "-o",
            "ControlMaster=auto",
            "-o",
            "ControlPersist=yes",
            *self.options,
            self.user_host,
            "true",
        ]

        try:
            # don't set a low timeout here, because otherwise we might get into a race condition
            # where ssh already forked off the persisted connection daemon, but gets killed here
            # before exiting itself. In that case, self._control_path will exist, and the retry
            # will hit the assert at the start of this function.
            self._exec(establish_cmd, check=True)
        except Exception:
            # cleanup the control socket to allow for retries
            if self._control_path.exists():
                self._control_path.unlink()
            raise

    def _check_liveness(self) -> int:
        """Checks whether the ControlPersist connection is still alive"""
        check_cmd = ["ssh", "-O", "check", *self.options, self.user_host]

        _, _, stderr = self._exec(check_cmd, check=True)

        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):
        """Closes the ControlPersist connection"""
        master_pid = self._check_liveness()

        stop_cmd = ["ssh", "-O", "stop", *self.options, self.user_host]

        _, _, stderr = self._exec(stop_cmd, check=True)

        assert "Stop listening request sent" in stderr

        try:
            with Timeout(5):
                utils.wait_process_termination(master_pid)

Review comment (suggestion):
Mh, if we do this then we can probably remove the assert altogether (and the comment, too). Although I think I slightly prefer not explicitly removing the path here. The idea behind this was that if the socket file ends up getting created, then we definitely have an SSH daemon process, and just retrying in this case will simply create more daemon processes. The assert prevent this from happening. Another option here would be to catch the exception and instead of unlinking, we do `self.close()`.

Code after:
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "ConnectTimeout=2",
            # Set up a persistent background connection
            "-o",
            "ControlMaster=auto",
            "-o",
            "ControlPersist=yes",
            *self.options,
            self.user_host,
            "true",
        ]

        try:
            # don't set a low timeout here, because otherwise we might get into a race condition
            # where ssh already forked off the persisted connection daemon, but gets killed here
            # before exiting itself. In that case, self._control_path will exist, and the retry
            # will hit the assert at the start of this function.
            self._exec(establish_cmd, check=True)
        except Exception:
            # if the control socket is present, then the daemon is running, and we should stop it
            # before retrying again
            if self._control_path.exists():
                self.close()
            raise

    def _check_liveness(self) -> int:
        """Checks whether the ControlPersist connection is still alive"""
        check_cmd = ["ssh", "-O", "check", *self.options, self.user_host]

        _, _, stderr = self._exec(check_cmd, check=True)

        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):
        """Closes the ControlPersist connection"""
        master_pid = self._check_liveness()

        stop_cmd = ["ssh", "-O", "stop", *self.options, self.user_host]

        _, _, stderr = self._exec(stop_cmd, check=True)

        assert "Stop listening request sent" in stderr

        try:
            with Timeout(5):