# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The race between a reader and a writer: update versus delete.

Two writers race on one key with "write the store, then update the cache"
and with "write the store, then delete the cache key"; then a reader that
misses and repopulates races a writer (the populate-after-invalidate race).
"""
import argparse, json, random, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--trials", type=int, default=10_000, help="Trials per experiment.")
p.add_argument("--max-sleep-ms", type=float, default=5.0, help="Largest random sleep between the steps of an operation.")
p.add_argument("--ttl", type=float, default=1.0, help="TTL bounding a stale entry, in seconds.")
p.add_argument("--ttl-checks", type=int, default=20, help="Stale entries followed until their TTL expires.")
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()
random.seed(args.seed)
pause = lambda: time.sleep(random.uniform(0, args.max_sleep_ms) / 1000)


def two_writers(policy):
    """One trial: W1 writes v1 and W2 writes v2 at the same time. Returns True if cache and store disagree."""
    store, cache = {"k": "v0"}, {"k": "v0"}
    def writer(v):
        store["k"] = v; pause()
        if policy == "update": cache["k"] = v
        else: cache.pop("k", None)
    a, b = threading.Thread(target=writer, args=("v1",)), threading.Thread(target=writer, args=("v2",))
    a.start(); b.start(); a.join(); b.join()
    if policy == "delete" and "k" not in cache:          # the next read repopulates from the store
        cache["k"] = store["k"]
    return cache["k"] != store["k"]


def reader_vs_writer():
    """A reader misses, reads the store, pauses, then populates; a writer writes and deletes meanwhile.
    Returns the cache dict (with its write time) if a stale value was left behind."""
    store, cache = {"k": "v0"}, {}
    def reader():
        v = store["k"]; pause()                         # miss -> read the store -> slow populate
        cache["k"] = (v, time.time())
    def writer():
        pause(); store["k"] = "v1"; cache.pop("k", None)
    a, b = threading.Thread(target=reader), threading.Thread(target=writer)
    a.start(); b.start(); a.join(); b.join()
    return (store, cache) if "k" in cache and cache["k"][0] != store["k"] else None


update = sum(two_writers("update") for _ in range(args.trials))
delete = sum(two_writers("delete") for _ in range(args.trials))
print(f"two writers, update the cache: {update:,} of {args.trials:,} trials disagree ({update / args.trials:.2%})")
print(f"two writers, delete the cache: {delete:,} of {args.trials:,} disagree after the next read")

def follow(store, cache):
    """Right after the race: a TTL-honouring reader reads every 10 ms; how long is it served the stale value?"""
    v, written = cache["k"]
    while True:
        entry = cache.get("k")
        if entry is None or time.time() - entry[1] >= args.ttl:      # expired: the next read goes to the store
            cache["k"] = (store["k"], time.time())
        if cache["k"][0] == store["k"]:
            return time.time() - written
        time.sleep(0.01)

stale, lifetimes = 0, []
for _ in range(args.trials):
    r = reader_vs_writer()
    if r:
        stale += 1
        if len(lifetimes) < args.ttl_checks: lifetimes.append(follow(*r))
print(f"reader against writer, delete: {stale:,} of {args.trials:,} trials leave a stale value "
      f"({stale / args.trials:.2%})")
print(f"with a {args.ttl:g} s TTL, the {len(lifetimes)} stale entries followed were served for at most {max(lifetimes):.2f} s")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {
            "races": [{"experiment": "two writers, update", "bad": update, "pct": update / args.trials * 100},
                      {"experiment": "two writers, delete", "bad": delete, "pct": delete / args.trials * 100},
                      {"experiment": "reader against writer, delete", "bad": stale, "pct": stale / args.trials * 100}],
            "summary": [{"update_pct": update / args.trials * 100, "delete_bad": delete,
                         "stale_pct": stale / args.trials * 100,
                         "max_stale_s": max(lifetimes) if lifetimes else 0, "ttl": args.ttl}]}}, f)
