[RFC] Use guest memory collection and other refactors

src/vmm/src/arch/aarch64/fdt.rs

@@ -9,7 +9,7 @@ use std::ffi::CString;
 use std::fmt::Debug;
 
 use vm_fdt::{Error as VmFdtError, FdtWriter, FdtWriterNode};
-use vm_memory::GuestMemoryError;
+use vm_memory::{GuestMemoryError, GuestMemoryRegion};
 
 use super::cache_info::{CacheEntry, read_cache_config};
 use super::gic::GICDevice;
@@ -22,7 +22,7 @@ use crate::device_manager::mmio::MMIODeviceInfo;
 use crate::device_manager::pci_mngr::PciDevices;
 use crate::devices::acpi::vmgenid::{VMGENID_MEM_SIZE, VmGenId};
 use crate::initrd::InitrdConfig;
-use crate::vstate::memory::{Address, GuestMemory, GuestMemoryMmap};
+use crate::vstate::memory::{Address, GuestMemory, GuestMemoryMmap, GuestRegionType};
 
 // This is a value for uniquely identifying the FDT node declaring the interrupt controller.
 const GIC_PHANDLE: u32 = 1;
@@ -227,14 +227,16 @@ fn create_memory_node(fdt: &mut FdtWriter, guest_mem: &GuestMemoryMmap) -> Resul
     // The reason we do this is that Linux does not allow remapping system memory. However, without
     // remap, kernel drivers cannot get virtual addresses to read data from device memory. Leaving
     // this memory region out allows Linux kernel modules to remap and thus read this region.
-    let mem_size = guest_mem.last_addr().raw_value()
-        - super::layout::DRAM_MEM_START
-        - super::layout::SYSTEM_MEM_SIZE
-        + 1;
-    let mem_reg_prop = &[
-        super::layout::DRAM_MEM_START + super::layout::SYSTEM_MEM_SIZE,
-        mem_size,
-    ];
+    // Pick the first (and only) memory region
+    let dram_region = guest_mem
+        .iter()
+        .find(|region| region.region_type == GuestRegionType::Dram)
+        .unwrap();
+    let start_addr = dram_region
+        .start_addr()
+        .unchecked_add(super::layout::SYSTEM_MEM_SIZE);
+    let mem_size = dram_region.last_addr().unchecked_offset_from(start_addr) + 1;
+    let mem_reg_prop = &[start_addr.raw_value(), mem_size];
     let mem = fdt.begin_node("memory@ram")?;
     fdt.property_string("device_type", "memory")?;
     fdt.property_array_u64("reg", mem_reg_prop)?;

src/vmm/src/arch/aarch64/mod.rs

@@ -22,15 +22,17 @@ use std::fs::File;
 
 use linux_loader::loader::pe::PE as Loader;
 use linux_loader::loader::{Cmdline, KernelLoader};
-use vm_memory::GuestMemoryError;
+use vm_memory::{GuestMemoryError, GuestMemoryRegion};
 
 use crate::arch::{BootProtocol, EntryPoint, arch_memory_regions_with_gap};
 use crate::cpu_config::aarch64::{CpuConfiguration, CpuConfigurationError};
 use crate::cpu_config::templates::CustomCpuTemplate;
 use crate::initrd::InitrdConfig;
 use crate::utils::{align_up, u64_to_usize, usize_to_u64};
 use crate::vmm_config::machine_config::MachineConfig;
-use crate::vstate::memory::{Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap};
+use crate::vstate::memory::{
+    Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap, GuestRegionType,
+};
 use crate::vstate::vcpu::KvmVcpuError;
 use crate::{DeviceManager, Kvm, Vcpu, VcpuConfig, Vm, logger};
 
@@ -151,31 +153,29 @@ pub fn initrd_load_addr(guest_mem: &GuestMemoryMmap, initrd_size: usize) -> Opti
         usize_to_u64(initrd_size),
         usize_to_u64(super::GUEST_PAGE_SIZE),
     );
-    match GuestAddress(get_fdt_addr(guest_mem)).checked_sub(rounded_size) {
-        Some(offset) => {
-            if guest_mem.address_in_range(offset) {
-                Some(offset.raw_value())
-            } else {
-                None
-            }
-        }
-        None => None,
-    }
+    GuestAddress(get_fdt_addr(guest_mem))
+        .checked_sub(rounded_size)
+        .filter(|&addr| guest_mem.address_in_range(addr))
+        .map(|addr| addr.raw_value())
 }
 
 // Auxiliary function to get the address where the device tree blob is loaded.
 fn get_fdt_addr(mem: &GuestMemoryMmap) -> u64 {
+    // Find the first (and only) DRAM region.
+    let dram_region = mem
+        .iter()
+        .find(|region| region.region_type == GuestRegionType::Dram)
+        .unwrap();
+
     // If the memory allocated is smaller than the size allocated for the FDT,
     // we return the start of the DRAM so that
     // we allow the code to try and load the FDT.
-
-    if let Some(addr) = mem.last_addr().checked_sub(layout::FDT_MAX_SIZE as u64 - 1)
-        && mem.address_in_range(addr)
-    {
-        return addr.raw_value();
-    }
-
-    layout::DRAM_MEM_START
+    dram_region
+        .last_addr()
+        .checked_sub(layout::FDT_MAX_SIZE as u64 - 1)
+        .filter(|&addr| mem.address_in_range(addr))
+        .map(|addr| addr.raw_value())
+        .unwrap_or(layout::DRAM_MEM_START)
 }
 
 /// Load linux kernel into guest memory.

src/vmm/src/arch/x86_64/mod.rs

@@ -31,12 +31,13 @@ pub mod xstate;
 #[allow(missing_docs)]
 pub mod generated;
 
+use std::cmp::max;
 use std::fs::File;
 
 use kvm::Kvm;
 use layout::{
-    CMDLINE_START, FIRST_ADDR_PAST_32BITS, FIRST_ADDR_PAST_64BITS_MMIO, MMIO32_MEM_SIZE,
-    MMIO32_MEM_START, MMIO64_MEM_SIZE, MMIO64_MEM_START, PCI_MMCONFIG_SIZE, PCI_MMCONFIG_START,
+    CMDLINE_START, MMIO32_MEM_SIZE, MMIO32_MEM_START, MMIO64_MEM_SIZE, MMIO64_MEM_START,
+    PCI_MMCONFIG_SIZE, PCI_MMCONFIG_START,
 };
 use linux_loader::configurator::linux::LinuxBootConfigurator;
 use linux_loader::configurator::pvh::PvhBootConfigurator;
@@ -59,7 +60,7 @@ use crate::initrd::InitrdConfig;
 use crate::utils::{align_down, u64_to_usize, usize_to_u64};
 use crate::vmm_config::machine_config::MachineConfig;
 use crate::vstate::memory::{
-    Address, GuestAddress, GuestMemory, GuestMemoryMmap, GuestMemoryRegion,
+    Address, GuestAddress, GuestMemory, GuestMemoryMmap, GuestMemoryRegion, GuestRegionType,
 };
 use crate::vstate::vcpu::KvmVcpuConfigureError;
 use crate::{Vcpu, VcpuConfig, Vm, logger};
@@ -253,10 +254,6 @@ fn configure_pvh(
     initrd: &Option<InitrdConfig>,
 ) -> Result<(), ConfigurationError> {
     const XEN_HVM_START_MAGIC_VALUE: u32 = 0x336e_c578;
-    let first_addr_past_32bits = GuestAddress(FIRST_ADDR_PAST_32BITS);
-    let end_32bit_gap_start = GuestAddress(MMIO32_MEM_START);
-    let first_addr_past_64bits = GuestAddress(FIRST_ADDR_PAST_64BITS_MMIO);
-    let end_64bit_gap_start = GuestAddress(MMIO64_MEM_START);
     let himem_start = GuestAddress(layout::HIMEM_START);
 
     // Vector to hold modules (currently either empty or holding initrd).
@@ -294,36 +291,21 @@ fn configure_pvh(
         type_: E820_RESERVED,
         ..Default::default()
     });
-    let last_addr = guest_mem.last_addr();
 
-    if last_addr > first_addr_past_64bits {
-        memmap.push(hvm_memmap_table_entry {
-            addr: first_addr_past_64bits.raw_value(),
-            size: last_addr.unchecked_offset_from(first_addr_past_64bits) + 1,
-            type_: MEMMAP_TYPE_RAM,
-            ..Default::default()
-        });
-    }
-
-    if last_addr > first_addr_past_32bits {
+    for region in guest_mem
+        .iter()
+        .filter(|region| region.region_type == GuestRegionType::Dram)
+    {
+        // the first 1MB is reserved for the kernel
+        let addr = max(himem_start, region.start_addr());
         memmap.push(hvm_memmap_table_entry {
-            addr: first_addr_past_32bits.raw_value(),
-            size: (end_64bit_gap_start.unchecked_offset_from(first_addr_past_32bits))
-                .min(last_addr.unchecked_offset_from(first_addr_past_32bits) + 1),
+            addr: addr.raw_value(),
+            size: region.last_addr().unchecked_offset_from(addr) + 1,
             type_: MEMMAP_TYPE_RAM,
             ..Default::default()
         });
     }
 
-    memmap.push(hvm_memmap_table_entry {
-        addr: himem_start.raw_value(),
-        size: end_32bit_gap_start
-            .unchecked_offset_from(himem_start)
-            .min(last_addr.unchecked_offset_from(himem_start) + 1),
-        type_: MEMMAP_TYPE_RAM,
-        ..Default::default()
-    });
-
     // Construct the hvm_start_info structure and serialize it into
     // boot_params.  This will be stored at PVH_INFO_START address, and %rbx
     // will be initialized to contain PVH_INFO_START prior to starting the
@@ -368,10 +350,6 @@ fn configure_64bit_boot(
     const KERNEL_HDR_MAGIC: u32 = 0x5372_6448;
     const KERNEL_LOADER_OTHER: u8 = 0xff;
     const KERNEL_MIN_ALIGNMENT_BYTES: u32 = 0x0100_0000; // Must be non-zero.
-    let first_addr_past_32bits = GuestAddress(FIRST_ADDR_PAST_32BITS);
-    let end_32bit_gap_start = GuestAddress(MMIO32_MEM_START);
-    let first_addr_past_64bits = GuestAddress(FIRST_ADDR_PAST_64BITS_MMIO);
-    let end_64bit_gap_start = GuestAddress(MMIO64_MEM_START);
 
     let himem_start = GuestAddress(layout::HIMEM_START);
 
@@ -409,35 +387,20 @@ fn configure_64bit_boot(
         E820_RESERVED,
     )?;
 
-    let last_addr = guest_mem.last_addr();
-
-    if last_addr > first_addr_past_64bits {
-        add_e820_entry(
-            &mut params,
-            first_addr_past_64bits.raw_value(),
-            last_addr.unchecked_offset_from(first_addr_past_64bits) + 1,
-            E820_RAM,
-        )?;
-    }
-
-    if last_addr > first_addr_past_32bits {
+    for region in guest_mem
+        .iter()
+        .filter(|region| region.region_type == GuestRegionType::Dram)
+    {
+        // the first 1MB is reserved for the kernel
+        let addr = max(himem_start, region.start_addr());
         add_e820_entry(
             &mut params,
-            first_addr_past_32bits.raw_value(),
-            (end_64bit_gap_start.unchecked_offset_from(first_addr_past_32bits))
-                .min(last_addr.unchecked_offset_from(first_addr_past_32bits) + 1),
+            addr.raw_value(),
+            region.last_addr().unchecked_offset_from(addr) + 1,
             E820_RAM,
         )?;
     }
 
-    add_e820_entry(
-        &mut params,
-        himem_start.raw_value(),
-        (last_addr.unchecked_offset_from(himem_start) + 1)
-            .min(end_32bit_gap_start.unchecked_offset_from(himem_start)),
-        E820_RAM,
-    )?;
-
     LinuxBootConfigurator::write_bootparams(
         &BootParams::new(&params, GuestAddress(layout::ZERO_PAGE_START)),
         guest_mem,
@@ -573,6 +536,7 @@ mod tests {
     use linux_loader::loader::bootparam::boot_e820_entry;
 
     use super::*;
+    use crate::arch::x86_64::layout::FIRST_ADDR_PAST_32BITS;
     use crate::test_utils::{arch_mem, single_region_mem};
     use crate::utils::mib_to_bytes;
     use crate::vstate::resources::ResourceAllocator;

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")]
@@ -462,7 +462,6 @@ pub fn build_microvm_from_snapshot(
         event_manager,
         vm_resources,
         instance_id: &instance_info.id,
-        restored_from_file: uffd.is_none(),
         vcpus_exit_evt: &vcpus_exit_evt,
     };
     #[allow(unused_mut)]

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/device_manager/mod.rs

@@ -439,7 +439,6 @@ pub struct DeviceRestoreArgs<'a> {
     pub vcpus_exit_evt: &'a EventFd,
     pub vm_resources: &'a mut VmResources,
     pub instance_id: &'a str,
-    pub restored_from_file: bool,
 }
 
 impl std::fmt::Debug for DeviceRestoreArgs<'_> {
@@ -449,7 +448,6 @@ impl std::fmt::Debug for DeviceRestoreArgs<'_> {
             .field("vm", &self.vm)
             .field("vm_resources", &self.vm_resources)
             .field("instance_id", &self.instance_id)
-            .field("restored_from_file", &self.restored_from_file)
             .finish()
     }
 }
@@ -487,7 +485,6 @@ impl<'a> Persist<'a> for DeviceManager {
             event_manager: constructor_args.event_manager,
             vm_resources: constructor_args.vm_resources,
             instance_id: constructor_args.instance_id,
-            restored_from_file: constructor_args.restored_from_file,
         };
         let mmio_devices = MMIODeviceManager::restore(mmio_ctor_args, &state.mmio_state)?;
 
@@ -505,7 +502,6 @@ impl<'a> Persist<'a> for DeviceManager {
             mem: constructor_args.mem,
             vm_resources: constructor_args.vm_resources,
             instance_id: constructor_args.instance_id,
-            restored_from_file: constructor_args.restored_from_file,
             event_manager: constructor_args.event_manager,
         };
         let pci_devices = PciDevices::restore(pci_ctor_args, &state.pci_state)?;

src/vmm/src/device_manager/pci_mngr.rs

@@ -247,7 +247,6 @@ pub struct PciDevicesConstructorArgs<'a> {
     pub mem: &'a GuestMemoryMmap,
     pub vm_resources: &'a mut VmResources,
     pub instance_id: &'a str,
-    pub restored_from_file: bool,
     pub event_manager: &'a mut EventManager,
 }
 
@@ -258,7 +257,6 @@ impl<'a> Debug for PciDevicesConstructorArgs<'a> {
             .field("mem", &self.mem)
             .field("vm_resources", &self.vm_resources)
             .field("instance_id", &self.instance_id)
-            .field("restored_from_file", &self.restored_from_file)
             .finish()
     }
 }
@@ -410,10 +408,7 @@ impl<'a> Persist<'a> for PciDevices {
         if let Some(balloon_state) = &state.balloon_device {
             let device = Arc::new(Mutex::new(
                 Balloon::restore(
-                    BalloonConstructorArgs {
-                        mem: mem.clone(),
-                        restored_from_file: constructor_args.restored_from_file,
-                    },
+                    BalloonConstructorArgs { mem: mem.clone() },
                     &balloon_state.device_state,
                 )
                 .unwrap(),
@@ -668,7 +663,6 @@ mod tests {
             mem: vmm.vm.guest_memory(),
             vm_resources,
             instance_id: "microvm-id",
-            restored_from_file: true,
             event_manager: &mut event_manager,
         };
         let _restored_dev_manager =