vm-migration: Don't store the id as part of Snapshot structure

The information about the identifier related to a Snapshot is only
relevant from the BTreeMap perspective, which is why we can get rid of
the duplicated identifier in every Snapshot structure.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>

devices/src/gic.rs

@@ -154,7 +154,7 @@ impl Snapshottable for Gic {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         let vgic = self.vgic.as_ref().unwrap().clone();
         let state = vgic.lock().unwrap().state().unwrap();
-        Snapshot::new_from_state(&self.id(), &state)
+        Snapshot::new_from_state(&state)
     }
 }
 

devices/src/ioapic.rs

@@ -438,7 +438,7 @@ impl Snapshottable for Ioapic {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

devices/src/legacy/gpio_pl061.rs

@@ -328,7 +328,7 @@ impl Snapshottable for Gpio {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

devices/src/legacy/serial.rs

@@ -334,7 +334,7 @@ impl Snapshottable for Serial {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

devices/src/legacy/uart_pl011.rs

@@ -454,7 +454,7 @@ impl Snapshottable for Pl011 {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

pci/src/configuration.rs

@@ -1072,7 +1072,7 @@ impl Snapshottable for PciConfiguration {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id(), &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

pci/src/msi.rs

@@ -288,6 +288,6 @@ impl Snapshottable for MsiConfig {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id(), &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }

pci/src/msix.rs

@@ -433,7 +433,7 @@ impl Snapshottable for MsixConfig {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id(), &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

pci/src/vfio.rs

@@ -1235,20 +1235,19 @@ impl Snapshottable for VfioCommon {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let mut vfio_common_snapshot =
+        let mut vfio_common_snapshot = Snapshot::new_from_versioned_state(&self.state())?;
-            Snapshot::new_from_versioned_state(&self.id(), &self.state())?;
 
         // Snapshot PciConfiguration
-        vfio_common_snapshot.add_snapshot(self.configuration.snapshot()?);
+        vfio_common_snapshot.add_snapshot(self.configuration.id(), self.configuration.snapshot()?);
 
         // Snapshot MSI
         if let Some(msi) = &mut self.interrupt.msi {
-            vfio_common_snapshot.add_snapshot(msi.cfg.snapshot()?);
+            vfio_common_snapshot.add_snapshot(msi.cfg.id(), msi.cfg.snapshot()?);
         }
 
         // Snapshot MSI-X
         if let Some(msix) = &mut self.interrupt.msix {
-            vfio_common_snapshot.add_snapshot(msix.bar.snapshot()?);
+            vfio_common_snapshot.add_snapshot(msix.bar.id(), msix.bar.snapshot()?);
         }
 
         Ok(vfio_common_snapshot)
@@ -1740,10 +1739,10 @@ impl Snapshottable for VfioPciDevice {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let mut vfio_pci_dev_snapshot = Snapshot::new(&self.id);
+        let mut vfio_pci_dev_snapshot = Snapshot::default();
 
         // Snapshot VfioCommon
-        vfio_pci_dev_snapshot.add_snapshot(self.common.snapshot()?);
+        vfio_pci_dev_snapshot.add_snapshot(self.common.id(), self.common.snapshot()?);
 
         Ok(vfio_pci_dev_snapshot)
     }

pci/src/vfio_user.rs

@@ -535,10 +535,10 @@ impl Snapshottable for VfioUserPciDevice {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let mut vfio_pci_dev_snapshot = Snapshot::new(&self.id);
+        let mut vfio_pci_dev_snapshot = Snapshot::default();
 
         // Snapshot VfioCommon
-        vfio_pci_dev_snapshot.add_snapshot(self.common.snapshot()?);
+        vfio_pci_dev_snapshot.add_snapshot(self.common.id(), self.common.snapshot()?);
 
         Ok(vfio_pci_dev_snapshot)
     }

virtio-devices/src/balloon.rs

@@ -585,7 +585,7 @@ impl Snapshottable for Balloon {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id(), &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 impl Transportable for Balloon {}

virtio-devices/src/block.rs

@@ -743,7 +743,7 @@ impl Snapshottable for Block {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id(), &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 impl Transportable for Block {}

virtio-devices/src/console.rs

@@ -819,7 +819,7 @@ impl Snapshottable for Console {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 impl Transportable for Console {}

virtio-devices/src/iommu.rs

@@ -1109,7 +1109,7 @@ impl Snapshottable for Iommu {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 impl Transportable for Iommu {}

virtio-devices/src/mem.rs

@@ -1014,7 +1014,7 @@ impl Snapshottable for Mem {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id(), &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 impl Transportable for Mem {}

virtio-devices/src/net.rs

@@ -857,7 +857,7 @@ impl Snapshottable for Net {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 impl Transportable for Net {}

virtio-devices/src/pmem.rs

@@ -467,7 +467,7 @@ impl Snapshottable for Pmem {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

virtio-devices/src/rng.rs

@@ -323,7 +323,7 @@ impl Snapshottable for Rng {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

virtio-devices/src/transport/pci_common_config.rs

@@ -321,7 +321,7 @@ impl Snapshottable for VirtioPciCommonConfig {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id(), &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

virtio-devices/src/transport/pci_device.rs

@@ -1272,18 +1272,20 @@ impl Snapshottable for VirtioPciDevice {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let mut virtio_pci_dev_snapshot =
+        let mut virtio_pci_dev_snapshot = Snapshot::new_from_versioned_state(&self.state())?;
-            Snapshot::new_from_versioned_state(&self.id, &self.state())?;
 
         // Snapshot PciConfiguration
-        virtio_pci_dev_snapshot.add_snapshot(self.configuration.snapshot()?);
+        virtio_pci_dev_snapshot
+            .add_snapshot(self.configuration.id(), self.configuration.snapshot()?);
 
         // Snapshot VirtioPciCommonConfig
-        virtio_pci_dev_snapshot.add_snapshot(self.common_config.snapshot()?);
+        virtio_pci_dev_snapshot
+            .add_snapshot(self.common_config.id(), self.common_config.snapshot()?);
 
         // Snapshot MSI-X
         if let Some(msix_config) = &self.msix_config {
-            virtio_pci_dev_snapshot.add_snapshot(msix_config.lock().unwrap().snapshot()?);
+            let mut msix_config = msix_config.lock().unwrap();
+            virtio_pci_dev_snapshot.add_snapshot(msix_config.id(), msix_config.snapshot()?);
         }
 
         Ok(virtio_pci_dev_snapshot)

virtio-devices/src/vdpa.rs

@@ -489,12 +489,9 @@ impl Snapshottable for Vdpa {
             )));
         }
 
-        let snapshot = Snapshot::new_from_versioned_state(
+        let snapshot = Snapshot::new_from_versioned_state(&self.state().map_err(|e| {
-            &self.id(),
+            MigratableError::Snapshot(anyhow!("Error snapshotting vDPA device: {:?}", e))
-            &self.state().map_err(|e| {
+        })?)?;
-                MigratableError::Snapshot(anyhow!("Error snapshotting vDPA device: {:?}", e))
-            })?,
-        )?;
 
         // Force the vhost handler to be dropped in order to close the vDPA
         // file. This will ensure the device can be accessed if the VM is

virtio-devices/src/vhost_user/blk.rs

@@ -381,7 +381,7 @@ impl Snapshottable for Blk {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        self.vu_common.snapshot(&self.id(), &self.state())
+        self.vu_common.snapshot(&self.state())
     }
 }
 impl Transportable for Blk {}

virtio-devices/src/vhost_user/fs.rs

@@ -651,7 +651,7 @@ impl Snapshottable for Fs {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        self.vu_common.snapshot(&self.id(), &self.state())
+        self.vu_common.snapshot(&self.state())
     }
 }
 impl Transportable for Fs {}

virtio-devices/src/vhost_user/mod.rs

@@ -430,15 +430,11 @@ impl VhostUserCommon {
         }
     }
 
-    pub fn snapshot<T>(
+    pub fn snapshot<T>(&mut self, state: &T) -> std::result::Result<Snapshot, MigratableError>
-        &mut self,
-        id: &str,
-        state: &T,
-    ) -> std::result::Result<Snapshot, MigratableError>
     where
         T: Versionize + VersionMapped,
     {
-        let snapshot = Snapshot::new_from_versioned_state(id, state)?;
+        let snapshot = Snapshot::new_from_versioned_state(state)?;
 
         if self.migration_started {
             self.shutdown();

virtio-devices/src/vhost_user/net.rs

@@ -418,7 +418,7 @@ impl Snapshottable for Net {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        self.vu_common.snapshot(&self.id(), &self.state())
+        self.vu_common.snapshot(&self.state())
     }
 }
 impl Transportable for Net {}

virtio-devices/src/vsock/device.rs

@@ -518,7 +518,7 @@ where
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 impl<B> Transportable for Vsock<B> where B: VsockBackend + Sync + 'static {}

virtio-devices/src/watchdog.rs

@@ -419,7 +419,7 @@ impl Snapshottable for Watchdog {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Snapshot::new_from_versioned_state(&self.id, &self.state())
+        Snapshot::new_from_versioned_state(&self.state())
     }
 }
 

vm-migration/src/lib.rs

@@ -139,9 +139,6 @@ impl SnapshotData {
 /// the actual device snapshot data.
 #[derive(Clone, Default, Deserialize, Serialize)]
 pub struct Snapshot {
-    /// The Snapshottable component id.
-    pub id: String,
-
     /// The Snapshottable component snapshots.
     pub snapshots: std::collections::BTreeMap<String, Snapshot>,
 
@@ -151,39 +148,31 @@ pub struct Snapshot {
 }
 
 impl Snapshot {
-    /// Create an empty Snapshot.
-    pub fn new(id: &str) -> Self {
-        Snapshot {
-            id: id.to_string(),
-            ..Default::default()
-        }
-    }
-
     /// Create from state that can be serialized
-    pub fn new_from_state<T>(id: &str, state: &T) -> Result<Self, MigratableError>
+    pub fn new_from_state<T>(state: &T) -> Result<Self, MigratableError>
     where
         T: Serialize,
     {
-        let mut snapshot_data = Snapshot::new(id);
+        let mut snapshot_data = Snapshot::default();
         snapshot_data.add_data_section(SnapshotData::new_from_state(state)?);
 
         Ok(snapshot_data)
     }
 
     /// Create from versioned state
-    pub fn new_from_versioned_state<T>(id: &str, state: &T) -> Result<Self, MigratableError>
+    pub fn new_from_versioned_state<T>(state: &T) -> Result<Self, MigratableError>
     where
         T: Versionize + VersionMapped,
     {
-        let mut snapshot_data = Snapshot::new(id);
+        let mut snapshot_data = Snapshot::default();
         snapshot_data.add_data_section(SnapshotData::new_from_versioned_state(state)?);
 
         Ok(snapshot_data)
     }
 
     /// Add a sub-component's Snapshot to the Snapshot.
-    pub fn add_snapshot(&mut self, snapshot: Snapshot) {
+    pub fn add_snapshot(&mut self, id: String, snapshot: Snapshot) {
-        self.snapshots.insert(snapshot.id.clone(), snapshot);
+        self.snapshots.insert(id, snapshot);
     }
 
     /// Add a SnapshotData to the component snapshot data.
@@ -240,7 +229,7 @@ pub trait Snapshottable: Pausable {
 
     /// Take a component snapshot.
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        Ok(Snapshot::new(""))
+        Ok(Snapshot::default())
     }
 }
 
@@ -271,7 +260,7 @@ pub trait Transportable: Pausable + Snapshottable {
     /// * `source_url` - The source URL to fetch the snapshot from. This could be an HTTP
     ///                  endpoint, a TCP address or a local file.
     fn recv(&self, _source_url: &str) -> std::result::Result<Snapshot, MigratableError> {
-        Ok(Snapshot::new(""))
+        Ok(Snapshot::default())
     }
 }
 

vmm/src/cpu.rs

@@ -387,11 +387,10 @@ impl Vcpu {
     }
 }
 
-const VCPU_SNAPSHOT_ID: &str = "vcpu";
 impl Pausable for Vcpu {}
 impl Snapshottable for Vcpu {
     fn id(&self) -> String {
-        VCPU_SNAPSHOT_ID.to_string()
+        self.id.to_string()
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
@@ -400,9 +399,7 @@ impl Snapshottable for Vcpu {
             .state()
             .map_err(|e| MigratableError::Pause(anyhow!("Could not get vCPU state {:?}", e)))?;
 
-        // TODO: The special format of the CPU id can be removed once ready to
+        let mut vcpu_snapshot = Snapshot::default();
-        // break live upgrade.
-        let mut vcpu_snapshot = Snapshot::new(&format!("{:03}", self.id));
         vcpu_snapshot.add_data_section(SnapshotData::new_from_state(&saved_state)?);
 
         self.saved_state = Some(saved_state);
@@ -776,7 +773,7 @@ impl CpuManager {
                 cpu_id,
                 // TODO: The special format of the CPU id can be removed once
                 // ready to break live upgrade.
-                snapshot_from_id(snapshot.as_ref(), &format!("{:03}", cpu_id)),
+                snapshot_from_id(snapshot.as_ref(), cpu_id.to_string().as_str()),
             )?);
         }
 
@@ -2076,12 +2073,12 @@ impl Snapshottable for CpuManager {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let mut cpu_manager_snapshot = Snapshot::new(CPU_MANAGER_SNAPSHOT_ID);
+        let mut cpu_manager_snapshot = Snapshot::default();
 
         // The CpuManager snapshot is a collection of all vCPUs snapshots.
         for vcpu in &self.vcpus {
-            let cpu_snapshot = vcpu.lock().unwrap().snapshot()?;
+            let mut vcpu = vcpu.lock().unwrap();
-            cpu_manager_snapshot.add_snapshot(cpu_snapshot);
+            cpu_manager_snapshot.add_snapshot(vcpu.id(), vcpu.snapshot()?);
         }
 
         Ok(cpu_manager_snapshot)

vmm/src/device_manager.rs

@@ -4445,13 +4445,13 @@ impl Snapshottable for DeviceManager {
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
-        let mut snapshot = Snapshot::new(DEVICE_MANAGER_SNAPSHOT_ID);
+        let mut snapshot = Snapshot::default();
 
         // We aggregate all devices snapshots.
         for (_, device_node) in self.device_tree.lock().unwrap().iter() {
             if let Some(migratable) = &device_node.migratable {
-                let device_snapshot = migratable.lock().unwrap().snapshot()?;
+                let mut migratable = migratable.lock().unwrap();
-                snapshot.add_snapshot(device_snapshot);
+                snapshot.add_snapshot(migratable.id(), migratable.snapshot()?);
             }
         }
 

vmm/src/memory_manager.rs

@@ -2428,7 +2428,7 @@ impl Snapshottable for MemoryManager {
     }
 
     fn snapshot(&mut self) -> result::Result<Snapshot, MigratableError> {
-        let mut memory_manager_snapshot = Snapshot::new(MEMORY_MANAGER_SNAPSHOT_ID);
+        let mut memory_manager_snapshot = Snapshot::default();
 
         let memory_ranges = self.memory_range_table(true)?;
 

vmm/src/vm.rs

@@ -2508,7 +2508,7 @@ impl Snapshottable for Vm {
             })?
         };
 
-        let mut vm_snapshot = Snapshot::new(VM_SNAPSHOT_ID);
+        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,10 +2517,21 @@ impl Snapshottable for Vm {
         })
         .map_err(|e| MigratableError::Snapshot(e.into()))?;
 
-        vm_snapshot.add_snapshot(self.cpu_manager.lock().unwrap().snapshot()?);
+        let (id, snapshot) = {
-        vm_snapshot.add_snapshot(self.memory_manager.lock().unwrap().snapshot()?);
+            let mut cpu_manager = self.cpu_manager.lock().unwrap();
-
+            (cpu_manager.id(), cpu_manager.snapshot()?)
-        vm_snapshot.add_snapshot(self.device_manager.lock().unwrap().snapshot()?);
+        };
+        vm_snapshot.add_snapshot(id, snapshot);
+        let (id, snapshot) = {
+            let mut memory_manager = self.memory_manager.lock().unwrap();
+            (memory_manager.id(), memory_manager.snapshot()?)
+        };
+        vm_snapshot.add_snapshot(id, snapshot);
+        let (id, snapshot) = {
+            let mut device_manager = self.device_manager.lock().unwrap();
+            (device_manager.id(), device_manager.snapshot()?)
+        };
+        vm_snapshot.add_snapshot(id, snapshot);
         vm_snapshot.add_data_section(SnapshotData(vm_snapshot_data));
 
         event!("vm", "snapshotted");