[PR #9679/3f2f4a7 backport][3.11] Return early in WebSocket receive if there is no processing to do (#9683)

Author: bdraco

CHANGES/9679.misc.rst

@@ -0,0 +1 @@
+Improved performance of calling ``receive`` for WebSockets for the most common message types -- by :user:`bdraco`.

aiohttp/client_ws.py

@@ -18,7 +18,7 @@
     WSMessage,
     WSMsgType,
 )
-from .http_websocket import WebSocketWriter  # WSMessage
+from .http_websocket import _INTERNAL_RECEIVE_TYPES, WebSocketWriter
 from .streams import EofStream, FlowControlDataQueue
 from .typedefs import (
     DEFAULT_JSON_DECODER,
@@ -359,6 +359,11 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
                 await self.close()
                 return WSMessage(WSMsgType.ERROR, exc, None)
 
+            if msg.type not in _INTERNAL_RECEIVE_TYPES:
+                # If its not a close/closing/ping/pong message
+                # we can return it immediately
+                return msg
+
             if msg.type is WSMsgType.CLOSE:
                 self._set_closing()
                 self._close_code = msg.data

aiohttp/http_websocket.py

@@ -13,6 +13,12 @@
 from ._websocket.reader import WebSocketReader
 from ._websocket.writer import WebSocketWriter
 
+# Messages that the WebSocketResponse.receive needs to handle internally
+_INTERNAL_RECEIVE_TYPES = frozenset(
+    (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.PING, WSMsgType.PONG)
+)
+
+
 __all__ = (
     "WS_CLOSED_MESSAGE",
     "WS_CLOSING_MESSAGE",

aiohttp/web_ws.py

@@ -27,6 +27,7 @@
     ws_ext_gen,
     ws_ext_parse,
 )
+from .http_websocket import _INTERNAL_RECEIVE_TYPES
 from .log import ws_logger
 from .streams import EofStream, FlowControlDataQueue
 from .typedefs import JSONDecoder, JSONEncoder
@@ -555,6 +556,11 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
                 await self.close()
                 return WSMessage(WSMsgType.ERROR, exc, None)
 
+            if msg.type not in _INTERNAL_RECEIVE_TYPES:
+                # If its not a close/closing/ping/pong message
+                # we can return it immediately
+                return msg
+
             if msg.type is WSMsgType.CLOSE:
                 self._set_closing(msg.data)
                 # Could be closed while awaiting reader.