From 800220acbb9b3a2ea186a84bbc9d114229d9886f Mon Sep 17 00:00:00 2001
From: Hui Zhu <teawater@antfin.com>
Date: Tue, 23 Jun 2020 17:52:30 +0800
Subject: [PATCH] virtio-balloon: Store the balloon size to support reboot

This commit store balloon size to MemoryConfig.
After reboot, virtio-balloon can use this size to inflate back to
the size before reboot.

Signed-off-by: Hui Zhu <teawater@antfin.com>
---
 src/main.rs                   |  1 +
 virtio-devices/src/balloon.rs |  5 +++--
 vmm/src/config.rs             |  5 +++++
 vmm/src/device_manager.rs     |  7 +++++--
 vmm/src/memory_manager.rs     | 13 ++++++-------
 vmm/src/vm.rs                 |  5 ++++-
 6 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 15a69e719..1d049c559 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -503,6 +503,7 @@ mod unit_tests {
                     shared: false,
                     hugepages: false,
                     balloon: false,
+                    balloon_size: 0,
                 },
                 kernel: Some(KernelConfig {
                     path: PathBuf::from("/path/to/kernel"),
diff --git a/virtio-devices/src/balloon.rs b/virtio-devices/src/balloon.rs
index 794f9ddd0..0a7b7d1be 100644
--- a/virtio-devices/src/balloon.rs
+++ b/virtio-devices/src/balloon.rs
@@ -410,10 +410,11 @@ pub struct Balloon {
 
 impl Balloon {
     // Create a new virtio-balloon.
-    pub fn new(id: String) -> io::Result<Self> {
+    pub fn new(id: String, size: u64) -> io::Result<Self> {
         let avail_features = 1u64 << VIRTIO_F_VERSION_1;
 
-        let config = VirtioBalloonConfig::default();
+        let mut config = VirtioBalloonConfig::default();
+        config.num_pages = (size >> PAGE_SHIFT) as u32;
 
         Ok(Balloon {
             id,
diff --git a/vmm/src/config.rs b/vmm/src/config.rs
index 15171074f..493a93fe7 100644
--- a/vmm/src/config.rs
+++ b/vmm/src/config.rs
@@ -340,6 +340,8 @@ pub struct MemoryConfig {
     pub hugepages: bool,
     #[serde(default)]
     pub balloon: bool,
+    #[serde(default)]
+    pub balloon_size: u64,
 }
 
 impl MemoryConfig {
@@ -400,6 +402,7 @@ impl MemoryConfig {
             shared,
             hugepages,
             balloon,
+            balloon_size: 0,
         })
     }
 }
@@ -415,6 +418,7 @@ impl Default for MemoryConfig {
             shared: false,
             hugepages: false,
             balloon: false,
+            balloon_size: 0,
         }
     }
 }
@@ -1873,6 +1877,7 @@ mod tests {
                 shared: false,
                 hugepages: false,
                 balloon: false,
+                balloon_size: 0,
             },
             kernel: Some(KernelConfig {
                 path: PathBuf::from("/path/to/kernel"),
diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index bfb20688b..02abad770 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -2309,8 +2309,11 @@ impl DeviceManager {
             let id = String::from(BALLOON_DEVICE_NAME);
 
             let virtio_balloon_device = Arc::new(Mutex::new(
-                virtio_devices::Balloon::new(id.clone())
-                    .map_err(DeviceManagerError::CreateVirtioBalloon)?,
+                virtio_devices::Balloon::new(
+                    id.clone(),
+                    self.config.lock().unwrap().memory.balloon_size,
+                )
+                .map_err(DeviceManagerError::CreateVirtioBalloon)?,
             ));
 
             self.memory_manager
diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index 6d0334f33..e9eb8fc12 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -799,13 +799,12 @@ impl MemoryManager {
         Ok(())
     }
 
-    pub fn balloon_resize(&mut self, expected_ram: u64) -> Result<(), Error> {
+    pub fn balloon_resize(&mut self, expected_ram: u64) -> Result<u64, Error> {
+        let mut balloon_size = 0;
         if let Some(balloon) = &self.balloon {
-            let balloon_size = if expected_ram < self.current_ram {
-                self.current_ram - expected_ram
-            } else {
-                0
-            };
+            if expected_ram < self.current_ram {
+                balloon_size = self.current_ram - expected_ram;
+            }
             balloon
                 .lock()
                 .unwrap()
@@ -813,7 +812,7 @@ impl MemoryManager {
                 .map_err(Error::VirtioBalloonResizeFail)?;
         }
 
-        Ok(())
+        Ok(balloon_size)
     }
 
     /// In case this function resulted in adding a new memory region to the
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 397ba6e1a..701b9b1e5 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -738,7 +738,10 @@ impl Vm {
         }
 
         if let Some(desired_ram_w_balloon) = desired_ram_w_balloon {
-            self.memory_manager
+            // update the configuration value for the balloon size to ensure
+            // a reboot would use the right value.
+            self.config.lock().unwrap().memory.balloon_size = self
+                .memory_manager
                 .lock()
                 .unwrap()
                 .balloon_resize(desired_ram_w_balloon)