Add regression tests and fixes for issue redis#1128

redis/asyncio/client.py

@@ -477,8 +477,17 @@ async def _send_command_parse_response(self, conn, command_name, *args, **option
  """
  Send a command and parse the response
  """
- await conn.send_command(*args)
- return await self.parse_response(conn, command_name, **options)
+ try:
+ await conn.send_command(*args)
+ return await self.parse_response(conn, command_name, **options)
+ except BaseException:
+ # An exception while reading or writing leaves the connection in
+ # an unknown state. There may be un-written or un-read data
+ # so we cannot re-use it for a subsequent command/response transaction.
+ # This may be a TimeoutError or any other error not handled by the
+ # Connection object itself.
+ await conn.disconnect(nowait=True)
+ raise
 
  async def _disconnect_raise(self, conn: Connection, error: Exception):
  """

redis/client.py

@@ -1231,8 +1231,18 @@ def _send_command_parse_response(self, conn, command_name, *args, **options):
  """
  Send a command and parse the response
  """
- conn.send_command(*args)
- return self.parse_response(conn, command_name, **options)
+ try:
+ conn.send_command(*args)
+ return self.parse_response(conn, command_name, **options)
+ except BaseException:
+ # An exception while reading or writing leaves the connection in
+ # an unknown state. There may be un-written or un-read data
+ # so we cannot re-use it for a subsequent command/response transaction.
+ # This can be any error not handled by the Connection itself, such
+ # as BaseExceptions may have been used to interrupt IO, when using
+ # gevent.
+ conn.disconnect()
+ raise
 
  def _disconnect_raise(self, conn, error):
  """

tests/test_asyncio/test_commands.py

@@ -1,10 +1,12 @@
 """
 Tests async overrides of commands from their mixins
 """
+import asyncio
 import binascii
 import datetime
 import re
 from string import ascii_letters
+import async_timeout
 
 import pytest
 import pytest_asyncio
@@ -2999,6 +3001,36 @@ async def test_module_list(self, r: redis.Redis):
  for x in await r.module_list():
  assert isinstance(x, dict)
 
+ @pytest.mark.onlynoncluster
+ async def test_interrupted_command(self, r: redis.Redis):
+ """
+ Regression test for issue #1128: An Un-handled BaseException
+ will leave the socket with un-read response to a previous
+ command.
+ """
+ ready = asyncio.Event()
+
+ async def helper():
+ with pytest.raises(asyncio.CancelledError):
+ # blocking pop
+ ready.set()
+ await r.brpop(["nonexist"])
+ # if all is well, we can continue. The following should not hang.
+ await r.set("status", "down")
+ return "done"
+
+ task = asyncio.create_task(helper())
+ await ready.wait()
+ await asyncio.sleep(0.01)
+ # the task is now sleeping, lets send it an exception
+ task.cancel()
+ try:
+ async with async_timeout.timeout(0.1):
+ assert await task == "done"
+ except asyncio.TimeoutError:
+ task.cancel()
+ await task
+
 
 @pytest.mark.onlynoncluster
 class TestBinarySave:

tests/test_commands.py

@@ -1,9 +1,13 @@
 import binascii
 import datetime
 import re
+import threading
 import time
 from string import ascii_letters
+from asyncio import CancelledError
 from unittest import mock
+from unittest.mock import patch
+import socket
 
 import pytest
 
@@ -4712,6 +4716,34 @@ def test_psync(self, r):
  res = r2.psync(r2.client_id(), 1)
  assert b"FULLRESYNC" in res
 
+ @pytest.mark.onlynoncluster
+ def test_interrupted_command(self, r: redis.Redis):
+ """
+ Regression test for issue #1128: An Un-handled BaseException
+ will leave the socket with un-read response to a previous
+ command.
+ """
+ print("start")
+
+ def helper():
+ try:
+ # blocking pop
+ with patch.object(
+ socket.socket, "recv_into", side_effect=CancelledError
+ ) as mock_recv:
+ r.brpop(["nonexist"])
+ except CancelledError:
+ print("canc")
+ pass # we got some BaseException.
+ # if all is well, we can continue.
+ r.set("status", "down") # should not hang
+ return "done"
+
+ thread = threading.Thread(target=helper)
+ thread.start()
+ thread.join(0.1)
+ assert not thread.is_alive()
+
 
 @pytest.mark.onlynoncluster
 class TestBinarySave: