[PR #13561/148205f2 backport][3.14] Add masked WebSocket read benchmarks and fix read benchmark backpressure

CHANGES/13561.contrib.rst

@@ -0,0 +1,3 @@
+Added benchmarks for reading masked WebSocket messages and fixed the
+existing read benchmarks, which stopped measuring the parser after the
+eighth large frame due to the queue limit -- by :user:`bdraco`.

tests/test_benchmarks_http_websocket.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from aiohttp._websocket.helpers import MSG_SIZE, PACK_LEN3
+from aiohttp._websocket.helpers import MSG_SIZE, PACK_LEN1, PACK_LEN3, websocket_mask
 from aiohttp._websocket.reader import WebSocketDataQueue
 from aiohttp.base_protocol import BaseProtocol
 from aiohttp.helpers import DEFAULT_CHUNK_SIZE
@@ -18,25 +18,51 @@
     BenchmarkFixture = pytest_codspeed.BenchmarkFixture
 
 
+# Large enough that a benchmark run never crosses the queue limit; hitting it
+# would engage read backpressure and silently stop the parser mid-benchmark.
+READ_QUEUE_LIMIT = 2**24
+MASK = b"\x9a\x3c\x71\xe5"
+
+
+def _make_reader(loop: asyncio.AbstractEventLoop) -> WebSocketReader:
+    protocol = BaseProtocol(loop)
+    # A WebSocket connection is always upgraded; without this, backpressure
+    # would hit ``assert self._parser is not None`` in pause_reading().
+    protocol._upgraded = True
+    queue = WebSocketDataQueue(protocol, READ_QUEUE_LIMIT, loop=loop)
+    return WebSocketReader(
+        queue, max_msg_size=DEFAULT_CHUNK_SIZE, compress=True, decode_text=True
+    )
+
+
+def _masked_frame(opcode: WSMsgType, payload: bytes) -> bytes:
+    """Build a client-to-server frame with a masked payload."""
+    masked = bytearray(payload)
+    websocket_mask(MASK, masked)
+    first_byte = 0x80 | opcode.value
+    length = len(payload)
+    assert length < 126 or length > 2**16
+    if length < 126:
+        header = PACK_LEN1(first_byte, 0x80 | length)
+    else:
+        header = PACK_LEN3(first_byte, 0x80 | 127, length)
+    return header + MASK + bytes(masked)
+
+
 def test_read_large_binary_websocket_messages(
     loop: asyncio.AbstractEventLoop, benchmark: BenchmarkFixture
 ) -> None:
     """Read one hundred large binary websocket messages."""
-    queue = WebSocketDataQueue(BaseProtocol(loop), DEFAULT_CHUNK_SIZE, loop=loop)
-    reader = WebSocketReader(
-        queue, max_msg_size=DEFAULT_CHUNK_SIZE, compress=True, decode_text=True
-    )
-
     # PACK3 has a minimum message length of 2**16 bytes.
     message = b"x" * ((2**16) + 1)
     msg_length = len(message)
     first_byte = 0x80 | 0 | WSMsgType.BINARY.value
     header = PACK_LEN3(first_byte, 127, msg_length)
     raw_message = header + message
-    feed_data = reader.feed_data
 
     @benchmark
     def _run() -> None:
+        feed_data = _make_reader(loop).feed_data
         for _ in range(100):
             feed_data(raw_message)
 
@@ -45,20 +71,42 @@ def test_read_one_hundred_websocket_text_messages(
     loop: asyncio.AbstractEventLoop, benchmark: BenchmarkFixture
 ) -> None:
     """Benchmark reading 100 WebSocket text messages."""
-    queue = WebSocketDataQueue(BaseProtocol(loop), DEFAULT_CHUNK_SIZE, loop=loop)  #
-    reader = WebSocketReader(
-        queue, max_msg_size=DEFAULT_CHUNK_SIZE, compress=True, decode_text=True
-    )  #
     raw_message = (
         b'\x81~\x01!{"id":1,"src":"shellyplugus-c049ef8c30e4","dst":"aios-1453812500'
         b'8","result":{"name":null,"id":"shellyplugus-c049ef8c30e4","mac":"C049EF8C30E'
         b'4","slot":1,"model":"SNPL-00116US","gen":2,"fw_id":"20231219-133953/1.1.0-g3'
         b'4b5d4f","ver":"1.1.0","app":"PlugUS","auth_en":false,"auth_domain":null}}'
     )
-    feed_data = reader.feed_data
 
     @benchmark
     def _run() -> None:
+        feed_data = _make_reader(loop).feed_data
+        for _ in range(100):
+            feed_data(raw_message)
+
+
+def test_read_one_hundred_masked_websocket_text_messages(
+    loop: asyncio.AbstractEventLoop, benchmark: BenchmarkFixture
+) -> None:
+    """Read 100 small masked text messages, as a server receives them."""
+    raw_message = _masked_frame(WSMsgType.TEXT, b'{"id":1,"type":"ping"}')
+
+    @benchmark
+    def _run() -> None:
+        feed_data = _make_reader(loop).feed_data
+        for _ in range(100):
+            feed_data(raw_message)
+
+
+def test_read_one_hundred_masked_large_binary_websocket_messages(
+    loop: asyncio.AbstractEventLoop, benchmark: BenchmarkFixture
+) -> None:
+    """Read 100 large masked binary messages, as a server receives them."""
+    raw_message = _masked_frame(WSMsgType.BINARY, b"x" * ((2**16) + 1))
+
+    @benchmark
+    def _run() -> None:
+        feed_data = _make_reader(loop).feed_data
         for _ in range(100):
             feed_data(raw_message)