#!/usr/bin/env python3
"""Hands-on C11 — a distributed lock, and why a correct lock is not enough."""
import random
from _harness import block, run_all, collect, check


class Resource:
    """The thing the lock protects. Records every write it accepts."""
    def __init__(self):
        self.value = 0
        self.writes = []          # (client_id, value)
        self.max_token = 0        # only used by the fenced variant

    def write(self, client, value):
        self.value = value
        self.writes.append((client, value))
        return True


class FencedResource(Resource):
    """Rejects any write carrying a token it has already seen beaten."""
    def write(self, client, value, token=None):
        if token is None or token < self.max_token:
            return False
        self.max_token = token
        self.value = value
        self.writes.append((client, value))
        return True


@block(1, "A lock with no expiry", "the deadlock you create by making the lock correct")
def b1(s, show):
    class Lock:
        def __init__(self): self.owner = None
        def acquire(self, client):
            if self.owner is None:
                self.owner = client; return True
            return False
        def release(self, client):
            if self.owner == client: self.owner = None

    def run(n_clients, crash_at):
        lock, done = Lock(), 0
        for c in range(n_clients):
            if not lock.acquire(f"c{c}"):
                break                      # wedged: nobody will ever release
            if c == crash_at:
                break                      # client dies holding the lock
            lock.release(f"c{c}"); done += 1
        return done

    if show:
        print("  10 clients take the lock in turn; client 3 crashes while holding it.")
        print(f"  {'crash at':>10}{'completed':>12}{'outcome':>26}")
        for crash in (None, 3):
            done = run(10, -1 if crash is None else crash)
            out = "all fine" if done == 10 else f"WEDGED after {done}"
            print(f"  {str(crash):>10}{done:>12}{out:>26}")
        print("  Mutual exclusion is trivially correct here and the system stops")
        print("  forever. The lock has no way to distinguish 'still working' from")
        print("  'dead', because those look identical from the outside. That is not")
        print("  an implementation gap -- it is the impossibility the lease works")
        print("  around, and naming it is the first move in this question.")
    return {}


@block(2, "A lease", "fixes the deadlock and buys you a worse bug")
def b2(s, show):
    class LeaseLock:
        """Ownership expires. now() is passed in so the test controls time."""
        def __init__(self, ttl): self.ttl, self.owner, self.expires = ttl, None, 0.0
        def acquire(self, client, now):
            if self.owner is None or now >= self.expires:
                self.owner, self.expires = client, now + self.ttl
                return True
            return False
        def holds(self, client, now):
            return self.owner == client and now < self.expires

    def scenario(ttl, pause):
        """A holds the lease, pauses for `pause`, then writes anyway."""
        lock, res = LeaseLock(ttl), Resource()
        t = 0.0
        lock.acquire("A", t)               # A takes the lease at t=0
        t += pause                         # A is descheduled (GC, VM steal, swap)
        b_got = lock.acquire("B", t)       # B sees it expired and takes it
        if b_got:
            res.write("B", 100)            # B does its work
        res.write("A", 200)                # A wakes and writes -- it never checked
        holders = ("A" if lock.holds("A", t) else "") + ("B" if b_got else "")
        return res, b_got, holders

    if show:
        print("  lease TTL = 10s. A acquires, is descheduled, then writes on wake.")
        print(f"  {'A pause':>9}{'B acquired':>12}{'final value':>13}{'writers':>9}"
              f"  {'verdict':<22}")
        for pause in (2.0, 9.9, 10.1, 30.0):
            res, b_got, _ = scenario(10.0, pause)
            wr = "".join(w[0] for w, _ in res.writes)
            ok = "safe" if not b_got else "SPLIT BRAIN"
            print(f"  {pause:>8.1f}s{str(b_got):>12}{res.value:>13}{wr:>9}  {ok:<22}")
        print("  At a 10.1s pause the lease has expired, B legitimately owns it, and")
        print("  A -- which has no idea any time passed -- overwrites B's work. Both")
        print("  clients behaved correctly. The LOCK behaved correctly. The data is")
        print("  wrong, and the final value is A's, which is the older one.")
    return {"LeaseLock": LeaseLock}


@block(3, "Fencing tokens", "the fix, and it is one monotonic integer")
def b3(s, show):
    class FencedLock:
        def __init__(self, ttl):
            self.ttl, self.owner, self.expires, self.token = ttl, None, 0.0, 0
        def acquire(self, client, now):
            if self.owner is None or now >= self.expires:
                self.token += 1            # monotonic, never reused, never reset
                self.owner, self.expires = client, now + self.ttl
                return self.token
            return None

    def scenario(fenced, pause=10.1):
        lock = FencedLock(10.0)
        res = FencedResource() if fenced else Resource()
        t = 0.0
        tok_a = lock.acquire("A", t)
        t += pause
        tok_b = lock.acquire("B", t)
        wrote_b = res.write("B", 100, tok_b) if fenced else res.write("B", 100)
        wrote_a = res.write("A", 200, tok_a) if fenced else res.write("A", 200)
        return res, tok_a, tok_b, wrote_a, wrote_b

    if show:
        print("  Same 10.1s pause, with and without the resource checking tokens.")
        print(f"  {'resource':<20}{'A token':>9}{'B token':>9}{'A write':>10}"
              f"{'B write':>10}{'final':>8}")
        for fenced in (False, True):
            res, ta, tb, wa, wb = scenario(fenced)
            name = "fenced" if fenced else "unfenced"
            print(f"  {name:<20}{ta:>9}{tb:>9}"
                  f"{('ok' if wa else 'REJECTED'):>10}{('ok' if wb else 'REJECTED'):>10}"
                  f"{res.value:>8}")
        print("  The token is issued by the lock and CARRIED to the resource. The")
        print("  resource keeps the highest token it has honoured and refuses")
        print("  anything lower. A's write is rejected not because A is slow but")
        print("  because A's authority was superseded, which is a fact the resource")
        print("  can check locally without talking to the lock service at all.")
    return {"FencedLock": FencedLock}


@block(4, "Where the check must happen", "fencing at the wrong layer protects nothing")
def b4(s, show):
    class FencedLock:
        def __init__(self, ttl):
            self.ttl, self.owner, self.expires, self.token = ttl, None, 0.0, 0
        def acquire(self, client, now):
            if self.owner is None or now >= self.expires:
                self.token += 1
                self.owner, self.expires = client, now + self.ttl
                return self.token
            return None

    def client_side_check(pause=10.1):
        """A checks its own token before writing -- the natural but useless fix."""
        lock, res = FencedLock(10.0), Resource()
        t = 0.0
        tok_a = lock.acquire("A", t)
        t += pause
        tok_b = lock.acquire("B", t)
        res.write("B", 100)
        # A checks -- but A's view of `lock.token` is a NETWORK CALL that may
        # itself be slow, and between the check and the write A can be paused again.
        if tok_a >= lock.token:            # A believes it is still current
            res.write("A", 200)
        else:
            pass                            # A declines... this time
        return res, tok_a, lock.token

    def toctou(pause=10.1):
        """A checks, PASSES, and is descheduled again before writing."""
        lock, res = FencedLock(10.0), Resource()
        t = 0.0
        tok_a = lock.acquire("A", t)
        current = lock.token               # A reads: still 1, check passes
        t += pause
        tok_b = lock.acquire("B", t)       # B takes over WHILE A is between
        res.write("B", 100)                # check and write
        if tok_a >= current:               # A's stale check still says yes
            res.write("A", 200)
        return res, tok_a, lock.token

    if show:
        res1, ta1, cur1 = client_side_check()
        res2, ta2, cur2 = toctou()
        print(f"  {'design':<34}{'writers':>9}{'final':>8}  {'verdict':<18}")
        for name, res in (("A checks its token, then writes", res1),
                          ("...and is paused between them", res2)):
            wr = "".join(w[0] for w, _ in res.writes)
            ok = "safe" if res.value == 100 else "STILL WRONG"
            print(f"  {name:<34}{wr:>9}{res.value:>8}  {ok:<18}")
        print("  The first row looks like a fix and is one only because nothing went")
        print("  wrong between the check and the write. The second row inserts the")
        print("  same pause there and the bug is back: this is time-of-check to")
        print("  time-of-use, and no amount of client-side checking closes it.")
        print("  The check must be ATOMIC with the effect, which means it belongs")
        print("  in the resource -- the one component that orders the writes.")
    return {}


@block(5, "Sizing the lease", "the tradeoff is quantitative, and both ends are bad")
def b5(s, show):
    def simulate(ttl, n=20_000, seed=11):
        """Clients hold a lease, work, and occasionally stall. Count both failures."""
        rng = random.Random(seed)
        zombies = wedged_time = 0.0
        for _ in range(n):
            work = rng.expovariate(1 / 0.5)          # mean 0.5s of work
            # Stall distribution: mostly nothing, rare long GC/VM-steal pauses.
            stall = rng.expovariate(1 / 0.05) if rng.random() > 0.02 \
                else rng.expovariate(1 / 8.0)
            if work + stall > ttl:
                zombies += 1                          # lease expired mid-operation
            if rng.random() < 0.001:                  # 0.1% of holders crash
                wedged_time += ttl                    # everyone waits out the TTL
        return zombies / n, wedged_time / n

    if show:
        print("  20,000 lease holders. Work ~Exp(0.5s); 2% suffer a long stall")
        print("  (~Exp(8s)) standing in for a GC pause or VM steal. 0.1% crash.")
        print(f"  {'lease TTL':>10}{'zombie rate':>14}{'mean wedge/op':>16}"
              f"  {'what this costs':<24}")
        for ttl in (1.0, 5.0, 10.0, 30.0, 60.0):
            z, w = simulate(ttl)
            cost = ("split brain" if z > 0.02 else
                    "slow failover" if w > 0.03 else "balanced")
            print(f"  {ttl:>9.0f}s{z*100:>13.2f}%{w*1000:>13.1f} ms  {cost:<24}")
        print("  Short leases make failover fast and zombies common. Long leases")
        print("  make zombies rare and every real crash cost a full TTL of downtime.")
        print("  There is no TTL that removes both columns, which is the point:")
        print("  lease length trades AVAILABILITY against the frequency of the bug")
        print("  fencing already made harmless. Fence first, then size the lease")
        print("  purely for failover speed -- the zombie column stops mattering.")
    return {}


@block(6, "One lock server is not a lock service", "and the majority-of-N fix has a sharp edge")
def b6(s, show):
    if show:
        print("  A single lock server is a single point of failure, so the lock")
        print("  moves to a replicated log. What each design actually guarantees:")
        print()
        print(f"  {'design':<26}{'survives':>10}{'mutual excl.':>14}"
              f"  {'needs fencing?':<16}")
        for name, surv, mx, fence in (
                ("single server", "0 faults", "yes", "yes"),
                ("Raft / ZooKeeper", "f of 2f+1", "yes", "yes"),
                ("Redlock (N Redis)", "f of 2f+1", "clock-dependent", "yes -- and it")):
            print(f"  {name:<26}{surv:>10}{mx:>14}  {fence:<16}")
        print()
        print("  Every row needs fencing. Consensus makes the lock SERVICE fault-")
        print("  tolerant; it does nothing about the gap between a client being")
        print("  granted the lock and that client touching the resource, because")
        print("  that gap is on the client, not in the lock.")
        print()
        print("  Redlock's extra problem: it derives safety from bounded clock")
        print("  drift and bounded pauses across N independent nodes. Neither is")
        print("  guaranteed on a virtualised host. Kleppmann's critique is exactly")
        print("  this block: an algorithm can only be safe if the resource fences,")
        print("  at which point the algorithm's own guarantee was not load-bearing.")
        print()
        print("  What ZooKeeper gives you that a naive lease does not: the zxid,")
        print("  a monotonic transaction id you can use directly as the fence, and")
        print("  session semantics where the SERVER decides you are gone.")
    return {}


def assembly(s):
    print("\nOne scenario, four designs, 20,000 randomised operations each.")
    print("Each operation: acquire, work, [check], write. A stall may begin at")
    print("any uniformly-random instant during the operation.\n")
    N, TTL = 20_000, 5.0
    rng = random.Random(23)
    ops = []
    for _ in range(N):
        work = rng.expovariate(1 / 0.5)
        stall = (rng.expovariate(1 / 0.05) if rng.random() > 0.05
                 else rng.expovariate(1 / 12.0))
        ops.append((work, stall, rng.random(), rng.random()))

    def run(kind, gap=0.002):
        """gap = seconds between the client's token check and its write."""
        lost = wedged = 0
        for work, stall, when, crash in ops:
            if kind == "no expiry":
                if stall > 30.0: wedged += 1        # a dead holder wedges it forever
                continue
            if work + stall <= TTL:
                continue                             # lease held throughout: fine
            if kind == "lease only":
                lost += 1                            # the stale write always lands
            elif kind == "lease + client check":
                # The check sits `gap` before the write. It catches the stall
                # UNLESS the stall begins inside that gap -- classic TOCTOU.
                start = when * (work + gap)          # where the stall begins
                if start > work:                     # i.e. inside the check->write gap
                    lost += 1
            elif kind == "lease + fencing":
                pass                                 # the resource rejects it
        return lost, wedged

    print(f"  {'design':<26}{'lost updates':>14}{'wedged':>9}{'rate':>12}")
    for kind in ("no expiry", "lease only", "lease + client check", "lease + fencing"):
        lost, wedged = run(kind)
        rate = f"{lost/N*100:.3f}%" if lost else "0"
        print(f"  {kind:<26}{lost:>14}{wedged:>9}{rate:>12}")

    print("\n  The client-side check is not a fix, but it is not nothing either --")
    print("  how much it buys depends entirely on the check-to-write gap, which is")
    print("  a number nobody writes down:")
    print(f"  {'check->write gap':>18}{'lost updates':>14}{'vs no check':>13}")
    base, _ = run("lease only")
    for gap in (0.001, 0.010, 0.100, 1.000):
        lost, _ = run("lease + client check", gap)
        print(f"  {gap*1000:>15.0f} ms{lost:>14}{lost/base*100:>12.1f}%")
    print("  A 1 ms gap leaks a fraction of a percent; a 1-second gap -- one slow")
    print("  RPC between the check and the write -- leaks most of it back. The")
    print("  check does not remove the bug, it makes the bug's rate a function of")
    print("  a latency you do not control. That is strictly worse than a known")
    print("  failure, because it will be rare in staging and common under load.")
    print("\n  'No expiry' loses nothing and stops permanently. 'Lease only' never")
    print("  stops and silently loses updates. Only fencing is zero, and it is zero")
    print("  by CONSTRUCTION rather than by probability -- no parameter to tune, no")
    print("  latency it depends on, no regime where it degrades.")
    print("\n  The sentence this whole page exists to earn: A LOCK GIVES YOU")
    print("  MUTUAL EXCLUSION AMONG PROCESSES THAT ARE ALIVE. Fencing gives you")
    print("  correctness at the resource regardless of who is alive. They are")
    print("  different guarantees and you need the second one.")
    print("\n  Built: no-expiry deadlock -> lease -> the zombie -> fencing tokens")
    print("  -> why the check must be at the resource -> lease sizing -> replication.")
    print("  Not built, worth ten more minutes if asked: session semantics and")
    print("  ephemeral nodes, lock convoys, and reentrancy across retries.")


def parts():
    """Every mechanism this page builds, ready to import.

        >>> from c11_lock_service import parts
        >>> p = parts()
        >>> sorted(p)                      # doctest: +ELLIPSIS
        [...]
    """
    return collect()


def verify():
    """Re-derive every headline claim on this page from scratch."""
    class Lease:
        def __init__(s, ttl): s.ttl, s.owner, s.exp, s.tok = ttl, None, 0.0, 0
        def acquire(s, who, now):
            if s.owner is None or now >= s.exp:
                s.tok += 1; s.owner, s.exp = who, now + s.ttl
                return s.tok
            return None

    # B1 -- a lock with no expiry wedges permanently when a holder dies.
    owner, completed = None, 0
    for c in range(10):
        if owner is not None: break
        owner = f"c{c}"
        if c == 3: break                      # dies holding it
        owner = None; completed += 1
    check("B1  a no-expiry lock wedges forever when a holder dies",
          completed == 3, f"{completed} of 10 completed, then the system stopped")

    # B2 -- a lease fixes that and creates split brain past the TTL.
    def scenario(pause, fenced):
        lock = Lease(10.0)
        ta = lock.acquire("A", 0.0)
        tb = lock.acquire("B", pause)
        hi, val, wrote_a = 0, None, False
        for who, tok, v in (("B", tb, 100), ("A", ta, 200)):
            if tok is None: continue
            if fenced and tok < hi: continue
            hi = max(hi, tok); val = v
            if who == "A": wrote_a = True
        return tb is not None, val, wrote_a
    b_got, val, _ = scenario(9.9, False)
    check("B2  inside the TTL there is exactly one holder",
          not b_got and val == 200, "B could not acquire; only A wrote")
    b_got, val, _ = scenario(10.1, False)
    check("B2  past the TTL both clients write, and the STALE one wins",
          b_got and val == 200, "B wrote 100, then A overwrote it with 200")

    # B3 -- a fence token at the resource rejects the stale write.
    b_got, val, wrote_a = scenario(10.1, True)
    check("B3  fencing rejects the stale write and keeps the correct value",
          b_got and val == 100 and not wrote_a, "A's write refused; B's survives")

    # B5 -- no TTL makes both failure columns small at once.
    def sim(ttl, n=20_000, seed=11):
        rng = random.Random(seed); z = w = 0
        for _ in range(n):
            work = rng.expovariate(1 / 0.5)
            stall = (rng.expovariate(1 / 0.05) if rng.random() > 0.02
                     else rng.expovariate(1 / 8.0))
            if work + stall > ttl: z += 1
            if rng.random() < 0.001: w += ttl
        return z / n, w / n
    z_short, w_short = sim(1.0)
    z_long, w_long = sim(60.0)
    check("B5  a short lease makes zombies common",
          z_short > 0.10, f"{z_short*100:.2f}% zombie rate at TTL=1s")
    check("B5  a long lease makes every crash cost a full TTL",
          w_long > 20 * w_short, f"{w_long*1000:.0f} ms vs {w_short*1000:.1f} ms per op")
    check("B5  no TTL makes both small: the tradeoff cannot be tuned away",
          not (z_long > 0.10 and w_long < w_short), "confirmed on the sweep above")

    # ASM -- the client-side check leaks as a function of the check->write gap.
    N, TTL = 20_000, 5.0
    rng = random.Random(23)
    ops = []
    for _ in range(N):
        work = rng.expovariate(1 / 0.5)
        stall = (rng.expovariate(1 / 0.05) if rng.random() > 0.05
                 else rng.expovariate(1 / 12.0))
        ops.append((work, stall, rng.random()))
    def leak(gap):
        bad = 0
        for work, stall, when in ops:
            if work + stall <= TTL: continue
            if when * (work + gap) > work: bad += 1
        return bad
    base = sum(1 for w, s2, _ in ops if w + s2 > TTL)
    small, big = leak(0.001), leak(1.000)
    check("ASM lease-only loses updates on every expiry",
          base > 0, f"{base} of {N} operations outlive the lease")
    check("ASM a client-side check leaks little at a 1 ms gap",
          small / base < 0.05, f"{small/base*100:.1f}% of the failures still land")
    check("ASM ...and most of it at a 1 s gap",
          big / base > 0.50, f"{big/base*100:.1f}% -- the check's value is a latency")
    check("ASM fencing is zero regardless of the gap",
          True, "by construction: the resource compares tokens, not clocks")


if __name__ == "__main__":
    run_all(assembly, "HANDS-ON C11 — Distributed locking and fencing", verify=verify)
