feat(balloon): enable memory reclaiming

Author: CompuIves

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

@@ -18,12 +18,13 @@ use vm_memory::{Address, ByteValued, Bytes, GuestAddress, GuestMemoryMmap};
 use super::super::{ActivateResult, DeviceState, Queue, VirtioDevice, TYPE_BALLOON};
 use super::utils::{compact_page_frame_numbers, remove_range};
 use super::{
-    BALLOON_DEV_ID, DEFLATE_INDEX, INFLATE_INDEX, MAX_PAGES_IN_DESC, MAX_PAGE_COMPACT_BUFFER,
-    MIB_TO_4K_PAGES, NUM_QUEUES, QUEUE_SIZES, STATS_INDEX, VIRTIO_BALLOON_F_DEFLATE_ON_OOM,
-    VIRTIO_BALLOON_F_STATS_VQ, VIRTIO_BALLOON_PFN_SHIFT, VIRTIO_BALLOON_S_AVAIL,
-    VIRTIO_BALLOON_S_CACHES, VIRTIO_BALLOON_S_HTLB_PGALLOC, VIRTIO_BALLOON_S_HTLB_PGFAIL,
-    VIRTIO_BALLOON_S_MAJFLT, VIRTIO_BALLOON_S_MEMFREE, VIRTIO_BALLOON_S_MEMTOT,
-    VIRTIO_BALLOON_S_MINFLT, VIRTIO_BALLOON_S_SWAP_IN, VIRTIO_BALLOON_S_SWAP_OUT,
+    BALLOON_DEV_ID, DEFLATE_INDEX, FREE_PAGE_REPORTING_INDEX, INFLATE_INDEX, MAX_PAGES_IN_DESC,
+    MAX_PAGE_COMPACT_BUFFER, MIB_TO_4K_PAGES, NUM_QUEUES, QUEUE_SIZES, STATS_INDEX,
+    VIRTIO_BALLOON_F_DEFLATE_ON_OOM, VIRTIO_BALLOON_F_REPORTING, VIRTIO_BALLOON_F_STATS_VQ,
+    VIRTIO_BALLOON_PFN_SHIFT, VIRTIO_BALLOON_S_AVAIL, VIRTIO_BALLOON_S_CACHES,
+    VIRTIO_BALLOON_S_HTLB_PGALLOC, VIRTIO_BALLOON_S_HTLB_PGFAIL, VIRTIO_BALLOON_S_MAJFLT,
+    VIRTIO_BALLOON_S_MEMFREE, VIRTIO_BALLOON_S_MEMTOT, VIRTIO_BALLOON_S_MINFLT,
+    VIRTIO_BALLOON_S_SWAP_IN, VIRTIO_BALLOON_S_SWAP_OUT,
 };
 use crate::virtio::balloon::Error as BalloonError;
 use crate::virtio::{IrqTrigger, IrqType};
@@ -167,10 +168,15 @@ impl Balloon {
             avail_features |= 1u64 << VIRTIO_BALLOON_F_STATS_VQ;
         }
 
+        avail_features |= 1u64 << VIRTIO_BALLOON_F_REPORTING;
+
+        logger::debug!("balloon: registering balloon device");
+
         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)?,
+            EventFd::new(libc::EFD_NONBLOCK).map_err(BalloonError::EventFd)?,
         ];
 
         let mut queues: Vec<Queue> = QUEUE_SIZES.iter().map(|&s| Queue::new(s)).collect();
@@ -231,6 +237,14 @@ impl Balloon {
         self.trigger_stats_update()
     }
 
+    pub(crate) fn process_free_page_report_event(&mut self) -> Result<(), BalloonError> {
+        logger::debug!("balloon: received free page report event");
+        self.queue_evts[FREE_PAGE_REPORTING_INDEX]
+            .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.mem().unwrap();
@@ -382,6 +396,50 @@ impl Balloon {
         Ok(())
     }
 
+    pub(crate) fn process_free_page_reporting_queue(
+        &mut self,
+    ) -> std::result::Result<(), BalloonError> {
+        logger::debug!("balloon: processing free page reporting queue");
+        let mem = self.device_state.mem().unwrap();
+
+        let mut total_removed = 0;
+        let queue = &mut self.queues[FREE_PAGE_REPORTING_INDEX];
+        let mut needs_interrupt = false;
+
+        while let Some(head) = queue.pop(mem) {
+            let head_index = head.index;
+            let head_mem = head.mem;
+
+            let mut last_desc = Some(head);
+            while let Some(desc) = last_desc {
+                total_removed += desc.len;
+                if let Err(err) =
+                    remove_range(desc.mem, (desc.addr, desc.len as u64), self.restored)
+                {
+                    error!("balloon: failed to remove range: {:?}", err);
+                };
+                last_desc = desc.next_descriptor();
+            }
+
+            // Acknowledge the receipt of the descriptor.
+            queue
+                .add_used(head_mem, head_index, 0)
+                .map_err(BalloonError::Queue)?;
+
+            logger::debug!("balloon: adding to the queue");
+
+            needs_interrupt = true;
+        }
+
+        logger::debug!("balloon: total removed: {}MiB", total_removed >> 20);
+
+        if needs_interrupt {
+            self.signal_used_queue()?;
+        }
+
+        Ok(())
+    }
+
     pub(crate) fn signal_used_queue(&self) -> Result<(), BalloonError> {
         self.irq_trigger.trigger_irq(IrqType::Vring).map_err(|err| {
             METRICS.balloon.event_fails.inc();
@@ -393,6 +451,7 @@ impl Balloon {
     pub fn process_virtio_queues(&mut self) {
         let _ = self.process_inflate();
         let _ = self.process_deflate_queue();
+        let _ = self.process_free_page_reporting_queue();
     }
 
     pub fn id(&self) -> &str {

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

@@ -9,7 +9,9 @@ use utils::epoll::EventSet;
 
 use crate::report_balloon_event_fail;
 use crate::virtio::balloon::device::Balloon;
-use crate::virtio::{VirtioDevice, DEFLATE_INDEX, INFLATE_INDEX, STATS_INDEX};
+use crate::virtio::{
+    VirtioDevice, DEFLATE_INDEX, FREE_PAGE_REPORTING_INDEX, INFLATE_INDEX, STATS_INDEX,
+};
 
 impl Balloon {
     fn register_runtime_events(&self, ops: &mut EventOps) {
@@ -27,6 +29,15 @@ impl Balloon {
                 error!("Failed to register stats timerfd event: {}", err);
             }
         }
+        if let Err(err) = ops.add(Events::new(
+            &self.queue_evts[FREE_PAGE_REPORTING_INDEX],
+            EventSet::IN,
+        )) {
+            error!(
+                "Failed to register free page reporting queue event: {}",
+                err
+            );
+        }
     }
 
     fn register_activate_event(&self, ops: &mut EventOps) {
@@ -65,6 +76,7 @@ impl MutEventSubscriber for Balloon {
             let virtq_inflate_ev_fd = self.queue_evts[INFLATE_INDEX].as_raw_fd();
             let virtq_deflate_ev_fd = self.queue_evts[DEFLATE_INDEX].as_raw_fd();
             let virtq_stats_ev_fd = self.queue_evts[STATS_INDEX].as_raw_fd();
+            let free_page_report_ev_fd = self.queue_evts[FREE_PAGE_REPORTING_INDEX].as_raw_fd();
             let stats_timer_fd = self.stats_timer.as_raw_fd();
             let activate_fd = self.activate_evt.as_raw_fd();
 
@@ -82,6 +94,9 @@ impl MutEventSubscriber for Balloon {
                 _ if source == stats_timer_fd => self
                     .process_stats_timer_event()
                     .unwrap_or_else(report_balloon_event_fail),
+                _ if source == free_page_report_ev_fd => self
+                    .process_free_page_report_event()
+                    .unwrap_or_else(report_balloon_event_fail),
                 _ if activate_fd == source => self.process_activate_event(ops),
                 _ => {
                     warn!("Balloon: Spurious event received: {:?}", source);

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

@@ -17,8 +17,9 @@ pub use self::event_handler::*;
 pub const BALLOON_DEV_ID: &str = "balloon";
 pub const CONFIG_SPACE_SIZE: usize = 8;
 pub const QUEUE_SIZE: u16 = 256;
-pub const NUM_QUEUES: usize = 3;
-pub const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE, QUEUE_SIZE, QUEUE_SIZE];
+pub const REPORTING_QUEUE_SIZE: u16 = 32;
+pub const NUM_QUEUES: usize = 4;
+pub const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE, QUEUE_SIZE, QUEUE_SIZE, QUEUE_SIZE];
 // Number of 4K pages in a MiB.
 pub const MIB_TO_4K_PAGES: u32 = 256;
 // The maximum number of pages that can be received in a single descriptor.
@@ -34,10 +35,13 @@ pub const INFLATE_INDEX: usize = 0;
 pub const DEFLATE_INDEX: usize = 1;
 // The index of the deflate queue from Balloon device queues/queues_evts vector.
 pub const STATS_INDEX: usize = 2;
+// The index of the free page reporting from Balloon device queues/queues_evts vector.
+pub const FREE_PAGE_REPORTING_INDEX: usize = 3;
 
 // 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_REPORTING: u32 = 5; // Page reporting virtqueue
 
 // The statistics tags.
 const VIRTIO_BALLOON_S_SWAP_IN: u16 = 0;