[snapshot-ga] Additional info added to snap state

Signed-off-by: Vibha Acharya <vibharya@amazon.co.uk>

Author: vibharya

src/vmm/benches/main.rs

@@ -16,10 +16,12 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use snapshot::Snapshot;
 use utils::tempfile::TempFile;
 use versionize::VersionMap;
-use vmm::persist::MicrovmState;
+use vmm::persist::{MicrovmState, VmInfo};
 use vmm::utilities::mock_resources::NOISY_KERNEL_IMAGE;
 use vmm::utilities::test_utils::create_vmm;
 use vmm::version_map::VERSION_MAP;
+use vmm::vmm_config::boot_source::BootSourceConfig;
+use vmm::vmm_config::machine_config::CpuFeaturesTemplate;
 use vmm::vmm_config::snapshot::{CreateSnapshotParams, SnapshotType};
 use vmm::{persist, FcExitCode};
 
@@ -78,10 +80,20 @@ fn create_microvm_state(is_diff: bool) -> MicrovmState {
         mem_file_path: memory_file.as_path().to_path_buf(),
         version: None,
     };
+    let mut vm_info = VmInfo {
+        mem_size_mib: 1u64,
+        ..Default::default()
+    };
 
     {
         let mut locked_vmm = vmm.lock().unwrap();
-        persist::create_snapshot(&mut locked_vmm, &snapshot_params, VERSION_MAP.clone()).unwrap();
+        persist::create_snapshot(
+            &mut locked_vmm,
+            &mut vm_info,
+            &snapshot_params,
+            VERSION_MAP.clone(),
+        )
+        .unwrap();
     }
 
     vmm.lock().unwrap().stop(FcExitCode::Ok);

src/vmm/src/builder.rs

@@ -51,7 +51,7 @@ use crate::vmm_config::machine_config::{VmConfigError, VmUpdateConfig};
 use crate::vstate::system::KvmContext;
 use crate::vstate::vcpu::{Vcpu, VcpuConfig};
 use crate::vstate::vm::Vm;
-use crate::{device_manager, mem_size_mib, Error, EventManager, Vmm, VmmEventsObserver};
+use crate::{device_manager, Error, EventManager, Vmm, VmmEventsObserver};
 
 /// Errors associated with starting the instance.
 #[derive(Debug)]
@@ -565,12 +565,15 @@ pub fn build_microvm_from_snapshot(
 
     vm_resources.update_vm_config(&VmUpdateConfig {
         vcpu_count: Some(vcpu_count),
-        mem_size_mib: Some(mem_size_mib(&guest_memory) as usize),
-        smt: Some(false),
-        cpu_template: None,
+        mem_size_mib: Some(microvm_state.vm_info.mem_size_mib as usize),
+        smt: Some(microvm_state.vm_info.smt),
+        cpu_template: Some(microvm_state.vm_info.cpu_template),
         track_dirty_pages: Some(track_dirty_pages),
     })?;
 
+    // Restore the boot source config paths.
+    vm_resources.set_boot_source_config(microvm_state.vm_info.boot_source);
+
     // Restore devices states.
     let mmio_ctor_args = MMIODevManagerConstructorArgs {
         mem: guest_memory,

src/vmm/src/lib.rs

@@ -461,7 +461,10 @@ impl Vmm {
     }
 
     /// Saves the state of a paused Microvm.
-    pub fn save_state(&mut self) -> std::result::Result<MicrovmState, MicrovmStateError> {
+    pub fn save_state(
+        &mut self,
+        vm_info: &VmInfo,
+    ) -> std::result::Result<MicrovmState, MicrovmStateError> {
         use self::MicrovmStateError::SaveVmState;
         let vcpu_states = self.save_vcpu_states()?;
         let vm_state = {
@@ -478,11 +481,10 @@ impl Vmm {
         };
         let device_states = self.mmio_device_manager.save();
 
-        let mem_size_mib = mem_size_mib(self.guest_memory());
         let memory_state = self.guest_memory().describe();
 
         Ok(MicrovmState {
-            vm_info: VmInfo { mem_size_mib },
+            vm_info: vm_info.clone(),
             memory_state,
             vm_state,
             vcpu_states,

src/vmm/src/persist.rs

@@ -16,7 +16,7 @@ use arch::regs::{get_manufacturer_id_from_host, get_manufacturer_id_from_state};
 #[cfg(target_arch = "x86_64")]
 use cpuid::common::{get_vendor_id_from_cpuid, get_vendor_id_from_host};
 use devices::virtio::TYPE_NET;
-use logger::{error, info};
+use logger::{error, info, warn};
 use seccompiler::BpfThreadMap;
 use serde::Serialize;
 use snapshot::Snapshot;
@@ -34,8 +34,9 @@ use crate::resources::VmResources;
 #[cfg(target_arch = "x86_64")]
 use crate::version_map::FC_V0_23_SNAP_VERSION;
 use crate::version_map::{FC_V1_0_SNAP_VERSION, FC_V1_1_SNAP_VERSION, FC_VERSION_TO_SNAP_VERSION};
+use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
-use crate::vmm_config::machine_config::MAX_SUPPORTED_VCPUS;
+use crate::vmm_config::machine_config::{CpuFeaturesTemplate, MAX_SUPPORTED_VCPUS};
 use crate::vmm_config::snapshot::{
     CreateSnapshotParams, LoadSnapshotParams, MemBackendType, SnapshotType,
 };
@@ -47,11 +48,59 @@ use crate::{mem_size_mib, memory_snapshot, vstate, Error as VmmError, EventManag
 const FC_V0_23_MAX_DEVICES: u32 = 11;
 
 /// Holds information related to the VM that is not part of VmState.
-#[derive(Debug, PartialEq, Versionize)]
+#[derive(Clone, Debug, Default, PartialEq, Versionize, Serialize)]
 // NOTICE: Any changes to this structure require a snapshot version bump.
 pub struct VmInfo {
     /// Guest memory size.
     pub mem_size_mib: u64,
+    /// smt information
+    #[version(start = 2, default_fn = "def_smt", ser_fn = "ser_smt")]
+    pub smt: bool,
+    /// CPU template type
+    #[version(
+        start = 2,
+        default_fn = "def_cpu_template",
+        ser_fn = "ser_cpu_template"
+    )]
+    pub cpu_template: CpuFeaturesTemplate,
+    /// Boot source information.
+    #[version(start = 2, default_fn = "def_boot_source", ser_fn = "ser_boot_source")]
+    pub boot_source: BootSourceConfig,
+}
+
+impl VmInfo {
+    fn def_smt(_: u16) -> bool {
+        warn!("SMT field not found in snapshot.");
+        false
+    }
+
+    fn ser_smt(&mut self, _target_version: u16) -> VersionizeResult<()> {
+        // v1.1 and older versions do not include smt info.
+        warn!("Saving to older snapshot version, SMT information will not be saved.");
+        Ok(())
+    }
+
+    fn def_cpu_template(_: u16) -> CpuFeaturesTemplate {
+        warn!("CPU template field not found in snapshot.");
+        CpuFeaturesTemplate::None
+    }
+
+    fn ser_cpu_template(&mut self, _target_version: u16) -> VersionizeResult<()> {
+        // v1.1 and older versions do not include cpu template info.
+        warn!("Saving to older snapshot version, CPU template information will not be saved.");
+        Ok(())
+    }
+
+    fn def_boot_source(_: u16) -> BootSourceConfig {
+        warn!("Boot source information not found in snapshot.");
+        BootSourceConfig::default()
+    }
+
+    fn ser_boot_source(&mut self, _target_version: u16) -> VersionizeResult<()> {
+        // v1.1 and older versions do not include boot source info.
+        warn!("Saving to older snapshot version, boot source information will not be saved.");
+        Ok(())
+    }
 }
 
 /// Contains the necesary state for saving/restoring a microVM.
@@ -195,14 +244,15 @@ impl Display for CreateSnapshotError {
 /// Creates a Microvm snapshot.
 pub fn create_snapshot(
     vmm: &mut Vmm,
+    vm_info: &VmInfo,
     params: &CreateSnapshotParams,
     version_map: VersionMap,
 ) -> std::result::Result<(), CreateSnapshotError> {
     // Fail early from invalid target version.
     let snapshot_data_version = get_snapshot_data_version(&params.version, &version_map, &vmm)?;
 
     let microvm_state = vmm
-        .save_state()
+        .save_state(vm_info)
         .map_err(CreateSnapshotError::MicrovmState)?;
 
     snapshot_state_to_file(
@@ -787,7 +837,10 @@ mod tests {
             device_states: states,
             memory_state,
             vcpu_states,
-            vm_info: VmInfo { mem_size_mib: 1u64 },
+            vm_info: VmInfo {
+                mem_size_mib: 1u64,
+                ..Default::default()
+            },
             #[cfg(target_arch = "aarch64")]
             vm_state: vmm.vm.save_state(&mpidrs).unwrap(),
             #[cfg(target_arch = "x86_64")]

src/vmm/src/rpc_interface.rs

@@ -22,7 +22,7 @@ use super::{
     resources::VmResources, Vmm,
 };
 use crate::builder::StartMicrovmError;
-use crate::persist::{CreateSnapshotError, RestoreFromSnapshotError};
+use crate::persist::{CreateSnapshotError, RestoreFromSnapshotError, VmInfo};
 use crate::resources::VmmConfig;
 use crate::version_map::VERSION_MAP;
 use crate::vmm_config::balloon::{
@@ -743,9 +743,21 @@ impl RuntimeApiController {
         }
 
         let mut locked_vmm = self.vmm.lock().unwrap();
+        let vm_cfg = self.vm_resources.vm_config();
+        let vm_info = VmInfo {
+            mem_size_mib: vm_cfg.mem_size_mib as u64,
+            smt: vm_cfg.smt,
+            cpu_template: vm_cfg.cpu_template,
+            boot_source: self.vm_resources.boot_source_config().clone(),
+        };
         let create_start_us = utils::time::get_time_us(utils::time::ClockType::Monotonic);
 
-        create_snapshot(&mut locked_vmm, create_params, VERSION_MAP.clone())?;
+        create_snapshot(
+            &mut locked_vmm,
+            &vm_info,
+            create_params,
+            VERSION_MAP.clone(),
+        )?;
 
         match create_params.snapshot_type {
             SnapshotType::Full => {
@@ -865,6 +877,7 @@ mod tests {
         pub vsock: VsockBuilder,
         balloon_config_called: bool,
         balloon_set: bool,
+        boot_src: BootSourceConfig,
         boot_cfg_set: bool,
         block_set: bool,
         vsock_set: bool,
@@ -927,17 +940,22 @@ mod tests {
 
         pub fn build_boot_source(
             &mut self,
-            _: BootSourceConfig,
+            boot_source: BootSourceConfig,
         ) -> Result<(), BootSourceConfigError> {
             if self.force_errors {
                 return Err(BootSourceConfigError::InvalidKernelPath(
                     std::io::Error::from_raw_os_error(0),
                 ));
             }
+            self.boot_src = boot_source;
             self.boot_cfg_set = true;
             Ok(())
         }
 
+        pub fn boot_source_config(&mut self) -> &BootSourceConfig {
+            &self.boot_src
+        }
+
         pub fn set_block_device(&mut self, _: BlockDeviceConfig) -> Result<(), DriveError> {
             if self.force_errors {
                 return Err(DriveError::RootBlockDeviceAlreadyAdded);
@@ -1140,6 +1158,7 @@ mod tests {
     // instead of our mocks.
     pub fn create_snapshot(
         _: &mut Vmm,
+        _: &VmInfo,
         _: &CreateSnapshotParams,
         _: versionize::VersionMap,
     ) -> std::result::Result<(), CreateSnapshotError> {

src/vmm/src/version_map.rs

@@ -11,6 +11,7 @@ use lazy_static::lazy_static;
 use versionize::{VersionMap, Versionize};
 
 use crate::device_manager::persist::DeviceStates;
+use crate::persist::VmInfo;
 #[cfg(target_arch = "x86_64")]
 use crate::vstate::vcpu::VcpuState;
 
@@ -25,6 +26,8 @@ pub const FC_V0_25_SNAP_VERSION: u16 = 3;
 pub const FC_V1_0_SNAP_VERSION: u16 = 4;
 /// Snap version for Firecracker v1.1
 pub const FC_V1_1_SNAP_VERSION: u16 = 5;
+/// Snap version for Firecracker v1.2
+pub const FC_V1_2_SNAP_VERSION: u16 = 6;
 
 lazy_static! {
     // Note: until we have a better design, this needs to be updated when the version changes.
@@ -48,6 +51,9 @@ lazy_static! {
         // v1.1 state change mappings.
         version_map.new_version().set_type_version(DeviceStates::type_id(), 3);
 
+        // v1.2 state change mappings.
+        version_map.new_version().set_type_version(VmInfo::type_id(), 2);
+
         version_map
     };
 
@@ -62,6 +68,7 @@ lazy_static! {
         mapping.insert(String::from("0.25.0"), FC_V0_25_SNAP_VERSION);
         mapping.insert(String::from("1.0.0"), FC_V1_0_SNAP_VERSION);
         mapping.insert(String::from("1.1.0"), FC_V1_1_SNAP_VERSION);
+        mapping.insert(String::from("1.2.0"), FC_V1_2_SNAP_VERSION);
 
         mapping
     };

src/vmm/src/vmm_config/boot_source.rs

@@ -6,6 +6,8 @@ use std::fs::File;
 use std::io;
 
 use serde::{Deserialize, Serialize};
+use versionize::{VersionMap, Versionize, VersionizeResult};
+use versionize_derive::Versionize;
 
 /// Default guest kernel command line:
 /// - `reboot=k` shut down the guest on reboot, instead of well... rebooting;
@@ -22,7 +24,7 @@ pub const DEFAULT_KERNEL_CMDLINE: &str = "reboot=k panic=1 pci=off nomodules 825
 
 /// Strongly typed data structure used to configure the boot source of the
 /// microvm.
-#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
+#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize, Versionize)]
 #[serde(deny_unknown_fields)]
 pub struct BootSourceConfig {
     /// Path of the kernel image.

src/vmm/src/vmm_config/machine_config.rs

@@ -1,9 +1,10 @@
 // Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-
 use std::fmt;
 
 use serde::{de, Deserialize, Serialize};
+use versionize::{VersionMap, Versionize, VersionizeError, VersionizeResult};
+use versionize_derive::Versionize;
 
 /// The default memory size of the VM, in MiB.
 pub const DEFAULT_MEM_SIZE_MIB: usize = 128;
@@ -237,7 +238,7 @@ where
 
 /// Template types available for configuring the CPU features that map
 /// to EC2 instances.
-#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Serialize)]
+#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Serialize, Versionize)]
 pub enum CpuFeaturesTemplate {
     /// C3 Template.
     C3,