feat: add bounded WebSocket server support

benchmarks/route_benchmark.py

@@ -0,0 +1,125 @@
+"""Repeatable routing comparison; run with ``python benchmarks/route_benchmark.py``."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import importlib.util
+import inspect
+import json
+from pathlib import Path
+import statistics
+import sys
+import time
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from smallserver import Headers, RegexRouteConfig, Request, Response, RouteMatchTimeout, SmallServer
+
+STATIC_DISPATCH_RATIO_FLOOR = 0.80
+
+
+async def _legacy_dispatch(routes, request):
+    """Model the pre-router static dictionary dispatch for same-run comparison."""
+    handler = routes[(request.method.upper(), request.path)]
+    result = handler(request)
+    if not inspect.isawaitable(result):
+        raise TypeError("benchmark handler must be awaitable")
+    response = await result
+    if not isinstance(response, Response):
+        raise TypeError("benchmark handler must return Response")
+    return response
+
+
+async def _measure(operation, iterations: int) -> float:
+    started = time.perf_counter()
+    for _ in range(iterations):
+        await operation()
+    return iterations / (time.perf_counter() - started)
+
+
+async def benchmark(iterations: int, rounds: int) -> dict[str, float | int | str | bool]:
+    app = SmallServer()
+
+    @app.get("/health")
+    async def health(request):
+        return Response()
+
+    request = Request("GET", "/health", Headers())
+    legacy_routes = {("GET", "/health"): health}
+
+    async def legacy_operation():
+        return await _legacy_dispatch(legacy_routes, request)
+
+    async def router_operation():
+        return await app.dispatch(request)
+
+    await _measure(legacy_operation, min(iterations, 1_000))
+    await _measure(router_operation, min(iterations, 1_000))
+    legacy_rates = []
+    router_rates = []
+    ratios = []
+    for round_number in range(rounds):
+        if round_number % 2:
+            router_rate = await _measure(router_operation, iterations)
+            legacy_rate = await _measure(legacy_operation, iterations)
+        else:
+            legacy_rate = await _measure(legacy_operation, iterations)
+            router_rate = await _measure(router_operation, iterations)
+        legacy_rates.append(legacy_rate)
+        router_rates.append(router_rate)
+        ratios.append(router_rate / legacy_rate)
+
+    median_ratio = statistics.median(ratios)
+    result: dict[str, float | int | str | bool] = {
+        "iterations": iterations,
+        "rounds": rounds,
+        "legacy_static_dispatches_per_second": round(statistics.median(legacy_rates), 2),
+        "router_static_dispatches_per_second": round(statistics.median(router_rates), 2),
+        "router_to_legacy_ratio": round(median_ratio, 4),
+        "static_dispatch_ratio_floor": STATIC_DISPATCH_RATIO_FLOOR,
+        "static_dispatch_floor_passed": median_ratio >= STATIC_DISPATCH_RATIO_FLOOR,
+    }
+
+    if importlib.util.find_spec("regex") is None:
+        result["regex"] = "skipped; install smallserver[regex-routes]"
+        return result
+
+    bounded = SmallServer(RegexRouteConfig(match_timeout=0.002, total_match_timeout=0.005))
+
+    @bounded.get_regex(r"/(a+)+$")
+    async def hostile(request):
+        return Response()
+
+    started = time.perf_counter()
+    try:
+        await bounded.dispatch(Request("GET", "/" + "a" * 5000 + "!", Headers()))
+    except RouteMatchTimeout:
+        pass
+    result["configured_regex_match_timeout_seconds"] = 0.002
+    result["observed_worst_case_regex_seconds"] = round(time.perf_counter() - started, 6)
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--iterations", type=int, default=25_000)
+    parser.add_argument("--rounds", type=int, default=5)
+    parser.add_argument("--release", action="store_true")
+    arguments = parser.parse_args()
+    if arguments.iterations <= 0:
+        parser.error("--iterations must be positive")
+    if arguments.rounds <= 0:
+        parser.error("--rounds must be positive")
+    if arguments.release and arguments.iterations < 10_000:
+        parser.error("--release requires at least 10000 iterations")
+    if arguments.release and arguments.rounds < 5:
+        parser.error("--release requires at least 5 rounds")
+    result = asyncio.run(benchmark(arguments.iterations, arguments.rounds))
+    print(json.dumps(result, indent=2, sort_keys=True))
+    if arguments.release and not result["static_dispatch_floor_passed"]:
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

demo.py

@@ -4,7 +4,7 @@
 
 import json
 
-from smallserver import HTTPError, Request, Response, SmallServer
+from smallserver import HTTPError, Request, Response, SmallServer, WebSocket
 
 
 app = SmallServer()
@@ -81,6 +81,16 @@ async def delete_task(request: Request) -> Response:
     return Response(status=204)
 
 
+@app.websocket("/ws")
+async def websocket_echo(socket: WebSocket) -> None:
+    await socket.accept()
+    async for message in socket:
+        if message.is_text:
+            await socket.send_text(message.text)
+        else:
+            await socket.send_bytes(message.bytes)
+
+
 if __name__ == "__main__":
     print("Starting SmallServer on http://127.0.0.1:8000")
     app.listen(host="127.0.0.1", port=8000)

examples/websocket_echo.py

@@ -0,0 +1,28 @@
+"""Run a bounded WebSocket echo endpoint on localhost:8000."""
+
+from smallserver import SmallServer, WebSocket, WebSocketConfig
+
+
+app = SmallServer(
+    websocket_config=WebSocketConfig(
+        max_frame_payload_bytes=64 * 1024,
+        max_message_bytes=256 * 1024,
+        max_inbound_messages=8,
+        max_outbound_commands=8,
+    )
+)
+
+
+@app.websocket("/echo")
+async def echo(socket: WebSocket) -> None:
+    await socket.accept()
+    async for message in socket:
+        if message.is_text:
+            await socket.send_text(message.text)
+        else:
+            await socket.send_bytes(message.bytes)
+
+
+if __name__ == "__main__":
+    print("Starting WebSocket echo server on ws://127.0.0.1:8000/echo")
+    app.listen(host="127.0.0.1", port=8000)

pyproject.toml

@@ -14,6 +14,8 @@ dependencies = []
 
 [project.optional-dependencies]
 dev = ["build>=1.2"]
+regex-routes = ["regex>=2023.10.3,<2027"]
+websocket = ["wsproto>=1.2,<2"]
 test = [
     "build>=1.2",
     "h2>=4,<5",

smallserver/__init__.py

@@ -10,8 +10,24 @@
     ServerStartupError,
 )
 from .http import Headers, Request, Response
+from .routing import (
+    RegexRouteConfig,
+    RegexRoutesUnavailable,
+    RouteErrorEvent,
+    RouteMatchTimeout,
+    RoutePathTooLarge,
+)
 from .runtime import ManagedRuntimeConfig
 from .server import ServerConfig, ServerHandle
+from .websocket import (
+    WebSocket,
+    WebSocketCapacityError,
+    WebSocketConfig,
+    WebSocketDisconnect,
+    WebSocketMessage,
+    WebSocketStateError,
+    WebSocketUnavailable,
+)
 
 if TYPE_CHECKING:
     from .adapters import AdapterRegistry, AdapterShutdownError, http_error_from_adapter
@@ -33,13 +49,25 @@ def __getattr__(name: str) -> Any:
     "Headers",
     "HTTPError",
     "ManagedRuntimeConfig",
+    "RegexRouteConfig",
+    "RegexRoutesUnavailable",
     "Request",
     "Response",
+    "RouteErrorEvent",
+    "RouteMatchTimeout",
+    "RoutePathTooLarge",
     "ServerConfig",
     "ServerConfigurationError",
     "ServerFinalizationError",
     "ServerHandle",
     "ServerStartupError",
     "SmallServer",
+    "WebSocket",
+    "WebSocketCapacityError",
+    "WebSocketConfig",
+    "WebSocketDisconnect",
+    "WebSocketMessage",
+    "WebSocketStateError",
+    "WebSocketUnavailable",
     "http_error_from_adapter",
 ]