WIP: Add partial PVTime snapshot persistence

- Changes to PVTime design (use base addr)
- Secondary constructor for PVTime using only base addr
- Implement Persist interface for PVTime, save base addr
- Add pvtime to struct MicrovmState and struct Vmm
- Restore PVTime device in build_microvm_from_snapshot

Signed-off-by: Dakshin Devanand <[email protected]>

Author: DakshinD

Diff for: src/vmm/src/arch/aarch64/pvtime/mod.rs

@@ -4,20 +4,24 @@ use std::collections::HashMap;
 
 use displaydoc::Display;
 use kvm_bindings::{KVM_ARM_VCPU_PVTIME_CTRL, KVM_ARM_VCPU_PVTIME_IPA};
-use log::{debug, info, warn};
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use vm_memory::GuestAddress;
 
 use crate::device_manager::resources::ResourceAllocator;
+use crate::snapshot::Persist;
 
 /// 64 bytes due to alignment requirement in 3.1 of https://www.kernel.org/doc/html/v5.8/virt/kvm/devices/vcpu.html#attribute-kvm-arm-vcpu-pvtime-ipa
 pub const STEALTIME_STRUCT_MEM_SIZE: u64 = 64;
 
 /// Represent PVTime device for ARM
+/// TODO: Decide whether we want to keep the hashmap OR the base IPA
 #[derive(Debug)]
 pub struct PVTime {
     /// Maps vCPU index to IPA location of stolen_time struct as defined in DEN0057A
     steal_time_regions: HashMap<u8, u64>,
+    /// The base IPA of the shared memory region
+    base_ipa: u64,
 }
 
 /// Errors associated with PVTime operations
@@ -32,31 +36,42 @@ pub enum PVTimeError {
 }
 
 impl PVTime {
-    /// Creates a new PVTime device by allocating system memory for all vCPUs
-    pub fn new(
-        resource_allocator: &mut ResourceAllocator,
-        vcpu_count: u8,
-    ) -> Result<Self, PVTimeError> {
-
-        // This returns the IPA of the start of our shared memory region for all vCPUs.
-        let base_addr: GuestAddress = GuestAddress(resource_allocator
-            .allocate_system_memory(
-                STEALTIME_STRUCT_MEM_SIZE * vcpu_count as u64,
-                64, 
-                vm_allocator::AllocPolicy::LastMatch,
-            )
-            .map_err(|e| PVTimeError::AllocationFailed(e.to_string()))?);
-
+    /// Create a new PVTime device given a base addr
+    /// - Assumes total shared memory region from base addr is already allocated
+    fn from_base(base_addr: GuestAddress, vcpu_count: u8) -> Self {
+        let base_ipa: u64 = base_addr.0;
 
         // Now we need to store the base IPA for each vCPU's steal_time struct.
         let mut steal_time_regions = HashMap::new();
         for i in 0..vcpu_count {
-            let ipa = base_addr.0 + (i as u64 * STEALTIME_STRUCT_MEM_SIZE);
+            let ipa = base_ipa + (i as u64 * STEALTIME_STRUCT_MEM_SIZE);
             steal_time_regions.insert(i, ipa);
         }
 
         // Return the PVTime device with the steal_time region IPAs mapped to vCPU indices.
-        Ok(PVTime { steal_time_regions })
+        PVTime {
+            steal_time_regions,
+            base_ipa,
+        }
+    }
+
+    /// Creates a new PVTime device by allocating new system memory for all vCPUs
+    pub fn new(
+        resource_allocator: &mut ResourceAllocator,
+        vcpu_count: u8,
+    ) -> Result<Self, PVTimeError> {
+        // This returns the IPA of the start of our shared memory region for all vCPUs.
+        let base_ipa: GuestAddress = GuestAddress(
+            resource_allocator
+                .allocate_system_memory(
+                    STEALTIME_STRUCT_MEM_SIZE * vcpu_count as u64,
+                    64,
+                    vm_allocator::AllocPolicy::LastMatch,
+                )
+                .map_err(|e| PVTimeError::AllocationFailed(e.to_string()))?,
+        );
+
+        Ok(Self::from_base(base_ipa, vcpu_count))
     }
 
     /// Register a vCPU with its pre-allocated steal time region
@@ -65,7 +80,6 @@ impl PVTime {
         vcpu_index: u8,
         vcpu_fd: &kvm_ioctls::VcpuFd,
     ) -> Result<(), PVTimeError> {
-
         // Get IPA of the steal_time region for this vCPU
         let ipa = self
             .steal_time_regions
@@ -80,15 +94,55 @@ impl PVTime {
             flags: 0,
         };
 
-        vcpu_fd.set_device_attr(&vcpu_device_attr).map_err(|err| {
-            PVTimeError::DeviceAttribute(err, true, KVM_ARM_VCPU_PVTIME_CTRL)
-        })?;
+        vcpu_fd
+            .set_device_attr(&vcpu_device_attr)
+            .map_err(|err| PVTimeError::DeviceAttribute(err, true, KVM_ARM_VCPU_PVTIME_CTRL))?;
 
         Ok(())
     }
 }
 
-// TODO/Q: Would we be correct in implementing Persist for PVTime? Some sort of persistence is
-// needed for snapshot capability. We would only need to store base_addr IPA of shared memory region
-// assuming the # of vCPUs is constant across snapshots. Also assuming we are correct that
-// kvm_set_device_attr takes an IPA and we get an IPA back from resource_allocator.
+/// Logic to save/restore the state of a PVTime device
+#[derive(Default, Debug, Clone, Serialize, Deserialize)]
+pub struct PVTimeState {
+    /// base IPA of the total shared memory region
+    pub base_ipa: u64,
+}
+
+#[derive(Debug)]
+pub struct PVTimeConstructorArgs<'a> {
+    pub resource_allocator: &'a mut ResourceAllocator,
+    pub vcpu_count: u8,
+}
+
+impl<'a> Persist<'a> for PVTime {
+    type State = PVTimeState;
+    type ConstructorArgs = PVTimeConstructorArgs<'a>;
+    type Error = PVTimeError;
+
+    /// Save base IPA of PVTime device for persistence
+    fn save(&self) -> Self::State {
+        PVTimeState {
+            base_ipa: self.base_ipa,
+        }
+    }
+
+    /// Restore state of PVTime device from given base IPA
+    fn restore(
+        constructor_args: Self::ConstructorArgs,
+        state: &Self::State,
+    ) -> std::result::Result<Self, PVTimeError> {
+        constructor_args
+            .resource_allocator
+            .allocate_system_memory(
+                STEALTIME_STRUCT_MEM_SIZE * constructor_args.vcpu_count as u64,
+                64,
+                vm_allocator::AllocPolicy::ExactMatch(state.base_ipa),
+            )
+            .map_err(|e| PVTimeError::AllocationFailed(e.to_string()))?;
+        Ok(Self::from_base(
+            GuestAddress(state.base_ipa),
+            constructor_args.vcpu_count,
+        ))
+    }
+}

Diff for: src/vmm/src/builder.rs

@@ -11,15 +11,15 @@ use std::sync::{Arc, Mutex};
 use event_manager::{MutEventSubscriber, SubscriberOps};
 use libc::EFD_NONBLOCK;
 use linux_loader::cmdline::Cmdline as LoaderKernelCmdline;
-use log::{info, warn};
+use log::warn;
 use userfaultfd::Uffd;
 use utils::time::TimestampUs;
 #[cfg(target_arch = "aarch64")]
 use vm_superio::Rtc;
 use vm_superio::Serial;
 use vmm_sys_util::eventfd::EventFd;
 
-use crate::arch::aarch64::pvtime::PVTimeError;
+use crate::arch::aarch64::pvtime::{PVTime, PVTimeConstructorArgs, PVTimeError};
 use crate::arch::{ConfigurationError, configure_system_for_boot, load_kernel};
 #[cfg(target_arch = "aarch64")]
 use crate::construct_kvm_mpidrs;
@@ -195,6 +195,7 @@ fn create_vmm_and_vcpus(
         #[cfg(target_arch = "x86_64")]
         pio_device_manager,
         acpi_device_manager,
+        pv_time: None,
     };
 
     Ok((vmm, vcpus))
@@ -420,6 +421,8 @@ pub enum BuildMicrovmFromSnapshotError {
     ACPIDeviManager(#[from] ACPIDeviceManagerRestoreError),
     /// VMGenID update failed: {0}
     VMGenIDUpdate(std::io::Error),
+    /// Failed to restore PVTime device: {0}
+    RestorePVTime(#[from] PVTimeError),
 }
 
 /// Builds and starts a microVM based on the provided MicrovmState.
@@ -470,6 +473,21 @@ pub fn build_microvm_from_snapshot(
             .map_err(BuildMicrovmFromSnapshotError::RestoreVcpus)?;
     }
 
+    // Restore the PVTime device
+    let pvtime_state = microvm_state.pvtime_state;
+    if let Some(pvtime_state) = pvtime_state {
+        let pvtime_ctor_args = PVTimeConstructorArgs {
+            resource_allocator: &mut vmm.resource_allocator,
+            vcpu_count: vcpus.len() as u8,
+        };
+        vmm.pv_time = Some(
+            PVTime::restore(pvtime_ctor_args, &pvtime_state)
+                .map_err(BuildMicrovmFromSnapshotError::RestorePVTime)?,
+        );
+        setup_pv_time_with(vmm.pv_time.as_ref().unwrap(), &mut vcpus) // We can force unwrap here b/c we just set it
+            .map_err(BuildMicrovmFromSnapshotError::RestorePVTime)?;
+    }
+
     #[cfg(target_arch = "aarch64")]
     {
         let mpidrs = construct_kvm_mpidrs(&microvm_state.vcpu_states);
@@ -569,7 +587,7 @@ pub fn setup_serial_device(
 #[cfg(target_arch = "aarch64")]
 fn setup_pv_time(vmm: &mut Vmm, vcpus: &mut Vec<Vcpu>) -> Result<(), StartMicrovmError> {
     use crate::arch::aarch64::pvtime::PVTime;
-    
+
     // Check if pvtime is enabled
     let pvtime_device_attr = kvm_bindings::kvm_device_attr {
         group: kvm_bindings::KVM_ARM_VCPU_PVTIME_CTRL,
@@ -590,15 +608,21 @@ fn setup_pv_time(vmm: &mut Vmm, vcpus: &mut Vec<Vcpu>) -> Result<(), StartMicrov
     let pv_time = PVTime::new(&mut vmm.resource_allocator, vcpus.len() as u8)
         .map_err(StartMicrovmError::CreatePVTime)?;
 
+    // Register all vcpus with pvtime device
+    setup_pv_time_with(&pv_time, vcpus).map_err(StartMicrovmError::CreatePVTime)?;
+
+    // TODO: Might be a better design to return the device instead of assigning it to vmm here.
+    vmm.pv_time = Some(pv_time);
+    Ok(())
+}
+
+/// Helper method to register all vcpus with pvtime device
+#[cfg(target_arch = "aarch64")]
+fn setup_pv_time_with(pv_time: &PVTime, vcpus: &mut Vec<Vcpu>) -> Result<(), PVTimeError> {
     // Register the vcpu with the pvtime device to map its steal time region
     for i in 0..vcpus.len() {
-        pv_time
-            .register_vcpu(i as u8, &vcpus[i].kvm_vcpu.fd)
-            .map_err(StartMicrovmError::CreatePVTime)?; // TODO: Change this to its own error
+        pv_time.register_vcpu(i as u8, &vcpus[i].kvm_vcpu.fd)?; // TODO: Change this to its own error
     }
-
-    // TODO: Store pv_time somewhere (in Vmm ?) for snapshotting later instead of dropping it
-
     Ok(())
 }
 

Diff for: src/vmm/src/lib.rs

@@ -157,7 +157,7 @@ use crate::vstate::memory::{
 use crate::vstate::vcpu::VcpuState;
 pub use crate::vstate::vcpu::{Vcpu, VcpuConfig, VcpuEvent, VcpuHandle, VcpuResponse};
 pub use crate::vstate::vm::Vm;
-
+use crate::arch::aarch64::pvtime::PVTime;
 /// Shorthand type for the EventManager flavour used by Firecracker.
 pub type EventManager = BaseEventManager<Arc<Mutex<dyn MutEventSubscriber>>>;
 
@@ -324,6 +324,9 @@ pub struct Vmm {
     #[cfg(target_arch = "x86_64")]
     pio_device_manager: PortIODeviceManager,
     acpi_device_manager: ACPIDeviceManager,
+    #[cfg(target_arch = "aarch64")]
+    pv_time: Option<PVTime>,
+
 }
 
 impl Vmm {
@@ -523,6 +526,7 @@ impl Vmm {
 
         let memory_state = self.guest_memory.describe();
         let acpi_dev_state = self.acpi_device_manager.save();
+        let pvtime_state = self.pv_time.as_ref().map(|pvtime| pvtime.save());
 
         Ok(MicrovmState {
             vm_info: vm_info.clone(),
@@ -532,6 +536,7 @@ impl Vmm {
             vcpu_states,
             device_states,
             acpi_dev_state,
+            pvtime_state,
         })
     }
 

Diff for: src/vmm/src/persist.rs

@@ -16,6 +16,7 @@ use serde::{Deserialize, Serialize};
 use userfaultfd::{FeatureFlags, Uffd, UffdBuilder};
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
 
+use crate::arch::aarch64::pvtime::PVTimeState;
 #[cfg(target_arch = "aarch64")]
 use crate::arch::aarch64::vcpu::get_manufacturer_id_from_host;
 use crate::builder::{self, BuildMicrovmFromSnapshotError};
@@ -88,6 +89,9 @@ pub struct MicrovmState {
     pub device_states: DeviceStates,
     /// ACPI devices state.
     pub acpi_dev_state: ACPIDeviceManagerState,
+    /// PVTime device state (optional for platforms that support it).
+    #[cfg(target_arch = "aarch64")]
+    pub pvtime_state: Option<PVTimeState>,
 }
 
 /// This describes the mapping between Firecracker base virtual address and