Merge branch 'ps_add_fail_over_events_handling' into hitless_handshake

Author: petyaslavova

redis/maintenance_events.py

@@ -13,10 +13,9 @@
 class MaintenanceState(enum.Enum):
     NONE = "none"
     MOVING = "moving"
-    MIGRATING = "migrating"
-    FAILING_OVER = "failing_over"
-
+    MAINTENANCE = "maintenance"
 
+    
 class EndpointType:
     """Constants for valid endpoint types used in CLIENT MAINT_NOTIFICATIONS command."""
 
@@ -26,6 +25,7 @@ class EndpointType:
     EXTERNAL_FQDN = "external-fqdn"
     NONE = "none"
 
+    
     @classmethod
     def get_valid_types(cls):
         """Return a set of all valid endpoint types."""
@@ -37,7 +37,7 @@ def get_valid_types(cls):
             cls.NONE,
         }
 
-
+      
 if TYPE_CHECKING:
     from redis.connection import (
         BlockingConnectionPool,
@@ -531,6 +531,7 @@ def __init__(
 
         self.endpoint_type = endpoint_type
 
+
     def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}("
@@ -748,30 +749,40 @@ def handle_node_moved_event(self, event: NodeMovingEvent):
 
 
 class MaintenanceEventConnectionHandler:
+    # 1 = "starting maintenance" events, 0 = "completed maintenance" events
+    _EVENT_TYPES: dict[type["MaintenanceEvent"], int] = {
+        NodeMigratingEvent: 1,
+        NodeFailingOverEvent: 1,
+        NodeMigratedEvent: 0,
+        NodeFailedOverEvent: 0,
+    }
+
     def __init__(
         self, connection: "ConnectionInterface", config: MaintenanceEventsConfig
     ) -> None:
         self.connection = connection
         self.config = config
 
     def handle_event(self, event: MaintenanceEvent):
-        if isinstance(event, NodeMigratingEvent):
-            return self.handle_maintenance_start_event(MaintenanceState.MIGRATING)
-        elif isinstance(event, NodeMigratedEvent):
-            return self.handle_maintenance_completed_event()
-        elif isinstance(event, NodeFailingOverEvent):
-            return self.handle_maintenance_start_event(MaintenanceState.FAILING_OVER)
-        elif isinstance(event, NodeFailedOverEvent):
-            return self.handle_maintenance_completed_event()
-        else:
+        # get the event type by checking its class in the _EVENT_TYPES dict
+        event_type = self._EVENT_TYPES.get(event.__class__, None)
+
+        if event_type is None:
             logging.error(f"Unhandled event type: {event}")
+            return
+
+        if event_type:
+            self.handle_maintenance_start_event(MaintenanceState.MAINTENANCE)
+        else:
+            self.handle_maintenance_completed_event()
 
     def handle_maintenance_start_event(self, maintenance_state: MaintenanceState):
         if (
             self.connection.maintenance_state == MaintenanceState.MOVING
             or not self.config.is_relax_timeouts_enabled()
         ):
             return
+
         self.connection.maintenance_state = maintenance_state
         self.connection.set_tmp_settings(tmp_relax_timeout=self.config.relax_timeout)
         # extend the timeout for all created connections

tests/test_maintenance_events.py

@@ -3,6 +3,7 @@
 import pytest
 
 from redis.connection import ConnectionInterface
+
 from redis.maintenance_events import (
     MaintenanceEvent,
     NodeMovingEvent,
@@ -564,7 +565,7 @@ def test_handle_event_migrating(self):
             self.handler, "handle_maintenance_start_event"
         ) as mock_handle:
             self.handler.handle_event(event)
-            mock_handle.assert_called_once_with(MaintenanceState.MIGRATING)
+            mock_handle.assert_called_once_with(MaintenanceState.MAINTENANCE)
 
     def test_handle_event_migrated(self):
         """Test handling of NodeMigratedEvent."""
@@ -584,7 +585,8 @@ def test_handle_event_failing_over(self):
             self.handler, "handle_maintenance_start_event"
         ) as mock_handle:
             self.handler.handle_event(event)
-            mock_handle.assert_called_once_with(MaintenanceState.FAILING_OVER)
+            mock_handle.assert_called_once_with(MaintenanceState.MAINTENANCE)
+
 
     def test_handle_event_failed_over(self):
         """Test handling of NodeFailedOverEvent."""
@@ -610,37 +612,28 @@ def test_handle_maintenance_start_event_disabled(self):
         config = MaintenanceEventsConfig(relax_timeout=-1)
         handler = MaintenanceEventConnectionHandler(self.mock_connection, config)
 
-        result = handler.handle_maintenance_start_event(MaintenanceState.MIGRATING)
+        result = handler.handle_maintenance_start_event(MaintenanceState.MAINTENANCE)
+
         assert result is None
         self.mock_connection.update_current_socket_timeout.assert_not_called()
 
     def test_handle_maintenance_start_event_moving_state(self):
         """Test maintenance start event handling when connection is in MOVING state."""
         self.mock_connection.maintenance_state = MaintenanceState.MOVING
 
-        result = self.handler.handle_maintenance_start_event(MaintenanceState.MIGRATING)
+        result = self.handler.handle_maintenance_start_event(
+            MaintenanceState.MAINTENANCE
+        )
         assert result is None
         self.mock_connection.update_current_socket_timeout.assert_not_called()
 
-    def test_handle_maintenance_start_event_migrating_success(self):
+    def test_handle_maintenance_start_event_success(self):
         """Test successful maintenance start event handling for migrating."""
         self.mock_connection.maintenance_state = MaintenanceState.NONE
 
-        self.handler.handle_maintenance_start_event(MaintenanceState.MIGRATING)
+        self.handler.handle_maintenance_start_event(MaintenanceState.MAINTENANCE)
 
-        assert self.mock_connection.maintenance_state == MaintenanceState.MIGRATING
-        self.mock_connection.update_current_socket_timeout.assert_called_once_with(20)
-        self.mock_connection.set_tmp_settings.assert_called_once_with(
-            tmp_relax_timeout=20
-        )
-
-    def test_handle_maintenance_start_event_failing_over_success(self):
-        """Test successful maintenance start event handling for failing over."""
-        self.mock_connection.maintenance_state = MaintenanceState.NONE
-
-        self.handler.handle_maintenance_start_event(MaintenanceState.FAILING_OVER)
-
-        assert self.mock_connection.maintenance_state == MaintenanceState.FAILING_OVER
+        assert self.mock_connection.maintenance_state == MaintenanceState.MAINTENANCE
         self.mock_connection.update_current_socket_timeout.assert_called_once_with(20)
         self.mock_connection.set_tmp_settings.assert_called_once_with(
             tmp_relax_timeout=20
@@ -665,11 +658,12 @@ def test_handle_maintenance_completed_event_moving_state(self):
 
     def test_handle_maintenance_completed_event_success(self):
         """Test successful maintenance completed event handling."""
-        self.mock_connection.maintenance_state = MaintenanceState.MIGRATING
+        self.mock_connection.maintenance_state = MaintenanceState.MAINTENANCE
 
         self.handler.handle_maintenance_completed_event()
 
         assert self.mock_connection.maintenance_state == MaintenanceState.NONE
+
         self.mock_connection.update_current_socket_timeout.assert_called_once_with(-1)
         self.mock_connection.reset_tmp_settings.assert_called_once_with(
             reset_relax_timeout=True

tests/test_maintenance_events_handling.py

@@ -1431,7 +1431,8 @@ def worker(idx):
     def test_moving_migrating_migrated_moved_state_transitions(self, pool_class):
         """
         Test moving configs are not lost if the per connection events get picked up after moving is handled.
-        MOVING → MIGRATING → MIGRATED → FAILING_OVER → FAILED_OVER → MOVED
+        Sequence of events: MOVING, MIGRATING, MIGRATED, FAILING_OVER, FAILED_OVER, MOVED.
+        Note: FAILING_OVER and FAILED_OVER events do not change the connection state when already in MOVING state.
         Checks the state after each event for all connections and for new connections created during each state.
         """
         # Setup
@@ -1541,25 +1542,6 @@ def test_moving_migrating_migrated_moved_state_transitions(self, pool_class):
             expected_current_peername=DEFAULT_ADDRESS.split(":")[0],
         )
 
-        # 4. FAILING_OVER event (simulate direct connection handler call)
-        for conn in in_use_connections:
-            conn._maintenance_event_connection_handler.handle_event(
-                NodeFailingOverEvent(id=3, ttl=1)
-            )
-        # State should not change for connections that are in MOVING state
-        Helpers.validate_in_use_connections_state(
-            in_use_connections,
-            expected_state=MaintenanceState.MOVING,
-            expected_host_address=tmp_address,
-            expected_socket_timeout=self.config.relax_timeout,
-            expected_socket_connect_timeout=self.config.relax_timeout,
-            expected_orig_host_address=DEFAULT_ADDRESS.split(":")[0],
-            expected_orig_socket_timeout=None,
-            expected_orig_socket_connect_timeout=None,
-            expected_current_socket_timeout=self.config.relax_timeout,
-            expected_current_peername=DEFAULT_ADDRESS.split(":")[0],
-        )
-
         # 5. FAILED_OVER event (simulate direct connection handler call)
         for conn in in_use_connections:
             conn._maintenance_event_connection_handler.handle_event(
@@ -1579,21 +1561,6 @@ def test_moving_migrating_migrated_moved_state_transitions(self, pool_class):
             expected_current_peername=DEFAULT_ADDRESS.split(":")[0],
         )
 
-        # 6. MOVED event (simulate timer expiry)
-        pool_handler.handle_node_moved_event(moving_event)
-        Helpers.validate_in_use_connections_state(
-            in_use_connections,
-            expected_state=MaintenanceState.NONE,
-            expected_host_address=DEFAULT_ADDRESS.split(":")[0],
-            expected_socket_timeout=None,
-            expected_socket_connect_timeout=None,
-            expected_orig_host_address=DEFAULT_ADDRESS.split(":")[0],
-            expected_orig_socket_timeout=None,
-            expected_orig_socket_connect_timeout=None,
-            expected_current_socket_timeout=None,
-            expected_current_peername=DEFAULT_ADDRESS.split(":")[0],
-        )
-
         # 6. MOVED event (simulate timer expiry)
         pool_handler.handle_node_moved_event(moving_event)
         Helpers.validate_in_use_connections_state(
@@ -1841,7 +1808,8 @@ def test_migrating_after_moving_multiple_proxies(self, pool_class):
         conn_event_handler = conn._maintenance_event_connection_handler
         conn_event_handler.handle_event(NodeMigratingEvent(id=3, ttl=1))
         # validate connection is in MIGRATING state
-        assert conn.maintenance_state == MaintenanceState.MIGRATING
+        assert conn.maintenance_state == MaintenanceState.MAINTENANCE
+
         assert conn.socket_timeout == self.config.relax_timeout
 
         # Send MIGRATED event to con with ip = key3