Add support for adjusting the server WebSocket writer limit (#9572) (#9573)

Author: bdraco

CHANGES/9572.feature.rst

@@ -0,0 +1 @@
+Added ``writer_limit`` to the :py:class:`~aiohttp.web.WebSocketResponse` to be able to adjust the limit before the writer forces the buffer to be drained -- by :user:`bdraco`.

aiohttp/web_ws.py

@@ -10,6 +10,7 @@
 from multidict import CIMultiDict
 
 from . import hdrs
+from ._websocket.writer import DEFAULT_LIMIT
 from .abc import AbstractStreamWriter
 from .helpers import calculate_timeout_when, set_exception, set_result
 from .http import (
@@ -70,6 +71,7 @@ def __init__(
         protocols: Iterable[str] = (),
         compress: bool = True,
         max_msg_size: int = 4 * 1024 * 1024,
+        writer_limit: int = DEFAULT_LIMIT,
     ) -> None:
         super().__init__(status=101)
         self._protocols = protocols
@@ -97,6 +99,7 @@ def __init__(
         self._compress = compress
         self._max_msg_size = max_msg_size
         self._ping_task: Optional[asyncio.Task[None]] = None
+        self._writer_limit = writer_limit
 
     def _cancel_heartbeat(self) -> None:
         self._cancel_pong_response_cb()
@@ -305,7 +308,11 @@ def _pre_start(self, request: BaseRequest) -> Tuple[str, WebSocketWriter]:
         transport = request._protocol.transport
         assert transport is not None
         writer = WebSocketWriter(
-            request._protocol, transport, compress=compress, notakeover=notakeover
+            request._protocol,
+            transport,
+            compress=compress,
+            notakeover=notakeover,
+            limit=self._writer_limit,
         )
 
         return protocol, writer

docs/web_reference.rst

@@ -954,7 +954,8 @@ and :ref:`aiohttp-web-signals` handlers::
 
 .. class:: WebSocketResponse(*, timeout=10.0, receive_timeout=None, \
                              autoclose=True, autoping=True, heartbeat=None, \
-                             protocols=(), compress=True, max_msg_size=4194304)
+                             protocols=(), compress=True, max_msg_size=4194304, \
+                             writer_limit=65536)
 
    Class for handling server-side websockets, inherited from
    :class:`StreamResponse`.
@@ -1005,6 +1006,11 @@ and :ref:`aiohttp-web-signals` handlers::
                            ``request.transport.close()`` to avoid
                            leaking resources.
 
+   :param int writer_limit: maximum size of write buffer, 64 KB by default.
+                            Once the buffer is full, the websocket will pause
+                            to drain the buffer.
+
+      .. versionadded:: 3.11
 
    The class supports ``async for`` statement for iterating over
    incoming messages::

tests/test_web_websocket.py

@@ -249,6 +249,18 @@ def test_closed_after_ctor() -> None:
     assert ws.close_code is None
 
 
+async def test_raise_writer_limit(make_request) -> None:
+    """Test the writer limit can be adjusted."""
+    req = make_request("GET", "/")
+    ws = WebSocketResponse(writer_limit=1234567)
+    await ws.prepare(req)
+    assert ws._reader is not None
+    assert ws._writer is not None
+    assert ws._writer._limit == 1234567
+    ws._reader.feed_data(WS_CLOSED_MESSAGE)
+    await ws.close()
+
+
 async def test_send_str_closed(make_request) -> None:
     req = make_request("GET", "/")
     ws = WebSocketResponse()