Snapshot restore for swiotlb regions

Support restoring VMs with swiotlb regions. For this, untangle the uffd
handshake from the actual restoration of memory regions, as we first
need to restore and register to the Vm _all_ memory regions, before we
can then send a single handshake containing both normal and swiotlb
regions to the Uffd handler.

While we're here, significantly simplify the jungle of error types.

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

Author: roypat

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

@@ -30,9 +30,7 @@ use crate::cpu_config::templates::CustomCpuTemplate;
 use crate::initrd::InitrdConfig;
 use crate::utils::{align_up, usize_to_u64};
 use crate::vmm_config::machine_config::MachineConfig;
-use crate::vstate::memory::{
-    Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap,
-};
+use crate::vstate::memory::{Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap};
 use crate::vstate::vcpu::KvmVcpuError;
 use crate::{Vcpu, VcpuConfig, Vmm};
 
@@ -118,7 +116,9 @@ pub fn configure_system_for_boot(
 
     let swiotlb_region = match vmm.vm.swiotlb_regions().num_regions() {
         0 | 1 => vmm.vm.swiotlb_regions().iter().next(),
-        _ => panic!("Firecracker tried to configure more than one swiotlb region. This is a logic bug.")
+        _ => panic!(
+            "Firecracker tried to configure more than one swiotlb region. This is a logic bug."
+        ),
     };
 
     let fdt = fdt::create_fdt(

src/vmm/src/builder.rs

@@ -4,16 +4,16 @@
 //! Enables pre-boot setup, instantiation and booting of a Firecracker VMM.
 
 use std::fmt::Debug;
-use std::io;
+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;
 use utils::time::TimestampUs;
+use vm_memory::GuestMemoryRegion;
 #[cfg(target_arch = "aarch64")]
 use vm_superio::Rtc;
 use vm_superio::Serial;
@@ -50,14 +50,17 @@ 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::{
+    MicrovmState, MicrovmStateError, RestoreMemoryError, SnapshotStateFromFileError,
+    restore_memory, 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::vcpu::{Vcpu, VcpuError};
 use crate::vstate::vm::Vm;
 use crate::{EventManager, Vmm, VmmError, device_manager};
@@ -393,6 +396,8 @@ pub enum BuildMicrovmFromSnapshotError {
     SetTsc(#[from] crate::arch::SetTscError),
     /// Failed to restore microVM state: {0}
     RestoreState(#[from] crate::vstate::vm::ArchVmError),
+    /// Failed to get snapshot state from file: {0}
+    LoadState(#[from] SnapshotStateFromFileError),
     /// Failed to update microVM configuration: {0}
     VmUpdateConfig(#[from] MachineConfigError),
     /// Failed to restore MMIO device: {0}
@@ -413,19 +418,19 @@ pub enum BuildMicrovmFromSnapshotError {
     ACPIDeviManager(#[from] ACPIDeviceManagerRestoreError),
     /// VMGenID update failed: {0}
     VMGenIDUpdate(std::io::Error),
+    /// Failed to restore guest memory: {0}
+    Memory(#[from] RestoreMemoryError),
 }
 
 /// Builds and starts a microVM based on the provided MicrovmState.
 ///
 /// 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> {
@@ -439,11 +444,63 @@ pub fn build_microvm_from_snapshot(
     )
     .map_err(StartMicrovmError::Internal)?;
 
+    let track_dirty_pages = vm_resources.machine_config.track_dirty_pages;
+    let huge_pages = vm_resources.machine_config.huge_pages;
+
+    let mem_backend_path = &mem_backend.backend_path;
+
+    let mem_file = match mem_backend.backend_type {
+        MemBackendType::File => Some(mem_backend_path),
+        MemBackendType::Uffd => None,
+    };
+
+    let guest_memory = restore_memory(
+        &microvm_state.vm_state.memory,
+        mem_file,
+        huge_pages,
+        track_dirty_pages,
+        0,
+    )?;
+    let io_memory = restore_memory(
+        &microvm_state.vm_state.io_memory,
+        mem_file,
+        huge_pages,
+        track_dirty_pages,
+        guest_memory.iter().map(|r| r.len()).sum(),
+    )?;
+
     vmm.vm
         .register_memory_regions(guest_memory)
         .map_err(VmmError::Vm)
         .map_err(StartMicrovmError::Internal)?;
-    vmm.uffd = uffd;
+
+    for region in io_memory {
+        vmm.vm
+            .register_swiotlb_region(region)
+            .map_err(VmmError::Vm)
+            .map_err(StartMicrovmError::Internal)?;
+    }
+
+    vmm.uffd = match mem_backend.backend_type {
+        MemBackendType::File => None,
+        MemBackendType::Uffd => {
+            let (uffd, mut mappings) = vmm
+                .vm
+                .create_uffd()
+                .map_err(RestoreMemoryError::UffdCreate)?;
+
+            #[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)
+                .map_err(RestoreMemoryError::UffdHandshake)?;
+
+            Some(uffd)
+        }
+    };
 
     #[cfg(target_arch = "x86_64")]
     {

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

@@ -478,6 +478,7 @@ pub(crate) mod tests {
                 libc::MAP_PRIVATE,
                 Some(file),
                 false,
+                0,
             )
             .unwrap()
             .into_iter()