1import enum
2import ipaddress
3import logging
4import re
5import threading
6import time
7from abc import ABC, abstractmethod
8from typing import (
9 TYPE_CHECKING,
10 Any,
11 Dict,
12 List,
13 Literal,
14 Mapping,
15 Optional,
16 Union,
17)
18
19from redis.observability.attributes import get_pool_name
20from redis.observability.recorder import (
21 record_connection_handoff,
22 record_connection_relaxed_timeout,
23 record_maint_notification_count,
24)
25from redis.typing import Number
26
27if TYPE_CHECKING:
28 from redis.cluster import MaintNotificationsAbstractRedisCluster
29
30logger = logging.getLogger(__name__)
31
32
33class MaintenanceState(enum.Enum):
34 NONE = "none"
35 MOVING = "moving"
36 MAINTENANCE = "maintenance"
37
38
39class EndpointType(enum.Enum):
40 """Valid endpoint types used in CLIENT MAINT_NOTIFICATIONS command."""
41
42 INTERNAL_IP = "internal-ip"
43 INTERNAL_FQDN = "internal-fqdn"
44 EXTERNAL_IP = "external-ip"
45 EXTERNAL_FQDN = "external-fqdn"
46 NONE = "none"
47
48 def __str__(self):
49 """Return the string value of the enum."""
50 return self.value
51
52
53if TYPE_CHECKING:
54 from redis.asyncio.connection import AsyncMaintNotificationsAbstractConnection
55 from redis.connection import (
56 MaintNotificationsAbstractConnection,
57 MaintNotificationsAbstractConnectionPool,
58 )
59
60
61class MaintenanceNotification(ABC):
62 """
63 Base class for maintenance notifications sent through push messages by Redis server.
64
65 This class provides common functionality for all maintenance notifications including
66 unique identification and TTL (Time-To-Live) functionality.
67
68 Attributes:
69 id (int): Unique identifier for this notification
70 ttl (int): Time-to-live in seconds for this notification
71 creation_time (float): Timestamp when the notification was created/read
72 """
73
74 def __init__(self, id: int, ttl: int):
75 """
76 Initialize a new MaintenanceNotification with unique ID and TTL functionality.
77
78 Args:
79 id (int): Unique identifier for this notification
80 ttl (int): Time-to-live in seconds for this notification
81 """
82 self.id = id
83 self.ttl = ttl
84 self.creation_time = time.monotonic()
85 self.expire_at = self.creation_time + self.ttl
86
87 def is_expired(self) -> bool:
88 """
89 Check if this notification has expired based on its TTL
90 and creation time.
91
92 Returns:
93 bool: True if the notification has expired, False otherwise
94 """
95 return time.monotonic() > (self.creation_time + self.ttl)
96
97 @abstractmethod
98 def __repr__(self) -> str:
99 """
100 Return a string representation of the maintenance notification.
101
102 This method must be implemented by all concrete subclasses.
103
104 Returns:
105 str: String representation of the notification
106 """
107 pass
108
109 @abstractmethod
110 def __eq__(self, other) -> bool:
111 """
112 Compare two maintenance notifications for equality.
113
114 This method must be implemented by all concrete subclasses.
115 Notifications are typically considered equal if they have the same id
116 and are of the same type.
117
118 Args:
119 other: The other object to compare with
120
121 Returns:
122 bool: True if the notifications are equal, False otherwise
123 """
124 pass
125
126 @abstractmethod
127 def __hash__(self) -> int:
128 """
129 Return a hash value for the maintenance notification.
130
131 This method must be implemented by all concrete subclasses to allow
132 instances to be used in sets and as dictionary keys.
133
134 Returns:
135 int: Hash value for the notification
136 """
137 pass
138
139
140class NodeMovingNotification(MaintenanceNotification):
141 """
142 This notification is received when a node is replaced with a new node
143 during cluster rebalancing or maintenance operations.
144 """
145
146 def __init__(
147 self,
148 id: int,
149 new_node_host: Optional[str],
150 new_node_port: Optional[int],
151 ttl: int,
152 ):
153 """
154 Initialize a new NodeMovingNotification.
155
156 Args:
157 id (int): Unique identifier for this notification
158 new_node_host (str): Hostname or IP address of the new replacement node
159 new_node_port (int): Port number of the new replacement node
160 ttl (int): Time-to-live in seconds for this notification
161 """
162 super().__init__(id, ttl)
163 self.new_node_host = new_node_host
164 self.new_node_port = new_node_port
165
166 def __repr__(self) -> str:
167 expiry_time = self.expire_at
168 remaining = max(0, expiry_time - time.monotonic())
169
170 return (
171 f"{self.__class__.__name__}("
172 f"id={self.id}, "
173 f"new_node_host='{self.new_node_host}', "
174 f"new_node_port={self.new_node_port}, "
175 f"ttl={self.ttl}, "
176 f"creation_time={self.creation_time}, "
177 f"expires_at={expiry_time}, "
178 f"remaining={remaining:.1f}s, "
179 f"expired={self.is_expired()}"
180 f")"
181 )
182
183 def __eq__(self, other) -> bool:
184 """
185 Two NodeMovingNotification notifications are considered equal if they have the same
186 id, new_node_host, and new_node_port.
187 """
188 if not isinstance(other, NodeMovingNotification):
189 return False
190 return (
191 self.id == other.id
192 and self.new_node_host == other.new_node_host
193 and self.new_node_port == other.new_node_port
194 )
195
196 def __hash__(self) -> int:
197 """
198 Return a hash value for the notification to allow
199 instances to be used in sets and as dictionary keys.
200
201 Returns:
202 int: Hash value based on notification type class name, id,
203 new_node_host and new_node_port
204 """
205 try:
206 node_port = int(self.new_node_port) if self.new_node_port else None
207 except ValueError:
208 node_port = 0
209
210 return hash(
211 (
212 self.__class__.__name__,
213 int(self.id),
214 str(self.new_node_host),
215 node_port,
216 )
217 )
218
219
220class NodeMigratingNotification(MaintenanceNotification):
221 """
222 Notification for when a Redis cluster node is in the process of migrating slots.
223
224 This notification is received when a node starts migrating its slots to another node
225 during cluster rebalancing or maintenance operations.
226
227 Args:
228 id (int): Unique identifier for this notification
229 ttl (int): Time-to-live in seconds for this notification
230 """
231
232 def __init__(self, id: int, ttl: int):
233 super().__init__(id, ttl)
234
235 def __repr__(self) -> str:
236 expiry_time = self.creation_time + self.ttl
237 remaining = max(0, expiry_time - time.monotonic())
238 return (
239 f"{self.__class__.__name__}("
240 f"id={self.id}, "
241 f"ttl={self.ttl}, "
242 f"creation_time={self.creation_time}, "
243 f"expires_at={expiry_time}, "
244 f"remaining={remaining:.1f}s, "
245 f"expired={self.is_expired()}"
246 f")"
247 )
248
249 def __eq__(self, other) -> bool:
250 """
251 Two NodeMigratingNotification notifications are considered equal if they have the same
252 id and are of the same type.
253 """
254 if not isinstance(other, NodeMigratingNotification):
255 return False
256 return self.id == other.id and type(self) is type(other)
257
258 def __hash__(self) -> int:
259 """
260 Return a hash value for the notification to allow
261 instances to be used in sets and as dictionary keys.
262
263 Returns:
264 int: Hash value based on notification type and id
265 """
266 return hash((self.__class__.__name__, int(self.id)))
267
268
269class NodeMigratedNotification(MaintenanceNotification):
270 """
271 Notification for when a Redis cluster node has completed migrating slots.
272
273 This notification is received when a node has finished migrating all its slots
274 to other nodes during cluster rebalancing or maintenance operations.
275
276 Args:
277 id (int): Unique identifier for this notification
278 """
279
280 DEFAULT_TTL = 5
281
282 def __init__(self, id: int):
283 super().__init__(id, NodeMigratedNotification.DEFAULT_TTL)
284
285 def __repr__(self) -> str:
286 expiry_time = self.creation_time + self.ttl
287 remaining = max(0, expiry_time - time.monotonic())
288 return (
289 f"{self.__class__.__name__}("
290 f"id={self.id}, "
291 f"ttl={self.ttl}, "
292 f"creation_time={self.creation_time}, "
293 f"expires_at={expiry_time}, "
294 f"remaining={remaining:.1f}s, "
295 f"expired={self.is_expired()}"
296 f")"
297 )
298
299 def __eq__(self, other) -> bool:
300 """
301 Two NodeMigratedNotification notifications are considered equal if they have the same
302 id and are of the same type.
303 """
304 if not isinstance(other, NodeMigratedNotification):
305 return False
306 return self.id == other.id and type(self) is type(other)
307
308 def __hash__(self) -> int:
309 """
310 Return a hash value for the notification to allow
311 instances to be used in sets and as dictionary keys.
312
313 Returns:
314 int: Hash value based on notification type and id
315 """
316 return hash((self.__class__.__name__, int(self.id)))
317
318
319class NodeFailingOverNotification(MaintenanceNotification):
320 """
321 Notification for when a Redis cluster node is in the process of failing over.
322
323 This notification is received when a node starts a failover process during
324 cluster maintenance operations or when handling node failures.
325
326 Args:
327 id (int): Unique identifier for this notification
328 ttl (int): Time-to-live in seconds for this notification
329 """
330
331 def __init__(self, id: int, ttl: int):
332 super().__init__(id, ttl)
333
334 def __repr__(self) -> str:
335 expiry_time = self.creation_time + self.ttl
336 remaining = max(0, expiry_time - time.monotonic())
337 return (
338 f"{self.__class__.__name__}("
339 f"id={self.id}, "
340 f"ttl={self.ttl}, "
341 f"creation_time={self.creation_time}, "
342 f"expires_at={expiry_time}, "
343 f"remaining={remaining:.1f}s, "
344 f"expired={self.is_expired()}"
345 f")"
346 )
347
348 def __eq__(self, other) -> bool:
349 """
350 Two NodeFailingOverNotification notifications are considered equal if they have the same
351 id and are of the same type.
352 """
353 if not isinstance(other, NodeFailingOverNotification):
354 return False
355 return self.id == other.id and type(self) is type(other)
356
357 def __hash__(self) -> int:
358 """
359 Return a hash value for the notification to allow
360 instances to be used in sets and as dictionary keys.
361
362 Returns:
363 int: Hash value based on notification type and id
364 """
365 return hash((self.__class__.__name__, int(self.id)))
366
367
368class NodeFailedOverNotification(MaintenanceNotification):
369 """
370 Notification for when a Redis cluster node has completed a failover.
371
372 This notification is received when a node has finished the failover process
373 during cluster maintenance operations or after handling node failures.
374
375 Args:
376 id (int): Unique identifier for this notification
377 """
378
379 DEFAULT_TTL = 5
380
381 def __init__(self, id: int):
382 super().__init__(id, NodeFailedOverNotification.DEFAULT_TTL)
383
384 def __repr__(self) -> str:
385 expiry_time = self.creation_time + self.ttl
386 remaining = max(0, expiry_time - time.monotonic())
387 return (
388 f"{self.__class__.__name__}("
389 f"id={self.id}, "
390 f"ttl={self.ttl}, "
391 f"creation_time={self.creation_time}, "
392 f"expires_at={expiry_time}, "
393 f"remaining={remaining:.1f}s, "
394 f"expired={self.is_expired()}"
395 f")"
396 )
397
398 def __eq__(self, other) -> bool:
399 """
400 Two NodeFailedOverNotification notifications are considered equal if they have the same
401 id and are of the same type.
402 """
403 if not isinstance(other, NodeFailedOverNotification):
404 return False
405 return self.id == other.id and type(self) is type(other)
406
407 def __hash__(self) -> int:
408 """
409 Return a hash value for the notification to allow
410 instances to be used in sets and as dictionary keys.
411
412 Returns:
413 int: Hash value based on notification type and id
414 """
415 return hash((self.__class__.__name__, int(self.id)))
416
417
418class OSSNodeMigratingNotification(MaintenanceNotification):
419 """
420 Notification for when a Redis OSS API client is used and a node is in the process of migrating slots.
421
422 This notification is received when a node starts migrating its slots to another node
423 during cluster rebalancing or maintenance operations.
424
425 Args:
426 id (int): Unique identifier for this notification
427 slots (Optional[List[int]]): List of slots being migrated
428 """
429
430 DEFAULT_TTL = 30
431
432 def __init__(
433 self,
434 id: int,
435 slots: Optional[str] = None,
436 ):
437 super().__init__(id, OSSNodeMigratingNotification.DEFAULT_TTL)
438 self.slots = slots
439
440 def __repr__(self) -> str:
441 expiry_time = self.creation_time + self.ttl
442 remaining = max(0, expiry_time - time.monotonic())
443 return (
444 f"{self.__class__.__name__}("
445 f"id={self.id}, "
446 f"slots={self.slots}, "
447 f"ttl={self.ttl}, "
448 f"creation_time={self.creation_time}, "
449 f"expires_at={expiry_time}, "
450 f"remaining={remaining:.1f}s, "
451 f"expired={self.is_expired()}"
452 f")"
453 )
454
455 def __eq__(self, other) -> bool:
456 """
457 Two OSSNodeMigratingNotification notifications are considered equal if they have the same
458 id and are of the same type.
459 """
460 if not isinstance(other, OSSNodeMigratingNotification):
461 return False
462 return self.id == other.id and type(self) is type(other)
463
464 def __hash__(self) -> int:
465 """
466 Return a hash value for the notification to allow
467 instances to be used in sets and as dictionary keys.
468
469 Returns:
470 int: Hash value based on notification type and id
471 """
472 return hash((self.__class__.__name__, int(self.id)))
473
474
475class OSSNodeMigratedNotification(MaintenanceNotification):
476 """
477 Notification for when a Redis OSS API client is used and a node has completed migrating slots.
478
479 This notification is received when a node has finished migrating all its slots
480 to other nodes during cluster rebalancing or maintenance operations.
481
482 Args:
483 id (int): Unique identifier for this notification
484 nodes_to_slots_mapping (Dict[str, List[Dict[str, str]]]): Map of source node address
485 to list of destination mappings. Each destination mapping is a dict with
486 the destination node address as key and the slot range as value.
487
488 Structure example:
489 {
490 "127.0.0.1:6379": [
491 {"127.0.0.1:6380": "1-100"},
492 {"127.0.0.1:6381": "101-200"}
493 ],
494 "127.0.0.1:6382": [
495 {"127.0.0.1:6383": "201-300"}
496 ]
497 }
498
499 Where:
500 - Key (str): Source node address in "host:port" format
501 - Value (List[Dict[str, str]]): List of destination mappings where each dict
502 contains destination node address as key and slot range as value
503 """
504
505 DEFAULT_TTL = 120
506
507 def __init__(
508 self,
509 id: int,
510 nodes_to_slots_mapping: Dict[str, List[Dict[str, str]]],
511 ):
512 super().__init__(id, OSSNodeMigratedNotification.DEFAULT_TTL)
513 self.nodes_to_slots_mapping = nodes_to_slots_mapping
514
515 def __repr__(self) -> str:
516 expiry_time = self.creation_time + self.ttl
517 remaining = max(0, expiry_time - time.monotonic())
518 return (
519 f"{self.__class__.__name__}("
520 f"id={self.id}, "
521 f"nodes_to_slots_mapping={self.nodes_to_slots_mapping}, "
522 f"ttl={self.ttl}, "
523 f"creation_time={self.creation_time}, "
524 f"expires_at={expiry_time}, "
525 f"remaining={remaining:.1f}s, "
526 f"expired={self.is_expired()}"
527 f")"
528 )
529
530 def __eq__(self, other) -> bool:
531 """
532 Two OSSNodeMigratedNotification notifications are considered equal if they have the same
533 id and are of the same type.
534 """
535 if not isinstance(other, OSSNodeMigratedNotification):
536 return False
537 return self.id == other.id and type(self) is type(other)
538
539 def __hash__(self) -> int:
540 """
541 Return a hash value for the notification to allow
542 instances to be used in sets and as dictionary keys.
543
544 Returns:
545 int: Hash value based on notification type and id
546 """
547 return hash((self.__class__.__name__, int(self.id)))
548
549
550def _is_private_fqdn(host: str) -> bool:
551 """
552 Determine if an FQDN is likely to be internal/private.
553
554 This uses heuristics based on RFC 952 and RFC 1123 standards:
555 - .local domains (RFC 6762 - Multicast DNS)
556 - .internal domains (common internal convention)
557 - Single-label hostnames (no dots)
558 - Common internal TLDs
559
560 Args:
561 host (str): The FQDN to check
562
563 Returns:
564 bool: True if the FQDN appears to be internal/private
565 """
566 host_lower = host.lower().rstrip(".")
567
568 # Single-label hostnames (no dots) are typically internal
569 if "." not in host_lower:
570 return True
571
572 # Common internal/private domain patterns
573 internal_patterns = [
574 r"\.local$", # mDNS/Bonjour domains
575 r"\.internal$", # Common internal convention
576 r"\.corp$", # Corporate domains
577 r"\.lan$", # Local area network
578 r"\.intranet$", # Intranet domains
579 r"\.private$", # Private domains
580 ]
581
582 for pattern in internal_patterns:
583 if re.search(pattern, host_lower):
584 return True
585
586 # If none of the internal patterns match, assume it's external
587 return False
588
589
590notification_types_mapping: dict[type[MaintenanceNotification], str] = {
591 NodeMovingNotification: "MOVING",
592 NodeMigratingNotification: "MIGRATING",
593 NodeMigratedNotification: "MIGRATED",
594 NodeFailingOverNotification: "FAILING_OVER",
595 NodeFailedOverNotification: "FAILED_OVER",
596 OSSNodeMigratingNotification: "SMIGRATING",
597 OSSNodeMigratedNotification: "SMIGRATED",
598}
599
600
601def add_debug_log_for_notification(
602 connection: "MaintNotificationsAbstractConnection",
603 notification: Union[str, MaintenanceNotification],
604):
605 if logger.isEnabledFor(logging.DEBUG):
606 socket_address = None
607 try:
608 socket_address = (
609 connection._sock.getsockname() if connection._sock else None
610 )
611 socket_address = socket_address[1] if socket_address else None
612 except (AttributeError, OSError):
613 pass
614
615 logger.debug(
616 f"Handling maintenance notification: {notification}, "
617 f"with connection: {connection}, connected to ip {connection.get_resolved_ip()}, "
618 f"local socket port: {socket_address}",
619 )
620
621
622class MaintNotificationsConfig:
623 """
624 Configuration class for maintenance notifications handling behaviour. Notifications are received through
625 push notifications.
626
627 This class defines how the Redis client should react to different push notifications
628 such as node moving, migrations, etc. in a Redis cluster.
629
630 """
631
632 def __init__(
633 self,
634 enabled: Union[bool, Literal["auto"]] = "auto",
635 proactive_reconnect: bool = True,
636 relaxed_timeout: Optional[Number] = 10,
637 endpoint_type: Optional[EndpointType] = None,
638 ):
639 """
640 Initialize a new MaintNotificationsConfig.
641
642 Args:
643 enabled (bool | "auto"): Controls maintenance notifications handling behavior.
644 - True: The CLIENT MAINT_NOTIFICATIONS command must succeed during connection setup,
645 otherwise a ResponseError is raised.
646 - "auto": The CLIENT MAINT_NOTIFICATIONS command is attempted but failures are
647 gracefully handled - a warning is logged and normal operation continues.
648 - False: Maintenance notifications are completely disabled.
649 Defaults to "auto".
650 proactive_reconnect (bool): Whether to proactively reconnect when a node is replaced.
651 Defaults to True.
652 relaxed_timeout (Number): The relaxed timeout to use for the connection during maintenance.
653 If -1 is provided - the relaxed timeout is disabled. Defaults to 20.
654 endpoint_type (Optional[EndpointType]): Override for the endpoint type to use in CLIENT MAINT_NOTIFICATIONS.
655 If None, the endpoint type will be automatically determined based on the host and TLS configuration.
656 Defaults to None.
657
658 Raises:
659 ValueError: If endpoint_type is provided but is not a valid endpoint type.
660 """
661 self.enabled = enabled
662 self.relaxed_timeout = relaxed_timeout
663 self.proactive_reconnect = proactive_reconnect
664 self.endpoint_type = endpoint_type
665
666 def __repr__(self) -> str:
667 return (
668 f"{self.__class__.__name__}("
669 f"enabled={self.enabled}, "
670 f"proactive_reconnect={self.proactive_reconnect}, "
671 f"relaxed_timeout={self.relaxed_timeout}, "
672 f"endpoint_type={self.endpoint_type!r}"
673 f")"
674 )
675
676 def is_relaxed_timeouts_enabled(self) -> bool:
677 """
678 Check if the relaxed_timeout is enabled. The '-1' value is used to disable the relaxed_timeout.
679 If relaxed_timeout is set to None, it will make the operation blocking
680 and waiting until any response is received.
681
682 Returns:
683 True if the relaxed_timeout is enabled, False otherwise.
684 """
685 return self.relaxed_timeout != -1
686
687 def get_endpoint_type(
688 self,
689 host: str,
690 connection: "MaintNotificationsAbstractConnection | AsyncMaintNotificationsAbstractConnection",
691 ) -> EndpointType:
692 """
693 Determine the appropriate endpoint type for CLIENT MAINT_NOTIFICATIONS command.
694
695 Logic:
696 1. If endpoint_type is explicitly set, use it
697 2. Otherwise, check the original host from connection.host:
698 - If host is an IP address, use it directly to determine internal-ip vs external-ip
699 - If host is an FQDN, get the resolved IP to determine internal-fqdn vs external-fqdn
700
701 Args:
702 host: User provided hostname to analyze
703 connection: The connection object to analyze for endpoint type determination
704
705 Returns:
706 """
707
708 # If endpoint_type is explicitly set, use it
709 if self.endpoint_type is not None:
710 return self.endpoint_type
711
712 # Check if the host is an IP address
713 try:
714 ip_addr = ipaddress.ip_address(host)
715 # Host is an IP address - use it directly
716 is_private = ip_addr.is_private
717 return EndpointType.INTERNAL_IP if is_private else EndpointType.EXTERNAL_IP
718 except ValueError:
719 # Host is an FQDN - need to check resolved IP to determine internal vs external
720 pass
721
722 # Host is an FQDN, get the resolved IP to determine if it's internal or external
723 resolved_ip = connection.get_resolved_ip()
724
725 if resolved_ip:
726 try:
727 ip_addr = ipaddress.ip_address(resolved_ip)
728 is_private = ip_addr.is_private
729 # Use FQDN types since the original host was an FQDN
730 return (
731 EndpointType.INTERNAL_FQDN
732 if is_private
733 else EndpointType.EXTERNAL_FQDN
734 )
735 except ValueError:
736 # This shouldn't happen since we got the IP from the socket, but fallback
737 pass
738
739 # Final fallback: use heuristics on the FQDN itself
740 is_private = _is_private_fqdn(host)
741 return EndpointType.INTERNAL_FQDN if is_private else EndpointType.EXTERNAL_FQDN
742
743
744_MAINTENANCE_START_NOTIFICATION_TYPES = (
745 NodeMigratingNotification,
746 NodeFailingOverNotification,
747 OSSNodeMigratingNotification,
748)
749_MAINTENANCE_COMPLETED_NOTIFICATION_TYPES = (
750 NodeMigratedNotification,
751 NodeFailedOverNotification,
752 OSSNodeMigratedNotification,
753)
754
755
756def _get_maintenance_notification_type(
757 notification: MaintenanceNotification,
758) -> Optional[int]:
759 if notification.__class__ in _MAINTENANCE_START_NOTIFICATION_TYPES:
760 return 1
761 if notification.__class__ in _MAINTENANCE_COMPLETED_NOTIFICATION_TYPES:
762 return 0
763 return None
764
765
766def _get_maintenance_notification_name(
767 notification: MaintenanceNotification,
768) -> str:
769 return notification_types_mapping.get(notification.__class__, "")
770
771
772def _should_skip_connection_timeout_update(
773 maintenance_state: MaintenanceState,
774 config: MaintNotificationsConfig,
775) -> bool:
776 return (
777 maintenance_state == MaintenanceState.MOVING
778 or not config.is_relaxed_timeouts_enabled()
779 )
780
781
782def _build_moving_connection_kwargs(
783 notification: NodeMovingNotification,
784 config: MaintNotificationsConfig,
785) -> dict[str, Any]:
786 kwargs: dict[str, Any] = {
787 "maintenance_state": MaintenanceState.MOVING,
788 "maintenance_notification_hash": hash(notification),
789 }
790 if notification.new_node_host is not None:
791 # the host is not updated if the new node host is None
792 # this happens when the MOVING push notification does not contain
793 # the new node host - in this case we only update the timeouts
794 kwargs["host"] = notification.new_node_host
795 if config.is_relaxed_timeouts_enabled():
796 kwargs.update(
797 {
798 "socket_timeout": config.relaxed_timeout,
799 "socket_connect_timeout": config.relaxed_timeout,
800 }
801 )
802 return kwargs
803
804
805def _build_moving_cleanup_connection_kwargs(
806 connection_kwargs: Mapping[str, Any],
807 notification_hash: int,
808) -> Optional[dict[str, Any]]:
809 # if the current maintenance_notification_hash in kwargs is not matching the notification
810 # it means there has been a new moving notification after this one
811 # and we don't need to revert the kwargs yet
812 if connection_kwargs.get("maintenance_notification_hash") != notification_hash:
813 return None
814
815 return {
816 "maintenance_state": MaintenanceState.NONE,
817 "maintenance_notification_hash": None,
818 "host": connection_kwargs.get("orig_host_address"),
819 "socket_timeout": connection_kwargs.get("orig_socket_timeout"),
820 "socket_connect_timeout": connection_kwargs.get("orig_socket_connect_timeout"),
821 }
822
823
824class MaintNotificationsPoolHandler:
825 def __init__(
826 self,
827 pool: "MaintNotificationsAbstractConnectionPool",
828 config: MaintNotificationsConfig,
829 ) -> None:
830 self.pool = pool
831 self.config = config
832 self._processed_notifications = set()
833 self._lock = threading.RLock()
834 self.connection = None
835
836 def set_connection(self, connection: "MaintNotificationsAbstractConnection"):
837 self.connection = connection
838
839 def get_handler_for_connection(self):
840 # Copy all data that should be shared between connections
841 # but each connection should have its own pool handler
842 # since each connection can be in a different state
843 copy = MaintNotificationsPoolHandler(self.pool, self.config)
844 copy._processed_notifications = self._processed_notifications
845 copy._lock = self._lock
846 copy.connection = None
847 return copy
848
849 def remove_expired_notifications(self):
850 with self._lock:
851 for notification in tuple(self._processed_notifications):
852 if notification.is_expired():
853 self._processed_notifications.remove(notification)
854
855 def handle_notification(self, notification: MaintenanceNotification):
856 self.remove_expired_notifications()
857
858 if isinstance(notification, NodeMovingNotification):
859 return self.handle_node_moving_notification(notification)
860 else:
861 logger.error(f"Unhandled notification type: {notification}")
862
863 def handle_node_moving_notification(self, notification: NodeMovingNotification):
864 if (
865 not self.config.proactive_reconnect
866 and not self.config.is_relaxed_timeouts_enabled()
867 ):
868 return
869
870 with self._lock:
871 if notification in self._processed_notifications:
872 # nothing to do in the connection pool handling
873 # the notification has already been handled or is expired
874 # just return
875 return
876
877 with self.pool._lock:
878 if logger.isEnabledFor(logging.DEBUG):
879 logger.debug(
880 f"Handling node MOVING notification: {notification}, "
881 f"with connection: {self.connection}, connected to ip "
882 f"{self.connection.get_resolved_ip() if self.connection else None}"
883 )
884 # Get the current connected address - if any
885 # This is the address that is being moved
886 # and we need to handle only connections
887 # connected to the same address
888 moving_address_src = (
889 self.connection.getpeername() if self.connection else None
890 )
891
892 if getattr(self.pool, "set_in_maintenance", False):
893 # Set pool in maintenance mode - executed only if
894 # BlockingConnectionPool is used
895 self.pool.set_in_maintenance(True)
896
897 # Update maintenance state, timeout and optionally host address
898 # connection settings for matching connections
899 self.pool.update_connections_settings(
900 state=MaintenanceState.MOVING,
901 maintenance_notification_hash=hash(notification),
902 relaxed_timeout=self.config.relaxed_timeout,
903 host_address=notification.new_node_host,
904 matching_address=moving_address_src,
905 matching_pattern="connected_address",
906 update_notification_hash=True,
907 include_free_connections=True,
908 )
909
910 if self.config.proactive_reconnect:
911 if notification.new_node_host is not None:
912 self.run_proactive_reconnect(moving_address_src)
913 else:
914 threading.Timer(
915 notification.ttl / 2,
916 self.run_proactive_reconnect,
917 args=(moving_address_src,),
918 ).start()
919
920 # Update config for new connections:
921 # Set state to MOVING
922 # update host
923 # if relax timeouts are enabled - update timeouts
924 self.pool.update_connection_kwargs(
925 **_build_moving_connection_kwargs(notification, self.config)
926 )
927
928 if getattr(self.pool, "set_in_maintenance", False):
929 self.pool.set_in_maintenance(False)
930
931 threading.Timer(
932 notification.ttl,
933 self.handle_node_moved_notification,
934 args=(notification,),
935 ).start()
936
937 record_connection_handoff(
938 pool_name=get_pool_name(self.pool),
939 )
940
941 self._processed_notifications.add(notification)
942
943 def run_proactive_reconnect(self, moving_address_src: Optional[str] = None):
944 """
945 Run proactive reconnect for the pool.
946 Active connections are marked for reconnect after they complete the current command.
947 Inactive connections are disconnected and will be connected on next use.
948 """
949 with self._lock:
950 with self.pool._lock:
951 # take care for the active connections in the pool
952 # mark them for reconnect after they complete the current command
953 self.pool.update_active_connections_for_reconnect(
954 moving_address_src=moving_address_src,
955 )
956 # take care for the inactive connections in the pool
957 # delete them and create new ones
958 self.pool.disconnect_free_connections(
959 moving_address_src=moving_address_src,
960 )
961
962 def handle_node_moved_notification(self, notification: NodeMovingNotification):
963 """
964 Handle the cleanup after a node moving notification expires.
965 """
966 notification_hash = hash(notification)
967
968 with self._lock:
969 if logger.isEnabledFor(logging.DEBUG):
970 logger.debug(
971 f"Reverting temporary changes related to notification: {notification}, "
972 f"with connection: {self.connection}, connected to ip "
973 f"{self.connection.get_resolved_ip() if self.connection else None}"
974 )
975 kwargs = _build_moving_cleanup_connection_kwargs(
976 self.pool.connection_kwargs, notification_hash
977 )
978 if kwargs is not None:
979 self.pool.update_connection_kwargs(**kwargs)
980
981 with self.pool._lock:
982 reset_relaxed_timeout = self.config.is_relaxed_timeouts_enabled()
983 reset_host_address = self.config.proactive_reconnect
984
985 self.pool.update_connections_settings(
986 relaxed_timeout=-1,
987 state=MaintenanceState.NONE,
988 maintenance_notification_hash=None,
989 matching_notification_hash=notification_hash,
990 matching_pattern="notification_hash",
991 update_notification_hash=True,
992 reset_relaxed_timeout=reset_relaxed_timeout,
993 reset_host_address=reset_host_address,
994 include_free_connections=True,
995 )
996
997
998class MaintNotificationsConnectionHandler:
999 # 1 = "starting maintenance" notifications, 0 = "completed maintenance" notifications
1000 _NOTIFICATION_TYPES: dict[type["MaintenanceNotification"], int] = {
1001 NodeMigratingNotification: 1,
1002 NodeFailingOverNotification: 1,
1003 OSSNodeMigratingNotification: 1,
1004 NodeMigratedNotification: 0,
1005 NodeFailedOverNotification: 0,
1006 OSSNodeMigratedNotification: 0,
1007 }
1008
1009 def __init__(
1010 self,
1011 connection: "MaintNotificationsAbstractConnection",
1012 config: MaintNotificationsConfig,
1013 ) -> None:
1014 self.connection = connection
1015 self.config = config
1016
1017 def _get_pool_name(self) -> str:
1018 """
1019 Get the pool name from the connection's pool handler.
1020 Falls back to connection representation if pool is not available.
1021 """
1022 pool_handler = getattr(
1023 self.connection, "_maint_notifications_pool_handler", None
1024 )
1025 if pool_handler and getattr(pool_handler, "pool", None):
1026 return get_pool_name(pool_handler.pool)
1027 # Fallback for standalone connections without a pool
1028 return repr(self.connection)
1029
1030 def handle_notification(self, notification: MaintenanceNotification):
1031 # get the notification type by checking its class
1032 # 1 for start, 0 for end notification type, None for unknown
1033 notification_type = _get_maintenance_notification_type(notification)
1034 maint_notification = _get_maintenance_notification_name(notification)
1035
1036 record_maint_notification_count(
1037 server_address=self.connection.host,
1038 server_port=self.connection.port,
1039 network_peer_address=self.connection.host,
1040 network_peer_port=self.connection.port,
1041 maint_notification=maint_notification,
1042 )
1043
1044 if notification_type is None:
1045 logger.error(f"Unhandled notification type: {notification}")
1046 return
1047
1048 if notification_type:
1049 self.handle_maintenance_start_notification(
1050 MaintenanceState.MAINTENANCE, notification
1051 )
1052 else:
1053 self.handle_maintenance_completed_notification(notification=notification)
1054
1055 def handle_maintenance_start_notification(
1056 self, maintenance_state: MaintenanceState, notification: MaintenanceNotification
1057 ):
1058 add_debug_log_for_notification(self.connection, notification)
1059
1060 if _should_skip_connection_timeout_update(
1061 self.connection.maintenance_state, self.config
1062 ):
1063 return
1064
1065 self.connection.maintenance_state = maintenance_state
1066 self.connection.set_tmp_settings(
1067 tmp_relaxed_timeout=self.config.relaxed_timeout
1068 )
1069 # extend the timeout for all created connections
1070 self.connection.update_current_socket_timeout(self.config.relaxed_timeout)
1071 if isinstance(notification, OSSNodeMigratingNotification):
1072 # add the notification id to the set of processed start maint notifications
1073 # this is used to skip the unrelaxing of the timeouts if we have received more than
1074 # one start notification before the the final end notification
1075 self.connection.add_maint_start_notification(notification.id)
1076
1077 maint_notification = _get_maintenance_notification_name(notification)
1078 record_connection_relaxed_timeout(
1079 connection_name=self._get_pool_name(),
1080 maint_notification=maint_notification,
1081 relaxed=True,
1082 )
1083
1084 def handle_maintenance_completed_notification(self, **kwargs: Any) -> None:
1085 # Only reset timeouts if state is not MOVING and relaxed timeouts are enabled
1086 if _should_skip_connection_timeout_update(
1087 self.connection.maintenance_state, self.config
1088 ):
1089 return
1090
1091 notification = None
1092 if kwargs.get("notification"):
1093 notification = kwargs["notification"]
1094 add_debug_log_for_notification(
1095 self.connection, notification if notification else "MAINTENANCE_COMPLETED"
1096 )
1097 self.connection.reset_tmp_settings(reset_relaxed_timeout=True)
1098 # Maintenance completed - reset the connection
1099 # timeouts by providing -1 as the relaxed timeout
1100 self.connection.update_current_socket_timeout(-1)
1101 self.connection.maintenance_state = MaintenanceState.NONE
1102 # reset the sets that keep track of received start maint
1103 # notifications and skipped end maint notifications
1104 self.connection.reset_received_notifications()
1105
1106 if notification:
1107 maint_notification = _get_maintenance_notification_name(notification)
1108 record_connection_relaxed_timeout(
1109 connection_name=self._get_pool_name(),
1110 maint_notification=maint_notification,
1111 relaxed=False,
1112 )
1113
1114
1115class OSSMaintNotificationsHandler:
1116 def __init__(
1117 self,
1118 cluster_client: "MaintNotificationsAbstractRedisCluster",
1119 config: MaintNotificationsConfig,
1120 ) -> None:
1121 self.cluster_client = cluster_client
1122 self.config = config
1123 self._processed_notifications = set()
1124 self._in_progress = set()
1125 self._lock = threading.RLock()
1126
1127 def remove_expired_notifications(self):
1128 with self._lock:
1129 for notification in tuple(self._processed_notifications):
1130 if notification.is_expired():
1131 self._processed_notifications.remove(notification)
1132
1133 def handle_notification(self, notification: MaintenanceNotification):
1134 if isinstance(notification, OSSNodeMigratedNotification):
1135 self.handle_oss_maintenance_completed_notification(notification)
1136 else:
1137 logger.error(f"Unhandled notification type: {notification}")
1138
1139 def handle_oss_maintenance_completed_notification(
1140 self, notification: OSSNodeMigratedNotification
1141 ):
1142 self.remove_expired_notifications()
1143
1144 with self._lock:
1145 if (
1146 notification in self._in_progress
1147 or notification in self._processed_notifications
1148 ):
1149 # we are already handling this notification or it has already been processed
1150 # we should skip in_progress notification since when we reinitialize the cluster
1151 # we execute a CLUSTER SLOTS command that can use a different connection
1152 # that has also has the notification and we don't want to
1153 # process the same notification twice
1154 return
1155
1156 if logger.isEnabledFor(logging.DEBUG):
1157 logger.debug(f"Handling SMIGRATED notification: {notification}")
1158 self._in_progress.add(notification)
1159
1160 try:
1161 # Extract the information about the src and destination nodes that are affected
1162 # by the maintenance. nodes_to_slots_mapping structure:
1163 # {
1164 # "src_host:port": [
1165 # {"dest_host:port": "slot_range"},
1166 # ...
1167 # ],
1168 # ...
1169 # }
1170 additional_startup_nodes_info = []
1171 affected_nodes = set()
1172 for (
1173 src_address,
1174 dest_mappings,
1175 ) in notification.nodes_to_slots_mapping.items():
1176 src_host, src_port = src_address.rsplit(":", 1)
1177 src_node = self.cluster_client.nodes_manager.get_node(
1178 host=src_host, port=int(src_port)
1179 )
1180 if src_node is not None:
1181 affected_nodes.add(src_node)
1182
1183 for dest_mapping in dest_mappings:
1184 for dest_address in dest_mapping.keys():
1185 dest_host, dest_port = dest_address.rsplit(":", 1)
1186 additional_startup_nodes_info.append(
1187 (dest_host, int(dest_port))
1188 )
1189
1190 # Updates the cluster slots cache with the new slots mapping
1191 # This will also update the nodes cache with the new nodes mapping
1192 self.cluster_client.nodes_manager.initialize(
1193 disconnect_startup_nodes_pools=False,
1194 additional_startup_nodes_info=additional_startup_nodes_info,
1195 )
1196
1197 all_nodes = set(affected_nodes)
1198 all_nodes = all_nodes.union(
1199 self.cluster_client.nodes_manager.nodes_cache.values()
1200 )
1201
1202 for current_node in all_nodes:
1203 if current_node.redis_connection is None:
1204 continue
1205 with current_node.redis_connection.connection_pool._lock:
1206 handoff_recorded = False
1207 if current_node in affected_nodes:
1208 # mark for reconnect all in use connections to the node - this will force them to
1209 # disconnect after they complete their current commands
1210 # Some of them might be used by sub sub and we don't know which ones - so we disconnect
1211 # all in flight connections after they are done with current command execution
1212 for conn in current_node.redis_connection.connection_pool._get_in_use_connections():
1213 add_debug_log_for_notification(
1214 conn, "SMIGRATED - mark for reconnect"
1215 )
1216 conn.mark_for_reconnect()
1217
1218 record_connection_handoff(
1219 pool_name=get_pool_name(
1220 current_node.redis_connection.connection_pool
1221 )
1222 )
1223 handoff_recorded = True
1224 else:
1225 if logger.isEnabledFor(logging.DEBUG):
1226 logger.debug(
1227 f"SMIGRATED: Node {current_node.name} not affected by maintenance, "
1228 f"skipping mark for reconnect"
1229 )
1230
1231 if (
1232 current_node
1233 not in self.cluster_client.nodes_manager.nodes_cache.values()
1234 ):
1235 # disconnect all free connections to the node - this node will be dropped
1236 # from the cluster, so we don't need to revert the timeouts
1237 for conn in current_node.redis_connection.connection_pool._get_free_connections():
1238 conn.disconnect()
1239
1240 # Only record handoff if not already recorded for this node
1241 if not handoff_recorded:
1242 record_connection_handoff(
1243 pool_name=get_pool_name(
1244 current_node.redis_connection.connection_pool
1245 )
1246 )
1247
1248 # mark the notification as processed
1249 self._processed_notifications.add(notification)
1250 finally:
1251 # Release the in-progress reservation. On success the notification
1252 # is also in _processed_notifications (so it won't be re-handled);
1253 # on failure it is not, allowing a later redelivery to retry.
1254 self._in_progress.discard(notification)