# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""A micro-benchmark that disappears, and a profiler's blind spot.

Part A: a C loop computes x * x for 1 billion values, first never using the
result, then summing it and printing the sum (compiled with -O2; needs cc).
Part B: a program sleeps 2 ms about once a second and does busy work otherwise,
while a timer-signal sampler records what it is doing at 99 Hz and at 999 Hz.
Counts the samples that land in the sleep, and the busy work lost to sampling.
"""
import argparse, json, os, random, signal, subprocess, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--iter", type=int, default=1_000_000_000, help="Part A: loop iterations.")
p.add_argument("--seconds", type=int, default=60, help="Part B: profiled seconds per sampling rate.")
p.add_argument("--sleep-ms", type=float, default=2, help="Part B: length of the once-a-second sleep.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

C = r"""
#include <stdio.h>
#include <time.h>
int main(void){
  struct timespec a,b; clock_gettime(CLOCK_MONOTONIC,&a);
  unsigned long sum = 0;
  for (unsigned long x = 0; x < ITER; x++) {
#ifdef OBSERVE
    sum += x * x;
    __asm__ volatile("" : "+r"(sum));          /* keeps the compiler from replacing the loop with a formula */
#else
    unsigned long y = x * x; (void)y;          /* never used */
#endif
  }
  clock_gettime(CLOCK_MONOTONIC,&b);
  double s = (b.tv_sec-a.tv_sec)+(b.tv_nsec-a.tv_nsec)/1e9;
  fprintf(stderr, "sum=%lu\n", sum);
  printf("%.9f\n", s);
}
"""
d = tempfile.mkdtemp(); src = os.path.join(d, "square.c"); open(src, "w").write(C); part_a = []
for name, flags in (("result never used", []), ("sum printed", ["-DOBSERVE"])):
    exe = os.path.join(d, "observe" if flags else "unused")
    subprocess.run(["cc", "-O2", f"-DITER={args.iter}UL", *flags, src, "-o", exe], check=True)
    sec = float(subprocess.run([exe], capture_output=True, text=True, check=True).stdout)
    part_a.append({"variant": name, "seconds": sec, "ns_per_iter": sec / args.iter * 1e9})

sleeping = False; samples = {"total": 0, "sleep": 0}; hit_events = set(); event = 0
def sample(signum, frame):
    samples["total"] += 1
    if sleeping: samples["sleep"] += 1; hit_events.add(event)

def profiled(hz):
    global sleeping, event
    samples.update(total=0, sleep=0); hit_events.clear(); event = 0; work = 0
    if hz: signal.signal(signal.SIGALRM, sample); signal.setitimer(signal.ITIMER_REAL, 1 / hz, 1 / hz)
    start = time.perf_counter(); next_sleep = start + 1
    while (now := time.perf_counter()) < start + args.seconds:
        if now >= next_sleep:
            event += 1; sleeping = True; time.sleep(args.sleep_ms / 1000); sleeping = False; next_sleep += random.uniform(0.95, 1.05)   # about once a second: a fixed 1 s would run in lockstep with the timer
        for _ in range(1000): work += 1                                    # busy work
    signal.setitimer(signal.ITIMER_REAL, 0)
    return {"rate": f"{hz} Hz" if hz else "no profiler", "hz": hz, "samples": samples["total"], "in_sleep": samples["sleep"],
            "expected_in_sleep": hz * event * args.sleep_ms / 1000, "sleeps": event, "sleeps_seen": len(hit_events), "work": work}

part_b = [profiled(0), profiled(99), profiled(999)]
for r in part_b: r["overhead_pct"] = (1 - r["work"] / part_b[0]["work"]) * 100

for r in part_a: print(f"A  {r['variant']:<18} {r['seconds']:.4f} s  ({r['ns_per_iter']:.4f} ns per iteration)")
for r in part_b[1:]:
    print(f"B  {r['rate']:>6}: {r['samples']} samples, {r['in_sleep']} in the sleep (expected {r['expected_in_sleep']:.1f}); "
          f"{r['sleeps_seen']} of {r['sleeps']} sleeps seen at all; busy work {r['overhead_pct']:.1f}% lower than unprofiled")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"part_a": part_a, "part_b": part_b}}, f)