From 5b177b205b46a442149e0ec85d7d7dd4f5c4763f Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Tue, 5 Oct 2021 14:20:19 +0200
Subject: [PATCH] arch, vmm: Extend the data being snapshot

Storing multiple data coming from the MemoryManager in order to be able
to restore without creating everything from scratch.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 arch/src/lib.rs           |  7 +++++-
 vmm/src/memory_manager.rs | 51 +++++++++++++++++++++++++++++++++++----
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/arch/src/lib.rs b/arch/src/lib.rs
index fb5c80705..a02a583a9 100644
--- a/arch/src/lib.rs
+++ b/arch/src/lib.rs
@@ -17,6 +17,9 @@ use std::collections::BTreeMap;
 use std::fmt;
 use std::result;
 use std::sync::Arc;
+use versionize::{VersionMap, Versionize, VersionizeError, VersionizeResult};
+use versionize_derive::Versionize;
+use vm_migration::VersionMapped;
 
 type GuestMemoryMmap = vm_memory::GuestMemoryMmap<vm_memory::bitmap::AtomicBitmap>;
 type GuestRegionMmap = vm_memory::GuestRegionMmap<vm_memory::bitmap::AtomicBitmap>;
@@ -54,7 +57,7 @@ pub enum Error {
 pub type Result<T> = result::Result<T, Error>;
 
 /// Type for memory region types.
-#[derive(PartialEq, Debug)]
+#[derive(Clone, Copy, PartialEq, Debug, Versionize)]
 pub enum RegionType {
     /// RAM type
     Ram,
@@ -72,6 +75,8 @@ pub enum RegionType {
     Reserved,
 }
 
+impl VersionMapped for RegionType {}
+
 /// Module for aarch64 related functionality.
 #[cfg(target_arch = "aarch64")]
 pub mod aarch64;
diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index 5ddb9970e..dbfc3e725 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -63,7 +63,7 @@ const MPOL_BIND: u32 = 2;
 const MPOL_MF_STRICT: u32 = 1;
 const MPOL_MF_MOVE: u32 = 1 << 1;
 
-#[derive(Default)]
+#[derive(Clone, Default, Versionize)]
 struct HotPlugState {
     base: u64,
     length: u64,
@@ -121,6 +121,7 @@ impl MemoryZone {
 
 pub type MemoryZones = HashMap<String, MemoryZone>;
 
+#[derive(Clone, Versionize)]
 struct GuestRamMapping {
     slot: u32,
     gpa: u64,
@@ -130,6 +131,13 @@ struct GuestRamMapping {
     file_offset: u64,
 }
 
+#[derive(Clone, Versionize)]
+struct ArchMemRegion {
+    base: u64,
+    size: usize,
+    r_type: RegionType,
+}
+
 pub struct MemoryManager {
     boot_guest_memory: GuestMemoryMmap,
     guest_memory: GuestMemoryAtomic<GuestMemoryMmap>,
@@ -155,7 +163,7 @@ pub struct MemoryManager {
     snapshot_memory_ranges: MemoryRangeTable,
     memory_zones: MemoryZones,
     log_dirty: bool, // Enable dirty logging for created RAM regions
-    arch_mem_regions: Vec<(GuestAddress, usize, RegionType)>,
+    arch_mem_regions: Vec<ArchMemRegion>,
 
     // Keep track of calls to create_userspace_mapping() for guest RAM.
     // This is useful for getting the dirty pages as we need to know the
@@ -784,7 +792,7 @@ impl MemoryManager {
 
         // Allocate SubRegion and Reserved address ranges.
         for region in self.arch_mem_regions.iter() {
-            if region.2 == RegionType::Ram {
+            if region.r_type == RegionType::Ram {
                 // Ignore the RAM type since ranges have already been allocated
                 // based on the GuestMemory regions.
                 continue;
@@ -792,7 +800,11 @@ impl MemoryManager {
             self.allocator
                 .lock()
                 .unwrap()
-                .allocate_mmio_addresses(Some(region.0), region.1 as GuestUsize, None)
+                .allocate_mmio_addresses(
+                    Some(GuestAddress(region.base)),
+                    region.size as GuestUsize,
+                    None,
+                )
                 .ok_or(Error::MemoryRangeAllocation)?;
         }
 
@@ -820,6 +832,15 @@ impl MemoryManager {
             .map(|r| (r.0, r.1))
             .collect();
 
+        let arch_mem_regions: Vec<ArchMemRegion> = arch_mem_regions
+            .iter()
+            .map(|(a, b, c)| ArchMemRegion {
+                base: a.0,
+                size: *b,
+                r_type: *c,
+            })
+            .collect();
+
         let (mem_regions, mut memory_zones) =
             Self::create_memory_regions_from_zones(&ram_regions, &zones, prefault)?;
 
@@ -2049,6 +2070,15 @@ impl Pausable for MemoryManager {}
 #[derive(Versionize)]
 pub struct MemoryManagerSnapshotData {
     memory_ranges: MemoryRangeTable,
+    guest_ram_mappings: Vec<GuestRamMapping>,
+    start_of_device_area: u64,
+    boot_ram: u64,
+    current_ram: u64,
+    arch_mem_regions: Vec<ArchMemRegion>,
+    hotplug_slots: Vec<HotPlugState>,
+    next_memory_slot: u32,
+    selected_slot: usize,
+    next_hotplug_slot: usize,
 }
 
 impl VersionMapped for MemoryManagerSnapshotData {}
@@ -2075,7 +2105,18 @@ impl Snapshottable for MemoryManager {
 
         memory_manager_snapshot.add_data_section(SnapshotDataSection::new_from_versioned_state(
             MEMORY_MANAGER_SNAPSHOT_ID,
-            &MemoryManagerSnapshotData { memory_ranges },
+            &MemoryManagerSnapshotData {
+                memory_ranges,
+                guest_ram_mappings: self.guest_ram_mappings.clone(),
+                start_of_device_area: self.start_of_device_area.0,
+                boot_ram: self.boot_ram,
+                current_ram: self.current_ram,
+                arch_mem_regions: self.arch_mem_regions.clone(),
+                hotplug_slots: self.hotplug_slots.clone(),
+                next_memory_slot: self.next_memory_slot,
+                selected_slot: self.selected_slot,
+                next_hotplug_slot: self.next_hotplug_slot,
+            },
         )?);
 
         Ok(memory_manager_snapshot)