Only copy bitmap if the snapshot fails

Only if the snapshot writing fails then store the kvm buffer inside the
firecracker internal bitmap. Also attempted to make the filling
algorithm in the test to be dynamic so it will work on ARM
chips

Signed-off-by: Jack Thomson <jackabt@amazon.com>

src/vmm/src/vstate/memory.rs

@@ -301,36 +301,21 @@ impl GuestMemoryExtension for GuestMemoryMmap {
         let mut writer_offset = 0;
         let page_size = get_page_size().map_err(MemoryError::PageSize)?;
 
-        self.iter().enumerate().for_each(|(slot, region)| {
-            let kvm_bitmap = dirty_bitmap.get(&slot).unwrap();
-            let firecracker_bitmap = region.bitmap();
-
-            for (i, v) in kvm_bitmap.iter().enumerate() {
-                for j in 0..64 {
-                    let is_kvm_page_dirty = ((v >> j) & 1u64) != 0u64;
-
-                    if is_kvm_page_dirty {
-                        let page_offset = ((i * 64) + j) * page_size;
-
-                        firecracker_bitmap.mark_dirty(page_offset, 1)
-                    }
-                }
-            }
-        });
-
-        self.iter()
+        let write_result = self.iter()
             .enumerate()
             .try_for_each(|(slot, region)| {
                 let kvm_bitmap = dirty_bitmap.get(&slot).unwrap();
                 let firecracker_bitmap = region.bitmap();
                 let mut write_size = 0;
                 let mut dirty_batch_start: u64 = 0;
 
-                for i in 0..kvm_bitmap.len() {
+                for (i, v) in kvm_bitmap.iter().enumerate() {
                     for j in 0..64 {
+                        let is_kvm_page_dirty = ((v >> j) & 1u64) != 0u64;
                         let page_offset = ((i * 64) + j) * page_size;
                         let is_firecracker_page_dirty = firecracker_bitmap.dirty_at(page_offset);
-                        if is_firecracker_page_dirty {
+
+                        if is_kvm_page_dirty || is_firecracker_page_dirty {
                             // We are at the start of a new batch of dirty pages.
                             if write_size == 0 {
                                 // Seek forward over the unmodified pages.
@@ -340,6 +325,7 @@ impl GuestMemoryExtension for GuestMemoryMmap {
                                 dirty_batch_start = page_offset as u64;
                             }
                             write_size += page_size;
+
                         } else if write_size > 0 {
                             // We are at the end of a batch of dirty pages.
                             writer.write_all_volatile(
@@ -362,12 +348,30 @@ impl GuestMemoryExtension for GuestMemoryMmap {
                 writer_offset += region.len();
 
                 Ok(())
-            })
-            .map_err(MemoryError::WriteMemory)?;
+            });
+
+        if write_result.is_err() {
+            self.iter().enumerate().for_each(|(slot, region)| {
+                let kvm_bitmap = dirty_bitmap.get(&slot).unwrap();
+                let firecracker_bitmap = region.bitmap();
+
+                for (i, v) in kvm_bitmap.iter().enumerate() {
+                    for j in 0..64 {
+                        let is_kvm_page_dirty = ((v >> j) & 1u64) != 0u64;
+
+                        if is_kvm_page_dirty {
+                            let page_offset = ((i * 64) + j) * page_size;
 
-        self.reset_dirty();
+                            firecracker_bitmap.mark_dirty(page_offset, 1)
+                        }
+                    }
+                }
+            });
+        } else {
+            self.reset_dirty();
+        }
 
-        Ok(())
+        write_result.map_err(MemoryError::WriteMemory)
     }
 
     /// Resets all the memory region bitmaps

tests/integration_tests/functional/test_snapshot_not_losing_dirty_pages.py

@@ -44,10 +44,12 @@ def test_diff_snapshot_works_after_error(
 
     chroot = Path(uvm.chroot())
 
-    # Create a large file, so we run out of space (ENOSPC) during the snapshot
-    # Assumes a Docker /srv tmpfs of 1G, derived by trial and error
+    # Create a large file dynamically based on available space
     fill = chroot / "fill"
-    subprocess.check_call(f"fallocate -l 330M {fill}", shell=True)
+    free_space = int(subprocess.check_output(f"df {chroot} | tail -1 | awk '{{print $4}}'", shell=True)) # in bytes
+    target_size = round(((free_space ) / 1024) * 0.9) # Attempt to fill 90% of free space
+
+    subprocess.check_call(f"fallocate -l {target_size}M {fill}", shell=True)
 
     try:
         uvm.snapshot_diff()