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

60 intervals that take 1 s each (a sleep). Run them sequentially, in parallel
on 4 workers, and on 4 workers with a depends_on_past rule (an interval may
start only after the previous one succeeded). Then make interval 10 fail and
count the intervals that run in each parallel mode.
"""
import argparse, json, threading, time
from concurrent.futures import ThreadPoolExecutor

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--intervals", type=int, default=60, help="Intervals.")
p.add_argument("--seconds", type=float, default=1.0, help="Seconds per interval.")
p.add_argument("--workers", type=int, default=4, help="Workers.")
p.add_argument("--fail", type=int, default=10, help="Interval that fails in the failure runs.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def work(i, fail):
    time.sleep(args.seconds)
    if i == fail: raise RuntimeError(f"interval {i} failed")

def run(mode, fail=-1):
    t0 = time.perf_counter(); ran = 0
    if mode == "sequential":
        for i in range(args.intervals):
            try: work(i, fail); ran += 1
            except RuntimeError: break
    elif mode == "parallel":
        with ThreadPoolExecutor(args.workers) as ex:
            futs = [ex.submit(work, i, fail) for i in range(args.intervals)]
            ran = sum(1 for f in futs if f.exception() is None)
    else:                                                              # depends_on_past: interval i waits for i - 1 to succeed
        ok = [threading.Event() for _ in range(args.intervals)]; failed = threading.Event(); count = [0]; lock = threading.Lock()
        def task(i):
            if i > 0:
                while not ok[i - 1].wait(0.01):
                    if failed.is_set(): return                         # blocked forever: give up
            try: work(i, fail); ok[i].set()
            except RuntimeError: failed.set(); return
            with lock: count[0] += 1
        with ThreadPoolExecutor(args.workers) as ex: list(ex.map(task, range(args.intervals)))
        ran = count[0]
    return {"mode": mode, "fail": "interval %d fails" % fail if fail >= 0 else "none", "wall_s": time.perf_counter() - t0, "succeeded": ran}

rows = [run("sequential"), run("parallel"), run("depends_on_past"), run("parallel", args.fail), run("depends_on_past", args.fail)]
for r in rows: print(f"{r['mode']:<16} failure: {r['fail']:<18} wall {r['wall_s']:5.1f} s   intervals succeeded {r['succeeded']}")

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