fix: saving/restoring async IO engine transport state

src/vmm/src/device_manager/pci_mngr.rs

@@ -283,9 +283,13 @@ impl<'a> Persist<'a> for PciDevices {
 
         for pci_dev in self.virtio_devices.values() {
             let locked_pci_dev = pci_dev.lock().expect("Poisoned lock");
-            let transport_state = locked_pci_dev.state();
             let virtio_dev = locked_pci_dev.virtio_device();
+            // We need to call `prepare_save()` on the device before saving the transport
+            // so that, if we modify the transport state while preparing the device, e.g. sending
+            // an interrupt to the guest, this is correctly captured in the saved transport state.
             let mut locked_virtio_dev = virtio_dev.lock().expect("Poisoned lock");
+            locked_virtio_dev.prepare_save();
+            let transport_state = locked_pci_dev.state();
 
             let pci_device_bdf = transport_state.pci_device_bdf.into();
 
@@ -316,7 +320,6 @@ impl<'a> Persist<'a> for PciDevices {
                              snapshotting yet"
                         );
                     } else {
-                        block_dev.prepare_save();
                         let device_state = block_dev.save();
                         state.block_devices.push(VirtioDeviceState {
                             device_id: block_dev.id().to_string(),
@@ -338,7 +341,6 @@ impl<'a> Persist<'a> for PciDevices {
                             imds_compat: mmds_guard.imds_compat(),
                         });
                     }
-                    net_dev.prepare_save();
                     let device_state = net_dev.save();
 
                     state.net_devices.push(VirtioDeviceState {

src/vmm/src/device_manager/persist.rs

@@ -227,11 +227,15 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
         let _: Result<(), ()> = self.for_each_virtio_device(|_, devid, device| {
             let mmio_transport_locked = device.inner.lock().expect("Poisoned lock");
+            let mut locked_device = mmio_transport_locked.locked_device();
+            // We need to call `prepare_save()` on the device before saving the transport
+            // so that, if we modify the transport state while preparing the device, e.g. sending
+            // an interrupt to the guest, this is correctly captured in the saved transport state.
+            locked_device.prepare_save();
             let transport_state = mmio_transport_locked.save();
             let device_info = device.resources;
             let device_id = devid.clone();
 
-            let mut locked_device = mmio_transport_locked.locked_device();
             match locked_device.device_type() {
                 virtio_ids::VIRTIO_ID_BALLOON => {
                     let device_state = locked_device
@@ -255,7 +259,6 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                              snapshotting yet"
                         );
                     } else {
-                        block.prepare_save();
                         let device_state = block.save();
                         states.block_devices.push(VirtioDeviceState {
                             device_id,
@@ -275,7 +278,6 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                         });
                     }
 
-                    net.prepare_save();
                     let device_state = net.save();
                     states.net_devices.push(VirtioDeviceState {
                         device_id,

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

@@ -82,13 +82,6 @@ impl Block {
         }
     }
 
-    pub fn prepare_save(&mut self) {
-        match self {
-            Self::Virtio(b) => b.prepare_save(),
-            Self::VhostUser(b) => b.prepare_save(),
-        }
-    }
-
     pub fn process_virtio_queues(&mut self) -> Result<(), InvalidAvailIdx> {
         match self {
             Self::Virtio(b) => b.process_virtio_queues(),
@@ -227,6 +220,13 @@ impl VirtioDevice for Block {
             self.process_virtio_queues();
         }
     }
+
+    fn prepare_save(&mut self) {
+        match self {
+            Self::Virtio(b) => b.prepare_save(),
+            Self::VhostUser(b) => b.prepare_save(),
+        }
+    }
 }
 
 impl MutEventSubscriber for Block {

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

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

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

@@ -938,26 +938,6 @@ impl Net {
 
         Ok(())
     }
-
-    /// Prepare saving state
-    pub fn prepare_save(&mut self) {
-        // We shouldn't be messing with the queue if the device is not activated.
-        // Anyways, if it isn't there's nothing to prepare; we haven't parsed any
-        // descriptors yet from it and we can't have a deferred frame.
-        if !self.is_activated() {
-            return;
-        }
-
-        // Give potential deferred RX frame to guest
-        self.rx_buffer.finish_frame(&mut self.queues[RX_INDEX]);
-        // Reset the parsed available descriptors, so we will re-parse them
-        self.queues[RX_INDEX].next_avail -=
-            Wrapping(u16::try_from(self.rx_buffer.parsed_descriptors.len()).unwrap());
-        self.rx_buffer.parsed_descriptors.clear();
-        self.rx_buffer.iovec.clear();
-        self.rx_buffer.used_bytes = 0;
-        self.rx_buffer.used_descriptors = 0;
-    }
 }
 
 impl VirtioDevice for Net {
@@ -1069,6 +1049,26 @@ impl VirtioDevice for Net {
             self.process_virtio_queues();
         }
     }
+
+    /// Prepare saving state
+    fn prepare_save(&mut self) {
+        // We shouldn't be messing with the queue if the device is not activated.
+        // Anyways, if it isn't there's nothing to prepare; we haven't parsed any
+        // descriptors yet from it and we can't have a deferred frame.
+        if !self.is_activated() {
+            return;
+        }
+
+        // Give potential deferred RX frame to guest
+        self.rx_buffer.finish_frame(&mut self.queues[RX_INDEX]);
+        // Reset the parsed available descriptors, so we will re-parse them
+        self.queues[RX_INDEX].next_avail -=
+            Wrapping(u16::try_from(self.rx_buffer.parsed_descriptors.len()).unwrap());
+        self.rx_buffer.parsed_descriptors.clear();
+        self.rx_buffer.iovec.clear();
+        self.rx_buffer.used_bytes = 0;
+        self.rx_buffer.used_descriptors = 0;
+    }
 }
 
 #[cfg(test)]

src/vmm/src/lib.rs

@@ -436,6 +436,12 @@ impl Vmm {
     /// Saves the state of a paused Microvm.
     pub fn save_state(&mut self, vm_info: &VmInfo) -> Result<MicrovmState, MicrovmStateError> {
         use self::MicrovmStateError::SaveVmState;
+        // We need to save device state before saving KVM state.
+        // Some devices, (at the time of writing this comment block device with async engine)
+        // might modify the VirtIO transport and send an interrupt to the guest. If we save KVM
+        // state before we save device state, that interrupt will never be delivered to the guest
+        // upon resuming from the snapshot.
+        let device_states = self.device_manager.save();
         let vcpu_states = self.save_vcpu_states()?;
         let kvm_state = self.kvm.save_state();
         let vm_state = {
@@ -450,7 +456,6 @@ impl Vmm {
                 self.vm.save_state(&mpidrs).map_err(SaveVmState)?
             }
         };
-        let device_states = self.device_manager.save();
 
         Ok(MicrovmState {
             vm_info: vm_info.clone(),