From ddff7b54b5db4ace94203dc141151b3f08060d2a Mon Sep 17 00:00:00 2001 From: Willian Moreira Date: Tue, 12 Mar 2024 08:14:25 -0300 Subject: [PATCH] Optimizing cluster initialization changing the checks for cluster-enabled flag (#3158) * change if the cluster-mode is enabled by trying run CLUSTER SLOT insted of INFO * fix typo * fixing cluster mode is not enabled on this node tests * remove changes on asyncio * rename mock flag to be more consistent * optimizing async cluster creation using CLUSTER SLOT command instead of INFO command * fixing test. Before INFO and CLUSTER_SLOT was used for performing the connection, now only the CLUSTER_SLOT, so the total commands is minus 1 * remove dot at the end of string * remove unecessary print from test * fix lint problems --------- Co-authored-by: Willian Moreira --- redis/asyncio/cluster.py | 7 +++---- redis/cluster.py | 5 +++-- tests/test_asyncio/test_cluster.py | 20 ++++++++++++++------ tests/test_cluster.py | 16 ++++++++++++---- 4 files changed, 32 insertions(+), 16 deletions(-) diff --git a/redis/asyncio/cluster.py b/redis/asyncio/cluster.py index 4fb2fc4647..11c423b848 100644 --- a/redis/asyncio/cluster.py +++ b/redis/asyncio/cluster.py @@ -1253,13 +1253,12 @@ async def initialize(self) -> None: for startup_node in self.startup_nodes.values(): try: # Make sure cluster mode is enabled on this node - if not (await startup_node.execute_command("INFO")).get( - "cluster_enabled" - ): + try: + cluster_slots = await startup_node.execute_command("CLUSTER SLOTS") + except ResponseError: raise RedisClusterException( "Cluster mode is not enabled on this node" ) - cluster_slots = await startup_node.execute_command("CLUSTER SLOTS") startup_nodes_reachable = True except Exception as e: # Try the next startup node. diff --git a/redis/cluster.py b/redis/cluster.py index ba25b92246..a9213f4235 100644 --- a/redis/cluster.py +++ b/redis/cluster.py @@ -1525,11 +1525,12 @@ def initialize(self): ) self.startup_nodes[startup_node.name].redis_connection = r # Make sure cluster mode is enabled on this node - if bool(r.info().get("cluster_enabled")) is False: + try: + cluster_slots = str_if_bytes(r.execute_command("CLUSTER SLOTS")) + except ResponseError: raise RedisClusterException( "Cluster mode is not enabled on this node" ) - cluster_slots = str_if_bytes(r.execute_command("CLUSTER SLOTS")) startup_nodes_reachable = True except Exception as e: # Try the next startup node. diff --git a/tests/test_asyncio/test_cluster.py b/tests/test_asyncio/test_cluster.py index a57d32f5d2..d7554b12a5 100644 --- a/tests/test_asyncio/test_cluster.py +++ b/tests/test_asyncio/test_cluster.py @@ -127,7 +127,9 @@ async def slowlog(r: RedisCluster) -> None: await r.config_set("slowlog-max-len", old_max_length_value) -async def get_mocked_redis_client(*args, **kwargs) -> RedisCluster: +async def get_mocked_redis_client( + cluster_slots_raise_error=False, *args, **kwargs +) -> RedisCluster: """ Return a stable RedisCluster object that have deterministic nodes and slots setup to remove the problem of different IP addresses @@ -139,9 +141,13 @@ async def get_mocked_redis_client(*args, **kwargs) -> RedisCluster: with mock.patch.object(ClusterNode, "execute_command") as execute_command_mock: async def execute_command(*_args, **_kwargs): + if _args[0] == "CLUSTER SLOTS": - mock_cluster_slots = cluster_slots - return mock_cluster_slots + if cluster_slots_raise_error: + raise ResponseError() + else: + mock_cluster_slots = cluster_slots + return mock_cluster_slots elif _args[0] == "COMMAND": return {"get": [], "set": []} elif _args[0] == "INFO": @@ -2458,7 +2464,10 @@ async def test_init_slots_cache_cluster_mode_disabled(self) -> None: """ with pytest.raises(RedisClusterException) as e: rc = await get_mocked_redis_client( - host=default_host, port=default_port, cluster_enabled=False + cluster_slots_raise_error=True, + host=default_host, + port=default_port, + cluster_enabled=False, ) await rc.aclose() assert "Cluster mode is not enabled on this node" in str(e.value) @@ -2719,10 +2728,9 @@ async def parse_response( async with r.pipeline() as pipe: with pytest.raises(ClusterDownError): await pipe.get(key).execute() - assert ( node.parse_response.await_count - == 4 * r.cluster_error_retry_attempts - 3 + == 3 * r.cluster_error_retry_attempts - 2 ) async def test_connection_error_not_raised(self, r: RedisCluster) -> None: diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 8a44d45ea3..1f505b816d 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -151,7 +151,9 @@ def cleanup(): r.config_set("slowlog-max-len", 128) -def get_mocked_redis_client(func=None, *args, **kwargs): +def get_mocked_redis_client( + func=None, cluster_slots_raise_error=False, *args, **kwargs +): """ Return a stable RedisCluster object that have deterministic nodes and slots setup to remove the problem of different IP addresses @@ -164,8 +166,11 @@ def get_mocked_redis_client(func=None, *args, **kwargs): def execute_command(*_args, **_kwargs): if _args[0] == "CLUSTER SLOTS": - mock_cluster_slots = cluster_slots - return mock_cluster_slots + if cluster_slots_raise_error: + raise ResponseError() + else: + mock_cluster_slots = cluster_slots + return mock_cluster_slots elif _args[0] == "COMMAND": return {"get": [], "set": []} elif _args[0] == "INFO": @@ -2654,7 +2659,10 @@ def test_init_slots_cache_cluster_mode_disabled(self): """ with pytest.raises(RedisClusterException) as e: get_mocked_redis_client( - host=default_host, port=default_port, cluster_enabled=False + cluster_slots_raise_error=True, + host=default_host, + port=default_port, + cluster_enabled=False, ) assert "Cluster mode is not enabled on this node" in str(e.value)