fix: only use mmap trick when restoring from file

When the balloon inflates, and the guest gives us back some pages of
memory, we need to free those pages. In booted VMs, we do this with
madvise(MADV_DONTNEED), and in restored VMs we do it by MAP_FIXED-ing a
new VMA on top of the range-to-free. This is because if guest memory is
a MAP_PRIVATE of a memory file, madvise has no effect.

However, we also do this MAP_FIXED trick if the snapshot is restored
with UFFD. In this case, its not needed (madvise works perfectly fine),
and in fact, its wrong: If we map over the memory range, UFFD will not
receive Remove events for the specified range.

Fix this by only using the mmap trick for file-based restored.

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

Author: roypat

src/vmm/src/builder.rs

@@ -531,6 +531,7 @@ pub fn build_microvm_from_snapshot(
         resource_allocator: &mut vmm.resource_allocator,
         vm_resources,
         instance_id: &instance_info.id,
+        restored_from_file: vmm.uffd.is_none(),
     };
 
     vmm.mmio_device_manager =

src/vmm/src/device_manager/persist.rs

@@ -214,6 +214,7 @@ pub struct MMIODevManagerConstructorArgs<'a> {
     pub resource_allocator: &'a mut ResourceAllocator,
     pub vm_resources: &'a mut VmResources,
     pub instance_id: &'a str,
+    pub restored_from_file: bool,
 }
 impl fmt::Debug for MMIODevManagerConstructorArgs<'_> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -512,7 +513,10 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
         if let Some(balloon_state) = &state.balloon_device {
             let device = Arc::new(Mutex::new(Balloon::restore(
-                BalloonConstructorArgs { mem: mem.clone() },
+                BalloonConstructorArgs {
+                    mem: mem.clone(),
+                    restored_from_file: constructor_args.restored_from_file,
+                },
                 &balloon_state.device_state,
             )?));
 
@@ -807,6 +811,7 @@ mod tests {
             resource_allocator: &mut resource_allocator,
             vm_resources,
             instance_id: "microvm-id",
+            restored_from_file: true,
         };
         let restored_dev_manager =
             MMIODeviceManager::restore(restore_args, &device_states).unwrap();

src/vmm/src/devices/virtio/balloon/device.rs

@@ -164,7 +164,7 @@ pub struct Balloon {
     pub(crate) irq_trigger: IrqTrigger,
 
     // Implementation specific fields.
-    pub(crate) restored: bool,
+    pub(crate) restored_from_file: bool,
     pub(crate) stats_polling_interval_s: u16,
     pub(crate) stats_timer: TimerFd,
     // The index of the previous stats descriptor is saved because
@@ -189,7 +189,7 @@ impl fmt::Debug for Balloon {
             .field("queue_evts", &self.queue_evts)
             .field("device_state", &self.device_state)
             .field("irq_trigger", &self.irq_trigger)
-            .field("restored", &self.restored)
+            .field("restored_from_file", &self.restored_from_file)
             .field("stats_polling_interval_s", &self.stats_polling_interval_s)
             .field("stats_desc_index", &self.stats_desc_index)
             .field("latest_stats", &self.latest_stats)
@@ -204,7 +204,7 @@ impl Balloon {
         amount_mib: u32,
         deflate_on_oom: bool,
         stats_polling_interval_s: u16,
-        restored: bool,
+        restored_from_file: bool,
     ) -> Result<Balloon, BalloonError> {
         let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
 
@@ -245,7 +245,7 @@ impl Balloon {
             irq_trigger: IrqTrigger::new().map_err(BalloonError::EventFd)?,
             device_state: DeviceState::Inactive,
             activate_evt: EventFd::new(libc::EFD_NONBLOCK).map_err(BalloonError::EventFd)?,
-            restored,
+            restored_from_file,
             stats_polling_interval_s,
             stats_timer,
             stats_desc_index: None,
@@ -355,7 +355,7 @@ impl Balloon {
                 if let Err(err) = remove_range(
                     mem,
                     (guest_addr, u64::from(range_len) << VIRTIO_BALLOON_PFN_SHIFT),
-                    self.restored,
+                    self.restored_from_file,
                 ) {
                     error!("Error removing memory range: {:?}", err);
                 }

src/vmm/src/devices/virtio/balloon/persist.rs

@@ -95,6 +95,7 @@ pub struct BalloonState {
 pub struct BalloonConstructorArgs {
     /// Pointer to guest memory.
     pub mem: GuestMemoryMmap,
+    pub restored_from_file: bool,
 }
 
 impl Persist<'_> for Balloon {
@@ -121,7 +122,12 @@ impl Persist<'_> for Balloon {
     ) -> Result<Self, Self::Error> {
         // We can safely create the balloon with arbitrary flags and
         // num_pages because we will overwrite them after.
-        let mut balloon = Balloon::new(0, false, state.stats_polling_interval_s, true)?;
+        let mut balloon = Balloon::new(
+            0,
+            false,
+            state.stats_polling_interval_s,
+            constructor_args.restored_from_file,
+        )?;
 
         let mut num_queues = BALLOON_NUM_QUEUES;
         // As per the virtio 1.1 specification, the statistics queue
@@ -192,13 +198,16 @@ mod tests {
 
         // Deserialize and restore the balloon device.
         let restored_balloon = Balloon::restore(
-            BalloonConstructorArgs { mem: guest_mem },
+            BalloonConstructorArgs {
+                mem: guest_mem,
+                restored_from_file: true,
+            },
             &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
         )
         .unwrap();
 
         assert_eq!(restored_balloon.device_type(), TYPE_BALLOON);
-        assert!(restored_balloon.restored);
+        assert!(restored_balloon.restored_from_file);
 
         assert_eq!(restored_balloon.acked_features, balloon.acked_features);
         assert_eq!(restored_balloon.avail_features, balloon.avail_features);

src/vmm/src/devices/virtio/balloon/util.rs

@@ -68,7 +68,7 @@ pub(crate) fn compact_page_frame_numbers(v: &mut [u32]) -> Vec<(u32, u32)> {
 pub(crate) fn remove_range(
     guest_memory: &GuestMemoryMmap,
     range: (GuestAddress, u64),
-    restored: bool,
+    restored_from_file: bool,
 ) -> Result<(), RemoveRegionError> {
     let (guest_address, range_len) = range;
 
@@ -83,7 +83,7 @@ pub(crate) fn remove_range(
         // Mmap a new anonymous region over the present one in order to create a hole.
         // This workaround is (only) needed after resuming from a snapshot because the guest memory
         // is mmaped from file as private and there is no `madvise` flag that works for this case.
-        if restored {
+        if restored_from_file {
             // SAFETY: The address and length are known to be valid.
             let ret = unsafe {
                 libc::mmap(

src/vmm/src/lib.rs

@@ -314,8 +314,6 @@ pub struct Vmm {
     vm: Vm,
     guest_memory: GuestMemoryMmap,
     // Save UFFD in order to keep it open in the Firecracker process, as well.
-    // Since this field is never read again, we need to allow `dead_code`.
-    #[allow(dead_code)]
     uffd: Option<Uffd>,
     vcpus_handles: Vec<VcpuHandle>,
     // Used by Vcpus and devices to initiate teardown; Vmm should never write here.