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

@@ -168,7 +168,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,
@@ -411,7 +411,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

@@ -586,7 +586,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();
@@ -632,7 +632,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();
@@ -685,7 +685,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,7 +228,10 @@ pub mod tests {
                 true,
                 HugePageConfig::None,
             )
-            .unwrap(),
+            .unwrap()
+            .into_iter()
+            .map(|region| GuestRegionMmapExt::dram_from_mmap_region(region, 0))
+            .collect(),
         )
         .unwrap()
     }

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

@@ -380,7 +380,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,
             };
@@ -480,7 +480,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(
@@ -490,7 +490,10 @@ pub(crate) mod tests {
                 Some(file),
                 false,
             )
-            .unwrap(),
+            .unwrap()
+            .into_iter()
+            .map(|region| GuestRegionMmapExt::dram_from_mmap_region(region, 0))
+            .collect(),
         )
         .unwrap()
     }
@@ -800,7 +803,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

@@ -12,11 +12,12 @@ use crate::builder::build_microvm_for_boot;
 use crate::resources::VmResources;
 use crate::seccomp::get_empty_filters;
 use crate::test_utils::mock_resources::{MockBootSourceConfig, MockVmConfig, MockVmResources};
+use crate::vm_memory_vendored::GuestRegionCollection;
 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 +44,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()
 }