Adding maintenance state to connections. Migrating and Migrated are not processed in in Moving state. Tests are updated

Author: petyaslavova

redis/connection.py

@@ -11,6 +11,7 @@
 from queue import Empty, Full, LifoQueue
 from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union
 from urllib.parse import parse_qs, unquote, urlparse
+import enum
 
 from redis.cache import (
     CacheEntry,
@@ -41,6 +42,7 @@
     MaintenanceEventConnectionHandler,
     MaintenanceEventPoolHandler,
     MaintenanceEventsConfig,
+    MaintenanceState,
 )
 from .retry import Retry
 from .utils import (
@@ -284,6 +286,7 @@ def __init__(
         maintenance_events_config: Optional[MaintenanceEventsConfig] = None,
         tmp_host_address: Optional[str] = None,
         tmp_relax_timeout: Optional[float] = -1,
+        maintenance_state: "MaintenanceState" = MaintenanceState.NONE,
     ):
         """
         Initialize a new Connection.
@@ -373,6 +376,7 @@ def __init__(
         self._should_reconnect = False
         self.tmp_host_address = tmp_host_address
         self.tmp_relax_timeout = tmp_relax_timeout
+        self.maintenance_state = maintenance_state
 
     def __repr__(self):
         repr_args = ",".join([f"{k}={v}" for k, v in self.repr_pieces()])
@@ -829,6 +833,9 @@ def update_tmp_settings(
         if tmp_relax_timeout is not SENTINEL:
             self.tmp_relax_timeout = tmp_relax_timeout
 
+    def set_maintenance_state(self, state: "MaintenanceState"):
+        self.maintenance_state = state
+
 
 class Connection(AbstractConnection):
     "Manages TCP communication to and from a Redis server"
@@ -1717,11 +1724,18 @@ def make_connection(self) -> "ConnectionInterface":
             raise ConnectionError("Too many connections")
         self._created_connections += 1
 
+        # Pass current maintenance_state to new connections
+        maintenance_state = self.connection_kwargs.get(
+            "maintenance_state", MaintenanceState.NONE
+        )
+        kwargs = dict(self.connection_kwargs)
+        kwargs["maintenance_state"] = maintenance_state
+
         if self.cache is not None:
             return CacheProxyConnection(
-                self.connection_class(**self.connection_kwargs), self.cache, self._lock
+                self.connection_class(**kwargs), self.cache, self._lock
             )
-        return self.connection_class(**self.connection_kwargs)
+        return self.connection_class(**kwargs)
 
     def release(self, connection: "Connection") -> None:
         "Releases the connection back to the pool"
@@ -1946,6 +1960,16 @@ async def _mock(self, error: RedisError):
         """
         pass
 
+    def set_maintenance_state_for_all(self, state: "MaintenanceState"):
+        with self._lock:
+            for conn in self._available_connections:
+                conn.set_maintenance_state(state)
+            for conn in self._in_use_connections:
+                conn.set_maintenance_state(state)
+
+    def set_maintenance_state_in_kwargs(self, state: "MaintenanceState"):
+        self.connection_kwargs["maintenance_state"] = state
+
 
 class BlockingConnectionPool(ConnectionPool):
     """
@@ -2040,15 +2064,20 @@ def make_connection(self):
             if self._in_maintenance:
                 self._lock.acquire()
                 self._locked = True
+            # Pass current maintenance_state to new connections
+            maintenance_state = self.connection_kwargs.get(
+                "maintenance_state", MaintenanceState.NONE
+            )
+            kwargs = dict(self.connection_kwargs)
+            kwargs["maintenance_state"] = maintenance_state
             if self.cache is not None:
                 connection = CacheProxyConnection(
-                    self.connection_class(**self.connection_kwargs),
+                    self.connection_class(**kwargs),
                     self.cache,
                     self._lock,
                 )
             else:
-                connection = self.connection_class(**self.connection_kwargs)
-
+                connection = self.connection_class(**kwargs)
                 self._connections.append(connection)
             return connection
         finally:
@@ -2259,3 +2288,12 @@ def _update_maintenance_events_configs_for_connections(
     def set_in_maintenance(self, in_maintenance: bool):
         """Set the maintenance mode for the connection pool."""
         self._in_maintenance = in_maintenance
+
+    def set_maintenance_state_for_all(self, state: "MaintenanceState"):
+        with self._lock:
+            for conn in getattr(self, "_connections", []):
+                if conn:
+                    conn.set_maintenance_state(state)
+
+    def set_maintenance_state_in_kwargs(self, state: "MaintenanceState"):
+        self.connection_kwargs["maintenance_state"] = state

redis/maintenance_events.py

@@ -3,6 +3,14 @@
 import time
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Optional, Union
+import enum
+
+
+class MaintenanceState(enum.Enum):
+    NONE = "none"
+    MOVING = "moving"
+    MIGRATING = "migrating"
+
 
 from redis.typing import Number
 
@@ -351,6 +359,9 @@ def handle_node_moving_event(self, event: NodeMovingEvent):
                 ):
                     if getattr(self.pool, "set_in_maintenance", False):
                         self.pool.set_in_maintenance(True)
+                    # Set state to MOVING for all connections and in kwargs (inside pool lock, after set_in_maintenance)
+                    self.pool.set_maintenance_state_for_all(MaintenanceState.MOVING)
+                    self.pool.set_maintenance_state_in_kwargs(MaintenanceState.MOVING)
                     # edit the config for new connections until the notification expires
                     self.pool.update_connection_kwargs_with_tmp_settings(
                         tmp_host_address=event.new_node_host,
@@ -368,7 +379,6 @@ def handle_node_moving_event(self, event: NodeMovingEvent):
                             tmp_host_address=event.new_node_host,
                             tmp_relax_timeout=self.config.relax_timeout,
                         )
-
                         # take care for the inactive connections in the pool
                         # delete them and create new ones
                         self.pool.disconnect_and_reconfigure_free_connections(
@@ -388,16 +398,19 @@ def handle_node_moved_event(self):
                 tmp_host_address=None,
                 tmp_relax_timeout=-1,
             )
+            # Clear state to NONE in kwargs immediately after updating tmp kwargs
+            self.pool.set_maintenance_state_in_kwargs(MaintenanceState.NONE)
             with self.pool._lock:
                 if self.config.is_relax_timeouts_enabled():
                     # reset the timeout for existing connections
                     self.pool.update_connections_current_timeout(
                         relax_timeout=-1, include_free_connections=True
                     )
-
                 self.pool.update_connections_tmp_settings(
                     tmp_host_address=None, tmp_relax_timeout=-1
                 )
+                # Clear state to NONE for all connections
+                self.pool.set_maintenance_state_for_all(MaintenanceState.NONE)
 
 
 class MaintenanceEventConnectionHandler:
@@ -416,17 +429,24 @@ def handle_event(self, event: MaintenanceEvent):
             logging.error(f"Unhandled event type: {event}")
 
     def handle_migrating_event(self, notification: NodeMigratingEvent):
-        if not self.config.is_relax_timeouts_enabled():
+        if (
+            self.connection.maintenance_state == MaintenanceState.MOVING
+            or not self.config.is_relax_timeouts_enabled()
+        ):
             return
-
+        self.connection.set_maintenance_state(MaintenanceState.MIGRATING)
         # extend the timeout for all created connections
         self.connection.update_current_socket_timeout(self.config.relax_timeout)
         self.connection.update_tmp_settings(tmp_relax_timeout=self.config.relax_timeout)
 
     def handle_migration_completed_event(self, notification: "NodeMigratedEvent"):
-        if not self.config.is_relax_timeouts_enabled():
+        # Only reset timeouts if state is not MOVING and relax timeouts are enabled
+        if (
+            self.connection.maintenance_state == MaintenanceState.MOVING
+            or not self.config.is_relax_timeouts_enabled()
+        ):
             return
-
+        self.connection.set_maintenance_state(MaintenanceState.NONE)
         # Node migration completed - reset the connection
         # timeouts by providing -1 as the relax timeout
         self.connection.update_current_socket_timeout(-1)

tests/test_connection_pool.py

@@ -9,6 +9,7 @@
 import redis
 from redis.cache import CacheConfig
 from redis.connection import CacheProxyConnection, Connection, to_bool
+from redis.maintenance_events import MaintenanceState
 from redis.utils import SSL_AVAILABLE
 
 from .conftest import (
@@ -53,10 +54,15 @@ def get_pool(
         return pool
 
     def test_connection_creation(self):
-        connection_kwargs = {"foo": "bar", "biz": "baz"}
+        connection_kwargs = {
+            "foo": "bar",
+            "biz": "baz",
+            "maintenance_state": MaintenanceState.NONE,
+        }
         pool = self.get_pool(
             connection_kwargs=connection_kwargs, connection_class=DummyConnection
         )
+
         connection = pool.get_connection()
         assert isinstance(connection, DummyConnection)
         assert connection.kwargs == connection_kwargs
@@ -150,7 +156,9 @@ def test_connection_creation(self, master_host):
             "host": master_host[0],
             "port": master_host[1],
         }
+
         pool = self.get_pool(connection_kwargs=connection_kwargs)
+        connection_kwargs["maintenance_state"] = MaintenanceState.NONE
         connection = pool.get_connection()
         assert isinstance(connection, DummyConnection)
         assert connection.kwargs == connection_kwargs