Merge branch 'main' into tracing-fix

tests/framework/microvm.py

@@ -38,6 +38,7 @@
 from framework.jailer import JailerContext
 from framework.microvm_helpers import MicrovmHelpers
 from framework.properties import global_props
+from framework.utils_drive import VhostUserBlkBackend, VhostUserBlkBackendType
 from host_tools.memory import MemoryMonitor
 
 LOG = logging.getLogger("microvm")
@@ -209,6 +210,7 @@ def __init__(
         # device dictionaries
         self.iface = {}
         self.disks = {}
+        self.disks_vhost_user = {}
         self.vcpus_count = None
         self.mem_size_bytes = None
 
@@ -229,6 +231,13 @@ def kill(self):
         """All clean up associated with this microVM should go here."""
         # pylint: disable=subprocess-run-check
 
+        # 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()
+
         if (
             self.expect_kill_by_signal is False
             and "Shutting down VM after intercepting signal" in self.log_data
@@ -466,6 +475,8 @@ def pin_api(self, cpu_id: int):
     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(
@@ -481,6 +492,8 @@ def pin_threads(self, first_cpu):
             first_cpu + self.vcpus_count + 1
         ), "Failed to pin fc_api thread."
 
+        return first_cpu + self.vcpus_count + 2
+
     def spawn(
         self,
         log_file="fc.log",
@@ -704,13 +717,28 @@ def add_drive(
     def add_vhost_user_drive(
         self,
         drive_id,
-        socket,
+        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,
@@ -719,6 +747,8 @@ def add_vhost_user_drive(
             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:

tests/framework/utils_drive.py

@@ -6,6 +6,7 @@
 import os
 import subprocess
 import time
+from abc import ABC, abstractmethod
 from enum import Enum
 from pathlib import Path
 from subprocess import check_output
@@ -46,78 +47,142 @@ def partuuid_and_disk_path(rootfs_ubuntu_22, disk_path):
     return (partuuid, disk_path)
 
 
-CROSVM_CTR_SOCKET_NAME = "crosvm_ctr.socket"
-
-
-def spawn_vhost_user_backend(
-    vm,
-    host_mem_path,
-    socket_path,
-    readonly=False,
-    backend=VhostUserBlkBackendType.CROSVM,
-):
-    """Spawn vhost-user-blk backend."""
-
-    uid = vm.jailer.uid
-    gid = vm.jailer.gid
-    host_vhost_user_socket_path = Path(vm.chroot()) / socket_path.strip("/")
-
-    if backend == VhostUserBlkBackendType.QEMU:
+class VhostUserBlkBackend(ABC):
+    """vhost-user-blk backend base class"""
+
+    @classmethod
+    def get_all_subclasses(cls):
+        """Get all subclasses of the class."""
+        subclasses = {}
+        for subclass in cls.__subclasses__():
+            subclasses[subclass.__name__] = subclass
+            subclasses.update(subclass.get_all_subclasses())
+        return subclasses
+
+    @classmethod
+    def with_backend(cls, backend: VhostUserBlkBackendType, *args, **kwargs):
+        """Get a backend of a specific type."""
+        subclasses = cls.get_all_subclasses()
+        return subclasses[backend.value + cls.__name__](*args, **kwargs)
+
+    def __init__(
+        self,
+        host_mem_path,
+        chroot,
+        backend_id,
+        readonly,
+    ):
+        self.host_mem_path = host_mem_path
+        self.socket_path = Path(chroot) / f"{backend_id}_vhost_user.sock"
+        self.readonly = readonly
+        self.proc = None
+
+    def spawn(self, uid, gid):
+        """
+        Spawn a backend.
+
+        Return socket path in the jail that can be used with FC API.
+        """
+        assert not self.proc, "backend already spawned"
+        args = self._spawn_cmd()
+        proc = subprocess.Popen(args)
+
+        # Give the backend time to initialise.
+        time.sleep(1)
+
+        assert proc is not None and proc.poll() is None, "backend is not up"
+        assert self.socket_path.exists()
+
+        os.chown(self.socket_path, uid, gid)
+
+        self.proc = proc
+
+        return str(Path("/") / os.path.basename(self.socket_path))
+
+    @abstractmethod
+    def _spawn_cmd(self):
+        """Return a spawn command for the backend"""
+        return ""
+
+    @abstractmethod
+    def resize(self, new_size):
+        """Resize the vhost-user-backed drive"""
+
+    def pin(self, cpu_id: int):
+        """Pin the vhost-user backend to a CPU list."""
+        return utils.ProcessManager.set_cpu_affinity(self.proc.pid, [cpu_id])
+
+    def kill(self):
+        """Kill the backend"""
+        if self.proc.poll() is None:
+            self.proc.terminate()
+            self.proc.wait()
+            os.remove(self.socket_path)
+        assert not os.path.exists(self.socket_path)
+
+
+class QemuVhostUserBlkBackend(VhostUserBlkBackend):
+    """vhost-user-blk backend implementaiton for Qemu backend"""
+
+    def _spawn_cmd(self):
         args = [
             "vhost-user-blk",
             "--socket-path",
-            host_vhost_user_socket_path,
+            self.socket_path,
             "--blk-file",
-            host_mem_path,
+            self.host_mem_path,
         ]
-        if readonly:
+        if self.readonly:
             args.append("--read-only")
-    elif backend == VhostUserBlkBackendType.CROSVM:
-        crosvm_ctr_socket_path = Path(vm.chroot()) / CROSVM_CTR_SOCKET_NAME.strip("/")
-        ro = ",ro" if readonly else ""
+        return args
+
+    def resize(self, new_size):
+        raise NotImplementedError("not supported for Qemu backend")
+
+
+class CrosvmVhostUserBlkBackend(VhostUserBlkBackend):
+    """vhost-user-blk backend implementaiton for crosvm backend"""
+
+    def __init__(
+        self,
+        host_mem_path,
+        chroot,
+        backend_id,
+        readonly=False,
+    ):
+        super().__init__(
+            host_mem_path,
+            chroot,
+            backend_id,
+            readonly,
+        )
+        self.ctr_socket_path = Path(chroot) / f"{backend_id}_ctr.sock"
+
+    def _spawn_cmd(self):
+        ro = ",ro" if self.readonly else ""
         args = [
             "crosvm",
             "--log-level",
             "off",
             "devices",
             "--disable-sandbox",
             "--control-socket",
-            crosvm_ctr_socket_path,
+            self.ctr_socket_path,
             "--block",
-            f"vhost={host_vhost_user_socket_path},path={host_mem_path}{ro}",
+            f"vhost={self.socket_path},path={self.host_mem_path}{ro}",
         ]
-        if os.path.exists(crosvm_ctr_socket_path):
-            os.remove(crosvm_ctr_socket_path)
-    else:
-        assert False, f"unknown vhost-user-blk backend `{backend}`"
-    proc = subprocess.Popen(args)
+        return args
 
-    # Give the backend time to initialise.
-    time.sleep(1)
+    def resize(self, new_size):
+        assert self.proc, "backend is not spawned"
+        assert self.ctr_socket_path.exists()
 
-    assert proc is not None and proc.poll() is None, "backend is not up"
+        utils.run_cmd(
+            f"crosvm disk resize 0 {new_size * 1024 * 1024} {self.ctr_socket_path}"
+        )
 
-    with utils.chroot(vm.chroot()):
-        # The backend will create the socket path with root rights.
-        # Change rights to the jailer's.
-        os.chown(socket_path, uid, gid)
-
-    return proc
-
-
-def resize_vhost_user_drive(vm, new_size):
-    """
-    Resize vhost-user-blk drive and send config change notification.
-
-    This only works with the crosvm vhost-user-blk backend.
-    New size is in MB.
-    """
-
-    crosvm_ctr_socket_path = Path(vm.chroot()) / CROSVM_CTR_SOCKET_NAME.strip("/")
-    assert os.path.exists(
-        crosvm_ctr_socket_path
-    ), "crosvm backend must be spawned first"
-
-    utils.run_cmd(
-        f"crosvm disk resize 0 {new_size * 1024 * 1024} {crosvm_ctr_socket_path}"
-    )
+    def kill(self):
+        super().kill()
+        assert self.proc.poll() is not None
+        os.remove(self.ctr_socket_path)
+        assert not os.path.exists(self.ctr_socket_path)

tests/integration_tests/functional/test_api.py

@@ -9,6 +9,7 @@
 import re
 import resource
 import time
+from pathlib import Path
 
 import packaging.version
 import pytest
@@ -18,7 +19,6 @@
 from framework import utils_cpuid
 from framework.utils import get_firecracker_version_from_toml, is_io_uring_supported
 from framework.utils_cpu_templates import nonci_on_arm
-from framework.utils_drive import spawn_vhost_user_backend
 
 MEM_LIMIT = 1000000000
 
@@ -749,12 +749,7 @@ def test_drive_patch(test_microvm_with_api):
     fs_vub = drive_tools.FilesystemFile(
         os.path.join(test_microvm.fsfiles, "scratch_vub")
     )
-    vhost_user_socket = "/vub.socket"
-    # Launching vhost-user-block backend
-    _backend = spawn_vhost_user_backend(
-        test_microvm, fs_vub.path, vhost_user_socket, False
-    )
-    test_microvm.add_vhost_user_drive("scratch_vub", vhost_user_socket)
+    test_microvm.add_vhost_user_drive("scratch_vub", fs_vub.path)
 
     # Patching drive before boot is not allowed.
     with pytest.raises(RuntimeError, match=NOT_SUPPORTED_BEFORE_START):
@@ -931,7 +926,10 @@ def _drive_patch(test_microvm):
             "path_on_host": None,
             "rate_limiter": None,
             "io_engine": None,
-            "socket": "/vub.socket",
+            "socket": str(
+                Path("/")
+                / test_microvm.disks_vhost_user["scratch_vub"].socket_path.name
+            ),
         },
     ]