refactor: reduce boilerplate to execute fn on virtio device

Make the with_virtio_device_with_id function more generic and introduce
its variant with error handling try_with_virtio_device_with_id. Then use
these new functions in vmm.rs whenever we need to perform an action on a
device.

To simplify the code, I also moved some balloon device error handling
back to the balloon device code and refactored it a little bit.

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

Author: Manciukic

src/vmm/src/device_manager/mod.rs

@@ -89,7 +89,7 @@ pub enum FindDeviceError {
     /// Device not found
     DeviceNotFound,
     /// Internal Device error: {0}
-    InternalDeviceError(String),
+    InternalDeviceError(anyhow::Error),
 }
 
 #[derive(Debug)]
@@ -355,27 +355,43 @@ impl DeviceManager {
     }
 
     /// Run fn `f()` for the virtio device matching `virtio_type` and `id`.
-    pub fn with_virtio_device_with_id<T, F>(
+    pub fn try_with_virtio_device_with_id<T, F, R, E>(
         &self,
         virtio_type: u32,
         id: &str,
         f: F,
-    ) -> Result<(), FindDeviceError>
+    ) -> Result<R, FindDeviceError>
     where
         T: VirtioDevice + 'static + Debug,
-        F: FnOnce(&mut T) -> Result<(), String>,
+        E: std::error::Error + 'static + Send + Sync,
+        F: FnOnce(&mut T) -> Result<R, E>,
     {
         if let Some(device) = self.get_virtio_device(virtio_type, id) {
             let mut dev = device.lock().expect("Poisoned lock");
             f(dev
                 .as_mut_any()
                 .downcast_mut::<T>()
                 .ok_or(FindDeviceError::InvalidDeviceType)?)
-            .map_err(FindDeviceError::InternalDeviceError)?;
+            .map_err(|e| FindDeviceError::InternalDeviceError(e.into()))
         } else {
-            return Err(FindDeviceError::DeviceNotFound);
+            Err(FindDeviceError::DeviceNotFound)
         }
-        Ok(())
+    }
+
+    /// Run fn `f()` for the virtio device matching `virtio_type` and `id`.
+    pub fn with_virtio_device_with_id<T, F, R>(
+        &self,
+        virtio_type: u32,
+        id: &str,
+        f: F,
+    ) -> Result<R, FindDeviceError>
+    where
+        T: VirtioDevice + 'static + Debug,
+        F: FnOnce(&mut T) -> R,
+    {
+        self.try_with_virtio_device_with_id(virtio_type, id, |dev: &mut T| {
+            Ok::<R, FindDeviceError>(f(dev))
+        })
     }
 }
 

src/vmm/src/devices/virtio/balloon/device.rs

@@ -31,6 +31,7 @@ use crate::devices::virtio::generated::virtio_ids::VIRTIO_ID_BALLOON;
 use crate::devices::virtio::queue::InvalidAvailIdx;
 use crate::devices::virtio::transport::{VirtioInterrupt, VirtioInterruptType};
 use crate::logger::IncMetric;
+use crate::mem_size_mib;
 use crate::utils::u64_to_usize;
 use crate::vstate::memory::{Address, ByteValued, Bytes, GuestAddress, GuestMemoryMmap};
 
@@ -450,6 +451,12 @@ impl Balloon {
     /// Update the target size of the balloon.
     pub fn update_size(&mut self, amount_mib: u32) -> Result<(), BalloonError> {
         if self.is_activated() {
+            let mem = &self.device_state.active_state().unwrap().mem;
+            // The balloon cannot have a target size greater than the size of
+            // the guest memory.
+            if u64::from(amount_mib) > mem_size_mib(mem) {
+                return Err(BalloonError::TooManyPagesRequested);
+            }
             self.config_space.num_pages = mib_to_pages(amount_mib)?;
             self.interrupt_trigger()
                 .trigger(VirtioInterruptType::Config)
@@ -504,15 +511,15 @@ impl Balloon {
     }
 
     /// Retrieve latest stats for the balloon device.
-    pub fn latest_stats(&mut self) -> Option<&BalloonStats> {
+    pub fn latest_stats(&mut self) -> Result<&BalloonStats, BalloonError> {
         if self.stats_enabled() {
             self.latest_stats.target_pages = self.config_space.num_pages;
             self.latest_stats.actual_pages = self.config_space.actual_pages;
             self.latest_stats.target_mib = pages_to_mib(self.latest_stats.target_pages);
             self.latest_stats.actual_mib = pages_to_mib(self.latest_stats.actual_pages);
-            Some(&self.latest_stats)
+            Ok(&self.latest_stats)
         } else {
-            None
+            Err(BalloonError::StatisticsDisabled)
         }
     }
 
@@ -1150,7 +1157,7 @@ pub(crate) mod tests {
         let mut balloon = Balloon::new(0, true, 0, false).unwrap();
         // Switch the state to active.
         balloon.device_state = DeviceState::Activated(ActiveState {
-            mem: single_region_mem(0x1),
+            mem: single_region_mem(32 << 20),
             interrupt: default_interrupt(),
         });
 

src/vmm/src/lib.rs

@@ -134,9 +134,7 @@ use vstate::kvm::Kvm;
 use vstate::vcpu::{self, StartThreadedError, VcpuSendEventError};
 
 use crate::cpu_config::templates::CpuConfiguration;
-use crate::devices::virtio::balloon::{
-    BALLOON_DEV_ID, Balloon, BalloonConfig, BalloonError, BalloonStats,
-};
+use crate::devices::virtio::balloon::{BALLOON_DEV_ID, Balloon, BalloonConfig, BalloonStats};
 use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::generated::virtio_ids;
 use crate::devices::virtio::mem::VirtioMemError;
@@ -525,14 +523,10 @@ impl Vmm {
         path_on_host: String,
     ) -> Result<(), VmmError> {
         self.device_manager
-            .with_virtio_device_with_id(
+            .try_with_virtio_device_with_id(
                 virtio_ids::VIRTIO_ID_BLOCK,
                 drive_id,
-                |block: &mut Block| {
-                    block
-                        .update_disk_image(path_on_host)
-                        .map_err(|err| err.to_string())
-                },
+                |block: &mut Block| block.update_disk_image(path_on_host),
             )
             .map_err(VmmError::FindDeviceError)
     }
@@ -545,25 +539,21 @@ impl Vmm {
         rl_ops: BucketUpdate,
     ) -> Result<(), VmmError> {
         self.device_manager
-            .with_virtio_device_with_id(
+            .try_with_virtio_device_with_id(
                 virtio_ids::VIRTIO_ID_BLOCK,
                 drive_id,
-                |block: &mut Block| {
-                    block
-                        .update_rate_limiter(rl_bytes, rl_ops)
-                        .map_err(|err| err.to_string())
-                },
+                |block: &mut Block| block.update_rate_limiter(rl_bytes, rl_ops),
             )
             .map_err(VmmError::FindDeviceError)
     }
 
     /// Updates the rate limiter parameters for block device with `drive_id` id.
     pub fn update_vhost_user_block_config(&mut self, drive_id: &str) -> Result<(), VmmError> {
         self.device_manager
-            .with_virtio_device_with_id(
+            .try_with_virtio_device_with_id(
                 virtio_ids::VIRTIO_ID_BLOCK,
                 drive_id,
-                |block: &mut Block| block.update_config().map_err(|err| err.to_string()),
+                |block: &mut Block| block.update_config(),
             )
             .map_err(VmmError::FindDeviceError)
     }
@@ -579,104 +569,56 @@ impl Vmm {
     ) -> Result<(), VmmError> {
         self.device_manager
             .with_virtio_device_with_id(virtio_ids::VIRTIO_ID_NET, net_id, |net: &mut Net| {
-                net.patch_rate_limiters(rx_bytes, rx_ops, tx_bytes, tx_ops);
-                Ok(())
+                net.patch_rate_limiters(rx_bytes, rx_ops, tx_bytes, tx_ops)
             })
             .map_err(VmmError::FindDeviceError)
     }
 
     /// Returns a reference to the balloon device if present.
-    pub fn balloon_config(&self) -> Result<BalloonConfig, BalloonError> {
-        if let Some(virtio_device) = self
-            .device_manager
-            .get_virtio_device(virtio_ids::VIRTIO_ID_BALLOON, BALLOON_DEV_ID)
-        {
-            let config = virtio_device
-                .lock()
-                .expect("Poisoned lock")
-                .as_mut_any()
-                .downcast_mut::<Balloon>()
-                .unwrap()
-                .config();
-
-            Ok(config)
-        } else {
-            Err(BalloonError::DeviceNotFound)
-        }
+    pub fn balloon_config(&self) -> Result<BalloonConfig, VmmError> {
+        self.device_manager
+            .with_virtio_device_with_id(
+                virtio_ids::VIRTIO_ID_BALLOON,
+                BALLOON_DEV_ID,
+                |dev: &mut Balloon| dev.config(),
+            )
+            .map_err(VmmError::FindDeviceError)
     }
 
     /// Returns the latest balloon statistics if they are enabled.
-    pub fn latest_balloon_stats(&self) -> Result<BalloonStats, BalloonError> {
-        if let Some(virtio_device) = self
-            .device_manager
-            .get_virtio_device(virtio_ids::VIRTIO_ID_BALLOON, BALLOON_DEV_ID)
-        {
-            let latest_stats = virtio_device
-                .lock()
-                .expect("Poisoned lock")
-                .as_mut_any()
-                .downcast_mut::<Balloon>()
-                .unwrap()
-                .latest_stats()
-                .ok_or(BalloonError::StatisticsDisabled)
-                .cloned()?;
-
-            Ok(latest_stats)
-        } else {
-            Err(BalloonError::DeviceNotFound)
-        }
+    pub fn latest_balloon_stats(&self) -> Result<BalloonStats, VmmError> {
+        self.device_manager
+            .try_with_virtio_device_with_id(
+                virtio_ids::VIRTIO_ID_BALLOON,
+                BALLOON_DEV_ID,
+                |dev: &mut Balloon| dev.latest_stats().cloned(),
+            )
+            .map_err(VmmError::FindDeviceError)
     }
 
     /// Updates configuration for the balloon device target size.
-    pub fn update_balloon_config(&mut self, amount_mib: u32) -> Result<(), BalloonError> {
-        // The balloon cannot have a target size greater than the size of
-        // the guest memory.
-        if u64::from(amount_mib) > mem_size_mib(self.vm.guest_memory()) {
-            return Err(BalloonError::TooManyPagesRequested);
-        }
-
-        if let Some(virtio_device) = self
-            .device_manager
-            .get_virtio_device(virtio_ids::VIRTIO_ID_BALLOON, BALLOON_DEV_ID)
-        {
-            {
-                virtio_device
-                    .lock()
-                    .expect("Poisoned lock")
-                    .as_mut_any()
-                    .downcast_mut::<Balloon>()
-                    .unwrap()
-                    .update_size(amount_mib)?;
-
-                Ok(())
-            }
-        } else {
-            Err(BalloonError::DeviceNotFound)
-        }
+    pub fn update_balloon_config(&mut self, amount_mib: u32) -> Result<(), VmmError> {
+        self.device_manager
+            .try_with_virtio_device_with_id(
+                virtio_ids::VIRTIO_ID_BALLOON,
+                BALLOON_DEV_ID,
+                |dev: &mut Balloon| dev.update_size(amount_mib),
+            )
+            .map_err(VmmError::FindDeviceError)
     }
 
     /// Updates configuration for the balloon device as described in `balloon_stats_update`.
     pub fn update_balloon_stats_config(
         &mut self,
         stats_polling_interval_s: u16,
-    ) -> Result<(), BalloonError> {
-        if let Some(virtio_device) = self
-            .device_manager
-            .get_virtio_device(virtio_ids::VIRTIO_ID_BALLOON, BALLOON_DEV_ID)
-        {
-            {
-                virtio_device
-                    .lock()
-                    .expect("Poisoned lock")
-                    .as_mut_any()
-                    .downcast_mut::<Balloon>()
-                    .unwrap()
-                    .update_stats_polling_interval(stats_polling_interval_s)?;
-            }
-            Ok(())
-        } else {
-            Err(BalloonError::DeviceNotFound)
-        }
+    ) -> Result<(), VmmError> {
+        self.device_manager
+            .try_with_virtio_device_with_id(
+                virtio_ids::VIRTIO_ID_BALLOON,
+                BALLOON_DEV_ID,
+                |dev: &mut Balloon| dev.update_stats_polling_interval(stats_polling_interval_s),
+            )
+            .map_err(VmmError::FindDeviceError)
     }
 
     /// Signals Vmm to stop and exit.

src/vmm/src/rpc_interface.rs

@@ -132,6 +132,8 @@ pub enum VmmAction {
 pub enum VmmActionError {
     /// Balloon config error: {0}
     BalloonConfig(#[from] BalloonConfigError),
+    /// Balloon update error: {0}
+    BalloonUpdate(VmmError),
     /// Boot source error: {0}
     BootSource(#[from] BootSourceConfigError),
     /// Create snapshot error: {0}
@@ -647,14 +649,14 @@ impl RuntimeApiController {
                 .expect("Poisoned lock")
                 .balloon_config()
                 .map(|state| VmmData::BalloonConfig(BalloonDeviceConfig::from(state)))
-                .map_err(|err| VmmActionError::BalloonConfig(BalloonConfigError::from(err))),
+                .map_err(VmmActionError::InternalVmm),
             GetBalloonStats => self
                 .vmm
                 .lock()
                 .expect("Poisoned lock")
                 .latest_balloon_stats()
                 .map(VmmData::BalloonStats)
-                .map_err(|err| VmmActionError::BalloonConfig(BalloonConfigError::from(err))),
+                .map_err(VmmActionError::InternalVmm),
             GetFullVmConfig => Ok(VmmData::FullVmConfig((&self.vm_resources).into())),
             GetMMDS => self.get_mmds(),
             GetVmMachineConfig => Ok(VmmData::MachineConfiguration(
@@ -678,14 +680,14 @@ impl RuntimeApiController {
                 .expect("Poisoned lock")
                 .update_balloon_config(balloon_update.amount_mib)
                 .map(|_| VmmData::Empty)
-                .map_err(|err| VmmActionError::BalloonConfig(BalloonConfigError::from(err))),
+                .map_err(VmmActionError::BalloonUpdate),
             UpdateBalloonStatistics(balloon_stats_update) => self
                 .vmm
                 .lock()
                 .expect("Poisoned lock")
                 .update_balloon_stats_config(balloon_stats_update.stats_polling_interval_s)
                 .map(|_| VmmData::Empty)
-                .map_err(|err| VmmActionError::BalloonConfig(BalloonConfigError::from(err))),
+                .map_err(VmmActionError::BalloonUpdate),
             UpdateBlockDevice(new_cfg) => self.update_block_device(new_cfg),
             UpdateNetworkInterface(netif_update) => self.update_net_rate_limiters(netif_update),