async_cluster: optimisations (#2205)

- return true from execute_pipeline if there are any errors
- use todo list to speedup retries
- store initialisation node in CommandsParser object
- add sync context manager for pipeline
- use if/else instead of try/except
- make command a function argument in _determine_nodes & _determine_slot

- add async cluster pipeline benchmark script

Author: utkarshgupta137

benchmarks/cluster_async.py

@@ -249,8 +249,8 @@ async def main(loop, gather=None):
     port = 16379
     password = None
 
-    count = 1000
-    size = 16
+    count = 10000
+    size = 256
 
     asyncio.run(main("asyncio"))
     asyncio.run(main("asyncio", gather=False))

benchmarks/cluster_async_pipeline.py

@@ -0,0 +1,107 @@
+import asyncio
+import functools
+import time
+
+import aioredis_cluster
+import aredis
+import uvloop
+
+import redis.asyncio as redispy
+
+
+def timer(func):
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        tic = time.perf_counter()
+        await func(*args, **kwargs)
+        toc = time.perf_counter()
+        return f"{toc - tic:.4f}"
+
+    return wrapper
+
+
+@timer
+async def warmup(client):
+    await asyncio.gather(
+        *(asyncio.create_task(client.exists(f"bench:warmup_{i}")) for i in range(100))
+    )
+
+
+@timer
+async def run(client):
+    data_str = "a" * size
+    data_int = int("1" * size)
+
+    for i in range(count):
+        with client.pipeline() as pipe:
+            await (
+                pipe.set(f"bench:str_{i}", data_str)
+                .set(f"bench:int_{i}", data_int)
+                .get(f"bench:str_{i}")
+                .get(f"bench:int_{i}")
+                .hset("bench:hset", str(i), data_str)
+                .hget("bench:hset", str(i))
+                .incr("bench:incr")
+                .lpush("bench:lpush", data_int)
+                .lrange("bench:lpush", 0, 300)
+                .lpop("bench:lpush")
+                .execute()
+            )
+
+
+async def main(loop):
+    arc = aredis.StrictRedisCluster(
+        host=host,
+        port=port,
+        password=password,
+        max_connections=2**31,
+        max_connections_per_node=2**31,
+        readonly=False,
+        reinitialize_steps=count,
+        skip_full_coverage_check=True,
+        decode_responses=False,
+        max_idle_time=count,
+        idle_check_interval=count,
+    )
+    print(f"{loop} {await warmup(arc)} aredis")
+    print(await run(arc))
+    arc.connection_pool.disconnect()
+
+    aiorc = await aioredis_cluster.create_redis_cluster(
+        [(host, port)],
+        password=password,
+        state_reload_interval=count,
+        idle_connection_timeout=count,
+        pool_maxsize=2**31,
+    )
+    print(f"{loop} {await warmup(aiorc)} aioredis-cluster")
+    print(await run(aiorc))
+    aiorc.close()
+    await aiorc.wait_closed()
+
+    async with redispy.RedisCluster(
+        host=host,
+        port=port,
+        password=password,
+        reinitialize_steps=count,
+        read_from_replicas=False,
+        decode_responses=False,
+        max_connections=2**31,
+    ) as rca:
+        print(f"{loop} {await warmup(rca)} redispy")
+        print(await run(rca))
+
+
+if __name__ == "__main__":
+    host = "localhost"
+    port = 16379
+    password = None
+
+    count = 10000
+    size = 256
+
+    asyncio.run(main("asyncio"))
+
+    uvloop.install()
+
+    asyncio.run(main("uvloop"))

redis/asyncio/cluster.py

@@ -468,9 +468,8 @@ def keyslot(self, key: EncodableT) -> int:
         return key_slot(k)
 
     async def _determine_nodes(
-        self, *args: Any, node_flag: Optional[str] = None
+        self, command: str, *args: Any, node_flag: Optional[str] = None
     ) -> List["ClusterNode"]:
-        command = args[0]
         if not node_flag:
             # get the nodes group for this command if it was predefined
             node_flag = self.command_flags.get(command)
@@ -495,16 +494,15 @@ async def _determine_nodes(
         # get the node that holds the key's slot
         return [
             self.nodes_manager.get_node_from_slot(
-                await self._determine_slot(*args),
+                await self._determine_slot(command, *args),
                 self.read_from_replicas and command in READ_COMMANDS,
             )
         ]
 
-    async def _determine_slot(self, *args: Any) -> int:
-        command = args[0]
+    async def _determine_slot(self, command: str, *args: Any) -> int:
         if self.command_flags.get(command) == SLOT_ID:
             # The command contains the slot ID
-            return int(args[1])
+            return int(args[0])
 
         # Get the keys in the command
 
@@ -516,19 +514,17 @@ async def _determine_slot(self, *args: Any) -> int:
         # - fix: https://github.com/redis/redis/pull/9733
         if command in ("EVAL", "EVALSHA"):
             # command syntax: EVAL "script body" num_keys ...
-            if len(args) <= 2:
-                raise RedisClusterException(f"Invalid args in command: {args}")
-            num_actual_keys = args[2]
-            eval_keys = args[3 : 3 + num_actual_keys]
+            if len(args) < 2:
+                raise RedisClusterException(
+                    f"Invalid args in command: {command, *args}"
+                )
+            keys = args[2 : 2 + args[1]]
             # if there are 0 keys, that means the script can be run on any node
             # so we can just return a random slot
-            if not eval_keys:
+            if not keys:
                 return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
-            keys = eval_keys
         else:
-            keys = await self.commands_parser.get_keys(
-                self.nodes_manager.default_node, *args
-            )
+            keys = await self.commands_parser.get_keys(command, *args)
             if not keys:
                 # FCALL can call a function with 0 keys, that means the function
                 #  can be run on any node so we can just return a random slot
@@ -848,13 +844,13 @@ def acquire_connection(self) -> Connection:
                 self._free.append(connection)
 
             return self._free.popleft()
-        else:
-            if len(self._connections) < self.max_connections:
-                connection = self.connection_class(**self.connection_kwargs)
-                self._connections.append(connection)
-                return connection
-            else:
-                raise ConnectionError("Too many connections")
+
+        if len(self._connections) < self.max_connections:
+            connection = self.connection_class(**self.connection_kwargs)
+            self._connections.append(connection)
+            return connection
+
+        raise ConnectionError("Too many connections")
 
     async def parse_response(
         self, connection: Connection, command: str, **kwargs: Any
@@ -872,10 +868,10 @@ async def parse_response(
             raise
 
         # Return response
-        try:
+        if command in self.response_callbacks:
             return self.response_callbacks[command](response, **kwargs)
-        except KeyError:
-            return response
+
+        return response
 
     async def execute_command(self, *args: Any, **kwargs: Any) -> Any:
         # Acquire connection
@@ -891,7 +887,7 @@ async def execute_command(self, *args: Any, **kwargs: Any) -> Any:
             # Release connection
             self._free.append(connection)
 
-    async def execute_pipeline(self) -> None:
+    async def execute_pipeline(self) -> bool:
         # Acquire connection
         connection = self.acquire_connection()
 
@@ -901,17 +897,20 @@ async def execute_pipeline(self) -> None:
         )
 
         # Read responses
-        try:
-            for cmd in self._command_stack:
-                try:
-                    cmd.result = await self.parse_response(
-                        connection, cmd.args[0], **cmd.kwargs
-                    )
-                except Exception as e:
-                    cmd.result = e
-        finally:
-            # Release connection
-            self._free.append(connection)
+        ret = False
+        for cmd in self._command_stack:
+            try:
+                cmd.result = await self.parse_response(
+                    connection, cmd.args[0], **cmd.kwargs
+                )
+            except Exception as e:
+                cmd.result = e
+                ret = True
+
+        # Release connection
+        self._free.append(connection)
+
+        return ret
 
 
 class NodesManager:
@@ -1257,6 +1256,13 @@ async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> N
     def __await__(self) -> Generator[Any, None, "ClusterPipeline"]:
         return self.initialize().__await__()
 
+    def __enter__(self) -> "ClusterPipeline":
+        self._command_stack = []
+        return self
+
+    def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
+        self._command_stack = []
+
     def __bool__(self) -> bool:
         return bool(self._command_stack)
 
@@ -1310,6 +1316,7 @@ async def execute(
 
                 try:
                     return await self._execute(
+                        self._client,
                         self._command_stack,
                         raise_on_error=raise_on_error,
                         allow_redirections=allow_redirections,
@@ -1331,60 +1338,60 @@ async def execute(
 
     async def _execute(
         self,
+        client: "RedisCluster",
         stack: List["PipelineCommand"],
         raise_on_error: bool = True,
         allow_redirections: bool = True,
     ) -> List[Any]:
-        client = self._client
+        todo = [
+            cmd for cmd in stack if not cmd.result or isinstance(cmd.result, Exception)
+        ]
+
         nodes = {}
-        for cmd in stack:
-            if not cmd.result or isinstance(cmd.result, Exception):
-                target_nodes = await client._determine_nodes(*cmd.args)
-                if not target_nodes:
-                    raise RedisClusterException(
-                        f"No targets were found to execute {cmd.args} command on"
-                    )
-                if len(target_nodes) > 1:
-                    raise RedisClusterException(
-                        f"Too many targets for command {cmd.args}"
-                    )
+        for cmd in todo:
+            target_nodes = await client._determine_nodes(*cmd.args)
+            if not target_nodes:
+                raise RedisClusterException(
+                    f"No targets were found to execute {cmd.args} command on"
+                )
+            if len(target_nodes) > 1:
+                raise RedisClusterException(f"Too many targets for command {cmd.args}")
 
-                node = target_nodes[0]
-                if node.name not in nodes:
-                    nodes[node.name] = node
-                    node._command_stack = []
-                node._command_stack.append(cmd)
+            node = target_nodes[0]
+            if node.name not in nodes:
+                nodes[node.name] = node
+                node._command_stack = []
+            node._command_stack.append(cmd)
 
-        await asyncio.gather(
+        errors = await asyncio.gather(
             *(asyncio.ensure_future(node.execute_pipeline()) for node in nodes.values())
         )
 
-        if allow_redirections:
-            # send each errored command individually
-            for cmd in stack:
-                if isinstance(cmd.result, (TryAgainError, MovedError, AskError)):
-                    try:
-                        cmd.result = await client.execute_command(
-                            *cmd.args, **cmd.kwargs
+        if any(errors):
+            if allow_redirections:
+                # send each errored command individually
+                for cmd in todo:
+                    if isinstance(cmd.result, (TryAgainError, MovedError, AskError)):
+                        try:
+                            cmd.result = await client.execute_command(
+                                *cmd.args, **cmd.kwargs
+                            )
+                        except Exception as e:
+                            cmd.result = e
+
+            if raise_on_error:
+                for cmd in todo:
+                    result = cmd.result
+                    if isinstance(result, Exception):
+                        command = " ".join(map(safe_str, cmd.args))
+                        msg = (
+                            f"Command # {cmd.position + 1} ({command}) of pipeline "
+                            f"caused error: {result.args}"
                         )
-                    except Exception as e:
-                        cmd.result = e
-
-        responses = [cmd.result for cmd in stack]
-
-        if raise_on_error:
-            for cmd in stack:
-                result = cmd.result
-                if isinstance(result, Exception):
-                    command = " ".join(map(safe_str, cmd.args))
-                    msg = (
-                        f"Command # {cmd.position + 1} ({command}) of pipeline "
-                        f"caused error: {result.args}"
-                    )
-                    result.args = (msg,) + result.args[1:]
-                    raise result
+                        result.args = (msg,) + result.args[1:]
+                        raise result
 
-        return responses
+        return [cmd.result for cmd in stack]
 
     def _split_command_across_slots(
         self, command: str, *keys: KeyT