vmm: x86_64: synchronize TSC offsets after restoring vCPUs

src/vmm/src/arch/x86_64/vcpu.rs

@@ -10,12 +10,16 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use kvm_bindings::{
-    CpuId, KVM_MAX_CPUID_ENTRIES, KVM_MAX_MSR_ENTRIES, Msrs, Xsave, kvm_debugregs, kvm_lapic_state,
-    kvm_mp_state, kvm_regs, kvm_sregs, kvm_vcpu_events, kvm_xcrs, kvm_xsave, kvm_xsave2,
+    CpuId, KVM_MAX_CPUID_ENTRIES, KVM_MAX_MSR_ENTRIES, KVM_VCPU_TSC_CTRL, KVM_VCPU_TSC_OFFSET,
+    KVMIO, Msrs, Xsave, kvm_debugregs, kvm_device_attr, kvm_lapic_state, kvm_mp_state, kvm_regs,
+    kvm_sregs, kvm_vcpu_events, kvm_xcrs, kvm_xsave, kvm_xsave2,
 };
 use kvm_ioctls::{VcpuExit, VcpuFd};
 use serde::{Deserialize, Serialize};
+use vmm_sys_util::errno;
 use vmm_sys_util::fam::{self, FamStruct};
+use vmm_sys_util::ioctl::ioctl_with_ref;
+use vmm_sys_util::ioctl_iow_nr;
 
 use crate::arch::EntryPoint;
 use crate::arch::x86_64::generated::msr_index::{MSR_IA32_TSC, MSR_IA32_TSC_DEADLINE};
@@ -36,6 +40,15 @@ use crate::vstate::vm::KvmVm;
 const TSC_KHZ_TOL_NUMERATOR: i64 = 250;
 const TSC_KHZ_TOL_DENOMINATOR: i64 = 1_000_000;
 
+// kvm-ioctls only exposes vCPU device-attribute accessors on aarch64.
+#[allow(missing_docs)]
+mod ioctls {
+    use super::*;
+    ioctl_iow_nr!(KVM_SET_DEVICE_ATTR, KVMIO, 0xe1, kvm_device_attr);
+    ioctl_iow_nr!(KVM_GET_DEVICE_ATTR, KVMIO, 0xe2, kvm_device_attr);
+    ioctl_iow_nr!(KVM_HAS_DEVICE_ATTR, KVMIO, 0xe3, kvm_device_attr);
+}
+
 /// A set of MSRs that should be restored separately after all other MSRs have already been restored
 const DEFERRED_MSRS: [u32; 1] = [
     // MSR_IA32_TSC_DEADLINE must be restored after MSR_IA32_TSC, otherwise we risk "losing" timer
@@ -390,6 +403,50 @@ impl KvmVcpu {
         Ok(res)
     }
 
+    /// Whether KVM supports directly accessing this vCPU's TSC offset.
+    pub fn supports_tsc_offset_attr(&self) -> bool {
+        let attr = kvm_device_attr {
+            group: KVM_VCPU_TSC_CTRL,
+            attr: u64::from(KVM_VCPU_TSC_OFFSET),
+            ..Default::default()
+        };
+        // SAFETY: The vCPU fd and attribute are valid, and HAS_DEVICE_ATTR ignores addr.
+        unsafe { ioctl_with_ref(&self.fd, ioctls::KVM_HAS_DEVICE_ATTR(), &attr) == 0 }
+    }
+
+    /// Read this vCPU's TSC offset relative to the host TSC.
+    pub fn get_tsc_offset(&self) -> Result<i64, errno::Error> {
+        let mut offset = 0_i64;
+        let attr = kvm_device_attr {
+            group: KVM_VCPU_TSC_CTRL,
+            attr: u64::from(KVM_VCPU_TSC_OFFSET),
+            addr: std::ptr::from_mut(&mut offset) as u64,
+            ..Default::default()
+        };
+        // SAFETY: The attribute points to a writable i64 that lives through the ioctl.
+        let ret = unsafe { ioctl_with_ref(&self.fd, ioctls::KVM_GET_DEVICE_ATTR(), &attr) };
+        if ret != 0 {
+            return Err(errno::Error::last());
+        }
+        Ok(offset)
+    }
+
+    /// Set this vCPU's TSC offset while the vCPU is not running.
+    pub fn set_tsc_offset(&self, offset: i64) -> Result<(), errno::Error> {
+        let attr = kvm_device_attr {
+            group: KVM_VCPU_TSC_CTRL,
+            attr: u64::from(KVM_VCPU_TSC_OFFSET),
+            addr: std::ptr::from_ref(&offset) as u64,
+            ..Default::default()
+        };
+        // SAFETY: The attribute points to a readable i64 that lives through the ioctl.
+        let ret = unsafe { ioctl_with_ref(&self.fd, ioctls::KVM_SET_DEVICE_ATTR(), &attr) };
+        if ret != 0 {
+            return Err(errno::Error::last());
+        }
+        Ok(())
+    }
+
     /// Get CPUID for this vCPU.
     ///
     /// Opposed to KVM_GET_SUPPORTED_CPUID, KVM_GET_CPUID2 does not update "nent" with valid number

src/vmm/src/builder.rs

@@ -419,6 +419,41 @@ pub enum BuildMicrovmFromSnapshotError {
     UnsupportedClockRealtime,
 }
 
+/// Align restored TSC offsets before starting any vCPU threads.
+#[cfg(target_arch = "x86_64")]
+fn synchronize_tsc_offsets(vcpus: &[crate::Vcpu]) {
+    let reference = vcpus
+        .first()
+        .expect("TSC synchronization requires at least one vCPU");
+    // KVM_VCPU_TSC_OFFSET is supported since Linux 5.16 (commit 828ca89628bf).
+    // Probe the attribute to keep restore best-effort on older hosts.
+    if !reference.kvm_vcpu.supports_tsc_offset_attr() {
+        crate::logger::debug!("KVM does not support TSC offset synchronization");
+        return;
+    }
+    let offset = match reference.kvm_vcpu.get_tsc_offset() {
+        Ok(offset) => offset,
+        Err(err) => {
+            crate::logger::warn!("Failed to read vCPU 0 TSC offset: {err}");
+            return;
+        }
+    };
+
+    let mut synchronized = true;
+    for vcpu in vcpus {
+        if let Err(err) = vcpu.kvm_vcpu.set_tsc_offset(offset) {
+            crate::logger::warn!(
+                "Failed to synchronize vCPU {} TSC offset: {err}",
+                vcpu.kvm_vcpu.index
+            );
+            synchronized = false;
+        }
+    }
+    if synchronized {
+        crate::logger::debug!("Synchronized all vCPU TSC offsets to {offset}");
+    }
+}
+
 /// Builds and starts a microVM based on the provided MicrovmState.
 ///
 /// An `Arc` reference of the built `Vmm` is also plugged in the `EventManager`, while another
@@ -472,6 +507,12 @@ pub fn build_microvm_from_snapshot(
             .map_err(BuildMicrovmFromSnapshotError::RestoreVcpus)?;
     }
 
+    // Restoring TSC MSRs separately can leave different offsets on each vCPU.
+    // Preserve vCPU0's restored timeline while preventing time from going backwards
+    // when the guest migrates between vCPUs.
+    #[cfg(target_arch = "x86_64")]
+    synchronize_tsc_offsets(&vcpus);
+
     #[cfg(target_arch = "aarch64")]
     {
         if clock_realtime {
@@ -821,6 +862,49 @@ pub(crate) mod tests {
         }
     }
 
+    #[cfg(target_arch = "x86_64")]
+    #[test]
+    fn test_synchronize_tsc_offsets() {
+        for offsets in [
+            &[-5_000_000_000_i64][..],
+            &[-5_000_000_000, 5_000_000_000][..],
+            &[5_000_000_000, 5_000_000_000][..],
+        ] {
+            let count = u8::try_from(offsets.len()).unwrap();
+            let mut source_vm = setup_vm_with_memory(0x1000);
+            let source_vcpus = source_vm.create_vcpus(count).unwrap();
+            if !source_vcpus[0].kvm_vcpu.supports_tsc_offset_attr() {
+                eprintln!("Skipping TSC offset synchronization: KVM attribute unavailable");
+                return;
+            }
+            let vm_state = source_vm.save_state().unwrap();
+            let states: Vec<_> = source_vcpus
+                .iter()
+                .map(|vcpu| vcpu.kvm_vcpu.save_state().unwrap())
+                .collect();
+
+            let mut vm = setup_vm_with_memory(0x1000);
+            let vcpus = vm.create_vcpus(count).unwrap();
+            for ((vcpu, state), &offset) in vcpus.iter().zip(&states).zip(offsets) {
+                vcpu.kvm_vcpu.restore_state(state).unwrap();
+                // Establish exact offsets without depending on KVM's MSR-write heuristics.
+                vcpu.kvm_vcpu.set_tsc_offset(offset).unwrap();
+                assert_eq!(vcpu.kvm_vcpu.get_tsc_offset().unwrap(), offset);
+            }
+
+            synchronize_tsc_offsets(&vcpus);
+            for vcpu in &vcpus {
+                assert_eq!(vcpu.kvm_vcpu.get_tsc_offset().unwrap(), offsets[0]);
+            }
+
+            // The subsequent VM clock restore must preserve the synchronized offsets.
+            vm.restore_state(&vm_state, false).unwrap();
+            for vcpu in &vcpus {
+                assert_eq!(vcpu.kvm_vcpu.get_tsc_offset().unwrap(), offsets[0]);
+            }
+        }
+    }
+
     fn cmdline_contains(cmdline: &Cmdline, slug: &str) -> bool {
         // The following unwraps can never fail; the only way any of these methods
         // would return an `Err` is if one of the following conditions is met:

tests/integration_tests/functional/test_snapshot_basic.py

@@ -9,6 +9,7 @@
 import platform
 import re
 import shutil
+import textwrap
 import time
 import uuid
 from pathlib import Path
@@ -743,6 +744,32 @@ def read_guest_clocksource(vm):
     return stdout.strip()
 
 
+def check_guest_monotonic_across_vcpus(vm):
+    """Check clock monotonicity as a task migrates between restored vCPUs."""
+    _, stdout, _ = vm.ssh.check_output(
+        textwrap.dedent("""\
+            python3 - <<'PY'
+            import os
+            import time
+
+            cpus = sorted(os.sched_getaffinity(0))
+            assert len(cpus) == 2, cpus
+            previous = time.monotonic_ns()
+            for _ in range(10_000):
+                for cpu in cpus:
+                    os.sched_setaffinity(0, {cpu})
+                    current = time.monotonic_ns()
+                    assert current >= previous, (cpu, previous, current)
+                    previous = current
+
+            print("20,000 clock samples across vCPU migrations: no regressions")
+            PY
+            """),
+        timeout=30,
+    )
+    print(stdout.strip())
+
+
 @pytest.mark.parametrize("clocksource", CLOCK_SOURCES)
 @pytest.mark.parametrize("clock_realtime", [False, True])
 def test_clocksource_snapshot_restore(
@@ -804,6 +831,8 @@ def test_clocksource_snapshot_restore(
     # If guest_delta is close to host_delta, the clock jumped forward
     # (suspend/resume behavior). If it's near 0, it resumed from where
     # it left off.
+    if not clock_realtime:
+        assert 0 <= guest_delta < 5.0, f"Unexpected clock delta: {guest_delta:.6f}s"
     jumped = abs(guest_delta - host_delta) < 5.0
 
     jumped_str = "JUMPED" if jumped else "RESUMED"
@@ -820,3 +849,7 @@ def test_clocksource_snapshot_restore(
     assert (
         jumped == clock_realtime
     ), f"Clock {jumped_str} but clock_realtime was {"not" if clock_realtime else ""} set."
+
+    if clocksource == "tsc":
+        check_guest_monotonic_across_vcpus(restored_vm)
+        assert read_guest_clocksource(restored_vm) == "tsc"