From e5b745945be0c6c0af0a230743b65d1f65c377fd Mon Sep 17 00:00:00 2001
From: Yijun Zhao <ariesdevil77@gmail.com>
Date: Sat, 5 Oct 2024 18:07:13 +0800
Subject: [PATCH] feat: Check if VM supports TSC

---
 src/tsc_now.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/tsc_now.rs b/src/tsc_now.rs
index 220c82d..216f074 100644
--- a/src/tsc_now.rs
+++ b/src/tsc_now.rs
@@ -4,6 +4,7 @@
 
 use std::cell::UnsafeCell;
 use std::fs::read_to_string;
+use std::io::ErrorKind;
 use std::time::Instant;
 
 static TSC_STATE: TSCState = TSCState {
@@ -92,10 +93,51 @@ impl TSCLevel {
 /// rely on the result to say tsc is stable so that no need to
 /// sync TSCs by ourselves.
 fn is_tsc_stable() -> bool {
-    let clock_source =
-        read_to_string("/sys/devices/system/clocksource/clocksource0/available_clocksource");
+    has_invariant_tsc() || clock_source_has_tsc()
+}
+
+fn clock_source_has_tsc() -> bool {
+    #[cfg(target_os = "linux")]
+    {
+        let clock_sources = [
+            "/sys/devices/system/clocksource/clocksource0/current_clocksource",
+            "/sys/devices/system/clocksource/clocksource0/available_clocksource",
+        ];
+
+        for &path in &clock_sources {
+            match read_to_string(path) {
+                Ok(content) => {
+                    if content.contains("tsc") {
+                        return true;
+                    }
+                }
+                Err(e) if e.kind() == ErrorKind::NotFound => continue,
+                Err(_) => return false,
+            }
+        }
+        false
+    }
+
+    #[cfg(not(target_os = "linux"))]
+    false
+}
+
+/// Invariant TSC could make sure TSC got synced among multi CPUs.
+/// They will be reset at same time, and run in same frequency.
+/// But in some VM, the max Extended Function in CPUID is < 0x80000007,
+/// we should enable TSC if the system clock source is TSC.
+#[inline]
+fn has_invariant_tsc() -> bool {
+    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+    unsafe {
+        use core::arch::x86_64::__cpuid;
+        let cpuid_invariant_tsc_bts = 1 << 8;
+        __cpuid(0x80000000).eax >= 0x80000007
+            && __cpuid(0x80000007).edx & cpuid_invariant_tsc_bts != 0
+    }
 
-    clock_source.map(|s| s.contains("tsc")).unwrap_or(false)
+    #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
+    false
 }
 
 /// Returns (1) cycles per second and (2) cycles from anchor.