Add trio message assembler.

Author: aaugustin

src/websockets/asyncio/messages.py

@@ -81,8 +81,7 @@ class Assembler:
 
     """
 
-    # coverage reports incorrectly: "line NN didn't jump to the function exit"
-    def __init__(  # pragma: no cover
+    def __init__(
         self,
         high: int | None = None,
         low: int | None = None,

src/websockets/trio/messages.py

@@ -0,0 +1,285 @@
+from __future__ import annotations
+
+import codecs
+import math
+from collections.abc import AsyncIterator
+from typing import Any, Callable, Literal, TypeVar, overload
+
+import trio
+
+from ..exceptions import ConcurrencyError
+from ..frames import OP_BINARY, OP_CONT, OP_TEXT, Frame
+from ..typing import Data
+
+
+__all__ = ["Assembler"]
+
+UTF8Decoder = codecs.getincrementaldecoder("utf-8")
+
+T = TypeVar("T")
+
+
+class Assembler:
+    """
+    Assemble messages from frames.
+
+    :class:`Assembler` expects only data frames. The stream of frames must
+    respect the protocol; if it doesn't, the behavior is undefined.
+
+    Args:
+        pause: Called when the buffer of frames goes above the high water mark;
+            should pause reading from the network.
+        resume: Called when the buffer of frames goes below the low water mark;
+            should resume reading from the network.
+
+    """
+
+    def __init__(
+        self,
+        high: int | None = None,
+        low: int | None = None,
+        pause: Callable[[], Any] = lambda: None,
+        resume: Callable[[], Any] = lambda: None,
+    ) -> None:
+        # Queue of incoming frames.
+        self.send_frames: trio.MemorySendChannel[Frame]
+        self.recv_frames: trio.MemoryReceiveChannel[Frame]
+        self.send_frames, self.recv_frames = trio.open_memory_channel(math.inf)
+
+        # We cannot put a hard limit on the size of the queue because a single
+        # call to Protocol.data_received() could produce thousands of frames,
+        # which must be buffered. Instead, we pause reading when the buffer goes
+        # above the high limit and we resume when it goes under the low limit.
+        if high is not None and low is None:
+            low = high // 4
+        if high is None and low is not None:
+            high = low * 4
+        if high is not None and low is not None:
+            if low < 0:
+                raise ValueError("low must be positive or equal to zero")
+            if high < low:
+                raise ValueError("high must be greater than or equal to low")
+        self.high, self.low = high, low
+        self.pause = pause
+        self.resume = resume
+        self.paused = False
+
+        # This flag prevents concurrent calls to get() by user code.
+        self.get_in_progress = False
+
+        # This flag marks the end of the connection.
+        self.closed = False
+
+    @overload
+    async def get(self, decode: Literal[True]) -> str: ...
+
+    @overload
+    async def get(self, decode: Literal[False]) -> bytes: ...
+
+    @overload
+    async def get(self, decode: bool | None = None) -> Data: ...
+
+    async def get(self, decode: bool | None = None) -> Data:
+        """
+        Read the next message.
+
+        :meth:`get` returns a single :class:`str` or :class:`bytes`.
+
+        If the message is fragmented, :meth:`get` waits until the last frame is
+        received, then it reassembles the message and returns it. To receive
+        messages frame by frame, use :meth:`get_iter` instead.
+
+        Args:
+            decode: :obj:`False` disables UTF-8 decoding of text frames and
+                returns :class:`bytes`. :obj:`True` forces UTF-8 decoding of
+                binary frames and returns :class:`str`.
+
+        Raises:
+            EOFError: If the stream of frames has ended.
+            UnicodeDecodeError: If a text frame contains invalid UTF-8.
+            ConcurrencyError: If two coroutines run :meth:`get` or
+                :meth:`get_iter` concurrently.
+
+        """
+        if self.get_in_progress:
+            raise ConcurrencyError("get() or get_iter() is already running")
+        self.get_in_progress = True
+
+        # Locking with get_in_progress prevents concurrent execution
+        # until get() fetches a complete message or is canceled.
+
+        try:
+            # First frame
+            try:
+                frame = await self.recv_frames.receive()
+            except trio.EndOfChannel:
+                raise EOFError("stream of frames ended")
+            self.maybe_resume()
+            assert frame.opcode is OP_TEXT or frame.opcode is OP_BINARY
+            if decode is None:
+                decode = frame.opcode is OP_TEXT
+            frames = [frame]
+
+            # Following frames, for fragmented messages
+            while not frame.fin:
+                try:
+                    frame = await self.recv_frames.receive()
+                except trio.Cancelled:
+                    # Put frames already received back into the queue
+                    # so that future calls to get() can return them.
+                    assert not self.send_frames._state.receive_tasks, (
+                        "no task should be waiting on receive()"
+                    )
+                    assert not self.send_frames._state.data, "queue should be empty"
+                    for frame in frames:
+                        self.send_frames.send_nowait(frame)
+                    raise
+                except trio.EndOfChannel:
+                    raise EOFError("stream of frames ended")
+                self.maybe_resume()
+                assert frame.opcode is OP_CONT
+                frames.append(frame)
+
+        finally:
+            self.get_in_progress = False
+
+        data = b"".join(frame.data for frame in frames)
+        if decode:
+            return data.decode()
+        else:
+            return data
+
+    @overload
+    def get_iter(self, decode: Literal[True]) -> AsyncIterator[str]: ...
+
+    @overload
+    def get_iter(self, decode: Literal[False]) -> AsyncIterator[bytes]: ...
+
+    @overload
+    def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]: ...
+
+    async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
+        """
+        Stream the next message.
+
+        Iterating the return value of :meth:`get_iter` asynchronously yields a
+        :class:`str` or :class:`bytes` for each frame in the message.
+
+        The iterator must be fully consumed before calling :meth:`get_iter` or
+        :meth:`get` again. Else, :exc:`ConcurrencyError` is raised.
+
+        This method only makes sense for fragmented messages. If messages aren't
+        fragmented, use :meth:`get` instead.
+
+        Args:
+            decode: :obj:`False` disables UTF-8 decoding of text frames and
+                returns :class:`bytes`. :obj:`True` forces UTF-8 decoding of
+                binary frames and returns :class:`str`.
+
+        Raises:
+            EOFError: If the stream of frames has ended.
+            UnicodeDecodeError: If a text frame contains invalid UTF-8.
+            ConcurrencyError: If two coroutines run :meth:`get` or
+                :meth:`get_iter` concurrently.
+
+        """
+        if self.get_in_progress:
+            raise ConcurrencyError("get() or get_iter() is already running")
+        self.get_in_progress = True
+
+        # Locking with get_in_progress prevents concurrent execution
+        # until get_iter() fetches a complete message or is canceled.
+
+        # If get_iter() raises an exception e.g. in decoder.decode(),
+        # get_in_progress remains set and the connection becomes unusable.
+
+        # First frame
+        try:
+            frame = await self.recv_frames.receive()
+        except trio.Cancelled:
+            self.get_in_progress = False
+            raise
+        except trio.EndOfChannel:
+            raise EOFError("stream of frames ended")
+        self.maybe_resume()
+        assert frame.opcode is OP_TEXT or frame.opcode is OP_BINARY
+        if decode is None:
+            decode = frame.opcode is OP_TEXT
+        if decode:
+            decoder = UTF8Decoder()
+            yield decoder.decode(frame.data, frame.fin)
+        else:
+            yield frame.data
+
+        # Following frames, for fragmented messages
+        while not frame.fin:
+            # We cannot handle trio.Cancelled because we don't buffer
+            # previous fragments — we're streaming them. Canceling get_iter()
+            # here will leave the assembler in a stuck state. Future calls to
+            # get() or get_iter() will raise ConcurrencyError.
+            try:
+                frame = await self.recv_frames.receive()
+            except trio.EndOfChannel:
+                raise EOFError("stream of frames ended")
+            self.maybe_resume()
+            assert frame.opcode is OP_CONT
+            if decode:
+                yield decoder.decode(frame.data, frame.fin)
+            else:
+                yield frame.data
+
+        self.get_in_progress = False
+
+    def put(self, frame: Frame) -> None:
+        """
+        Add ``frame`` to the next message.
+
+        Raises:
+            EOFError: If the stream of frames has ended.
+
+        """
+        if self.closed:
+            raise EOFError("stream of frames ended")
+
+        self.send_frames.send_nowait(frame)
+        self.maybe_pause()
+
+    def maybe_pause(self) -> None:
+        """Pause the writer if queue is above the high water mark."""
+        # Skip if flow control is disabled
+        if self.high is None:
+            return
+
+        # Bypass the statistics() method for performance reasons.
+        # Check for "> high" to support high = 0
+        if len(self.send_frames._state.data) > self.high and not self.paused:
+            self.paused = True
+            self.pause()
+
+    def maybe_resume(self) -> None:
+        """Resume the writer if queue is below the low water mark."""
+        # Skip if flow control is disabled
+        if self.low is None:
+            return
+
+        # Bypass the statistics() method for performance reasons.
+        # Check for "<= low" to support low = 0
+        if len(self.send_frames._state.data) <= self.low and self.paused:
+            self.paused = False
+            self.resume()
+
+    def close(self) -> None:
+        """
+        End the stream of frames.
+
+        Calling :meth:`close` concurrently with :meth:`get`, :meth:`get_iter`,
+        or :meth:`put` is safe. They will raise :exc:`trio.EndOfChannel`.
+
+        """
+        if self.closed:
+            return
+
+        self.closed = True
+
+        # Unblock get() or get_iter().
+        self.send_frames.close()

tests/asyncio/test_messages.py

@@ -267,6 +267,7 @@ async def test_get_iter_fragmented_text_message_not_received_yet(self):
         self.assertEqual(await anext(iterator), "f")
         self.assembler.put(Frame(OP_CONT, b"\xa9"))
         self.assertEqual(await anext(iterator), "é")
+        await iterator.aclose()
 
     async def test_get_iter_fragmented_binary_message_not_received_yet(self):
         """get_iter yields a fragmented binary message when it is received."""
@@ -277,6 +278,7 @@ async def test_get_iter_fragmented_binary_message_not_received_yet(self):
         self.assertEqual(await anext(iterator), b"e")
         self.assembler.put(Frame(OP_CONT, b"a"))
         self.assertEqual(await anext(iterator), b"a")
+        await iterator.aclose()
 
     async def test_get_iter_fragmented_text_message_being_received(self):
         """get_iter yields a fragmented text message that is partially received."""
@@ -287,6 +289,7 @@ async def test_get_iter_fragmented_text_message_being_received(self):
         self.assertEqual(await anext(iterator), "f")
         self.assembler.put(Frame(OP_CONT, b"\xa9"))
         self.assertEqual(await anext(iterator), "é")
+        await iterator.aclose()
 
     async def test_get_iter_fragmented_binary_message_being_received(self):
         """get_iter yields a fragmented binary message that is partially received."""
@@ -297,6 +300,7 @@ async def test_get_iter_fragmented_binary_message_being_received(self):
         self.assertEqual(await anext(iterator), b"e")
         self.assembler.put(Frame(OP_CONT, b"a"))
         self.assertEqual(await anext(iterator), b"a")
+        await iterator.aclose()
 
     async def test_get_iter_encoded_text_message(self):
         """get_iter yields a text message without UTF-8 decoding."""
@@ -334,6 +338,8 @@ async def test_get_iter_resumes_reading(self):
         await anext(iterator)
         self.resume.assert_called_once_with()
 
+        await iterator.aclose()
+
     async def test_get_iter_does_not_resume_reading(self):
         """get_iter does not resume reading when the low-water mark is unset."""
         self.assembler.low = None
@@ -345,6 +351,7 @@ async def test_get_iter_does_not_resume_reading(self):
         await anext(iterator)
         await anext(iterator)
         await anext(iterator)
+        await iterator.aclose()
 
         self.resume.assert_not_called()
 
@@ -467,7 +474,7 @@ async def test_get_iter_queued_fragmented_message_after_close(self):
         self.assertEqual(fragments, [b"t", b"e", b"a"])
 
     async def test_get_partially_queued_fragmented_message_after_close(self):
-        """get raises EOF on a partial fragmented message after close is called."""
+        """get raises EOFError on a partial fragmented message after close is called."""
         self.assembler.put(Frame(OP_BINARY, b"t", fin=False))
         self.assembler.put(Frame(OP_CONT, b"e", fin=False))
         self.assembler.close()