From 3931b99d4e45d062c53f585ea41f86cd2746f3a1 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Fri, 2 Dec 2022 16:03:50 +0100
Subject: [PATCH] vm-migration: Introduce new constructor for Snapshot

This simplifies the Snapshot creation as we expect a SnapshotData to be
provided most of the time.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vm-migration/src/lib.rs   | 24 +++++++++++-------------
 vmm/src/cpu.rs            |  9 ++++-----
 vmm/src/device_manager.rs |  5 +----
 vmm/src/memory_manager.rs |  8 ++------
 vmm/src/vm.rs             |  4 ++--
 5 files changed, 20 insertions(+), 30 deletions(-)

diff --git a/vm-migration/src/lib.rs b/vm-migration/src/lib.rs
index b426b6e14..367ad49c0 100644
--- a/vm-migration/src/lib.rs
+++ b/vm-migration/src/lib.rs
@@ -148,15 +148,19 @@ pub struct Snapshot {
 }
 
 impl Snapshot {
+    pub fn from_data(data: SnapshotData) -> Self {
+        Snapshot {
+            snapshot_data: Some(data),
+            ..Default::default()
+        }
+    }
+
     /// Create from state that can be serialized
     pub fn new_from_state<T>(state: &T) -> Result<Self, MigratableError>
     where
         T: Serialize,
     {
-        let mut snapshot_data = Snapshot::default();
-        snapshot_data.add_data(SnapshotData::new_from_state(state)?);
-
-        Ok(snapshot_data)
+        Ok(Snapshot::from_data(SnapshotData::new_from_state(state)?))
     }
 
     /// Create from versioned state
@@ -164,10 +168,9 @@ impl Snapshot {
     where
         T: Versionize + VersionMapped,
     {
-        let mut snapshot_data = Snapshot::default();
-        snapshot_data.add_data(SnapshotData::new_from_versioned_state(state)?);
-
-        Ok(snapshot_data)
+        Ok(Snapshot::from_data(SnapshotData::new_from_versioned_state(
+            state,
+        )?))
     }
 
     /// Add a sub-component's Snapshot to the Snapshot.
@@ -175,11 +178,6 @@ impl Snapshot {
         self.snapshots.insert(id, snapshot);
     }
 
-    /// Add a SnapshotData to the component snapshot data.
-    pub fn add_data(&mut self, section: SnapshotData) {
-        self.snapshot_data = Some(section);
-    }
-
     /// Generate the state data from the snapshot
     pub fn to_state<'a, T>(&'a self) -> Result<T, MigratableError>
     where
diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index 3db6b8334..7231d33eb 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -399,12 +399,11 @@ impl Snapshottable for Vcpu {
             .state()
             .map_err(|e| MigratableError::Pause(anyhow!("Could not get vCPU state {:?}", e)))?;
 
-        let mut vcpu_snapshot = Snapshot::default();
-        vcpu_snapshot.add_data(SnapshotData::new_from_state(&saved_state)?);
+        self.saved_state = Some(saved_state.clone());
 
-        self.saved_state = Some(saved_state);
-
-        Ok(vcpu_snapshot)
+        Ok(Snapshot::from_data(SnapshotData::new_from_state(
+            &saved_state,
+        )?))
     }
 }
 
diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index fe74c11e1..c14c8c6ee 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -4445,7 +4445,7 @@ impl Snapshottable for DeviceManager {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let mut snapshot = Snapshot::default();
+        let mut snapshot = Snapshot::from_data(SnapshotData::new_from_state(&self.state())?);
 
         // We aggregate all devices snapshots.
         for (_, device_node) in self.device_tree.lock().unwrap().iter() {
@@ -4455,9 +4455,6 @@ impl Snapshottable for DeviceManager {
             }
         }
 
-        // Then we store the DeviceManager state.
-        snapshot.add_data(SnapshotData::new_from_state(&self.state())?);
-
         Ok(snapshot)
     }
 }
diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index 96c369da9..da42bbf97 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -2428,8 +2428,6 @@ impl Snapshottable for MemoryManager {
     }
 
     fn snapshot(&mut self) -> result::Result<Snapshot, MigratableError> {
-        let mut memory_manager_snapshot = Snapshot::default();
-
         let memory_ranges = self.memory_range_table(true)?;
 
         // Store locally this list of ranges as it will be used through the
@@ -2442,11 +2440,9 @@ impl Snapshottable for MemoryManager {
         // memory range content for the ranges requiring it.
         self.snapshot_memory_ranges = memory_ranges;
 
-        memory_manager_snapshot.add_data(SnapshotData::new_from_versioned_state(
+        Ok(Snapshot::from_data(SnapshotData::new_from_versioned_state(
             &self.snapshot_data(),
-        )?);
-
-        Ok(memory_manager_snapshot)
+        )?))
     }
 }
 
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 2e90010e6..0c8f2f951 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -2508,7 +2508,6 @@ impl Snapshottable for Vm {
             })?
         };
 
-        let mut vm_snapshot = Snapshot::default();
         let vm_snapshot_data = serde_json::to_vec(&VmSnapshot {
             #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
             clock: self.saved_clock,
@@ -2517,6 +2516,8 @@ impl Snapshottable for Vm {
         })
         .map_err(|e| MigratableError::Snapshot(e.into()))?;
 
+        let mut vm_snapshot = Snapshot::from_data(SnapshotData(vm_snapshot_data));
+
         let (id, snapshot) = {
             let mut cpu_manager = self.cpu_manager.lock().unwrap();
             (cpu_manager.id(), cpu_manager.snapshot()?)
@@ -2532,7 +2533,6 @@ impl Snapshottable for Vm {
             (device_manager.id(), device_manager.snapshot()?)
         };
         vm_snapshot.add_snapshot(id, snapshot);
-        vm_snapshot.add_data(SnapshotData(vm_snapshot_data));
 
         event!("vm", "snapshotted");
         Ok(vm_snapshot)