Ensure that all async generators are explicitly closed

Author: agronholm

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [Unreleased]
+
+- Explicitly close all async generators to ensure predictable behavior
+
 ## Version 1.0.9 (April 24th, 2025)
 
 - Resolve https://github.com/advisories/GHSA-vqfr-h8mv-ghfj with h11 dependency update. (#1008)

httpcore/_async/connection_pool.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from collections.abc import AsyncGenerator
+
 import ssl
 import sys
 import types
@@ -10,9 +12,13 @@
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
 from .._models import Origin, Proxy, Request, Response
 from .._synchronization import AsyncEvent, AsyncShieldCancellation, AsyncThreadLock
+from .._utils import aclosing
 from .connection import AsyncHTTPConnection
 from .interfaces import AsyncConnectionInterface, AsyncRequestInterface
 
+if typing.TYPE_CHECKING:
+    from .http11 import HTTP11ConnectionByteStream
+    from .http2 import HTTP2ConnectionByteStream
 
 class AsyncPoolRequest:
     def __init__(self, request: Request) -> None:
@@ -389,7 +395,7 @@ def __repr__(self) -> str:
 class PoolByteStream:
     def __init__(
         self,
-        stream: typing.AsyncIterable[bytes],
+        stream: HTTP11ConnectionByteStream | HTTP2ConnectionByteStream,
         pool_request: AsyncPoolRequest,
         pool: AsyncConnectionPool,
     ) -> None:
@@ -398,20 +404,16 @@ def __init__(
         self._pool = pool
         self._closed = False
 
-    async def __aiter__(self) -> typing.AsyncIterator[bytes]:
-        try:
-            async for part in self._stream:
-                yield part
-        except BaseException as exc:
-            await self.aclose()
-            raise exc from None
+    async def __aiter__(self) -> AsyncGenerator[bytes]:
+        async with aclosing(self._stream.__aiter__()) as iterator:
+            async for chunk in iterator:
+                yield chunk
 
     async def aclose(self) -> None:
         if not self._closed:
             self._closed = True
             with AsyncShieldCancellation():
-                if hasattr(self._stream, "aclose"):
-                    await self._stream.aclose()
+                await self._stream.aclose()
 
             with self._pool._optional_thread_lock:
                 self._pool._requests.remove(self._pool_request)

httpcore/_async/http11.py

@@ -6,6 +6,7 @@
 import time
 import types
 import typing
+from collections.abc import AsyncGenerator
 
 import h11
 
@@ -21,6 +22,7 @@
 from .._synchronization import AsyncLock, AsyncShieldCancellation
 from .._trace import Trace
 from .interfaces import AsyncConnectionInterface
+from .._utils import aclosing
 
 logger = logging.getLogger("httpcore.http11")
 
@@ -193,9 +195,7 @@ async def _receive_response_headers(
 
         return http_version, event.status_code, event.reason, headers, trailing_data
 
-    async def _receive_response_body(
-        self, request: Request
-    ) -> typing.AsyncIterator[bytes]:
+    async def _receive_response_body(self, request: Request) -> AsyncGenerator[bytes]:
         timeouts = request.extensions.get("timeout", {})
         timeout = timeouts.get("read", None)
 
@@ -327,12 +327,13 @@ def __init__(self, connection: AsyncHTTP11Connection, request: Request) -> None:
         self._request = request
         self._closed = False
 
-    async def __aiter__(self) -> typing.AsyncIterator[bytes]:
+    async def __aiter__(self) -> AsyncGenerator[bytes]:
         kwargs = {"request": self._request}
         try:
             async with Trace("receive_response_body", logger, self._request, kwargs):
-                async for chunk in self._connection._receive_response_body(**kwargs):
-                    yield chunk
+                async with aclosing(self._connection._receive_response_body(**kwargs)) as body:
+                    async for chunk in body:
+                        yield chunk
         except BaseException as exc:
             # If we get an exception while streaming the response,
             # we want to close the response (and possibly the connection)

httpcore/_async/http2.py

@@ -5,6 +5,7 @@
 import time
 import types
 import typing
+from collections.abc import AsyncGenerator
 
 import h2.config
 import h2.connection
@@ -308,7 +309,7 @@ async def _receive_response(
 
     async def _receive_response_body(
         self, request: Request, stream_id: int
-    ) -> typing.AsyncIterator[bytes]:
+    ) -> AsyncGenerator[bytes]:
         """
         Iterator that returns the bytes of the response body for a given stream ID.
         """
@@ -568,7 +569,7 @@ def __init__(
         self._stream_id = stream_id
         self._closed = False
 
-    async def __aiter__(self) -> typing.AsyncIterator[bytes]:
+    async def __aiter__(self) -> AsyncGenerator[bytes]:
         kwargs = {"request": self._request, "stream_id": self._stream_id}
         try:
             async with Trace("receive_response_body", logger, self._request, kwargs):

httpcore/_async/interfaces.py

@@ -2,6 +2,7 @@
 
 import contextlib
 import typing
+from collections.abc import AsyncGenerator
 
 from .._models import (
     URL,
@@ -56,9 +57,9 @@ async def stream(
         url: URL | bytes | str,
         *,
         headers: HeaderTypes = None,
-        content: bytes | typing.AsyncIterator[bytes] | None = None,
+        content: bytes | AsyncGenerator[bytes] | None = None,
         extensions: Extensions | None = None,
-    ) -> typing.AsyncIterator[Response]:
+    ) -> AsyncGenerator[Response]:
         # Strict type checking on our parameters.
         method = enforce_bytes(method, name="method")
         url = enforce_url(url, name="url")

httpcore/_models.py

@@ -4,6 +4,9 @@
 import ssl
 import typing
 import urllib.parse
+from collections.abc import AsyncGenerator
+
+from ._utils import aclosing
 
 # Functions for typechecking...
 
@@ -151,7 +154,7 @@ def __init__(self, content: bytes) -> None:
     def __iter__(self) -> typing.Iterator[bytes]:
         yield self._content
 
-    async def __aiter__(self) -> typing.AsyncIterator[bytes]:
+    async def __aiter__(self) -> AsyncGenerator[bytes]:
         yield self._content
 
     def __repr__(self) -> str:
@@ -463,10 +466,11 @@ async def aread(self) -> bytes:
                 "You should use 'response.read()' instead."
             )
         if not hasattr(self, "_content"):
-            self._content = b"".join([part async for part in self.aiter_stream()])
+            async with aclosing(self.aiter_stream()) as parts:
+                self._content = b"".join([part async for part in parts])
         return self._content
 
-    async def aiter_stream(self) -> typing.AsyncIterator[bytes]:
+    async def aiter_stream(self) -> AsyncGenerator[bytes]:
         if not isinstance(self.stream, typing.AsyncIterable):  # pragma: nocover
             raise RuntimeError(
                 "Attempted to stream an synchronous response using 'async for ... in "
@@ -479,8 +483,9 @@ async def aiter_stream(self) -> typing.AsyncIterator[bytes]:
                 "more than once."
             )
         self._stream_consumed = True
-        async for chunk in self.stream:
-            yield chunk
+        async with aclosing(self.stream) as parts:
+            async for chunk in parts:
+                yield chunk
 
     async def aclose(self) -> None:
         if not isinstance(self.stream, typing.AsyncIterable):  # pragma: nocover

httpcore/_utils.py

@@ -4,6 +4,19 @@
 import socket
 import sys
 
+if sys.version_info >= (3, 10):
+    from contextlib import aclosing as aclosing
+else:
+    class aclosing(AbstractAsyncContextManager):
+        def __init__(self, thing):
+            self.thing = thing
+
+        async def __aenter__(self):
+            return self.thing
+
+        async def __aexit__(self, *exc_info):
+            await self.thing.aclose()
+
 
 def is_socket_readable(sock: socket.socket | None) -> bool:
     """

tests/test_cancellations.py

@@ -171,7 +171,6 @@ async def test_h11_timeout_during_response():
         assert conn.is_closed()
 
 
-@pytest.mark.xfail
 @pytest.mark.anyio
 async def test_h2_timeout_during_handshake():
     """
@@ -186,7 +185,6 @@ async def test_h2_timeout_during_handshake():
         assert conn.is_closed()
 
 
-@pytest.mark.xfail
 @pytest.mark.anyio
 async def test_h2_timeout_during_request():
     """
@@ -207,7 +205,6 @@ async def test_h2_timeout_during_request():
         assert conn.is_idle()
 
 
-@pytest.mark.xfail
 @pytest.mark.anyio
 async def test_h2_timeout_during_response():
     """

tests/test_models.py

@@ -163,6 +163,9 @@ async def __aiter__(self) -> typing.AsyncIterator[bytes]:
         for chunk in self._chunks:
             yield chunk
 
+    async def aclose(self) -> None:
+        pass
+
 
 @pytest.mark.trio
 async def test_response_async_read():