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