/*
 * vcpu_steal
 *
 * measures actual vcpu steal time from a host that's running as a
 * virtual guest inside of kvm, like a vps from most common providers.
 * we do this by busy-looping a cycle counter on a singular pinned vcpu,
 * trying to find gaps between any iteration that count as a preemption
 * by the hypervisor.
 *
 * usually the actual steal time is not exposed on cheap vps providers,
 * and even more expensive providers, and the value in `/proc/stat`
 * which should represent it is often just empty.
 * since we fully pin a single cpu and run a busy loop, the only events
 * that can interrupt us are guest-kernel preemption (timer ticks, irqs)
 * and the host stealing our vcpu. the guest-kernel events are short and
 * usually sub-microsecond or just a little higher, so our 5us threshold
 * for found gaps filters them out as noise, leaving the host preemption
 * as the primary thing producing the longer gaps we measure.
 * this results in an error of +1-2% at the absolute most.
 *
 * reports the cycle counter gap value, alongside the kernel wall vs.
 * cpu time for comparison.
 * note that the guest kernel doesn't really know what's going on and
 * thinks that any time stolen from a task by the hypervisor belongs
 * to the running task, thus this is only here to contrast the tested
 * value we compute.
 *
 * usage: ./vcpu_steal [seconds] [cpu]
 *
 * Copyright (c) 2026, Mel G. <mel@rnrd.eu>
 *
 * SPDX-License-Identifier: MPL-2.0
 */

#define _GNU_SOURCE

#include <sched.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <unistd.h>

// recognition threshold for a preemption gap in nanoseconds.
#define GAP_THRESHOLD_NS 5000.0

// gap-size histogram buckets in microseconds.
// last bucket collects all higher values.
#define HISTOGRAM_BUCKET_COUNT 8
static const double HISTOGRAM_BUCKETS_US[HISTOGRAM_BUCKET_COUNT - 1] = {
    10, 50, 100, 500, 1000, 5000, 10000,
};
static const char* HISTOGRAM_LABELS[HISTOGRAM_BUCKET_COUNT] = {
    "<10us", "10-50us", "50-100us", "100-500us",
    "500us-1ms", "1-5ms", "5-10ms", ">=10ms",
};

// reads a monotonic cycle counter that keeps advancing while the host
// runs someone else on our vCPU -- which is exactly what we want for
// steal detection. on x86 we use rdtscp because it serializes prior
// loads and stores past itself; on aarch64 we use cntvct_el0 with an
// isb in front for the same reason.
// reads monotonic cycle counter for our current vcpu,
// depending on the architecture.
#if defined(__x86_64__) || defined(__i386__)
static inline uint64_t
cycles_now(void)
{
    uint32_t lo, hi, aux;
    // serializes prior loads and stores past itself
    __asm__ volatile("rdtscp" : "=a"(lo), "=d"(hi), "=c"(aux));
    return ((uint64_t)hi << 32) | lo;
}
#elif defined(__aarch64__)
static inline uint64_t
cycles_now(void)
{
    uint64_t v;
    // w/ isb to match x86 behavior of prior synchronization
    __asm__ volatile("isb; mrs %0, cntvct_el0" : "=r"(v));
    return v;
}
#else
#error "no cycle counter on this architecture, sorry :("
#endif

double
timespec_diff_ns(struct timespec a, struct timespec b)
{
    return (b.tv_sec - a.tv_sec) * 1e9 + (b.tv_nsec - a.tv_nsec);
}

struct Calibration {
    double cycles_per_second;
};

// calibrates cycle counter to find out rough time of cycles per second.
struct Calibration
calibrate(void)
{
    struct timespec t0, t1;
    struct timespec hundred_ms = { .tv_sec = 0, .tv_nsec = 100 * 1000 * 1000 };

    // both cycle and monotonic counters tick during preemption,
    // making calibration work even if the vcpu gets stolen in the meantime!
    clock_gettime(CLOCK_MONOTONIC, &t0);
    uint64_t c0 = cycles_now();
    nanosleep(&hundred_ms, NULL);
    clock_gettime(CLOCK_MONOTONIC, &t1);
    uint64_t c1 = cycles_now();

    return (struct Calibration){
        .cycles_per_second = (double)(c1 - c0) / (timespec_diff_ns(t0, t1) / 1e9),
    };
}

void
pin_to_cpu(int cpu)
{
    cpu_set_t set;
    CPU_ZERO(&set);
    CPU_SET(cpu, &set);
    if (sched_setaffinity(0, sizeof(set), &set) != 0) {
        perror("couldn't pin to cpu, sorry :(");
        exit(EXIT_FAILURE);
    }
}

struct Measurement {
    // kernel values for cpu time and wall time
    // if kernel can't access kvm steal-time msr, the difference should be near zero.
    double wall_ns;
    double cpu_ns;

    // computed steal, as detected by cycle count gap detection
    uint64_t gap_count;
    double gap_total_ns;
    double gap_max_ns;
    uint64_t gap_histogram[HISTOGRAM_BUCKET_COUNT];

    uint64_t iterations;
};

// finds bucket for a cycle count cap.
uint32_t
histogram_bucket_for(double gap_us)
{
    // start from largest catch-all and find last fitting bucket
    for (uint32_t b = 0; b < HISTOGRAM_BUCKET_COUNT - 1; b++)
        if (gap_us < HISTOGRAM_BUCKETS_US[b]) return b;
    return HISTOGRAM_BUCKET_COUNT - 1;
}

// runs gap detection for a given amount of seconds.
struct Measurement
measure(double seconds, struct Calibration calibration)
{
    struct Measurement m = { 0 };

    struct timespec wall_start, wall_end, cpu_start, cpu_end;
    clock_gettime(CLOCK_MONOTONIC, &wall_start);
    clock_gettime(CLOCK_THREAD_CPUTIME_ID, &cpu_start);

    uint64_t start_cycles = cycles_now();
    uint64_t end_cycles = start_cycles + (uint64_t)(calibration.cycles_per_second * seconds);
    uint64_t prev = start_cycles;

    while (1) {
        uint64_t now = cycles_now();
        double gap_ns = (double)(now - prev) * 1e9 / calibration.cycles_per_second;
        // any gap larger than threshold counts as stolen time
        if (gap_ns > GAP_THRESHOLD_NS) {
            m.gap_count++;
            m.gap_total_ns += gap_ns;
            if (gap_ns > m.gap_max_ns) m.gap_max_ns = gap_ns;
            m.gap_histogram[histogram_bucket_for(gap_ns / 1000.0)]++;
        }
        prev = now;
        m.iterations++;
        if (now >= end_cycles) break;
    }

    clock_gettime(CLOCK_MONOTONIC, &wall_end);
    clock_gettime(CLOCK_THREAD_CPUTIME_ID, &cpu_end);

    m.wall_ns = timespec_diff_ns(wall_start, wall_end);
    m.cpu_ns = timespec_diff_ns(cpu_start, cpu_end);
    return m;
}

void
report(struct Measurement m, struct Calibration cal)
{
    printf("ran at %.3f GHz, for %lu iterations!\n\n",
           cal.cycles_per_second / 1e9, m.iterations);

    // what the kernel thinks
    // only works if kernel has access to kvm steal-time msr.
    double kernel_steal_ns = m.wall_ns - m.cpu_ns;
    printf("kernel says:   %.3f%% stolen (wall=%.2fms, cpu=%.2fms, diff=%.2fms)\n",
           kernel_steal_ns * 100.0 / m.wall_ns,
           m.wall_ns / 1e6,
           m.cpu_ns / 1e6,
           kernel_steal_ns / 1e6);

    // what we have detected
    // how much the cycle counter advanced without us seeing progress on our loop.
    printf("we say:        %.3f%% stolen (gaps=%lu, total=%.2fms, max=%.3fms)\n",
           m.gap_total_ns * 100.0 / m.wall_ns,
           m.gap_count,
           m.gap_total_ns / 1e6,
           m.gap_max_ns / 1e6);

    printf("\nhistogram:\n");
    for (uint32_t b = 0; b < HISTOGRAM_BUCKET_COUNT; b++)
        printf("    %-9s = %6lu\n", HISTOGRAM_LABELS[b], m.gap_histogram[b]);
    printf("\n");
}

struct Arguments {
    double seconds;
    int cpu;
};

struct Arguments
parse_arguments(int argc, char* argv[])
{
    struct Arguments a = { .seconds = 10.0, .cpu = 0 };
    if (argc > 1) a.seconds = atof(argv[1]);
    if (argc > 2) a.cpu = atoi(argv[2]);
    return a;
}

void
usage(const char* program)
{
    fprintf(stderr,
            "usage: %s [seconds] [cpu]\n"
            "\n"
            "    seconds  how long to probe for (default: 10.0)\n"
            "    cpu      which cpu to pin to     (default: 0)\n",
            program);
}

int
main(int argc, char* argv[])
{
    if (argc > 1 && (argv[1][0] == '-' || argv[1][0] == 'h')) {
        usage(argv[0]);
        return EXIT_SUCCESS;
    }

    struct Arguments args = parse_arguments(argc, argv);
    pin_to_cpu(args.cpu);

    struct Calibration calibration = calibrate();
    struct Measurement measurement = measure(args.seconds, calibration);
    report(measurement, calibration);

    return EXIT_SUCCESS;
}