Refactor PVTime design

- Moved most PVTime logic into pvtime.rs
- Fixed includes for only aarch64
- PVTime only uses base_addr now
- Made clippy happy

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

Author: DakshinD

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

@@ -1,56 +1,55 @@
 // Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-use std::collections::HashMap;
-
 use displaydoc::Display;
 use kvm_bindings::{KVM_ARM_VCPU_PVTIME_CTRL, KVM_ARM_VCPU_PVTIME_IPA};
+use kvm_ioctls::VcpuFd;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use vm_memory::GuestAddress;
 
+use crate::Vcpu;
 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>,
+    /// Number of vCPUs
+    vcpu_count: u8,
     /// The base IPA of the shared memory region
-    base_ipa: u64,
+    base_ipa: GuestAddress,
 }
 
 /// Errors associated with PVTime operations
 #[derive(Debug, Error, Display, PartialEq, Eq)]
 pub enum PVTimeError {
     /// Failed to allocate memory region: {0}
-    AllocationFailed(String),
+    AllocationFailed(vm_allocator::Error),
     /// Invalid VCPU ID: {0}
     InvalidVcpuIndex(u8),
-    /// Error while setting or getting device attributes for vCPU: {0}, {1}, {2}
-    DeviceAttribute(kvm_ioctls::Error, bool, u32),
+    /// Error while setting or getting device attributes for vCPU: {0}
+    DeviceAttribute(kvm_ioctls::Error),
 }
 
 impl PVTime {
-    /// 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_ipa + (i as u64 * STEALTIME_STRUCT_MEM_SIZE);
-            steal_time_regions.insert(i, ipa);
+    /// Helper function to get the IPA of the steal_time region for a given vCPU
+    fn get_steal_time_region_addr(&self, vcpu_index: u8) -> Result<GuestAddress, PVTimeError> {
+        if vcpu_index >= self.vcpu_count {
+            return Err(PVTimeError::InvalidVcpuIndex(vcpu_index));
         }
+        Ok(GuestAddress(
+            self.base_ipa.0 + (vcpu_index as u64 * STEALTIME_STRUCT_MEM_SIZE),
+        ))
+    }
 
-        // Return the PVTime device with the steal_time region IPAs mapped to vCPU indices.
+    /// Create a new PVTime device given a base addr
+    /// - Assumes total shared memory region from base addr is already allocated
+    fn from_base(base_ipa: GuestAddress, vcpu_count: u8) -> Self {
         PVTime {
-            steal_time_regions,
+            vcpu_count,
             base_ipa,
         }
     }
@@ -65,41 +64,61 @@ impl PVTime {
             resource_allocator
                 .allocate_system_memory(
                     STEALTIME_STRUCT_MEM_SIZE * vcpu_count as u64,
-                    64,
+                    STEALTIME_STRUCT_MEM_SIZE,
                     vm_allocator::AllocPolicy::LastMatch,
                 )
-                .map_err(|e| PVTimeError::AllocationFailed(e.to_string()))?,
+                .map_err(PVTimeError::AllocationFailed)?,
         );
-
         Ok(Self::from_base(base_ipa, vcpu_count))
     }
 
+    /// Check if PVTime is supported on vcpu
+    pub fn is_supported(vcpu_fd: &VcpuFd) -> bool {
+        // Check if pvtime is enabled
+        let pvtime_device_attr = kvm_bindings::kvm_device_attr {
+            group: kvm_bindings::KVM_ARM_VCPU_PVTIME_CTRL,
+            attr: kvm_bindings::KVM_ARM_VCPU_PVTIME_IPA as u64,
+            addr: 0,
+            flags: 0,
+        };
+
+        // Use kvm_has_device_attr to check if PVTime is supported
+        vcpu_fd.has_device_attr(&pvtime_device_attr).is_ok()
+    }
+
     /// Register a vCPU with its pre-allocated steal time region
-    pub fn register_vcpu(
+    fn register_vcpu(
         &self,
         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
-            .get(&vcpu_index)
-            .ok_or(PVTimeError::InvalidVcpuIndex(vcpu_index))?;
+        let ipa = self.get_steal_time_region_addr(vcpu_index)?;
 
         // Use KVM syscall (kvm_set_device_attr) to register the vCPU with the steal_time region
         let vcpu_device_attr = kvm_bindings::kvm_device_attr {
             group: KVM_ARM_VCPU_PVTIME_CTRL,
             attr: KVM_ARM_VCPU_PVTIME_IPA as u64,
-            addr: ipa as *const u64 as u64, // userspace address of attr data
+            addr: &ipa.0 as *const u64 as u64, // userspace address of attr data
             flags: 0,
         };
 
         vcpu_fd
             .set_device_attr(&vcpu_device_attr)
-            .map_err(|err| PVTimeError::DeviceAttribute(err, true, KVM_ARM_VCPU_PVTIME_CTRL))?;
+            .map_err(PVTimeError::DeviceAttribute)?;
 
         Ok(())
     }
+
+    /// Register all vCPUs with their pre-allocated steal time regions
+    pub fn register_all_vcpus(&self, vcpus: &mut [Vcpu]) -> Result<(), PVTimeError> {
+        // Register the vcpu with the pvtime device to map its steal time region
+        for (i, vcpu) in vcpus.iter().enumerate() {
+            #[allow(clippy::cast_possible_truncation)] // We know vcpu_count is u8 according to VcpuConfig
+            self.register_vcpu(i as u8, &vcpu.kvm_vcpu.fd)?;
+        }
+        Ok(())
+    }
 }
 
 /// Logic to save/restore the state of a PVTime device
@@ -109,9 +128,12 @@ pub struct PVTimeState {
     pub base_ipa: u64,
 }
 
+/// Arguments to restore a PVTime device from PVTimeState
 #[derive(Debug)]
 pub struct PVTimeConstructorArgs<'a> {
+    /// For steal_time shared memory region
     pub resource_allocator: &'a mut ResourceAllocator,
+    /// Number of vCPUs (should be consistent with pre-snapshot state)
     pub vcpu_count: u8,
 }
 
@@ -123,7 +145,7 @@ impl<'a> Persist<'a> for PVTime {
     /// Save base IPA of PVTime device for persistence
     fn save(&self) -> Self::State {
         PVTimeState {
-            base_ipa: self.base_ipa,
+            base_ipa: self.base_ipa.0,
         }
     }
 
@@ -136,10 +158,10 @@ impl<'a> Persist<'a> for PVTime {
             .resource_allocator
             .allocate_system_memory(
                 STEALTIME_STRUCT_MEM_SIZE * constructor_args.vcpu_count as u64,
-                64,
+                STEALTIME_STRUCT_MEM_SIZE,
                 vm_allocator::AllocPolicy::ExactMatch(state.base_ipa),
             )
-            .map_err(|e| PVTimeError::AllocationFailed(e.to_string()))?;
+            .map_err(PVTimeError::AllocationFailed)?;
         Ok(Self::from_base(
             GuestAddress(state.base_ipa),
             constructor_args.vcpu_count,

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

@@ -19,6 +19,7 @@ use vm_superio::Rtc;
 use vm_superio::Serial;
 use vmm_sys_util::eventfd::EventFd;
 
+#[cfg(target_arch = "aarch64")]
 use crate::arch::aarch64::pvtime::{PVTime, PVTimeConstructorArgs, PVTimeError};
 use crate::arch::{ConfigurationError, configure_system_for_boot, load_kernel};
 #[cfg(target_arch = "aarch64")]
@@ -70,9 +71,6 @@ pub enum StartMicrovmError {
     AttachBlockDevice(io::Error),
     /// Unable to attach the VMGenID device: {0}
     AttachVmgenidDevice(kvm_ioctls::Error),
-    /// PVTime not supported: {0}
-    #[cfg(target_arch = "aarch64")]
-    PVTimeNotSupported(kvm_ioctls::Error),
     /// System configuration error: {0}
     ConfigureSystem(#[from] ConfigurationError),
     /// Failed to create guest config: {0}
@@ -195,6 +193,7 @@ fn create_vmm_and_vcpus(
         #[cfg(target_arch = "x86_64")]
         pio_device_manager,
         acpi_device_manager,
+        #[cfg(target_arch = "aarch64")]
         pv_time: None,
     };
 
@@ -298,13 +297,13 @@ pub fn build_microvm_for_boot(
 
     // Attempt to setup PVTime, continue if not supported
     #[cfg(target_arch = "aarch64")]
-    if let Err(e) = setup_pv_time(&mut vmm, vcpus.as_mut()) {
-        match e {
-            StartMicrovmError::PVTimeNotSupported(e) => {
-                warn!("PVTime not supported: {}", e);
-            }
-            other => return Err(other),
-        }
+    {
+        vmm.pv_time = if PVTime::is_supported(&vcpus[0].kvm_vcpu.fd) {
+            Some(setup_pv_time(&mut vmm, vcpus.as_mut())?)
+        } else {
+            warn!("PVTime is not supported by KVM. Steal time will not be reported to the guest.");
+            None
+        };
     }
 
     configure_system_for_boot(
@@ -474,18 +473,25 @@ pub fn build_microvm_from_snapshot(
     }
 
     // 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 pvtime_state = microvm_state.pvtime_state;
+        if let Some(pvtime_state) = pvtime_state {
+            #[allow(clippy::cast_possible_truncation)] // We know vcpu_count is u8 according to VcpuConfig
+            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)?,
+            );
+            vmm.pv_time
+                .as_ref()
+                .unwrap()
+                .register_all_vcpus(&mut vcpus) // We can safely unwrap here
+                .map_err(StartMicrovmError::CreatePVTime)?;
+        }
     }
 
     #[cfg(target_arch = "aarch64")]
@@ -585,45 +591,20 @@ pub fn setup_serial_device(
 
 /// Sets up the pvtime device.
 #[cfg(target_arch = "aarch64")]
-fn setup_pv_time(vmm: &mut Vmm, vcpus: &mut Vec<Vcpu>) -> Result<(), StartMicrovmError> {
+fn setup_pv_time(vmm: &mut Vmm, vcpus: &mut [Vcpu]) -> Result<PVTime, 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,
-        attr: kvm_bindings::KVM_ARM_VCPU_PVTIME_IPA as u64,
-        addr: 0,
-        flags: 0,
-    };
-
-    // Use kvm_has_device_attr to check if PVTime is supported
-    // TODO: Flesh out feature detection. Either throw error and handle, or
-    // even return silently and just log "no support", or something else.
-    let vcpu_fd = &vcpus[0].kvm_vcpu.fd;
-    vcpu_fd
-        .has_device_attr(&pvtime_device_attr)
-        .map_err(StartMicrovmError::PVTimeNotSupported)?;
-
     // Create the pvtime device
+    #[allow(clippy::cast_possible_truncation)] // We know vcpu_count is u8 according to VcpuConfig
     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(())
-}
+    pv_time
+        .register_all_vcpus(vcpus)
+        .map_err(StartMicrovmError::CreatePVTime)?;
 
-/// 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)?; // TODO: Change this to its own error
-    }
-    Ok(())
+    Ok(pv_time)
 }
 
 #[cfg(target_arch = "aarch64")]
@@ -929,6 +910,8 @@ pub(crate) mod tests {
             #[cfg(target_arch = "x86_64")]
             pio_device_manager,
             acpi_device_manager,
+            #[cfg(target_arch = "aarch64")]
+            pv_time: None,
         }
     }
 

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

@@ -134,6 +134,8 @@ use vstate::kvm::Kvm;
 use vstate::vcpu::{self, StartThreadedError, VcpuSendEventError};
 
 use crate::arch::DeviceType;
+#[cfg(target_arch = "aarch64")]
+use crate::arch::aarch64::pvtime::PVTime;
 use crate::cpu_config::templates::CpuConfiguration;
 #[cfg(target_arch = "x86_64")]
 use crate::device_manager::legacy::PortIODeviceManager;
@@ -157,7 +159,6 @@ 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>>>;
 
@@ -326,7 +327,6 @@ pub struct Vmm {
     acpi_device_manager: ACPIDeviceManager,
     #[cfg(target_arch = "aarch64")]
     pv_time: Option<PVTime>,
-
 }
 
 impl Vmm {
@@ -536,6 +536,7 @@ impl Vmm {
             vcpu_states,
             device_states,
             acpi_dev_state,
+            #[cfg(target_arch = "aarch64")]
             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;
 
+#[cfg(target_arch = "aarch64")]
 use crate::arch::aarch64::pvtime::PVTimeState;
 #[cfg(target_arch = "aarch64")]
 use crate::arch::aarch64::vcpu::get_manufacturer_id_from_host;
@@ -769,6 +770,8 @@ mod tests {
             #[cfg(target_arch = "x86_64")]
             vm_state: vmm.vm.save_state().unwrap(),
             acpi_dev_state: vmm.acpi_device_manager.save(),
+            #[cfg(target_arch = "aarch64")] 
+            pvtime_state: None,
         };
 
         let mut buf = vec![0; 10000];