diff --git a/.env.example b/.env.example index 4358863..3323345 100644 --- a/.env.example +++ b/.env.example @@ -9,3 +9,25 @@ ALLOWED_ORIGINS= AI_AGENT=model LOG_LEVEL=INFO + +# --- Resource limits (hard caps per container) ------------------------------- +SERVER_CPUS=1.0 +SERVER_MEMORY=512M +AI_CPUS=1.0 +AI_MEMORY=1G +WEBSITE_CPUS=0.25 +WEBSITE_MEMORY=64M + +# --- Capacity ----------------------------------------------------------------- +# Games running at the same time; more game starts are refused ("server busy"). +MAX_ACTIVE_GAMES=50 +# Open websocket connections in total / per client IP. +MAX_CONNECTIONS=500 +MAX_CONNECTIONS_PER_IP=10 + +# --- AI service --------------------------------------------------------------- +# Decisions computed in parallel, waiting requests before refusing, and queue +# depth from which the cheap heuristic agent is used instead of the model. +AI_WORKERS=1 +AI_MAX_QUEUE=64 +AI_DEGRADE_QUEUE=8 diff --git a/compose.yaml b/compose.yaml index dd9d582..ea24013 100644 --- a/compose.yaml +++ b/compose.yaml @@ -4,6 +4,12 @@ services: wizard_website: build: ./wizard_website restart: unless-stopped + deploy: + resources: + limits: + cpus: "${WEBSITE_CPUS:-0.25}" + memory: ${WEBSITE_MEMORY:-64M} + pids_limit: 64 networks: - wizard_website @@ -14,6 +20,16 @@ services: AI_SHARED_SECRET: ${AI_SHARED_SECRET:?set AI_SHARED_SECRET in .env} ALLOWED_ORIGINS: ${ALLOWED_ORIGINS:-} LOG_LEVEL: ${LOG_LEVEL:-INFO} + MAX_ACTIVE_GAMES: ${MAX_ACTIVE_GAMES:-50} + MAX_CONNECTIONS: ${MAX_CONNECTIONS:-500} + MAX_CONNECTIONS_PER_IP: ${MAX_CONNECTIONS_PER_IP:-10} + # Hard caps so Wizard can never starve other services on the host. + deploy: + resources: + limits: + cpus: "${SERVER_CPUS:-1.0}" + memory: ${SERVER_MEMORY:-512M} + pids_limit: 256 networks: - wizard_api - wizard_internal @@ -31,6 +47,18 @@ services: AGENT: ${AI_AGENT:-model} MODEL_PATH: /srv/models/latest LOG_LEVEL: ${LOG_LEVEL:-INFO} + AI_WORKERS: ${AI_WORKERS:-1} + AI_MAX_QUEUE: ${AI_MAX_QUEUE:-64} + AI_DEGRADE_QUEUE: ${AI_DEGRADE_QUEUE:-8} + # Numeric libraries (numpy/torch) must not start one thread per host core. + OMP_NUM_THREADS: ${AI_WORKERS:-1} + MKL_NUM_THREADS: ${AI_WORKERS:-1} + deploy: + resources: + limits: + cpus: "${AI_CPUS:-1.0}" + memory: ${AI_MEMORY:-1G} + pids_limit: 256 volumes: # Trained checkpoints; replace without rebuilding the image. - ./wizard_ai/models:/srv/models:ro diff --git a/docs/04-server.md b/docs/04-server.md index 000a129..d7f082f 100644 --- a/docs/04-server.md +++ b/docs/04-server.md @@ -26,7 +26,9 @@ flowchart TB | `service.py` | `WizardService` | the "controller": handles a browser connection, dispatches messages, manages sessions and cleanup | | `sessions.py` | `Session`, `Connection`, `SessionRegistry` | player identity across reconnects; ordered sending | | `lobby.py` | `LobbyManager`, `Lobby`, `Seat`, `LobbyError` | lobby membership, host actions, running the game, AI turns, broadcasting | -| `ai_bridge.py` | `AIBridge`, `AIConnection` | talks to the AI service; random fallback | +| `ai_bridge.py` | `AIBridge`, `AIConnection` | talks to the AI service; random fallback, in-flight cap, circuit breaker | +| `limits.py` | `TokenBucket`, `KeyedRateLimiter`, `ConnectionLimiter` | abuse limits (message rate, connections, lobby creation) | +| `load.py` | `LoadMonitor`, `next_level`, cgroup helpers | detects when the server runs out of compute | ## Routes (`main.py`) @@ -249,7 +251,102 @@ The AI service keeps a websocket open to `/internal/ai`. `AIBridge`: the action is malformed or illegal (logged as warning, except the "no service connected" case which is logged at debug level); - calls `on_change` when a service connects/disconnects, and the service then - re-broadcasts all lobbies so the lobby screen shows whether the AI is online. + re-broadcasts all lobbies so the lobby screen shows whether the AI is online; +- sends `timeout` with every request so the AI service can drop requests the + server has already given up on; +- **in-flight cap**: if `AI_MAX_INFLIGHT` requests are already waiting on a + connection, the move is played by the fallback without asking; +- **circuit breaker**: after `AI_BREAKER_FAILURES` failures in a row (timeouts, + errors, illegal answers) the AI is not asked for `AI_BREAKER_COOLDOWN` + seconds, so a struggling AI service is not buried under more requests. + +## Limits and overload handling + +The server protects the host and its players in three layers. The guiding rule +is: **running games always continue; new work is refused when there is no +capacity; abuse is rejected immediately.** + +```mermaid +flowchart TD + C[new websocket] --> L1{"connection limit
total / per IP"} + L1 -- over --> R1["error busy + close 1013"] + L1 -- ok --> M[message] + M --> L2{"message rate
token bucket"} + L2 -- over --> R2["dropped + error rate_limited
(close 1008 when flooding)"] + L2 -- ok --> D{"create_lobby /
start_game?"} + D -- no --> OK[handled normally] + D -- yes --> L3{"accepting_games?
not overloaded and
running games < MAX_ACTIVE_GAMES"} + L3 -- no --> R3["error busy"] + L3 -- yes --> OK +``` + +### 1. Hard limits (Docker and uvicorn) + +- `compose.yaml` caps every container (`deploy.resources.limits` + `pids_limit`): + server 1 CPU / 512 MB, AI 1 CPU / 1 GB, website 0.25 CPU / 64 MB by default, + overridable in `.env` (`SERVER_CPUS`, `SERVER_MEMORY`, …). +- uvicorn (`wizard_server/Dockerfile`): `--ws-max-size 65536` (bigger messages + close the connection with 1009), `--ws-max-queue 16`, + `--limit-concurrency 1000` (HTTP 503 beyond that — a last resort above the + app's own connection limit). + +### 2. Abuse limits (`limits.py`, applied in `WizardService`) + +| Limit | Default | When exceeded | +|-------|---------|---------------| +| open connections total / per client IP | 500 / 10 | `error` (code `busy`), then close **1013** "Try Again Later" | +| messages per connection (token bucket) | 10/s, burst 30 | message dropped, `error` (code `rate_limited`, not repeated for every message); after 50 drops the connection is closed with 1008 | +| lobbies created per client IP | 10 per minute | `error` (code `rate_limited`) | + +The client IP is `websocket.client.host`. Behind Caddy, uvicorn's +`--proxy-headers` sets it from `X-Forwarded-For` (Caddy does not pass on +client-supplied values, so it cannot be spoofed). + +### 3. Load monitor and admission control (`load.py`) + +`LoadMonitor.run()` (started in the app lifespan) samples once per second: + +- **event loop lag** — how late the loop wakes up; the most direct "players + notice slowness" signal; +- **CPU** — process CPU time relative to the container's CPU quota (read from + the process's own cgroup, `/proc/self/cgroup` → `…/cpu.max`), at most one + core because game logic runs on one core; +- **throttling** — share of CPU periods in which the container hit its limit + (`cpu.stat`: `nr_throttled / nr_periods`). + +Each sample is classified (`classify`) and `next_level` derives the level with +hysteresis: + +| Level | Entered when | Left when | +|-------|--------------|-----------| +| `busy` | one sample with lag ≥ 0.1 s or CPU ≥ 70 % | 10 normal samples in a row | +| `overloaded` | 3 samples in a row with lag ≥ 0.25 s, CPU ≥ 90 % or throttled ≥ 50 % — or **one** sample with lag ≥ 1 s | 10 samples in a row that are not overloaded | + +`busy` is only informational. While **overloaded**, or while +`MAX_ACTIVE_GAMES` (default 50) games are running, `accepting_games` is false: + +| Action | Refused when not accepting games | +|--------|----------------------------------| +| `create_lobby` | yes (`error` code `busy`) | +| `start_game` | yes (`error` code `busy`) | +| `join_lobby`, moves, `return_to_lobby`, reconnects | never — they are cheap and keep running games alive | + +Refusing new games sheds load by itself: finished games are not replaced, so +the load drops and the server recovers. + +Every connected client receives a `server_status` message after `welcome` and +whenever the level or `accepting_games` changes (the website then disables +"Create lobby" / "Start game" and shows a notice). `/health` reports `load`, +`loop_lag`, `cpu`, `throttled`, `accepting_games`, `active_games`, +`connections` and AI bridge stats, but its `status` **stays "ok"** under load: +an overloaded server must never be restarted by Docker, which would end all +running games. + +Measured on a server capped at 15 % of a CPU with 30 games played at full +speed: overloaded after ~3 s, new lobbies refused, running games kept going, +back to normal ~10 s after the load dropped. + +For manual testing, `LOAD_FORCE_LEVEL=overloaded` pins the level. ## Configuration (`config.py`) @@ -265,27 +362,51 @@ All values are read from environment variables when the module is imported. | `MAX_LOBBIES` | 1000 | limit of simultaneously open lobbies | | `ALLOWED_ORIGINS` | empty (= all) | comma separated allowed `Origin` headers for `/ws` | | `LOG_LEVEL` | INFO | Python logging level | +| `MAX_CONNECTIONS` / `MAX_CONNECTIONS_PER_IP` | 500 / 10 | open websocket connections | +| `MESSAGE_RATE` / `MESSAGE_BURST` | 10 / 30 | messages per second per connection | +| `MAX_RATE_VIOLATIONS` | 50 | dropped messages before a flooding connection is closed | +| `LOBBY_CREATE_RATE` | 10 | lobbies per IP per minute | +| `MAX_ACTIVE_GAMES` | 50 | games running at the same time | +| `LOAD_SAMPLE_INTERVAL` | 1.0 s | load sampling period | +| `LOAD_LAG_BUSY` / `LOAD_LAG_OVERLOAD` / `LOAD_LAG_SEVERE` | 0.1 / 0.25 / 1.0 s | event loop lag thresholds | +| `LOAD_CPU_BUSY` / `LOAD_CPU_OVERLOAD` | 0.7 / 0.9 | CPU share of the quota | +| `LOAD_THROTTLE_OVERLOAD` | 0.5 | throttled share of CPU periods | +| `LOAD_OVERLOAD_SAMPLES` / `LOAD_RECOVER_SAMPLES` | 3 / 10 | hysteresis | +| `LOAD_FORCE_LEVEL` | empty | pin the level (testing) | +| `AI_MAX_INFLIGHT` | 32 | waiting AI requests per AI connection | +| `AI_BREAKER_FAILURES` / `AI_BREAKER_COOLDOWN` | 5 / 30 s | AI circuit breaker | ## Housekeeping `WizardService.cleanup_loop()` runs every 30 s (started in the app lifespan) -and removes expired sessions (leaving their lobby first). +and removes expired sessions (leaving their lobby first) and forgets idle +per-IP lobby-creation buckets. -## Tests (`wizard_server/tests/test_server.py`) +## Tests (`wizard_server/tests/`) -Uses FastAPI's `TestClient` websockets with zero delays. Covers: health, -hello requirement, validation errors, the full lobby flow including a complete -20-round game between two humans and an AI, reconnect by token, leaving -mid-game (seat becomes AI, host moves), the AI endpoint's secret and proxy -checks, and a fake AI service answering decide requests. +Uses FastAPI's `TestClient` websockets with zero delays and a raised message +rate limit (`helpers.TEST_SETTINGS`). -Note for writing new tests: enter websocket sessions with -`client.sockets.enter_context(client.websocket_connect(...))` (see the fixture), -otherwise the test client hangs on shutdown. +- `test_server.py` — health, hello requirement, validation errors, the full + lobby flow including a complete 20-round game between two humans and an AI, + reconnect by token, leaving mid-game (seat becomes AI, host moves), the AI + endpoint's secret and proxy checks, and a fake AI service. +- `test_protection.py` — message flood, connections per IP, lobby creation + rate, forced overload, `MAX_ACTIVE_GAMES`, a running game continuing while + overloaded, `server_status` messages. +- `test_limits.py` — unit tests of the token bucket, limiters, load levels, + cgroup parsing and the AI circuit breaker / in-flight cap. + +Fixtures (`conftest.py`): `client` (default test settings) and +`make_client(**overrides)` for custom settings. Enter websocket sessions with +`client.sockets.enter_context(client.websocket_connect(...))` (as +`helpers.connect` does), otherwise the test client hangs on shutdown. ## Limitations to be aware of - In-memory state: restarting the server ends all lobbies and games. - Must run as a single process (no `--workers`). - No spectators: only seated players receive game state. -- No rate limiting on websocket messages. +- A running game whose humans all disconnected keeps its `MAX_ACTIVE_GAMES` + slot until their sessions expire (`SESSION_TTL`), because they may still + reconnect. It costs no CPU while waiting. diff --git a/docs/05-protocol.md b/docs/05-protocol.md index 9b0860e..3618d33 100644 --- a/docs/05-protocol.md +++ b/docs/05-protocol.md @@ -93,25 +93,51 @@ contains the **complete** current state. - `ai_available` is `false` when no AI service is connected; AI seats then play random moves. +#### `server_status` + +Sent right after `welcome` and to every connected client whenever the server's +load level or capacity changes. + +```json +{"type": "server_status", "load": "overloaded", "accepting_games": false} +``` + +- `load`: `normal`, `busy` (informational) or `overloaded`. +- `accepting_games`: `false` while the server refuses `create_lobby` and + `start_game` (overloaded, or `MAX_ACTIVE_GAMES` games running). Everything + else, including running games, keeps working. + #### `error` ```json {"type": "error", "message": "Only the host can do that"} +{"type": "error", "message": "The server is busy right now, …", "code": "busy"} ``` +`code` is present for errors the client may react to: + +| `code` | Meaning | +|--------|---------| +| `busy` | no capacity: connection limit reached (followed by close 1013), or `create_lobby` / `start_game` refused because the server does not accept new games | +| `rate_limited` | too many messages on this connection, or too many lobbies created from this IP | + Sent only to the client that caused it. Typical messages: `Send 'hello' first`, `Invalid message: …` (followed by the pydantic reason), `Lobby ABC123 does not exist`, `The game in this lobby has already started`, `This lobby is full`, `Only the host can do that`, `Wizard needs 3 to 6 players`, `Illegal action Bid(value=2) for player 0`, `It is not player 1's turn`, `You were removed from the lobby`, -`This session was opened in another window`. +`This session was opened in another window`, `Too many messages, slow down.`, +`You are creating lobbies too quickly, please wait a moment.` ### Close codes | Code | Meaning | Frontend reaction | |------|---------|-------------------| | (handshake rejected) | origin not in `ALLOWED_ORIGINS`; the server closes with 1008 before accepting, so the browser sees a failed connection | reconnects with backoff | +| 1008 | message flood (too many dropped messages) | reconnects with backoff | +| 1009 | message larger than 64 KB (uvicorn `--ws-max-size`) | reconnects with backoff | +| 1013 | "Try Again Later": connection limit reached (preceded by an `error` with code `busy`) | banner "The server is full", retries slowly (15 s, 30 s, 60 s …) | | 4000 | the session was opened in another tab | stops reconnecting, shows a banner | | other | network problem, server restart | reconnects (1 s, 2 s, 4 s … max 10 s) | @@ -153,9 +179,9 @@ headers (see "Routes" in [the server chapter](04-server.md)). | Direction | Message | Meaning | |-----------|---------|---------| | AI → server | `{"type":"hello","agent":"heuristic"}` | informational, logged | -| server → AI | `{"type":"decide","request_id":"","observation":{…}}` | choose an action for `observation.player` | -| AI → server | `{"type":"action","request_id":"","action":{"type":"play","card":17}}` | the chosen action | -| AI → server | `{"type":"error","request_id":"","message":"…"}` | the AI failed; the server plays a random move | +| server → AI | `{"type":"decide","request_id":"","timeout":5.0,"observation":{…}}` | choose an action for `observation.player`; the server waits at most `timeout` seconds | +| AI → server | `{"type":"action","request_id":"","action":{"type":"play","card":17},"agent":"heuristic"}` | the chosen action; `agent` names the agent that decided (logged) | +| AI → server | `{"type":"error","request_id":"","message":"…"}` | the AI failed or refused (`overloaded`, `expired`); the server plays a random move | Rules: @@ -163,6 +189,11 @@ Rules: (matched by `request_id`) and may handle several concurrently. - The answer must be one of `observation.legal_actions`; otherwise the server plays a random legal move. -- No answer within `AI_DECISION_TIMEOUT` (5 s) → random move. +- No answer within `AI_DECISION_TIMEOUT` (5 s) → random move. The AI service + drops requests whose `timeout` passed before it could start them (`expired`). +- The server stops asking an AI connection that already has `AI_MAX_INFLIGHT` + unanswered requests, and pauses the AI after repeated failures (circuit + breaker); see "Limits and overload handling" in + [the server chapter](04-server.md). - Several AI service instances may connect; the server distributes requests round robin. diff --git a/docs/06-frontend.md b/docs/06-frontend.md index 86bbd29..9c5befd 100644 --- a/docs/06-frontend.md +++ b/docs/06-frontend.md @@ -53,7 +53,8 @@ The only stateful piece of the app. | Value | Meaning | |-------|---------| -| `status` | `connecting` \| `open` \| `closed` (will retry) \| `replaced` (opened in another tab, no retry) | +| `status` | `connecting` \| `open` \| `closed` (will retry) \| `replaced` (opened in another tab, no retry) \| `busy` (server full, retries slowly) | +| `serverStatus` | last `server_status` message: `{ load, accepting_games }` | | `playerId` | from `welcome` | | `lobby`, `game` | exactly what the last `state` message contained | | `completedTrick` | the last `trick_completed` event, kept for 2.2 s so the trick stays visible | @@ -73,6 +74,10 @@ The only stateful piece of the app. second tab is a different player. - On close it reconnects with exponential backoff (1, 2, 4, 8, 10, 10 … s), except after close code 4000. +- If the server refuses the connection because it is full (an `error` with + code `busy` before `welcome`, then close code 1013), the status becomes + `busy` and it retries slowly (15 s, 30 s, 60 s, …) so waiting clients don't + add load. - `send()` shows a toast if the socket is not open. ## Screen selection (`App.tsx`) @@ -96,6 +101,8 @@ shows the connection status when it is not `open`. - Code input + "Join" → `actions.joinLobby(code, name)`. The code is prefilled from `?lobby=` so invite links work. - Buttons are disabled while disconnected or without a name. +- While `serverStatus.accepting_games` is false, "Create lobby" is disabled and + a notice explains that the server is busy (joining still works). ### LobbyScreen @@ -107,6 +114,8 @@ shows the connection status when it is not `open`. seats are filled). Others see the options read-only and "Waiting for the host…". - Shows a warning if AI players are present but the AI service is offline. +- While `serverStatus.accepting_games` is false, "Start game" shows + "Server busy" (disabled) with a notice; the lobby stays open. - Explains the fixed bidding rule (bids may not add up to the cards in hand). ### GameScreen (`screens/GameScreen/`) diff --git a/docs/07-ai.md b/docs/07-ai.md index cc6e4fd..ad56cf3 100644 --- a/docs/07-ai.md +++ b/docs/07-ai.md @@ -194,8 +194,11 @@ registered with `@register_policy`, and calls its `load`. the container (`compose.yaml`); restart `wizard_ai`. With `AGENT=model` (default) it is loaded automatically. -Model size matters: the service answers every AI move in a worker thread and -the server waits at most `AI_DECISION_TIMEOUT` (5 s). +Model size matters: the service answers every AI move in a worker thread, the +container is capped at `AI_CPUS` (1 CPU by default) and the server waits at +most `AI_DECISION_TIMEOUT` (5 s). If the model is too slow the service +automatically answers with the heuristic instead (see below), so measure the +decision time before deploying. ## Training (`training/`) @@ -244,14 +247,28 @@ count as a win for everyone tied). Seeded for fair comparisons. ## Production service (`service/client.py`) -`AIService(server_url, agent, secret).run_forever()`: +`AIService(server_url, agent, secret, fallback=…, workers=…, max_queue=…, degrade_queue=…).run_forever()`: - connects to `SERVER_URL` with `Authorization: Bearer `, sends `hello`, then answers `decide` requests (protocol in [chapter 5](05-protocol.md#ai-protocol)); -- each request runs `agent.act` in a worker thread (`asyncio.to_thread`) behind - a lock, so a slow model does not block the websocket and non-thread-safe - agents stay consistent; +- each request runs `agent.act` in a worker thread (`asyncio.to_thread`), so a + slow model does not block the websocket; +- **bounded work** (`handle_request`): + - at most `AI_WORKERS` (default 1) decisions are computed at the same time — + agents must be thread safe if you raise it; + - at most `AI_MAX_QUEUE` (64) requests may be in progress; more are answered + with `error` "overloaded" immediately; + - a request whose server `timeout` passed while it waited for a worker is + answered with `error` "expired" without computing it; +- **degradation**: when `AI_DEGRADE_QUEUE` (8) or more other requests are + waiting, or the primary agent's average decision time is more than half of + the remaining time, the cheap fallback agent (`HeuristicAgent`) decides + instead. `serve` sets the heuristic as fallback automatically unless the + primary agent already is the heuristic. The reply's `agent` field says which + one decided; +- `OMP_NUM_THREADS` / `MKL_NUM_THREADS` are set to `AI_WORKERS` in + `compose.yaml`, so numeric libraries don't start a thread per host core; - errors while deciding are reported back as `error` (the server then plays a random move) instead of crashing; - reconnects forever with exponential backoff (1 s … 30 s); logs a hint if the @@ -271,10 +288,13 @@ Environment variables (`config.py`): `SERVER_URL` (default `ws://localhost:8000/internal/ai`; in Docker `ws://wizard_server:8083/internal/ai`), `AI_SHARED_SECRET`, `AGENT` (default `model`), `MODEL_PATH` (default `models/latest`, in Docker -`/srv/models/latest`), `LOG_LEVEL`. +`/srv/models/latest`), `LOG_LEVEL`, `AI_WORKERS` (1), `AI_MAX_QUEUE` (64), +`AI_DEGRADE_QUEUE` (8). ## Tests (`wizard_ai/tests/test_ai.py`) Feature length and masks for 3–6 players, action-space round trip, all agents play legal full games, env rewards add up to scores, trainer runs and writes a -loadable checkpoint, `make_agent` fallback, heuristic beats random. +loadable checkpoint, `make_agent` fallback, heuristic beats random, and the +service limits (queue full → overloaded, expired requests, degradation to the +heuristic when busy or when the model is too slow). diff --git a/docs/08-operations.md b/docs/08-operations.md index 1ac3327..eb81a43 100644 --- a/docs/08-operations.md +++ b/docs/08-operations.md @@ -44,8 +44,16 @@ not read by anything. | `ALLOWED_ORIGINS` | server | no | e.g. `https://game.marvinkrausser.com` | | `AI_AGENT` | AI | no (default `model`) | `model`, `heuristic` or `random` | | `LOG_LEVEL` | server + AI | no | `INFO`, `DEBUG`, … | +| `SERVER_CPUS` / `SERVER_MEMORY` | compose | no (1.0 / 512M) | hard caps of the server container | +| `AI_CPUS` / `AI_MEMORY` | compose | no (1.0 / 1G) | hard caps of the AI container | +| `WEBSITE_CPUS` / `WEBSITE_MEMORY` | compose | no (0.25 / 64M) | hard caps of the website container | +| `MAX_ACTIVE_GAMES` | server | no (50) | games running at the same time; more starts are refused ("server busy") | +| `MAX_CONNECTIONS` / `MAX_CONNECTIONS_PER_IP` | server | no (500 / 10) | open websocket connections | +| `AI_WORKERS` / `AI_MAX_QUEUE` / `AI_DEGRADE_QUEUE` | AI | no (1 / 64 / 8) | parallel decisions, waiting requests before refusing, queue depth from which the heuristic is used | -Further server settings (`AI_MOVE_DELAY`, `TRICK_PAUSE`, …, see +How the limits work together is described in "Limits and overload handling" +in [the server chapter](04-server.md). Further server settings +(`AI_MOVE_DELAY`, `TRICK_PAUSE`, load thresholds, …, see "Configuration" in [the server chapter](04-server.md)) can be added to the `environment:` section of `wizard_server` in `compose.yaml`. @@ -147,8 +155,8 @@ cd wizard_website && npm run lint && npm run build | Suite | Count | Covers | |-------|-------|--------| | game engine | 52 | rules, full random games for 3–6 players, invariants, JSON | -| server | 9 | websocket flows incl. a complete game, reconnect, AI endpoint security | -| AI | 13 | encoding, agents, env, trainer, checkpoints | +| server | 32 | websocket flows incl. a complete game, reconnect, AI endpoint security, abuse limits, overload handling, load levels | +| AI | 17 | encoding, agents, env, trainer, checkpoints, service limits | ## Troubleshooting @@ -161,6 +169,10 @@ cd wizard_website && npm run lint && npm run build | AI players only make weird/random moves | Same as above (random fallback), or the AI answers too slowly (`AI_DECISION_TIMEOUT`), see server logs "Playing random move for player …". | | Game waits and nothing happens | It is a disconnected human's turn (seat shows "offline"). The game waits until they reconnect or their session expires (`SESSION_TTL`, then an AI takes over). | | All lobbies gone | The server restarted — state is in memory only. | +| Start page / lobby says "The server is busy", Create/Start disabled | The server is overloaded or `MAX_ACTIVE_GAMES` games are running. Check `https://…/wizard/api/health` (`load`, `cpu`, `throttled`, `active_games`) and `docker stats`. Running games are not affected. If it happens often, raise `SERVER_CPUS` or `MAX_ACTIVE_GAMES`. | +| Banner "The server is full right now" | Connection limit reached (`MAX_CONNECTIONS`, or `MAX_CONNECTIONS_PER_IP` — e.g. many tabs behind one NAT). The page retries on its own every 15–60 s. | +| Error toast "Too many messages, slow down." | A client exceeded the message rate limit (10/s). Normal play never gets near it; a script or broken client might. | +| AI plays noticeably worse under load | The AI service switched to the heuristic because it was busy (`AI_DEGRADE_QUEUE`) or the model was too slow; see `docker compose logs wizard_ai` ("Busy … using heuristic agent"). Raise `AI_CPUS`/`AI_WORKERS` or use a smaller model. | | Tests pass but behaviour does not change after editing `wizard_game` | A non-editable copy of `wizard_game` is installed in your virtualenv and shadows the source. `pip uninstall wizard-game && pip install -e wizard_server/game_logic`. Check with `python -c "import wizard_game; print(wizard_game.__file__)"`. | | Server tests hang | A test opened a websocket without a context manager; use `client.sockets.enter_context(client.websocket_connect(...))`. | @@ -190,6 +202,17 @@ cd wizard_website && npm run lint && npm run build `actions`. 4. Add a test in `wizard_server/tests/test_server.py`. +### Check resource usage + +```bash +docker stats wizard-wizard_server-1 wizard-wizard_ai-1 wizard-wizard_website-1 +curl -s https://game.marvinkrausser.com/wizard/api/health +``` + +To test the overload behaviour without real load, start the server with +`LOAD_FORCE_LEVEL=overloaded` (new games are refused, the UI shows the busy +notice). + ### Change how fast AI players move Set `AI_MOVE_DELAY` (per move) and `TRICK_PAUSE` (after a trick) for diff --git a/wizard_ai/tests/test_ai.py b/wizard_ai/tests/test_ai.py index 5ebf6da..153d6e2 100644 --- a/wizard_ai/tests/test_ai.py +++ b/wizard_ai/tests/test_ai.py @@ -78,3 +78,82 @@ def test_make_agent_falls_back_without_checkpoint(tmp_path): def test_heuristic_beats_random(): result = evaluate([HeuristicAgent(), RandomAgent(1), RandomAgent(2), RandomAgent(3)], games=40, seed=0) assert result.mean_scores[0] > max(result.mean_scores[1:]) + + +# --- Service limits ------------------------------------------------------------- + +import asyncio # noqa: E402 +import threading # noqa: E402 + +from wizard_ai.service import AIService # noqa: E402 + + +class BlockingAgent(RandomAgent): + """Blocks in act() until released, to simulate a slow model.""" + + name = "slow-model" + + def __init__(self): + super().__init__(seed=0) + self.release = threading.Event() + + def act(self, observation): + self.release.wait(5) + return super().act(observation) + + +def decide_message(request_id, timeout=5.0): + game = WizardGame(3, seed=1) + return {"type": "decide", "request_id": request_id, "timeout": timeout, + "observation": game.observation(game.current_player).to_dict()} + + +def test_service_answers_with_agent_name(): + service = AIService("ws://unused", HeuristicAgent()) + reply = asyncio.run(service.handle_request(decide_message("r1"))) + assert reply["type"] == "action" and reply["request_id"] == "r1" and reply["agent"] == "heuristic" + + +def test_service_refuses_when_queue_is_full_and_degrades_when_busy(): + async def scenario(): + slow = BlockingAgent() + service = AIService("ws://unused", slow, fallback=HeuristicAgent(), workers=1, max_queue=3, degrade_queue=1) + first = asyncio.create_task(service.handle_request(decide_message("a"))) + await asyncio.sleep(0.05) # "a" is computing with the slow agent + second = asyncio.create_task(service.handle_request(decide_message("b"))) + third = asyncio.create_task(service.handle_request(decide_message("c"))) + await asyncio.sleep(0.05) + refused = await service.handle_request(decide_message("d")) + assert refused == {"type": "error", "request_id": "d", "message": "overloaded"} + slow.release.set() + replies = await asyncio.gather(first, second, third) + # "b" starts while "c" is still waiting -> degraded; "c" runs when nothing waits anymore. + assert [r["agent"] for r in replies] == ["slow-model", "heuristic", "slow-model"] + assert service.in_progress == 0 + + asyncio.run(scenario()) + + +def test_service_drops_expired_requests(): + async def scenario(): + slow = BlockingAgent() + service = AIService("ws://unused", slow, workers=1) + first = asyncio.create_task(service.handle_request(decide_message("a"))) + await asyncio.sleep(0.05) + waiting = asyncio.create_task(service.handle_request(decide_message("b", timeout=0.1))) + await asyncio.sleep(0.2) # "b" expires while "a" blocks the only worker + slow.release.set() + await first + assert (await waiting)["message"] == "expired" + + asyncio.run(scenario()) + + +def test_service_uses_fallback_when_primary_is_too_slow(): + async def scenario(): + service = AIService("ws://unused", RandomAgent(1), fallback=HeuristicAgent()) + service._primary_latency = 3.0 # model needs 3 s, server waits 5 s -> more than half + reply = await service.handle_request(decide_message("a", timeout=5.0)) + assert reply["agent"] == "heuristic" + + asyncio.run(scenario()) diff --git a/wizard_ai/wizard_ai/__main__.py b/wizard_ai/wizard_ai/__main__.py index 85de975..6900bd9 100644 --- a/wizard_ai/wizard_ai/__main__.py +++ b/wizard_ai/wizard_ai/__main__.py @@ -49,10 +49,20 @@ def main() -> None: logging.basicConfig(level=args.log_level.upper(), format="%(asctime)s %(levelname)s %(name)s: %(message)s") if args.command == "serve": + from .agents import HeuristicAgent from .service import AIService agent = make_agent(args.agent, args.model_path) - asyncio.run(AIService(args.server_url, agent, config.secret).run_forever()) + # Cheap agent used when the service is busy (pointless if it already is the heuristic). + fallback = None if isinstance(agent, HeuristicAgent) else HeuristicAgent() + service = AIService( + args.server_url, agent, config.secret, + fallback=fallback, + workers=config.workers, + max_queue=config.max_queue, + degrade_queue=config.degrade_queue, + ) + asyncio.run(service.run_forever()) elif args.command == "train": from .model import UniformPolicy, checkpoint_exists, load_policy diff --git a/wizard_ai/wizard_ai/config.py b/wizard_ai/wizard_ai/config.py index 68f6a12..36ab6fe 100644 --- a/wizard_ai/wizard_ai/config.py +++ b/wizard_ai/wizard_ai/config.py @@ -17,3 +17,9 @@ class Config: # Checkpoint directory of the trained policy. model_path: str = os.environ.get("MODEL_PATH", "models/latest") log_level: str = os.environ.get("LOG_LEVEL", "INFO") + # Decisions computed in parallel (agents must be thread safe if > 1). + workers: int = int(os.environ.get("AI_WORKERS", 1)) + # Requests allowed to wait; more are refused with "overloaded". + max_queue: int = int(os.environ.get("AI_MAX_QUEUE", 64)) + # Waiting requests from which the heuristic agent is used instead of the model. + degrade_queue: int = int(os.environ.get("AI_DEGRADE_QUEUE", 8)) diff --git a/wizard_ai/wizard_ai/service/client.py b/wizard_ai/wizard_ai/service/client.py index f5e12d9..d09ee02 100644 --- a/wizard_ai/wizard_ai/service/client.py +++ b/wizard_ai/wizard_ai/service/client.py @@ -3,6 +3,18 @@ The service connects to the server's ``/internal/ai`` endpoint and answers ``decide`` requests (see ``wizard_server/app/ai_bridge.py``). It reconnects automatically if the connection drops. + +Work is bounded so a busy service degrades instead of falling over: + +- at most ``workers`` decisions are computed at the same time; +- at most ``max_queue`` requests may wait; more are answered with the error + "overloaded" right away (the server then plays a fallback move); +- a request that waited longer than the server's ``timeout`` is answered with + "expired" without computing it (the server has already moved on); +- when ``degrade_queue`` or more other requests are waiting, or the primary + agent typically needs more than half of the remaining time, the cheap + ``fallback`` agent (heuristic) decides instead of the primary agent (e.g. a + neural network). """ from __future__ import annotations @@ -10,6 +22,8 @@ from __future__ import annotations import asyncio import json import logging +import time +from collections.abc import Callable from typing import Any from websockets.asyncio.client import ClientConnection, connect @@ -22,16 +36,36 @@ from ..agents import Agent logger = logging.getLogger(__name__) MAX_BACKOFF = 30.0 +# Weight of the newest measurement in the average decision time. +LATENCY_SMOOTHING = 0.2 class AIService: - def __init__(self, server_url: str, agent: Agent, secret: str = "") -> None: + def __init__( + self, + server_url: str, + agent: Agent, + secret: str = "", + *, + fallback: Agent | None = None, + workers: int = 1, + max_queue: int = 64, + degrade_queue: int = 8, + clock: Callable[[], float] = time.monotonic, + ) -> None: self.server_url = server_url self.agent = agent + self.fallback = fallback self.secret = secret - # Inference runs in worker threads; the lock keeps agents that are not - # thread safe (e.g. stateful models) consistent. - self._agent_lock = asyncio.Lock() + self.max_queue = max_queue + self.degrade_queue = degrade_queue + self._clock = clock + # Agents run in worker threads; with workers > 1 they must be thread safe. + self._workers = asyncio.Semaphore(max(1, workers)) + # Requests received and not answered yet (waiting + computing). + self.in_progress = 0 + # Smoothed seconds the primary agent needs per decision. + self._primary_latency = 0.0 async def run_forever(self) -> None: backoff = 1.0 @@ -55,21 +89,62 @@ class AIService: async for raw in websocket: message = json.loads(raw) if message.get("type") == "decide": - task = asyncio.create_task(self._decide(websocket, message)) + task = asyncio.create_task(self._answer(websocket, message)) tasks.add(task) task.add_done_callback(tasks.discard) - async def _decide(self, websocket: ClientConnection, message: dict[str, Any]) -> None: - request_id = message.get("request_id") - try: - observation = Observation.from_dict(message["observation"]) - async with self._agent_lock: - action = await asyncio.to_thread(self.agent.act, observation) - reply = {"type": "action", "request_id": request_id, "action": action.to_dict()} - except Exception as exc: # report instead of crashing; the server falls back - logger.exception("Failed to decide request %s", request_id) - reply = {"type": "error", "request_id": request_id, "message": str(exc)} + async def _answer(self, websocket: ClientConnection, message: dict[str, Any]) -> None: + reply = await self.handle_request(message) try: await websocket.send(json.dumps(reply)) except ConnectionClosed: pass + + async def handle_request(self, message: dict[str, Any]) -> dict[str, Any]: + """Answer one ``decide`` request (independent of the websocket, for tests).""" + request_id = message.get("request_id") + received = self._clock() + timeout = message.get("timeout") + deadline = received + timeout if isinstance(timeout, (int, float)) and timeout > 0 else None + + if self.in_progress >= self.max_queue: + logger.warning("Overloaded: %d requests in progress, refusing %s", self.in_progress, request_id) + return self._error(request_id, "overloaded") + + self.in_progress += 1 + try: + async with self._workers: + now = self._clock() + if deadline is not None and now >= deadline: + logger.info("Dropping expired request %s (waited %.2fs)", request_id, now - received) + return self._error(request_id, "expired") + agent = self._choose_agent(deadline, now) + observation = Observation.from_dict(message["observation"]) + started = self._clock() + action = await asyncio.to_thread(agent.act, observation) + if agent is self.agent: + elapsed = self._clock() - started + self._primary_latency += LATENCY_SMOOTHING * (elapsed - self._primary_latency) + return {"type": "action", "request_id": request_id, "action": action.to_dict(), "agent": agent.name} + except Exception as exc: # report instead of crashing; the server falls back + logger.exception("Failed to decide request %s", request_id) + return self._error(request_id, str(exc) or type(exc).__name__) + finally: + self.in_progress -= 1 + + def _choose_agent(self, deadline: float | None, now: float) -> Agent: + """The primary agent, or the cheap fallback when the service is busy.""" + if self.fallback is None: + return self.agent + waiting = self.in_progress - 1 # requests besides this one + if waiting >= self.degrade_queue: + logger.info("Busy (%d waiting): using %s agent", waiting, self.fallback.name) + return self.fallback + if deadline is not None and self._primary_latency > (deadline - now) / 2: + logger.info("Slow (%.2fs per decision): using %s agent", self._primary_latency, self.fallback.name) + return self.fallback + return self.agent + + @staticmethod + def _error(request_id: Any, message: str) -> dict[str, Any]: + return {"type": "error", "request_id": request_id, "message": message} diff --git a/wizard_server/Dockerfile b/wizard_server/Dockerfile index 5abd4c9..def8dc0 100644 --- a/wizard_server/Dockerfile +++ b/wizard_server/Dockerfile @@ -23,4 +23,8 @@ EXPOSE 8083 HEALTHCHECK --interval=30s --timeout=3s --start-period=5s \ CMD python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8083/health', timeout=2)" -CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8083", "--proxy-headers", "--forwarded-allow-ips", "*"] +# Transport limits as a last line of defence (the app enforces friendlier limits first): +# messages are < 1 KB, so 64 KB is plenty; at most 1000 concurrent connections. +CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8083", \ + "--proxy-headers", "--forwarded-allow-ips", "*", \ + "--ws-max-size", "65536", "--ws-max-queue", "16", "--limit-concurrency", "1000"] diff --git a/wizard_server/app/ai_bridge.py b/wizard_server/app/ai_bridge.py index e81cac6..e66a47b 100644 --- a/wizard_server/app/ai_bridge.py +++ b/wizard_server/app/ai_bridge.py @@ -4,14 +4,20 @@ The AI service (``wizard_ai``) connects to ``/internal/ai`` via websocket. Whenever an AI seat has to act, the server sends it the seat's observation and waits for an action: - server -> ai: {"type": "decide", "request_id": "...", "observation": {...}} - ai -> server: {"type": "action", "request_id": "...", "action": {...}} + server -> ai: {"type": "decide", "request_id": "...", "timeout": 5.0, "observation": {...}} + ai -> server: {"type": "action", "request_id": "...", "action": {...}, "agent": "heuristic"} ai -> server: {"type": "error", "request_id": "...", "message": "..."} -Several AI service instances may connect; requests are distributed round -robin. If no instance is connected, it does not answer in time, or it answers -with an illegal action, a random legal action is played instead so a game -never gets stuck. +``timeout`` tells the AI service how long the server waits, so it can drop +requests that are already too old. Several AI service instances may connect; +requests are distributed round robin. + +A random legal action is played instead (so a game never gets stuck) when no +instance is connected, it does not answer in time, it answers with an error or +an illegal action, too many requests are already waiting for it +(``max_inflight``), or the circuit breaker is open: after ``breaker_failures`` +failures in a row the AI is not asked for ``breaker_cooldown`` seconds, so a +struggling AI service is not buried under even more requests. """ from __future__ import annotations @@ -55,6 +61,9 @@ class AIBridge: timeout: float, rng: random.Random | None = None, on_change: Callable[[], None] | None = None, + max_inflight: int = 32, + breaker_failures: int = 5, + breaker_cooldown: float = 30.0, ) -> None: self.timeout = timeout # Called whenever an AI service connects or disconnects. @@ -62,6 +71,11 @@ class AIBridge: self._connections: list[AIConnection] = [] self._round_robin = itertools.count() self._rng = rng or random.Random() + self._max_inflight = max_inflight + self._breaker_failures = breaker_failures + self._breaker_cooldown = breaker_cooldown + self._consecutive_failures = 0 + self._breaker_open_until = 0.0 @property def available(self) -> bool: @@ -76,32 +90,49 @@ class AIBridge: connection = self._pick_connection() if connection is None: return self._fallback(observation, "no AI service connected", logging.DEBUG) + loop = asyncio.get_running_loop() + if loop.time() < self._breaker_open_until: + return self._fallback(observation, "AI paused after repeated failures", logging.DEBUG) + if len(connection.pending) >= self._max_inflight: + return self._failure(observation, f"{len(connection.pending)} AI requests already waiting") request_id = uuid.uuid4().hex - future: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future() + future: asyncio.Future[dict[str, Any]] = loop.create_future() connection.pending[request_id] = future try: await connection.send({ "type": "decide", "request_id": request_id, + "timeout": self.timeout, "observation": observation.to_dict(), }) reply = await asyncio.wait_for(future, self.timeout) except (TimeoutError, ConnectionError, RuntimeError, WebSocketDisconnect) as exc: - return self._fallback(observation, f"AI request failed: {exc!r}") + return self._failure(observation, f"AI request failed: {exc!r}") finally: connection.pending.pop(request_id, None) if reply.get("type") != "action": - return self._fallback(observation, f"AI error: {reply.get('message')}") + return self._failure(observation, f"AI error: {reply.get('message')}") try: action = action_from_dict(reply.get("action", {})) except ValueError as exc: - return self._fallback(observation, f"malformed AI action: {exc}") + return self._failure(observation, f"malformed AI action: {exc}") if action not in legal: - return self._fallback(observation, f"illegal AI action {action}") + return self._failure(observation, f"illegal AI action {action}") + self._consecutive_failures = 0 + if reply.get("agent"): + logger.debug("AI decision by %s agent", reply["agent"]) return action + def stats(self) -> dict[str, Any]: + loop_time = asyncio.get_running_loop().time() + return { + "connections": len(self._connections), + "pending": sum(len(c.pending) for c in self._connections), + "paused": loop_time < self._breaker_open_until, + } + async def serve(self, websocket: WebSocket) -> None: """Handle one AI service connection until it disconnects.""" connection = AIConnection(websocket, name=websocket.client.host if websocket.client else "?") @@ -136,6 +167,16 @@ class AIBridge: return None return self._connections[next(self._round_robin) % len(self._connections)] + def _failure(self, observation: Observation, reason: str) -> Action: + """A failed AI request: count it for the circuit breaker, then fall back.""" + self._consecutive_failures += 1 + if self._consecutive_failures >= self._breaker_failures: + self._consecutive_failures = 0 + self._breaker_open_until = asyncio.get_running_loop().time() + self._breaker_cooldown + logger.warning("AI failed %d times in a row, pausing AI requests for %.0fs", + self._breaker_failures, self._breaker_cooldown) + return self._fallback(observation, reason) + def _fallback(self, observation: Observation, reason: str, level: int = logging.WARNING) -> Action: logger.log(level, "Playing random move for player %d: %s", observation.player, reason) return self._rng.choice(observation.legal_actions) diff --git a/wizard_server/app/config.py b/wizard_server/app/config.py index 2426c28..9bd1de8 100644 --- a/wizard_server/app/config.py +++ b/wizard_server/app/config.py @@ -35,5 +35,44 @@ class Settings: ) log_level: str = os.environ.get("LOG_LEVEL", "INFO") + # --- Abuse limits ------------------------------------------------------- + # Open websocket connections in total and per client IP. + max_connections: int = _int("MAX_CONNECTIONS", 500) + max_connections_per_ip: int = _int("MAX_CONNECTIONS_PER_IP", 10) + # Messages per second a single connection may send (token bucket). + message_rate: float = _float("MESSAGE_RATE", 10.0) + message_burst: int = _int("MESSAGE_BURST", 30) + # Dropped messages after which a flooding connection is closed. + max_rate_violations: int = _int("MAX_RATE_VIOLATIONS", 50) + # Lobbies a single IP may create per minute. + lobby_create_rate: int = _int("LOBBY_CREATE_RATE", 10) + + # --- Load control ------------------------------------------------------- + # Games running at the same time; further game starts are refused. + max_active_games: int = _int("MAX_ACTIVE_GAMES", 50) + load_sample_interval: float = _float("LOAD_SAMPLE_INTERVAL", 1.0) + # Event loop lag (seconds) / CPU share of the available quota (0..1). + load_lag_busy: float = _float("LOAD_LAG_BUSY", 0.1) + load_lag_overload: float = _float("LOAD_LAG_OVERLOAD", 0.25) + # A single sample with this much lag means overload right away: when the + # server is that slow, even the load sampler runs rarely. + load_lag_severe: float = _float("LOAD_LAG_SEVERE", 1.0) + load_cpu_busy: float = _float("LOAD_CPU_BUSY", 0.7) + load_cpu_overload: float = _float("LOAD_CPU_OVERLOAD", 0.9) + # Share of CPU periods in which the container was throttled (cgroup v2). + load_throttle_overload: float = _float("LOAD_THROTTLE_OVERLOAD", 0.5) + # Consecutive overloaded samples to enter / good samples to leave overload. + load_overload_samples: int = _int("LOAD_OVERLOAD_SAMPLES", 3) + load_recover_samples: int = _int("LOAD_RECOVER_SAMPLES", 10) + # "normal", "busy" or "overloaded" to pin the level (testing only). + load_force_level: str = os.environ.get("LOAD_FORCE_LEVEL", "") + + # --- AI service protection ---------------------------------------------- + # Unanswered decide requests per AI connection before falling back. + ai_max_inflight: int = _int("AI_MAX_INFLIGHT", 32) + # Consecutive AI failures that pause AI requests for ai_breaker_cooldown seconds. + ai_breaker_failures: int = _int("AI_BREAKER_FAILURES", 5) + ai_breaker_cooldown: float = _float("AI_BREAKER_COOLDOWN", 30.0) + settings = Settings() diff --git a/wizard_server/app/limits.py b/wizard_server/app/limits.py new file mode 100644 index 0000000..248f485 --- /dev/null +++ b/wizard_server/app/limits.py @@ -0,0 +1,90 @@ +"""Rate and connection limits that protect the server from floods and abuse.""" + +from __future__ import annotations + +import time +from collections import Counter +from collections.abc import Callable, Hashable + +Clock = Callable[[], float] + + +class TokenBucket: + """Allows ``rate`` actions per second on average with bursts up to ``burst``.""" + + def __init__(self, rate: float, burst: float, clock: Clock = time.monotonic) -> None: + self.rate = rate + self.burst = burst + self._clock = clock + self._tokens = float(burst) + self._updated = clock() + + def allow(self, cost: float = 1.0) -> bool: + self._refill() + if self._tokens >= cost: + self._tokens -= cost + return True + return False + + @property + def full(self) -> bool: + self._refill() + return self._tokens >= self.burst + + def _refill(self) -> None: + now = self._clock() + self._tokens = min(self.burst, self._tokens + (now - self._updated) * self.rate) + self._updated = now + + +class KeyedRateLimiter: + """One token bucket per key (e.g. per client IP).""" + + def __init__(self, rate: float, burst: float, clock: Clock = time.monotonic) -> None: + self.rate = rate + self.burst = burst + self._clock = clock + self._buckets: dict[Hashable, TokenBucket] = {} + + def allow(self, key: Hashable) -> bool: + bucket = self._buckets.get(key) + if bucket is None: + bucket = self._buckets[key] = TokenBucket(self.rate, self.burst, self._clock) + return bucket.allow() + + def prune(self) -> None: + """Forget keys whose bucket has refilled completely (they behave like new keys).""" + for key in [key for key, bucket in self._buckets.items() if bucket.full]: + del self._buckets[key] + + def __len__(self) -> int: + return len(self._buckets) + + +class ConnectionLimiter: + """Counts open connections in total and per client IP.""" + + def __init__(self, max_total: int, max_per_ip: int) -> None: + self.max_total = max_total + self.max_per_ip = max_per_ip + self._per_ip: Counter[str] = Counter() + self._total = 0 + + @property + def total(self) -> int: + return self._total + + def acquire(self, ip: str) -> bool: + if self._total >= self.max_total or self._per_ip[ip] >= self.max_per_ip: + return False + self._per_ip[ip] += 1 + self._total += 1 + return True + + def release(self, ip: str) -> None: + if self._per_ip[ip] <= 0: + return + self._per_ip[ip] -= 1 + self._total -= 1 + if self._per_ip[ip] == 0: + del self._per_ip[ip] diff --git a/wizard_server/app/load.py b/wizard_server/app/load.py new file mode 100644 index 0000000..59fadb7 --- /dev/null +++ b/wizard_server/app/load.py @@ -0,0 +1,248 @@ +"""Load monitoring: detects when the server runs out of compute. + +Every ``load_sample_interval`` seconds the monitor measures + +- **event loop lag**: how much later than scheduled the loop woke up. This is + the most direct signal that the server is slow for players; +- **CPU usage** of the process relative to what it may use (the container's + CPU quota, at most one core because Python runs game logic on one core); +- **CPU throttling** from cgroup v2 (``/sys/fs/cgroup/cpu.stat``): the share + of scheduling periods in which the container hit its CPU limit. + +Each sample is classified as normal / busy / overloaded, and ``next_level`` +turns the recent classifications into the current level with hysteresis, so a +single slow second does not flip the state back and forth. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from collections import deque +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from enum import Enum +from pathlib import Path + +from .config import Settings + +logger = logging.getLogger(__name__) + +CGROUP_DIR = Path("/sys/fs/cgroup") + + +class LoadLevel(str, Enum): + NORMAL = "normal" + BUSY = "busy" + OVERLOADED = "overloaded" + + +@dataclass(frozen=True, slots=True) +class LoadSample: + loop_lag: float + # Share of the available CPU used (None if unknown). + cpu: float | None = None + # Share of CPU periods in which the container was throttled (None if unknown). + throttled: float | None = None + + +@dataclass(frozen=True, slots=True) +class LoadThresholds: + lag_busy: float + lag_overload: float + lag_severe: float + cpu_busy: float + cpu_overload: float + throttle_overload: float + overload_samples: int + recover_samples: int + + @classmethod + def from_settings(cls, settings: Settings) -> LoadThresholds: + return cls( + lag_busy=settings.load_lag_busy, + lag_overload=settings.load_lag_overload, + lag_severe=settings.load_lag_severe, + cpu_busy=settings.load_cpu_busy, + cpu_overload=settings.load_cpu_overload, + throttle_overload=settings.load_throttle_overload, + overload_samples=max(1, settings.load_overload_samples), + recover_samples=max(1, settings.load_recover_samples), + ) + + +def classify(sample: LoadSample, t: LoadThresholds) -> LoadLevel: + """Level of a single sample.""" + if ( + sample.loop_lag >= t.lag_overload + or (sample.cpu is not None and sample.cpu >= t.cpu_overload) + or (sample.throttled is not None and sample.throttled >= t.throttle_overload) + ): + return LoadLevel.OVERLOADED + if sample.loop_lag >= t.lag_busy or (sample.cpu is not None and sample.cpu >= t.cpu_busy): + return LoadLevel.BUSY + return LoadLevel.NORMAL + + +def next_level( + current: LoadLevel, recent: Sequence[LoadLevel], t: LoadThresholds, severe: bool = False +) -> LoadLevel: + """Current level from the recent sample levels (oldest first), with hysteresis. + + - Overloaded after ``overload_samples`` overloaded samples in a row, or at + once when the latest sample was ``severe`` (extreme event loop lag). + - Leaving overload (and busy) requires ``recover_samples`` samples in a row + below that level. + - Busy is entered immediately (it is only a hint for players). + """ + if severe: + return LoadLevel.OVERLOADED + if not recent: + return current + last_bad = recent[-t.overload_samples:] + if len(last_bad) == t.overload_samples and all(level is LoadLevel.OVERLOADED for level in last_bad): + return LoadLevel.OVERLOADED + + window = recent[-t.recover_samples:] + full_window = len(window) == t.recover_samples + if current is LoadLevel.OVERLOADED and not ( + full_window and all(level is not LoadLevel.OVERLOADED for level in window) + ): + return LoadLevel.OVERLOADED + if recent[-1] is not LoadLevel.NORMAL: + return LoadLevel.BUSY + if current is not LoadLevel.NORMAL and not (full_window and all(level is LoadLevel.NORMAL for level in window)): + return LoadLevel.BUSY + return LoadLevel.NORMAL + + +# --- cgroup v2 helpers ------------------------------------------------------ + + +def parse_cpu_max(text: str) -> float | None: + """CPU quota in cores from ``cpu.max`` ("max 100000" = unlimited).""" + parts = text.split() + if len(parts) != 2 or parts[0] == "max": + return None + try: + return int(parts[0]) / int(parts[1]) + except (ValueError, ZeroDivisionError): + return None + + +def parse_cpu_stat(text: str) -> dict[str, int]: + stats: dict[str, int] = {} + for line in text.splitlines(): + key, _, value = line.partition(" ") + if value.strip().isdigit(): + stats[key] = int(value) + return stats + + +def own_cgroup_dir(root: Path = CGROUP_DIR, proc_cgroup: Path = Path("/proc/self/cgroup")) -> Path: + """The cgroup v2 directory of this process. + + In a container this is ``/sys/fs/cgroup`` itself (cgroup namespace); on a + host it is a sub directory such as ``/sys/fs/cgroup/user.slice/…scope``. + """ + for line in (_read(proc_cgroup) or "").splitlines(): + if line.startswith("0::"): + candidate = root / line[3:].strip().lstrip("/") + if candidate.is_dir(): + return candidate + return root + + +def _read(path: Path) -> str | None: + try: + return path.read_text() + except OSError: + return None + + +class LoadMonitor: + def __init__( + self, + settings: Settings, + on_change: Callable[[LoadLevel], None] | None = None, + cgroup_dir: Path | None = None, + ) -> None: + self._settings = settings + self._thresholds = LoadThresholds.from_settings(settings) + self._on_change = on_change + self._cgroup_dir = cgroup_dir if cgroup_dir is not None else own_cgroup_dir() + self._forced = LoadLevel(settings.load_force_level) if settings.load_force_level else None + self._level = LoadLevel.NORMAL + self._recent: deque[LoadLevel] = deque( + maxlen=max(self._thresholds.overload_samples, self._thresholds.recover_samples) + ) + self.last_sample = LoadSample(loop_lag=0.0) + + quota = parse_cpu_max(_read(self._cgroup_dir / "cpu.max") or "") + # Game logic runs on one core, so more than one core cannot be used. + self._cpu_capacity = min(quota, 1.0) if quota else 1.0 + self._last_cpu_time = time.process_time() + self._last_wall = time.monotonic() + self._last_cgroup = self._cgroup_stats() + + @property + def level(self) -> LoadLevel: + return self._forced or self._level + + def record(self, sample: LoadSample) -> None: + """Add a sample and update the level (called by ``run``; public for tests).""" + self.last_sample = sample + self._recent.append(classify(sample, self._thresholds)) + severe = sample.loop_lag >= self._thresholds.lag_severe + new_level = next_level(self._level, list(self._recent), self._thresholds, severe) + if new_level is self._level: + return + logger.log( + logging.WARNING if new_level is LoadLevel.OVERLOADED else logging.INFO, + "Load level %s -> %s (lag %.3fs, cpu %s, throttled %s)", + self._level.value, new_level.value, sample.loop_lag, sample.cpu, sample.throttled, + ) + self._level = new_level + if self._on_change is not None and self._forced is None: + self._on_change(new_level) + + async def run(self) -> None: + loop = asyncio.get_running_loop() + interval = self._settings.load_sample_interval + while True: + started = loop.time() + await asyncio.sleep(interval) + lag = max(0.0, loop.time() - started - interval) + try: + self.record(self._measure(lag)) + except Exception: + logger.exception("Load measurement failed") + + def stats(self) -> dict[str, object]: + sample = self.last_sample + return { + "load": self.level.value, + "loop_lag": round(sample.loop_lag, 4), + "cpu": None if sample.cpu is None else round(sample.cpu, 3), + "throttled": None if sample.throttled is None else round(sample.throttled, 3), + } + + def _measure(self, lag: float) -> LoadSample: + now_cpu, now_wall = time.process_time(), time.monotonic() + wall = now_wall - self._last_wall + cpu = (now_cpu - self._last_cpu_time) / wall / self._cpu_capacity if wall > 0 else None + self._last_cpu_time, self._last_wall = now_cpu, now_wall + + throttled = None + stats = self._cgroup_stats() + if stats and self._last_cgroup: + periods = stats.get("nr_periods", 0) - self._last_cgroup.get("nr_periods", 0) + if periods > 0: + throttled = (stats.get("nr_throttled", 0) - self._last_cgroup.get("nr_throttled", 0)) / periods + self._last_cgroup = stats + return LoadSample(loop_lag=lag, cpu=cpu, throttled=throttled) + + def _cgroup_stats(self) -> dict[str, int]: + text = _read(self._cgroup_dir / "cpu.stat") + return parse_cpu_stat(text) if text else {} diff --git a/wizard_server/app/lobby.py b/wizard_server/app/lobby.py index 4498337..31799d4 100644 --- a/wizard_server/app/lobby.py +++ b/wizard_server/app/lobby.py @@ -5,7 +5,7 @@ from __future__ import annotations import asyncio import logging import secrets -from collections.abc import Iterable +from collections.abc import Callable, Iterable from dataclasses import dataclass from enum import Enum from typing import Any @@ -24,6 +24,7 @@ from wizard_game import ( from .ai_bridge import AIBridge from .config import Settings +from .protocol import ERROR_BUSY from .sessions import Session logger = logging.getLogger(__name__) @@ -40,6 +41,11 @@ AI_NAMES = ( class LobbyError(Exception): """An error that is reported to the client that caused it.""" + def __init__(self, message: str, code: str | None = None) -> None: + super().__init__(message) + # Optional machine readable code, see ``protocol.ERROR_*``. + self.code = code + class LobbyStatus(str, Enum): WAITING = "waiting" @@ -60,8 +66,16 @@ class Seat: class Lobby: - def __init__(self, code: str, bridge: AIBridge, settings: Settings) -> None: + def __init__( + self, + code: str, + bridge: AIBridge, + settings: Settings, + on_status_change: Callable[[], None] | None = None, + ) -> None: self.code = code + # Called when a game starts or ends (the number of running games changes). + self._on_status_change = on_status_change self.seats: list[Seat] = [] self.host_id: str | None = None self.options = GameOptions() @@ -188,6 +202,7 @@ class Lobby: self.status = LobbyStatus.PLAYING self._pause_until = 0.0 logger.info("Lobby %s: game started with %d players", self.code, len(self.seats)) + self._notify_status_change() self._schedule_ai() def return_to_lobby(self, by: Session) -> None: @@ -224,6 +239,9 @@ class Lobby: self._closed = True if self._ai_task and not self._ai_task.done(): self._ai_task.cancel() + if self.status is LobbyStatus.PLAYING: + self.status = LobbyStatus.FINISHED + self._notify_status_change() # --- Internals ---------------------------------------------------------- @@ -247,12 +265,19 @@ class Lobby: loop = asyncio.get_running_loop() if any(isinstance(e, TrickCompleted) for e in events): self._pause_until = loop.time() + self._settings.trick_pause - if any(isinstance(e, GameFinished) for e in events): + finished = any(isinstance(e, GameFinished) for e in events) + if finished: self.status = LobbyStatus.FINISHED logger.info("Lobby %s: game finished", self.code) self.broadcast(events) + if finished: + self._notify_status_change() self._schedule_ai() + def _notify_status_change(self) -> None: + if self._on_status_change is not None: + self._on_status_change() + def _current_ai_seat(self) -> int | None: game = self.game if self._closed or self.status is not LobbyStatus.PLAYING or game is None: @@ -292,9 +317,15 @@ class Lobby: class LobbyManager: - def __init__(self, bridge: AIBridge, settings: Settings) -> None: + def __init__( + self, + bridge: AIBridge, + settings: Settings, + on_games_changed: Callable[[], None] | None = None, + ) -> None: self._bridge = bridge self._settings = settings + self._on_games_changed = on_games_changed self._lobbies: dict[str, Lobby] = {} def __len__(self) -> int: @@ -303,14 +334,17 @@ class LobbyManager: def __iter__(self): return iter(list(self._lobbies.values())) + def active_games(self) -> int: + return sum(lobby.status is LobbyStatus.PLAYING for lobby in self._lobbies.values()) + def create(self) -> Lobby: if len(self._lobbies) >= self._settings.max_lobbies: - raise LobbyError("Too many open lobbies, please try again later") + raise LobbyError("Too many open lobbies, please try again later", ERROR_BUSY) while True: code = "".join(secrets.choice(LOBBY_CODE_ALPHABET) for _ in range(LOBBY_CODE_LENGTH)) if code not in self._lobbies: break - lobby = Lobby(code, self._bridge, self._settings) + lobby = Lobby(code, self._bridge, self._settings, self._on_games_changed) self._lobbies[code] = lobby logger.info("Lobby %s created (%d open)", code, len(self._lobbies)) return lobby diff --git a/wizard_server/app/main.py b/wizard_server/app/main.py index 8860403..6bab8ed 100644 --- a/wizard_server/app/main.py +++ b/wizard_server/app/main.py @@ -30,23 +30,22 @@ def create_app(config: Settings = settings) -> FastAPI: async def lifespan(_: FastAPI) -> AsyncIterator[None]: if not config.ai_shared_secret: logger.warning("AI_SHARED_SECRET is not set, /internal/ai is unauthenticated") - cleanup = asyncio.create_task(service.cleanup_loop()) + tasks = [ + asyncio.create_task(service.cleanup_loop()), + asyncio.create_task(service.load.run()), + ] try: yield finally: - cleanup.cancel() + for task in tasks: + task.cancel() app = FastAPI(title="Wizard game server", lifespan=lifespan, docs_url=None, redoc_url=None) app.state.service = service @app.get("/health") async def health() -> dict: - return { - "status": "ok", - "lobbies": len(service.lobbies), - "sessions": len(service.sessions), - "ai_available": service.bridge.available, - } + return service.health() @app.websocket("/ws") async def client_socket(websocket: WebSocket) -> None: diff --git a/wizard_server/app/protocol.py b/wizard_server/app/protocol.py index 0c51a4c..c5c326e 100644 --- a/wizard_server/app/protocol.py +++ b/wizard_server/app/protocol.py @@ -5,7 +5,8 @@ messages are plain dicts: - ``{"type": "welcome", "player_id", "token", "name"}`` - ``{"type": "state", "lobby": LobbyView | null, "game": GameView | null, "events": [...]}`` -- ``{"type": "error", "message"}`` +- ``{"type": "error", "message", "code"?}`` +- ``{"type": "server_status", "load", "accepting_games"}`` """ from __future__ import annotations @@ -75,5 +76,16 @@ ClientMessage = Annotated[ client_message_adapter: TypeAdapter[ClientMessage] = TypeAdapter(ClientMessage) -def error(message: str) -> dict: - return {"type": "error", "message": message} +# Machine readable error codes (the frontend reacts to them). +ERROR_BUSY = "busy" # server has no capacity for this right now +ERROR_RATE_LIMITED = "rate_limited" # client sends too many messages / creates too many lobbies + +# Websocket close code "Try Again Later" (RFC 6455 registry), used when full. +CLOSE_TRY_AGAIN_LATER = 1013 + + +def error(message: str, code: str | None = None) -> dict: + message_dict: dict = {"type": "error", "message": message} + if code: + message_dict["code"] = code + return message_dict diff --git a/wizard_server/app/service.py b/wizard_server/app/service.py index ce4fcb8..38867ec 100644 --- a/wizard_server/app/service.py +++ b/wizard_server/app/service.py @@ -15,6 +15,8 @@ from wizard_game import GameOptions from . import protocol from .ai_bridge import AIBridge from .config import Settings +from .limits import ConnectionLimiter, KeyedRateLimiter, TokenBucket +from .load import LoadLevel, LoadMonitor from .lobby import Lobby, LobbyError, LobbyManager from .sessions import Connection, Session, SessionRegistry @@ -27,8 +29,53 @@ class WizardService: def __init__(self, settings: Settings) -> None: self.settings = settings self.sessions = SessionRegistry() - self.bridge = AIBridge(settings.ai_decision_timeout, on_change=self._on_ai_availability_changed) - self.lobbies = LobbyManager(self.bridge, settings) + self.bridge = AIBridge( + settings.ai_decision_timeout, + on_change=self._on_ai_availability_changed, + max_inflight=settings.ai_max_inflight, + breaker_failures=settings.ai_breaker_failures, + breaker_cooldown=settings.ai_breaker_cooldown, + ) + self.lobbies = LobbyManager(self.bridge, settings, on_games_changed=self._broadcast_server_status) + self.connections = ConnectionLimiter(settings.max_connections, settings.max_connections_per_ip) + self.lobby_creation = KeyedRateLimiter(settings.lobby_create_rate / 60, settings.lobby_create_rate) + self.load = LoadMonitor(settings, on_change=lambda _level: self._broadcast_server_status()) + self._last_server_status: dict | None = None + + # --- Capacity ----------------------------------------------------------- + + @property + def accepting_games(self) -> bool: + """Whether new lobbies may be created and games started. + + Running games are never affected; only new work is refused. + """ + return ( + self.load.level is not LoadLevel.OVERLOADED + and self.lobbies.active_games() < self.settings.max_active_games + ) + + def server_status(self) -> dict: + return { + "type": "server_status", + "load": self.load.level.value, + "accepting_games": self.accepting_games, + } + + def _broadcast_server_status(self) -> None: + status = self.server_status() + if status == self._last_server_status: + return + self._last_server_status = status + for session in self.sessions.connected(): + session.send(status) + + def _require_capacity(self) -> None: + if not self.accepting_games: + raise LobbyError( + "The server is busy right now, new games can't be started. Please try again in a moment.", + protocol.ERROR_BUSY, + ) # --- Client connections ------------------------------------------------- @@ -41,10 +88,26 @@ class WizardService: await websocket.accept() connection = Connection(websocket) connection.start() + if not self.connections.acquire(connection.ip): + logger.warning("Refusing connection from %s: connection limit reached", connection.ip) + connection.send(protocol.error("The server is full right now, please try again later.", protocol.ERROR_BUSY)) + await connection.close(code=protocol.CLOSE_TRY_AGAIN_LATER) + return + + bucket = TokenBucket(self.settings.message_rate, self.settings.message_burst) + violations = 0 session: Session | None = None try: while True: raw = await websocket.receive_text() + if not bucket.allow(): + violations += 1 + if violations >= self.settings.max_rate_violations: + logger.warning("Closing connection from %s: message flood", connection.ip) + break + if violations == 1 or violations % 10 == 0: + connection.send(protocol.error("Too many messages, slow down.", protocol.ERROR_RATE_LIMITED)) + continue try: message = protocol.client_message_adapter.validate_json(raw) except ValidationError as exc: @@ -60,16 +123,17 @@ class WizardService: try: self._dispatch(session, message) except LobbyError as exc: - session.send(protocol.error(str(exc))) + session.send(protocol.error(str(exc), exc.code)) except WebSocketDisconnect: pass finally: + self.connections.release(connection.ip) if session is not None and session.connection is connection: session.connection = None session.last_seen = time.monotonic() if lobby := self.lobbies.get(session.lobby_code): lobby.broadcast() - await connection.close() + await connection.close(code=1008 if violations >= self.settings.max_rate_violations else 1000) def _hello(self, connection: Connection, token: str | None) -> Session: session = self.sessions.get(token) @@ -87,6 +151,7 @@ class WizardService: "token": session.token, "name": session.name, }) + session.send(self.server_status()) lobby = self.lobbies.get(session.lobby_code) if lobby is not None and lobby.seat_index(session.player_id) is not None: lobby.add_human(session) # re-binds the seat to this session @@ -99,6 +164,11 @@ class WizardService: def _dispatch(self, session: Session, message: object) -> None: match message: case protocol.CreateLobby(name=name): + self._require_capacity() + ip = session.connection.ip if session.connection else "unknown" + if not self.lobby_creation.allow(ip): + raise LobbyError("You are creating lobbies too quickly, please wait a moment.", + protocol.ERROR_RATE_LIMITED) self._leave_current_lobby(session) session.name = name lobby = self.lobbies.create() @@ -132,6 +202,7 @@ class WizardService: lobby.broadcast() case protocol.StartGame(): lobby = self._lobby_of(session) + self._require_capacity() lobby.start(session) lobby.broadcast() case protocol.ReturnToLobby(): @@ -200,6 +271,20 @@ class WizardService: for session in self.sessions.expired(self.settings.session_ttl): self._leave_current_lobby(session) self.sessions.remove(session) + self.lobby_creation.prune() + + def health(self) -> dict: + return { + "status": "ok", # stays ok under load: an overloaded server must not be restarted + **self.load.stats(), + "accepting_games": self.accepting_games, + "lobbies": len(self.lobbies), + "active_games": self.lobbies.active_games(), + "sessions": len(self.sessions), + "connections": self.connections.total, + "ai_available": self.bridge.available, + "ai": self.bridge.stats(), + } def _validation_message(exc: ValidationError) -> str: diff --git a/wizard_server/app/sessions.py b/wizard_server/app/sessions.py index 2756096..42e8d90 100644 --- a/wizard_server/app/sessions.py +++ b/wizard_server/app/sessions.py @@ -24,6 +24,8 @@ logger = logging.getLogger(__name__) class Connection: def __init__(self, websocket: WebSocket) -> None: self.websocket = websocket + # Client IP; behind the reverse proxy uvicorn takes it from X-Forwarded-For. + self.ip = websocket.client.host if websocket.client else "unknown" self._queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue() self._writer: asyncio.Task | None = None @@ -85,6 +87,9 @@ class SessionRegistry: def remove(self, session: Session) -> None: self._by_token.pop(session.token, None) + def connected(self) -> list[Session]: + return [s for s in self._by_token.values() if s.connected] + def expired(self, ttl: float) -> list[Session]: now = time.monotonic() return [s for s in self._by_token.values() if not s.connected and now - s.last_seen > ttl] diff --git a/wizard_server/tests/conftest.py b/wizard_server/tests/conftest.py new file mode 100644 index 0000000..fbb2586 --- /dev/null +++ b/wizard_server/tests/conftest.py @@ -0,0 +1,26 @@ +import contextlib + +import pytest +from fastapi.testclient import TestClient + +from app.config import Settings +from app.main import create_app +from helpers import TEST_SETTINGS + + +@pytest.fixture +def make_client(): + """Factory: ``make_client(**settings_overrides)`` returns a started TestClient.""" + with contextlib.ExitStack() as stack: + def factory(**overrides): + settings = Settings(**{**TEST_SETTINGS, **overrides}) + test_client = stack.enter_context(TestClient(create_app(settings))) + # Websocket sessions are closed before the app stops (entered later, exited first). + test_client.sockets = stack.enter_context(contextlib.ExitStack()) + return test_client + yield factory + + +@pytest.fixture +def client(make_client): + return make_client() diff --git a/wizard_server/tests/helpers.py b/wizard_server/tests/helpers.py new file mode 100644 index 0000000..3923a9a --- /dev/null +++ b/wizard_server/tests/helpers.py @@ -0,0 +1,33 @@ +"""Helpers shared by the server tests.""" + +SECRET = "test-secret" + +# Tests play much faster than humans, so the message rate limit is raised. +TEST_SETTINGS = dict( + ai_shared_secret=SECRET, + ai_move_delay=0, + trick_pause=0, + ai_decision_timeout=2, + message_rate=10_000, + message_burst=10_000, +) + + +def receive_until(ws, predicate, limit=5000): + for _ in range(limit): + message = ws.receive_json() + if predicate(message): + return message + raise AssertionError("expected message not received") + + +def is_state(message): + return message["type"] == "state" + + +def connect(client, token=None): + ws = client.sockets.enter_context(client.websocket_connect("/ws")) + ws.send_json({"type": "hello", "token": token}) + welcome = receive_until(ws, lambda m: m["type"] == "welcome") + receive_until(ws, is_state) + return ws, welcome diff --git a/wizard_server/tests/test_limits.py b/wizard_server/tests/test_limits.py new file mode 100644 index 0000000..de21235 --- /dev/null +++ b/wizard_server/tests/test_limits.py @@ -0,0 +1,214 @@ +"""Unit tests of the limiter, load-monitor and AI-bridge protection building blocks.""" + +import asyncio +import random + +from wizard_game import WizardGame + +from app.ai_bridge import AIBridge, AIConnection +from app.config import Settings +from app.limits import ConnectionLimiter, KeyedRateLimiter, TokenBucket +from app.load import ( + LoadLevel, + LoadMonitor, + LoadSample, + LoadThresholds, + classify, + next_level, + own_cgroup_dir, + parse_cpu_max, + parse_cpu_stat, +) + +N, B, O = LoadLevel.NORMAL, LoadLevel.BUSY, LoadLevel.OVERLOADED + + +class FakeClock: + def __init__(self): + self.now = 0.0 + + def __call__(self): + return self.now + + +# --- TokenBucket / KeyedRateLimiter / ConnectionLimiter ----------------------- + + +def test_token_bucket_burst_and_refill(): + clock = FakeClock() + bucket = TokenBucket(rate=2, burst=3, clock=clock) + assert [bucket.allow() for _ in range(4)] == [True, True, True, False] + clock.now = 0.5 # one token back + assert bucket.allow() and not bucket.allow() + clock.now = 100 # never more than the burst + assert sum(bucket.allow() for _ in range(10)) == 3 + + +def test_keyed_rate_limiter_is_per_key_and_prunes(): + clock = FakeClock() + limiter = KeyedRateLimiter(rate=1, burst=1, clock=clock) + assert limiter.allow("a") and not limiter.allow("a") + assert limiter.allow("b") + clock.now = 10 + limiter.prune() + assert len(limiter) == 0 + + +def test_connection_limiter(): + limiter = ConnectionLimiter(max_total=3, max_per_ip=2) + assert limiter.acquire("1.1.1.1") and limiter.acquire("1.1.1.1") + assert not limiter.acquire("1.1.1.1") + assert limiter.acquire("2.2.2.2") + assert not limiter.acquire("3.3.3.3") # total reached + limiter.release("1.1.1.1") + assert limiter.total == 2 + assert limiter.acquire("3.3.3.3") + limiter.release("9.9.9.9") # unknown ip is ignored + assert limiter.total == 3 + + +# --- Load levels -------------------------------------------------------------- + +T = LoadThresholds( + lag_busy=0.1, lag_overload=0.25, lag_severe=1.0, cpu_busy=0.7, cpu_overload=0.9, + throttle_overload=0.5, overload_samples=3, recover_samples=4, +) + + +def test_classify(): + assert classify(LoadSample(0.0, 0.1, 0.0), T) is N + assert classify(LoadSample(0.15), T) is B + assert classify(LoadSample(0.0, 0.75), T) is B + assert classify(LoadSample(0.3), T) is O + assert classify(LoadSample(0.0, 0.95), T) is O + assert classify(LoadSample(0.0, 0.2, 0.6), T) is O + + +def run_levels(samples): + level, recent, levels = N, [], [] + for sample in samples: + recent.append(sample) + level = next_level(level, recent, T) + levels.append(level) + return levels + + +def test_overload_needs_consecutive_samples(): + assert run_levels([O, O, N]) == [B, B, B] + assert run_levels([O, O, O]) == [B, B, O] + + +def test_overload_recovery_needs_a_calm_window(): + levels = run_levels([O, O, O, N, N, N, B, N, N, N, N]) + assert levels[3:6] == [O, O, O] # still overloaded while recovering + assert levels[6] is B # 4 samples without overload -> busy (last sample busy) + assert levels[-1] is N # 4 normal samples -> normal + + +def test_severe_lag_overloads_immediately(): + assert next_level(N, [O], T, severe=True) is O + monitor = LoadMonitor(Settings()) + monitor.record(LoadSample(loop_lag=2.0)) + assert monitor.level is O + + +def test_busy_is_immediate_and_sticky(): + assert run_levels([N, B, N, N, N, N]) == [N, B, B, B, B, N] + + +# --- cgroup parsing / LoadMonitor --------------------------------------------- + + +def test_parse_cgroup_files(): + assert parse_cpu_max("max 100000\n") is None + assert parse_cpu_max("50000 100000\n") == 0.5 + assert parse_cpu_max("") is None + stats = parse_cpu_stat("usage_usec 10\nnr_periods 200\nnr_throttled 50\n") + assert stats["nr_periods"] == 200 and stats["nr_throttled"] == 50 + + +def test_own_cgroup_dir(tmp_path): + (tmp_path / "user.slice" / "app.scope").mkdir(parents=True) + proc = tmp_path / "cgroup" + proc.write_text("0::/user.slice/app.scope\n") + assert own_cgroup_dir(tmp_path, proc) == tmp_path / "user.slice" / "app.scope" + proc.write_text("0::/\n") # inside a container + assert own_cgroup_dir(tmp_path, proc) == tmp_path + assert own_cgroup_dir(tmp_path, tmp_path / "missing") == tmp_path + + +def test_monitor_reads_cgroup_throttling(tmp_path): + (tmp_path / "cpu.max").write_text("50000 100000\n") + (tmp_path / "cpu.stat").write_text("nr_periods 100\nnr_throttled 0\n") + monitor = LoadMonitor(Settings(), cgroup_dir=tmp_path) + (tmp_path / "cpu.stat").write_text("nr_periods 200\nnr_throttled 80\n") + sample = monitor._measure(lag=0.0) + assert sample.throttled == 0.8 + + +def test_monitor_levels_and_callback(): + changes = [] + monitor = LoadMonitor(Settings(load_overload_samples=2, load_recover_samples=2), on_change=changes.append) + for _ in range(2): + monitor.record(LoadSample(loop_lag=0.5)) # overloaded, but not severe + assert monitor.level is O + for _ in range(2): + monitor.record(LoadSample(loop_lag=0.0)) + assert monitor.level is N + assert changes == [B, O, N] + assert monitor.stats()["load"] == "normal" + + +def test_forced_level(): + monitor = LoadMonitor(Settings(load_force_level="overloaded")) + assert monitor.level is O + monitor.record(LoadSample(loop_lag=0.0)) + assert monitor.level is O + + +# --- AI bridge protection ----------------------------------------------------- + + +class SilentWebSocket: + """Accepts requests but never answers (a hanging AI service).""" + + def __init__(self): + self.sent = [] + + async def send_json(self, message): + self.sent.append(message) + + +def observation(): + game = WizardGame(3, seed=1) + return game.observation(game.current_player) + + +def test_ai_circuit_breaker(): + async def scenario(): + bridge = AIBridge(timeout=0.01, rng=random.Random(0), breaker_failures=2, breaker_cooldown=60) + ws = SilentWebSocket() + bridge._connections.append(AIConnection(ws, "fake")) + obs = observation() + for _ in range(4): + action = await bridge.decide(obs) + assert action in obs.legal_actions # random fallback + assert len(ws.sent) == 2 # after 2 timeouts the AI is not asked anymore + assert ws.sent[0]["timeout"] == 0.01 + assert bridge.stats()["paused"] is True + + asyncio.run(scenario()) + + +def test_ai_inflight_limit(): + async def scenario(): + bridge = AIBridge(timeout=1, rng=random.Random(0), max_inflight=1) + ws = SilentWebSocket() + connection = AIConnection(ws, "fake") + connection.pending["busy"] = asyncio.get_running_loop().create_future() + bridge._connections.append(connection) + obs = observation() + assert await bridge.decide(obs) in obs.legal_actions + assert ws.sent == [] # not even sent: too many requests already waiting + + asyncio.run(scenario()) diff --git a/wizard_server/tests/test_protection.py b/wizard_server/tests/test_protection.py new file mode 100644 index 0000000..239bb81 --- /dev/null +++ b/wizard_server/tests/test_protection.py @@ -0,0 +1,137 @@ +"""Abuse limits, load-based admission control and server_status messages.""" + +import pytest +from starlette.websockets import WebSocketDisconnect + +from app.load import LoadSample + +from helpers import connect, is_state, receive_until + + +def is_error(code): + return lambda m: m["type"] == "error" and m.get("code") == code + + +def test_welcome_is_followed_by_server_status(client): + ws = client.sockets.enter_context(client.websocket_connect("/ws")) + ws.send_json({"type": "hello", "token": None}) + assert ws.receive_json()["type"] == "welcome" + status = ws.receive_json() + assert status == {"type": "server_status", "load": "normal", "accepting_games": True} + + +def test_message_flood_is_limited_and_closed(make_client): + client = make_client(message_rate=0.001, message_burst=3, max_rate_violations=5) + ws, _ = connect(client) # uses 1 message + for _ in range(2): # uses the remaining burst + ws.send_json({"type": "leave_lobby"}) + ws.send_json({"type": "leave_lobby"}) # over the limit + assert receive_until(ws, is_error("rate_limited")) + with pytest.raises(WebSocketDisconnect) as closed: + for _ in range(10): + ws.send_json({"type": "leave_lobby"}) + for _ in range(20): + ws.receive_json() + assert closed.value.code == 1008 + + +def test_connections_per_ip_are_limited(make_client): + client = make_client(max_connections_per_ip=2) + connect(client) + connect(client) + ws = client.sockets.enter_context(client.websocket_connect("/ws")) + assert ws.receive_json()["code"] == "busy" + with pytest.raises(WebSocketDisconnect) as closed: + ws.receive_json() + assert closed.value.code == 1013 + + +def test_closed_connections_free_their_slot(make_client): + client = make_client(max_connections_per_ip=1) + with client.websocket_connect("/ws") as ws: + ws.send_json({"type": "hello", "token": None}) + receive_until(ws, is_state) + ws, _ = connect(client) # works again after the first one closed + + +def test_lobby_creation_rate_per_ip(make_client): + client = make_client(lobby_create_rate=2) + ws, _ = connect(client) + for _ in range(2): + ws.send_json({"type": "create_lobby", "name": "Ann"}) + receive_until(ws, lambda m: is_state(m) and m["lobby"]) + ws.send_json({"type": "create_lobby", "name": "Ann"}) + assert receive_until(ws, is_error("rate_limited")) + + +def test_overloaded_server_refuses_new_games(make_client): + client = make_client(load_force_level="overloaded") + ws = client.sockets.enter_context(client.websocket_connect("/ws")) + ws.send_json({"type": "hello", "token": None}) + status = receive_until(ws, lambda m: m["type"] == "server_status") + assert status["load"] == "overloaded" and not status["accepting_games"] + ws.send_json({"type": "create_lobby", "name": "Ann"}) + assert receive_until(ws, is_error("busy")) + assert client.get("/health").json()["accepting_games"] is False + + +def start_game_with_ais(client, name): + ws, _ = connect(client) + ws.send_json({"type": "create_lobby", "name": name}) + receive_until(ws, lambda m: is_state(m) and m["lobby"]) + ws.send_json({"type": "add_ai"}) + ws.send_json({"type": "add_ai"}) + return ws + + +def test_max_active_games(make_client): + client = make_client(max_active_games=1) + first = start_game_with_ais(client, "Ann") + second = start_game_with_ais(client, "Ben") + + first.send_json({"type": "start_game"}) + receive_until(first, lambda m: is_state(m) and m["game"]) + # Everybody learns that no further games can start. + status = receive_until(second, lambda m: m["type"] == "server_status") + assert status["accepting_games"] is False + + second.send_json({"type": "start_game"}) + assert receive_until(second, is_error("busy")) + + # Leaving the running game ends it (only AIs left) and frees the slot. + first.send_json({"type": "leave_lobby"}) + status = receive_until(second, lambda m: m["type"] == "server_status") + assert status["accepting_games"] is True + second.send_json({"type": "start_game"}) + assert receive_until(second, lambda m: is_state(m) and m["game"]) + + +def test_running_game_continues_when_overloaded(client): + ws = start_game_with_ais(client, "Ann") + ws.send_json({"type": "start_game"}) + latest = receive_until(ws, lambda m: is_state(m) and m["game"]) + + service = client.app.state.service + for _ in range(3): # three overloaded samples in a row + client.portal.call(service.load.record, LoadSample(loop_lag=1.0)) + + # normal -> busy (first bad sample) -> overloaded (third in a row). + # Keep track of the latest game state seen on the way. + while True: + message = ws.receive_json() + if is_state(message): + latest = message + if message["type"] == "server_status" and message["load"] == "overloaded": + assert message["accepting_games"] is False + break + + # The running game still accepts moves. + def my_turn(state): + obs = state["game"]["observation"] + return obs["current_player"] == state["game"]["you"] and obs["legal_actions"] + + while not my_turn(latest): + latest = receive_until(ws, is_state) + ws.send_json({"type": "action", "action": latest["game"]["observation"]["legal_actions"][0]}) + reply = ws.receive_json() + assert reply["type"] == "state", reply diff --git a/wizard_server/tests/test_server.py b/wizard_server/tests/test_server.py index 87a3b23..f685d53 100644 --- a/wizard_server/tests/test_server.py +++ b/wizard_server/tests/test_server.py @@ -1,42 +1,9 @@ -import contextlib import random import threading import pytest -from fastapi.testclient import TestClient -from app.config import Settings -from app.main import create_app - -SECRET = "test-secret" - - -@pytest.fixture -def client(): - settings = Settings(ai_shared_secret=SECRET, ai_move_delay=0, trick_pause=0, ai_decision_timeout=2) - with TestClient(create_app(settings)) as test_client, contextlib.ExitStack() as sockets: - test_client.sockets = sockets # websocket sessions are closed before the app stops - yield test_client - - -def receive_until(ws, predicate, limit=5000): - for _ in range(limit): - message = ws.receive_json() - if predicate(message): - return message - raise AssertionError("expected message not received") - - -def is_state(message): - return message["type"] == "state" - - -def connect(client, token=None): - ws = client.sockets.enter_context(client.websocket_connect("/ws")) - ws.send_json({"type": "hello", "token": token}) - welcome = receive_until(ws, lambda m: m["type"] == "welcome") - receive_until(ws, is_state) - return ws, welcome +from helpers import SECRET, connect, is_state, receive_until def test_health(client): diff --git a/wizard_website/src/App.tsx b/wizard_website/src/App.tsx index 9b314a9..597a325 100644 --- a/wizard_website/src/App.tsx +++ b/wizard_website/src/App.tsx @@ -10,10 +10,12 @@ const STATUS_MESSAGES = { connecting: 'Connecting to the server…', closed: 'Connection lost. Reconnecting…', replaced: 'This game was opened in another window. Reload the page to continue here.', + busy: 'The server is full right now. Retrying in a moment…', } as const function App() { - const { status, playerId, lobby, game, completedTrick, toasts, dismissToast, actions } = useWizard() + const { status, playerId, lobby, game, completedTrick, serverStatus, toasts, dismissToast, actions } = + useWizard() const lobbyCode = lobby?.code // Keep the lobby code in the address bar so the URL can be shared as an invite link. @@ -27,9 +29,22 @@ function App() { let screen if (!lobby || !playerId) { - screen = + screen = ( + + ) } else if (lobby.status === 'waiting' || !game) { - screen = + screen = ( + + ) } else { screen = ( (null) const [toasts, setToasts] = useState([]) + const [serverStatus, setServerStatus] = useState({ load: 'normal', accepting_games: true }) const socketRef = useRef(null) const toastId = useRef(0) @@ -76,6 +82,7 @@ export function useWizard() { useEffect(() => { let disposed = false let attempt = 0 + let busyAttempt = 0 let retryTimer: ReturnType | undefined let trickTimer: ReturnType | undefined @@ -97,6 +104,9 @@ export function useWizard() { if (!message.game) setCompletedTrick(null) break } + case 'server_status': + setServerStatus({ load: message.load, accepting_games: message.accepting_games }) + break case 'error': showError(message.message) break @@ -107,6 +117,9 @@ export function useWizard() { const socket = new WebSocket(websocketUrl()) socketRef.current = socket + // A "busy" error before the welcome means the server refused this connection. + let welcomed = false + let refused = false socket.onopen = () => { attempt = 0 setStatus('open') @@ -114,7 +127,18 @@ export function useWizard() { } socket.onmessage = (event) => { try { - handleMessage(JSON.parse(event.data as string) as ServerMessage) + const message = JSON.parse(event.data as string) as ServerMessage + if (message.type === 'welcome') { + welcomed = true + busyAttempt = 0 + } + // The "server is full" error arrives right before the 1013 close; the + // banner explains it, so no toast. + if (message.type === 'error' && message.code === 'busy' && !welcomed) { + refused = true + return + } + handleMessage(message) } catch (error) { console.error('Invalid message from server', error) } @@ -126,9 +150,17 @@ export function useWizard() { setStatus('replaced') return } - setStatus('closed') - const delay = Math.min(1000 * 2 ** attempt, MAX_RETRY_DELAY_MS) - attempt += 1 + let delay: number + if (event.code === CLOSE_TRY_AGAIN_LATER || refused) { + setStatus('busy') + delay = BUSY_RETRY_DELAYS_MS[Math.min(busyAttempt, BUSY_RETRY_DELAYS_MS.length - 1)] + busyAttempt += 1 + } else { + setStatus('closed') + busyAttempt = 0 + delay = Math.min(1000 * 2 ** attempt, MAX_RETRY_DELAY_MS) + attempt += 1 + } retryTimer = setTimeout(() => { setStatus('connecting') connect() @@ -173,7 +205,7 @@ export function useWizard() { [send], ) - return { status, playerId, lobby, game, completedTrick, toasts, dismissToast, actions } + return { status, playerId, lobby, game, completedTrick, serverStatus, toasts, dismissToast, actions } } export type WizardActions = ReturnType['actions'] diff --git a/wizard_website/src/screens/HomeScreen/HomeScreen.module.css b/wizard_website/src/screens/HomeScreen/HomeScreen.module.css index 1f8be76..588b80a 100644 --- a/wizard_website/src/screens/HomeScreen/HomeScreen.module.css +++ b/wizard_website/src/screens/HomeScreen/HomeScreen.module.css @@ -104,3 +104,13 @@ color: var(--text-muted); font-size: 13px; } + +.busy { + margin: 0; + padding: 10px 12px; + border: 1px solid rgb(255 209 102 / 40%); + border-radius: 10px; + background: rgb(255 209 102 / 8%); + color: var(--gold); + font-size: 14px; +} diff --git a/wizard_website/src/screens/HomeScreen/HomeScreen.tsx b/wizard_website/src/screens/HomeScreen/HomeScreen.tsx index 079ca6d..4e78a9a 100644 --- a/wizard_website/src/screens/HomeScreen/HomeScreen.tsx +++ b/wizard_website/src/screens/HomeScreen/HomeScreen.tsx @@ -20,10 +20,12 @@ function initialCode(): string { interface HomeScreenProps { connected: boolean + // False while the server refuses new games (overloaded or full). + acceptingGames: boolean actions: WizardActions } -export function HomeScreen({ connected, actions }: HomeScreenProps) { +export function HomeScreen({ connected, acceptingGames, actions }: HomeScreenProps) { const [name, setName] = useState(initialName) const [code, setCode] = useState(initialCode) const trimmedName = name.trim() @@ -73,9 +75,20 @@ export function HomeScreen({ connected, actions }: HomeScreenProps) { onChange={(event) => setName(event.target.value)} /> - + {connected && !acceptingGames && ( +

+ The server is busy right now, so new games can't be created. You can still join an existing + lobby — please try again in a moment. +

+ )}
or join a lobby diff --git a/wizard_website/src/screens/LobbyScreen/LobbyScreen.module.css b/wizard_website/src/screens/LobbyScreen/LobbyScreen.module.css index fdee31e..d93e05c 100644 --- a/wizard_website/src/screens/LobbyScreen/LobbyScreen.module.css +++ b/wizard_website/src/screens/LobbyScreen/LobbyScreen.module.css @@ -211,3 +211,13 @@ color: var(--text-muted); font-size: 14px; } + +.busy { + margin: 0; + padding: 12px 16px; + border: 1px solid rgb(255 209 102 / 40%); + border-radius: 12px; + background: rgb(255 209 102 / 8%); + color: var(--gold); + font-size: 14px; +} diff --git a/wizard_website/src/screens/LobbyScreen/LobbyScreen.tsx b/wizard_website/src/screens/LobbyScreen/LobbyScreen.tsx index 0748f41..3e08525 100644 --- a/wizard_website/src/screens/LobbyScreen/LobbyScreen.tsx +++ b/wizard_website/src/screens/LobbyScreen/LobbyScreen.tsx @@ -7,6 +7,8 @@ import styles from './LobbyScreen.module.css' interface LobbyScreenProps { lobby: Lobby playerId: string + // False while the server refuses new games (overloaded or full). + acceptingGames: boolean actions: WizardActions } @@ -17,7 +19,7 @@ function inviteLink(code: string): string { return url.toString() } -export function LobbyScreen({ lobby, playerId, actions }: LobbyScreenProps) { +export function LobbyScreen({ lobby, playerId, acceptingGames, actions }: LobbyScreenProps) { const [copied, setCopied] = useState<'code' | 'link' | null>(null) const isHost = lobby.host_id === playerId const isFull = lobby.players.length >= lobby.max_players @@ -116,13 +118,24 @@ export function LobbyScreen({ lobby, playerId, actions }: LobbyScreenProps) { + {!acceptingGames && ( +

+ The server is busy right now, so games can't be started. Your lobby stays open — try again in a + moment. +

+ )} +