# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The pointer swap: two writers, one table.

A table commit on the local file system: each version is a file vN.json and
the "pointer" is the highest version number. Creating v(N+1).json with the
exclusive-create flag acts as a compare-and-swap. Two writer threads commit 50
changes each and retry when they lose. Then the same with a plain write
(create-or-overwrite) in place of the exclusive create.
"""
import argparse, json, os, tempfile, threading

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--commits", type=int, default=50, help="Commits per writer.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def commit(table_dir, base_version, change, exclusive):
    new_version = base_version + 1
    path = os.path.join(table_dir, f"v{new_version}.json")
    flags = os.O_CREAT | os.O_WRONLY | (os.O_EXCL if exclusive else os.O_TRUNC)
    try:
        fd = os.open(path, flags)                                          # with O_EXCL: fails if the file exists
    except FileExistsError:
        return False                                                       # lost the race
    with os.fdopen(fd, "w") as f:
        json.dump({"version": new_version, "change": change}, f)
    return True

def latest(table_dir):
    return max(int(n[1:-5]) for n in os.listdir(table_dir) if n.startswith("v") and n.endswith(".json"))

def writer(table_dir, name, results, retries, exclusive):
    for i in range(args.commits):
        while True:
            base = latest(table_dir)
            if commit(table_dir, base, f"{name}-{i}", exclusive):
                results[name] += 1
                break                                                      # else: retry on top of the new latest
            retries[name] += 1

rows = []
for exclusive, label in ((True, "exclusive create (compare-and-swap)"), (False, "plain write (create or overwrite)")):
    d = tempfile.mkdtemp(); open(os.path.join(d, "v0.json"), "w").write("{}")
    results, retries = {"A": 0, "B": 0}, {"A": 0, "B": 0}
    ts = [threading.Thread(target=writer, args=(d, n, results, retries, exclusive)) for n in "AB"]
    [t.start() for t in ts]; [t.join() for t in ts]
    stored = []
    for n in os.listdir(d):
        if n != "v0.json":
            try: stored.append(json.load(open(os.path.join(d, n)))["change"])
            except (json.JSONDecodeError, KeyError): stored.append("<torn>")
    claimed = results["A"] + results["B"]
    rows.append({"commit": label, "claimed": claimed, "latest": latest(d), "stored": len(set(stored) - {"<torn>"}),
                 "lost": claimed - len(set(stored) - {"<torn>"}), "retries": retries["A"] + retries["B"]})
for r in rows: print(f"{r['commit']:<38} writers think they committed {r['claimed']}; latest = v{r['latest']}; changes stored {r['stored']}; lost {r['lost']}; retries {r['retries']}")

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