Skip to content

Commit f7cac7e

Browse files
Reduce WebSocket buffer slicing overhead (#10601)
<!-- Thank you for your contribution! --> ## What do these changes do? Use a `const unsigned char *` for the buffer (Cython will automatically extract is using `__Pyx_PyBytes_AsUString`) as its a lot faster than copying around `PyBytes` objects. We do need to be careful that all slices are bounded and we bound check everything to make sure we do not do an out of bounds read since Cython does not bounds check C strings. I checked that all accesses to `buf_cstr` are proceeded by a bounds check but it would be good to get another set of eyes on that to verify in the `self._state == READ_PAYLOAD` block that we will never try to read out of bounds. <img width="376" alt="Screenshot 2025-03-19 at 10 21 54 AM" src="https://github.com/user-attachments/assets/a340ffa2-f09b-4aff-a4f7-c487dae186c8" /> ## Are there changes in behavior for the user? performance improvement ## Is it a substantial burden for the maintainers to support this? no There is a small risk that someone could remove a bounds check in the future and create a memory safety issue, however in this case its likely we would already be trying to read data that wasn't there if we are missing the bounds checking so the pure python version would throw if we are testing properly. --------- Co-authored-by: Sam Bull <[email protected]>
1 parent 8ac4830 commit f7cac7e

File tree

3 files changed

+13
-9
lines changed

3 files changed

+13
-9
lines changed

CHANGES/10601.misc.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Improved performance of WebSocket buffer handling -- by :user:`bdraco`.

aiohttp/_websocket/reader_c.pxd

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ cdef class WebSocketReader:
9999
chunk_size="unsigned int",
100100
chunk_len="unsigned int",
101101
buf_length="unsigned int",
102+
buf_cstr="const unsigned char *",
102103
first_byte="unsigned char",
103104
second_byte="unsigned char",
104105
end_pos="unsigned int",

aiohttp/_websocket/reader_py.py

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -333,14 +333,15 @@ def parse_frame(
333333

334334
start_pos: int = 0
335335
buf_length = len(buf)
336+
buf_cstr = buf
336337

337338
while True:
338339
# read header
339340
if self._state == READ_HEADER:
340341
if buf_length - start_pos < 2:
341342
break
342-
first_byte = buf[start_pos]
343-
second_byte = buf[start_pos + 1]
343+
first_byte = buf_cstr[start_pos]
344+
second_byte = buf_cstr[start_pos + 1]
344345
start_pos += 2
345346

346347
fin = (first_byte >> 7) & 1
@@ -405,14 +406,14 @@ def parse_frame(
405406
if length_flag == 126:
406407
if buf_length - start_pos < 2:
407408
break
408-
first_byte = buf[start_pos]
409-
second_byte = buf[start_pos + 1]
409+
first_byte = buf_cstr[start_pos]
410+
second_byte = buf_cstr[start_pos + 1]
410411
start_pos += 2
411412
self._payload_length = first_byte << 8 | second_byte
412413
elif length_flag > 126:
413414
if buf_length - start_pos < 8:
414415
break
415-
data = buf[start_pos : start_pos + 8]
416+
data = buf_cstr[start_pos : start_pos + 8]
416417
start_pos += 8
417418
self._payload_length = UNPACK_LEN3(data)[0]
418419
else:
@@ -424,7 +425,7 @@ def parse_frame(
424425
if self._state == READ_PAYLOAD_MASK:
425426
if buf_length - start_pos < 4:
426427
break
427-
self._frame_mask = buf[start_pos : start_pos + 4]
428+
self._frame_mask = buf_cstr[start_pos : start_pos + 4]
428429
start_pos += 4
429430
self._state = READ_PAYLOAD
430431

@@ -440,10 +441,10 @@ def parse_frame(
440441
if self._frame_payload_len:
441442
if type(self._frame_payload) is not bytearray:
442443
self._frame_payload = bytearray(self._frame_payload)
443-
self._frame_payload += buf[start_pos:end_pos]
444+
self._frame_payload += buf_cstr[start_pos:end_pos]
444445
else:
445446
# Fast path for the first frame
446-
self._frame_payload = buf[start_pos:end_pos]
447+
self._frame_payload = buf_cstr[start_pos:end_pos]
447448

448449
self._frame_payload_len += end_pos - start_pos
449450
start_pos = end_pos
@@ -469,6 +470,7 @@ def parse_frame(
469470
self._frame_payload_len = 0
470471
self._state = READ_HEADER
471472

472-
self._tail = buf[start_pos:] if start_pos < buf_length else b""
473+
# XXX: Cython needs slices to be bounded, so we can't omit the slice end here.
474+
self._tail = buf_cstr[start_pos:buf_length] if start_pos < buf_length else b""
473475

474476
return frames

0 commit comments

Comments
 (0)