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

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:
@@ -179,6 +181,25 @@ async def read_response(
 class MaintenanceNotificationsParser:
     """Protocol defining maintenance push notification parsing functionality"""
 
+    @staticmethod
+    def parse_oss_maintenance_start_msg(response):
+        # TODO This format is not the final - will be changed later
+        # Expected message format is: SMIGRATING <seq_number> <src_node> <dest_node> <slots>
+        id = response[1]
+        src_node = response[2]
+        dest_node = response[3]
+        slots = response[4]
+        return OSSNodeMigratingNotification(id, src_node, dest_node, slots)
+
+    @staticmethod
+    def parse_oss_maintenance_completed_msg(response):
+        # TODO This format is not the final - will be changed later
+        # Expected message format is: SMIGRATED <seq_number> <node_address> <slots>
+        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>
@@ -215,12 +236,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[
@@ -246,6 +270,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,
+    ),
 }
 
 
@@ -256,6 +288,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_maintenance_push_handler_func: Optional[Callable] = None
 
     def handle_pubsub_push_response(self, response):
         """Handle pubsub push responses"""
@@ -270,6 +303,7 @@ def handle_push_response(self, response, **kwargs):
             _INVALIDATION_MESSAGE,
             *_MAINTENANCE_MESSAGES,
             _MOVING_MESSAGE,
+            _SMIGRATED_MESSAGE,
         ):
             return self.pubsub_push_handler_func(response)
 
@@ -292,13 +326,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_maintenance_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_maintenance_push_handler_func(notification)
         except Exception as e:
             logger.error(
                 "Error handling {} message ({}): {}".format(msg_type, response, e)
@@ -318,6 +366,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_maintenance_push_handler(self, oss_maintenance_push_handler_func):
+        self.oss_maintenance_push_handler_func = oss_maintenance_push_handler_func
+
 
 class AsyncPushNotificationsParser(Protocol):
     """Protocol defining async RESP3-specific parsing functionality"""
@@ -326,6 +377,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_maintenance_push_handler_func: Optional[Callable[..., Awaitable[None]]] = None
 
     async def handle_pubsub_push_response(self, response):
         """Handle pubsub push responses asynchronously"""
@@ -342,6 +394,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)
 
@@ -366,13 +419,28 @@ 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_maintenance_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_maintenance_push_handler_func(
+                            notification
+                        )
         except Exception as e:
             logger.error(
                 "Error handling {} message ({}): {}".format(msg_type, response, e)
@@ -394,6 +462,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_maintenance_push_handler(self, oss_maintenance_push_handler_func):
+        self.oss_maintenance_push_handler_func = oss_maintenance_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_maintenance_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_maintenance_push_handler_func = None
         self.invalidation_push_handler_func = None
 
     def handle_pubsub_push_response(self, response):

redis/maint_notifications.py

@@ -7,6 +7,7 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, List, Literal, Optional, Union
 
+from redis.cluster import NodesManager
 from redis.typing import Number
 
 
@@ -886,6 +887,7 @@ class MaintNotificationsConnectionHandler:
     _NOTIFICATION_TYPES: dict[type["MaintenanceNotification"], int] = {
         NodeMigratingNotification: 1,
         NodeFailingOverNotification: 1,
+        OSSNodeMigratingNotification: 1,
         NodeMigratedNotification: 0,
         NodeFailedOverNotification: 0,
     }
@@ -939,3 +941,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}")