redis 6.3.0__py3-none-any.whl → 7.0.0__py3-none-any.whl

redis/maint_notifications.py

@@ -0,0 +1,810 @@
+import enum
+import ipaddress
+import logging
+import re
+import threading
+import time
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Literal, Optional, Union
+
+from redis.typing import Number
+
+
+class MaintenanceState(enum.Enum):
+    NONE = "none"
+    MOVING = "moving"
+    MAINTENANCE = "maintenance"
+
+
+class EndpointType(enum.Enum):
+    """Valid endpoint types used in CLIENT MAINT_NOTIFICATIONS command."""
+
+    INTERNAL_IP = "internal-ip"
+    INTERNAL_FQDN = "internal-fqdn"
+    EXTERNAL_IP = "external-ip"
+    EXTERNAL_FQDN = "external-fqdn"
+    NONE = "none"
+
+    def __str__(self):
+        """Return the string value of the enum."""
+        return self.value
+
+
+if TYPE_CHECKING:
+    from redis.connection import (
+        MaintNotificationsAbstractConnection,
+        MaintNotificationsAbstractConnectionPool,
+    )
+
+
+class MaintenanceNotification(ABC):
+    """
+    Base class for maintenance notifications sent through push messages by Redis server.
+
+    This class provides common functionality for all maintenance notifications including
+    unique identification and TTL (Time-To-Live) functionality.
+
+    Attributes:
+        id (int): Unique identifier for this notification
+        ttl (int): Time-to-live in seconds for this notification
+        creation_time (float): Timestamp when the notification was created/read
+    """
+
+    def __init__(self, id: int, ttl: int):
+        """
+        Initialize a new MaintenanceNotification with unique ID and TTL functionality.
+
+        Args:
+            id (int): Unique identifier for this notification
+            ttl (int): Time-to-live in seconds for this notification
+        """
+        self.id = id
+        self.ttl = ttl
+        self.creation_time = time.monotonic()
+        self.expire_at = self.creation_time + self.ttl
+
+    def is_expired(self) -> bool:
+        """
+        Check if this notification has expired based on its TTL
+        and creation time.
+
+        Returns:
+            bool: True if the notification has expired, False otherwise
+        """
+        return time.monotonic() > (self.creation_time + self.ttl)
+
+    @abstractmethod
+    def __repr__(self) -> str:
+        """
+        Return a string representation of the maintenance notification.
+
+        This method must be implemented by all concrete subclasses.
+
+        Returns:
+            str: String representation of the notification
+        """
+        pass
+
+    @abstractmethod
+    def __eq__(self, other) -> bool:
+        """
+        Compare two maintenance notifications for equality.
+
+        This method must be implemented by all concrete subclasses.
+        Notifications are typically considered equal if they have the same id
+        and are of the same type.
+
+        Args:
+            other: The other object to compare with
+
+        Returns:
+            bool: True if the notifications are equal, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    def __hash__(self) -> int:
+        """
+        Return a hash value for the maintenance notification.
+
+        This method must be implemented by all concrete subclasses to allow
+        instances to be used in sets and as dictionary keys.
+
+        Returns:
+            int: Hash value for the notification
+        """
+        pass
+
+
+class NodeMovingNotification(MaintenanceNotification):
+    """
+    This notification is received when a node is replaced with a new node
+    during cluster rebalancing or maintenance operations.
+    """
+
+    def __init__(
+        self,
+        id: int,
+        new_node_host: Optional[str],
+        new_node_port: Optional[int],
+        ttl: int,
+    ):
+        """
+        Initialize a new NodeMovingNotification.
+
+        Args:
+            id (int): Unique identifier for this notification
+            new_node_host (str): Hostname or IP address of the new replacement node
+            new_node_port (int): Port number of the new replacement node
+            ttl (int): Time-to-live in seconds for this notification
+        """
+        super().__init__(id, ttl)
+        self.new_node_host = new_node_host
+        self.new_node_port = new_node_port
+
+    def __repr__(self) -> str:
+        expiry_time = self.expire_at
+        remaining = max(0, expiry_time - time.monotonic())
+
+        return (
+            f"{self.__class__.__name__}("
+            f"id={self.id}, "
+            f"new_node_host='{self.new_node_host}', "
+            f"new_node_port={self.new_node_port}, "
+            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 NodeMovingNotification notifications are considered equal if they have the same
+        id, new_node_host, and new_node_port.
+        """
+        if not isinstance(other, NodeMovingNotification):
+            return False
+        return (
+            self.id == other.id
+            and self.new_node_host == other.new_node_host
+            and self.new_node_port == other.new_node_port
+        )
+
+    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 class name, id,
+            new_node_host and new_node_port
+        """
+        try:
+            node_port = int(self.new_node_port) if self.new_node_port else None
+        except ValueError:
+            node_port = 0
+
+        return hash(
+            (
+                self.__class__.__name__,
+                int(self.id),
+                str(self.new_node_host),
+                node_port,
+            )
+        )
+
+
+class NodeMigratingNotification(MaintenanceNotification):
+    """
+    Notification for when a Redis cluster 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
+        ttl (int): Time-to-live in seconds for this notification
+    """
+
+    def __init__(self, id: int, ttl: int):
+        super().__init__(id, ttl)
+
+    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"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 NodeMigratingNotification notifications are considered equal if they have the same
+        id and are of the same type.
+        """
+        if not isinstance(other, NodeMigratingNotification):
+            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 NodeMigratedNotification(MaintenanceNotification):
+    """
+    Notification for when a Redis cluster 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
+    """
+
+    DEFAULT_TTL = 5
+
+    def __init__(self, id: int):
+        super().__init__(id, NodeMigratedNotification.DEFAULT_TTL)
+
+    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"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 NodeMigratedNotification notifications are considered equal if they have the same
+        id and are of the same type.
+        """
+        if not isinstance(other, NodeMigratedNotification):
+            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 NodeFailingOverNotification(MaintenanceNotification):
+    """
+    Notification for when a Redis cluster node is in the process of failing over.
+
+    This notification is received when a node starts a failover process during
+    cluster maintenance operations or when handling node failures.
+
+    Args:
+        id (int): Unique identifier for this notification
+        ttl (int): Time-to-live in seconds for this notification
+    """
+
+    def __init__(self, id: int, ttl: int):
+        super().__init__(id, ttl)
+
+    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"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 NodeFailingOverNotification notifications are considered equal if they have the same
+        id and are of the same type.
+        """
+        if not isinstance(other, NodeFailingOverNotification):
+            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 NodeFailedOverNotification(MaintenanceNotification):
+    """
+    Notification for when a Redis cluster node has completed a failover.
+
+    This notification is received when a node has finished the failover process
+    during cluster maintenance operations or after handling node failures.
+
+    Args:
+        id (int): Unique identifier for this notification
+    """
+
+    DEFAULT_TTL = 5
+
+    def __init__(self, id: int):
+        super().__init__(id, NodeFailedOverNotification.DEFAULT_TTL)
+
+    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"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 NodeFailedOverNotification notifications are considered equal if they have the same
+        id and are of the same type.
+        """
+        if not isinstance(other, NodeFailedOverNotification):
+            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.
+
+    This uses heuristics based on RFC 952 and RFC 1123 standards:
+    - .local domains (RFC 6762 - Multicast DNS)
+    - .internal domains (common internal convention)
+    - Single-label hostnames (no dots)
+    - Common internal TLDs
+
+    Args:
+        host (str): The FQDN to check
+
+    Returns:
+        bool: True if the FQDN appears to be internal/private
+    """
+    host_lower = host.lower().rstrip(".")
+
+    # Single-label hostnames (no dots) are typically internal
+    if "." not in host_lower:
+        return True
+
+    # Common internal/private domain patterns
+    internal_patterns = [
+        r"\.local$",  # mDNS/Bonjour domains
+        r"\.internal$",  # Common internal convention
+        r"\.corp$",  # Corporate domains
+        r"\.lan$",  # Local area network
+        r"\.intranet$",  # Intranet domains
+        r"\.private$",  # Private domains
+    ]
+
+    for pattern in internal_patterns:
+        if re.search(pattern, host_lower):
+            return True
+
+    # If none of the internal patterns match, assume it's external
+    return False
+
+
+class MaintNotificationsConfig:
+    """
+    Configuration class for maintenance notifications handling behaviour. Notifications are received through
+    push notifications.
+
+    This class defines how the Redis client should react to different push notifications
+    such as node moving, migrations, etc. in a Redis cluster.
+
+    """
+
+    def __init__(
+        self,
+        enabled: Union[bool, Literal["auto"]] = "auto",
+        proactive_reconnect: bool = True,
+        relaxed_timeout: Optional[Number] = 10,
+        endpoint_type: Optional[EndpointType] = None,
+    ):
+        """
+        Initialize a new MaintNotificationsConfig.
+
+        Args:
+            enabled (bool | "auto"): Controls maintenance notifications handling behavior.
+                - True: The CLIENT MAINT_NOTIFICATIONS command must succeed during connection setup,
+                otherwise a ResponseError is raised.
+                - "auto": The CLIENT MAINT_NOTIFICATIONS command is attempted but failures are
+                gracefully handled - a warning is logged and normal operation continues.
+                - False: Maintenance notifications are completely disabled.
+                Defaults to "auto".
+            proactive_reconnect (bool): Whether to proactively reconnect when a node is replaced.
+                Defaults to True.
+            relaxed_timeout (Number): The relaxed timeout to use for the connection during maintenance.
+                If -1 is provided - the relaxed timeout is disabled. Defaults to 20.
+            endpoint_type (Optional[EndpointType]): Override for the endpoint type to use in CLIENT MAINT_NOTIFICATIONS.
+                If None, the endpoint type will be automatically determined based on the host and TLS configuration.
+                Defaults to None.
+
+        Raises:
+            ValueError: If endpoint_type is provided but is not a valid endpoint type.
+        """
+        self.enabled = enabled
+        self.relaxed_timeout = relaxed_timeout
+        self.proactive_reconnect = proactive_reconnect
+        self.endpoint_type = endpoint_type
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"enabled={self.enabled}, "
+            f"proactive_reconnect={self.proactive_reconnect}, "
+            f"relaxed_timeout={self.relaxed_timeout}, "
+            f"endpoint_type={self.endpoint_type!r}"
+            f")"
+        )
+
+    def is_relaxed_timeouts_enabled(self) -> bool:
+        """
+        Check if the relaxed_timeout is enabled. The '-1' value is used to disable the relaxed_timeout.
+        If relaxed_timeout is set to None, it will make the operation blocking
+        and waiting until any response is received.
+
+        Returns:
+            True if the relaxed_timeout is enabled, False otherwise.
+        """
+        return self.relaxed_timeout != -1
+
+    def get_endpoint_type(
+        self, host: str, connection: "MaintNotificationsAbstractConnection"
+    ) -> EndpointType:
+        """
+        Determine the appropriate endpoint type for CLIENT MAINT_NOTIFICATIONS command.
+
+        Logic:
+        1. If endpoint_type is explicitly set, use it
+        2. Otherwise, check the original host from connection.host:
+           - If host is an IP address, use it directly to determine internal-ip vs external-ip
+           - If host is an FQDN, get the resolved IP to determine internal-fqdn vs external-fqdn
+
+        Args:
+            host: User provided hostname to analyze
+            connection: The connection object to analyze for endpoint type determination
+
+        Returns:
+        """
+
+        # If endpoint_type is explicitly set, use it
+        if self.endpoint_type is not None:
+            return self.endpoint_type
+
+        # Check if the host is an IP address
+        try:
+            ip_addr = ipaddress.ip_address(host)
+            # Host is an IP address - use it directly
+            is_private = ip_addr.is_private
+            return EndpointType.INTERNAL_IP if is_private else EndpointType.EXTERNAL_IP
+        except ValueError:
+            # Host is an FQDN - need to check resolved IP to determine internal vs external
+            pass
+
+        # Host is an FQDN, get the resolved IP to determine if it's internal or external
+        resolved_ip = connection.get_resolved_ip()
+
+        if resolved_ip:
+            try:
+                ip_addr = ipaddress.ip_address(resolved_ip)
+                is_private = ip_addr.is_private
+                # Use FQDN types since the original host was an FQDN
+                return (
+                    EndpointType.INTERNAL_FQDN
+                    if is_private
+                    else EndpointType.EXTERNAL_FQDN
+                )
+            except ValueError:
+                # This shouldn't happen since we got the IP from the socket, but fallback
+                pass
+
+        # Final fallback: use heuristics on the FQDN itself
+        is_private = _is_private_fqdn(host)
+        return EndpointType.INTERNAL_FQDN if is_private else EndpointType.EXTERNAL_FQDN
+
+
+class MaintNotificationsPoolHandler:
+    def __init__(
+        self,
+        pool: "MaintNotificationsAbstractConnectionPool",
+        config: MaintNotificationsConfig,
+    ) -> None:
+        self.pool = pool
+        self.config = config
+        self._processed_notifications = 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
+        # but each connection should have its own pool handler
+        # since each connection can be in a different state
+        copy = MaintNotificationsPoolHandler(self.pool, self.config)
+        copy._processed_notifications = self._processed_notifications
+        copy._lock = self._lock
+        copy.connection = None
+        return copy
+
+    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):
+        self.remove_expired_notifications()
+
+        if isinstance(notification, NodeMovingNotification):
+            return self.handle_node_moving_notification(notification)
+        else:
+            logging.error(f"Unhandled notification type: {notification}")
+
+    def handle_node_moving_notification(self, notification: NodeMovingNotification):
+        if (
+            not self.config.proactive_reconnect
+            and not self.config.is_relaxed_timeouts_enabled()
+        ):
+            return
+        with self._lock:
+            if notification in self._processed_notifications:
+                # nothing to do in the connection pool handling
+                # the notification has already been handled or is expired
+                # just return
+                return
+
+            with self.pool._lock:
+                if (
+                    self.config.proactive_reconnect
+                    or self.config.is_relaxed_timeouts_enabled()
+                ):
+                    # Get the current connected address - if any
+                    # This is the address that is being moved
+                    # and we need to handle only connections
+                    # connected to the same address
+                    moving_address_src = (
+                        self.connection.getpeername() if self.connection else None
+                    )
+
+                    if getattr(self.pool, "set_in_maintenance", False):
+                        # Set pool in maintenance mode - executed only if
+                        # BlockingConnectionPool is used
+                        self.pool.set_in_maintenance(True)
+
+                    # Update maintenance state, timeout and optionally host address
+                    # connection settings for matching connections
+                    self.pool.update_connections_settings(
+                        state=MaintenanceState.MOVING,
+                        maintenance_notification_hash=hash(notification),
+                        relaxed_timeout=self.config.relaxed_timeout,
+                        host_address=notification.new_node_host,
+                        matching_address=moving_address_src,
+                        matching_pattern="connected_address",
+                        update_notification_hash=True,
+                        include_free_connections=True,
+                    )
+
+                    if self.config.proactive_reconnect:
+                        if notification.new_node_host is not None:
+                            self.run_proactive_reconnect(moving_address_src)
+                        else:
+                            threading.Timer(
+                                notification.ttl / 2,
+                                self.run_proactive_reconnect,
+                                args=(moving_address_src,),
+                            ).start()
+
+                    # Update config for new connections:
+                    # Set state to MOVING
+                    # update host
+                    # if relax timeouts are enabled - update timeouts
+                    kwargs: dict = {
+                        "maintenance_state": MaintenanceState.MOVING,
+                        "maintenance_notification_hash": hash(notification),
+                    }
+                    if notification.new_node_host is not None:
+                        # the host is not updated if the new node host is None
+                        # this happens when the MOVING push notification does not contain
+                        # the new node host - in this case we only update the timeouts
+                        kwargs.update(
+                            {
+                                "host": notification.new_node_host,
+                            }
+                        )
+                    if self.config.is_relaxed_timeouts_enabled():
+                        kwargs.update(
+                            {
+                                "socket_timeout": self.config.relaxed_timeout,
+                                "socket_connect_timeout": self.config.relaxed_timeout,
+                            }
+                        )
+                    self.pool.update_connection_kwargs(**kwargs)
+
+                    if getattr(self.pool, "set_in_maintenance", False):
+                        self.pool.set_in_maintenance(False)
+
+            threading.Timer(
+                notification.ttl,
+                self.handle_node_moved_notification,
+                args=(notification,),
+            ).start()
+
+            self._processed_notifications.add(notification)
+
+    def run_proactive_reconnect(self, moving_address_src: Optional[str] = None):
+        """
+        Run proactive reconnect for the pool.
+        Active connections are marked for reconnect after they complete the current command.
+        Inactive connections are disconnected and will be connected on next use.
+        """
+        with self._lock:
+            with self.pool._lock:
+                # take care for the active connections in the pool
+                # mark them for reconnect after they complete the current command
+                self.pool.update_active_connections_for_reconnect(
+                    moving_address_src=moving_address_src,
+                )
+                # take care for the inactive connections in the pool
+                # delete them and create new ones
+                self.pool.disconnect_free_connections(
+                    moving_address_src=moving_address_src,
+                )
+
+    def handle_node_moved_notification(self, notification: NodeMovingNotification):
+        """
+        Handle the cleanup after a node moving notification expires.
+        """
+        notification_hash = hash(notification)
+
+        with self._lock:
+            # if the current maintenance_notification_hash in kwargs is not matching the notification
+            # it means there has been a new moving notification after this one
+            # and we don't need to revert the kwargs yet
+            if (
+                self.pool.connection_kwargs.get("maintenance_notification_hash")
+                == notification_hash
+            ):
+                orig_host = self.pool.connection_kwargs.get("orig_host_address")
+                orig_socket_timeout = self.pool.connection_kwargs.get(
+                    "orig_socket_timeout"
+                )
+                orig_connect_timeout = self.pool.connection_kwargs.get(
+                    "orig_socket_connect_timeout"
+                )
+                kwargs: dict = {
+                    "maintenance_state": MaintenanceState.NONE,
+                    "maintenance_notification_hash": None,
+                    "host": orig_host,
+                    "socket_timeout": orig_socket_timeout,
+                    "socket_connect_timeout": orig_connect_timeout,
+                }
+                self.pool.update_connection_kwargs(**kwargs)
+
+            with self.pool._lock:
+                reset_relaxed_timeout = self.config.is_relaxed_timeouts_enabled()
+                reset_host_address = self.config.proactive_reconnect
+
+                self.pool.update_connections_settings(
+                    relaxed_timeout=-1,
+                    state=MaintenanceState.NONE,
+                    maintenance_notification_hash=None,
+                    matching_notification_hash=notification_hash,
+                    matching_pattern="notification_hash",
+                    update_notification_hash=True,
+                    reset_relaxed_timeout=reset_relaxed_timeout,
+                    reset_host_address=reset_host_address,
+                    include_free_connections=True,
+                )
+
+
+class MaintNotificationsConnectionHandler:
+    # 1 = "starting maintenance" notifications, 0 = "completed maintenance" notifications
+    _NOTIFICATION_TYPES: dict[type["MaintenanceNotification"], int] = {
+        NodeMigratingNotification: 1,
+        NodeFailingOverNotification: 1,
+        NodeMigratedNotification: 0,
+        NodeFailedOverNotification: 0,
+    }
+
+    def __init__(
+        self,
+        connection: "MaintNotificationsAbstractConnection",
+        config: MaintNotificationsConfig,
+    ) -> None:
+        self.connection = connection
+        self.config = config
+
+    def handle_notification(self, notification: MaintenanceNotification):
+        # get the notification type by checking its class in the _NOTIFICATION_TYPES dict
+        notification_type = self._NOTIFICATION_TYPES.get(notification.__class__, None)
+
+        if notification_type is None:
+            logging.error(f"Unhandled notification type: {notification}")
+            return
+
+        if notification_type:
+            self.handle_maintenance_start_notification(MaintenanceState.MAINTENANCE)
+        else:
+            self.handle_maintenance_completed_notification()
+
+    def handle_maintenance_start_notification(
+        self, maintenance_state: MaintenanceState
+    ):
+        if (
+            self.connection.maintenance_state == MaintenanceState.MOVING
+            or not self.config.is_relaxed_timeouts_enabled()
+        ):
+            return
+
+        self.connection.maintenance_state = maintenance_state
+        self.connection.set_tmp_settings(
+            tmp_relaxed_timeout=self.config.relaxed_timeout
+        )
+        # extend the timeout for all created connections
+        self.connection.update_current_socket_timeout(self.config.relaxed_timeout)
+
+    def handle_maintenance_completed_notification(self):
+        # Only reset timeouts if state is not MOVING and relaxed timeouts are enabled
+        if (
+            self.connection.maintenance_state == MaintenanceState.MOVING
+            or not self.config.is_relaxed_timeouts_enabled()
+        ):
+            return
+        self.connection.reset_tmp_settings(reset_relaxed_timeout=True)
+        # Maintenance completed - reset the connection
+        # timeouts by providing -1 as the relaxed timeout
+        self.connection.update_current_socket_timeout(-1)
+        self.connection.maintenance_state = MaintenanceState.NONE