From f48b05eee6f4b946146aa07790e9e1ed3aa64d7a Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Tue, 7 Jun 2022 16:32:57 +0200
Subject: [PATCH] pci: vfio_user: Implement Migratable for VfioUserPciDevice

Based on the VfioCommon implementation, the VfioUserPciDevice now
implements the Migratable trait.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 pci/src/vfio_user.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/pci/src/vfio_user.rs b/pci/src/vfio_user.rs
index 669223da9..7a9de41b8 100644
--- a/pci/src/vfio_user.rs
+++ b/pci/src/vfio_user.rs
@@ -26,6 +26,7 @@ use vm_memory::bitmap::AtomicBitmap;
 use vm_memory::{
     Address, GuestAddress, GuestAddressSpace, GuestMemory, GuestMemoryRegion, GuestRegionMmap,
 };
+use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
 use vmm_sys_util::eventfd::EventFd;
 
 pub struct VfioUserPciDevice {
@@ -551,6 +552,34 @@ impl Drop for VfioUserPciDevice {
     }
 }
 
+impl Pausable for VfioUserPciDevice {}
+
+impl Snapshottable for VfioUserPciDevice {
+    fn id(&self) -> String {
+        self.id.clone()
+    }
+
+    fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
+        let mut vfio_pci_dev_snapshot = Snapshot::new(&self.id);
+
+        // Snapshot VfioCommon
+        vfio_pci_dev_snapshot.add_snapshot(self.common.snapshot()?);
+
+        Ok(vfio_pci_dev_snapshot)
+    }
+
+    fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
+        // Restore VfioCommon
+        if let Some(vfio_common_snapshot) = snapshot.snapshots.get(&self.common.id()) {
+            self.common.restore(*vfio_common_snapshot.clone())?;
+        }
+
+        Ok(())
+    }
+}
+impl Transportable for VfioUserPciDevice {}
+impl Migratable for VfioUserPciDevice {}
+
 pub struct VfioUserDmaMapping<M: GuestAddressSpace> {
     client: Arc<Mutex<Client>>,
     memory: Arc<M>,