Fix async clients safety when used as an async context manager (#3512)

* Fix async safety when Redis client is used as an async context manager

When the async Redis client is used as an async context manager and
called from different corotuines, one coroutine can exit, shutting
down the client's connection pool, while another coroutine is
attempting to use a connection. This results in a connection error,
such as:

redis.exceptions.ConnectionError: Connection closed by server.

Additional locking in `ConnectionPool` resolves the problem but
introduces extreme latency due to the locking. Instead, this PR
implements a shielded counter that increments as callers enter the async
context manager and decrements when they exit. The client then closes
its connection pool only after all active contexts exit.

Performance is on par with use of the client without a context manager.

Author: abrookins

CONTRIBUTING.md

@@ -81,6 +81,10 @@ using `invoke standalone-tests`; similarly, RedisCluster tests can be run by usi
 Each run of tests starts and stops the various dockers required. Sometimes
 things get stuck, an `invoke clean` can help.
 
+## Linting
+
+Call `invoke linters` to run linters without also running tests.
+
 ## Documentation
 
 If relevant, update the code documentation, via docstrings, or in `/docs`.

redis/asyncio/client.py

@@ -387,6 +387,12 @@ def __init__(
         # on a set of redis commands
         self._single_conn_lock = asyncio.Lock()
 
+        # 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()
+
     def __repr__(self):
         return (
             f"<{self.__class__.__module__}.{self.__class__.__name__}"
@@ -594,10 +600,47 @@ def client(self) -> "Redis":
         )
 
     async def __aenter__(self: _RedisT) -> _RedisT:
-        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.
+        """
+        await self._increment_usage()
+        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
+            await self._decrement_usage()
+            raise
+
+    async def _increment_usage(self) -> int:
+        """
+        Helper coroutine to increment 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 _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):
-        await self.aclose()
+        """
+        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())
 
     _DEL_MESSAGE = "Unclosed Redis client"
 

redis/asyncio/cluster.py

@@ -431,6 +431,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:
@@ -467,10 +473,47 @@ 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.
+        """
+        await self._increment_usage()
+        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
+            await self._decrement_usage()
+            raise
 
-    async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
-        await self.aclose()
+    async def _increment_usage(self) -> int:
+        """
+        Helper coroutine to increment 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 _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

@@ -0,0 +1,16 @@
+import asyncio
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_usage_counter(r):
+    async def dummy_task():
+        async with r:
+            await asyncio.sleep(0.01)
+
+    tasks = [dummy_task() for _ in range(20)]
+    await asyncio.gather(*tasks)
+
+    # After all tasks have completed, the usage counter should be back to zero.
+    assert r._usage_counter == 0