device_manager: save resource allocator in snapshot

vm-allocator now allows us to (De)serialize IdAllocator and
AddressAllocator types. Add ResourceAllocator in DeviceManager snapshot
state and restore it when loading a snapshot. Like this we can avoid
doing the ExactMatch allocations during snapshot resumes for reserving
the exact same MMIO ranges.

Moreover, change DeviceManager and PciDevices to provide save/restore
functionality via the Persist trait. Like that we can avoid first
creating the objects and then restoring their state, overwriting their
fields.

Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

Cargo.lock

@@ -52,7 +52,7 @@ userfaultfd = "0.8.1"
 utils = { path = "../utils" }
 uuid = "1.16.0"
 vhost = { version = "0.14.0", features = ["vhost-user-frontend"] }
-vm-allocator = "0.1.2"
+vm-allocator = { version = "0.1.2", features = ["serde"] }
 vm-device = { path = "../vm-device" }
 vm-memory = { version = "0.16.2", features = [
   "backend-mmap",

src/vmm/Cargo.toml

@@ -41,6 +41,7 @@ use crate::logger::debug;
 use crate::persist::{MicrovmState, MicrovmStateError};
 use crate::resources::VmResources;
 use crate::seccomp::BpfThreadMap;
+use crate::snapshot::Persist;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::MachineConfigError;
 use crate::vstate::kvm::Kvm;
@@ -410,12 +411,25 @@ pub fn build_microvm_from_snapshot(
         .create_vcpus(vm_resources.machine_config.vcpu_count)
         .unwrap();
 
-    let mut device_manager = DeviceManager::new(event_manager, &vcpus_exit_evt, &vm).unwrap();
-
     vm.register_memory_regions(guest_memory)
         .map_err(VmmError::Vm)
         .map_err(StartMicrovmError::Internal)?;
 
+    // Restore devices states.
+    let device_ctor_args = DeviceRestoreArgs {
+        mem: vm.guest_memory(),
+        vm: &vm,
+        event_manager,
+        vm_resources,
+        instance_id: &instance_info.id,
+        restored_from_file: uffd.is_none(),
+        vcpus_exit_evt: &vcpus_exit_evt,
+    };
+
+    #[allow(unused_mut)]
+    let mut device_manager =
+        DeviceManager::restore(device_ctor_args, &microvm_state.device_states)?;
+
     #[cfg(target_arch = "x86_64")]
     {
         // Scale TSC to match, extract the TSC freq from the state if specified
@@ -430,16 +444,6 @@ pub fn build_microvm_from_snapshot(
         }
     }
 
-    // Restore allocator state
-    #[cfg(target_arch = "aarch64")]
-    if let Some(pvtime_ipa) = vcpus[0].kvm_vcpu.pvtime_ipa {
-        allocate_pvtime_region(
-            &mut device_manager,
-            vcpus.len(),
-            vm_allocator::AllocPolicy::ExactMatch(pvtime_ipa.0),
-        )?;
-    }
-
     // Restore vcpus kvm state.
     for (vcpu, state) in vcpus.iter_mut().zip(microvm_state.vcpu_states.iter()) {
         vcpu.kvm_vcpu
@@ -462,18 +466,6 @@ pub fn build_microvm_from_snapshot(
     // Restore the boot source config paths.
     vm_resources.boot_source.config = microvm_state.vm_info.boot_source;
 
-    // Restore devices states.
-    let device_ctor_args = DeviceRestoreArgs {
-        mem: vm.guest_memory(),
-        vm: &vm,
-        event_manager,
-        vm_resources,
-        instance_id: &instance_info.id,
-        restored_from_file: uffd.is_none(),
-    };
-
-    device_manager.restore(&microvm_state.device_states, device_ctor_args)?;
-
     let mut vmm = Vmm {
         events_observer: Some(std::io::stdin()),
         instance_info: instance_info.clone(),

src/vmm/src/builder.rs

@@ -15,7 +15,7 @@ use legacy::{LegacyDeviceError, PortIODeviceManager};
 use linux_loader::loader::Cmdline;
 use log::error;
 use mmio::{MMIODeviceManager, MmioError};
-use pci_mngr::{PciDevices, PciManagerError};
+use pci_mngr::{PciDevices, PciDevicesConstructorArgs, PciManagerError};
 use persist::{ACPIDeviceManagerConstructorArgs, MMIODevManagerConstructorArgs};
 use resources::ResourceAllocator;
 use serde::{Deserialize, Serialize};
@@ -127,30 +127,39 @@ impl DeviceManager {
         Ok(serial)
     }
 
+    #[cfg(target_arch = "x86_64")]
+    fn create_legacy_devices(
+        event_manager: &mut EventManager,
+        vcpus_exit_evt: &EventFd,
+        vm: &Vm,
+        resource_allocator: &ResourceAllocator,
+    ) -> Result<PortIODeviceManager, DeviceManagerCreateError> {
+        Self::set_stdout_nonblocking();
+
+        // Create serial device
+        let serial = Self::setup_serial_device(event_manager)?;
+        let reset_evt = vcpus_exit_evt
+            .try_clone()
+            .map_err(DeviceManagerCreateError::EventFd)?;
+        // Create keyboard emulator for reset event
+        let i8042 = Arc::new(Mutex::new(I8042Device::new(reset_evt)?));
+
+        // create pio dev manager with legacy devices
+        let mut legacy_devices = PortIODeviceManager::new(serial, i8042)?;
+        legacy_devices.register_devices(&resource_allocator.pio_bus, vm)?;
+        Ok(legacy_devices)
+    }
+
     #[cfg_attr(target_arch = "aarch64", allow(unused))]
     pub fn new(
         event_manager: &mut EventManager,
-        vcpu_exit_evt: &EventFd,
+        vcpus_exit_evt: &EventFd,
         vm: &Vm,
     ) -> Result<Self, DeviceManagerCreateError> {
         let resource_allocator = Arc::new(ResourceAllocator::new()?);
         #[cfg(target_arch = "x86_64")]
-        let legacy_devices = {
-            Self::set_stdout_nonblocking();
-
-            // Create serial device
-            let serial = Self::setup_serial_device(event_manager)?;
-            let reset_evt = vcpu_exit_evt
-                .try_clone()
-                .map_err(DeviceManagerCreateError::EventFd)?;
-            // Create keyboard emulator for reset event
-            let i8042 = Arc::new(Mutex::new(I8042Device::new(reset_evt)?));
-
-            // create pio dev manager with legacy devices
-            let mut legacy_devices = PortIODeviceManager::new(serial, i8042)?;
-            legacy_devices.register_devices(&resource_allocator.pio_bus, vm)?;
-            legacy_devices
-        };
+        let legacy_devices =
+            Self::create_legacy_devices(event_manager, vcpus_exit_evt, vm, &resource_allocator)?;
 
         Ok(DeviceManager {
             resource_allocator,
@@ -270,6 +279,8 @@ impl DeviceManager {
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]
 /// State of devices in the system
 pub struct DevicesState {
+    /// Resource allocator state
+    pub resource_allocator_state: resources::ResourceAllocatorState,
     /// MMIO devices state
     pub mmio_state: persist::DeviceStates,
     /// ACPI devices state
@@ -292,12 +303,15 @@ pub enum DevicePersistError {
     SerialRestore(#[from] EmulateSerialInitError),
     /// Error inserting device in bus: {0}
     Bus(#[from] vm_device::BusError),
+    /// Error creating DeviceManager: {0}
+    DeviceManager(#[from] DeviceManagerCreateError),
 }
 
 pub struct DeviceRestoreArgs<'a> {
     pub mem: &'a GuestMemoryMmap,
     pub vm: &'a Vm,
     pub event_manager: &'a mut EventManager,
+    pub vcpus_exit_evt: &'a EventFd,
     pub vm_resources: &'a mut VmResources,
     pub instance_id: &'a str,
     pub restored_from_file: bool,
@@ -315,15 +329,82 @@ impl std::fmt::Debug for DeviceRestoreArgs<'_> {
     }
 }
 
-impl DeviceManager {
-    pub fn save(&self) -> DevicesState {
+impl<'a> Persist<'a> for DeviceManager {
+    type State = DevicesState;
+    type ConstructorArgs = DeviceRestoreArgs<'a>;
+    type Error = DevicePersistError;
+
+    fn save(&self) -> Self::State {
         DevicesState {
+            resource_allocator_state: self.resource_allocator.save(),
             mmio_state: self.mmio_devices.save(),
             acpi_state: self.acpi_devices.save(),
             pci_state: self.pci_devices.save(),
         }
     }
 
+    fn restore(
+        constructor_args: Self::ConstructorArgs,
+        state: &Self::State,
+    ) -> std::result::Result<Self, Self::Error> {
+        // Safe to unwrap here. ResourceAllocator restoring cannot fail.
+        let resource_allocator =
+            Arc::new(ResourceAllocator::restore((), &state.resource_allocator_state).unwrap());
+
+        // Restore MMIO devices
+        let mmio_ctor_args = MMIODevManagerConstructorArgs {
+            mem: constructor_args.mem,
+            vm: constructor_args.vm,
+            event_manager: constructor_args.event_manager,
+            resource_allocator: &resource_allocator,
+            vm_resources: constructor_args.vm_resources,
+            instance_id: constructor_args.instance_id,
+            restored_from_file: constructor_args.restored_from_file,
+        };
+        let mmio_devices = MMIODeviceManager::restore(mmio_ctor_args, &state.mmio_state)?;
+
+        // Restore ACPI devices
+        let acpi_ctor_args = ACPIDeviceManagerConstructorArgs {
+            mem: constructor_args.mem,
+            resource_allocator: &resource_allocator,
+            vm: constructor_args.vm,
+        };
+        let mut acpi_devices = ACPIDeviceManager::restore(acpi_ctor_args, &state.acpi_state)?;
+        acpi_devices.notify_vmgenid()?;
+
+        // Restore PCI devices
+        let pci_ctor_args = PciDevicesConstructorArgs {
+            resource_allocator: &resource_allocator,
+        };
+        let pci_devices = PciDevices::restore(pci_ctor_args, &state.pci_state)?;
+
+        // Setup legacy devices in case of x86
+        #[cfg(target_arch = "x86_64")]
+        let legacy_devices = Self::create_legacy_devices(
+            constructor_args.event_manager,
+            constructor_args.vcpus_exit_evt,
+            constructor_args.vm,
+            &resource_allocator,
+        )?;
+
+        let device_manager = DeviceManager {
+            resource_allocator,
+            mmio_devices,
+            #[cfg(target_arch = "x86_64")]
+            legacy_devices,
+            acpi_devices,
+            pci_devices,
+        };
+
+        // Restore serial.
+        // We need to do that after we restore mmio devices, otherwise it won't succeed in Aarch64
+        device_manager.emulate_serial_init()?;
+
+        Ok(device_manager)
+    }
+}
+
+impl DeviceManager {
     /// Sets RDA bit in serial console
     pub fn emulate_serial_init(&self) -> Result<(), EmulateSerialInitError> {
         // When restoring from a previously saved state, there is no serial
@@ -361,43 +442,6 @@ impl DeviceManager {
             Ok(())
         }
     }
-
-    pub fn restore(
-        &mut self,
-        state: &DevicesState,
-        restore_args: DeviceRestoreArgs,
-    ) -> Result<(), DevicePersistError> {
-        // Restore MMIO devices
-        let mmio_ctor_args = MMIODevManagerConstructorArgs {
-            mem: restore_args.mem,
-            vm: restore_args.vm,
-            event_manager: restore_args.event_manager,
-            resource_allocator: &self.resource_allocator,
-            vm_resources: restore_args.vm_resources,
-            instance_id: restore_args.instance_id,
-            restored_from_file: restore_args.restored_from_file,
-        };
-        self.mmio_devices = MMIODeviceManager::restore(mmio_ctor_args, &state.mmio_state)?;
-
-        // Restore serial.
-        // We need to do that after we restore mmio devices, otherwise it won't succeed in Aarch64
-        self.emulate_serial_init()?;
-
-        // Restore ACPI devices
-        let acpi_ctor_args = ACPIDeviceManagerConstructorArgs {
-            mem: restore_args.mem,
-            resource_allocator: &self.resource_allocator,
-            vm: restore_args.vm,
-        };
-        self.acpi_devices = ACPIDeviceManager::restore(acpi_ctor_args, &state.acpi_state)?;
-        self.acpi_devices.notify_vmgenid()?;
-
-        // Restore PCI devices
-        self.pci_devices
-            .restore(&state.pci_state, &self.resource_allocator)?;
-
-        Ok(())
-    }
 }
 
 #[cfg(test)]

src/vmm/src/device_manager/mod.rs

@@ -16,6 +16,7 @@ use crate::device_manager::resources::ResourceAllocator;
 use crate::devices::pci::PciSegment;
 use crate::devices::virtio::device::VirtioDevice;
 use crate::devices::virtio::transport::pci::device::{VirtioPciDevice, VirtioPciDeviceError};
+use crate::snapshot::Persist;
 use crate::vstate::vm::InterruptError;
 
 #[derive(Debug, Default)]
@@ -63,24 +64,6 @@ impl PciDevices {
         Ok(())
     }
 
-    pub fn save(&self) -> PciDevicesState {
-        PciDevicesState {
-            pci_enabled: self.pci_segment.is_some(),
-        }
-    }
-
-    pub fn restore(
-        &mut self,
-        state: &PciDevicesState,
-        resource_allocator: &Arc<ResourceAllocator>,
-    ) -> Result<(), PciManagerError> {
-        if state.pci_enabled {
-            self.attach_pci_segment(resource_allocator)?;
-        }
-
-        Ok(())
-    }
-
     pub(crate) fn attach_pci_virtio_device<
         T: 'static + VirtioDevice + MutEventSubscriber + Debug,
     >(
@@ -186,3 +169,33 @@ impl PciDevices {
 pub struct PciDevicesState {
     pci_enabled: bool,
 }
+
+#[derive(Debug)]
+pub struct PciDevicesConstructorArgs<'a> {
+    pub resource_allocator: &'a Arc<ResourceAllocator>,
+}
+
+impl<'a> Persist<'a> for PciDevices {
+    type State = PciDevicesState;
+    type ConstructorArgs = PciDevicesConstructorArgs<'a>;
+    type Error = PciManagerError;
+
+    fn save(&self) -> Self::State {
+        PciDevicesState {
+            pci_enabled: self.pci_segment.is_some(),
+        }
+    }
+
+    fn restore(
+        constructor_args: Self::ConstructorArgs,
+        state: &Self::State,
+    ) -> std::result::Result<Self, Self::Error> {
+        let mut pci_devices = PciDevices::new();
+
+        if state.pci_enabled {
+            pci_devices.attach_pci_segment(constructor_args.resource_allocator)?;
+        }
+
+        Ok(pci_devices)
+    }
+}