Optimize WebSocketReader for 2 byte length case

aio-libs · Oct 28, 2024 · 0b57e5f · 0b57e5f
1 parent 10f07ea
commit 0b57e5f
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 5 deletions.
diff --git a/CHANGES/9554.feature.rst b/CHANGES/9554.feature.rst
@@ -0,0 +1 @@
+9543.feature.rst
diff --git a/aiohttp/_websocket/helpers.py b/aiohttp/_websocket/helpers.py
@@ -8,7 +8,6 @@
 from ..helpers import NO_EXTENSIONS
 from .models import WSHandshakeError
 
-UNPACK_LEN2 = Struct("!H").unpack_from
 UNPACK_LEN3 = Struct("!Q").unpack_from
 UNPACK_CLOSE_CODE = Struct("!H").unpack
 PACK_LEN1 = Struct("!BB").pack

diff --git a/aiohttp/_websocket/reader_c.pxd b/aiohttp/_websocket/reader_c.pxd
@@ -15,7 +15,6 @@ cdef unsigned int OP_CODE_CLOSE
 cdef unsigned int OP_CODE_PING
 cdef unsigned int OP_CODE_PONG
 
-cdef object UNPACK_LEN2
 cdef object UNPACK_LEN3
 cdef object UNPACK_CLOSE_CODE
 cdef object TUPLE_NEW

diff --git a/aiohttp/_websocket/reader_py.py b/aiohttp/_websocket/reader_py.py
@@ -5,7 +5,7 @@
 from ..compression_utils import ZLibDecompressor
 from ..helpers import set_exception
 from ..streams import DataQueue
-from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN2, UNPACK_LEN3, websocket_mask
+from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask
 from .models import (
     WS_DEFLATE_TRAILING,
     WebSocketError,
@@ -313,9 +313,10 @@ def parse_frame(
                 if length_flag == 126:
                     if buf_length - start_pos < 2:
                         break
-                    data = buf[start_pos : start_pos + 2]
+                    first_byte = buf[start_pos]
+                    second_byte = buf[start_pos + 1]
                     start_pos += 2
-                    self._payload_length = UNPACK_LEN2(data)[0]
+                    self._payload_length = first_byte << 8 | second_byte
                 elif length_flag > 126:
                     if buf_length - start_pos < 8:
                         break