Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add the ability to specify the buffer size. #186

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Add the ability to specify the buffer size. #186

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4bb02a4
Add the ability to specify the buffer size.
Nov 16, 2023
e988b56
add type annotation
Nov 16, 2023
a6d4816
properly add receive_buffer_size everywhere
Nov 23, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
11 changes: 8 additions & 3 deletions trio_websocket/_impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
CONN_TIMEOUT = 60 # default connect & disconnect timeout, in seconds
MESSAGE_QUEUE_SIZE = 1
MAX_MESSAGE_SIZE = 2 ** 20 # 1 MiB
RECEIVE_BYTES = 4 * 2 ** 10 # 4 KiB
DEFAULT_RECEIVE_BYTES = 4 * 2 ** 10 # 4 KiB
palkeo marked this conversation as resolved.
Show resolved Hide resolved
logger = logging.getLogger('trio-websocket')


Expand Down Expand Up @@ -687,7 +687,8 @@ class WebSocketConnection(trio.abc.AsyncResource):
def __init__(self, stream, ws_connection, *, host=None, path=None,
client_subprotocols=None, client_extra_headers=None,
message_queue_size=MESSAGE_QUEUE_SIZE,
max_message_size=MAX_MESSAGE_SIZE):
max_message_size=MAX_MESSAGE_SIZE,
receive_buffer_size=DEFAULT_RECEIVE_BYTES):
palkeo marked this conversation as resolved.
Show resolved Hide resolved
'''
Constructor.

Expand All @@ -713,6 +714,9 @@ def __init__(self, stream, ws_connection, *, host=None, path=None,
:param int max_message_size: The maximum message size as measured by
``len()``. If a message is received that is larger than this size,
then the connection is closed with code 1009 (Message Too Big).
:param Optional[int] receive_buffer_size: The buffer size we use to
receive messages internally. None to let trio choose. Defaults
to 4 KiB.
'''
# NOTE: The implementation uses _close_reason for more than an advisory
# purpose. It's critical internal state, indicating when the
Expand All @@ -725,6 +729,7 @@ def __init__(self, stream, ws_connection, *, host=None, path=None,
self._message_size = 0
self._message_parts: List[Union[bytes, str]] = []
self._max_message_size = max_message_size
self._receive_buffer_size: Optional[int] = receive_buffer_size
self._reader_running = True
if ws_connection.client:
self._initial_request: Optional[Request] = Request(host=host, target=path,
Expand Down Expand Up @@ -1232,7 +1237,7 @@ async def _reader_task(self):

# Get network data.
try:
data = await self._stream.receive_some(RECEIVE_BYTES)
data = await self._stream.receive_some(self._receive_buffer_size)
except (trio.BrokenResourceError, trio.ClosedResourceError):
await self._abort_web_socket()
break
Expand Down