fix: only reset dirty bitmap if all regions were dumped successfully

When taking a diff snapshot, Firecracker resets the dirty bitmap of each
guest memory region if the region itself was successfully dumped to the
snapshot memory file. This is wrong, because a later region might fail,
and then some regions will have their dirty bitmap reset, yet no
snapshot was successfully taken, meaning a failure during diff snapshot
taking will forever corrupt the dirty bitmap.

Fix this by only reset the dirty bitmap of all regions after (and if!)
all regions were successfully dumped to the file.

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

Author: roypat

CHANGELOG.md

@@ -44,6 +44,9 @@ and this project adheres to
   the `SendCtrlAltDel` command not working for ACPI-enabled guest kernels, by
   dropping the i8042.nopnp argument from the default kernel command line
   Firecracker constructs.
+- #[???](???): Fix failure during diff snapshot creation corrupting internal
+  state about which pages were dirtied since the last snapshot, meaning no
+  further diff snapshots could ever be taken, even if the error is transient.
 
 ## [1.11.0]
 

src/vmm/src/vstate/memory.rs

@@ -273,8 +273,6 @@ impl GuestMemoryExtension for GuestMemoryMmap {
 
         if write_result.is_err() {
             self.store_dirty_bitmap(dirty_bitmap, page_size);
-        } else {
-            self.reset_dirty();
         }
 
         write_result.map_err(MemoryError::WriteMemory)
@@ -622,6 +620,7 @@ mod tests {
 
         let mut file = TempFile::new().unwrap().into_file();
         guest_memory.dump_dirty(&mut file, &dirty_bitmap).unwrap();
+        guest_memory.reset_dirty();
 
         // We can restore from this because this is the first dirty dump.
         let restored_guest_memory = GuestMemoryMmap::from_regions(
@@ -656,6 +655,7 @@ mod tests {
             .unwrap();
 
         guest_memory.dump_dirty(&mut reader, &dirty_bitmap).unwrap();
+        guest_memory.reset_dirty();
 
         // Check that only the dirty regions are dumped.
         let mut diff_file_content = Vec::new();

src/vmm/src/vstate/vm.rs

@@ -262,18 +262,15 @@ impl Vm {
         match snapshot_type {
             SnapshotType::Diff => {
                 let dirty_bitmap = self.get_dirty_bitmap().map_err(DirtyBitmap)?;
-                self.guest_memory().dump_dirty(&mut file, &dirty_bitmap)
+                self.guest_memory().dump_dirty(&mut file, &dirty_bitmap)?;
             }
             SnapshotType::Full => {
-                let dump_res = self.guest_memory().dump(&mut file);
-                if dump_res.is_ok() {
-                    self.reset_dirty_bitmap();
-                    self.guest_memory().reset_dirty();
-                }
-
-                dump_res
+                self.guest_memory().dump(&mut file)?;
+                self.reset_dirty_bitmap();
             }
-        }?;
+        };
+
+        self.guest_memory().reset_dirty();
 
         file.flush()
             .map_err(|err| MemoryBackingFile("flush", err))?;