refactor(vm-memory): use GuestMemoryCollection to store KVM slot

This refactor enables us to add more information to the GuestMemory.
To start off, I am adding KVM slot number and type of memory. So far we
only support Dram (normal memory), but we will also add Hotpluggable
region type in the future.

Signed-off-by: Riccardo Mancini <[email protected]>

Author: Manciukic

src/vmm/src/builder.rs

@@ -171,7 +171,7 @@ pub fn build_microvm_for_boot(
     // Build custom CPU config if a custom template is provided.
     let mut vm = Vm::new(&kvm)?;
     let (mut vcpus, vcpus_exit_evt) = vm.create_vcpus(vm_resources.machine_config.vcpu_count)?;
-    vm.register_memory_regions(guest_memory)?;
+    vm.register_dram_memory_regions(guest_memory)?;
 
     let mut device_manager = DeviceManager::new(
         event_manager,
@@ -421,7 +421,7 @@ pub fn build_microvm_from_snapshot(
         .create_vcpus(vm_resources.machine_config.vcpu_count)
         .map_err(StartMicrovmError::Vm)?;
 
-    vm.register_memory_regions(guest_memory)
+    vm.restore_memory_regions(guest_memory, &microvm_state.vm_state.memory)
         .map_err(StartMicrovmError::Vm)?;
 
     #[cfg(target_arch = "x86_64")]

src/vmm/src/device_manager/mmio.rs

@@ -587,7 +587,7 @@ pub(crate) mod tests {
         let guest_mem = multi_region_mem_raw(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
         let kvm = Kvm::new(vec![]).expect("Cannot create Kvm");
         let mut vm = Vm::new(&kvm).unwrap();
-        vm.register_memory_regions(guest_mem).unwrap();
+        vm.register_dram_memory_regions(guest_mem).unwrap();
         let mut device_manager = MMIODeviceManager::new();
 
         let mut cmdline = kernel_cmdline::Cmdline::new(4096).unwrap();
@@ -633,7 +633,7 @@ pub(crate) mod tests {
         let guest_mem = multi_region_mem_raw(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
         let kvm = Kvm::new(vec![]).expect("Cannot create Kvm");
         let mut vm = Vm::new(&kvm).unwrap();
-        vm.register_memory_regions(guest_mem).unwrap();
+        vm.register_dram_memory_regions(guest_mem).unwrap();
         let mut device_manager = MMIODeviceManager::new();
 
         let mut cmdline = kernel_cmdline::Cmdline::new(4096).unwrap();
@@ -686,7 +686,7 @@ pub(crate) mod tests {
         let guest_mem = multi_region_mem_raw(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
         let kvm = Kvm::new(vec![]).expect("Cannot create Kvm");
         let mut vm = Vm::new(&kvm).unwrap();
-        vm.register_memory_regions(guest_mem).unwrap();
+        vm.register_dram_memory_regions(guest_mem).unwrap();
 
         #[cfg(target_arch = "x86_64")]
         vm.setup_irqchip().unwrap();

src/vmm/src/devices/virtio/block/virtio/io/mod.rs

@@ -179,14 +179,15 @@ pub mod tests {
     #![allow(clippy::undocumented_unsafe_blocks)]
     use std::os::unix::ffi::OsStrExt;
 
+    use vm_memory::GuestMemoryRegion;
     use vmm_sys_util::tempfile::TempFile;
 
     use super::*;
     use crate::devices::virtio::block::virtio::device::FileEngineType;
     use crate::utils::u64_to_usize;
     use crate::vmm_config::machine_config::HugePageConfig;
     use crate::vstate::memory;
-    use crate::vstate::memory::{Bitmap, Bytes, GuestMemory};
+    use crate::vstate::memory::{Bitmap, Bytes, GuestMemory, GuestRegionMmapExt};
 
     const FILE_LEN: u32 = 1024;
     // 2 pages of memory should be enough to test read/write ops and also dirty tracking.
@@ -227,20 +228,23 @@ pub mod tests {
                 true,
                 HugePageConfig::None,
             )
-            .unwrap(),
+            .unwrap()
+            .into_iter()
+            .map(|region| GuestRegionMmapExt::dram_from_mmap_region(region, 0))
+            .collect(),
         )
         .unwrap()
     }
 
     fn check_dirty_mem(mem: &GuestMemoryMmap, addr: GuestAddress, len: u32) {
-        let bitmap = mem.find_region(addr).unwrap().bitmap().as_ref().unwrap();
+        let bitmap = mem.find_region(addr).unwrap().bitmap();
         for offset in addr.0..addr.0 + u64::from(len) {
             assert!(bitmap.dirty_at(u64_to_usize(offset)));
         }
     }
 
     fn check_clean_mem(mem: &GuestMemoryMmap, addr: GuestAddress, len: u32) {
-        let bitmap = mem.find_region(addr).unwrap().bitmap().as_ref().unwrap();
+        let bitmap = mem.find_region(addr).unwrap().bitmap();
         for offset in addr.0..addr.0 + u64::from(len) {
             assert!(!bitmap.dirty_at(u64_to_usize(offset)));
         }

src/vmm/src/devices/virtio/vhost_user.rs

@@ -378,7 +378,7 @@ impl<T: VhostUserHandleBackend> VhostUserHandleImpl<T> {
             let vhost_user_net_reg = VhostUserMemoryRegionInfo {
                 guest_phys_addr: region.start_addr().raw_value(),
                 memory_size: region.len(),
-                userspace_addr: region.as_ptr() as u64,
+                userspace_addr: region.inner.as_ptr() as u64,
                 mmap_offset,
                 mmap_handle,
             };
@@ -478,7 +478,7 @@ pub(crate) mod tests {
     use crate::devices::virtio::test_utils::default_interrupt;
     use crate::test_utils::create_tmp_socket;
     use crate::vstate::memory;
-    use crate::vstate::memory::GuestAddress;
+    use crate::vstate::memory::{GuestAddress, GuestRegionMmapExt};
 
     pub(crate) fn create_mem(file: File, regions: &[(GuestAddress, usize)]) -> GuestMemoryMmap {
         GuestMemoryMmap::from_regions(
@@ -488,7 +488,10 @@ pub(crate) mod tests {
                 Some(file),
                 false,
             )
-            .unwrap(),
+            .unwrap()
+            .into_iter()
+            .map(|region| GuestRegionMmapExt::dram_from_mmap_region(region, 0))
+            .collect(),
         )
         .unwrap()
     }
@@ -798,7 +801,7 @@ pub(crate) mod tests {
             .map(|region| VhostUserMemoryRegionInfo {
                 guest_phys_addr: region.start_addr().raw_value(),
                 memory_size: region.len(),
-                userspace_addr: region.as_ptr() as u64,
+                userspace_addr: region.inner.as_ptr() as u64,
                 mmap_offset: region.file_offset().unwrap().start(),
                 mmap_handle: region.file_offset().unwrap().file().as_raw_fd(),
             })

src/vmm/src/persist.rs

@@ -588,7 +588,7 @@ mod tests {
     use crate::vmm_config::balloon::BalloonDeviceConfig;
     use crate::vmm_config::net::NetworkInterfaceConfig;
     use crate::vmm_config::vsock::tests::default_config;
-    use crate::vstate::memory::GuestMemoryRegionState;
+    use crate::vstate::memory::{GuestMemoryRegionState, GuestRegionType};
 
     fn default_vmm_with_devices() -> Vmm {
         let mut event_manager = EventManager::new().expect("Cannot create EventManager");
@@ -693,6 +693,7 @@ mod tests {
             regions: vec![GuestMemoryRegionState {
                 base_address: 0,
                 size: 0x20000,
+                region_type: GuestRegionType::Dram,
             }],
         };
 

src/vmm/src/test_utils/mod.rs

@@ -5,7 +5,7 @@
 
 use std::sync::{Arc, Mutex};
 
-use vm_memory::GuestAddress;
+use vm_memory::{GuestAddress, GuestRegionCollection};
 use vmm_sys_util::tempdir::TempDir;
 
 use crate::builder::build_microvm_for_boot;
@@ -16,7 +16,7 @@ use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::HugePageConfig;
 use crate::vstate::memory;
-use crate::vstate::memory::{GuestMemoryMmap, GuestRegionMmap};
+use crate::vstate::memory::{GuestMemoryMmap, GuestRegionMmap, GuestRegionMmapExt};
 use crate::{EventManager, Vmm};
 
 pub mod mock_resources;
@@ -43,9 +43,12 @@ pub fn single_region_mem_at_raw(at: u64, size: usize) -> Vec<GuestRegionMmap> {
 
 /// Creates a [`GuestMemoryMmap`] with multiple regions and without dirty page tracking.
 pub fn multi_region_mem(regions: &[(GuestAddress, usize)]) -> GuestMemoryMmap {
-    GuestMemoryMmap::from_regions(
+    GuestRegionCollection::from_regions(
         memory::anonymous(regions.iter().copied(), false, HugePageConfig::None)
-            .expect("Cannot initialize memory"),
+            .expect("Cannot initialize memory")
+            .into_iter()
+            .map(|region| GuestRegionMmapExt::dram_from_mmap_region(region, 0))
+            .collect(),
     )
     .unwrap()
 }