refactor: cleanup vmm::snapshot module

Implement the suggestions from firecracker-microvm#4523 to give a cleaner API for the
snapshot module. Conceptually, the change is that we now store the thing
we are snapshotting inside of the `Snapshot` struct.
Implementation-wise, this allows us to deserialize the snapshot data in
a single pass (instead of having to read it twice, once for
deserialization and once for CRC validation), which means we do not have
to first query the snapshot file for the size of the snapshot.

Closes firecracker-microvm#4691, firecracker-microvm#4523, firecracker-microvm#5195

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

Author: roypat

src/firecracker/src/main.rs

@@ -28,7 +28,7 @@ use vmm::persist::SNAPSHOT_VERSION;
 use vmm::resources::VmResources;
 use vmm::seccomp::BpfThreadMap;
 use vmm::signal_handler::register_signal_handlers;
-use vmm::snapshot::{Snapshot, SnapshotError};
+use vmm::snapshot::{SnapshotError, get_format_version};
 use vmm::vmm_config::instance_info::{InstanceInfo, VmState};
 use vmm::vmm_config::metrics::{MetricsConfig, MetricsConfigError, init_metrics};
 use vmm::{EventManager, FcExitCode, HTTP_MAX_PAYLOAD_SIZE};
@@ -538,8 +538,8 @@ fn print_snapshot_data_format(snapshot_path: &str) -> Result<(), SnapshotVersion
     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 data_format_version =
+        get_format_version(&mut snapshot_reader).map_err(SnapshotVersionError::SnapshotVersion)?;
 
     println!("v{}", data_format_version);
     Ok(())

src/snapshot-editor/src/edit_vmstate.rs

@@ -51,8 +51,8 @@ fn edit(
     output_path: &PathBuf,
     f: impl Fn(MicrovmState) -> Result<MicrovmState, EditVmStateError>,
 ) -> Result<(), EditVmStateError> {
-    let (microvm_state, _) = open_vmstate(vmstate_path)?;
-    let microvm_state = f(microvm_state)?;
+    let snapshot = open_vmstate(vmstate_path)?;
+    let microvm_state = f(snapshot.data)?;
     save_vmstate(microvm_state, output_path)?;
     Ok(())
 }

src/snapshot-editor/src/info.rs

@@ -4,8 +4,8 @@
 use std::path::PathBuf;
 
 use clap::Subcommand;
-use semver::Version;
 use vmm::persist::MicrovmState;
+use vmm::snapshot::Snapshot;
 
 use crate::utils::*;
 
@@ -50,27 +50,27 @@ pub fn info_vmstate_command(command: InfoVmStateSubCommand) -> Result<(), InfoVm
 
 fn info(
     vmstate_path: &PathBuf,
-    f: impl Fn(&MicrovmState, Version) -> Result<(), InfoVmStateError>,
+    f: impl Fn(&Snapshot<MicrovmState>) -> Result<(), InfoVmStateError>,
 ) -> Result<(), InfoVmStateError> {
-    let (vmstate, version) = open_vmstate(vmstate_path)?;
-    f(&vmstate, version)?;
+    let snapshot = open_vmstate(vmstate_path)?;
+    f(&snapshot)?;
     Ok(())
 }
 
-fn info_version(_: &MicrovmState, version: Version) -> Result<(), InfoVmStateError> {
-    println!("v{version}");
+fn info_version(snapshot: &Snapshot<MicrovmState>) -> Result<(), InfoVmStateError> {
+    println!("v{}", snapshot.version());
     Ok(())
 }
 
-fn info_vcpu_states(state: &MicrovmState, _: Version) -> Result<(), InfoVmStateError> {
-    for (i, state) in state.vcpu_states.iter().enumerate() {
+fn info_vcpu_states(snapshot: &Snapshot<MicrovmState>) -> Result<(), InfoVmStateError> {
+    for (i, state) in snapshot.data.vcpu_states.iter().enumerate() {
         println!("vcpu {i}:");
         println!("{state:#?}");
     }
     Ok(())
 }
 
-fn info_vmstate(vmstate: &MicrovmState, _version: Version) -> Result<(), InfoVmStateError> {
-    println!("{vmstate:#?}");
+fn info_vmstate(snapshot: &Snapshot<MicrovmState>) -> Result<(), InfoVmStateError> {
+    println!("{:#?}", snapshot.data);
     Ok(())
 }

src/snapshot-editor/src/utils.rs

@@ -4,19 +4,15 @@
 use std::fs::{File, OpenOptions};
 use std::path::PathBuf;
 
-use semver::Version;
-use vmm::persist::{MicrovmState, SNAPSHOT_VERSION};
+use vmm::persist::MicrovmState;
 use vmm::snapshot::Snapshot;
-use vmm::utils::u64_to_usize;
 
 // Some errors are only used in aarch64 code
 #[allow(unused)]
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum UtilsError {
     /// Can not open snapshot file: {0}
     VmStateFileOpen(std::io::Error),
-    /// Can not retrieve metadata for snapshot file: {0}
-    VmStateFileMeta(std::io::Error),
     /// Can not load snapshot: {0}
     VmStateLoad(vmm::snapshot::SnapshotError),
     /// Can not open output file: {0}
@@ -26,11 +22,9 @@ pub enum UtilsError {
 }
 
 #[allow(unused)]
-pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<(MicrovmState, Version), UtilsError> {
+pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<Snapshot<MicrovmState>, UtilsError> {
     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)
+    Snapshot::load(&mut snapshot_reader).map_err(UtilsError::VmStateLoad)
 }
 
 // This method is used only in aarch64 code so far
@@ -42,9 +36,9 @@ pub fn save_vmstate(microvm_state: MicrovmState, output_path: &PathBuf) -> Resul
         .truncate(true)
         .open(output_path)
         .map_err(UtilsError::OutputFileOpen)?;
-    let mut snapshot = Snapshot::new(SNAPSHOT_VERSION);
+    let mut snapshot = Snapshot::new(microvm_state);
     snapshot
-        .save(&mut output_file, &microvm_state)
+        .save(&mut output_file)
         .map_err(UtilsError::VmStateSave)?;
     Ok(())
 }

src/vmm/src/arch/aarch64/regs.rs

@@ -546,8 +546,8 @@ mod tests {
 
         let mut buf = vec![0; 10000];
 
-        Snapshot::serialize(&mut buf.as_mut_slice(), &v).unwrap();
-        let restored: Aarch64RegisterVec = Snapshot::deserialize(&mut buf.as_slice()).unwrap();
+        Snapshot::new(&v).save(&mut buf.as_mut_slice()).unwrap();
+        let restored: Aarch64RegisterVec = Snapshot::load(&mut buf.as_slice()).unwrap().data;
 
         for (old, new) in v.iter().zip(restored.iter()) {
             assert_eq!(old, new);
@@ -572,11 +572,11 @@ mod tests {
 
         let mut buf = vec![0; 10000];
 
-        Snapshot::serialize(&mut buf.as_mut_slice(), &v).unwrap();
+        Snapshot::new(&v).save(&mut buf.as_mut_slice()).unwrap();
 
         // Total size of registers according IDs are 16 + 16 = 32,
         // but actual data size is 8 + 16 = 24.
-        Snapshot::deserialize::<_, Aarch64RegisterVec>(&mut buf.as_slice()).unwrap_err();
+        Snapshot::<Aarch64RegisterVec>::load(&mut buf.as_slice()).unwrap_err();
     }
 
     #[test]
@@ -595,10 +595,10 @@ mod tests {
 
         let mut buf = vec![0; 10000];
 
-        Snapshot::serialize(&mut buf.as_mut_slice(), &v).unwrap();
+        Snapshot::new(v).save(&mut buf.as_mut_slice()).unwrap();
 
         // 4096 bit wide registers are not supported.
-        Snapshot::deserialize::<_, Aarch64RegisterVec>(&mut buf.as_slice()).unwrap_err();
+        Snapshot::<Aarch64RegisterVec>::load(&mut buf.as_slice()).unwrap_err();
     }
 
     #[test]

src/vmm/src/arch/x86_64/vm.rs

@@ -306,8 +306,10 @@ mod tests {
         let (_, mut vm) = setup_vm_with_memory(0x1000);
         vm.setup_irqchip().unwrap();
         let state = vm.save_state().unwrap();
-        Snapshot::serialize(&mut snapshot_data.as_mut_slice(), &state).unwrap();
-        let restored_state: VmState = Snapshot::deserialize(&mut snapshot_data.as_slice()).unwrap();
+        Snapshot::new(state)
+            .save(&mut snapshot_data.as_mut_slice())
+            .unwrap();
+        let restored_state: VmState = Snapshot::load(&mut snapshot_data.as_slice()).unwrap().data;
 
         vm.restore_state(&restored_state).unwrap();
     }

src/vmm/src/device_manager/persist.rs

@@ -768,7 +768,9 @@ mod tests {
             let entropy_config = EntropyDeviceConfig::default();
             insert_entropy_device(&mut vmm, &mut cmdline, &mut event_manager, entropy_config);
 
-            Snapshot::serialize(&mut buf.as_mut_slice(), &vmm.mmio_device_manager.save()).unwrap();
+            Snapshot::new(vmm.mmio_device_manager.save())
+                .save(&mut buf.as_mut_slice())
+                .unwrap();
 
             // We only want to keep the device map from the original MmioDeviceManager.
             vmm.mmio_device_manager.soft_clone()
@@ -777,7 +779,7 @@ mod tests {
 
         let mut event_manager = EventManager::new().expect("Unable to create EventManager");
         let vmm = default_vmm();
-        let device_states: DeviceStates = Snapshot::deserialize(&mut buf.as_slice()).unwrap();
+        let device_states: DeviceStates = Snapshot::load(&mut buf.as_slice()).unwrap().data;
         let vm_resources = &mut VmResources::default();
         let restore_args = MMIODevManagerConstructorArgs {
             mem: vmm.vm.guest_memory(),

src/vmm/src/devices/virtio/balloon/persist.rs

@@ -194,15 +194,17 @@ mod tests {
         // Create and save the balloon device.
         let balloon = Balloon::new(0x42, false, 2, false).unwrap();
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &balloon.save()).unwrap();
+        Snapshot::new(balloon.save())
+            .save(&mut mem.as_mut_slice())
+            .unwrap();
 
         // Deserialize and restore the balloon device.
         let restored_balloon = Balloon::restore(
             BalloonConstructorArgs {
                 mem: guest_mem,
                 restored_from_file: true,
             },
-            &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
+            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
         )
         .unwrap();
 

src/vmm/src/devices/virtio/block/virtio/persist.rs

@@ -186,7 +186,9 @@ mod tests {
         // Save the block device.
         let mut mem = vec![0; 4096];
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &block.save()).unwrap();
+        Snapshot::new(block.save())
+            .save(&mut mem.as_mut_slice())
+            .unwrap();
     }
 
     #[test]
@@ -229,12 +231,14 @@ mod tests {
         // Save the block device.
         let mut mem = vec![0; 4096];
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &block.save()).unwrap();
+        Snapshot::new(block.save())
+            .save(&mut mem.as_mut_slice())
+            .unwrap();
 
         // Restore the block device.
         let restored_block = VirtioBlock::restore(
             BlockConstructorArgs { mem: guest_mem },
-            &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
+            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
         )
         .unwrap();
 

src/vmm/src/devices/virtio/net/persist.rs

@@ -195,7 +195,9 @@ mod tests {
 
         // Create and save the net device.
         {
-            Snapshot::serialize(&mut mem.as_mut_slice(), &net.save()).unwrap();
+            Snapshot::new(net.save())
+                .save(&mut mem.as_mut_slice())
+                .unwrap();
 
             // Save some fields that we want to check later.
             id = net.id.clone();
@@ -215,7 +217,7 @@ mod tests {
                     mem: guest_mem,
                     mmds: mmds_ds,
                 },
-                &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
+                &Snapshot::load(&mut mem.as_slice()).unwrap().data,
             ) {
                 Ok(restored_net) => {
                     // Test that virtio specific fields are the same.