refactor: hoist up hugepages/file-based restore incompatibility check

Move this check out of the memory allocation code and up into the
snapshot restoration code. This allows us to encode the impossibility of
mmap-ing a snapshot file with hugepages at the type level (by simply not
giving an option to pass a `HugePageConfig` to `GuestMemoryMmap::create`
if `MemoryKind::Snapshot(_)` is passed).

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

Author: roypat

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

@@ -379,7 +379,6 @@ 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::vmm_config::machine_config::HugePageConfig;
     use crate::vstate::memory::{GuestMemoryExtension, MemoryKind};
 
     #[test]
@@ -784,7 +783,6 @@ mod tests {
             [(GuestAddress(0), region_size)].into_iter(),
             MemoryKind::FilePrivate(file),
             false,
-            HugePageConfig::None,
         )
         .unwrap();
 

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

@@ -232,9 +232,8 @@ pub mod tests {
     fn create_mem() -> GuestMemoryMmap {
         GuestMemoryMmap::create(
             [(GuestAddress(0), MEM_LEN)].into_iter(),
-            MemoryKind::MapAnonymous,
+            MemoryKind::MapAnonymous(HugePageConfig::None),
             true,
-            HugePageConfig::None,
         )
         .unwrap()
     }

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

@@ -466,7 +466,6 @@ mod tests {
 
     use super::*;
     use crate::test_utils::create_tmp_socket;
-    use crate::vmm_config::machine_config::HugePageConfig;
     use crate::vstate::memory::{GuestAddress, GuestMemoryExtension, MemoryKind};
 
     #[test]
@@ -764,13 +763,9 @@ mod tests {
             (GuestAddress(0x10000), region_size),
         ];
 
-        let guest_memory = GuestMemoryMmap::create(
-            regions.into_iter(),
-            MemoryKind::FilePrivate(file),
-            false,
-            HugePageConfig::None,
-        )
-        .unwrap();
+        let guest_memory =
+            GuestMemoryMmap::create(regions.into_iter(), MemoryKind::FilePrivate(file), false)
+                .unwrap();
 
         vuh.update_mem_table(&guest_memory).unwrap();
 
@@ -887,7 +882,6 @@ mod tests {
             [(GuestAddress(0), region_size)].into_iter(),
             MemoryKind::FilePrivate(file),
             false,
-            HugePageConfig::None,
         )
         .unwrap();
 

src/vmm/src/persist.rs

@@ -449,16 +449,20 @@ pub fn restore_from_snapshot(
     let mem_state = &microvm_state.memory_state;
 
     let (guest_memory, uffd) = match params.mem_backend.backend_type {
-        MemBackendType::File => (
-            guest_memory_from_file(
-                mem_backend_path,
-                mem_state,
-                track_dirty_pages,
-                vm_resources.machine_config.huge_pages,
+        MemBackendType::File => {
+            if vm_resources.machine_config.huge_pages.is_hugetlbfs() {
+                return Err(RestoreFromSnapshotGuestMemoryError::File(
+                    GuestMemoryFromFileError::HugetlbfsSnapshot,
+                )
+                .into());
+            }
+
+            (
+                guest_memory_from_file(mem_backend_path, mem_state, track_dirty_pages)
+                    .map_err(RestoreFromSnapshotGuestMemoryError::File)?,
+                None,
             )
-            .map_err(RestoreFromSnapshotGuestMemoryError::File)?,
-            None,
-        ),
+        }
         MemBackendType::Uffd => guest_memory_from_uffd(
             mem_backend_path,
             mem_state,
@@ -514,14 +518,15 @@ pub enum GuestMemoryFromFileError {
     File(#[from] std::io::Error),
     /// Failed to restore guest memory: {0}
     Restore(#[from] MemoryError),
+    /// Cannot restore hugetlbfs backed snapshot by mapping the memory file. Please use uffd.
+    HugetlbfsSnapshot,
 }
 
 /// 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::create(
@@ -531,7 +536,6 @@ pub fn guest_memory_from_file(
             .map(|region| (GuestAddress(region.base_address), region.size)),
         MemoryKind::FilePrivate(mem_file),
         track_dirty_pages,
-        huge_pages,
     )?;
     Ok(guest_mem)
 }
@@ -596,9 +600,8 @@ fn create_guest_memory(
             .regions
             .iter()
             .map(|region| (GuestAddress(region.base_address), region.size)),
-        MemoryKind::MapAnonymous,
+        MemoryKind::MapAnonymous(huge_pages),
         track_dirty_pages,
-        huge_pages,
     )?;
     let mut backend_mappings = Vec::with_capacity(guest_memory.num_regions());
     for (mem_region, state_region) in guest_memory.iter().zip(mem_state.regions.iter()) {

src/vmm/src/resources.rs

@@ -471,14 +471,13 @@ impl VmResources {
                 self.machine_config.huge_pages,
             )?
         } else {
-            MemoryKind::MapAnonymous
+            MemoryKind::MapAnonymous(self.machine_config.huge_pages)
         };
 
         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

@@ -36,9 +36,8 @@ pub fn single_region_mem_at(at: u64, size: usize) -> GuestMemoryMmap {
 pub fn multi_region_mem(regions: &[(GuestAddress, usize)]) -> GuestMemoryMmap {
     GuestMemoryMmap::create(
         regions.iter().copied(),
-        MemoryKind::MapAnonymous,
+        MemoryKind::MapAnonymous(HugePageConfig::None),
         false,
-        HugePageConfig::None,
     )
     .expect("Cannot initialize memory")
 }