Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import asyncio
2import collections
3import logging
4import random
5import socket
6import threading
7import time
8import warnings
9import weakref
10from abc import ABC, abstractmethod
11from collections import defaultdict
12from copy import copy
13from itertools import chain
14from types import MethodType
15from typing import (
16 TYPE_CHECKING,
17 Any,
18 Callable,
19 Coroutine,
20 Deque,
21 Dict,
22 Generator,
23 List,
24 Literal,
25 Mapping,
26 Optional,
27 Set,
28 Tuple,
29 Type,
30 TypeVar,
31 Union,
32)
34if TYPE_CHECKING:
35 from redis.asyncio.keyspace_notifications import (
36 AsyncClusterKeyspaceNotifications,
37 )
39from redis._defaults import (
40 DEFAULT_RETRY_BASE,
41 DEFAULT_RETRY_CAP,
42 DEFAULT_RETRY_COUNT,
43 DEFAULT_SOCKET_CONNECT_TIMEOUT,
44 DEFAULT_SOCKET_READ_SIZE,
45 DEFAULT_SOCKET_TIMEOUT,
46)
47from redis._parsers import AsyncCommandsParser, Encoder
48from redis._parsers.commands import CommandPolicies, RequestPolicy, ResponsePolicy
49from redis._parsers.helpers import get_response_callbacks
50from redis.asyncio.client import PubSub, ResponseCallbackT
51from redis.asyncio.connection import (
52 AbstractConnection,
53 Connection,
54 ConnectionPoolInterface,
55 SSLConnection,
56 parse_url,
57)
58from redis.asyncio.lock import Lock
59from redis.asyncio.observability.recorder import (
60 record_error_count,
61 record_operation_duration,
62)
63from redis.asyncio.retry import Retry
64from redis.auth.token import TokenInterface
65from redis.backoff import ExponentialWithJitterBackoff, NoBackoff
66from redis.client import EMPTY_RESPONSE, NEVER_DECODE, AbstractRedis
67from redis.cluster import (
68 PIPELINE_BLOCKED_COMMANDS,
69 PRIMARY,
70 REPLICA,
71 SLOT_ID,
72 AbstractRedisCluster,
73 LoadBalancer,
74 LoadBalancingStrategy,
75 block_pipeline_command,
76 get_node_name,
77 parse_cluster_shards,
78 parse_cluster_shards_unified,
79 parse_cluster_shards_with_str_keys,
80 parse_cluster_slots,
81)
82from redis.commands import READ_COMMANDS, AsyncRedisClusterCommands
83from redis.commands.helpers import list_or_args, parse_pubsub_subscriptions
84from redis.commands.policies import AsyncPolicyResolver, AsyncStaticPolicyResolver
85from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
86from redis.credentials import CredentialProvider
87from redis.driver_info import DriverInfo, resolve_driver_info
88from redis.event import (
89 AfterAsyncClusterInstantiationEvent,
90 AsyncAfterSlotsCacheRefreshEvent,
91 AsyncEventListenerInterface,
92 EventDispatcher,
93)
94from redis.exceptions import (
95 AskError,
96 BusyLoadingError,
97 ClusterDownError,
98 ClusterError,
99 ConnectionError,
100 CrossSlotTransactionError,
101 DataError,
102 ExecAbortError,
103 InvalidPipelineStack,
104 MaxConnectionsError,
105 MovedError,
106 RedisClusterException,
107 RedisError,
108 ResponseError,
109 SlotNotCoveredError,
110 TimeoutError,
111 TryAgainError,
112 WatchError,
113)
114from redis.typing import (
115 AnyKeyT,
116 ChannelT,
117 EncodableT,
118 KeyT,
119 PubSubHandler,
120 Subscription,
121)
122from redis.utils import (
123 SENTINEL,
124 SSL_AVAILABLE,
125 deprecated_args,
126 deprecated_function,
127 safe_str,
128 str_if_bytes,
129 truncate_text,
130)
132if SSL_AVAILABLE:
133 from ssl import TLSVersion, VerifyFlags, VerifyMode
134else:
135 TLSVersion = None
136 VerifyMode = None
137 VerifyFlags = None
139logger = logging.getLogger(__name__)
141TargetNodesT = TypeVar(
142 "TargetNodesT", str, "ClusterNode", List["ClusterNode"], Dict[Any, "ClusterNode"]
143)
146class RedisCluster(AbstractRedis, AbstractRedisCluster, AsyncRedisClusterCommands):
147 """
148 Create a new RedisCluster client.
150 Pass one of parameters:
152 - `host` & `port`
153 - `startup_nodes`
155 | Use ``await`` :meth:`initialize` to find cluster nodes & create connections.
156 | Use ``await`` :meth:`close` to disconnect connections & close client.
158 Many commands support the target_nodes kwarg. It can be one of the
159 :attr:`NODE_FLAGS`:
161 - :attr:`PRIMARIES`
162 - :attr:`REPLICAS`
163 - :attr:`ALL_NODES`
164 - :attr:`RANDOM`
165 - :attr:`DEFAULT_NODE`
167 Note: This client is not thread/process/fork safe.
169 :param host:
170 | Can be used to point to a startup node
171 :param port:
172 | Port used if **host** is provided
173 :param startup_nodes:
174 | :class:`~.ClusterNode` to used as a startup node
175 :param require_full_coverage:
176 | When set to ``False``: the client will not require a full coverage of
177 the slots. However, if not all slots are covered, and at least one node
178 has ``cluster-require-full-coverage`` set to ``yes``, the server will throw
179 a :class:`~.ClusterDownError` for some key-based commands.
180 | When set to ``True``: all slots must be covered to construct the cluster
181 client. If not all slots are covered, :class:`~.RedisClusterException` will be
182 thrown.
183 | See:
184 https://redis.io/docs/manual/scaling/#redis-cluster-configuration-parameters
185 :param read_from_replicas:
186 | @deprecated - please use load_balancing_strategy instead
187 | Enable read from replicas in READONLY mode.
188 When set to true, read commands will be assigned between the primary and
189 its replications in a Round-Robin manner.
190 The data read from replicas is eventually consistent with the data in primary nodes.
191 :param load_balancing_strategy:
192 | Enable read from replicas in READONLY mode and defines the load balancing
193 strategy that will be used for cluster node selection.
194 The data read from replicas is eventually consistent with the data in primary nodes.
195 :param dynamic_startup_nodes:
196 | Set the RedisCluster's startup nodes to all the discovered nodes.
197 If true (default value), the cluster's discovered nodes will be used to
198 determine the cluster nodes-slots mapping in the next topology refresh.
199 It will remove the initial passed startup nodes if their endpoints aren't
200 listed in the CLUSTER SLOTS output.
201 If you use dynamic DNS endpoints for startup nodes but CLUSTER SLOTS lists
202 specific IP addresses, it is best to set it to false.
203 :param reinitialize_steps:
204 | Specifies the number of MOVED errors that need to occur before reinitializing
205 the whole cluster topology. If a MOVED error occurs and the cluster does not
206 need to be reinitialized on this current error handling, only the MOVED slot
207 will be patched with the redirected node.
208 To reinitialize the cluster on every MOVED error, set reinitialize_steps to 1.
209 To avoid reinitializing the cluster on moved errors, set reinitialize_steps to
210 0.
211 :param cluster_error_retry_attempts:
212 | @deprecated - Please configure the 'retry' object instead
213 In case 'retry' object is set - this argument is ignored!
215 Number of times to retry before raising an error when :class:`~.TimeoutError`,
216 :class:`~.ConnectionError`, :class:`~.SlotNotCoveredError`
217 or :class:`~.ClusterDownError` are encountered
218 :param retry:
219 | A retry object that defines the retry strategy and the number of
220 retries for the cluster client.
221 In current implementation for the cluster client (starting form redis-py version 6.0.0)
222 the retry object is not yet fully utilized, instead it is used just to determine
223 the number of retries for the cluster client.
224 In the future releases the retry object will be used to handle the cluster client retries!
225 :param max_connections:
226 | Maximum number of connections per node. If there are no free connections & the
227 maximum number of connections are already created, a
228 :class:`~.MaxConnectionsError` is raised.
229 :param socket_keepalive:
230 | If ``True``, TCP keepalive is enabled for TCP socket connections.
231 :param socket_keepalive_options:
232 | Mapping of TCP keepalive socket option constants to values, for
233 example ``{socket.TCP_KEEPIDLE: 30}``. If left unspecified, redis-py
234 uses TCP keepalive defaults when ``socket_keepalive`` is enabled:
235 idle 30 seconds, interval 5 seconds, and 3 probes.
236 Platform-specific options that are not available are skipped.
237 Pass ``None`` or ``{}`` to avoid setting additional TCP keepalive
238 options.
239 :param address_remap:
240 | An optional callable which, when provided with an internal network
241 address of a node, e.g. a `(host, port)` tuple, will return the address
242 where the node is reachable. This can be used to map the addresses at
243 which the nodes _think_ they are, to addresses at which a client may
244 reach them, such as when they sit behind a proxy.
246 | Rest of the arguments will be passed to the
247 :class:`~redis.asyncio.connection.Connection` instances when created
249 :raises RedisClusterException:
250 if any arguments are invalid or unknown. Eg:
252 - `db` != 0 or None
253 - `path` argument for unix socket connection
254 - none of the `host`/`port` & `startup_nodes` were provided
256 """
258 @classmethod
259 def from_url(cls, url: str, **kwargs: Any) -> "RedisCluster":
260 """
261 Return a Redis client object configured from the given URL.
263 For example::
265 redis://[[username]:[password]]@localhost:6379/0
266 rediss://[[username]:[password]]@localhost:6379/0
268 Three URL schemes are supported:
270 - `redis://` creates a TCP socket connection. See more at:
271 <https://www.iana.org/assignments/uri-schemes/prov/redis>
272 - `rediss://` creates a SSL wrapped TCP socket connection. See more at:
273 <https://www.iana.org/assignments/uri-schemes/prov/rediss>
275 The username, password, hostname, path and all querystring values are passed
276 through ``urllib.parse.unquote`` in order to replace any percent-encoded values
277 with their corresponding characters.
279 All querystring options are cast to their appropriate Python types. Boolean
280 arguments can be specified with string values "True"/"False" or "Yes"/"No".
281 Values that cannot be properly cast cause a ``ValueError`` to be raised. Once
282 parsed, the querystring arguments and keyword arguments are passed to
283 :class:`~redis.asyncio.connection.Connection` when created.
284 In the case of conflicting arguments, querystring arguments are used.
285 """
286 kwargs.update(parse_url(url))
287 if kwargs.pop("connection_class", None) is SSLConnection:
288 kwargs["ssl"] = True
289 return cls(**kwargs)
291 # Type discrimination marker for @overload self-type pattern
292 _is_async_client: Literal[True] = True
294 __slots__ = (
295 "_initialize",
296 "_lock",
297 "retry",
298 "command_flags",
299 "commands_parser",
300 "connection_kwargs",
301 "encoder",
302 "node_flags",
303 "nodes_manager",
304 "read_from_replicas",
305 "reinitialize_counter",
306 "reinitialize_steps",
307 "response_callbacks",
308 "result_callbacks",
309 )
311 @deprecated_args(
312 args_to_warn=["read_from_replicas"],
313 reason="Please configure the 'load_balancing_strategy' instead",
314 version="5.3.0",
315 )
316 @deprecated_args(
317 args_to_warn=[
318 "cluster_error_retry_attempts",
319 ],
320 reason="Please configure the 'retry' object instead",
321 version="6.0.0",
322 )
323 @deprecated_args(
324 args_to_warn=["lib_name", "lib_version"],
325 reason="Use 'driver_info' parameter instead. "
326 "lib_name and lib_version will be removed in a future version.",
327 )
328 def __init__(
329 self,
330 host: str | None = None,
331 port: str | int = 6379,
332 # Cluster related kwargs
333 startup_nodes: List["ClusterNode"] | None = None,
334 require_full_coverage: bool = True,
335 read_from_replicas: bool = False,
336 load_balancing_strategy: LoadBalancingStrategy | None = None,
337 dynamic_startup_nodes: bool = True,
338 reinitialize_steps: int = 5,
339 cluster_error_retry_attempts: int = DEFAULT_RETRY_COUNT,
340 max_connections: int = 100,
341 retry: Retry | None = None,
342 retry_on_error: List[Type[Exception]] | None = None,
343 # Client related kwargs
344 db: str | int = 0,
345 path: str | None = None,
346 credential_provider: CredentialProvider | None = None,
347 username: str | None = None,
348 password: str | None = None,
349 client_name: str | None = None,
350 lib_name: str | object | None = SENTINEL,
351 lib_version: str | object | None = SENTINEL,
352 driver_info: DriverInfo | object | None = SENTINEL,
353 # Encoding related kwargs
354 encoding: str = "utf-8",
355 encoding_errors: str = "strict",
356 decode_responses: bool = False,
357 # Connection related kwargs
358 health_check_interval: float = 0,
359 socket_timeout: float | None = DEFAULT_SOCKET_TIMEOUT,
360 socket_connect_timeout: float | None = DEFAULT_SOCKET_CONNECT_TIMEOUT,
361 socket_read_size: int = DEFAULT_SOCKET_READ_SIZE,
362 socket_keepalive: bool = True,
363 socket_keepalive_options: Mapping[int, int | bytes] | object | None = SENTINEL,
364 # SSL related kwargs
365 ssl: bool = False,
366 ssl_ca_certs: str | None = None,
367 ssl_ca_data: str | None = None,
368 ssl_cert_reqs: "str | VerifyMode" = "required",
369 ssl_include_verify_flags: List["VerifyFlags"] | None = None,
370 ssl_exclude_verify_flags: List["VerifyFlags"] | None = None,
371 ssl_certfile: str | None = None,
372 ssl_check_hostname: bool = True,
373 ssl_keyfile: str | None = None,
374 ssl_min_version: "TLSVersion | None" = None,
375 ssl_ciphers: str | None = None,
376 protocol: int | None = None,
377 legacy_responses: bool = True,
378 address_remap: Callable[[Tuple[str, int]], Tuple[str, int]] | None = None,
379 event_dispatcher: EventDispatcher | None = None,
380 policy_resolver: AsyncPolicyResolver = AsyncStaticPolicyResolver(),
381 ) -> None:
382 if db:
383 raise RedisClusterException(
384 "Argument 'db' must be 0 or None in cluster mode"
385 )
387 if path:
388 raise RedisClusterException(
389 "Unix domain socket is not supported in cluster mode"
390 )
392 if (not host or not port) and not startup_nodes:
393 raise RedisClusterException(
394 "RedisCluster requires at least one node to discover the cluster.\n"
395 "Please provide one of the following or use RedisCluster.from_url:\n"
396 ' - host and port: RedisCluster(host="localhost", port=6379)\n'
397 " - startup_nodes: RedisCluster(startup_nodes=["
398 'ClusterNode("localhost", 6379), ClusterNode("localhost", 6380)])'
399 )
401 computed_driver_info = resolve_driver_info(driver_info, lib_name, lib_version)
403 kwargs: Dict[str, Any] = {
404 "max_connections": max_connections,
405 "connection_class": Connection,
406 # Client related kwargs
407 "credential_provider": credential_provider,
408 "username": username,
409 "password": password,
410 "client_name": client_name,
411 "driver_info": computed_driver_info,
412 # Encoding related kwargs
413 "encoding": encoding,
414 "encoding_errors": encoding_errors,
415 "decode_responses": decode_responses,
416 # Connection related kwargs
417 "health_check_interval": health_check_interval,
418 "socket_connect_timeout": socket_connect_timeout,
419 "socket_keepalive": socket_keepalive,
420 "socket_keepalive_options": socket_keepalive_options,
421 "socket_read_size": socket_read_size,
422 "socket_timeout": socket_timeout,
423 "protocol": protocol,
424 "legacy_responses": legacy_responses,
425 }
427 if ssl:
428 # SSL related kwargs
429 kwargs.update(
430 {
431 "connection_class": SSLConnection,
432 "ssl_ca_certs": ssl_ca_certs,
433 "ssl_ca_data": ssl_ca_data,
434 "ssl_cert_reqs": ssl_cert_reqs,
435 "ssl_include_verify_flags": ssl_include_verify_flags,
436 "ssl_exclude_verify_flags": ssl_exclude_verify_flags,
437 "ssl_certfile": ssl_certfile,
438 "ssl_check_hostname": ssl_check_hostname,
439 "ssl_keyfile": ssl_keyfile,
440 "ssl_min_version": ssl_min_version,
441 "ssl_ciphers": ssl_ciphers,
442 }
443 )
445 if read_from_replicas or load_balancing_strategy:
446 # Call our on_connect function to configure READONLY mode
447 kwargs["redis_connect_func"] = self.on_connect
449 if retry:
450 self.retry = retry
451 else:
452 self.retry = Retry(
453 backoff=ExponentialWithJitterBackoff(
454 base=DEFAULT_RETRY_BASE, cap=DEFAULT_RETRY_CAP
455 ),
456 retries=cluster_error_retry_attempts,
457 )
458 if retry_on_error:
459 self.retry.update_supported_errors(retry_on_error)
461 kwargs["response_callbacks"] = get_response_callbacks(
462 user_protocol=kwargs.get("protocol"),
463 legacy_responses=kwargs.get("legacy_responses", True),
464 )
465 if not kwargs.get("legacy_responses", True):
466 kwargs["response_callbacks"]["CLUSTER SHARDS"] = (
467 parse_cluster_shards_unified
468 )
469 elif kwargs.get("protocol") is None:
470 kwargs["response_callbacks"]["CLUSTER SHARDS"] = (
471 parse_cluster_shards_with_str_keys
472 )
473 else:
474 kwargs["response_callbacks"]["CLUSTER SHARDS"] = parse_cluster_shards
475 self.connection_kwargs = kwargs
477 if startup_nodes:
478 passed_nodes = []
479 for node in startup_nodes:
480 passed_nodes.append(
481 ClusterNode(node.host, node.port, **self.connection_kwargs)
482 )
483 startup_nodes = passed_nodes
484 else:
485 startup_nodes = []
486 if host and port:
487 startup_nodes.append(ClusterNode(host, port, **self.connection_kwargs))
489 if event_dispatcher is None:
490 self._event_dispatcher = EventDispatcher()
491 else:
492 self._event_dispatcher = event_dispatcher
494 self.startup_nodes = startup_nodes
495 self.nodes_manager = NodesManager(
496 startup_nodes,
497 require_full_coverage,
498 kwargs,
499 dynamic_startup_nodes=dynamic_startup_nodes,
500 address_remap=address_remap,
501 event_dispatcher=self._event_dispatcher,
502 )
503 self.encoder = Encoder(encoding, encoding_errors, decode_responses)
504 self.read_from_replicas = read_from_replicas
505 self.load_balancing_strategy = load_balancing_strategy
506 self.reinitialize_steps = reinitialize_steps
507 self.reinitialize_counter = 0
509 # For backward compatibility, mapping from existing policies to new one
510 self._command_flags_mapping: dict[str, Union[RequestPolicy, ResponsePolicy]] = {
511 self.__class__.RANDOM: RequestPolicy.DEFAULT_KEYLESS,
512 self.__class__.PRIMARIES: RequestPolicy.ALL_SHARDS,
513 self.__class__.ALL_NODES: RequestPolicy.ALL_NODES,
514 self.__class__.REPLICAS: RequestPolicy.ALL_REPLICAS,
515 self.__class__.DEFAULT_NODE: RequestPolicy.DEFAULT_NODE,
516 SLOT_ID: RequestPolicy.DEFAULT_KEYED,
517 }
519 self._policies_callback_mapping: dict[
520 Union[RequestPolicy, ResponsePolicy], Callable
521 ] = {
522 RequestPolicy.DEFAULT_KEYLESS: lambda command_name: [
523 self.get_random_primary_or_all_nodes(command_name)
524 ],
525 RequestPolicy.DEFAULT_KEYED: self.get_nodes_from_slot,
526 RequestPolicy.DEFAULT_NODE: lambda: [self.get_default_node()],
527 RequestPolicy.ALL_SHARDS: self.get_primaries,
528 RequestPolicy.ALL_NODES: self.get_nodes,
529 RequestPolicy.ALL_REPLICAS: self.get_replicas,
530 RequestPolicy.SPECIAL: self.get_special_nodes,
531 ResponsePolicy.DEFAULT_KEYLESS: lambda res: res,
532 ResponsePolicy.DEFAULT_KEYED: lambda res: res,
533 }
535 self._policy_resolver = policy_resolver
536 self.commands_parser = AsyncCommandsParser()
537 self._aggregate_nodes = None
538 self.node_flags = self.__class__.NODE_FLAGS.copy()
539 self.command_flags = self.__class__.COMMAND_FLAGS.copy()
540 self.response_callbacks = kwargs["response_callbacks"]
541 self.result_callbacks = self.__class__.RESULT_CALLBACKS.copy()
542 self.result_callbacks["CLUSTER SLOTS"] = (
543 lambda cmd, res, **kwargs: parse_cluster_slots(
544 list(res.values())[0], **kwargs
545 )
546 )
548 self._initialize = True
549 self._lock: Optional[asyncio.Lock] = None
551 # When used as an async context manager, we need to increment and decrement
552 # a usage counter so that we can close the connection pool when no one is
553 # using the client.
554 self._usage_counter = 0
555 self._usage_lock = asyncio.Lock()
557 async def initialize(
558 self,
559 additional_startup_nodes_info: Optional[List[Tuple[str, int]]] = None,
560 last_failed_node_name: Optional[str] = None,
561 ) -> "RedisCluster":
562 """Get all nodes from startup nodes & creates connections if not initialized."""
563 if self._initialize:
564 if not self._lock:
565 self._lock = asyncio.Lock()
566 async with self._lock:
567 if self._initialize:
568 try:
569 await self.nodes_manager.initialize(
570 additional_startup_nodes_info=additional_startup_nodes_info,
571 last_failed_node_name=last_failed_node_name,
572 )
573 await self.commands_parser.initialize(
574 self.nodes_manager.default_node
575 )
576 self._initialize = False
577 except BaseException:
578 await self.nodes_manager.aclose()
579 await self.nodes_manager.aclose("startup_nodes")
580 raise
581 return self
583 async def aclose(self) -> None:
584 """Close all connections & client if initialized."""
585 if not self._initialize:
586 if not self._lock:
587 self._lock = asyncio.Lock()
588 async with self._lock:
589 if not self._initialize:
590 self._initialize = True
591 await self.nodes_manager.aclose()
592 await self.nodes_manager.aclose("startup_nodes")
594 @deprecated_function(version="5.0.0", reason="Use aclose() instead", name="close")
595 async def close(self) -> None:
596 """alias for aclose() for backwards compatibility"""
597 await self.aclose()
599 async def __aenter__(self) -> "RedisCluster":
600 """
601 Async context manager entry. Increments a usage counter so that the
602 connection pool is only closed (via aclose()) when no context is using
603 the client.
604 """
605 await self._increment_usage()
606 try:
607 # Initialize the client (i.e. establish connection, etc.)
608 return await self.initialize()
609 except Exception:
610 # If initialization fails, decrement the counter to keep it in sync
611 await self._decrement_usage()
612 raise
614 async def _increment_usage(self) -> int:
615 """
616 Helper coroutine to increment the usage counter while holding the lock.
617 Returns the new value of the usage counter.
618 """
619 async with self._usage_lock:
620 self._usage_counter += 1
621 return self._usage_counter
623 async def _decrement_usage(self) -> int:
624 """
625 Helper coroutine to decrement the usage counter while holding the lock.
626 Returns the new value of the usage counter.
627 """
628 async with self._usage_lock:
629 self._usage_counter -= 1
630 return self._usage_counter
632 async def __aexit__(self, exc_type, exc_value, traceback):
633 """
634 Async context manager exit. Decrements a usage counter. If this is the
635 last exit (counter becomes zero), the client closes its connection pool.
636 """
637 current_usage = await asyncio.shield(self._decrement_usage())
638 if current_usage == 0:
639 # This was the last active context, so disconnect the pool.
640 await asyncio.shield(self.aclose())
642 def __await__(self) -> Generator[Any, None, "RedisCluster"]:
643 return self.initialize().__await__()
645 _DEL_MESSAGE = "Unclosed RedisCluster client"
647 def __del__(
648 self,
649 _warn: Any = warnings.warn,
650 _grl: Any = asyncio.get_running_loop,
651 ) -> None:
652 if hasattr(self, "_initialize") and not self._initialize:
653 _warn(f"{self._DEL_MESSAGE} {self!r}", ResourceWarning, source=self)
654 try:
655 context = {"client": self, "message": self._DEL_MESSAGE}
656 _grl().call_exception_handler(context)
657 except RuntimeError:
658 pass
660 async def on_connect(self, connection: Connection) -> None:
661 await connection.on_connect()
663 # Sending READONLY command to server to configure connection as
664 # readonly. Since each cluster node may change its server type due
665 # to a failover, we should establish a READONLY connection
666 # regardless of the server type. If this is a primary connection,
667 # READONLY would not affect executing write commands.
668 await connection.send_command("READONLY")
669 if str_if_bytes(await connection.read_response()) != "OK":
670 raise ConnectionError("READONLY command failed")
672 def get_nodes(self) -> List["ClusterNode"]:
673 """Get all nodes of the cluster."""
674 return list(self.nodes_manager.nodes_cache.values())
676 def get_primaries(self) -> List["ClusterNode"]:
677 """Get the primary nodes of the cluster."""
678 return self.nodes_manager.get_nodes_by_server_type(PRIMARY)
680 def get_replicas(self) -> List["ClusterNode"]:
681 """Get the replica nodes of the cluster."""
682 return self.nodes_manager.get_nodes_by_server_type(REPLICA)
684 def get_random_node(self) -> "ClusterNode":
685 """Get a random node of the cluster."""
686 return random.choice(list(self.nodes_manager.nodes_cache.values()))
688 def get_default_node(self) -> "ClusterNode":
689 """Get the default node of the client."""
690 return self.nodes_manager.default_node
692 def set_default_node(self, node: "ClusterNode") -> None:
693 """
694 Set the default node of the client.
696 :raises DataError: if None is passed or node does not exist in cluster.
697 """
698 if not node or not self.get_node(node_name=node.name):
699 raise DataError("The requested node does not exist in the cluster.")
701 self.nodes_manager.default_node = node
703 def get_node(
704 self,
705 host: Optional[str] = None,
706 port: Optional[int] = None,
707 node_name: Optional[str] = None,
708 ) -> Optional["ClusterNode"]:
709 """Get node by (host, port) or node_name."""
710 return self.nodes_manager.get_node(host, port, node_name)
712 def get_node_from_key(
713 self, key: str, replica: bool = False
714 ) -> Optional["ClusterNode"]:
715 """
716 Get the cluster node corresponding to the provided key.
718 :param key:
719 :param replica:
720 | Indicates if a replica should be returned
721 |
722 None will returned if no replica holds this key
724 :raises SlotNotCoveredError: if the key is not covered by any slot.
725 """
726 slot = self.keyslot(key)
727 slot_cache = self.nodes_manager.slots_cache.get(slot)
728 if not slot_cache:
729 raise SlotNotCoveredError(f'Slot "{slot}" is not covered by the cluster.')
731 if replica:
732 if len(self.nodes_manager.slots_cache[slot]) < 2:
733 return None
734 node_idx = 1
735 else:
736 node_idx = 0
738 return slot_cache[node_idx]
740 def get_random_primary_or_all_nodes(self, command_name):
741 """
742 Returns random primary or all nodes depends on READONLY mode.
743 """
744 if self.read_from_replicas and command_name in READ_COMMANDS:
745 return self.get_random_node()
747 return self.get_random_primary_node()
749 def get_random_primary_node(self) -> "ClusterNode":
750 """
751 Returns a random primary node
752 """
753 return random.choice(self.get_primaries())
755 async def get_nodes_from_slot(self, command: str, *args):
756 """
757 Returns a list of nodes that hold the specified keys' slots.
758 """
759 # get the node that holds the key's slot
760 return [
761 self.nodes_manager.get_node_from_slot(
762 await self._determine_slot(command, *args),
763 self.read_from_replicas and command in READ_COMMANDS,
764 self.load_balancing_strategy if command in READ_COMMANDS else None,
765 )
766 ]
768 def get_special_nodes(self) -> Optional[list["ClusterNode"]]:
769 """
770 Returns a list of nodes for commands with a special policy.
771 """
772 if not self._aggregate_nodes:
773 raise RedisClusterException(
774 "Cannot execute FT.CURSOR commands without FT.AGGREGATE"
775 )
777 return self._aggregate_nodes
779 def keyslot(self, key: EncodableT) -> int:
780 """
781 Find the keyslot for a given key.
783 See: https://redis.io/docs/manual/scaling/#redis-cluster-data-sharding
784 """
785 return key_slot(self.encoder.encode(key))
787 def get_encoder(self) -> Encoder:
788 """Get the encoder object of the client."""
789 return self.encoder
791 def get_connection_kwargs(self) -> Dict[str, Optional[Any]]:
792 """Get the kwargs passed to :class:`~redis.asyncio.connection.Connection`."""
793 return self.connection_kwargs
795 def set_retry(self, retry: Retry) -> None:
796 self.retry = retry
798 def set_response_callback(self, command: str, callback: ResponseCallbackT) -> None:
799 """Set a custom response callback."""
800 self.response_callbacks[command] = callback
802 async def _determine_nodes(
803 self,
804 command: str,
805 *args: Any,
806 request_policy: RequestPolicy,
807 node_flag: Optional[str] = None,
808 ) -> List["ClusterNode"]:
809 # Determine which nodes should be executed the command on.
810 # Returns a list of target nodes.
811 if not node_flag:
812 # get the nodes group for this command if it was predefined
813 node_flag = self.command_flags.get(command)
815 if node_flag in self._command_flags_mapping:
816 request_policy = self._command_flags_mapping[node_flag]
818 policy_callback = self._policies_callback_mapping[request_policy]
820 if request_policy == RequestPolicy.DEFAULT_KEYED:
821 nodes = await policy_callback(command, *args)
822 elif request_policy == RequestPolicy.DEFAULT_KEYLESS:
823 nodes = policy_callback(command)
824 else:
825 nodes = policy_callback()
827 if command.lower() == "ft.aggregate":
828 self._aggregate_nodes = nodes
830 return nodes
832 async def _determine_slot(self, command: str, *args: Any) -> int:
833 if self.command_flags.get(command) == SLOT_ID:
834 # The command contains the slot ID
835 return int(args[0])
837 # Get the keys in the command
839 # EVAL and EVALSHA are common enough that it's wasteful to go to the
840 # redis server to parse the keys. Besides, there is a bug in redis<7.0
841 # where `self._get_command_keys()` fails anyway. So, we special case
842 # EVAL/EVALSHA.
843 # - issue: https://github.com/redis/redis/issues/9493
844 # - fix: https://github.com/redis/redis/pull/9733
845 if command.upper() in ("EVAL", "EVALSHA"):
846 # command syntax: EVAL "script body" num_keys ...
847 if len(args) < 2:
848 raise RedisClusterException(
849 f"Invalid args in command: {command, *args}"
850 )
851 keys = args[2 : 2 + int(args[1])]
852 # if there are 0 keys, that means the script can be run on any node
853 # so we can just return a random slot
854 if not keys:
855 return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
856 else:
857 keys = await self.commands_parser.get_keys(command, *args)
858 if not keys:
859 # FCALL can call a function with 0 keys, that means the function
860 # can be run on any node so we can just return a random slot
861 if command.upper() in ("FCALL", "FCALL_RO"):
862 return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
863 raise RedisClusterException(
864 "No way to dispatch this command to Redis Cluster. "
865 "Missing key.\nYou can execute the command by specifying "
866 f"target nodes.\nCommand: {args}"
867 )
869 # single key command
870 if len(keys) == 1:
871 return self.keyslot(keys[0])
873 # multi-key command; we need to make sure all keys are mapped to
874 # the same slot
875 slots = {self.keyslot(key) for key in keys}
876 if len(slots) != 1:
877 raise RedisClusterException(
878 f"{command} - all keys must map to the same key slot"
879 )
881 return slots.pop()
883 def _is_node_flag(self, target_nodes: Any) -> bool:
884 return isinstance(target_nodes, str) and target_nodes in self.node_flags
886 def _parse_target_nodes(self, target_nodes: Any) -> List["ClusterNode"]:
887 if isinstance(target_nodes, list):
888 nodes = target_nodes
889 elif isinstance(target_nodes, ClusterNode):
890 # Supports passing a single ClusterNode as a variable
891 nodes = [target_nodes]
892 elif isinstance(target_nodes, dict):
893 # Supports dictionaries of the format {node_name: node}.
894 # It enables to execute commands with multi nodes as follows:
895 # rc.cluster_save_config(rc.get_primaries())
896 nodes = list(target_nodes.values())
897 else:
898 raise TypeError(
899 "target_nodes type can be one of the following: "
900 "node_flag (PRIMARIES, REPLICAS, RANDOM, ALL_NODES),"
901 "ClusterNode, list<ClusterNode>, or dict<any, ClusterNode>. "
902 f"The passed type is {type(target_nodes)}"
903 )
904 return nodes
906 async def _record_error_metric(
907 self,
908 error: Exception,
909 connection: Union[Connection, "ClusterNode"],
910 is_internal: bool = True,
911 retry_attempts: Optional[int] = None,
912 ):
913 """
914 Records error count metric directly.
915 Accepts either a Connection or ClusterNode object.
916 """
917 await record_error_count(
918 server_address=connection.host,
919 server_port=connection.port,
920 network_peer_address=connection.host,
921 network_peer_port=connection.port,
922 error_type=error,
923 retry_attempts=retry_attempts if retry_attempts is not None else 0,
924 is_internal=is_internal,
925 )
927 async def _record_command_metric(
928 self,
929 command_name: str,
930 duration_seconds: float,
931 connection: Union[Connection, "ClusterNode"],
932 error: Optional[Exception] = None,
933 ):
934 """
935 Records operation duration metric directly.
936 Accepts either a Connection or ClusterNode object.
937 """
938 # Connection has db attribute, ClusterNode has connection_kwargs
939 if hasattr(connection, "db"):
940 db = connection.db
941 else:
942 db = connection.connection_kwargs.get("db", 0)
943 await record_operation_duration(
944 command_name=command_name,
945 duration_seconds=duration_seconds,
946 server_address=connection.host,
947 server_port=connection.port,
948 db_namespace=str(db) if db is not None else None,
949 error=error,
950 )
952 async def execute_command(self, *args: EncodableT, **kwargs: Any) -> Any:
953 """
954 Execute a raw command on the appropriate cluster node or target_nodes.
956 It will retry the command as specified by the retries property of
957 the :attr:`retry` & then raise an exception.
959 :param args:
960 | Raw command args
961 :param kwargs:
963 - target_nodes: :attr:`NODE_FLAGS` or :class:`~.ClusterNode`
964 or List[:class:`~.ClusterNode`] or Dict[Any, :class:`~.ClusterNode`]
965 - Rest of the kwargs are passed to the Redis connection
967 :raises RedisClusterException: if target_nodes is not provided & the command
968 can't be mapped to a slot
969 """
970 command = args[0]
971 target_nodes = []
972 target_nodes_specified = False
973 retry_attempts = self.retry.get_retries()
975 passed_targets = kwargs.pop("target_nodes", None)
976 if passed_targets and not self._is_node_flag(passed_targets):
977 target_nodes = self._parse_target_nodes(passed_targets)
978 target_nodes_specified = True
979 retry_attempts = 0
981 command_policies = await self._policy_resolver.resolve(args[0].lower())
983 if not command_policies and not target_nodes_specified:
984 command_flag = self.command_flags.get(command)
985 if not command_flag:
986 # Fallback to default policy
987 if not self.get_default_node():
988 slot = None
989 else:
990 slot = await self._determine_slot(*args)
991 if slot is None:
992 command_policies = CommandPolicies()
993 else:
994 command_policies = CommandPolicies(
995 request_policy=RequestPolicy.DEFAULT_KEYED,
996 response_policy=ResponsePolicy.DEFAULT_KEYED,
997 )
998 else:
999 if command_flag in self._command_flags_mapping:
1000 command_policies = CommandPolicies(
1001 request_policy=self._command_flags_mapping[command_flag]
1002 )
1003 else:
1004 command_policies = CommandPolicies()
1005 elif not command_policies and target_nodes_specified:
1006 command_policies = CommandPolicies()
1008 # Add one for the first execution
1009 execute_attempts = 1 + retry_attempts
1010 failure_count = 0
1012 # Start timing for observability
1013 start_time = time.monotonic()
1014 last_failed_node_name = None
1016 for _ in range(execute_attempts):
1017 if self._initialize:
1018 await self.initialize(last_failed_node_name=last_failed_node_name)
1019 last_failed_node_name = None
1020 if (
1021 len(target_nodes) == 1
1022 and target_nodes[0] == self.get_default_node()
1023 ):
1024 # Replace the default cluster node
1025 self.replace_default_node()
1026 try:
1027 if not target_nodes_specified:
1028 # Determine the nodes to execute the command on
1029 target_nodes = await self._determine_nodes(
1030 *args,
1031 request_policy=command_policies.request_policy,
1032 node_flag=passed_targets,
1033 )
1034 if not target_nodes:
1035 raise RedisClusterException(
1036 f"No targets were found to execute {args} command on"
1037 )
1039 if len(target_nodes) == 1:
1040 # Return the processed result
1041 ret = await self._execute_command(target_nodes[0], *args, **kwargs)
1042 if command in self.result_callbacks:
1043 ret = self.result_callbacks[command](
1044 command, {target_nodes[0].name: ret}, **kwargs
1045 )
1046 return self._policies_callback_mapping[
1047 command_policies.response_policy
1048 ](ret)
1049 else:
1050 keys = [node.name for node in target_nodes]
1051 values = await asyncio.gather(
1052 *(
1053 asyncio.create_task(
1054 self._execute_command(node, *args, **kwargs)
1055 )
1056 for node in target_nodes
1057 )
1058 )
1059 if command in self.result_callbacks:
1060 return self.result_callbacks[command](
1061 command, dict(zip(keys, values)), **kwargs
1062 )
1063 return self._policies_callback_mapping[
1064 command_policies.response_policy
1065 ](dict(zip(keys, values)))
1066 except Exception as e:
1067 if retry_attempts > 0 and type(e) in self.__class__.ERRORS_ALLOW_RETRY:
1068 # The nodes and slots cache were should be reinitialized.
1069 # Try again with the new cluster setup.
1070 retry_attempts -= 1
1071 failure_count += 1
1072 last_failed_node_name = getattr(e, "last_failed_node_name", None)
1074 if hasattr(e, "connection"):
1075 await self._record_command_metric(
1076 command_name=command,
1077 duration_seconds=time.monotonic() - start_time,
1078 connection=e.connection,
1079 error=e,
1080 )
1081 await self._record_error_metric(
1082 error=e,
1083 connection=e.connection,
1084 retry_attempts=failure_count,
1085 )
1086 continue
1087 else:
1088 # raise the exception
1089 if hasattr(e, "connection"):
1090 await self._record_error_metric(
1091 error=e,
1092 connection=e.connection,
1093 retry_attempts=failure_count,
1094 is_internal=False,
1095 )
1096 raise e
1098 async def _execute_command(
1099 self, target_node: "ClusterNode", *args: Union[KeyT, EncodableT], **kwargs: Any
1100 ) -> Any:
1101 asking = moved = False
1102 redirect_addr = None
1103 ttl = self.RedisClusterRequestTTL
1104 command = args[0]
1105 start_time = time.monotonic()
1107 while ttl > 0:
1108 ttl -= 1
1109 try:
1110 if asking:
1111 target_node = self.get_node(node_name=redirect_addr)
1112 await target_node.execute_command("ASKING")
1113 asking = False
1114 elif moved:
1115 # MOVED occurred and the slots cache was updated,
1116 # refresh the target node
1117 slot = await self._determine_slot(*args)
1118 target_node = self.nodes_manager.get_node_from_slot(
1119 slot,
1120 self.read_from_replicas and args[0] in READ_COMMANDS,
1121 self.load_balancing_strategy
1122 if args[0] in READ_COMMANDS
1123 else None,
1124 )
1125 moved = False
1127 response = await target_node.execute_command(*args, **kwargs)
1128 await self._record_command_metric(
1129 command_name=command,
1130 duration_seconds=time.monotonic() - start_time,
1131 connection=target_node,
1132 )
1133 return response
1134 except BusyLoadingError as e:
1135 e.connection = target_node
1136 await self._record_command_metric(
1137 command_name=command,
1138 duration_seconds=time.monotonic() - start_time,
1139 connection=target_node,
1140 error=e,
1141 )
1142 raise
1143 except MaxConnectionsError as e:
1144 # MaxConnectionsError indicates client-side resource exhaustion
1145 # (too many connections in the pool), not a node failure.
1146 # Don't treat this as a node failure - just re-raise the error
1147 # without reinitializing the cluster.
1148 e.connection = target_node
1149 await self._record_command_metric(
1150 command_name=command,
1151 duration_seconds=time.monotonic() - start_time,
1152 connection=target_node,
1153 error=e,
1154 )
1155 raise
1156 except (ConnectionError, TimeoutError) as e:
1157 # Connection retries are being handled in the node's
1158 # Retry object.
1159 # Mark active connections for reconnect and disconnect free ones
1160 # This handles connection state (like READONLY) that may be stale
1161 target_node.update_active_connections_for_reconnect()
1162 await target_node.disconnect_free_connections()
1164 # Move the failed node to the end of the cached nodes list
1165 # so it's tried last during reinitialization
1166 self.nodes_manager.move_node_to_end_of_cached_nodes(target_node.name)
1167 e.last_failed_node_name = target_node.name
1169 # Signal that reinitialization is needed
1170 # The retry loop will handle initialize() AND replace_default_node()
1171 self._initialize = True
1172 e.connection = target_node
1173 await self._record_command_metric(
1174 command_name=command,
1175 duration_seconds=time.monotonic() - start_time,
1176 connection=target_node,
1177 error=e,
1178 )
1179 raise
1180 except (ClusterDownError, SlotNotCoveredError) as e:
1181 # ClusterDownError can occur during a failover and to get
1182 # self-healed, we will try to reinitialize the cluster layout
1183 # and retry executing the command
1185 # SlotNotCoveredError can occur when the cluster is not fully
1186 # initialized or can be temporary issue.
1187 # We will try to reinitialize the cluster topology
1188 # and retry executing the command
1190 await self.aclose()
1191 await asyncio.sleep(0.25)
1192 e.connection = target_node
1193 await self._record_command_metric(
1194 command_name=command,
1195 duration_seconds=time.monotonic() - start_time,
1196 connection=target_node,
1197 error=e,
1198 )
1199 raise
1200 except MovedError as e:
1201 # First, we will try to patch the slots/nodes cache with the
1202 # redirected node output and try again. If MovedError exceeds
1203 # 'reinitialize_steps' number of times, we will force
1204 # reinitializing the tables, and then try again.
1205 # 'reinitialize_steps' counter will increase faster when
1206 # the same client object is shared between multiple threads. To
1207 # reduce the frequency you can set this variable in the
1208 # RedisCluster constructor.
1209 self.reinitialize_counter += 1
1210 if (
1211 self.reinitialize_steps
1212 and self.reinitialize_counter % self.reinitialize_steps == 0
1213 ):
1214 await self.aclose()
1215 # Reset the counter
1216 self.reinitialize_counter = 0
1217 else:
1218 await self.nodes_manager.move_slot(e)
1219 moved = True
1220 await self._record_command_metric(
1221 command_name=command,
1222 duration_seconds=time.monotonic() - start_time,
1223 connection=target_node,
1224 error=e,
1225 )
1226 await self._record_error_metric(
1227 error=e,
1228 connection=target_node,
1229 )
1230 except AskError as e:
1231 redirect_addr = get_node_name(host=e.host, port=e.port)
1232 asking = True
1233 await self._record_command_metric(
1234 command_name=command,
1235 duration_seconds=time.monotonic() - start_time,
1236 connection=target_node,
1237 error=e,
1238 )
1239 await self._record_error_metric(
1240 error=e,
1241 connection=target_node,
1242 )
1243 except TryAgainError as e:
1244 if ttl < self.RedisClusterRequestTTL / 2:
1245 await asyncio.sleep(0.05)
1246 await self._record_command_metric(
1247 command_name=command,
1248 duration_seconds=time.monotonic() - start_time,
1249 connection=target_node,
1250 error=e,
1251 )
1252 await self._record_error_metric(
1253 error=e,
1254 connection=target_node,
1255 )
1256 except ResponseError as e:
1257 e.connection = target_node
1258 await self._record_command_metric(
1259 command_name=command,
1260 duration_seconds=time.monotonic() - start_time,
1261 connection=target_node,
1262 error=e,
1263 )
1264 raise
1265 except Exception as e:
1266 e.connection = target_node
1267 await self._record_command_metric(
1268 command_name=command,
1269 duration_seconds=time.monotonic() - start_time,
1270 connection=target_node,
1271 error=e,
1272 )
1273 raise
1275 e = ClusterError("TTL exhausted.")
1276 e.connection = target_node
1277 await self._record_command_metric(
1278 command_name=command,
1279 duration_seconds=time.monotonic() - start_time,
1280 connection=target_node,
1281 error=e,
1282 )
1283 raise e
1285 def pipeline(
1286 self, transaction: Optional[Any] = None, shard_hint: Optional[Any] = None
1287 ) -> "ClusterPipeline":
1288 """
1289 Create & return a new :class:`~.ClusterPipeline` object.
1291 Cluster implementation of pipeline does not support transaction or shard_hint.
1293 :raises RedisClusterException: if transaction or shard_hint are truthy values
1294 """
1295 if shard_hint:
1296 raise RedisClusterException("shard_hint is deprecated in cluster mode")
1298 return ClusterPipeline(self, transaction)
1300 def pubsub(
1301 self,
1302 node: Optional["ClusterNode"] = None,
1303 host: Optional[str] = None,
1304 port: Optional[int] = None,
1305 **kwargs: Any,
1306 ) -> "ClusterPubSub":
1307 """
1308 Create and return a ClusterPubSub instance.
1310 Allows passing a ClusterNode, or host&port, to get a pubsub instance
1311 connected to the specified node
1313 :param node: ClusterNode to connect to
1314 :param host: Host of the node to connect to
1315 :param port: Port of the node to connect to
1316 :param kwargs: Additional keyword arguments
1317 :return: ClusterPubSub instance
1318 """
1319 return ClusterPubSub(self, node=node, host=host, port=port, **kwargs)
1321 def keyspace_notifications(
1322 self,
1323 key_prefix: Union[str, bytes, None] = None,
1324 ignore_subscribe_messages: bool = True,
1325 ) -> "AsyncClusterKeyspaceNotifications":
1326 """
1327 Return an
1328 :class:`~redis.asyncio.keyspace_notifications.AsyncClusterKeyspaceNotifications`
1329 object for subscribing to keyspace and keyevent notifications across
1330 all primary nodes in the cluster.
1332 Note: Keyspace notifications must be enabled on all Redis cluster nodes
1333 via the ``notify-keyspace-events`` configuration option.
1335 Args:
1336 key_prefix: Optional prefix to filter and strip from keys in
1337 notifications.
1338 ignore_subscribe_messages: If True, subscribe/unsubscribe
1339 confirmations are not returned by
1340 get_message/listen.
1341 """
1342 from redis.asyncio.keyspace_notifications import (
1343 AsyncClusterKeyspaceNotifications,
1344 )
1346 return AsyncClusterKeyspaceNotifications(
1347 self,
1348 key_prefix=key_prefix,
1349 ignore_subscribe_messages=ignore_subscribe_messages,
1350 )
1352 def lock(
1353 self,
1354 name: KeyT,
1355 timeout: Optional[float] = None,
1356 sleep: float = 0.1,
1357 blocking: bool = True,
1358 blocking_timeout: Optional[float] = None,
1359 lock_class: Optional[Type[Lock]] = None,
1360 thread_local: bool = True,
1361 raise_on_release_error: bool = True,
1362 ) -> Lock:
1363 """
1364 Return a new Lock object using key ``name`` that mimics
1365 the behavior of threading.Lock.
1367 If specified, ``timeout`` indicates a maximum life for the lock.
1368 By default, it will remain locked until release() is called.
1370 ``sleep`` indicates the amount of time to sleep per loop iteration
1371 when the lock is in blocking mode and another client is currently
1372 holding the lock.
1374 ``blocking`` indicates whether calling ``acquire`` should block until
1375 the lock has been acquired or to fail immediately, causing ``acquire``
1376 to return False and the lock not being acquired. Defaults to True.
1377 Note this value can be overridden by passing a ``blocking``
1378 argument to ``acquire``.
1380 ``blocking_timeout`` indicates the maximum amount of time in seconds to
1381 spend trying to acquire the lock. A value of ``None`` indicates
1382 continue trying forever. ``blocking_timeout`` can be specified as a
1383 float or integer, both representing the number of seconds to wait.
1385 ``lock_class`` forces the specified lock implementation. Note that as
1386 of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
1387 a Lua-based lock). So, it's unlikely you'll need this parameter, unless
1388 you have created your own custom lock class.
1390 ``thread_local`` indicates whether the lock token is placed in
1391 thread-local storage. By default, the token is placed in thread local
1392 storage so that a thread only sees its token, not a token set by
1393 another thread. Consider the following timeline:
1395 time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
1396 thread-1 sets the token to "abc"
1397 time: 1, thread-2 blocks trying to acquire `my-lock` using the
1398 Lock instance.
1399 time: 5, thread-1 has not yet completed. redis expires the lock
1400 key.
1401 time: 5, thread-2 acquired `my-lock` now that it's available.
1402 thread-2 sets the token to "xyz"
1403 time: 6, thread-1 finishes its work and calls release(). if the
1404 token is *not* stored in thread local storage, then
1405 thread-1 would see the token value as "xyz" and would be
1406 able to successfully release the thread-2's lock.
1408 ``raise_on_release_error`` indicates whether to raise an exception when
1409 the lock is no longer owned when exiting the context manager. By default,
1410 this is True, meaning an exception will be raised. If False, the warning
1411 will be logged and the exception will be suppressed.
1413 In some use cases it's necessary to disable thread local storage. For
1414 example, if you have code where one thread acquires a lock and passes
1415 that lock instance to a worker thread to release later. If thread
1416 local storage isn't disabled in this case, the worker thread won't see
1417 the token set by the thread that acquired the lock. Our assumption
1418 is that these cases aren't common and as such default to using
1419 thread local storage."""
1420 if lock_class is None:
1421 lock_class = Lock
1422 return lock_class(
1423 self,
1424 name,
1425 timeout=timeout,
1426 sleep=sleep,
1427 blocking=blocking,
1428 blocking_timeout=blocking_timeout,
1429 thread_local=thread_local,
1430 raise_on_release_error=raise_on_release_error,
1431 )
1433 async def transaction(
1434 self, func: Coroutine[None, "ClusterPipeline", Any], *watches, **kwargs
1435 ):
1436 """
1437 Convenience method for executing the callable `func` as a transaction
1438 while watching all keys specified in `watches`. The 'func' callable
1439 should expect a single argument which is a Pipeline object.
1440 """
1441 shard_hint = kwargs.pop("shard_hint", None)
1442 value_from_callable = kwargs.pop("value_from_callable", False)
1443 watch_delay = kwargs.pop("watch_delay", None)
1444 async with self.pipeline(True, shard_hint) as pipe:
1445 while True:
1446 try:
1447 if watches:
1448 await pipe.watch(*watches)
1449 func_value = await func(pipe)
1450 exec_value = await pipe.execute()
1451 return func_value if value_from_callable else exec_value
1452 except WatchError:
1453 if watch_delay is not None and watch_delay > 0:
1454 time.sleep(watch_delay)
1455 continue
1458class ClusterNode:
1459 """
1460 Create a new ClusterNode.
1462 Each ClusterNode manages multiple :class:`~redis.asyncio.connection.Connection`
1463 objects for the (host, port).
1464 """
1466 __slots__ = (
1467 "_background_tasks",
1468 "_connections",
1469 "_free",
1470 "_lock",
1471 "_event_dispatcher",
1472 "connection_class",
1473 "connection_kwargs",
1474 "host",
1475 "max_connections",
1476 "name",
1477 "port",
1478 "response_callbacks",
1479 "server_type",
1480 )
1482 def __init__(
1483 self,
1484 host: str,
1485 port: Union[str, int],
1486 server_type: Optional[str] = None,
1487 *,
1488 max_connections: int = 100,
1489 connection_class: Type[Connection] = Connection,
1490 **connection_kwargs: Any,
1491 ) -> None:
1492 if host == "localhost":
1493 host = socket.gethostbyname(host)
1495 connection_kwargs["host"] = host
1496 connection_kwargs["port"] = port
1497 self.host = host
1498 self.port = port
1499 self.name = get_node_name(host, port)
1500 self.server_type = server_type
1502 self.max_connections = max_connections
1503 self.connection_class = connection_class
1504 self.connection_kwargs = connection_kwargs
1505 self.response_callbacks = connection_kwargs.pop("response_callbacks", {})
1507 self._connections: List[Connection] = []
1508 self._free: Deque[Connection] = collections.deque(maxlen=self.max_connections)
1509 self._background_tasks: Set[asyncio.Task] = set()
1510 self._event_dispatcher = self.connection_kwargs.get("event_dispatcher", None)
1511 if self._event_dispatcher is None:
1512 self._event_dispatcher = EventDispatcher()
1514 def __repr__(self) -> str:
1515 return (
1516 f"[host={self.host}, port={self.port}, "
1517 f"name={self.name}, server_type={self.server_type}]"
1518 )
1520 def __eq__(self, obj: Any) -> bool:
1521 return isinstance(obj, ClusterNode) and obj.name == self.name
1523 def __hash__(self) -> int:
1524 return hash(self.name)
1526 _DEL_MESSAGE = "Unclosed ClusterNode object"
1528 def __del__(
1529 self,
1530 _warn: Any = warnings.warn,
1531 _grl: Any = asyncio.get_running_loop,
1532 ) -> None:
1533 for connection in self._connections:
1534 if connection.is_connected:
1535 _warn(f"{self._DEL_MESSAGE} {self!r}", ResourceWarning, source=self)
1537 try:
1538 context = {"client": self, "message": self._DEL_MESSAGE}
1539 _grl().call_exception_handler(context)
1540 except RuntimeError:
1541 pass
1542 break
1544 async def disconnect(self) -> None:
1545 ret = await asyncio.gather(
1546 *(
1547 asyncio.create_task(connection.disconnect())
1548 for connection in self._connections
1549 ),
1550 return_exceptions=True,
1551 )
1552 exc = next((res for res in ret if isinstance(res, Exception)), None)
1553 if exc:
1554 raise exc
1556 def acquire_connection(self) -> Connection:
1557 try:
1558 return self._free.popleft()
1559 except IndexError:
1560 if len(self._connections) < self.max_connections:
1561 # We are configuring the connection pool not to retry
1562 # connections on lower level clients to avoid retrying
1563 # connections to nodes that are not reachable
1564 # and to avoid blocking the connection pool.
1565 # The only error that will have some handling in the lower
1566 # level clients is ConnectionError which will trigger disconnection
1567 # of the socket.
1568 # The retries will be handled on cluster client level
1569 # where we will have proper handling of the cluster topology
1570 retry = Retry(
1571 backoff=NoBackoff(),
1572 retries=0,
1573 supported_errors=(ConnectionError,),
1574 )
1575 connection_kwargs = self.connection_kwargs.copy()
1576 connection_kwargs["retry"] = retry
1577 connection = self.connection_class(**connection_kwargs)
1578 self._connections.append(connection)
1579 return connection
1581 raise MaxConnectionsError()
1583 async def disconnect_if_needed(self, connection: Connection) -> None:
1584 """
1585 Disconnect a connection if it's marked for reconnect.
1586 This implements lazy disconnection to avoid race conditions.
1587 The connection will auto-reconnect on next use.
1588 """
1589 if connection.should_reconnect():
1590 await connection.disconnect()
1592 def release(self, connection: Connection) -> None:
1593 """
1594 Release connection back to free queue.
1595 If the connection is marked for reconnect, disconnect it before
1596 returning it to the free queue.
1597 """
1598 if connection.should_reconnect():
1599 task = asyncio.create_task(self._disconnect_and_release(connection))
1600 self._background_tasks.add(task)
1601 task.add_done_callback(self._background_tasks.discard)
1602 return
1603 self._free.append(connection)
1605 async def _disconnect_and_release(self, connection: Connection) -> None:
1606 try:
1607 await connection.disconnect()
1608 except Exception as exc:
1609 logger.debug(
1610 "disconnecting released cluster connection failed: %r",
1611 exc,
1612 exc_info=True,
1613 )
1614 try:
1615 self._connections.remove(connection)
1616 except ValueError:
1617 pass
1618 return
1620 self._free.append(connection)
1622 def get_encoder(self) -> Encoder:
1623 """Return an :class:`Encoder` derived from this node's connection kwargs."""
1624 kwargs = self.connection_kwargs
1625 encoder_class = kwargs.get("encoder_class", Encoder)
1626 return encoder_class(
1627 encoding=kwargs.get("encoding", "utf-8"),
1628 encoding_errors=kwargs.get("encoding_errors", "strict"),
1629 decode_responses=kwargs.get("decode_responses", False),
1630 )
1632 def update_active_connections_for_reconnect(self) -> None:
1633 """
1634 Mark all in-use (active) connections for reconnect.
1635 In-use connections are those in _connections but not currently in _free.
1636 They will be disconnected after their current operation completes.
1637 """
1638 free_set = set(self._free)
1639 for connection in self._connections:
1640 if connection not in free_set:
1641 connection.mark_for_reconnect()
1643 async def disconnect_free_connections(self) -> None:
1644 """
1645 Disconnect all free/idle connections in the pool.
1646 This is useful after topology changes (e.g., failover) to clear
1647 stale connection state like READONLY mode.
1648 The connections remain in the pool and will reconnect on next use.
1649 """
1650 if self._free:
1651 # Take a snapshot to avoid issues if _free changes during await
1652 await asyncio.gather(
1653 *(connection.disconnect() for connection in tuple(self._free)),
1654 return_exceptions=True,
1655 )
1657 async def parse_response(
1658 self, connection: Connection, command: str, **kwargs: Any
1659 ) -> Any:
1660 try:
1661 if NEVER_DECODE in kwargs:
1662 response = await connection.read_response(disable_decoding=True)
1663 kwargs.pop(NEVER_DECODE)
1664 else:
1665 response = await connection.read_response()
1666 except ResponseError:
1667 if EMPTY_RESPONSE in kwargs:
1668 return kwargs[EMPTY_RESPONSE]
1669 raise
1671 if EMPTY_RESPONSE in kwargs:
1672 kwargs.pop(EMPTY_RESPONSE)
1674 # Remove keys entry, it needs only for cache.
1675 kwargs.pop("keys", None)
1677 # Return response
1678 if command in self.response_callbacks:
1679 return self.response_callbacks[command](response, **kwargs)
1681 return response
1683 async def execute_command(self, *args: Any, **kwargs: Any) -> Any:
1684 # Acquire connection
1685 connection = self.acquire_connection()
1686 try:
1687 # Handle lazy disconnect for connections marked for reconnect
1688 await self.disconnect_if_needed(connection)
1690 # Execute command
1691 await connection.send_packed_command(connection.pack_command(*args))
1693 # Read response
1694 return await self.parse_response(connection, args[0], **kwargs)
1695 finally:
1696 try:
1697 await self.disconnect_if_needed(connection)
1698 finally:
1699 # Release connection
1700 self.release(connection)
1702 async def execute_pipeline(self, commands: List["PipelineCommand"]) -> bool:
1703 # Acquire connection
1704 connection = self.acquire_connection()
1705 try:
1706 # Handle lazy disconnect for connections marked for reconnect
1707 await self.disconnect_if_needed(connection)
1709 # Execute command
1710 await connection.send_packed_command(
1711 connection.pack_commands(cmd.args for cmd in commands)
1712 )
1714 # Read responses
1715 ret = False
1716 for cmd in commands:
1717 try:
1718 cmd.result = await self.parse_response(
1719 connection, cmd.args[0], **cmd.kwargs
1720 )
1721 except Exception as e:
1722 cmd.result = e
1723 ret = True
1725 return ret
1726 finally:
1727 try:
1728 await self.disconnect_if_needed(connection)
1729 finally:
1730 # Release connection
1731 self.release(connection)
1733 async def re_auth_callback(self, token: TokenInterface):
1734 tmp_queue = collections.deque()
1735 while self._free:
1736 conn = self._free.popleft()
1737 await conn.retry.call_with_retry(
1738 lambda: conn.send_command(
1739 "AUTH", token.try_get("oid"), token.get_value()
1740 ),
1741 lambda error: self._mock(error),
1742 )
1743 await conn.retry.call_with_retry(
1744 lambda: conn.read_response(), lambda error: self._mock(error)
1745 )
1746 tmp_queue.append(conn)
1748 while tmp_queue:
1749 conn = tmp_queue.popleft()
1750 self._free.append(conn)
1752 async def _mock(self, error: RedisError):
1753 """
1754 Dummy functions, needs to be passed as error callback to retry object.
1755 :param error:
1756 :return:
1757 """
1758 pass
1761class NodesManager:
1762 __slots__ = (
1763 "_dynamic_startup_nodes",
1764 "_event_dispatcher",
1765 "_background_tasks",
1766 "connection_kwargs",
1767 "default_node",
1768 "nodes_cache",
1769 "_epoch",
1770 "read_load_balancer",
1771 "_initialize_lock",
1772 "require_full_coverage",
1773 "slots_cache",
1774 "startup_nodes",
1775 "address_remap",
1776 )
1778 def __init__(
1779 self,
1780 startup_nodes: List["ClusterNode"],
1781 require_full_coverage: bool,
1782 connection_kwargs: Dict[str, Any],
1783 dynamic_startup_nodes: bool = True,
1784 address_remap: Optional[Callable[[Tuple[str, int]], Tuple[str, int]]] = None,
1785 event_dispatcher: Optional[EventDispatcher] = None,
1786 ) -> None:
1787 self.startup_nodes = {node.name: node for node in startup_nodes}
1788 self.require_full_coverage = require_full_coverage
1789 self.connection_kwargs = connection_kwargs
1790 self.address_remap = address_remap
1792 self.default_node: "ClusterNode" = None
1793 self.nodes_cache: Dict[str, "ClusterNode"] = {}
1794 self.slots_cache: Dict[int, List["ClusterNode"]] = {}
1795 self._epoch: int = 0
1796 self.read_load_balancer = LoadBalancer()
1797 self._initialize_lock: asyncio.Lock = asyncio.Lock()
1799 self._background_tasks: Set[asyncio.Task] = set()
1800 self._dynamic_startup_nodes: bool = dynamic_startup_nodes
1801 if event_dispatcher is None:
1802 self._event_dispatcher = EventDispatcher()
1803 else:
1804 self._event_dispatcher = event_dispatcher
1806 def get_node(
1807 self,
1808 host: Optional[str] = None,
1809 port: Optional[int] = None,
1810 node_name: Optional[str] = None,
1811 ) -> Optional["ClusterNode"]:
1812 if host and port:
1813 # the user passed host and port
1814 if host == "localhost":
1815 host = socket.gethostbyname(host)
1816 return self.nodes_cache.get(get_node_name(host=host, port=port))
1817 elif node_name:
1818 return self.nodes_cache.get(node_name)
1819 else:
1820 raise DataError(
1821 "get_node requires one of the following: 1. node name 2. host and port"
1822 )
1824 def set_nodes(
1825 self,
1826 old: Dict[str, "ClusterNode"],
1827 new: Dict[str, "ClusterNode"],
1828 remove_old: bool = False,
1829 ) -> None:
1830 if remove_old:
1831 for name in list(old.keys()):
1832 if name not in new:
1833 # Node is removed from cache before disconnect starts,
1834 # so it won't be found in lookups during disconnect
1835 # Mark active connections so in-flight commands can
1836 # finish, then disconnect them when their current
1837 # operation completes. Free connections can be
1838 # disconnected immediately.
1839 removed_node = old.pop(name)
1840 removed_node.update_active_connections_for_reconnect()
1841 task = asyncio.create_task(
1842 removed_node.disconnect_free_connections()
1843 )
1844 self._background_tasks.add(task)
1845 task.add_done_callback(self._background_tasks.discard)
1847 for name, node in new.items():
1848 if name in old:
1849 # Preserve the existing node but mark connections for reconnect.
1850 # This method is sync so we can't call disconnect_free_connections()
1851 # which is async. Instead, we mark free connections for reconnect
1852 # and they will be lazily disconnected when acquired via
1853 # disconnect_if_needed() to avoid race conditions.
1854 # TODO: Make this method async in the next major release to allow
1855 # immediate disconnection of free connections.
1856 existing_node = old[name]
1857 existing_node.server_type = node.server_type
1858 existing_node.update_active_connections_for_reconnect()
1859 for conn in existing_node._free:
1860 conn.mark_for_reconnect()
1861 continue
1862 # New node is detected and should be added to the pool
1863 old[name] = node
1865 def move_node_to_end_of_cached_nodes(self, node_name: str) -> None:
1866 """
1867 Move a failing node to the end of startup_nodes and nodes_cache so it's
1868 tried last during reinitialization and when selecting the default node.
1869 If the node is not in the respective list, nothing is done.
1870 """
1871 # Move in startup_nodes
1872 if node_name in self.startup_nodes and len(self.startup_nodes) > 1:
1873 node = self.startup_nodes.pop(node_name)
1874 self.startup_nodes[node_name] = node # Re-insert at end
1876 # Move in nodes_cache - this affects get_nodes_by_server_type ordering
1877 # which is used to select the default_node during initialize()
1878 if node_name in self.nodes_cache and len(self.nodes_cache) > 1:
1879 node = self.nodes_cache.pop(node_name)
1880 self.nodes_cache[node_name] = node # Re-insert at end
1882 async def move_slot(self, e: AskError | MovedError):
1883 node_changed = False
1884 redirected_node = self.get_node(host=e.host, port=e.port)
1885 if redirected_node:
1886 # The node already exists
1887 if redirected_node.server_type != PRIMARY:
1888 # Update the node's server type
1889 redirected_node.server_type = PRIMARY
1890 else:
1891 # This is a new node, we will add it to the nodes cache
1892 redirected_node = ClusterNode(
1893 e.host, e.port, PRIMARY, **self.connection_kwargs
1894 )
1895 self.set_nodes(self.nodes_cache, {redirected_node.name: redirected_node})
1896 slot_nodes = self.slots_cache[e.slot_id]
1897 if redirected_node not in slot_nodes:
1898 # The new slot owner is a new server, or a server from a different
1899 # shard. We need to remove all current nodes from the slot's list
1900 # (including replications) and add just the new node.
1901 self.slots_cache[e.slot_id] = [redirected_node]
1902 node_changed = True
1903 elif redirected_node is not slot_nodes[0]:
1904 # The MOVED error resulted from a failover, and the new slot owner
1905 # had previously been a replica.
1906 old_primary = slot_nodes[0]
1907 # Update the old primary to be a replica and add it to the end of
1908 # the slot's node list
1909 old_primary.server_type = REPLICA
1910 slot_nodes.append(old_primary)
1911 # Remove the old replica, which is now a primary, from the slot's
1912 # node list
1913 slot_nodes.remove(redirected_node)
1914 # Override the old primary with the new one
1915 slot_nodes[0] = redirected_node
1916 if self.default_node == old_primary:
1917 # Update the default node with the new primary
1918 self.default_node = redirected_node
1919 node_changed = True
1920 # else: circular MOVED to current primary -> no-op
1921 # Dispatch so listeners can run shard-pubsub reconciliation; skipped on
1922 # the no-op branch to avoid needless walks under MOVED storms. A
1923 # listener must not break slots-cache refresh; log and continue so a
1924 # single buggy listener cannot starve the rest.
1925 if node_changed:
1926 try:
1927 await self._event_dispatcher.dispatch_async(
1928 AsyncAfterSlotsCacheRefreshEvent()
1929 )
1930 except Exception as exc:
1931 # Don't shadow the method parameter ``e``: ``except as`` binds
1932 # the listener exception in the function scope and ``del``s
1933 # the name on block exit (PEP 3134), which would also wipe
1934 # out the original AskError/MovedError parameter.
1935 logger.exception(
1936 "listener raised during slots-cache refresh: %s: %s",
1937 type(exc).__name__,
1938 exc,
1939 )
1941 def get_node_from_slot(
1942 self,
1943 slot: int,
1944 read_from_replicas: bool = False,
1945 load_balancing_strategy=None,
1946 ) -> "ClusterNode":
1947 if read_from_replicas is True and load_balancing_strategy is None:
1948 load_balancing_strategy = LoadBalancingStrategy.ROUND_ROBIN
1950 try:
1951 if len(self.slots_cache[slot]) > 1 and load_balancing_strategy:
1952 # get the server index using the strategy defined in load_balancing_strategy
1953 primary_name = self.slots_cache[slot][0].name
1954 node_idx = self.read_load_balancer.get_server_index(
1955 primary_name, len(self.slots_cache[slot]), load_balancing_strategy
1956 )
1957 return self.slots_cache[slot][node_idx]
1958 return self.slots_cache[slot][0]
1959 except (IndexError, TypeError):
1960 raise SlotNotCoveredError(
1961 f'Slot "{slot}" not covered by the cluster. '
1962 f'"require_full_coverage={self.require_full_coverage}"'
1963 )
1965 def get_nodes_by_server_type(self, server_type: str) -> List["ClusterNode"]:
1966 return [
1967 node
1968 for node in self.nodes_cache.values()
1969 if node.server_type == server_type
1970 ]
1972 async def initialize(
1973 self,
1974 additional_startup_nodes_info: Optional[List[Tuple[str, int]]] = None,
1975 last_failed_node_name: Optional[str] = None,
1976 ) -> None:
1977 self.read_load_balancer.reset()
1978 tmp_nodes_cache: Dict[str, "ClusterNode"] = {}
1979 tmp_slots: Dict[int, List["ClusterNode"]] = {}
1980 disagreements = []
1981 startup_nodes_reachable = False
1982 fully_covered = False
1983 exception = None
1984 epoch = self._epoch
1985 if additional_startup_nodes_info is None:
1986 additional_startup_nodes_info = []
1988 async with self._initialize_lock:
1989 if self._epoch != epoch:
1990 # another initialize call has already reinitialized the
1991 # nodes since we started waiting for the lock;
1992 # we don't need to do it again.
1993 return
1995 # Copy to a list to prevent RuntimeError if self.startup_nodes
1996 # is modified during iteration, then shuffle the iteration order.
1997 startup_nodes = list(self.startup_nodes.values())
1998 deferred_failed_nodes = []
1999 if last_failed_node_name is not None:
2000 for index, node in enumerate(startup_nodes):
2001 if node.name == last_failed_node_name:
2002 deferred_failed_nodes.append(startup_nodes.pop(index))
2003 break
2004 if len(startup_nodes) > 1:
2005 # Vary which startup node is queried first so clients do not
2006 # all reinitialize through the same node.
2007 random.shuffle(startup_nodes)
2008 additional_startup_nodes = [
2009 ClusterNode(host, port, **self.connection_kwargs)
2010 for host, port in additional_startup_nodes_info
2011 ]
2012 if last_failed_node_name is not None:
2013 for index, node in enumerate(additional_startup_nodes):
2014 if node.name == last_failed_node_name:
2015 if not deferred_failed_nodes:
2016 deferred_failed_nodes.append(node)
2017 additional_startup_nodes.pop(index)
2018 break
2019 for startup_node in chain(
2020 startup_nodes,
2021 additional_startup_nodes,
2022 deferred_failed_nodes,
2023 ):
2024 try:
2025 # Make sure cluster mode is enabled on this node
2026 try:
2027 self._event_dispatcher.dispatch(
2028 AfterAsyncClusterInstantiationEvent(
2029 self.nodes_cache,
2030 self.connection_kwargs.get("credential_provider", None),
2031 )
2032 )
2033 cluster_slots = await startup_node.execute_command(
2034 "CLUSTER SLOTS"
2035 )
2036 except ResponseError:
2037 raise RedisClusterException(
2038 "Cluster mode is not enabled on this node"
2039 )
2040 startup_nodes_reachable = True
2041 except Exception as e:
2042 # Try the next startup node.
2043 # The exception is saved and raised only if we have no more nodes.
2044 exception = e
2045 continue
2047 # CLUSTER SLOTS command results in the following output:
2048 # [[slot_section[from_slot,to_slot,master,replica1,...,replicaN]]]
2049 # where each node contains the following list: [IP, port, node_id]
2050 # Therefore, cluster_slots[0][2][0] will be the IP address of the
2051 # primary node of the first slot section.
2052 # If there's only one server in the cluster, its ``host`` is ''
2053 # Fix it to the host in startup_nodes
2054 if (
2055 len(cluster_slots) == 1
2056 and not cluster_slots[0][2][0]
2057 and len(self.startup_nodes) == 1
2058 ):
2059 cluster_slots[0][2][0] = startup_node.host
2061 for slot in cluster_slots:
2062 for i in range(2, len(slot)):
2063 slot[i] = [str_if_bytes(val) for val in slot[i]]
2064 primary_node = slot[2]
2065 host = primary_node[0]
2066 if host == "":
2067 host = startup_node.host
2068 port = int(primary_node[1])
2069 host, port = self.remap_host_port(host, port)
2071 nodes_for_slot = []
2073 target_node = tmp_nodes_cache.get(get_node_name(host, port))
2074 if not target_node:
2075 target_node = ClusterNode(
2076 host, port, PRIMARY, **self.connection_kwargs
2077 )
2078 # add this node to the nodes cache
2079 tmp_nodes_cache[target_node.name] = target_node
2080 nodes_for_slot.append(target_node)
2082 replica_nodes = slot[3:]
2083 for replica_node in replica_nodes:
2084 host = replica_node[0]
2085 port = replica_node[1]
2086 host, port = self.remap_host_port(host, port)
2088 target_replica_node = tmp_nodes_cache.get(
2089 get_node_name(host, port)
2090 )
2091 if not target_replica_node:
2092 target_replica_node = ClusterNode(
2093 host, port, REPLICA, **self.connection_kwargs
2094 )
2095 # add this node to the nodes cache
2096 tmp_nodes_cache[target_replica_node.name] = target_replica_node
2097 nodes_for_slot.append(target_replica_node)
2099 for i in range(int(slot[0]), int(slot[1]) + 1):
2100 if i not in tmp_slots:
2101 tmp_slots[i] = nodes_for_slot
2102 else:
2103 # Validate that 2 nodes want to use the same slot cache
2104 # setup
2105 tmp_slot = tmp_slots[i][0]
2106 if tmp_slot.name != target_node.name:
2107 disagreements.append(
2108 f"{tmp_slot.name} vs {target_node.name} on slot: {i}"
2109 )
2111 if len(disagreements) > 5:
2112 raise RedisClusterException(
2113 f"startup_nodes could not agree on a valid "
2114 f"slots cache: {', '.join(disagreements)}"
2115 )
2117 # Validate if all slots are covered or if we should try next startup node
2118 fully_covered = True
2119 for i in range(REDIS_CLUSTER_HASH_SLOTS):
2120 if i not in tmp_slots:
2121 fully_covered = False
2122 break
2123 if fully_covered:
2124 break
2126 if not startup_nodes_reachable:
2127 raise RedisClusterException(
2128 f"Redis Cluster cannot be connected. Please provide at least "
2129 f"one reachable node: {str(exception)}"
2130 ) from exception
2132 # Check if the slots are not fully covered
2133 if not fully_covered and self.require_full_coverage:
2134 # Despite the requirement that the slots be covered, there
2135 # isn't a full coverage
2136 raise RedisClusterException(
2137 f"All slots are not covered after query all startup_nodes. "
2138 f"{len(tmp_slots)} of {REDIS_CLUSTER_HASH_SLOTS} "
2139 f"covered..."
2140 )
2142 # Set the tmp variables to the real variables
2143 self.set_nodes(self.nodes_cache, tmp_nodes_cache, remove_old=True)
2144 # tmp_slots was built from CLUSTER SLOTS responses and can contain
2145 # newly-created ClusterNode objects for nodes we already know about.
2146 # Rebuild the slots cache with the preserved nodes_cache instances
2147 # so existing per-node connection pools stay in use after refresh.
2148 # Keep the shared node-list-per-slot-range shape from tmp_slots to
2149 # avoid allocating a separate list for every slot.
2150 node_lists_by_id: Dict[int, List["ClusterNode"]] = {}
2151 new_slots_cache: Dict[int, List["ClusterNode"]] = {}
2152 for slot, nodes in tmp_slots.items():
2153 node_list_id = id(nodes)
2154 slot_nodes = node_lists_by_id.get(node_list_id)
2155 if slot_nodes is None:
2156 slot_nodes = [self.nodes_cache[node.name] for node in nodes]
2157 node_lists_by_id[node_list_id] = slot_nodes
2158 new_slots_cache[slot] = slot_nodes
2159 self.slots_cache = new_slots_cache
2161 if self._dynamic_startup_nodes:
2162 # Populate the startup nodes with all discovered nodes
2163 self.set_nodes(self.startup_nodes, self.nodes_cache, remove_old=True)
2165 # Set the default node
2166 self.default_node = self.get_nodes_by_server_type(PRIMARY)[0]
2167 self._epoch += 1
2168 # Dispatch so listeners (e.g. ClusterPubSub) can reconcile per-node
2169 # state after slot ownership may have changed. A listener must not
2170 # break slots-cache refresh; log and continue so a single buggy
2171 # listener cannot starve the rest.
2172 try:
2173 await self._event_dispatcher.dispatch_async(
2174 AsyncAfterSlotsCacheRefreshEvent()
2175 )
2176 except Exception as e:
2177 logger.exception(
2178 "listener raised during slots-cache refresh: %s: %s",
2179 type(e).__name__,
2180 e,
2181 )
2183 async def aclose(self, attr: str = "nodes_cache") -> None:
2184 self.default_node = None
2185 await asyncio.gather(
2186 *(
2187 asyncio.create_task(node.disconnect())
2188 for node in getattr(self, attr).values()
2189 )
2190 )
2192 def remap_host_port(self, host: str, port: int) -> Tuple[str, int]:
2193 """
2194 Remap the host and port returned from the cluster to a different
2195 internal value. Useful if the client is not connecting directly
2196 to the cluster.
2197 """
2198 if self.address_remap:
2199 return self.address_remap((host, port))
2200 return host, port
2203class ClusterPipeline(AbstractRedis, AbstractRedisCluster, AsyncRedisClusterCommands):
2204 """
2205 Create a new ClusterPipeline object.
2207 Usage::
2209 result = await (
2210 rc.pipeline()
2211 .set("A", 1)
2212 .get("A")
2213 .hset("K", "F", "V")
2214 .hgetall("K")
2215 .mset_nonatomic({"A": 2, "B": 3})
2216 .get("A")
2217 .get("B")
2218 .delete("A", "B", "K")
2219 .execute()
2220 )
2221 # result = [True, "1", 1, {"F": "V"}, True, True, "2", "3", 1, 1, 1]
2223 Note: For commands `DELETE`, `EXISTS`, `TOUCH`, `UNLINK`, `mset_nonatomic`, which
2224 are split across multiple nodes, you'll get multiple results for them in the array.
2226 Retryable errors:
2227 - :class:`~.ClusterDownError`
2228 - :class:`~.ConnectionError`
2229 - :class:`~.TimeoutError`
2231 Redirection errors:
2232 - :class:`~.TryAgainError`
2233 - :class:`~.MovedError`
2234 - :class:`~.AskError`
2236 :param client:
2237 | Existing :class:`~.RedisCluster` client
2238 """
2240 __slots__ = (
2241 "cluster_client",
2242 "_transaction",
2243 "_execution_strategy",
2244 )
2246 # Type discrimination marker for @overload self-type pattern
2247 _is_async_client: Literal[True] = True
2249 def __init__(
2250 self, client: RedisCluster, transaction: Optional[bool] = None
2251 ) -> None:
2252 self.cluster_client = client
2253 self._transaction = transaction
2254 self._execution_strategy: ExecutionStrategy = (
2255 PipelineStrategy(self)
2256 if not self._transaction
2257 else TransactionStrategy(self)
2258 )
2260 @property
2261 def nodes_manager(self) -> "NodesManager":
2262 """Get the nodes manager from the cluster client."""
2263 return self.cluster_client.nodes_manager
2265 def set_response_callback(self, command: str, callback: ResponseCallbackT) -> None:
2266 """Set a custom response callback on the cluster client."""
2267 self.cluster_client.set_response_callback(command, callback)
2269 async def initialize(self) -> "ClusterPipeline":
2270 await self._execution_strategy.initialize()
2271 return self
2273 async def __aenter__(self) -> "ClusterPipeline":
2274 return await self.initialize()
2276 async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
2277 await self.reset()
2279 def __await__(self) -> Generator[Any, None, "ClusterPipeline"]:
2280 return self.initialize().__await__()
2282 def __bool__(self) -> bool:
2283 "Pipeline instances should always evaluate to True on Python 3+"
2284 return True
2286 def __len__(self) -> int:
2287 return len(self._execution_strategy)
2289 def execute_command(
2290 self, *args: Union[KeyT, EncodableT], **kwargs: Any
2291 ) -> "ClusterPipeline":
2292 """
2293 Append a raw command to the pipeline.
2295 :param args:
2296 | Raw command args
2297 :param kwargs:
2299 - target_nodes: :attr:`NODE_FLAGS` or :class:`~.ClusterNode`
2300 or List[:class:`~.ClusterNode`] or Dict[Any, :class:`~.ClusterNode`]
2301 - Rest of the kwargs are passed to the Redis connection
2302 """
2303 return self._execution_strategy.execute_command(*args, **kwargs)
2305 async def execute(
2306 self, raise_on_error: bool = True, allow_redirections: bool = True
2307 ) -> List[Any]:
2308 """
2309 Execute the pipeline.
2311 It will retry the commands as specified by retries specified in :attr:`retry`
2312 & then raise an exception.
2314 :param raise_on_error:
2315 | Raise the first error if there are any errors
2316 :param allow_redirections:
2317 | Whether to retry each failed command individually in case of redirection
2318 errors
2320 :raises RedisClusterException: if target_nodes is not provided & the command
2321 can't be mapped to a slot
2322 """
2323 try:
2324 return await self._execution_strategy.execute(
2325 raise_on_error, allow_redirections
2326 )
2327 finally:
2328 await self.reset()
2330 def _split_command_across_slots(
2331 self, command: str, *keys: KeyT
2332 ) -> "ClusterPipeline":
2333 for slot_keys in self.cluster_client._partition_keys_by_slot(keys).values():
2334 self.execute_command(command, *slot_keys)
2336 return self
2338 async def reset(self):
2339 """
2340 Reset back to empty pipeline.
2341 """
2342 await self._execution_strategy.reset()
2344 def multi(self):
2345 """
2346 Start a transactional block of the pipeline after WATCH commands
2347 are issued. End the transactional block with `execute`.
2348 """
2349 self._execution_strategy.multi()
2351 async def discard(self):
2352 """ """
2353 await self._execution_strategy.discard()
2355 async def watch(self, *names):
2356 """Watches the values at keys ``names``"""
2357 await self._execution_strategy.watch(*names)
2359 async def unwatch(self):
2360 """Unwatches all previously specified keys"""
2361 await self._execution_strategy.unwatch()
2363 async def unlink(self, *names):
2364 await self._execution_strategy.unlink(*names)
2366 def mset_nonatomic(
2367 self, mapping: Mapping[AnyKeyT, EncodableT]
2368 ) -> "ClusterPipeline":
2369 return self._execution_strategy.mset_nonatomic(mapping)
2372for command in PIPELINE_BLOCKED_COMMANDS:
2373 command = command.replace(" ", "_").lower()
2374 if command == "mset_nonatomic":
2375 continue
2377 setattr(ClusterPipeline, command, block_pipeline_command(command))
2380class PipelineCommand:
2381 def __init__(self, position: int, *args: Any, **kwargs: Any) -> None:
2382 self.args = args
2383 self.kwargs = kwargs
2384 self.position = position
2385 self.result: Union[Any, Exception] = None
2386 self.command_policies: Optional[CommandPolicies] = None
2388 def __repr__(self) -> str:
2389 return f"[{self.position}] {self.args} ({self.kwargs})"
2392class ExecutionStrategy(ABC):
2393 @abstractmethod
2394 async def initialize(self) -> "ClusterPipeline":
2395 """
2396 Initialize the execution strategy.
2398 See ClusterPipeline.initialize()
2399 """
2400 pass
2402 @abstractmethod
2403 def execute_command(
2404 self, *args: Union[KeyT, EncodableT], **kwargs: Any
2405 ) -> "ClusterPipeline":
2406 """
2407 Append a raw command to the pipeline.
2409 See ClusterPipeline.execute_command()
2410 """
2411 pass
2413 @abstractmethod
2414 async def execute(
2415 self, raise_on_error: bool = True, allow_redirections: bool = True
2416 ) -> List[Any]:
2417 """
2418 Execute the pipeline.
2420 It will retry the commands as specified by retries specified in :attr:`retry`
2421 & then raise an exception.
2423 See ClusterPipeline.execute()
2424 """
2425 pass
2427 @abstractmethod
2428 def mset_nonatomic(
2429 self, mapping: Mapping[AnyKeyT, EncodableT]
2430 ) -> "ClusterPipeline":
2431 """
2432 Executes multiple MSET commands according to the provided slot/pairs mapping.
2434 See ClusterPipeline.mset_nonatomic()
2435 """
2436 pass
2438 @abstractmethod
2439 async def reset(self):
2440 """
2441 Resets current execution strategy.
2443 See: ClusterPipeline.reset()
2444 """
2445 pass
2447 @abstractmethod
2448 def multi(self):
2449 """
2450 Starts transactional context.
2452 See: ClusterPipeline.multi()
2453 """
2454 pass
2456 @abstractmethod
2457 async def watch(self, *names):
2458 """
2459 Watch given keys.
2461 See: ClusterPipeline.watch()
2462 """
2463 pass
2465 @abstractmethod
2466 async def unwatch(self):
2467 """
2468 Unwatches all previously specified keys
2470 See: ClusterPipeline.unwatch()
2471 """
2472 pass
2474 @abstractmethod
2475 async def discard(self):
2476 pass
2478 @abstractmethod
2479 async def unlink(self, *names):
2480 """
2481 "Unlink a key specified by ``names``"
2483 See: ClusterPipeline.unlink()
2484 """
2485 pass
2487 @abstractmethod
2488 def __len__(self) -> int:
2489 pass
2492class AbstractStrategy(ExecutionStrategy):
2493 def __init__(self, pipe: ClusterPipeline) -> None:
2494 self._pipe: ClusterPipeline = pipe
2495 self._command_queue: List["PipelineCommand"] = []
2497 async def initialize(self) -> "ClusterPipeline":
2498 if self._pipe.cluster_client._initialize:
2499 await self._pipe.cluster_client.initialize()
2500 self._command_queue = []
2501 return self._pipe
2503 def execute_command(
2504 self, *args: Union[KeyT, EncodableT], **kwargs: Any
2505 ) -> "ClusterPipeline":
2506 self._command_queue.append(
2507 PipelineCommand(len(self._command_queue), *args, **kwargs)
2508 )
2509 return self._pipe
2511 def _annotate_exception(self, exception, number, command):
2512 """
2513 Provides extra context to the exception prior to it being handled
2514 """
2515 cmd = " ".join(map(safe_str, command))
2516 msg = (
2517 f"Command # {number} ({truncate_text(cmd)}) of pipeline "
2518 f"caused error: {exception.args[0]}"
2519 )
2520 exception.args = (msg,) + exception.args[1:]
2522 @abstractmethod
2523 def mset_nonatomic(
2524 self, mapping: Mapping[AnyKeyT, EncodableT]
2525 ) -> "ClusterPipeline":
2526 pass
2528 @abstractmethod
2529 async def execute(
2530 self, raise_on_error: bool = True, allow_redirections: bool = True
2531 ) -> List[Any]:
2532 pass
2534 @abstractmethod
2535 async def reset(self):
2536 pass
2538 @abstractmethod
2539 def multi(self):
2540 pass
2542 @abstractmethod
2543 async def watch(self, *names):
2544 pass
2546 @abstractmethod
2547 async def unwatch(self):
2548 pass
2550 @abstractmethod
2551 async def discard(self):
2552 pass
2554 @abstractmethod
2555 async def unlink(self, *names):
2556 pass
2558 def __len__(self) -> int:
2559 return len(self._command_queue)
2562class PipelineStrategy(AbstractStrategy):
2563 def __init__(self, pipe: ClusterPipeline) -> None:
2564 super().__init__(pipe)
2566 def mset_nonatomic(
2567 self, mapping: Mapping[AnyKeyT, EncodableT]
2568 ) -> "ClusterPipeline":
2569 encoder = self._pipe.cluster_client.encoder
2571 slots_pairs = {}
2572 for pair in mapping.items():
2573 slot = key_slot(encoder.encode(pair[0]))
2574 slots_pairs.setdefault(slot, []).extend(pair)
2576 for pairs in slots_pairs.values():
2577 self.execute_command("MSET", *pairs)
2579 return self._pipe
2581 async def execute(
2582 self, raise_on_error: bool = True, allow_redirections: bool = True
2583 ) -> List[Any]:
2584 if not self._command_queue:
2585 return []
2587 try:
2588 retry_attempts = self._pipe.cluster_client.retry.get_retries()
2589 while True:
2590 try:
2591 if self._pipe.cluster_client._initialize:
2592 await self._pipe.cluster_client.initialize()
2593 return await self._execute(
2594 self._pipe.cluster_client,
2595 self._command_queue,
2596 raise_on_error=raise_on_error,
2597 allow_redirections=allow_redirections,
2598 )
2600 except RedisCluster.ERRORS_ALLOW_RETRY as e:
2601 if retry_attempts > 0:
2602 # Try again with the new cluster setup. All other errors
2603 # should be raised.
2604 retry_attempts -= 1
2605 await self._pipe.cluster_client.aclose()
2606 await asyncio.sleep(0.25)
2607 else:
2608 # All other errors should be raised.
2609 raise e
2610 finally:
2611 await self.reset()
2613 async def _execute(
2614 self,
2615 client: "RedisCluster",
2616 stack: List["PipelineCommand"],
2617 raise_on_error: bool = True,
2618 allow_redirections: bool = True,
2619 ) -> List[Any]:
2620 todo = [
2621 cmd for cmd in stack if not cmd.result or isinstance(cmd.result, Exception)
2622 ]
2624 nodes = {}
2625 for cmd in todo:
2626 passed_targets = cmd.kwargs.pop("target_nodes", None)
2627 command_policies = await client._policy_resolver.resolve(
2628 cmd.args[0].lower()
2629 )
2631 if passed_targets and not client._is_node_flag(passed_targets):
2632 target_nodes = client._parse_target_nodes(passed_targets)
2634 if not command_policies:
2635 command_policies = CommandPolicies()
2636 else:
2637 if not command_policies:
2638 command_flag = client.command_flags.get(cmd.args[0])
2639 if not command_flag:
2640 # Fallback to default policy
2641 if not client.get_default_node():
2642 slot = None
2643 else:
2644 slot = await client._determine_slot(*cmd.args)
2645 if slot is None:
2646 command_policies = CommandPolicies()
2647 else:
2648 command_policies = CommandPolicies(
2649 request_policy=RequestPolicy.DEFAULT_KEYED,
2650 response_policy=ResponsePolicy.DEFAULT_KEYED,
2651 )
2652 else:
2653 if command_flag in client._command_flags_mapping:
2654 command_policies = CommandPolicies(
2655 request_policy=client._command_flags_mapping[
2656 command_flag
2657 ]
2658 )
2659 else:
2660 command_policies = CommandPolicies()
2662 target_nodes = await client._determine_nodes(
2663 *cmd.args,
2664 request_policy=command_policies.request_policy,
2665 node_flag=passed_targets,
2666 )
2667 if not target_nodes:
2668 raise RedisClusterException(
2669 f"No targets were found to execute {cmd.args} command on"
2670 )
2671 cmd.command_policies = command_policies
2672 if len(target_nodes) > 1:
2673 raise RedisClusterException(f"Too many targets for command {cmd.args}")
2674 node = target_nodes[0]
2675 if node.name not in nodes:
2676 nodes[node.name] = (node, [])
2677 nodes[node.name][1].append(cmd)
2679 # Start timing for observability
2680 start_time = time.monotonic()
2682 errors = await asyncio.gather(
2683 *(
2684 asyncio.create_task(node[0].execute_pipeline(node[1]))
2685 for node in nodes.values()
2686 )
2687 )
2689 # Record operation duration for each node
2690 for node_name, (node, commands) in nodes.items():
2691 # Find the first error in this node's commands, if any
2692 node_error = None
2693 for cmd in commands:
2694 if isinstance(cmd.result, Exception):
2695 node_error = cmd.result
2696 break
2698 db = node.connection_kwargs.get("db", 0)
2699 await record_operation_duration(
2700 command_name="PIPELINE",
2701 duration_seconds=time.monotonic() - start_time,
2702 server_address=node.host,
2703 server_port=node.port,
2704 db_namespace=str(db) if db is not None else None,
2705 error=node_error,
2706 )
2708 if any(errors):
2709 if allow_redirections:
2710 # send each errored command individually
2711 for cmd in todo:
2712 if isinstance(cmd.result, (TryAgainError, MovedError, AskError)):
2713 try:
2714 cmd.result = client._policies_callback_mapping[
2715 cmd.command_policies.response_policy
2716 ](await client.execute_command(*cmd.args, **cmd.kwargs))
2717 except Exception as e:
2718 cmd.result = e
2720 if raise_on_error:
2721 for cmd in todo:
2722 result = cmd.result
2723 if isinstance(result, Exception):
2724 command = " ".join(map(safe_str, cmd.args))
2725 msg = (
2726 f"Command # {cmd.position + 1} "
2727 f"({truncate_text(command)}) "
2728 f"of pipeline caused error: {result.args}"
2729 )
2730 result.args = (msg,) + result.args[1:]
2731 raise result
2733 default_cluster_node = client.get_default_node()
2735 # Check whether the default node was used. In some cases,
2736 # 'client.get_default_node()' may return None. The check below
2737 # prevents a potential AttributeError.
2738 if default_cluster_node is not None:
2739 default_node = nodes.get(default_cluster_node.name)
2740 if default_node is not None:
2741 # This pipeline execution used the default node, check if we need
2742 # to replace it.
2743 # Note: when the error is raised we'll reset the default node in the
2744 # caller function.
2745 for cmd in default_node[1]:
2746 # Check if it has a command that failed with a relevant
2747 # exception
2748 if type(cmd.result) in RedisCluster.ERRORS_ALLOW_RETRY:
2749 client.replace_default_node()
2750 break
2752 return [cmd.result for cmd in stack]
2754 async def reset(self):
2755 """
2756 Reset back to empty pipeline.
2757 """
2758 self._command_queue = []
2760 def multi(self):
2761 raise RedisClusterException(
2762 "method multi() is not supported outside of transactional context"
2763 )
2765 async def watch(self, *names):
2766 raise RedisClusterException(
2767 "method watch() is not supported outside of transactional context"
2768 )
2770 async def unwatch(self):
2771 raise RedisClusterException(
2772 "method unwatch() is not supported outside of transactional context"
2773 )
2775 async def discard(self):
2776 raise RedisClusterException(
2777 "method discard() is not supported outside of transactional context"
2778 )
2780 async def unlink(self, *names):
2781 if len(names) != 1:
2782 raise RedisClusterException(
2783 "unlinking multiple keys is not implemented in pipeline command"
2784 )
2786 return self.execute_command("UNLINK", names[0])
2789class TransactionStrategy(AbstractStrategy):
2790 NO_SLOTS_COMMANDS = {"UNWATCH"}
2791 IMMEDIATE_EXECUTE_COMMANDS = {"WATCH", "UNWATCH"}
2792 UNWATCH_COMMANDS = {"DISCARD", "EXEC", "UNWATCH"}
2793 SLOT_REDIRECT_ERRORS = (AskError, MovedError)
2794 CONNECTION_ERRORS = (
2795 ConnectionError,
2796 OSError,
2797 ClusterDownError,
2798 SlotNotCoveredError,
2799 )
2801 def __init__(self, pipe: ClusterPipeline) -> None:
2802 super().__init__(pipe)
2803 self._explicit_transaction = False
2804 self._watching = False
2805 self._pipeline_slots: Set[int] = set()
2806 self._transaction_node: Optional[ClusterNode] = None
2807 self._transaction_connection: Optional[Connection] = None
2808 self._executing = False
2809 self._retry = copy(self._pipe.cluster_client.retry)
2810 self._retry.update_supported_errors(
2811 RedisCluster.ERRORS_ALLOW_RETRY + self.SLOT_REDIRECT_ERRORS
2812 )
2814 def _get_client_and_connection_for_transaction(
2815 self,
2816 ) -> Tuple[ClusterNode, Connection]:
2817 """
2818 Find a connection for a pipeline transaction.
2820 For running an atomic transaction, watch keys ensure that contents have not been
2821 altered as long as the watch commands for those keys were sent over the same
2822 connection. So once we start watching a key, we fetch a connection to the
2823 node that owns that slot and reuse it.
2824 """
2825 if not self._pipeline_slots:
2826 raise RedisClusterException(
2827 "At least a command with a key is needed to identify a node"
2828 )
2830 node: ClusterNode = self._pipe.cluster_client.nodes_manager.get_node_from_slot(
2831 list(self._pipeline_slots)[0], False
2832 )
2833 self._transaction_node = node
2835 if not self._transaction_connection:
2836 connection: Connection = self._transaction_node.acquire_connection()
2837 self._transaction_connection = connection
2839 return self._transaction_node, self._transaction_connection
2841 def execute_command(self, *args: Union[KeyT, EncodableT], **kwargs: Any) -> "Any":
2842 # Given the limitation of ClusterPipeline sync API, we have to run it in thread.
2843 response = None
2844 error = None
2846 def runner():
2847 nonlocal response
2848 nonlocal error
2849 try:
2850 response = asyncio.run(self._execute_command(*args, **kwargs))
2851 except Exception as e:
2852 error = e
2854 thread = threading.Thread(target=runner)
2855 thread.start()
2856 thread.join()
2858 if error:
2859 raise error
2861 return response
2863 async def _execute_command(
2864 self, *args: Union[KeyT, EncodableT], **kwargs: Any
2865 ) -> Any:
2866 if self._pipe.cluster_client._initialize:
2867 await self._pipe.cluster_client.initialize()
2869 slot_number: Optional[int] = None
2870 if args[0] not in self.NO_SLOTS_COMMANDS:
2871 slot_number = await self._pipe.cluster_client._determine_slot(*args)
2873 if (
2874 self._watching or args[0] in self.IMMEDIATE_EXECUTE_COMMANDS
2875 ) and not self._explicit_transaction:
2876 if args[0] == "WATCH":
2877 self._validate_watch()
2879 if slot_number is not None:
2880 if self._pipeline_slots and slot_number not in self._pipeline_slots:
2881 raise CrossSlotTransactionError(
2882 "Cannot watch or send commands on different slots"
2883 )
2885 self._pipeline_slots.add(slot_number)
2886 elif args[0] not in self.NO_SLOTS_COMMANDS:
2887 raise RedisClusterException(
2888 f"Cannot identify slot number for command: {args[0]},"
2889 "it cannot be triggered in a transaction"
2890 )
2892 return self._immediate_execute_command(*args, **kwargs)
2893 else:
2894 if slot_number is not None:
2895 self._pipeline_slots.add(slot_number)
2897 return super().execute_command(*args, **kwargs)
2899 def _validate_watch(self):
2900 if self._explicit_transaction:
2901 raise RedisError("Cannot issue a WATCH after a MULTI")
2903 self._watching = True
2905 async def _immediate_execute_command(self, *args, **options):
2906 return await self._retry.call_with_retry(
2907 lambda: self._get_connection_and_send_command(*args, **options),
2908 self._reinitialize_on_error,
2909 with_failure_count=True,
2910 )
2912 async def _get_connection_and_send_command(self, *args, **options):
2913 redis_node, connection = self._get_client_and_connection_for_transaction()
2914 # Only disconnect if not watching - disconnecting would lose WATCH state
2915 if not self._watching:
2916 await redis_node.disconnect_if_needed(connection)
2918 # Start timing for observability
2919 start_time = time.monotonic()
2921 try:
2922 response = await self._send_command_parse_response(
2923 connection, redis_node, args[0], *args, **options
2924 )
2926 await record_operation_duration(
2927 command_name=args[0],
2928 duration_seconds=time.monotonic() - start_time,
2929 server_address=connection.host,
2930 server_port=connection.port,
2931 db_namespace=str(connection.db),
2932 )
2934 return response
2935 except Exception as e:
2936 e.connection = connection
2937 await record_operation_duration(
2938 command_name=args[0],
2939 duration_seconds=time.monotonic() - start_time,
2940 server_address=connection.host,
2941 server_port=connection.port,
2942 db_namespace=str(connection.db),
2943 error=e,
2944 )
2945 raise
2947 async def _send_command_parse_response(
2948 self,
2949 connection: Connection,
2950 redis_node: ClusterNode,
2951 command_name,
2952 *args,
2953 **options,
2954 ):
2955 """
2956 Send a command and parse the response
2957 """
2959 await connection.send_command(*args)
2960 output = await redis_node.parse_response(connection, command_name, **options)
2962 if command_name in self.UNWATCH_COMMANDS:
2963 self._watching = False
2964 return output
2966 async def _reinitialize_on_error(self, error, failure_count):
2967 if hasattr(error, "connection"):
2968 await record_error_count(
2969 server_address=error.connection.host,
2970 server_port=error.connection.port,
2971 network_peer_address=error.connection.host,
2972 network_peer_port=error.connection.port,
2973 error_type=error,
2974 retry_attempts=failure_count,
2975 is_internal=True,
2976 )
2978 if self._watching:
2979 if type(error) in self.SLOT_REDIRECT_ERRORS and self._executing:
2980 raise WatchError("Slot rebalancing occurred while watching keys")
2982 if (
2983 type(error) in self.SLOT_REDIRECT_ERRORS
2984 or type(error) in self.CONNECTION_ERRORS
2985 ):
2986 if self._transaction_connection and self._transaction_node:
2987 # Disconnect and release back to pool
2988 await self._transaction_connection.disconnect()
2989 self._transaction_node.release(self._transaction_connection)
2990 self._transaction_connection = None
2992 self._pipe.cluster_client.reinitialize_counter += 1
2993 if (
2994 self._pipe.cluster_client.reinitialize_steps
2995 and self._pipe.cluster_client.reinitialize_counter
2996 % self._pipe.cluster_client.reinitialize_steps
2997 == 0
2998 ):
2999 await self._pipe.cluster_client.nodes_manager.initialize()
3000 self.reinitialize_counter = 0
3001 else:
3002 if isinstance(error, AskError):
3003 await self._pipe.cluster_client.nodes_manager.move_slot(error)
3005 self._executing = False
3007 async def _raise_first_error(self, responses, stack, start_time):
3008 """
3009 Raise the first exception on the stack
3010 """
3011 for r, cmd in zip(responses, stack):
3012 if isinstance(r, Exception):
3013 self._annotate_exception(r, cmd.position + 1, cmd.args)
3015 await record_operation_duration(
3016 command_name="TRANSACTION",
3017 duration_seconds=time.monotonic() - start_time,
3018 server_address=self._transaction_connection.host,
3019 server_port=self._transaction_connection.port,
3020 db_namespace=str(self._transaction_connection.db),
3021 error=r,
3022 )
3024 raise r
3026 def mset_nonatomic(
3027 self, mapping: Mapping[AnyKeyT, EncodableT]
3028 ) -> "ClusterPipeline":
3029 raise NotImplementedError("Method is not supported in transactional context.")
3031 async def execute(
3032 self, raise_on_error: bool = True, allow_redirections: bool = True
3033 ) -> List[Any]:
3034 stack = self._command_queue
3035 if not stack and (not self._watching or not self._pipeline_slots):
3036 return []
3038 return await self._execute_transaction_with_retries(stack, raise_on_error)
3040 async def _execute_transaction_with_retries(
3041 self, stack: List["PipelineCommand"], raise_on_error: bool
3042 ):
3043 return await self._retry.call_with_retry(
3044 lambda: self._execute_transaction(stack, raise_on_error),
3045 lambda error, failure_count: self._reinitialize_on_error(
3046 error, failure_count
3047 ),
3048 with_failure_count=True,
3049 )
3051 async def _execute_transaction(
3052 self, stack: List["PipelineCommand"], raise_on_error: bool
3053 ):
3054 if len(self._pipeline_slots) > 1:
3055 raise CrossSlotTransactionError(
3056 "All keys involved in a cluster transaction must map to the same slot"
3057 )
3059 self._executing = True
3061 redis_node, connection = self._get_client_and_connection_for_transaction()
3062 # Only disconnect if not watching - disconnecting would lose WATCH state
3063 if not self._watching:
3064 await redis_node.disconnect_if_needed(connection)
3066 stack = chain(
3067 [PipelineCommand(0, "MULTI")],
3068 stack,
3069 [PipelineCommand(0, "EXEC")],
3070 )
3071 commands = [c.args for c in stack if EMPTY_RESPONSE not in c.kwargs]
3072 packed_commands = connection.pack_commands(commands)
3074 # Start timing for observability
3075 start_time = time.monotonic()
3077 await connection.send_packed_command(packed_commands)
3078 errors = []
3080 # parse off the response for MULTI
3081 # NOTE: we need to handle ResponseErrors here and continue
3082 # so that we read all the additional command messages from
3083 # the socket
3084 try:
3085 await redis_node.parse_response(connection, "MULTI")
3086 except ResponseError as e:
3087 self._annotate_exception(e, 0, "MULTI")
3088 errors.append(e)
3089 except self.CONNECTION_ERRORS as cluster_error:
3090 self._annotate_exception(cluster_error, 0, "MULTI")
3091 cluster_error.connection = connection
3092 raise
3094 # and all the other commands
3095 for i, command in enumerate(self._command_queue):
3096 if EMPTY_RESPONSE in command.kwargs:
3097 errors.append((i, command.kwargs[EMPTY_RESPONSE]))
3098 else:
3099 try:
3100 _ = await redis_node.parse_response(connection, "_")
3101 except self.SLOT_REDIRECT_ERRORS as slot_error:
3102 self._annotate_exception(slot_error, i + 1, command.args)
3103 errors.append(slot_error)
3104 except self.CONNECTION_ERRORS as cluster_error:
3105 self._annotate_exception(cluster_error, i + 1, command.args)
3106 cluster_error.connection = connection
3107 raise
3108 except ResponseError as e:
3109 self._annotate_exception(e, i + 1, command.args)
3110 errors.append(e)
3112 response = None
3113 # parse the EXEC.
3114 try:
3115 response = await redis_node.parse_response(connection, "EXEC")
3116 except ExecAbortError:
3117 if errors:
3118 raise errors[0]
3119 raise
3121 self._executing = False
3123 # EXEC clears any watched keys
3124 self._watching = False
3126 if response is None:
3127 raise WatchError("Watched variable changed.")
3129 # put any parse errors into the response
3130 for i, e in errors:
3131 response.insert(i, e)
3133 if len(response) != len(self._command_queue):
3134 raise InvalidPipelineStack(
3135 "Unexpected response length for cluster pipeline EXEC."
3136 " Command stack was {} but response had length {}".format(
3137 [c.args[0] for c in self._command_queue], len(response)
3138 )
3139 )
3141 # find any errors in the response and raise if necessary
3142 if raise_on_error or len(errors) > 0:
3143 await self._raise_first_error(
3144 response,
3145 self._command_queue,
3146 start_time,
3147 )
3149 # We have to run response callbacks manually
3150 data = []
3151 for r, cmd in zip(response, self._command_queue):
3152 if not isinstance(r, Exception):
3153 command_name = cmd.args[0]
3154 if command_name in self._pipe.cluster_client.response_callbacks:
3155 r = self._pipe.cluster_client.response_callbacks[command_name](
3156 r, **cmd.kwargs
3157 )
3158 data.append(r)
3160 await record_operation_duration(
3161 command_name="TRANSACTION",
3162 duration_seconds=time.monotonic() - start_time,
3163 server_address=connection.host,
3164 server_port=connection.port,
3165 db_namespace=str(connection.db),
3166 )
3168 return data
3170 async def reset(self):
3171 self._command_queue = []
3173 try:
3174 # make sure to reset the connection state in the event that we
3175 # were watching something
3176 if self._transaction_connection:
3177 try:
3178 if self._watching:
3179 # call this manually since our unwatch or
3180 # immediate_execute_command methods can call reset()
3181 await self._transaction_connection.send_command("UNWATCH")
3182 await self._transaction_connection.read_response()
3183 except self.CONNECTION_ERRORS:
3184 # disconnect will also remove any previous WATCHes
3185 if self._transaction_connection:
3186 await self._transaction_connection.disconnect()
3187 except asyncio.CancelledError:
3188 # Disconnect so any unread UNWATCH reply does not get
3189 # served to the next caller that takes the connection.
3190 if self._transaction_connection:
3191 await self._transaction_connection.disconnect()
3192 raise
3193 else:
3194 # On the happy path, honor lazy reconnect before release.
3195 await self._transaction_node.disconnect_if_needed(
3196 self._transaction_connection
3197 )
3198 finally:
3199 # Always return the connection to the node's free queue, even on
3200 # cancellation, so cancelled resets do not leak pooled
3201 # connections. Detach the reference before releasing so the
3202 # strategy never holds a pointer to a returned connection.
3203 # ClusterNode.release is synchronous, so no shield is required.
3204 if self._transaction_connection and self._transaction_node:
3205 connection, self._transaction_connection = (
3206 self._transaction_connection,
3207 None,
3208 )
3209 self._transaction_node.release(connection)
3210 # clean up the other instance attributes
3211 self._transaction_connection = None
3212 self._transaction_node = None
3213 self._watching = False
3214 self._explicit_transaction = False
3215 self._pipeline_slots = set()
3216 self._executing = False
3218 def multi(self):
3219 if self._explicit_transaction:
3220 raise RedisError("Cannot issue nested calls to MULTI")
3221 if self._command_queue:
3222 raise RedisError(
3223 "Commands without an initial WATCH have already been issued"
3224 )
3225 self._explicit_transaction = True
3227 async def watch(self, *names):
3228 if self._explicit_transaction:
3229 raise RedisError("Cannot issue a WATCH after a MULTI")
3231 return await self.execute_command("WATCH", *names)
3233 async def unwatch(self):
3234 if self._watching:
3235 return await self.execute_command("UNWATCH")
3237 return True
3239 async def discard(self):
3240 await self.reset()
3242 async def unlink(self, *names):
3243 return self.execute_command("UNLINK", *names)
3246class _ClusterNodePoolAdapter(ConnectionPoolInterface):
3247 """Thin adapter exposing the :class:`ConnectionPoolInterface` that
3248 :class:`PubSub` requires, backed by a :class:`ClusterNode`'s own
3249 connection pool.
3251 Connections are acquired from the node via
3252 :meth:`ClusterNode.acquire_connection` and returned via
3253 :meth:`ClusterNode.release`. :meth:`PubSub.aclose` already
3254 disconnects the connection *before* calling :meth:`release`, so the
3255 connection is returned to the node's free-queue in a disconnected
3256 state — guaranteeing that a subscribed socket is never silently
3257 reused for regular commands.
3259 Methods that do not apply to this adapter (the underlying node's
3260 lifecycle is managed by the cluster, not by individual PubSub
3261 instances) are implemented as no-ops so the adapter remains a valid
3262 :class:`ConnectionPoolInterface`.
3263 """
3265 def __init__(self, node: "ClusterNode") -> None:
3266 self._node = node
3267 self.connection_kwargs = node.connection_kwargs
3269 # -- methods used by PubSub ------------------------------------------------
3271 def get_encoder(self) -> Encoder:
3272 return self._node.get_encoder()
3274 async def get_connection(
3275 self, command_name: Optional[str] = None, *keys: Any, **options: Any
3276 ) -> AbstractConnection:
3277 connection = self._node.acquire_connection()
3278 try:
3279 await connection.connect()
3280 except BaseException:
3281 # connect() may fail mid-handshake (e.g. after the TCP socket
3282 # is established but before AUTH/HELLO completes) leaving the
3283 # connection in a partially-connected state. Disconnect before
3284 # returning it to the node's free queue so it is not reused.
3285 await connection.disconnect()
3286 self._node.release(connection)
3287 raise
3288 return connection
3290 async def release(self, connection: AbstractConnection) -> None:
3291 # PubSub.aclose() disconnects the connection before calling
3292 # release(), so it is safe to put it back in the node's free
3293 # queue – it will reconnect lazily on next use.
3294 await self._node.disconnect_if_needed(connection)
3295 self._node.release(connection)
3297 # -- no-op stubs for the rest of ConnectionPoolInterface -------------------
3298 # The node's connections are shared with regular cluster traffic and its
3299 # lifecycle is managed by RedisCluster / NodesManager, so the adapter must
3300 # not reset, disconnect, retry-configure or re-auth them on behalf of a
3301 # single PubSub instance.
3303 def get_protocol(self):
3304 return self.connection_kwargs.get("protocol", None)
3306 def reset(self) -> None:
3307 pass
3309 async def disconnect(self, inuse_connections: bool = True) -> None:
3310 pass
3312 async def aclose(self) -> None:
3313 pass
3315 def set_retry(self, retry: "Retry") -> None:
3316 pass
3318 async def re_auth_callback(self, token: TokenInterface) -> None:
3319 pass
3321 def get_connection_count(self) -> List[Tuple[int, dict]]:
3322 return []
3325def _unregister_slots_cache_listener(
3326 dispatcher_ref: "weakref.ref[EventDispatcher]",
3327 listener: AsyncEventListenerInterface,
3328 event_type: Type[object],
3329) -> None:
3330 # Module-level finalizer callback. Kept free of strong references to the
3331 # owning ClusterPubSub so attaching it via weakref.finalize does not
3332 # extend the pubsub's lifetime.
3333 dispatcher = dispatcher_ref()
3334 if dispatcher is not None:
3335 dispatcher.unregister_listeners({event_type: [listener]})
3338class ClusterPubSubSlotsCacheListener(AsyncEventListenerInterface):
3339 """
3340 Async listener that forwards AsyncAfterSlotsCacheRefreshEvent to a
3341 ClusterPubSub.
3343 Holds a weak reference to the pubsub so it does not keep the instance
3344 alive. Deterministic cleanup of the dispatcher's strong reference to this
3345 listener is performed by a ``weakref.finalize`` attached to the owning
3346 ClusterPubSub in ``ClusterPubSub.__init__``.
3347 """
3349 def __init__(self, pubsub: "ClusterPubSub") -> None:
3350 self._pubsub_ref: "weakref.ref[ClusterPubSub]" = weakref.ref(pubsub)
3352 async def listen(self, event: object) -> None:
3353 pubsub = self._pubsub_ref()
3354 if pubsub is None:
3355 # Race window between pubsub GC and the finalizer running; safe
3356 # no-op, finalizer will remove this listener shortly.
3357 return
3358 try:
3359 await pubsub.on_slots_changed()
3360 except Exception as e:
3361 # Listeners must not break slots-cache refresh; log and continue so
3362 # a single buggy pubsub cannot starve the rest.
3363 logger.exception(
3364 "pubsub %r raised during slots-cache change: %s: %s",
3365 pubsub,
3366 type(e).__name__,
3367 e,
3368 )
3371class ClusterPubSub(PubSub):
3372 """
3373 Async cluster implementation for pub/sub.
3375 IMPORTANT: before using ClusterPubSub, read about the known limitations
3376 with pubsub in Cluster mode and learn how to workaround them:
3377 https://redis.readthedocs.io/en/stable/clustering.html#known-pubsub-limitations
3378 """
3380 def __init__(
3381 self,
3382 redis_cluster: "RedisCluster",
3383 node: Optional["ClusterNode"] = None,
3384 host: Optional[str] = None,
3385 port: Optional[int] = None,
3386 push_handler_func: Optional[Callable] = None,
3387 event_dispatcher: Optional[EventDispatcher] = None,
3388 **kwargs: Any,
3389 ) -> None:
3390 """
3391 When a pubsub instance is created without specifying a node, a single
3392 node will be transparently chosen for the pubsub connection on the
3393 first command execution. The node will be determined by:
3394 1. Hashing the channel name in the request to find its keyslot
3395 2. Selecting a node that handles the keyslot: If read_from_replicas is
3396 set to true or load_balancing_strategy is set, a replica can be selected.
3398 :param redis_cluster: RedisCluster instance
3399 :param node: ClusterNode to connect to
3400 :param host: Host of the node to connect to
3401 :param port: Port of the node to connect to
3402 :param push_handler_func: Optional push handler function
3403 :param event_dispatcher: Optional event dispatcher
3404 :param kwargs: Additional keyword arguments
3405 """
3406 self.node = None
3407 self.set_pubsub_node(redis_cluster, node, host, port)
3409 # Borrow the node's own connection pool via an adapter rather than
3410 # creating a second, detached ConnectionPool for pubsub.
3411 if self.node is not None:
3412 connection_pool = _ClusterNodePoolAdapter(self.node)
3413 else:
3414 connection_pool = None
3416 self.cluster = redis_cluster
3417 self.node_pubsub_mapping: Dict[str, PubSub] = {}
3418 # Reverse index: shard channel (normalized) -> owning node.name. Used to
3419 # route sunsubscribe calls and reconcile subscriptions after slot
3420 # migration / failover.
3421 self._shard_channel_to_node: Dict[Any, str] = {}
3422 # Dedicated lock for shard-subscription bookkeeping. Distinct from
3423 # PubSub.self._lock (which serializes wire I/O on the cluster-level
3424 # connection used by aclose / send_command / regular subscribe) so
3425 # that reconciliation cannot starve those unrelated coroutines
3426 # during long per-channel migrations.
3427 self._shard_state_lock: asyncio.Lock = asyncio.Lock()
3428 # Background tasks created by on_slots_changed; kept to prevent GC.
3429 self._reconcile_tasks: Set[asyncio.Task] = set()
3430 self._pubsubs_generator = self._pubsubs_generator()
3431 if event_dispatcher is None:
3432 self._event_dispatcher = EventDispatcher()
3433 else:
3434 self._event_dispatcher = event_dispatcher
3435 super().__init__(
3436 connection_pool=connection_pool,
3437 encoder=redis_cluster.encoder,
3438 push_handler_func=push_handler_func,
3439 event_dispatcher=self._event_dispatcher,
3440 **kwargs,
3441 )
3442 # Subscribe to slots-cache change notifications so shard subscriptions
3443 # can be reconciled automatically after topology refreshes.
3444 nm_dispatcher = redis_cluster.nodes_manager._event_dispatcher
3445 self._slots_cache_listener = ClusterPubSubSlotsCacheListener(self)
3446 nm_dispatcher.register_listeners(
3447 {AsyncAfterSlotsCacheRefreshEvent: [self._slots_cache_listener]}
3448 )
3449 # Deterministic GC-time cleanup so short-lived pubsubs do not leak
3450 # listeners in the dispatcher when no slots-refresh event ever fires.
3451 weakref.finalize(
3452 self,
3453 _unregister_slots_cache_listener,
3454 weakref.ref(nm_dispatcher),
3455 self._slots_cache_listener,
3456 AsyncAfterSlotsCacheRefreshEvent,
3457 )
3459 def set_pubsub_node(
3460 self,
3461 cluster: "RedisCluster",
3462 node: Optional["ClusterNode"] = None,
3463 host: Optional[str] = None,
3464 port: Optional[int] = None,
3465 ) -> None:
3466 """
3467 The pubsub node will be set according to the passed node, host and port
3468 When none of the node, host, or port are specified - the node is set
3469 to None and will be determined by the keyslot of the channel in the
3470 first command to be executed.
3471 RedisClusterException will be thrown if the passed node does not exist
3472 in the cluster.
3473 If host is passed without port, or vice versa, a DataError will be
3474 thrown.
3475 """
3476 if node is not None:
3477 # node is passed by the user
3478 self._raise_on_invalid_node(cluster, node, node.host, node.port)
3479 pubsub_node = node
3480 elif host is not None and port is not None:
3481 # host and port passed by the user
3482 node = cluster.get_node(host=host, port=port)
3483 self._raise_on_invalid_node(cluster, node, host, port)
3484 pubsub_node = node
3485 elif host is not None or port is not None:
3486 # only one of host and port is specified
3487 raise DataError("Specify both host and port")
3488 else:
3489 # nothing specified by the user
3490 pubsub_node = None
3491 self.node = pubsub_node
3493 def get_pubsub_node(self) -> Optional["ClusterNode"]:
3494 """
3495 Get the node that is being used as the pubsub connection.
3497 :return: The ClusterNode being used for pubsub, or None if not yet determined
3498 """
3499 return self.node
3501 async def _resubscribe_shard_channels(self) -> None:
3502 # A single node can own multiple slot ranges, so a batched
3503 # ``SSUBSCRIBE`` covering every tracked channel would be rejected by
3504 # Redis with a ``CROSSSLOT`` error. Group by hash slot and emit one
3505 # ``SSUBSCRIBE`` per slot.
3506 by_slot: defaultdict[int, dict] = defaultdict(dict)
3507 for k, v in self.shard_channels.items():
3508 by_slot[key_slot(self.encoder.encode(k))][k] = v
3509 for subscriptions in by_slot.values():
3510 await self._resubscribe(subscriptions, self.ssubscribe)
3512 def _get_node_pubsub(self, node: "ClusterNode") -> PubSub:
3513 """Get or create a PubSub instance for the given node."""
3514 try:
3515 return self.node_pubsub_mapping[node.name]
3516 except KeyError:
3517 pubsub = PubSub(
3518 connection_pool=_ClusterNodePoolAdapter(node),
3519 encoder=self.cluster.encoder,
3520 push_handler_func=self.push_handler_func,
3521 event_dispatcher=self._event_dispatcher,
3522 )
3523 # Replay shard subscriptions on reconnect with slot-aware grouping
3524 # so that channels spanning multiple slots owned by this node do
3525 # not trigger a CROSSSLOT error.
3526 pubsub._resubscribe_shard_channels = MethodType(
3527 ClusterPubSub._resubscribe_shard_channels, pubsub
3528 )
3529 self.node_pubsub_mapping[node.name] = pubsub
3530 return pubsub
3532 def _find_node_name_for_pubsub(self, pubsub: PubSub) -> Optional[str]:
3533 for name, candidate in self.node_pubsub_mapping.items():
3534 if candidate is pubsub:
3535 return name
3536 return None
3538 async def _sharded_message_generator(
3539 self, timeout: float = 0.0
3540 ) -> Tuple[Optional[PubSub], Optional[Dict[str, Any]]]:
3541 """Generate messages from shard channels across all nodes."""
3542 for _ in range(len(self.node_pubsub_mapping)):
3543 pubsub = next(self._pubsubs_generator)
3544 # Don't pass ignore_subscribe_messages here - let get_sharded_message
3545 # handle the filtering after processing subscription state changes
3546 message = await pubsub.get_message(
3547 ignore_subscribe_messages=False, timeout=timeout
3548 )
3549 if message is not None:
3550 return pubsub, message
3551 return None, None
3553 def _pubsubs_generator(self) -> Generator[PubSub, None, None]:
3554 """Generator that yields PubSub instances in round-robin fashion."""
3555 while True:
3556 current_nodes = list(self.node_pubsub_mapping.values())
3557 if not current_nodes:
3558 return # Avoid infinite loop when no subscriptions exist
3559 yield from current_nodes
3561 async def get_sharded_message(
3562 self,
3563 ignore_subscribe_messages: bool = False,
3564 timeout: float = 0.0,
3565 target_node: Optional["ClusterNode"] = None,
3566 ) -> Optional[Dict[str, Any]]:
3567 """
3568 Get a message from shard channels.
3570 :param ignore_subscribe_messages: Whether to ignore subscribe messages
3571 :param timeout: Timeout for message retrieval
3572 :param target_node: Specific node to get message from
3573 :return: Message dictionary or None
3574 """
3575 pubsub: Optional[PubSub]
3576 if target_node:
3577 pubsub = self.node_pubsub_mapping.get(target_node.name)
3578 if pubsub:
3579 # Don't pass ignore_subscribe_messages here - let get_sharded_message
3580 # handle the filtering after processing subscription state changes
3581 message = await pubsub.get_message(
3582 ignore_subscribe_messages=False, timeout=timeout
3583 )
3584 else:
3585 message = None
3586 else:
3587 pubsub, message = await self._sharded_message_generator(timeout=timeout)
3589 if message is None:
3590 return None
3591 # Only sunsubscribe mutates cluster-level shard state; bypassing the
3592 # lock on the data-message hot path keeps smessage delivery from
3593 # competing with the reconciliation task for _shard_state_lock.
3594 if str_if_bytes(message["type"]) == "sunsubscribe":
3595 # Serialize state mutation against reinitialize_shard_subscriptions
3596 # (background task). The blocking get_message above intentionally
3597 # runs outside the lock so reconciliation is not stalled by long
3598 # polls.
3599 async with self._shard_state_lock:
3600 if message["channel"] in self.pending_unsubscribe_shard_channels:
3601 # User-initiated sunsubscribe: drop from cluster-level tracking.
3602 self.pending_unsubscribe_shard_channels.remove(message["channel"])
3603 self.shard_channels.pop(message["channel"], None)
3604 self._shard_channel_to_node.pop(message["channel"], None)
3605 # Drop the per-node pubsub that delivered the confirmation once
3606 # it no longer holds any shard subscriptions, regardless of
3607 # whether the sunsubscribe was user-initiated or driven by
3608 # slot-migration reconciliation (_migrate_shard_channel, which
3609 # intentionally does not add the channel to
3610 # pending_unsubscribe_shard_channels). This releases the
3611 # dedicated connection that would otherwise linger.
3612 # Identifying the receiving pubsub directly (rather than via
3613 # the cluster's current slot map) is required after slot
3614 # migration, where the channel's owner is no longer the node
3615 # that received our original SSUBSCRIBE.
3616 if pubsub is not None and not pubsub.subscribed:
3617 name = self._find_node_name_for_pubsub(pubsub)
3618 if name is not None:
3619 try:
3620 await pubsub.aclose()
3621 except Exception:
3622 pass
3623 self.node_pubsub_mapping.pop(name, None)
3625 # Only suppress subscribe/unsubscribe messages, not data messages (smessage)
3626 if str_if_bytes(message["type"]) in ("ssubscribe", "sunsubscribe"):
3627 if self.ignore_subscribe_messages or ignore_subscribe_messages:
3628 return None
3629 return message
3631 async def ssubscribe(
3632 self, *args: ChannelT | Subscription, **kwargs: PubSubHandler
3633 ) -> None:
3634 """
3635 Subscribe to shard channels.
3637 :param args: Channel names or ``Subscription`` objects
3638 :param kwargs: Channel names with handlers
3639 """
3640 s_channels = parse_pubsub_subscriptions(args, kwargs)
3642 # Serialize against reinitialize_shard_subscriptions (background
3643 # task) so the reverse index, shard_channels, and node_pubsub_mapping
3644 # are not mutated concurrently. _migrate_shard_channel below does not
3645 # re-acquire this lock (asyncio.Lock is non-reentrant).
3646 async with self._shard_state_lock:
3647 for s_channel, handler in s_channels.items():
3648 node = self.cluster.get_node_from_key(s_channel)
3649 if not node:
3650 continue
3651 # Lazy re-route: if this channel is already tracked against a
3652 # different node (e.g. after a slot migration), migrate it now
3653 # so the caller's intent is applied on the current owner.
3654 normalized_key = next(iter(self._normalize_keys({s_channel: None})))
3655 old_name = self._shard_channel_to_node.get(normalized_key)
3656 if old_name and old_name != node.name:
3657 # Match PubSub.ssubscribe() dict.update() semantics: the
3658 # caller's newly supplied handler (including None) always
3659 # overrides any previously registered handler.
3660 await self._migrate_shard_channel(
3661 normalized_key,
3662 handler,
3663 old_name,
3664 node,
3665 )
3666 continue
3667 pubsub = self._get_node_pubsub(node)
3668 if handler:
3669 await pubsub.ssubscribe(Subscription(s_channel, handler))
3670 else:
3671 await pubsub.ssubscribe(s_channel)
3672 self.shard_channels.update(pubsub.shard_channels)
3673 self._shard_channel_to_node[normalized_key] = node.name
3674 self.pending_unsubscribe_shard_channels.difference_update(
3675 self._normalize_keys({s_channel: None})
3676 )
3678 async def sunsubscribe(self, *args: Any) -> None:
3679 """
3680 Unsubscribe from shard channels.
3682 :param args: Channel names to unsubscribe from. If empty, unsubscribe from all.
3683 """
3684 if args:
3685 args = list_or_args(args[0], args[1:])
3686 else:
3687 args = list(self.shard_channels.keys())
3689 # Serialize against reinitialize_shard_subscriptions: the reverse
3690 # index and node_pubsub_mapping must not change between the lookup
3691 # and the per-node sunsubscribe call below.
3692 async with self._shard_state_lock:
3693 for s_channel in args:
3694 normalized_key = next(iter(self._normalize_keys({s_channel: None})))
3695 # Route via the reverse index so we unsubscribe on the node
3696 # that actually holds the subscription. After a slot migration
3697 # the cluster's current owner may no longer be that node.
3698 name = self._shard_channel_to_node.get(normalized_key)
3699 if name and name in self.node_pubsub_mapping:
3700 pubsub = self.node_pubsub_mapping[name]
3701 else:
3702 node = self.cluster.get_node_from_key(s_channel)
3703 if not node or node.name not in self.node_pubsub_mapping:
3704 continue
3705 pubsub = self.node_pubsub_mapping[node.name]
3706 await pubsub.sunsubscribe(s_channel)
3707 self.pending_unsubscribe_shard_channels.update(
3708 pubsub.pending_unsubscribe_shard_channels
3709 )
3711 async def reinitialize_shard_subscriptions(self) -> None:
3712 """
3713 Reconcile per-node shard subscriptions against the cluster's current
3714 slot ownership map. For each tracked shard channel whose owning node
3715 has changed (e.g. after CLUSTER SETSLOT / failover), sunsubscribe on
3716 the old node's pubsub and ssubscribe on the new owner's pubsub,
3717 preserving any registered handler.
3718 """
3719 uncovered: list = []
3720 made_progress = False
3721 first_migrate_error: Optional[BaseException] = None
3722 async with self._shard_state_lock:
3723 for channel, handler in list(self.shard_channels.items()):
3724 try:
3725 new_node = self.cluster.get_node_from_key(channel)
3726 except SlotNotCoveredError:
3727 # Slot is transiently uncovered (mid-migration / partial
3728 # topology refresh). Defer this channel so coverable
3729 # siblings still reconcile this pass; we surface the
3730 # error below so the caller (and logs) know not every
3731 # channel was reconciled. Retry happens on the next
3732 # slots-cache change notification.
3733 uncovered.append(channel)
3734 continue
3735 old_name = self._shard_channel_to_node.get(channel)
3736 if old_name == new_node.name:
3737 continue
3738 try:
3739 await self._migrate_shard_channel(
3740 channel, handler, old_name, new_node
3741 )
3742 made_progress = True
3743 except (ConnectionError, TimeoutError, OSError) as e:
3744 # Transient connectivity error while subscribing on the
3745 # new owner (or unsubscribing on the old owner if its
3746 # handler chose to re-raise). Do not abort reconciliation
3747 # for sibling channels: _shard_channel_to_node was not
3748 # advanced for this channel, so the next slots-cache
3749 # change notification will retry it.
3750 logger.warning(
3751 "shard channel %r migration deferred: %s: %s",
3752 channel,
3753 type(e).__name__,
3754 e,
3755 )
3756 if first_migrate_error is None:
3757 first_migrate_error = e
3758 continue
3759 # Garbage-collect per-node pubsubs that no longer hold any
3760 # subscription so their connections are released.
3761 for name, pubsub in list(self.node_pubsub_mapping.items()):
3762 if not pubsub.subscribed:
3763 try:
3764 await pubsub.aclose()
3765 except Exception:
3766 pass
3767 self.node_pubsub_mapping.pop(name, None)
3768 if uncovered:
3769 # Surface the uncovered channels so the caller (and observer
3770 # notification path) knows reconciliation was incomplete. All
3771 # coverable siblings have already been migrated above.
3772 raise SlotNotCoveredError(
3773 f"{len(uncovered)} shard channel(s) left unreconciled; "
3774 f"slot(s) not covered by the cluster: {uncovered!r}"
3775 )
3776 if first_migrate_error is not None and not made_progress:
3777 # Every migration attempted in this pass failed transiently and
3778 # nothing else made progress. Re-raise the first caught error
3779 # (typically the root cause; later failures are often downstream
3780 # symptoms of the same unreachable node) so the task's done-
3781 # callback surfaces a single representative failure through the
3782 # same logger channel used for SlotNotCoveredError. Per-channel
3783 # WARNINGs above preserve the full forensic detail.
3784 raise first_migrate_error
3786 async def _migrate_shard_channel(
3787 self,
3788 channel: Any,
3789 handler: Optional[Callable],
3790 old_name: Optional[str],
3791 new_node: "ClusterNode",
3792 ) -> None:
3793 # Detach from the old per-node pubsub, best-effort: the old node may
3794 # already be unreachable during migration / failover.
3795 if old_name and old_name in self.node_pubsub_mapping:
3796 old_pubsub = self.node_pubsub_mapping[old_name]
3797 try:
3798 await old_pubsub.sunsubscribe(channel)
3799 except (ConnectionError, TimeoutError, OSError):
3800 # redis-py's Connection has already called ``disconnect()``
3801 # before raising (see Connection.read_response /
3802 # send_packed_command with ``disconnect_on_error=True``),
3803 # so ``old_pubsub``'s dedicated socket is gone. Two cases:
3804 #
3805 # 1. The old node is no longer in the cluster topology
3806 # (e.g. removed by failover / topology refresh): no
3807 # reconnect target exists, so ``old_pubsub.subscribed``
3808 # would stay True forever and the end-of-pass GC block
3809 # would skip it. Drop it eagerly so the round-robin
3810 # generator does not keep yielding a dead pubsub that
3811 # produces periodic errors from ``get_sharded_message``.
3812 # 2. The old node is still known (transiently slow /
3813 # unreachable): ``PubSub._execute`` auto-reconnects and
3814 # ``on_connect`` re-subscribes to remaining channels,
3815 # so other subscriptions on the same pubsub recover
3816 # naturally. Leave it alone.
3817 if self.cluster.get_node(node_name=old_name) is None:
3818 try:
3819 await old_pubsub.aclose()
3820 except Exception:
3821 pass
3822 self.node_pubsub_mapping.pop(old_name, None)
3823 # Attach to the new per-node pubsub, preserving the handler. Decode to
3824 # a text key only when we must pass it as a kwarg (handler present).
3825 new_pubsub = self._get_node_pubsub(new_node)
3826 if handler:
3827 await new_pubsub.ssubscribe(Subscription(channel, handler))
3828 else:
3829 await new_pubsub.ssubscribe(channel)
3830 self.shard_channels.update(new_pubsub.shard_channels)
3831 normalized_key = next(iter(self._normalize_keys({channel: None})))
3832 self._shard_channel_to_node[normalized_key] = new_node.name
3833 self.pending_unsubscribe_shard_channels.difference_update(
3834 self._normalize_keys({channel: None})
3835 )
3837 async def on_slots_changed(self) -> None:
3838 # Observer hook invoked by NodesManager after a slots-cache refresh.
3839 # Schedule reconciliation as a separate task so the caller's code
3840 # path (typically MovedError handling in _execute_command) is not
3841 # blocked on the network I/O performed by reinitialize_shard_
3842 # subscriptions. No-op when there are no shard subscriptions to
3843 # reconcile.
3844 if not self.shard_channels:
3845 return
3846 task = asyncio.create_task(self.reinitialize_shard_subscriptions())
3847 self._reconcile_tasks.add(task)
3848 task.add_done_callback(self._reconcile_tasks.discard)
3849 # Consume the task's exception (if any) so Python does not emit a
3850 # "Task exception was never retrieved" warning. reinitialize_shard_
3851 # subscriptions surfaces SlotNotCoveredError when a slot is still
3852 # transiently uncovered; route it through the same logger channel
3853 # as sync ClusterPubSubSlotsCacheListener for consistent observability.
3854 task.add_done_callback(self._log_reconcile_task_exception)
3856 @staticmethod
3857 def _log_reconcile_task_exception(task: "asyncio.Task") -> None:
3858 if task.cancelled():
3859 return
3860 exc = task.exception()
3861 if exc is not None:
3862 logger.error(
3863 "shard subscription reconciliation failed: %r", exc, exc_info=exc
3864 )
3866 def get_redis_connection(self) -> Optional["AbstractConnection"]:
3867 """
3868 Get the Redis connection of the pubsub connected node.
3870 Returns the pubsub's dedicated connection (acquired from its own
3871 connection pool), not from the ClusterNode's connection pool.
3872 This avoids the connection pool resource leak that would occur
3873 if we called node.acquire_connection() without releasing.
3874 """
3875 # Return the pubsub's own dedicated connection, which is acquired
3876 # from self.connection_pool when executing pubsub commands.
3877 # This is safe because it's the connection dedicated to this pubsub
3878 # instance, not a shared pool connection from the ClusterNode.
3879 return self.connection
3881 async def aclose(self) -> None:
3882 """
3883 Disconnect the pubsub connection.
3884 """
3885 # Cancel and gather in-flight reconciliation tasks BEFORE acquiring
3886 # _shard_state_lock. The tasks themselves take that lock inside
3887 # reinitialize_shard_subscriptions; since asyncio.Lock is non-
3888 # reentrant, gathering while holding it would deadlock. Awaiting
3889 # each task with suppressed CancelledError also avoids unhandled-
3890 # exception warnings if the task was created but not yet scheduled.
3891 if self._reconcile_tasks:
3892 tasks = list(self._reconcile_tasks)
3893 for task in tasks:
3894 task.cancel()
3895 await asyncio.gather(*tasks, return_exceptions=True)
3896 # Hold _shard_state_lock across the rest of the teardown so it
3897 # observes the same mutual-exclusion discipline as ssubscribe /
3898 # sunsubscribe / get_sharded_message / reinitialize_shard_
3899 # subscriptions, which all mutate shard_channels,
3900 # _shard_channel_to_node, and node_pubsub_mapping under this lock.
3901 # Without it, super().aclose() rebinds shard_channels and
3902 # pending_unsubscribe_shard_channels in parallel with a concurrent
3903 # user-coroutine mutation that resumes during one of the awaits
3904 # below, silently dropping subscription intent.
3905 async with self._shard_state_lock:
3906 self._reconcile_tasks.clear()
3907 # Close all shard pubsub instances first
3908 for pubsub in self.node_pubsub_mapping.values():
3909 await pubsub.aclose()
3910 # Drop the now-dead per-node pubsubs from the mapping so the
3911 # round-robin in _pubsubs_generator / _sharded_message_generator
3912 # cannot yield them between teardown and re-subscription.
3913 self.node_pubsub_mapping.clear()
3914 # _pubsubs_generator captures node_pubsub_mapping.values() into
3915 # a local list inside ``yield from``; clearing the mapping does
3916 # not reach references already held by that captured snapshot,
3917 # so a generator suspended mid-yield-from would still surface
3918 # the now-aclose()'d per-node pubsubs after re-subscription.
3919 # Recreate it to drop the captured list. type(self) bypasses
3920 # the instance-level self-shadow established at __init__
3921 # (self._pubsubs_generator = self._pubsubs_generator()).
3922 self._pubsubs_generator = type(self)._pubsubs_generator( # type: ignore[method-assign]
3923 self
3924 )
3925 # Let parent handle self.connection disconnect under the lock
3926 # (includes disconnect, release to pool, and clearing
3927 # self.connection)
3928 await super().aclose()
3929 # Clear the reverse index so a reused instance doesn't route
3930 # against stale mappings. super().aclose() has already cleared
3931 # shard_channels.
3932 self._shard_channel_to_node.clear()
3934 def _raise_on_invalid_node(
3935 self,
3936 redis_cluster: "RedisCluster",
3937 node: Optional["ClusterNode"],
3938 host: Optional[str],
3939 port: Optional[int],
3940 ) -> None:
3941 """
3942 Raise a RedisClusterException if the node is None or doesn't exist in
3943 the cluster.
3944 """
3945 if node is None or redis_cluster.get_node(node_name=node.name) is None:
3946 raise RedisClusterException(
3947 f"Node {host}:{port} doesn't exist in the cluster"
3948 )
3950 async def execute_command(self, *args: Any, **kwargs: Any) -> Any:
3951 """
3952 Execute a command on the appropriate cluster node.
3954 Taken code from redis-py and tweaked to make it work within a cluster.
3955 """
3956 # NOTE: don't parse the response in this function -- it could pull a
3957 # legitimate message off the stack if the connection is already
3958 # subscribed to one or more channels
3960 # For shard commands, route to appropriate node
3961 command = args[0].upper() if args else ""
3962 if command in ("SSUBSCRIBE", "SUNSUBSCRIBE", "SPUBLISH"):
3963 if len(args) > 1:
3964 channel = args[1]
3965 node = self.cluster.get_node_from_key(channel)
3966 if node:
3967 pubsub = self._get_node_pubsub(node)
3968 return await pubsub.execute_command(*args, **kwargs)
3970 # For other commands, use the set node or lazily discover one
3971 if self.connection is None:
3972 if self.connection_pool is None:
3973 if len(args) > 1:
3974 # Hash the first channel and get one of the nodes holding
3975 # this slot
3976 channel = args[1]
3977 slot = self.cluster.keyslot(channel)
3978 node = self.cluster.nodes_manager.get_node_from_slot(
3979 slot,
3980 self.cluster.read_from_replicas,
3981 self.cluster.load_balancing_strategy,
3982 )
3983 else:
3984 # Get a random node
3985 node = self.cluster.get_random_node()
3986 self.node = node
3987 self.connection_pool = _ClusterNodePoolAdapter(node)
3989 # Now we have a connection_pool, use parent's execute_command
3990 return await super().execute_command(*args, **kwargs)