refactor: further simplify with_virtio_device functions

Replace 2 functions: `try_with_virtio_device_with_id`
and `with_virtio_device_with_id`, with single
`with_virtio_device` that handles any return type of the
passed function.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

Author: ShadowCurse

src/vmm/src/device_manager/mod.rs

@@ -368,15 +368,10 @@ impl DeviceManager {
     }
 
     /// Run fn `f()` for the virtio device matching `virtio_type` and `id`.
-    pub fn try_with_virtio_device_with_id<T, F, R, E>(
-        &self,
-        id: &str,
-        f: F,
-    ) -> Result<Result<R, E>, FindDeviceError>
+    pub fn with_virtio_device<T, F, R>(&self, id: &str, f: F) -> Result<R, FindDeviceError>
     where
         T: VirtioDevice + 'static + Debug,
-        E: std::error::Error + 'static + Send + Sync,
-        F: FnOnce(&mut T) -> Result<R, E>,
+        F: FnOnce(&mut T) -> R,
     {
         if let Some(device) = self.get_virtio_device(T::const_device_type(), id) {
             let mut dev = device.lock().expect("Poisoned lock");
@@ -388,15 +383,6 @@ impl DeviceManager {
             Err(FindDeviceError::DeviceNotFound)
         }
     }
-
-    /// Run fn `f()` for the virtio device matching `virtio_type` and `id`.
-    pub fn with_virtio_device_with_id<T, F, R>(&self, id: &str, f: F) -> Result<R, FindDeviceError>
-    where
-        T: VirtioDevice + 'static + Debug,
-        F: FnOnce(&mut T) -> R,
-    {
-        self.try_with_virtio_device_with_id(id, |dev: &mut T| Ok::<R, FindDeviceError>(f(dev)))?
-    }
 }
 
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]

src/vmm/src/lib.rs

@@ -526,10 +526,9 @@ impl Vmm {
         path_on_host: String,
     ) -> Result<(), VmmError> {
         self.device_manager
-            .try_with_virtio_device_with_id(drive_id, |block: &mut Block| {
+            .with_virtio_device(drive_id, |block: &mut Block| {
                 block.update_disk_image(path_on_host)
-            })
-            .map_err(VmmError::FindDeviceError)??;
+            })??;
         Ok(())
     }
 
@@ -541,18 +540,16 @@ impl Vmm {
         rl_ops: BucketUpdate,
     ) -> Result<(), VmmError> {
         self.device_manager
-            .try_with_virtio_device_with_id(drive_id, |block: &mut Block| {
+            .with_virtio_device(drive_id, |block: &mut Block| {
                 block.update_rate_limiter(rl_bytes, rl_ops)
-            })
-            .map_err(VmmError::FindDeviceError)??;
+            })??;
         Ok(())
     }
 
     /// 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
-            .try_with_virtio_device_with_id(drive_id, |block: &mut Block| block.update_config())
-            .map_err(VmmError::FindDeviceError)??;
+            .with_virtio_device(drive_id, |block: &mut Block| block.update_config())??;
         Ok(())
     }
 
@@ -566,35 +563,34 @@ impl Vmm {
         tx_ops: BucketUpdate,
     ) -> Result<(), VmmError> {
         self.device_manager
-            .with_virtio_device_with_id(net_id, |net: &mut Net| {
+            .with_virtio_device(net_id, |net: &mut Net| {
                 net.patch_rate_limiters(rx_bytes, rx_ops, tx_bytes, tx_ops)
-            })
-            .map_err(VmmError::FindDeviceError)
+            })?;
+        Ok(())
     }
 
     /// Returns a reference to the balloon device if present.
     pub fn balloon_config(&self) -> Result<BalloonConfig, VmmError> {
-        self.device_manager
-            .with_virtio_device_with_id(BALLOON_DEV_ID, |dev: &mut Balloon| dev.config())
-            .map_err(VmmError::FindDeviceError)
+        let config = self
+            .device_manager
+            .with_virtio_device(BALLOON_DEV_ID, |dev: &mut Balloon| dev.config())?;
+        Ok(config)
     }
 
     /// Returns the latest balloon statistics if they are enabled.
     pub fn latest_balloon_stats(&self) -> Result<BalloonStats, VmmError> {
         let stats = self
             .device_manager
-            .try_with_virtio_device_with_id(BALLOON_DEV_ID, |dev: &mut Balloon| dev.latest_stats())
-            .map_err(VmmError::FindDeviceError)??;
+            .with_virtio_device(BALLOON_DEV_ID, |dev: &mut Balloon| dev.latest_stats())??;
         Ok(stats)
     }
 
     /// Updates configuration for the balloon device target size.
     pub fn update_balloon_config(&mut self, amount_mib: u32) -> Result<(), VmmError> {
         self.device_manager
-            .try_with_virtio_device_with_id(BALLOON_DEV_ID, |dev: &mut Balloon| {
+            .with_virtio_device(BALLOON_DEV_ID, |dev: &mut Balloon| {
                 dev.update_size(amount_mib)
-            })
-            .map_err(VmmError::FindDeviceError)??;
+            })??;
         Ok(())
     }
 
@@ -604,10 +600,9 @@ impl Vmm {
         stats_polling_interval_s: u16,
     ) -> Result<(), VmmError> {
         self.device_manager
-            .try_with_virtio_device_with_id(BALLOON_DEV_ID, |dev: &mut Balloon| {
+            .with_virtio_device(BALLOON_DEV_ID, |dev: &mut Balloon| {
                 dev.update_stats_polling_interval(stats_polling_interval_s)
-            })
-            .map_err(VmmError::FindDeviceError)??;
+            })??;
         Ok(())
     }