Turn acquire_connection into a context manager

CHANGES

@@ -67,6 +67,7 @@
     * Close Unix sockets if the connection attempt fails. This prevents `ResourceWarning`s. (#3314)
     * Close SSL sockets if the connection attempt fails, or if validations fail. (#3317)
     * Eliminate mutable default arguments in the `redis.commands.core.Script` class. (#3332)
+    * Turn `acquire_connection` into a context manager (#3435)
 
 * 4.1.3 (Feb 8, 2022)
   * Fix flushdb and flushall (#1926)

redis/asyncio/cluster.py

@@ -4,6 +4,7 @@
 import socket
 import ssl
 import warnings
+from contextlib import contextmanager
 from typing import (
     Any,
     Callable,
@@ -1006,16 +1007,23 @@ async def disconnect(self) -> None:
         if exc:
             raise exc
 
+    @contextmanager
     def acquire_connection(self) -> Connection:
+        """Context manager acquiring a connection on enter and automatically
+        freeing it on exit."""
         try:
-            return self._free.popleft()
+            connection = self._free.popleft()
         except IndexError:
             if len(self._connections) < self.max_connections:
                 connection = self.connection_class(**self.connection_kwargs)
                 self._connections.append(connection)
-                return connection
+            else:
+                raise MaxConnectionsError()
 
-            raise MaxConnectionsError()
+        try:
+            yield connection
+        finally:
+            self._free.append(connection)
 
     async def parse_response(
         self, connection: Connection, command: str, **kwargs: Any
@@ -1045,42 +1053,35 @@ async def parse_response(
 
     async def execute_command(self, *args: Any, **kwargs: Any) -> Any:
         # Acquire connection
-        connection = self.acquire_connection()
+        with self.acquire_connection() as connection:
 
-        # Execute command
-        await connection.send_packed_command(connection.pack_command(*args), False)
+            # Execute command
+            await connection.send_packed_command(connection.pack_command(*args), False)
 
-        # Read response
-        try:
+            # Read response
             return await self.parse_response(connection, args[0], **kwargs)
-        finally:
-            # Release connection
-            self._free.append(connection)
 
     async def execute_pipeline(self, commands: List["PipelineCommand"]) -> bool:
         # Acquire connection
-        connection = self.acquire_connection()
+        with self.acquire_connection() as connection:
 
-        # Execute command
-        await connection.send_packed_command(
-            connection.pack_commands(cmd.args for cmd in commands), False
-        )
-
-        # Read responses
-        ret = False
-        for cmd in commands:
-            try:
-                cmd.result = await self.parse_response(
-                    connection, cmd.args[0], **cmd.kwargs
-                )
-            except Exception as e:
-                cmd.result = e
-                ret = True
+            # Execute command
+            await connection.send_packed_command(
+                connection.pack_commands(cmd.args for cmd in commands), False
+            )
 
-        # Release connection
-        self._free.append(connection)
+            # Read responses
+            ret = False
+            for cmd in commands:
+                try:
+                    cmd.result = await self.parse_response(
+                        connection, cmd.args[0], **cmd.kwargs
+                    )
+                except Exception as e:
+                    cmd.result = e
+                    ret = True
 
-        return ret
+            return ret
 
 
 class NodesManager:

tests/test_asyncio/test_cluster.py

@@ -360,20 +360,20 @@ async def test_cluster_set_get_retry_object(self, request: FixtureRequest):
                 assert n_retry._retries == retry._retries
                 assert isinstance(n_retry._backoff, NoBackoff)
             rand_cluster_node = r.get_random_node()
-            existing_conn = rand_cluster_node.acquire_connection()
-            # Change retry policy
-            new_retry = Retry(ExponentialBackoff(), 3)
-            r.set_retry(new_retry)
-            assert r.get_retry()._retries == new_retry._retries
-            assert isinstance(r.get_retry()._backoff, ExponentialBackoff)
-            for node in r.get_nodes():
-                n_retry = node.connection_kwargs.get("retry")
-                assert n_retry is not None
-                assert n_retry._retries == new_retry._retries
-                assert isinstance(n_retry._backoff, ExponentialBackoff)
-            assert existing_conn.retry._retries == new_retry._retries
-            new_conn = rand_cluster_node.acquire_connection()
-            assert new_conn.retry._retries == new_retry._retries
+            with rand_cluster_node.acquire_connection() as existing_conn:
+                # Change retry policy
+                new_retry = Retry(ExponentialBackoff(), 3)
+                r.set_retry(new_retry)
+                assert r.get_retry()._retries == new_retry._retries
+                assert isinstance(r.get_retry()._backoff, ExponentialBackoff)
+                for node in r.get_nodes():
+                    n_retry = node.connection_kwargs.get("retry")
+                    assert n_retry is not None
+                    assert n_retry._retries == new_retry._retries
+                    assert isinstance(n_retry._backoff, ExponentialBackoff)
+                assert existing_conn.retry._retries == new_retry._retries
+                with rand_cluster_node.acquire_connection() as new_conn:
+                    assert new_conn.retry._retries == new_retry._retries
 
     async def test_cluster_retry_object(self, request: FixtureRequest) -> None:
         url = request.config.getoption("--redis-url")