Cluster code micro-optimization

benchmarks/command_packer_benchmark.py

@@ -9,6 +9,8 @@ def send_packed_command(self, command, check_health=True):
         if not self._sock:
             self.connect()
         try:
+            if isinstance(command, list):
+                command = SYM_EMPTY.join(command)
             self._sock.sendall(command)
         except OSError as e:
             self.disconnect()

redis/_parsers/commands.py

@@ -114,6 +114,37 @@ def initialize(self, r):
         for cmd in uppercase_commands:
             commands[cmd.lower()] = commands.pop(cmd)
         self.commands = commands
+        for command in self.commands.values():
+            first_key_pos = command["first_key_pos"]
+            last_key_pos = command["last_key_pos"]
+            flags = command["flags"]
+            if (
+                first_key_pos > 0
+                and first_key_pos == last_key_pos
+                and "movablekeys" not in flags
+                and "pubsub" not in flags
+                and command["name"] != "pubsub"
+            ):
+                command["_single_key_pos"] = first_key_pos
+
+    def _is_keyed_command(self, *args):
+        """
+        Determines whether the command is always keyed, never keyless.
+
+        Must have been initialized, won't automatically do so.
+        """
+        if len(args) < 2:
+            # The command has no keys in it
+            return False
+
+        cmd_name = args[0].lower()
+        commands = self.commands
+        command = commands.get(cmd_name)
+        if command is None:
+            return False
+
+        single_pos = command.get("_single_key_pos")
+        return single_pos is not None
 
     # As soon as this PR is merged into Redis, we should reimplement
     # our logic to use COMMAND INFO changes to determine the key positions
@@ -134,54 +165,65 @@ def get_keys(self, redis_conn, *args):
             return None
 
         cmd_name = args[0].lower()
-        if cmd_name not in self.commands:
+        commands = self.commands
+        command = commands.get(cmd_name)
+        if command is None:
             # try to split the command name and to take only the main command,
             # e.g. 'memory' for 'memory usage'
             cmd_name_split = cmd_name.split()
             cmd_name = cmd_name_split[0]
-            if cmd_name in self.commands:
+            if cmd_name in commands:
                 # save the split command to args
                 args = cmd_name_split + list(args[1:])
             else:
                 # We'll try to reinitialize the commands cache, if the engine
                 # version has changed, the commands may not be current
                 self.initialize(redis_conn)
-                if cmd_name not in self.commands:
+                commands = self.commands
+                if cmd_name not in commands:
                     raise RedisError(
                         f"{cmd_name.upper()} command doesn't exist in Redis commands"
                     )
+            command = commands.get(cmd_name)
 
-        command = self.commands.get(cmd_name)
-        if "movablekeys" in command["flags"]:
+        single_pos = command.get("_single_key_pos")
+        if single_pos is not None:
+            return [args[single_pos]]
+
+        flags = command["flags"]
+        if "movablekeys" in flags:
             keys = self._get_moveable_keys(redis_conn, *args)
-        elif "pubsub" in command["flags"] or command["name"] == "pubsub":
+        elif "pubsub" in flags or command["name"] == "pubsub":
             keys = self._get_pubsub_keys(*args)
         else:
-            if (
-                command["step_count"] == 0
-                and command["first_key_pos"] == 0
-                and command["last_key_pos"] == 0
-            ):
+            step_count = command["step_count"]
+            first_key_pos = command["first_key_pos"]
+            last_key_pos = command["last_key_pos"]
+            if step_count == 0 and first_key_pos == 0 and last_key_pos == 0:
                 is_subcmd = False
                 if "subcommands" in command:
                     subcmd_name = f"{cmd_name}|{args[1].lower()}"
                     for subcmd in command["subcommands"]:
                         if str_if_bytes(subcmd[0]) == subcmd_name:
                             command = self.parse_subcommand(subcmd)
 
-                            if command["first_key_pos"] > 0:
+                            step_count = command["step_count"]
+                            first_key_pos = command["first_key_pos"]
+                            last_key_pos = command["last_key_pos"]
+
+                            if first_key_pos > 0:
                                 is_subcmd = True
 
                 # The command doesn't have keys in it
                 if not is_subcmd:
                     return None
-            last_key_pos = command["last_key_pos"]
             if last_key_pos < 0:
-                last_key_pos = len(args) - abs(last_key_pos)
-            keys_pos = list(
-                range(command["first_key_pos"], last_key_pos + 1, command["step_count"])
+                last_key_pos += len(args)
+            keys = list(
+                map(
+                    args.__getitem__, range(first_key_pos, last_key_pos + 1, step_count)
+                )
             )
-            keys = [args[pos] for pos in keys_pos]
 
         return keys
 
@@ -421,7 +463,22 @@ async def initialize(self, node: Optional["ClusterNode"] = None) -> None:
             self.node = node
 
         commands = await self.node.execute_command("COMMAND")
-        self.commands = {cmd.lower(): command for cmd, command in commands.items()}
+        commands = {cmd.lower(): command for cmd, command in commands.items()}
+
+        for command in commands.values():
+            first_key_pos = command["first_key_pos"]
+            last_key_pos = command["last_key_pos"]
+            flags = command["flags"]
+            if (
+                first_key_pos > 0
+                and first_key_pos == last_key_pos
+                and "movablekeys" not in flags
+                and "pubsub" not in flags
+                and command["name"] != "pubsub"
+            ):
+                command["_single_key_pos"] = first_key_pos
+
+        self.commands = commands
 
     # As soon as this PR is merged into Redis, we should reimplement
     # our logic to use COMMAND INFO changes to determine the key positions
@@ -442,54 +499,62 @@ async def get_keys(self, *args: Any) -> Optional[Tuple[str, ...]]:
             return None
 
         cmd_name = args[0].lower()
-        if cmd_name not in self.commands:
+        commands = self.commands
+        command = commands.get(cmd_name)
+        if command is None:
             # try to split the command name and to take only the main command,
             # e.g. 'memory' for 'memory usage'
             cmd_name_split = cmd_name.split()
             cmd_name = cmd_name_split[0]
-            if cmd_name in self.commands:
+            if cmd_name in commands:
                 # save the split command to args
                 args = cmd_name_split + list(args[1:])
             else:
                 # We'll try to reinitialize the commands cache, if the engine
                 # version has changed, the commands may not be current
                 await self.initialize()
-                if cmd_name not in self.commands:
+                commands = self.commands
+                if cmd_name not in commands:
                     raise RedisError(
                         f"{cmd_name.upper()} command doesn't exist in Redis commands"
                     )
+            command = commands.get(cmd_name)
+
+        single_pos = command.get("_single_key_pos")
+        if single_pos is not None:
+            return [args[single_pos]]
 
-        command = self.commands.get(cmd_name)
-        if "movablekeys" in command["flags"]:
+        flags = command["flags"]
+        if "movablekeys" in flags:
             keys = await self._get_moveable_keys(*args)
-        elif "pubsub" in command["flags"] or command["name"] == "pubsub":
+        elif "pubsub" in flags or command["name"] == "pubsub":
             keys = self._get_pubsub_keys(*args)
         else:
-            if (
-                command["step_count"] == 0
-                and command["first_key_pos"] == 0
-                and command["last_key_pos"] == 0
-            ):
+            step_count = command["step_count"]
+            first_key_pos = command["first_key_pos"]
+            last_key_pos = command["last_key_pos"]
+            if step_count == 0 and first_key_pos == 0 and last_key_pos == 0:
                 is_subcmd = False
                 if "subcommands" in command:
                     subcmd_name = f"{cmd_name}|{args[1].lower()}"
                     for subcmd in command["subcommands"]:
                         if str_if_bytes(subcmd[0]) == subcmd_name:
                             command = self.parse_subcommand(subcmd)
 
-                            if command["first_key_pos"] > 0:
+                            first_key_pos = command["first_key_pos"]
+                            last_key_pos = command["last_key_pos"]
+                            step_count = command["step_count"]
+                            if first_key_pos > 0:
                                 is_subcmd = True
 
                 # The command doesn't have keys in it
                 if not is_subcmd:
                     return None
-            last_key_pos = command["last_key_pos"]
             if last_key_pos < 0:
-                last_key_pos = len(args) - abs(last_key_pos)
-            keys_pos = list(
-                range(command["first_key_pos"], last_key_pos + 1, command["step_count"])
-            )
-            keys = [args[pos] for pos in keys_pos]
+                last_key_pos += len(args)
+            keys = [
+                args[pos] for pos in range(first_key_pos, last_key_pos + 1, step_count)
+            ]
 
         return keys
 

redis/_parsers/encoders.py

@@ -15,24 +15,23 @@ def encode(self, value):
         "Return a bytestring or bytes-like representation of the value"
         if isinstance(value, (bytes, bytearray, memoryview)):
             return value
+        elif isinstance(value, str):
+            return value.encode(self.encoding, self.encoding_errors)
         elif isinstance(value, bool):
             # special case bool since it is a subclass of int
             raise DataError(
                 "Invalid input of type: 'bool'. Convert to a "
                 "bytes, string, int or float first."
             )
         elif isinstance(value, (int, float)):
-            value = repr(value).encode()
-        elif not isinstance(value, str):
+            return repr(value).encode()
+        else:
             # a value we don't know how to deal with. throw an error
             typename = type(value).__name__
             raise DataError(
                 f"Invalid input of type: '{typename}'. "
                 f"Convert to a bytes, string, int or float first."
             )
-        if isinstance(value, str):
-            value = value.encode(self.encoding, self.encoding_errors)
-        return value
 
     def decode(self, value, force=False):
         "Return a unicode string from the bytes-like representation"

redis/_parsers/helpers.py

@@ -1,4 +1,4 @@
-import datetime
+from datetime import datetime
 
 from redis.utils import str_if_bytes
 
@@ -11,7 +11,7 @@ def timestamp_to_datetime(response):
         response = int(response)
     except ValueError:
         return None
-    return datetime.datetime.fromtimestamp(response)
+    return datetime.fromtimestamp(response)
 
 
 def parse_debug_object(response):
@@ -875,7 +875,7 @@ def float_or_none(response):
 
 
 def bool_ok(response, **options):
-    return str_if_bytes(response) == "OK"
+    return response in (b"OK", "OK")
 
 
 def parse_zadd(response, **options):
@@ -1259,7 +1259,7 @@ def parse_pubsub_numsub(response, **options):
 def parse_client_kill(response, **options):
     if isinstance(response, int):
         return response
-    return str_if_bytes(response) == "OK"
+    return bool_ok(response)
 
 
 def parse_acl_getuser(response, **options):
@@ -1451,7 +1451,7 @@ def parse_set_result(response, **options):
         # Redis will return a getCommand result.
         # See `setGenericCommand` in t_string.c
         return response
-    return response and str_if_bytes(response) == "OK"
+    return response and bool_ok(response)
 
 
 def parse_function_list_unified(response, **options):

redis/_parsers/hiredis.py

@@ -349,17 +349,18 @@ async def read_response(
         if not self._connected:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
 
+        reader = self._reader
         if disable_decoding:
-            response = self._reader.gets(False)
+            response = reader.gets(False)
         else:
-            response = self._reader.gets()
+            response = reader.gets()
 
         while response is NOT_ENOUGH_DATA:
             await self.read_from_socket()
             if disable_decoding:
-                response = self._reader.gets(False)
+                response = reader.gets(False)
             else:
-                response = self._reader.gets()
+                response = reader.gets()
 
         # if the response is a ConnectionError or the response is a list and
         # the first item is a ConnectionError, raise it as something bad

redis/client.py

@@ -927,22 +927,26 @@ def failure_callback(error, failure_count):
 
     def parse_response(self, connection, command_name, **options):
         """Parses a response from the Redis server"""
-        try:
-            if NEVER_DECODE in options:
-                response = connection.read_response(disable_decoding=True)
-                options.pop(NEVER_DECODE)
-            else:
-                response = connection.read_response()
-        except ResponseError:
-            if EMPTY_RESPONSE in options:
-                return options[EMPTY_RESPONSE]
-            raise
+        if not options:
+            # Fast-path for the common case of no options
+            response = connection.read_response()
+        else:
+            try:
+                if NEVER_DECODE in options:
+                    response = connection.read_response(disable_decoding=True)
+                    options.pop(NEVER_DECODE)
+                else:
+                    response = connection.read_response()
+            except ResponseError:
+                if EMPTY_RESPONSE in options:
+                    return options[EMPTY_RESPONSE]
+                raise
 
-        if EMPTY_RESPONSE in options:
-            options.pop(EMPTY_RESPONSE)
+            if EMPTY_RESPONSE in options:
+                options.pop(EMPTY_RESPONSE)
 
-        # Remove keys entry, it needs only for cache.
-        options.pop("keys", None)
+            # Remove keys entry, it needs only for cache.
+            options.pop("keys", None)
 
         if command_name in self.response_callbacks:
             return self.response_callbacks[command_name](response, **options)