hypervisor: aarch64: capture the guest counter for snapshot/restore

Unlike x86, ARM64 has no kvmclock support to sync guest time upon
required. However, the guest reads the architected virtual timer
(CNTVCT_EL0) directly which can be modified by the VMM to update the
time after snapshot restore. Since the CNTVCT is in ticks, we also need
to read CNTFRQ (via mrs due to lack of ONEREG interface) to compute the
ticks from wall clock difference.

Because the counter is a vCPU register, the capture must run with the
vCPUs quiesced, so the VMM now captures the clock just after
cpu_manager.pause() through the boot vCPU. This is behaviorally
identical for x86, whose clock is VM-wide. There is no restore/advance
yet, so aarch64 guests still resume behind real time until the following
commit.

Signed-off-by: Atish Patra <atishp@meta.com>

hypervisor/src/arch/aarch64/mod.rs

@@ -7,6 +7,23 @@ pub mod regs;
 
 use serde::{Deserialize, Serialize};
 
+/// Reads the architected counter frequency (`CNTFRQ_EL0`, Hz) from the host: KVM
+/// does not expose it through ONE_REG and the guest counter runs at host frequency.
+pub fn get_cntfrq() -> u64 {
+    use std::arch::asm;
+    let cntfrq: u64;
+    // SAFETY: `mrs cntfrq_el0` only reads a read-only system register and
+    // touches no memory (nomem, nostack, preserves_flags).
+    unsafe {
+        asm!(
+            "mrs {}, cntfrq_el0",
+            out(reg) cntfrq,
+            options(nomem, nostack, preserves_flags),
+        );
+    }
+    cntfrq
+}
+
 #[derive(Clone, Serialize, Deserialize)]
 pub struct ExtendedReg {
     pub id: u64,

hypervisor/src/cpu.rs

@@ -511,6 +511,11 @@ pub trait Vcpu: Send + Sync {
     #[cfg(target_arch = "aarch64")]
     fn get_sys_reg(&self, sys_reg: u32) -> Result<u64>;
     ///
+    /// Gets the guest virtual counter (`CNTVCT_EL0`) via `KVM_REG_ARM_TIMER_CNT`.
+    ///
+    #[cfg(all(target_arch = "aarch64", feature = "kvm"))]
+    fn get_cntvct(&self) -> Result<u64>;
+    ///
     /// Gets the value of a non-core register on RISC-V 64-bit
     ///
     #[cfg(target_arch = "riscv64")]

hypervisor/src/kvm/mod.rs

@@ -27,6 +27,8 @@ use std::sync::Mutex;
 #[cfg(target_arch = "x86_64")]
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, RwLock};
+#[cfg(target_arch = "aarch64")]
+use std::time::{SystemTime, UNIX_EPOCH};
 
 #[cfg(target_arch = "x86_64")]
 use anyhow::Context;
@@ -78,7 +80,9 @@ use x86_64::check_required_kvm_extensions;
 pub use x86_64::{CpuId, ExtendedControlRegisters, MsrEntries, VcpuKvmState};
 
 #[cfg(target_arch = "x86_64")]
-use crate::{ClockData, ClockState};
+use crate::ClockData;
+#[cfg(any(target_arch = "x86_64", all(target_arch = "aarch64", feature = "kvm")))]
+use crate::ClockState;
 #[cfg(target_arch = "x86_64")]
 use crate::arch::x86::{
     CpuIdEntry, FpuState, LapicState, MTRR_MSR_INDICES, MsrEntry, NUM_IOAPIC_PINS,
@@ -119,9 +123,11 @@ pub use kvm_bindings::{
 #[cfg(target_arch = "aarch64")]
 use kvm_bindings::{
     KVM_GUESTDBG_USE_HW, KVM_NR_SPSR, KVM_REG_ARM_COPROC_MASK, KVM_REG_ARM_CORE, KVM_REG_ARM64,
-    KVM_REG_ARM64_SYSREG, KVM_REG_ARM64_SYSREG_CRM_MASK, KVM_REG_ARM64_SYSREG_CRN_MASK,
-    KVM_REG_ARM64_SYSREG_OP0_MASK, KVM_REG_ARM64_SYSREG_OP1_MASK, KVM_REG_ARM64_SYSREG_OP2_MASK,
-    KVM_REG_SIZE_U32, KVM_REG_SIZE_U64, KVM_REG_SIZE_U128, kvm_regs, user_pt_regs,
+    KVM_REG_ARM64_SYSREG, KVM_REG_ARM64_SYSREG_CRM_MASK, KVM_REG_ARM64_SYSREG_CRM_SHIFT,
+    KVM_REG_ARM64_SYSREG_CRN_MASK, KVM_REG_ARM64_SYSREG_CRN_SHIFT, KVM_REG_ARM64_SYSREG_OP0_MASK,
+    KVM_REG_ARM64_SYSREG_OP0_SHIFT, KVM_REG_ARM64_SYSREG_OP1_MASK, KVM_REG_ARM64_SYSREG_OP1_SHIFT,
+    KVM_REG_ARM64_SYSREG_OP2_MASK, KVM_REG_SIZE_U32, KVM_REG_SIZE_U64, KVM_REG_SIZE_U128, kvm_regs,
+    user_pt_regs,
 };
 #[cfg(target_arch = "riscv64")]
 use kvm_bindings::{KVM_REG_RISCV_CORE, KVM_REG_RISCV_TIMER, kvm_riscv_core};
@@ -141,10 +147,27 @@ use vmm_sys_util::{ioctl::ioctl_with_val, ioctl_iowr_nr};
 #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
 use crate::RegList;
 #[cfg(target_arch = "aarch64")]
-use crate::arch::aarch64::regs;
+use crate::TimerState;
+#[cfg(target_arch = "aarch64")]
+use crate::arch::aarch64::{get_cntfrq, regs};
 #[cfg(target_arch = "x86_64")]
 use crate::kvm::x86_64::XsaveStateError;
 
+// `KVM_REG_ARM_TIMER_CNT`, the timer-counter firmware register the kernel defines
+// as `ARM64_SYS_REG(3, 3, 14, 3, 2)`. kvm-bindings exposes no constant for it, so
+// build the id the way the kernel's `ARM64_SYS_REG` macro would. The kernel pins
+// the EL0 virtual-timer reg encodings by value (CVAL/CNT were historically
+// swapped), so CNT must be exactly op0=3, op1=3, crn=14, crm=3, op2=2.
+#[cfg(target_arch = "aarch64")]
+const KVM_REG_ARM_TIMER_CNT: u64 = KVM_REG_ARM64
+    | KVM_REG_SIZE_U64
+    | KVM_REG_ARM64_SYSREG as u64
+    | ((3_u64 << KVM_REG_ARM64_SYSREG_OP0_SHIFT) & KVM_REG_ARM64_SYSREG_OP0_MASK as u64)
+    | ((3_u64 << KVM_REG_ARM64_SYSREG_OP1_SHIFT) & KVM_REG_ARM64_SYSREG_OP1_MASK as u64)
+    | ((14_u64 << KVM_REG_ARM64_SYSREG_CRN_SHIFT) & KVM_REG_ARM64_SYSREG_CRN_MASK as u64)
+    | ((3_u64 << KVM_REG_ARM64_SYSREG_CRM_SHIFT) & KVM_REG_ARM64_SYSREG_CRM_MASK as u64)
+    | (2_u64 & KVM_REG_ARM64_SYSREG_OP2_MASK as u64);
+
 #[cfg(target_arch = "x86_64")]
 ioctl_io_nr!(KVM_NMI, kvm_bindings::KVMIO, 0x9a);
 
@@ -1266,10 +1289,30 @@ impl vm::Vm for KvmVm {
 
     /// Capture kvmclock (filling realtime) for snapshot/migration.
     #[cfg(target_arch = "x86_64")]
-    fn snapshot_clock(&self) -> vm::Result<Option<ClockState>> {
+    fn snapshot_clock(&self, _boot_vcpu: &dyn cpu::Vcpu) -> vm::Result<Option<ClockState>> {
         Ok(Some(self.get_clock()?.with_realtime_filled()))
     }
 
+    /// Capture the guest virtual counter and host wall clock for
+    /// snapshot/migration. The vCPUs must be paused.
+    #[cfg(target_arch = "aarch64")]
+    fn snapshot_clock(&self, boot_vcpu: &dyn cpu::Vcpu) -> vm::Result<Option<ClockState>> {
+        // cntvct and the host wall clock must be sampled back-to-back; cntfrq is
+        // static, so read it last.
+        let cntvct = boot_vcpu
+            .get_cntvct()
+            .map_err(|e| vm::HypervisorVmError::CaptureTimerState(e.into()))?;
+        let host_realtime_ns = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .map_err(|e| vm::HypervisorVmError::CaptureTimerState(e.into()))?
+            .as_nanos() as u64;
+        Ok(Some(TimerState {
+            cntvct,
+            host_realtime_ns,
+            cntfrq: get_cntfrq(),
+        }))
+    }
+
     /// Restore kvmclock before the vCPUs resume.
     #[cfg(target_arch = "x86_64")]
     fn restore_clock(&self, state: &ClockState) -> vm::Result<()> {
@@ -2734,6 +2777,18 @@ impl cpu::Vcpu for KvmVcpu {
         Ok(u64::from_le_bytes(bytes))
     }
 
+    ///
+    /// Gets the guest virtual counter via `KVM_REG_ARM_TIMER_CNT`.
+    ///
+    #[cfg(target_arch = "aarch64")]
+    fn get_cntvct(&self) -> cpu::Result<u64> {
+        let mut bytes = [0_u8; 8];
+        self.fd
+            .get_one_reg(KVM_REG_ARM_TIMER_CNT, &mut bytes)
+            .map_err(|e| cpu::HypervisorCpuError::GetSysRegister(e.into()))?;
+        Ok(u64::from_le_bytes(bytes))
+    }
+
     ///
     /// Gets the value of a non-core register
     ///

hypervisor/src/lib.rs

@@ -226,9 +226,28 @@ impl ClockData {
 }
 
 /// Guest clock state preserved across pause/resume and snapshot/restore
-/// (`ClockData` on x86).
+/// (`ClockData` on x86, `TimerState` on aarch64+kvm.
+/// The platform where it has not guest clock, `Option<ClockState>` will be None.
 #[cfg(target_arch = "x86_64")]
 pub type ClockState = ClockData;
+#[cfg(all(target_arch = "aarch64", feature = "kvm"))]
+pub type ClockState = TimerState;
+#[cfg(not(any(target_arch = "x86_64", all(target_arch = "aarch64", feature = "kvm"))))]
+#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
+pub enum ClockState {}
+
+/// Guest timer state captured on aarch64 for snapshot/migration: the guest
+/// virtual counter (`CNTVCT_EL0`) plus the host wall clock and counter
+/// frequency needed to advance it to current wall time on restore. aarch64 has
+/// no `KVM_SET_CLOCK`/`KVM_CLOCK_REALTIME`, so the VMM records these and does the
+/// advance itself.
+#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
+#[cfg(all(target_arch = "aarch64", feature = "kvm"))]
+pub struct TimerState {
+    pub cntvct: u64,
+    pub host_realtime_ns: u64,
+    pub cntfrq: u64,
+}
 
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
 pub struct HypervisorVmConfig {

hypervisor/src/mshv/mod.rs

@@ -1324,6 +1324,11 @@ impl cpu::Vcpu for MshvVcpu {
         Ok(res)
     }
 
+    #[cfg(all(target_arch = "aarch64", feature = "kvm"))]
+    fn get_cntvct(&self) -> cpu::Result<u64> {
+        unimplemented!()
+    }
+
     #[cfg(target_arch = "aarch64")]
     fn get_reg_list(&self, _reg_list: &mut crate::RegList) -> cpu::Result<()> {
         unimplemented!()
@@ -2202,7 +2207,7 @@ impl vm::Vm for MshvVm {
 
     /// Capture the partition reference time for snapshot/migration.
     #[cfg(target_arch = "x86_64")]
-    fn snapshot_clock(&self) -> vm::Result<Option<ClockState>> {
+    fn snapshot_clock(&self, _boot_vcpu: &dyn cpu::Vcpu) -> vm::Result<Option<ClockState>> {
         Ok(Some(self.get_clock()?.with_realtime_filled()))
     }
 

hypervisor/src/vm.rs

@@ -23,7 +23,8 @@ use thiserror::Error;
 use vmm_sys_util::eventfd::EventFd;
 
 #[cfg(target_arch = "x86_64")]
-use crate::{ClockData, ClockState};
+use crate::ClockData;
+use crate::ClockState;
 #[cfg(target_arch = "aarch64")]
 use crate::arch::aarch64::gic::{Vgic, VgicConfig};
 #[cfg(target_arch = "riscv64")]
@@ -140,6 +141,12 @@ pub enum HypervisorVmError {
     #[error("Failed to set clock")]
     SetClock(#[source] anyhow::Error),
     ///
+    /// Capture guest timer state error (aarch64)
+    ///
+    #[cfg(all(target_arch = "aarch64", feature = "kvm"))]
+    #[error("Failed to capture the guest timer state")]
+    CaptureTimerState(#[source] anyhow::Error),
+    ///
     /// Create passthrough device
     ///
     #[error("Failed to create passthrough device")]
@@ -385,9 +392,10 @@ pub trait Vm: Send + Sync + Any {
     #[cfg(target_arch = "x86_64")]
     fn set_clock(&self, data: &ClockData) -> Result<()>;
     /// Capture the guest clock for snapshot/migration while the VM is paused.
-    /// `Ok(None)` means this backend has no clock to preserve.
-    #[cfg(target_arch = "x86_64")]
-    fn snapshot_clock(&self) -> Result<Option<ClockState>> {
+    /// `Ok(None)` means this backend has no clock to preserve. `boot_vcpu` is the
+    /// boot vCPU, used by backends whose clock is per-vCPU state (e.g the aarch64
+    /// counter);
+    fn snapshot_clock(&self, _boot_vcpu: &dyn Vcpu) -> Result<Option<ClockState>> {
         Ok(None)
     }
     /// Re-establish the guest clock before the vCPUs resume.

vmm/src/cpu.rs

@@ -499,6 +499,12 @@ pub struct Vcpu {
 }
 
 impl Vcpu {
+    /// Borrow the underlying hypervisor vCPU, for guest-clock capture/restore
+    /// because the aarch64 counter is per-vCPU state
+    pub fn hypervisor_vcpu(&self) -> &dyn hypervisor::Vcpu {
+        self.vcpu.as_ref()
+    }
+
     /// Constructs a new VCPU for `vm`.
     ///
     /// # Arguments
@@ -1691,6 +1697,11 @@ impl CpuManager {
             .collect()
     }
 
+    /// The boot vCPU (vCPU 0), or `None` before the vCPUs are created.
+    pub fn boot_vcpu(&self) -> Option<Arc<Mutex<Vcpu>>> {
+        self.vcpus.first().cloned()
+    }
+
     #[cfg(target_arch = "aarch64")]
     pub fn get_saved_states(&self) -> Vec<CpuState> {
         self.vcpus

vmm/src/vm.rs

@@ -100,11 +100,9 @@ use crate::landlock::LandlockError;
 use crate::memory_manager::{
     Error as MemoryManagerError, MemoryManager, MemoryManagerSnapshotData,
 };
-#[cfg(target_arch = "x86_64")]
-use crate::migration::get_vm_snapshot;
 #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
 use crate::migration::url_to_file;
-use crate::migration::{SNAPSHOT_CONFIG_FILE, SNAPSHOT_STATE_FILE, url_to_path};
+use crate::migration::{SNAPSHOT_CONFIG_FILE, SNAPSHOT_STATE_FILE, get_vm_snapshot, url_to_path};
 #[cfg(all(feature = "kvm", feature = "sev_snp", feature = "fw_cfg"))]
 use crate::sev::MeasuredBootInfo;
 #[cfg(feature = "fw_cfg")]
@@ -528,7 +526,6 @@ pub struct Vm {
     #[cfg_attr(any(not(feature = "kvm"), target_arch = "aarch64"), allow(dead_code))]
     // The hypervisor abstracted virtual machine.
     vm: Arc<dyn hypervisor::Vm>,
-    #[cfg(target_arch = "x86_64")]
     saved_clock: Option<hypervisor::ClockState>,
     #[cfg(not(target_arch = "riscv64"))]
     numa_nodes: NumaNodes,
@@ -682,7 +679,6 @@ impl Vm {
             .transpose()
             .map_err(Error::InitramfsFile)?;
 
-        #[cfg(target_arch = "x86_64")]
         let saved_clock = if let Some(snapshot) = snapshot.as_ref() {
             let vm_snapshot = get_vm_snapshot(snapshot).map_err(Error::Restore)?;
             vm_snapshot.clock
@@ -707,7 +703,6 @@ impl Vm {
             cpu_manager,
             memory_manager,
             vm,
-            #[cfg(target_arch = "x86_64")]
             saved_clock,
             #[cfg(not(target_arch = "riscv64"))]
             numa_nodes,
@@ -3214,14 +3209,6 @@ impl Pausable for Vm {
             .valid_transition(new_state)
             .map_err(|e| MigratableError::Pause(anyhow!("Invalid transition: {e:?}")))?;
 
-        #[cfg(target_arch = "x86_64")]
-        {
-            self.saved_clock = self
-                .vm
-                .snapshot_clock()
-                .map_err(|e| MigratableError::Pause(anyhow!("Could not capture guest clock: {e}")))?;
-        }
-
         // Before pausing the vCPUs activate any pending virtio devices that might
         // need activation between starting the pause (or e.g. a migration it's part of)
         self.activate_virtio_devices().map_err(|e| {
@@ -3229,6 +3216,16 @@ impl Pausable for Vm {
         })?;
 
         self.cpu_manager.lock().unwrap().pause()?;
+
+        // Capture the guest clock once the vCPUs are paused: aarch64 reads the
+        // boot vCPU's virtual counter, which requires the vCPU to be quiesced.
+        if let Some(boot_vcpu) = self.cpu_manager.lock().unwrap().boot_vcpu() {
+            let boot_vcpu = boot_vcpu.lock().unwrap();
+            self.saved_clock = self
+                .vm
+                .snapshot_clock(boot_vcpu.hypervisor_vcpu())
+                .map_err(|e| MigratableError::Pause(anyhow!("Could not capture guest clock: {e}")))?;
+        }
         self.device_manager.lock().unwrap().pause()?;
 
         self.vm
@@ -3279,7 +3276,7 @@ impl Pausable for Vm {
 
 #[derive(Serialize, Deserialize)]
 pub struct VmSnapshot {
-    #[cfg(target_arch = "x86_64")]
+    #[serde(default)]
     pub clock: Option<hypervisor::ClockState>,
     #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
     pub common_cpuid: Vec<hypervisor::arch::x86::CpuIdEntry>,
@@ -3341,7 +3338,6 @@ impl Snapshottable for Vm {
         };
 
         let vm_snapshot_state = VmSnapshot {
-            #[cfg(target_arch = "x86_64")]
             clock: self.saved_clock,
             #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
             common_cpuid,