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>

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,