refactor(test): spawn uffd handler inside restore_from_snapshot

Drop the requirement of having to call spawn_pf_handler() before the
call to microvm.restore_from_snapshot() in integration tests, and
instead just have restore_from_snapshot() take the name of the uffd
handler as an optional keyword argument, and do the spawning. This
removes the weird calling convention of calling restore_from_snapshot()
without a snapshot parameter (because it was inferred from uffd).

Signed-off-by: Patrick Roy <[email protected]>

Author: roypat

tests/framework/microvm.py

@@ -975,19 +975,22 @@ def snapshot_full(self, *, mem_path: str = "mem", vmstate_path="vmstate"):
 
     def restore_from_snapshot(
         self,
-        snapshot: Snapshot = None,
+        snapshot: Snapshot,
         resume: bool = False,
         rename_interfaces: dict = None,
+        *,
+        uffd_handler_name: str = None,
     ):
         """Restore a snapshot"""
-        if self.uffd_handler is None:
-            assert (
-                snapshot is not None
-            ), "snapshot file must be provided if no uffd handler is attached!"
 
-            jailed_snapshot = snapshot.copy_to_chroot(Path(self.chroot()))
-        else:
-            jailed_snapshot = self.uffd_handler.snapshot
+        jailed_snapshot = snapshot.copy_to_chroot(Path(self.chroot()))
+
+        if uffd_handler_name:
+            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
@@ -1180,14 +1183,9 @@ def build_n_from_snapshot(
             microvm = self.build()
             microvm.spawn()
 
-            if uffd_handler_name is not None:
-                spawn_pf_handler(
-                    microvm,
-                    uffd_handler(uffd_handler_name, binary_dir=self.binary_path),
-                    current_snapshot,
-                )
-
-            snapshot_copy = microvm.restore_from_snapshot(current_snapshot, resume=True)
+            snapshot_copy = microvm.restore_from_snapshot(
+                current_snapshot, resume=True, uffd_handler_name=uffd_handler_name
+            )
 
             yield microvm
 

tests/framework/utils_uffd.py

@@ -83,10 +83,9 @@ def __del__(self):
             self.proc.kill()
 
 
-def spawn_pf_handler(vm, handler_path, snapshot):
+def spawn_pf_handler(vm, handler_path, jailed_snapshot):
     """Spawn page fault handler process."""
     # Copy snapshot memory file into chroot of microVM.
-    jailed_snapshot = snapshot.copy_to_chroot(Path(vm.chroot()))
     # Copy the valid page fault binary into chroot of microVM.
     jailed_handler = vm.create_jailed_resource(handler_path)
     handler_name = os.path.basename(jailed_handler)

tests/integration_tests/functional/test_uffd.py

@@ -9,7 +9,6 @@
 import requests
 
 from framework.utils import Timeout, check_output
-from framework.utils_uffd import spawn_pf_handler, uffd_handler
 
 
 @pytest.fixture(scope="function", name="snapshot")
@@ -90,11 +89,7 @@ def test_valid_handler(uvm_plain, snapshot):
     vm = uvm_plain
     vm.memory_monitor = None
     vm.spawn()
-
-    # Spawn page fault handler process.
-    spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot)
-
-    vm.restore_from_snapshot(resume=True)
+    vm.restore_from_snapshot(snapshot, resume=True, uffd_handler_name="on_demand")
 
     # Inflate balloon.
     vm.api.balloon.patch(amount_mib=200)
@@ -124,14 +119,13 @@ def test_malicious_handler(uvm_plain, snapshot):
     vm.memory_monitor = None
     vm.spawn()
 
-    # Spawn page fault handler process.
-    spawn_pf_handler(vm, uffd_handler("malicious"), snapshot)
-
     # We expect Firecracker to freeze while resuming from a snapshot
     # due to the malicious handler's unavailability.
     try:
         with Timeout(seconds=30):
-            vm.restore_from_snapshot(resume=True)
+            vm.restore_from_snapshot(
+                snapshot, resume=True, uffd_handler_name="malicious"
+            )
             assert False, "Firecracker should freeze"
     except (TimeoutError, requests.exceptions.ReadTimeout):
         pass

tests/integration_tests/performance/test_huge_pages.py

@@ -10,7 +10,6 @@
 from framework.microvm import HugePagesConfig
 from framework.properties import global_props
 from framework.utils_ftrace import ftrace_events
-from framework.utils_uffd import spawn_pf_handler, uffd_handler
 
 
 def check_hugetlbfs_in_use(pid: int, allocation_name: str):
@@ -91,11 +90,7 @@ def test_hugetlbfs_snapshot(microvm_factory, guest_kernel_linux_5_10, rootfs):
     ### Restore Snapshot ###
     vm = microvm_factory.build()
     vm.spawn()
-
-    # Spawn page fault handler process.
-    spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot)
-
-    vm.restore_from_snapshot(resume=True)
+    vm.restore_from_snapshot(snapshot, resume=True, uffd_handler_name="on_demand")
 
     check_hugetlbfs_in_use(vm.firecracker_pid, "/anon_hugepage")
 
@@ -133,11 +128,9 @@ def test_hugetlbfs_diff_snapshot(microvm_factory, uvm_plain):
 
     vm = microvm_factory.build()
     vm.spawn()
-
-    # Spawn page fault handler process.
-    spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot_merged)
-
-    vm.restore_from_snapshot(resume=True)
+    vm.restore_from_snapshot(
+        snapshot_merged, resume=True, uffd_handler_name="on_demand"
+    )
 
     # Verify if the restored microvm works.
 
@@ -192,11 +185,8 @@ def test_ept_violation_count(
     vm.jailer.extra_args.update({"no-seccomp": None})
     vm.spawn()
 
-    # Spawn page fault handler process.
-    spawn_pf_handler(vm, uffd_handler("fault_all"), snapshot)
-
     with ftrace_events("kvm:*"):
-        vm.restore_from_snapshot(resume=True)
+        vm.restore_from_snapshot(snapshot, resume=True, uffd_handler_name="fault_all")
 
         # Verify if guest can run commands, and also wake up the fast page fault helper to trigger page faults.
         vm.ssh.check_output(f"kill -s {signal.SIGUSR1} {pid}")