vmm: Prefault snapshot pages in background

Userfaultfd is a great mechanism for providing fast restore to Cloud
Hypervisor VMs. But that means the price to pay for bringing pages in
happens at runtime, which might slow down the guest when it's touching
pages which haven't been brought in yet.

By prefaulting the pages in the background, we're trying to get the best
of both worlds. That means we still get a very fast restore with the
uffd handler, but within a few seconds (depending on VM's RAM size), we
also get the pages fully faulted and we can stop the uffd handler thread
at that point.

Signed-off-by: Sebastien Boeuf <sboeuf@meta.com>
Assisted-by: Claude:claude-opus-4-7

vmm/src/memory_manager.rs

@@ -7,7 +7,8 @@
 use std::collections::BTreeMap;
 use std::collections::HashMap;
 use std::fs::{File, OpenOptions};
-use std::io::{self, Read as _, Seek, SeekFrom};
+use std::io::{self, Seek, SeekFrom};
+use std::num::NonZeroUsize;
 use std::ops::{BitAnd, Not, Sub};
 use std::os::fd::{AsFd, BorrowedFd, OwnedFd};
 use std::os::unix::fs::FileExt;
@@ -107,6 +108,27 @@ struct UffdRange {
     page_size: u64,
 }
 
+impl UffdRange {
+    fn num_pages(&self) -> u64 {
+        self.length.div_ceil(self.page_size)
+    }
+
+    fn page_addr(&self, page_idx: u64) -> u64 {
+        self.host_addr + page_idx * self.page_size
+    }
+
+    fn file_pos(&self, page_idx: u64) -> u64 {
+        self.file_offset + page_idx * self.page_size
+    }
+
+    /// Returns the page index containing `addr` if it falls within this range.
+    fn page_index_of(&self, addr: u64) -> Option<u64> {
+        let page_addr = addr & !(self.page_size - 1);
+        (page_addr >= self.host_addr && page_addr < self.host_addr + self.length)
+            .then(|| (page_addr - self.host_addr) / self.page_size)
+    }
+}
+
 pub const MEMORY_MANAGER_ACPI_SIZE: usize = 0x18;
 
 const DEFAULT_MEMORY_ZONE: &str = "mem0";
@@ -1101,16 +1123,19 @@ impl MemoryManager {
         }
     }
 
-    /// Poll the UFFD fd and serve page faults from the snapshot file.
+    /// Poll the UFFD fd and serve page faults from the snapshot file, while
+    /// opportunistically prefaulting the remaining pages in between faults.
     ///
-    /// 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.
+    /// Runs until the fd is closed (EPOLLHUP), `stop_event` fires, or an
+    /// unrecoverable error occurs. Each fault triggers a read from the
+    /// snapshot file followed by a `UFFDIO_COPY` to resolve the fault and
+    /// wake the faulting thread. When no fault is pending, one prefault
+    /// page is copied per loop iteration.
     #[allow(clippy::needless_pass_by_value)]
     fn uffd_handler_loop(
         uffd_fd: OwnedFd,
         stop_event: EventFd,
-        mut snapshot_file: File,
+        snapshot_file: File,
         ranges: &[UffdRange],
         page_size: u64,
         ready_tx: &SyncSender<()>,
@@ -1118,8 +1143,28 @@ impl MemoryManager {
         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 total_pages: u64 = ranges.iter().map(UffdRange::num_pages).sum();
         let mut pages_served: u64 = 0;
+        let mut pages_prefaulted: u64 = 0;
+
+        // Per-range bitmap tracking which pages have been populated by the
+        // on-demand fault handler. Lets the prefault cursor skip them without
+        // doing a wasted file read + UFFDIO_COPY.
+        let served_bitmap: Vec<AtomicBitmap> = ranges
+            .iter()
+            .map(|r| {
+                AtomicBitmap::new(
+                    r.length as usize,
+                    NonZeroUsize::new(r.page_size as usize).unwrap(),
+                )
+            })
+            .collect();
+
+        // Prefault cursor: (range index, page index within range). `None`
+        // means prefault was given up due to an error (natural completion
+        // returns from the function instead).
+        let mut prefault_cursor: Option<(usize, u64)> = (!ranges.is_empty()).then_some((0, 0));
+        let prefault_start = std::time::Instant::now();
 
         const EVENT_STOP: u64 = 0;
         const EVENT_UFFD: u64 = 1;
@@ -1148,7 +1193,10 @@ impl MemoryManager {
 
         let mut events = vec![epoll::Event::new(epoll::Events::empty(), 0); 2];
         loop {
-            let num_events = match epoll::wait(epoll_fd, -1, &mut events) {
+            // Block only when prefault is done; otherwise poll non-blocking
+            // so we can advance prefault between faults.
+            let timeout = if prefault_cursor.is_some() { 0 } else { -1 };
+            let num_events = match epoll::wait(epoll_fd, timeout, &mut events) {
                 Ok(n) => n,
                 Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
                 Err(e) => return Err(e),
@@ -1178,54 +1226,51 @@ impl MemoryManager {
                 }
             }
 
-            if !got_uffd_data {
-                continue;
-            }
+            if got_uffd_data {
+                // 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,
+                        (&raw mut msg).cast(),
+                        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",
+                    ));
+                }
 
-            // 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,
-                    (&raw mut msg).cast(),
-                    std::mem::size_of::<uffd::UffdMsg>(),
-                )
-            };
-            if n < 0 {
-                let err = io::Error::last_os_error();
-                if err.kind() == io::ErrorKind::WouldBlock {
+                if msg.event != crate::userfaultfd::UFFD_EVENT_PAGEFAULT {
                     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 fault_addr = msg.pf_address;
+                let mut served = false;
+                for (range_idx, range) in ranges.iter().enumerate() {
+                    let Some(page_idx) = range.page_index_of(fault_addr) else {
+                        continue;
+                    };
+                    let page_addr = range.page_addr(page_idx);
+                    let file_pos = range.file_pos(page_idx);
 
-            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])?;
+                    snapshot_file
+                        .read_exact_at(&mut page_buf[..range.page_size as usize], file_pos)?;
 
                     loop {
                         match uffd::copy(
@@ -1236,6 +1281,7 @@ impl MemoryManager {
                         ) {
                             Ok(()) => {
                                 pages_served += 1;
+                                served_bitmap[range_idx].set_bit(page_idx as usize);
                                 break;
                             }
                             Err(e) if e.raw_os_error() == Some(libc::EEXIST) => {
@@ -1244,6 +1290,7 @@ impl MemoryManager {
                                 {
                                     warn!("UFFDIO_WAKE failed at {page_addr:#x}: {e}");
                                 }
+                                served_bitmap[range_idx].set_bit(page_idx as usize);
                                 break;
                             }
                             Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => {
@@ -1257,18 +1304,104 @@ impl MemoryManager {
                     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",
+                    )));
+                }
+
+                continue;
             }
 
-            if !served {
-                return Err(io::Error::other(format!(
-                    "UFFD handler: fault at {fault_addr:#x} does not belong to any registered range",
-                )));
+            // No fault pending — advance the prefault cursor past served and
+            // end-of-range pages, then prefault one fresh page below.
+            let cursor = loop {
+                let Some((range_idx, page_idx)) = prefault_cursor else {
+                    break None;
+                };
+                if page_idx >= ranges[range_idx].num_pages() {
+                    if range_idx + 1 < ranges.len() {
+                        prefault_cursor = Some((range_idx + 1, 0));
+                        continue;
+                    }
+                    // Reached the end of the last range — every page is
+                    // mapped, so no future faults can occur. Exit.
+                    let elapsed = prefault_start.elapsed();
+                    info!(
+                        "UFFD handler: prefault done in {elapsed:.3?} — \
+                         prefaulted={pages_prefaulted} served={pages_served} \
+                         total={total_pages}"
+                    );
+                    return Ok(());
+                }
+                if served_bitmap[range_idx].is_bit_set(page_idx as usize) {
+                    prefault_cursor = Some((range_idx, page_idx + 1));
+                    continue;
+                }
+                break Some((range_idx, page_idx));
+            };
+
+            let Some((range_idx, page_idx)) = cursor else {
+                // Prefault was given up earlier (or the range list was
+                // empty). Keep serving on-demand faults.
+                continue;
+            };
+
+            let range = &ranges[range_idx];
+            let file_pos = range.file_pos(page_idx);
+            let page_addr = range.page_addr(page_idx);
+            let len = range.page_size as usize;
+
+            let advance = match snapshot_file.read_exact_at(&mut page_buf[..len], file_pos) {
+                Ok(()) => match uffd::copy(
+                    uffd_fd.as_fd(),
+                    page_addr,
+                    page_buf.as_ptr(),
+                    range.page_size,
+                ) {
+                    Ok(()) => {
+                        pages_prefaulted += 1;
+                        true
+                    }
+                    Err(e) if e.raw_os_error() == Some(libc::EEXIST) => {
+                        // Should be unreachable: in single-thread mode with
+                        // MISSING-only registration, the only installer is us,
+                        // and the bitmap check above already filtered served
+                        // pages. Treat as a bug signal but keep going.
+                        warn!(
+                            "UFFD prefault: unexpected EEXIST at {page_addr:#x} \
+                             (bitmap/handler bug?)"
+                        );
+                        true
+                    }
+                    Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => {
+                        // Unlike the on-demand handler (which must retry to
+                        // wake the faulting thread), prefault can safely skip:
+                        // any future guest access will simply page-fault and
+                        // be served by the on-demand path.
+                        true
+                    }
+                    Err(e) => {
+                        warn!("UFFD prefault: UFFDIO_COPY error at {page_addr:#x}: {e}");
+                        false
+                    }
+                },
+                Err(e) => {
+                    warn!("UFFD prefault: read error at {file_pos:#x}: {e}");
+                    false
+                }
+            };
+
+            if !advance {
+                // Prefault hit an unrecoverable error; give up but keep
+                // serving on-demand faults.
+                warn!("UFFD prefault: abandoning background prefault after error");
+                prefault_cursor = None;
+                continue;
             }
 
-            if pages_served == total_pages {
-                info!("UFFD handler: all {pages_served} pages served, exiting");
-                return Ok(());
-            }
+            prefault_cursor = Some((range_idx, page_idx + 1));
         }
     }