microvm.py

# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

"""Classes for working with microVMs.

This module defines `Microvm`, which can be used to create, test drive, and
destroy microvms.

- Use the Firecracker Open API spec to populate Microvm API resource URLs.
"""

# pylint:disable=too-many-lines

import json
import logging
import os
import re
import select
import shutil
import signal
import time
import uuid
from collections import namedtuple
from dataclasses import dataclass
from enum import Enum, auto
from functools import cached_property, lru_cache
from pathlib import Path
from typing import Optional

import psutil
from tenacity import Retrying, retry, stop_after_attempt, wait_fixed

import host_tools.cargo_build as build_tools
import host_tools.network as net_tools
from framework import utils
from framework.defs import DEFAULT_BINARY_DIR, MAX_API_CALL_DURATION_MS
from framework.guest import GuestDistro
from framework.http_api import Api
from framework.jailer import JailerContext
from framework.microvm_helpers import MicrovmHelpers
from framework.properties import global_props
from framework.utils_cpu_templates import get_cpu_template_name
from framework.utils_drive import VhostUserBlkBackend, VhostUserBlkBackendType
from framework.utils_uffd import spawn_pf_handler, uffd_handler
from host_tools.fcmetrics import FCMetricsMonitor
from host_tools.memory import MemoryMonitor

LOG = logging.getLogger("microvm")


class SnapshotType(Enum):
    """Supported snapshot types."""

    FULL = auto()
    DIFF = auto()
    DIFF_MINCORE = auto()

    def __repr__(self):
        cls_name = self.__class__.__name__
        return f"{cls_name}.{self.name}"

    @property
    def needs_rebase(self) -> bool:
        """Does this snapshot type need rebasing on top of a base snapshot before restoration?"""
        return self in [SnapshotType.DIFF, SnapshotType.DIFF_MINCORE]

    @property
    def needs_dirty_page_tracking(self) -> bool:
        """Does taking this snapshot type require dirty page tracking to be enabled?"""
        return self == SnapshotType.DIFF

    @property
    def api_type(self) -> str:
        """Converts this `SnapshotType` to the string value expected by the Firecracker API"""
        match self:
            case SnapshotType.FULL:
                return "Full"
            case SnapshotType.DIFF | SnapshotType.DIFF_MINCORE:
                return "Diff"


def hardlink_or_copy(src, dst):
    """If src and dst are in the same device, hardlink. Otherwise, copy."""
    dst.touch(exist_ok=False)
    if dst.stat().st_dev == src.stat().st_dev:
        dst.unlink()
        dst.hardlink_to(src)
    else:
        shutil.copyfile(src, dst)


@dataclass(frozen=True, repr=True)
class Snapshot:
    """A Firecracker snapshot"""

    vmstate: Path
    mem: Path
    net_ifaces: list
    disks: dict
    ssh_key: Path
    snapshot_type: SnapshotType
    meta: dict

    def rebase_snapshot(
        self, base, use_snapshot_editor=False, binary_dir=DEFAULT_BINARY_DIR
    ):
        """Rebases current incremental snapshot onto a specified base layer."""
        if not self.snapshot_type.needs_rebase:
            raise ValueError(f"Cannot rebase {self.snapshot_type}")
        if use_snapshot_editor:
            build_tools.run_snap_editor_rebase(
                base.mem, self.mem, binary_dir=binary_dir
            )
        else:
            build_tools.run_rebase_snap_bin(base.mem, self.mem)

        new_args = self.__dict__ | {"mem": base.mem}
        return Snapshot(**new_args)

    def copy_to_chroot(self, chroot) -> "Snapshot":
        """
        Move all the snapshot files into the microvm jail.
        Use different names so a snapshot doesn't overwrite our original snapshot.
        """
        mem_src = chroot / self.mem.with_suffix(".src").name
        hardlink_or_copy(self.mem, mem_src)
        vmstate_src = chroot / self.vmstate.with_suffix(".src").name
        hardlink_or_copy(self.vmstate, vmstate_src)

        return Snapshot(
            vmstate=vmstate_src,
            mem=mem_src,
            net_ifaces=self.net_ifaces,
            disks=self.disks,
            ssh_key=self.ssh_key,
            snapshot_type=self.snapshot_type,
            meta=self.meta,
        )

    @classmethod
    # TBD when Python 3.11: -> Self
    def load_from(cls, src: Path) -> "Snapshot":
        """Load a snapshot saved with `save_to`"""
        snap_json = src / "snapshot.json"
        obj = json.loads(snap_json.read_text())
        return cls(
            vmstate=src / obj["vmstate"],
            mem=src / obj["mem"],
            net_ifaces=[net_tools.NetIfaceConfig(**d) for d in obj["net_ifaces"]],
            disks={dsk: src / p for dsk, p in obj["disks"].items()},
            ssh_key=src / obj["ssh_key"],
            snapshot_type=SnapshotType(obj["snapshot_type"]),
            meta=obj["meta"],
        )

    def save_to(self, dst: Path):
        """Serialize snapshot details to `dst`

        Deserialize the snapshot with `load_from`
        """
        for path in [self.vmstate, self.mem, self.ssh_key]:
            new_path = dst / path.name
            hardlink_or_copy(path, new_path)
        new_disks = {}
        for disk_id, path in self.disks.items():
            new_path = dst / path.name
            hardlink_or_copy(path, new_path)
            new_disks[disk_id] = new_path.name
        obj = {
            "vmstate": self.vmstate.name,
            "mem": self.mem.name,
            "net_ifaces": [x.__dict__ for x in self.net_ifaces],
            "disks": new_disks,
            "ssh_key": self.ssh_key.name,
            "snapshot_type": self.snapshot_type.value,
            "meta": self.meta,
        }
        snap_json = dst / "snapshot.json"
        snap_json.write_text(json.dumps(obj))

    def delete(self):
        """Delete the backing files from disk."""
        self.mem.unlink()
        self.vmstate.unlink()


class HugePagesConfig(str, Enum):
    """Enum describing the huge pages configurations supported Firecracker"""

    NONE = "None"
    HUGETLBFS_2MB = "2M"


# pylint: disable=R0904
class Microvm:
    """Class to represent a Firecracker microvm.

    A microvm is described by a unique identifier, a path to all the resources
    it needs in order to be able to start and the binaries used to spawn it.
    Besides keeping track of microvm resources and exposing microvm API
    methods, `spawn()` and `kill()` can be used to start/end the microvm
    process.
    """

    def __init__(
        self,
        microvm_id: str,
        fc_binary_path: Path,
        jailer_binary_path: Path,
        netns: net_tools.NetNs,
        monitor_memory: bool = True,
        jailer_kwargs: Optional[dict] = None,
        numa_node=None,
        custom_cpu_template: Path = None,
        pci: bool = False,
    ):
        """Set up microVM attributes, paths, and data structures."""
        # pylint: disable=too-many-statements
        # Unique identifier for this machine.
        assert microvm_id is not None
        self._microvm_id = microvm_id

        self.kernel_file = None
        self.rootfs_file = None
        self.distro = None
        self.ssh_key = None
        self.initrd_file = None
        self.boot_args = None
        self.uffd_handler = None

        self.fc_binary_path = Path(fc_binary_path)
        assert fc_binary_path.exists()
        self.jailer_binary_path = Path(jailer_binary_path)
        assert jailer_binary_path.exists()

        jailer_kwargs = jailer_kwargs or {}
        self.netns = netns
        # Create the jailer context associated with this microvm.
        self.jailer = JailerContext(
            jailer_id=self._microvm_id,
            exec_file=self.fc_binary_path,
            netns=netns,
            new_pid_ns=True,
            **jailer_kwargs,
        )

        self.pci_enabled = pci
        if pci:
            self.jailer.extra_args["enable-pci"] = None

        # Copy the /etc/localtime file in the jailer root
        self.jailer.jailed_path("/etc/localtime", subdir="etc")

        self._screen_pid = None

        self.time_api_requests = global_props.host_linux_version != "6.1"
        # disable the HTTP API timings as they cause a lot of false positives
        if int(os.environ.get("PYTEST_XDIST_WORKER_COUNT", 1)) > 1:
            self.time_api_requests = False

        self.monitors = []
        self.memory_monitor = None
        if monitor_memory:
            self.memory_monitor = MemoryMonitor(self)
            self.monitors.append(self.memory_monitor)

        self.api = None
        self.log_file = None
        self.serial_out_path = None
        self.metrics_file = None
        self._spawned = False
        self._killed = False

        # device dictionaries
        self.iface = {}
        self.disks = {}
        self.disks_vhost_user = {}
        self.vcpus_count = None
        self.mem_size_bytes = None
        self.cpu_template_name = "None"
        # The given custom CPU template will be set in basic_config() but could
        # be overwritten via set_cpu_template().
        self.custom_cpu_template = custom_cpu_template

        self._connections = []

        self._pre_cmd = []
        if numa_node:
            node_str = str(numa_node)
            self.add_pre_cmd([["numactl", "-N", node_str, "-m", node_str]])

        # MMDS content from file
        self.metadata_file = None

        self.help = MicrovmHelpers(self)

        self.gdb_socket = None

    def __repr__(self):
        return f"<Microvm id={self.id}>"

    def mark_killed(self):
        """
        Marks this `Microvm` as killed, meaning test tear down should not try to kill it

        raises an exception if the Firecracker process managing this VM is not actually dead
        """
        if self.firecracker_pid is not None:
            utils.wait_process_termination(self.firecracker_pid)

        self._killed = True

    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
        if self._killed:
            return

        # Stop any registered monitors
        for monitor in self.monitors:
            monitor.stop()

        # Kill all background SSH connections
        for connection in self._connections:
            connection.close(strict=not might_be_dead)

        # We start with vhost-user backends,
        # because if we stop Firecracker first, the backend will want
        # to exit as well and this will cause a race condition.
        for backend in self.disks_vhost_user.values():
            backend.kill()
        self.disks_vhost_user.clear()

        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:
            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)

                raise

        # if microvm was spawned then check if it gets killed
        if self._spawned:
            # Wait until the Firecracker process is actually dead
            utils.wait_process_termination(self.firecracker_pid)

            # The following logic guards us against the case where `firecracker_pid` for some
            # reason is the wrong PID, e.g. this is a regression test for
            # https://github.com/firecracker-microvm/firecracker/pull/4442/commits/d63eb7a65ffaaae0409d15ed55d99ecbd29bc572

            # filter ps results for the jailer's unique id
            _, stdout, stderr = utils.run_cmd(
                f"ps ax -o pid,cmd -ww | grep {self.jailer.jailer_id}"
            )

            assert not stderr, f"error querying processes using `ps`: {stderr}"

            offenders = []
            for proc in stdout.splitlines():
                _, cmd = proc.lower().split(maxsplit=1)
                if "firecracker" in proc and not cmd.startswith("screen"):
                    offenders.append(proc)

            # make sure firecracker was killed
            assert not offenders, (
                f"Firecracker reported its pid {self.firecracker_pid}, which was killed, but there still exist processes using the supposedly dead Firecracker's jailer_id: \n"
                + "\n".join(offenders)
            )

        if self.uffd_handler and self.uffd_handler.is_running():
            self.uffd_handler.kill()

        # Mark the microVM as not spawned, so we avoid trying to kill twice.
        self._spawned = False
        self._killed = True

        if self.time_api_requests:
            self._validate_api_response_times()

        if self.memory_monitor:
            self.memory_monitor.check_samples()

    def _validate_api_response_times(self):
        """
        Parses the firecracker logs for information regarding api server request processing times, and asserts they
        are within acceptable bounds.
        """
        # Log messages are either
        # 2023-06-16T07:45:41.767987318 [fc44b23e-ce47-4635-9549-5779a6bd9cee:fc_api] The API server received a Get request on "/mmds".
        # or
        # 2023-06-16T07:47:31.204704732 [2f2427c7-e4de-4226-90e6-e3556402be84:fc_api] The API server received a Put request on "/actions" with body "{\"action_type\": \"InstanceStart\"}".
        api_request_regex = re.compile(
            r"\] The API server received a (?P<method>\w+) request on \"(?P<url>(/(\w|-)*)+)\"( with body (?P<body>.*))?\."
        )
        api_request_times_regex = re.compile(
            r"\] Total previous API call duration: (?P<execution_time>\d+) us.$"
        )

        # Note: Processing of api requests is synchronous, so these messages cannot be torn by concurrency effects
        log_lines = self.log_data.split("\n")

        ApiCall = namedtuple("ApiCall", "method url body")

        current_call = None

        for log_line in log_lines:
            match = api_request_regex.search(log_line)

            if match:
                if current_call is not None:
                    raise Exception(
                        f"API call duration log entry for {current_call.method} {current_call.url} with body {current_call.body} is missing!"
                    )

                current_call = ApiCall(
                    match.group("method"), match.group("url"), match.group("body")
                )

            match = api_request_times_regex.search(log_line)

            if match:
                if current_call is None:
                    raise Exception(
                        "Got API call duration log entry before request entry"
                    )

                if current_call.url not in ["/snapshot/create", "/snapshot/load"]:
                    exec_time = float(match.group("execution_time")) / 1000.0

                    assert (
                        exec_time <= MAX_API_CALL_DURATION_MS
                    ), f"{current_call.method} {current_call.url} API call exceeded maximum duration: {exec_time} ms. Body: {current_call.body}"

                current_call = None

    @property
    def firecracker_version(self):
        """Return the version of the Firecracker executable."""
        _, stdout, _ = utils.check_output(f"{self.fc_binary_path} --version")
        return re.match(r"^Firecracker v(.+)", stdout.partition("\n")[0]).group(1)

    @property
    def path(self):
        """Return the path on disk used that represents this microVM."""
        return self.jailer.chroot_base_with_id()

    # some functions use this
    fsfiles = path

    @property
    def id(self):
        """Return the unique identifier of this microVM."""
        return self._microvm_id

    @property
    def log_data(self):
        """Return the log data."""
        if self.log_file is None:
            return ""
        return self.log_file.read_text()

    @property
    def state(self):
        """Get the InstanceInfo property and return the state field."""
        return self.api.describe.get().json()["state"]

    @cached_property
    def firecracker_pid(self):
        """Return Firecracker's PID

        Reads the pid from a file created by jailer.
        """
        if not self._spawned:
            return None

        # Read the PID from Firecracker's pidfile. Retry if
        # file doesn't exist yet, or doesn't yet contain an integer
        for attempt in Retrying(
            stop=stop_after_attempt(5),
            wait=wait_fixed(0.1),
            reraise=True,
        ):
            with attempt:
                return int(self.jailer.pid_file.read_text(encoding="ascii"))

    @cached_property
    def ps(self):
        """Returns a handle to the psutil.Process for this VM"""
        return psutil.Process(self.firecracker_pid)

    @property
    def dimensions(self):
        """Gets a default set of cloudwatch dimensions describing the configuration of this microvm"""
        return {
            "instance": global_props.instance,
            "cpu_model": global_props.cpu_model,
            "host_kernel": f"linux-{global_props.host_linux_version}",
            "guest_kernel": self.kernel_file.stem[2:],
            "rootfs": self.rootfs_file.name,
            "vcpus": str(self.vcpus_count),
            "guest_memory": f"{self.mem_size_bytes / (1024 * 1024)}MB",
            "pci": f"{self.pci_enabled}",
        }

    @property
    def guest_kernel_version(self):
        """Get the guest kernel version from the filename

        It won't work if the file name does not like name-X.Y.Z
        """
        splits = self.kernel_file.name.split("-")
        if len(splits) < 2:
            return None
        return tuple(int(x) for x in splits[1].split("."))

    def get_metrics(self):
        """Return iterator to metric data points written by FC"""
        with self.metrics_file.open() as fd:
            for line in fd:
                if not line.endswith("}\n"):
                    LOG.warning("Line is not a proper JSON object. Partial write?")
                    continue
                yield json.loads(line)

    def get_all_metrics(self):
        """Return all metric data points written by FC."""
        return list(self.get_metrics())

    def flush_metrics(self):
        """Flush the microvm metrics and get the latest datapoint"""
        self.api.actions.put(action_type="FlushMetrics")
        # get the latest metrics
        return self.get_all_metrics()[-1]

    def create_jailed_resource(self, path):
        """Create a hard link to some resource inside this microvm."""
        return self.jailer.jailed_path(path, create=True)

    def get_jailed_resource(self, path):
        """Get the relative jailed path to a resource."""
        return self.jailer.jailed_path(path, create=False)

    def chroot(self):
        """Get the chroot of this microVM."""
        return self.jailer.chroot_path()

    @property
    def screen_session(self):
        """The screen session name

        The id of this microVM, which should be unique.
        """
        return self.id

    @property
    def screen_log(self):
        """Get the screen log file."""
        return f"/tmp/screen-{self.screen_session}.log"

    @property
    def screen_pid(self) -> Optional[int]:
        """Get the screen PID."""
        if self._screen_pid:
            return int(self._screen_pid)
        return None

    def pin_vmm(self, cpu_id: int) -> bool:
        """Pin the firecracker process VMM thread to a cpu list."""
        if self.firecracker_pid:
            for thread_name, thread_pids in utils.get_threads(
                self.firecracker_pid
            ).items():
                # the firecracker thread should start with firecracker...
                if thread_name.startswith("firecracker"):
                    for pid in thread_pids:
                        utils.set_cpu_affinity(pid, [cpu_id])
                return True
        return False

    def pin_vcpu(self, vcpu_id: int, cpu_id: int):
        """Pin the firecracker vcpu thread to a cpu list."""
        if self.firecracker_pid:
            for thread in utils.get_threads(self.firecracker_pid)[f"fc_vcpu {vcpu_id}"]:
                utils.set_cpu_affinity(thread, [cpu_id])
            return True
        return False

    def pin_api(self, cpu_id: int):
        """Pin the firecracker process API server thread to a cpu list."""
        if self.firecracker_pid:
            for thread in utils.get_threads(self.firecracker_pid)["fc_api"]:
                utils.set_cpu_affinity(thread, [cpu_id])
            return True
        return False

    def pin_threads(self, first_cpu):
        """
        Pins all microvm threads (VMM, API and vCPUs) to consecutive physical cpu core, starting with "first_cpu"

        Return next "free" cpu core.
        """
        for vcpu, pcpu in enumerate(range(first_cpu, first_cpu + self.vcpus_count)):
            assert self.pin_vcpu(
                vcpu, pcpu
            ), f"Failed to pin fc_vcpu {vcpu} thread to core {pcpu}."
        # The cores first_cpu,...,first_cpu + self.vcpus_count - 1 are assigned to the individual vCPU threads,
        # So the remaining two threads (VMM and API) get first_cpu + self.vcpus_count
        # and first_cpu + self.vcpus_count + 1
        assert self.pin_vmm(
            first_cpu + self.vcpus_count
        ), "Failed to pin firecracker thread."
        assert self.pin_api(
            first_cpu + self.vcpus_count + 1
        ), "Failed to pin fc_api thread."

        return first_cpu + self.vcpus_count + 2

    def add_pre_cmd(self, pre_cmd):
        """Prepends commands to the command line to launch the microVM

        For example, this can be used to pin the VM to a NUMA node or to trace the VM with strace.
        """
        self._pre_cmd = pre_cmd + self._pre_cmd

    def spawn(
        self,
        log_file="fc.log",
        serial_out_path="serial.log",
        log_level="Debug",
        log_show_level=False,
        log_show_origin=False,
        metrics_path="fc.ndjson",
        emit_metrics: bool = False,
        validate_api: bool = True,
    ):
        """Start a microVM as a daemon or in a screen session."""
        # pylint: disable=subprocess-run-check
        # pylint: disable=too-many-branches
        self.jailer.setup()
        self.api = Api(
            self.jailer.api_socket_path(),
            validate=validate_api,
            on_error=lambda verb, uri, err_msg: self._dump_debug_information(
                f"Error during {verb} {uri}: {err_msg}"
            ),
        )

        if log_file is not None:
            self.log_file = Path(self.path) / log_file
            self.log_file.touch()
            self.create_jailed_resource(self.log_file)
            # The default value for `level`, when configuring the logger via cmd
            # line, is `Info`. We set the level to `Debug` to also have the boot
            # time printed in the log.
            self.jailer.extra_args.update({"log-path": log_file, "level": log_level})
            if log_show_level:
                self.jailer.extra_args["show-level"] = None
            if log_show_origin:
                self.jailer.extra_args["show-log-origin"] = None

        if serial_out_path is not None:
            self.serial_out_path = Path(self.path) / serial_out_path
            self.serial_out_path.touch()
            self.create_jailed_resource(self.serial_out_path)

        if metrics_path is not None:
            self.metrics_file = Path(self.path) / metrics_path
            self.metrics_file.touch()
            self.create_jailed_resource(self.metrics_file)
            self.jailer.extra_args.update({"metrics-path": self.metrics_file.name})
        else:
            assert not emit_metrics

        if self.metadata_file:
            if os.path.exists(self.metadata_file):
                LOG.debug("metadata file exists, adding as a jailed resource")
                self.create_jailed_resource(self.metadata_file)
            self.jailer.extra_args.update(
                {"metadata": os.path.basename(self.metadata_file)}
            )

        if log_level != "Debug":
            # Checking the timings requires DEBUG level log messages
            self.time_api_requests = False

        cmd = [
            *self._pre_cmd,
            str(self.jailer_binary_path),
            *self.jailer.construct_param_list(),
        ]

        # When the daemonize flag is on, we want to clone-exec into the
        # jailer rather than executing it via spawning a shell.
        if self.jailer.daemonize:
            utils.check_output(cmd, shell=False)
        else:
            # Run Firecracker under screen. This is used when we want to access
            # the serial console. The file will collect the output from
            # 'screen'ed Firecracker.
            screen_pid = utils.start_screen_process(
                self.screen_log,
                self.screen_session,
                cmd[0],
                cmd[1:],
            )
            self._screen_pid = screen_pid

        # If `--new-pid-ns` is used, the Firecracker process will detach from
        # the screen and the screen process will exit. We do not want to
        # attempt to kill it in that case to avoid a race condition.
        if self.jailer.new_pid_ns:
            self._screen_pid = None

        self._spawned = True

        if emit_metrics:
            self.monitors.append(FCMetricsMonitor(self))

        # Ensure Firecracker is in as good a state as possible wrts guest
        # responsiveness / API availability.
        # If we are using a config file and it has a network device specified,
        # use SSH to wait until guest userspace is available. If we are
        # using the API, wait until the log message indicating the API server
        # has finished initializing is printed (if logging is enabled), or
        # until the API socket file has been created.
        # If none of these apply, do a last ditch effort to make sure the
        # Firecracker process itself at least came up by checking
        # for the startup log message. Otherwise, you're on your own kid.
        if "config-file" in self.jailer.extra_args and self.iface:
            assert not serial_out_path
            self.wait_for_ssh_up()
        elif "no-api" not in self.jailer.extra_args:
            if self.log_file and log_level in ("Trace", "Debug", "Info"):
                self.check_log_message("API server started.")
            else:
                self._wait_for_api_socket()

            if serial_out_path is not None:
                self.api.serial.put(serial_out_path=serial_out_path)
        elif self.log_file and log_level in ("Trace", "Debug", "Info"):
            assert not serial_out_path
            self.check_log_message("Running Firecracker")

    @retry(wait=wait_fixed(0.2), stop=stop_after_attempt(5), reraise=True)
    def _wait_for_api_socket(self):
        """Wait until the API socket and chroot folder are available."""

        # We expect the jailer to start within 80 ms. However, we wait for
        # 1 sec since we are rechecking the existence of the socket 5 times
        # and leave 0.2 delay between them.
        os.stat(self.jailer.api_socket_path())

    @retry(wait=wait_fixed(0.2), stop=stop_after_attempt(5), reraise=True)
    def check_log_message(self, message):
        """Wait until `message` appears in logging output."""
        assert (
            message in self.log_data
        ), f'Message ("{message}") not found in log data ("{self.log_data}").'

    @retry(wait=wait_fixed(0.2), stop=stop_after_attempt(5), reraise=True)
    def get_exit_code(self):
        """Get exit code from logging output"""
        exit_msg_pattern = (
            r"Firecracker exiting (with error|successfully). exit_code=(\d+)"
        )
        match = re.search(exit_msg_pattern, self.log_data)
        if match:
            exit_code = int(match.group(2))
            return exit_code
        raise AssertionError(f"unable to find exit code from the log: {self.log_data}")

    @retry(wait=wait_fixed(0.2), stop=stop_after_attempt(5), reraise=True)
    def check_any_log_message(self, messages):
        """Wait until any message in `messages` appears in logging output."""
        for message in messages:
            if message in self.log_data:
                return
        raise AssertionError(
            f"`{messages}` were not found in this log: {self.log_data}"
        )

    def serial_input(self, input_string):
        """Send a string to the Firecracker serial console via screen."""
        input_cmd = f'screen -S {self.screen_session} -p 0 -X stuff "{input_string}"'
        return utils.check_output(input_cmd)

    def basic_config(
        self,
        vcpu_count: int = 2,
        smt: bool = None,
        mem_size_mib: int = 256,
        add_root_device: bool = True,
        boot_args: str = None,
        use_initrd: bool = False,
        track_dirty_pages: bool = False,
        huge_pages: HugePagesConfig = None,
        rootfs_io_engine=None,
        cpu_template: Optional[str] = None,
        enable_entropy_device=False,
    ):
        """Shortcut for quickly configuring a microVM.

        It handles:
        - CPU and memory.
        - Kernel image (will load the one in the microVM allocated path).
        - Root File System (will use the one in the microVM allocated path).
        - Does not start the microvm.

        The function checks the response status code and asserts that
        the response is within the interval [200, 300).

        If boot_args is None, the default boot_args used in tests is
            reboot=k panic=1 nomodule swiotlb=noforce console=ttyS0 [pci=off]
        which differs from Firecracker's default only in the enabling of the serial console.
        Reference: file:../../src/vmm/src/vmm_config/boot_source.rs::DEFAULT_KERNEL_CMDLINE
        """
        self.api.machine_config.put(
            vcpu_count=vcpu_count,
            smt=smt,
            mem_size_mib=mem_size_mib,
            track_dirty_pages=track_dirty_pages,
            huge_pages=huge_pages,
        )
        self.vcpus_count = vcpu_count
        self.mem_size_bytes = mem_size_mib * 2**20

        if self.custom_cpu_template is not None:
            self.set_cpu_template(self.custom_cpu_template)

        if cpu_template is not None:
            self.set_cpu_template(cpu_template)

        if self.memory_monitor:
            self.memory_monitor.start()

        if boot_args is not None:
            self.boot_args = boot_args
        else:
            self.boot_args = "reboot=k panic=1 nomodule swiotlb=noforce console=ttyS0"
            if not self.pci_enabled:
                self.boot_args += " pci=off"
        boot_source_args = {
            "kernel_image_path": self.create_jailed_resource(self.kernel_file),
            "boot_args": self.boot_args,
        }

        if use_initrd and self.initrd_file is not None:
            boot_source_args.update(
                initrd_path=self.create_jailed_resource(self.initrd_file)
            )

        self.api.boot.put(**boot_source_args)

        if add_root_device and self.rootfs_file is not None:
            read_only = self.rootfs_file.suffix == ".squashfs"

            # Add the root file system
            self.add_drive(
                drive_id="rootfs",
                path_on_host=self.rootfs_file,
                is_root_device=True,
                is_read_only=read_only,
                io_engine=rootfs_io_engine,
            )

        if enable_entropy_device:
            self.enable_entropy_device()

    def set_cpu_template(self, cpu_template):
        """Set guest CPU template."""
        self.cpu_template_name = get_cpu_template_name(cpu_template)
        if cpu_template is None:
            return
        # static CPU template
        if isinstance(cpu_template, str):
            self.api.machine_config.patch(cpu_template=cpu_template)
        # custom CPU template
        elif isinstance(cpu_template, dict):
            self.api.cpu_config.put(**cpu_template["template"])

    def add_drive(
        self,
        drive_id,
        path_on_host,
        is_root_device=False,
        is_read_only=False,
        partuuid=None,
        cache_type=None,
        io_engine=None,
    ):
        """Add a block device."""

        path_on_jail = self.create_jailed_resource(path_on_host)
        self.api.drive.put(
            drive_id=drive_id,
            path_on_host=path_on_jail,
            is_root_device=is_root_device,
            is_read_only=is_read_only,
            partuuid=partuuid,
            cache_type=cache_type,
            io_engine=io_engine,
        )
        self.disks[drive_id] = path_on_host

    def add_vhost_user_drive(
        self,
        drive_id,
        path_on_host,
        partuuid=None,
        is_root_device=False,
        is_read_only=False,
        cache_type=None,
        backend_type=VhostUserBlkBackendType.CROSVM,
    ):
        """Add a vhost-user block device."""

        # It is possible that the user adds another drive
        # with the same ID. In that case, we should clean
        # the previous backend up first.
        prev = self.disks_vhost_user.pop(drive_id, None)
        if prev:
            prev.kill()

        backend = VhostUserBlkBackend.with_backend(
            backend_type, path_on_host, self.chroot(), drive_id, is_read_only
        )

        socket = backend.spawn(self.jailer.uid, self.jailer.gid)

        self.api.drive.put(
            drive_id=drive_id,
            socket=socket,
            partuuid=partuuid,
            is_root_device=is_root_device,
            cache_type=cache_type,
        )

        self.disks_vhost_user[drive_id] = backend

    def patch_drive(self, drive_id, file=None):
        """Modify/patch an existing block device."""
        if file:
            self.api.drive.patch(
                drive_id=drive_id,
                path_on_host=self.create_jailed_resource(file.path),
            )
            self.disks[drive_id] = Path(file.path)
        else:
            self.api.drive.patch(drive_id=drive_id)

    def add_net_iface(self, iface=None, api=True, **kwargs):
        """Add a network interface"""
        if iface is None:
            iface = net_tools.NetIfaceConfig.with_id(len(self.iface))
        tap = self.netns.add_tap(
            iface.tap_name, ip=f"{iface.host_ip}/{iface.netmask_len}"
        )
        self.iface[iface.dev_name] = {
            "iface": iface,
            "tap": tap,
        }

        # If api, call it... there may be cases when we don't want it, for
        # example during restore
        if api:
            self.api.network.put(
                iface_id=iface.dev_name,
                host_dev_name=iface.tap_name,
                guest_mac=iface.guest_mac,
                **kwargs,
            )

        return iface

    def add_pmem(
        self,
        pmem_id,
        path_on_host,
        root_device=False,
        read_only=False,
    ):
        """Add a pmem device."""

        path_on_jail = self.create_jailed_resource(path_on_host)
        self.api.pmem.put(
            id=pmem_id,
            path_on_host=path_on_jail,
            root_device=root_device,
            read_only=read_only,
        )
        self.disks[pmem_id] = path_on_host

    def start(self):
        """Start the microvm.

        This function validates that the microvm boot succeeds.
        """
        # Check that the VM has not started yet
        assert self.state == "Not started"

        self.api.actions.put(action_type="InstanceStart")

        # Check that the VM has started
        assert self.state == "Running"

        if self.iface:
            self.wait_for_ssh_up()

    def pause(self):
        """Pauses the microVM"""
        self.api.vm.patch(state="Paused")

    def resume(self):
        """Resume the microVM"""
        self.api.vm.patch(state="Resumed")

    def make_snapshot(
        self,
        snapshot_type: SnapshotType,
        *,
        mem_path: str = "mem",
        vmstate_path="vmstate",
    ):
        """Create a Snapshot object from a microvm.

        The snapshot's memory and vstate files will be saved at the specified paths
        relative to the Microvm's chroot.

        It pauses the microvm before taking the snapshot.
        """
        self.pause()
        # Notify monitor that snapshot is being created
        if self.memory_monitor:
            self.memory_monitor.set_threshold_for_snapshot()
        self.api.snapshot_create.put(
            mem_file_path=str(mem_path),
            snapshot_path=str(vmstate_path),
            snapshot_type=snapshot_type.api_type,
        )
        root = Path(self.chroot())
        return Snapshot(
            vmstate=root / vmstate_path,
            mem=root / mem_path,
            disks=self.disks,
            net_ifaces=[x["iface"] for ifname, x in self.iface.items()],
            ssh_key=self.ssh_key,
            snapshot_type=snapshot_type,
            meta={
                "kernel_file": str(self.kernel_file),
                "rootfs_file": str(self.rootfs_file) if self.rootfs_file else None,
                "vcpus_count": self.vcpus_count,
            },
        )

    def snapshot_diff(self, *, mem_path: str = "mem", vmstate_path="vmstate"):
        """Make a Diff snapshot"""
        return self.make_snapshot(
            SnapshotType.DIFF, mem_path=mem_path, vmstate_path=vmstate_path
        )

    def snapshot_full(self, *, mem_path: str = "mem", vmstate_path="vmstate"):
        """Make a Full snapshot"""
        return self.make_snapshot(
            SnapshotType.FULL, mem_path=mem_path, vmstate_path=vmstate_path
        )

    def restore_from_snapshot(
        self,
        snapshot: Snapshot,
        resume: bool = False,
        rename_interfaces: dict = None,
        vsock_override: str = None,
        clock_realtime: bool = False,
        *,
        uffd_handler_name: str = None,
    ):
        """Restore a snapshot"""

        jailed_snapshot = snapshot.copy_to_chroot(Path(self.chroot()))

        if uffd_handler_name:
            self.uffd_handler = spawn_pf_handler(
                self,
                uffd_handler(uffd_handler_name, binary_dir=self.fc_binary_path.parent),
                jailed_snapshot,
            )

        jailed_mem = Path("/") / jailed_snapshot.mem.name
        jailed_vmstate = Path("/") / jailed_snapshot.vmstate.name

        snapshot_disks = [v for k, v in jailed_snapshot.disks.items()]
        assert len(snapshot_disks) > 0, "Snapshot requires at least one disk."
        jailed_disks = []
        for disk in snapshot_disks:
            jailed_disks.append(self.create_jailed_resource(disk))
        self.disks = jailed_snapshot.disks
        self.ssh_key = jailed_snapshot.ssh_key

        # Create network interfaces.
        for iface in jailed_snapshot.net_ifaces:
            self.add_net_iface(iface, api=False)

        mem_backend = {"backend_type": "File", "backend_path": str(jailed_mem)}
        if self.uffd_handler is not None:
            mem_backend = {
                "backend_type": "Uffd",
                "backend_path": str(self.uffd_handler.socket_path),
            }

        for key, value in jailed_snapshot.meta.items():
            setattr(self, key, value)
        # Adjust things just in case
        self.kernel_file = Path(self.kernel_file)
        if self.rootfs_file:
            self.rootfs_file = Path(self.rootfs_file)
            self.distro = GuestDistro(self.rootfs_file)

        iface_overrides = []
        if rename_interfaces is not None:
            iface_overrides = [
                {"iface_id": k, "host_dev_name": v}
                for k, v in rename_interfaces.items()
            ]

        optional_kwargs = {}
        if iface_overrides:
            # For backwards compatibility ab testing we want to avoid adding
            # new parameters until we have a release baseline with the new
            # parameter. Once the release baseline has moved, this assignment
            # can be inline in the snapshot_load command below
            optional_kwargs["network_overrides"] = iface_overrides

        if vsock_override is not None:
            optional_kwargs["vsock_override"] = {"uds_path": vsock_override}

        if clock_realtime:
            optional_kwargs["clock_realtime"] = clock_realtime

        self.api.snapshot_load.put(
            mem_backend=mem_backend,
            snapshot_path=str(jailed_vmstate),
            enable_diff_snapshots=jailed_snapshot.snapshot_type.needs_dirty_page_tracking,
            resume_vm=resume,
            **optional_kwargs,
        )

        if self.memory_monitor:
            response = self.api.machine_config.get()
            self.mem_size_bytes = int(response.json()["mem_size_mib"]) * 2**20
            # Notify monitor that this is a restored VM
            self.memory_monitor.set_threshold_for_restored_vm()
            self.memory_monitor.start()

        # This is not a "wait for boot", but rather a "VM still works after restoration"
        if jailed_snapshot.net_ifaces and resume:
            self.wait_for_ssh_up()
        return jailed_snapshot

    def enable_entropy_device(self):
        """Enable entropy device for microVM"""
        self.api.entropy.put()

    def restore_from_path(self, snap_dir: Path, **kwargs):
        """Restore snapshot from a path"""
        return self.restore_from_snapshot(Snapshot.load_from(snap_dir), **kwargs)

    @lru_cache
    def ssh_iface(self, iface_idx=0):
        """Return a cached SSH connection on a given interface id."""
        guest_ip = list(self.iface.values())[iface_idx]["iface"].guest_ip
        self.ssh_key = Path(self.ssh_key)
        connection = net_tools.SSHConnection(
            netns=self.netns.id,
            ssh_key=self.ssh_key,
            user="root",
            host=guest_ip,
            control_path=Path(self.chroot()) / f"ssh-{iface_idx}.sock",
            on_error=lambda exc: self._dump_debug_information(
                f"Failure executing command via SSH in microVM: {exc}"
            ),
        )
        self._connections.append(connection)
        return connection

    @property
    def ssh(self):
        """Return a cached SSH connection on the 1st interface"""
        return self.ssh_iface(0)

    @property
    def thread_backtraces(self):
        """Return backtraces of all threads"""
        backtraces = []
        for thread_name, thread_pids in utils.get_threads(self.firecracker_pid).items():
            for pid in thread_pids:
                try:
                    stack = Path(f"/proc/{pid}/stack").read_text("UTF-8")
                except FileNotFoundError:
                    continue  # process might've gone away between get_threads() call and here

                backtraces.append(f"{thread_name} ({pid=}):\n{stack}")
        return "\n".join(backtraces)

    def _dump_debug_information(self, what: str):
        """
        Dumps debug information about this microvm

        Used for example when running a command inside the guest via `SSHConnection.check_output` fails.
        """
        LOG.error(what)
        LOG.error("Firecracker logs:\n%s", self.log_data)
        if self.uffd_handler:
            LOG.error("Uffd logs:\n%s", self.uffd_handler.log_data)
        if not self._killed:
            LOG.error("Thread backtraces:\n%s", self.thread_backtraces)

    def wait_for_ssh_up(self):
        """Wait for guest running inside the microVM to come up and respond."""
        # Ensure that we have an initialized SSH connection to the guest that can
        # run commands. The actual connection retry loop happens in SSHConnection._init_connection
        _ = self.ssh_iface(0)

    def enable_gdb(self):
        """Enables GDB debugging"""
        self.gdb_socket = "gdb.socket"
        self.api.machine_config.patch(gdb_socket_path=self.gdb_socket)

    def hotplug_memory(
        self, requested_size_mib: int, timeout: int = 60, poll: float = 0.1
    ):
        """Send a hot(un)plug request and wait up to timeout seconds for completion polling every poll seconds

        Returns: api latency (secs), total latency (secs)
        """
        api_start = time.time()
        self.api.memory_hotplug.patch(requested_size_mib=requested_size_mib)
        api_end = time.time()
        # Wait for the hotplug to complete
        deadline = time.time() + timeout
        while time.time() < deadline:
            if (
                self.api.memory_hotplug.get().json()["plugged_size_mib"]
                == requested_size_mib
            ):
                plug_end = time.time()
                return api_end - api_start, plug_end - api_start
            time.sleep(poll)
        raise TimeoutError(f"Hotplug did not complete within {timeout} seconds")


class MicroVMFactory:
    """MicroVM factory"""

    def __init__(self, binary_path: Path, **kwargs):
        self.vms = []
        self.binary_path = binary_path
        self.netns_factory = kwargs.pop("netns_factory", net_tools.NetNs)
        self.kwargs = kwargs

        assert self.fc_binary_path.exists(), "missing firecracker binary"
        assert self.jailer_binary_path.exists(), "missing jailer binary"

    @property
    def fc_binary_path(self):
        """The path to the firecracker binary from which this factory will build VMs"""
        return self.binary_path / "firecracker"

    @property
    def jailer_binary_path(self):
        """The path to the jailer binary using which this factory will build VMs"""
        return self.binary_path / "jailer"

    def build(self, kernel=None, rootfs=None, **kwargs):
        """Build a microvm"""
        kwargs = self.kwargs | kwargs
        microvm_id = kwargs.pop("microvm_id", str(uuid.uuid4()))
        vm = Microvm(
            microvm_id=microvm_id,
            fc_binary_path=kwargs.pop("fc_binary_path", self.fc_binary_path),
            jailer_binary_path=kwargs.pop(
                "jailer_binary_path", self.jailer_binary_path
            ),
            netns=kwargs.pop("netns", self.netns_factory(microvm_id)),
            **kwargs,
        )
        vm.netns.setup()
        self.vms.append(vm)
        if kernel is not None:
            vm.kernel_file = kernel
        if rootfs is not None:
            ssh_key = rootfs.with_suffix(".id_rsa")
            # copy only iff not a read-only rootfs
            rootfs_path = rootfs
            if rootfs_path.suffix != ".squashfs":
                rootfs_path = Path(vm.path) / rootfs.name
                shutil.copyfile(rootfs, rootfs_path)
            vm.rootfs_file = rootfs_path
            vm.distro = GuestDistro(rootfs_path)
            vm.ssh_key = ssh_key
        return vm

    def build_from_snapshot(
        self, snapshot: Snapshot, uffd_handler_name=None, clock_realtime=False
    ):
        """Build a microvm from a snapshot"""
        vm = self.build()
        vm.spawn()
        vm.restore_from_snapshot(
            snapshot,
            resume=True,
            uffd_handler_name=uffd_handler_name,
            clock_realtime=clock_realtime,
        )
        return vm

    def build_n_from_snapshot(
        self,
        current_snapshot,
        nr_vms,
        *,
        uffd_handler_name=None,
        incremental=False,
        use_snapshot_editor=True,
        no_netns_reuse=False,
    ):
        """A generator of `n` microvms restored, either all restored from the same given snapshot
        (incremental=False), or created by taking successive snapshots of restored VMs
        """
        last_snapshot = None
        for _ in range(nr_vms):
            microvm = self.build(
                **(
                    {"netns": net_tools.NetNs(str(uuid.uuid4()))}
                    if no_netns_reuse
                    else {}
                )
            )
            microvm.spawn()

            snapshot_copy = microvm.restore_from_snapshot(
                current_snapshot, resume=True, uffd_handler_name=uffd_handler_name
            )

            yield microvm

            if incremental:
                # When doing diff snapshots, we continuously overwrite the same base snapshot file from the first
                # iteration in-place with successive snapshots, so don't delete it!
                if (
                    last_snapshot is not None
                    and not last_snapshot.snapshot_type.needs_rebase
                ):
                    last_snapshot.delete()

                next_snapshot = microvm.make_snapshot(current_snapshot.snapshot_type)

                if current_snapshot.snapshot_type.needs_rebase:
                    next_snapshot = next_snapshot.rebase_snapshot(
                        current_snapshot,
                        use_snapshot_editor,
                        binary_dir=microvm.fc_binary_path.parent,
                    )

                last_snapshot = current_snapshot
                current_snapshot = next_snapshot

            microvm.kill()
            snapshot_copy.delete()

        if last_snapshot is not None and not last_snapshot.snapshot_type.needs_rebase:
            last_snapshot.delete()
        current_snapshot.delete()

    def kill(self):
        """Clean up all built VMs"""
        for vm in self.vms:
            vm.kill()
            vm.jailer.cleanup()
            chroot_base_with_id = vm.jailer.chroot_base_with_id()
            if len(vm.jailer.jailer_id) > 0 and chroot_base_with_id.exists():
                shutil.rmtree(chroot_base_with_id)
            vm.netns.cleanup()

        self.vms.clear()


class Serial:
    """Class for serial console communication with a Microvm."""

    RX_TIMEOUT_S = 60

    def __init__(self, vm):
        """Initialize a new Serial object."""
        self._poller = None
        self._vm = vm

    def open(self):
        """Open a serial connection."""
        # Open the screen log file.
        if self._poller is not None:
            # serial already opened
            return

        attempt = 0
        while not Path(self._vm.screen_log).exists() and attempt < 5:
            time.sleep(0.2)
            attempt += 1

        serial_log_fd = os.open(self._vm.screen_log, os.O_RDONLY)
        self._poller = select.poll()
        self._poller.register(serial_log_fd, select.POLLIN | select.POLLHUP)

    def tx(self, input_string, end="\n"):
        # pylint: disable=invalid-name
        # No need to have a snake_case naming style for a single word.
        r"""Send a string terminated by an end token (defaulting to "\n")."""
        self._vm.serial_input(input_string + end)

    def rx_char(self):
        """Read a single character."""
        result = self._poller.poll(0.1)

        for fd, flag in result:
            if flag & select.POLLHUP:
                assert False, "Oh! The console vanished before test completed."

            if flag & select.POLLIN:
                output_char = str(os.read(fd, 1), encoding="utf-8", errors="ignore")
                return output_char

        return ""

    def rx(self, token="\n"):
        # pylint: disable=invalid-name
        # No need to have a snake_case naming style for a single word.
        r"""Read a string delimited by an end token (defaults to "\n")."""
        rx_str = ""
        start = time.time()
        while True:
            rx_str += self.rx_char()
            if rx_str.endswith(token):
                break
            if (time.time() - start) >= self.RX_TIMEOUT_S:
                self._vm.kill()
                assert False

        return rx_str