Add WebSocket protocol benchmarks for wsproto and websockets-sansio (#2849)

Benchmark handshake and text frame sending using the same mock
transport approach as HTTP benchmarks. The legacy websockets
implementation is excluded as it manages its own internal tasks.

tests/benchmarks/test_ws.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import importlib.util
+from typing import TYPE_CHECKING
+
+import pytest
+
+from tests.benchmarks.ws import WS_UPGRADE, get_connected_ws_protocol
+from uvicorn._types import ASGIReceiveCallable, ASGISendCallable, Scope
+
+if TYPE_CHECKING:
+    from tests.benchmarks.ws import WSProtocolClass
+
+pytestmark = [pytest.mark.anyio, pytest.mark.benchmark]
+
+
+@pytest.fixture(
+    params=[
+        pytest.param(
+            "wsproto",
+            marks=pytest.mark.skipif(not importlib.util.find_spec("wsproto"), reason="wsproto not installed."),
+            id="wsproto",
+        ),
+        pytest.param("websockets-sansio", id="websockets-sansio"),
+    ]
+)
+def ws_cls(request: pytest.FixtureRequest) -> WSProtocolClass:
+    if request.param == "wsproto":
+        from uvicorn.protocols.websockets.wsproto_impl import WSProtocol
+
+        return WSProtocol
+    from uvicorn.protocols.websockets.websockets_sansio_impl import WebSocketsSansIOProtocol
+
+    return WebSocketsSansIOProtocol
+
+
+async def _ws_accept_close_app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable) -> None:
+    await receive()
+    await send({"type": "websocket.accept"})
+    await send({"type": "websocket.close", "code": 1000})
+
+
+async def _ws_send_text_app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable) -> None:
+    await receive()
+    await send({"type": "websocket.accept"})
+    await send({"type": "websocket.send", "text": "Hello, world!"})
+    await send({"type": "websocket.close", "code": 1000})
+
+
+async def test_bench_ws_handshake(ws_cls: WSProtocolClass) -> None:
+    protocol = get_connected_ws_protocol(_ws_accept_close_app, ws_cls)
+    protocol.data_received(WS_UPGRADE)
+    await protocol.loop.run_one()
+
+
+async def test_bench_ws_send_text(ws_cls: WSProtocolClass) -> None:
+    protocol = get_connected_ws_protocol(_ws_send_text_app, ws_cls)
+    protocol.data_received(WS_UPGRADE)
+    await protocol.loop.run_one()

tests/benchmarks/ws.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, TypeAlias
+
+from tests.benchmarks.http import MockLoop, MockTransport
+from uvicorn._types import ASGIApplication
+from uvicorn.config import Config
+from uvicorn.lifespan.off import LifespanOff
+from uvicorn.server import ServerState
+
+if TYPE_CHECKING:
+    from uvicorn.protocols.websockets.websockets_sansio_impl import WebSocketsSansIOProtocol
+    from uvicorn.protocols.websockets.wsproto_impl import WSProtocol
+
+    WSProtocolClass: TypeAlias = type[WSProtocol] | type[WebSocketsSansIOProtocol]
+
+WS_UPGRADE = (
+    b"GET / HTTP/1.1\r\n"
+    b"Host: example.org\r\n"
+    b"Upgrade: websocket\r\n"
+    b"Connection: Upgrade\r\n"
+    b"Sec-WebSocket-Key: YmVuY2htYXJra2V5MTIzNA==\r\n"
+    b"Sec-WebSocket-Version: 13\r\n"
+    b"\r\n"
+)
+
+# Masked text frame: "Hello, world!" (13 bytes) with zero mask key
+WS_TEXT_FRAME = b"\x81\x8d\x00\x00\x00\x00Hello, world!"
+
+# Masked close frame: code 1000 with zero mask key
+WS_CLOSE_FRAME = b"\x88\x82\x00\x00\x00\x00\x03\xe8"
+
+
+def get_connected_ws_protocol(app: ASGIApplication, ws_protocol_cls: WSProtocolClass, **kwargs: Any) -> Any:
+    loop = MockLoop()
+    transport = MockTransport()
+    config = Config(app=app, access_log=False, **kwargs)
+    lifespan = LifespanOff(config)
+    server_state = ServerState()
+    protocol = ws_protocol_cls(config=config, server_state=server_state, app_state=lifespan.state, _loop=loop)  # type: ignore[arg-type]
+    protocol.connection_made(transport)  # type: ignore[arg-type]
+    return protocol