fix(test): cleanup snapshot files better

tests/framework/microvm.py

@@ -971,33 +971,43 @@ def snapshot_full(self, *, mem_path: str = "mem", vmstate_path="vmstate"):
 
     def restore_from_snapshot(
         self,
-        snapshot: Snapshot,
+        snapshot: Snapshot = None,
         resume: bool = False,
-        uffd_path: Path = None,
         rename_interfaces: dict = None,
     ):
         """Restore a snapshot"""
-        jailed_snapshot = snapshot.copy_to_chroot(Path(self.chroot()))
+        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_mem = Path("/") / jailed_snapshot.mem.name
         jailed_vmstate = Path("/") / jailed_snapshot.vmstate.name
 
-        snapshot_disks = [v for k, v in snapshot.disks.items()]
+        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 = snapshot.disks
-        self.ssh_key = snapshot.ssh_key
+        self.disks = jailed_snapshot.disks
+        self.ssh_key = jailed_snapshot.ssh_key
 
         # Create network interfaces.
-        for iface in snapshot.net_ifaces:
+        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 uffd_path is not None:
-            mem_backend = {"backend_type": "Uffd", "backend_path": str(uffd_path)}
+        if self.uffd_handler is not None:
+            mem_backend = {
+                "backend_type": "Uffd",
+                "backend_path": str(self.uffd_handler.socket_path),
+            }
 
-        for key, value in snapshot.meta.items():
+        for key, value in jailed_snapshot.meta.items():
             setattr(self, key, value)
         # Adjust things just in case
         self.kernel_file = Path(self.kernel_file)
@@ -1020,12 +1030,12 @@ def restore_from_snapshot(
         self.api.snapshot_load.put(
             mem_backend=mem_backend,
             snapshot_path=str(jailed_vmstate),
-            enable_diff_snapshots=snapshot.is_diff,
+            enable_diff_snapshots=jailed_snapshot.is_diff,
             resume_vm=resume,
             **optional_kwargs,
         )
         # This is not a "wait for boot", but rather a "VM still works after restoration"
-        if snapshot.net_ifaces and resume:
+        if jailed_snapshot.net_ifaces and resume:
             self.wait_for_ssh_up()
         return jailed_snapshot
 
@@ -1148,7 +1158,7 @@ def build_from_snapshot(self, snapshot: Snapshot):
 
     def build_n_from_snapshot(
         self,
-        snapshot,
+        current_snapshot,
         nr_vms,
         *,
         uffd_handler_name=None,
@@ -1158,39 +1168,44 @@ def build_n_from_snapshot(
         """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()
             microvm.spawn()
 
-            uffd_path = None
             if uffd_handler_name is not None:
-                pf_handler = spawn_pf_handler(
+                spawn_pf_handler(
                     microvm,
                     uffd_handler(uffd_handler_name, binary_dir=self.binary_path),
-                    snapshot.mem,
+                    current_snapshot,
                 )
-                uffd_path = pf_handler.socket_path
 
-            snapshot_copy = microvm.restore_from_snapshot(
-                snapshot, resume=True, uffd_path=uffd_path
-            )
+            snapshot_copy = microvm.restore_from_snapshot(current_snapshot, resume=True)
 
             yield microvm
 
             if incremental:
-                new_snapshot = microvm.make_snapshot(snapshot.snapshot_type)
+                # 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.is_diff:
+                    last_snapshot.delete()
+
+                next_snapshot = microvm.make_snapshot(current_snapshot.snapshot_type)
 
-                if snapshot.is_diff:
-                    new_snapshot = new_snapshot.rebase_snapshot(
-                        snapshot, use_snapshot_editor
+                if current_snapshot.is_diff:
+                    next_snapshot = next_snapshot.rebase_snapshot(
+                        current_snapshot, use_snapshot_editor
                     )
 
-                snapshot = new_snapshot
+                last_snapshot = current_snapshot
+                current_snapshot = next_snapshot
 
             microvm.kill()
             snapshot_copy.delete()
 
-        snapshot.delete()
+        if last_snapshot is not None and not last_snapshot.is_diff:
+            last_snapshot.delete()
+        current_snapshot.delete()
 
     def kill(self):
         """Clean up all built VMs"""

tests/framework/utils_uffd.py

@@ -17,12 +17,14 @@
 class UffdHandler:
     """Describe the UFFD page fault handler process."""
 
-    def __init__(self, name, socket_path, mem_file, chroot_path, log_file_name):
+    def __init__(
+        self, name, socket_path, snapshot: "Snapshot", chroot_path, log_file_name
+    ):
         """Instantiate the handler process with arguments."""
         self._proc = None
         self._handler_name = name
         self.socket_path = socket_path
-        self._mem_file = mem_file
+        self.snapshot = snapshot
         self._chroot = chroot_path
         self._log_file = log_file_name
 
@@ -35,7 +37,11 @@ def spawn(self, uid, gid):
 
             chroot_log_file = Path("/") / self._log_file
             with open(chroot_log_file, "w", encoding="utf-8") as logfile:
-                args = [f"/{self._handler_name}", self.socket_path, self._mem_file]
+                args = [
+                    f"/{self._handler_name}",
+                    self.socket_path,
+                    self.snapshot.mem.name,
+                ]
                 self._proc = subprocess.Popen(
                     args, stdout=logfile, stderr=subprocess.STDOUT
                 )
@@ -77,16 +83,16 @@ def __del__(self):
             self.proc.kill()
 
 
-def spawn_pf_handler(vm, handler_path, mem_path):
+def spawn_pf_handler(vm, handler_path, snapshot):
     """Spawn page fault handler process."""
     # Copy snapshot memory file into chroot of microVM.
-    jailed_mem = vm.create_jailed_resource(mem_path)
+    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)
 
     uffd_handler = UffdHandler(
-        handler_name, SOCKET_PATH, jailed_mem, vm.chroot(), "uffd.log"
+        handler_name, SOCKET_PATH, jailed_snapshot, vm.chroot(), "uffd.log"
     )
     uffd_handler.spawn(vm.jailer.uid, vm.jailer.gid)
     vm.uffd_handler = uffd_handler

tests/integration_tests/functional/test_uffd.py

@@ -92,9 +92,9 @@ def test_valid_handler(uvm_plain, snapshot):
     vm.spawn()
 
     # Spawn page fault handler process.
-    pf_handler = spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot.mem)
+    spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot)
 
-    vm.restore_from_snapshot(snapshot, resume=True, uffd_path=pf_handler.socket_path)
+    vm.restore_from_snapshot(resume=True)
 
     # Inflate balloon.
     vm.api.balloon.patch(amount_mib=200)
@@ -125,15 +125,13 @@ def test_malicious_handler(uvm_plain, snapshot):
     vm.spawn()
 
     # Spawn page fault handler process.
-    pf_handler = spawn_pf_handler(vm, uffd_handler("malicious"), snapshot.mem)
+    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(
-                snapshot, resume=True, uffd_path=pf_handler.socket_path
-            )
+            vm.restore_from_snapshot(resume=True)
             assert False, "Firecracker should freeze"
     except (TimeoutError, requests.exceptions.ReadTimeout):
         pass

tests/integration_tests/performance/test_huge_pages.py

@@ -93,9 +93,9 @@ def test_hugetlbfs_snapshot(microvm_factory, guest_kernel_linux_5_10, rootfs):
     vm.spawn()
 
     # Spawn page fault handler process.
-    pf_handler = spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot.mem)
+    spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot)
 
-    vm.restore_from_snapshot(snapshot, resume=True, uffd_path=pf_handler.socket_path)
+    vm.restore_from_snapshot(resume=True)
 
     check_hugetlbfs_in_use(vm.firecracker_pid, "/anon_hugepage")
 
@@ -135,11 +135,9 @@ def test_hugetlbfs_diff_snapshot(microvm_factory, uvm_plain):
     vm.spawn()
 
     # Spawn page fault handler process.
-    pf_handler = spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot_merged.mem)
+    spawn_pf_handler(vm, uffd_handler("on_demand"), snapshot_merged)
 
-    vm.restore_from_snapshot(
-        snapshot_merged, resume=True, uffd_path=pf_handler.socket_path
-    )
+    vm.restore_from_snapshot(resume=True)
 
     # Verify if the restored microvm works.
 
@@ -195,12 +193,10 @@ def test_ept_violation_count(
     vm.spawn()
 
     # Spawn page fault handler process.
-    pf_handler = spawn_pf_handler(vm, uffd_handler("fault_all"), snapshot.mem)
+    spawn_pf_handler(vm, uffd_handler("fault_all"), snapshot)
 
     with ftrace_events("kvm:*"):
-        vm.restore_from_snapshot(
-            snapshot, resume=True, uffd_path=pf_handler.socket_path
-        )
+        vm.restore_from_snapshot(resume=True)
 
         # 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}")