fix(async): avoid asyncio.timeout on PubSub poll path (#3748)

redis/_parsers/base.py

@@ -72,6 +72,10 @@
 logger = logging.getLogger(__name__)
 
 
+class BufferedResponseIncomplete(Exception):
+    """Raised when a buffered-only read needs more socket data."""
+
+
 class BaseParser(ABC):
     EXCEPTION_CLASSES = {
         "ERR": {
@@ -562,7 +566,7 @@ async def can_read(self) -> bool:
         # parser and fail loudly if the private buffer API changes.
         return bool(self._stream._buffer)
 
-    async def _read(self, length: int) -> bytes:
+    async def _read(self, length: int, *, buffered_only: bool = False) -> bytes:
         """
         Read `length` bytes of data.  These are assumed to be followed
         by a '\r\n' terminator which is subsequently discarded.
@@ -572,6 +576,8 @@ async def _read(self, length: int) -> bytes:
         if len(self._buffer) >= end:
             result = self._buffer[self._pos : end - 2]
         else:
+            if buffered_only:
+                raise BufferedResponseIncomplete()
             tail = self._buffer[self._pos :]
             try:
                 data = await self._stream.readexactly(want - len(tail))
@@ -582,7 +588,7 @@ async def _read(self, length: int) -> bytes:
         self._pos += want
         return result
 
-    async def _readline(self) -> bytes:
+    async def _readline(self, *, buffered_only: bool = False) -> bytes:
         """
         read an unknown number of bytes up to the next '\r\n'
         line separator, which is discarded.
@@ -591,6 +597,8 @@ async def _readline(self) -> bytes:
         if found >= 0:
             result = self._buffer[self._pos : found]
         else:
+            if buffered_only:
+                raise BufferedResponseIncomplete()
             tail = self._buffer[self._pos :]
             data = await self._stream.readline()
             if not data.endswith(b"\r\n"):

redis/_parsers/hiredis.py

@@ -348,7 +348,11 @@ async def read_from_socket(self):
         return True
 
     async def read_response(
-        self, disable_decoding: bool = False, push_request: bool = False
+        self,
+        disable_decoding: bool = False,
+        push_request: bool = False,
+        *,
+        buffered_only: bool = False,
     ) -> Union[EncodableT, List[EncodableT]]:
         # If `on_disconnect()` has been called, prohibit any more reads
         # even if they could happen because data might be present.
@@ -362,6 +366,8 @@ async def read_response(
             response = self._reader.gets()
 
         while response is NOT_ENOUGH_DATA:
+            if buffered_only:
+                return None
             await self.read_from_socket()
             if disable_decoding:
                 response = self._reader.gets(False)
@@ -379,7 +385,9 @@ async def read_response(
             response = await self.handle_push_response(response)
             if not push_request:
                 return await self.read_response(
-                    disable_decoding=disable_decoding, push_request=push_request
+                    disable_decoding=disable_decoding,
+                    push_request=push_request,
+                    buffered_only=buffered_only,
                 )
             else:
                 return response

redis/_parsers/resp2.py

@@ -3,7 +3,7 @@
 from ..exceptions import ConnectionError, InvalidResponse, ResponseError
 from ..typing import EncodableT
 from ..utils import SENTINEL
-from .base import _AsyncRESPBase, _RESPBase
+from .base import BufferedResponseIncomplete, _AsyncRESPBase, _RESPBase
 from .socket import SERVER_CLOSED_CONNECTION_ERROR
 
 
@@ -78,23 +78,35 @@ def _read_response(
 class _AsyncRESP2Parser(_AsyncRESPBase):
     """Async class for the RESP2 protocol"""
 
-    async def read_response(self, disable_decoding: bool = False):
+    async def read_response(
+        self, disable_decoding: bool = False, *, buffered_only: bool = False
+    ):
         if not self._connected:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
         if self._chunks:
             # augment parsing buffer with previously read data
             self._buffer += b"".join(self._chunks)
             self._chunks.clear()
         self._pos = 0
-        response = await self._read_response(disable_decoding=disable_decoding)
+        try:
+            response = await self._read_response(
+                disable_decoding=disable_decoding, buffered_only=buffered_only
+            )
+        except BufferedResponseIncomplete:
+            if buffered_only:
+                return None
+            raise
         # Successfully parsing a response allows us to clear our parsing buffer
         self._clear()
         return response
 
     async def _read_response(
-        self, disable_decoding: bool = False
+        self,
+        disable_decoding: bool = False,
+        *,
+        buffered_only: bool = False,
     ) -> Union[EncodableT, ResponseError, None]:
-        raw = await self._readline()
+        raw = await self._readline(buffered_only=buffered_only)
         response: Any
         byte, response = raw[:1], raw[1:]
 
@@ -122,13 +134,17 @@ async def _read_response(
         elif byte == b"$" and response == b"-1":
             return None
         elif byte == b"$":
-            response = await self._read(int(response))
+            response = await self._read(int(response), buffered_only=buffered_only)
         # multi-bulk response
         elif byte == b"*" and response == b"-1":
             return None
         elif byte == b"*":
             response = [
-                (await self._read_response(disable_decoding))
+                (
+                    await self._read_response(
+                        disable_decoding, buffered_only=buffered_only
+                    )
+                )
                 for _ in range(int(response))  # noqa
             ]
         else:

redis/_parsers/resp3.py

@@ -6,6 +6,7 @@
 from ..utils import SENTINEL
 from .base import (
     AsyncPushNotificationsParser,
+    BufferedResponseIncomplete,
     PushNotificationsParser,
     _AsyncRESPBase,
     _RESPBase,
@@ -175,7 +176,11 @@ async def handle_pubsub_push_response(self, response):
         return response
 
     async def read_response(
-        self, disable_decoding: bool = False, push_request: bool = False
+        self,
+        disable_decoding: bool = False,
+        push_request: bool = False,
+        *,
+        buffered_only: bool = False,
     ):
         if not self._connected:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
@@ -184,19 +189,30 @@ async def read_response(
             self._buffer += b"".join(self._chunks)
             self._chunks.clear()
         self._pos = 0
-        response = await self._read_response(
-            disable_decoding=disable_decoding, push_request=push_request
-        )
+        try:
+            response = await self._read_response(
+                disable_decoding=disable_decoding,
+                push_request=push_request,
+                buffered_only=buffered_only,
+            )
+        except BufferedResponseIncomplete:
+            if buffered_only:
+                return None
+            raise
         # Successfully parsing a response allows us to clear our parsing buffer
         self._clear()
         return response
 
     async def _read_response(
-        self, disable_decoding: bool = False, push_request: bool = False
+        self,
+        disable_decoding: bool = False,
+        push_request: bool = False,
+        *,
+        buffered_only: bool = False,
     ) -> Union[EncodableT, ResponseError, None]:
         if not self._stream or not self.encoder:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
-        raw = await self._readline()
+        raw = await self._readline(buffered_only=buffered_only)
         response: Any
         byte, response = raw[:1], raw[1:]
 
@@ -206,7 +222,7 @@ async def _read_response(
         # server returned an error
         if byte in (b"-", b"!"):
             if byte == b"!":
-                response = await self._read(int(response))
+                response = await self._read(int(response), buffered_only=buffered_only)
             response = response.decode("utf-8", errors="replace")
             error = self.parse_error(response)
             # if the error is a ConnectionError, raise immediately so the user
@@ -236,22 +252,34 @@ async def _read_response(
             return response == b"t"
         # bulk response
         elif byte == b"$":
-            response = await self._read(int(response))
+            response = await self._read(int(response), buffered_only=buffered_only)
         # verbatim string response
         elif byte == b"=":
-            response = (await self._read(int(response)))[4:]
+            response = (await self._read(int(response), buffered_only=buffered_only))[
+                4:
+            ]
         # array response
         elif byte == b"*":
             response = [
-                (await self._read_response(disable_decoding=disable_decoding))
+                (
+                    await self._read_response(
+                        disable_decoding=disable_decoding,
+                        buffered_only=buffered_only,
+                    )
+                )
                 for _ in range(int(response))
             ]
         # set response
         elif byte == b"~":
             # redis can return unhashable types (like dict) in a set,
             # so we always convert to a list, to have predictable return types
             response = [
-                (await self._read_response(disable_decoding=disable_decoding))
+                (
+                    await self._read_response(
+                        disable_decoding=disable_decoding,
+                        buffered_only=buffered_only,
+                    )
+                )
                 for _ in range(int(response))
             ]
         # map response
@@ -261,25 +289,34 @@ async def _read_response(
             # became defined to be left-right in version 3.8
             resp_dict = {}
             for _ in range(int(response)):
-                key = await self._read_response(disable_decoding=disable_decoding)
+                key = await self._read_response(
+                    disable_decoding=disable_decoding,
+                    buffered_only=buffered_only,
+                )
                 resp_dict[key] = await self._read_response(
-                    disable_decoding=disable_decoding, push_request=push_request
+                    disable_decoding=disable_decoding,
+                    push_request=push_request,
+                    buffered_only=buffered_only,
                 )
             response = resp_dict
         # push response
         elif byte == b">":
             response = [
                 (
                     await self._read_response(
-                        disable_decoding=disable_decoding, push_request=push_request
+                        disable_decoding=disable_decoding,
+                        push_request=push_request,
+                        buffered_only=buffered_only,
                     )
                 )
                 for _ in range(int(response))
             ]
             response = await self.handle_push_response(response)
             if not push_request:
                 return await self._read_response(
-                    disable_decoding=disable_decoding, push_request=push_request
+                    disable_decoding=disable_decoding,
+                    push_request=push_request,
+                    buffered_only=buffered_only,
                 )
             else:
                 return response

redis/asyncio/client.py

@@ -1138,6 +1138,15 @@ async def listen(self) -> AsyncIterator[MonitorCommandInfo]:
             yield await self.next_command()
 
 
+async def _await_connection_can_read(conn: "Connection", timeout: float = 0) -> bool:
+    """Call ``can_read`` on a connection, tolerating legacy no-arg overrides."""
+    can_read = conn.can_read
+    try:
+        return await can_read(timeout=timeout)
+    except TypeError:
+        return await can_read()
+
+
 class PubSub:
     """
     PubSub provides publish, subscribe and listen support to Redis channels.
@@ -1445,14 +1454,24 @@ async def parse_response(self, block: bool = True, timeout: float = 0):
         # removed. That swap is a breaking change to the
         # Connection.read_response signature so it must wait for a
         # major release.
-        read_timeout = math.inf if block else timeout
-        response = await self._execute(
-            conn,
-            conn.read_response,
-            timeout=read_timeout,
-            disconnect_on_error=False,
-            push_request=True,
-        )
+        async def try_read():
+            if not block:
+                if not await _await_connection_can_read(conn, timeout):
+                    if timeout == 0:
+                        # Match PubSub.run(): yield so other tasks make progress
+                        # when polling with timeout=0 and no buffered data.
+                        await asyncio.sleep(0)
+                    return None
+                read_timeout = timeout
+            else:
+                read_timeout = math.inf
+            return await conn.read_response(
+                timeout=read_timeout,
+                disconnect_on_error=False,
+                push_request=True,
+            )
+
+        response = await self._execute(conn, try_read)
 
         if conn.health_check_interval and response in self.health_check_response:
             # ignore the health check message as user might not expect it