# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Differential profile of two builds.

A small program parses JSON records and inserts them into a hash map. Build B
makes one small change: a lock around the record counter. Both are profiled
for the same time under the same load by a sampling profiler (a CPU-time
timer signal that records the Python stack), then diffed frame by frame. The
diff is repeated with B profiled while CPU-heavy processes run alongside.
Writes folded stacks (the input format of flamegraph.pl) next to --json.
"""
import argparse, collections, json, multiprocessing as mp, os, signal, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--seconds", type=float, default=30, help="Profiled seconds per run.")
p.add_argument("--hz", type=int, default=999, help="Sampling rate.")
p.add_argument("--load-procs", type=int, default=8, help="CPU-heavy processes in the different-load run.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

RECORDS = [json.dumps({"id": i, "user": f"u{i % 997}", "amount": i * 1.5, "tags": ["a", "b", str(i % 7)]}) for i in range(20_000)]
lock = threading.Lock(); n = 0; LOCKED = False

def parse(s): return json.loads(s)
def insert(table, rec): table[rec["user"]] = table.get(rec["user"], 0) + rec["amount"]
def count():
    global n
    if LOCKED:
        with lock: n += 1                                          # build B: the one change
    else: n += 1

def workload(seconds):
    table = {}; end = time.process_time() + seconds
    while time.process_time() < end:
        for s in RECORDS[:2000]:
            insert(table, parse(s)); count()

stacks = collections.Counter()
def sample(signum, frame):
    names = []
    while frame is not None:
        names.append(frame.f_code.co_name); frame = frame.f_back
    stacks[";".join(reversed(names))] += 1

def burn(stop):
    x = 0
    while not stop.is_set(): x += 1

def profile(locked, load=0):
    global LOCKED
    LOCKED = locked; stacks.clear(); stop = mp.Event(); procs = [mp.Process(target=burn, args=(stop,)) for _ in range(load)]
    for q in procs: q.start()
    signal.signal(signal.SIGPROF, sample); signal.setitimer(signal.ITIMER_PROF, 1 / args.hz, 1 / args.hz)
    workload(args.seconds)
    signal.setitimer(signal.ITIMER_PROF, 0); stop.set()
    for q in procs: q.join()
    return collections.Counter(stacks)

def leaf_shares(folded):
    """Share of samples per frame (self time): the width of each frame's top edge in a flamegraph."""
    total = sum(folded.values()); c = collections.Counter()
    for st, n in folded.items(): c[st.split(";")[-1]] += n
    return {k: v / total * 100 for k, v in c.items()}

def diff(a, b, name):
    sa, sb = leaf_shares(a), leaf_shares(b)
    rows = [{"diff": name, "frame": f, "share_a": sa.get(f, 0), "share_b": sb.get(f, 0), "delta_pts": sb.get(f, 0) - sa.get(f, 0)}
            for f in set(sa) | set(sb)]
    return sorted(rows, key=lambda r: -abs(r["delta_pts"]))[:8]

if __name__ == "__main__":
    a = profile(False); b = profile(True); a2 = profile(False); b_load = profile(True, args.load_procs)
    same, noise, loaded = diff(a, b, "A vs B, same load"), diff(a, a2, "A vs A (no change)"), diff(a, b_load, "A vs B, B under load")
    for title, rows in (("A vs B, same load", same), ("A vs A, no change", noise), ("A vs B, B under extra load", loaded)):
        print(title)
        for r in rows[:5]: print(f"  {r['frame']:<18} {r['share_a']:6.2f}% -> {r['share_b']:6.2f}%  ({r['delta_pts']:+.2f} pts)")
    if args.json:
        base = os.path.splitext(args.json)[0]
        for tag, folded in (("a", a), ("b", b), ("b-loaded", b_load)):
            with open(f"{base}.{tag}.folded", "w") as f: f.writelines(f"{k} {v}\n" for k, v in folded.items())
        with open(args.json, "w") as f:
            json.dump({"tables": {"same": same, "noise": noise, "loaded": loaded}}, f)