Add graceful clean up for responses

Author: karpetrosyan

httpcore/_async/connection_pool.py

@@ -342,6 +342,12 @@ def __init__(
         self._pool = pool
         self._status = status
 
+    @AutoBackend.graceful_call
+    async def _force_clean_up(self) -> None:
+        if hasattr(self._stream, "aclose"):
+            await self._stream.aclose()
+        await self._pool.response_closed(self._status)
+
     async def __aiter__(self) -> AsyncIterator[bytes]:
         async for part in self._stream:
             yield part
@@ -351,4 +357,4 @@ async def aclose(self) -> None:
             if hasattr(self._stream, "aclose"):
                 await self._stream.aclose()
         finally:
-            await self._pool.response_closed(self._status)
+            await self._force_clean_up()

httpcore/_sync/connection_pool.py

@@ -342,6 +342,12 @@ def __init__(
         self._pool = pool
         self._status = status
 
+    @SyncBackend.graceful_call
+    def _force_clean_up(self) -> None:
+        if hasattr(self._stream, "close"):
+            self._stream.close()
+        self._pool.response_closed(self._status)
+
     def __iter__(self) -> Iterator[bytes]:
         for part in self._stream:
             yield part
@@ -351,4 +357,4 @@ def close(self) -> None:
             if hasattr(self._stream, "close"):
                 self._stream.close()
         finally:
-            self._pool.response_closed(self._status)
+            self._force_clean_up()

httpcore/backends/auto.py

@@ -1,10 +1,14 @@
 import typing
 from typing import Optional
 
+import anyio
 import sniffio
 
 from .base import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
 
+_TP = typing.ParamSpec("_TP")
+_R = typing.TypeVar("_R")
+
 
 class AutoBackend(AsyncNetworkBackend):
     async def _init_backend(self) -> None:
@@ -50,3 +54,13 @@ async def connect_unix_socket(
     async def sleep(self, seconds: float) -> None:  # pragma: nocover
         await self._init_backend()
         return await self._backend.sleep(seconds)
+
+    @staticmethod
+    def graceful_call(
+        fnc: typing.Callable[_TP, typing.Awaitable[_R]]
+    ) -> typing.Callable[_TP, typing.Awaitable[_R]]:
+        async def inner(*args: _TP.args, **kwargs: _TP.kwargs) -> _R:
+            with anyio.CancelScope(shield=True):
+                return await fnc(*args, **kwargs)
+
+        return inner

httpcore/backends/sync.py

@@ -16,6 +16,9 @@
 from .._utils import is_socket_readable
 from .base import SOCKET_OPTION, NetworkBackend, NetworkStream
 
+_TP = typing.ParamSpec("_TP")
+_R = typing.TypeVar("_R")
+
 
 class SyncStream(NetworkStream):
     def __init__(self, sock: socket.socket) -> None:
@@ -131,3 +134,7 @@ def connect_unix_socket(
             sock.settimeout(timeout)
             sock.connect(path)
         return SyncStream(sock)
+
+    @staticmethod
+    def graceful_call(fnc: typing.Callable[_TP, _R]) -> typing.Callable[_TP, _R]:
+        return fnc