Add resource limits, abuse protection and overload handling
Running games always continue; new work is refused when there is no
capacity; abuse is rejected immediately.
- Docker: CPU/memory/pids caps per container (overridable in .env);
uvicorn message size, queue and concurrency limits
- Server abuse limits: connections total/per IP (close 1013), message
rate per connection (token bucket, close 1008 on floods), lobby
creation rate per IP
- Load monitor: event loop lag, CPU of the own cgroup quota and
throttling, with hysteresis (normal/busy/overloaded, severe lag
overloads immediately)
- Admission control: create_lobby and start_game refused while
overloaded or MAX_ACTIVE_GAMES are running; server_status message
broadcast on changes; /health reports load but stays ok
- AI bridge: timeout sent with requests, in-flight cap, circuit breaker
- AI service: bounded workers and queue ("overloaded"), drops expired
requests, degrades to the heuristic agent when busy or too slow
- Website: busy notices, Create/Start disabled when not accepting
games, slow retry when the server is full
- Tests for limiters, load levels, admission control and AI service
limits; docs updated
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 5.5
@@ -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"]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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]
|
||||
@@ -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 {}
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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())
|
||||
@@ -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
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user