fix test, apply logic to async cluster client

Author: abrookins

redis/asyncio/client.py

@@ -570,7 +570,8 @@ def client(self) -> "Redis":
     async def __aenter__(self: _RedisT) -> _RedisT:
         """
         Async context manager entry. Increments a usage counter so that the
-        connection pool is only closed (via aclose()) when no one is using the client.
+        connection pool is only closed (via aclose()) when no context is using
+        the client.
         """
         async with self._usage_lock:
             self._usage_counter += 1

redis/asyncio/cluster.py

@@ -379,6 +379,12 @@ def __init__(
         self._initialize = True
         self._lock: Optional[asyncio.Lock] = None
 
+        # When used as an async context manager, we need to increment and decrement
+        # a usage counter so that we can close the connection pool when no one is
+        # using the client.
+        self._usage_counter = 0
+        self._usage_lock = asyncio.Lock()
+
     async def initialize(self) -> "RedisCluster":
         """Get all nodes from startup nodes & creates connections if not initialized."""
         if self._initialize:
@@ -415,10 +421,40 @@ async def close(self) -> None:
         await self.aclose()
 
     async def __aenter__(self) -> "RedisCluster":
-        return await self.initialize()
+        """
+        Async context manager entry. Increments a usage counter so that the
+        connection pool is only closed (via aclose()) when no context is using
+        the client.
+        """
+        async with self._usage_lock:
+            self._usage_counter += 1
+        try:
+            # Initialize the client (i.e. establish connection, etc.)
+            return await self.initialize()
+        except Exception:
+            # If initialization fails, decrement the counter to keep it in sync
+            async with self._usage_lock:
+                self._usage_counter -= 1
+            raise
 
-    async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
-        await self.aclose()
+    async def _decrement_usage(self) -> int:
+        """
+        Helper coroutine to decrement the usage counter while holding the lock.
+        Returns the new value of the usage counter.
+        """
+        async with self._usage_lock:
+            self._usage_counter -= 1
+            return self._usage_counter
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        """
+        Async context manager exit. Decrements a usage counter. If this is the
+        last exit (counter becomes zero), the client closes its connection pool.
+        """
+        current_usage = await asyncio.shield(self._decrement_usage())
+        if current_usage == 0:
+            # This was the last active context, so disconnect the pool.
+            await asyncio.shield(self.aclose())
 
     def __await__(self) -> Generator[Any, None, "RedisCluster"]:
         return self.initialize().__await__()

tests/test_asyncio/test_usage_counter.py

@@ -1,12 +1,11 @@
 import asyncio
 
 import pytest
+import redis
 
 
 @pytest.mark.asyncio
-async def test_usage_counter(create_redis):
-    r = await create_redis(decode_responses=True)
-
+async def test_usage_counter(r):
     async def dummy_task():
         async with r:
             await asyncio.sleep(0.01)