Refactor Snapshot Module

src/firecracker/src/main.rs

@@ -529,17 +529,16 @@ fn warn_deprecated_parameters() {}
 enum SnapshotVersionError {
     /// Unable to open snapshot state file: {0}
     OpenSnapshot(io::Error),
-    /// Invalid data format version of snapshot file: {0}
-    SnapshotVersion(SnapshotError),
+    /// Invalid data format of snapshot file: {0}
+    LoadSnapshot(SnapshotError),
 }
 
 // Print data format of provided snapshot state file.
 fn print_snapshot_data_format(snapshot_path: &str) -> Result<(), SnapshotVersionError> {
-    let mut snapshot_reader =
-        File::open(snapshot_path).map_err(SnapshotVersionError::OpenSnapshot)?;
+    let mut snapshot_reader = File::open(snapshot_path).map_err(SnapshotVersionError::OpenSnapshot);
 
-    let data_format_version = Snapshot::get_format_version(&mut snapshot_reader)
-        .map_err(SnapshotVersionError::SnapshotVersion)?;
+    let snapshot = Snapshot::load(snapshot_reader).map_err();
+    let data_format_version = snapshot.version(SnapshotVersionError::LoadSnapshot);
 
     println!("v{}", data_format_version);
     Ok(())

src/snapshot-editor/src/utils.rs

@@ -30,7 +30,18 @@ pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<(MicrovmState, Version),
     let mut snapshot_reader = File::open(snapshot_path).map_err(UtilsError::VmStateFileOpen)?;
     let metadata = std::fs::metadata(snapshot_path).map_err(UtilsError::VmStateFileMeta)?;
     let snapshot_len = u64_to_usize(metadata.len());
-    Snapshot::load(&mut snapshot_reader, snapshot_len).map_err(UtilsError::VmStateLoad)
+
+    let snapshot: Result<Snapshot<MicrovmState>, UtilsError> = Snapshot::load(&mut snapshot_reader, snapshot_len).map_err(UtilsError::VmStateLoad);
+    match snapshot {
+        Ok(snapshot) => {
+            let version = snapshot.version();
+            Ok((snapshot.data, version.to_owned()))
+        }
+        Err(e) => {
+            return Err(e);
+        }
+    }
+
 }
 
 // This method is used only in aarch64 code so far
@@ -46,9 +57,10 @@ pub fn save_vmstate(
         .truncate(true)
         .open(output_path)
         .map_err(UtilsError::OutputFileOpen)?;
-    let mut snapshot = Snapshot::new(version);
+    let snapshot_hdr = SnapshotHdr::new(version);
+    let snapshot = Snapshot::new(snapshot_hdr, microvm_state);
     snapshot
-        .save(&mut output_file, &microvm_state)
+        .save(&mut output_file)
         .map_err(UtilsError::VmStateSave)?;
     Ok(())
 }

src/vmm/src/persist.rs

@@ -425,15 +425,12 @@ pub enum SnapshotStateFromFileError {
 fn snapshot_state_from_file(
     snapshot_path: &Path,
 ) -> Result<MicrovmState, SnapshotStateFromFileError> {
-    let snapshot = Snapshot::new(SNAPSHOT_VERSION);
     let mut snapshot_reader =
         File::open(snapshot_path).map_err(SnapshotStateFromFileError::Open)?;
-    let metadata = std::fs::metadata(snapshot_path).map_err(SnapshotStateFromFileError::Meta)?;
-    let snapshot_len = u64_to_usize(metadata.len());
-    let state: MicrovmState = snapshot
-        .load_with_version_check(&mut snapshot_reader, snapshot_len)
+    let state: Snapshot<MicrovmState> = Snapshot::load_with_version_check(&mut snapshot_reader, SNAPSHOT_VERSION)
         .map_err(SnapshotStateFromFileError::Load)?;
-    Ok(state)
+
+    Ok(state.data)
 }
 
 /// Error type for [`guest_memory_from_file`].
@@ -686,17 +683,22 @@ mod tests {
             vm_state: vmm.vm.save_state().unwrap(),
             acpi_dev_state: vmm.acpi_device_manager.save(),
         };
+        let vm_info = microvm_state.vm_info.clone();
+        let device_states = microvm_state.device_states.clone();
 
         let mut buf = vec![0; 10000];
-        Snapshot::serialize(&mut buf.as_mut_slice(), &microvm_state).unwrap();
 
-        let restored_microvm_state: MicrovmState =
-            Snapshot::deserialize(&mut buf.as_slice()).unwrap();
+        let snapshot = Snapshot::new(Version::new(1, 0, 42), microvm_state);
+        snapshot.save(&mut buf.as_mut_slice()).unwrap();
+
+        let restored_snapshot: Snapshot<MicrovmState> =
+            Snapshot::load(&mut buf.as_slice());
+        let restored_microvm_state = restored_snapshot.data;
 
-        assert_eq!(restored_microvm_state.vm_info, microvm_state.vm_info);
+        assert_eq!(restored_microvm_state.vm_info, vm_info);
         assert_eq!(
             restored_microvm_state.device_states,
-            microvm_state.device_states
+            device_states
         )
     }