Fix memory consumption on empty/small messages (aio-libs#13393)

Co-authored-by: J. Nick Koston <nick@koston.org>

Author: Dreamsorcerer

CHANGES/13393.breaking.rst

@@ -0,0 +1 @@
+The WebSocket receive queue now only holds a weak reference to the ``WebSocketReader`` while parsing is stalled; code constructing a reader directly and passing it to ``set_parser()`` must keep its own strong reference to it, or frames the reader stopped short of parsing are lost with it -- by :user:`Dreamsorcerer`.

CHANGES/13393.bugfix.rst

@@ -0,0 +1 @@
+Fixed excessive memory consumption with small WebSocket messages -- by :user:`Dreamsorcerer`.

THREAT_MODEL.md

@@ -582,6 +582,21 @@ client-side, the writer adds masks to outgoing frames.
   fragment count (`max(1024, max_msg_size // 256)`) and pauses reading for
   backpressure once exceeded, mirroring the HTTP chunk-splits limit in
   `StreamReader` (PR #11894).
+- **PR #13393** — queue accounting only counted payload bytes, so a flood
+  of empty/tiny frames could pin unbounded per-message object overhead in
+  `WebSocketDataQueue` before the `_limit` high-water mark fired. Each
+  queued message is now charged a 128-byte accounting overhead
+  (`MSG_SIZE_OVERHEAD`) and the reader stalls at a frame boundary once the
+  queue is over the mark: the unparsed remainder stays compressed in
+  `_tail`, the queue keeps only a weakref to the stalled reader, and
+  drains re-drive it at a `_limit // 2` low-water mark (each resume
+  re-slices the tail, so per-pop resumes would make draining a tiny-frame
+  burst quadratic CPU), and the transport stays paused until the stash is
+  exhausted (resuming earlier would admit a fresh socket read into `_tail`
+  per couple of drained messages, relocating the memory bound there).
+  Callers constructing a `WebSocketReader` for
+  `set_parser()` must hold a strong reference to it (both in-tree response
+  classes do, via `_parser`).
 
 ---
 

aiohttp/_websocket/reader_c.pxd

@@ -44,18 +44,21 @@ cdef set MESSAGE_TYPES_WITH_CONTENT
 cdef tuple EMPTY_FRAME
 cdef tuple EMPTY_FRAME_ERROR
 
+cdef unsigned int MSG_SIZE_OVERHEAD
+
 cdef class WebSocketDataQueue:
 
-    cdef unsigned int _size
+    cdef readonly unsigned int _size
     cdef public object _protocol
-    cdef unsigned int _limit
+    cdef readonly unsigned int _limit
     cdef object _loop
     cdef bint _eof
     cdef object _waiter
     cdef object _exception
     cdef public object _buffer
     cdef object _get_buffer
     cdef object _put_buffer
+    cdef readonly object _stalled_reader
 
     cdef void _release_waiter(self)
 
@@ -67,6 +70,11 @@ cdef class WebSocketDataQueue:
 
 cdef class WebSocketReader:
 
+    # cdef classes are not weak-referenceable without this; the queue parks a
+    # weakref here while parsing is stalled.
+    cdef object __weakref__
+    cdef object _weak_self
+
     cdef WebSocketDataQueue queue
     cdef unsigned int _max_msg_size
     cdef bint _decode_text

aiohttp/_websocket/reader_py.py

@@ -3,11 +3,14 @@
 import asyncio
 import builtins
 import sys
+import weakref
 from collections import deque
+from typing import Final
 
 from ..base_protocol import BaseProtocol
 from ..compression_utils import TooManyMembersError, ZLibDecompressor
 from ..helpers import _EXC_SENTINEL, set_exception
+from ..log import ws_logger
 from ..streams import EofStream
 from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask
 from .models import (
@@ -60,6 +63,15 @@
 
 TUPLE_NEW = tuple.__new__
 
+# Overhead added to each message to ensure that tiny messages can't use
+# unreasonable amounts of memory.
+MSG_SIZE_OVERHEAD: Final[int] = 128
+
+STALLED_READER_COLLECTED: Final[str] = (
+    "WebSocketReader was garbage collected while stalled; "
+    "callers of set_parser() must hold a strong reference"
+)
+
 cython_int = int  # Typed to int in Python, but cython with use a signed int in the pxd
 
 
@@ -82,6 +94,7 @@ def __init__(
         self._buffer: deque[WSMessage] = deque()
         self._get_buffer = self._buffer.popleft
         self._put_buffer = self._buffer.append
+        self._stalled_reader: "weakref.ref[WebSocketReader] | None" = None
 
     def is_eof(self) -> bool:
         return self._eof
@@ -113,8 +126,10 @@ def feed_eof(self) -> None:
         self._exception = None  # Break cyclic references
 
     def feed_data(self, data: "WSMessage") -> None:
+        # Unbox into the typed local before adding, so Cython keeps the sum in
+        # C instead of boxing MSG_SIZE_OVERHEAD for a Python-level add.
         size = data.size
-        self._size += size
+        self._size += size + MSG_SIZE_OVERHEAD
         self._put_buffer(data)
         self._release_waiter()
         if self._size > self._limit and not self._protocol._reading_paused:
@@ -135,8 +150,33 @@ def _read_from_buffer(self) -> WSMessage:
         if self._buffer:
             data = self._get_buffer()
             size = data.size
-            self._size -= size
-            if self._size < self._limit and self._protocol._reading_paused:
+            self._size -= size + MSG_SIZE_OVERHEAD
+            if self._stalled_reader is not None and self._size <= self._limit // 2:
+                # Resume parsing once the queue drains to the low-water mark.
+                # Each resume re-slices the parser's unparsed tail, so waiting
+                # for headroom makes a drain cost one copy per batch of
+                # messages instead of one per message.
+                if (reader := self._stalled_reader()) is not None:
+                    reader.feed_data(b"")
+                else:
+                    # The stash died with the reader. Deliver what was already
+                    # queued, then surface the contract violation on the next
+                    # read instead of hanging. Log as well, since a caller that
+                    # stops reading early never sees the deferred exception. A
+                    # real failure that was already recorded stays the reported
+                    # cause.
+                    self._stalled_reader = None
+                    ws_logger.warning(STALLED_READER_COLLECTED)
+                    if self._exception is None:
+                        self.set_exception(RuntimeError(STALLED_READER_COLLECTED))
+            # Resuming the transport while a stash remains would admit
+            # another socket read into the tail for every couple of messages
+            # drained, moving the memory bound from the queue into the tail.
+            if (
+                self._stalled_reader is None
+                and self._size < self._limit
+                and self._protocol._reading_paused
+            ):
                 self._protocol.resume_reading()
             return data
         if self._exception is not None:
@@ -155,6 +195,9 @@ def __init__(
         self.queue = queue
         self._max_msg_size = max_msg_size
         self._decode_text = decode_text
+        # Parked on the queue while parsing is stalled; created once so
+        # stalling does not allocate.
+        self._weak_self = weakref.ref(self)
 
         self._exc: Exception | None = None
         self._partial = bytearray()
@@ -349,6 +392,7 @@ def _handle_frame(
 
     def _feed_data(self, data: bytes) -> None:
         """Return the next frame from the socket."""
+        self.queue._stalled_reader = None
         if self._tail:
             data, self._tail = self._tail + data, b""
 
@@ -357,6 +401,14 @@ def _feed_data(self, data: bytes) -> None:
         data_cstr = data
 
         while True:
+            if start_pos < data_len and self.queue._size > self.queue._limit:
+                # Over the high-water mark with unparsed bytes left: stash the
+                # remainder and stall. Gating on unparsed bytes keeps a read
+                # that ended on a frame boundary from arming an empty stall,
+                # which would hold the transport paused with nothing to drain.
+                self.queue._stalled_reader = self._weak_self
+                break
+
             # read header
             if self._state == READ_HEADER:
                 if data_len - start_pos < 2:

aiohttp/client.py

@@ -1210,14 +1210,14 @@ async def _ws_connect(
                 compress=compress,
                 client_notakeover=notakeover,
             )
-            parser = WebSocketReader(
+            ws_resp._parser = WebSocketReader(
                 reader,
                 max_msg_size,
                 compress=bool(compress),
                 decode_text=decode_text,
             )
             cb = None if heartbeat is None else ws_resp._on_data_received
-            conn_proto.set_parser(parser, reader, data_received_cb=cb)
+            conn_proto.set_parser(ws_resp._parser, reader, data_received_cb=cb)
             return ws_resp
 
     def _prepare_headers(self, headers: LooseHeaders | None) -> "CIMultiDict[str]":

aiohttp/client_ws.py

@@ -7,7 +7,7 @@
 from types import TracebackType
 from typing import Any, Final, Generic, Literal, overload
 
-from ._websocket.reader import WebSocketDataQueue
+from ._websocket.reader import WebSocketDataQueue, WebSocketReader
 from .client_exceptions import ClientError, ServerTimeoutError, WSMessageTypeError
 from .client_reqrep import ClientResponse
 from .helpers import calculate_timeout_when, frozen_dataclass_decorator, set_result
@@ -79,6 +79,10 @@ def __init__(
 
         self._writer = writer
         self._reader = reader
+        # Set by ClientSession._ws_connect; owns the parser so a stalled
+        # reader parked on the queue by weakref stays alive while this
+        # response can still be drained.
+        self._parser: WebSocketReader | None = None
         self._protocol = protocol
         self._closed = False
         self._closing = False
@@ -220,6 +224,8 @@ def _handle_ping_pong_exception(self, exc: BaseException) -> None:
         if self._closed:
             return
         self._set_closed()
+        # close() is never reached after this; release the parser here.
+        self._parser = None
         self._close_code = WSCloseCode.ABNORMAL_CLOSURE
         self._exception = exc
         self._response.close()
@@ -337,25 +343,8 @@ async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bo
 
         self._set_closed()
         try:
-            await self._writer.close(code, message)
-        except asyncio.CancelledError:
-            self._close_code = WSCloseCode.ABNORMAL_CLOSURE
-            self._response.close()
-            raise
-        except Exception as exc:
-            self._close_code = WSCloseCode.ABNORMAL_CLOSURE
-            self._exception = exc
-            self._response.close()
-            return True
-
-        if self._close_code:
-            self._response.close()
-            return True
-
-        while True:
             try:
-                async with async_timeout.timeout(self._timeout.ws_close):
-                    msg = await self._reader.read()
+                await self._writer.close(code, message)
             except asyncio.CancelledError:
                 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 self._response.close()
@@ -366,11 +355,33 @@ async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bo
                 self._response.close()
                 return True
 
-            if msg.type is WSMsgType.CLOSE:
-                self._close_code = msg.data
+            if self._close_code:
                 self._response.close()
                 return True
 
+            while True:
+                try:
+                    async with async_timeout.timeout(self._timeout.ws_close):
+                        msg = await self._reader.read()
+                except asyncio.CancelledError:
+                    self._close_code = WSCloseCode.ABNORMAL_CLOSURE
+                    self._response.close()
+                    raise
+                except Exception as exc:
+                    self._close_code = WSCloseCode.ABNORMAL_CLOSURE
+                    self._exception = exc
+                    self._response.close()
+                    return True
+
+                if msg.type is WSMsgType.CLOSE:
+                    self._close_code = msg.data
+                    self._response.close()
+                    return True
+        finally:
+            # Once closed the response can no longer be drained; release the
+            # parser and the stash it retains.
+            self._parser = None
+
     @overload
     async def receive(
         self: "ClientWebSocketResponse[Literal[True]]", timeout: float | None = None
@@ -425,7 +436,10 @@ async def receive(
                 await self.close()
                 return WS_CLOSED_MESSAGE
             except ClientError:
-                # Likely ServerDisconnectedError when connection is lost
+                # Likely ServerDisconnectedError when connection is lost.
+                # close() is not called on this path, so release the parser
+                # and the stash it retains here.
+                self._parser = None
                 self._set_closed()
                 self._close_code = WSCloseCode.ABNORMAL_CLOSURE
                 return WS_CLOSED_MESSAGE