diff --git a/CHANGES/13561.contrib.rst b/CHANGES/13561.contrib.rst new file mode 100644 index 00000000000..7066f0b6516 --- /dev/null +++ b/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`. diff --git a/tests/test_benchmarks_http_websocket.py b/tests/test_benchmarks_http_websocket.py index 8a64d5585ca..19362ff2807 100644 --- a/tests/test_benchmarks_http_websocket.py +++ b/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)