Fix WebSocket reader flow control calculations

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/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_py.py

@@ -1,10 +1,12 @@
 """Reader for WebSocket protocol versions 13 and 8."""
 
+import asyncio
 from typing import Final, List, Optional, Set, Tuple, Union
 
+from ..base_protocol import BaseProtocol
 from ..compression_utils import ZLibDecompressor
 from ..helpers import set_exception
-from ..streams import DataQueue
+from ..streams import DataQueue, EofStream
 from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask
 from .models import (
     WS_DEFLATE_TRAILING,
@@ -45,9 +47,47 @@
 TUPLE_NEW = tuple.__new__
 
 
+class WebSocketDataQueue(DataQueue[WSMessage]):
+    """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:
+        super().__init__(loop=loop)
+        self._size = 0
+        self._protocol = protocol
+        self._limit = limit * 2
+
+    def feed_data(self, data: WSMessage) -> None:
+        self._size += len(data[0])
+        self._buffer.append(data)
+        if (waiter := self._waiter) is not None:
+            self._waiter = None
+            if not waiter.done():
+                waiter.set_result(None)
+        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:
+            await self._wait_for_data()
+        if self._buffer:
+            data = self._buffer.popleft()
+            self._size -= len(data)
+            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

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,
@@ -102,7 +103,6 @@
 )
 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 .tracing import Trace, TraceConfig
 from .typedefs import JSONEncoder, LooseCookies, LooseHeaders, Query, StrOrURL
 
@@ -1035,7 +1035,7 @@ async def _ws_connect(
 
             transport = conn.transport
             assert transport is not None
-            reader: FlowControlDataQueue[WSMessage] = FlowControlDataQueue(
+            reader: WebSocketDataQueue[WSMessage] = WebSocketDataQueue(
                 conn_proto, 2**16, loop=self._loop
             )
             conn_proto.set_parser(WebSocketReader(reader, max_msg_size), reader)

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[WSMessage]",
         writer: WebSocketWriter,
         protocol: Optional[str],
         response: ClientResponse,

aiohttp/streams.py

@@ -30,11 +30,9 @@
     "EofStream",
     "StreamReader",
     "DataQueue",
-    "FlowControlDataQueue",
 )
 
 _T = TypeVar("_T")
-_SizedT = TypeVar("_SizedT", bound=collections.abc.Sized)
 
 
 class EofStream(Exception):
@@ -600,15 +598,15 @@ def read_nowait(self, n: int = -1) -> bytes:
 EMPTY_PAYLOAD: Final[StreamReader] = EmptyStreamReader()
 
 
-class DataQueue(Generic[_SizedT]):
+class DataQueue(Generic[_T]):
     """DataQueue is a general-purpose blocking queue with one reader."""
 
     def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
         self._loop = loop
         self._eof = False
         self._waiter: Optional[asyncio.Future[None]] = None
         self._exception: Union[Type[BaseException], BaseException, None] = None
-        self._buffer: Deque[_SizedT] = collections.deque()
+        self._buffer: Deque[_T] = collections.deque()
 
     def __len__(self) -> int:
         return len(self._buffer)
@@ -633,7 +631,7 @@ def set_exception(
             self._waiter = None
             set_exception(waiter, exc, exc_cause)
 
-    def feed_data(self, data: _SizedT) -> None:
+    def feed_data(self, data: _T) -> None:
         self._buffer.append(data)
         if (waiter := self._waiter) is not None:
             self._waiter = None
@@ -654,7 +652,7 @@ async def _wait_for_data(self) -> None:
             self._waiter = None
             raise
 
-    async def read(self) -> _SizedT:
+    async def read(self) -> _T:
         if not self._buffer and not self._eof:
             await self._wait_for_data()
         if self._buffer:
@@ -663,42 +661,5 @@ async def read(self) -> _SizedT:
             raise self._exception
         raise EofStream
 
-    def __aiter__(self) -> AsyncStreamIterator[_SizedT]:
+    def __aiter__(self) -> AsyncStreamIterator[_T]:
         return AsyncStreamIterator(self.read)
-
-
-class FlowControlDataQueue(DataQueue[_SizedT]):
-    """FlowControlDataQueue resumes and pauses an underlying stream.
-
-    It is a destination for parsed data.
-    """
-
-    def __init__(
-        self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop
-    ) -> None:
-        super().__init__(loop=loop)
-        self._size = 0
-        self._protocol = protocol
-        self._limit = limit * 2
-
-    def feed_data(self, data: _SizedT) -> None:
-        self._size += len(data)
-        self._buffer.append(data)
-        if (waiter := self._waiter) is not None:
-            self._waiter = None
-            set_result(waiter, None)
-        if self._size > self._limit and not self._protocol._reading_paused:
-            self._protocol.pause_reading()
-
-    async def read(self) -> _SizedT:
-        if not self._buffer and not self._eof:
-            await self._wait_for_data()
-        if self._buffer:
-            data = self._buffer.popleft()
-            self._size -= len(data)
-            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

aiohttp/web_ws.py

@@ -10,6 +10,7 @@
 from multidict import CIMultiDict
 
 from . import hdrs
+from ._websocket.reader import WebSocketDataQueue
 from ._websocket.writer import DEFAULT_LIMIT
 from .abc import AbstractStreamWriter
 from .client_exceptions import WSMessageTypeError
@@ -29,7 +30,7 @@
 )
 from .http_websocket import _INTERNAL_RECEIVE_TYPES, WSMessageError
 from .log import ws_logger
-from .streams import EofStream, FlowControlDataQueue
+from .streams import EofStream
 from .typedefs import JSONDecoder, JSONEncoder
 from .web_exceptions import HTTPBadRequest, HTTPException
 from .web_request import BaseRequest
@@ -105,7 +106,7 @@ def __init__(
         self._protocols = protocols
         self._ws_protocol: Optional[str] = None
         self._writer: Optional[WebSocketWriter] = None
-        self._reader: Optional[FlowControlDataQueue[WSMessage]] = None
+        self._reader: Optional[WebSocketDataQueue[WSMessage]] = None
         self._closed = False
         self._closing = False
         self._conn_lost = 0
@@ -355,7 +356,7 @@ def _post_start(
 
         loop = self._loop
         assert loop is not None
-        self._reader = FlowControlDataQueue(request._protocol, 2**16, loop=loop)
+        self._reader = WebSocketDataQueue(request._protocol, 2**16, loop=loop)
         request.protocol.set_parser(
             WebSocketReader(self._reader, self._max_msg_size, compress=self._compress)
         )

tests/test_flowcontrol_streams.py

@@ -19,13 +19,6 @@ def stream(
     return streams.StreamReader(protocol, limit=1, loop=loop)
 
 
-@pytest.fixture
-def buffer(
-    loop: asyncio.AbstractEventLoop, protocol: BaseProtocol
-) -> streams.FlowControlDataQueue[str]:
-    return streams.FlowControlDataQueue[str](protocol, limit=1, loop=loop)
-
-
 class TestFlowControlStreamReader:
     async def test_read(self, stream: streams.StreamReader) -> None:
         stream.feed_data(b"da")
@@ -114,20 +107,3 @@ async def test_read_nowait(self, stream: streams.StreamReader) -> None:
         res = stream.read_nowait(5)
         assert res == b""
         assert stream._protocol.resume_reading.call_count == 1  # type: ignore[attr-defined]
-
-
-class TestFlowControlDataQueue:
-    def test_feed_pause(self, buffer: streams.FlowControlDataQueue[str]) -> None:
-        buffer._protocol._reading_paused = False
-        buffer.feed_data("x" * 100)
-
-        assert buffer._protocol.pause_reading.called  # type: ignore[attr-defined]
-
-    async def test_resume_on_read(
-        self, buffer: streams.FlowControlDataQueue[str]
-    ) -> None:
-        buffer.feed_data("x" * 100)
-
-        buffer._protocol._reading_paused = True
-        await buffer.read()
-        assert buffer._protocol.resume_reading.called  # type: ignore[attr-defined]

tests/test_websocket_data_queue.py

@@ -0,0 +1,36 @@
+import asyncio
+from unittest import mock
+
+import pytest
+
+from aiohttp._websocket.models import WSMessageBinary
+from aiohttp._websocket.reader import WebSocketDataQueue
+from aiohttp.base_protocol import BaseProtocol
+
+
+@pytest.fixture
+def protocol() -> BaseProtocol:
+    return mock.create_autospec(BaseProtocol, spec_set=True, instance=True, _reading_paused=False)  # type: ignore[no-any-return]
+
+
+@pytest.fixture
+def buffer(
+    loop: asyncio.AbstractEventLoop, protocol: BaseProtocol
+) -> WebSocketDataQueue:
+    return WebSocketDataQueue(protocol, limit=1, loop=loop)
+
+
+class TestWebSocketDataQueue:
+    def test_feed_pause(self, buffer: WebSocketDataQueue) -> None:
+        buffer._protocol._reading_paused = False
+        for _ in range(3):
+            buffer.feed_data(WSMessageBinary(b"x"))
+
+        assert buffer._protocol.pause_reading.called  # type: ignore[attr-defined]
+
+    async def test_resume_on_read(self, buffer: WebSocketDataQueue) -> None:
+        buffer.feed_data(WSMessageBinary(b"x"))
+
+        buffer._protocol._reading_paused = True
+        await buffer.read()
+        assert buffer._protocol.resume_reading.called  # type: ignore[attr-defined]