refactor: Hoist update_vm_config call to begin of restore path

Ensure that the vm_config field of the `VmResources` struct is
initialized from snapshot data as early as possible. This means that
further steps in the snapshot restore procedure can rely on the
vm_resources object for vm configuration information, instead of needing
to access the microvm state (or passing various parameters down to ever
snapshot restore function aslong with the uninitialized `VmResources`
struct, as was done for dirty page tracking). This change also ensures
that we fail early if the snapshot is malformed in some way, as opposed
to failing after various KVM resources have already (potentially
incorrectly) been initialized.

Signed-off-by: Patrick Roy <roypat@amazon.co.uk>

Author: roypat

src/vmm/src/builder.rs

@@ -56,7 +56,7 @@ use crate::snapshot::Persist;
 use crate::vmm_config::boot_source::BootConfig;
 use crate::vmm_config::drive::BlockDeviceType;
 use crate::vmm_config::instance_info::InstanceInfo;
-use crate::vmm_config::machine_config::{MachineConfigUpdate, VmConfig, VmConfigError};
+use crate::vmm_config::machine_config::{VmConfig, VmConfigError};
 use crate::vstate::memory::{GuestAddress, GuestMemory, GuestMemoryExtension, GuestMemoryMmap};
 use crate::vstate::vcpu::{Vcpu, VcpuConfig, VcpuError};
 use crate::vstate::vm::Vm;
@@ -363,8 +363,6 @@ pub fn build_and_boot_microvm(
 pub enum BuildMicrovmFromSnapshotError {
     /// Failed to create microVM and vCPUs: {0}
     CreateMicrovmAndVcpus(#[from] StartMicrovmError),
-    /// Only 255 vCPU state are supported, but {0} states where given.
-    TooManyVCPUs(usize),
     /// Could not access KVM: {0}
     KvmAccess(#[from] utils::errno::Error),
     /// Error configuring the TSC, frequency not present in the given snapshot.
@@ -406,23 +404,18 @@ pub fn build_microvm_from_snapshot(
     microvm_state: MicrovmState,
     guest_memory: GuestMemoryMmap,
     uffd: Option<Uffd>,
-    track_dirty_pages: bool,
     seccomp_filters: &BpfThreadMap,
     vm_resources: &mut VmResources,
 ) -> Result<Arc<Mutex<Vmm>>, BuildMicrovmFromSnapshotError> {
-    let vcpu_count = u8::try_from(microvm_state.vcpu_states.len()).map_err(|_| {
-        BuildMicrovmFromSnapshotError::TooManyVCPUs(microvm_state.vcpu_states.len())
-    })?;
-
     // Build Vmm.
     debug!("event_start: build microvm from snapshot");
     let (mut vmm, mut vcpus) = create_vmm_and_vcpus(
         instance_info,
         event_manager,
         guest_memory.clone(),
         uffd,
-        track_dirty_pages,
-        vcpu_count,
+        vm_resources.vm_config.track_dirty_pages,
+        vm_resources.vm_config.vcpu_count,
         microvm_state.vm_state.kvm_cap_modifiers.clone(),
     )?;
 
@@ -466,14 +459,6 @@ pub fn build_microvm_from_snapshot(
     #[cfg(target_arch = "x86_64")]
     vmm.vm.restore_state(&microvm_state.vm_state)?;
 
-    vm_resources.update_vm_config(&MachineConfigUpdate {
-        vcpu_count: Some(vcpu_count),
-        mem_size_mib: Some(u64_to_usize(microvm_state.vm_info.mem_size_mib)),
-        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);
 

src/vmm/src/persist.rs

@@ -32,7 +32,7 @@ use crate::resources::VmResources;
 use crate::snapshot::Snapshot;
 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::{MachineConfigUpdate, VmConfigError, MAX_SUPPORTED_VCPUS};
 use crate::vmm_config::snapshot::{
     CreateSnapshotParams, LoadSnapshotParams, MemBackendType, SnapshotType,
 };
@@ -392,13 +392,30 @@ pub fn restore_from_snapshot(
     vm_resources: &mut VmResources,
 ) -> Result<Arc<Mutex<Vmm>>, RestoreFromSnapshotError> {
     let microvm_state = snapshot_state_from_file(&params.snapshot_path)?;
+    let track_dirty_pages = params.enable_diff_snapshots;
+
+    let vcpu_count = microvm_state
+        .vcpu_states
+        .len()
+        .try_into()
+        .map_err(|_| VmConfigError::InvalidVcpuCount)
+        .map_err(BuildMicrovmFromSnapshotError::VmUpdateConfig)?;
+
+    vm_resources
+        .update_vm_config(&MachineConfigUpdate {
+            vcpu_count: Some(vcpu_count),
+            mem_size_mib: Some(u64_to_usize(microvm_state.vm_info.mem_size_mib)),
+            smt: Some(microvm_state.vm_info.smt),
+            cpu_template: Some(microvm_state.vm_info.cpu_template),
+            track_dirty_pages: Some(track_dirty_pages),
+        })
+        .map_err(BuildMicrovmFromSnapshotError::VmUpdateConfig)?;
 
     // Some sanity checks before building the microvm.
     snapshot_state_sanity_check(&microvm_state)?;
 
     let mem_backend_path = &params.mem_backend.backend_path;
     let mem_state = &microvm_state.memory_state;
-    let track_dirty_pages = params.enable_diff_snapshots;
 
     let (guest_memory, uffd) = match params.mem_backend.backend_type {
         MemBackendType::File => (
@@ -422,7 +439,6 @@ pub fn restore_from_snapshot(
         microvm_state,
         guest_memory,
         uffd,
-        track_dirty_pages,
         seccomp_filters,
         vm_resources,
     )

src/vmm/tests/integration_tests.rs

@@ -254,7 +254,6 @@ fn verify_load_snapshot(snapshot_file: TempFile, memory_file: TempFile) {
         microvm_state,
         mem,
         None,
-        false,
         &empty_seccomp_filters,
         vm_resources,
     )