Reset dirty bitmap on Full Snapshot

src/vmm/src/lib.rs

@@ -588,6 +588,19 @@ impl Vmm {
         Ok(cpu_configs)
     }
 
+    /// Retrieves the KVM dirty bitmap for each of the guest's memory regions.
+    pub fn reset_dirty_bitmap(&self) {
+        self.guest_memory
+            .iter()
+            .enumerate()
+            .for_each(|(slot, region)| {
+                let _ = self
+                    .vm
+                    .fd()
+                    .get_dirty_log(u32::try_from(slot).unwrap(), u64_to_usize(region.len()));
+            });
+    }
+
     /// Retrieves the KVM dirty bitmap for each of the guest's memory regions.
     pub fn get_dirty_bitmap(&self) -> Result<DirtyBitmap, VmmError> {
         let mut bitmap: DirtyBitmap = HashMap::new();

src/vmm/src/persist.rs

@@ -254,7 +254,15 @@ fn snapshot_memory_to_file(
                 .dump_dirty(&mut file, &dirty_bitmap)
                 .map_err(Memory)
         }
-        SnapshotType::Full => vmm.guest_memory().dump(&mut file).map_err(Memory),
+        SnapshotType::Full => {
+            let dump_res = vmm.guest_memory().dump(&mut file).map_err(Memory);
+            if dump_res.is_ok() {
+                vmm.reset_dirty_bitmap();
+                vmm.guest_memory().reset_dirty();
+            }
+
+            dump_res
+        }
     }?;
     file.flush()
         .map_err(|err| MemoryBackingFile("flush", err))?;

src/vmm/src/vstate/memory.rs

@@ -104,6 +104,9 @@ where
         writer: &mut T,
         dirty_bitmap: &DirtyBitmap,
     ) -> Result<(), MemoryError>;
+
+    /// Resets all the memory region bitmaps
+    fn reset_dirty(&self);
 }
 
 /// State of a guest memory region saved to file/buffer.
@@ -349,6 +352,15 @@ impl GuestMemoryExtension for GuestMemoryMmap {
             })
             .map_err(MemoryError::WriteMemory)
     }
+
+    /// Resets all the memory region bitmaps
+    fn reset_dirty(&self) {
+        self.iter().for_each(|region| {
+            if let Some(bitmap) = region.bitmap() {
+                bitmap.reset();
+            }
+        })
+    }
 }
 
 fn create_memfd(

tests/integration_tests/functional/test_dirty_pages_in_full_snapshot.py

@@ -0,0 +1,30 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Test scenario for reseting dirty pages after making a full snapshot."""
+
+
+def test_dirty_pages_after_full_snapshot(uvm_plain):
+    """
+    Test if dirty pages are erased after making a full snapshot of a VM
+    """
+
+    vm_mem_size = 128
+    uvm = uvm_plain
+    uvm.spawn()
+    uvm.basic_config(mem_size_mib=vm_mem_size, track_dirty_pages=True)
+    uvm.add_net_iface()
+    uvm.start()
+    uvm.ssh.run("true")
+
+    snap_full = uvm.snapshot_full(vmstate_path="vmstate_full", mem_path="mem_full")
+    snap_diff = uvm.snapshot_diff(vmstate_path="vmstate_diff", mem_path="mem_diff")
+    snap_diff2 = uvm.snapshot_diff(vmstate_path="vmstate_diff2", mem_path="mem_diff2")
+
+    # file size is the same, but the `diff` snapshot is actually a sparse file
+    assert snap_full.mem.stat().st_size == snap_diff.mem.stat().st_size
+
+    # diff -> diff there should be no differences
+    assert snap_diff2.mem.stat().st_blocks == 0
+
+    # full -> diff there should be no differences
+    assert snap_diff.mem.stat().st_blocks == 0

tests/integration_tests/functional/test_snapshot_basic.py

@@ -252,14 +252,14 @@ def test_cmp_full_and_first_diff_mem(microvm_factory, guest_kernel, rootfs):
     exit_code, _, _ = vm.ssh.run("sync")
     assert exit_code == 0
 
-    logger.info("Create full snapshot.")
-    # Create full snapshot.
-    full_snapshot = vm.snapshot_full(mem_path="mem_full")
-
     logger.info("Create diff snapshot.")
     # Create diff snapshot.
     diff_snapshot = vm.snapshot_diff()
 
+    logger.info("Create full snapshot.")
+    # Create full snapshot.
+    full_snapshot = vm.snapshot_full(mem_path="mem_full")
+
     assert full_snapshot.mem != diff_snapshot.mem
     assert filecmp.cmp(full_snapshot.mem, diff_snapshot.mem, shallow=False)