async_cluster: add concurrent test & use read_response/_update_moved_slots without lock

Author: utkarshgupta137

redis/asyncio/cluster.py

@@ -2,7 +2,6 @@
 import collections
 import random
 import socket
-import threading
 import warnings
 from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union
 
@@ -339,7 +338,7 @@ async def on_connect(self, connection: Connection) -> None:
             # regardless of the server type. If this is a primary connection,
             # READONLY would not affect executing write commands.
             await connection.send_command("READONLY")
-            if str_if_bytes(await connection.read_response()) != "OK":
+            if str_if_bytes(await connection.read_response_without_lock()) != "OK":
                 raise ConnectionError("READONLY command failed")
 
     def get_node(
@@ -789,10 +788,10 @@ async def execute_command(self, *args, **kwargs) -> Any:
         connection = None
         if self._free:
             for _ in range(len(self._free)):
-                if self._free[0].is_connected:
-                    connection = self._free.popleft()
+                connection = self._free.popleft()
+                if connection.is_connected:
                     break
-                self._free.rotate(-1)
+                self._free.append(connection)
             else:
                 connection = self._free.popleft()
         else:
@@ -807,9 +806,11 @@ async def execute_command(self, *args, **kwargs) -> Any:
         await connection.send_packed_command(command, False)
         try:
             if NEVER_DECODE in kwargs:
-                response = await connection.read_response(disable_decoding=True)
+                response = await connection.read_response_without_lock(
+                    disable_decoding=True
+                )
             else:
-                response = await connection.read_response()
+                response = await connection.read_response_without_lock()
         except ResponseError:
             if EMPTY_RESPONSE in kwargs:
                 return kwargs[EMPTY_RESPONSE]
@@ -827,7 +828,6 @@ async def execute_command(self, *args, **kwargs) -> Any:
 
 class NodesManager:
     __slots__ = (
-        "_lock",
         "_moved_exception",
         "_require_full_coverage",
         "connection_kwargs",
@@ -852,7 +852,6 @@ def __init__(
         self._moved_exception = None
         self.connection_kwargs = kwargs
         self.read_load_balancer = LoadBalancer()
-        self._lock = threading.Lock()
 
     def get_node(
         self,
@@ -936,9 +935,7 @@ def get_node_from_slot(
         self, slot: int, read_from_replicas: bool = False
     ) -> "ClusterNode":
         if self._moved_exception:
-            with self._lock:
-                if self._moved_exception:
-                    self._update_moved_slots()
+            self._update_moved_slots()
 
         try:
             if read_from_replicas:

redis/asyncio/connection.py

@@ -939,6 +939,41 @@ async def read_response(self, disable_decoding: bool = False):
             raise response from None
         return response
 
+    async def read_response_without_lock(self, disable_decoding: bool = False):
+        """Read the response from a previously sent command"""
+        try:
+            if self.socket_timeout:
+                async with async_timeout.timeout(self.socket_timeout):
+                    response = await self._parser.read_response(
+                        disable_decoding=disable_decoding
+                    )
+            else:
+                response = await self._parser.read_response(
+                    disable_decoding=disable_decoding
+                )
+        except asyncio.TimeoutError:
+            await self.disconnect()
+            raise TimeoutError(f"Timeout reading from {self.host}:{self.port}")
+        except OSError as e:
+            await self.disconnect()
+            raise ConnectionError(
+                f"Error while reading from {self.host}:{self.port} : {e.args}"
+            )
+        except BaseException:
+            await self.disconnect()
+            raise
+
+        if self.health_check_interval:
+            if sys.version_info[0:2] == (3, 6):
+                func = asyncio.get_event_loop
+            else:
+                func = asyncio.get_running_loop
+            self.next_health_check = func().time() + self.health_check_interval
+
+        if isinstance(response, ResponseError):
+            raise response from None
+        return response
+
     def pack_command(self, *args: EncodableT) -> List[bytes]:
         """Pack a series of arguments into the Redis protocol"""
         output = []

tests/test_asyncio/test_cluster.py

@@ -134,7 +134,7 @@ def mock_node_resp(
 ) -> ClusterNode:
     connection = mock.AsyncMock()
     connection.is_connected = True
-    connection.read_response.return_value = response
+    connection.read_response_without_lock.return_value = response
     while node._free:
         node._free.pop()
     node._free.append(connection)
@@ -288,10 +288,10 @@ async def test_execute_command_node_flag_primaries(self, r: RedisCluster) -> Non
         assert await r.ping(target_nodes=RedisCluster.PRIMARIES) is True
         for primary in primaries:
             conn = primary._free.pop()
-            assert conn.read_response.called is True
+            assert conn.read_response_without_lock.called is True
         for replica in replicas:
             conn = replica._free.pop()
-            assert conn.read_response.called is not True
+            assert conn.read_response_without_lock.called is not True
 
     async def test_execute_command_node_flag_replicas(self, r: RedisCluster) -> None:
         """
@@ -305,10 +305,10 @@ async def test_execute_command_node_flag_replicas(self, r: RedisCluster) -> None
         assert await r.ping(target_nodes=RedisCluster.REPLICAS) is True
         for replica in replicas:
             conn = replica._free.pop()
-            assert conn.read_response.called is True
+            assert conn.read_response_without_lock.called is True
         for primary in primaries:
             conn = primary._free.pop()
-            assert conn.read_response.called is not True
+            assert conn.read_response_without_lock.called is not True
 
         await r.close()
 
@@ -320,7 +320,7 @@ async def test_execute_command_node_flag_all_nodes(self, r: RedisCluster) -> Non
         assert await r.ping(target_nodes=RedisCluster.ALL_NODES) is True
         for node in r.get_nodes():
             conn = node._free.pop()
-            assert conn.read_response.called is True
+            assert conn.read_response_without_lock.called is True
 
     async def test_execute_command_node_flag_random(self, r: RedisCluster) -> None:
         """
@@ -331,7 +331,7 @@ async def test_execute_command_node_flag_random(self, r: RedisCluster) -> None:
         called_count = 0
         for node in r.get_nodes():
             conn = node._free.pop()
-            if conn.read_response.called is True:
+            if conn.read_response_without_lock.called is True:
                 called_count += 1
         assert called_count == 1
 
@@ -344,7 +344,7 @@ async def test_execute_command_default_node(self, r: RedisCluster) -> None:
         mock_node_resp(def_node, "PONG")
         assert await r.ping() is True
         conn = def_node._free.pop()
-        assert conn.read_response.called
+        assert conn.read_response_without_lock.called
 
     async def test_ask_redirection(self, r: RedisCluster) -> None:
         """
@@ -488,7 +488,7 @@ async def test_reading_from_replicas_in_round_robin(self) -> None:
         with mock.patch.multiple(
             Connection,
             send_command=mock.DEFAULT,
-            read_response=mock.DEFAULT,
+            read_response_without_lock=mock.DEFAULT,
             _connect=mock.DEFAULT,
             can_read=mock.DEFAULT,
             on_connect=mock.DEFAULT,
@@ -520,7 +520,7 @@ def execute_command_mock_third(self, *args, **options):
                 # so we'll mock some of the Connection's functions to allow it
                 execute_command.side_effect = execute_command_mock_first
                 mocks["send_command"].return_value = True
-                mocks["read_response"].return_value = "OK"
+                mocks["read_response_without_lock"].return_value = "OK"
                 mocks["_connect"].return_value = True
                 mocks["can_read"].return_value = False
                 mocks["on_connect"].return_value = True
@@ -682,6 +682,14 @@ async def test_not_require_full_coverage_cluster_down_error(
                 else:
                     raise e
 
+    async def test_can_run_concurrent_commands(self, r: RedisCluster) -> None:
+        assert await r.ping(target_nodes=RedisCluster.ALL_NODES) is True
+        assert all(
+            await asyncio.gather(
+                *(r.ping(target_nodes=RedisCluster.ALL_NODES) for _ in range(100))
+            )
+        )
+
 
 @pytest.mark.onlycluster
 class TestClusterRedisCommands:
@@ -792,7 +800,7 @@ async def test_cluster_addslots(self, r: RedisCluster) -> None:
 
     @skip_if_server_version_lt("7.0.0")
     @skip_if_redis_enterprise()
-    async def test_cluster_addslotsrange(self, r):
+    async def test_cluster_addslotsrange(self, r: RedisCluster):
         node = r.get_random_node()
         mock_node_resp(node, "OK")
         assert await r.cluster_addslotsrange(node, 1, 5)
@@ -820,14 +828,14 @@ async def test_cluster_delslots(self) -> None:
         node0 = r.get_node(default_host, 7000)
         node1 = r.get_node(default_host, 7001)
         assert await r.cluster_delslots(0, 8192) == [True, True]
-        assert node0._free.pop().read_response.called
-        assert node1._free.pop().read_response.called
+        assert node0._free.pop().read_response_without_lock.called
+        assert node1._free.pop().read_response_without_lock.called
 
         await r.close()
 
     @skip_if_server_version_lt("7.0.0")
     @skip_if_redis_enterprise()
-    async def test_cluster_delslotsrange(self, r):
+    async def test_cluster_delslotsrange(self, r: RedisCluster):
         node = r.get_random_node()
         mock_node_resp(node, "OK")
         await r.cluster_addslots(node, 1, 2, 3, 4, 5)
@@ -990,7 +998,7 @@ async def test_cluster_setslot_stable(self, r: RedisCluster) -> None:
         node = r.nodes_manager.get_node_from_slot(12182)
         mock_node_resp(node, "OK")
         assert await r.cluster_setslot_stable(12182) is True
-        assert node._free.pop().read_response.called
+        assert node._free.pop().read_response_without_lock.called
 
     @skip_if_redis_enterprise()
     async def test_cluster_replicas(self, r: RedisCluster) -> None:
@@ -1014,7 +1022,7 @@ async def test_cluster_replicas(self, r: RedisCluster) -> None:
         )
 
     @skip_if_server_version_lt("7.0.0")
-    async def test_cluster_links(self, r):
+    async def test_cluster_links(self, r: RedisCluster):
         node = r.get_random_node()
         res = await r.cluster_links(node)
         links_to = sum(x.count("to") for x in res)
@@ -1032,7 +1040,7 @@ async def test_readonly(self) -> None:
         for res in all_replicas_results.values():
             assert res is True
         for replica in r.get_replicas():
-            assert replica._free.pop().read_response.called
+            assert replica._free.pop().read_response_without_lock.called
 
         await r.close()
 
@@ -1045,7 +1053,7 @@ async def test_readwrite(self) -> None:
         for res in all_replicas_results.values():
             assert res is True
         for replica in r.get_replicas():
-            assert replica._free.pop().read_response.called
+            assert replica._free.pop().read_response_without_lock.called
 
         await r.close()