device-manager: kick devices on snapshot restore

Signed-off-by: Adrian Catangiu <[email protected]>

Author: acatangiu

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

@@ -196,7 +196,14 @@ impl Block {
             METRICS.block.event_fails.inc();
         } else if self.rate_limiter.is_blocked() {
             METRICS.block.rate_limiter_throttled_events.inc();
-        } else if self.process_queue(0) {
+        } else {
+            self.process_virtio_queues();
+        }
+    }
+
+    /// Process device virtio queue(s).
+    pub fn process_virtio_queues(&mut self) {
+        if self.process_queue(0) {
             let _ = self.signal_used_queue();
         }
     }

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

@@ -704,6 +704,12 @@ impl Net {
             }
         }
     }
+
+    /// Process device virtio queue(s).
+    pub fn process_virtio_queues(&mut self) {
+        let _ = self.resume_rx();
+        let _ = self.process_tx();
+    }
 }
 
 impl VirtioDevice for Net {

src/vmm/src/device_manager/persist.rs

@@ -216,12 +216,21 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             restore_helper(
                 device.clone(),
-                device,
+                device.clone(),
                 &block_state.device_id,
                 &block_state.transport_state,
                 &block_state.mmio_slot,
                 constructor_args.event_manager,
             )?;
+
+            // 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.
+            let mut dev = device.lock().expect("Poisoned lock");
+            if dev.is_activated() {
+                dev.process_virtio_queues();
+            }
         }
         for net_state in &state.net_devices {
             let device = Arc::new(Mutex::new(
@@ -234,12 +243,21 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             restore_helper(
                 device.clone(),
-                device,
+                device.clone(),
                 &net_state.device_id,
                 &net_state.transport_state,
                 &net_state.mmio_slot,
                 constructor_args.event_manager,
             )?;
+
+            // 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.
+            let mut dev = device.lock().expect("Poisoned lock");
+            if dev.is_activated() {
+                dev.process_virtio_queues();
+            }
         }
         if let Some(vsock_state) = &state.vsock_device {
             let ctor_args = VsockUdsConstructorArgs {
@@ -266,6 +284,10 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                 &vsock_state.mmio_slot,
                 constructor_args.event_manager,
             )?;
+
+            // 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'.
         }
 
         Ok(dev_manager)