Update vcpu index to u64 and fix tests

- Change PVTime::vcpu_count to u64 to avoid casts
- Fix integration tests to check steal_time increase

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

Author: DakshinD

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),
             )

src/vmm/src/builder.rs

@@ -11,8 +11,6 @@ use std::sync::{Arc, Mutex};
 use event_manager::{MutEventSubscriber, SubscriberOps};
 use libc::EFD_NONBLOCK;
 use linux_loader::cmdline::Cmdline as LoaderKernelCmdline;
-#[cfg(target_arch = "aarch64")]
-use log::warn;
 use userfaultfd::Uffd;
 use utils::time::TimestampUs;
 #[cfg(target_arch = "aarch64")]
@@ -303,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
         };
     }
@@ -484,11 +484,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)
@@ -603,8 +601,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

tests/integration_tests/functional/test_pvtime.py renamed to tests/integration_tests/functional/test_steal_time.py

@@ -3,14 +3,21 @@
 
 """Tests for verifying the PVTime device behavior under contention and across snapshots."""
 
-import re
 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"
 )
@@ -43,35 +50,20 @@ def test_pvtime_steal_time_increases(uvm_plain):
     # Pin both vCPUs to the same physical CPU to induce contention
     vm.pin_vcpu(0, 0)
     vm.pin_vcpu(1, 0)
-    vm.pin_vmm(1)
-    vm.pin_api(2)
 
     # 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)
-
-    # Measure steal time before
-    _, out_before, _ = vm.ssh.run("grep '^cpu[0-9]' /proc/stat")
-    steal_before = sum(
-        int(re.split(r"\s+", line.strip())[8])
-        for line in out_before.strip().splitlines()
-    )
-
-    time.sleep(2)
-
-    # Measure steal time after
-    _, out_after, _ = vm.ssh.run("grep '^cpu[0-9]' /proc/stat")
-    steal_after = sum(
-        int(re.split(r"\s+", line.strip())[8])
-        for line in out_after.strip().splitlines()
-    )
+    steal_after = get_steal_time_ms(vm)
 
     # Require increase in steal time
     assert (
-        steal_after - steal_before >= 200
+        steal_after > steal_before
     ), f"Steal time did not increase as expected. Before: {steal_before}, After: {steal_after}"
 
 
@@ -88,21 +80,13 @@ def test_pvtime_snapshot(uvm_plain, microvm_factory):
 
     vm.pin_vcpu(0, 0)
     vm.pin_vcpu(1, 0)
-    vm.pin_vmm(1)
-    vm.pin_api(2)
 
     hog_cmd = "nohup bash -c 'while true; do :; done' >/dev/null 2>&1 &"
     vm.ssh.run(hog_cmd)
     vm.ssh.run(hog_cmd)
 
-    time.sleep(1)
-
     # Snapshot pre-steal time
-    _, out_before_snap, _ = vm.ssh.run("grep '^cpu[0-9]' /proc/stat")
-    steal_before = [
-        int(re.split(r"\s+", line.strip())[8])
-        for line in out_before_snap.strip().splitlines()
-    ]
+    steal_before = get_steal_time_ms(vm)
 
     snapshot = vm.snapshot_full()
     vm.kill()
@@ -115,31 +99,20 @@ def test_pvtime_snapshot(uvm_plain, microvm_factory):
 
     restored_vm.pin_vcpu(0, 0)
     restored_vm.pin_vcpu(1, 0)
-    restored_vm.pin_vmm(1)
-    restored_vm.pin_api(2)
     restored_vm.resume()
 
-    time.sleep(1)
-
     # Steal time just after restoring
-    _, out_after_snap, _ = restored_vm.ssh.run("grep '^cpu[0-9]' /proc/stat")
-    steal_after_snap = [
-        int(re.split(r"\s+", line.strip())[8])
-        for line in out_after_snap.strip().splitlines()
-    ]
+    steal_after_snap = get_steal_time_ms(restored_vm)
 
     time.sleep(2)
 
     # Steal time after running resumed VM
-    _, out_after_resume, _ = restored_vm.ssh.run("grep '^cpu[0-9]' /proc/stat")
-    steal_after_resume = [
-        int(re.split(r"\s+", line.strip())[8])
-        for line in out_after_resume.strip().splitlines()
-    ]
+    steal_after_resume = get_steal_time_ms(restored_vm)
 
     # Ensure steal time persisted and continued increasing
-    persisted = sum(steal_before) + 100 <= sum(steal_after_snap)
-    increased = sum(steal_after_resume) > sum(steal_after_snap)
+    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