Use retry mechanism in async version of Connection objects

Author: szumka

redis/asyncio/connection.py

@@ -637,6 +637,8 @@ def __init__(
  retry_on_error = []
  if retry_on_timeout:
  retry_on_error.append(TimeoutError)
+ retry_on_error.append(socket.timeout)
+ retry_on_error.append(asyncio.TimeoutError)
  self.retry_on_error = retry_on_error
  if retry_on_error:
  if not retry:
@@ -706,7 +708,9 @@ async def connect(self):
  if self.is_connected:
  return
  try:
- await self._connect()
+ await self.retry.call_with_retry(
+ lambda: self._connect(), lambda error: self.disconnect(error)
+ )
  except asyncio.CancelledError:
  raise
  except (socket.timeout, asyncio.TimeoutError):
@@ -816,7 +820,7 @@ async def on_connect(self) -> None:
  if str_if_bytes(await self.read_response()) != "OK":
  raise ConnectionError("Invalid Database")
 
- async def disconnect(self) -> None:
+ async def disconnect(self, *args) -> None:
  """Disconnects from the Redis server"""
  try:
  async with async_timeout.timeout(self.socket_connect_timeout):

redis/asyncio/sentinel.py

@@ -44,7 +44,7 @@ async def connect_to(self, address):
  if str_if_bytes(await self.read_response()) != "PONG":
  raise ConnectionError("PING failed")
 
- async def connect(self):
+ async def _connect_retry(self):
  if self._reader:
  return # already connected
  if self.connection_pool.is_master:
@@ -57,6 +57,12 @@ async def connect(self):
  continue
  raise SlaveNotFoundError # Never be here
 
+ async def connect(self):
+ return await self.retry.call_with_retry(
+ self._connect_retry,
+ lambda error: asyncio.sleep(0),
+ )
+
  async def read_response(self, disable_decoding: bool = False):
  try:
  return await super().read_response(disable_decoding=disable_decoding)

tests/test_asyncio/test_connection.py

@@ -1,10 +1,18 @@
 import asyncio
+import socket
 import types
+from unittest.mock import patch
 
 import pytest
 
-from redis.asyncio.connection import PythonParser, UnixDomainSocketConnection
-from redis.exceptions import InvalidResponse
+from redis.asyncio.connection import (
+ Connection,
+ PythonParser,
+ UnixDomainSocketConnection,
+)
+from redis.asyncio.retry import Retry
+from redis.backoff import NoBackoff
+from redis.exceptions import ConnectionError, InvalidResponse, TimeoutError
 from redis.utils import HIREDIS_AVAILABLE
 from tests.conftest import skip_if_server_version_lt
 
@@ -60,3 +68,44 @@ async def test_socket_param_regression(r):
 async def test_can_run_concurrent_commands(r):
  assert await r.ping() is True
  assert all(await asyncio.gather(*(r.ping() for _ in range(10))))
+
+
+async def test_connect_retry_on_timeout_error():
+ """Test that the _connect function is retried in case of a timeout"""
+ conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 3))
+ origin_connect = conn._connect
+ conn._connect = mock.AsyncMock()
+
+ async def mock_connect():
+ # connect only on the last retry
+ if conn._connect.call_count <= 2:
+ raise socket.timeout
+ else:
+ return await origin_connect()
+
+ conn._connect.side_effect = mock_connect
+ await conn.connect()
+ assert conn._connect.call_count == 3
+
+
+async def test_connect_without_retry_on_os_error():
+ """Test that the _connect function is not being retried in case of a OSError"""
+ with patch.object(Connection, "_connect") as _connect:
+ _connect.side_effect = OSError("")
+ conn = Connection(retry_on_timeout=True, retry=Retry(NoBackoff(), 2))
+ with pytest.raises(ConnectionError):
+ await conn.connect()
+ assert _connect.call_count == 1
+
+
+async def test_connect_timeout_error_without_retry():
+ """Test that the _connect function is not being retried if retry_on_timeout is
+ set to False"""
+ conn = Connection(retry_on_timeout=False)
+ conn._connect = mock.AsyncMock()
+ conn._connect.side_effect = socket.timeout
+
+ with pytest.raises(TimeoutError) as e:
+ await conn.connect()
+ assert conn._connect.call_count == 1
+ assert str(e.value) == "Timeout connecting to server"

tests/test_asyncio/test_sentinel_managed_connection.py

@@ -0,0 +1,37 @@
+import socket
+
+import pytest
+
+from redis.asyncio.retry import Retry
+from redis.asyncio.sentinel import SentinelManagedConnection
+from redis.backoff import NoBackoff
+
+from .compat import mock
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_connect_retry_on_timeout_error():
+ """Test that the _connect function is retried in case of a timeout"""
+ connection_pool = mock.AsyncMock()
+ connection_pool.get_master_address = mock.AsyncMock(
+ return_value=("localhost", 6379)
+ )
+ conn = SentinelManagedConnection(
+ retry_on_timeout=True,
+ retry=Retry(NoBackoff(), 3),
+ connection_pool=connection_pool,
+ )
+ origin_connect = conn._connect
+ conn._connect = mock.AsyncMock()
+
+ async def mock_connect():
+ # connect only on the last retry
+ if conn._connect.call_count <= 2:
+ raise socket.timeout
+ else:
+ return await origin_connect()
+
+ conn._connect.side_effect = mock_connect
+ await conn.connect()
+ assert conn._connect.call_count == 3