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