Added support for async cluster

Author: vladvildanov

redis/asyncio/cluster.py

@@ -3,6 +3,7 @@
 import random
 import socket
 import ssl
+import threading
 import warnings
 from typing import (
     Any,
@@ -29,6 +30,7 @@
 from redis.asyncio.connection import Connection, DefaultParser, SSLConnection, parse_url
 from redis.asyncio.lock import Lock
 from redis.asyncio.retry import Retry
+from redis.auth.token import TokenInterface
 from redis.backoff import default_backoff
 from redis.client import EMPTY_RESPONSE, NEVER_DECODE, AbstractRedis
 from redis.cluster import (
@@ -45,6 +47,7 @@
 from redis.commands import READ_COMMANDS, AsyncRedisClusterCommands
 from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
 from redis.credentials import CredentialProvider
+from redis.event import EventDispatcher, AsyncAfterConnectionReleasedEvent, AfterAsyncClusterInstantiationEvent
 from redis.exceptions import (
     AskError,
     BusyLoadingError,
@@ -60,7 +63,7 @@
     ResponseError,
     SlotNotCoveredError,
     TimeoutError,
-    TryAgainError,
+    TryAgainError, RedisError,
 )
 from redis.typing import AnyKeyT, EncodableT, KeyT
 from redis.utils import (
@@ -270,6 +273,7 @@ def __init__(
         ssl_ciphers: Optional[str] = None,
         protocol: Optional[int] = 2,
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher(),
     ) -> None:
         if db:
             raise RedisClusterException(
@@ -371,6 +375,7 @@ def __init__(
             require_full_coverage,
             kwargs,
             address_remap=address_remap,
+            event_dispatcher=event_dispatcher,
         )
         self.encoder = Encoder(encoding, encoding_errors, decode_responses)
         self.read_from_replicas = read_from_replicas
@@ -929,6 +934,8 @@ class ClusterNode:
     __slots__ = (
         "_connections",
         "_free",
+        "_lock",
+        "_event_dispatcher",
         "connection_class",
         "connection_kwargs",
         "host",
@@ -966,6 +973,9 @@ def __init__(
 
         self._connections: List[Connection] = []
         self._free: Deque[Connection] = collections.deque(maxlen=self.max_connections)
+        self._event_dispatcher = self.connection_kwargs.get("event_dispatcher", None)
+        if self._event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
 
     def __repr__(self) -> str:
         return (
@@ -1082,10 +1092,37 @@ async def execute_pipeline(self, commands: List["PipelineCommand"]) -> bool:
 
         return ret
 
+    async def re_auth_callback(self, token: TokenInterface):
+        tmp_queue = collections.deque()
+        while self._free:
+            conn = self._free.popleft()
+            await conn.retry.call_with_retry(
+                lambda: conn.send_command('AUTH', token.try_get('oid'), token.get_value()),
+                lambda error: self._mock(error)
+            )
+            await conn.retry.call_with_retry(
+                lambda: conn.read_response(),
+                lambda error: self._mock(error)
+            )
+            tmp_queue.append(conn)
+
+        while tmp_queue:
+            conn = tmp_queue.popleft()
+            self._free.append(conn)
+
+    async def _mock(self, error: RedisError):
+        """
+        Dummy functions, needs to be passed as error callback to retry object.
+        :param error:
+        :return:
+        """
+        pass
+
 
 class NodesManager:
     __slots__ = (
         "_moved_exception",
+        "_event_dispatcher",
         "connection_kwargs",
         "default_node",
         "nodes_cache",
@@ -1102,6 +1139,7 @@ def __init__(
         require_full_coverage: bool,
         connection_kwargs: Dict[str, Any],
         address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
+        event_dispatcher: Optional[EventDispatcher] = EventDispatcher(),
     ) -> None:
         self.startup_nodes = {node.name: node for node in startup_nodes}
         self.require_full_coverage = require_full_coverage
@@ -1113,6 +1151,7 @@ def __init__(
         self.slots_cache: Dict[int, List["ClusterNode"]] = {}
         self.read_load_balancer = LoadBalancer()
         self._moved_exception: MovedError = None
+        self._event_dispatcher = event_dispatcher
 
     def get_node(
         self,
@@ -1230,6 +1269,11 @@ async def initialize(self) -> None:
             try:
                 # Make sure cluster mode is enabled on this node
                 try:
+                    self._event_dispatcher.dispatch(
+                        AfterAsyncClusterInstantiationEvent(
+                            self.nodes_cache,
+                            self.connection_kwargs.get("credential_provider", None))
+                    )
                     cluster_slots = await startup_node.execute_command("CLUSTER SLOTS")
                 except ResponseError:
                     raise RedisClusterException(

redis/event.py

@@ -4,6 +4,7 @@
 from enum import Enum
 from typing import List, Union, Optional
 
+from redis.auth.token import TokenInterface
 from redis.credentials import StreamingCredentialProvider, CredentialProvider
 
 
@@ -57,6 +58,9 @@ def __init__(self):
             AfterSingleConnectionInstantiationEvent: [
                 RegisterReAuthForSingleConnection()
             ],
+            AfterAsyncClusterInstantiationEvent: [
+                RegisterReAuthForAsyncClusterNodes()
+            ],
             AsyncAfterConnectionReleasedEvent: [
                 AsyncReAuthConnectionListener(),
             ],
@@ -154,6 +158,29 @@ def connection_lock(self) -> Union[threading.Lock, asyncio.Lock]:
         return self._connection_lock
 
 
+class AfterAsyncClusterInstantiationEvent:
+    """
+    Event that will be fired after async cluster instance was created.
+
+    Async cluster doesn't use connection pools, instead ClusterNode object manages connections.
+    """
+    def __init__(
+            self,
+            nodes: dict,
+            credential_provider: Optional[CredentialProvider] = None,
+    ):
+        self._nodes = nodes
+        self._credential_provider = credential_provider
+
+    @property
+    def nodes(self) -> dict:
+        return self._nodes
+
+    @property
+    def credential_provider(self) -> Union[CredentialProvider, None]:
+        return self._credential_provider
+
+
 class ReAuthConnectionListener(EventListenerInterface):
     """
     Listener that performs re-authentication of given connection.
@@ -225,3 +252,17 @@ async def _re_auth_async(self, token):
         async with self._event.connection_lock:
             await self._event.connection.send_command('AUTH', token.try_get('oid'), token.get_value())
             await self._event.connection.read_response()
+
+
+class RegisterReAuthForAsyncClusterNodes(EventListenerInterface):
+    def __init__(self):
+        self._event = None
+
+    def listen(self, event: AfterAsyncClusterInstantiationEvent):
+        if isinstance(event.credential_provider, StreamingCredentialProvider):
+            self._event = event
+            event.credential_provider.on_next(self._re_auth)
+
+    async def _re_auth(self, token: TokenInterface):
+        for key in self._event.nodes:
+            await self._event.nodes[key].re_auth_callback(token)