virtio: add kick() method in VirtioDevice trait

So that we don't have to downcast VirtioDevice trait objects to the
actual device type before calling the logic to process events for each
device.

Signed-off-by: Babis Chalios <bchalios@amazon.es>

Author: bchalios

src/vmm/src/device_manager/mod.rs

@@ -30,14 +30,8 @@ use crate::devices::legacy::RTCDevice;
 use crate::devices::legacy::serial::SerialOut;
 use crate::devices::legacy::{IER_RDA_BIT, IER_RDA_OFFSET, SerialDevice};
 use crate::devices::pseudo::BootTimer;
-use crate::devices::virtio::balloon::Balloon;
-use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::device::VirtioDevice;
-use crate::devices::virtio::net::Net;
-use crate::devices::virtio::rng::Entropy;
 use crate::devices::virtio::transport::mmio::{IrqTrigger, MmioTransport};
-use crate::devices::virtio::vsock::{TYPE_VSOCK, Vsock, VsockUnixBackend};
-use crate::devices::virtio::{TYPE_BALLOON, TYPE_BLOCK, TYPE_NET, TYPE_RNG};
 use crate::resources::VmResources;
 use crate::snapshot::Persist;
 use crate::vstate::memory::GuestMemoryMmap;
@@ -265,85 +259,28 @@ impl DeviceManager {
         self.pci_devices.attach_pci_segment(vm)
     }
 
-    fn do_kick_device(virtio_device: Arc<Mutex<dyn VirtioDevice>>) {
-        let mut device = virtio_device.lock().expect("Poisoned lock");
-        match device.device_type() {
-            TYPE_BALLOON => {
-                let balloon = device.as_mut_any().downcast_mut::<Balloon>().unwrap();
-                // If device is activated, kick the balloon queue(s) to make up for any
-                // pending or in-flight epoll events we may have not captured in snapshot.
-                // Stats queue doesn't need kicking as it is notified via a `timer_fd`.
-                if balloon.is_activated() {
-                    info!("kick balloon {}.", balloon.id());
-                    balloon.process_virtio_queues();
-                }
-            }
-            TYPE_BLOCK => {
-                // We only care about kicking virtio block.
-                // If we need to kick vhost-user-block we can do nothing.
-                if let Some(block) = device.as_mut_any().downcast_mut::<Block>() {
-                    // If device is activated, kick the block queue(s) to make up for any
-                    // pending or in-flight epoll events we may have not captured in
-                    // snapshot. No need to kick Ratelimiters
-                    // because they are restored 'unblocked' so
-                    // any inflight `timer_fd` events can be safely discarded.
-                    if block.is_activated() {
-                        info!("kick block {}.", block.id());
-                        block.process_virtio_queues();
-                    }
-                }
-            }
-            TYPE_NET => {
-                let net = device.as_mut_any().downcast_mut::<Net>().unwrap();
-                // If device is activated, kick the net queue(s) to make up for any
-                // pending or in-flight epoll events we may have not captured in snapshot.
-                // No need to kick Ratelimiters because they are restored 'unblocked' so
-                // any inflight `timer_fd` events can be safely discarded.
-                if net.is_activated() {
-                    info!("kick net {}.", net.id());
-                    net.process_virtio_queues();
-                }
-            }
-            TYPE_VSOCK => {
-                // Vsock has complicated protocol that isn't resilient to any packet loss,
-                // so for Vsock we don't support connection persistence through snapshot.
-                // Any in-flight packets or events are simply lost.
-                // Vsock is restored 'empty'.
-                // The only reason we still `kick` it is to make guest process
-                // `TRANSPORT_RESET_EVENT` event we sent during snapshot creation.
-                let vsock = device
-                    .as_mut_any()
-                    .downcast_mut::<Vsock<VsockUnixBackend>>()
-                    .unwrap();
-                if vsock.is_activated() {
-                    info!("kick vsock {}.", vsock.id());
-                    vsock.signal_used_queue(0).unwrap();
-                }
-            }
-            TYPE_RNG => {
-                let entropy = device.as_mut_any().downcast_mut::<Entropy>().unwrap();
-                if entropy.is_activated() {
-                    info!("kick entropy {}.", entropy.id());
-                    entropy.process_virtio_queues();
-                }
-            }
-            _ => (),
-        }
-    }
-
     /// Artificially kick VirtIO devices as if they had external events.
     pub fn kick_virtio_devices(&self) {
         info!("Artificially kick devices");
         // Go through MMIO VirtIO devices
         let _: Result<(), MmioError> = self.mmio_devices.for_each_virtio_device(|_, _, device| {
             let mmio_transport_locked = device.inner.lock().expect("Poisoned lock");
-            Self::do_kick_device(mmio_transport_locked.device());
+            mmio_transport_locked
+                .device()
+                .lock()
+                .expect("Poisoned lock")
+                .kick();
             Ok(())
         });
         // Go through PCI VirtIO devices
-        for device in self.pci_devices.virtio_devices.values() {
-            let virtio_device = device.lock().expect("Poisoned lock").virtio_device();
-            Self::do_kick_device(virtio_device);
+        for virtio_pci_device in self.pci_devices.virtio_devices.values() {
+            virtio_pci_device
+                .lock()
+                .expect("Poisoned lock")
+                .virtio_device()
+                .lock()
+                .expect("Poisoned lock")
+                .kick();
         }
     }
 

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

@@ -5,7 +5,7 @@ use std::ops::Deref;
 use std::sync::Arc;
 use std::time::Duration;
 
-use log::error;
+use log::{error, info};
 use serde::Serialize;
 use timerfd::{ClockId, SetTimeFlags, TimerFd, TimerState};
 use vmm_sys_util::eventfd::EventFd;
@@ -615,6 +615,16 @@ impl VirtioDevice for Balloon {
     fn is_activated(&self) -> bool {
         self.device_state.is_activated()
     }
+
+    fn kick(&mut self) {
+        // If device is activated, kick the balloon queue(s) to make up for any
+        // pending or in-flight epoll events we may have not captured in snapshot.
+        // Stats queue doesn't need kicking as it is notified via a `timer_fd`.
+        if self.is_activated() {
+            info!("kick balloon {}.", self.id());
+            self.process_virtio_queues();
+        }
+    }
 }
 
 #[cfg(test)]

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

@@ -4,6 +4,7 @@
 use std::sync::Arc;
 
 use event_manager::{EventOps, Events, MutEventSubscriber};
+use log::info;
 use vmm_sys_util::eventfd::EventFd;
 
 use super::BlockError;
@@ -214,6 +215,18 @@ impl VirtioDevice for Block {
             Self::VhostUser(b) => b.device_state.is_activated(),
         }
     }
+
+    fn kick(&mut self) {
+        // If device is activated, kick the block queue(s) to make up for any
+        // pending or in-flight epoll events we may have not captured in
+        // snapshot. No need to kick Ratelimiters
+        // because they are restored 'unblocked' so
+        // any inflight `timer_fd` events can be safely discarded.
+        if self.is_activated() {
+            info!("kick block {}.", self.id());
+            self.process_virtio_queues();
+        }
+    }
 }
 
 impl MutEventSubscriber for Block {

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

@@ -159,6 +159,9 @@ pub trait VirtioDevice: AsAny + Send {
         }
         Ok(())
     }
+
+    /// Kick the device, as if it had received external events.
+    fn kick(&mut self) {}
 }
 
 impl fmt::Debug for dyn VirtioDevice {

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

@@ -13,7 +13,7 @@ use std::ops::Deref;
 use std::sync::{Arc, Mutex};
 
 use libc::{EAGAIN, iovec};
-use log::error;
+use log::{error, info};
 use vmm_sys_util::eventfd::EventFd;
 
 use super::NET_QUEUE_MAX_SIZE;
@@ -1053,6 +1053,17 @@ impl VirtioDevice for Net {
     fn is_activated(&self) -> bool {
         self.device_state.is_activated()
     }
+
+    fn kick(&mut self) {
+        // If device is activated, kick the net queue(s) to make up for any
+        // pending or in-flight epoll events we may have not captured in snapshot.
+        // No need to kick Ratelimiters because they are restored 'unblocked' so
+        // any inflight `timer_fd` events can be safely discarded.
+        if self.is_activated() {
+            info!("kick net {}.", self.id());
+            self.process_virtio_queues();
+        }
+    }
 }
 
 #[cfg(test)]

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

@@ -6,6 +6,7 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use aws_lc_rs::rand;
+use log::info;
 use vm_memory::GuestMemoryError;
 use vmm_sys_util::eventfd::EventFd;
 
@@ -309,6 +310,13 @@ impl VirtioDevice for Entropy {
         self.device_state = DeviceState::Activated(ActiveState { mem, interrupt });
         Ok(())
     }
+
+    fn kick(&mut self) {
+        if self.is_activated() {
+            info!("kick entropy {}.", self.id());
+            self.process_virtio_queues();
+        }
+    }
 }
 
 #[cfg(test)]

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

@@ -24,7 +24,7 @@ use std::fmt::Debug;
 use std::ops::Deref;
 use std::sync::Arc;
 
-use log::{error, warn};
+use log::{error, info, warn};
 use vmm_sys_util::eventfd::EventFd;
 
 use super::super::super::DeviceError;
@@ -370,6 +370,19 @@ where
     fn is_activated(&self) -> bool {
         self.device_state.is_activated()
     }
+
+    fn kick(&mut self) {
+        // Vsock has complicated protocol that isn't resilient to any packet loss,
+        // so for Vsock we don't support connection persistence through snapshot.
+        // Any in-flight packets or events are simply lost.
+        // Vsock is restored 'empty'.
+        // The only reason we still `kick` it is to make guest process
+        // `TRANSPORT_RESET_EVENT` event we sent during snapshot creation.
+        if self.is_activated() {
+            info!("kick vsock {}.", self.id());
+            self.signal_used_queue(0).unwrap();
+        }
+    }
 }
 
 #[cfg(test)]