Watermark with an idle partition

# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""A watermark with an idle partition.

Three partitions send one event per second each, with up to 0.2 minutes of
skew. Partition 3 goes quiet at minute 10; at minute 20 it sends the events it
held back (event times from minute 10 to 20), then carries on. The global
watermark is the minimum of the partition watermarks (max event time seen
minus a delay). Records how late each 1-minute window closes, without an
idleness timeout and with one, and counts events that arrive after their
window closed.
"""
import argparse, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--minutes", type=int, default=30, help="Simulated minutes.")
p.add_argument("--idle-at", type=int, default=10, help="Minute at which partition 3 goes quiet.")
p.add_argument("--resume-at", type=int, default=20, help="Minute at which partition 3 sends its backlog.")
p.add_argument("--timeout", type=float, default=2.0, help="Idleness timeout (minutes).")
p.add_argument("--delay", type=float, default=0.5, help="Watermark delay (minutes).")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def run(timeout):
    rng = random.Random(args.seed)
    seen = {0: 0.0, 1: 0.0, 2: 0.0}; last = {0: 0.0, 1: 0.0, 2: 0.0}
    closed = 0; late = 0; delays = []; backlog_sent = False
    for s in range(args.minutes * 60):
        now = s / 60; arrivals = []
        for part in (0, 1, 2):
            if part == 2 and args.idle_at <= now < args.resume_at: continue                 # quiet
            arrivals.append((part, max(0.0, now - rng.uniform(0, 0.2))))
        if now >= args.resume_at and not backlog_sent:                                       # the held-back events, all at once
            arrivals += [(2, args.idle_at + i / 60) for i in range((args.resume_at - args.idle_at) * 60)]; backlog_sent = True
        for part, ev in arrivals:
            if ev < closed: late += 1                                                        # its window has already closed
            seen[part] = max(seen[part], ev); last[part] = now
        active = [seen[p_] for p_ in seen if timeout is None or now - last[p_] <= timeout]
        wm = min(active) - args.delay
        while closed + 1 <= wm: closed += 1; delays.append(now - closed)                     # window [closed-1, closed) closes
    return {"idleness_timeout": f"{timeout:g} min" if timeout is not None else "none", "windows_closed": closed,
            "longest_close_delay_min": max(delays), "late_events": late}
rows = [run(None), run(args.timeout)]
for r in rows: print(f"timeout {r['idleness_timeout']:<6}: {r['windows_closed']} windows closed, longest wait past a window's end {r['longest_close_delay_min']:.1f} min, late events {r['late_events']}")

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