vmm: Move balloon code from MemoryManager to DeviceManager

Now that we have a new dedicated way of asking for a balloon through the
CLI and the REST API, we can move all the balloon code to the device
manager. This allows us to simplify the memory manager, which is already
quite complex.

It also simplifies the behavior of the balloon resizing command. Instead
of providing the expected size for the RAM, which is complex when memory
zones are involved, it now expects the balloon size. This is a much more
straightforward behavior as it really resizes the balloon to the desired
size. Additionally to the simplication, the benefit of this approach is
that it does not need to be tied to the memory manager at all.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>

src/bin/ch-remote.rs

@@ -226,7 +226,7 @@ fn resize_api_command(
         None
     };
 
-    let desired_ram_w_balloon: Option<u64> = if let Some(balloon) = balloon {
+    let desired_balloon: Option<u64> = if let Some(balloon) = balloon {
         Some(
             balloon
                 .parse::<ByteSized>()
@@ -240,7 +240,7 @@ fn resize_api_command(
     let resize = vmm::api::VmResizeData {
         desired_vcpus,
         desired_ram,
-        desired_ram_w_balloon,
+        desired_balloon,
     };
 
     simple_api_command(
@@ -577,7 +577,7 @@ fn main() {
                 .arg(
                     Arg::with_name("balloon")
                         .long("balloon")
-                        .help("New memory with balloon size in bytes (supports K/M/G suffix)")
+                        .help("New balloon size in bytes (supports K/M/G suffix)")
                         .takes_value(true)
                         .number_of_values(1),
                 ),

vmm/src/api/mod.rs

@@ -157,7 +157,7 @@ pub struct VmmPingResponse {
 pub struct VmResizeData {
     pub desired_vcpus: Option<u8>,
     pub desired_ram: Option<u64>,
-    pub desired_ram_w_balloon: Option<u64>,
+    pub desired_balloon: Option<u64>,
 }
 
 #[derive(Clone, Deserialize, Serialize, Default)]

vmm/src/device_manager.rs

@@ -358,6 +358,12 @@ pub enum DeviceManagerError {
 
     /// No support for device passthrough
     NoDevicePassthroughSupport,
+
+    /// Failed to resize virtio-balloon
+    VirtioBalloonResize(virtio_devices::balloon::Error),
+
+    /// Missing virtio-balloon, can't proceed as expected.
+    MissingVirtioBalloon,
 }
 pub type DeviceManagerResult<T> = result::Result<T, DeviceManagerError>;
 
@@ -780,6 +786,9 @@ pub struct DeviceManager {
     // List of guest NUMA nodes.
     #[cfg(feature = "acpi")]
     numa_nodes: NumaNodes,
+
+    // Possible handle to the virtio-balloon device
+    balloon: Option<Arc<Mutex<virtio_devices::Balloon>>>,
 }
 
 impl DeviceManager {
@@ -850,6 +859,7 @@ impl DeviceManager {
             seccomp_action,
             #[cfg(feature = "acpi")]
             numa_nodes,
+            balloon: None,
         };
 
         #[cfg(feature = "acpi")]
@@ -2440,22 +2450,19 @@ impl DeviceManager {
     ) -> DeviceManagerResult<Vec<(VirtioDeviceArc, bool, String)>> {
         let mut devices = Vec::new();
 
-        if self.config.lock().unwrap().memory.balloon {
+        if let Some(balloon_config) = &self.config.lock().unwrap().balloon {
             let id = String::from(BALLOON_DEVICE_NAME);
 
             let virtio_balloon_device = Arc::new(Mutex::new(
                 virtio_devices::Balloon::new(
                     id.clone(),
-                    self.config.lock().unwrap().memory.balloon_size,
+                    balloon_config.size,
                     self.seccomp_action.clone(),
                 )
                 .map_err(DeviceManagerError::CreateVirtioBalloon)?,
             ));
 
-            self.memory_manager
+            self.balloon = Some(virtio_balloon_device.clone());
-                .lock()
-                .unwrap()
-                .set_balloon(virtio_balloon_device.clone());
 
             devices.push((
                 Arc::clone(&virtio_balloon_device) as VirtioDeviceArc,
@@ -3197,6 +3204,27 @@ impl DeviceManager {
 
         counters
     }
+
+    pub fn resize_balloon(&mut self, size: u64) -> DeviceManagerResult<()> {
+        if let Some(balloon) = &self.balloon {
+            return balloon
+                .lock()
+                .unwrap()
+                .resize(size)
+                .map_err(DeviceManagerError::VirtioBalloonResize);
+        }
+
+        warn!("No balloon setup: Can't resize the balloon");
+        Err(DeviceManagerError::MissingVirtioBalloon)
+    }
+
+    pub fn balloon_size(&self) -> u64 {
+        if let Some(balloon) = &self.balloon {
+            return balloon.lock().unwrap().get_actual();
+        }
+
+        0
+    }
 }
 
 #[cfg(feature = "acpi")]

vmm/src/lib.rs

@@ -469,7 +469,7 @@ impl Vmm {
 
                 let mut memory_actual_size = config.lock().unwrap().memory.total_size();
                 if let Some(vm) = &self.vm {
-                    memory_actual_size -= vm.get_balloon_actual();
+                    memory_actual_size -= vm.balloon_size();
                 }
 
                 Ok(VmInfo {
@@ -511,10 +511,10 @@ impl Vmm {
         &mut self,
         desired_vcpus: Option<u8>,
         desired_ram: Option<u64>,
-        desired_ram_w_balloon: Option<u64>,
+        desired_balloon: Option<u64>,
     ) -> result::Result<(), VmError> {
         if let Some(ref mut vm) = self.vm {
-            if let Err(e) = vm.resize(desired_vcpus, desired_ram, desired_ram_w_balloon) {
+            if let Err(e) = vm.resize(desired_vcpus, desired_ram, desired_balloon) {
                 error!("Error when resizing VM: {:?}", e);
                 Err(e)
             } else {
@@ -801,7 +801,7 @@ impl Vmm {
                                         .vm_resize(
                                             resize_data.desired_vcpus,
                                             resize_data.desired_ram,
-                                            resize_data.desired_ram_w_balloon,
+                                            resize_data.desired_balloon,
                                         )
                                         .map_err(ApiError::VmResize)
                                         .map(|_| ApiResponsePayload::Empty);

vmm/src/memory_manager.rs

@@ -117,7 +117,6 @@ pub struct MemoryManager {
     snapshot: Mutex<Option<GuestMemoryLoadGuard<GuestMemoryMmap>>>,
     shared: bool,
     hugepages: bool,
-    balloon: Option<Arc<Mutex<virtio_devices::Balloon>>>,
     #[cfg(target_arch = "x86_64")]
     sgx_epc_region: Option<SgxEpcRegion>,
     user_provided_zones: bool,
@@ -176,9 +175,6 @@ pub enum Error {
     /// memory regions.
     InvalidAmountExternalBackingFiles,
 
-    /// Failed to virtio-balloon resize
-    VirtioBalloonResizeFail(virtio_devices::balloon::Error),
-
     /// Invalid SGX EPC section size
     #[cfg(target_arch = "x86_64")]
     EpcSectionSizeInvalid,
@@ -688,7 +684,6 @@ impl MemoryManager {
             snapshot: Mutex::new(None),
             shared: config.shared,
             hugepages: config.hugepages,
-            balloon: None,
             #[cfg(target_arch = "x86_64")]
             sgx_epc_region: None,
             user_provided_zones,
@@ -1085,10 +1080,6 @@ 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()
     }
@@ -1242,30 +1233,6 @@ impl MemoryManager {
         Err(Error::UnknownMemoryZone)
     }
 
-    pub fn balloon_resize(&mut self, expected_ram: u64) -> Result<u64, Error> {
-        let mut balloon_size = 0;
-        if let Some(balloon) = &self.balloon {
-            if expected_ram < self.current_ram {
-                balloon_size = self.current_ram - expected_ram;
-            }
-            balloon
-                .lock()
-                .unwrap()
-                .resize(balloon_size)
-                .map_err(Error::VirtioBalloonResizeFail)?;
-        }
-
-        Ok(balloon_size)
-    }
-
-    pub fn get_balloon_actual(&self) -> u64 {
-        if let Some(balloon) = &self.balloon {
-            return balloon.lock().unwrap().get_actual();
-        }
-
-        0
-    }
-
     /// 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

vmm/src/vm.rs

@@ -1040,7 +1040,7 @@ impl Vm {
         &mut self,
         desired_vcpus: Option<u8>,
         desired_memory: Option<u64>,
-        desired_ram_w_balloon: Option<u64>,
+        desired_balloon: Option<u64>,
     ) -> Result<()> {
         if let Some(desired_vcpus) = desired_vcpus {
             if self
@@ -1103,15 +1103,18 @@ impl Vm {
             }
         }
 
-        if let Some(desired_ram_w_balloon) = desired_ram_w_balloon {
+        if let Some(desired_balloon) = desired_balloon {
-            // update the configuration value for the balloon size to ensure
+            self.device_manager
-            // 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)
+                .resize_balloon(desired_balloon)
-                .map_err(Error::MemoryManager)?;
+                .map_err(Error::DeviceManager)?;
+
+            // Update the configuration value for the balloon size to ensure
+            // a reboot would use the right value.
+            if let Some(balloon_config) = &mut self.config.lock().unwrap().balloon {
+                balloon_config.size = desired_balloon;
+            }
         }
 
         Ok(())
@@ -1607,8 +1610,8 @@ impl Vm {
     }
 
     /// Gets the actual size of the balloon.
-    pub fn get_balloon_actual(&self) -> u64 {
+    pub fn balloon_size(&self) -> u64 {
-        self.memory_manager.lock().unwrap().get_balloon_actual()
+        self.device_manager.lock().unwrap().balloon_size()
     }
 }