fix(balloon): fix balloon not found when pci is enabled

The code managing the balloon logic is only looking at the mmio device
manager. Make it use the generic device manager to find the device.

Signed-off-by: Riccardo Mancini <[email protected]>
Signed-off-by: Babis Chalios <[email protected]>

Author: Manciukic

src/vmm/src/lib.rs

@@ -123,7 +123,6 @@ use std::time::Duration;
 
 use device_manager::DeviceManager;
 use devices::acpi::vmgenid::VmGenIdError;
-use devices::virtio::device::VirtioDevice;
 use event_manager::{EventManager as BaseEventManager, EventOps, Events, MutEventSubscriber};
 use seccomp::BpfProgram;
 use snapshot::Persist;
@@ -329,20 +328,6 @@ impl Vmm {
         self.shutdown_exit_code
     }
 
-    /// Gets the specified bus device.
-    pub fn get_virtio_device(
-        &self,
-        device_type: u32,
-        device_id: &str,
-    ) -> Option<Arc<Mutex<dyn VirtioDevice>>> {
-        let device = self
-            .device_manager
-            .mmio_devices
-            .get_virtio_device(device_type, device_id)?;
-
-        Some(device.inner.lock().expect("Poisoned lock").device().clone())
-    }
-
     /// Starts the microVM vcpus.
     ///
     /// # Errors
@@ -589,7 +574,10 @@ impl Vmm {
 
     /// Returns a reference to the balloon device if present.
     pub fn balloon_config(&self) -> Result<BalloonConfig, BalloonError> {
-        if let Some(virtio_device) = self.get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID) {
+        if let Some(virtio_device) = self
+            .device_manager
+            .get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID)
+        {
             let config = virtio_device
                 .lock()
                 .expect("Poisoned lock")
@@ -606,7 +594,10 @@ impl Vmm {
 
     /// Returns the latest balloon statistics if they are enabled.
     pub fn latest_balloon_stats(&self) -> Result<BalloonStats, BalloonError> {
-        if let Some(virtio_device) = self.get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID) {
+        if let Some(virtio_device) = self
+            .device_manager
+            .get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID)
+        {
             let latest_stats = virtio_device
                 .lock()
                 .expect("Poisoned lock")
@@ -631,7 +622,10 @@ impl Vmm {
             return Err(BalloonError::TooManyPagesRequested);
         }
 
-        if let Some(virtio_device) = self.get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID) {
+        if let Some(virtio_device) = self
+            .device_manager
+            .get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID)
+        {
             {
                 virtio_device
                     .lock()
@@ -653,7 +647,10 @@ impl Vmm {
         &mut self,
         stats_polling_interval_s: u16,
     ) -> Result<(), BalloonError> {
-        if let Some(virtio_device) = self.get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID) {
+        if let Some(virtio_device) = self
+            .device_manager
+            .get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID)
+        {
             {
                 virtio_device
                     .lock()