# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""What a rebalance repeats: eager, cooperative and static membership.

A consumer group reads a partitioned topic and commits each partition's offset on
a fixed interval. One consumer is redeployed. A partition that changes owner restarts
from its last committed offset, so everything processed since that commit is processed
again. Eager rebalancing revokes every partition; cooperative revokes only the ones
that move; static membership with a restart inside the session timeout moves nothing.
Optionally, offsets are committed when partitions are revoked. Each deploy happens at a
random moment in the commit cycle; the script repeats it many times.
"""
import argparse, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--consumers", type=int, default=12, help="Consumers in the group.")
p.add_argument("--partitions", type=int, default=24, help="Partitions of the topic.")
p.add_argument("--rate", type=float, default=400, help="Records per second per partition.")
p.add_argument("--commit", type=float, default=5.0, help="Seconds between offset commits.")
p.add_argument("--deploys", type=int, default=10000, help="Simulated deploys (each at a random moment).")
p.add_argument("--seed", type=int, default=1, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

per_consumer = args.partitions // args.consumers


def deploy(rnd, protocol, commit_on_revoke):
    # each consumer commits its partitions together, at its own phase of the commit cycle
    since_commit = [rnd.uniform(0, args.commit) for _ in range(args.consumers)]
    if protocol == "static membership":
        return 0, 0
    moved = range(args.consumers) if protocol == "eager" else [rnd.randrange(args.consumers)]
    paused = per_consumer * len(moved)
    if commit_on_revoke:
        # the redeployed consumer is being stopped cleanly too, so it commits before leaving
        return 0, paused
    return sum(since_commit[c] * args.rate * per_consumer for c in moved), paused


rows = []
for revoke in (False, True):
    for protocol in ("eager", "cooperative", "static membership"):
        rnd = random.Random(args.seed)
        res = [deploy(rnd, protocol, revoke) for _ in range(args.deploys)]
        reps = sorted(r[0] for r in res)
        rows.append({"protocol": protocol, "commit_on_revoke": "yes" if revoke else "no",
                     "partitions_paused": res[0][1], "mean_repeats": sum(reps) / len(reps),
                     "p99_repeats": reps[int(len(reps) * 0.99)], "max_possible": (args.partitions if protocol == "eager" else per_consumer if protocol == "cooperative" else 0) * args.rate * args.commit * (0 if revoke else 1)})
for r in rows:
    print(f"{r['protocol']:<18} commit on revoke {r['commit_on_revoke']:<3}  partitions paused {r['partitions_paused']:>2}  "
          f"repeats per deploy: mean {r['mean_repeats']:8,.0f}  p99 {r['p99_repeats']:8,.0f}  bound {r['max_possible']:8,.0f}")

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