undercooked/tests/integration/test_ws_budgets.py

"""The budget suite (phase 06 task 06, `slow`): PLAN §10's
performance numbers, measured on a **live** 4-player room — the mean
tick cost over 400 real 20 Hz ticks (≤ 2 ms, the p99 ≤ 8 ms), the
`snap` frame size (≤ 1.5 KB on the 4-player wire shape), the
downlink rate (≤ 60 KB/s over a 5 s window on one client) and the
input→snapshot echo (≤ 120 ms — the send to the first `snap` that
arrives after it: at 20 Hz the server answers within one tick —
wall-tick jitter re-encodes the same sim tick on some ticks, so the
*distinct*-tick definition is the 2-tick one, and the budget is 2.4
ticks).

The tick cost is measured by wrapping `room._tick` (the bound method
the room's own tick task drives — no `FakeClock`, no replay: the
wall-clock 20 Hz loop, the substep dance and all). The sim is kept
working by three driving sockets (a `input` frame each ~10 s⁻¹, a
random unit direction and an occasional press — the same choreography
the two-client round test uses, in miniature); the fourth cook stays
still and is the measurement socket (its `input` seq counter is not
interleaved with the drivers' — a stale-seq refusal would poison
both).

`docs/PROTOCOL.md` "Testing the protocol" carries the measured
values this suite bounds (the bound is the budget; the measurement
is what a regression moves).
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import math
import random
import time
from pathlib import Path
from typing import Any

import httpx
import pytest
import websockets
from websockets.exceptions import ConnectionClosed

pytestmark = [pytest.mark.integration, pytest.mark.slow]

PLAYABLE_LEVELS_DIR = (
    Path(__file__).resolve().parents[1] / "fixtures" / "levels_playable"
)

#: The tick budget (PLAN §10): the mean over 400 real ticks, and the
#: p99 (a single slow tick may not hide in the mean).
TICKS_TO_MEASURE = 400
MEAN_TICK_BUDGET_S = 0.002
P99_TICK_BUDGET_S = 0.008
#: The `snap` size budget (the 4p wire shape, `test_encoding.py`'s
#: 1.5 KB).
SNAPSHOT_BUDGET_BYTES = 1500
#: The downlink budget (bytes per second, a 5 s window on one client).
DOWNLINK_WINDOW_S = 5.0
DOWNLINK_BUDGET_BYTES_PER_S = 60 * 1024
#: The input→snapshot echo budget (send → the first `snap` that
#: arrives after the send; at 20 Hz the server answers within one
#: tick — 50 ms — and the budget is 2.4 ticks).
ECHO_BUDGET_S = 0.120


# ------------------------------------------------------------- the room


def _create_room(url: str, token: str) -> str:
    response = httpx.post(
        f"{url}/api/rooms", json={"token": token}, timeout=10.0
    )
    assert response.status_code == 201, response.text
    return response.json()["room_code"]


def _join(url: str, code: str, token: str) -> None:
    response = httpx.post(
        f"{url}/api/rooms/{code}/join", json={"token": token}, timeout=10.0
    )
    assert 200 <= response.status_code < 300, response.text


async def _connect(url: str, code: str, token: str, name: str) -> Any:
    ws = await websockets.connect(
        url.replace("http://", "ws://", 1) + f"/ws/{code}", open_timeout=10.0
    )
    await ws.send(
        json.dumps({"type": "hello", "p": 1, "token": token, "name": name})
    )
    welcome = json.loads(await ws.recv())
    assert welcome["type"] == "welcome", welcome
    return ws


async def _wait_for_async(predicate: Any, timeout_s: float) -> None:
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        if predicate():
            return
        await asyncio.sleep(0.05)
    msg = f"timed out after {timeout_s:.0f} s"
    raise TimeoutError(msg)


def _input_frame(rng: random.Random, seq: int) -> dict[str, Any]:
    """One `input` frame: a random unit movement intent (the sim
    clamps any magnitude) and an occasional press."""
    angle = rng.uniform(0.0, 2.0 * math.pi)
    return {
        "type": "input",
        "seq": seq,
        "move": {
            "x": round(math.cos(angle), 6),
            "y": round(math.sin(angle), 6),
        },
        "act": rng.random() < 0.1,
    }


# --------------------------------------------------- the three phases


async def _measure_tick_cost(room: Any) -> tuple[float, float, int]:
    """The mean / p99 / n of the room's real tick, over
    `TICKS_TO_MEASURE` ticks (the `room._tick` wrap — the room's own
    tick task drives it; `tick_once`'s crash isolation stays outside
    the wrap)."""
    durations: list[float] = []
    original_tick = room._tick

    def _timed_tick(dt: float, when: float) -> None:
        started = time.perf_counter()
        original_tick(dt, when)
        durations.append(time.perf_counter() - started)

    room._tick = _timed_tick  # type: ignore[method-assign]
    try:
        await _wait_for_async(
            lambda: len(durations) >= TICKS_TO_MEASURE, timeout_s=60.0
        )
    finally:
        room._tick = original_tick  # type: ignore[method-assign]
    mean_tick = sum(durations) / len(durations)
    p99 = sorted(durations)[int(0.99 * len(durations)) - 1]
    return mean_tick, p99, len(durations)


async def _measure_downlink(ws: Any) -> tuple[float, int, int]:
    """One client's downlink over `DOWNLINK_WINDOW_S`: the bytes per
    second, the largest `snap` frame and the snap count. The kernel
    TCP backlog the tick phase left on this (unread) socket is
    drained first (2 s — the window's bound, and loopback delivers
    the backlog far faster): without it the window would measure the
    catch-up burst, not the steady state a continuously-reading
    client sees."""
    drain_deadline = time.monotonic() + 2.0
    while time.monotonic() < drain_deadline:
        with contextlib.suppress(
            TimeoutError, asyncio.TimeoutError, ConnectionClosed
        ):
            await asyncio.wait_for(ws.recv(decode=False), 0.5)
    total_bytes = 0
    snapshot_sizes: list[int] = []
    deadline = time.monotonic() + DOWNLINK_WINDOW_S
    while time.monotonic() < deadline:
        try:
            raw = await asyncio.wait_for(ws.recv(decode=False), 1.0)
        except TimeoutError:
            continue
        except ConnectionClosed:
            break
        total_bytes += len(raw)
        if json.loads(raw)["type"] == "snap":
            snapshot_sizes.append(len(raw))
    downlink_bytes_per_s = total_bytes / DOWNLINK_WINDOW_S
    max_snapshot = max(snapshot_sizes) if snapshot_sizes else 0
    return downlink_bytes_per_s, max_snapshot, len(snapshot_sizes)


async def _measure_echo(ws: Any, rng: random.Random) -> list[float]:
    """Three input→snapshot echoes (~1 s apart): the send to the first
    `snap` that arrives after it. The echo budget is the bound — a
    sample that misses it fails, and a snap already in the socket's
    buffer (the tick between the send and the read) reads 0 ms."""
    elapsed_s: list[float] = []
    for sample in range(3):
        frame = _input_frame(rng, 100_000 + sample)
        sent_at = time.monotonic()
        with contextlib.suppress(Exception):
            await ws.send(json.dumps(frame))
        saw_echo = False
        sample_elapsed: float | None = None
        deadline = sent_at + ECHO_BUDGET_S
        while time.monotonic() < deadline:
            try:
                raw = await asyncio.wait_for(ws.recv(decode=False), 1.0)
            except TimeoutError:
                break  # the echo budget is the bound; it expires
            except ConnectionClosed:
                break
            if json.loads(raw)["type"] == "snap":
                elapsed = time.monotonic() - sent_at
                assert elapsed <= ECHO_BUDGET_S, (
                    f"the input→snapshot echo is "
                    f"{elapsed * 1000:.0f} ms "
                    f"(budget {ECHO_BUDGET_S * 1000:.0f} ms)"
                )
                saw_echo = True
                sample_elapsed = elapsed
                break
        assert saw_echo and sample_elapsed is not None, (
            f"sample {sample}: no snap frame arrived after the input "
            f"within the {ECHO_BUDGET_S * 1000:.0f} ms echo budget"
        )
        elapsed_s.append(sample_elapsed)
        await asyncio.sleep(1.0)
    return elapsed_s


# ------------------------------------------------------------- the test


async def test_the_4_player_room_holds_the_tick_snapshot_downlink_and_echo_budgets(
    ws_server: Any,
) -> None:
    with ws_server(levels_dir=str(PLAYABLE_LEVELS_DIR)) as (url, app):
        code = _create_room(url, "tok-0")
        for seat in (1, 2, 3):
            _join(url, code, f"tok-{seat}")
        clients = [
            await _connect(url, code, f"tok-{seat}", f"Cook{seat}")
            for seat in range(4)
        ]
        room = app.state.supervisor.get(code)
        assert room is not None

        # Ready-up (every seat — the start guard), then the start.
        for ws in clients:
            await ws.send(json.dumps({"type": "ready"}))
        await clients[0].send(json.dumps({"type": "start"}))
        await _wait_for_async(
            lambda: room.state.value == "playing"
            and room.sim_state is not None,
            timeout_s=30.0,
        )

        # The drivers: three cooks wander on a ~10 s⁻¹ cadence (a
        # random unit direction, an occasional press). The sim is the
        # loaded machine the budget applies to. The fourth (the
        # measurement socket) stays still — its `input` seq counter
        # is not interleaved with the drivers' (a stale-seq refusal
        # would poison both). The seeded generator is the test's own
        # determinism (a fixed choreography, never a credential —
        # `random` is the sim's only RNG, and this is not the sim).
        rng = random.Random(20240606)  # noqa: S311

        async def _drive(ws: Any) -> None:
            seq = 1
            while True:
                frame = _input_frame(rng, seq)
                seq += 1
                with contextlib.suppress(Exception):
                    await ws.send(json.dumps(frame))
                await asyncio.sleep(0.1)

        drivers = [
            asyncio.create_task(_drive(ws)) for ws in clients[:3]
        ]

        # The measurement socket's keep-alive (the documented client
        # behaviour — a silent socket for `HEARTBEAT_DEAD_S` (12 s) is
        # a dead client, and the test's window outlives the silence):
        # a `ping` every 2 s keeps it alive for the full ~30 s.
        async def _keepalive(ws: Any) -> None:
            n = 0
            while True:
                await asyncio.sleep(2.0)
                n += 1
                with contextlib.suppress(Exception):
                    await ws.send(json.dumps({"type": "ping", "t": n}))

        keepalive = asyncio.create_task(_keepalive(clients[3]))
        try:
            # ---- the tick cost (400 real 20 Hz ticks, ~20 s of wall)
            mean_tick, p99, n_ticks = await _measure_tick_cost(room)
            assert mean_tick <= MEAN_TICK_BUDGET_S, (
                f"the mean tick cost is {mean_tick * 1000:.3f} ms over "
                f"{n_ticks} ticks (budget "
                f"{MEAN_TICK_BUDGET_S * 1000:.1f} ms)"
            )
            assert p99 <= P99_TICK_BUDGET_S, (
                f"the p99 tick cost is {p99 * 1000:.3f} ms "
                f"(budget {P99_TICK_BUDGET_S * 1000:.1f} ms)"
            )

            # ---- the downlink + the snapshot size (one client, 5 s)
            bytes_per_s, max_snapshot, n_snaps = await _measure_downlink(
                clients[3]
            )
            assert n_snaps > 0, "no snap frames in the downlink window"
            assert bytes_per_s <= DOWNLINK_BUDGET_BYTES_PER_S, (
                f"the downlink is {bytes_per_s / 1024:.1f} KB/s "
                f"(budget {DOWNLINK_BUDGET_BYTES_PER_S / 1024:.0f} KB/s)"
            )
            assert max_snapshot <= SNAPSHOT_BUDGET_BYTES, (
                f"a snap frame is {max_snapshot} bytes "
                f"(budget {SNAPSHOT_BUDGET_BYTES})"
            )

            # ---- the input→snapshot echo (3 samples, ~1 s apart)
            await _measure_echo(clients[3], rng)
        finally:
            for task in drivers:
                task.cancel()
            keepalive.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await asyncio.gather(*drivers, keepalive)