Refactoring the SMIGRATED flow - the notification is changed to contain the src node address for each slot range movement.

Author: petyaslavova

redis/_parsers/base.py

@@ -11,6 +11,7 @@
     NodeMigratedNotification,
     NodeMigratingNotification,
     NodeMovingNotification,
+    NodesToSlotsMapping,
     OSSNodeMigratedNotification,
     OSSNodeMigratingNotification,
 )
@@ -195,9 +196,18 @@ def parse_oss_maintenance_completed_msg(response):
         # SMIGRATED <seq_number> [<host:port> <slot, range1-range2,...>, ...]
         id = response[1]
         nodes_to_slots_mapping_data = response[2]
-        nodes_to_slots_mapping = {}
-        for node, slots in nodes_to_slots_mapping_data:
-            nodes_to_slots_mapping[safe_str(node)] = safe_str(slots)
+        nodes_to_slots_mapping = []
+        for src_node, node, slots in nodes_to_slots_mapping_data:
+            # Parse the node address to extract host and port
+            src_node_str = safe_str(src_node)
+            node_str = safe_str(node)
+            slots_str = safe_str(slots)
+            # The src_node_address is not provided in the SMIGRATED message,
+            # so we use an empty string as a placeholder
+            mapping = NodesToSlotsMapping(
+                src_node_address=src_node_str, node_address=node_str, slots=slots_str
+            )
+            nodes_to_slots_mapping.append(mapping)
 
         return OSSNodeMigratedNotification(id, nodes_to_slots_mapping)
 

redis/asyncio/cluster.py

@@ -1212,6 +1212,9 @@ def __repr__(self) -> str:
     def __eq__(self, obj: Any) -> bool:
         return isinstance(obj, ClusterNode) and obj.name == self.name
 
+    def __hash__(self) -> int:
+        return hash(self.name)
+
     _DEL_MESSAGE = "Unclosed ClusterNode object"
 
     def __del__(

redis/cluster.py

@@ -1676,6 +1676,9 @@ def __repr__(self):
     def __eq__(self, obj):
         return isinstance(obj, ClusterNode) and obj.name == self.name
 
+    def __hash__(self):
+        return hash(self.name)
+
 
 class LoadBalancingStrategy(Enum):
     ROUND_ROBIN = "round_robin"

redis/connection.py

@@ -440,23 +440,17 @@ def _configure_maintenance_notifications(
         else:
             self._maint_notifications_pool_handler = None
 
+        self._maint_notifications_connection_handler = (
+            MaintNotificationsConnectionHandler(self, self.maint_notifications_config)
+        )
+
         if oss_cluster_maint_notifications_handler:
-            # Extract a reference to a new handler that copies all properties
-            # of the original one and has a different connection reference
-            # This is needed because when we attach the handler to the parser
-            # we need to make sure that the handler has a reference to the
-            # connection that the parser is attached to.
             self._oss_cluster_maint_notifications_handler = (
-                oss_cluster_maint_notifications_handler.get_handler_for_connection()
+                oss_cluster_maint_notifications_handler
             )
-            self._oss_cluster_maint_notifications_handler.set_connection(self)
         else:
             self._oss_cluster_maint_notifications_handler = None
 
-        self._maint_notifications_connection_handler = (
-            MaintNotificationsConnectionHandler(self, self.maint_notifications_config)
-        )
-
         # Set up OSS cluster handler to parser if available
         if self._oss_cluster_maint_notifications_handler:
             parser.set_oss_cluster_maint_push_handler(
@@ -521,21 +515,12 @@ def set_maint_notifications_pool_handler_for_connection(
     def set_maint_notifications_cluster_handler_for_connection(
         self, oss_cluster_maint_notifications_handler: OSSMaintNotificationsHandler
     ):
-        # Deep copy the cluster handler to avoid sharing the same handler
-        # between multiple connections, because otherwise each connection will override
-        # the connection reference and the handler will only hold a reference
-        # to the last connection that was set.
-        maint_notifications_cluster_handler_copy = (
-            oss_cluster_maint_notifications_handler.get_handler_for_connection()
-        )
-
-        maint_notifications_cluster_handler_copy.set_connection(self)
         self._get_parser().set_oss_cluster_maint_push_handler(
-            maint_notifications_cluster_handler_copy.handle_notification
+            oss_cluster_maint_notifications_handler.handle_notification
         )
 
         self._oss_cluster_maint_notifications_handler = (
-            maint_notifications_cluster_handler_copy
+            oss_cluster_maint_notifications_handler
         )
 
         # Update maintenance notification connection handler if it doesn't exist

redis/maint_notifications.py

@@ -5,7 +5,8 @@
 import threading
 import time
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Literal, Optional, Union
 
 from redis.typing import Number
 
@@ -454,6 +455,13 @@ def __hash__(self) -> int:
         return hash((self.__class__.__name__, int(self.id)))
 
 
+@dataclass
+class NodesToSlotsMapping:
+    src_node_address: str
+    node_address: str
+    slots: str
+
+
 class OSSNodeMigratedNotification(MaintenanceNotification):
     """
     Notification for when a Redis OSS API client is used and a node has completed migrating slots.
@@ -463,15 +471,15 @@ class OSSNodeMigratedNotification(MaintenanceNotification):
 
     Args:
         id (int): Unique identifier for this notification
-        nodes_to_slots_mapping (Dict[str, str]): Mapping of node addresses to slots
+        nodes_to_slots_mapping (List[NodesToSlotsMapping]): List of node-to-slots mappings
     """
 
     DEFAULT_TTL = 30
 
     def __init__(
         self,
         id: int,
-        nodes_to_slots_mapping: Dict[str, str],
+        nodes_to_slots_mapping: List[NodesToSlotsMapping],
     ):
         super().__init__(id, OSSNodeMigratedNotification.DEFAULT_TTL)
         self.nodes_to_slots_mapping = nodes_to_slots_mapping
@@ -967,10 +975,6 @@ def __init__(
         self._processed_notifications = set()
         self._in_progress = set()
         self._lock = threading.RLock()
-        self.connection = None
-
-    def set_connection(self, connection: "MaintNotificationsAbstractConnection"):
-        self.connection = connection
 
     def get_handler_for_connection(self):
         # Copy all data that should be shared between connections
@@ -980,7 +984,6 @@ def get_handler_for_connection(self):
         copy._processed_notifications = self._processed_notifications
         copy._in_progress = self._in_progress
         copy._lock = self._lock
-        copy.connection = None
         return copy
 
     def remove_expired_notifications(self):
@@ -1011,55 +1014,56 @@ def handle_oss_maintenance_completed_notification(
                 # that has also has the notification and we don't want to
                 # process the same notification twice
                 return
-            if self.connection is None:
-                logging.error(
-                    "Connection is not set for OSSMaintNotificationsHandler. "
-                    f"Failed to handle notification: {notification}"
-                )
-                return
 
-            logging.debug(
-                f"Handling SMIGRATED notification: {notification} with connection: {self.connection}, connected to ip {self.connection.get_resolved_ip()}"
-            )
+            logging.debug(f"Handling SMIGRATED notification: {notification}")
             self._in_progress.add(notification)
 
-            # get the node to which the connection is connected
-            # before refreshing the cluster topology
-            current_node = self.cluster_client.nodes_manager.get_node(
-                host=self.connection.host, port=self.connection.port
-            )
-
-            # Updates the cluster slots cache with the new slots mapping
-            # This will also update the nodes cache with the new nodes mapping
+            # Extract the information about the src and destination nodes that are affected by the maintenance
             additional_startup_nodes_info = []
-            for node_address, _ in notification.nodes_to_slots_mapping.items():
-                new_node_host, new_node_port = node_address.split(":")
+            affected_nodes = set()
+            for mapping in notification.nodes_to_slots_mapping:
+                new_node_host, new_node_port = mapping.node_address.split(":")
+                src_host, src_port = mapping.src_node_address.split(":")
+                src_node = self.cluster_client.nodes_manager.get_node(
+                    host=src_host, port=src_port
+                )
+                if src_node is not None:
+                    affected_nodes.add(src_node)
+
                 additional_startup_nodes_info.append(
                     (new_node_host, int(new_node_port))
                 )
 
+            # Updates the cluster slots cache with the new slots mapping
+            # This will also update the nodes cache with the new nodes mapping
             self.cluster_client.nodes_manager.initialize(
                 disconnect_startup_nodes_pools=False,
                 additional_startup_nodes_info=additional_startup_nodes_info,
             )
 
-            with current_node.redis_connection.connection_pool._lock:
-                # mark for reconnect all in use connections to the node - this will force them to
-                # disconnect after they complete their current commands
-                # Some of them might be used by sub sub and we don't know which ones - so we disconnect
-                # all in flight connections after they are done with current command execution
-                for conn in current_node.redis_connection.connection_pool._get_in_use_connections():
-                    conn.mark_for_reconnect()
+            all_nodes = self.cluster_client.nodes_manager.nodes_cache.values()
+
+            for current_node in all_nodes:
+                if current_node.redis_connection is None:
+                    continue
+                with current_node.redis_connection.connection_pool._lock:
+                    if current_node in affected_nodes:
+                        # mark for reconnect all in use connections to the node - this will force them to
+                        # disconnect after they complete their current commands
+                        # Some of them might be used by sub sub and we don't know which ones - so we disconnect
+                        # all in flight connections after they are done with current command execution
+                        for conn in current_node.redis_connection.connection_pool._get_in_use_connections():
+                            conn.mark_for_reconnect()
+
+                    # if (
+                    #     current_node
+                    #     not in self.cluster_client.nodes_manager.nodes_cache.values()
+                    # ):
+                    #     # disconnect all free connections to the node - this node will be dropped
+                    #     # from the cluster, so we don't need to revert the timeouts
+                    #     for conn in current_node.redis_connection.connection_pool._get_free_connections():
+                    #         conn.disconnect()
 
-                if (
-                    current_node
-                    not in self.cluster_client.nodes_manager.nodes_cache.values()
-                ):
-                    # disconnect all free connections to the node - this node will be dropped
-                    # from the cluster, so we don't need to revert the timeouts
-                    for conn in current_node.redis_connection.connection_pool._get_free_connections():
-                        conn.disconnect()
-                else:
                     if self.config.is_relaxed_timeouts_enabled():
                         # reset the timeouts for the node to which the connection is connected
                         # Perform check if other maintenance ops are in progress for the same node

tests/maint_notifications/proxy_server_helpers.py

@@ -48,12 +48,13 @@ def oss_maint_notification_to_resp(txt: str) -> str:
                 ">3\r\n"  # Push message with 3 elements
                 f"+{notification}\r\n"  # Element 1: Command
                 f":{seq_id}\r\n"  # Element 2: SeqID
-                f"*{len(hosts_and_slots) // 2}\r\n"  # Element 3: Array of host:port, slots pairs
+                f"*{len(hosts_and_slots) // 3}\r\n"  # Element 3: Array of src_host:src_port, dest_host:dest_port, slots pairs
             )
-            for i in range(0, len(hosts_and_slots), 2):
-                resp += "*2\r\n"
+            for i in range(0, len(hosts_and_slots), 3):
+                resp += "*3\r\n"
                 resp += f"+{hosts_and_slots[i]}\r\n"
                 resp += f"+{hosts_and_slots[i + 1]}\r\n"
+                resp += f"+{hosts_and_slots[i + 2]}\r\n"
         else:
             # SMIGRATING
             # Format: SMIGRATING SeqID slot,range1-range2
@@ -211,20 +212,12 @@ def get_connections(self) -> dict:
 
     def send_notification(
         self,
-        connected_to_port: Union[int, str],
         notification: str,
     ) -> dict:
         """
-        Send a notification to all connections connected to
-        a specific node(identified by port number).
-
-        This method:
-        1. Fetches stats from the interceptor server
-        2. Finds all connection IDs connected to the specified node
-        3. Sends the notification to each connection
+        Send a notification to all connections.
 
         Args:
-            connected_to_port: Port number of the node to send the notification to
             notification: The notification message to send (in RESP format)
 
         Returns:
@@ -233,32 +226,12 @@ def send_notification(
         Example:
             interceptor = ProxyInterceptorHelper(None, "http://localhost:4000")
             result = interceptor.send_notification(
-                "6379",
                 "KjENCiQ0DQpQSU5HDQo="  # PING command in base64
             )
         """
-        # Get stats to find connection IDs for the node
-        stats = self.get_stats()
-
-        # Extract connection IDs for the specified node
-        conn_ids = []
-        for node_key, node_info in stats.items():
-            node_port = node_key.split("@")[1]
-            if int(node_port) == int(connected_to_port):
-                for conn in node_info.get("connections", []):
-                    conn_ids.append(conn["id"])
-
-        if not conn_ids:
-            raise RuntimeError(
-                f"No connections found for node {node_port}. "
-                f"Available nodes: {list(set(c.get('node') for c in stats.get('connections', {}).values()))}"
-            )
-
-        # Send notification to each connection
+        # Send notification to all connections
         results = {}
-        logging.info(f"Sending notification to {len(conn_ids)} connections: {conn_ids}")
-        connections_query = f"connectionIds={','.join(conn_ids)}"
-        url = f"{self.server_url}/send-to-clients?{connections_query}&encoding=base64"
+        url = f"{self.server_url}/send-to-all-clients?encoding=base64"
         # Encode notification to base64
         data = base64.b64encode(notification.encode("utf-8"))
 
@@ -271,8 +244,6 @@ def send_notification(
             results = {"error": str(e)}
 
         return {
-            "node_address": node_port,
-            "connection_ids": conn_ids,
             "results": results,
         }