Do not clear the redis-reader's state when we disconnect so that it can finish reading the final message

Author: kristjanvalur

redis/asyncio/connection.py

@@ -350,13 +350,14 @@ async def _readline(self) -> bytes:
 class HiredisParser(BaseParser):
     """Parser class for connections using Hiredis"""
 
-    __slots__ = BaseParser.__slots__ + ("_reader",)
+    __slots__ = BaseParser.__slots__ + ("_reader", "_connected")
 
     def __init__(self, socket_read_size: int):
         if not HIREDIS_AVAILABLE:
             raise RedisError("Hiredis is not available.")
         super().__init__(socket_read_size=socket_read_size)
         self._reader: Optional[hiredis.Reader] = None
+        self._connected: bool = False
 
     def on_connect(self, connection: "Connection"):
         self._stream = connection._reader
@@ -369,13 +370,13 @@ def on_connect(self, connection: "Connection"):
             kwargs["errors"] = connection.encoder.encoding_errors
 
         self._reader = hiredis.Reader(**kwargs)
+        self._connected = False
 
     def on_disconnect(self):
-        self._stream = None
-        self._reader = None
+        self._connected = False
 
     async def can_read_destructive(self):
-        if not self._stream or not self._reader:
+        if not self._connected:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
         if self._reader.gets():
             return True
@@ -397,8 +398,10 @@ async def read_from_socket(self):
     async def read_response(
         self, disable_decoding: bool = False
     ) -> Union[EncodableT, List[EncodableT]]:
-        if not self._stream or not self._reader:
-            self.on_disconnect()
+        # If `on_disconnect()` has been called, prohibit any more reads
+        # even if they could happen because data might be present.
+        # We still allow reads in progress to finish
+        if not self._connected:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
 
         response = self._reader.gets()

tests/test_asyncio/test_connection.py

@@ -150,15 +150,18 @@ async def test_connection_parse_response_resume(r: redis.Redis):
     assert i > 0
 
 
-@pytest.mark.xfail
 @pytest.mark.onlynoncluster
 async def test_connection_hiredis_disconect_race():
     """
     This test reproduces the case in issue #2349
-    where a connection is closed while the parser is reading to feed the internal buffer.
-    The stremam read() will succeed, but when it returns, another task has already called
-    `disconnect()` and is waiting for close to finish.  When it attempts to feed the
-    buffer, it will fail, since the buffer is no longer there.
+    where a connection is closed while the parser is reading to feed the
+    internal buffer.The stremam read() will succeed, but when it returns,
+    another task hasalready called `disconnect()` and is waiting for
+    close to finish.  When it attempts to feed the buffer, it will fail
+    since the buffer is no longer there.
+
+    This test verifies that a read in progress can finish even
+    if the `disconnect()` method is called.
     """
     if not HIREDIS_AVAILABLE:
         pytest.skip("Hiredis not available)")
@@ -194,7 +197,8 @@ async def do_close():
         await conn.disconnect()
 
     async def do_read():
-        await conn.read_response()
+        with pytest.raises(InvalidResponse):
+            await conn.read_response()
 
     reader = AsyncMock()
     writer = AsyncMock()