Add snapshot save and restore functionality to CpuContainer

Implement Persist and the required functions for saving and restoring
snapshots

Note: Currently when saving a snapshot, custom CPU templates are not
saved. This is because, up until now, it has only been needed at
initialization. With the implementation of hotplugging however, CPU
templates and the vCPU config are now required for the lifetime of the
VMM. As a result of this, hotplugging will throw an error after
restoring from snapshot.

Signed-off-by: James Curtis <jxcurtis@amazon.co.uk>

Author: JamesC1305

src/vmm/src/builder.rs

@@ -29,6 +29,8 @@ use vm_superio::Serial;
 
 #[cfg(target_arch = "x86_64")]
 use crate::acpi;
+#[cfg(target_arch = "x86_64")]
+use crate::arch::DeviceType;
 use crate::arch::InitrdConfig;
 #[cfg(target_arch = "aarch64")]
 use crate::construct_kvm_mpidrs;
@@ -58,6 +60,8 @@ use crate::devices::virtio::mmio::MmioTransport;
 use crate::devices::virtio::net::Net;
 use crate::devices::virtio::rng::Entropy;
 use crate::devices::virtio::vsock::{Vsock, VsockUnixBackend};
+#[cfg(target_arch = "x86_64")]
+use crate::devices::BusDevice;
 use crate::logger::{debug, error};
 use crate::persist::{MicrovmState, MicrovmStateError};
 use crate::resources::VmResources;
@@ -451,6 +455,8 @@ pub enum BuildMicrovmFromSnapshotError {
     StartVcpus(#[from] crate::StartVcpusError),
     /// Failed to restore vCPUs: {0}
     RestoreVcpus(#[from] VcpuError),
+    /// Failed to restore CPUID: {0}
+    RestoreCpuId(#[from] GuestConfigError),
     /// Failed to apply VMM secccomp filter as none found.
     MissingVmmSeccompFilters,
     /// Failed to apply VMM secccomp filter: {0}
@@ -522,6 +528,31 @@ pub fn build_microvm_from_snapshot(
     #[cfg(target_arch = "x86_64")]
     vmm.vm.restore_state(&microvm_state.vm_state)?;
 
+    #[cfg(target_arch = "x86_64")]
+    {
+        // FIXME: Custom templates are not stored when snapshots are taken, as they were previously
+        // only needed at boottime. With hotplugging, the template could be necessary at any time,
+        // so snapshot must be modified in order to save custom templates
+
+        // let cpuid = Cpuid::try_from(vmm.vm.supported_cpuid().clone())
+        //     .map_err(GuestConfigError::CpuidFromKvmCpuid)?;
+
+        // let_cpu_template = &microvm_state.vm_info.cpu_template;
+        //
+        // let msrs = vcpus[0]
+        //     .kvm_vcpu
+        //     .get_msrs(&[])
+        //     .map_err(GuestConfigError::VcpuIoctl)?;
+        //
+        // let cpu_config = CpuConfiguration { cpuid, msrs };
+        //
+        // let vcpu_config = VcpuConfig {
+        //     vcpu_count: vcpus.len().try_into().unwrap(),
+        //     smt: microvm_state.vm_info.smt,
+        //     cpu_config,
+        // };
+    }
+
     // Restore the boot source config paths.
     vm_resources.set_boot_source_config(microvm_state.vm_info.boot_source);
 
@@ -550,6 +581,22 @@ pub fn build_microvm_from_snapshot(
         vmm.acpi_device_manager =
             ACPIDeviceManager::restore(acpi_ctor_args, &microvm_state.acpi_dev_state)?;
 
+        #[cfg(target_arch = "x86_64")]
+        if let Some(BusDevice::CpuContainer(container)) = vmm.get_bus_device(
+            DeviceType::CpuContainer,
+            &DeviceType::CpuContainer.to_string(),
+        ) {
+            vmm.acpi_device_manager
+                .attach_cpu_container(container.clone(), vmm.vm.fd())
+                .map_err(|err| {
+                    BuildMicrovmFromSnapshotError::ACPIDeviManager(
+                        ACPIDeviceManagerRestoreError::CpuContainer(
+                            CpuContainerError::RegisterIrqFd(err),
+                        ),
+                    )
+                })?;
+        }
+
         // Inject the notification to VMGenID that we have resumed from a snapshot.
         // This needs to happen before we resume vCPUs, so that we minimize the time between vCPUs
         // resuming and notification being handled by the driver.

src/vmm/src/device_manager/persist.rs

@@ -17,6 +17,10 @@ use super::mmio::*;
 use super::resources::ResourceAllocator;
 #[cfg(target_arch = "aarch64")]
 use crate::arch::DeviceType;
+#[cfg(target_arch = "x86_64")]
+use crate::devices::acpi::cpu_container::{
+    CpuContainer, CpuContainerConstructorArgs, CpuContainerError, CpuContainerState,
+};
 use crate::devices::acpi::vmgenid::{VMGenIDState, VMGenIdConstructorArgs, VmGenId, VmGenIdError};
 use crate::devices::virtio::balloon::persist::{BalloonConstructorArgs, BalloonState};
 use crate::devices::virtio::balloon::{Balloon, BalloonError};
@@ -55,6 +59,9 @@ pub enum DevicePersistError {
     Balloon(#[from] BalloonError),
     /// Block: {0}
     Block(#[from] BlockError),
+    /// CpuContainer: {0}
+    #[cfg(target_arch = "x86_64")]
+    CpuContainer(#[from] CpuContainerError),
     /// Device manager: {0}
     DeviceManager(#[from] super::mmio::MmioError),
     /// Mmio transport
@@ -151,6 +158,13 @@ pub struct ConnectedLegacyState {
     pub device_info: MMIODeviceInfo,
 }
 
+#[cfg(target_arch = "x86_64")]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ConnectedCpuContainerState {
+    pub device_state: CpuContainerState,
+    pub device_info: MMIODeviceInfo,
+}
+
 /// Holds the MMDS data store version.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub enum MmdsVersionState {
@@ -194,6 +208,9 @@ pub struct DeviceStates {
     pub mmds_version: Option<MmdsVersionState>,
     /// Entropy device state.
     pub entropy_device: Option<ConnectedEntropyState>,
+    /// Cpu Container device state.
+    #[cfg(target_arch = "x86_64")]
+    pub cpu_container: Option<ConnectedCpuContainerState>,
 }
 
 /// A type used to extract the concrete `Arc<Mutex<T>>` for each of the device
@@ -245,6 +262,9 @@ pub enum ACPIDeviceManagerRestoreError {
     Interrupt(#[from] kvm_ioctls::Error),
     /// Could not create VMGenID device: {0}
     VMGenID(#[from] VmGenIdError),
+    /// Could not create CpuContainer device: {0}
+    #[cfg(target_arch = "x86_64")]
+    CpuContainer(#[from] CpuContainerError),
 }
 
 impl<'a> Persist<'a> for ACPIDeviceManager {
@@ -283,11 +303,21 @@ impl<'a> Persist<'a> for MMIODeviceManager {
     type Error = DevicePersistError;
 
     fn save(&self) -> Self::State {
+        use crate::devices::bus::BusDevice::*;
         let mut states = DeviceStates::default();
+        #[allow(unused_variables)]
         let _: Result<(), ()> = self.for_each_device(|devtype, devid, device_info, bus_dev| {
-            if *devtype == crate::arch::DeviceType::BootTimer {
-                // No need to save BootTimer state.
-                return Ok(());
+            match bus_dev {
+                BootTimer(_) => return Ok(()),
+                #[cfg(target_arch = "x86_64")]
+                CpuContainer(x) => {
+                    states.cpu_container = Some(ConnectedCpuContainerState {
+                        device_state: x.lock().expect("Poisoned lock").save(),
+                        device_info: device_info.clone(),
+                    });
+                    return Ok(());
+                }
+                _ => (),
             }
 
             #[cfg(target_arch = "aarch64")]
@@ -644,6 +674,17 @@ impl<'a> Persist<'a> for MMIODeviceManager {
             )?;
         }
 
+        #[cfg(target_arch = "x86_64")]
+        if let Some(container) = &state.cpu_container {
+            let ctor_args = CpuContainerConstructorArgs { vm };
+            let device = Arc::new(Mutex::new(CpuContainer::restore(
+                ctor_args,
+                &container.device_state,
+            )?));
+
+            dev_manager.register_mmio_cpu_container(vm, device, &container.device_info)?;
+        }
+
         Ok(dev_manager)
     }
 }

src/vmm/src/devices/acpi/cpu_container.rs

@@ -5,6 +5,7 @@
 
 use acpi_tables::madt::LocalAPIC;
 use acpi_tables::{aml, Aml};
+use kvm_ioctls::VmFd;
 use log::{debug, error};
 use serde::{Deserialize, Serialize};
 use utils::eventfd::EventFd;
@@ -15,6 +16,7 @@ use zerocopy::AsBytes;
 use crate::device_manager::mmio::MMIO_LEN;
 use crate::device_manager::resources::ResourceAllocator;
 use crate::devices::legacy::EventFdTrigger;
+use crate::snapshot::Persist;
 use crate::vmm_config::machine_config::MAX_SUPPORTED_VCPUS;
 
 #[derive(Debug)]
@@ -150,6 +152,33 @@ impl CpuContainer {
     }
 }
 
+impl<'a> Persist<'a> for CpuContainer {
+    type State = CpuContainerState;
+    type ConstructorArgs = CpuContainerConstructorArgs<'a>;
+    type Error = CpuContainerError;
+
+    fn save(&self) -> Self::State {
+        Self::State {
+            mmio_address: self.mmio_address.0,
+            gsi: self.gsi,
+            cpu_devices: self.cpu_devices.clone(),
+            selected_cpu: self.selected_cpu,
+        }
+    }
+
+    fn restore(
+        _constructor_args: Self::ConstructorArgs,
+        state: &Self::State,
+    ) -> std::result::Result<Self, Self::Error> {
+        Self::from_parts(
+            GuestAddress(state.mmio_address),
+            state.gsi,
+            state.selected_cpu,
+            state.cpu_devices.clone(),
+        )
+    }
+}
+
 impl Aml for CpuContainer {
     fn append_aml_bytes(&self, v: &mut Vec<u8>) {
         // CPU hotplug controller
@@ -220,6 +249,19 @@ impl Aml for CpuContainer {
     }
 }
 
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
+pub struct CpuContainerState {
+    pub mmio_address: u64,
+    pub gsi: u32,
+    pub cpu_devices: Vec<CpuDevice>,
+    pub selected_cpu: u8,
+}
+
+#[derive(Debug, Clone)]
+pub struct CpuContainerConstructorArgs<'a> {
+    pub vm: &'a VmFd,
+}
+
 #[derive(Default, Debug, Clone, Serialize, Deserialize)]
 pub struct CpuDevice {
     /// The ID of this CPU, matches APICID

src/vmm/src/vmm_config/hotplug.rs

@@ -38,6 +38,8 @@ pub enum HotplugVcpuError {
     VcpuStart(StartVcpusError),
     /// No seccomp filter for thread category: {0}
     MissingSeccompFilters(String),
+    /// Cannot hotplug vCPUs after restoring from snapshot
+    RestoredFromSnapshot,
 }
 
 /// Config for hotplugging vCPUS

tests/integration_tests/functional/test_api.py

@@ -1146,6 +1146,7 @@ def test_get_full_config_after_restoring_snapshot(microvm_factory, uvm_nano):
 
     # Start the microvm.
     uvm_nano.start()
+    uvm_nano.wait_for_up()
 
     # Add a tx rate limiter to the net device.
     tx_rl = {