# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Aligned against unaligned checkpoints (a model).

Not a real stream engine: a tick-by-tick model of a source, a network buffer
of bounded size (backpressure) and a sink, with checkpoint barriers injected
every 10 seconds. An aligned barrier travels in order behind the records
already buffered, so the checkpoint completes when the sink has drained them;
an unaligned barrier overtakes the buffer and the checkpoint stores the
buffered (in-flight) records instead. Checkpoint duration and size as the
sink slows down.
"""
import argparse, json
from collections import deque

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--source-rate", type=int, default=1000, help="Records per second the source can emit.")
p.add_argument("--sink-rates", default="2000,1000,500,250,100", help="Records per second the sink can write.")
p.add_argument("--buffer", type=int, default=20_000, help="In-flight records the network buffers hold.")
p.add_argument("--state-kb", type=float, default=512, help="Operator state size (KB).")
p.add_argument("--record-bytes", type=int, default=200, help="Bytes per in-flight record.")
p.add_argument("--seconds", type=int, default=300, help="Simulated seconds.")
p.add_argument("--interval", type=int, default=10, help="Checkpoint interval (s).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

TICK = 10                                                                   # ticks per second
def sim(sink_rate, unaligned):
    q = deque(); durations, sizes = [], []; pending = []                    # pending: (inject time, records ahead)
    src, snk = args.source_rate / TICK, sink_rate / TICK; acc_src = acc_snk = 0.0
    for tick in range(args.seconds * TICK):
        t = tick / TICK
        if tick % (args.interval * TICK) == 0 and tick:
            if unaligned:                                                    # overtake: snapshot the buffered records, done at once
                durations.append(1 / TICK); sizes.append(args.state_kb + len(q) * args.record_bytes / 1024)
            else:
                pending.append([t, len(q)])                                  # wait behind everything already buffered
        acc_src += src
        while acc_src >= 1 and len(q) < args.buffer: q.append(1); acc_src -= 1
        acc_src = min(acc_src, 1.0)                                          # backpressure: the source waits
        acc_snk += snk; drained = 0
        while acc_snk >= 1 and q: q.popleft(); acc_snk -= 1; drained += 1
        acc_snk = min(acc_snk, 1.0)
        for p_ in pending: p_[1] -= drained
        for p_ in [p_ for p_ in pending if p_[1] <= 0]:
            durations.append(t + 1 / TICK - p_[0]); sizes.append(args.state_kb); pending.remove(p_)
    tail = durations[len(durations) // 2:] or [float("nan")]; stail = sizes[len(sizes) // 2:] or [float("nan")]
    return sum(tail) / len(tail), sum(stail) / len(stail), len(durations)
rows = []
for s in map(int, args.sink_rates.split(",")):
    for unaligned in (False, True):
        d, size, n = sim(s, unaligned)
        rows.append({"sink_rate": s, "mode": "unaligned" if unaligned else "aligned", "duration_s": d, "size_kb": size, "completed": n})
for r in rows: print(f"sink {r['sink_rate']:>5}/s  {r['mode']:<9} checkpoint takes {r['duration_s']:6.1f} s, size {r['size_kb']:8.0f} KB  ({r['completed']} completed)")

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