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

redis/asyncio/client.py

@@ -362,6 +362,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__}"
@@ -562,10 +568,40 @@ 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 one is using the client.
+        """
+        async with self._usage_lock:
+            self._usage_counter += 1
+            current_usage = self._usage_counter
+        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 _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"
 
@@ -1347,9 +1383,7 @@ async def _disconnect_reset_raise(self, conn, error):
         # indicates the user should retry this transaction.
         if self.watching:
             await self.aclose()
-            raise WatchError(
-                "A ConnectionError occurred on while watching one or more keys"
-            )
+            raise
         # if retry_on_error is not set or the error is not one
         # of the specified error types, raise it
         if (

redis/asyncio/connection.py

@@ -1157,19 +1157,20 @@ async def disconnect(self, inuse_connections: bool = True):
         current in use, potentially by other tasks. Otherwise only disconnect
         connections that are idle in the pool.
         """
-        if inuse_connections:
-            connections: Iterable[AbstractConnection] = chain(
-                self._available_connections, self._in_use_connections
+        async with self._lock:
+            if inuse_connections:
+                connections: Iterable[AbstractConnection] = chain(
+                    self._available_connections, self._in_use_connections
+                )
+            else:
+                connections = self._available_connections
+            resp = await asyncio.gather(
+                *(connection.disconnect() for connection in connections),
+                return_exceptions=True,
             )
-        else:
-            connections = self._available_connections
-        resp = await asyncio.gather(
-            *(connection.disconnect() for connection in connections),
-            return_exceptions=True,
-        )
-        exc = next((r for r in resp if isinstance(r, BaseException)), None)
-        if exc:
-            raise exc
+            exc = next((r for r in resp if isinstance(r, BaseException)), None)
+            if exc:
+                raise exc
 
     async def aclose(self) -> None:
         """Close the pool, disconnecting all connections"""

tests/test_asyncio/test_usage_counter.py

@@ -0,0 +1,18 @@
+import asyncio
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_usage_counter(create_redis):
+    r = await create_redis(decode_responses=True)
+
+    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