set fix address as a property of SentinelManagedConnection

Author: agnesnatasya

redis/asyncio/connection.py

@@ -1118,6 +1118,7 @@ def __init__(
         self,
         connection_class: Type[AbstractConnection] = Connection,
         max_connections: Optional[int] = None,
+        index_available_connections: bool = False,
         **connection_kwargs,
     ):
         max_connections = max_connections or 2**31
@@ -1130,6 +1131,7 @@ def __init__(
 
         self._available_connections: ConnectionsIndexer = ConnectionsIndexer()
         self._in_use_connections: Set[AbstractConnection] = set()
+        self._index_available_connections= index_available_connections
         self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder)
 
     def __repr__(self):
@@ -1139,7 +1141,9 @@ def __repr__(self):
         )
 
     def reset(self):
-        self._available_connections = []
+        self._available_connections: ConnectionsIndexer | list = (
+            ConnectionsIndexer() if self._index_available_connections else []
+        )
         self._in_use_connections = weakref.WeakSet()
 
     def can_get_connection(self) -> bool:

redis/asyncio/sentinel.py

@@ -35,6 +35,11 @@ class SlaveNotFoundError(ConnectionError):
 class SentinelManagedConnection(Connection):
     def __init__(self, **kwargs):
         self.connection_pool = kwargs.pop("connection_pool")
+        # To be set to True if we want to prevent
+        # the sentinel managed connection to connect
+        # to the most relevant sentinel in the pool and just
+        # connect to the current self.host and self.port
+        self._is_address_fixed = False
         super().__init__(**kwargs)
 
     def __repr__(self):
@@ -48,6 +53,10 @@ def __repr__(self):
             s += host_info
         return s + ")>"
 
+    def fix_address(self, address):
+        self.host, self.port = address
+        self._is_address_fixed = True
+
     async def connect_to(self, address):
         self.host, self.port = address
         await super().connect()
@@ -56,12 +65,12 @@ async def connect_to(self, address):
             if str_if_bytes(await self.read_response()) != "PONG":
                 raise ConnectionError("PING failed")
 
-    async def _connect_retry(self, same_address: bool = False):
+    async def _connect_retry(self):
         if self._reader:
             return  # already connected
-        # If same_server is True, it means that the connection
+        # If address is fixed, it means that the connection
         # is not rotating to the next slave (if the connection pool is in replica mode)
-        if same_address:
+        if self._is_address_fixed:
             self.connect_to(self.host, self.port)
             return
         # If same_server is False, connnect to master in master mode
@@ -82,24 +91,6 @@ async def connect(self):
             lambda error: asyncio.sleep(0),
         )
 
-    async def connect_to_same_address(self):
-        """
-        Similar to connect, but instead of rotating to the next slave (in replica mode),
-        it just connects to the same address of the connection object.
-        """
-        return await self.retry.call_with_retry(
-            lambda: self._connect_retry(same_address=True),
-            lambda error: asyncio.sleep(0),
-        )
-
-    async def connect_to_address(self, address):
-        """
-        Similar to connect, but instead of rotating to the next slave (in replica mode),
-        it just connects to the same address of the connection object.
-        """
-        self.host, self.port = address
-        return await self.connect_to_same_address()
-
     async def read_response(
         self,
         disable_decoding: bool = False,
@@ -202,34 +193,6 @@ async def rotate_slaves(self) -> AsyncIterator:
             pass
         raise SlaveNotFoundError(f"No slave found for {self.service_name!r}")
 
-    async def ensure_connection_connected_to_address(
-        self, connection: SentinelManagedConnection
-    ):
-        """
-        Ensure the connection is already connected to the server that this connection
-        object wants to connect to
-
-        Similar to self.ensure_connection, but calling connection.connect()
-        in SentinelManagedConnection (replica mode) will cause the
-        connection object to connect to the next replica in rotation,
-        and we don't wnat behavior. Look at get_connection inline docs for details.
-
-        Here, we just try to make sure that the connection is already connected
-        to the replica we wanted it to.
-        """
-        await connection.connect_to_same_address()
-        try:
-            if (
-                await connection.can_read_destructive()
-                and connection.client_cache is None
-            ):
-                raise ConnectionError("Connection has data")
-        except (ConnectionError, OSError):
-            await connection.disconnect()
-            await connection.connect_to_same_address()
-            if await connection.can_read_destructive():
-                raise ConnectionError("Connection has data")
-
     def cleanup(self, **options):
         """
         Remove the SCAN ITER family command's request id from the dictionary
@@ -291,8 +254,8 @@ async def get_connection(
             # connect to the previous replica.
             # This will connect to the host and port of the replica
             else:
-                await connection.connect_to_address((server_host, server_port))
-            await self.ensure_connection_connected_to_address(connection)
+                connection.fix_address((server_host, server_port))
+            await self.ensure_connection(connection)
         except BaseException:
             # Release the connection back to the pool so that we don't
             # leak it

redis/connection.py

@@ -1135,6 +1135,7 @@ def __init__(
         self,
         connection_class=Connection,
         max_connections: Optional[int] = None,
+        index_available_connections: bool = False,
         **connection_kwargs,
     ):
         max_connections = max_connections or 2**31
@@ -1154,6 +1155,7 @@ def __init__(
         # will notice the first thread already did the work and simply
         # release the lock.
         self._fork_lock = threading.Lock()
+        self._index_available_connections= index_available_connections
         self.reset()
 
     def __repr__(self) -> (str, str):
@@ -1175,7 +1177,9 @@ def cleanup(self, **options):
     def reset(self) -> None:
         self._lock = threading.Lock()
         self._created_connections = 0
-        self._available_connections = ConnectionsIndexer()
+        self._available_connections: ConnectionsIndexer | list = (
+            ConnectionsIndexer() if self._index_available_connections else []
+        )
         self._in_use_connections = set()
 
         # this must be the last operation in this method. while reset() is

redis/sentinel.py

@@ -20,6 +20,11 @@ class SlaveNotFoundError(ConnectionError):
 class SentinelManagedConnection(Connection):
     def __init__(self, **kwargs):
         self.connection_pool = kwargs.pop("connection_pool")
+        # To be set to True if we want to prevent
+        # the sentinel managed connection to connect
+        # to the most relevant sentinel in the pool and just
+        # connect to the current self.host and self.port
+        self._is_address_fixed = False
         super().__init__(**kwargs)
 
     def __repr__(self):
@@ -33,6 +38,10 @@ def __repr__(self):
             s = s % host_info
         return s
 
+    def fix_address(self, address):
+        self.host, self.port = address
+        self._is_address_fixed = True
+
     def connect_to(self, address):
         self.host, self.port = address
         super().connect()
@@ -41,12 +50,12 @@ def connect_to(self, address):
             if str_if_bytes(self.read_response()) != "PONG":
                 raise ConnectionError("PING failed")
 
-    def _connect_retry(self, same_address: bool = False):
+    def _connect_retry(self):
         if self._sock:
             return  # already connected
-        # If same_server is True, it means that the connection
+        # If address is fixed, it means that the connection
         # is not rotating to the next slave (if the connection pool is not master)
-        if same_address:
+        if self._is_address_fixed:
             self.connect_to((self.host, self.port))
             return
         # If same_server is False, connnect to master in master mode
@@ -66,40 +75,6 @@ def connect(self):
             lambda: self._connect_retry(), lambda error: None
         )
 
-    def connect_to_same_address(self):
-        """
-        Similar to connect, but instead of rotating to the next slave (in replica mode),
-        it just connects to the same address of the connection object.
-        """
-        return self.retry.call_with_retry(
-            lambda: self._connect_retry(same_address=True), lambda error: None
-        )
-
-    def connect_to_address(self, address):
-        """
-        Similar to connect, but instead of rotating to the next slave (in replica mode),
-        it just connects to the address supplied.
-        """
-        self.host, self.port = address
-        return self.connect_to_same_address()
-
-    def can_read_same_address(self, timeout=0):
-        """
-        Similar to can_read_same_address, but calls
-        connect_to_same_address instead of connect
-        """
-        sock = self._sock
-        if not sock:
-            self.connect_to_same_address()
-
-        host_error = self._host_error()
-
-        try:
-            return self._parser.can_read(timeout)
-        except OSError as e:
-            self.disconnect()
-            raise ConnectionError(f"Error while reading from {host_error}: {e.args}")
-
     def read_response(
         self,
         disable_decoding=False,
@@ -203,7 +178,7 @@ def __init__(self, service_name, sentinel_manager, **kwargs):
             service_name=service_name,
             sentinel_manager=sentinel_manager,
         )
-        super().__init__(**kwargs)
+        super().__init__(index_available_connections=True, **kwargs)
         self.connection_kwargs["connection_pool"] = self.proxy
         self.service_name = service_name
         self.sentinel_manager = sentinel_manager
@@ -238,31 +213,6 @@ def rotate_slaves(self):
         "Round-robin slave balancer"
         return self.proxy.rotate_slaves()
 
-    def ensure_connection_connected_to_address(
-        self, connection: SentinelManagedConnection
-    ):
-        """
-        Ensure the connection is already connected to the server that this connection
-        object wants to connect to.
-
-        Similar to self.ensure_connection, but calling connection.connect()
-        in SentinelManagedConnection (replica mode) will cause the
-        connection object to connect to the next replica in rotation,
-        and we don't wnat behavior. Look at get_connection inline docs for details.
-
-        Here, we just try to make sure that the connection is already connected
-        to the replica we wanted it to.
-        """
-        connection.connect_to_same_address()
-        try:
-            if connection.can_read_same_address() and connection.client_cache is None:
-                raise ConnectionError("Connection has data")
-        except (ConnectionError, OSError):
-            connection.disconnect()
-            connection.connect_to_same_address()
-            if connection.can_read():
-                raise ConnectionError("Connection has data")
-
     def cleanup(self, **options):
         """
         Remove the SCAN ITER family command's request id from the dictionary
@@ -309,7 +259,7 @@ def get_connection(
                 host=server_host, port=server_port
             )
             # If not, make a new dummy connection object, and set its host and port
-            # to the one that we want later in the call to ``connect_to_same_address``
+            # to the one that we want later in the call to ``fix_address``
             if not connection:
                 connection = self.make_connection()
         assert connection
@@ -324,8 +274,8 @@ def get_connection(
             # connect to the previous replica.
             # This will connect to the host and port of the replica
             else:
-                connection.connect_to_address((server_host, server_port))
-            self.ensure_connection_connected_to_address(connection)
+                connection.fix_address((server_host, server_port))
+            self.ensure_connection(connection)
         except BaseException:
             # Release the connection back to the pool so that we don't
             # leak it

tests/test_connection.py

@@ -298,6 +298,7 @@ def mock_disconnect(_):
     assert called == 1
     pool.disconnect()
 
+
 @pytest.mark.parametrize(
     "conn, error, expected_message",
     [
@@ -347,6 +348,7 @@ def test_unix_socket_connection_failure():
         == "Error 2 connecting to unix:///tmp/a.sock. No such file or directory."
     )
 
+
 def test_connections_indexer_operations():
     ci = ConnectionsIndexer()
     c1 = Connection(host="1", port=2)