fix(async): apply socket_timeout per read in async parsers

redis/_parsers/base.py

@@ -1,8 +1,14 @@
 import logging
+import sys
 from abc import ABC, abstractmethod
 from asyncio import IncompleteReadError, StreamReader
 from typing import Awaitable, Callable, List, Optional, Protocol, Union
 
+if sys.version_info >= (3, 11, 3):
+    from asyncio import timeout as async_timeout
+else:
+    from async_timeout import timeout as async_timeout
+
 from redis.maint_notifications import (
     MaintenanceNotification,
     NodeFailedOverNotification,
@@ -13,7 +19,6 @@
     OSSNodeMigratedNotification,
     OSSNodeMigratingNotification,
 )
-from redis.utils import deprecated_function, safe_str
 
 from ..exceptions import (
     AskError,
@@ -36,6 +41,7 @@
     TryAgainError,
 )
 from ..typing import EncodableT
+from ..utils import SENTINEL, deprecated_function, safe_str
 from .encoders import Encoder
 from .socket import SERVER_CLOSED_CONNECTION_ERROR, SocketBuffer
 
@@ -183,7 +189,10 @@ async def can_read(self) -> bool:
         pass
 
     async def read_response(
-        self, disable_decoding: bool = False
+        self,
+        disable_decoding: bool = False,
+        push_request: bool = False,
+        timeout: Union[float, object] = SENTINEL,
     ) -> Union[EncodableT, ResponseError, None, List[EncodableT]]:
         raise NotImplementedError()
 
@@ -509,7 +518,9 @@ def __init__(self, socket_read_size: int):
         self._pos = 0
 
     def _clear(self):
-        self._buffer = b""
+        """Clear parsed data but preserve unconsumed pipelined bytes."""
+        self._buffer = self._buffer[self._pos :]
+        self._pos = 0
         self._chunks.clear()
 
     def on_connect(self, connection):
@@ -518,7 +529,10 @@ def on_connect(self, connection):
         if self._stream is None:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
         self.encoder = connection.encoder
-        self._clear()
+        # Full reset on (re)connect — discard stale data from old connection
+        self._buffer = b""
+        self._pos = 0
+        self._chunks.clear()
         self._connected = True
 
     def on_disconnect(self):
@@ -545,40 +559,65 @@ async def can_read(self) -> bool:
         # parser and fail loudly if the private buffer API changes.
         return bool(self._stream._buffer) or self._stream.at_eof()
 
-    async def _read(self, length: int) -> bytes:
+    async def _read_from_stream(
+        self, timeout: Union[float, object] = SENTINEL, max_bytes: int = 0
+    ) -> bytes:
+        """
+        Read the next chunk from the underlying stream with an optional
+        per-read timeout.  This mirrors the sync client's per-recv timeout
+        semantics: each individual socket read gets its own timeout window.
+
+        ``max_bytes`` limits how many bytes may be returned. When 0, the
+        parser's ``_read_size`` is used.
+        """
+        stream = self._stream
+        if stream is None:
+            raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
+        size = max_bytes if max_bytes > 0 else self._read_size
+        if timeout is not SENTINEL and isinstance(timeout, (int, float)):
+            async with async_timeout(timeout):
+                return await stream.read(size)
+        return await stream.read(size)
+
+    async def _read(
+        self, length: int, timeout: Union[float, object] = SENTINEL
+    ) -> bytes:
         """
         Read `length` bytes of data.  These are assumed to be followed
         by a '\r\n' terminator which is subsequently discarded.
         """
         want = length + 2
         end = self._pos + want
-        if len(self._buffer) >= end:
-            result = self._buffer[self._pos : end - 2]
-        else:
-            tail = self._buffer[self._pos :]
+        while len(self._buffer) < end:
+            need = end - len(self._buffer)
             try:
-                data = await self._stream.readexactly(want - len(tail))
+                chunk = await self._read_from_stream(
+                    timeout=timeout, max_bytes=min(need, self._read_size)
+                )
             except IncompleteReadError as error:
                 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from error
-            result = (tail + data)[:-2]
-            self._chunks.append(data)
-        self._pos += want
+            if not chunk:
+                raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
+            self._buffer += chunk
+        result = self._buffer[self._pos : end - 2]
+        self._pos = end
         return result
 
-    async def _readline(self) -> bytes:
+    async def _readline(self, timeout: Union[float, object] = SENTINEL) -> bytes:
         """
         read an unknown number of bytes up to the next '\r\n'
         line separator, which is discarded.
         """
-        found = self._buffer.find(b"\r\n", self._pos)
-        if found >= 0:
-            result = self._buffer[self._pos : found]
-        else:
-            tail = self._buffer[self._pos :]
-            data = await self._stream.readline()
-            if not data.endswith(b"\r\n"):
+        while True:
+            found = self._buffer.find(b"\r\n", self._pos)
+            if found >= 0:
+                result = self._buffer[self._pos : found]
+                self._pos = found + 2
+                return result
+            try:
+                chunk = await self._read_from_stream(timeout=timeout)
+            except IncompleteReadError as error:
+                raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from error
+            if not chunk:
                 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
-            result = (tail + data)[:-2]
-            self._chunks.append(data)
-        self._pos += len(result) + 2
-        return result
+            self._buffer += chunk

redis/_parsers/hiredis.py

@@ -1,9 +1,15 @@
 import select
 import selectors
 import socket
+import sys
 from logging import getLogger
 from typing import Callable, List, Optional, TypedDict, Union
 
+if sys.version_info >= (3, 11, 3):
+    from asyncio import timeout as async_timeout
+else:
+    from async_timeout import timeout as async_timeout
+
 from ..exceptions import ConnectionError, InvalidResponse, RedisError, TimeoutError
 from ..typing import EncodableT
 from ..utils import HIREDIS_AVAILABLE, SENTINEL, deprecated_function
@@ -265,8 +271,12 @@ async def can_read(self) -> bool:
         # with a real StreamReader guard this private buffer API in CI.
         return bool(self._stream._buffer)
 
-    async def read_from_socket(self):
-        buffer = await self._stream.read(self._read_size)
+    async def read_from_socket(self, timeout: Union[float, object] = SENTINEL):
+        if timeout is not SENTINEL:
+            async with async_timeout(timeout):
+                buffer = await self._stream.read(self._read_size)
+        else:
+            buffer = await self._stream.read(self._read_size)
         if not buffer or not isinstance(buffer, bytes):
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
         self._reader.feed(buffer)
@@ -275,7 +285,10 @@ 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,
+        timeout: Union[float, object] = SENTINEL,
     ) -> Union[EncodableT, List[EncodableT]]:
         # If `on_disconnect()` has been called, prohibit any more reads
         # even if they could happen because data might be present.
@@ -289,7 +302,7 @@ async def read_response(
             response = self._reader.gets()
 
         while response is NOT_ENOUGH_DATA:
-            await self.read_from_socket()
+            await self.read_from_socket(timeout=timeout)
             if disable_decoding:
                 response = self._reader.gets(False)
             else:
@@ -306,7 +319,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,
+                    timeout=timeout,
                 )
             else:
                 return response

redis/_parsers/resp2.py

@@ -78,23 +78,27 @@ 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, timeout: Union[float, object] = SENTINEL
+    ):
         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)
+        response = await self._read_response(
+            disable_decoding=disable_decoding, timeout=timeout
+        )
         # 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, timeout: Union[float, object] = SENTINEL
     ) -> Union[EncodableT, ResponseError, None]:
-        raw = await self._readline()
+        raw = await self._readline(timeout=timeout)
         response: Any
         byte, response = raw[:1], raw[1:]
 
@@ -122,13 +126,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), timeout=timeout)
         # 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=disable_decoding, timeout=timeout
+                    )
+                )
                 for _ in range(int(response))  # noqa
             ]
         else:

redis/_parsers/resp3.py

@@ -153,6 +153,7 @@ def _read_response(
             return self._read_response(
                 disable_decoding=disable_decoding,
                 push_request=push_request,
+                timeout=timeout,
             )
         else:
             raise InvalidResponse(f"Protocol Error: {raw!r}")
@@ -175,26 +176,34 @@ 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,
+        timeout: Union[float, object] = SENTINEL,
     ):
         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, push_request=push_request
+            disable_decoding=disable_decoding,
+            push_request=push_request,
+            timeout=timeout,
         )
         # 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,
+        timeout: Union[float, object] = SENTINEL,
     ) -> 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(timeout=timeout)
         response: Any
         byte, response = raw[:1], raw[1:]
 
@@ -204,7 +213,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), timeout=timeout)
             response = response.decode("utf-8", errors="replace")
             error = self.parse_error(response)
             # if the error is a ConnectionError, raise immediately so the user
@@ -234,22 +243,30 @@ 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), timeout=timeout)
         # verbatim string response
         elif byte == b"=":
-            response = (await self._read(int(response)))[4:]
+            response = (await self._read(int(response), timeout=timeout))[4:]
         # array response
         elif byte == b"*":
             response = [
-                (await self._read_response(disable_decoding=disable_decoding))
+                (
+                    await self._read_response(
+                        disable_decoding=disable_decoding, timeout=timeout
+                    )
+                )
                 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, timeout=timeout
+                    )
+                )
                 for _ in range(int(response))
             ]
         # map response
@@ -259,25 +276,33 @@ 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, timeout=timeout
+                )
                 resp_dict[key] = await self._read_response(
-                    disable_decoding=disable_decoding, push_request=push_request
+                    disable_decoding=disable_decoding,
+                    push_request=push_request,
+                    timeout=timeout,
                 )
             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,
+                        timeout=timeout,
                     )
                 )
                 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,
+                    timeout=timeout,
                 )
             else:
                 return response