Wire up huge pages support with snapshot feature

We store the huge pages configuration in the snapshot's vmstate file,
and enforce that a snapshot gets restored with the same hugepages
configuration with which it was taken (for simplicity reasons).

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

Author: roypat

src/vmm/src/persist.rs

@@ -32,7 +32,7 @@ use crate::logger::{info, warn};
 use crate::resources::VmResources;
 use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
-use crate::vmm_config::machine_config::{MachineConfigUpdate, VmConfigError};
+use crate::vmm_config::machine_config::{HugePageConfig, MachineConfigUpdate, VmConfigError};
 use crate::vmm_config::snapshot::{
     CreateSnapshotParams, LoadSnapshotParams, MemBackendType, SnapshotType,
 };
@@ -54,6 +54,8 @@ pub struct VmInfo {
     pub cpu_template: StaticCpuTemplate,
     /// Boot source information.
     pub boot_source: BootSourceConfig,
+    /// Huge page configuration
+    pub huge_pages: HugePageConfig,
 }
 
 impl From<&VmResources> for VmInfo {
@@ -63,6 +65,7 @@ impl From<&VmResources> for VmInfo {
             smt: value.vm_config.smt,
             cpu_template: StaticCpuTemplate::from(&value.vm_config.cpu_template),
             boot_source: value.boot_source_config().clone(),
+            huge_pages: value.vm_config.huge_pages,
         }
     }
 }
@@ -399,7 +402,7 @@ pub fn restore_from_snapshot(
             smt: Some(microvm_state.vm_info.smt),
             cpu_template: Some(microvm_state.vm_info.cpu_template),
             track_dirty_pages: Some(track_dirty_pages),
-            huge_pages: None, // TODO: snapshot integration
+            huge_pages: Some(microvm_state.vm_info.huge_pages),
         })
         .map_err(BuildMicrovmFromSnapshotError::VmUpdateConfig)?;
 
@@ -411,8 +414,13 @@ pub fn restore_from_snapshot(
 
     let (guest_memory, uffd) = match params.mem_backend.backend_type {
         MemBackendType::File => (
-            guest_memory_from_file(mem_backend_path, mem_state, track_dirty_pages)
-                .map_err(RestoreFromSnapshotGuestMemoryError::File)?,
+            guest_memory_from_file(
+                mem_backend_path,
+                mem_state,
+                track_dirty_pages,
+                vm_resources.vm_config.huge_pages,
+            )
+            .map_err(RestoreFromSnapshotGuestMemoryError::File)?,
             None,
         ),
         MemBackendType::Uffd => guest_memory_from_uffd(
@@ -422,6 +430,7 @@ pub fn restore_from_snapshot(
             // We enable the UFFD_FEATURE_EVENT_REMOVE feature only if a balloon device
             // is present in the microVM state.
             microvm_state.device_states.balloon_device.is_some(),
+            vm_resources.vm_config.huge_pages,
         )
         .map_err(RestoreFromSnapshotGuestMemoryError::Uffd)?,
     };
@@ -475,9 +484,11 @@ 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)?;
+    let guest_mem =
+        GuestMemoryMmap::from_state(Some(&mem_file), mem_state, track_dirty_pages, huge_pages)?;
     Ok(guest_mem)
 }
 
@@ -501,8 +512,9 @@ fn guest_memory_from_uffd(
     mem_state: &GuestMemoryState,
     track_dirty_pages: bool,
     enable_balloon: bool,
+    huge_pages: HugePageConfig,
 ) -> Result<(GuestMemoryMmap, Option<Uffd>), GuestMemoryFromUffdError> {
-    let guest_memory = GuestMemoryMmap::from_state(None, mem_state, track_dirty_pages)?;
+    let guest_memory = GuestMemoryMmap::from_state(None, mem_state, track_dirty_pages, huge_pages)?;
 
     let mut uffd_builder = UffdBuilder::new();
 

src/vmm/src/rpc_interface.rs

@@ -1069,6 +1069,7 @@ mod tests {
                 smt: value.vm_config.smt,
                 cpu_template: StaticCpuTemplate::from(&value.vm_config.cpu_template),
                 boot_source: value.boot_source_config().clone(),
+                huge_pages: value.vm_config.huge_pages,
             }
         }
     }

src/vmm/src/vstate/memory.rs

@@ -52,6 +52,8 @@ pub enum MemoryError {
     MemfdSetLen(std::io::Error),
     /// Failed to use madvise to enable THP: {0:?}
     Madvise(std::io::Error),
+    /// Cannot restore hugetlbfs backed snapshot by maping the memory file. Please use uffd.
+    HugetlbfsSnapshot,
 }
 
 /// Defines the interface for snapshotting memory.
@@ -87,6 +89,7 @@ where
         file: Option<&File>,
         state: &GuestMemoryState,
         track_dirty_pages: bool,
+        huge_pages: HugePageConfig,
     ) -> Result<Self, MemoryError>;
 
     /// Describes GuestMemoryMmap through a GuestMemoryState struct.
@@ -227,9 +230,14 @@ impl GuestMemoryExtension for GuestMemoryMmap {
         file: Option<&File>,
         state: &GuestMemoryState,
         track_dirty_pages: bool,
+        huge_pages: HugePageConfig,
     ) -> Result<Self, MemoryError> {
         match file {
             Some(f) => {
+                if huge_pages.is_hugetlbfs() {
+                    return Err(MemoryError::HugetlbfsSnapshot);
+                }
+
                 let regions = state
                     .regions
                     .iter()
@@ -242,15 +250,20 @@ impl GuestMemoryExtension for GuestMemoryMmap {
                     .collect::<Result<Vec<_>, std::io::Error>>()
                     .map_err(MemoryError::FileError)?;
 
-                Self::from_raw_regions_file(regions, track_dirty_pages, false, false)
+                Self::from_raw_regions_file(
+                    regions,
+                    track_dirty_pages,
+                    false,
+                    huge_pages == HugePageConfig::Thp,
+                )
             }
             None => {
                 let regions = state
                     .regions
                     .iter()
                     .map(|r| (GuestAddress(r.base_address), r.size))
                     .collect::<Vec<_>>();
-                Self::from_raw_regions(&regions, track_dirty_pages, HugePageConfig::None)
+                Self::from_raw_regions(&regions, track_dirty_pages, huge_pages)
             }
         }
     }
@@ -512,6 +525,25 @@ pub mod tests {
         }
     }
 
+    #[test]
+    fn test_from_state() {
+        let state = GuestMemoryState {
+            regions: vec![GuestMemoryRegionState {
+                base_address: 0,
+                size: 4096,
+                offset: 0,
+            }],
+        };
+        let file = TempFile::new().unwrap().into_file();
+
+        // No mapping of snapshots that were taken with hugetlbfs enabled
+        let err =
+            GuestMemoryMmap::from_state(Some(&file), &state, false, HugePageConfig::Hugetlbfs2M)
+                .unwrap_err();
+
+        assert!(matches!(err, MemoryError::HugetlbfsSnapshot), "{:?}", err);
+    }
+
     #[test]
     fn test_mark_dirty() {
         let page_size = get_page_size().unwrap();
@@ -691,8 +723,13 @@ pub mod tests {
         let mut memory_file = TempFile::new().unwrap().into_file();
         guest_memory.dump(&mut memory_file).unwrap();
 
-        let restored_guest_memory =
-            GuestMemoryMmap::from_state(Some(&memory_file), &memory_state, false).unwrap();
+        let restored_guest_memory = GuestMemoryMmap::from_state(
+            Some(&memory_file),
+            &memory_state,
+            false,
+            HugePageConfig::None,
+        )
+        .unwrap();
 
         // Check that the region contents are the same.
         let mut restored_region = vec![0u8; page_size * 2];
@@ -750,7 +787,8 @@ pub mod tests {
 
         // We can restore from this because this is the first dirty dump.
         let restored_guest_memory =
-            GuestMemoryMmap::from_state(Some(&file), &memory_state, false).unwrap();
+            GuestMemoryMmap::from_state(Some(&file), &memory_state, false, HugePageConfig::None)
+                .unwrap();
 
         // Check that the region contents are the same.
         let mut restored_region = vec![0u8; region_size];

src/vmm/tests/integration_tests.rs

@@ -16,6 +16,7 @@ use vmm::utilities::mock_resources::{MockVmResources, NOISY_KERNEL_IMAGE};
 use vmm::utilities::test_utils::dirty_tracking_vmm;
 use vmm::utilities::test_utils::{create_vmm, default_vmm, default_vmm_no_boot};
 use vmm::vmm_config::instance_info::{InstanceInfo, VmState};
+use vmm::vmm_config::machine_config::HugePageConfig;
 use vmm::vmm_config::snapshot::{CreateSnapshotParams, SnapshotType};
 use vmm::{DumpCpuConfigError, EventManager, FcExitCode};
 
@@ -242,6 +243,7 @@ fn verify_load_snapshot(snapshot_file: TempFile, memory_file: TempFile) {
         Some(memory_file.as_file()),
         &microvm_state.memory_state,
         false,
+        HugePageConfig::None,
     )
     .unwrap();