test(pvtime): Add integration/unit tests and fix x86

- Added integration tests to validate PVTime steal time
  increases and persists after snapshot
- Added unit tests for PVTime internals
- Fixed build errors on x86_64
- Updated PVTime::vcpu_count to u64 for clarity and correctness
- Bumped snapshot version and updated changelog

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

Author: DakshinD

CHANGELOG.md

@@ -10,6 +10,9 @@ and this project adheres to
 
 ### Added
 
+- [#5139](https://github.com/firecracker-microvm/firecracker/pull/5139): Added
+  support for [PVTime](https://docs.kernel.org/virt/kvm/arm/pvtime.html). This is 
+  used to support steal time on ARM machines.
 - [#5048](https://github.com/firecracker-microvm/firecracker/pull/5048): Added
   support for [PVH boot mode](docs/pvh.md). This is used when an x86 kernel
   provides the appropriate ELF Note to indicate that PVH boot mode is supported.

src/vmm/src/arch/aarch64/pvtime.rs

@@ -18,7 +18,7 @@ pub const STEALTIME_STRUCT_MEM_SIZE: u64 = 64;
 #[derive(Debug)]
 pub struct PVTime {
     /// Number of vCPUs
-    vcpu_count: u8,
+    vcpu_count: u64,
     /// The base IPA of the shared memory region
     base_ipa: GuestAddress,
 }
@@ -29,25 +29,25 @@ pub enum PVTimeError {
     /// Failed to allocate memory region: {0}
     AllocationFailed(vm_allocator::Error),
     /// Invalid VCPU ID: {0}
-    InvalidVcpuIndex(u8),
+    InvalidVcpuIndex(u64),
     /// Error while setting or getting device attributes for vCPU: {0}
     DeviceAttribute(kvm_ioctls::Error),
 }
 
 impl PVTime {
     /// 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> {
+    fn get_steal_time_region_addr(&self, vcpu_index: u64) -> 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),
+            self.base_ipa.0 + (vcpu_index * STEALTIME_STRUCT_MEM_SIZE),
         ))
     }
 
     /// 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 {
+    fn from_base(base_ipa: GuestAddress, vcpu_count: u64) -> Self {
         PVTime {
             vcpu_count,
             base_ipa,
@@ -57,13 +57,13 @@ impl PVTime {
     /// Creates a new PVTime device by allocating new system memory for all vCPUs
     pub fn new(
         resource_allocator: &mut ResourceAllocator,
-        vcpu_count: u8,
+        vcpu_count: u64,
     ) -> 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,
+                    STEALTIME_STRUCT_MEM_SIZE * vcpu_count,
                     STEALTIME_STRUCT_MEM_SIZE,
                     vm_allocator::AllocPolicy::LastMatch,
                 )
@@ -89,7 +89,7 @@ impl PVTime {
     /// Register a vCPU with its pre-allocated steal time region
     fn register_vcpu(
         &self,
-        vcpu_index: u8,
+        vcpu_index: u64,
         vcpu_fd: &kvm_ioctls::VcpuFd,
     ) -> Result<(), PVTimeError> {
         // Get IPA of the steal_time region for this vCPU
@@ -114,9 +114,7 @@ impl PVTime {
     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)?;
+            self.register_vcpu(i as u64, &vcpu.kvm_vcpu.fd)?;
         }
         Ok(())
     }
@@ -135,7 +133,7 @@ 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,
+    pub vcpu_count: u64,
 }
 
 impl<'a> Persist<'a> for PVTime {
@@ -158,7 +156,7 @@ impl<'a> Persist<'a> for PVTime {
         constructor_args
             .resource_allocator
             .allocate_system_memory(
-                STEALTIME_STRUCT_MEM_SIZE * constructor_args.vcpu_count as u64,
+                STEALTIME_STRUCT_MEM_SIZE * constructor_args.vcpu_count,
                 STEALTIME_STRUCT_MEM_SIZE,
                 vm_allocator::AllocPolicy::ExactMatch(state.base_ipa),
             )
@@ -169,3 +167,41 @@ impl<'a> Persist<'a> for PVTime {
         ))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::device_manager::resources::ResourceAllocator;
+
+    #[test]
+    fn test_get_steal_time_region_addr() {
+        let base = GuestAddress(0x1000);
+        let dev = PVTime::from_base(base, 3);
+
+        assert_eq!(
+            dev.get_steal_time_region_addr(0).unwrap(),
+            GuestAddress(0x1000)
+        );
+        assert_eq!(
+            dev.get_steal_time_region_addr(1).unwrap(),
+            GuestAddress(0x1000 + STEALTIME_STRUCT_MEM_SIZE)
+        );
+        assert!(matches!(
+            dev.get_steal_time_region_addr(3),
+            Err(PVTimeError::InvalidVcpuIndex(3))
+        ));
+    }
+
+    #[test]
+    fn test_new_pvtime_allocates_correctly() {
+        let mut allocator = ResourceAllocator::new().unwrap();
+        let dev = PVTime::new(&mut allocator, 2).unwrap();
+
+        assert_eq!(dev.vcpu_count, 2);
+        assert_eq!(
+            dev.base_ipa.0 % STEALTIME_STRUCT_MEM_SIZE,
+            0,
+            "Base IPA should be aligned"
+        );
+    }
+}

src/vmm/src/builder.rs

@@ -11,7 +11,6 @@ use std::sync::{Arc, Mutex};
 use event_manager::{MutEventSubscriber, SubscriberOps};
 use libc::EFD_NONBLOCK;
 use linux_loader::cmdline::Cmdline as LoaderKernelCmdline;
-use log::warn;
 use userfaultfd::Uffd;
 use utils::time::TimestampUs;
 #[cfg(target_arch = "aarch64")]
@@ -85,6 +84,7 @@ pub enum StartMicrovmError {
     /// Error creating VMGenID device: {0}
     CreateVMGenID(VmGenIdError),
     /// Error creating PVTime device: {0}
+    #[cfg(target_arch = "aarch64")]
     CreatePVTime(PVTimeError),
     /// Invalid Memory Configuration: {0}
     GuestMemory(crate::vstate::memory::MemoryError),
@@ -301,7 +301,9 @@ pub fn build_microvm_for_boot(
         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.");
+            log::warn!(
+                "PVTime is not supported by KVM. Steal time will not be reported to the guest."
+            );
             None
         };
     }
@@ -420,6 +422,7 @@ pub enum BuildMicrovmFromSnapshotError {
     ACPIDeviManager(#[from] ACPIDeviceManagerRestoreError),
     /// VMGenID update failed: {0}
     VMGenIDUpdate(std::io::Error),
+    #[cfg(target_arch = "aarch64")]
     /// Failed to restore PVTime device: {0}
     RestorePVTime(#[from] PVTimeError),
 }
@@ -477,11 +480,9 @@ pub fn build_microvm_from_snapshot(
     {
         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,
+                vcpu_count: vcpus.len() as u64,
             };
             vmm.pv_time = Some(
                 PVTime::restore(pvtime_ctor_args, &pvtime_state)
@@ -596,8 +597,7 @@ fn setup_pv_time(vmm: &mut Vmm, vcpus: &mut [Vcpu]) -> Result<PVTime, StartMicro
     use crate::arch::aarch64::pvtime::PVTime;
 
     // 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)
+    let pv_time = PVTime::new(&mut vmm.resource_allocator, vcpus.len() as u64)
         .map_err(StartMicrovmError::CreatePVTime)?;
 
     // Register all vcpus with pvtime device

src/vmm/src/persist.rs

@@ -157,7 +157,7 @@ pub enum CreateSnapshotError {
 }
 
 /// Snapshot version
-pub const SNAPSHOT_VERSION: Version = Version::new(5, 0, 0);
+pub const SNAPSHOT_VERSION: Version = Version::new(6, 1, 0);
 
 /// Creates a Microvm snapshot.
 pub fn create_snapshot(

tests/integration_tests/functional/test_steal_time.py

@@ -0,0 +1,119 @@
+# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Tests for verifying the PVTime device behavior under contention and across snapshots."""
+
+import time
+
+import pytest
+
+from framework.properties import global_props
+
+
+def get_steal_time_ms(vm):
+    """Returns total steal time of vCPUs in VM in milliseconds"""
+    _, out, _ = vm.ssh.run("grep -w '^cpu' /proc/stat")
+    steal_time_tck = int(out.strip().split()[8])
+    clk_tck = int(vm.ssh.run("getconf CLK_TCK").stdout)
+    return steal_time_tck / clk_tck * 1000
+
+
+@pytest.mark.skipif(
+    global_props.cpu_architecture != "aarch64", reason="Only run in aarch64"
+)
+def test_guest_has_pvtime_enabled(uvm_plain):
+    """
+    Check that the guest kernel has enabled PV steal time.
+    """
+    vm = uvm_plain
+    vm.spawn()
+    vm.basic_config()
+    vm.add_net_iface()
+    vm.start()
+
+    _, stdout, _ = vm.ssh.run("dmesg | grep 'stolen time PV'")
+    assert (
+        "stolen time PV" in stdout
+    ), "Guest kernel did not report PV steal time enabled"
+
+
+def test_pvtime_steal_time_increases(uvm_plain):
+    """
+    Test that PVTime steal time increases when both vCPUs are contended on the same pCPU.
+    """
+    vm = uvm_plain
+    vm.spawn()
+    vm.basic_config()
+    vm.add_net_iface()
+    vm.start()
+
+    # Pin both vCPUs to the same physical CPU to induce contention
+    vm.pin_vcpu(0, 0)
+    vm.pin_vcpu(1, 0)
+
+    # Start two infinite loops to hog CPU time
+    hog_cmd = "nohup bash -c 'while true; do :; done' >/dev/null 2>&1 &"
+    vm.ssh.run(hog_cmd)
+    vm.ssh.run(hog_cmd)
+
+    # Measure before and after steal time
+    steal_before = get_steal_time_ms(vm)
+    time.sleep(2)
+    steal_after = get_steal_time_ms(vm)
+
+    # Require increase in steal time
+    assert (
+        steal_after > steal_before
+    ), f"Steal time did not increase as expected. Before: {steal_before}, After: {steal_after}"
+
+
+def test_pvtime_snapshot(uvm_plain, microvm_factory):
+    """
+    Test that PVTime steal time is preserved across snapshot/restore
+    and continues increasing post-resume.
+    """
+    vm = uvm_plain
+    vm.spawn()
+    vm.basic_config()
+    vm.add_net_iface()
+    vm.start()
+
+    vm.pin_vcpu(0, 0)
+    vm.pin_vcpu(1, 0)
+
+    hog_cmd = "nohup bash -c 'while true; do :; done' >/dev/null 2>&1 &"
+    vm.ssh.run(hog_cmd)
+    vm.ssh.run(hog_cmd)
+
+    # Snapshot pre-steal time
+    steal_before = get_steal_time_ms(vm)
+
+    snapshot = vm.snapshot_full()
+    vm.kill()
+
+    # Restore microVM from snapshot and resume
+    restored_vm = microvm_factory.build()
+    restored_vm.spawn()
+    restored_vm.restore_from_snapshot(snapshot, resume=False)
+    snapshot.delete()
+
+    restored_vm.pin_vcpu(0, 0)
+    restored_vm.pin_vcpu(1, 0)
+    restored_vm.resume()
+
+    # Steal time just after restoring
+    steal_after_snap = get_steal_time_ms(restored_vm)
+
+    time.sleep(2)
+
+    # Steal time after running resumed VM
+    steal_after_resume = get_steal_time_ms(restored_vm)
+
+    # Ensure steal time persisted and continued increasing
+    tolerance = 1500  # 1.5 seconds tolerance for persistence check
+    persisted = steal_after_snap - steal_before < tolerance
+    increased = steal_after_resume > steal_after_snap
+
+    assert (
+        persisted and increased
+    ), "Steal time did not persist through snapshot or failed to increase after resume"