refactor(memory): GuestMemoryExtension introduction

Changed `SnapshotMemory` trait to `GuestMemoryExtension`
and moved all public methods that create GuestMemoryMmap into it.
This removed a need for a public `test_utils` module and it makes
all methods that directly act on `GuestMemoryMmap` to be associated
with it.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

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

@@ -426,6 +426,7 @@ mod tests {
     use super::*;
     use crate::arch::aarch64::gic::create_gic;
     use crate::arch::aarch64::{arch_memory_regions, layout};
+    use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryMmap};
 
     const LEN: u64 = 4096;
 
@@ -458,8 +459,8 @@ mod tests {
     #[test]
     fn test_create_fdt_with_devices() {
         let regions = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000);
-        let mem = crate::vstate::memory::test_utils::create_anon_guest_memory(&regions, false)
-            .expect("Cannot initialize memory");
+        let mem =
+            GuestMemoryMmap::from_raw_regions(&regions, false).expect("Cannot initialize memory");
 
         let dev_info: HashMap<(DeviceType, std::string::String), MMIODeviceInfo> = [
             (
@@ -498,8 +499,8 @@ mod tests {
     #[test]
     fn test_create_fdt() {
         let regions = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000);
-        let mem = crate::vstate::memory::test_utils::create_anon_guest_memory(&regions, false)
-            .expect("Cannot initialize memory");
+        let mem =
+            GuestMemoryMmap::from_raw_regions(&regions, false).expect("Cannot initialize memory");
         let kvm = Kvm::new().unwrap();
         let vm = kvm.create_vm().unwrap();
         let gic = create_gic(&vm, 1, None).unwrap();
@@ -556,8 +557,8 @@ mod tests {
     #[test]
     fn test_create_fdt_with_initrd() {
         let regions = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000);
-        let mem = crate::vstate::memory::test_utils::create_anon_guest_memory(&regions, false)
-            .expect("Cannot initialize memory");
+        let mem =
+            GuestMemoryMmap::from_raw_regions(&regions, false).expect("Cannot initialize memory");
         let kvm = Kvm::new().unwrap();
         let vm = kvm.create_vm().unwrap();
         let gic = create_gic(&vm, 1, None).unwrap();

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

@@ -114,6 +114,7 @@ fn get_fdt_addr(mem: &GuestMemoryMmap) -> u64 {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::vstate::memory::GuestMemoryExtension;
 
     #[test]
     fn test_regions_lt_1024gb() {
@@ -134,18 +135,18 @@ mod tests {
     #[test]
     fn test_get_fdt_addr() {
         let regions = arch_memory_regions(layout::FDT_MAX_SIZE - 0x1000);
-        let mem = crate::vstate::memory::test_utils::create_anon_guest_memory(&regions, false)
-            .expect("Cannot initialize memory");
+        let mem =
+            GuestMemoryMmap::from_raw_regions(&regions, false).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), layout::DRAM_MEM_START);
 
         let regions = arch_memory_regions(layout::FDT_MAX_SIZE);
-        let mem = crate::vstate::memory::test_utils::create_anon_guest_memory(&regions, false)
-            .expect("Cannot initialize memory");
+        let mem =
+            GuestMemoryMmap::from_raw_regions(&regions, false).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), layout::DRAM_MEM_START);
 
         let regions = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000);
-        let mem = crate::vstate::memory::test_utils::create_anon_guest_memory(&regions, false)
-            .expect("Cannot initialize memory");
+        let mem =
+            GuestMemoryMmap::from_raw_regions(&regions, false).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), 0x1000 + layout::DRAM_MEM_START);
     }
 }

src/vmm/src/arch/aarch64/vcpu.rs

@@ -201,15 +201,16 @@ mod tests {
 
     use super::*;
     use crate::arch::aarch64::{arch_memory_regions, layout};
+    use crate::vstate::memory::GuestMemoryExtension;
 
     #[test]
     fn test_setup_regs() {
         let kvm = Kvm::new().unwrap();
         let vm = kvm.create_vm().unwrap();
         let vcpu = vm.create_vcpu(0).unwrap();
         let regions = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000);
-        let mem = crate::vstate::memory::test_utils::create_anon_guest_memory(&regions, false)
-            .expect("Cannot initialize memory");
+        let mem =
+            GuestMemoryMmap::from_raw_regions(&regions, false).expect("Cannot initialize memory");
 
         let res = setup_boot_regs(&vcpu, 0, 0x0, &mem);
         assert!(matches!(

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

@@ -210,6 +210,7 @@ mod tests {
     use linux_loader::loader::bootparam::boot_e820_entry;
 
     use super::*;
+    use crate::vstate::memory::GuestMemoryExtension;
 
     #[test]
     fn regions_lt_4gb() {
@@ -230,11 +231,7 @@ mod tests {
     #[test]
     fn test_system_configuration() {
         let no_vcpus = 4;
-        let gm = crate::vstate::memory::test_utils::create_anon_guest_memory(
-            &[(GuestAddress(0), 0x10000)],
-            false,
-        )
-        .unwrap();
+        let gm = GuestMemoryMmap::from_raw_regions(&[(GuestAddress(0), 0x10000)], false).unwrap();
         let config_err = configure_system(&gm, GuestAddress(0), 0, &None, 1);
         assert!(config_err.is_err());
         assert_eq!(
@@ -245,25 +242,19 @@ mod tests {
         // Now assigning some memory that falls before the 32bit memory hole.
         let mem_size = 128 << 20;
         let arch_mem_regions = arch_memory_regions(mem_size);
-        let gm =
-            crate::vstate::memory::test_utils::create_anon_guest_memory(&arch_mem_regions, false)
-                .unwrap();
+        let gm = GuestMemoryMmap::from_raw_regions(&arch_mem_regions, false).unwrap();
         configure_system(&gm, GuestAddress(0), 0, &None, no_vcpus).unwrap();
 
         // Now assigning some memory that is equal to the start of the 32bit memory hole.
         let mem_size = 3328 << 20;
         let arch_mem_regions = arch_memory_regions(mem_size);
-        let gm =
-            crate::vstate::memory::test_utils::create_anon_guest_memory(&arch_mem_regions, false)
-                .unwrap();
+        let gm = GuestMemoryMmap::from_raw_regions(&arch_mem_regions, false).unwrap();
         configure_system(&gm, GuestAddress(0), 0, &None, no_vcpus).unwrap();
 
         // Now assigning some memory that falls after the 32bit memory hole.
         let mem_size = 3330 << 20;
         let arch_mem_regions = arch_memory_regions(mem_size);
-        let gm =
-            crate::vstate::memory::test_utils::create_anon_guest_memory(&arch_mem_regions, false)
-                .unwrap();
+        let gm = GuestMemoryMmap::from_raw_regions(&arch_mem_regions, false).unwrap();
         configure_system(&gm, GuestAddress(0), 0, &None, no_vcpus).unwrap();
     }
 

src/vmm/src/arch/x86_64/mptable.rs

@@ -288,8 +288,9 @@ pub fn setup_mptable(mem: &GuestMemoryMmap, num_cpus: u8) -> Result<(), MptableE
 
 #[cfg(test)]
 mod tests {
+
     use super::*;
-    use crate::vstate::memory::Bytes;
+    use crate::vstate::memory::{Bytes, GuestMemoryExtension};
 
     fn table_entry_size(type_: u8) -> usize {
         match u32::from(type_) {
@@ -305,7 +306,7 @@ mod tests {
     #[test]
     fn bounds_check() {
         let num_cpus = 4;
-        let mem = crate::vstate::memory::test_utils::create_guest_memory_unguarded(
+        let mem = GuestMemoryMmap::from_raw_regions_unguarded(
             &[(GuestAddress(MPTABLE_START), compute_mp_size(num_cpus))],
             false,
         )
@@ -317,7 +318,7 @@ mod tests {
     #[test]
     fn bounds_check_fails() {
         let num_cpus = 4;
-        let mem = crate::vstate::memory::test_utils::create_guest_memory_unguarded(
+        let mem = GuestMemoryMmap::from_raw_regions_unguarded(
             &[(GuestAddress(MPTABLE_START), compute_mp_size(num_cpus) - 1)],
             false,
         )
@@ -329,7 +330,7 @@ mod tests {
     #[test]
     fn mpf_intel_checksum() {
         let num_cpus = 1;
-        let mem = crate::vstate::memory::test_utils::create_guest_memory_unguarded(
+        let mem = GuestMemoryMmap::from_raw_regions_unguarded(
             &[(GuestAddress(MPTABLE_START), compute_mp_size(num_cpus))],
             false,
         )
@@ -345,7 +346,7 @@ mod tests {
     #[test]
     fn mpc_table_checksum() {
         let num_cpus = 4;
-        let mem = crate::vstate::memory::test_utils::create_guest_memory_unguarded(
+        let mem = GuestMemoryMmap::from_raw_regions_unguarded(
             &[(GuestAddress(MPTABLE_START), compute_mp_size(num_cpus))],
             false,
         )
@@ -379,7 +380,7 @@ mod tests {
 
     #[test]
     fn cpu_entry_count() {
-        let mem = crate::vstate::memory::test_utils::create_guest_memory_unguarded(
+        let mem = GuestMemoryMmap::from_raw_regions_unguarded(
             &[(
                 GuestAddress(MPTABLE_START),
                 compute_mp_size(MAX_SUPPORTED_CPUS),
@@ -417,8 +418,8 @@ mod tests {
     #[test]
     fn cpu_entry_count_max() {
         let cpus = MAX_SUPPORTED_CPUS + 1;
-        let mem = crate::vstate::memory::test_utils::create_guest_memory_unguarded(
-            &[(GuestAddress(MPTABLE_START), compute_mp_size(cpus as u8))],
+        let mem = GuestMemoryMmap::from_raw_regions_unguarded(
+            &[(GuestAddress(MPTABLE_START), compute_mp_size(cpus))],
             false,
         )
         .unwrap();

src/vmm/src/arch/x86_64/regs.rs

@@ -242,23 +242,16 @@ mod tests {
     use utils::u64_to_usize;
 
     use super::*;
-    use crate::vstate::memory::{Bytes, GuestAddress, GuestMemoryMmap};
+    use crate::vstate::memory::{Bytes, GuestAddress, GuestMemoryExtension, GuestMemoryMmap};
 
     fn create_guest_mem(mem_size: Option<u64>) -> GuestMemoryMmap {
         let page_size = 0x10000usize;
         let mem_size = u64_to_usize(mem_size.unwrap_or(page_size as u64));
         if mem_size % page_size == 0 {
-            crate::vstate::memory::test_utils::create_anon_guest_memory(
-                &[(GuestAddress(0), mem_size)],
-                false,
-            )
-            .unwrap()
+            GuestMemoryMmap::from_raw_regions(&[(GuestAddress(0), mem_size)], false).unwrap()
         } else {
-            crate::vstate::memory::test_utils::create_guest_memory_unguarded(
-                &[(GuestAddress(0), mem_size)],
-                false,
-            )
-            .unwrap()
+            GuestMemoryMmap::from_raw_regions_unguarded(&[(GuestAddress(0), mem_size)], false)
+                .unwrap()
         }
     }
 

src/vmm/src/builder.rs

@@ -55,7 +55,7 @@ use crate::vmm_config::boot_source::BootConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::{MachineConfigUpdate, VmConfig, VmConfigError};
 use crate::volatile::ReadVolatile;
-use crate::vstate::memory::{GuestAddress, GuestMemory, GuestMemoryMmap};
+use crate::vstate::memory::{GuestAddress, GuestMemory, GuestMemoryExtension, GuestMemoryMmap};
 use crate::vstate::vcpu::{Vcpu, VcpuConfig};
 use crate::vstate::vm::Vm;
 use crate::{device_manager, EventManager, RestoreVcpusError, Vmm, VmmError};
@@ -265,7 +265,9 @@ pub fn build_microvm_for_boot(
         .ok_or(MissingKernelConfig)?;
 
     let track_dirty_pages = vm_resources.track_dirty_pages();
-    let guest_memory = create_guest_memory(vm_resources.vm_config.mem_size_mib, track_dirty_pages)?;
+    let guest_memory =
+        GuestMemoryMmap::with_size(vm_resources.vm_config.mem_size_mib, track_dirty_pages)
+            .map_err(StartMicrovmError::GuestMemory)?;
     let entry_addr = load_kernel(boot_config, &guest_memory)?;
     let initrd = load_initrd_from_config(boot_config, &guest_memory)?;
     // Clone the command-line so that a failed boot doesn't pollute the original.
@@ -540,24 +542,6 @@ pub fn build_microvm_from_snapshot(
     Ok(vmm)
 }
 
-/// Creates GuestMemory of `mem_size_mib` MiB in size.
-pub fn create_guest_memory(
-    mem_size_mib: usize,
-    track_dirty_pages: bool,
-) -> Result<GuestMemoryMmap, StartMicrovmError> {
-    let mem_size = mem_size_mib << 20;
-    let arch_mem_regions = crate::arch::arch_memory_regions(mem_size);
-
-    crate::vstate::memory::create_guest_memory(
-        &arch_mem_regions
-            .iter()
-            .map(|(addr, size)| (None, *addr, *size))
-            .collect::<Vec<_>>()[..],
-        track_dirty_pages,
-    )
-    .map_err(StartMicrovmError::GuestMemory)
-}
-
 fn load_kernel(
     boot_config: &BootConfig,
     guest_memory: &GuestMemoryMmap,
@@ -986,7 +970,6 @@ pub mod tests {
     use crate::vmm_config::net::{NetBuilder, NetworkInterfaceConfig};
     use crate::vmm_config::vsock::tests::default_config;
     use crate::vmm_config::vsock::{VsockBuilder, VsockDeviceConfig};
-    use crate::vstate::memory::GuestMemory;
 
     #[derive(Debug)]
     pub(crate) struct CustomBlockConfig {
@@ -1051,7 +1034,7 @@ pub mod tests {
     }
 
     pub(crate) fn default_vmm() -> Vmm {
-        let guest_memory = create_guest_memory(128, false).unwrap();
+        let guest_memory = GuestMemoryMmap::with_size(128, false).unwrap();
 
         let vcpus_exit_evt = EventFd::new(libc::EFD_NONBLOCK)
             .map_err(VmmError::EventFd)
@@ -1228,8 +1211,7 @@ pub mod tests {
     }
 
     fn create_guest_mem_at(at: GuestAddress, size: usize) -> GuestMemoryMmap {
-        crate::vstate::memory::test_utils::create_guest_memory_unguarded(&[(at, size)], false)
-            .unwrap()
+        GuestMemoryMmap::from_raw_regions_unguarded(&[(at, size)], false).unwrap()
     }
 
     pub(crate) fn create_guest_mem_with_size(size: usize) -> GuestMemoryMmap {
@@ -1305,21 +1287,21 @@ pub mod tests {
 
         // Case 1: create guest memory without dirty page tracking
         {
-            let guest_memory = create_guest_memory(mem_size, false).unwrap();
+            let guest_memory = GuestMemoryMmap::with_size(mem_size, false).unwrap();
             assert!(!is_dirty_tracking_enabled(&guest_memory));
         }
 
         // Case 2: create guest memory with dirty page tracking
         {
-            let guest_memory = create_guest_memory(mem_size, true).unwrap();
+            let guest_memory = GuestMemoryMmap::with_size(mem_size, true).unwrap();
             assert!(is_dirty_tracking_enabled(&guest_memory));
         }
     }
 
     #[test]
     fn test_create_vcpus() {
         let vcpu_count = 2;
-        let guest_memory = create_guest_memory(128, false).unwrap();
+        let guest_memory = GuestMemoryMmap::with_size(128, false).unwrap();
 
         #[allow(unused_mut)]
         let mut vm = Vm::new(vec![]).unwrap();

src/vmm/src/device_manager/legacy.rs

@@ -172,16 +172,13 @@ impl PortIODeviceManager {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::vstate::memory::GuestAddress;
+    use crate::vstate::memory::{GuestAddress, GuestMemoryExtension, GuestMemoryMmap};
     use crate::Vm;
 
     #[test]
     fn test_register_legacy_devices() {
-        let guest_mem = crate::vstate::memory::test_utils::create_anon_guest_memory(
-            &[(GuestAddress(0x0), 0x1000)],
-            false,
-        )
-        .unwrap();
+        let guest_mem =
+            GuestMemoryMmap::from_raw_regions(&[(GuestAddress(0x0), 0x1000)], false).unwrap();
         let mut vm = Vm::new(vec![]).unwrap();
         vm.memory_init(&guest_mem, false).unwrap();
         crate::builder::setup_interrupt_controller(&mut vm).unwrap();

src/vmm/src/device_manager/mmio.rs

@@ -469,7 +469,7 @@ mod tests {
 
     use super::*;
     use crate::devices::virtio::{ActivateError, Queue, VirtioDevice};
-    use crate::vstate::memory::{GuestAddress, GuestMemoryMmap};
+    use crate::vstate::memory::{GuestAddress, GuestMemoryExtension, GuestMemoryMmap};
     use crate::{builder, Vm};
 
     const QUEUE_SIZES: &[u16] = &[64];
@@ -573,7 +573,7 @@ mod tests {
     fn test_register_virtio_device() {
         let start_addr1 = GuestAddress(0x0);
         let start_addr2 = GuestAddress(0x1000);
-        let guest_mem = crate::vstate::memory::test_utils::create_anon_guest_memory(
+        let guest_mem = GuestMemoryMmap::from_raw_regions(
             &[(start_addr1, 0x1000), (start_addr2, 0x1000)],
             false,
         )
@@ -603,7 +603,7 @@ mod tests {
     fn test_register_too_many_devices() {
         let start_addr1 = GuestAddress(0x0);
         let start_addr2 = GuestAddress(0x1000);
-        let guest_mem = crate::vstate::memory::test_utils::create_anon_guest_memory(
+        let guest_mem = GuestMemoryMmap::from_raw_regions(
             &[(start_addr1, 0x1000), (start_addr2, 0x1000)],
             false,
         )
@@ -663,7 +663,7 @@ mod tests {
     fn test_device_info() {
         let start_addr1 = GuestAddress(0x0);
         let start_addr2 = GuestAddress(0x1000);
-        let guest_mem = crate::vstate::memory::test_utils::create_anon_guest_memory(
+        let guest_mem = GuestMemoryMmap::from_raw_regions(
             &[(start_addr1, 0x1000), (start_addr2, 0x1000)],
             false,
         )