From 3a9de0c1457e04bbe81acfefd031ff436c1da98d Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 8 Aug 2024 19:17:35 -0500
Subject: [PATCH] [PR #8660/14d5295 backport][3.10] Improve performance of
 WebSockets when there is no timeout (#8663)

---
 CHANGES/8660.misc.rst |  3 +++
 aiohttp/client_ws.py  | 11 ++++++++++-
 aiohttp/web_ws.py     | 10 +++++++++-
 3 files changed, 22 insertions(+), 2 deletions(-)
 create mode 100644 CHANGES/8660.misc.rst

diff --git a/CHANGES/8660.misc.rst b/CHANGES/8660.misc.rst
new file mode 100644
index 00000000000..8710063329e
--- /dev/null
+++ b/CHANGES/8660.misc.rst
@@ -0,0 +1,3 @@
+Improved performance of :py:meth:`~aiohttp.ClientWebSocketResponse.receive` and :py:meth:`~aiohttp.web.WebSocketResponse.receive` when there is no timeout. -- by :user:`bdraco`.
+
+The timeout context manager is now avoided when there is no timeout as it accounted for up to 50% of the time spent in the :py:meth:`~aiohttp.ClientWebSocketResponse.receive` and :py:meth:`~aiohttp.web.WebSocketResponse.receive` methods.
diff --git a/aiohttp/client_ws.py b/aiohttp/client_ws.py
index 247f62c758e..7fd141248bd 100644
--- a/aiohttp/client_ws.py
+++ b/aiohttp/client_ws.py
@@ -281,6 +281,8 @@ async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bo
             return False
 
     async def receive(self, timeout: Optional[float] = None) -> WSMessage:
+        receive_timeout = timeout or self._receive_timeout
+
         while True:
             if self._waiting:
                 raise RuntimeError("Concurrent call to receive() is not allowed")
@@ -294,7 +296,14 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
             try:
                 self._waiting = True
                 try:
-                    async with async_timeout.timeout(timeout or self._receive_timeout):
+                    if receive_timeout:
+                        # Entering the context manager and creating
+                        # Timeout() object can take almost 50% of the
+                        # run time in this loop so we avoid it if
+                        # there is no read timeout.
+                        async with async_timeout.timeout(receive_timeout):
+                            msg = await self._reader.read()
+                    else:
                         msg = await self._reader.read()
                     self._reset_heartbeat()
                 finally:
diff --git a/aiohttp/web_ws.py b/aiohttp/web_ws.py
index ba3332715a6..fe8f537dc76 100644
--- a/aiohttp/web_ws.py
+++ b/aiohttp/web_ws.py
@@ -484,6 +484,7 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
 
         loop = self._loop
         assert loop is not None
+        receive_timeout = timeout or self._receive_timeout
         while True:
             if self._waiting:
                 raise RuntimeError("Concurrent call to receive() is not allowed")
@@ -499,7 +500,14 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
             try:
                 self._waiting = True
                 try:
-                    async with async_timeout.timeout(timeout or self._receive_timeout):
+                    if receive_timeout:
+                        # Entering the context manager and creating
+                        # Timeout() object can take almost 50% of the
+                        # run time in this loop so we avoid it if
+                        # there is no read timeout.
+                        async with async_timeout.timeout(receive_timeout):
+                            msg = await self._reader.read()
+                    else:
                         msg = await self._reader.read()
                     self._reset_heartbeat()
                 finally: