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

@@ -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"),

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,

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"),

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

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

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)