Applying moving/moved only on connections to the same proxy.

Author: petyaslavova

redis/_parsers/base.py

@@ -202,17 +202,20 @@ def handle_push_response(self, response, **kwargs):
         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:
@@ -260,17 +263,20 @@ async def handle_push_response(self, response, **kwargs):
             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)
@@ -283,7 +289,7 @@ def set_invalidation_push_handler(self, invalidation_push_handler_func):
         """Set the invalidation push handler function"""
         self.invalidation_push_handler_func = invalidation_push_handler_func
 
-    def set_node_moving_push_handler_func(self, node_moving_push_handler_func):
+    def set_node_moving_push_handler(self, node_moving_push_handler_func):
         self.node_moving_push_handler_func = node_moving_push_handler_func
 
     def set_maintenance_push_handler(self, maintenance_push_handler_func):

redis/connection.py

@@ -8,7 +8,7 @@
 from abc import abstractmethod
 from itertools import chain
 from queue import Empty, Full, LifoQueue
-from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, Type, TypeVar, Union
 from urllib.parse import parse_qs, unquote, urlparse
 
 from redis.cache import (
@@ -249,6 +249,13 @@ def maintenance_state(self, state: "MaintenanceState"):
         """
         pass
 
+    @abstractmethod
+    def getpeername(self):
+        """
+        Returns the peer name of the connection.
+        """
+        pass
+
     @abstractmethod
     def mark_for_reconnect(self):
         """
@@ -402,6 +409,7 @@ def __init__(
 
         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
                 )
@@ -484,6 +492,7 @@ def set_parser(self, parser_class):
     def set_maintenance_event_pool_handler(
         self, maintenance_event_pool_handler: MaintenanceEventPoolHandler
     ):
+        maintenance_event_pool_handler.set_connection(self)
         self._parser.set_node_moving_push_handler(
             maintenance_event_pool_handler.handle_event
         )
@@ -867,6 +876,11 @@ def maintenance_state(self) -> MaintenanceState:
     def maintenance_state(self, state: "MaintenanceState"):
         self._maintenance_state = state
 
+    def getpeername(self):
+        if not self._sock:
+            return None
+        return self._sock.getpeername()[0]
+
     def mark_for_reconnect(self):
         self._should_reconnect = True
 
@@ -1892,10 +1906,27 @@ def re_auth_callback(self, token: TokenInterface):
             for conn in self._in_use_connections:
                 conn.set_re_auth_token(token)
 
-    def set_maintenance_state_for_all_connections(self, state: "MaintenanceState"):
+    def set_maintenance_state_for_connections(
+        self,
+        state: "MaintenanceState",
+        matching_address: Optional[str] = None,
+        address_type_to_match: Literal["connected", "configured"] = "connected",
+    ):
         for conn in self._available_connections:
+            if address_type_to_match == "connected":
+                if matching_address and conn.getpeername() != matching_address:
+                    continue
+            else:
+                if matching_address and conn.host != matching_address:
+                    continue
             conn.maintenance_state = state
         for conn in self._in_use_connections:
+            if address_type_to_match == "connected":
+                if matching_address and conn.getpeername() != matching_address:
+                    continue
+            else:
+                if matching_address and conn.host != matching_address:
+                    continue
             conn.maintenance_state = state
 
     def set_maintenance_state_in_connection_kwargs(self, state: "MaintenanceState"):
@@ -1963,7 +1994,12 @@ def remove_tmp_config_from_connection_kwargs(self):
             }
         )
 
-    def reset_connections_tmp_settings(self):
+    def reset_connections_tmp_settings(
+        self,
+        moving_address: Optional[str] = None,
+        reset_host_address: bool = False,
+        reset_relax_timeout: bool = False,
+    ):
         """
         Restore original settings from temporary configuration for all connections in the pool.
 
@@ -1978,16 +2014,25 @@ def reset_connections_tmp_settings(self):
         """
         with self._lock:
             for conn in self._available_connections:
+                if moving_address and conn.host != moving_address:
+                    continue
                 conn.reset_tmp_settings(
-                    reset_host_address=True, reset_relax_timeout=True
+                    reset_host_address=reset_host_address,
+                    reset_relax_timeout=reset_relax_timeout,
                 )
             for conn in self._in_use_connections:
+                if moving_address and conn.host != moving_address:
+                    continue
                 conn.reset_tmp_settings(
-                    reset_host_address=True, reset_relax_timeout=True
+                    reset_host_address=reset_host_address,
+                    reset_relax_timeout=reset_relax_timeout,
                 )
 
     def update_active_connections_for_reconnect(
-        self, tmp_host_address: str, tmp_relax_timeout: Optional[float] = None
+        self,
+        tmp_host_address: str,
+        tmp_relax_timeout: Optional[float] = None,
+        moving_address_src: Optional[str] = None,
     ):
         """
         Mark all active connections for reconnect.
@@ -1999,6 +2044,8 @@ def update_active_connections_for_reconnect(
         :param tmp_relax_timeout: The relax timeout to use for the connection.
         """
         for conn in self._in_use_connections:
+            if moving_address_src and conn.getpeername() != moving_address_src:
+                continue
             self._update_connection_for_reconnect(
                 conn, tmp_host_address, tmp_relax_timeout
             )
@@ -2007,6 +2054,7 @@ def disconnect_and_reconfigure_free_connections(
         self,
         tmp_host_address: str,
         tmp_relax_timeout: Optional[float] = None,
+        moving_address_src: Optional[str] = None,
     ):
         """
         Disconnect all free/available connections.
@@ -2019,13 +2067,17 @@ def disconnect_and_reconfigure_free_connections(
         """
 
         for conn in self._available_connections:
+            if moving_address_src and conn.getpeername() != moving_address_src:
+                continue
             self._disconnect_and_update_connection_for_reconnect(
                 conn, tmp_host_address, tmp_relax_timeout
             )
 
     def update_connections_current_timeout(
         self,
         relax_timeout: Optional[float],
+        matching_address: Optional[str] = None,
+        address_type_to_match: Literal["connected", "configured"] = "connected",
         include_free_connections: bool = False,
     ):
         """
@@ -2039,10 +2091,22 @@ def update_connections_current_timeout(
         :param include_available_connections: Whether to include available connections in the update.
         """
         for conn in self._in_use_connections:
+            if address_type_to_match == "connected":
+                if matching_address and conn.getpeername() != matching_address:
+                    continue
+            else:
+                if matching_address and conn.host != matching_address:
+                    continue
             conn.update_current_socket_timeout(relax_timeout)
 
         if include_free_connections:
             for conn in self._available_connections:
+                if address_type_to_match == "connected":
+                    if matching_address and conn.getpeername() != matching_address:
+                        continue
+                else:
+                    if matching_address and conn.host != matching_address:
+                        continue
                 conn.update_current_socket_timeout(relax_timeout)
 
     def _update_connection_for_reconnect(
@@ -2308,7 +2372,10 @@ def disconnect(self):
                 self._locked = False
 
     def update_active_connections_for_reconnect(
-        self, tmp_host_address: str, tmp_relax_timeout: Optional[float] = None
+        self,
+        tmp_host_address: str,
+        tmp_relax_timeout: Optional[float] = None,
+        moving_address_src: Optional[str] = None,
     ):
         """
         Mark all active connections for reconnect.
@@ -2323,6 +2390,8 @@ def update_active_connections_for_reconnect(
             connections_in_queue = {conn for conn in self.pool.queue if conn}
             for conn in self._connections:
                 if conn not in connections_in_queue:
+                    if moving_address_src and conn.getpeername() != moving_address_src:
+                        continue
                     self._update_connection_for_reconnect(
                         conn, tmp_host_address, tmp_relax_timeout
                     )
@@ -2331,6 +2400,7 @@ def disconnect_and_reconfigure_free_connections(
         self,
         tmp_host_address: str,
         tmp_relax_timeout: Optional[Number] = None,
+        moving_address_src: Optional[str] = None,
     ):
         """
         Disconnect all free/available connections.
@@ -2345,13 +2415,17 @@ def disconnect_and_reconfigure_free_connections(
 
         for conn in existing_connections:
             if conn:
+                if moving_address_src and conn.getpeername() != moving_address_src:
+                    continue
                 self._disconnect_and_update_connection_for_reconnect(
                     conn, tmp_host_address, tmp_relax_timeout
                 )
 
     def update_connections_current_timeout(
         self,
         relax_timeout: Optional[float] = None,
+        matching_address: Optional[str] = None,
+        address_type_to_match: Literal["connected", "configured"] = "connected",
         include_free_connections: bool = False,
     ):
         """
@@ -2365,11 +2439,23 @@ def update_connections_current_timeout(
         """
         if include_free_connections:
             for conn in tuple(self._connections):
+                if address_type_to_match == "connected":
+                    if matching_address and conn.getpeername() != matching_address:
+                        continue
+                else:
+                    if matching_address and conn.host != matching_address:
+                        continue
                 conn.update_current_socket_timeout(relax_timeout)
         else:
             connections_in_queue = {conn for conn in self.pool.queue if conn}
             for conn in self._connections:
                 if conn not in connections_in_queue:
+                    if address_type_to_match == "connected":
+                        if matching_address and conn.getpeername() != matching_address:
+                            continue
+                    else:
+                        if matching_address and conn.host != matching_address:
+                            continue
                     conn.update_current_socket_timeout(relax_timeout)
 
     def _update_maintenance_events_config_for_connections(
@@ -2387,14 +2473,24 @@ def _update_maintenance_events_configs_for_connections(
                 conn.set_maintenance_event_pool_handler(maintenance_events_pool_handler)
                 conn.maintenance_events_config = maintenance_events_pool_handler.config
 
-    def reset_connections_tmp_settings(self):
+    def reset_connections_tmp_settings(
+        self,
+        moving_address: Optional[str] = None,
+        reset_host_address: bool = False,
+        reset_relax_timeout: bool = False,
+    ):
         """
         Override base class method to work with BlockingConnectionPool's structure.
 
         Restore original settings from temporary configuration for all connections in the pool.
         """
         for conn in tuple(self._connections):
-            conn.reset_tmp_settings(reset_host_address=True, reset_relax_timeout=True)
+            if moving_address and conn.host != moving_address:
+                continue
+            conn.reset_tmp_settings(
+                reset_host_address=reset_host_address,
+                reset_relax_timeout=reset_relax_timeout,
+            )
 
     def set_in_maintenance(self, in_maintenance: bool):
         """
@@ -2405,6 +2501,18 @@ def set_in_maintenance(self, in_maintenance: bool):
         """
         self._in_maintenance = in_maintenance
 
-    def set_maintenance_state_for_all_connections(self, state: "MaintenanceState"):
+    def set_maintenance_state_for_connections(
+        self,
+        state: "MaintenanceState",
+        matching_address: Optional[str] = None,
+        address_type_to_match: Literal["connected", "configured"] = "connected",
+    ):
         for conn in self._connections:
+            if address_type_to_match == "connected":
+                if matching_address and conn.getpeername() != matching_address:
+                    continue
+            else:
+                if matching_address and conn.host != matching_address:
+                    continue
+
             conn.maintenance_state = state