Adding OSS Notification Classes for SMIGRATING and SMIGRATED. Handling of SMIGRATING is completed and covered with tests. (#3849)

* Adding maintenance notifications for OSS API enabled connections

* Adding oss maint notifications handler configurations to parsers. Placeholder for smigrated handler in OSSMaintNotificationsHandler class

* Applying review comments. Finilized the notifications format.

Author: petyaslavova

redis/_parsers/base.py

@@ -11,6 +11,8 @@
     NodeMigratedNotification,
     NodeMigratingNotification,
     NodeMovingNotification,
+    OSSNodeMigratedNotification,
+    OSSNodeMigratingNotification,
 )
 
 if sys.version_info.major >= 3 and sys.version_info.minor >= 11:
@@ -178,16 +180,39 @@ async def read_response(
 class MaintenanceNotificationsParser:
     """Protocol defining maintenance push notification parsing functionality"""
 
+    @staticmethod
+    def parse_oss_maintenance_start_msg(response):
+        # Expected message format is:
+        # SMIGRATING <seq_number> <slot, range1-range2,...>
+        id = response[1]
+        slots = response[2]
+        return OSSNodeMigratingNotification(id, slots)
+
+    @staticmethod
+    def parse_oss_maintenance_completed_msg(response):
+        # Expected message format is:
+        # SMIGRATED <seq_number> <host:port> <slot, range1-range2,...>
+        id = response[1]
+        node_address = response[2]
+        slots = response[3]
+        return OSSNodeMigratedNotification(id, node_address, slots)
+
     @staticmethod
     def parse_maintenance_start_msg(response, notification_type):
         # Expected message format is: <notification_type> <seq_number> <time>
+        # Examples:
+        # MIGRATING 1 10
+        # FAILING_OVER 2 20
         id = response[1]
         ttl = response[2]
         return notification_type(id, ttl)
 
     @staticmethod
     def parse_maintenance_completed_msg(response, notification_type):
         # Expected message format is: <notification_type> <seq_number>
+        # Examples:
+        # MIGRATED 1
+        # FAILED_OVER 2
         id = response[1]
         return notification_type(id)
 
@@ -214,12 +239,15 @@ def parse_moving_msg(response):
 _MIGRATED_MESSAGE = "MIGRATED"
 _FAILING_OVER_MESSAGE = "FAILING_OVER"
 _FAILED_OVER_MESSAGE = "FAILED_OVER"
+_SMIGRATING_MESSAGE = "SMIGRATING"
+_SMIGRATED_MESSAGE = "SMIGRATED"
 
 _MAINTENANCE_MESSAGES = (
     _MIGRATING_MESSAGE,
     _MIGRATED_MESSAGE,
     _FAILING_OVER_MESSAGE,
     _FAILED_OVER_MESSAGE,
+    _SMIGRATING_MESSAGE,
 )
 
 MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING: dict[
@@ -245,6 +273,14 @@ def parse_moving_msg(response):
         NodeMovingNotification,
         MaintenanceNotificationsParser.parse_moving_msg,
     ),
+    _SMIGRATING_MESSAGE: (
+        OSSNodeMigratingNotification,
+        MaintenanceNotificationsParser.parse_oss_maintenance_start_msg,
+    ),
+    _SMIGRATED_MESSAGE: (
+        OSSNodeMigratedNotification,
+        MaintenanceNotificationsParser.parse_oss_maintenance_completed_msg,
+    ),
 }
 
 
@@ -255,6 +291,7 @@ class PushNotificationsParser(Protocol):
     invalidation_push_handler_func: Optional[Callable] = None
     node_moving_push_handler_func: Optional[Callable] = None
     maintenance_push_handler_func: Optional[Callable] = None
+    oss_cluster_maint_push_handler_func: Optional[Callable] = None
 
     def handle_pubsub_push_response(self, response):
         """Handle pubsub push responses"""
@@ -269,6 +306,7 @@ def handle_push_response(self, response, **kwargs):
             _INVALIDATION_MESSAGE,
             *_MAINTENANCE_MESSAGES,
             _MOVING_MESSAGE,
+            _SMIGRATED_MESSAGE,
         ):
             return self.pubsub_push_handler_func(response)
 
@@ -291,13 +329,27 @@ def handle_push_response(self, response, **kwargs):
                 parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
                     msg_type
                 ][1]
-                notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
-                    msg_type
-                ][0]
-                notification = parser_function(response, notification_type)
+                if msg_type == _SMIGRATING_MESSAGE:
+                    notification = parser_function(response)
+                else:
+                    notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
+                        msg_type
+                    ][0]
+                    notification = parser_function(response, notification_type)
 
                 if notification is not None:
                     return self.maintenance_push_handler_func(notification)
+            if (
+                msg_type == _SMIGRATED_MESSAGE
+                and self.oss_cluster_maint_push_handler_func
+            ):
+                parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
+                    msg_type
+                ][1]
+                notification = parser_function(response)
+
+                if notification is not None:
+                    return self.oss_cluster_maint_push_handler_func(notification)
         except Exception as e:
             logger.error(
                 "Error handling {} message ({}): {}".format(msg_type, response, e)
@@ -317,6 +369,9 @@ def set_node_moving_push_handler(self, node_moving_push_handler_func):
     def set_maintenance_push_handler(self, maintenance_push_handler_func):
         self.maintenance_push_handler_func = maintenance_push_handler_func
 
+    def set_oss_cluster_maint_push_handler(self, oss_cluster_maint_push_handler_func):
+        self.oss_cluster_maint_push_handler_func = oss_cluster_maint_push_handler_func
+
 
 class AsyncPushNotificationsParser(Protocol):
     """Protocol defining async RESP3-specific parsing functionality"""
@@ -325,6 +380,7 @@ class AsyncPushNotificationsParser(Protocol):
     invalidation_push_handler_func: Optional[Callable] = None
     node_moving_push_handler_func: Optional[Callable[..., Awaitable[None]]] = None
     maintenance_push_handler_func: Optional[Callable[..., Awaitable[None]]] = None
+    oss_cluster_maint_push_handler_func: Optional[Callable[..., Awaitable[None]]] = None
 
     async def handle_pubsub_push_response(self, response):
         """Handle pubsub push responses asynchronously"""
@@ -341,6 +397,7 @@ async def handle_push_response(self, response, **kwargs):
             _INVALIDATION_MESSAGE,
             *_MAINTENANCE_MESSAGES,
             _MOVING_MESSAGE,
+            _SMIGRATED_MESSAGE,
         ):
             return await self.pubsub_push_handler_func(response)
 
@@ -365,13 +422,26 @@ async def handle_push_response(self, response, **kwargs):
                 parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
                     msg_type
                 ][1]
-                notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
-                    msg_type
-                ][0]
-                notification = parser_function(response, notification_type)
+                if msg_type == _SMIGRATING_MESSAGE:
+                    notification = parser_function(response)
+                else:
+                    notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
+                        msg_type
+                    ][0]
+                    notification = parser_function(response, notification_type)
 
                 if notification is not None:
                     return await self.maintenance_push_handler_func(notification)
+            if (
+                msg_type == _SMIGRATED_MESSAGE
+                and self.oss_cluster_maint_push_handler_func
+            ):
+                parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
+                    msg_type
+                ][1]
+                notification = parser_function(response)
+                if notification is not None:
+                    return await self.oss_cluster_maint_push_handler_func(notification)
         except Exception as e:
             logger.error(
                 "Error handling {} message ({}): {}".format(msg_type, response, e)
@@ -393,6 +463,9 @@ def set_node_moving_push_handler(self, node_moving_push_handler_func):
     def set_maintenance_push_handler(self, maintenance_push_handler_func):
         self.maintenance_push_handler_func = maintenance_push_handler_func
 
+    def set_oss_cluster_maint_push_handler(self, oss_cluster_maint_push_handler_func):
+        self.oss_cluster_maint_push_handler_func = oss_cluster_maint_push_handler_func
+
 
 class _AsyncRESPBase(AsyncBaseParser):
     """Base class for async resp parsing"""

redis/_parsers/hiredis.py

@@ -49,6 +49,7 @@ def __init__(self, socket_read_size):
         self.pubsub_push_handler_func = self.handle_pubsub_push_response
         self.node_moving_push_handler_func = None
         self.maintenance_push_handler_func = None
+        self.oss_cluster_maint_push_handler_func = None
         self.invalidation_push_handler_func = None
         self._hiredis_PushNotificationType = None
 

redis/_parsers/resp3.py

@@ -20,6 +20,7 @@ def __init__(self, socket_read_size):
         self.pubsub_push_handler_func = self.handle_pubsub_push_response
         self.node_moving_push_handler_func = None
         self.maintenance_push_handler_func = None
+        self.oss_cluster_maint_push_handler_func = None
         self.invalidation_push_handler_func = None
 
     def handle_pubsub_push_response(self, response):

redis/maint_notifications.py

@@ -5,10 +5,13 @@
 import threading
 import time
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Literal, Optional, Union
+from typing import TYPE_CHECKING, List, Literal, Optional, Union
 
 from redis.typing import Number
 
+if TYPE_CHECKING:
+    from redis.cluster import NodesManager
+
 
 class MaintenanceState(enum.Enum):
     NONE = "none"
@@ -394,6 +397,125 @@ def __hash__(self) -> int:
         return hash((self.__class__.__name__, int(self.id)))
 
 
+class OSSNodeMigratingNotification(MaintenanceNotification):
+    """
+    Notification for when a Redis OSS API client is used and a node is in the process of migrating slots.
+
+    This notification is received when a node starts migrating its slots to another node
+    during cluster rebalancing or maintenance operations.
+
+    Args:
+        id (int): Unique identifier for this notification
+        slots (Optional[List[int]]): List of slots being migrated
+    """
+
+    DEFAULT_TTL = 30
+
+    def __init__(
+        self,
+        id: int,
+        slots: Optional[List[int]] = None,
+    ):
+        super().__init__(id, OSSNodeMigratingNotification.DEFAULT_TTL)
+        self.slots = slots
+
+    def __repr__(self) -> str:
+        expiry_time = self.creation_time + self.ttl
+        remaining = max(0, expiry_time - time.monotonic())
+        return (
+            f"{self.__class__.__name__}("
+            f"id={self.id}, "
+            f"slots={self.slots}, "
+            f"ttl={self.ttl}, "
+            f"creation_time={self.creation_time}, "
+            f"expires_at={expiry_time}, "
+            f"remaining={remaining:.1f}s, "
+            f"expired={self.is_expired()}"
+            f")"
+        )
+
+    def __eq__(self, other) -> bool:
+        """
+        Two OSSNodeMigratingNotification notifications are considered equal if they have the same
+        id and are of the same type.
+        """
+        if not isinstance(other, OSSNodeMigratingNotification):
+            return False
+        return self.id == other.id and type(self) is type(other)
+
+    def __hash__(self) -> int:
+        """
+        Return a hash value for the notification to allow
+        instances to be used in sets and as dictionary keys.
+
+        Returns:
+            int: Hash value based on notification type and id
+        """
+        return hash((self.__class__.__name__, int(self.id)))
+
+
+class OSSNodeMigratedNotification(MaintenanceNotification):
+    """
+    Notification for when a Redis OSS API client is used and a node has completed migrating slots.
+
+    This notification is received when a node has finished migrating all its slots
+    to other nodes during cluster rebalancing or maintenance operations.
+
+    Args:
+        id (int): Unique identifier for this notification
+        node_address (Optional[str]): Address of the node that has
+                                      completed migration - this is the destination node.
+        slots (Optional[List[int]]): List of slots that have been migrated
+    """
+
+    DEFAULT_TTL = 30
+
+    def __init__(
+        self,
+        id: int,
+        node_address: Optional[str] = None,
+        slots: Optional[List[int]] = None,
+    ):
+        super().__init__(id, OSSNodeMigratedNotification.DEFAULT_TTL)
+        self.node_address = node_address
+        self.slots = slots
+
+    def __repr__(self) -> str:
+        expiry_time = self.creation_time + self.ttl
+        remaining = max(0, expiry_time - time.monotonic())
+        return (
+            f"{self.__class__.__name__}("
+            f"id={self.id}, "
+            f"node_address={self.node_address}, "
+            f"slots={self.slots}, "
+            f"ttl={self.ttl}, "
+            f"creation_time={self.creation_time}, "
+            f"expires_at={expiry_time}, "
+            f"remaining={remaining:.1f}s, "
+            f"expired={self.is_expired()}"
+            f")"
+        )
+
+    def __eq__(self, other) -> bool:
+        """
+        Two OSSNodeMigratedNotification notifications are considered equal if they have the same
+        id and are of the same type.
+        """
+        if not isinstance(other, OSSNodeMigratedNotification):
+            return False
+        return self.id == other.id and type(self) is type(other)
+
+    def __hash__(self) -> int:
+        """
+        Return a hash value for the notification to allow
+        instances to be used in sets and as dictionary keys.
+
+        Returns:
+            int: Hash value based on notification type and id
+        """
+        return hash((self.__class__.__name__, int(self.id)))
+
+
 def _is_private_fqdn(host: str) -> bool:
     """
     Determine if an FQDN is likely to be internal/private.
@@ -755,6 +877,7 @@ class MaintNotificationsConnectionHandler:
     _NOTIFICATION_TYPES: dict[type["MaintenanceNotification"], int] = {
         NodeMigratingNotification: 1,
         NodeFailingOverNotification: 1,
+        OSSNodeMigratingNotification: 1,
         NodeMigratedNotification: 0,
         NodeFailedOverNotification: 0,
     }
@@ -808,3 +931,31 @@ def handle_maintenance_completed_notification(self):
         # timeouts by providing -1 as the relaxed timeout
         self.connection.update_current_socket_timeout(-1)
         self.connection.maintenance_state = MaintenanceState.NONE
+
+
+class OSSMaintNotificationsHandler:
+    def __init__(
+        self, nodes_manager: "NodesManager", config: MaintNotificationsConfig
+    ) -> None:
+        self.nodes_manager = nodes_manager
+        self.config = config
+        self._processed_notifications = set()
+        self._lock = threading.RLock()
+
+    def remove_expired_notifications(self):
+        with self._lock:
+            for notification in tuple(self._processed_notifications):
+                if notification.is_expired():
+                    self._processed_notifications.remove(notification)
+
+    def handle_notification(self, notification: MaintenanceNotification):
+        if isinstance(notification, OSSNodeMigratedNotification):
+            self.handle_oss_maintenance_completed_notification(notification)
+        else:
+            logging.error(f"Unhandled notification type: {notification}")
+
+    def handle_oss_maintenance_completed_notification(
+        self, notification: OSSNodeMigratedNotification
+    ):
+        self.remove_expired_notifications()
+        logging.info(f"Received OSS maintenance completed notification: {notification}")