From 030d9763b2fc3b04a45ae764821914e9369e0491 Mon Sep 17 00:00:00 2001
From: zhouli57 <zhouli57@huawei.com>
Date: Wed, 5 Jan 2022 09:48:27 +0800
Subject: [PATCH 03/10] migration: use device id as snapshot id

At present, snapshots use the device initialization sequence as the key,
but after the introduction of the device hot-plug mechanism, due to the
dynamic addition of devices, the sequence of snapshot restoration may be
in consistent with the original sequence, resulting in abnormal device
state restoration. Therefore, a new interface is provided, which can
specify the device id during registration, so as to ensure that the
snapshot is restored to match the original device.
In addition, an unregister interface has been added for device
destruction to clean up related resources.

Signed-off-by: zhouli57 <zhouli57@huawei.com>
---
 Cargo.lock                                    |   1 +
 address_space/src/state.rs                    |   4 +-
 .../src/interrupt_controller/aarch64/gicv3.rs |  10 +-
 hypervisor/src/kvm/mod.rs                     |   6 +-
 machine/src/lib.rs                            |  17 +-
 machine/src/standard_vm/mod.rs                |  44 +++--
 migration/Cargo.toml                          |   1 +
 migration/src/device_state.rs                 |   7 -
 migration/src/lib.rs                          |   4 +-
 migration/src/manager.rs                      | 181 +++++++++++++-----
 migration/src/snapshot.rs                     |  63 ++++--
 migration_derive/src/struct_parser.rs         |  11 +-
 pci/src/msix.rs                               |  10 +-
 pci/src/root_port.rs                          |   9 +-
 virtio/src/block.rs                           |   4 +
 virtio/src/net.rs                             |   4 +
 virtio/src/virtio_pci.rs                      |  20 +-
 17 files changed, 287 insertions(+), 109 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index df5dc88..215f1d5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -327,6 +327,7 @@ version = "2.1.0"
 dependencies = [
  "error-chain",
  "kvm-ioctls",
+ "log",
  "migration_derive",
  "once_cell",
  "serde",
diff --git a/address_space/src/state.rs b/address_space/src/state.rs
index 2347378..eb34e91 100644
--- a/address_space/src/state.rs
+++ b/address_space/src/state.rs
@@ -17,7 +17,7 @@ use std::sync::Arc;
 
 use crate::{AddressSpace, FileBackend, GuestAddress, HostMemMapping, Region};
 use migration::errors::{ErrorKind, Result, ResultExt};
-use migration::{DeviceStateDesc, FieldDesc, MigrationHook, MigrationManager, StateTransfer};
+use migration::{DeviceStateDesc, FieldDesc, MigrationHook, StateTransfer};
 use util::byte_code::ByteCode;
 use util::unix::host_page_size;
 
@@ -76,7 +76,7 @@ impl StateTransfer for AddressSpace {
 }
 
 impl MigrationHook for AddressSpace {
-    fn pre_save(&self, _id: u64, writer: &mut dyn Write) -> Result<()> {
+    fn pre_save(&self, _id: &str, writer: &mut dyn Write) -> Result<()> {
         let ram_state = self.get_state_vec()?;
         writer.write_all(&ram_state)?;
         let padding_buffer =
diff --git a/devices/src/interrupt_controller/aarch64/gicv3.rs b/devices/src/interrupt_controller/aarch64/gicv3.rs
index 21008a6..79ebb27 100644
--- a/devices/src/interrupt_controller/aarch64/gicv3.rs
+++ b/devices/src/interrupt_controller/aarch64/gicv3.rs
@@ -21,7 +21,7 @@ use super::{
 use crate::interrupt_controller::errors::{ErrorKind, Result, ResultExt};
 use hypervisor::kvm::KVM_FDS;
 use machine_manager::machine::{KvmVmState, MachineLifecycle};
-use migration::MigrationManager;
+use migration::{MigrationManager, MigrationRestoreOrder};
 use util::device_tree::{self, FdtBuilder};
 
 // See arch/arm64/include/uapi/asm/kvm.h file from the linux kernel.
@@ -393,10 +393,14 @@ impl GICDevice for GICv3 {
             MigrationManager::register_device_instance(
                 GICv3ItsState::descriptor(),
                 gicv3.its_dev.as_ref().unwrap().clone(),
-                true,
+                MigrationRestoreOrder::Gicv3Its,
             );
         }
-        MigrationManager::register_device_instance(GICv3State::descriptor(), gicv3.clone(), true);
+        MigrationManager::register_device_instance(
+            GICv3State::descriptor(),
+            gicv3.clone(),
+            MigrationRestoreOrder::Gicv3,
+        );
 
         Ok(gicv3)
     }
diff --git a/hypervisor/src/kvm/mod.rs b/hypervisor/src/kvm/mod.rs
index a318a65..80012bb 100644
--- a/hypervisor/src/kvm/mod.rs
+++ b/hypervisor/src/kvm/mod.rs
@@ -27,6 +27,8 @@ use once_cell::sync::Lazy;
 use vmm_sys_util::eventfd::EventFd;
 
 use crate::errors::{Result, ResultExt};
+#[cfg(target_arch = "x86_64")]
+use migration::{MigrationManager, MigrationRestoreOrder};
 
 // See: https://elixir.bootlin.com/linux/v4.19.123/source/include/uapi/asm-generic/kvm.h
 pub const KVM_SET_DEVICE_ATTR: u32 = 0x4018_aee1;
@@ -124,10 +126,10 @@ impl KVMFds {
         };
 
         #[cfg(target_arch = "x86_64")]
-        migration::MigrationManager::register_device_instance(
+        MigrationManager::register_device_instance(
             state::KvmDeviceState::descriptor(),
             Arc::new(state::KvmDevice {}),
-            false,
+            MigrationRestoreOrder::Default,
         );
 
         kvm_fds
diff --git a/machine/src/lib.rs b/machine/src/lib.rs
index 4421deb..7f88b22 100644
--- a/machine/src/lib.rs
+++ b/machine/src/lib.rs
@@ -132,7 +132,7 @@ use machine_manager::config::{
 };
 use machine_manager::event_loop::EventLoop;
 use machine_manager::machine::{KvmVmState, MachineInterface};
-use migration::MigrationManager;
+use migration::{MigrationManager, MigrationRestoreOrder};
 use util::loop_context::{EventNotifier, NotifierCallback, NotifierOperation};
 use util::seccomp::{BpfRule, SeccompOpt, SyscallFilter};
 use vfio::{VfioDevice, VfioPciDevice};
@@ -243,7 +243,11 @@ pub trait MachineOps {
             ));
             cpus.push(cpu.clone());
 
-            MigrationManager::register_device_instance(cpu::ArchCPU::descriptor(), cpu, false);
+            MigrationManager::register_device_instance(
+                cpu::ArchCPU::descriptor(),
+                cpu,
+                MigrationRestoreOrder::Default,
+            );
         }
 
         if let Some(boot_config) = boot_cfg {
@@ -486,7 +490,11 @@ pub trait MachineOps {
         let device_cfg = parse_blk(vm_config, cfg_args)?;
         let device = Arc::new(Mutex::new(Block::new(device_cfg.clone())));
         self.add_virtio_pci_device(&device_cfg.id, &bdf, device.clone(), multi_func)?;
-        MigrationManager::register_device_instance_mutex(BlockState::descriptor(), device);
+        MigrationManager::register_device_instance_mutex_with_id(
+            BlockState::descriptor(),
+            device,
+            &device_cfg.id,
+        );
         self.reset_bus(&device_cfg.id)?;
         Ok(())
     }
@@ -502,9 +510,10 @@ pub trait MachineOps {
             )))
         } else {
             let device = Arc::new(Mutex::new(virtio::Net::new(device_cfg.clone())));
-            MigrationManager::register_device_instance_mutex(
+            MigrationManager::register_device_instance_mutex_with_id(
                 VirtioNetState::descriptor(),
                 device.clone(),
+                &device_cfg.id,
             );
             device
         };
diff --git a/machine/src/standard_vm/mod.rs b/machine/src/standard_vm/mod.rs
index 1fca3bf..c96f89a 100644
--- a/machine/src/standard_vm/mod.rs
+++ b/machine/src/standard_vm/mod.rs
@@ -80,10 +80,11 @@ use machine_manager::config::{
 };
 use machine_manager::machine::{DeviceInterface, KvmVmState};
 use machine_manager::qmp::{qmp_schema, QmpChannel, Response};
+use migration::MigrationManager;
 use pci::hotplug::{handle_plug, handle_unplug_request};
 use pci::PciBus;
 use util::byte_code::ByteCode;
-use virtio::{qmp_balloon, qmp_query_balloon, Block, VhostKern, VirtioDevice};
+use virtio::{qmp_balloon, qmp_query_balloon, Block, BlockState, VhostKern, VirtioNetState};
 
 #[cfg(target_arch = "aarch64")]
 use aarch64::{LayoutEntryType, MEM_LAYOUT};
@@ -549,7 +550,7 @@ impl StdMachine {
 
         let blk = if let Some(conf) = self.get_vm_config().lock().unwrap().drives.get(drive) {
             let dev = BlkDevConfig {
-                id: conf.id.clone(),
+                id: args.id.clone(),
                 path_on_host: conf.path_on_host.clone(),
                 read_only: conf.read_only,
                 direct: conf.direct,
@@ -558,13 +559,22 @@ impl StdMachine {
                 iops: conf.iops,
             };
             dev.check()?;
-            Arc::new(Mutex::new(Block::new(dev)))
+            dev
         } else {
             bail!("Drive not found");
         };
 
-        self.add_virtio_pci_device(&args.id, pci_bdf, blk, multifunction)
-            .chain_err(|| "Failed to add virtio pci block device")
+        let blk_id = blk.id.clone();
+        let blk = Arc::new(Mutex::new(Block::new(blk)));
+        self.add_virtio_pci_device(&args.id, pci_bdf, blk.clone(), multifunction)
+            .chain_err(|| "Failed to add virtio pci block device")?;
+
+        MigrationManager::register_device_instance_mutex_with_id(
+            BlockState::descriptor(),
+            blk,
+            &blk_id,
+        );
+        Ok(())
     }
 
     fn plug_virtio_pci_net(
@@ -581,7 +591,7 @@ impl StdMachine {
 
         let dev = if let Some(conf) = self.get_vm_config().lock().unwrap().netdevs.get(netdev) {
             let dev = NetworkInterfaceConfig {
-                id: conf.id.clone(),
+                id: args.id.clone(),
                 host_dev_name: conf.ifname.clone(),
                 mac: args.mac.clone(),
                 tap_fd: conf.tap_fd,
@@ -595,14 +605,22 @@ impl StdMachine {
             bail!("Netdev not found");
         };
 
-        let net: Arc<Mutex<dyn VirtioDevice>> = if dev.vhost_type.is_some() {
-            Arc::new(Mutex::new(VhostKern::Net::new(&dev, self.get_sys_mem())))
+        if dev.vhost_type.is_some() {
+            let net = Arc::new(Mutex::new(VhostKern::Net::new(&dev, self.get_sys_mem())));
+            self.add_virtio_pci_device(&args.id, &pci_bdf, net, multifunction)
+                .chain_err(|| "Failed to add virtio net device")?;
         } else {
-            Arc::new(Mutex::new(virtio::Net::new(dev)))
-        };
-
-        self.add_virtio_pci_device(&args.id, &pci_bdf, net, multifunction)
-            .chain_err(|| "Failed to add virtio pci net device")
+            let net_id = dev.id.clone();
+            let net = Arc::new(Mutex::new(virtio::Net::new(dev)));
+            self.add_virtio_pci_device(&args.id, &pci_bdf, net.clone(), multifunction)
+                .chain_err(|| "Failed to add virtio net device")?;
+            MigrationManager::register_device_instance_mutex_with_id(
+                VirtioNetState::descriptor(),
+                net,
+                &net_id,
+            );
+        }
+        Ok(())
     }
 
     fn plug_vfio_pci_device(
diff --git a/migration/Cargo.toml b/migration/Cargo.toml
index fc877ad..6991804 100644
--- a/migration/Cargo.toml
+++ b/migration/Cargo.toml
@@ -11,6 +11,7 @@ kvm-ioctls = "0.6.0"
 serde = { version = ">=1.0.114", features = ["derive"] }
 serde_json = "1.0.55"
 once_cell = "1.9.0"
+log = "0.4.8"
 
 [dev-dependencies]
 migration_derive = { path = "../migration_derive" }
diff --git a/migration/src/device_state.rs b/migration/src/device_state.rs
index 75bf3b9..de9c16c 100644
--- a/migration/src/device_state.rs
+++ b/migration/src/device_state.rs
@@ -171,13 +171,6 @@ pub mod tests {
     use super::{DeviceStateDesc, FieldDesc, StateTransfer, VersionCheck};
     use util::byte_code::ByteCode;
 
-    struct MigrationManager {}
-    impl MigrationManager {
-        fn desc_db_len() -> u64 {
-            0
-        }
-    }
-
     #[derive(Default)]
     // A simple device version 1.
     pub struct DeviceV1 {
diff --git a/migration/src/lib.rs b/migration/src/lib.rs
index 9751fd3..9ccab5f 100644
--- a/migration/src/lib.rs
+++ b/migration/src/lib.rs
@@ -19,6 +19,8 @@ extern crate error_chain;
 #[cfg(test)]
 #[macro_use]
 extern crate migration_derive;
+#[macro_use]
+extern crate log;
 
 mod device_state;
 mod header;
@@ -27,7 +29,7 @@ mod snapshot;
 mod status;
 
 pub use device_state::{DeviceStateDesc, FieldDesc, StateTransfer};
-pub use manager::{MigrationHook, MigrationManager};
+pub use manager::{MigrationHook, MigrationManager, MigrationRestoreOrder};
 pub use status::MigrationStatus;
 
 pub mod errors {
diff --git a/migration/src/manager.rs b/migration/src/manager.rs
index f5d52b1..ef903f9 100644
--- a/migration/src/manager.rs
+++ b/migration/src/manager.rs
@@ -10,8 +10,11 @@
 // NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 // See the Mulan PSL v2 for more details.
 
-use std::collections::{BTreeMap, HashMap};
+use std::cmp;
+use std::collections::hash_map::DefaultHasher;
+use std::collections::HashMap;
 use std::fs::File;
+use std::hash::{Hash, Hasher};
 use std::io::{Read, Write};
 use std::sync::{Arc, Mutex, RwLock};
 
@@ -23,11 +26,27 @@ use util::byte_code::ByteCode;
 
 /// Glocal MigrationManager to manage all migration combined interface.
 pub(crate) static MIGRATION_MANAGER: Lazy<MigrationManager> = Lazy::new(|| MigrationManager {
-    entry: Arc::new(RwLock::new(BTreeMap::<u64, MigrationEntry>::new())),
+    entry: Arc::new(RwLock::new([
+        Vec::<(String, MigrationEntry)>::new(),
+        Vec::<(String, MigrationEntry)>::new(),
+        Vec::<(String, MigrationEntry)>::new(),
+    ])),
     desc_db: Arc::new(RwLock::new(HashMap::<String, DeviceStateDesc>::new())),
     status: Arc::new(RwLock::new(MigrationStatus::None)),
 });
 
+/// Used to map Device id from String to u64 only.
+/// Because instance_id in InstanceId can't be String for it has no Copy trait.
+///
+/// # Arguments
+///
+/// * `dev_id` - The device id.
+pub fn id_remap(dev_id: &str) -> u64 {
+    let mut hash = DefaultHasher::new();
+    dev_id.hash(&mut hash);
+    hash.finish()
+}
+
 /// A hook for `Device` to save device state to `Write` object and load device
 /// from `[u8]` slice.
 ///
@@ -44,7 +63,7 @@ pub trait MigrationHook: StateTransfer {
     /// * `id` - This unique id to represent a single device. It can be treated
     ///          as `object_id` in `InstanceId`.
     /// * `writer` - The `Write` trait object to store or receive data.
-    fn pre_save(&self, id: u64, writer: &mut dyn Write) -> Result<()> {
+    fn pre_save(&self, id: &str, writer: &mut dyn Write) -> Result<()> {
         let state_data = self
             .get_state_vec()
             .chain_err(|| "Failed to get device state")?;
@@ -52,7 +71,7 @@ pub trait MigrationHook: StateTransfer {
         let device_alias = self.get_device_alias();
         let instance_id = InstanceId {
             object_type: device_alias,
-            object_id: id,
+            object_id: id_remap(&id),
         };
 
         writer
@@ -131,11 +150,36 @@ pub enum MigrationEntry {
     Memory(Arc<dyn MigrationHook + Send + Sync>),
 }
 
+/// Ensure the recovery sequence of different devices based on priorities.
+/// At present, we need to ensure that the state recovery of the gic device
+/// must be after the cpu, so different priorities are defined.
+#[derive(Debug)]
+pub enum MigrationRestoreOrder {
+    Default = 0,
+    Gicv3 = 1,
+    Gicv3Its = 2,
+    Max = 3,
+}
+
+impl From<MigrationRestoreOrder> for u16 {
+    fn from(order: MigrationRestoreOrder) -> u16 {
+        match order {
+            MigrationRestoreOrder::Default => 0,
+            MigrationRestoreOrder::Gicv3 => 1,
+            MigrationRestoreOrder::Gicv3Its => 2,
+            _ => 3,
+        }
+    }
+}
+
+/// The entry list size is the same as the MigrationRestoreOrder number
+type MigrationEntryList = [Vec<(String, MigrationEntry)>; 3];
+
 /// This structure is to manage all resource during migration.
 /// It is also the only way to call on `MIGRATION_MANAGER`.
 pub struct MigrationManager {
     /// The map offers the device_id and combined migratable device entry.
-    pub(crate) entry: Arc<RwLock<BTreeMap<u64, MigrationEntry>>>,
+    pub(crate) entry: Arc<RwLock<MigrationEntryList>>,
     /// The map offers the device type and its device state describe structure.
     pub(crate) desc_db: Arc<RwLock<HashMap<String, DeviceStateDesc>>>,
     /// The status of migration work.
@@ -161,28 +205,23 @@ impl MigrationManager {
     ///
     /// * `device_desc` - The `DeviceStateDesc` of device instance.
     /// * `entry` - Device instance with migratable interface.
-    /// * `reverse` - Register device in order or in the reverse order.
+    /// * `restore_order` - device restore order.
     pub fn register_device_instance<T>(
         device_desc: DeviceStateDesc,
         device_entry: Arc<T>,
-        reverse: bool,
+        restore_order: MigrationRestoreOrder,
     ) where
         T: MigrationHook + Sync + Send + 'static,
     {
+        let name = device_desc.name.clone();
         Self::register_device_desc(device_desc);
 
         let entry = MigrationEntry::Safe(device_entry);
-        let nr_entry = if reverse {
-            !0 - Self::entry_db_len()
-        } else {
-            Self::entry_db_len()
-        };
-
-        MIGRATION_MANAGER
-            .entry
-            .write()
-            .unwrap()
-            .insert(nr_entry, entry);
+        info!(
+            "Register device instance: id {} order {:?}",
+            &name, &restore_order
+        );
+        MigrationManager::insert_entry(name, restore_order.into(), entry, true);
     }
 
     /// Register mutex device instance to entry hashmap with instance_id.
@@ -197,16 +236,34 @@ impl MigrationManager {
     ) where
         T: MigrationHook + Sync + Send + 'static,
     {
+        let name = device_desc.name.clone();
+        let order = MigrationRestoreOrder::Default.into();
         Self::register_device_desc(device_desc);
 
         let entry = MigrationEntry::Mutex(device_entry);
-        let nr_entry = Self::entry_db_len();
+        info!("Register device instance mutex: id {}", &name);
+        MigrationManager::insert_entry(name, order, entry, true);
+    }
+
+    pub fn register_device_instance_mutex_with_id<T>(
+        device_desc: DeviceStateDesc,
+        device_entry: Arc<Mutex<T>>,
+        id: &str,
+    ) where
+        T: MigrationHook + Sync + Send + 'static,
+    {
+        let name = device_desc.name.clone() + "/" + id;
+        let order = MigrationRestoreOrder::Default.into();
+        Self::register_device_desc(device_desc);
+        let entry = MigrationEntry::Mutex(device_entry);
+        info!("Register device instance with id: id {}", &name);
+        MigrationManager::insert_entry(name, order, entry, false);
+    }
 
-        MIGRATION_MANAGER
-            .entry
-            .write()
-            .unwrap()
-            .insert(nr_entry, entry);
+    pub fn unregister_device_instance_mutex_by_id(device_desc: DeviceStateDesc, id: &str) {
+        let name = device_desc.name + "/" + id;
+        info!("Unregister device instance: id {}", &name);
+        MigrationManager::remove_entry(&name);
     }
 
     /// Register memory instance.
@@ -219,23 +276,55 @@ impl MigrationManager {
         T: MigrationHook + Sync + Send + 'static,
     {
         let entry = MigrationEntry::Memory(entry);
-        let nr_entry = Self::entry_db_len();
-
-        MIGRATION_MANAGER
-            .entry
-            .write()
-            .unwrap()
-            .insert(nr_entry, entry);
+        info!("Register memory instance");
+        MigrationManager::insert_entry(String::from("MemoryState/Memory"), 0, entry, true);
     }
 
-    /// Get entry_db's length.
-    pub fn entry_db_len() -> u64 {
-        MIGRATION_MANAGER.entry.read().unwrap().len() as u64
+    /// Insert entry. If the name is duplicated, you should set gen_instance_id to true to
+    /// generated instance id to ensure that the id is unique.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - Entry name.
+    /// * `order` - Restore order.
+    /// * `entry` - Instance with migratable interface.
+    /// * `gen_instance_id` - If auto-generated instance id.
+    fn insert_entry(name: String, order: u16, entry: MigrationEntry, gen_instance_id: bool) {
+        let mut entrys = MIGRATION_MANAGER.entry.write().unwrap();
+        let mut index = 0;
+        if gen_instance_id {
+            for (key, _) in &entrys[order as usize] {
+                if let Some(pos) = key.rfind(':') {
+                    let (tmp_id, num_id) = key.split_at(pos);
+                    if tmp_id == name {
+                        let num = num_id.strip_prefix(':').unwrap();
+                        index = cmp::max(index, num.parse::<u16>().unwrap() + 1);
+                    }
+                }
+            }
+        }
+        // ID is format as "{name}:{instance_id}"
+        let id = format!("{}:{}", name, index);
+        debug!("Insert entry: id {}", &id);
+        entrys[order as usize].push((id, entry));
     }
 
-    /// Get desc_db's length.
-    pub fn desc_db_len() -> u64 {
-        MIGRATION_MANAGER.desc_db.read().unwrap().len() as u64
+    /// Remove entry by the unique name. Not support to remove the entry with instance id.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - Entry name.
+    fn remove_entry(name: &str) {
+        let eid = format!("{}:0", name);
+        let mut entrys = MIGRATION_MANAGER.entry.write().unwrap();
+        for (i, item) in entrys.iter().enumerate() {
+            let pos = item.iter().position(|(key, _)| key == &eid);
+            if let Some(index) = pos {
+                debug!("Remove entry: eid {}", &eid);
+                entrys[i].remove(index);
+                return;
+            }
+        }
     }
 
     /// Get `Device`'s alias from device type string.
@@ -244,12 +333,7 @@ impl MigrationManager {
     ///
     /// * `device_type` - The type string of device instance.
     pub fn get_desc_alias(device_type: &str) -> Option<u64> {
-        MIGRATION_MANAGER
-            .desc_db
-            .read()
-            .unwrap()
-            .get(device_type)
-            .map(|desc| desc.alias)
+        Some(id_remap(device_type))
     }
 
     /// Return `desc_db` value len(0 restored as `serde_json`)
@@ -340,23 +424,26 @@ mod tests {
         let device_v2 = Arc::new(DeviceV2::default());
         let device_v2_mutex = Arc::new(Mutex::new(DeviceV2::default()));
 
-        MigrationManager::register_device_instance(DeviceV1State::descriptor(), device_v1, false);
+        MigrationManager::register_device_instance(
+            DeviceV1State::descriptor(),
+            device_v1,
+            MigrationRestoreOrder::Default,
+        );
         MigrationManager::register_memory_instance(device_v2);
         MigrationManager::register_device_instance_mutex(
             DeviceV2State::descriptor(),
             device_v2_mutex,
         );
 
-        assert_eq!(MigrationManager::desc_db_len(), 2);
         assert!(MigrationManager::get_desc_alias("DeviceV1State").is_some());
         assert_eq!(
             MigrationManager::get_desc_alias("DeviceV1State").unwrap(),
-            0
+            id_remap("DeviceV1State")
         );
         assert!(MigrationManager::get_desc_alias("DeviceV2State").is_some());
         assert_eq!(
             MigrationManager::get_desc_alias("DeviceV2State").unwrap(),
-            0
+            id_remap("DeviceV2State")
         );
     }
 }
diff --git a/migration/src/snapshot.rs b/migration/src/snapshot.rs
index edea8ce..30ee13c 100644
--- a/migration/src/snapshot.rs
+++ b/migration/src/snapshot.rs
@@ -23,7 +23,7 @@ use util::unix::host_page_size;
 use crate::device_state::{DeviceStateDesc, VersionCheck};
 use crate::errors::{ErrorKind, Result, ResultExt};
 use crate::header::{FileFormat, MigrationHeader};
-use crate::manager::{InstanceId, MigrationEntry, MigrationManager, MIGRATION_MANAGER};
+use crate::manager::{id_remap, InstanceId, MigrationEntry, MigrationManager, MIGRATION_MANAGER};
 use crate::status::MigrationStatus;
 
 /// The length of `MigrationHeader` part occupies bytes in snapshot file.
@@ -187,10 +187,13 @@ impl MigrationManager {
     ///
     /// * `writer` - The `Write` trait object.
     fn save_memory(writer: &mut dyn Write) -> Result<()> {
-        for (id, entry) in MIGRATION_MANAGER.entry.read().unwrap().iter() {
-            if let MigrationEntry::Memory(i) = entry {
-                i.pre_save(*id, writer)
-                    .chain_err(|| "Failed to save vm memory")?;
+        let entry = MIGRATION_MANAGER.entry.read().unwrap();
+        for item in entry.iter() {
+            for (id, entry) in item.iter() {
+                if let MigrationEntry::Memory(i) = entry {
+                    i.pre_save(id, writer)
+                        .chain_err(|| "Failed to save vm memory")?;
+                }
             }
         }
 
@@ -205,10 +208,13 @@ impl MigrationManager {
     fn load_memory(file: &mut File) -> Result<()> {
         let mut state_bytes = [0_u8].repeat((host_page_size() as usize) * 2 - HEADER_LENGTH);
         file.read_exact(&mut state_bytes)?;
-        for (_, entry) in MIGRATION_MANAGER.entry.read().unwrap().iter() {
-            if let MigrationEntry::Memory(i) = entry {
-                i.pre_load(&state_bytes, Some(file))
-                    .chain_err(|| "Failed to load vm memory")?;
+        let entry = MIGRATION_MANAGER.entry.read().unwrap();
+        for item in entry.iter() {
+            for (_, entry) in item.iter() {
+                if let MigrationEntry::Memory(i) = entry {
+                    i.pre_load(&state_bytes, Some(file))
+                        .chain_err(|| "Failed to load vm memory")?;
+                }
             }
         }
 
@@ -221,11 +227,14 @@ impl MigrationManager {
     ///
     /// * `writer` - The `Write` trait object.
     fn save_device_state(writer: &mut dyn Write) -> Result<()> {
-        for (device_id, entry) in MIGRATION_MANAGER.entry.read().unwrap().iter() {
-            match entry {
-                MigrationEntry::Safe(i) => i.pre_save(*device_id, writer)?,
-                MigrationEntry::Mutex(i) => i.lock().unwrap().pre_save(*device_id, writer)?,
-                _ => {}
+        let entry = MIGRATION_MANAGER.entry.read().unwrap();
+        for item in entry.iter() {
+            for (id, entry) in item.iter() {
+                match entry {
+                    MigrationEntry::Safe(i) => i.pre_save(id, writer)?,
+                    MigrationEntry::Mutex(i) => i.lock().unwrap().pre_save(id, writer)?,
+                    _ => {}
+                }
             }
         }
 
@@ -275,10 +284,19 @@ impl MigrationManager {
                 }
             }
 
-            match device_entry.get(&instance_id.object_id).unwrap() {
-                MigrationEntry::Safe(i) => i.pre_load(&state_data, None)?,
-                MigrationEntry::Mutex(i) => i.lock().unwrap().pre_load_mut(&state_data, None)?,
-                _ => {}
+            for item in device_entry.iter() {
+                for (key, state) in item {
+                    if id_remap(key) == instance_id.object_id {
+                        info!("Load VM state: key {}", key);
+                        match state {
+                            MigrationEntry::Safe(i) => i.pre_load(&state_data, None)?,
+                            MigrationEntry::Mutex(i) => {
+                                i.lock().unwrap().pre_load_mut(&state_data, None)?
+                            }
+                            _ => {}
+                        }
+                    }
+                }
             }
         }
 
@@ -288,9 +306,12 @@ impl MigrationManager {
     /// Resume recovered device.
     /// This function will be called after restore device state.
     fn resume() -> Result<()> {
-        for (_, entry) in MIGRATION_MANAGER.entry.read().unwrap().iter() {
-            if let MigrationEntry::Mutex(i) = entry {
-                i.lock().unwrap().resume()?
+        let entry = MIGRATION_MANAGER.entry.read().unwrap();
+        for item in entry.iter() {
+            for (_, state) in item {
+                if let MigrationEntry::Mutex(i) = state {
+                    i.lock().unwrap().resume()?
+                }
             }
         }
         Ok(())
diff --git a/migration_derive/src/struct_parser.rs b/migration_derive/src/struct_parser.rs
index bc7d9d2..4e66d7e 100644
--- a/migration_derive/src/struct_parser.rs
+++ b/migration_derive/src/struct_parser.rs
@@ -26,10 +26,19 @@ pub fn parse_struct(
 
     let fields = parse_fields(&input.fields, ident);
 
+    use std::collections::hash_map::DefaultHasher;
+    use std::hash::{Hash, Hasher};
+
+    let id_remap = |s: &str| -> u64 {
+        let mut hash = DefaultHasher::new();
+        s.hash(&mut hash);
+        hash.finish()
+    };
+    let alias = id_remap(&name);
     quote! {
         #struct_ident {
             name: #name.to_string(),
-            alias: MigrationManager::desc_db_len(),
+            alias: #alias,
             size: std::mem::size_of::<#ident>() as u32,
             current_version: #current_version,
             compat_version: #compat_version,
diff --git a/pci/src/msix.rs b/pci/src/msix.rs
index 83b3d05..71d172e 100644
--- a/pci/src/msix.rs
+++ b/pci/src/msix.rs
@@ -405,6 +405,7 @@ pub fn init_msix(
     vector_nr: u32,
     config: &mut PciConfig,
     dev_id: Arc<AtomicU16>,
+    id: &str,
 ) -> Result<()> {
     if vector_nr > MSIX_TABLE_SIZE_MAX as u32 + 1 {
         bail!("Too many msix vectors.");
@@ -439,7 +440,7 @@ pub fn init_msix(
     config.msix = Some(msix.clone());
 
     #[cfg(not(test))]
-    MigrationManager::register_device_instance_mutex(MsixState::descriptor(), msix);
+    MigrationManager::register_device_instance_mutex_with_id(MsixState::descriptor(), msix, id);
 
     Ok(())
 }
@@ -469,11 +470,12 @@ mod tests {
             0,
             MSIX_TABLE_SIZE_MAX as u32 + 2,
             &mut pci_config,
-            Arc::new(AtomicU16::new(0))
+            Arc::new(AtomicU16::new(0)),
+            "msix"
         )
         .is_err());
 
-        init_msix(1, 2, &mut pci_config, Arc::new(AtomicU16::new(0))).unwrap();
+        init_msix(1, 2, &mut pci_config, Arc::new(AtomicU16::new(0)), "msix").unwrap();
         let msix_cap_start = 64_u8;
         assert_eq!(pci_config.last_cap_end, 64 + MSIX_CAP_SIZE as u16);
         // Capabilities pointer
@@ -538,7 +540,7 @@ mod tests {
     #[test]
     fn test_write_config() {
         let mut pci_config = PciConfig::new(PCI_CONFIG_SPACE_SIZE, 2);
-        init_msix(0, 2, &mut pci_config, Arc::new(AtomicU16::new(0))).unwrap();
+        init_msix(0, 2, &mut pci_config, Arc::new(AtomicU16::new(0)), "msix").unwrap();
         let msix = pci_config.msix.as_ref().unwrap();
         let mut locked_msix = msix.lock().unwrap();
         locked_msix.enabled = false;
diff --git a/pci/src/root_port.rs b/pci/src/root_port.rs
index b00c14a..286e92d 100644
--- a/pci/src/root_port.rs
+++ b/pci/src/root_port.rs
@@ -289,7 +289,7 @@ impl PciDevOps for RootPort {
             .add_pcie_cap(self.devfn, self.port_num, PcieDevType::RootPort as u8)?;
 
         self.dev_id.store(self.devfn as u16, Ordering::SeqCst);
-        init_msix(0, 1, &mut self.config, self.dev_id.clone())?;
+        init_msix(0, 1, &mut self.config, self.dev_id.clone(), &self.name)?;
 
         let parent_bus = self.parent_bus.upgrade().unwrap();
         let mut locked_parent_bus = parent_bus.lock().unwrap();
@@ -303,6 +303,7 @@ impl PciDevOps for RootPort {
             .add_subregion(self.sec_bus.lock().unwrap().mem_region.clone(), 0)
             .chain_err(|| "Failed to register subregion in memory space.")?;
 
+        let name = self.name.clone();
         let root_port = Arc::new(Mutex::new(self));
         #[allow(unused_mut)]
         let mut locked_root_port = root_port.lock().unwrap();
@@ -327,7 +328,11 @@ impl PciDevOps for RootPort {
         }
         // Need to drop locked_root_port in order to register root_port instance.
         drop(locked_root_port);
-        MigrationManager::register_device_instance_mutex(RootPortState::descriptor(), root_port);
+        MigrationManager::register_device_instance_mutex_with_id(
+            RootPortState::descriptor(),
+            root_port,
+            &name,
+        );
 
         Ok(())
     }
diff --git a/virtio/src/block.rs b/virtio/src/block.rs
index e0ced06..a2e35e8 100644
--- a/virtio/src/block.rs
+++ b/virtio/src/block.rs
@@ -967,6 +967,10 @@ impl VirtioDevice for Block {
     }
 
     fn unrealize(&mut self) -> Result<()> {
+        MigrationManager::unregister_device_instance_mutex_by_id(
+            BlockState::descriptor(),
+            &self.blk_cfg.id,
+        );
         Ok(())
     }
 
diff --git a/virtio/src/net.rs b/virtio/src/net.rs
index bbb1cc7..096121b 100644
--- a/virtio/src/net.rs
+++ b/virtio/src/net.rs
@@ -620,6 +620,10 @@ impl VirtioDevice for Net {
     }
 
     fn unrealize(&mut self) -> Result<()> {
+        MigrationManager::unregister_device_instance_mutex_by_id(
+            VirtioNetState::descriptor(),
+            &self.net_cfg.id,
+        );
         Ok(())
     }
 
diff --git a/virtio/src/virtio_pci.rs b/virtio/src/virtio_pci.rs
index bf1f58a..9eae777 100644
--- a/virtio/src/virtio_pci.rs
+++ b/virtio/src/virtio_pci.rs
@@ -24,7 +24,7 @@ use pci::config::{
     VENDOR_ID,
 };
 use pci::errors::{ErrorKind, Result as PciResult, ResultExt};
-use pci::msix::update_dev_id;
+use pci::msix::{update_dev_id, MsixState};
 use pci::{
     config::PciConfig, init_msix, init_multifunction, le_write_u16, ranges_overlap, PciBus,
     PciDevOps,
@@ -939,6 +939,7 @@ impl PciDevOps for VirtioPciDevice {
             nvectors as u32,
             &mut self.config,
             self.dev_id.clone(),
+            &self.name,
         )?;
 
         self.assign_interrupt_cb();
@@ -964,6 +965,7 @@ impl PciDevOps for VirtioPciDevice {
             .realize()
             .chain_err(|| "Failed to realize virtio device")?;
 
+        let name = self.name.clone();
         let devfn = self.devfn;
         let dev = Arc::new(Mutex::new(self));
         let pci_bus = dev.lock().unwrap().parent_bus.upgrade().unwrap();
@@ -978,7 +980,11 @@ impl PciDevOps for VirtioPciDevice {
                 pci_device.unwrap().lock().unwrap().name()
             );
         }
-        MigrationManager::register_device_instance_mutex(VirtioPciState::descriptor(), dev);
+        MigrationManager::register_device_instance_mutex_with_id(
+            VirtioPciState::descriptor(),
+            dev,
+            &name,
+        );
 
         Ok(())
     }
@@ -992,6 +998,15 @@ impl PciDevOps for VirtioPciDevice {
 
         let bus = self.parent_bus.upgrade().unwrap();
         self.config.unregister_bars(&bus)?;
+
+        MigrationManager::unregister_device_instance_mutex_by_id(
+            MsixState::descriptor(),
+            &self.name,
+        );
+        MigrationManager::unregister_device_instance_mutex_by_id(
+            VirtioPciState::descriptor(),
+            &self.name,
+        );
         Ok(())
     }
 
@@ -1491,6 +1506,7 @@ mod tests {
             virtio_pci.device.lock().unwrap().queue_num() as u32 + 1,
             &mut virtio_pci.config,
             virtio_pci.dev_id.clone(),
+            &virtio_pci.name,
         )
         .unwrap();
         // Prepare valid queue config
-- 
2.25.1

