From f2dd1f9f7934b275f05823626fc9ad20aef23c67 Mon Sep 17 00:00:00 2001
From: Changyuan Lyu <changyuanl@google.com>
Date: Sun, 20 Oct 2024 17:55:15 -0700
Subject: [PATCH] feat(balloon): implement inflate queue

Signed-off-by: Changyuan Lyu <changyuanl@google.com>
---
 alioth/src/mem/mapped.rs         | 14 +++++++++++---
 alioth/src/virtio/dev/balloon.rs | 10 +++++++---
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/alioth/src/mem/mapped.rs b/alioth/src/mem/mapped.rs
index 55c957e..07846e8 100644
--- a/alioth/src/mem/mapped.rs
+++ b/alioth/src/mem/mapped.rs
@@ -26,11 +26,11 @@ use std::ptr::{null_mut, NonNull};
 use std::sync::Arc;
 
 use libc::{
-    c_void, mmap, msync, munmap, MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, MAP_SHARED, MS_ASYNC,
-    PROT_EXEC, PROT_READ, PROT_WRITE,
+    c_void, madvise, mmap, msync, munmap, MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, MAP_SHARED,
+    MS_ASYNC, PROT_EXEC, PROT_READ, PROT_WRITE,
 };
 #[cfg(target_os = "linux")]
-use libc::{madvise, MADV_HUGEPAGE, MFD_CLOEXEC};
+use libc::{MADV_HUGEPAGE, MFD_CLOEXEC};
 use parking_lot::{RwLock, RwLockReadGuard};
 use snafu::ResultExt;
 use zerocopy::{FromBytes, Immutable, IntoBytes};
@@ -404,6 +404,14 @@ impl Ram {
     pub fn iter(&self) -> impl DoubleEndedIterator<Item = (u64, &ArcMemPages)> {
         self.inner.iter()
     }
+
+    pub fn madvise(&self, gpa: u64, size: u64, advice: i32) -> Result<()> {
+        for r in self.slice_iter_mut(gpa, size) {
+            let s = r?;
+            ffi!(unsafe { madvise(s.as_mut_ptr() as _, s.len(), advice) })?;
+        }
+        Ok(())
+    }
 }
 
 impl Default for RamBus {
diff --git a/alioth/src/virtio/dev/balloon.rs b/alioth/src/virtio/dev/balloon.rs
index 73b06d2..f96a76e 100644
--- a/alioth/src/virtio/dev/balloon.rs
+++ b/alioth/src/virtio/dev/balloon.rs
@@ -143,9 +143,9 @@ impl Balloon {
         })
     }
 
-    fn inflate(&self, desc: &[IoSlice], _ram: &Ram) {
+    fn inflate(&self, desc: &[IoSlice], ram: &Ram) {
         for buf in desc {
-            for bytes in buf.chunks(4) {
+            for bytes in buf.chunks(size_of::<u32>()) {
                 let Ok(page_num) = u32::read_from_bytes(bytes) else {
                     log::error!(
                         "{}: inflate: invalid page_num bytes: {bytes:02x?}",
@@ -154,7 +154,11 @@ impl Balloon {
                     continue;
                 };
                 let gpa = (page_num as u64) << 12;
-                log::info!("{}: TODO: inflating GPA {gpa:#x}", self.name);
+                if let Err(e) = ram.madvise(gpa, 1 << 12, libc::MADV_DONTNEED) {
+                    log::error!("{}: inflate at GPA {gpa:#x}: {e:?}", self.name);
+                } else {
+                    log::trace!("{}: freed GPA {gpa:#x}", self.name);
+                }
             }
         }
     }