wip: refactor: break apart guest mem restoration

Do it in two steps: First restore the mappings, and then in a second
step after registering all the mapping to the VM, create the uffd and
send the handshake.

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

Author: roypat

src/vmm/src/builder.rs

@@ -4,15 +4,17 @@
 //! Enables pre-boot setup, instantiation and booting of a Firecracker VMM.
 
 use std::fmt::Debug;
-use std::io;
+use std::fs::File;
+use std::io::{self};
 #[cfg(feature = "gdb")]
 use std::sync::mpsc;
 use std::sync::{Arc, Mutex};
 
 use event_manager::{MutEventSubscriber, SubscriberOps};
 use libc::EFD_NONBLOCK;
 use linux_loader::cmdline::Cmdline as LoaderKernelCmdline;
-use userfaultfd::Uffd;
+#[cfg(target_arch = "aarch64")]
+use linux_loader::loader::pe::PE as Loader;
 use utils::time::TimestampUs;
 #[cfg(target_arch = "aarch64")]
 use vm_superio::Rtc;
@@ -50,14 +52,19 @@ use crate::devices::virtio::vsock::{Vsock, VsockUnixBackend};
 use crate::gdb;
 use crate::initrd::{InitrdConfig, InitrdError};
 use crate::logger::{debug, error};
-use crate::persist::{MicrovmState, MicrovmStateError};
+use crate::persist::{
+    GuestMemoryFromFileError, MicrovmState, MicrovmStateError, RestoreFromSnapshotGuestMemoryError,
+    send_uffd_handshake,
+};
 use crate::resources::VmResources;
 use crate::seccomp::BpfThreadMap;
 use crate::snapshot::Persist;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::MachineConfigError;
+use crate::vmm_config::snapshot::{MemBackendConfig, MemBackendType};
 use crate::vstate::kvm::Kvm;
-use crate::vstate::memory::GuestRegionMmap;
+use crate::vstate::memory;
+use crate::vstate::memory::GuestMemory;
 use crate::vstate::vcpu::{Vcpu, VcpuError};
 use crate::vstate::vm::Vm;
 use crate::{EventManager, Vmm, VmmError, device_manager};
@@ -409,13 +416,11 @@ pub enum BuildMicrovmFromSnapshotError {
 ///
 /// An `Arc` reference of the built `Vmm` is also plugged in the `EventManager`, while another
 /// is returned.
-#[allow(clippy::too_many_arguments)]
 pub fn build_microvm_from_snapshot(
     instance_info: &InstanceInfo,
     event_manager: &mut EventManager,
     microvm_state: MicrovmState,
-    guest_memory: Vec<GuestRegionMmap>,
-    uffd: Option<Uffd>,
+    mem_backend: &MemBackendConfig,
     seccomp_filters: &BpfThreadMap,
     vm_resources: &mut VmResources,
 ) -> Result<Arc<Mutex<Vmm>>, BuildMicrovmFromSnapshotError> {
@@ -429,11 +434,54 @@ pub fn build_microvm_from_snapshot(
     )
     .map_err(StartMicrovmError::Internal)?;
 
+    let track_dirty_pages = vm_resources.machine_config.track_dirty_pages;
+
+    let mem_backend_path = &mem_backend.backend_path;
+    let mem_state = &microvm_state.vm_state.memory;
+
+    let guest_memory = match mem_backend.backend_type {
+        MemBackendType::File => {
+            if vm_resources.machine_config.huge_pages.is_hugetlbfs() {
+                panic!(
+                    "{:?}",
+                    RestoreFromSnapshotGuestMemoryError::File(
+                        GuestMemoryFromFileError::HugetlbfsSnapshot,
+                    )
+                );
+            }
+
+            let mem_file = File::open(mem_backend_path).unwrap();
+            memory::snapshot_file(mem_file, mem_state.regions(), track_dirty_pages).unwrap()
+        }
+        MemBackendType::Uffd => memory::anonymous(
+            mem_state.regions(),
+            track_dirty_pages,
+            vm_resources.machine_config.huge_pages,
+        )
+        .unwrap(),
+    };
+
     vmm.vm
         .register_memory_regions(guest_memory)
         .map_err(VmmError::Vm)
         .map_err(StartMicrovmError::Internal)?;
-    vmm.uffd = uffd;
+
+    vmm.uffd = match mem_backend.backend_type {
+        MemBackendType::File => None,
+        MemBackendType::Uffd => {
+            let (uffd, mut mappings) = vmm.vm.create_uffd().unwrap();
+
+            #[allow(deprecated)]
+            mappings.iter_mut().for_each(|mapping| {
+                mapping.page_size = vm_resources.machine_config.huge_pages.page_size();
+                mapping.page_size_kib = vm_resources.machine_config.huge_pages.page_size();
+            });
+
+            send_uffd_handshake(mem_backend_path, &mappings, &uffd).unwrap();
+
+            Some(uffd)
+        }
+    };
 
     #[cfg(target_arch = "x86_64")]
     {

src/vmm/src/persist.rs

@@ -13,7 +13,6 @@ use std::sync::{Arc, Mutex};
 
 use semver::Version;
 use serde::{Deserialize, Serialize};
-use userfaultfd::{FeatureFlags, Uffd, UffdBuilder};
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
 
 #[cfg(target_arch = "aarch64")]
@@ -33,10 +32,9 @@ use crate::utils::u64_to_usize;
 use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::{HugePageConfig, MachineConfigError, MachineConfigUpdate};
-use crate::vmm_config::snapshot::{CreateSnapshotParams, LoadSnapshotParams, MemBackendType};
+use crate::vmm_config::snapshot::{CreateSnapshotParams, LoadSnapshotParams};
 use crate::vstate::kvm::KvmState;
-use crate::vstate::memory;
-use crate::vstate::memory::{GuestMemoryState, GuestRegionMmap, MemoryError};
+use crate::vstate::memory::MemoryError;
 use crate::vstate::vcpu::{VcpuSendEventError, VcpuState};
 use crate::vstate::vm::VmState;
 use crate::{EventManager, Vmm, vstate};
@@ -93,7 +91,7 @@ pub struct MicrovmState {
 /// E.g. Guest memory contents for a region of `size` bytes can be found in the
 /// backend at `offset` bytes from the beginning, and should be copied/populated
 /// into `base_host_address`.
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
+#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
 pub struct GuestRegionUffdMapping {
     /// Base host virtual address where the guest memory contents for this
     /// region should be copied/populated.
@@ -373,37 +371,11 @@ pub fn restore_from_snapshot(
     // Some sanity checks before building the microvm.
     snapshot_state_sanity_check(&microvm_state)?;
 
-    let mem_backend_path = &params.mem_backend.backend_path;
-    let mem_state = &microvm_state.vm_state.memory;
-
-    let (guest_memory, uffd) = match params.mem_backend.backend_type {
-        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,
-            )
-        }
-        MemBackendType::Uffd => guest_memory_from_uffd(
-            mem_backend_path,
-            mem_state,
-            track_dirty_pages,
-            vm_resources.machine_config.huge_pages,
-        )
-        .map_err(RestoreFromSnapshotGuestMemoryError::Uffd)?,
-    };
     builder::build_microvm_from_snapshot(
         instance_info,
         event_manager,
         microvm_state,
-        guest_memory,
-        uffd,
+        &params.mem_backend,
         seccomp_filters,
         vm_resources,
     )
@@ -448,16 +420,6 @@ pub enum GuestMemoryFromFileError {
     HugetlbfsSnapshot,
 }
 
-fn guest_memory_from_file(
-    mem_file_path: &Path,
-    mem_state: &GuestMemoryState,
-    track_dirty_pages: bool,
-) -> Result<Vec<GuestRegionMmap>, GuestMemoryFromFileError> {
-    let mem_file = File::open(mem_file_path)?;
-    let guest_mem = memory::snapshot_file(mem_file, mem_state.regions(), track_dirty_pages)?;
-    Ok(guest_mem)
-}
-
 /// Error type for [`guest_memory_from_uffd`]
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum GuestMemoryFromUffdError {
@@ -473,60 +435,7 @@ pub enum GuestMemoryFromUffdError {
     Send(#[from] vmm_sys_util::errno::Error),
 }
 
-fn guest_memory_from_uffd(
-    mem_uds_path: &Path,
-    mem_state: &GuestMemoryState,
-    track_dirty_pages: bool,
-    huge_pages: HugePageConfig,
-) -> Result<(Vec<GuestRegionMmap>, Option<Uffd>), GuestMemoryFromUffdError> {
-    let (guest_memory, backend_mappings) =
-        create_guest_memory(mem_state, track_dirty_pages, huge_pages)?;
-
-    let mut uffd_builder = UffdBuilder::new();
-
-    uffd_builder.require_features(FeatureFlags::EVENT_REMOVE);
-
-    let uffd = uffd_builder
-        .close_on_exec(true)
-        .non_blocking(true)
-        .user_mode_only(false)
-        .create()
-        .map_err(GuestMemoryFromUffdError::Create)?;
-
-    for mem_region in guest_memory.iter() {
-        uffd.register(mem_region.as_ptr().cast(), mem_region.size() as _)
-            .map_err(GuestMemoryFromUffdError::Register)?;
-    }
-
-    send_uffd_handshake(mem_uds_path, &backend_mappings, &uffd)?;
-
-    Ok((guest_memory, Some(uffd)))
-}
-
-fn create_guest_memory(
-    mem_state: &GuestMemoryState,
-    track_dirty_pages: bool,
-    huge_pages: HugePageConfig,
-) -> Result<(Vec<GuestRegionMmap>, Vec<GuestRegionUffdMapping>), GuestMemoryFromUffdError> {
-    let guest_memory = memory::anonymous(mem_state.regions(), track_dirty_pages, huge_pages)?;
-    let mut backend_mappings = Vec::with_capacity(guest_memory.len());
-    let mut offset = 0;
-    for mem_region in guest_memory.iter() {
-        #[allow(deprecated)]
-        backend_mappings.push(GuestRegionUffdMapping {
-            base_host_virt_addr: mem_region.as_ptr() as u64,
-            size: mem_region.size(),
-            offset,
-            page_size: huge_pages.page_size(),
-            page_size_kib: huge_pages.page_size(),
-        });
-        offset += mem_region.size() as u64;
-    }
-
-    Ok((guest_memory, backend_mappings))
-}
-
-fn send_uffd_handshake(
+pub(crate) fn send_uffd_handshake(
     mem_uds_path: &Path,
     backend_mappings: &[GuestRegionUffdMapping],
     uffd: &impl AsRawFd,
@@ -595,7 +504,6 @@ mod tests {
     use crate::vmm_config::balloon::BalloonDeviceConfig;
     use crate::vmm_config::net::NetworkInterfaceConfig;
     use crate::vmm_config::vsock::tests::default_config;
-    use crate::vstate::memory::GuestMemoryRegionState;
 
     fn default_vmm_with_devices() -> Vmm {
         let mut event_manager = EventManager::new().expect("Cannot create EventManager");
@@ -692,24 +600,6 @@ mod tests {
         )
     }
 
-    #[test]
-    fn test_create_guest_memory() {
-        let mem_state = GuestMemoryState {
-            regions: vec![GuestMemoryRegionState {
-                base_address: 0,
-                size: 0x20000,
-            }],
-        };
-
-        let (_, uffd_regions) =
-            create_guest_memory(&mem_state, false, HugePageConfig::None).unwrap();
-
-        assert_eq!(uffd_regions.len(), 1);
-        assert_eq!(uffd_regions[0].size, 0x20000);
-        assert_eq!(uffd_regions[0].offset, 0);
-        assert_eq!(uffd_regions[0].page_size, HugePageConfig::None.page_size());
-    }
-
     #[test]
     fn test_send_uffd_handshake() {
         #[allow(deprecated)]

src/vmm/src/vstate/vm.rs

@@ -12,13 +12,15 @@ use std::path::Path;
 use std::sync::Arc;
 
 use kvm_ioctls::VmFd;
+use userfaultfd::{FeatureFlags, Uffd, UffdBuilder};
 use vmm_sys_util::eventfd::EventFd;
 
 pub use crate::arch::{ArchVm as Vm, ArchVmError, VmState};
 use crate::logger::info;
 use crate::persist::CreateSnapshotError;
 use crate::utils::u64_to_usize;
 use crate::vmm_config::snapshot::SnapshotType;
+use crate::persist::GuestRegionUffdMapping;
 use crate::vstate::memory::{
     GuestMemory, GuestMemoryExtension, GuestMemoryMmap, GuestMemoryRegion, GuestRegionMmap,
     KvmRegion,
@@ -185,6 +187,37 @@ impl Vm {
         &self.common.guest_memory
     }
 
+    pub(crate) fn create_uffd(
+        &self,
+    ) -> Result<(Uffd, Vec<GuestRegionUffdMapping>), userfaultfd::Error> {
+        let mut uffd_builder = UffdBuilder::new();
+        let mut mappings = Vec::new();
+
+        uffd_builder.require_features(FeatureFlags::EVENT_REMOVE);
+
+        let uffd = uffd_builder
+            .close_on_exec(true)
+            .non_blocking(true)
+            .user_mode_only(false)
+            .create()?;
+
+        let mut offset = 0;
+
+        for mem_region in self.common.guest_memory.iter() {
+            uffd.register(mem_region.as_ptr().cast(), mem_region.size() as _)?;
+            mappings.push(GuestRegionUffdMapping {
+                base_host_virt_addr: mem_region.as_ptr() as u64,
+                size: mem_region.size(),
+                offset,
+                ..Default::default()
+            });
+
+            offset += mem_region.size() as u64;
+        }
+
+        Ok((uffd, mappings))
+    }
+
     /// Resets the KVM dirty bitmap for each of the guest's memory regions.
     pub fn reset_dirty_bitmap(&self) {
         self.guest_memory().iter().for_each(|region| {