[hyperlight-host/trace] Fixup: update the way the TSC frequency is calculated

- Rely on TSC readings on the host from when the guest is started and
when the first records are received. Use these two moments in time and TSC
readings to calculate the TSC frequency.
- Afterwards we store the TSC reading from inside the guest and calculate all
incoming TSC readings relative to the first one and using the frequency we can
calculate the amount of time it took.

Signed-off-by: Doru Blânzeanu <[email protected]>

Author: dblnz

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -763,8 +763,10 @@ impl Hypervisor for HypervLinuxDriver {
         } else {
             #[cfg(feature = "trace_guest")]
             if self.trace_info.guest_start_epoch.is_none() {
-                // Set the guest start epoch to the current time, before running the vcpu
-                crate::debug!("HyperV - Guest Start Epoch set");
+                // Store the guest start epoch and cycles to trace the guest execution time
+                crate::debug!("MSHV - Guest Start Epoch set");
+                self.trace_info.guest_start_tsc =
+                    Some(hyperlight_guest_tracing::invariant_tsc::read_tsc());
                 self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
             }
             // Note: if a `InterruptHandle::kill()` called while this thread is **here**

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -734,8 +734,10 @@ impl Hypervisor for HypervWindowsDriver {
         } else {
             #[cfg(feature = "trace_guest")]
             if self.trace_info.guest_start_epoch.is_none() {
-                // Set the guest start epoch to the current time, before running the vcpu
-                crate::debug!("MSHV - Guest Start Epoch set");
+                // Store the guest start epoch and cycles to trace the guest execution time
+                crate::debug!("HyperV - Guest Start Epoch set");
+                self.trace_info.guest_start_tsc =
+                    Some(hyperlight_guest_tracing::invariant_tsc::read_tsc());
                 self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
             }
             self.processor.run()?

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -645,9 +645,11 @@ impl Hypervisor for KVMDriver {
         } else {
             #[cfg(feature = "trace_guest")]
             if self.trace_info.guest_start_epoch.is_none() {
-                // Set the guest start epoch to the current time, before running the vcpu
+                // Store the guest start epoch and cycles to trace the guest execution time
                 crate::debug!("KVM - Guest Start Epoch set");
                 self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
+                self.trace_info.guest_start_tsc =
+                    Some(hyperlight_guest_tracing::invariant_tsc::read_tsc());
             }
 
             // Note: if a `InterruptHandle::kill()` called while this thread is **here**

src/hyperlight_host/src/sandbox/mod.rs

@@ -90,7 +90,7 @@ pub type ExtraAllowedSyscall = i64;
 /// Determine whether a suitable hypervisor is available to run
 /// this sandbox.
 ///
-//  Returns a boolean indicating whether a suitable hypervisor is present.
+///  Returns a boolean indicating whether a suitable hypervisor is present.
 #[instrument(skip_all, parent = Span::current())]
 pub fn is_hypervisor_present() -> bool {
     hypervisor::get_available_hypervisor().is_some()
@@ -103,18 +103,23 @@ pub fn is_hypervisor_present() -> bool {
 pub(crate) struct TraceInfo {
     /// The epoch against which trace events are timed; at least as
     /// early as the creation of the sandbox being traced.
-    #[allow(dead_code)]
     pub epoch: std::time::Instant,
     /// The frequency of the timestamp counter.
-    #[allow(dead_code)]
-    pub tsc_freq: u64,
+    pub tsc_freq: Option<u64>,
     /// The epoch at which the guest started, if it has started.
     /// This is used to calculate the time spent in the guest relative to the
-    /// time of the host.
-    #[allow(dead_code)]
+    /// time when the host started.
     pub guest_start_epoch: Option<std::time::Instant>,
-    /// The start guest time, in TSC cycles, for the current guest.
-    #[allow(dead_code)]
+    /// The start guest time, in TSC cycles, for the current guest has a double purpose.
+    /// This field is used in two ways:
+    /// 1. It contains the TSC value recorded on the host when the guest started.
+    ///    This is used to calculate the TSC frequency which is the same on the host and guest.
+    ///    The TSC frequency is used to convert TSC values to timestamps in the trace.
+    ///    **NOTE**: This is only used until the TSC frequency is calculated, when the first
+    ///    records are received.
+    /// 2. To store the TSC value at recorded on the guest when the guest started (first record
+    ///    received)
+    ///    This is used to calculate the records timestamps relative to when guest started.
     pub guest_start_tsc: Option<u64>,
     /// The file to which the trace is being written
     #[allow(dead_code)]
@@ -139,8 +144,17 @@ impl TraceInfo {
     ) -> crate::Result<Self> {
         let mut path = std::env::current_dir()?;
         path.push("trace");
+
+        // create directory if it does not exist
+        if !path.exists() {
+            std::fs::create_dir(&path)?;
+        }
         path.push(uuid::Uuid::new_v4().to_string());
         path.set_extension("trace");
+
+        log::info!("Creating trace file at: {}", path.display());
+        println!("Creating trace file at: {}", path.display());
+
         #[cfg(feature = "unwind_guest")]
         let hash = unwind_module.hash();
         #[cfg(feature = "unwind_guest")]
@@ -150,11 +164,17 @@ impl TraceInfo {
             let cache = framehop::x86_64::CacheX86_64::new();
             (unwinder, Arc::new(Mutex::new(cache)))
         };
-        let tsc_freq = Self::calculate_tsc_freq()?;
+        if !hyperlight_guest_tracing::invariant_tsc::has_invariant_tsc() {
+            // If the platform does not support invariant TSC, warn the user.
+            // On Azure nested virtualization, the TSC invariant bit is not correctly reported, this is a known issue.
+            log::warn!(
+                "Invariant TSC is not supported on this platform, trace timestamps may be inaccurate"
+            );
+        }
 
         let ret = Self {
             epoch: std::time::Instant::now(),
-            tsc_freq,
+            tsc_freq: None,
             guest_start_epoch: None,
             guest_start_tsc: None,
             file: Arc::new(Mutex::new(std::fs::File::create_new(path)?)),
@@ -173,22 +193,40 @@ impl TraceInfo {
         Ok(ret)
     }
 
-    /// Calculate the TSC frequency based on the RDTSC instruction.
-    fn calculate_tsc_freq() -> crate::Result<u64> {
-        if !hyperlight_guest_tracing::invariant_tsc::has_invariant_tsc() {
-            return Err(crate::new_error!(
-                "Invariant TSC is not supported on this platform"
-            ));
-        }
-        let start = hyperlight_guest_tracing::invariant_tsc::read_tsc();
-        let start_time = std::time::Instant::now();
-        // Sleep for 1 second to get a good sample
-        std::thread::sleep(std::time::Duration::from_secs(1));
-        let end = hyperlight_guest_tracing::invariant_tsc::read_tsc();
+    /// Calculate the TSC frequency based on the RDTSC instruction on the host.
+    fn calculate_tsc_freq(&mut self) -> crate::Result<()> {
+        let (start, start_time) = match (
+            self.guest_start_tsc.as_ref(),
+            self.guest_start_epoch.as_ref(),
+        ) {
+            (Some(start), Some(start_time)) => (*start, *start_time),
+            _ => {
+                // If the guest start TSC and time are not set, we use the current time and TSC.
+                // This is not ideal, but it allows us to calculate the TSC frequency without
+                // failing.
+                // This is a fallback mechanism to ensure that we can still calculate, however it
+                // should be noted that this may lead to inaccuracies in the TSC frequency.
+                // The start time should be already set before running the guest for each sandbox.
+                log::error!(
+                    "Guest start TSC and time are not set. Calculating TSC frequency will use current time and TSC."
+                );
+                (
+                    hyperlight_guest_tracing::invariant_tsc::read_tsc(),
+                    std::time::Instant::now(),
+                )
+            }
+        };
+
         let end_time = std::time::Instant::now();
+        let end = hyperlight_guest_tracing::invariant_tsc::read_tsc();
+
         let elapsed = end_time.duration_since(start_time).as_secs_f64();
+        let tsc_freq = ((end - start) as f64 / elapsed) as u64;
+
+        log::info!("Calculated TSC frequency: {} Hz", tsc_freq);
+        self.tsc_freq = Some(tsc_freq);
 
-        Ok(((end - start) as f64 / elapsed) as u64)
+        Ok(())
     }
 }
 

src/hyperlight_host/src/sandbox/outb.rs

@@ -232,8 +232,13 @@ pub(super) fn record_guest_trace_frame<F: FnOnce(&mut std::fs::File)>(
             .guest_start_tsc
             .as_ref()
             .map_or_else(|| 0, |c| *c);
+
     // Convert cycles to microseconds based on the TSC frequency
-    let micros = cycles_spent as f64 / trace_info.tsc_freq as f64 * 1_000_000f64;
+    let tsc_freq = trace_info
+        .tsc_freq
+        .as_ref()
+        .ok_or_else(|| new_error!("TSC frequency not set in TraceInfo"))?;
+    let micros = cycles_spent as f64 / *tsc_freq as f64 * 1_000_000f64;
 
     // Convert to a Duration
     let guest_duration = std::time::Duration::from_micros(micros as u64);
@@ -353,7 +358,6 @@ fn handle_outb_impl(
                 .as_ref()
                 .shared_mem
                 .copy_to_slice(buffer, ptr as usize - SandboxMemoryLayout::BASE_ADDRESS)
-                // .read::<u64>((addr - SandboxMemoryLayout::BASE_ADDRESS as u64) as usize)
                 .map_err(|e| {
                     new_error!(
                         "Failed to copy trace records from guest memory to host: {:?}",
@@ -367,11 +371,24 @@ fn handle_outb_impl(
 
             {
                 let trace_info = _hv.trace_info_as_mut();
-                // Store the start guest cycles if not already set
-                // This is the `entrypoint` of the guest execution
-                // This should be set only once, at the start of the guest execution
-                if trace_info.guest_start_tsc.is_none() && !traces.is_empty() {
-                    trace_info.guest_start_tsc = Some(traces[0].cycles);
+
+                // Calculate the TSC frequency based on the current TSC reading
+                // This is done only once, when the first trace record is received
+                // Ideally, we should use a timer or a clock to measure the time elapsed,
+                // but that adds delays.
+                // To avoid that we store the TSC value and a timestamp right
+                // before starting the guest execution and then calculate the TSC frequency when
+                // the first trace record is received, based on the current TSC value and clock.
+                if trace_info.tsc_freq.is_none() {
+                    trace_info.calculate_tsc_freq()?;
+
+                    // After the TSC frequency is calculated, we no longer need the value of TSC
+                    // recorded on the host when the guest started, so we can set the guest_start_tsc field
+                    // to store the TSC value recorded on the guest when the guest started executing.
+                    // This is used to calculate the records timestamps relative to the first trace record.
+                    if !traces.is_empty() {
+                        trace_info.guest_start_tsc = Some(traces[0].cycles);
+                    }
                 }
             }