vmm: Add virtio-balloon support

This commit adds new option balloon to memory config.
Set it to on will open the balloon function.

Signed-off-by: Hui Zhu <teawater@antfin.com>

src/main.rs

@@ -502,6 +502,7 @@ mod unit_tests {
                     hotplug_size: None,
                     shared: false,
                     hugepages: false,
+                    balloon: false,
                 },
                 kernel: Some(KernelConfig {
                     path: PathBuf::from("/path/to/kernel"),

vmm/src/config.rs

@@ -338,6 +338,8 @@ pub struct MemoryConfig {
     pub shared: bool,
     #[serde(default)]
     pub hugepages: bool,
+    #[serde(default)]
+    pub balloon: bool,
 }
 
 impl MemoryConfig {
@@ -350,7 +352,8 @@ impl MemoryConfig {
             .add("hotplug_method")
             .add("hotplug_size")
             .add("shared")
-            .add("hugepages");
+            .add("hugepages")
+            .add("balloon");
         parser.parse(memory).map_err(Error::ParseMemory)?;
 
         let size = parser
@@ -382,6 +385,11 @@ impl MemoryConfig {
             .map_err(Error::ParseMemory)?
             .unwrap_or(Toggle(false))
             .0;
+        let balloon = parser
+            .convert::<Toggle>("balloon")
+            .map_err(Error::ParseMemory)?
+            .unwrap_or(Toggle(false))
+            .0;
 
         Ok(MemoryConfig {
             size,
@@ -391,6 +399,7 @@ impl MemoryConfig {
             hotplug_size,
             shared,
             hugepages,
+            balloon,
         })
     }
 }
@@ -405,6 +414,7 @@ impl Default for MemoryConfig {
             hotplug_size: None,
             shared: false,
             hugepages: false,
+            balloon: false,
         }
     }
 }
@@ -1862,6 +1872,7 @@ mod tests {
                 hotplug_size: None,
                 shared: false,
                 hugepages: false,
+                balloon: false,
             },
             kernel: Some(KernelConfig {
                 path: PathBuf::from("/path/to/kernel"),

vmm/src/device_manager.rs

@@ -104,6 +104,7 @@ const CONSOLE_DEVICE_NAME: &str = "_console";
 const DISK_DEVICE_NAME_PREFIX: &str = "_disk";
 const FS_DEVICE_NAME_PREFIX: &str = "_fs";
 const MEM_DEVICE_NAME: &str = "_mem";
+const BALLOON_DEVICE_NAME: &str = "_balloon";
 const NET_DEVICE_NAME_PREFIX: &str = "_net";
 const PMEM_DEVICE_NAME_PREFIX: &str = "_pmem";
 const RNG_DEVICE_NAME: &str = "_rng";
@@ -165,6 +166,9 @@ pub enum DeviceManagerError {
     /// Cannot create virtio-iommu device
     CreateVirtioIommu(io::Error),
 
+    /// Cannot create virtio-balloon device
+    CreateVirtioBalloon(io::Error),
+
     /// Failed parsing disk image format
     DetectImageType(qcow::Error),
 
@@ -1500,6 +1504,9 @@ impl DeviceManager {
 
         devices.append(&mut self.make_virtio_mem_devices()?);
 
+        // Add virtio-balloon if required
+        devices.append(&mut self.make_virtio_balloon_devices()?);
+
         Ok(devices)
     }
 
@@ -2293,6 +2300,39 @@ impl DeviceManager {
         Ok(devices)
     }
 
+    fn make_virtio_balloon_devices(
+        &mut self,
+    ) -> DeviceManagerResult<Vec<(VirtioDeviceArc, bool, String)>> {
+        let mut devices = Vec::new();
+
+        if self.config.lock().unwrap().memory.balloon {
+            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)?,
+            ));
+
+            self.memory_manager
+                .lock()
+                .unwrap()
+                .set_balloon(virtio_balloon_device.clone());
+
+            devices.push((
+                Arc::clone(&virtio_balloon_device) as VirtioDeviceArc,
+                false,
+                id.clone(),
+            ));
+
+            self.device_tree
+                .lock()
+                .unwrap()
+                .insert(id.clone(), device_node!(id, virtio_balloon_device));
+        }
+
+        Ok(devices)
+    }
+
     #[cfg(feature = "pci_support")]
     fn create_kvm_device(vm: &Arc<dyn hypervisor::Vm>) -> DeviceManagerResult<DeviceFd> {
         let mut vfio_dev = hypervisor::kvm::kvm_create_device {

vmm/src/memory_manager.rs

@@ -70,6 +70,7 @@ pub struct MemoryManager {
     snapshot: Mutex<Option<GuestMemoryLoadGuard<GuestMemoryMmap>>>,
     shared: bool,
     hugepages: bool,
+    balloon: Option<Arc<Mutex<virtio_devices::Balloon>>>,
 }
 
 #[derive(Debug)]
@@ -122,6 +123,9 @@ pub enum Error {
     /// The number of external backing files doesn't match the number of
     /// memory regions.
     InvalidAmountExternalBackingFiles,
+
+    /// Failed to virtio-balloon resize
+    VirtioBalloonResizeFail(virtio_devices::balloon::Error),
 }
 
 const ENABLE_FLAG: usize = 0;
@@ -338,6 +342,7 @@ impl MemoryManager {
             snapshot: Mutex::new(None),
             shared: config.shared,
             hugepages: config.hugepages,
+            balloon: None,
         }));
 
         guest_memory.memory().with_regions(|_, region| {
@@ -645,6 +650,10 @@ impl MemoryManager {
         Ok(region)
     }
 
+    pub fn set_balloon(&mut self, balloon: Arc<Mutex<virtio_devices::Balloon>>) {
+        self.balloon = Some(balloon);
+    }
+
     pub fn guest_memory(&self) -> GuestMemoryAtomic<GuestMemoryMmap> {
         self.guest_memory.clone()
     }
@@ -790,6 +799,23 @@ impl MemoryManager {
         Ok(())
     }
 
+    pub fn balloon_resize(&mut self, expected_ram: u64) -> Result<(), Error> {
+        if let Some(balloon) = &self.balloon {
+            let balloon_size = if expected_ram < self.current_ram {
+                self.current_ram - expected_ram
+            } else {
+                0
+            };
+            balloon
+                .lock()
+                .unwrap()
+                .resize(balloon_size)
+                .map_err(Error::VirtioBalloonResizeFail)?;
+        }
+
+        Ok(())
+    }
+
     /// In case this function resulted in adding a new memory region to the
     /// guest memory, the new region is returned to the caller. The virtio-mem
     /// use case never adds a new region as the whole hotpluggable memory has