refactor: add trait function const_device_type to use in Sized contexts

By defining an associated function to the trait, we can simplify the
logic to execute code on a specific device from the VMM, while also
statically guaranteeing we are passing the right constant.

The downside is that we need both a sized and a dyn implementation for
the function. To ensure devices implement them correctly, a utility
macro is provided. We cannot do it as a default trait impl as the const_
variant is only defined on Sized.

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

Author: Manciukic

src/vmm/src/device_manager/mmio.rs

@@ -468,7 +468,7 @@ pub(crate) mod tests {
     use crate::test_utils::multi_region_mem_raw;
     use crate::vstate::kvm::Kvm;
     use crate::vstate::memory::{GuestAddress, GuestMemoryMmap};
-    use crate::{Vm, arch};
+    use crate::{Vm, arch, impl_device_type};
 
     const QUEUE_SIZES: &[u16] = &[64];
 
@@ -522,6 +522,8 @@ pub(crate) mod tests {
     }
 
     impl VirtioDevice for DummyDevice {
+        impl_device_type!(0);
+
         fn avail_features(&self) -> u64 {
             0
         }
@@ -532,10 +534,6 @@ pub(crate) mod tests {
 
         fn set_acked_features(&mut self, _: u64) {}
 
-        fn device_type(&self) -> u32 {
-            0
-        }
-
         fn queues(&self) -> &[Queue] {
             &self.queues
         }

src/vmm/src/device_manager/mod.rs

@@ -373,7 +373,6 @@ 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,
-        virtio_type: u32,
         id: &str,
         f: F,
     ) -> Result<R, FindDeviceError>
@@ -382,7 +381,7 @@ impl DeviceManager {
         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) {
+        if let Some(device) = self.get_virtio_device(T::const_device_type(), id) {
             let mut dev = device.lock().expect("Poisoned lock");
             f(dev
                 .as_mut_any()
@@ -395,19 +394,12 @@ impl DeviceManager {
     }
 
     /// 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>
+    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(virtio_type, id, |dev: &mut T| {
-            Ok::<R, FindDeviceError>(f(dev))
-        })
+        self.try_with_virtio_device_with_id(id, |dev: &mut T| Ok::<R, FindDeviceError>(f(dev)))
     }
 }
 

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

@@ -31,9 +31,9 @@ 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};
+use crate::{impl_device_type, mem_size_mib};
 
 const SIZE_OF_U32: usize = std::mem::size_of::<u32>();
 const SIZE_OF_STAT: usize = std::mem::size_of::<BalloonStat>();
@@ -544,6 +544,8 @@ impl Balloon {
 }
 
 impl VirtioDevice for Balloon {
+    impl_device_type!(VIRTIO_ID_BALLOON);
+
     fn avail_features(&self) -> u64 {
         self.avail_features
     }
@@ -556,10 +558,6 @@ impl VirtioDevice for Balloon {
         self.acked_features = acked_features;
     }
 
-    fn device_type(&self) -> u32 {
-        VIRTIO_ID_BALLOON
-    }
-
     fn queues(&self) -> &[Queue] {
         &self.queues
     }

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

@@ -16,6 +16,7 @@ use crate::devices::virtio::device::VirtioDevice;
 use crate::devices::virtio::generated::virtio_ids::VIRTIO_ID_BLOCK;
 use crate::devices::virtio::queue::{InvalidAvailIdx, Queue};
 use crate::devices::virtio::transport::VirtioInterrupt;
+use crate::impl_device_type;
 use crate::rate_limiter::BucketUpdate;
 use crate::snapshot::Persist;
 use crate::vmm_config::drive::BlockDeviceConfig;
@@ -132,6 +133,8 @@ impl Block {
 }
 
 impl VirtioDevice for Block {
+    impl_device_type!(VIRTIO_ID_BLOCK);
+
     fn avail_features(&self) -> u64 {
         match self {
             Self::Virtio(b) => b.avail_features,
@@ -153,10 +156,6 @@ impl VirtioDevice for Block {
         }
     }
 
-    fn device_type(&self) -> u32 {
-        VIRTIO_ID_BLOCK
-    }
-
     fn queues(&self) -> &[Queue] {
         match self {
             Self::Virtio(b) => &b.queues,

src/vmm/src/devices/virtio/block/vhost_user/device.rs

@@ -27,6 +27,7 @@ use crate::devices::virtio::vhost_user::{VhostUserHandleBackend, VhostUserHandle
 use crate::devices::virtio::vhost_user_metrics::{
     VhostUserDeviceMetrics, VhostUserMetricsPerDevice,
 };
+use crate::impl_device_type;
 use crate::logger::{IncMetric, StoreMetric, log_dev_preview_warning};
 use crate::utils::u64_to_usize;
 use crate::vmm_config::drive::BlockDeviceConfig;
@@ -287,6 +288,8 @@ impl<T: VhostUserHandleBackend> VhostUserBlockImpl<T> {
 }
 
 impl<T: VhostUserHandleBackend + Send + 'static> VirtioDevice for VhostUserBlockImpl<T> {
+    impl_device_type!(VIRTIO_ID_BLOCK);
+
     fn avail_features(&self) -> u64 {
         self.avail_features
     }
@@ -299,10 +302,6 @@ impl<T: VhostUserHandleBackend + Send + 'static> VirtioDevice for VhostUserBlock
         self.acked_features = acked_features;
     }
 
-    fn device_type(&self) -> u32 {
-        VIRTIO_ID_BLOCK
-    }
-
     fn queues(&self) -> &[Queue] {
         &self.queues
     }

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

@@ -34,6 +34,7 @@ use crate::devices::virtio::generated::virtio_ids::VIRTIO_ID_BLOCK;
 use crate::devices::virtio::generated::virtio_ring::VIRTIO_RING_F_EVENT_IDX;
 use crate::devices::virtio::queue::{InvalidAvailIdx, Queue};
 use crate::devices::virtio::transport::{VirtioInterrupt, VirtioInterruptType};
+use crate::impl_device_type;
 use crate::logger::{IncMetric, error, warn};
 use crate::rate_limiter::{BucketUpdate, RateLimiter};
 use crate::utils::u64_to_usize;
@@ -582,6 +583,8 @@ impl VirtioBlock {
 }
 
 impl VirtioDevice for VirtioBlock {
+    impl_device_type!(VIRTIO_ID_BLOCK);
+
     fn avail_features(&self) -> u64 {
         self.avail_features
     }
@@ -594,10 +597,6 @@ impl VirtioDevice for VirtioBlock {
         self.acked_features = acked_features;
     }
 
-    fn device_type(&self) -> u32 {
-        VIRTIO_ID_BLOCK
-    }
-
     fn queues(&self) -> &[Queue] {
         &self.queues
     }

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

@@ -74,7 +74,14 @@ pub trait VirtioDevice: AsAny + Send {
         (self.acked_features() & (1 << feature)) != 0
     }
 
+    /// The virtio device type (as a constant of the struct).
+    fn const_device_type() -> u32
+    where
+        Self: Sized;
+
     /// The virtio device type.
+    ///
+    /// It should be the same as returned by Self::const_device_type().
     fn device_type(&self) -> u32;
 
     /// Returns the device queues.
@@ -170,6 +177,20 @@ impl fmt::Debug for dyn VirtioDevice {
     }
 }
 
+/// Utility to define both const_device_type and device_type with a u32 constant
+#[macro_export]
+macro_rules! impl_device_type {
+    ($const_type:expr) => {
+        fn const_device_type() -> u32 {
+            $const_type
+        }
+
+        fn device_type(&self) -> u32 {
+            Self::const_device_type()
+        }
+    };
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
     use super::*;
@@ -180,6 +201,8 @@ pub(crate) mod tests {
     }
 
     impl VirtioDevice for MockVirtioDevice {
+        impl_device_type!(0);
+
         fn avail_features(&self) -> u64 {
             todo!()
         }
@@ -192,10 +215,6 @@ pub(crate) mod tests {
             todo!()
         }
 
-        fn device_type(&self) -> u32 {
-            todo!()
-        }
-
         fn queues(&self) -> &[Queue] {
             todo!()
         }

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

@@ -40,6 +40,7 @@ use crate::devices::virtio::transport::{VirtioInterrupt, VirtioInterruptType};
 use crate::devices::{DeviceError, report_net_event_fail};
 use crate::dumbo::pdu::arp::ETH_IPV4_FRAME_LEN;
 use crate::dumbo::pdu::ethernet::{EthernetFrame, PAYLOAD_OFFSET};
+use crate::impl_device_type;
 use crate::logger::{IncMetric, METRICS};
 use crate::mmds::data_store::Mmds;
 use crate::mmds::ns::MmdsNetworkStack;
@@ -961,6 +962,8 @@ impl Net {
 }
 
 impl VirtioDevice for Net {
+    impl_device_type!(VIRTIO_ID_NET);
+
     fn avail_features(&self) -> u64 {
         self.avail_features
     }
@@ -973,10 +976,6 @@ impl VirtioDevice for Net {
         self.acked_features = acked_features;
     }
 
-    fn device_type(&self) -> u32 {
-        VIRTIO_ID_NET
-    }
-
     fn queues(&self) -> &[Queue] {
         &self.queues
     }

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

@@ -21,6 +21,7 @@ use crate::devices::virtio::iov_deque::IovDequeError;
 use crate::devices::virtio::iovec::IoVecBufferMut;
 use crate::devices::virtio::queue::{FIRECRACKER_MAX_QUEUE_SIZE, InvalidAvailIdx, Queue};
 use crate::devices::virtio::transport::{VirtioInterrupt, VirtioInterruptType};
+use crate::impl_device_type;
 use crate::logger::{IncMetric, debug, error};
 use crate::rate_limiter::{RateLimiter, TokenType};
 use crate::vstate::memory::GuestMemoryMmap;
@@ -253,9 +254,7 @@ impl Entropy {
 }
 
 impl VirtioDevice for Entropy {
-    fn device_type(&self) -> u32 {
-        VIRTIO_ID_RNG
-    }
+    impl_device_type!(VIRTIO_ID_RNG);
 
     fn queues(&self) -> &[Queue] {
         &self.queues

src/vmm/src/devices/virtio/transport/mmio.rs

@@ -476,6 +476,7 @@ pub(crate) mod tests {
     use super::*;
     use crate::devices::virtio::ActivateError;
     use crate::devices::virtio::device_status::DEVICE_NEEDS_RESET;
+    use crate::impl_device_type;
     use crate::test_utils::single_region_mem;
     use crate::utils::byte_order::{read_le_u32, write_le_u32};
     use crate::utils::u64_to_usize;
@@ -516,6 +517,8 @@ pub(crate) mod tests {
     }
 
     impl VirtioDevice for DummyDevice {
+        impl_device_type!(123);
+
         fn avail_features(&self) -> u64 {
             self.avail_features
         }
@@ -528,10 +531,6 @@ pub(crate) mod tests {
             self.acked_features = acked_features;
         }
 
-        fn device_type(&self) -> u32 {
-            123
-        }
-
         fn queues(&self) -> &[Queue] {
             &self.queues
         }