Apply routing and response aggregation policies in OSS Cluster mode (#3834)

* Added policies parsing and policies resolving (#3805) * Added poilcy resolution method * Moved main command proceessing on top * Fixed return type and keyless detection * Added Dynamic and Static policies * Added coverage for policy resolvers * Removed all policies except search (phase 1) * Applied comments * Apply policies for normal and pipeline mode (#3818) * Added poilcy resolution method * Moved main command proceessing on top * Fixed return type and keyless detection * Added Dynamic and Static policies * Added coverage for policy resolvers * Applied request policies * Removed all policies except search (phase 1) * Policy applied for normal and pipeline mode * Added assertion with core command * Applied comments * Renamed method * Added async implementation for Request-Response policies (#3824) * Added poilcy resolution method * Moved main command proceessing on top * Fixed return type and keyless detection * Added Dynamic and Static policies * Added coverage for policy resolvers * Applied request policies * Removed all policies except search (phase 1) * Policy applied for normal and pipeline mode * Added assertion with core command * Applied comments * Added async support for Request/Responce policies * Fixed method name * Updated async to sync * Codestyle changes * Codestyle changes * Codestyle changes * Revert changes * Marked tests with Redis 8.0 * Added timeouts for index creation * Codestyle changes * Fixed RESP3 responses * Codestyle changes * Added additional option for cluster

Author: vladvildanov

docker-compose.yml

@@ -4,7 +4,7 @@ x-client-libs-stack-image: &client-libs-stack-image
  image: "redislabs/client-libs-test:${CLIENT_LIBS_TEST_STACK_IMAGE_TAG:-8.2}"
 
 x-client-libs-image: &client-libs-image
- image: "redislabs/client-libs-test:${CLIENT_LIBS_TEST_IMAGE_TAG:-8.2}"
+ image: "redislabs/client-libs-test:${CLIENT_LIBS_TEST_IMAGE_TAG:-8.4-RC1-pre.2}"
 
 services:
 
@@ -58,7 +58,7 @@ services:
  - TLS_ENABLED=yes
  - PORT=16379
  - TLS_PORT=27379
- command: ${REDIS_EXTRA_ARGS:---enable-debug-command yes --enable-module-command yes --tls-auth-clients optional --save ""}
+ command: ${REDIS_EXTRA_ARGS:---enable-debug-command yes --enable-module-command yes --tls-auth-clients optional --save "" --tls-cluster yes}
  ports:
  - "16379-16384:16379-16384"
  - "27379-27384:27379-27384"

redis/_parsers/commands.py

@@ -26,6 +26,7 @@
 )
 
 from redis._parsers import AsyncCommandsParser, Encoder
+from redis._parsers.commands import CommandPolicies, RequestPolicy, ResponsePolicy
 from redis._parsers.helpers import (
  _RedisCallbacks,
  _RedisCallbacksRESP2,
@@ -51,6 +52,7 @@
  parse_cluster_slots,
 )
 from redis.commands import READ_COMMANDS, AsyncRedisClusterCommands
+from redis.commands.policies import AsyncPolicyResolver, AsyncStaticPolicyResolver
 from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
 from redis.credentials import CredentialProvider
 from redis.event import AfterAsyncClusterInstantiationEvent, EventDispatcher
@@ -310,6 +312,7 @@ def __init__(
  protocol: Optional[int] = 2,
  address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
  event_dispatcher: Optional[EventDispatcher] = None,
+ policy_resolver: AsyncPolicyResolver = AsyncStaticPolicyResolver(),
  ) -> None:
  if db:
  raise RedisClusterException(
@@ -423,7 +426,36 @@ def __init__(
  self.load_balancing_strategy = load_balancing_strategy
  self.reinitialize_steps = reinitialize_steps
  self.reinitialize_counter = 0
+
+ # For backward compatibility, mapping from existing policies to new one
+ self._command_flags_mapping: dict[str, Union[RequestPolicy, ResponsePolicy]] = {
+ self.__class__.RANDOM: RequestPolicy.DEFAULT_KEYLESS,
+ self.__class__.PRIMARIES: RequestPolicy.ALL_SHARDS,
+ self.__class__.ALL_NODES: RequestPolicy.ALL_NODES,
+ self.__class__.REPLICAS: RequestPolicy.ALL_REPLICAS,
+ self.__class__.DEFAULT_NODE: RequestPolicy.DEFAULT_NODE,
+ SLOT_ID: RequestPolicy.DEFAULT_KEYED,
+ }
+
+ self._policies_callback_mapping: dict[
+ Union[RequestPolicy, ResponsePolicy], Callable
+ ] = {
+ RequestPolicy.DEFAULT_KEYLESS: lambda command_name: [
+ self.get_random_primary_or_all_nodes(command_name)
+ ],
+ RequestPolicy.DEFAULT_KEYED: self.get_nodes_from_slot,
+ RequestPolicy.DEFAULT_NODE: lambda: [self.get_default_node()],
+ RequestPolicy.ALL_SHARDS: self.get_primaries,
+ RequestPolicy.ALL_NODES: self.get_nodes,
+ RequestPolicy.ALL_REPLICAS: self.get_replicas,
+ RequestPolicy.SPECIAL: self.get_special_nodes,
+ ResponsePolicy.DEFAULT_KEYLESS: lambda res: res,
+ ResponsePolicy.DEFAULT_KEYED: lambda res: res,
+ }
+
+ self._policy_resolver = policy_resolver
  self.commands_parser = AsyncCommandsParser()
+ self._aggregate_nodes = None
  self.node_flags = self.__class__.NODE_FLAGS.copy()
  self.command_flags = self.__class__.COMMAND_FLAGS.copy()
  self.response_callbacks = kwargs["response_callbacks"]
@@ -619,6 +651,45 @@ def get_node_from_key(
 
  return slot_cache[node_idx]
 
+ def get_random_primary_or_all_nodes(self, command_name):
+ """
+ Returns random primary or all nodes depends on READONLY mode.
+ """
+ if self.read_from_replicas and command_name in READ_COMMANDS:
+ return self.get_random_node()
+
+ return self.get_random_primary_node()
+
+ def get_random_primary_node(self) -> "ClusterNode":
+ """
+ Returns a random primary node
+ """
+ return random.choice(self.get_primaries())
+
+ async def get_nodes_from_slot(self, command: str, *args):
+ """
+ Returns a list of nodes that hold the specified keys' slots.
+ """
+ # get the node that holds the key's slot
+ return [
+ self.nodes_manager.get_node_from_slot(
+ await self._determine_slot(command, *args),
+ self.read_from_replicas and command in READ_COMMANDS,
+ self.load_balancing_strategy if command in READ_COMMANDS else None,
+ )
+ ]
+
+ def get_special_nodes(self) -> Optional[list["ClusterNode"]]:
+ """
+ Returns a list of nodes for commands with a special policy.
+ """
+ if not self._aggregate_nodes:
+ raise RedisClusterException(
+ "Cannot execute FT.CURSOR commands without FT.AGGREGATE"
+ )
+
+ return self._aggregate_nodes
+
  def keyslot(self, key: EncodableT) -> int:
  """
  Find the keyslot for a given key.
@@ -643,39 +714,34 @@ def set_response_callback(self, command: str, callback: ResponseCallbackT) -> No
  self.response_callbacks[command] = callback
 
  async def _determine_nodes(
- self, command: str, *args: Any, node_flag: Optional[str] = None
+ self,
+ command: str,
+ *args: Any,
+ request_policy: RequestPolicy,
+ node_flag: Optional[str] = None,
  ) -> List["ClusterNode"]:
  # Determine which nodes should be executed the command on.
  # Returns a list of target nodes.
  if not node_flag:
  # get the nodes group for this command if it was predefined
  node_flag = self.command_flags.get(command)
 
- if node_flag in self.node_flags:
- if node_flag == self.__class__.DEFAULT_NODE:
- # return the cluster's default node
- return [self.nodes_manager.default_node]
- if node_flag == self.__class__.PRIMARIES:
- # return all primaries
- return self.nodes_manager.get_nodes_by_server_type(PRIMARY)
- if node_flag == self.__class__.REPLICAS:
- # return all replicas
- return self.nodes_manager.get_nodes_by_server_type(REPLICA)
- if node_flag == self.__class__.ALL_NODES:
- # return all nodes
- return list(self.nodes_manager.nodes_cache.values())
- if node_flag == self.__class__.RANDOM:
- # return a random node
- return [random.choice(list(self.nodes_manager.nodes_cache.values()))]
+ if node_flag in self._command_flags_mapping:
+ request_policy = self._command_flags_mapping[node_flag]
 
- # get the node that holds the key's slot
- return [
- self.nodes_manager.get_node_from_slot(
- await self._determine_slot(command, *args),
- self.read_from_replicas and command in READ_COMMANDS,
- self.load_balancing_strategy if command in READ_COMMANDS else None,
- )
- ]
+ policy_callback = self._policies_callback_mapping[request_policy]
+
+ if request_policy == RequestPolicy.DEFAULT_KEYED:
+ nodes = await policy_callback(command, *args)
+ elif request_policy == RequestPolicy.DEFAULT_KEYLESS:
+ nodes = policy_callback(command)
+ else:
+ nodes = policy_callback()
+
+ if command.lower() == "ft.aggregate":
+ self._aggregate_nodes = nodes
+
+ return nodes
 
  async def _determine_slot(self, command: str, *args: Any) -> int:
  if self.command_flags.get(command) == SLOT_ID:
@@ -780,6 +846,33 @@ async def execute_command(self, *args: EncodableT, **kwargs: Any) -> Any:
  target_nodes_specified = True
  retry_attempts = 0
 
+ command_policies = await self._policy_resolver.resolve(args[0].lower())
+
+ if not command_policies and not target_nodes_specified:
+ command_flag = self.command_flags.get(command)
+ if not command_flag:
+ # Fallback to default policy
+ if not self.get_default_node():
+ slot = None
+ else:
+ slot = await self._determine_slot(*args)
+ if not slot:
+ command_policies = CommandPolicies()
+ else:
+ command_policies = CommandPolicies(
+ request_policy=RequestPolicy.DEFAULT_KEYED,
+ response_policy=ResponsePolicy.DEFAULT_KEYED,
+ )
+ else:
+ if command_flag in self._command_flags_mapping:
+ command_policies = CommandPolicies(
+ request_policy=self._command_flags_mapping[command_flag]
+ )
+ else:
+ command_policies = CommandPolicies()
+ elif not command_policies and target_nodes_specified:
+ command_policies = CommandPolicies()
+
  # Add one for the first execution
  execute_attempts = 1 + retry_attempts
  for _ in range(execute_attempts):
@@ -795,7 +888,9 @@ async def execute_command(self, *args: EncodableT, **kwargs: Any) -> Any:
  if not target_nodes_specified:
  # Determine the nodes to execute the command on
  target_nodes = await self._determine_nodes(
- *args, node_flag=passed_targets
+ *args,
+ request_policy=command_policies.request_policy,
+ node_flag=passed_targets,
  )
  if not target_nodes:
  raise RedisClusterException(
@@ -806,10 +901,12 @@ async def execute_command(self, *args: EncodableT, **kwargs: Any) -> Any:
  # Return the processed result
  ret = await self._execute_command(target_nodes[0], *args, **kwargs)
  if command in self.result_callbacks:
- return self.result_callbacks[command](
+ ret = self.result_callbacks[command](
  command, {target_nodes[0].name: ret}, **kwargs
  )
- return ret
+ return self._policies_callback_mapping[
+ command_policies.response_policy
+ ](ret)
  else:
  keys = [node.name for node in target_nodes]
  values = await asyncio.gather(
@@ -824,7 +921,9 @@ async def execute_command(self, *args: EncodableT, **kwargs: Any) -> Any:
  return self.result_callbacks[command](
  command, dict(zip(keys, values)), **kwargs
  )
- return dict(zip(keys, values))
+ return self._policies_callback_mapping[
+ command_policies.response_policy
+ ](dict(zip(keys, values)))
  except Exception as e:
  if retry_attempts > 0 and type(e) in self.__class__.ERRORS_ALLOW_RETRY:
  # The nodes and slots cache were should be reinitialized.
@@ -1740,6 +1839,7 @@ def __init__(self, position: int, *args: Any, **kwargs: Any) -> None:
  self.kwargs = kwargs
  self.position = position
  self.result: Union[Any, Exception] = None
+ self.command_policies: Optional[CommandPolicies] = None
 
  def __repr__(self) -> str:
  return f"[{self.position}] {self.args} ({self.kwargs})"
@@ -1980,16 +2080,51 @@ async def _execute(
  nodes = {}
  for cmd in todo:
  passed_targets = cmd.kwargs.pop("target_nodes", None)
+ command_policies = await client._policy_resolver.resolve(
+ cmd.args[0].lower()
+ )
+
  if passed_targets and not client._is_node_flag(passed_targets):
  target_nodes = client._parse_target_nodes(passed_targets)
+
+ if not command_policies:
+ command_policies = CommandPolicies()
  else:
+ if not command_policies:
+ command_flag = client.command_flags.get(cmd.args[0])
+ if not command_flag:
+ # Fallback to default policy
+ if not client.get_default_node():
+ slot = None
+ else:
+ slot = await client._determine_slot(*cmd.args)
+ if not slot:
+ command_policies = CommandPolicies()
+ else:
+ command_policies = CommandPolicies(
+ request_policy=RequestPolicy.DEFAULT_KEYED,
+ response_policy=ResponsePolicy.DEFAULT_KEYED,
+ )
+ else:
+ if command_flag in client._command_flags_mapping:
+ command_policies = CommandPolicies(
+ request_policy=client._command_flags_mapping[
+ command_flag
+ ]
+ )
+ else:
+ command_policies = CommandPolicies()
+
  target_nodes = await client._determine_nodes(
- *cmd.args, node_flag=passed_targets
+ *cmd.args,
+ request_policy=command_policies.request_policy,
+ node_flag=passed_targets,
  )
  if not target_nodes:
  raise RedisClusterException(
  f"No targets were found to execute {cmd.args} command on"
  )
+ cmd.command_policies = command_policies
  if len(target_nodes) > 1:
  raise RedisClusterException(f"Too many targets for command {cmd.args}")
  node = target_nodes[0]
@@ -2010,9 +2145,9 @@ async def _execute(
  for cmd in todo:
  if isinstance(cmd.result, (TryAgainError, MovedError, AskError)):
  try:
- cmd.result = await client.execute_command(
- *cmd.args, **cmd.kwargs
- )
+ cmd.result = client._policies_callback_mapping[
+ cmd.command_policies.response_policy
+ ](await client.execute_command(*cmd.args, **cmd.kwargs))
  except Exception as e:
  cmd.result = e