balloon: let balloon inflation works on page size other than 4k

Currently, virtio-balloon can't work well with page size other than 4k.
The virtio-balloon always works in units of 4kiB (BALLOON_PAGE_SIZE), but
we can only actually discard memory in units of the host page size.

We get some idea from [1] to solve this issue.

What has been done in this commit:

For balloon inflation:

A bitmap is employed to track the memory range to be released in 4k
granularity. Once it accumulates to one host page size, the corresponding
page is released, and the bitmap is cleared to handle the next record.
This process continues until all the memory range is managed. Memory will
only be released when a consecutive set of balloon request entries from
the same host page reaches the full host page size. If a balloon request
entry from a different host page is encountered, the bitmap and the base
host page address will be reset. Consequently, memory is released in
units of the page size, ensuring efficient memory management. That's say
if memory range length to be released smaller than page size or if the
guest scatters requests each of whose size is smaller than page size
across different host pages no memory will be released.

[1] https://patchwork.kernel.org/project/qemu-devel/patch/20190214043916.22128-6-david@gibson.dropbear.id.au/

Fixes: cloud-hypervisor#5369
Signed-off-by: Jianyong Wu <jianyong.wu@arm.com>

virtio-devices/src/balloon.rs

@@ -29,6 +29,7 @@ use thiserror::Error;
 use versionize::{VersionMap, Versionize, VersionizeResult};
 use versionize_derive::Versionize;
 use virtio_queue::{Queue, QueueT};
+use vm_allocator::page_size::{align_page_size_down, get_page_size};
 use vm_memory::{
     Address, ByteValued, Bytes, GuestAddress, GuestAddressSpace, GuestMemory, GuestMemoryAtomic,
     GuestMemoryError, GuestMemoryRegion,
@@ -94,6 +95,50 @@ pub struct VirtioBalloonConfig {
     actual: u32,
 }
 
+#[derive(Clone, Debug)]
+struct PartiallyBalloonedPage {
+    addr: u64,
+    bitmap: Vec<u64>,
+    page_size: u64,
+}
+
+impl PartiallyBalloonedPage {
+    fn new() -> Self {
+        let page_size = get_page_size();
+        let len = ((page_size >> VIRTIO_BALLOON_PFN_SHIFT) + 63) / 64;
+        // Initial each padding bit as 1 in bitmap.
+        let mut bitmap = vec![0_u64; len as usize];
+        let pad_num = len * 64 - (page_size >> VIRTIO_BALLOON_PFN_SHIFT);
+        bitmap[(len - 1) as usize] = !((1 << (64 - pad_num)) - 1);
+        Self {
+            addr: 0,
+            bitmap,
+            page_size,
+        }
+    }
+
+    fn pfn_match(&self, addr: u64) -> bool {
+        self.addr == addr & !(self.page_size - 1)
+    }
+
+    fn bitmap_full(&self) -> bool {
+        self.bitmap.iter().all(|b| *b == u64::MAX)
+    }
+
+    fn set_bit(&mut self, addr: u64) {
+        let addr_offset = (addr % self.page_size) >> VIRTIO_BALLOON_PFN_SHIFT;
+        self.bitmap[(addr_offset / 64) as usize] |= 1 << (addr_offset % 64);
+    }
+
+    fn reset(&mut self) {
+        let len = ((self.page_size >> VIRTIO_BALLOON_PFN_SHIFT) + 63) / 64;
+        self.addr = 0;
+        self.bitmap = vec![0; len as usize];
+        let pad_num = len * 64 - (self.page_size >> VIRTIO_BALLOON_PFN_SHIFT);
+        self.bitmap[(len - 1) as usize] = !((1 << (64 - pad_num)) - 1);
+    }
+}
+
 const CONFIG_ACTUAL_OFFSET: u64 = 4;
 const CONFIG_ACTUAL_SIZE: usize = 4;
 
@@ -109,6 +154,7 @@ struct BalloonEpollHandler {
     reporting_queue_evt: Option<EventFd>,
     kill_evt: EventFd,
     pause_evt: EventFd,
+    pbp: Option<PartiallyBalloonedPage>,
 }
 
 impl BalloonEpollHandler {
@@ -165,6 +211,43 @@ impl BalloonEpollHandler {
         Self::advise_memory_range(memory, range_base, range_len, libc::MADV_DONTNEED)
     }
 
+    fn release_memory_range_4k(
+        pbp: &mut Option<PartiallyBalloonedPage>,
+        memory: &GuestMemoryMmap,
+        pfn: u32,
+    ) -> result::Result<(), Error> {
+        let range_base = GuestAddress((pfn as u64) << VIRTIO_BALLOON_PFN_SHIFT);
+        let range_len = 1 << VIRTIO_BALLOON_PFN_SHIFT;
+
+        let page_size: u64 = get_page_size();
+        if page_size == 1 << VIRTIO_BALLOON_PFN_SHIFT {
+            return Self::release_memory_range(memory, range_base, range_len);
+        }
+
+        if pbp.is_none() {
+            *pbp = Some(PartiallyBalloonedPage::new());
+        }
+
+        if !pbp.as_ref().unwrap().pfn_match(range_base.0) {
+            // We are trying to free memory region in a different pfn with current pbp. Flush pbp.
+            pbp.as_mut().unwrap().reset();
+            pbp.as_mut().unwrap().addr = align_page_size_down(range_base.0);
+        }
+
+        pbp.as_mut().unwrap().set_bit(range_base.0);
+        if pbp.as_ref().unwrap().bitmap_full() {
+            Self::release_memory_range(
+                memory,
+                vm_memory::GuestAddress(pbp.as_ref().unwrap().addr),
+                page_size as usize,
+            )?;
+
+            pbp.as_mut().unwrap().reset();
+        }
+
+        Ok(())
+    }
+
     fn process_queue(&mut self, queue_index: usize) -> result::Result<(), Error> {
         let mut used_descs = false;
         while let Some(mut desc_chain) =
@@ -198,7 +281,7 @@ impl BalloonEpollHandler {
 
                 match queue_index {
                     0 => {
-                        Self::release_memory_range(desc_chain.memory(), range_base, range_len)?;
+                        Self::release_memory_range_4k(&mut self.pbp, desc_chain.memory(), pfn)?;
                     }
                     1 => {
                         Self::advise_memory_range(
@@ -543,6 +626,7 @@ impl VirtioDevice for Balloon {
             reporting_queue_evt,
             kill_evt,
             pause_evt,
+            pbp: None,
         };
 
         let paused = self.common.paused.clone();