vmm: memory_manager: add on-demand snapshot restore via userfaultfd

When memory_restore_mode=ondemand is specified on the restore command,
the memory manager creates a userfaultfd descriptor, registers each
guest RAM range for missing-page fault interception, and spawns a
handler thread that serves page faults from the snapshot file using
UFFDIO_COPY. This avoids reading the entire memory-ranges file into
guest RAM before restore completes.

The handler uses epoll to multiplex the userfaultfd and a stop eventfd
for clean shutdown. Concurrent faults from multiple vCPUs are handled
by treating EEXIST as a benign race and waking blocked threads with
UFFDIO_WAKE. Once all pages have been served the handler exits
automatically. If the handler thread panics the VMM is signalled to
exit since the VM cannot continue without page fault service.

MemoryZone gains a backing_page_size field so the handler resolves
fault granularity from the zone rather than the top-level config.

Errors from the UFFD setup path use a structured UffdError enum
and a new MigratableError::OnDemandRestore variant, with a From
impl to keep call sites concise.

The seccomp filter is updated to allow the userfaultfd syscall and
the four uffd ioctls (UFFDIO_API, UFFDIO_COPY, UFFDIO_REGISTER,
UFFDIO_WAKE) under the VMM thread profile.

Signed-off-by: Shayon Mukherjee <shayonj@gmail.com>

vm-migration/src/lib.rs

@@ -14,6 +14,38 @@ mod bitpos_iterator;
 mod context;
 pub mod protocol;
 
+#[derive(Error, Debug)]
+pub enum UffdError {
+    #[error("Snapshot ranges are not page-aligned")]
+    UnalignedRanges,
+
+    #[error("Failed to create userfaultfd")]
+    Create(#[source] std::io::Error),
+
+    #[error("Cannot translate GPA {gpa:#x} to host address")]
+    GpaTranslation { gpa: u64 },
+
+    #[error("Failed to register region at {addr:#x}+{len:#x}")]
+    Register {
+        addr: u64,
+        len: u64,
+        #[source]
+        source: std::io::Error,
+    },
+
+    #[error("Region at {addr:#x}+{len:#x} missing COPY/WAKE support")]
+    MissingIoctlSupport { addr: u64, len: u64 },
+
+    #[error("Failed to spawn handler thread")]
+    SpawnThread(#[source] std::io::Error),
+
+    #[error("Handler terminated before startup completed")]
+    HandlerStartup,
+
+    #[error("Handler failed after startup")]
+    HandlerFailed(#[source] std::io::Error),
+}
+
 #[derive(Error, Debug)]
 pub enum MigratableError {
     #[error("Failed to pause migratable component")]
@@ -34,6 +66,9 @@ pub enum MigratableError {
     #[error("Failed to receive migratable component snapshot")]
     MigrateReceive(#[source] anyhow::Error),
 
+    #[error("On-demand restore failed")]
+    OnDemandRestore(#[source] UffdError),
+
     #[error("Socket error")]
     MigrateSocket(#[source] std::io::Error),
 

vmm/src/lib.rs

@@ -51,7 +51,7 @@ use crate::api::{
     ApiRequest, ApiResponse, RequestHandler, VmInfoResponse, VmReceiveMigrationData,
     VmSendMigrationData, VmmPingResponse,
 };
-use crate::config::{RestoreConfig, add_to_config};
+use crate::config::{MemoryRestoreMode, RestoreConfig, add_to_config};
 #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
 use crate::coredump::GuestDebuggable;
 use crate::landlock::Landlock;
@@ -88,6 +88,8 @@ mod pci_segment;
 pub mod seccomp_filters;
 mod serial_manager;
 mod sigwinch_listener;
+mod uffd;
+mod userfaultfd;
 pub mod vm;
 pub mod vm_config;
 
@@ -1506,6 +1508,7 @@ impl Vmm {
         source_url: &str,
         vm_config: Arc<Mutex<VmConfig>>,
         prefault: bool,
+        memory_restore_mode: MemoryRestoreMode,
     ) -> std::result::Result<(), VmError> {
         let snapshot = recv_vm_state(source_url).map_err(VmError::Restore)?;
         #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
@@ -1548,6 +1551,7 @@ impl Vmm {
             Some(&snapshot),
             Some(source_url),
             Some(prefault),
+            Some(memory_restore_mode),
         )?;
         self.vm = Some(vm);
 
@@ -1754,6 +1758,7 @@ impl RequestHandler for Vmm {
                         None,
                         None,
                         None,
+                        None,
                     )?;
 
                     self.vm = Some(vm);
@@ -1838,17 +1843,22 @@ impl RequestHandler for Vmm {
             }
         }
 
-        self.vm_restore(source_url, vm_config, restore_cfg.prefault)
-            .map_err(|vm_restore_err| {
-                error!("VM Restore failed: {vm_restore_err:?}");
+        self.vm_restore(
+            source_url,
+            vm_config,
+            restore_cfg.prefault,
+            restore_cfg.memory_restore_mode,
+        )
+        .map_err(|vm_restore_err| {
+            error!("VM Restore failed: {vm_restore_err:?}");
 
-                // Cleanup the VM being created while vm restore
-                if let Err(e) = self.vm_delete() {
-                    return e;
-                }
+            // Cleanup the VM being created while vm restore
+            if let Err(e) = self.vm_delete() {
+                return e;
+            }
 
-                vm_restore_err
-            })
+            vm_restore_err
+        })
     }
 
     #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
@@ -1930,6 +1940,7 @@ impl RequestHandler for Vmm {
             None,
             None,
             None,
+            None,
         )?;
 
         // And we boot it

vmm/src/memory_manager.rs

@@ -7,13 +7,13 @@
 use std::collections::BTreeMap;
 use std::collections::HashMap;
 use std::fs::{File, OpenOptions};
-use std::io::{self};
+use std::io::{self, Read as _, Seek, SeekFrom};
 use std::ops::{BitAnd, Not, Sub};
-#[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
-use std::os::fd::AsFd;
+use std::os::fd::{AsFd, OwnedFd};
 use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicU32, Ordering};
+use std::sync::mpsc::{self, Receiver, SyncSender};
 use std::sync::{Arc, Barrier, Mutex};
 use std::{ffi, result, thread};
 
@@ -44,15 +44,31 @@ use vm_memory::{
 use vm_migration::protocol::{MemoryRange, MemoryRangeTable};
 use vm_migration::{
     Migratable, MigratableError, Pausable, Snapshot, SnapshotData, Snapshottable, Transportable,
+    UffdError,
 };
+use vmm_sys_util::eventfd::EventFd;
 
+use crate::config::MemoryRestoreMode;
 #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
 use crate::coredump::{
     CoredumpMemoryRegion, CoredumpMemoryRegions, DumpState, GuestDebuggableError,
 };
 use crate::migration::url_to_path;
 use crate::vm_config::{HotplugMethod, MemoryConfig, MemoryZoneConfig};
-use crate::{GuestMemoryMmap, GuestRegionMmap, MEMORY_MANAGER_SNAPSHOT_ID};
+use crate::{GuestMemoryMmap, GuestRegionMmap, MEMORY_MANAGER_SNAPSHOT_ID, uffd};
+
+struct UffdHandler {
+    stop_event: EventFd,
+    result_rx: Receiver<Result<(), io::Error>>,
+    handle: thread::JoinHandle<()>,
+}
+
+struct UffdRange {
+    host_addr: u64,
+    length: u64,
+    file_offset: u64,
+    page_size: u64,
+}
 
 pub const MEMORY_MANAGER_ACPI_SIZE: usize = 0x18;
 
@@ -116,13 +132,25 @@ impl VirtioMemZone {
     }
 }
 
-#[derive(Default)]
 pub struct MemoryZone {
     regions: Vec<Arc<GuestRegionMmap>>,
     virtio_mem_zone: Option<VirtioMemZone>,
+    shared: bool,
+    hugepages: bool,
+    backing_page_size: u64,
 }
 
 impl MemoryZone {
+    fn new(shared: bool, hugepages: bool, backing_page_size: u64) -> Self {
+        Self {
+            regions: Vec::new(),
+            virtio_mem_zone: None,
+            shared,
+            hugepages,
+            backing_page_size,
+        }
+    }
+
     pub fn regions(&self) -> &Vec<Arc<GuestRegionMmap>> {
         &self.regions
     }
@@ -132,6 +160,21 @@ impl MemoryZone {
     pub fn virtio_mem_zone_mut(&mut self) -> Option<&mut VirtioMemZone> {
         self.virtio_mem_zone.as_mut()
     }
+
+    fn backing_page_size_for_gpa(&self, gpa: u64) -> Option<u64> {
+        if self.regions.iter().any(|region| {
+            let start = region.start_addr().raw_value();
+            gpa >= start && gpa < start + region.len()
+        }) {
+            return Some(self.backing_page_size);
+        }
+
+        self.virtio_mem_zone.as_ref().and_then(|virtio_mem_zone| {
+            let start = virtio_mem_zone.region.start_addr().raw_value();
+            (gpa >= start && gpa < start + virtio_mem_zone.region.len())
+                .then_some(self.backing_page_size)
+        })
+    }
 }
 
 pub type MemoryZones = HashMap<String, MemoryZone>;
@@ -187,6 +230,7 @@ pub struct MemoryManager {
     // This is useful for getting the dirty pages as we need to know the
     // slots that the mapping is created in.
     guest_ram_mappings: Vec<GuestRamMapping>,
+    uffd_handler: Option<UffdHandler>,
 
     pub acpi_address: Option<GuestAddress>,
     #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
@@ -350,6 +394,12 @@ pub enum Error {
     MisalignedMemorySize,
 }
 
+impl From<UffdError> for Error {
+    fn from(e: UffdError) -> Self {
+        Error::Restore(MigratableError::OnDemandRestore(e))
+    }
+}
+
 const ENABLE_FLAG: usize = 0;
 const INSERTING_FLAG: usize = 1;
 const REMOVING_FLAG: usize = 2;
@@ -551,7 +601,10 @@ impl MemoryManager {
         }
 
         // Add zone id to the list of memory zones.
-        memory_zones.insert(zone.id.clone(), MemoryZone::default());
+        memory_zones.insert(
+            zone.id.clone(),
+            MemoryZone::new(zone.shared, zone.hugepages, zone_align_size),
+        );
 
         for ram_region in ram_regions.iter() {
             let mut ram_region_offset = 0;
@@ -642,7 +695,10 @@ impl MemoryManager {
                         );
                         return Err(Error::DuplicateZoneId);
                     }
-                    memory_zones.insert(zone.id.clone(), MemoryZone::default());
+                    memory_zones.insert(
+                        zone.id.clone(),
+                        MemoryZone::new(zone.shared, zone.hugepages, zone_align_size),
+                    );
                 }
 
                 if ram_region_consumed {
@@ -670,7 +726,11 @@ impl MemoryManager {
         let mut memory_zones = HashMap::new();
 
         for zone_config in zones_config {
-            memory_zones.insert(zone_config.id.clone(), MemoryZone::default());
+            let zone_page_size = memory_zone_get_align_size(zone_config)?;
+            memory_zones.insert(
+                zone_config.id.clone(),
+                MemoryZone::new(zone_config.shared, zone_config.hugepages, zone_page_size),
+            );
         }
 
         for guest_ram_mapping in guest_ram_mappings {
@@ -760,6 +820,360 @@ impl MemoryManager {
         Ok(())
     }
 
+    /// Restore guest memory using userfaultfd for lazy demand paging.
+    ///
+    /// Instead of reading the entire snapshot into guest RAM upfront (which
+    /// blocks restore for hundreds of milliseconds at multi-GB sizes), this
+    /// registers the guest memory regions with a userfaultfd. A background
+    /// thread handles page faults by reading the corresponding page from the
+    /// snapshot file and copying it into guest memory via `UFFDIO_COPY`.
+    ///
+    /// This preserves the original memory mapping type (anonymous or shared),
+    /// making it compatible with VFIO device passthrough and shared-memory
+    /// guest RAM.
+    ///
+    /// Fails the restore if UFFD setup cannot be completed successfully.
+    ///
+    /// The handler thread keeps the snapshot file open while lazy restore
+    /// is active. The file must remain available until the VM is shut down or
+    /// all faulted pages have been served.
+    fn restore_by_uffd(
+        &mut self,
+        file_path: &Path,
+        saved_regions: &MemoryRangeTable,
+        exit_evt: &EventFd,
+    ) -> Result<(), Error> {
+        if saved_regions.is_empty() {
+            return Ok(());
+        }
+
+        let guest_memory = self.guest_memory.memory();
+        let required_uffd_features = self.required_uffd_features();
+
+        // SAFETY: FFI call. Trivially safe.
+        let base_page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as u64;
+
+        info!(
+            "UFFD restore: attempting demand-paged restore for {} region(s)",
+            saved_regions.regions().len()
+        );
+
+        if saved_regions
+            .regions()
+            .iter()
+            .any(|range| range.gpa % base_page_size != 0 || range.length % base_page_size != 0)
+        {
+            return Err(UffdError::UnalignedRanges.into());
+        }
+
+        let snapshot_file = File::open(file_path).map_err(Error::SnapshotOpen)?;
+
+        let uffd_fd = uffd::create(required_uffd_features).map_err(UffdError::Create)?;
+
+        let mut handler_ranges: Vec<UffdRange> = Vec::new();
+        let mut file_offset: u64 = 0;
+
+        for range in saved_regions.regions() {
+            let host_addr = guest_memory
+                .get_host_address(GuestAddress(range.gpa))
+                .map_err(|_| UffdError::GpaTranslation { gpa: range.gpa })?
+                as u64;
+
+            let ioctls = uffd::register(uffd_fd.as_fd(), host_addr, range.length).map_err(|e| {
+                UffdError::Register {
+                    addr: host_addr,
+                    len: range.length,
+                    source: e,
+                }
+            })?;
+
+            if ioctls & crate::userfaultfd::UFFD_API_RANGE_IOCTLS_BASIC
+                != crate::userfaultfd::UFFD_API_RANGE_IOCTLS_BASIC
+            {
+                return Err(UffdError::MissingIoctlSupport {
+                    addr: host_addr,
+                    len: range.length,
+                }
+                .into());
+            }
+
+            let range_page_size = self
+                .memory_zones
+                .values()
+                .find_map(|zone| zone.backing_page_size_for_gpa(range.gpa))
+                .unwrap_or(base_page_size);
+
+            handler_ranges.push(UffdRange {
+                host_addr,
+                length: range.length,
+                file_offset,
+                page_size: range_page_size,
+            });
+
+            file_offset += range.length;
+        }
+
+        info!(
+            "UFFD restore: registered {} region(s), {} total bytes, spawning handler",
+            handler_ranges.len(),
+            file_offset
+        );
+
+        let stop_event = EventFd::new(libc::EFD_NONBLOCK).map_err(Error::EventFdFail)?;
+        let thread_stop_event = stop_event.try_clone().map_err(Error::EventFdFail)?;
+        let thread_exit_evt = exit_evt.try_clone().map_err(Error::EventFdFail)?;
+        let (ready_tx, ready_rx) = mpsc::sync_channel(1);
+        let (result_tx, result_rx) = mpsc::sync_channel(1);
+        let handle = thread::Builder::new()
+            .name("uffd-handler".to_string())
+            .spawn(move || {
+                std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
+                    let max_page_size = handler_ranges
+                        .iter()
+                        .map(|r| r.page_size)
+                        .max()
+                        .unwrap_or(base_page_size);
+                    let result = Self::uffd_handler_loop(
+                        uffd_fd,
+                        thread_stop_event,
+                        snapshot_file,
+                        &handler_ranges,
+                        max_page_size,
+                        &ready_tx,
+                    );
+
+                    if let Err(e) = &result {
+                        error!("UFFD handler exited with error: {e}");
+                    }
+
+                    result_tx.send(result).ok();
+                }))
+                .map_err(|_| {
+                    error!("uffd-handler thread panicked");
+                    thread_exit_evt.write(1).ok();
+                })
+                .ok();
+            })
+            .map_err(UffdError::SpawnThread)?;
+
+        if ready_rx.recv().is_err() {
+            handle.join().ok();
+            return Err(UffdError::HandlerStartup.into());
+        }
+
+        if let Ok(Err(e)) = result_rx.try_recv() {
+            handle.join().ok();
+            return Err(UffdError::HandlerFailed(e).into());
+        }
+
+        self.uffd_handler = Some(UffdHandler {
+            stop_event,
+            result_rx,
+            handle,
+        });
+
+        info!("UFFD restore: demand-paged restore enabled");
+
+        Ok(())
+    }
+
+    fn required_uffd_features(&self) -> u64 {
+        let mut features = 0u64;
+        if self.memory_zones.values().any(|z| z.shared && !z.hugepages) {
+            features |= crate::userfaultfd::UFFD_FEATURE_MISSING_SHMEM;
+        }
+        if self.memory_zones.values().any(|z| z.hugepages) {
+            features |= crate::userfaultfd::UFFD_FEATURE_MISSING_HUGETLBFS;
+        }
+        features
+    }
+
+    fn stop_uffd_handler(&mut self) {
+        if let Some(uffd_handler) = self.uffd_handler.take() {
+            uffd_handler.stop_event.write(1).ok();
+            uffd_handler.handle.join().ok();
+
+            match uffd_handler.result_rx.try_recv() {
+                Ok(Err(e)) => error!("UFFD handler terminated with error: {e}"),
+                Err(mpsc::TryRecvError::Disconnected) => {
+                    warn!("UFFD handler terminated unexpectedly (possible panic)");
+                }
+                _ => {}
+            }
+        }
+    }
+
+    /// Poll the UFFD fd and serve page faults from the snapshot file.
+    ///
+    /// Runs until the fd is closed (EPOLLHUP) or an unrecoverable error occurs.
+    /// Each fault triggers a seek + read from the snapshot file followed by a
+    /// `UFFDIO_COPY` to resolve the fault and wake the faulting thread.
+    #[allow(clippy::needless_pass_by_value)]
+    fn uffd_handler_loop(
+        uffd_fd: OwnedFd,
+        stop_event: EventFd,
+        mut snapshot_file: File,
+        ranges: &[UffdRange],
+        page_size: u64,
+        ready_tx: &SyncSender<()>,
+    ) -> Result<(), io::Error> {
+        let uffd_raw_fd = uffd_fd.as_raw_fd();
+        let mut page_buf = vec![0u8; page_size as usize];
+
+        let total_pages: u64 = ranges.iter().map(|r| r.length.div_ceil(r.page_size)).sum();
+        let mut pages_served: u64 = 0;
+
+        const EVENT_STOP: u64 = 0;
+        const EVENT_UFFD: u64 = 1;
+
+        let epoll_fd = epoll::create(true).map_err(io::Error::other)?;
+        // SAFETY: epoll_fd is valid and owned by this scope.
+        let _epoll_file = unsafe { File::from_raw_fd(epoll_fd) };
+
+        epoll::ctl(
+            epoll_fd,
+            epoll::ControlOptions::EPOLL_CTL_ADD,
+            stop_event.as_raw_fd(),
+            epoll::Event::new(epoll::Events::EPOLLIN, EVENT_STOP),
+        )
+        .map_err(io::Error::other)?;
+
+        epoll::ctl(
+            epoll_fd,
+            epoll::ControlOptions::EPOLL_CTL_ADD,
+            uffd_raw_fd,
+            epoll::Event::new(epoll::Events::EPOLLIN | epoll::Events::EPOLLHUP, EVENT_UFFD),
+        )
+        .map_err(io::Error::other)?;
+
+        ready_tx.send(()).ok();
+
+        let mut events = vec![epoll::Event::new(epoll::Events::empty(), 0); 2];
+        loop {
+            let num_events = match epoll::wait(epoll_fd, -1, &mut events) {
+                Ok(n) => n,
+                Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
+                Err(e) => return Err(e),
+            };
+
+            let mut got_uffd_data = false;
+            for event in events.iter().take(num_events) {
+                let token = event.data;
+                let evt_flags = event.events;
+
+                if token == EVENT_STOP {
+                    stop_event.read().ok();
+                    info!("UFFD handler: received stop event, exiting");
+                    return Ok(());
+                }
+
+                if token == EVENT_UFFD
+                    && (evt_flags & epoll::Events::EPOLLHUP.bits()) != 0
+                    && (evt_flags & epoll::Events::EPOLLIN.bits()) == 0
+                {
+                    info!("UFFD handler: fd closed (EPOLLHUP), exiting");
+                    return Ok(());
+                }
+
+                if token == EVENT_UFFD && (evt_flags & epoll::Events::EPOLLIN.bits()) != 0 {
+                    got_uffd_data = true;
+                }
+            }
+
+            if !got_uffd_data {
+                continue;
+            }
+
+            // SAFETY: UffdMsg is a plain repr(C) struct, safe to zero-init.
+            let mut msg: uffd::UffdMsg = unsafe { std::mem::zeroed() };
+            // SAFETY: reading a uffd_msg-sized struct from the valid uffd fd.
+            let n = unsafe {
+                libc::read(
+                    uffd_raw_fd,
+                    &mut msg as *mut uffd::UffdMsg as *mut libc::c_void,
+                    std::mem::size_of::<uffd::UffdMsg>(),
+                )
+            };
+            if n < 0 {
+                let err = io::Error::last_os_error();
+                if err.kind() == io::ErrorKind::WouldBlock {
+                    continue;
+                }
+                return Err(err);
+            }
+            if n == 0 {
+                info!("UFFD handler: EOF on fd, exiting");
+                return Ok(());
+            }
+            if n as usize != std::mem::size_of::<uffd::UffdMsg>() {
+                return Err(io::Error::new(
+                    io::ErrorKind::UnexpectedEof,
+                    "Short read from userfaultfd",
+                ));
+            }
+
+            if msg.event != crate::userfaultfd::UFFD_EVENT_PAGEFAULT {
+                continue;
+            }
+
+            let fault_addr = msg.pf_address;
+
+            let mut served = false;
+            for range in ranges {
+                // Round down to the page boundary containing the faulted address.
+                let page_addr = fault_addr & !(range.page_size - 1);
+                if page_addr >= range.host_addr && page_addr < range.host_addr + range.length {
+                    let offset_in_range = page_addr - range.host_addr;
+                    let file_pos = range.file_offset + offset_in_range;
+
+                    snapshot_file.seek(SeekFrom::Start(file_pos))?;
+                    snapshot_file.read_exact(&mut page_buf[..range.page_size as usize])?;
+
+                    loop {
+                        match uffd::copy(
+                            uffd_fd.as_fd(),
+                            page_addr,
+                            page_buf.as_ptr(),
+                            range.page_size,
+                        ) {
+                            Ok(()) => {
+                                pages_served += 1;
+                                break;
+                            }
+                            Err(e) if e.raw_os_error() == Some(libc::EEXIST) => {
+                                if let Err(e) =
+                                    uffd::wake(uffd_fd.as_fd(), page_addr, range.page_size)
+                                {
+                                    warn!("UFFDIO_WAKE failed at {page_addr:#x}: {e}");
+                                }
+                                break;
+                            }
+                            Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => {
+                                // The kernel can report a transient EAGAIN while the fault
+                                // is being resolved; yield and retry instead of aborting restore.
+                                thread::yield_now();
+                            }
+                            Err(e) => return Err(e),
+                        }
+                    }
+                    served = true;
+                    break;
+                }
+            }
+
+            if !served {
+                return Err(io::Error::other(format!(
+                    "UFFD handler: fault at {fault_addr:#x} does not belong to any registered range",
+                )));
+            }
+
+            if pages_served == total_pages {
+                info!("UFFD handler: all {pages_served} pages served, exiting");
+                return Ok(());
+            }
+        }
+    }
+
     fn validate_memory_config(
         config: &MemoryConfig,
         user_provided_zones: bool,
@@ -1227,6 +1641,7 @@ impl MemoryManager {
             snapshot_memory_ranges: MemoryRangeTable::default(),
             memory_zones,
             guest_ram_mappings: Vec::new(),
+            uffd_handler: None,
             acpi_address,
             log_dirty: dynamic, // Cannot log dirty pages on a TD
             arch_mem_regions,
@@ -1240,13 +1655,16 @@ impl MemoryManager {
         Ok(Arc::new(Mutex::new(memory_manager)))
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn new_from_snapshot(
         snapshot: &Snapshot,
         vm: Arc<dyn hypervisor::Vm>,
         config: &MemoryConfig,
         source_url: Option<&str>,
         prefault: bool,
+        memory_restore_mode: MemoryRestoreMode,
         phys_bits: u8,
+        exit_evt: &EventFd,
     ) -> Result<Arc<Mutex<MemoryManager>>, Error> {
         if let Some(source_url) = source_url {
             let mut memory_file_path = url_to_path(source_url).map_err(Error::Restore)?;
@@ -1266,9 +1684,17 @@ impl MemoryManager {
                 Default::default(),
             )?;
 
-            mm.lock()
-                .unwrap()
-                .fill_saved_regions(memory_file_path, &mem_snapshot.memory_ranges)?;
+            if memory_restore_mode == MemoryRestoreMode::OnDemand {
+                mm.lock().unwrap().restore_by_uffd(
+                    &memory_file_path,
+                    &mem_snapshot.memory_ranges,
+                    exit_evt,
+                )?;
+            } else {
+                mm.lock()
+                    .unwrap()
+                    .fill_saved_regions(memory_file_path, &mem_snapshot.memory_ranges)?;
+            }
 
             Ok(mm)
         } else {
@@ -2530,6 +2956,12 @@ impl Aml for MemoryManager {
 
 impl Pausable for MemoryManager {}
 
+impl Drop for MemoryManager {
+    fn drop(&mut self) {
+        self.stop_uffd_handler();
+    }
+}
+
 #[derive(Clone, Serialize, Deserialize)]
 pub struct MemoryManagerSnapshotData {
     memory_ranges: MemoryRangeTable,

vmm/src/seccomp_filters.rs

@@ -25,6 +25,8 @@ use vhost::vhost_kern::vhost_binding::{
     VHOST_VDPA_SET_STATUS, VHOST_VDPA_SET_VRING_ENABLE, VHOST_VDPA_SUSPEND,
 };
 
+use crate::userfaultfd::{UFFDIO_API, UFFDIO_COPY, UFFDIO_REGISTER, UFFDIO_WAKE};
+
 #[derive(Copy, Clone)]
 pub enum Thread {
     HttpApi,
@@ -362,6 +364,10 @@ fn create_vmm_ioctl_seccomp_rule_common(
             VHOST_VDPA_GET_CONFIG_SIZE()
         )?],
         and![Cond::new(1, ArgLen::Dword, Eq, VHOST_VDPA_SUSPEND())?],
+        and![Cond::new(1, ArgLen::Dword, Eq, UFFDIO_API)?],
+        and![Cond::new(1, ArgLen::Dword, Eq, UFFDIO_COPY)?],
+        and![Cond::new(1, ArgLen::Dword, Eq, UFFDIO_REGISTER)?],
+        and![Cond::new(1, ArgLen::Dword, Eq, UFFDIO_WAKE)?],
     ];
 
     let hypervisor_rules = create_vmm_ioctl_seccomp_rule_hypervisor(hypervisor_type)?;
@@ -691,6 +697,7 @@ fn vmm_thread_rules(
         (libc::SYS_unlink, vec![]),
         #[cfg(target_arch = "aarch64")]
         (libc::SYS_unlinkat, vec![]),
+        (libc::SYS_userfaultfd, vec![]),
         (libc::SYS_wait4, vec![]),
         (libc::SYS_write, vec![]),
         (libc::SYS_writev, vec![]),

vmm/src/vm.rs

@@ -76,7 +76,7 @@ use vm_migration::{
 use vmm_sys_util::eventfd::EventFd;
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
 
-use crate::config::{ValidationError, add_to_config};
+use crate::config::{MemoryRestoreMode, ValidationError, add_to_config};
 use crate::console_devices::{ConsoleDeviceError, ConsoleInfo};
 #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
 use crate::coredump::{
@@ -1265,6 +1265,7 @@ impl Vm {
         snapshot: Option<&Snapshot>,
         source_url: Option<&str>,
         prefault: Option<bool>,
+        memory_restore_mode: Option<MemoryRestoreMode>,
     ) -> Result<Self> {
         trace_scoped!("Vm::new");
 
@@ -1300,8 +1301,10 @@ impl Vm {
                     vm.clone(),
                     &vm_config.lock().unwrap().memory.clone(),
                     source_url,
-                    prefault.unwrap(),
+                    prefault.unwrap_or(false),
+                    memory_restore_mode.unwrap_or_default(),
                     phys_bits,
+                    &exit_evt,
                 )
                 .map_err(Error::MemoryManager)?
             } else {