test: VirtIO PCI device create and restoration

Refactor the test code that inserts VirtIO devices in a Vmm object and
then add a test which creates a Vmm with PCI devices and then serializes
and deserializes the device manager and ensures that everything is as
restored as expected.

Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

src/vmm/src/builder.rs

@@ -848,7 +848,6 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .mmio_devices
                 .get_virtio_device(TYPE_VSOCK, &vsock_dev_id)
                 .is_some()
         );
@@ -874,7 +873,6 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .mmio_devices
                 .get_virtio_device(TYPE_RNG, ENTROPY_DEV_ID)
                 .is_some()
         );
@@ -909,7 +907,6 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .mmio_devices
                 .get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID)
                 .is_some()
         );
@@ -961,7 +958,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda ro"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -983,7 +979,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 rw"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -1006,7 +1001,6 @@ pub(crate) mod tests {
             assert!(!cmdline_contains(&cmdline, "root=/dev/vda"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -1044,19 +1038,16 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 rw"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, "root")
                     .is_some()
             );
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, "secondary")
                     .is_some()
             );
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, "third")
                     .is_some()
             );
@@ -1086,7 +1077,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda rw"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -1108,7 +1098,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 ro"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -1130,7 +1119,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda rw"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );

src/vmm/src/device_manager/mod.rs

@@ -375,6 +375,36 @@ impl DeviceManager {
             Self::do_mark_virtio_queue_memory_dirty(virtio_device, mem);
         }
     }
+
+    /// Get a VirtIO device of type `virtio_type` with ID `device_id`
+    pub fn get_virtio_device(
+        &self,
+        virtio_type: u32,
+        device_id: &str,
+    ) -> Option<Arc<Mutex<dyn VirtioDevice>>> {
+        if self.pci_devices.pci_segment.is_some() {
+            let pci_device = self.pci_devices.get_virtio_device(virtio_type, device_id)?;
+            Some(
+                pci_device
+                    .lock()
+                    .expect("Poisoned lock")
+                    .virtio_device()
+                    .clone(),
+            )
+        } else {
+            let mmio_device = self
+                .mmio_devices
+                .get_virtio_device(virtio_type, device_id)?;
+            Some(
+                mmio_device
+                    .inner
+                    .lock()
+                    .expect("Poisoned lock")
+                    .device()
+                    .clone(),
+            )
+        }
+    }
 }
 
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]

src/vmm/src/device_manager/pci_mngr.rs

@@ -259,6 +259,16 @@ impl PciDevices {
         event_manager.add_subscriber(device);
         Ok(())
     }
+
+    /// Gets the specified device.
+    pub fn get_virtio_device(
+        &self,
+        device_type: u32,
+        device_id: &str,
+    ) -> Option<&Arc<Mutex<VirtioPciDevice>>> {
+        self.virtio_devices
+            .get(&(device_type, device_id.to_string()))
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -617,3 +627,187 @@ impl<'a> Persist<'a> for PciDevices {
         Ok(pci_devices)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vmm_sys_util::tempfile::TempFile;
+
+    use super::*;
+    use crate::builder::tests::*;
+    use crate::device_manager;
+    use crate::devices::virtio::block::CacheType;
+    use crate::mmds::data_store::MmdsVersion;
+    use crate::resources::VmmConfig;
+    use crate::snapshot::Snapshot;
+    use crate::vmm_config::balloon::BalloonDeviceConfig;
+    use crate::vmm_config::entropy::EntropyDeviceConfig;
+    use crate::vmm_config::net::NetworkInterfaceConfig;
+    use crate::vmm_config::vsock::VsockDeviceConfig;
+
+    #[test]
+    fn test_device_manager_persistence() {
+        let mut buf = vec![0; 65536];
+        // These need to survive so the restored blocks find them.
+        let _block_files;
+        let mut tmp_sock_file = TempFile::new().unwrap();
+        tmp_sock_file.remove().unwrap();
+        // Set up a vmm with one of each device, and get the serialized DeviceStates.
+        {
+            let mut event_manager = EventManager::new().expect("Unable to create EventManager");
+            let mut vmm = default_vmm();
+            vmm.device_manager.enable_pci(&vmm.vm).unwrap();
+            let mut cmdline = default_kernel_cmdline();
+
+            // Add a balloon device.
+            let balloon_cfg = BalloonDeviceConfig {
+                amount_mib: 123,
+                deflate_on_oom: false,
+                stats_polling_interval_s: 1,
+            };
+            insert_balloon_device(&mut vmm, &mut cmdline, &mut event_manager, balloon_cfg);
+            // Add a block device.
+            let drive_id = String::from("root");
+            let block_configs = vec![CustomBlockConfig::new(
+                drive_id,
+                true,
+                None,
+                true,
+                CacheType::Unsafe,
+            )];
+            _block_files =
+                insert_block_devices(&mut vmm, &mut cmdline, &mut event_manager, block_configs);
+            // Add a net device.
+            let network_interface = NetworkInterfaceConfig {
+                iface_id: String::from("netif"),
+                host_dev_name: String::from("hostname"),
+                guest_mac: None,
+                rx_rate_limiter: None,
+                tx_rate_limiter: None,
+            };
+            insert_net_device_with_mmds(
+                &mut vmm,
+                &mut cmdline,
+                &mut event_manager,
+                network_interface,
+                MmdsVersion::V2,
+            );
+            // Add a vsock device.
+            let vsock_dev_id = "vsock";
+            let vsock_config = VsockDeviceConfig {
+                vsock_id: Some(vsock_dev_id.to_string()),
+                guest_cid: 3,
+                uds_path: tmp_sock_file.as_path().to_str().unwrap().to_string(),
+            };
+            insert_vsock_device(&mut vmm, &mut cmdline, &mut event_manager, vsock_config);
+            // Add an entropy device.
+            let entropy_config = EntropyDeviceConfig::default();
+            insert_entropy_device(&mut vmm, &mut cmdline, &mut event_manager, entropy_config);
+
+            Snapshot::serialize(&mut buf.as_mut_slice(), &vmm.device_manager.save()).unwrap();
+        }
+
+        tmp_sock_file.remove().unwrap();
+
+        let mut event_manager = EventManager::new().expect("Unable to create EventManager");
+        // Keep in mind we are re-creating here an empty DeviceManager. Restoring later on
+        // will create a new PciDevices manager different than vmm.pci_devices. We're doing
+        // this to avoid restoring the whole Vmm, since what we really need from Vmm is the Vm
+        // object and calling default_vmm() is the easiest way to create one.
+        let vmm = default_vmm();
+        let device_manager_state: device_manager::DevicesState =
+            Snapshot::deserialize(&mut buf.as_slice()).unwrap();
+        let vm_resources = &mut VmResources::default();
+        let restore_args = PciDevicesConstructorArgs {
+            vm: vmm.vm.clone(),
+            mem: vmm.vm.guest_memory(),
+            vm_resources,
+            instance_id: "microvm-id",
+            restored_from_file: true,
+            event_manager: &mut event_manager,
+        };
+        let _restored_dev_manager =
+            PciDevices::restore(restore_args, &device_manager_state.pci_state).unwrap();
+
+        let expected_vm_resources = format!(
+            r#"{{
+  "balloon": {{
+    "amount_mib": 123,
+    "deflate_on_oom": false,
+    "stats_polling_interval_s": 1
+  }},
+  "drives": [
+    {{
+      "drive_id": "root",
+      "partuuid": null,
+      "is_root_device": true,
+      "cache_type": "Unsafe",
+      "is_read_only": true,
+      "path_on_host": "{}",
+      "rate_limiter": null,
+      "io_engine": "Sync",
+      "socket": null
+    }}
+  ],
+  "boot-source": {{
+    "kernel_image_path": "",
+    "initrd_path": null,
+    "boot_args": null
+  }},
+  "cpu-config": null,
+  "logger": null,
+  "machine-config": {{
+    "vcpu_count": 1,
+    "mem_size_mib": 128,
+    "smt": false,
+    "track_dirty_pages": false,
+    "huge_pages": "None"
+  }},
+  "metrics": null,
+  "mmds-config": {{
+    "version": "V2",
+    "network_interfaces": [
+      "netif"
+    ],
+    "ipv4_address": "169.254.169.254"
+  }},
+  "network-interfaces": [
+    {{
+      "iface_id": "netif",
+      "host_dev_name": "hostname",
+      "guest_mac": null,
+      "rx_rate_limiter": null,
+      "tx_rate_limiter": null
+    }}
+  ],
+  "vsock": {{
+    "guest_cid": 3,
+    "uds_path": "{}"
+  }},
+  "entropy": {{
+    "rate_limiter": null
+  }}
+}}"#,
+            _block_files.last().unwrap().as_path().to_str().unwrap(),
+            tmp_sock_file.as_path().to_str().unwrap()
+        );
+
+        assert_eq!(
+            vm_resources
+                .mmds
+                .as_ref()
+                .unwrap()
+                .lock()
+                .unwrap()
+                .version(),
+            MmdsVersion::V2
+        );
+        assert_eq!(
+            device_manager_state.pci_state.mmds_version.unwrap(),
+            MmdsVersion::V2.into()
+        );
+        assert_eq!(
+            expected_vm_resources,
+            serde_json::to_string_pretty(&VmmConfig::from(&*vm_resources)).unwrap()
+        );
+    }
+}