smithy-http: stream aiohttp response bodies

packages/smithy-http/.changes/next-release/smithy-http-bugfix-b0c1d627a5c34c5b953a785d049c0bb6.json

@@ -0,0 +1,4 @@
+{
+  "type": "bugfix",
+  "description": "Stream aiohttp response bodies incrementally to enable true output event streams."
+}

packages/smithy-http/src/smithy_http/aio/aiohttp.py

@@ -1,5 +1,6 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
+from collections.abc import AsyncIterable, AsyncIterator
 from copy import copy, deepcopy
 from itertools import chain
 from typing import TYPE_CHECKING, Any, Self
@@ -22,7 +23,6 @@
 
 from smithy_core.aio.interfaces import StreamingBlob
 from smithy_core.aio.types import AsyncBytesReader
-from smithy_core.aio.utils import async_list
 from smithy_core.exceptions import MissingDependencyError
 from smithy_core.interfaces import URI
 
@@ -49,13 +49,28 @@ def __post_init__(self) -> None:
         _assert_aiohttp()
 
 
+class _AIOHTTPStreamingBody(AsyncIterable[bytes]):
+    """Streams a response body, releasing the response once it is done."""
+
+    def __init__(self, response: "aiohttp.ClientResponse") -> None:
+        self._response = response
+
+    def __aiter__(self) -> AsyncIterator[bytes]:
+        # The reader iterates by line, so use iter_any to get raw chunks.
+        return self._response.content.iter_any()
+
+    async def close(self) -> None:
+        # Pools the connection if the body was read to completion, closes it if not.
+        self._response.release()
+
+
 class AIOHTTPClient(HTTPClient):
     """Implementation of :py:class:`.interfaces.HTTPClient` using aiohttp."""
 
     TIMEOUT_EXCEPTIONS = (TimeoutError,)
 
-    # aiohttp has no HTTP/2 support and this client fully buffers the response
-    # before returning, so it can never interleave request and response data.
+    # aiohttp has no HTTP/2 support, so it can never interleave request and
+    # response data.
     SUPPORTS_DUPLEX_STREAMING = False
 
     def __init__(
@@ -113,15 +128,19 @@ async def send(
         # The typing on `params` is incorrect, it'll happily accept a mapping whose
         # values are lists (or tuples) and produce expected values.
         # See: https://github.com/aio-libs/aiohttp/issues/8563
-        async with self._session.request(
+        resp = await self._session.request(
             method=request.method,
             url=self._serialize_uri_without_query(request.destination),
             params=parse_qs(request.destination.query),  # type: ignore
             headers=headers_list,
             data=body,
             allow_redirects=False,
-        ) as resp:
-            return await self._marshal_response(resp)
+        )
+        try:
+            return self._marshal_response(resp)
+        except BaseException:
+            resp.release()
+            raise
 
     async def close(self) -> None:
         """Close the underlying aiohttp session and its connection pool."""
@@ -163,7 +182,7 @@ def _serialize_uri_without_query(self, uri: URI) -> yarl.URL:
             encoded=True,
         )
 
-    async def _marshal_response(
+    def _marshal_response(
         self, aiohttp_resp: "aiohttp.ClientResponse"
     ) -> HTTPResponseInterface:
         """Convert a ``aiohttp.ClientResponse`` to a ``smithy_http.aio.HTTPResponse``"""
@@ -181,7 +200,7 @@ async def _marshal_response(
         return HTTPResponse(
             status=aiohttp_resp.status,
             fields=headers,
-            body=async_list([await aiohttp_resp.read()]),
+            body=_AIOHTTPStreamingBody(aiohttp_resp),
             reason=aiohttp_resp.reason,
         )
 

packages/smithy-http/tests/unit/aio/test_aiohttp.py

@@ -3,14 +3,14 @@
 #  pyright: reportPrivateUsage=false
 from collections.abc import AsyncIterator
 from typing import Any, cast
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from smithy_core import URI
 from smithy_core.aio.types import AsyncBytesReader
 from smithy_http import Field, Fields
 from smithy_http.aio import HTTPRequest
-from smithy_http.aio.aiohttp import AIOHTTPClient
+from smithy_http.aio.aiohttp import AIOHTTPClient, _AIOHTTPStreamingBody
 from smithy_http.exceptions import SmithyHTTPError
 
 
@@ -20,11 +20,19 @@ def _create_client() -> tuple[AIOHTTPClient, MagicMock]:
 
     session = MagicMock()
     session.close = AsyncMock()
-    session.request.return_value.__aenter__ = AsyncMock(return_value=response)
-    session.request.return_value.__aexit__ = AsyncMock(return_value=None)
+    session.request = AsyncMock(return_value=response)
     return AIOHTTPClient(_session=cast(Any, session)), session
 
 
+def _create_request() -> HTTPRequest:
+    return HTTPRequest(
+        method="GET",
+        destination=URI(scheme="https", host="example.com", path="/"),
+        body=AsyncBytesReader(b""),
+        fields=Fields(),
+    )
+
+
 async def test_close_closes_session() -> None:
     client, session = _create_client()
 
@@ -57,14 +65,8 @@ async def test_context_manager_closes_session() -> None:
 
 async def test_send_omits_empty_async_reader_body() -> None:
     client, session = _create_client()
-    request = HTTPRequest(
-        method="GET",
-        destination=URI(scheme="https", host="example.com", path="/"),
-        body=AsyncBytesReader(b""),
-        fields=Fields(),
-    )
 
-    await client.send(request)
+    await client.send(_create_request())
 
     assert session.request.call_args.kwargs["data"] is None
 
@@ -86,18 +88,86 @@ async def test_send_preserves_explicitly_framed_empty_body() -> None:
 
 async def test_send_disables_redirects() -> None:
     client, session = _create_client()
-    request = HTTPRequest(
-        method="GET",
-        destination=URI(scheme="https", host="example.com", path="/"),
-        body=AsyncBytesReader(b""),
-        fields=Fields(),
-    )
 
-    await client.send(request)
+    await client.send(_create_request())
 
     assert session.request.call_args.kwargs["allow_redirects"] is False
 
 
+async def test_send_streams_response_body_and_releases_it_on_close() -> None:
+    async def chunks() -> AsyncIterator[bytes]:
+        yield b"first"
+        yield b"second"
+
+    client, session = _create_client()
+    aiohttp_response = session.request.return_value
+    aiohttp_response.content.iter_any.return_value = chunks()
+
+    response = await client.send(_create_request())
+
+    aiohttp_response.read.assert_not_awaited()
+    aiohttp_response.content.iter_any.assert_not_called()
+    assert isinstance(response.body, _AIOHTTPStreamingBody)
+    assert [chunk async for chunk in response.body] == [
+        b"first",
+        b"second",
+    ]
+    aiohttp_response.content.iter_any.assert_called_once_with()
+    aiohttp_response.release.assert_not_called()
+
+    await response.body.close()
+    aiohttp_response.release.assert_called_once_with()
+
+
+async def test_non_streaming_response_body_can_be_consumed() -> None:
+    async def chunks() -> AsyncIterator[bytes]:
+        yield b'{"message":'
+        yield b'"hello"}'
+
+    client, session = _create_client()
+    aiohttp_response = session.request.return_value
+    aiohttp_response.content.iter_any.return_value = chunks()
+
+    response = await client.send(_create_request())
+
+    assert await response.consume_body_async() == b'{"message":"hello"}'
+    aiohttp_response.content.iter_any.assert_called_once_with()
+
+
+async def test_response_body_close_releases_partially_consumed_response() -> None:
+    async def chunks() -> AsyncIterator[bytes]:
+        yield b"first"
+        yield b"second"
+
+    client, session = _create_client()
+    aiohttp_response = session.request.return_value
+    aiohttp_response.content.iter_any.return_value = chunks()
+
+    response = await client.send(_create_request())
+    assert isinstance(response.body, _AIOHTTPStreamingBody)
+    body_iterator = aiter(response.body)
+    assert await anext(body_iterator) == b"first"
+
+    await response.body.close()
+
+    aiohttp_response.release.assert_called_once_with()
+
+
+async def test_send_releases_response_when_marshaling_fails() -> None:
+    client, session = _create_client()
+    aiohttp_response = session.request.return_value
+
+    with (
+        patch.object(
+            client, "_marshal_response", side_effect=ValueError("invalid response")
+        ),
+        pytest.raises(ValueError, match="invalid response"),
+    ):
+        await client.send(_create_request())
+
+    aiohttp_response.release.assert_called_once_with()
+
+
 async def test_prepare_body_preserves_nonempty_reader_position() -> None:
     client, _ = _create_client()
     body = AsyncBytesReader(b"request body")