Skip to content

Commit

Permalink
Optimizing cluster initialization changing the checks for cluster-ena…
Browse files Browse the repository at this point in the history
…bled 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 <willian.moreira@ifood.com.br>
  • Loading branch information
2 people authored and gerzse committed Jul 11, 2024
1 parent 49cbc44 commit 33e76d8
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 16 deletions.
7 changes: 3 additions & 4 deletions redis/asyncio/cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -1226,13 +1226,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.
Expand Down
5 changes: 3 additions & 2 deletions redis/cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -1513,11 +1513,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.
Expand Down
20 changes: 14 additions & 6 deletions tests/test_asyncio/test_cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,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
Expand All @@ -140,9 +142,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":
Expand Down Expand Up @@ -2457,7 +2463,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)
Expand Down Expand Up @@ -2718,10 +2727,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:
Expand Down
16 changes: 12 additions & 4 deletions tests/test_cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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":
Expand Down Expand Up @@ -2651,7 +2656,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)

Expand Down

0 comments on commit 33e76d8

Please sign in to comment.