virtio-balloon: Add free page reporting

Free page reporting is a mechanism in which the guest will notify the
host of pages which are not currently in use. This feature can only be
configured on boot and will continue to report continuously.

With free page reporting firecracker will `MADV_DONTNEED` on the ranges
reported. This allows the host to free up memory and reduce the RSS of
the VM. With UFFD this is sent as the `UFFD_EVENT_REMOVE` after the call
with `MADV_DONTNEED`.

Signed-off-by: Jack Thomson <jackabt@amazon.com>

Author: JackThomson2

src/vmm/src/builder.rs

@@ -1231,6 +1231,7 @@ pub(crate) mod tests {
             amount_mib: 0,
             deflate_on_oom: false,
             stats_polling_interval_s: 0,
+            free_page_reporting: false,
         };
 
         let mut cmdline = default_kernel_cmdline();

src/vmm/src/device_manager/pci_mngr.rs

@@ -645,6 +645,7 @@ mod tests {
                 amount_mib: 123,
                 deflate_on_oom: false,
                 stats_polling_interval_s: 1,
+                free_page_reporting: false,
             };
             insert_balloon_device(&mut vmm, &mut cmdline, &mut event_manager, balloon_cfg);
             // Add a block device.
@@ -728,7 +729,8 @@ mod tests {
   "balloon": {{
     "amount_mib": 123,
     "deflate_on_oom": false,
-    "stats_polling_interval_s": 1
+    "stats_polling_interval_s": 1,
+    "free_page_reporting": false
   }},
   "drives": [
     {{

src/vmm/src/device_manager/persist.rs

@@ -661,6 +661,7 @@ mod tests {
                 amount_mib: 123,
                 deflate_on_oom: false,
                 stats_polling_interval_s: 1,
+                free_page_reporting: false,
             };
             insert_balloon_device(&mut vmm, &mut cmdline, &mut event_manager, balloon_cfg);
             // Add a block device.
@@ -740,7 +741,8 @@ mod tests {
   "balloon": {{
     "amount_mib": 123,
     "deflate_on_oom": false,
-    "stats_polling_interval_s": 1
+    "stats_polling_interval_s": 1,
+    "free_page_reporting": false
   }},
   "drives": [
     {{

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

@@ -24,7 +24,7 @@ use super::{
     VIRTIO_BALLOON_S_MEMTOT, VIRTIO_BALLOON_S_MINFLT, VIRTIO_BALLOON_S_SWAP_IN,
     VIRTIO_BALLOON_S_SWAP_OUT,
 };
-use crate::devices::virtio::balloon::BalloonError;
+use crate::devices::virtio::balloon::{BalloonError, VIRTIO_BALLOON_F_FREE_PAGE_REPORTING};
 use crate::devices::virtio::device::ActiveState;
 use crate::devices::virtio::generated::virtio_config::VIRTIO_F_VERSION_1;
 use crate::devices::virtio::generated::virtio_ids::VIRTIO_ID_BALLOON;
@@ -83,6 +83,9 @@ pub struct BalloonConfig {
     pub deflate_on_oom: bool,
     /// Interval of time in seconds at which the balloon statistics are updated.
     pub stats_polling_interval_s: u16,
+    /// Free page reporting enabled
+    #[serde(default)]
+    pub free_page_reporting: bool,
 }
 
 /// BalloonStats holds statistics returned from the stats_queue.
@@ -169,7 +172,7 @@ pub struct Balloon {
 
     // Transport related fields.
     pub(crate) queues: Vec<Queue>,
-    pub(crate) queue_evts: [EventFd; BALLOON_NUM_QUEUES],
+    pub(crate) queue_evts: Vec<EventFd>,
     pub(crate) device_state: DeviceState,
 
     // Implementation specific fields.
@@ -189,6 +192,7 @@ impl Balloon {
         amount_mib: u32,
         deflate_on_oom: bool,
         stats_polling_interval_s: u16,
+        free_page_reporting: bool,
     ) -> Result<Balloon, BalloonError> {
         let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
 
@@ -200,20 +204,26 @@ impl Balloon {
             avail_features |= 1u64 << VIRTIO_BALLOON_F_STATS_VQ;
         }
 
-        let queue_evts = [
-            EventFd::new(libc::EFD_NONBLOCK).map_err(BalloonError::EventFd)?,
-            EventFd::new(libc::EFD_NONBLOCK).map_err(BalloonError::EventFd)?,
-            EventFd::new(libc::EFD_NONBLOCK).map_err(BalloonError::EventFd)?,
-        ];
-
-        let mut queues: Vec<Queue> = BALLOON_QUEUE_SIZES.iter().map(|&s| Queue::new(s)).collect();
-
         // The VirtIO specification states that the statistics queue should
         // not be present at all if the statistics are not enabled.
+        let mut queue_count = BALLOON_QUEUE_SIZES.len();
         if stats_polling_interval_s == 0 {
-            let _ = queues.remove(STATS_INDEX);
+            queue_count -= 1;
+        }
+
+        if free_page_reporting {
+            avail_features |= 1u64 << VIRTIO_BALLOON_F_FREE_PAGE_REPORTING;
+        } else {
+            queue_count -= 1;
         }
 
+        let queues: Vec<Queue> = (0..queue_count)
+            .map(|s| Queue::new(BALLOON_QUEUE_SIZES[s]))
+            .collect();
+        let queue_evts = (0..queue_count)
+            .map(|_| EventFd::new(libc::EFD_NONBLOCK).map_err(BalloonError::EventFd))
+            .collect::<Result<Vec<_>, _>>()?;
+
         let stats_timer =
             TimerFd::new_custom(ClockId::Monotonic, true, true).map_err(BalloonError::Timer)?;
 
@@ -262,9 +272,20 @@ impl Balloon {
         self.trigger_stats_update()
     }
 
+    pub(crate) fn process_free_page_reporting_queue_event(&mut self) -> Result<(), BalloonError> {
+        self.queue_evts[self.free_page_reporting_idx()]
+            .read()
+            .map_err(BalloonError::EventFd)?;
+        self.process_free_page_reporting_queue()
+    }
+
     pub(crate) fn process_inflate(&mut self) -> Result<(), BalloonError> {
         // This is safe since we checked in the event handler that the device is activated.
-        let mem = &self.device_state.active_state().unwrap().mem;
+        let mem = &self
+            .device_state
+            .active_state()
+            .ok_or(BalloonError::DeviceNotActive)?
+            .mem;
         METRICS.inflate_count.inc();
 
         let queue = &mut self.queues[INFLATE_INDEX];
@@ -406,6 +427,37 @@ impl Balloon {
         Ok(())
     }
 
+    pub(crate) fn process_free_page_reporting_queue(&mut self) -> Result<(), BalloonError> {
+        let mem = &self.device_state.active_state().unwrap().mem;
+
+        let idx = self.free_page_reporting_idx();
+        let queue = &mut self.queues[idx];
+        let mut needs_interrupt = false;
+
+        while let Some(head) = queue.pop()? {
+            let head_index = head.index;
+
+            let mut last_desc = Some(head);
+            while let Some(desc) = last_desc {
+                if let Err(err) = mem.discard_range(desc.addr, desc.len as usize) {
+                    error!("balloon: failed to remove range: {err:?}");
+                }
+                last_desc = desc.next_descriptor();
+            }
+
+            queue.add_used(head.index, 0)?;
+            needs_interrupt = true;
+        }
+
+        queue.advance_used_ring_idx();
+
+        if needs_interrupt {
+            self.signal_used_queue(idx)?;
+        }
+
+        Ok(())
+    }
+
     pub(crate) fn signal_used_queue(&self, qidx: usize) -> Result<(), BalloonError> {
         self.interrupt_trigger()
             .trigger(VirtioInterruptType::Queue(
@@ -427,6 +479,13 @@ impl Balloon {
             return Err(err);
         }
 
+        if self.free_page_reporting()
+            && let Err(BalloonError::InvalidAvailIdx(err)) =
+                self.process_free_page_reporting_queue()
+        {
+            return Err(err);
+        }
+
         Ok(())
     }
 
@@ -466,6 +525,20 @@ impl Balloon {
         }
     }
 
+    pub fn free_page_reporting(&self) -> bool {
+        self.avail_features & (1u64 << VIRTIO_BALLOON_F_FREE_PAGE_REPORTING) != 0
+    }
+
+    pub fn free_page_reporting_idx(&self) -> usize {
+        let mut idx = STATS_INDEX;
+
+        if self.stats_polling_interval_s > 0 {
+            idx += 1;
+        }
+
+        idx
+    }
+
     /// Update the statistics polling interval.
     pub fn update_stats_polling_interval(&mut self, interval_s: u16) -> Result<(), BalloonError> {
         if self.stats_polling_interval_s == interval_s {
@@ -529,6 +602,7 @@ impl Balloon {
             amount_mib: self.size_mb(),
             deflate_on_oom: self.deflate_on_oom(),
             stats_polling_interval_s: self.stats_polling_interval_s(),
+            free_page_reporting: self.free_page_reporting(),
         }
     }
 
@@ -737,7 +811,7 @@ pub(crate) mod tests {
         // Test all feature combinations.
         for deflate_on_oom in [true, false].iter() {
             for stats_interval in [0, 1].iter() {
-                let mut balloon = Balloon::new(0, *deflate_on_oom, *stats_interval).unwrap();
+                let mut balloon = Balloon::new(0, *deflate_on_oom, *stats_interval, false).unwrap();
                 assert_eq!(balloon.device_type(), VIRTIO_ID_BALLOON);
 
                 let features: u64 = (1u64 << VIRTIO_F_VERSION_1)
@@ -764,12 +838,13 @@ pub(crate) mod tests {
 
     #[test]
     fn test_virtio_read_config() {
-        let balloon = Balloon::new(0x10, true, 0).unwrap();
+        let balloon = Balloon::new(0x10, true, 0, false).unwrap();
 
         let cfg = BalloonConfig {
             amount_mib: 16,
             deflate_on_oom: true,
             stats_polling_interval_s: 0,
+            free_page_reporting: false,
         };
         assert_eq!(balloon.config(), cfg);
 
@@ -798,7 +873,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_virtio_write_config() {
-        let mut balloon = Balloon::new(0, true, 0).unwrap();
+        let mut balloon = Balloon::new(0, true, 0, false).unwrap();
 
         let expected_config_space: [u8; BALLOON_CONFIG_SPACE_SIZE] =
             [0x00, 0x50, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
@@ -824,7 +899,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_invalid_request() {
-        let mut balloon = Balloon::new(0, true, 0).unwrap();
+        let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         let mem = default_mem();
         let interrupt = default_interrupt();
         // Only initialize the inflate queue to demonstrate invalid request handling.
@@ -885,7 +960,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_inflate() {
-        let mut balloon = Balloon::new(0, true, 0).unwrap();
+        let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         let mem = default_mem();
         let interrupt = default_interrupt();
         let infq = VirtQueue::new(GuestAddress(0), &mem, 16);
@@ -957,7 +1032,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_deflate() {
-        let mut balloon = Balloon::new(0, true, 0).unwrap();
+        let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         let mem = default_mem();
         let interrupt = default_interrupt();
         let defq = VirtQueue::new(GuestAddress(0), &mem, 16);
@@ -1007,7 +1082,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_stats() {
-        let mut balloon = Balloon::new(0, true, 1).unwrap();
+        let mut balloon = Balloon::new(0, true, 1, false).unwrap();
         let mem = default_mem();
         let interrupt = default_interrupt();
         let statsq = VirtQueue::new(GuestAddress(0), &mem, 16);
@@ -1099,7 +1174,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_process_balloon_queues() {
-        let mut balloon = Balloon::new(0x10, true, 0).unwrap();
+        let mut balloon = Balloon::new(0x10, true, 0, false).unwrap();
         let mem = default_mem();
         let interrupt = default_interrupt();
         let infq = VirtQueue::new(GuestAddress(0), &mem, 16);
@@ -1114,7 +1189,7 @@ pub(crate) mod tests {
 
     #[test]
     fn test_update_stats_interval() {
-        let mut balloon = Balloon::new(0, true, 0).unwrap();
+        let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         let mem = default_mem();
         let q = VirtQueue::new(GuestAddress(0), &mem, 16);
         balloon.set_queue(INFLATE_INDEX, q.create_queue());
@@ -1127,7 +1202,7 @@ pub(crate) mod tests {
         );
         balloon.update_stats_polling_interval(0).unwrap();
 
-        let mut balloon = Balloon::new(0, true, 1).unwrap();
+        let mut balloon = Balloon::new(0, true, 1, false).unwrap();
         let mem = default_mem();
         let q = VirtQueue::new(GuestAddress(0), &mem, 16);
         balloon.set_queue(INFLATE_INDEX, q.create_queue());
@@ -1145,14 +1220,14 @@ pub(crate) mod tests {
 
     #[test]
     fn test_cannot_update_inactive_device() {
-        let mut balloon = Balloon::new(0, true, 0).unwrap();
+        let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         // Assert that we can't update an inactive device.
         balloon.update_size(1).unwrap_err();
     }
 
     #[test]
     fn test_num_pages() {
-        let mut balloon = Balloon::new(0, true, 0).unwrap();
+        let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         // Switch the state to active.
         balloon.device_state = DeviceState::Activated(ActiveState {
             mem: single_region_mem(32 << 20),

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

@@ -15,6 +15,7 @@ impl Balloon {
     const PROCESS_VIRTQ_DEFLATE: u32 = 2;
     const PROCESS_VIRTQ_STATS: u32 = 3;
     const PROCESS_STATS_TIMER: u32 = 4;
+    const PROCESS_VIRTQ_FREE_PAGE_REPORTING: u32 = 6;
 
     fn register_runtime_events(&self, ops: &mut EventOps) {
         if let Err(err) = ops.add(Events::with_data(
@@ -47,6 +48,19 @@ impl Balloon {
                 error!("Failed to register stats timerfd event: {}", err);
             }
         }
+
+        if self.free_page_reporting()
+            && let Err(err) = ops.add(Events::with_data(
+                &self.queue_evts[self.free_page_reporting_idx()],
+                Self::PROCESS_VIRTQ_FREE_PAGE_REPORTING,
+                EventSet::IN,
+            ))
+        {
+            error!(
+                "Failed to register free page reporting queue event: {}",
+                err
+            );
+        }
     }
 
     fn register_activate_event(&self, ops: &mut EventOps) {
@@ -103,6 +117,9 @@ impl MutEventSubscriber for Balloon {
                 Self::PROCESS_STATS_TIMER => self
                     .process_stats_timer_event()
                     .unwrap_or_else(report_balloon_event_fail),
+                Self::PROCESS_VIRTQ_FREE_PAGE_REPORTING => self
+                    .process_free_page_reporting_queue_event()
+                    .unwrap_or_else(report_balloon_event_fail),
                 _ => {
                     warn!("Balloon: Spurious event received: {:?}", source);
                 }
@@ -142,7 +159,7 @@ pub mod tests {
     #[test]
     fn test_event_handler() {
         let mut event_manager = EventManager::new().unwrap();
-        let mut balloon = Balloon::new(0, true, 10).unwrap();
+        let mut balloon = Balloon::new(0, true, 10, false).unwrap();
         let mem = default_mem();
         let interrupt = default_interrupt();
         let infq = VirtQueue::new(GuestAddress(0), &mem, 16);