refactor: move GuestMemoryState into VmState

Since the guest memory is now stored in `struct Vm`, we can more easily
store the guest memory state inside the VmState.

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

Author: roypat

src/vmm/src/arch/aarch64/vm.rs

@@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize};
 use crate::Kvm;
 use crate::arch::aarch64::gic::GicState;
 use crate::vstate::vm::{VmCommon, VmError};
+use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryState};
 
 /// Structure representing the current architecture's understand of what a "virtual machine" is.
 #[derive(Debug)]
@@ -68,6 +69,7 @@ impl ArchVm {
     /// Saves and returns the Kvm Vm state.
     pub fn save_state(&self, mpidrs: &[u64]) -> Result<VmState, ArchVmError> {
         Ok(VmState {
+            memory: self.common.guest_memory.describe(),
             gic: self
                 .get_irqchip()
                 .save_device(mpidrs)
@@ -84,13 +86,16 @@ impl ArchVm {
         self.get_irqchip()
             .restore_device(mpidrs, &state.gic)
             .map_err(ArchVmError::RestoreGic)?;
+
         Ok(())
     }
 }
 
 /// Structure holding an general specific VM state.
 #[derive(Debug, Default, Serialize, Deserialize)]
 pub struct VmState {
+    /// Guest memory state
+    pub memory: GuestMemoryState,
     /// GIC state.
     pub gic: GicState,
 }

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

@@ -12,6 +12,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::arch::x86_64::msr::MsrError;
 use crate::utils::u64_to_usize;
+use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryState};
 use crate::vstate::vm::{VmCommon, VmError};
 
 /// Error type for [`Vm::restore_state`]
@@ -185,6 +186,7 @@ impl ArchVm {
             .map_err(ArchVmError::VmGetIrqChip)?;
 
         Ok(VmState {
+            memory: self.common.guest_memory.describe(),
             pitstate,
             clock,
             pic_master,
@@ -207,6 +209,8 @@ impl ArchVm {
 #[derive(Default, Deserialize, Serialize)]
 /// Structure holding VM kvm state.
 pub struct VmState {
+    /// guest memory state
+    pub memory: GuestMemoryState,
     pitstate: kvm_pit_state2,
     clock: kvm_clock_data,
     // TODO: rename this field to adopt inclusive language once Linux updates it, too.

src/vmm/src/lib.rs

@@ -148,9 +148,7 @@ use crate::rate_limiter::BucketUpdate;
 use crate::snapshot::Persist;
 use crate::utils::u64_to_usize;
 use crate::vmm_config::instance_info::{InstanceInfo, VmState};
-use crate::vstate::memory::{
-    GuestMemory, GuestMemoryExtension, GuestMemoryMmap, GuestMemoryRegion,
-};
+use crate::vstate::memory::{GuestMemory, GuestMemoryMmap, GuestMemoryRegion};
 use crate::vstate::vcpu::VcpuState;
 pub use crate::vstate::vcpu::{Vcpu, VcpuConfig, VcpuEvent, VcpuHandle, VcpuResponse};
 pub use crate::vstate::vm::Vm;
@@ -529,12 +527,10 @@ impl Vmm {
         };
         let device_states = self.mmio_device_manager.save();
 
-        let memory_state = self.guest_memory().describe();
         let acpi_dev_state = self.acpi_device_manager.save();
 
         Ok(MicrovmState {
             vm_info: vm_info.clone(),
-            memory_state,
             kvm_state,
             vm_state,
             vcpu_states,

src/vmm/src/persist.rs

@@ -75,8 +75,6 @@ impl From<&VmResources> for VmInfo {
 pub struct MicrovmState {
     /// Miscellaneous VM info.
     pub vm_info: VmInfo,
-    /// Memory state.
-    pub memory_state: GuestMemoryState,
     /// KVM KVM state.
     pub kvm_state: KvmState,
     /// VM KVM state.
@@ -367,7 +365,7 @@ pub fn snapshot_state_sanity_check(
     // Check if the snapshot contains at least 1 mem region.
     // Upper bound check will be done when creating guest memory by comparing against
     // KVM max supported value kvm_context.max_memslots().
-    if microvm_state.memory_state.regions.is_empty() {
+    if microvm_state.vm_state.memory.regions.is_empty() {
         return Err(SnapShotStateSanityCheckError::NoMemory);
     }
 
@@ -450,7 +448,7 @@ pub fn restore_from_snapshot(
     snapshot_state_sanity_check(&microvm_state)?;
 
     let mem_backend_path = &params.mem_backend.backend_path;
-    let mem_state = &microvm_state.memory_state;
+    let mem_state = &microvm_state.vm_state.memory;
 
     let (guest_memory, uffd) = match params.mem_backend.backend_type {
         MemBackendType::File => {
@@ -745,13 +743,11 @@ mod tests {
         assert!(states.vsock_device.is_some());
         assert!(states.balloon_device.is_some());
 
-        let memory_state = vmm.guest_memory().describe();
         let vcpu_states = vec![VcpuState::default()];
         #[cfg(target_arch = "aarch64")]
         let mpidrs = construct_kvm_mpidrs(&vcpu_states);
         let microvm_state = MicrovmState {
             device_states: states,
-            memory_state,
             vcpu_states,
             kvm_state: Default::default(),
             vm_info: VmInfo {

src/vmm/src/vstate/vm.rs

@@ -23,7 +23,8 @@ pub struct VmCommon {
     /// The KVM file descriptor used to access this Vm.
     pub fd: VmFd,
     max_memslots: usize,
-    guest_memory: GuestMemoryMmap,
+    /// The guest memory of this Vm.
+    pub guest_memory: GuestMemoryMmap,
 }
 
 /// Errors associated with the wrappers over KVM ioctls.

src/vmm/tests/integration_tests.rs

@@ -299,7 +299,7 @@ fn test_snapshot_load_sanity_checks() {
     snapshot_state_sanity_check(&microvm_state).unwrap();
 
     // Remove memory regions.
-    microvm_state.memory_state.regions.clear();
+    microvm_state.vm_state.memory.regions.clear();
 
     // Validate sanity checks fail because there is no mem region in state.
     assert_eq!(