refactor: have memory.rs hand out Vec<GuestRegionMmap>

This way, we can accumulate different memory regions into one
`GuestMemoryMmap` over this inside `struct Vm`. Sadly, vm-memory does
not offer APIs for merging `GuestMemoryMmap`s, so there can only be one
of these.

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

Author: roypat

src/vmm/benches/memory_access.rs

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
-use vm_memory::GuestMemory;
 use vmm::resources::VmResources;
 use vmm::vmm_config::machine_config::{HugePageConfig, MachineConfig};
 
@@ -14,7 +13,7 @@ fn bench_single_page_fault(c: &mut Criterion, configuration: VmResources) {
                 // Get a pointer to the first memory region (cannot do `.get_slice(GuestAddress(0),
                 // 1)`, because on ARM64 guest memory does not start at physical
                 // address 0).
-                let ptr = memory.iter().next().unwrap().as_ptr();
+                let ptr = memory.first().unwrap().as_ptr();
 
                 // fine to return both here, because ptr is not a reference into `memory` (e.g. no
                 // self-referential structs are happening here)

src/vmm/src/acpi/mod.rs

@@ -188,7 +188,7 @@ mod tests {
     use crate::acpi::{AcpiError, AcpiTableWriter};
     use crate::arch::x86_64::layout::{SYSTEM_MEM_SIZE, SYSTEM_MEM_START};
     use crate::builder::tests::default_vmm;
-    use crate::test_utils::arch_mem;
+    use crate::test_utils::arch_mem_raw;
 
     struct MockSdt(Vec<u8>);
 
@@ -265,7 +265,7 @@ mod tests {
     fn test_write_acpi_table_small_memory() {
         let mut vmm = default_vmm();
         vmm.vm
-            .memory_init(arch_mem(
+            .register_memory_regions(arch_mem_raw(
                 (SYSTEM_MEM_START + SYSTEM_MEM_SIZE - 4096)
                     .try_into()
                     .unwrap(),

src/vmm/src/builder.rs

@@ -72,7 +72,7 @@ use crate::vmm_config::boot_source::BootConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::{MachineConfig, MachineConfigError};
 use crate::vstate::kvm::Kvm;
-use crate::vstate::memory::{GuestAddress, GuestMemory, GuestMemoryMmap};
+use crate::vstate::memory::{GuestAddress, GuestMemory, GuestMemoryMmap, GuestRegionMmap};
 use crate::vstate::vcpu::{Vcpu, VcpuConfig, VcpuError};
 use crate::vstate::vm::Vm;
 use crate::{EventManager, Vmm, VmmError, device_manager};
@@ -154,7 +154,6 @@ impl std::convert::From<linux_loader::cmdline::Error> for StartMicrovmError {
 fn create_vmm_and_vcpus(
     instance_info: &InstanceInfo,
     event_manager: &mut EventManager,
-    guest_memory: GuestMemoryMmap,
     uffd: Option<Uffd>,
     vcpu_count: u8,
     kvm_capabilities: Vec<KvmCapability>,
@@ -163,7 +162,6 @@ fn create_vmm_and_vcpus(
     // Set up Kvm Vm and register memory regions.
     // Build custom CPU config if a custom template is provided.
     let mut vm = Vm::new(&kvm)?;
-    vm.memory_init(guest_memory)?;
 
     let resource_allocator = ResourceAllocator::new()?;
 
@@ -250,12 +248,18 @@ pub fn build_microvm_for_boot(
     let (mut vmm, mut vcpus) = create_vmm_and_vcpus(
         instance_info,
         event_manager,
-        guest_memory,
         None,
         vm_resources.machine_config.vcpu_count,
         cpu_template.kvm_capabilities.clone(),
     )?;
 
+    for region in guest_memory {
+        vmm.vm
+            .register_memory_region(region)
+            .map_err(VmmError::Vm)
+            .map_err(StartMicrovmError::Internal)?;
+    }
+
     let entry_point = load_kernel(boot_config, vmm.guest_memory())?;
     let initrd = load_initrd_from_config(boot_config, vmm.guest_memory())?;
 
@@ -433,7 +437,7 @@ pub fn build_microvm_from_snapshot(
     instance_info: &InstanceInfo,
     event_manager: &mut EventManager,
     microvm_state: MicrovmState,
-    guest_memory: GuestMemoryMmap,
+    guest_memory: Vec<GuestRegionMmap>,
     uffd: Option<Uffd>,
     seccomp_filters: &BpfThreadMap,
     vm_resources: &mut VmResources,
@@ -443,13 +447,17 @@ pub fn build_microvm_from_snapshot(
     let (mut vmm, mut vcpus) = create_vmm_and_vcpus(
         instance_info,
         event_manager,
-        guest_memory,
         uffd,
         vm_resources.machine_config.vcpu_count,
         microvm_state.kvm_state.kvm_cap_modifiers.clone(),
     )
     .map_err(Internal)?;
 
+    vmm.vm
+        .register_memory_regions(guest_memory)
+        .map_err(VmmError::Vm)
+        .map_err(StartMicrovmError::Internal)?;
+
     #[cfg(target_arch = "x86_64")]
     {
         // Scale TSC to match, extract the TSC freq from the state if specified

src/vmm/src/device_manager/mmio.rs

@@ -532,7 +532,7 @@ mod tests {
     use crate::devices::virtio::ActivateError;
     use crate::devices::virtio::device::{IrqTrigger, VirtioDevice};
     use crate::devices::virtio::queue::Queue;
-    use crate::test_utils::multi_region_mem;
+    use crate::test_utils::multi_region_mem_raw;
     use crate::vstate::kvm::Kvm;
     use crate::vstate::memory::{GuestAddress, GuestMemoryMmap};
 
@@ -649,10 +649,10 @@ mod tests {
     fn test_register_virtio_device() {
         let start_addr1 = GuestAddress(0x0);
         let start_addr2 = GuestAddress(0x1000);
-        let guest_mem = multi_region_mem(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
+        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.memory_init(guest_mem).unwrap();
+        vm.register_memory_regions(guest_mem).unwrap();
         let mut device_manager = MMIODeviceManager::new();
         let mut resource_allocator = ResourceAllocator::new().unwrap();
 
@@ -680,10 +680,10 @@ mod tests {
     fn test_register_too_many_devices() {
         let start_addr1 = GuestAddress(0x0);
         let start_addr2 = GuestAddress(0x1000);
-        let guest_mem = multi_region_mem(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
+        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.memory_init(guest_mem).unwrap();
+        vm.register_memory_regions(guest_mem).unwrap();
         let mut device_manager = MMIODeviceManager::new();
         let mut resource_allocator = ResourceAllocator::new().unwrap();
 
@@ -736,10 +736,10 @@ mod tests {
     fn test_device_info() {
         let start_addr1 = GuestAddress(0x0);
         let start_addr2 = GuestAddress(0x1000);
-        let guest_mem = multi_region_mem(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
+        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.memory_init(guest_mem).unwrap();
+        vm.register_memory_regions(guest_mem).unwrap();
 
         #[cfg(target_arch = "x86_64")]
         vm.setup_irqchip().unwrap();

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

@@ -378,7 +378,8 @@ mod tests {
     use crate::devices::virtio::block::virtio::device::FileEngineType;
     use crate::devices::virtio::mmio::VIRTIO_MMIO_INT_CONFIG;
     use crate::test_utils::create_tmp_socket;
-    use crate::vstate::memory::{GuestAddress, GuestMemoryExtension};
+    use crate::vstate::memory;
+    use crate::vstate::memory::GuestAddress;
 
     #[test]
     fn test_from_config() {
@@ -779,9 +780,10 @@ mod tests {
         let file = TempFile::new().unwrap().into_file();
         file.set_len(region_size as u64).unwrap();
         let regions = vec![(GuestAddress(0x0), region_size)];
-        let guest_memory =
-            GuestMemoryMmap::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false)
-                .unwrap();
+        let guest_memory = GuestMemoryMmap::from_regions(
+            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false).unwrap(),
+        )
+        .unwrap();
 
         // During actiavion of the device features, memory and queues should be set and activated.
         vhost_block.activate(guest_memory).unwrap();

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

@@ -194,7 +194,8 @@ pub mod tests {
     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::{Bitmap, Bytes, GuestMemory, GuestMemoryExtension};
+    use crate::vstate::memory;
+    use crate::vstate::memory::{Bitmap, Bytes, GuestMemory};
 
     const FILE_LEN: u32 = 1024;
     // 2 pages of memory should be enough to test read/write ops and also dirty tracking.
@@ -230,10 +231,13 @@ pub mod tests {
     }
 
     fn create_mem() -> GuestMemoryMmap {
-        GuestMemoryMmap::anonymous(
-            [(GuestAddress(0), MEM_LEN)].into_iter(),
-            true,
-            HugePageConfig::None,
+        GuestMemoryMmap::from_regions(
+            memory::anonymous(
+                [(GuestAddress(0), MEM_LEN)].into_iter(),
+                true,
+                HugePageConfig::None,
+            )
+            .unwrap(),
         )
         .unwrap()
     }

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

@@ -466,7 +466,8 @@ mod tests {
 
     use super::*;
     use crate::test_utils::create_tmp_socket;
-    use crate::vstate::memory::{GuestAddress, GuestMemoryExtension};
+    use crate::vstate::memory;
+    use crate::vstate::memory::GuestAddress;
 
     #[test]
     fn test_new() {
@@ -763,9 +764,10 @@ mod tests {
             (GuestAddress(0x10000), region_size),
         ];
 
-        let guest_memory =
-            GuestMemoryMmap::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false)
-                .unwrap();
+        let guest_memory = GuestMemoryMmap::from_regions(
+            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false).unwrap(),
+        )
+        .unwrap();
 
         vuh.update_mem_table(&guest_memory).unwrap();
 
@@ -879,9 +881,10 @@ mod tests {
         file.set_len(region_size as u64).unwrap();
         let regions = vec![(GuestAddress(0x0), region_size)];
 
-        let guest_memory =
-            GuestMemoryMmap::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false)
-                .unwrap();
+        let guest_memory = GuestMemoryMmap::from_regions(
+            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false).unwrap(),
+        )
+        .unwrap();
 
         let mut queue = Queue::new(69);
         queue.initialize(&guest_memory).unwrap();

src/vmm/src/persist.rs

@@ -37,9 +37,8 @@ use crate::vmm_config::snapshot::{
     CreateSnapshotParams, LoadSnapshotParams, MemBackendType, SnapshotType,
 };
 use crate::vstate::kvm::KvmState;
-use crate::vstate::memory::{
-    GuestMemory, GuestMemoryExtension, GuestMemoryMmap, GuestMemoryState, MemoryError,
-};
+use crate::vstate::memory;
+use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryState, GuestRegionMmap, MemoryError};
 use crate::vstate::vcpu::{VcpuSendEventError, VcpuState};
 use crate::vstate::vm::VmState;
 use crate::{EventManager, Vmm, VmmError, mem_size_mib, vstate};
@@ -532,10 +531,9 @@ fn guest_memory_from_file(
     mem_file_path: &Path,
     mem_state: &GuestMemoryState,
     track_dirty_pages: bool,
-) -> Result<GuestMemoryMmap, GuestMemoryFromFileError> {
+) -> Result<Vec<GuestRegionMmap>, GuestMemoryFromFileError> {
     let mem_file = File::open(mem_file_path)?;
-    let guest_mem =
-        GuestMemoryMmap::snapshot_file(mem_file, mem_state.regions(), track_dirty_pages)?;
+    let guest_mem = memory::snapshot_file(mem_file, mem_state.regions(), track_dirty_pages)?;
     Ok(guest_mem)
 }
 
@@ -560,7 +558,7 @@ fn guest_memory_from_uffd(
     track_dirty_pages: bool,
     enable_balloon: bool,
     huge_pages: HugePageConfig,
-) -> Result<(GuestMemoryMmap, Option<Uffd>), GuestMemoryFromUffdError> {
+) -> Result<(Vec<GuestRegionMmap>, Option<Uffd>), GuestMemoryFromUffdError> {
     let (guest_memory, backend_mappings) =
         create_guest_memory(mem_state, track_dirty_pages, huge_pages)?;
 
@@ -593,10 +591,9 @@ fn create_guest_memory(
     mem_state: &GuestMemoryState,
     track_dirty_pages: bool,
     huge_pages: HugePageConfig,
-) -> Result<(GuestMemoryMmap, Vec<GuestRegionUffdMapping>), GuestMemoryFromUffdError> {
-    let guest_memory =
-        GuestMemoryMmap::anonymous(mem_state.regions(), track_dirty_pages, huge_pages)?;
-    let mut backend_mappings = Vec::with_capacity(guest_memory.num_regions());
+) -> Result<(Vec<GuestRegionMmap>, Vec<GuestRegionUffdMapping>), GuestMemoryFromUffdError> {
+    let guest_memory = memory::anonymous(mem_state.regions(), track_dirty_pages, huge_pages)?;
+    let mut backend_mappings = Vec::with_capacity(guest_memory.len());
     let mut offset = 0;
     for mem_region in guest_memory.iter() {
         #[allow(deprecated)]

src/vmm/src/resources.rs

@@ -29,7 +29,8 @@ use crate::vmm_config::metrics::{MetricsConfig, MetricsConfigError, init_metrics
 use crate::vmm_config::mmds::{MmdsConfig, MmdsConfigError};
 use crate::vmm_config::net::*;
 use crate::vmm_config::vsock::*;
-use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryMmap, MemoryError};
+use crate::vstate::memory;
+use crate::vstate::memory::{GuestRegionMmap, MemoryError};
 
 /// Errors encountered when configuring microVM resources.
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
@@ -440,7 +441,7 @@ impl VmResources {
     ///
     /// If vhost-user-blk devices are in use, allocates memfd-backed shared memory, otherwise
     /// prefers anonymous memory for performance reasons.
-    pub fn allocate_guest_memory(&self) -> Result<GuestMemoryMmap, MemoryError> {
+    pub fn allocate_guest_memory(&self) -> Result<Vec<GuestRegionMmap>, MemoryError> {
         let vhost_user_device_used = self
             .block
             .devices
@@ -460,13 +461,13 @@ impl VmResources {
             self.machine_config.mem_size_mib << MIB_TO_BYTES_SHIFT,
         );
         if vhost_user_device_used {
-            GuestMemoryMmap::memfd_backed(
+            memory::memfd_backed(
                 regions.as_ref(),
                 self.machine_config.track_dirty_pages,
                 self.machine_config.huge_pages,
             )
         } else {
-            GuestMemoryMmap::anonymous(
+            memory::anonymous(
                 regions.into_iter(),
                 self.machine_config.track_dirty_pages,
                 self.machine_config.huge_pages,

src/vmm/src/test_utils/mod.rs

@@ -15,7 +15,8 @@ use crate::test_utils::mock_resources::{MockBootSourceConfig, MockVmConfig, Mock
 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::{GuestMemoryExtension, GuestMemoryMmap};
+use crate::vstate::memory;
+use crate::vstate::memory::{GuestMemoryMmap, GuestRegionMmap};
 use crate::{EventManager, Vmm};
 
 pub mod mock_resources;
@@ -26,15 +27,31 @@ pub fn single_region_mem(region_size: usize) -> GuestMemoryMmap {
     single_region_mem_at(0, region_size)
 }
 
+pub fn single_region_mem_raw(region_size: usize) -> Vec<GuestRegionMmap> {
+    single_region_mem_at_raw(0, region_size)
+}
+
 /// Creates a [`GuestMemoryMmap`] with a single region of the given size starting at the given
 /// guest physical address `at` and without dirty tracking.
 pub fn single_region_mem_at(at: u64, size: usize) -> GuestMemoryMmap {
     multi_region_mem(&[(GuestAddress(at), size)])
 }
 
+pub fn single_region_mem_at_raw(at: u64, size: usize) -> Vec<GuestRegionMmap> {
+    multi_region_mem_raw(&[(GuestAddress(at), size)])
+}
+
 /// Creates a [`GuestMemoryMmap`] with multiple regions and without dirty page tracking.
 pub fn multi_region_mem(regions: &[(GuestAddress, usize)]) -> GuestMemoryMmap {
-    GuestMemoryMmap::anonymous(regions.iter().copied(), false, HugePageConfig::None)
+    GuestMemoryMmap::from_regions(
+        memory::anonymous(regions.iter().copied(), false, HugePageConfig::None)
+            .expect("Cannot initialize memory"),
+    )
+    .unwrap()
+}
+
+pub fn multi_region_mem_raw(regions: &[(GuestAddress, usize)]) -> Vec<GuestRegionMmap> {
+    memory::anonymous(regions.iter().copied(), false, HugePageConfig::None)
         .expect("Cannot initialize memory")
 }
 
@@ -44,6 +61,10 @@ pub fn arch_mem(mem_size_bytes: usize) -> GuestMemoryMmap {
     multi_region_mem(&crate::arch::arch_memory_regions(mem_size_bytes))
 }
 
+pub fn arch_mem_raw(mem_size_bytes: usize) -> Vec<GuestRegionMmap> {
+    multi_region_mem_raw(&crate::arch::arch_memory_regions(mem_size_bytes))
+}
+
 pub fn create_vmm(
     _kernel_image: Option<&str>,
     is_diff: bool,