virtio: restore the PCI config space before reinstalling the doorbells

vm/devices/virtio/virtio/src/tests.rs

@@ -136,6 +136,27 @@ async fn yield_and_poll_device(dev: &mut impl PollDevice) {
 struct VirtioTestMemoryAccess {
     memory_map: Mutex<MemoryMap>,
     doorbell_count: AtomicUsize,
+    /// The doorbells that are currently INSTALLED, i.e. whose registration
+    /// object is still alive.
+    ///
+    /// Kept alongside `doorbell_count` because a count of registration calls
+    /// cannot answer the question a restore poses: it says a registration
+    /// happened, not that a doorbell exists NOW at the address the guest kicks.
+    /// Behind an `Arc` so each handed-out registration can remove its own entry
+    /// when it is dropped.
+    live_doorbells: Arc<Mutex<Vec<DoorbellSpec>>>,
+}
+
+/// A doorbell registration as the transport asked for it.
+///
+/// All queues of one virtio PCI device register at the SAME address (BAR0's
+/// notify register) and are told apart by `value`, the datamatch. So an
+/// assertion on addresses alone cannot see a missing queue.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
+struct DoorbellSpec {
+    address: u64,
+    value: Option<u64>,
+    length: Option<u32>,
 }
 
 #[derive(Default)]
@@ -199,6 +220,13 @@ impl VirtioTestMemoryAccess {
         Default::default()
     }
 
+    /// The doorbells installed right now, in a stable order for comparison.
+    fn installed_doorbells(&self) -> Vec<DoorbellSpec> {
+        let mut specs = self.live_doorbells.lock().clone();
+        specs.sort_unstable();
+        specs
+    }
+
     fn modify_memory_map(&self, address: u64, data: &[u8], writeable: bool) {
         self.memory_map.lock().insert(address, data, writeable);
     }
@@ -284,18 +312,41 @@ unsafe impl GuestMemoryAccess for VirtioTestMemoryAccess {
     }
 }
 
-struct DoorbellEntry;
+/// The handle a doorbell registration hands back. Dropping it is what
+/// UNINSTALLS the doorbell, so it de-registers itself from `live_doorbells`.
+struct DoorbellEntry {
+    live: Arc<Mutex<Vec<DoorbellSpec>>>,
+    spec: DoorbellSpec,
+}
+
+impl Drop for DoorbellEntry {
+    fn drop(&mut self) {
+        let mut live = self.live.lock();
+        if let Some(i) = live.iter().position(|s| *s == self.spec) {
+            live.remove(i);
+        }
+    }
+}
 
 impl DoorbellRegistration for VirtioTestMemoryAccess {
     fn register_doorbell(
         &self,
-        _: u64,
-        _: Option<u64>,
-        _: Option<u32>,
+        address: u64,
+        value: Option<u64>,
+        length: Option<u32>,
         _: &Event,
     ) -> io::Result<Box<dyn Send + Sync>> {
         self.doorbell_count.fetch_add(1, Ordering::Relaxed);
-        Ok(Box::new(DoorbellEntry))
+        let spec = DoorbellSpec {
+            address,
+            value,
+            length,
+        };
+        self.live_doorbells.lock().push(spec);
+        Ok(Box::new(DoorbellEntry {
+            live: self.live_doorbells.clone(),
+            spec,
+        }))
     }
 }
 
@@ -4896,63 +4947,67 @@ async fn mmio_save_not_supported_device(_driver: DefaultDriver) {
 async fn pci_restore_reinstalls_doorbells(driver: DefaultDriver) {
     use vmcore::save_restore::SaveRestore;
 
+    const QUEUE_COUNT: u16 = 2;
+    // What `setup_pci_device` programs BAR0 to.
+    const BAR0_BASE: u64 = 0x10000000000;
+    // BAR0's base plus its notify offset. Every queue's doorbell lands on this one
+    // address; the queue index is the datamatch value.
+    let notify_address: u64 = BAR0_BASE + VIRTIO_PCI_COMMON_CFG_SIZE as u64;
+
     let test_mem = VirtioTestMemoryAccess::new();
 
-    let guest = VirtioTestGuest::new_split(&driver, &test_mem, 1, 4, true);
+    let guest = VirtioTestGuest::new_split(&driver, &test_mem, QUEUE_COUNT, 4, true);
 
-    let mut dev = VirtioPciTestDevice::new(&driver, 1, &test_mem, None);
+    let mut dev = VirtioPciTestDevice::new(&driver, QUEUE_COUNT, &test_mem, None);
     guest
         .setup_pci_device(&mut dev, guest.queue_features())
         .await;
 
-    // After setup, doorbells should be registered.
-    let doorbells_after_setup = test_mem.doorbell_count.load(Ordering::Relaxed);
-    assert!(
-        doorbells_after_setup > 0,
-        "doorbells should be registered after setup"
+    // What the running device serves, and what the restore has to come back
+    // with. Spelled out rather than just counted so a doorbell installed at the
+    // wrong address, or one queue short, cannot pass.
+    let expected: Vec<DoorbellSpec> = (0..QUEUE_COUNT)
+        .map(|i| DoorbellSpec {
+            address: notify_address,
+            value: Some(i as u64),
+            length: Some(2),
+        })
+        .collect();
+    assert_eq!(
+        test_mem.installed_doorbells(),
+        expected,
+        "the live device should serve one doorbell per queue at BAR0's notify register"
     );
 
-    // Stop, save, restore into a new device.
     dev.pci_device.stop().await;
     let saved = dev.pci_device.save().expect("save should succeed");
 
-    let mut dev2 = VirtioPciTestDevice::new(&driver, 1, &test_mem, None);
-    // Configure BARs on the target device so doorbells can be registered.
-    let bar_address1: u64 = 0x10000000000;
-    dev2.pci_device
-        .pci_cfg_write(
-            0x14,
-            ByteEnabledDwordWrite::with_all_bytes_enabled((bar_address1 >> 32) as u32),
-        )
-        .unwrap();
-    dev2.pci_device
-        .pci_cfg_write(
-            0x10,
-            ByteEnabledDwordWrite::with_all_bytes_enabled(bar_address1 as u32),
-        )
-        .unwrap();
-    dev2.pci_device
-        .pci_cfg_write(
-            0x4,
-            ByteEnabledDwordWrite::new(
-                cfg_space::Command::new()
-                    .with_mmio_enabled(true)
-                    .into_bits() as u32,
-                PciConfigByteEnable::LOW_WORD,
-            ),
-        )
-        .unwrap();
-    // Reset counter to isolate restore behavior.
-    test_mem.doorbell_count.store(0, Ordering::Relaxed);
+    // Drop the saved device before restoring, so what the assertions below see
+    // can only be the RESTORED device's doorbells and not the original's.
+    drop(dev);
+    assert!(
+        test_mem.installed_doorbells().is_empty(),
+        "dropping the saved device should uninstall its doorbells"
+    );
+
+    // Restore into a device that is FRESH, with nothing having programmed its
+    // BARs - which is the real restore: the guest is resumed mid-flight and never
+    // writes its BARs or DRIVER_OK again, so the config space can only come back
+    // from the saved state. This test used to program the BARs here by hand
+    // before restoring, and that is precisely what hid #159: the doorbells were
+    // installed from an address the production path does not have yet at that
+    // point, so the assertion held while every restored VM ran without a single
+    // doorbell.
+    let mut dev2 = VirtioPciTestDevice::new(&driver, QUEUE_COUNT, &test_mem, None);
     dev2.pci_device
         .restore(saved)
         .expect("restore should succeed");
 
-    // Doorbells must be reinstalled during restore.
-    let doorbells_after_restore = test_mem.doorbell_count.load(Ordering::Relaxed);
-    assert!(
-        doorbells_after_restore > 0,
-        "doorbells should be reinstalled after restore, got {doorbells_after_restore}"
+    assert_eq!(
+        test_mem.installed_doorbells(),
+        expected,
+        "restore must reinstall every queue's doorbell at BAR0's notify register; \
+         without them each kick is an MMIO exit to userspace instead of the ioeventfd"
     );
 
     dev2.pci_device.stop().await;

vm/devices/virtio/virtio/src/transport/pci.rs

@@ -881,6 +881,22 @@ mod saved_state {
 
         fn restore(&mut self, state: Self::SavedState) -> Result<(), RestoreError> {
             let saved_queue_count = state.queues.len();
+
+            // Restore the PCI config space (BARs, command register, etc.) so
+            // that host pre-assigned BARs are preserved across the cycle.
+            //
+            // This has to come before `restore_common`, which reinstalls the
+            // queue doorbells: `doorbell_region` reads `bar_address(0)`, and a
+            // BAR address only becomes active once the config space brings it
+            // back. If restored the other way round, `doorbell_region` answers
+            // `None` at that moment and no doorbell is installed. The only other
+            // place they go in is the guest writing DRIVER_OK, which a guest
+            // resumed mid-flight never does again, so the omission lasts for the
+            // life of the VM and every queue kick takes an MMIO exit instead of
+            // the registered doorbell. Nothing in `restore_common` reads the
+            // interrupt state restored below.
+            self.pci.config_space.restore(state.cfg_space)?;
+
             self.core.restore_common(
                 &mut self.pci,
                 &state.common,
@@ -898,10 +914,6 @@ mod saved_state {
             }
             self.pci.msix_config_vector = state.msix_config_vector;
 
-            // Restore the PCI config space (BARs, command register, etc.) so
-            // that host pre-assigned BARs are preserved across the cycle.
-            self.pci.config_space.restore(state.cfg_space)?;
-
             Ok(())
         }
     }