Topic 9: Module 9 lab

"""Module 9 lab: run the notes server as a small production fleet and exercise it end to end.

Two replicas with a shared requestState key, a per-client rate limit, and checkpointed
imports, behind a sticky round-robin balancer. Every step prints what it observed.
"""
from __future__ import annotations

import secrets
import shutil
import signal
from pathlib import Path

import anyio
import httpx2

from examples.m09_eras import one_call
from examples.m09_fleet import fleet
from examples.m09_import_client import main as import_notes
from examples.m09_loadtest import HEADERS, body
from examples.m09_registry_check import check

WORK = Path("/tmp/m09-lab")


async def burst(url: str, count: int) -> dict[int, int]:
    statuses: dict[int, int] = {}
    async with httpx2.AsyncClient() as http:
        for n in range(count):
            status = (await http.post(url, json=body(n), headers=HEADERS)).status_code
            statuses[status] = statuses.get(status, 0) + 1
    return statuses


async def lab(f) -> None:
    print("1. health")
    async with httpx2.AsyncClient() as http:
        for url in f.replica_urls:
            print("  ", (await http.get(url + "/health")).json())

    print("2. both protocol eras through the balancer")
    for mode in ("auto", "legacy"):
        print("  ", await one_call(f.url, mode))

    print("3. multi-round-trip import across replicas, then the same key again")
    await import_notes(f.url, "lab-batch-0001", 3)
    await import_notes(f.url, "lab-batch-0001", 3)

    print("4. burst of 40 calls from one client (limit: 10 per replica, refill 2/s)")
    await anyio.sleep(5)  # let the buckets refill after the steps above
    print("   statuses:", await burst(f.url, 40))

    print("5. SIGTERM replica-b, then call again")
    f.processes["replica-b"].send_signal(signal.SIGTERM)
    f.processes["replica-b"].wait(timeout=30)
    await anyio.sleep(5)
    for mode in ("auto", "legacy"):
        print("  ", await one_call(f.url, mode))

    print("6. registry metadata")
    problems = check(Path("registry/server.json"))
    print("   server.json:", "valid" if not problems else problems)


def main() -> None:
    shutil.rmtree(WORK, ignore_errors=True)
    shutil.copytree("notes", WORK / "notes")
    env = {"NOTES_DIR": str(WORK / "notes"), "NOTES_CHECKPOINT_DIR": str(WORK / "checkpoints"),
           "NOTES_REQUEST_STATE_KEY": secrets.token_hex(32), "NOTES_RATE_CAPACITY": "10", "NOTES_RATE_REFILL": "2"}
    with fleet(replicas=2, sticky=True, env=env) as f:
        anyio.run(lab, f)
        print("\nbalancer warnings:")
        print("\n".join(line for line in f.log("balancer").splitlines() if "refused" in line) or "   (none)")


if __name__ == "__main__":
    main()