Skip to content

Commit

Permalink
Add a PXD only Cython extension for the WebSocket reader to improve p…
Browse files Browse the repository at this point in the history
…erformance (#9543)
  • Loading branch information
bdraco authored Oct 27, 2024
1 parent dd9a1fd commit 6ae2f05
Show file tree
Hide file tree
Showing 10 changed files with 516 additions and 357 deletions.
1 change: 1 addition & 0 deletions CHANGES/9543.feature.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Improved performance of reading WebSocket messages with a Cython implementation -- by :user:`bdraco`.
8 changes: 7 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,11 @@ endif
aiohttp/_find_header.c: $(call to-hash,aiohttp/hdrs.py ./tools/gen.py)
./tools/gen.py

# Special case for reader since we want to be able to disable
# the extension with AIOHTTP_NO_EXTENSIONS
aiohttp/_websocket/reader_c.c: aiohttp/_websocket/reader_c.py
cython -3 -o $@ $< -I aiohttp -Werror

# _find_headers generator creates _headers.pyi as well
aiohttp/%.c: aiohttp/%.pyx $(call to-hash,$(CYS)) aiohttp/_find_header.c
cython -3 -o $@ $< -I aiohttp -Werror
Expand All @@ -74,7 +79,7 @@ vendor/llhttp/node_modules: vendor/llhttp/package.json
generate-llhttp: .llhttp-gen

.PHONY: cythonize
cythonize: .install-cython $(PYXS:.pyx=.c)
cythonize: .install-cython $(PYXS:.pyx=.c) aiohttp/_websocket/reader_c.c

.install-deps: .install-cython $(PYXS:.pyx=.c) $(call to-hash,$(CYS) $(REQS))
@python -m pip install -r requirements/dev.in -c requirements/dev.txt
Expand Down Expand Up @@ -157,6 +162,7 @@ clean:
@rm -f aiohttp/_http_parser.c
@rm -f aiohttp/_http_writer.c
@rm -f aiohttp/_websocket.c
@rm -f aiohttp/_websocket/reader_c.c
@rm -rf .tox
@rm -f .develop
@rm -f .flake
Expand Down
3 changes: 3 additions & 0 deletions aiohttp/_websocket/mask.pxd
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
"""Cython declarations for websocket masking."""

cpdef void _websocket_mask_cython(bytes mask, bytearray data)
10 changes: 1 addition & 9 deletions aiohttp/_websocket/mask.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ cdef extern from "Python.h":
from libc.stdint cimport uint32_t, uint64_t, uintmax_t


def _websocket_mask_cython(object mask, object data):
cpdef void _websocket_mask_cython(bytes mask, bytearray data):
"""Note, this function mutates its `data` argument
"""
cdef:
Expand All @@ -21,14 +21,6 @@ def _websocket_mask_cython(object mask, object data):

assert len(mask) == 4

if not isinstance(mask, bytes):
mask = bytes(mask)

if isinstance(data, bytearray):
data = <bytearray>data
else:
data = bytearray(data)

data_len = len(data)
in_buf = <unsigned char*>PyByteArray_AsString(data)
mask_buf = <const unsigned char*>PyBytes_AsString(mask)
Expand Down
Loading

0 comments on commit 6ae2f05

Please sign in to comment.