Add maintenance notifications support for OSS API cluster clients

.gitignore

@@ -27,3 +27,4 @@ docker/stunnel/keys
 /dockers/replica/
 /dockers/sentinel/
 /dockers/redis-stack/
+/experimenting/

docker-compose.yml

@@ -6,6 +6,9 @@ x-client-libs-stack-image: &client-libs-stack-image
 x-client-libs-image: &client-libs-image
   image: "redislabs/client-libs-test:${CLIENT_LIBS_TEST_IMAGE_TAG:-8.4.0}"
 
+networks:
+  redis-net:
+    driver: bridge
 services:
 
   redis:
@@ -108,3 +111,44 @@ services:
       - standalone
       - all-stack
       - all
+
+  redis-proxied:
+    <<: *client-libs-image
+    container_name: redis-proxied
+    ports:
+      - "3000:3000"
+    networks:
+      - redis-net
+    healthcheck:
+      test: ["CMD", "redis-cli", "-p", "3000", "PING"]
+      interval: 10s
+      timeout: 3s
+      retries: 5
+
+  resp-proxy:
+    image: redislabs/client-resp-proxy
+    container_name: resp-proxy
+    environment:
+      LISTEN_HOST: "0.0.0.0"
+      LISTEN_PORT: "15379,15380,15381"
+      TARGET_HOST: "redis-proxied"
+      TARGET_PORT: "3000"
+      API_PORT: "4000"
+      ENABLE_LOGGING: true
+      SIMULATE_CLUSTER: true
+      DEFAULT_INTERCEPTORS: "cluster,hitless,logger"
+
+    ports:
+      - "15379:15379"
+      - "15380:15380"
+      - "15381:15381"
+      - "4000:4000"
+    depends_on:
+      - redis-proxied
+    networks:
+      - redis-net
+    healthcheck:
+      test: ["CMD", "sh", "-c", "wget -qO- http://localhost:4000/stats || exit 1"]
+      interval: 10s
+      timeout: 3s
+      retries: 5

redis/_parsers/base.py

@@ -11,7 +11,10 @@
     NodeMigratedNotification,
     NodeMigratingNotification,
     NodeMovingNotification,
+    OSSNodeMigratedNotification,
+    OSSNodeMigratingNotification,
 )
+from redis.utils import safe_str
 
 if sys.version_info.major >= 3 and sys.version_info.minor >= 11:
     from asyncio import timeout as async_timeout
@@ -178,16 +181,56 @@ 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> [[<src_host:port> <dest_host:port> <slot_range>], ...]
+        id = response[1]
+        nodes_to_slots_mapping_data = response[2]
+        # Build the nodes_to_slots_mapping dict structure:
+        # {
+        #     "src_host:port": [
+        #         {"dest_host:port": "slot_range"},
+        #         ...
+        #     ],
+        #     ...
+        # }
+        nodes_to_slots_mapping = {}
+        for src_node, dest_node, slots in nodes_to_slots_mapping_data:
+            src_node_str = safe_str(src_node)
+            dest_node_str = safe_str(dest_node)
+            slots_str = safe_str(slots)
+
+            if src_node_str not in nodes_to_slots_mapping:
+                nodes_to_slots_mapping[src_node_str] = []
+            nodes_to_slots_mapping[src_node_str].append({dest_node_str: slots_str})
+
+        return OSSNodeMigratedNotification(id, nodes_to_slots_mapping)
+
     @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)
 
@@ -199,9 +242,7 @@ def parse_moving_msg(response):
         if response[3] is None:
             host, port = None, None
         else:
-            value = response[3]
-            if isinstance(value, bytes):
-                value = value.decode()
+            value = safe_str(response[3])
             host, port = value.split(":")
             port = int(port) if port is not None else None
 
@@ -214,12 +255,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 +289,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 +307,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 +322,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 +345,30 @@ 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
+                or self.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:
+                    if self.maintenance_push_handler_func:
+                        self.maintenance_push_handler_func(notification)
+                    if self.oss_cluster_maint_push_handler_func:
+                        self.oss_cluster_maint_push_handler_func(notification)
         except Exception as e:
             logger.error(
                 "Error handling {} message ({}): {}".format(msg_type, response, e)
@@ -317,6 +388,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 +399,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 +416,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 +441,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 +482,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/_parsers/socket.py

@@ -62,7 +62,7 @@ def _read_from_socket(
             sock.settimeout(timeout)
         try:
             while True:
-                data = self._sock.recv(socket_read_size)
+                data = sock.recv(socket_read_size)
                 # an empty string indicates the server shutdown the socket
                 if isinstance(data, bytes) and len(data) == 0:
                     raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)

redis/asyncio/cluster.py

@@ -856,7 +856,7 @@ async def execute_command(self, *args: EncodableT, **kwargs: Any) -> Any:
                     slot = None
                 else:
                     slot = await self._determine_slot(*args)
-                if not slot:
+                if slot is None:
                     command_policies = CommandPolicies()
                 else:
                     command_policies = CommandPolicies(
@@ -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__(
@@ -2188,7 +2191,7 @@ async def _execute(
                             slot = None
                         else:
                             slot = await client._determine_slot(*cmd.args)
-                        if not slot:
+                        if slot is None:
                             command_policies = CommandPolicies()
                         else:
                             command_policies = CommandPolicies(