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

@@ -669,6 +669,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.
@@ -748,7 +749,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.
@@ -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;
 
@@ -204,16 +208,26 @@ impl Balloon {
             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)?,
+            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 dropped_queue_count = 0;
         if stats_polling_interval_s == 0 {
-            let _ = queues.remove(STATS_INDEX);
+            dropped_queue_count += 1;
+        }
+
+        if free_page_reporting {
+            avail_features |= 1u64 << VIRTIO_BALLOON_F_FREE_PAGE_REPORTING;
+        } else {
+            dropped_queue_count += 1;
         }
 
+        queues.truncate(queues.len() - dropped_queue_count);
+
         let stats_timer =
             TimerFd::new_custom(ClockId::Monotonic, true, true).map_err(BalloonError::Timer)?;
 
@@ -262,9 +276,16 @@ 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,41 @@ 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 +483,11 @@ 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 +527,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 +604,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 +813,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 +840,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 +875,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 +901,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 +962,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 +1034,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 +1084,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 +1176,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 +1191,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 +1204,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 +1222,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,17 @@ 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 +115,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 +157,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);

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

@@ -25,13 +25,11 @@ pub const BALLOON_DEV_ID: &str = "balloon";
 /// The size of the config space.
 pub const BALLOON_CONFIG_SPACE_SIZE: usize = 8;
 /// Number of virtio queues.
-pub const BALLOON_NUM_QUEUES: usize = 3;
+pub const BALLOON_NUM_QUEUES: usize = 4;
 /// Virtio queue sizes, in number of descriptor chain heads.
-//  There are 3 queues for a virtio device (in this order): RX, TX, Event
+//  There are 4 queues for a virtio device (in this order): RX, TX, Event, Reporting
 pub const BALLOON_QUEUE_SIZES: [u16; BALLOON_NUM_QUEUES] = [
-    FIRECRACKER_MAX_QUEUE_SIZE,
-    FIRECRACKER_MAX_QUEUE_SIZE,
-    FIRECRACKER_MAX_QUEUE_SIZE,
+    FIRECRACKER_MAX_QUEUE_SIZE; BALLOON_NUM_QUEUES
 ];
 // Number of 4K pages in a MiB.
 pub const MIB_TO_4K_PAGES: u32 = 256;
@@ -42,16 +40,17 @@ pub const MAX_PAGES_IN_DESC: usize = 256;
 pub const MAX_PAGE_COMPACT_BUFFER: usize = 2048;
 /// The addresses given by the driver are divided by 4096.
 pub const VIRTIO_BALLOON_PFN_SHIFT: u32 = 12;
-/// The index of the deflate queue from Balloon device queues/queues_evts vector.
+/// The index of the inflate queue from Balloon device queues/queues_evts vector.
 pub const INFLATE_INDEX: usize = 0;
 /// The index of the deflate queue from Balloon device queues/queues_evts vector.
 pub const DEFLATE_INDEX: usize = 1;
-/// The index of the deflate queue from Balloon device queues/queues_evts vector.
+/// The index of the stats queue from Balloon device queues/queues_evts vector.
 pub const STATS_INDEX: usize = 2;
 
 // The feature bitmap for virtio balloon.
 const VIRTIO_BALLOON_F_STATS_VQ: u32 = 1; // Enable statistics.
 const VIRTIO_BALLOON_F_DEFLATE_ON_OOM: u32 = 2; // Deflate balloon on OOM.
+const VIRTIO_BALLOON_F_FREE_PAGE_REPORTING: u32 = 5; // Enable free page reportin
 
 // The statistics tags.
 const VIRTIO_BALLOON_S_SWAP_IN: u16 = 0;