Resolve "Default max_queue blocks websocket cancellation with high traffic"

Author: btschwertfeger

src/websockets/asyncio/connection.py

@@ -913,6 +913,9 @@ async def send_context(
         if wait_for_close:
             try:
                 async with asyncio_timeout_at(self.close_deadline):
+                    self.recv_messages.cancelling = True
+                    if self.recv_messages.paused:
+                        self.recv_messages.resume()
                     await asyncio.shield(self.connection_lost_waiter)
             except TimeoutError:
                 # There's no risk to overwrite another error because

src/websockets/asyncio/messages.py

@@ -113,9 +113,13 @@ def __init__(  # pragma: no cover
         # This flag prevents concurrent calls to get() by user code.
         self.get_in_progress = False
 
+        # This flag marks a soon cancellation
+        self.cancelling = False
+
         # This flag marks the end of the connection.
         self.closed = False
 
+
     async def get(self, decode: bool | None = None) -> Data:
         """
         Read the next message.
@@ -138,6 +142,8 @@ async def get(self, decode: bool | None = None) -> Data:
                 :meth:`get_iter` concurrently.
 
         """
+        if self.cancelling:
+            return
         if self.get_in_progress:
             raise ConcurrencyError("get() or get_iter() is already running")
         self.get_in_progress = True
@@ -201,6 +207,8 @@ async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
                 :meth:`get_iter` concurrently.
 
         """
+        if self.cancelling:
+            return
         if self.get_in_progress:
             raise ConcurrencyError("get() or get_iter() is already running")
         self.get_in_progress = True
@@ -251,6 +259,8 @@ def put(self, frame: Frame) -> None:
             EOFError: If the stream of frames has ended.
 
         """
+        if self.cancelling:
+            return
         if self.closed:
             raise EOFError("stream of frames ended")
 
@@ -283,7 +293,7 @@ def close(self) -> None:
         """
         End the stream of frames.
 
-        Callling :meth:`close` concurrently with :meth:`get`, :meth:`get_iter`,
+        Calling :meth:`close` concurrently with :meth:`get`, :meth:`get_iter`,
         or :meth:`put` is safe. They will raise :exc:`EOFError`.
 
         """