End CRT stream when source is empty

Author: JordonPhillips

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

@@ -337,6 +337,7 @@ async def _consume_body_async(
     ) -> None:
         async for chunk in source:
             dest.write(chunk)
+        dest.end_stream()
 
     def __deepcopy__(self, memo: Any) -> "AWSCRTHTTPClient":
         return AWSCRTHTTPClient(
@@ -354,6 +355,7 @@ def __init__(self) -> None:
         # will be much more efficient than a list.
         self._chunks: deque[bytes] = deque()
         self._closed = False
+        self._done = False
 
     def read(self, size: int | None = -1) -> bytes:
         if self._closed:
@@ -363,21 +365,24 @@ def read(self, size: int | None = -1) -> bytes:
             # When the CRT recieves this, it'll try again later.
             raise BlockingIOError("read")
 
-        chunk = self._chunks.popleft()
-        if size is None or size < 1:
-            # We could compile all the chunks here instead of just returning
-            # the one, BUT the CRT will keep calling read until empty bytes
-            # are returned. So it's actually better to just return one chunk
-            # since combining them would have some potentially bad memory
-            # usage issues.
-            return chunk
-        else:
-            result = chunk[:size]
-            remainder = chunk[size:]
+        # We could compile all the chunks here instead of just returning
+        # the one, BUT the CRT will keep calling read until empty bytes
+        # are returned. So it's actually better to just return one chunk
+        # since combining them would have some potentially bad memory
+        # usage issues.
+        result = self._chunks.popleft()
+        if size is not None and size > 0:
+            remainder = result[size:]
+            result = result[:size]
             if remainder:
                 self._chunks.appendleft(remainder)
             return result
 
+        if self._done and len(self._chunks) == 0:
+            self.close()
+
+        return result
+
     def read1(self, size: int = -1) -> bytes:
         return self.read(size)
 
@@ -410,6 +415,11 @@ def closed(self) -> bool:
 
     def close(self) -> None:
         self._closed = True
+        self._done = True
 
         # Clear out the remaining chunks so that they don't sit around in memory.
         self._chunks.clear()
+
+    def end_stream(self) -> None:
+        """End the stream, letting any remaining chunks be read before it is closed."""
+        self._done = True

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

@@ -91,3 +91,13 @@ def test_stream_readinto_bytearray() -> None:
     stream.write(b"foobar")
     stream.readinto(buffer)
     assert bytes(buffer) == b"foo"
+
+
+def test_end_stream() -> None:
+    stream = BufferableByteStream()
+    stream.write(b"foo")
+    stream.end_stream()
+
+    assert not stream.closed
+    assert stream.read() == b"foo"
+    assert stream.closed