Wrap completion future in crt response body

Author: JordonPhillips

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

@@ -3,6 +3,7 @@
 #  pyright: reportMissingTypeStubs=false,reportUnknownMemberType=false
 #  flake8: noqa: F811
 import asyncio
+from asyncio import Future as AsyncFuture
 from concurrent.futures import Future as ConcurrentFuture
 from collections import deque
 from collections.abc import AsyncGenerator, AsyncIterable
@@ -105,6 +106,7 @@ def __repr__(self) -> str:
 class CRTResponseBody:
     def __init__(self) -> None:
         self._stream: "crt_http.HttpClientStream | None" = None
+        self._completion_future: AsyncFuture[int] | None = None
         self._chunk_futures: deque[ConcurrentFuture[bytes]] = deque()
 
         # deque is thread safe and the crt is only going to be writing
@@ -117,7 +119,9 @@ def set_stream(self, stream: "crt_http.HttpClientStream") -> None:
         if self._stream is not None:
             raise SmithyHTTPException("Stream already set on AWSCRTHTTPResponse object")
         self._stream = stream
-        self._stream.completion_future.add_done_callback(self._on_complete)
+        concurrent_future: ConcurrentFuture[int] = stream.completion_future
+        self._completion_future = asyncio.wrap_future(concurrent_future)
+        self._completion_future.add_done_callback(self._on_complete)
         self._stream.activate()
 
     def on_body(self, chunk: bytes, **kwargs: Any) -> None:  # pragma: crt-callback
@@ -129,21 +133,21 @@ def on_body(self, chunk: bytes, **kwargs: Any) -> None:  # pragma: crt-callback
             self._received_chunks.append(chunk)
 
     async def next(self) -> bytes:
-        if self._stream is None:
+        if self._completion_future is None:
             raise SmithyHTTPException("Stream not set")
 
         # TODO: update backpressure window once CRT supports it
         if self._received_chunks:
             return self._received_chunks.popleft()
-        elif self._stream.completion_future.done():
+        elif self._completion_future.done():
             return b""
         else:
             future = ConcurrentFuture[bytes]()
             self._chunk_futures.append(future)
             return await asyncio.wrap_future(future)
 
     def _on_complete(
-        self, completion_future: ConcurrentFuture[int]
+        self, completion_future: AsyncFuture[int]
     ) -> None:  # pragma: crt-callback
         for future in self._chunk_futures:
             future.set_result(b"")

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

@@ -1,14 +1,18 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
+import asyncio
 from copy import deepcopy
 from io import BytesIO
+from unittest.mock import Mock
+from concurrent.futures import Future as ConcurrentFuture
 
 import pytest
+from awscrt.http import HttpClientStream  # type: ignore
 
 from smithy_core import URI
 from smithy_http import Fields
 from smithy_http.aio import HTTPRequest
-from smithy_http.aio.crt import AWSCRTHTTPClient, BufferableByteStream
+from smithy_http.aio.crt import AWSCRTHTTPClient, BufferableByteStream, CRTResponseBody
 
 
 def test_deepcopy_client() -> None:
@@ -136,3 +140,77 @@ def test_end_stream() -> None:
     assert not stream.closed
     assert stream.read() == b"foo"
     assert stream.closed
+
+
+async def test_response_body_completed_stream() -> None:
+    completion_future = ConcurrentFuture[int]()
+    mock_stream = Mock(spec=HttpClientStream)
+    mock_stream.completion_future = completion_future
+
+    response_body = CRTResponseBody()
+    response_body.set_stream(mock_stream)
+    completion_future.set_result(200)
+
+    assert await response_body.next() == b""
+
+
+async def test_response_body_empty_stream() -> None:
+    completion_future = ConcurrentFuture[int]()
+    mock_stream = Mock(spec=HttpClientStream)
+    mock_stream.completion_future = completion_future
+
+    response_body = CRTResponseBody()
+    response_body.set_stream(mock_stream)
+
+    read_task = asyncio.create_task(response_body.next())
+
+    # Sleep briefly so the read task gets priority. It should
+    # add a chunk future and then await it.
+    await asyncio.sleep(0.01)
+
+    assert len(response_body._chunk_futures) == 1  # type: ignore
+    response_body.on_body(b"foo")
+    assert await read_task == b"foo"
+
+
+async def test_response_body_stream_completion_clears_buffer() -> None:
+    completion_future = ConcurrentFuture[int]()
+    mock_stream = Mock(spec=HttpClientStream)
+    mock_stream.completion_future = completion_future
+
+    response_body = CRTResponseBody()
+    response_body.set_stream(mock_stream)
+
+    read_tasks = (
+        asyncio.create_task(response_body.next()),
+        asyncio.create_task(response_body.next()),
+        asyncio.create_task(response_body.next()),
+        asyncio.create_task(response_body.next()),
+    )
+
+    # Sleep briefly so the read tasks gets priority. It should
+    # add a chunk future and then await it.
+    await asyncio.sleep(0.01)
+
+    assert len(response_body._chunk_futures) == 4  # type: ignore
+    completion_future.set_result(200)
+    await asyncio.sleep(0.01)
+
+    # Tasks should have been drained
+    assert len(response_body._chunk_futures) == 0  # type: ignore
+
+    # Tasks should still be awaited, and should all return empty
+    results = asyncio.gather(*read_tasks)
+    assert results.result() == [b"", b"", b"", b""]
+
+
+async def test_response_body_non_empty_stream() -> None:
+    completion_future = ConcurrentFuture[int]()
+    mock_stream = Mock(spec=HttpClientStream)
+    mock_stream.completion_future = completion_future
+
+    response_body = CRTResponseBody()
+    response_body.set_stream(mock_stream)
+    response_body.on_body(b"foo")
+
+    assert await response_body.next() == b"foo"