kala 2.0.0

#![cfg(target_arch = "x86_64")]

use crate::arch::TscFrequency;
use std::sync::LazyLock;

pub fn tick_frequency() -> TscFrequency {
  static FREQUENCY: LazyLock<TscFrequency> =
    LazyLock::new(freq::calculate_frequency);
  *FREQUENCY
}

pub fn is_hw_supported() -> bool {
  static SUPPORTED: LazyLock<bool> = LazyLock::new(|| {
    // SSE2 is part of x86-64-v1 that rustc targets by default:
    //   https://en.wikipedia.org/wiki/X86-64#Microarchitecture_levels
    // RDTSC is part of x86-64-v1.

    // We verify that the CPU's TSC frequency is _invariant_, i.e. does not
    // dynamically change as the core clock changes.
    raw_cpuid::CpuId::new()
      .get_advanced_power_mgmt_info()
      .is_some_and(|apm| apm.has_invariant_tsc())
  });

  *SUPPORTED
}

#[inline(always)]
pub fn clock_fence() {
  #[target_feature(enable = "sse2")]
  #[inline]
  fn inner() {
    core::arch::x86_64::_mm_lfence();
  }
  // SAFETY: Rustc targets x86-64-v1 which includes SSE2
  unsafe { inner() }
}

#[inline(always)]
fn ticks() -> u64 {
  // SAFETY: Rustc targets x86-64-v1 which includes RDTSC
  unsafe { core::arch::x86_64::_rdtsc() }
}

#[inline(always)]
pub fn now() -> crate::time::Instant {
  crate::time::Instant::Ticks(ticks())
}

mod freq {
  use crate::arch::TscFrequency;

  use super::*;
  use std::thread::sleep;
  use std::time::{Duration, Instant};

  #[derive(Debug, Copy, Clone)]
  struct Time {
    ticks: u64,         // The (estimated) TSC tick value for "now"
    runtime_ticks: u64, // The number of ticks it took to get the current time
    instant: Instant,   // The reference clock value for "now"
  }

  // Returns the current time from both the TSC clock and the reference
  // `Instant::now()`.
  fn get_time() -> Time {
    clock_fence();
    let first = ticks();
    let instant = Instant::now();
    let second = ticks();
    clock_fence();

    // `saturating_sub` in case of weird thread migration
    let runtime_ticks = second.saturating_sub(first);

    // We get the in-between value for ticks to account for the `Instant::now()`
    // call overhead. It's not perfect, but better than not doing it.
    let avg_ticks = first + runtime_ticks / 2; // calculate avg without overflow

    Time {
      ticks: avg_ticks,
      runtime_ticks,
      instant,
    }
  }

  // Runs `get_time()` in a short loop to get the shortest runtime and thus most
  // precise number of ticks for `get_time()`.
  //
  // Useful because `Time::runtime_ticks` can vary by 10x between calls to
  // `get_time()`. The first call right after `sleep()` is quite slow, likely
  // due to cold CPU cache.
  //
  // This approach also handles the rare edge case where the thread is
  // pre-empted inside the timed section of `get_time()` and a bonkers value is
  // returned.
  fn best_time() -> Time {
    const MAX_ITERS: usize = 10;
    let mut best = get_time();

    // Start at 1 because we called `get_time()` once already.
    for _ in 1..MAX_ITERS {
      let current = get_time();
      if current.runtime_ticks > 0 && current.runtime_ticks < best.runtime_ticks
      {
        best = current;
      }
    }
    best
  }

  // Estimate the frequency (in Hz) of the TSC clock. Keep in mind that we
  // insist on CPUs with _invariant_ (constant) TSC clocks in
  // `is_hw_supported()`.
  //
  // We only ever call this function if the TSC frequency metadata in CPUID is
  // missing. Modern Intel CPUs should have the metadata, but it's ~always
  // missing on AMD.
  fn estimate_frequency() -> TscFrequency {
    let start = best_time();
    // Longer sleep amounts improve estimate precision.
    sleep(Duration::from_millis(100));
    let end = best_time();

    // I can't see how this subtraction could ever underflow, but I'm using
    // `strict_sub` just in case; better to surface a panic than incorrect data.
    // Please file a bug report with detailed info if it ever happens in the
    // wild.
    let delta_ticks = end.ticks.strict_sub(start.ticks);
    let delta_secs = (end.instant - start.instant).as_secs_f64();
    TscFrequency::Estimate((delta_ticks as f64 / delta_secs).round() as u64)
  }

  // Returns tick (TSC) frequency in Hz. Checks CPUID first, then falls back to
  // estimating the frequency.
  pub fn calculate_frequency() -> TscFrequency {
    raw_cpuid::CpuId::new()
      .get_tsc_info()
      .and_then(|tsc| tsc.tsc_frequency()) // NOTE: Only works on Intel!
      .map(TscFrequency::CpuMetadata)
      .unwrap_or_else(estimate_frequency) // AMD path.
  }

  #[cfg(test)]
  mod tests {
    use super::*;
    use crate::test_utils::assert_roughly_eq;

    #[test]
    fn estimated_frequency_matches_cpu_metadata() {
      // When `calculate_frequency()` returns `TscFrequency::Estimate`, then
      // `estimate_frequency()` is tested through `Stopwatch` tests.
      // When it returns `TscFrequency::CpuMetadata`, this test covers
      // `estimate_frequency()`. Thus we always have test coverage for it,
      // regardless of the host CPU.

      let TscFrequency::CpuMetadata(official_freq) = calculate_frequency()
      else {
        return; // Skip the test. (See big comment above.)
      };

      let TscFrequency::Estimate(estimated_freq) = estimate_frequency() else {
        unreachable!("must be an estimated freq");
      };

      // Our estimate should match the official frequency to within 1%. That's a
      // loose bound to prevent flakiness; in reality our estimate should be
      // _much_ tighter than that (likely to ~100ppm).
      assert_roughly_eq(estimated_freq as f64, official_freq as f64, 0.01);
    }
  }
}