From 52c80cfcf53c9769af88e54bbf09ea2e133fabbf Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Tue, 28 Apr 2020 12:08:51 +0200
Subject: [PATCH] vmm: Snapshot and restore DeviceManager state

The DeviceManager itself must be snapshotted in order to store the
information regarding the devices associated with it, which effectively
means we need to store the device tree.

The mechanics to snapshot and restore the DeviceManagerState are added
to the existing snapshot() and restore() implementations.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vm-device/src/lib.rs      |  6 +++--
 vmm/src/device_manager.rs | 56 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/vm-device/src/lib.rs b/vm-device/src/lib.rs
index 457efb9de..f7914f552 100644
--- a/vm-device/src/lib.rs
+++ b/vm-device/src/lib.rs
@@ -1,3 +1,5 @@
+#[macro_use]
+extern crate serde_derive;
 extern crate vm_memory;
 
 pub mod interrupt;
@@ -8,7 +10,7 @@ use vm_memory::{
 };
 
 /// Type of Message Singaled Interrupt
-#[derive(Copy, Clone, PartialEq)]
+#[derive(Copy, Clone, PartialEq, Serialize, Deserialize)]
 pub enum MsiIrqType {
     /// PCI MSI IRQ numbers.
     PciMsi,
@@ -20,7 +22,7 @@ pub enum MsiIrqType {
 
 /// Enumeration for device resources.
 #[allow(missing_docs)]
-#[derive(Clone)]
+#[derive(Clone, Serialize, Deserialize)]
 pub enum Resource {
     /// IO Port address range.
     PioAddressRange { base: u16, size: u16 },
diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index 93e02dd62..08bcee456 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -55,7 +55,10 @@ use vm_memory::guest_memory::FileOffset;
 use vm_memory::{
     Address, GuestAddress, GuestAddressSpace, GuestRegionMmap, GuestUsize, MmapRegion,
 };
-use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
+use vm_migration::{
+    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
+    Transportable,
+};
 #[cfg(feature = "pci_support")]
 use vm_virtio::transport::VirtioPciDevice;
 use vm_virtio::transport::VirtioTransport;
@@ -545,13 +548,19 @@ impl Drop for ActivatedBackend {
 }
 
 #[allow(unused)]
-#[derive(Default)]
+#[derive(Clone, Default, Serialize, Deserialize)]
 struct Node {
     resources: Vec<Resource>,
     parent: Option<String>,
     child: Option<String>,
 }
 
+#[derive(Serialize, Deserialize)]
+struct DeviceManagerState {
+    device_tree: HashMap<String, Node>,
+    device_id_cnt: Wrapping<usize>,
+}
+
 pub struct DeviceManager {
     // Manage address space related to devices
     address_manager: Arc<AddressManager>,
@@ -795,6 +804,20 @@ impl DeviceManager {
         Ok(device_manager)
     }
 
+    fn state(&self) -> DeviceManagerState {
+        DeviceManagerState {
+            device_tree: self.device_tree.clone(),
+            device_id_cnt: self.device_id_cnt,
+        }
+    }
+
+    fn set_state(&mut self, state: &DeviceManagerState) -> DeviceManagerResult<()> {
+        self.device_tree = state.device_tree.clone();
+        self.device_id_cnt = state.device_id_cnt;
+
+        Ok(())
+    }
+
     fn add_migratable_device(&mut self, migratable_device: Arc<Mutex<dyn Migratable>>) {
         let id = migratable_device.lock().unwrap().id();
         self.migratable_devices.push((id, migratable_device));
@@ -2792,16 +2815,43 @@ impl Snapshottable for DeviceManager {
     fn snapshot(&self) -> std::result::Result<Snapshot, MigratableError> {
         let mut snapshot = Snapshot::new(DEVICE_MANAGER_SNAPSHOT_ID);
 
-        // We aggregate all devices snapshot.
+        // We aggregate all devices snapshots.
         for (_, dev) in self.migratable_devices.iter() {
             let device_snapshot = dev.lock().unwrap().snapshot()?;
             snapshot.add_snapshot(device_snapshot);
         }
 
+        // Then we store the DeviceManager state.
+        snapshot.add_data_section(SnapshotDataSection {
+            id: format!("{}-section", DEVICE_MANAGER_SNAPSHOT_ID),
+            snapshot: serde_json::to_vec(&self.state())
+                .map_err(|e| MigratableError::Snapshot(e.into()))?,
+        });
+
         Ok(snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
+        // Let's first restore the DeviceManager.
+        if let Some(device_manager_section) = snapshot
+            .snapshot_data
+            .get(&format!("{}-section", DEVICE_MANAGER_SNAPSHOT_ID))
+        {
+            let device_manager_state = serde_json::from_slice(&device_manager_section.snapshot)
+                .map_err(|e| {
+                    MigratableError::Restore(anyhow!("Could not deserialize DeviceManager {}", e))
+                })?;
+
+            self.set_state(&device_manager_state).map_err(|e| {
+                MigratableError::Restore(anyhow!("Could not restore DeviceManager state {:?}", e))
+            })?;
+        } else {
+            return Err(MigratableError::Restore(anyhow!(
+                "Could not find DeviceManager snapshot section"
+            )));
+        }
+
+        // Then restore all devices associated with the DeviceManager.
         for (id, dev) in self.migratable_devices.iter() {
             debug!("Restoring {} from DeviceManager", id);
             if let Some(snapshot) = snapshot.snapshots.get(id) {