Allow immediate merging of diff snapshots during creation

Allows the automatic merging of diff snapshots with a base snapshot by
directly writing the diff snapshot on top of a base snapshot file of
matching size.

Signed-off-by: Patrick Roy <[email protected]>
Co-authored-by: Ives van Hoorne <[email protected]>

Author: roypat

CHANGELOG.md

@@ -59,6 +59,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Removed support for creating Firecracker snapshots targeting older versions
   of Firecracker. With this change, running 'firecracker --version' will not
   print the supported snapshot versions.
+- [#4301](https://github.com/firecracker-microvm/firecracker/pull/4301):
+  Allow merging of diff snapshots into base snapshots by directly writing
+  the diff snapshot on top of the base snapshot's memory file. This can be
+  done by setting the `mem_file_path` to the path of the pre-existing full
+  snapshot.
 
 ### Deprecated
 
@@ -87,6 +92,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Fixed a bug in the asynchronous virtio-block engine that rendered the device
   non-functional after a PATCH request was issued to Firecracker for updating
   the path to the host-side backing file of the device.
+- [#4301](https://github.com/firecracker-microvm/firecracker/pull/4301):
+  Fixed a bug where if Firecracker was instructed to take a snapshot of a
+  microvm which itself was restored from a snapshot, specifying `mem_file_path`
+  to be the path of the memory file from which the microvm was restored would
+  result in both the microvm and the snapshot being corrupted. It now instead
+  performs a "write-back" of all memory that was updated since the snapshot
+  was originally loaded.
 
 ## [1.5.0]
 

src/vmm/src/persist.rs

@@ -174,6 +174,9 @@ pub enum CreateSnapshotError {
     SerializeMicrovmState(snapshot::Error),
     /// Cannot perform {0} on the snapshot backing file: {1}
     SnapshotBackingFile(&'static str, io::Error),
+    #[rustfmt::skip]
+    #[doc = "Size mismatch when writing diff snapshot on top of base layer: base layer size is {0} but diff layer is size {1}."]
+    SnapshotBackingFileLengthMismatch(u64, u64),
     #[cfg(target_arch = "x86_64")]
     #[rustfmt::skip]
     #[doc = "Too many devices attached: {0}. The maximum number allowed for the snapshot data version requested is {FC_V0_23_MAX_DEVICES:}."]
@@ -231,23 +234,54 @@ fn snapshot_state_to_file(
         .map_err(|err| SnapshotBackingFile("sync_all", err))
 }
 
+/// Takes a snapshot of the virtual machine running inside the given [`Vmm`] and saves it to
+/// `mem_file_path`.
+///
+/// If `snapshot_type` is [`SnapshotType::Diff`], and `mem_file_path` exists and is a snapshot file
+/// of matching size, then the diff snapshot will be directly merged into the existing snapshot.
+/// Otherwise, existing files are simply overwritten.
 fn snapshot_memory_to_file(
     vmm: &Vmm,
     mem_file_path: &Path,
     snapshot_type: SnapshotType,
 ) -> Result<(), CreateSnapshotError> {
     use self::CreateSnapshotError::*;
+
+    // Need to check this here, as we create the file in the line below
+    let file_existed = mem_file_path.exists();
+
     let mut file = OpenOptions::new()
         .write(true)
         .create(true)
-        .truncate(true)
         .open(mem_file_path)
         .map_err(|err| MemoryBackingFile("open", err))?;
 
-    // Set the length of the file to the full size of the memory area.
+    // Determine what size our total memory area is.
     let mem_size_mib = mem_size_mib(vmm.guest_memory());
-    file.set_len(mem_size_mib * 1024 * 1024)
-        .map_err(|err| MemoryBackingFile("set_length", err))?;
+    let expected_size = mem_size_mib * 1024 * 1024;
+
+    if file_existed {
+        let file_size = file
+            .metadata()
+            .map_err(|e| MemoryBackingFile("get_metadata", e))?
+            .len();
+
+        // Here we only truncate the file if the size mismatches.
+        // - For full snapshots, the entire file's contents will be overwritten anyway. We have to
+        //   avoid truncating here to deal with the edge case where it represents the snapshot file
+        //   from which this very microVM was loaded (as modifying the memory file would be
+        //   reflected in the mmap of the file, meaning a truncate operation would zero out guest
+        //   memory, and thus corrupt the VM).
+        // - For diff snapshots, we want to merge the diff layer directly into the file.
+        if file_size != expected_size {
+            file.set_len(0)
+                .map_err(|err| MemoryBackingFile("truncate", err))?;
+        }
+    }
+
+    // Set the length of the file to the full size of the memory area.
+    file.set_len(expected_size)
+        .map_err(|e| MemoryBackingFile("set_length", e))?;
 
     match snapshot_type {
         SnapshotType::Diff => {