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
@@ -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