vmm: Add postcopy support to receive-migration

Plumb the SocketUffdMemorySource into the receiving side of live
migration. When memory_mode=postcopy is requested, the destination
brings up a dedicated fault connection, registers userfaultfd on the
restored memory regions, and serves guest pages on demand over that
connection while the VM resumes early.

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

vm-migration/src/lib.rs

@@ -47,6 +47,9 @@ pub enum UffdError {
     #[error("Region at {addr:#x}+{len:#x} missing COPY/WAKE support")]
     MissingIoctlSupport { addr: u64, len: u64 },
 
+    #[error("Failed to configure socket")]
+    SetSocket(#[source] io::Error),
+
     #[error("Failed to spawn handler thread")]
     SpawnThread(#[source] io::Error),
 

vmm/src/lib.rs

@@ -13,7 +13,7 @@ use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 use std::panic::AssertUnwindSafe;
 #[cfg(feature = "guest_debug")]
 use std::path::PathBuf;
-use std::sync::mpsc::{Receiver, RecvError, SendError, Sender};
+use std::sync::mpsc::{Receiver, RecvError, SendError, Sender, channel};
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 use std::{io, mem, result, thread};
@@ -49,7 +49,7 @@ use vmm_sys_util::signal::unblock_signal;
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
 
 use crate::api::{
-    ApiRequest, ApiResponse, RequestHandler, TimeoutStrategy, VmInfoResponse,
+    ApiRequest, ApiResponse, MigrationMode, RequestHandler, TimeoutStrategy, VmInfoResponse,
     VmReceiveMigrationData, VmSendMigrationData, VmmPingResponse,
 };
 use crate::config::{MemoryRestoreMode, RestoreConfig, add_to_config};
@@ -685,13 +685,19 @@ pub struct Vmm {
     check_migration_evt: EventFd,
 }
 
+/// Time before aborting on the page fault connection.
+const FAULT_CONNECTION_ACCEPT_TIMEOUT: Duration = Duration::from_secs(30);
+
 /// Just a wrapper for the data that goes into
 /// [`ReceiveMigrationState::Configured`]
 struct ReceiveMigrationConfiguredData {
     memory_manager: Arc<Mutex<MemoryManager>>,
     guest_memory: GuestMemoryAtomic<GuestMemoryMmap>,
     connections: ReceiveAdditionalConnections,
+    shared_backing: bool,
+    fault_rx: Receiver<SocketStream>,
 }
+
 /// The receiver's state machine behind the migration protocol.
 enum ReceiveMigrationState {
     /// The connection is established and we haven't received any commands yet.
@@ -938,7 +944,7 @@ impl Vmm {
         listener: &ReceiveListener,
         state: ReceiveMigrationState,
         req: &Request,
-        _receive_data_migration: &VmReceiveMigrationData,
+        receive_data_migration: &VmReceiveMigrationData,
     ) -> std::result::Result<ReceiveMigrationState, MigratableError> {
         use ReceiveMigrationState::*;
 
@@ -948,23 +954,29 @@ impl Vmm {
             )))
         };
 
+        let mode = receive_data_migration.memory_mode;
         let mut configure_vm =
             |socket: &mut SocketStream,
              memory_files: HashMap<u32, File>|
              -> std::result::Result<ReceiveMigrationConfiguredData, MigratableError> {
-                let memory_manager = self.vm_receive_config(req, socket, memory_files)?;
+                let shared_backing = !memory_files.is_empty();
+                let memory_manager = self.vm_receive_config(req, socket, memory_files, mode)?;
                 let guest_memory = memory_manager.lock().unwrap().guest_memory();
                 // Create the additional-connection receiver even in the single-connection case.
                 // At this point the receiver does not know whether the sender will use extra TCP
                 // connections. If it does not, no worker connections are accepted and memory
                 // requests continue to arrive on the main connection.
-                let connections = listener
-                    .try_clone()
-                    .and_then(|l| ReceiveAdditionalConnections::new(l, guest_memory.clone()))?;
+                // The accept thread hands the page fault connection back via this channel.
+                let (fault_tx, fault_rx) = channel();
+                let connections = listener.try_clone().and_then(|l| {
+                    ReceiveAdditionalConnections::new(l, guest_memory.clone(), fault_tx)
+                })?;
                 Ok(ReceiveMigrationConfiguredData {
                     memory_manager,
                     guest_memory,
                     connections,
+                    shared_backing,
+                    fault_rx,
                 })
             };
 
@@ -1024,18 +1036,7 @@ impl Vmm {
                     Ok(Configured(config_data))
                 }
                 Command::State => {
-                    let state_receive_begin = Instant::now();
-                    config_data.connections.cleanup()?;
-                    let (recv_state_dur, restore_vm_dur) =
-                        self.vm_receive_state(req, socket, config_data.memory_manager)?;
-                    debug!(
-                        "Migration (incoming): recv_snapshot:{}ms restore:{}ms",
-                        recv_state_dur.as_millis(),
-                        restore_vm_dur.as_millis(),
-                    );
-                    Ok(StateReceived {
-                        state_receive_begin,
-                    })
+                    self.vm_receive_state_command(req, socket, config_data, receive_data_migration)
                 }
                 c => invalid_command(state_name, c),
             },
@@ -1076,11 +1077,65 @@ impl Vmm {
         }
     }
 
+    fn vm_receive_state_command(
+        &mut self,
+        req: &Request,
+        socket: &mut SocketStream,
+        mut config_data: ReceiveMigrationConfiguredData,
+        receive_data_migration: &VmReceiveMigrationData,
+    ) -> std::result::Result<ReceiveMigrationState, MigratableError> {
+        let state_receive_begin = Instant::now();
+
+        // Serve faults before restore so accesses during restore resolve on demand.
+        if matches!(receive_data_migration.memory_mode, MigrationMode::Postcopy) {
+            let shared_backing = config_data.shared_backing;
+            let fault_stream = config_data
+                .fault_rx
+                .recv_timeout(FAULT_CONNECTION_ACCEPT_TIMEOUT)
+                .map_err(|e| {
+                    config_data.connections.cleanup().ok();
+                    MigratableError::MigrateReceive(anyhow!(
+                        "Timed out waiting for postcopy fault connection: {e}"
+                    ))
+                })?;
+            let mm = config_data.memory_manager.clone();
+            let saved_regions = mm.lock().unwrap().memory_range_table(false)?;
+            mm.lock()
+                .unwrap()
+                .start_postcopy_serving(
+                    &saved_regions,
+                    shared_backing,
+                    fault_stream,
+                    &self.exit_evt,
+                )
+                .map_err(|e| {
+                    config_data.connections.cleanup().ok();
+                    MigratableError::MigrateReceive(anyhow!("start_postcopy_serving: {e:?}"))
+                })?;
+        }
+
+        // The fault connection is in hand, so stop the accept thread.
+        config_data.connections.cleanup()?;
+
+        let (recv_state_dur, restore_vm_dur) =
+            self.vm_receive_state(req, socket, config_data.memory_manager)?;
+        debug!(
+            "Migration (incoming): recv_snapshot:{}ms restore:{}ms",
+            recv_state_dur.as_millis(),
+            restore_vm_dur.as_millis(),
+        );
+
+        Ok(ReceiveMigrationState::StateReceived {
+            state_receive_begin,
+        })
+    }
+
     fn vm_receive_config<T>(
         &mut self,
         req: &Request,
         socket: &mut T,
         existing_memory_files: HashMap<u32, File>,
+        mode: MigrationMode,
     ) -> std::result::Result<Arc<Mutex<MemoryManager>>, MigratableError>
     where
         T: Read,
@@ -1097,6 +1152,24 @@ impl Vmm {
                 MigratableError::MigrateReceive(anyhow!("Error deserialising config: {e}"))
             })?;
 
+        // Eager prefault populates memory before UFFD is registered, so those
+        // pages never fault and are never served. Reject postcopy+prefault
+        // rather than serve stale data.
+        if matches!(mode, MigrationMode::Postcopy) {
+            let memory = &vm_migration_config.vm_config.lock().unwrap().memory;
+            let prefault_enabled = memory.prefault
+                || memory
+                    .zones
+                    .as_ref()
+                    .is_some_and(|zones| zones.iter().any(|zone| zone.prefault));
+            if prefault_enabled {
+                return Err(MigratableError::MigrateReceive(anyhow!(
+                    "postcopy migration is incompatible with memory prefault; \
+                     the source VM must not be configured with prefault=on"
+                )));
+            }
+        }
+
         #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
         self.vm_check_cpuid_compatibility(
             &vm_migration_config.vm_config,
@@ -1551,7 +1624,7 @@ impl Vmm {
                 // No memory was transferred
                 MemoryMigrationContext::empty_finalized(),
             )
-            .expect("migration context should transition to VmPaused for local migration");
+            .expect("migration context should transition to VmPaused for local/postcopy migration");
         } else {
             let mut mem_send = transport::SendAdditionalConnections::new(
                 &send_data_migration.destination_url,
@@ -2822,12 +2895,17 @@ impl RequestHandler for Vmm {
             response.write_to(&mut socket)?;
         }
 
-        if let ReceiveMigrationState::Aborted = state {
-            event!("vm", "migration-receive-failed");
-            self.vm = VmOwnership::None;
-            self.vm_config = None;
-        } else {
-            event!("vm", "migration-receive-finished");
+        match state {
+            ReceiveMigrationState::Aborted => {
+                event!("vm", "migration-receive-failed");
+                self.vm = VmOwnership::None;
+                self.vm_config = None;
+            }
+            ReceiveMigrationState::Completed => {
+                // Serving and resume already happened in the protocol loop.
+                event!("vm", "migration-receive-finished");
+            }
+            _ => unreachable!("loop only exits in Completed or Aborted"),
         }
 
         Ok(())

vmm/src/memory_manager.rs

@@ -54,9 +54,13 @@ use crate::config::MemoryRestoreMode;
 use crate::coredump::{
     CoredumpMemoryRegion, CoredumpMemoryRegions, DumpState, GuestDebuggableError,
 };
+use crate::migration::transport::SocketStream;
 use crate::migration::url_to_path;
 use crate::sparse::{next_data_extent, write_region_sparse};
-use crate::uffd::{self, FaultResolution, FileUffdMemorySource, UffdMemorySource, UffdRange};
+use crate::uffd::{
+    self, FaultResolution, FileUffdMemorySource, SocketUffdMemorySource, UffdMemorySource,
+    UffdRange,
+};
 use crate::vm_config::{HotplugMethod, MemoryConfig, MemoryZoneConfig};
 use crate::{GuestMemoryMmap, GuestRegionMmap, MEMORY_MANAGER_SNAPSHOT_ID};
 
@@ -64,6 +68,7 @@ struct UffdHandler {
     stop_event: EventFd,
     result_rx: Receiver<Result<(), io::Error>>,
     handle: thread::JoinHandle<()>,
+    fault_socket_fd: Option<OwnedFd>,
 }
 
 pub const MEMORY_MANAGER_ACPI_SIZE: usize = 0x18;
@@ -912,23 +917,59 @@ impl MemoryManager {
         saved_regions: &MemoryRangeTable,
         exit_evt: &EventFd,
     ) -> Result<(), Error> {
-        if saved_regions.is_empty() {
+        let mut file_offset: u64 = 0;
+        let Some((uffd_fd, ranges)) = self.prepare_uffd(saved_regions, |r| {
+            let o = file_offset;
+            file_offset += r.length;
+            o
+        })?
+        else {
             return Ok(());
-        }
-
+        };
         let snapshot_file = File::open(file_path).map_err(Error::SnapshotOpen)?;
         let source: Box<dyn UffdMemorySource> = Box::new(FileUffdMemorySource::new(snapshot_file));
-        self.register_uffd_handler(saved_regions, exit_evt, source)
+        self.spawn_uffd_handler(uffd_fd, None, ranges, source, exit_evt)?;
+        info!("UFFD restore: demand-paged restore enabled");
+        Ok(())
     }
 
-    /// Register every range against userfaultfd, then spawn the
-    /// handler thread that resolves faults through the memory source.
-    fn register_uffd_handler(
+    /// Register UFFD and spawn the handler that serves faults over `socket`.
+    /// No-op when there are no regions.
+    pub(crate) fn start_postcopy_serving(
         &mut self,
         saved_regions: &MemoryRangeTable,
+        shared_backing: bool,
+        socket: SocketStream,
         exit_evt: &EventFd,
-        source: Box<dyn UffdMemorySource>,
     ) -> Result<(), Error> {
+        // PageFault uses the GPA as the page identifier on the wire.
+        let Some((uffd_fd, ranges)) = self.prepare_uffd(saved_regions, |r| r.gpa)? else {
+            return Ok(());
+        };
+        // Make every fault a small request/response round-trip.
+        socket.set_nodelay(true).map_err(UffdError::SetSocket)?;
+        let socket_fd = socket
+            .as_fd()
+            .try_clone_to_owned()
+            .map_err(UffdError::SetSocket)?;
+        let source: Box<dyn UffdMemorySource> =
+            Box::new(SocketUffdMemorySource::new(socket, shared_backing));
+        self.spawn_uffd_handler(uffd_fd, Some(socket_fd), ranges, source, exit_evt)
+    }
+
+    /// Create a UFFD fd and register every range.
+    fn prepare_uffd<F>(
+        &mut self,
+        saved_regions: &MemoryRangeTable,
+        mut source_offset_for: F,
+    ) -> Result<Option<(OwnedFd, Vec<UffdRange>)>, Error>
+    where
+        F: FnMut(&MemoryRange) -> u64,
+    {
+        if saved_regions.is_empty() {
+            return Ok(None);
+        }
+
         let guest_memory = self.guest_memory.memory();
         let required_uffd_features = self.required_uffd_features();
 
@@ -936,7 +977,7 @@ impl MemoryManager {
         let base_page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as u64;
 
         info!(
-            "UFFD restore: attempting demand-paged restore for {} region(s)",
+            "UFFD: registering {} region(s) for demand paging",
             saved_regions.regions().len()
         );
 
@@ -951,7 +992,6 @@ impl MemoryManager {
         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
@@ -988,22 +1028,32 @@ impl MemoryManager {
             handler_ranges.push(UffdRange {
                 host_addr,
                 length: range.length,
-                source_offset: file_offset,
+                source_offset: source_offset_for(range),
                 page_size: range_page_size,
             });
-
-            file_offset += range.length;
         }
 
+        Ok(Some((uffd_fd, handler_ranges)))
+    }
+
+    /// Spawn the UFFD handler thread that resolves faults through `source`.
+    fn spawn_uffd_handler(
+        &mut self,
+        uffd_fd: OwnedFd,
+        fault_socket_fd: Option<OwnedFd>,
+        handler_ranges: Vec<UffdRange>,
+        source: Box<dyn UffdMemorySource>,
+        exit_evt: &EventFd,
+    ) -> Result<(), Error> {
         info!(
-            "UFFD restore: registered {} region(s), {} total bytes, spawning handler",
-            handler_ranges.len(),
-            file_offset
+            "UFFD: spawning handler for {} region(s)",
+            handler_ranges.len()
         );
 
         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 panic_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()
@@ -1020,13 +1070,14 @@ impl MemoryManager {
 
                     if let Err(e) = &result {
                         error!("UFFD handler exited with error: {e}");
+                        thread_exit_evt.write(1).ok();
                     }
 
                     result_tx.send(result).ok();
                 }))
                 .map_err(|_| {
                     error!("uffd-handler thread panicked");
-                    thread_exit_evt.write(1).ok();
+                    panic_exit_evt.write(1).ok();
                 })
                 .ok();
             })
@@ -1046,16 +1097,15 @@ impl MemoryManager {
             stop_event,
             result_rx,
             handle,
+            fault_socket_fd,
         });
 
-        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) {
+        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) {
@@ -1067,6 +1117,10 @@ impl MemoryManager {
     fn stop_uffd_handler(&mut self) {
         if let Some(uffd_handler) = self.uffd_handler.take() {
             uffd_handler.stop_event.write(1).ok();
+            if let Some(fd) = &uffd_handler.fault_socket_fd {
+                // SAFETY: fd is a valid owned fd for the duration of this call.
+                unsafe { libc::shutdown(fd.as_raw_fd(), libc::SHUT_RDWR) };
+            }
             uffd_handler.handle.join().ok();
 
             match uffd_handler.result_rx.try_recv() {

vmm/src/migration/transport.rs

@@ -6,16 +6,15 @@
 use std::io::{self, ErrorKind, Read, Write};
 use std::net::{TcpListener, TcpStream};
 use std::num::{NonZeroU32, ParseIntError};
-use std::os::fd::{AsFd, BorrowedFd};
-use std::os::unix::io::AsRawFd;
+use std::os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd};
 use std::os::unix::net::{UnixListener, UnixStream};
 use std::path::{Path, PathBuf};
 use std::result::Result;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc::{Receiver, Sender, SyncSender, TrySendError, channel, sync_channel};
 use std::sync::{Arc, Mutex};
-use std::thread;
 use std::time::Duration;
+use std::{mem, thread};
 
 use anyhow::{Context, anyhow};
 use log::{debug, error, info, warn};
@@ -26,7 +25,7 @@ use vm_memory::{
     Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, ReadVolatile, VolatileMemoryError,
     VolatileSlice, WriteVolatile,
 };
-use vm_migration::protocol::{Command, MemoryRangeTable, Request, Response};
+use vm_migration::protocol::{Command, ConnectionRole, MemoryRangeTable, Request, Response};
 use vm_migration::tls::{TlsServerConfig, TlsStream};
 use vm_migration::{MigratableError, Snapshot};
 use vmm_sys_util::eventfd::EventFd;
@@ -159,6 +158,40 @@ impl Write for SocketStream {
     }
 }
 
+impl SocketStream {
+    pub(crate) fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
+        match self {
+            SocketStream::Unix(s) => s.set_read_timeout(dur),
+            SocketStream::Tcp(s) => s.set_read_timeout(dur),
+            SocketStream::Tls(s) => {
+                let fd = s.as_fd().as_raw_fd();
+                // SAFETY: fd is borrowed from the TLS stream and forgotten
+                // immediately. The TLS stream retains ownership.
+                let tcp = unsafe { TcpStream::from_raw_fd(fd) };
+                let r = tcp.set_read_timeout(dur);
+                mem::forget(tcp);
+                r
+            }
+        }
+    }
+
+    pub(crate) fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
+        match self {
+            SocketStream::Unix(_) => Ok(()),
+            SocketStream::Tcp(s) => s.set_nodelay(nodelay),
+            SocketStream::Tls(s) => {
+                let fd = s.as_fd().as_raw_fd();
+                // SAFETY: fd is borrowed from the TLS stream and forgotten
+                // immediately. The TLS stream retains ownership.
+                let tcp = unsafe { TcpStream::from_raw_fd(fd) };
+                let r = tcp.set_nodelay(nodelay);
+                mem::forget(tcp);
+                r
+            }
+        }
+    }
+}
+
 impl AsFd for SocketStream {
     fn as_fd(&self) -> BorrowedFd<'_> {
         match self {
@@ -263,6 +296,7 @@ impl ReceiveAdditionalConnections {
     pub(crate) fn new(
         listener: ReceiveListener,
         guest_memory: GuestMemoryAtomic<GuestMemoryMmap>,
+        fault_tx: Sender<SocketStream>,
     ) -> Result<Self, MigratableError> {
         let event_fd = EventFd::new(0)
             .context("Error creating terminate fd")
@@ -275,7 +309,9 @@ impl ReceiveAdditionalConnections {
 
         let accept_thread = thread::Builder::new()
             .name("migrate-receive-accept-connections".to_owned())
-            .spawn(move || Self::accept_connections(listener, &terminate_fd, &guest_memory))
+            .spawn(move || {
+                Self::accept_connections(listener, &terminate_fd, &guest_memory, &fault_tx)
+            })
             .context("Error creating connection accept thread")
             .map_err(MigratableError::MigrateReceive)?;
 
@@ -289,6 +325,7 @@ impl ReceiveAdditionalConnections {
         mut listener: ReceiveListener,
         terminate_fd: &EventFd,
         guest_memory: &GuestMemoryAtomic<GuestMemoryMmap>,
+        fault_tx: &Sender<SocketStream>,
     ) -> Result<(), MigratableError> {
         let mut threads: Vec<thread::JoinHandle<Result<(), MigratableError>>> = Vec::new();
         let mut first_err = loop {
@@ -306,6 +343,39 @@ impl ReceiveAdditionalConnections {
                 )));
             }
 
+            // Read the role header with a timeout so a stalled peer can't block
+            // acceptance of other connections.
+            socket
+                .set_read_timeout(Some(Duration::from_secs(10)))
+                .map_err(MigratableError::MigrateSocket)?;
+            let role = match ConnectionRole::read_from(&mut socket) {
+                Ok(role) => role,
+                Err(e) => {
+                    warn!("Dropping connection: failed to read connection role: {e}");
+                    continue;
+                }
+            };
+            socket
+                .set_read_timeout(None)
+                .map_err(MigratableError::MigrateSocket)?;
+
+            match role {
+                ConnectionRole::Invalid => {
+                    warn!("Dropping connection: invalid connection role");
+                    continue;
+                }
+                ConnectionRole::Fault => {
+                    // Hand the fault connection to the main thread. No memory worker for it.
+                    if let Err(e) = fault_tx.send(socket) {
+                        break Err(MigratableError::MigrateReceive(anyhow!(
+                            "Failed to hand off fault connection: {e}"
+                        )));
+                    }
+                    continue;
+                }
+                ConnectionRole::PrecopyMemory => {}
+            }
+
             let guest_memory = guest_memory.clone();
             let terminate_fd = match terminate_fd
                 .try_clone()
@@ -539,6 +609,7 @@ impl SendAdditionalConnections {
         // this case we create one additional thread for each connection.
         for n in 0..configured_connections {
             let mut socket = send_migration_socket(destination, tls_dir)?;
+            ConnectionRole::PrecopyMemory.write_to(&mut socket)?;
             let guest_memory = guest_memory.clone();
             let message_rx = message_rx.clone();
             let worker_error = worker_error.clone();

vmm/src/uffd.rs

@@ -278,15 +278,12 @@ impl UffdMemorySource for FileUffdMemorySource {
 }
 
 /// Memory source that provides pages content over a socket.
-// Wired into the receive-migration path in the next commit.
-#[allow(dead_code)]
 pub(crate) struct SocketUffdMemorySource {
     stream: SocketStream,
     shared_backing: bool,
     buf: Vec<u8>,
 }
 
-#[allow(dead_code)]
 impl SocketUffdMemorySource {
     pub fn new(stream: SocketStream, shared_backing: bool) -> Self {
         Self {