From 4da3bdcd6ecb6a81cc8363b593c8f94a3c4cb215 Mon Sep 17 00:00:00 2001
From: Henry Wang <Henry.Wang@arm.com>
Date: Sun, 6 Jun 2021 16:13:59 +0800
Subject: [PATCH] vmm: Split restore device_manager and devices

Signed-off-by: Henry Wang <Henry.Wang@arm.com>
---
 vmm/src/device_manager.rs | 59 ++++++++++++++++++++++-----------------
 vmm/src/vm.rs             | 11 ++++++++
 2 files changed, 44 insertions(+), 26 deletions(-)

diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index 3553578b8..a6d5af20a 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -3464,6 +3464,39 @@ impl DeviceManager {
         self.device_tree.clone()
     }
 
+    pub fn restore_devices(
+        &mut self,
+        snapshot: Snapshot,
+    ) -> std::result::Result<(), MigratableError> {
+        // Finally, restore all devices associated with the DeviceManager.
+        // It's important to restore devices in the right order, that's why
+        // the device tree is the right way to ensure we restore a child before
+        // its parent node.
+        for node in self
+            .device_tree
+            .lock()
+            .unwrap()
+            .breadth_first_traversal()
+            .rev()
+        {
+            // Restore the node
+            if let Some(migratable) = &node.migratable {
+                debug!("Restoring {} from DeviceManager", node.id);
+                if let Some(snapshot) = snapshot.snapshots.get(&node.id) {
+                    migratable.lock().unwrap().pause()?;
+                    migratable.lock().unwrap().restore(*snapshot.clone())?;
+                } else {
+                    return Err(MigratableError::Restore(anyhow!(
+                        "Missing device {}",
+                        node.id
+                    )));
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     #[cfg(feature = "acpi")]
     #[cfg(target_arch = "x86_64")]
     pub fn notify_power_button(&self) -> DeviceManagerResult<()> {
@@ -3969,32 +4002,6 @@ impl Snapshottable for DeviceManager {
         self.create_devices(None, None)
             .map_err(|e| MigratableError::Restore(anyhow!("Could not create devices {:?}", e)))?;
 
-        // Finally, restore all devices associated with the DeviceManager.
-        // It's important to restore devices in the right order, that's why
-        // the device tree is the right way to ensure we restore a child before
-        // its parent node.
-        for node in self
-            .device_tree
-            .lock()
-            .unwrap()
-            .breadth_first_traversal()
-            .rev()
-        {
-            // Restore the node
-            if let Some(migratable) = &node.migratable {
-                debug!("Restoring {} from DeviceManager", node.id);
-                if let Some(snapshot) = snapshot.snapshots.get(&node.id) {
-                    migratable.lock().unwrap().pause()?;
-                    migratable.lock().unwrap().restore(*snapshot.clone())?;
-                } else {
-                    return Err(MigratableError::Restore(anyhow!(
-                        "Missing device {}",
-                        node.id
-                    )));
-                }
-            }
-        }
-
         Ok(())
     }
 }
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index b23093d45..a5dd446b5 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -2277,6 +2277,17 @@ impl Snapshottable for Vm {
         #[cfg(target_arch = "aarch64")]
         self.restore_vgic_and_enable_interrupt(&snapshot)?;
 
+        if let Some(device_manager_snapshot) = snapshot.snapshots.get(DEVICE_MANAGER_SNAPSHOT_ID) {
+            self.device_manager
+                .lock()
+                .unwrap()
+                .restore_devices(*device_manager_snapshot.clone())?;
+        } else {
+            return Err(MigratableError::Restore(anyhow!(
+                "Missing device manager snapshot"
+            )));
+        }
+
         // Now we can start all vCPUs from here.
         self.cpu_manager
             .lock()