more features for ConnectionsINdexer

Author: agnesnatasya

redis/asyncio/sentinel.py

@@ -257,7 +257,7 @@ async def get_connection(
             # This will connect to the host and port of the replica
             else:
                 await connection.connect_to_address(server_host, server_port)
-            self.ensure_connection(connection)
+            await self.ensure_connection(connection)
         except BaseException:
             # Release the connection back to the pool so that we don't
             # leak it

redis/connection.py

@@ -737,39 +737,53 @@ def _host_error(self):
 
 class ConnectionsIndexer(Iterable):
     """
-    Data structure that manages a list of available connections which
-    is also indexed based on the address (ip and port) of the connection
+    Data structure that simulates a list of available connections. 
+    Instead of list, we keep 2 additional DS to support O(1) operations
+    on all of the class' methods.
+    The first DS is indexed on the connection object's ID. 
+    The second DS is indexed on the address (ip and port) of the connection.
     """
 
     def __init__(self):
-        self._connections = []
+        # Map the id to the connection object
+        self._id_to_connection = {}
         # Map the address to a dictionary of connections
         # The inner dictionary is a map between the object id to the object itself
-        # This is to support O(1) operations on all of the class' methods
+        # Both of these DS support O(1) operations on all of the class' methods
         self._address_to_connections = defaultdict(dict)
 
     def pop(self):
-        connection = self._connections.pop()
-        del self._address_to_connections[(connection.host, connection.port)][
-            id(connection)
-        ]
+        try:
+            _, connection = self._id_to_connection.popitem()
+            del self._address_to_connections[(connection.host, connection.port)][
+                id(connection)
+            ]
+        except KeyError:
+            # We are simulating a list, hence we raise IndexError
+            # when there's no item in the dictionary
+            raise IndexError()
         return connection
 
     def append(self, connection: Connection):
-        self._connections.append(connection)
+        self._id_to_connection[id(connection)] = connection
         self._address_to_connections[(connection.host, connection.port)][
             id(connection)
         ] = connection
 
     def get_connection(self, host: str, port: int):
         try:
-            connection = self._address_to_connections[(host, port)].popitem()
+            _, connection = self._address_to_connections[(host, port)].popitem()
+            del self._id_to_connection[id(connection)]
         except KeyError:
             return None
         return connection
 
     def __iter__(self):
-        return iter(self._connections)
+        # This is an O(1) operation in python3.7 and later
+        return iter(self._id_to_connection.values())
+
+    def __len__(self):
+        return len(self._id_to_connection)
 
 
 class SSLConnection(Connection):

redis/sentinel.py

@@ -41,9 +41,16 @@ def connect_to(self, address):
             if str_if_bytes(self.read_response()) != "PONG":
                 raise ConnectionError("PING failed")
 
-    def _connect_retry(self):
+    def _connect_retry(self, same_server: bool = False):
         if self._sock:
             return  # already connected
+        # If same_server is True, it means that the connection
+        # is not rotating to the next slave (if the connection pool is not master)
+        if same_server:
+            self.connect_to(self.host, self.port)
+            return
+        # If same_server is False, connnect to master in master mode 
+        # and rotate to the next slave in slave mode
         if self.connection_pool.is_master:
             self.connect_to(self.connection_pool.get_master_address())
         else:
@@ -54,8 +61,11 @@ def _connect_retry(self):
                     continue
             raise SlaveNotFoundError  # Never be here
 
-    def connect(self):
-        return self.retry.call_with_retry(self._connect_retry, lambda error: None)
+    def connect(self, same_server: bool = False):
+        return self.retry.call_with_retry(
+                lambda: self._connect_retry(same_server), 
+                lambda error: None
+        )
 
     def read_response(
         self,
@@ -195,6 +205,29 @@ 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(same_address=True)
+        try:
+            if connection.can_read(same_address=True) and connection.client_cache is None:
+                raise ConnectionError("Connection has data")
+        except (ConnectionError, OSError):
+            connection.disconnect()
+            connection.connect(same_address=True)
+            if connection.can_read():
+                raise ConnectionError("Connection has data")
+
     def cleanup_scan(self, **options):
         """
         Remove the SCAN ITER family command's request id from the dictionary
@@ -248,6 +281,7 @@ def get_connection(
                 connection = self.make_connection()
         assert connection
         self._in_use_connections.add(connection)
+        breakpoint()
         try:
             # Ensure this connection is connected to Redis
             # If this is the first scan request, it will
@@ -259,7 +293,9 @@ def get_connection(
             # This will connect to the host and port of the replica
             else:
                 connection.connect_to_address(server_host, server_port)
-            self.ensure_connection(connection)
+            breakpoint()
+            self.ensure_connection_connected_to_address(connection)
+            breakpoint()
         except BaseException:
             # Release the connection back to the pool so that we don't
             # leak it
@@ -270,6 +306,7 @@ def get_connection(
             connection.host,
             connection.port,
         )
+        breakpoint()
         return connection