implement in sync client

Author: agnesnatasya

redis/asyncio/client.py

@@ -651,6 +651,8 @@ async def execute_command(self, *args, **options):
         finally:
             if not self.connection:
                 await pool.release(conn)
+                if "ITER" in command_name.upper():
+                    pool.cleanup_scan(iter_req_id=options.get("_iter_req_id", None))
 
     async def parse_response(
         self, connection: Connection, command_name: Union[str, bytes], **options

redis/asyncio/connection.py

@@ -38,7 +38,7 @@
 
 from redis.asyncio.retry import Retry
 from redis.backoff import NoBackoff
-from redis.connection import DEFAULT_RESP_VERSION
+from redis.connection import DEFAULT_RESP_VERSION, ConnectionsIndexer
 from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider
 from redis.exceptions import (
     AuthenticationError,
@@ -1057,6 +1057,13 @@ class ConnectionPool:
     ``connection_class``.
     """
 
+    @abstractmethod
+    def cleanup_scan(self, **options):
+        """
+        Additional cleanup operations that the connection pool might
+        need to do after a SCAN ITER family command is executed
+        """
+
     @classmethod
     def from_url(cls: Type[_CP], url: str, **kwargs) -> _CP:
         """
@@ -1118,7 +1125,7 @@ def __init__(
         self.connection_kwargs = connection_kwargs
         self.max_connections = max_connections
 
-        self._available_connections: List[AbstractConnection] = []
+        self._available_connections: ConnectionsIndexer = ConnectionsIndexer()
         self._in_use_connections: Set[AbstractConnection] = set()
         self.encoder_class = self.connection_kwargs.get("encoder_class", Encoder)
 
@@ -1324,3 +1331,6 @@ async def release(self, connection: AbstractConnection):
         async with self._condition:
             await super().release(connection)
             self._condition.notify()
+
+    def cleanup_scan(self, **options):
+        pass

redis/asyncio/sentinel.py

@@ -147,7 +147,7 @@ def __init__(self, service_name, sentinel_manager, **kwargs):
         self.sentinel_manager = sentinel_manager
         self.master_address = None
         self.slave_rr_counter = None
-        self._request_id_to_replica_address = {}
+        self._iter_req_id_to_replica_address = {}
 
     def __repr__(self):
         return (
@@ -193,6 +193,14 @@ async def rotate_slaves(self) -> AsyncIterator:
             pass
         raise SlaveNotFoundError(f"No slave found for {self.service_name!r}")
 
+    def cleanup_scan(self, **options):
+        """
+        Remove the SCAN ITER family command's request id from the dictionary
+        """
+        self._iter_req_id_to_replica_address.pop(
+            options.get("_iter_req_id", None), None
+        )
+
     async def get_connection(
         self, command_name: str, *keys: Any, **options: Any
     ) -> SentinelManagedConnection:
@@ -217,7 +225,7 @@ async def get_connection(
         (
             server_host,
             server_port,
-        ) = self._request_id_to_replica_address.get(iter_req_id, (None, None))
+        ) = self._iter_req_id_to_replica_address.get(iter_req_id, (None, None))
         connection = None
         # If this is the first scan request of the iter command,
         # get a connection from the pool
@@ -228,18 +236,12 @@ async def get_connection(
                 connection = self.make_connection()
         # If this is not the first scan request of the iter command
         else:
-            # Check from the available connections, if any of the connection
-            # is connected to the host and port that we want
-            for available_connection in self._available_connections.copy():
-                # if yes, use that connection
-                if (
-                    available_connection.host == server_host
-                    and available_connection.port == server_port
-                ):
-                    self._available_connections.remove(available_connection)
-                    connection = available_connection
-            # 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_address``
+            # Get the connection that has the same host and port
+            connection = self._available_connections.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_address``
             if not connection:
                 connection = self.make_connection()
         assert connection
@@ -255,25 +257,14 @@ 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)
-            # Connections that the pool provides should be ready to send
-            # a command. If not, the connection was either returned to the
-            # pool before all data has been read or the socket has been
-            # closed. Either way, reconnect and verify everything is good.
-            try:
-                if await connection.can_read_destructive():
-                    raise ConnectionError("Connection has data") from None
-            except (ConnectionError, OSError):
-                await connection.disconnect()
-                await connection.connect()
-                if await connection.can_read_destructive():
-                    raise ConnectionError("Connection not ready") from None
+            self.ensure_connection(connection)
         except BaseException:
             # Release the connection back to the pool so that we don't
             # leak it
             await self.release(connection)
             raise
         # Store the connection to the dictionary
-        self._request_id_to_replica_address[iter_req_id] = (
+        self._iter_req_id_to_replica_address[iter_req_id] = (
             connection.host,
             connection.port,
         )

redis/commands/core.py

@@ -3041,9 +3041,15 @@ def scan_iter(
             Additionally, Redis modules can expose other types as well.
         """
         cursor = "0"
+        iter_req_id = uuid.uuid4()
         while cursor != 0:
             cursor, data = self.scan(
-                cursor=cursor, match=match, count=count, _type=_type, **kwargs
+                cursor=cursor,
+                match=match,
+                count=count,
+                _type=_type,
+                _iter_req_id=iter_req_id,
+                **kwargs,
             )
             yield from data
 
@@ -3087,8 +3093,11 @@ def sscan_iter(
         ``count`` allows for hint the minimum number of returns
         """
         cursor = "0"
+        iter_req_id = uuid.uuid4()
         while cursor != 0:
-            cursor, data = self.sscan(name, cursor=cursor, match=match, count=count)
+            cursor, data = self.sscan(
+                name, cursor=cursor, match=match, count=count, _iter_req_id=iter_req_id
+            )
             yield from data
 
     def hscan(
@@ -3139,9 +3148,15 @@ def hscan_iter(
         ``no_values`` indicates to return only the keys, without values
         """
         cursor = "0"
+        iter_req_id = uuid.uuid4()
         while cursor != 0:
             cursor, data = self.hscan(
-                name, cursor=cursor, match=match, count=count, no_values=no_values
+                name,
+                cursor=cursor,
+                match=match,
+                count=count,
+                no_values=no_values,
+                _iter_req_id=iter_req_id,
             )
             if no_values:
                 yield from data
@@ -3195,13 +3210,15 @@ def zscan_iter(
         ``score_cast_func`` a callable used to cast the score return value
         """
         cursor = "0"
+        iter_req_id = uuid.uuid4()
         while cursor != 0:
             cursor, data = self.zscan(
                 name,
                 cursor=cursor,
                 match=match,
                 count=count,
                 score_cast_func=score_cast_func,
+                _iter_req_id=iter_req_id,
             )
             yield from data
 

redis/connection.py

@@ -6,10 +6,11 @@
 import threading
 import weakref
 from abc import abstractmethod
+from collections import defaultdict
 from itertools import chain
 from queue import Empty, Full, LifoQueue
 from time import time
-from typing import Any, Callable, List, Optional, Sequence, Type, Union
+from typing import Any, Callable, Iterable, List, Optional, Sequence, Type, Union
 from urllib.parse import parse_qs, unquote, urlparse
 
 from ._cache import (
@@ -734,6 +735,43 @@ def _host_error(self):
         return f"{self.host}:{self.port}"
 
 
+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
+    """
+
+    def __init__(self):
+        self._connections = []
+        # 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
+        self._address_to_connections = defaultdict(dict)
+
+    def pop(self):
+        connection = self._connections.pop()
+        del self._address_to_connections[(connection.host, connection.port)][
+            id(connection)
+        ]
+        return connection
+
+    def append(self, connection: Connection):
+        self._connections.append(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()
+        except KeyError:
+            return None
+        return connection
+
+    def __iter__(self):
+        return iter(self._connections)
+
+
 class SSLConnection(Connection):
     """Manages SSL connections to and from the Redis server(s).
     This class extends the Connection class, adding SSL functionality, and making
@@ -1107,7 +1145,7 @@ def __repr__(self) -> (str, str):
     def reset(self) -> None:
         self._lock = threading.Lock()
         self._created_connections = 0
-        self._available_connections = []
+        self._available_connections = ConnectionsIndexer()
         self._in_use_connections = set()
 
         # this must be the last operation in this method. while reset() is
@@ -1168,6 +1206,25 @@ def _checkpid(self) -> None:
             finally:
                 self._fork_lock.release()
 
+    def ensure_connection(self, connection: AbstractConnection):
+        # ensure this connection is connected to Redis
+        connection.connect()
+        # if client caching is not enabled connections that the pool
+        # provides should be ready to send a command.
+        # if not, the connection was either returned to the
+        # pool before all data has been read or the socket has been
+        # closed. either way, reconnect and verify everything is good.
+        # (if caching enabled the connection will not always be ready
+        # to send a command because it may contain invalidation messages)
+        try:
+            if connection.can_read() and connection.client_cache is None:
+                raise ConnectionError("Connection has data")
+        except (ConnectionError, OSError):
+            connection.disconnect()
+            connection.connect()
+            if connection.can_read():
+                raise ConnectionError("Connection not ready")
+
     def get_connection(self, command_name: str, *keys, **options) -> "Connection":
         "Get a connection from the pool"
         self._checkpid()
@@ -1179,23 +1236,7 @@ def get_connection(self, command_name: str, *keys, **options) -> "Connection":
             self._in_use_connections.add(connection)
 
         try:
-            # ensure this connection is connected to Redis
-            connection.connect()
-            # if client caching is not enabled connections that the pool
-            # provides should be ready to send a command.
-            # if not, the connection was either returned to the
-            # pool before all data has been read or the socket has been
-            # closed. either way, reconnect and verify everything is good.
-            # (if caching enabled the connection will not always be ready
-            # to send a command because it may contain invalidation messages)
-            try:
-                if connection.can_read() and connection.client_cache is None:
-                    raise ConnectionError("Connection has data")
-            except (ConnectionError, OSError):
-                connection.disconnect()
-                connection.connect()
-                if connection.can_read():
-                    raise ConnectionError("Connection not ready")
+            self.ensure_connection(connection)
         except BaseException:
             # release the connection back to the pool so that we don't
             # leak it
@@ -1408,20 +1449,7 @@ def get_connection(self, command_name, *keys, **options):
             connection = self.make_connection()
 
         try:
-            # ensure this connection is connected to Redis
-            connection.connect()
-            # connections that the pool provides should be ready to send
-            # a command. if not, the connection was either returned to the
-            # pool before all data has been read or the socket has been
-            # closed. either way, reconnect and verify everything is good.
-            try:
-                if connection.can_read():
-                    raise ConnectionError("Connection has data")
-            except (ConnectionError, OSError):
-                connection.disconnect()
-                connection.connect()
-                if connection.can_read():
-                    raise ConnectionError("Connection not ready")
+            self.ensure_connection(connection)
         except BaseException:
             # release the connection back to the pool so that we don't leak it
             self.release(connection)