pci: Implement Migratable for VfioCommon

Introduces the common code to handle one aspect of the migration
support. Particularly, the ability to store VMM internal states related
to such device. The internal state of the device will happen later in a
dedicated patchset that will implement the VFIO migration API.

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

pci/src/vfio.rs

@@ -4,10 +4,11 @@
 //
 
 use crate::{
-    msi_num_enabled_vectors, BarReprogrammingParams, MsiConfig, MsixCap, MsixConfig,
+    msi_num_enabled_vectors, BarReprogrammingParams, MsiCap, MsiConfig, MsixCap, MsixConfig,
     PciBarConfiguration, PciBarRegionType, PciBdf, PciCapabilityId, PciClassCode, PciConfiguration,
     PciDevice, PciDeviceError, PciHeaderType, PciSubclass, MSIX_TABLE_ENTRY_SIZE,
 };
+use anyhow::anyhow;
 use byteorder::{ByteOrder, LittleEndian};
 use hypervisor::HypervisorVmError;
 use std::any::Any;
@@ -17,6 +18,8 @@ use std::os::unix::io::AsRawFd;
 use std::ptr::null_mut;
 use std::sync::{Arc, Barrier, Mutex};
 use thiserror::Error;
+use versionize::{VersionMap, Versionize, VersionizeResult};
+use versionize_derive::Versionize;
 use vfio_bindings::bindings::vfio::*;
 use vfio_ioctls::{
     VfioContainer, VfioDevice, VfioIrq, VfioRegionInfoCap, VfioRegionSparseMmapArea,
@@ -27,6 +30,7 @@ use vm_device::interrupt::{
 };
 use vm_device::{BusDevice, Resource};
 use vm_memory::{Address, GuestAddress, GuestUsize};
+use vm_migration::{MigratableError, Pausable, Snapshot, Snapshottable, VersionMapped};
 use vmm_sys_util::eventfd::EventFd;
 
 #[derive(Debug, Error)]
@@ -71,11 +75,22 @@ enum InterruptUpdateAction {
     DisableMsix,
 }
 
+#[derive(Versionize)]
+struct IntxState {
+    enabled: bool,
+}
+
 pub(crate) struct VfioIntx {
     interrupt_source_group: Arc<dyn InterruptSourceGroup>,
     enabled: bool,
 }
 
+#[derive(Versionize)]
+struct MsiState {
+    cap: MsiCap,
+    cap_offset: u32,
+}
+
 pub(crate) struct VfioMsi {
     pub(crate) cfg: MsiConfig,
     cap_offset: u32,
@@ -102,6 +117,13 @@ impl VfioMsi {
     }
 }
 
+#[derive(Versionize)]
+struct MsixState {
+    cap: MsixCap,
+    cap_offset: u32,
+    bdf: u32,
+}
+
 pub(crate) struct VfioMsix {
     pub(crate) bar: MsixConfig,
     cap: MsixCap,
@@ -356,6 +378,15 @@ impl Vfio for VfioDeviceWrapper {
     }
 }
 
+#[derive(Versionize)]
+struct VfioCommonState {
+    intx_state: Option<IntxState>,
+    msi_state: Option<MsiState>,
+    msix_state: Option<MsixState>,
+}
+
+impl VersionMapped for VfioCommonState {}
+
 pub(crate) struct VfioCommon {
     pub(crate) configuration: PciConfiguration,
     pub(crate) mmio_regions: Vec<MmioRegion>,
@@ -970,6 +1001,127 @@ impl VfioCommon {
         // The config register read comes from the VFIO device itself.
         self.vfio_wrapper.read_config_dword((reg_idx * 4) as u32) & mask
     }
+
+    fn state(&self) -> VfioCommonState {
+        let intx_state = self.interrupt.intx.as_ref().map(|intx| IntxState {
+            enabled: intx.enabled,
+        });
+
+        let msi_state = self.interrupt.msi.as_ref().map(|msi| MsiState {
+            cap: msi.cfg.cap,
+            cap_offset: msi.cap_offset,
+        });
+
+        let msix_state = self.interrupt.msix.as_ref().map(|msix| MsixState {
+            cap: msix.cap,
+            cap_offset: msix.cap_offset,
+            bdf: msix.bar.devid,
+        });
+
+        VfioCommonState {
+            intx_state,
+            msi_state,
+            msix_state,
+        }
+    }
+
+    fn set_state(&mut self, state: &VfioCommonState) -> Result<(), VfioPciError> {
+        if let (Some(intx), Some(interrupt_source_group)) =
+            (&state.intx_state, self.legacy_interrupt_group.clone())
+        {
+            self.interrupt.intx = Some(VfioIntx {
+                interrupt_source_group,
+                enabled: false,
+            });
+
+            if intx.enabled {
+                self.enable_intx()?;
+            }
+        }
+
+        if let Some(msi) = &state.msi_state {
+            self.initialize_msi(msi.cap.msg_ctl, msi.cap_offset);
+        }
+
+        if let Some(msix) = &state.msix_state {
+            self.initialize_msix(msix.cap, msix.cap_offset, msix.bdf.into());
+        }
+
+        Ok(())
+    }
+}
+
+impl Pausable for VfioCommon {}
+
+impl Snapshottable for VfioCommon {
+    fn id(&self) -> String {
+        String::from("vfio_common")
+    }
+
+    fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
+        let mut vfio_common_snapshot =
+            Snapshot::new_from_versioned_state(&self.id(), &self.state())?;
+
+        // Snapshot PciConfiguration
+        vfio_common_snapshot.add_snapshot(self.configuration.snapshot()?);
+
+        // Snapshot MSI
+        if let Some(msi) = &mut self.interrupt.msi {
+            vfio_common_snapshot.add_snapshot(msi.cfg.snapshot()?);
+        }
+
+        // Snapshot MSI-X
+        if let Some(msix) = &mut self.interrupt.msix {
+            vfio_common_snapshot.add_snapshot(msix.bar.snapshot()?);
+        }
+
+        Ok(vfio_common_snapshot)
+    }
+
+    fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
+        if let Some(vfio_common_section) = snapshot
+            .snapshot_data
+            .get(&format!("{}-section", self.id()))
+        {
+            // It has to be invoked first as we want Interrupt to be initialized
+            // correctly before we try to restore MSI and MSI-X configurations.
+            self.set_state(&vfio_common_section.to_versioned_state()?)
+                .map_err(|e| {
+                    MigratableError::Restore(anyhow!("Could not restore VFIO_COMMON state {:?}", e))
+                })?;
+
+            // Restore PciConfiguration
+            if let Some(pci_config_snapshot) = snapshot.snapshots.get(&self.configuration.id()) {
+                self.configuration.restore(*pci_config_snapshot.clone())?;
+            }
+
+            // Restore MSI
+            if let Some(msi) = &mut self.interrupt.msi {
+                if let Some(msi_snapshot) = snapshot.snapshots.get(&msi.cfg.id()) {
+                    msi.cfg.restore(*msi_snapshot.clone())?;
+                }
+                if msi.cfg.enabled() {
+                    self.enable_msi().unwrap();
+                }
+            }
+
+            // Restore MSI-X
+            if let Some(msix) = &mut self.interrupt.msix {
+                if let Some(msix_snapshot) = snapshot.snapshots.get(&msix.bar.id()) {
+                    msix.bar.restore(*msix_snapshot.clone())?;
+                }
+                if msix.bar.enabled() {
+                    self.enable_msix().unwrap();
+                }
+            }
+
+            return Ok(());
+        }
+
+        Err(MigratableError::Restore(anyhow!(
+            "Could not find VFIO_COMMON snapshot section"
+        )))
+    }
 }
 
 /// VfioPciDevice represents a VFIO PCI device.