From b886731bd81dbbe1a6e6b92072583a9f572de05f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= Date: Wed, 11 Oct 2023 13:22:09 +0000 Subject: [PATCH] Refactor how asyncio.BlockingConnectionPool gets connection. --- redis/asyncio/connection.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index 1ef9960ff3..519c1bc071 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -1028,7 +1028,7 @@ def can_get_connection(self) -> bool: async def get_connection(self, command_name, *keys, **options): """Get a connected connection from the pool""" - connection = self.get_available_connection() + connection = await self.get_available_connection() try: await self.ensure_connection(connection) except BaseException: @@ -1037,7 +1037,7 @@ async def get_connection(self, command_name, *keys, **options): return connection - def get_available_connection(self): + async def get_available_connection(self): """Get a connection from the pool, without making sure it is connected""" try: connection = self._available_connections.pop() @@ -1167,24 +1167,16 @@ def __init__( self._condition = asyncio.Condition() self.timeout = timeout - async def get_connection(self, command_name, *keys, **options): + async def get_available_connection(self): """Gets a connection from the pool, blocking until one is available""" try: async with self._condition: async with async_timeout(self.timeout): await self._condition.wait_for(self.can_get_connection) - connection = super().get_available_connection() + return await super().get_available_connection() except asyncio.TimeoutError as err: raise ConnectionError("No connection available.") from err - # We now perform the connection check outside of the lock. - try: - await self.ensure_connection(connection) - return connection - except BaseException: - await self.release(connection) - raise - async def release(self, connection: AbstractConnection): """Releases the connection back to the pool.""" async with self._condition: