fix: use temporary buffer for loading vmstate file

roypat · roypat · commit feb4bc1675cf · 2025-08-20T10:03:14.000-04:00
In commit e7504ae ("refactor: cleanup vmm::snapshot module"), firecracker started reading the snapshot vmstate file in a single pass instead of first loading it into a Vec and then deserializing. This seems to have caused some performance regression due to the deserializer doing many successive reads, resulting in many read(2) syscalls. Fix this by going back to first reading the snapshot file into a buffer, and then deserializing from slice instead. Signed-off-by: Patrick Roy <roypat@amazon.co.uk>
diff --git a/src/vmm/src/arch/aarch64/regs.rs b/src/vmm/src/arch/aarch64/regs.rs
@@ -547,7 +547,9 @@ mod tests {
         let mut buf = vec![0; 10000];
 
         Snapshot::new(&v).save(&mut buf.as_mut_slice()).unwrap();
-        let restored: Aarch64RegisterVec = Snapshot::load(&mut buf.as_slice()).unwrap().data;
+        let restored: Aarch64RegisterVec = Snapshot::load_without_crc_check(buf.as_slice())
+            .unwrap()
+            .data;
 
         for (old, new) in v.iter().zip(restored.iter()) {
             assert_eq!(old, new);
@@ -576,7 +578,7 @@ mod tests {
 
         // Total size of registers according IDs are 16 + 16 = 32,
         // but actual data size is 8 + 16 = 24.
-        Snapshot::<Aarch64RegisterVec>::load(&mut buf.as_slice()).unwrap_err();
+        Snapshot::<Aarch64RegisterVec>::load_without_crc_check(buf.as_slice()).unwrap_err();
     }
 
     #[test]
@@ -598,7 +600,7 @@ mod tests {
         Snapshot::new(v).save(&mut buf.as_mut_slice()).unwrap();
 
         // 4096 bit wide registers are not supported.
-        Snapshot::<Aarch64RegisterVec>::load(&mut buf.as_slice()).unwrap_err();
+        Snapshot::<Aarch64RegisterVec>::load_without_crc_check(buf.as_slice()).unwrap_err();
     }
 
     #[test]
diff --git a/src/vmm/src/arch/x86_64/vm.rs b/src/vmm/src/arch/x86_64/vm.rs
@@ -321,7 +321,9 @@ mod tests {
         Snapshot::new(state)
             .save(&mut snapshot_data.as_mut_slice())
             .unwrap();
-        let restored_state: VmState = Snapshot::load(&mut snapshot_data.as_slice()).unwrap().data;
+        let restored_state: VmState = Snapshot::load_without_crc_check(snapshot_data.as_slice())
+            .unwrap()
+            .data;
 
         vm.restore_state(&restored_state).unwrap();
     }
diff --git a/src/vmm/src/device_manager/pci_mngr.rs b/src/vmm/src/device_manager/pci_mngr.rs
@@ -659,7 +659,9 @@ mod tests {
         // object and calling default_vmm() is the easiest way to create one.
         let vmm = default_vmm();
         let device_manager_state: device_manager::DevicesState =
-            Snapshot::load(&mut buf.as_slice()).unwrap().data;
+            Snapshot::load_without_crc_check(buf.as_slice())
+                .unwrap()
+                .data;
         let vm_resources = &mut VmResources::default();
         let restore_args = PciDevicesConstructorArgs {
             vm: vmm.vm.clone(),
diff --git a/src/vmm/src/device_manager/persist.rs b/src/vmm/src/device_manager/persist.rs
@@ -680,7 +680,9 @@ mod tests {
         let mut event_manager = EventManager::new().expect("Unable to create EventManager");
         let vmm = default_vmm();
         let device_manager_state: device_manager::DevicesState =
-            Snapshot::load(&mut buf.as_slice()).unwrap().data;
+            Snapshot::load_without_crc_check(buf.as_slice())
+                .unwrap()
+                .data;
         let vm_resources = &mut VmResources::default();
         let restore_args = MMIODevManagerConstructorArgs {
             mem: vmm.vm.guest_memory(),
diff --git a/src/vmm/src/devices/virtio/balloon/persist.rs b/src/vmm/src/devices/virtio/balloon/persist.rs
@@ -196,7 +196,9 @@ mod tests {
                 mem: guest_mem,
                 restored_from_file: true,
             },
-            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 
diff --git a/src/vmm/src/devices/virtio/block/virtio/persist.rs b/src/vmm/src/devices/virtio/block/virtio/persist.rs
@@ -223,7 +223,9 @@ mod tests {
         // Restore the block device.
         let restored_block = VirtioBlock::restore(
             BlockConstructorArgs { mem: guest_mem },
-            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 
diff --git a/src/vmm/src/devices/virtio/net/persist.rs b/src/vmm/src/devices/virtio/net/persist.rs
@@ -175,7 +175,9 @@ mod tests {
                     mem: guest_mem,
                     mmds: mmds_ds,
                 },
-                &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+                &Snapshot::load_without_crc_check(mem.as_slice())
+                    .unwrap()
+                    .data,
             ) {
                 Ok(restored_net) => {
                     // Test that virtio specific fields are the same.
diff --git a/src/vmm/src/devices/virtio/persist.rs b/src/vmm/src/devices/virtio/persist.rs
@@ -366,8 +366,13 @@ mod tests {
             mem,
             is_activated: true,
         };
-        let restored_queue =
-            Queue::restore(ca, &Snapshot::load(&mut bytes.as_slice()).unwrap().data).unwrap();
+        let restored_queue = Queue::restore(
+            ca,
+            &Snapshot::load_without_crc_check(bytes.as_slice())
+                .unwrap()
+                .data,
+        )
+        .unwrap();
 
         assert_eq!(restored_queue, queue);
     }
@@ -380,7 +385,9 @@ mod tests {
         let state = VirtioDeviceState::from_device(&dummy);
         Snapshot::new(&state).save(&mut mem.as_mut_slice()).unwrap();
 
-        let restored_state: VirtioDeviceState = Snapshot::load(&mut mem.as_slice()).unwrap().data;
+        let restored_state: VirtioDeviceState = Snapshot::load_without_crc_check(mem.as_slice())
+            .unwrap()
+            .data;
         assert_eq!(restored_state, state);
     }
 
@@ -419,7 +426,9 @@ mod tests {
         };
         let restored_mmio_transport = MmioTransport::restore(
             restore_args,
-            &Snapshot::load(&mut buf.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(buf.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 
diff --git a/src/vmm/src/devices/virtio/rng/persist.rs b/src/vmm/src/devices/virtio/rng/persist.rs
@@ -92,7 +92,9 @@ mod tests {
         let guest_mem = create_virtio_mem();
         let restored = Entropy::restore(
             EntropyConstructorArgs { mem: guest_mem },
-            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 
diff --git a/src/vmm/src/devices/virtio/vsock/persist.rs b/src/vmm/src/devices/virtio/vsock/persist.rs
@@ -180,7 +180,9 @@ pub(crate) mod tests {
 
         Snapshot::new(&state).save(&mut mem.as_mut_slice()).unwrap();
 
-        let restored_state: VsockState = Snapshot::load(&mut mem.as_slice()).unwrap().data;
+        let restored_state: VsockState = Snapshot::load_without_crc_check(mem.as_slice())
+            .unwrap()
+            .data;
         let mut restored_device = Vsock::restore(
             VsockConstructorArgs {
                 mem: ctx.mem.clone(),
diff --git a/src/vmm/src/mmds/persist.rs b/src/vmm/src/mmds/persist.rs
@@ -68,7 +68,9 @@ mod tests {
 
         let restored_ns = MmdsNetworkStack::restore(
             Arc::new(Mutex::new(Mmds::default())),
-            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 
diff --git a/src/vmm/src/persist.rs b/src/vmm/src/persist.rs
@@ -676,8 +676,9 @@ mod tests {
             .save(&mut buf.as_mut_slice())
             .unwrap();
 
-        let restored_microvm_state: MicrovmState =
-            Snapshot::load(&mut buf.as_slice()).unwrap().data;
+        let restored_microvm_state: MicrovmState = Snapshot::load_without_crc_check(buf.as_slice())
+            .unwrap()
+            .data;
 
         assert_eq!(restored_microvm_state.vm_info, microvm_state.vm_info);
         assert_eq!(
diff --git a/src/vmm/src/rate_limiter/persist.rs b/src/vmm/src/rate_limiter/persist.rs
@@ -120,8 +120,13 @@ mod tests {
             .save(&mut mem.as_mut_slice())
             .unwrap();
 
-        let restored_tb =
-            TokenBucket::restore((), &Snapshot::load(&mut mem.as_slice()).unwrap().data).unwrap();
+        let restored_tb = TokenBucket::restore(
+            (),
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
+        )
+        .unwrap();
         assert!(tb.partial_eq(&restored_tb));
     }
 
@@ -197,8 +202,13 @@ mod tests {
         Snapshot::new(rate_limiter.save())
             .save(&mut mem.as_mut_slice())
             .unwrap();
-        let restored_rate_limiter =
-            RateLimiter::restore((), &Snapshot::load(&mut mem.as_slice()).unwrap().data).unwrap();
+        let restored_rate_limiter = RateLimiter::restore(
+            (),
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
+        )
+        .unwrap();
 
         assert!(
             rate_limiter
diff --git a/src/vmm/src/snapshot/mod.rs b/src/vmm/src/snapshot/mod.rs
@@ -26,7 +26,7 @@
 pub mod crc;
 mod persist;
 use std::fmt::Debug;
-use std::io::{Read, Write};
+use std::io::{Read, Seek, SeekFrom, Write};
 
 use bincode::config;
 use bincode::config::{Configuration, Fixint, Limit, LittleEndian};
@@ -38,7 +38,7 @@ use serde::{Deserialize, Serialize};
 use crate::persist::SNAPSHOT_VERSION;
 use crate::snapshot::crc::{CRC64Reader, CRC64Writer};
 pub use crate::snapshot::persist::Persist;
-use crate::utils::mib_to_bytes;
+use crate::utils::{mib_to_bytes, u64_to_usize};
 
 #[cfg(target_arch = "x86_64")]
 const SNAPSHOT_MAGIC_ID: u64 = 0x0710_1984_8664_0000u64;
@@ -64,10 +64,14 @@ pub enum SnapshotError {
     InvalidFormatVersion(Version),
     /// Magic value does not match arch: {0}
     InvalidMagic(u64),
+    /// Snapshot file is not long enough to even contain the CRC
+    TooShort,
     /// An error occured during bincode encoding: {0}
     Encode(#[from] EncodeError),
     /// An error occured during bincode decoding: {0}
     Decode(#[from] DecodeError),
+    /// IO Error: {0}
+    Io(#[from] std::io::Error),
 }
 
 fn serialize<S: Serialize, W: Write>(data: &S, write: &mut W) -> Result<(), SnapshotError> {
@@ -86,8 +90,8 @@ struct SnapshotHdr {
 }
 
 impl SnapshotHdr {
-    fn load<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
-        let hdr: SnapshotHdr = bincode::serde::decode_from_std_read(reader, BINCODE_CONFIG)?;
+    fn load(buf: &mut &[u8]) -> Result<Self, SnapshotError> {
+        let (hdr, bytes_read) = bincode::serde::decode_from_slice::<Self, _>(buf, BINCODE_CONFIG)?;
 
         if hdr.magic != SNAPSHOT_MAGIC_ID {
             return Err(SnapshotError::InvalidMagic(hdr.magic));
@@ -98,6 +102,8 @@ impl SnapshotHdr {
             return Err(SnapshotError::InvalidFormatVersion(hdr.version));
         }
 
+        *buf = &buf[bytes_read..];
+
         Ok(hdr)
     }
 }
@@ -138,17 +144,27 @@ impl<Data> Snapshot<Data> {
 }
 
 impl<Data: DeserializeOwned> Snapshot<Data> {
-    fn load_without_crc_check<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
-        let header = SnapshotHdr::load(reader)?;
-        let data = bincode::serde::decode_from_std_read(reader, BINCODE_CONFIG)?;
+    pub(crate) fn load_without_crc_check(mut buf: &[u8]) -> Result<Self, SnapshotError> {
+        let header = SnapshotHdr::load(&mut buf)?;
+        let data = bincode::serde::decode_from_slice(buf, BINCODE_CONFIG)?.0;
         Ok(Self { header, data })
     }
 
     /// Loads a snapshot from the given [`Read`] instance, performing all validations
     /// (CRC, snapshot magic value, snapshot version).
-    pub fn load<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
+    pub fn load<R: Read + Seek>(reader: &mut R) -> Result<Self, SnapshotError> {
+        let snapshot_size = reader.seek(SeekFrom::End(0))?;
+        reader.seek(SeekFrom::Start(0))?;
+        // dont read the CRC yet.
+        let mut buf = vec![
+            0;
+            u64_to_usize(snapshot_size)
+                .checked_sub(size_of::<u64>())
+                .ok_or(SnapshotError::TooShort)?
+        ];
         let mut crc_reader = CRC64Reader::new(reader);
-        let snapshot = Self::load_without_crc_check(&mut crc_reader)?;
+        crc_reader.read_exact(buf.as_mut_slice())?;
+        let snapshot = Self::load_without_crc_check(buf.as_slice())?;
         let computed_checksum = crc_reader.checksum();
         let stored_checksum: u64 =
             bincode::serde::decode_from_std_read(&mut crc_reader.reader, BINCODE_CONFIG)?;
@@ -199,10 +215,16 @@ mod tests {
             }
         }
 
+        impl Seek for BadReader {
+            fn seek(&mut self, _pos: SeekFrom) -> std::io::Result<u64> {
+                Ok(9) // needs to be long enough to prevent to have a CRC
+            }
+        }
+
         let mut reader = BadReader {};
 
         assert!(
-            matches!(Snapshot::<()>::load(&mut reader), Err(SnapshotError::Decode(DecodeError::Io {inner, ..})) if inner.kind() == std::io::ErrorKind::InvalidInput)
+            matches!(Snapshot::<()>::load(&mut reader), Err(SnapshotError::Io(inner)) if inner.kind() == std::io::ErrorKind::InvalidInput)
         );
     }
 
@@ -239,7 +261,7 @@ mod tests {
         serialize(&snapshot, &mut data.as_mut_slice()).unwrap();
 
         assert!(matches!(
-            Snapshot::<()>::load(&mut data.as_slice()),
+            Snapshot::<()>::load(&mut std::io::Cursor::new(data.as_slice())),
             Err(SnapshotError::Crc64(_))
         ));
     }
@@ -254,7 +276,7 @@ mod tests {
         snapshot.save(&mut data.as_mut_slice()).unwrap();
 
         assert!(matches!(
-            Snapshot::<()>::load(&mut data.as_slice()),
+            Snapshot::<()>::load_without_crc_check(data.as_slice()),
             Err(SnapshotError::InvalidFormatVersion(v)) if v.major == SNAPSHOT_VERSION.major + 1
         ));
 
@@ -263,36 +285,36 @@ mod tests {
         snapshot.header.version.minor = SNAPSHOT_VERSION.minor + 1;
         snapshot.save(&mut data.as_mut_slice()).unwrap();
         assert!(matches!(
-            Snapshot::<()>::load(&mut data.as_slice()),
+            Snapshot::<()>::load_without_crc_check(data.as_slice()),
             Err(SnapshotError::InvalidFormatVersion(v)) if v.minor == SNAPSHOT_VERSION.minor + 1
         ));
 
         // But we can support minor versions smaller or equal to ours. We also support
         // all patch versions within our supported major.minor version.
         let snapshot = Snapshot::new(());
         snapshot.save(&mut data.as_mut_slice()).unwrap();
-        Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+        Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
 
         if SNAPSHOT_VERSION.minor != 0 {
             let mut snapshot = Snapshot::new(());
             snapshot.header.version.minor = SNAPSHOT_VERSION.minor - 1;
             snapshot.save(&mut data.as_mut_slice()).unwrap();
-            Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+            Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
         }
 
         let mut snapshot = Snapshot::new(());
         snapshot.header.version.patch = 0;
         snapshot.save(&mut data.as_mut_slice()).unwrap();
-        Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+        Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
 
         let mut snapshot = Snapshot::new(());
         snapshot.header.version.patch = SNAPSHOT_VERSION.patch + 1;
         snapshot.save(&mut data.as_mut_slice()).unwrap();
-        Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+        Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
 
         let mut snapshot = Snapshot::new(());
         snapshot.header.version.patch = 1024;
         snapshot.save(&mut data.as_mut_slice()).unwrap();
-        Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+        Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
     }
 }
diff --git a/src/vmm/src/vmm_config/boot_source.rs b/src/vmm/src/vmm_config/boot_source.rs
@@ -132,7 +132,9 @@ pub(crate) mod tests {
         Snapshot::new(&boot_src_cfg)
             .save(&mut snapshot_data.as_mut_slice())
             .unwrap();
-        let restored_boot_cfg = Snapshot::load(&mut snapshot_data.as_slice()).unwrap().data;
+        let restored_boot_cfg = Snapshot::load_without_crc_check(snapshot_data.as_slice())
+            .unwrap()
+            .data;
         assert_eq!(boot_src_cfg, restored_boot_cfg);
     }
 }
diff --git a/src/vmm/src/vstate/memory.rs b/src/vmm/src/vstate/memory.rs
@@ -436,7 +436,9 @@ mod tests {
         Snapshot::new(&original_state)
             .save(&mut snapshot_data.as_mut_slice())
             .unwrap();
-        let restored_state = Snapshot::load(&mut snapshot_data.as_slice()).unwrap().data;
+        let restored_state = Snapshot::load_without_crc_check(snapshot_data.as_slice())
+            .unwrap()
+            .data;
         assert_eq!(original_state, restored_state);
     }
 
diff --git a/src/vmm/src/vstate/resources.rs b/src/vmm/src/vstate/resources.rs
diff --git a/src/vmm/src/vstate/vm.rs b/src/vmm/src/vstate/vm.rs
diff --git a/tests/integration_tests/functional/test_snapshot_basic.py b/tests/integration_tests/functional/test_snapshot_basic.py
