polymorphism for reset available connections instead

Author: agnesnatasya

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, ConnectionsIndexer
+from redis.connection import DEFAULT_RESP_VERSION
 from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider
 from redis.exceptions import (
     AuthenticationError,
@@ -1118,7 +1118,6 @@ 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
@@ -1129,11 +1128,14 @@ def __init__(
         self.connection_kwargs = connection_kwargs
         self.max_connections = max_connections
 
+<<<<<<< HEAD
         self._available_connections: ConnectionsIndexer = (
             ConnectionsIndexer() if index_available_connections else []
         )
+=======
+        self._available_connections = self.reset_available_connections()
+>>>>>>> bafbc03 (polymorphism for reset available connections instead)
         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):
@@ -1143,11 +1145,12 @@ def __repr__(self):
         )
 
     def reset(self):
-        self._available_connections: ConnectionsIndexer | list = (
-            ConnectionsIndexer() if self._index_available_connections else []
-        )
+        self._available_connections = self.reset_available_connections()
         self._in_use_connections = weakref.WeakSet()
 
+    def reset_available_connections(self):
+        return []
+
     def can_get_connection(self) -> bool:
         """Return True if a connection can be retrieved from the pool."""
         return (

redis/asyncio/sentinel.py

@@ -21,6 +21,7 @@
 )
 from redis.commands import AsyncSentinelCommands
 from redis.exceptions import ConnectionError, ReadOnlyError, ResponseError, TimeoutError
+from redis.sentinel import ConnectionsIndexer
 from redis.utils import str_if_bytes
 
 
@@ -163,6 +164,9 @@ def reset(self):
         self.master_address = None
         self.slave_rr_counter = None
 
+    def reset_available_connections(self):
+        return ConnectionsIndexer()
+
     def owns_connection(self, connection: Connection):
         check = not self.is_master or (
             self.is_master and self.master_address == (connection.host, connection.port)

redis/connection.py

@@ -6,11 +6,10 @@
 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, Iterable, List, Optional, Sequence, Type, Union
+from typing import Any, Callable, List, Optional, Sequence, Type, Union
 from urllib.parse import parse_qs, unquote, urlparse
 
 from ._cache import (
@@ -735,57 +734,6 @@ def _host_error(self):
         return f"{self.host}:{self.port}"
 
 
-class ConnectionsIndexer(Iterable):
-    """
-    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):
-        # 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
-        # Both of these DS support O(1) operations on all of the class' methods
-        self._address_to_connections = defaultdict(dict)
-
-    def pop(self):
-        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._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()
-            del self._id_to_connection[id(connection)]
-        except KeyError:
-            return None
-        return connection
-
-    def __iter__(self):
-        # 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):
     """Manages SSL connections to and from the Redis server(s).
     This class extends the Connection class, adding SSL functionality, and making
@@ -1135,7 +1083,6 @@ 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
@@ -1155,7 +1102,6 @@ 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):
@@ -1177,9 +1123,7 @@ def cleanup(self, **options):
     def reset(self) -> None:
         self._lock = threading.Lock()
         self._created_connections = 0
-        self._available_connections: ConnectionsIndexer | list = (
-            ConnectionsIndexer() if self._index_available_connections else []
-        )
+        self._available_connections = self.reset_available_connections()
         self._in_use_connections = set()
 
         # this must be the last operation in this method. while reset() is
@@ -1193,6 +1137,9 @@ def reset(self) -> None:
         # reset() and they will immediately release _fork_lock and continue on.
         self.pid = os.getpid()
 
+    def reset_available_connections(self):
+        return []
+
     def _checkpid(self) -> None:
         # _checkpid() attempts to keep ConnectionPool fork-safe on modern
         # systems. this is called by all ConnectionPool methods that

redis/sentinel.py

@@ -1,6 +1,7 @@
 import random
 import weakref
-from typing import Any, Optional
+from typing import Any, Iterable, Optional
+from collections import defaultdict
 
 from redis.client import Redis
 from redis.commands import SentinelCommands
@@ -155,6 +156,57 @@ def rotate_slaves(self):
         raise SlaveNotFoundError(f"No slave found for {self.service_name!r}")
 
 
+class ConnectionsIndexer(Iterable):
+    """
+    Data structure that stores available connections in a pool.
+    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):
+        # 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
+        # Both of these DS support O(1) operations on all of the class' methods
+        self._address_to_connections = defaultdict(dict)
+
+    def pop(self):
+        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._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()
+            del self._id_to_connection[id(connection)]
+        except KeyError:
+            return None
+        return connection
+
+    def __iter__(self):
+        # 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 SentinelConnectionPool(ConnectionPool):
     """
     Sentinel backed connection pool.
@@ -181,7 +233,7 @@ def __init__(self, service_name, sentinel_manager, **kwargs):
             service_name=service_name,
             sentinel_manager=sentinel_manager,
         )
-        super().__init__(index_available_connections=True, **kwargs)
+        super().__init__(**kwargs)
         self.connection_kwargs["connection_pool"] = self.proxy
         self.service_name = service_name
         self.sentinel_manager = sentinel_manager
@@ -198,6 +250,9 @@ def reset(self):
         super().reset()
         self.proxy.reset()
 
+    def reset_available_connections(self):
+        return ConnectionsIndexer()
+
     @property
     def master_address(self):
         return self.proxy.master_address

tests/test_connection.py

@@ -10,7 +10,6 @@
 from redis.backoff import NoBackoff
 from redis.connection import (
     Connection,
-    ConnectionsIndexer,
     SSLConnection,
     UnixDomainSocketConnection,
     parse_url,
@@ -348,17 +347,3 @@ 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)
-    ci.append(c1)
-    assert list(ci) == [c1]
-    assert ci.pop() == c1
-
-    c2 = Connection(host="3", port=4)
-    ci.append(c1)
-    ci.append(c2)
-    assert ci.get_connection("3", 4) == c2
-    assert not ci.get_connection("5", 6)
-    assert list(ci) == [c1]

tests/test_sentinel.py

@@ -5,6 +5,7 @@
 import redis.sentinel
 from redis import exceptions
 from redis.sentinel import (
+    ConnectionsIndexer,
     MasterNotFoundError,
     Sentinel,
     SentinelConnectionPool,
@@ -266,3 +267,18 @@ def mock_disconnect():
 
     assert calls == 1
     pool.disconnect()
+
+
+def test_connections_indexer_operations():
+    ci = ConnectionsIndexer()
+    c1 = Connection(host="1", port=2)
+    ci.append(c1)
+    assert list(ci) == [c1]
+    assert ci.pop() == c1
+
+    c2 = Connection(host="3", port=4)
+    ci.append(c1)
+    ci.append(c2)
+    assert ci.get_connection("3", 4) == c2
+    assert not ci.get_connection("5", 6)
+    assert list(ci) == [c1]