Simplify BlockingConnectionPool

Author: kristjanvalur

redis/asyncio/connection.py

@@ -1015,6 +1015,13 @@ def reset(self):
         self._available_connections = []
         self._in_use_connections = set()
 
+    def can_get_connection(self) -> bool:
+        """Return True if a connection can be retrieved from the pool."""
+        return (
+            self._available_connections
+            or self._created_connections < self.max_connections
+        )
+
     async def get_connection(self, command_name, *keys, **options):
         """Get a connection from the pool"""
         try:
@@ -1102,19 +1109,19 @@ class BlockingConnectionPool(ConnectionPool):
     """
     A blocking connection pool::
 
-        >>> from redis.client import Redis
+        >>> from redis.asyncio.client import Redis
         >>> client = Redis(connection_pool=BlockingConnectionPool())
 
     It performs the same function as the default
-    :py:class:`~redis.ConnectionPool` implementation, in that,
+    :py:class:`~redis.asyncio.ConnectionPool` implementation, in that,
     it maintains a pool of reusable connections that can be shared by
     multiple async redis clients.
 
     The difference is that, in the event that a client tries to get a
     connection from the pool when all of connections are in use, rather than
     raising a :py:class:`~redis.ConnectionError` (as the default
-    :py:class:`~redis.ConnectionPool` implementation does), it
-    makes the client wait ("blocks") for a specified number of seconds until
+    :py:class:`~redis.asyncio.ConnectionPool` implementation does), it
+    makes blocks the current `Task` for a specified number of seconds until
     a connection becomes available.
 
     Use ``max_connections`` to increase / decrease the pool size::
@@ -1137,107 +1144,29 @@ def __init__(
         max_connections: int = 50,
         timeout: Optional[int] = 20,
         connection_class: Type[AbstractConnection] = Connection,
-        queue_class: Type[asyncio.Queue] = asyncio.LifoQueue,
+        queue_class: Type[asyncio.Queue] = asyncio.LifoQueue,  # deprecated
         **connection_kwargs,
     ):
 
-        self.queue_class = queue_class
-        self.timeout = timeout
-        self._connections: List[AbstractConnection]
         super().__init__(
             connection_class=connection_class,
             max_connections=max_connections,
             **connection_kwargs,
         )
-        self._lock = asyncio.Lock()
-
-    def reset(self):
-        # Create and fill up a queue with ``None`` values.
-        self.pool = self.queue_class(self.max_connections)
-        while True:
-            try:
-                self.pool.put_nowait(None)
-            except asyncio.QueueFull:
-                break
-
-        # Keep a list of actual connection instances so that we can
-        # disconnect them later.
-        self._connections = []
-
-    def make_connection(self):
-        """Make a fresh connection."""
-        connection = self.connection_class(**self.connection_kwargs)
-        self._connections.append(connection)
-        return connection
+        self._condition = asyncio.Condition()
+        self.timeout = timeout
 
     async def get_connection(self, command_name, *keys, **options):
-        """
-        Get a connection, blocking for ``self.timeout`` until a connection
-        is available from the pool.
-
-        If the connection returned is ``None`` then creates a new connection.
-        Because we use a last-in first-out queue, the existing connections
-        (having been returned to the pool after the initial ``None`` values
-        were added) will be returned before ``None`` values. This means we only
-        create new connections when we need to, i.e.: the actual number of
-        connections will only increase in response to demand.
-        """
-
-        # Try and get a connection from the pool. If one isn't available within
-        # self.timeout then raise a ``ConnectionError``.
-        connection = None
         try:
             async with async_timeout(self.timeout):
-                connection = await self.pool.get()
-        except (asyncio.QueueEmpty, asyncio.TimeoutError):
-            # Note that this is not caught by the redis client and will be
-            # raised unless handled by application code. If you want never to
-            raise ConnectionError("No connection available.")
-
-        # If the ``connection`` is actually ``None`` then that's a cue to make
-        # a new connection to add to the pool.
-        if connection is None:
-            connection = self.make_connection()
-
-        try:
-            # ensure this connection is connected to Redis
-            await 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 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
-        except BaseException:
-            # release the connection back to the pool so that we don't leak it
-            await self.release(connection)
-            raise
-
-        return connection
+                async with self._condition:
+                    await self._condition.wait_for(self.can_get_connection)
+                    return await super().get_connection(command_name, *keys, **options)
+        except asyncio.TimeoutError as err:
+            raise ConnectionError("No connection available.") from err
 
     async def release(self, connection: AbstractConnection):
         """Releases the connection back to the pool."""
-        # Put the connection back into the pool.
-        try:
-            self.pool.put_nowait(connection)
-        except asyncio.QueueFull:
-            # perhaps the pool has been reset() after a fork? regardless,
-            # we don't want this connection
-            pass
-
-    async def disconnect(self, inuse_connections: bool = True):
-        """Disconnects all connections in the pool."""
-        async with self._lock:
-            resp = await asyncio.gather(
-                *(connection.disconnect() for connection in self._connections),
-                return_exceptions=True,
-            )
-            exc = next((r for r in resp if isinstance(r, BaseException)), None)
-            if exc:
-                raise exc
+        async with self._condition:
+            await super().release(connection)
+            self._condition.notify()

tests/test_asyncio/test_connection_pool.py

@@ -94,6 +94,9 @@ class DummyConnection(Connection):
     def __init__(self, **kwargs):
         self.kwargs = kwargs
 
+    def repr_pieces(self):
+        return [("id", id(self)), ("kwargs", self.kwargs)]
+
     async def connect(self):
         pass