Safe async cancellations. (#880)

* Connection pool work

* Connection pool work

* Connection pool work

* Connection pool work

* Comments

* Comments

* Connection pool work

* Reraise

* Lookin sharp

* nocover directive

* Safe cancellations

* Update CHANGELOG

Author: tomchristie

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## Unreleased
 
+- Fix support for async cancellations. (#880)
 - Fix trace extension when used with socks proxy. (#849)
 - Fix SSL context for connections using the "wss" scheme (#869)
 

httpcore/_async/connection.py

@@ -6,7 +6,7 @@
 
 from .._backends.auto import AutoBackend
 from .._backends.base import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
-from .._exceptions import ConnectError, ConnectionNotAvailable, ConnectTimeout
+from .._exceptions import ConnectError, ConnectTimeout
 from .._models import Origin, Request, Response
 from .._ssl import default_ssl_context
 from .._synchronization import AsyncLock
@@ -70,9 +70,9 @@ async def handle_async_request(self, request: Request) -> Response:
                 f"Attempted to send request to {request.url.origin} on connection to {self._origin}"
             )
 
-        async with self._request_lock:
-            if self._connection is None:
-                try:
+        try:
+            async with self._request_lock:
+                if self._connection is None:
                     stream = await self._connect(request)
 
                     ssl_object = stream.get_extra_info("ssl_object")
@@ -94,11 +94,9 @@ async def handle_async_request(self, request: Request) -> Response:
                             stream=stream,
                             keepalive_expiry=self._keepalive_expiry,
                         )
-                except Exception as exc:
-                    self._connect_failed = True
-                    raise exc
-            elif not self._connection.is_available():
-                raise ConnectionNotAvailable()
+        except BaseException as exc:
+            self._connect_failed = True
+            raise exc
 
         return await self._connection.handle_async_request(request)