Fix WebSocket reader flow control calculations

CHANGES/9685.breaking.rst

@@ -0,0 +1 @@
+The ``FlowControlDataQueue`` has been replaced with the ``WebSocketDataQueue`` -- by :user:`bdraco`.

aiohttp/__init__.py

@@ -84,13 +84,7 @@
     payload_type,
 )
 from .resolver import AsyncResolver, DefaultResolver, ThreadedResolver
-from .streams import (
-    EMPTY_PAYLOAD,
-    DataQueue,
-    EofStream,
-    FlowControlDataQueue,
-    StreamReader,
-)
+from .streams import EMPTY_PAYLOAD, DataQueue, EofStream, StreamReader
 from .tracing import (
     TraceConfig,
     TraceConnectionCreateEndParams,
@@ -206,7 +200,6 @@
     "DataQueue",
     "EMPTY_PAYLOAD",
     "EofStream",
-    "FlowControlDataQueue",
     "StreamReader",
     # tracing
     "TraceConfig",

aiohttp/_websocket/models.py

@@ -40,12 +40,14 @@ class WSMsgType(IntEnum):
 
 class WSMessageContinuation(NamedTuple):
     data: bytes
+    size: int
     extra: Optional[str] = None
     type: Literal[WSMsgType.CONTINUATION] = WSMsgType.CONTINUATION
 
 
 class WSMessageText(NamedTuple):
     data: str
+    size: int
     extra: Optional[str] = None
     type: Literal[WSMsgType.TEXT] = WSMsgType.TEXT
 
@@ -58,6 +60,7 @@ def json(
 
 class WSMessageBinary(NamedTuple):
     data: bytes
+    size: int
     extra: Optional[str] = None
     type: Literal[WSMsgType.BINARY] = WSMsgType.BINARY
 
@@ -70,36 +73,42 @@ def json(
 
 class WSMessagePing(NamedTuple):
     data: bytes
+    size: int
     extra: Optional[str] = None
     type: Literal[WSMsgType.PING] = WSMsgType.PING
 
 
 class WSMessagePong(NamedTuple):
     data: bytes
+    size: int
     extra: Optional[str] = None
     type: Literal[WSMsgType.PONG] = WSMsgType.PONG
 
 
 class WSMessageClose(NamedTuple):
     data: int
+    size: int
     extra: Optional[str] = None
     type: Literal[WSMsgType.CLOSE] = WSMsgType.CLOSE
 
 
 class WSMessageClosing(NamedTuple):
     data: None = None
+    size: int = 0
     extra: Optional[str] = None
     type: Literal[WSMsgType.CLOSING] = WSMsgType.CLOSING
 
 
 class WSMessageClosed(NamedTuple):
     data: None = None
+    size: int = 0
     extra: Optional[str] = None
     type: Literal[WSMsgType.CLOSED] = WSMsgType.CLOSED
 
 
 class WSMessageError(NamedTuple):
     data: BaseException
+    size: int = 0
     extra: Optional[str] = None
     type: Literal[WSMsgType.ERROR] = WSMsgType.ERROR
 

aiohttp/_websocket/reader.py

@@ -5,17 +5,27 @@
 from ..helpers import NO_EXTENSIONS
 
 if TYPE_CHECKING or NO_EXTENSIONS:  # pragma: no cover
-    from .reader_py import WebSocketReader as WebSocketReaderPython
+    from .reader_py import (
+        WebSocketDataQueue as WebSocketDataQueuePython,
+        WebSocketReader as WebSocketReaderPython,
+    )
 
     WebSocketReader = WebSocketReaderPython
+    WebSocketDataQueue = WebSocketDataQueuePython
 else:
     try:
         from .reader_c import (  # type: ignore[import-not-found]
+            WebSocketDataQueue as WebSocketDataQueueCython,
             WebSocketReader as WebSocketReaderCython,
         )
 
         WebSocketReader = WebSocketReaderCython
+        WebSocketDataQueue = WebSocketDataQueueCython
     except ImportError:  # pragma: no cover
-        from .reader_py import WebSocketReader as WebSocketReaderPython
+        from .reader_py import (
+            WebSocketDataQueue as WebSocketDataQueuePython,
+            WebSocketReader as WebSocketReaderPython,
+        )
 
         WebSocketReader = WebSocketReaderPython
+        WebSocketDataQueue = WebSocketDataQueuePython

aiohttp/_websocket/reader_c.pxd

@@ -36,11 +36,30 @@ cdef set MESSAGE_TYPES_WITH_CONTENT
 cdef tuple EMPTY_FRAME
 cdef tuple EMPTY_FRAME_ERROR
 
+cdef class WebSocketDataQueue:
+
+    cdef unsigned int _size
+    cdef public object _protocol
+    cdef 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 void _release_waiter(self)
+
+    @cython.locals(size="unsigned int")
+    cpdef void feed_data(self, object data)
+
+    @cython.locals(size="unsigned int")
+    cdef _read_from_buffer(self)
 
 cdef class WebSocketReader:
 
-    cdef object queue
-    cdef object _queue_feed_data
+    cdef WebSocketDataQueue queue
     cdef unsigned int _max_msg_size
 
     cdef Exception _exc

aiohttp/_websocket/reader_py.py

@@ -1,10 +1,14 @@
 """Reader for WebSocket protocol versions 13 and 8."""
 
-from typing import Final, List, Optional, Set, Tuple, Union
+import asyncio
+import builtins
+from collections import deque
+from typing import Deque, Final, List, Optional, Set, Tuple, Type, Union
 
+from ..base_protocol import BaseProtocol
 from ..compression_utils import ZLibDecompressor
-from ..helpers import set_exception
-from ..streams import DataQueue
+from ..helpers import _EXC_SENTINEL, set_exception
+from ..streams import EofStream
 from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask
 from .models import (
     WS_DEFLATE_TRAILING,
@@ -45,12 +49,88 @@
 TUPLE_NEW = tuple.__new__
 
 
+class WebSocketDataQueue:
+    """WebSocketDataQueue resumes and pauses an underlying stream.
+
+    It is a destination for WebSocket data.
+    """
+
+    def __init__(
+        self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop
+    ) -> None:
+        self._size = 0
+        self._protocol = protocol
+        self._limit = limit * 2
+        self._loop = loop
+        self._eof = False
+        self._waiter: Optional[asyncio.Future[None]] = None
+        self._exception: Union[Type[BaseException], BaseException, None] = None
+        self._buffer: Deque[WSMessage] = deque()
+        self._get_buffer = self._buffer.popleft
+        self._put_buffer = self._buffer.append
+
+    def exception(self) -> Optional[Union[Type[BaseException], BaseException]]:
+        return self._exception
+
+    def set_exception(
+        self,
+        exc: Union[Type[BaseException], BaseException],
+        exc_cause: builtins.BaseException = _EXC_SENTINEL,
+    ) -> None:
+        self._eof = True
+        self._exception = exc
+        if (waiter := self._waiter) is not None:
+            self._waiter = None
+            set_exception(waiter, exc, exc_cause)
+
+    def _release_waiter(self) -> None:
+        if (waiter := self._waiter) is None:
+            return
+        self._waiter = None
+        if not waiter.done():
+            waiter.set_result(None)
+
+    def feed_eof(self) -> None:
+        self._eof = True
+        self._release_waiter()
+
+    def feed_data(self, data: "WSMessage") -> None:
+        size = data.size
+        self._size += size
+        self._put_buffer(data)
+        self._release_waiter()
+        if self._size > self._limit and not self._protocol._reading_paused:
+            self._protocol.pause_reading()
+
+    async def read(self) -> WSMessage:
+        if not self._buffer and not self._eof:
+            assert not self._waiter
+            self._waiter = self._loop.create_future()
+            try:
+                await self._waiter
+            except (asyncio.CancelledError, asyncio.TimeoutError):
+                self._waiter = None
+                raise
+        return self._read_from_buffer()
+
+    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._protocol.resume_reading()
+            return data
+        if self._exception is not None:
+            raise self._exception
+        raise EofStream
+
+
 class WebSocketReader:
     def __init__(
-        self, queue: DataQueue[WSMessage], max_msg_size: int, compress: bool = True
+        self, queue: WebSocketDataQueue, max_msg_size: int, compress: bool = True
     ) -> None:
         self.queue = queue
-        self._queue_feed_data = queue.feed_data
         self._max_msg_size = max_msg_size
 
         self._exc: Optional[Exception] = None
@@ -177,6 +257,7 @@ def _feed_data(self, data: bytes) -> None:
                 else:
                     payload_merged = bytes(assembled_payload)
 
+                size = len(payload_merged)
                 if opcode == OP_CODE_TEXT:
                     try:
                         text = payload_merged.decode("utf-8")
@@ -189,15 +270,16 @@ def _feed_data(self, data: bytes) -> None:
                     # bottleneck, so we use tuple.__new__ to improve performance.
                     # This is not type safe, but many tests should fail in
                     # test_client_ws_functional.py if this is wrong.
-                    msg = TUPLE_NEW(WSMessageText, (text, "", WS_MSG_TYPE_TEXT))
+                    msg = TUPLE_NEW(WSMessageText, (text, size, "", WS_MSG_TYPE_TEXT))
                 else:
                     msg = TUPLE_NEW(
-                        WSMessageBinary, (payload_merged, "", WS_MSG_TYPE_BINARY)
+                        WSMessageBinary, (payload_merged, size, "", WS_MSG_TYPE_BINARY)
                     )
 
-                self._queue_feed_data(msg)
+                self.queue.feed_data(msg)
             elif opcode == OP_CODE_CLOSE:
-                if len(payload) >= 2:
+                payload_len = len(payload)
+                if payload_len >= 2:
                     close_code = UNPACK_CLOSE_CODE(payload[:2])[0]
                     if close_code < 3000 and close_code not in ALLOWED_CLOSE_CODES:
                         raise WebSocketError(
@@ -210,25 +292,26 @@ def _feed_data(self, data: bytes) -> None:
                         raise WebSocketError(
                             WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
                         ) from exc
-                    msg = WSMessageClose(data=close_code, extra=close_message)
+                    msg = WSMessageClose(
+                        data=close_code, size=payload_len, extra=close_message
+                    )
                 elif payload:
                     raise WebSocketError(
                         WSCloseCode.PROTOCOL_ERROR,
                         f"Invalid close frame: {fin} {opcode} {payload!r}",
                     )
                 else:
-                    msg = WSMessageClose(data=0, extra="")
-
-                self._queue_feed_data(msg)
+                    msg = WSMessageClose(data=0, size=payload_len, extra="")
 
+                self.queue.feed_data(msg)
             elif opcode == OP_CODE_PING:
-                msg = WSMessagePing(data=payload, extra="")
-                self._queue_feed_data(msg)
-
+                self.queue.feed_data(
+                    WSMessagePing(data=payload, size=len(payload), extra="")
+                )
             elif opcode == OP_CODE_PONG:
-                msg = WSMessagePong(data=payload, extra="")
-                self._queue_feed_data(msg)
-
+                self.queue.feed_data(
+                    WSMessagePong(data=payload, size=len(payload), extra="")
+                )
             else:
                 raise WebSocketError(
                     WSCloseCode.PROTOCOL_ERROR, f"Unexpected opcode={opcode!r}"

aiohttp/client.py

@@ -39,6 +39,7 @@
 from yarl import URL
 
 from . import hdrs, http, payload
+from ._websocket.reader import WebSocketDataQueue
 from .abc import AbstractCookieJar
 from .client_exceptions import (
     ClientConnectionError,
@@ -101,8 +102,7 @@
     strip_auth_from_url,
 )
 from .http import WS_KEY, HttpVersion, WebSocketReader, WebSocketWriter
-from .http_websocket import WSHandshakeError, WSMessage, ws_ext_gen, ws_ext_parse
-from .streams import FlowControlDataQueue
+from .http_websocket import WSHandshakeError, ws_ext_gen, ws_ext_parse
 from .tracing import Trace, TraceConfig
 from .typedefs import JSONEncoder, LooseCookies, LooseHeaders, Query, StrOrURL
 
@@ -1035,9 +1035,7 @@ async def _ws_connect(
 
             transport = conn.transport
             assert transport is not None
-            reader: FlowControlDataQueue[WSMessage] = FlowControlDataQueue(
-                conn_proto, 2**16, loop=self._loop
-            )
+            reader = WebSocketDataQueue(conn_proto, 2**16, loop=self._loop)
             conn_proto.set_parser(WebSocketReader(reader, max_msg_size), reader)
             writer = WebSocketWriter(
                 conn_proto,

aiohttp/client_proto.py

@@ -167,7 +167,7 @@ def set_exception(
     def set_parser(self, parser: Any, payload: Any) -> None:
         # TODO: actual types are:
         #   parser: WebSocketReader
-        #   payload: FlowControlDataQueue
+        #   payload: WebSocketDataQueue
         # but they are not generi enough
         # Need an ABC for both types
         self._payload = payload

aiohttp/client_ws.py

@@ -6,6 +6,7 @@
 from types import TracebackType
 from typing import Any, Final, Optional, Type
 
+from ._websocket.reader import WebSocketDataQueue
 from .client_exceptions import ClientError, ServerTimeoutError, WSMessageTypeError
 from .client_reqrep import ClientResponse
 from .helpers import calculate_timeout_when, set_result
@@ -18,7 +19,7 @@
     WSMsgType,
 )
 from .http_websocket import _INTERNAL_RECEIVE_TYPES, WebSocketWriter, WSMessageError
-from .streams import EofStream, FlowControlDataQueue
+from .streams import EofStream
 from .typedefs import (
     DEFAULT_JSON_DECODER,
     DEFAULT_JSON_ENCODER,
@@ -46,7 +47,7 @@ class ClientWSTimeout:
 class ClientWebSocketResponse:
     def __init__(
         self,
-        reader: "FlowControlDataQueue[WSMessage]",
+        reader: WebSocketDataQueue,
         writer: WebSocketWriter,
         protocol: Optional[str],
         response: ClientResponse,