From ad9968ce2d796f17c5e044df0ae9b6e39be0f68f Mon Sep 17 00:00:00 2001
From: Hui Zhu <teawater@antgroup.com>
Date: Tue, 3 Sep 2024 17:56:09 +0800
Subject: [PATCH] runtime-rs: Add enable_balloon_f_reporting for dragonball

Under normal circumstances, the virtual machine only requests memory
from the host and does not actively release it back to host when it is
no longer needed, leading to a waste of memory resources.

Free page reporting is a sub-feature of virtio-balloon. When this
feature is enabled, the Linux guest kernel will send information about
released pages to dragonball via virtio-balloon, and dragonball will
then release these pages.

This commit adds an option enable_balloon_f_reporting to runtime-rs.
When this option is enabled, runtime-rs will insert a virtio-balloon
device with the f_reporting option enabled during the Dragonball virtual
machine startup.

Signed-off-by: Hui Zhu <teawater@antgroup.com>
---
 .../kata-types/src/config/hypervisor/mod.rs   |  6 +++++
 .../crates/hypervisor/src/dragonball/inner.rs | 23 +++++++++++++++----
 .../crates/hypervisor/src/dragonball/mod.rs   | 21 ++++++++++++++++-
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/libs/kata-types/src/config/hypervisor/mod.rs b/src/libs/kata-types/src/config/hypervisor/mod.rs
index f17cc87239..f5fff4ba06 100644
--- a/src/libs/kata-types/src/config/hypervisor/mod.rs
+++ b/src/libs/kata-types/src/config/hypervisor/mod.rs
@@ -491,6 +491,12 @@ pub struct DeviceInfo {
     /// Enabling this will result in the VM device having iommu_platform=on set
     #[serde(default)]
     pub enable_iommu_platform: bool,
+
+    /// Enable balloon f_reporting, default false
+    ///
+    /// Enabling this will result in the VM balloon device having f_reporting=on set
+    #[serde(default)]
+    pub enable_balloon_f_reporting: bool,
 }
 
 impl DeviceInfo {
diff --git a/src/runtime-rs/crates/hypervisor/src/dragonball/inner.rs b/src/runtime-rs/crates/hypervisor/src/dragonball/inner.rs
index ab51f8dcc4..57613813dd 100644
--- a/src/runtime-rs/crates/hypervisor/src/dragonball/inner.rs
+++ b/src/runtime-rs/crates/hypervisor/src/dragonball/inner.rs
@@ -6,8 +6,9 @@
 
 use super::vmm_instance::VmmInstance;
 use crate::{
-    device::DeviceType, hypervisor_persist::HypervisorState, kernel_param::KernelParams,
-    MemoryConfig, VmmState, DEV_HUGEPAGES, HUGETLBFS, HUGE_SHMEM, HYPERVISOR_DRAGONBALL, SHMEM,
+    device::DeviceType, firecracker::sl, hypervisor_persist::HypervisorState,
+    kernel_param::KernelParams, MemoryConfig, VmmState, DEV_HUGEPAGES, HUGETLBFS, HUGE_SHMEM,
+    HYPERVISOR_DRAGONBALL, SHMEM,
 };
 use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
@@ -118,6 +119,20 @@ impl DragonballInner {
         }
     }
 
+    pub(crate) async fn try_insert_balloon_f_reporting(&mut self) {
+        let balloon_config = BalloonDeviceConfigInfo {
+            balloon_id: BALLOON_DEVICE_ID.to_owned(),
+            size_mib: 0,
+            use_shared_irq: None,
+            use_generic_irq: None,
+            f_deflate_on_oom: false,
+            f_reporting: true,
+        };
+        if let Err(e) = self.vmm_instance.insert_balloon_device(balloon_config) {
+            error!(sl(), "failed to insert f_reporting balloon device: {:?}", e);
+        }
+    }
+
     pub(crate) async fn cold_start_vm(&mut self, timeout: i32) -> Result<()> {
         info!(sl!(), "start sandbox cold");
 
@@ -415,7 +430,7 @@ impl DragonballInner {
                         use_shared_irq: None,
                         use_generic_irq: None,
                         f_deflate_on_oom: false,
-                        f_reporting: false,
+                        f_reporting: self.config.device_info.enable_balloon_f_reporting,
                     };
                     self.vmm_instance
                         .insert_balloon_device(balloon_config)
@@ -447,7 +462,7 @@ impl DragonballInner {
                     use_shared_irq: None,
                     use_generic_irq: None,
                     f_deflate_on_oom: false,
-                    f_reporting: false,
+                    f_reporting: self.config.device_info.enable_balloon_f_reporting,
                 };
                 self.balloon_size = had_mem_mb - new_mem_mb;
                 self.vmm_instance
diff --git a/src/runtime-rs/crates/hypervisor/src/dragonball/mod.rs b/src/runtime-rs/crates/hypervisor/src/dragonball/mod.rs
index 9313854714..36ac15f932 100644
--- a/src/runtime-rs/crates/hypervisor/src/dragonball/mod.rs
+++ b/src/runtime-rs/crates/hypervisor/src/dragonball/mod.rs
@@ -77,7 +77,26 @@ impl Hypervisor for Dragonball {
     #[instrument]
     async fn start_vm(&self, timeout: i32) -> Result<()> {
         let mut inner = self.inner.write().await;
-        inner.start_vm(timeout).await
+        let ret = inner.start_vm(timeout).await;
+
+        if ret.is_ok() && inner.config.device_info.enable_balloon_f_reporting {
+            // The virtio-balloon device must be inserted into dragonball and
+            // recognized by the guest kernel only after the dragonball upcall is ready.
+            // The dragonball upcall is not ready immediately after the VM starts,
+            // so here we create an asynchronous task that waits for 5 seconds before
+            // inserting the virtio-balloon device.
+            let inner_clone = self.inner.clone();
+            tokio::spawn(async move {
+                tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+                inner_clone
+                    .write()
+                    .await
+                    .try_insert_balloon_f_reporting()
+                    .await;
+            });
+        }
+
+        ret
     }
 
     async fn stop_vm(&self) -> Result<()> {