Hitless upgrade: Adding handshake command to enable the notifications after connection is established (#3735)

Author: elena-kolevska

redis/_parsers/base.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 from abc import ABC
 from asyncio import IncompleteReadError, StreamReader, TimeoutError
@@ -56,6 +57,8 @@
     "Client sent AUTH, but no password is set": AuthenticationError,
 }
 
+logger = logging.getLogger(__name__)
+
 
 class BaseParser(ABC):
     EXCEPTION_CLASSES = {
@@ -199,31 +202,42 @@ def handle_push_response(self, response, **kwargs):
             *_MOVING_MESSAGE,
         ):
             return self.pubsub_push_handler_func(response)
-        if msg_type in _INVALIDATION_MESSAGE and self.invalidation_push_handler_func:
-            return self.invalidation_push_handler_func(response)
-        if msg_type in _MOVING_MESSAGE and self.node_moving_push_handler_func:
-            # TODO: PARSE latest format when available
-            host, port = response[2].decode().split(":")
-            ttl = response[1]
-            id = 1  # Hardcoded value until the notification starts including the id
-            notification = NodeMovingEvent(id, host, port, ttl)
-            return self.node_moving_push_handler_func(notification)
-        if msg_type in _MAINTENANCE_MESSAGES and self.maintenance_push_handler_func:
-            if msg_type in _MIGRATING_MESSAGE:
-                # TODO: PARSE latest format when available
-                ttl = response[1]
-                id = 2  # Hardcoded value until the notification starts including the id
-                notification = NodeMigratingEvent(id, ttl)
-            elif msg_type in _MIGRATED_MESSAGE:
-                # TODO: PARSE latest format when available
-                id = 3  # Hardcoded value until the notification starts including the id
-                notification = NodeMigratedEvent(id)
-            else:
+
+        try:
+            if (
+                msg_type in _INVALIDATION_MESSAGE
+                and self.invalidation_push_handler_func
+            ):
+                return self.invalidation_push_handler_func(response)
+
+            if msg_type in _MOVING_MESSAGE and self.node_moving_push_handler_func:
+                # Expected message format is: MOVING <seq_number> <time> <endpoint>
+                id = response[1]
+                ttl = response[2]
+                host, port = response[3].decode().split(":")
+                notification = NodeMovingEvent(id, host, port, ttl)
+                return self.node_moving_push_handler_func(notification)
+
+            if msg_type in _MAINTENANCE_MESSAGES and self.maintenance_push_handler_func:
                 notification = None
-            if notification is not None:
-                return self.maintenance_push_handler_func(notification)
-            else:
-                return None
+
+                if msg_type in _MIGRATING_MESSAGE:
+                    # Expected message format is: MIGRATING <seq_number> <time> <shard_id-s>
+                    id = response[1]
+                    ttl = response[2]
+                    notification = NodeMigratingEvent(id, ttl)
+                elif msg_type in _MIGRATED_MESSAGE:
+                    id = response[1]
+                    notification = NodeMigratedEvent(id)
+
+                if notification is not None:
+                    return self.maintenance_push_handler_func(notification)
+        except Exception as e:
+            logger.error(
+                "Error handling {} message ({}): {}".format(msg_type, response, e)
+            )
+
+        return None
 
     def set_pubsub_push_handler(self, pubsub_push_handler_func):
         self.pubsub_push_handler_func = pubsub_push_handler_func
@@ -252,34 +266,49 @@ async def handle_pubsub_push_response(self, response):
 
     async def handle_push_response(self, response, **kwargs):
         """Handle push responses asynchronously"""
+
         msg_type = response[0]
         if msg_type not in (
             *_INVALIDATION_MESSAGE,
             *_MAINTENANCE_MESSAGES,
             *_MOVING_MESSAGE,
         ):
             return await self.pubsub_push_handler_func(response)
-        if msg_type in _INVALIDATION_MESSAGE and self.invalidation_push_handler_func:
-            return await self.invalidation_push_handler_func(response)
-        if msg_type in _MOVING_MESSAGE and self.node_moving_push_handler_func:
-            # push notification from enterprise cluster for node moving
-            # TODO: PARSE latest format when available
-            host, port = response[2].split(":")
-            ttl = response[1]
-            id = 1  # Hardcoded value for async parser
-            notification = NodeMovingEvent(id, host, port, ttl)
-            return await self.node_moving_push_handler_func(notification)
-        if msg_type in _MAINTENANCE_MESSAGES and self.maintenance_push_handler_func:
-            if msg_type in _MIGRATING_MESSAGE:
-                # TODO: PARSE latest format when available
-                ttl = response[1]
-                id = 2  # Hardcoded value for async parser
-                notification = NodeMigratingEvent(id, ttl)
-            elif msg_type in _MIGRATED_MESSAGE:
-                # TODO: PARSE latest format when available
-                id = 3  # Hardcoded value for async parser
-                notification = NodeMigratedEvent(id)
-            return await self.maintenance_push_handler_func(notification)
+
+        try:
+            if (
+                msg_type in _INVALIDATION_MESSAGE
+                and self.invalidation_push_handler_func
+            ):
+                return await self.invalidation_push_handler_func(response)
+
+            if msg_type in _MOVING_MESSAGE and self.node_moving_push_handler_func:
+                # push notification from enterprise cluster for node moving
+                id = response[1]
+                ttl = response[2]
+                host, port = response[3].split(":")
+                notification = NodeMovingEvent(id, host, port, ttl)
+                return await self.node_moving_push_handler_func(notification)
+
+            if msg_type in _MAINTENANCE_MESSAGES and self.maintenance_push_handler_func:
+                notification = None
+
+                if msg_type in _MIGRATING_MESSAGE:
+                    id = response[1]
+                    ttl = response[2]
+                    notification = NodeMigratingEvent(id, ttl)
+                elif msg_type in _MIGRATED_MESSAGE:
+                    id = response[1]
+                    notification = NodeMigratedEvent(id)
+
+                if notification is not None:
+                    return await self.maintenance_push_handler_func(notification)
+        except Exception as e:
+            logger.error(
+                "Error handling {} message ({}): {}".format(msg_type, response, e)
+            )
+
+        return None
 
     def set_pubsub_push_handler(self, pubsub_push_handler_func):
         """Set the pubsub push handler function"""

redis/connection.py

@@ -282,6 +282,13 @@ def should_reconnect(self):
         """
         pass
 
+    @abstractmethod
+    def get_resolved_ip(self):
+        """
+        Get resolved ip address for the connection.
+        """
+        pass
+
     @abstractmethod
     def update_current_socket_timeout(self, relax_timeout: Optional[float] = None):
         """
@@ -421,32 +428,16 @@ def __init__(
             parser_class = _RESP3Parser
         self.set_parser(parser_class)
 
-        if maintenance_events_config and maintenance_events_config.enabled:
-            if maintenance_events_pool_handler:
-                maintenance_events_pool_handler.set_connection(self)
-                self._parser.set_node_moving_push_handler(
-                    maintenance_events_pool_handler.handle_event
-                )
-            self._maintenance_event_connection_handler = (
-                MaintenanceEventConnectionHandler(self, maintenance_events_config)
-            )
-            self._parser.set_maintenance_push_handler(
-                self._maintenance_event_connection_handler.handle_event
-            )
+        self.maintenance_events_config = maintenance_events_config
+
+        # Set up maintenance events if enabled
+        self._configure_maintenance_events(
+            maintenance_events_pool_handler,
+            orig_host_address,
+            orig_socket_timeout,
+            orig_socket_connect_timeout,
+        )
 
-            self.orig_host_address = (
-                orig_host_address if orig_host_address else self.host
-            )
-            self.orig_socket_timeout = (
-                orig_socket_timeout if orig_socket_timeout else self.socket_timeout
-            )
-            self.orig_socket_connect_timeout = (
-                orig_socket_connect_timeout
-                if orig_socket_connect_timeout
-                else self.socket_connect_timeout
-            )
-        else:
-            self._maintenance_event_connection_handler = None
         self._should_reconnect = False
         self.maintenance_state = maintenance_state
 
@@ -505,6 +496,46 @@ def set_parser(self, parser_class):
         """
         self._parser = parser_class(socket_read_size=self._socket_read_size)
 
+    def _configure_maintenance_events(
+        self,
+        maintenance_events_pool_handler=None,
+        orig_host_address=None,
+        orig_socket_timeout=None,
+        orig_socket_connect_timeout=None,
+    ):
+        """Enable maintenance events by setting up handlers and storing original connection parameters."""
+        if (
+            not self.maintenance_events_config
+            or not self.maintenance_events_config.enabled
+        ):
+            self._maintenance_event_connection_handler = None
+            return
+
+        # Set up pool handler if available
+        if maintenance_events_pool_handler:
+            self._parser.set_node_moving_push_handler(
+                maintenance_events_pool_handler.handle_event
+            )
+
+        # Set up connection handler
+        self._maintenance_event_connection_handler = MaintenanceEventConnectionHandler(
+            self, self.maintenance_events_config
+        )
+        self._parser.set_maintenance_push_handler(
+            self._maintenance_event_connection_handler.handle_event
+        )
+
+        # Store original connection parameters
+        self.orig_host_address = orig_host_address if orig_host_address else self.host
+        self.orig_socket_timeout = (
+            orig_socket_timeout if orig_socket_timeout else self.socket_timeout
+        )
+        self.orig_socket_connect_timeout = (
+            orig_socket_connect_timeout
+            if orig_socket_connect_timeout
+            else self.socket_connect_timeout
+        )
+
     def set_maintenance_event_pool_handler(
         self, maintenance_event_pool_handler: MaintenanceEventPoolHandler
     ):
@@ -652,6 +683,39 @@ def on_connect_check_health(self, check_health: bool = True):
             ):
                 raise ConnectionError("Invalid RESP version")
 
+        # Send maintenance notifications handshake if RESP3 is active and maintenance events are enabled
+        # and we have a host to determine the endpoint type from
+        if (
+            self.protocol not in [2, "2"]
+            and self.maintenance_events_config
+            and self.maintenance_events_config.enabled
+            and self._maintenance_event_connection_handler
+            and hasattr(self, "host")
+        ):
+            try:
+                endpoint_type = self.maintenance_events_config.get_endpoint_type(
+                    self.host, self
+                )
+                self.send_command(
+                    "CLIENT",
+                    "MAINT_NOTIFICATIONS",
+                    "ON",
+                    "moving-endpoint-type",
+                    endpoint_type.value,
+                    check_health=check_health,
+                )
+                response = self.read_response()
+                if str_if_bytes(response) != "OK":
+                    raise ConnectionError(
+                        "The server doesn't support maintenance notifications"
+                    )
+            except Exception as e:
+                # Log warning but don't fail the connection
+                import logging
+
+                logger = logging.getLogger(__name__)
+                logger.warning(f"Failed to enable maintenance notifications: {e}")
+
         # if a client_name is given, set it
         if self.client_name:
             self.send_command(
@@ -888,6 +952,56 @@ def re_auth(self):
             self.read_response()
             self._re_auth_token = None
 
+    def get_resolved_ip(self) -> Optional[str]:
+        """
+        Extract the resolved IP address from an
+        established connection or resolve it from the host.
+
+        First tries to get the actual IP from the socket (most accurate),
+        then falls back to DNS resolution if needed.
+
+        Args:
+            connection: The connection object to extract the IP from
+
+        Returns:
+            str: The resolved IP address, or None if it cannot be determined
+        """
+
+        # Method 1: Try to get the actual IP from the established socket connection
+        # This is most accurate as it shows the exact IP being used
+        try:
+            if self._sock is not None:
+                peer_addr = self._sock.getpeername()
+                if peer_addr and len(peer_addr) >= 1:
+                    # For TCP sockets, peer_addr is typically (host, port) tuple
+                    # Return just the host part
+                    return peer_addr[0]
+        except (AttributeError, OSError):
+            # Socket might not be connected or getpeername() might fail
+            pass
+
+        # Method 2: Fallback to DNS resolution of the host
+        # This is less accurate but works when socket is not available
+        try:
+            host = getattr(self, "host", "localhost")
+            port = getattr(self, "port", 6379)
+            if host:
+                # Use getaddrinfo to resolve the hostname to IP
+                # This mimics what the connection would do during _connect()
+                addr_info = socket.getaddrinfo(
+                    host, port, socket.AF_UNSPEC, socket.SOCK_STREAM
+                )
+                if addr_info:
+                    # Return the IP from the first result
+                    # addr_info[0] is (family, socktype, proto, canonname, sockaddr)
+                    # sockaddr[0] is the IP address
+                    return addr_info[0][4][0]
+        except (AttributeError, OSError, socket.gaierror):
+            # DNS resolution might fail
+            pass
+
+        return None
+
     @property
     def maintenance_state(self) -> MaintenanceState:
         return self._maintenance_state