Fix backpressure in the Trio implementation.

Previously, saturating the buffer would crash, due to attempting
1/ to re-acquire a lock and 2/ to release it from another task.

Test backpressure in the asyncio and threading implementations.

Author: aaugustin

docs/project/changelog.rst

@@ -38,6 +38,9 @@ Bug fixes
 * Restored compatibility of :meth:`~asyncio.server.Server.serve_forever` in the
   :mod:`asyncio` implementation with third-party event loops such as uvloop_.
 
+* Prevented the Trio implementation from crashing when backpressure kicks in,
+  i.e. when receiving data faster than the application can process it.
+
 .. _uvloop: https://uvloop.readthedocs.io/
 
 .. _17.0:

src/websockets/trio/connection.py

@@ -84,14 +84,14 @@ def __init__(
         """Opening handshake response."""
 
         # Lock stopping reads when the assembler buffer is full.
-        self.recv_flow_control = trio.Lock()
+        self.recv_flow_control: trio.Event | None = None
 
         # Assembler turning frames into messages and serializing reads.
         self.recv_messages = Assembler(
             max_queue_high,
             max_queue_low,
-            pause=self.recv_flow_control.acquire_nowait,
-            resume=self.recv_flow_control.release,
+            pause=self.pause_reading,
+            resume=self.resume_reading,
         )
 
         # Deadline for the closing handshake.
@@ -882,8 +882,8 @@ async def recv_events(self) -> None:
             while True:
                 try:
                     # If the assembler buffer is full, block until it drains.
-                    async with self.recv_flow_control:
-                        pass
+                    if self.recv_flow_control is not None:
+                        await self.recv_flow_control.wait()
                     data = await self.stream.receive_some()
                 except Exception as exc:
                     if self.debug:
@@ -966,6 +966,16 @@ async def recv_events(self) -> None:
             # This isn't expected to raise an exception.
             await self.close_stream()
 
+    def pause_reading(self) -> None:
+        """Pause recv_events() until resume_reading() is called."""
+        self.recv_flow_control = trio.Event()
+
+    def resume_reading(self) -> None:
+        """Resume recv_events() after pause_reading() was called."""
+        assert self.recv_flow_control is not None
+        self.recv_flow_control.set()
+        self.recv_flow_control = None
+
     @contextlib.asynccontextmanager
     async def send_context(
         self,

tests/asyncio/test_connection.py

@@ -1255,6 +1255,38 @@ async def test_close_reason(self):
         """Connection has a close_reason attribute."""
         self.assertIsNone(self.connection.close_reason)
 
+    # Test backpressure.
+
+    async def test_backpressure(self):
+        """Connection stops reading from the network when the buffer fills up."""
+        self.connection.recv_messages.high = 1
+        self.connection.recv_messages.low = 0
+        frames_buffer = self.connection.recv_messages.frames
+
+        await self.remote_connection.send("A")
+        await self.remote_connection.send("B")
+
+        # Assembler buffer is above the high water mark. Backpressure kicks in.
+        await self.wait_for_remote_side()
+        self.assertFalse(self.transport.is_reading())
+        self.assertEqual(len(frames_buffer), 2)
+
+        await self.remote_connection.send("C")
+
+        # A third frame is sent by the peer but the connection doesn't read it.
+        await self.wait_for_remote_side()
+        self.assertEqual(len(frames_buffer), 2)
+
+        self.assertEqual(await self.connection.recv(), "A")
+        self.assertEqual(await self.connection.recv(), "B")
+
+        # Draining the assembler buffer below the low water mark resumes reading.
+        await self.wait_for_remote_side()
+        self.assertTrue(self.transport.is_reading())
+        self.assertEqual(len(frames_buffer), 1)
+
+        self.assertEqual(await self.connection.recv(), "C")
+
     # Test reporting of network errors.
 
     async def test_writing_in_data_received_fails(self):

tests/sync/test_connection.py

@@ -949,6 +949,38 @@ def test_close_reason(self):
         """Connection has a close_reason attribute."""
         self.assertIsNone(self.connection.close_reason)
 
+    # Test backpressure.
+
+    def test_backpressure(self):
+        """Connection stops reading from the network when the buffer fills up."""
+        self.connection.recv_messages.high = 1
+        self.connection.recv_messages.low = 0
+        frames_buffer = self.connection.recv_messages.frames
+
+        self.remote_connection.send("A")
+        self.remote_connection.send("B")
+
+        # Assembler buffer is above the high water mark. Backpressure kicks in.
+        self.wait_for_remote_side()
+        self.assertTrue(self.connection.recv_flow_control.locked())
+        self.assertEqual(frames_buffer.qsize(), 2)
+
+        self.remote_connection.send("C")
+
+        # A third frame is sent by the peer but the connection doesn't read it.
+        self.wait_for_remote_side()
+        self.assertEqual(frames_buffer.qsize(), 2)
+
+        self.assertEqual(self.connection.recv(), "A")
+        self.assertEqual(self.connection.recv(), "B")
+
+        # Draining the assembler buffer below the low water mark resumes reading.
+        self.wait_for_remote_side()
+        self.assertFalse(self.connection.recv_flow_control.locked())
+        self.assertEqual(frames_buffer.qsize(), 1)
+
+        self.assertEqual(self.connection.recv(), "C")
+
     # Test reporting of network errors.
 
     def test_writing_in_recv_events_fails(self):

tests/trio/test_connection.py

@@ -1210,6 +1210,39 @@ async def test_close_reason(self):
         """Connection has a close_reason attribute."""
         self.assertIsNone(self.connection.close_reason)
 
+    # Test backpressure.
+
+    async def test_backpressure(self):
+        """Connection stops reading from the network when the buffer fills up."""
+        self.connection.recv_messages.high = 1
+        self.connection.recv_messages.low = 0
+        frames_buffer = self.connection.recv_messages.recv_frames
+
+        await self.remote_connection.send("A")
+        await self.remote_connection.send("B")
+
+        # Assembler buffer is above the high water mark. Backpressure kicks in.
+        await trio.testing.wait_all_tasks_blocked()
+        self.assertIsNotNone(self.connection.recv_flow_control)
+        self.assertFalse(self.connection.recv_flow_control.is_set())
+        self.assertEqual(frames_buffer.statistics().current_buffer_used, 2)
+
+        await self.remote_connection.send("C")
+
+        # A third frame is sent by the peer but the connection doesn't read it.
+        await trio.testing.wait_all_tasks_blocked()
+        self.assertEqual(frames_buffer.statistics().current_buffer_used, 2)
+
+        self.assertEqual(await self.connection.recv(), "A")
+        self.assertEqual(await self.connection.recv(), "B")
+
+        # Draining the assembler buffer below the low water mark resumes reading.
+        await trio.testing.wait_all_tasks_blocked()
+        self.assertIsNone(self.connection.recv_flow_control)
+        self.assertEqual(frames_buffer.statistics().current_buffer_used, 1)
+
+        self.assertEqual(await self.connection.recv(), "C")
+
     # Test reporting of network errors.
 
     async def test_writing_in_recv_events_fails(self):