Updating retry object for existing and new connections

Author: barshaul

redis/asyncio/client.py

@@ -278,6 +278,7 @@ def get_retry(self) -> Optional["Retry"]:
 
     def set_retry(self, retry: "Retry") -> None:
         self.get_connection_kwargs().update({"retry": retry})
+        self.connection_pool.set_retry(retry)
 
     def load_external_module(self, funcname, func):
         """

redis/asyncio/cluster.py

@@ -207,7 +207,6 @@ def from_url(cls, url: str, **kwargs: Any) -> "RedisCluster":
         "reinitialize_steps",
         "response_callbacks",
         "result_callbacks",
-        "retry",
     )
 
     def __init__(
@@ -504,6 +503,8 @@ def set_retry(self, retry: "Retry") -> None:
         self.retry = retry
         for node in self.get_nodes():
             node.connection_kwargs.update({"retry": retry})
+            for conn in node._connections:
+                conn.retry = retry
 
     def set_response_callback(self, command: str, callback: ResponseCallbackT) -> None:
         """Set a custom response callback."""

redis/asyncio/connection.py

@@ -1522,6 +1522,12 @@ async def disconnect(self, inuse_connections: bool = True):
             if exc:
                 raise exc
 
+    def set_retry(self, retry: "Retry") -> None:
+        for conn in self._available_connections:
+            conn.retry = retry
+        for conn in self._in_use_connections:
+            conn.retry = retry
+
 
 class BlockingConnectionPool(ConnectionPool):
     """

redis/client.py

@@ -1050,6 +1050,7 @@ def get_retry(self) -> Optional["Retry"]:
 
     def set_retry(self, retry: "Retry") -> None:
         self.get_connection_kwargs().update({"retry": retry})
+        self.connection_pool.set_retry(retry)
 
     def set_response_callback(self, command, callback):
         """Set a custom Response Callback"""

redis/connection.py

@@ -529,7 +529,7 @@ def __init__(
             # Add TimeoutError to the errors list to retry on
             retry_on_error.append(TimeoutError)
         self.retry_on_error = retry_on_error
-        if retry_on_error:
+        if retry or retry_on_error:
             if retry is None:
                 self.retry = Retry(NoBackoff(), 1)
             else:
@@ -1446,6 +1446,13 @@ def disconnect(self, inuse_connections=True):
             for connection in connections:
                 connection.disconnect()
 
+    def set_retry(self, retry: "Retry") -> None:
+        self.connection_kwargs.update({"retry": retry})
+        for conn in self._available_connections:
+            conn.retry = retry
+        for conn in self._in_use_connections:
+            conn.retry = retry
+
 
 class BlockingConnectionPool(ConnectionPool):
     """

tests/test_asyncio/test_cluster.py

@@ -260,6 +260,8 @@ async def test_cluster_set_get_retry_object(self, request: FixtureRequest):
                 assert n_retry is not None
                 assert n_retry._retries == retry._retries
                 assert isinstance(n_retry._backoff, NoBackoff)
+            rand_cluster_node = r.get_random_node()
+            existing_conn = rand_cluster_node.acquire_connection()
             # Change retry policy
             new_retry = Retry(ExponentialBackoff(), 3)
             r.set_retry(new_retry)
@@ -270,6 +272,9 @@ async def test_cluster_set_get_retry_object(self, request: FixtureRequest):
                 assert n_retry is not None
                 assert n_retry._retries == new_retry._retries
                 assert isinstance(n_retry._backoff, ExponentialBackoff)
+            assert existing_conn.retry._retries == new_retry._retries
+            new_conn = rand_cluster_node.acquire_connection()
+            assert new_conn.retry._retries == new_retry._retries
 
     async def test_cluster_retry_object(self, request: FixtureRequest) -> None:
         url = request.config.getoption("--redis-url")

tests/test_asyncio/test_retry.py

@@ -127,6 +127,10 @@ async def test_get_set_retry_object(self, request):
         assert r.get_retry()._retries == retry._retries
         assert isinstance(r.get_retry()._backoff, NoBackoff)
         new_retry_policy = Retry(ExponentialBackoff(), 3)
+        exiting_conn = await r.connection_pool.get_connection("_")
         r.set_retry(new_retry_policy)
         assert r.get_retry()._retries == new_retry_policy._retries
         assert isinstance(r.get_retry()._backoff, ExponentialBackoff)
+        assert exiting_conn.retry._retries == new_retry_policy._retries
+        new_conn = await r.connection_pool.get_connection("_")
+        assert new_conn.retry._retries == new_retry_policy._retries

tests/test_cluster.py

@@ -749,12 +749,14 @@ def moved_redirect_effect(connection, *args, **options):
 
     def test_cluster_get_set_retry_object(self, request):
         retry = Retry(NoBackoff(), 2)
-        r = _get_client(Redis, request, retry=retry)
+        r = _get_client(RedisCluster, request, retry=retry)
         assert r.get_retry()._retries == retry._retries
         assert isinstance(r.get_retry()._backoff, NoBackoff)
         for node in r.get_nodes():
             assert node.redis_connection.get_retry()._retries == retry._retries
             assert isinstance(node.redis_connection.get_retry()._backoff, NoBackoff)
+        rand_node = r.get_random_node()
+        existing_conn = rand_node.redis_connection.connection_pool.get_connection("_")
         # Change retry policy
         new_retry = Retry(ExponentialBackoff(), 3)
         r.set_retry(new_retry)
@@ -765,6 +767,9 @@ def test_cluster_get_set_retry_object(self, request):
             assert isinstance(
                 node.redis_connection.get_retry()._backoff, ExponentialBackoff
             )
+        assert existing_conn.retry._retries == new_retry._retries
+        new_conn = rand_node.redis_connection.connection_pool.get_connection("_")
+        assert new_conn.retry._retries == new_retry._retries
 
     def test_cluster_retry_object(self, r) -> None:
         # Test default retry

tests/test_retry.py

@@ -207,9 +207,13 @@ def test_client_retry_on_timeout(self, request):
     def test_get_set_retry_object(self, request):
         retry = Retry(NoBackoff(), 2)
         r = _get_client(Redis, request, retry_on_timeout=True, retry=retry)
+        exist_conn = r.connection_pool.get_connection("_")
         assert r.get_retry()._retries == retry._retries
         assert isinstance(r.get_retry()._backoff, NoBackoff)
         new_retry_policy = Retry(ExponentialBackoff(), 3)
         r.set_retry(new_retry_policy)
         assert r.get_retry()._retries == new_retry_policy._retries
         assert isinstance(r.get_retry()._backoff, ExponentialBackoff)
+        assert exist_conn.retry._retries == new_retry_policy._retries
+        new_conn = r.connection_pool.get_connection("_")
+        assert new_conn.retry._retries == new_retry_policy._retries