Retry response cleanup after cancellation

httpcore/_async/connection_pool.py

@@ -9,7 +9,12 @@
 from .._backends.base import SOCKET_OPTION, AsyncNetworkBackend
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
 from .._models import Origin, Proxy, Request, Response
-from .._synchronization import AsyncEvent, AsyncShieldCancellation, AsyncThreadLock
+from .._synchronization import (
+    AsyncEvent,
+    AsyncLock,
+    AsyncShieldCancellation,
+    AsyncThreadLock,
+)
 from .connection import AsyncHTTPConnection
 from .interfaces import AsyncConnectionInterface, AsyncRequestInterface
 
@@ -397,6 +402,8 @@ def __init__(
         self._pool_request = pool_request
         self._pool = pool
         self._closed = False
+        self._close_lock = AsyncLock()
+        self._closing: list[AsyncConnectionInterface] | None = None
 
     async def __aiter__(self) -> typing.AsyncIterator[bytes]:
         try:
@@ -407,14 +414,20 @@ async def __aiter__(self) -> typing.AsyncIterator[bytes]:
             raise exc from None
 
     async def aclose(self) -> None:
-        if not self._closed:
-            self._closed = True
-            with AsyncShieldCancellation():
-                if hasattr(self._stream, "aclose"):
-                    await self._stream.aclose()
-
-            with self._pool._optional_thread_lock:
-                self._pool._requests.remove(self._pool_request)
-                closing = self._pool._assign_requests_to_connections()
-
-            await self._pool._close_connections(closing)
+        with AsyncShieldCancellation():
+            async with self._close_lock:
+                if not self._closed:
+                    if self._closing is None:
+                        if hasattr(self._stream, "aclose"):
+                            await self._stream.aclose()
+
+                        with self._pool._optional_thread_lock:
+                            self._pool._requests.remove(self._pool_request)
+                            self._closing = self._pool._assign_requests_to_connections()
+
+                    # Retain detached connections until each close succeeds, so a
+                    # retry neither loses cleanup nor removes the request twice.
+                    while self._closing:
+                        await self._pool._close_connections(self._closing[:1])
+                        del self._closing[0]
+                    self._closed = True

httpcore/_async/http11.py

@@ -326,6 +326,7 @@ def __init__(self, connection: AsyncHTTP11Connection, request: Request) -> None:
         self._connection = connection
         self._request = request
         self._closed = False
+        self._close_lock = AsyncLock()
 
     async def __aiter__(self) -> typing.AsyncIterator[bytes]:
         kwargs = {"request": self._request}
@@ -342,10 +343,11 @@ async def __aiter__(self) -> typing.AsyncIterator[bytes]:
             raise exc
 
     async def aclose(self) -> None:
-        if not self._closed:
-            self._closed = True
-            async with Trace("response_closed", logger, self._request):
-                await self._connection._response_closed()
+        async with self._close_lock:
+            if not self._closed:
+                async with Trace("response_closed", logger, self._request):
+                    await self._connection._response_closed()
+                self._closed = True
 
 
 class AsyncHTTP11UpgradeStream(AsyncNetworkStream):

httpcore/_sync/connection_pool.py

@@ -9,7 +9,12 @@
 from .._backends.base import SOCKET_OPTION, NetworkBackend
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
 from .._models import Origin, Proxy, Request, Response
-from .._synchronization import Event, ShieldCancellation, ThreadLock
+from .._synchronization import (
+    Event,
+    Lock,
+    ShieldCancellation,
+    ThreadLock,
+)
 from .connection import HTTPConnection
 from .interfaces import ConnectionInterface, RequestInterface
 
@@ -397,6 +402,8 @@ def __init__(
         self._pool_request = pool_request
         self._pool = pool
         self._closed = False
+        self._close_lock = Lock()
+        self._closing: list[ConnectionInterface] | None = None
 
     def __iter__(self) -> typing.Iterator[bytes]:
         try:
@@ -407,14 +414,20 @@ def __iter__(self) -> typing.Iterator[bytes]:
             raise exc from None
 
     def close(self) -> None:
-        if not self._closed:
-            self._closed = True
-            with ShieldCancellation():
-                if hasattr(self._stream, "close"):
-                    self._stream.close()
-
-            with self._pool._optional_thread_lock:
-                self._pool._requests.remove(self._pool_request)
-                closing = self._pool._assign_requests_to_connections()
-
-            self._pool._close_connections(closing)
+        with ShieldCancellation():
+            with self._close_lock:
+                if not self._closed:
+                    if self._closing is None:
+                        if hasattr(self._stream, "close"):
+                            self._stream.close()
+
+                        with self._pool._optional_thread_lock:
+                            self._pool._requests.remove(self._pool_request)
+                            self._closing = self._pool._assign_requests_to_connections()
+
+                    # Retain detached connections until each close succeeds, so a
+                    # retry neither loses cleanup nor removes the request twice.
+                    while self._closing:
+                        self._pool._close_connections(self._closing[:1])
+                        del self._closing[0]
+                    self._closed = True

httpcore/_sync/http11.py

@@ -326,6 +326,7 @@ def __init__(self, connection: HTTP11Connection, request: Request) -> None:
         self._connection = connection
         self._request = request
         self._closed = False
+        self._close_lock = Lock()
 
     def __iter__(self) -> typing.Iterator[bytes]:
         kwargs = {"request": self._request}
@@ -342,10 +343,11 @@ def __iter__(self) -> typing.Iterator[bytes]:
             raise exc
 
     def close(self) -> None:
-        if not self._closed:
-            self._closed = True
-            with Trace("response_closed", logger, self._request):
-                self._connection._response_closed()
+        with self._close_lock:
+            if not self._closed:
+                with Trace("response_closed", logger, self._request):
+                    self._connection._response_closed()
+                self._closed = True
 
 
 class HTTP11UpgradeStream(NetworkStream):

pyproject.toml

@@ -56,13 +56,18 @@ Source = "https://github.com/encode/httpcore"
 path = "httpcore/__init__.py"
 
 [tool.hatch.build.targets.sdist]
+# Keep distribution metadata compatible with the Twine version used by CI.
+core-metadata-version = "2.4"
 include = [
     "/httpcore",
     "/CHANGELOG.md",
     "/README.md",
     "/tests"
 ]
 
+[tool.hatch.build.targets.wheel]
+core-metadata-version = "2.4"
+
 [tool.hatch.metadata.hooks.fancy-pypi-readme]
 content-type = "text/markdown"
 

tests/test_cancellations.py

@@ -1,3 +1,4 @@
+import asyncio
 import typing
 
 import anyio
@@ -95,6 +96,68 @@ async def connect_tcp(
         return SlowReadStream(self._buffer)
 
 
+class InterruptibleCloseStream(SlowReadStream):
+    def __init__(self, buffer: typing.List[bytes]):
+        super().__init__(buffer)
+        self.close_calls = 0
+
+    async def aclose(self) -> None:
+        self.close_calls += 1
+        if self.close_calls == 1:
+            raise asyncio.CancelledError
+
+
+class InterruptibleCloseBackend(httpcore.AsyncNetworkBackend):
+    def __init__(self, buffer: typing.List[bytes]):
+        self.stream = InterruptibleCloseStream(buffer)
+
+    async def connect_tcp(
+        self,
+        host: str,
+        port: int,
+        timeout: typing.Optional[float] = None,
+        local_address: typing.Optional[str] = None,
+        socket_options: typing.Optional[typing.Iterable[httpcore.SOCKET_OPTION]] = None,
+    ) -> httpcore.AsyncNetworkStream:
+        return self.stream
+
+
+@pytest.mark.parametrize("read_body", [False, True])
+def test_connection_pool_retries_interrupted_response_close(read_body):
+    async def run_test() -> None:
+        network_backend = InterruptibleCloseBackend(
+            [
+                b"HTTP/1.1 200 OK\r\n",
+                b"Content-Length: 13\r\n",
+                b"\r\n",
+                b"Hello, world!",
+            ]
+        )
+        async with httpcore.AsyncConnectionPool(
+            network_backend=network_backend, max_keepalive_connections=0
+        ) as pool:
+            response = await pool.handle_async_request(
+                httpcore.Request(
+                    "GET", "http://example.com", headers={"Host": "example.com"}
+                )
+            )
+
+            if read_body:
+                await response.aread()
+
+            with pytest.raises(asyncio.CancelledError):
+                await response.aclose()
+
+            assert bool(pool.connections) is not read_body
+            await response.aclose()
+            assert not pool.connections
+            assert network_backend.stream.close_calls == 2
+            await response.aclose()
+            assert network_backend.stream.close_calls == 2
+
+    asyncio.run(run_test())
+
+
 @pytest.mark.anyio
 async def test_connection_pool_timeout_during_request():
     """