From 1592a9292f7651f5da27ba54633c36503932472e Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Mon, 27 Apr 2020 13:36:41 +0200
Subject: [PATCH] vm-virtio: pmem: Expect an identifier upon device creation

This identifier is chosen from the DeviceManager so that it will manage
all identifiers across the VM, which will ensure uniqueness.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vm-virtio/src/pmem.rs     | 15 +++++++--------
 vmm/src/device_manager.rs | 25 +++++++++++++++++--------
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/vm-virtio/src/pmem.rs b/vm-virtio/src/pmem.rs
index d916ed739..d14ab0a82 100644
--- a/vm-virtio/src/pmem.rs
+++ b/vm-virtio/src/pmem.rs
@@ -337,6 +337,7 @@ impl PmemEpollHandler {
 }
 
 pub struct Pmem {
+    id: String,
     kill_evt: Option<EventFd>,
     pause_evt: Option<EventFd>,
     disk: Option<File>,
@@ -363,6 +364,7 @@ pub struct PmemState {
 
 impl Pmem {
     pub fn new(
+        id: String,
         disk: File,
         addr: GuestAddress,
         mapping: UserspaceMapping,
@@ -381,6 +383,7 @@ impl Pmem {
         }
 
         Ok(Pmem {
+            id,
             kill_evt: None,
             pause_evt: None,
             disk: Some(disk),
@@ -571,19 +574,18 @@ impl VirtioDevice for Pmem {
 }
 
 virtio_pausable!(Pmem);
-const PMEM_SNAPSHOT_ID: &str = "virtio-pmem";
 impl Snapshottable for Pmem {
     fn id(&self) -> String {
-        PMEM_SNAPSHOT_ID.to_string()
+        self.id.clone()
     }
 
     fn snapshot(&self) -> std::result::Result<Snapshot, MigratableError> {
         let snapshot =
             serde_json::to_vec(&self.state()).map_err(|e| MigratableError::Snapshot(e.into()))?;
 
-        let mut pmem_snapshot = Snapshot::new(PMEM_SNAPSHOT_ID);
+        let mut pmem_snapshot = Snapshot::new(self.id.as_str());
         pmem_snapshot.add_data_section(SnapshotDataSection {
-            id: format!("{}-section", PMEM_SNAPSHOT_ID),
+            id: format!("{}-section", self.id),
             snapshot,
         });
 
@@ -591,10 +593,7 @@ impl Snapshottable for Pmem {
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        if let Some(pmem_section) = snapshot
-            .snapshot_data
-            .get(&format!("{}-section", PMEM_SNAPSHOT_ID))
-        {
+        if let Some(pmem_section) = snapshot.snapshot_data.get(&format!("{}-section", self.id)) {
             let pmem_state = match serde_json::from_slice(&pmem_section.snapshot) {
                 Ok(state) => state,
                 Err(error) => {
diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index 3d472543c..79ac772db 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -1560,9 +1560,13 @@ impl DeviceManager {
         &mut self,
         pmem_cfg: &mut PmemConfig,
     ) -> DeviceManagerResult<(VirtioDeviceArc, bool, Option<String>)> {
-        if pmem_cfg.id.is_none() {
-            pmem_cfg.id = Some(self.next_device_name(PMEM_DEVICE_NAME_PREFIX)?);
-        }
+        let id = if let Some(id) = &pmem_cfg.id {
+            id.clone()
+        } else {
+            let id = self.next_device_name(PMEM_DEVICE_NAME_PREFIX)?;
+            pmem_cfg.id = Some(id.clone());
+            id
+        };
 
         let (custom_flags, set_len) = if pmem_cfg.file.is_dir() {
             if pmem_cfg.size.is_none() {
@@ -1646,13 +1650,18 @@ impl DeviceManager {
         };
 
         let virtio_pmem_device = Arc::new(Mutex::new(
-            vm_virtio::Pmem::new(file, pmem_guest_addr, mapping, mmap_region, pmem_cfg.iommu)
-                .map_err(DeviceManagerError::CreateVirtioPmem)?,
+            vm_virtio::Pmem::new(
+                id,
+                file,
+                pmem_guest_addr,
+                mapping,
+                mmap_region,
+                pmem_cfg.iommu,
+            )
+            .map_err(DeviceManagerError::CreateVirtioPmem)?,
         ));
 
-        let migratable = Arc::clone(&virtio_pmem_device) as Arc<Mutex<dyn Migratable>>;
-        let id = migratable.lock().unwrap().id();
-        self.migratable_devices.push((id, migratable));
+        self.add_migratable_device(Arc::clone(&virtio_pmem_device) as Arc<Mutex<dyn Migratable>>);
 
         Ok((
             Arc::clone(&virtio_pmem_device) as VirtioDeviceArc,