refactor: centralize GuestMemoryMmap creation

In this day and age, Firecracker supports theoretically 4 different ways
of backing guest memory:

1. Normal MAP_ANONYMOUS | MAP_PRIVATE memory
2. memfd backed memory, mapped as shared
3. direct mapping of a snapshot file
4. MAP_ANONYMOUS again, but this time regions are described by snapshot
   file.

We have 3 different functions for creating these different backing
stores, which then call each other and vm_memory's APIs. Clean this up
by consolidating these into just one function that can be called with a
description of the memory regions, plus an enum argument stating how
each region should be backed.

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

Author: roypat

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

@@ -372,13 +372,15 @@ mod tests {
     use std::sync::atomic::Ordering;
 
     use vhost::{VhostUserMemoryRegionInfo, VringConfigData};
+    use vm_memory::GuestAddress;
     use vmm_sys_util::tempfile::TempFile;
 
     use super::*;
     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::{FileOffset, GuestAddress, GuestMemoryExtension};
+    use crate::vmm_config::machine_config::HugePageConfig;
+    use crate::vstate::memory::{GuestMemoryExtension, MemoryKind};
 
     #[test]
     fn test_from_config() {
@@ -778,12 +780,13 @@ mod tests {
         let region_size = 0x10000;
         let file = TempFile::new().unwrap().into_file();
         file.set_len(region_size as u64).unwrap();
-        let regions = vec![(
-            FileOffset::new(file.try_clone().unwrap(), 0x0),
-            GuestAddress(0x0),
-            region_size,
-        )];
-        let guest_memory = GuestMemoryMmap::from_raw_regions_file(regions, false, false).unwrap();
+        let guest_memory = GuestMemoryMmap::create(
+            [(GuestAddress(0), region_size)].into_iter(),
+            MemoryKind::FilePrivate(file),
+            false,
+            HugePageConfig::None,
+        )
+        .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,7 @@ 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::{Bitmap, Bytes, GuestMemory, GuestMemoryExtension, MemoryKind};
 
     const FILE_LEN: u32 = 1024;
     // 2 pages of memory should be enough to test read/write ops and also dirty tracking.
@@ -230,8 +230,13 @@ pub mod tests {
     }
 
     fn create_mem() -> GuestMemoryMmap {
-        GuestMemoryMmap::from_raw_regions(&[(GuestAddress(0), MEM_LEN)], true, HugePageConfig::None)
-            .unwrap()
+        GuestMemoryMmap::create(
+            [(GuestAddress(0), MEM_LEN)].into_iter(),
+            MemoryKind::MapAnonymous,
+            true,
+            HugePageConfig::None,
+        )
+        .unwrap()
     }
 
     fn check_dirty_mem(mem: &GuestMemoryMmap, addr: GuestAddress, len: u32) {

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::{FileOffset, GuestAddress, GuestMemoryExtension};
+    use crate::vmm_config::machine_config::HugePageConfig;
+    use crate::vstate::memory::{GuestAddress, GuestMemoryExtension, MemoryKind};
 
     #[test]
     fn test_new() {
@@ -758,20 +759,18 @@ mod tests {
         let file = TempFile::new().unwrap().into_file();
         let file_size = 2 * region_size;
         file.set_len(file_size as u64).unwrap();
-        let regions = vec![
-            (
-                FileOffset::new(file.try_clone().unwrap(), 0x0),
-                GuestAddress(0x0),
-                region_size,
-            ),
-            (
-                FileOffset::new(file.try_clone().unwrap(), 0x10000),
-                GuestAddress(0x10000),
-                region_size,
-            ),
+        let regions = [
+            (GuestAddress(0x0), region_size),
+            (GuestAddress(0x10000), region_size),
         ];
 
-        let guest_memory = GuestMemoryMmap::from_raw_regions_file(regions, false, false).unwrap();
+        let guest_memory = GuestMemoryMmap::create(
+            regions.into_iter(),
+            MemoryKind::FilePrivate(file),
+            false,
+            HugePageConfig::None,
+        )
+        .unwrap();
 
         vuh.update_mem_table(&guest_memory).unwrap();
 
@@ -883,13 +882,14 @@ mod tests {
         let region_size = 0x10000;
         let file = TempFile::new().unwrap().into_file();
         file.set_len(region_size as u64).unwrap();
-        let regions = vec![(
-            FileOffset::new(file.try_clone().unwrap(), 0x0),
-            GuestAddress(0x0),
-            region_size,
-        )];
 
-        let guest_memory = GuestMemoryMmap::from_raw_regions_file(regions, false, false).unwrap();
+        let guest_memory = GuestMemoryMmap::create(
+            [(GuestAddress(0), region_size)].into_iter(),
+            MemoryKind::FilePrivate(file),
+            false,
+            HugePageConfig::None,
+        )
+        .unwrap();
 
         let mut queue = Queue::new(69);
         queue.initialize(&guest_memory).unwrap();

src/vmm/src/persist.rs

@@ -14,6 +14,7 @@ use std::sync::{Arc, Mutex};
 use semver::Version;
 use serde::{Deserialize, Serialize};
 use userfaultfd::{FeatureFlags, Uffd, UffdBuilder};
+use vm_memory::GuestAddress;
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
 
 #[cfg(target_arch = "aarch64")]
@@ -38,7 +39,7 @@ use crate::vmm_config::snapshot::{
 };
 use crate::vstate::kvm::KvmState;
 use crate::vstate::memory::{
-    GuestMemory, GuestMemoryExtension, GuestMemoryMmap, GuestMemoryState, MemoryError,
+    GuestMemory, GuestMemoryExtension, GuestMemoryMmap, GuestMemoryState, MemoryError, MemoryKind,
 };
 use crate::vstate::vcpu::{VcpuSendEventError, VcpuState};
 use crate::vstate::vm::VmState;
@@ -515,15 +516,23 @@ pub enum GuestMemoryFromFileError {
     Restore(#[from] MemoryError),
 }
 
-fn guest_memory_from_file(
+/// Creates a [`GuestMemoryMmap`] from a snapshot file
+pub fn guest_memory_from_file(
     mem_file_path: &Path,
     mem_state: &GuestMemoryState,
     track_dirty_pages: bool,
     huge_pages: HugePageConfig,
 ) -> Result<GuestMemoryMmap, GuestMemoryFromFileError> {
     let mem_file = File::open(mem_file_path)?;
-    let guest_mem =
-        GuestMemoryMmap::from_state(Some(&mem_file), mem_state, track_dirty_pages, huge_pages)?;
+    let guest_mem = GuestMemoryMmap::create(
+        mem_state
+            .regions
+            .iter()
+            .map(|region| (GuestAddress(region.base_address), region.size)),
+        MemoryKind::FilePrivate(mem_file),
+        track_dirty_pages,
+        huge_pages,
+    )?;
     Ok(guest_mem)
 }
 
@@ -582,7 +591,15 @@ fn create_guest_memory(
     track_dirty_pages: bool,
     huge_pages: HugePageConfig,
 ) -> Result<(GuestMemoryMmap, Vec<GuestRegionUffdMapping>), GuestMemoryFromUffdError> {
-    let guest_memory = GuestMemoryMmap::from_state(None, mem_state, track_dirty_pages, huge_pages)?;
+    let guest_memory = GuestMemoryMmap::create(
+        mem_state
+            .regions
+            .iter()
+            .map(|region| (GuestAddress(region.base_address), region.size)),
+        MemoryKind::MapAnonymous,
+        track_dirty_pages,
+        huge_pages,
+    )?;
     let mut backend_mappings = Vec::with_capacity(guest_memory.num_regions());
     let mut offset = 0;
     for mem_region in guest_memory.iter() {

src/vmm/src/resources.rs

@@ -28,7 +28,7 @@ use crate::vmm_config::metrics::{init_metrics, MetricsConfig, MetricsConfigError
 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::{GuestMemoryExtension, GuestMemoryMmap, MemoryError, MemoryKind};
 
 /// Errors encountered when configuring microVM resources.
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
@@ -464,20 +464,22 @@ impl VmResources {
         // because that would require running a backend process. If in the future we converge to
         // a single way of backing guest memory for vhost-user and non-vhost-user cases,
         // that would not be worth the effort.
-        if vhost_user_device_used {
-            GuestMemoryMmap::memfd_backed(
+
+        let memory_kind = if vhost_user_device_used {
+            MemoryKind::memfd(
                 self.machine_config.mem_size_mib,
-                self.machine_config.track_dirty_pages,
                 self.machine_config.huge_pages,
-            )
+            )?
         } else {
-            let regions = crate::arch::arch_memory_regions(self.machine_config.mem_size_mib << 20);
-            GuestMemoryMmap::from_raw_regions(
-                &regions,
-                self.machine_config.track_dirty_pages,
-                self.machine_config.huge_pages,
-            )
-        }
+            MemoryKind::MapAnonymous
+        };
+
+        GuestMemoryMmap::create(
+            crate::arch::arch_memory_regions(self.machine_config.mem_size_mib << 20).into_iter(),
+            memory_kind,
+            self.machine_config.track_dirty_pages,
+            self.machine_config.huge_pages,
+        )
     }
 }
 

src/vmm/src/test_utils/mod.rs

@@ -15,7 +15,7 @@ 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::{GuestMemoryExtension, GuestMemoryMmap, MemoryKind};
 use crate::{EventManager, Vmm};
 
 pub mod mock_resources;
@@ -34,8 +34,13 @@ pub fn single_region_mem_at(at: u64, size: usize) -> GuestMemoryMmap {
 
 /// Creates a [`GuestMemoryMmap`] with multiple regions and without dirty page tracking.
 pub fn multi_region_mem(regions: &[(GuestAddress, usize)]) -> GuestMemoryMmap {
-    GuestMemoryMmap::from_raw_regions(regions, false, HugePageConfig::None)
-        .expect("Cannot initialize memory")
+    GuestMemoryMmap::create(
+        regions.iter().copied(),
+        MemoryKind::MapAnonymous,
+        false,
+        HugePageConfig::None,
+    )
+    .expect("Cannot initialize memory")
 }
 
 /// Creates a [`GuestMemoryMmap`] of the given size with the contained regions laid out in