Ensure streams get closed

Author: JordonPhillips

packages/aws-event-stream/src/aws_event_stream/_private/deserializers.py

@@ -40,9 +40,18 @@ def __init__(
         self._source = source
         self._is_client_mode = is_client_mode
         self._deserializer = deserializer
+        self._closed = False
 
     async def receive(self) -> E | None:
-        event = await Event.decode_async(self._source)
+        if self._closed:
+            return None
+
+        try:
+            event = await Event.decode_async(self._source)
+        except Exception as e:
+            await self.close()
+            raise IOError("Failed to read from stream.") from e
+
         if event is None:
             return None
 
@@ -57,10 +66,18 @@ async def receive(self) -> E | None:
         return result
 
     async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+
         if (close := getattr(self._source, "close", None)) is not None:
             if asyncio.iscoroutine(result := close()):
                 await result
 
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
 
 class EventDeserializer(SpecificShapeDeserializer):
     def __init__(

packages/aws-event-stream/src/aws_event_stream/_private/serializers.py

@@ -50,8 +50,11 @@ def __init__(
         self._serializer = EventSerializer(
             payload_codec=payload_codec, is_client_mode=is_client_mode
         )
+        self._closed = False
 
     async def send(self, event: E) -> None:
+        if self._closed:
+            raise IOError("Attempted to write to closed stream.")
         event.serialize(self._serializer)
         result = self._serializer.get_result()
         if result is None:
@@ -60,13 +63,27 @@ async def send(self, event: E) -> None:
             )
         if self._signer is not None:
             result = self._signer(result)
-        await self._writer.write(result.encode())
+
+        encoded_result = result.encode()
+        try:
+            await self._writer.write(encoded_result)
+        except Exception as e:
+            await self.close()
+            raise IOError("Failed to write to stream.") from e
 
     async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+
         if (close := getattr(self._writer, "close", None)) is not None:
             if asyncio.iscoroutine(result := close()):
                 await result
 
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
 
 class EventSerializer(SpecificShapeSerializer):
     def __init__(

packages/aws-event-stream/src/aws_event_stream/aio/__init__.py

@@ -75,31 +75,35 @@ def __init__(
         self.response: R | None = None
 
     async def await_output(self) -> tuple[R, AsyncEventReceiver[O]]:
-        async_reader = AsyncBytesReader((await self._awaitable_response).body)
-        if self.output_stream is None:
-            self.output_stream = _AWSEventReceiver[O](
-                payload_codec=self._payload_codec,
-                source=async_reader,
-                deserializer=self._deserializer,
-                is_client_mode=self._is_client_mode,
-            )
-
-        if self.response is None:
-            if self._deserializeable_response is None:
-                initial_response = await self._awaitable_output
-            else:
-                initial_response_stream = _AWSEventReceiver(
+        try:
+            async_reader = AsyncBytesReader((await self._awaitable_response).body)
+            if self.output_stream is None:
+                self.output_stream = _AWSEventReceiver[O](
                     payload_codec=self._payload_codec,
                     source=async_reader,
-                    deserializer=self._deserializeable_response.deserialize,
+                    deserializer=self._deserializer,
                     is_client_mode=self._is_client_mode,
                 )
-                initial_response = await initial_response_stream.receive()
-                if initial_response is None:
-                    raise MissingInitialResponse()
-                self.response = initial_response
-        else:
-            initial_response = self.response
+
+            if self.response is None:
+                if self._deserializeable_response is None:
+                    initial_response = await self._awaitable_output
+                else:
+                    initial_response_stream = _AWSEventReceiver(
+                        payload_codec=self._payload_codec,
+                        source=async_reader,
+                        deserializer=self._deserializeable_response.deserialize,
+                        is_client_mode=self._is_client_mode,
+                    )
+                    initial_response = await initial_response_stream.receive()
+                    if initial_response is None:
+                        raise MissingInitialResponse()
+                    self.response = initial_response
+            else:
+                initial_response = self.response
+        except Exception:
+            await self.input_stream.close()
+            raise
 
         return initial_response, self.output_stream
 
@@ -137,7 +141,11 @@ def __init__(
 
     async def await_output(self) -> R:
         if self.response is None:
-            self.response = await self._awaitable_response
+            try:
+                self.response = await self._awaitable_response
+            except Exception:
+                await self.input_stream.close()
+                raise
         return self.response
 
 

packages/aws-event-stream/tests/unit/_private/test_deserializers.py

@@ -89,3 +89,43 @@ def test_deserialize_unmodeled_error():
 
     with pytest.raises(UnmodeledEventError, match="InternalError"):
         EventStreamOperationInputOutput.deserialize(deserializer)
+
+
+async def test_receiver_closes_source() -> None:
+    source = AsyncBytesReader(b"")
+    deserializer = EventStreamDeserializer()
+    receiver = AWSAsyncEventReceiver[Any](
+        payload_codec=JSONCodec(), source=source, deserializer=deserializer.deserialize
+    )
+    assert not receiver.closed
+    assert not source.closed
+    await receiver.close()
+    assert receiver.closed
+    assert source.closed
+
+
+async def test_read_closed_receiver() -> None:
+    source = AsyncBytesReader(b"")
+    deserializer = EventStreamDeserializer()
+    receiver = AWSAsyncEventReceiver[Any](
+        payload_codec=JSONCodec(), source=source, deserializer=deserializer.deserialize
+    )
+
+    await receiver.close()
+    assert receiver.closed
+    assert await receiver.receive() is None
+
+
+async def test_read_closed_receiver_source() -> None:
+    source = AsyncBytesReader(b"")
+    deserializer = EventStreamDeserializer()
+    receiver = AWSAsyncEventReceiver[Any](
+        payload_codec=JSONCodec(), source=source, deserializer=deserializer.deserialize
+    )
+
+    await source.close()
+    assert source.closed
+    assert not receiver.closed
+    with pytest.raises(IOError):
+        await receiver.receive()
+    assert receiver.closed

packages/aws-event-stream/tests/unit/_private/test_serializers.py

@@ -1,10 +1,16 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
+from typing import Any
+
 import pytest
 from smithy_core.serializers import SerializeableShape
+from smithy_core.aio.types import AsyncBytesProvider
 from smithy_json import JSONCodec
 
-from aws_event_stream._private.serializers import EventSerializer
+from aws_event_stream._private.serializers import (
+    EventSerializer,
+    AWSAsyncEventPublisher,
+)
 from aws_event_stream.events import EventMessage
 
 from . import EVENT_STREAM_SERDE_CASES, INITIAL_REQUEST_CASE, INITIAL_RESPONSE_CASE
@@ -36,3 +42,41 @@ def test_serialize_initial_request():
 
 def test_serialize_initial_response():
     test_event_serializer_server_mode(*INITIAL_RESPONSE_CASE)
+
+
+async def test_publisher_closes_reader():
+    writer = AsyncBytesProvider()
+    publisher: AWSAsyncEventPublisher[Any] = AWSAsyncEventPublisher(
+        payload_codec=JSONCodec(), async_writer=writer
+    )
+
+    assert not publisher.closed
+    assert not writer.closed
+    await publisher.close()
+    assert publisher.closed
+    assert writer.closed
+
+
+async def test_send_after_close():
+    writer = AsyncBytesProvider()
+    publisher: AWSAsyncEventPublisher[Any] = AWSAsyncEventPublisher(
+        payload_codec=JSONCodec(), async_writer=writer
+    )
+
+    await publisher.close()
+    assert publisher.closed
+    with pytest.raises(IOError):
+        await publisher.send(EVENT_STREAM_SERDE_CASES[0][0])
+
+
+async def test_send_to_closed_writer():
+    writer = AsyncBytesProvider()
+    publisher: AWSAsyncEventPublisher[Any] = AWSAsyncEventPublisher(
+        payload_codec=JSONCodec(), async_writer=writer
+    )
+
+    await writer.close()
+    with pytest.raises(IOError):
+        await publisher.send(EVENT_STREAM_SERDE_CASES[0][0])
+
+    assert publisher.closed

packages/smithy-core/src/smithy_core/aio/types.py

@@ -114,12 +114,12 @@ def closed(self) -> bool:
 
     async def close(self) -> None:
         """Closes the stream, as well as the underlying stream where possible."""
+        self._closed = True
         if (close := getattr(self._data, "close", None)) is not None:
             if asyncio.iscoroutine(result := close()):
                 await result
 
         self._data = None
-        self._closed = True
 
 
 class SeekableAsyncBytesReader:

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

@@ -7,6 +7,7 @@
 from collections import deque
 from collections.abc import AsyncGenerator, AsyncIterable
 from copy import deepcopy
+from functools import partial
 from io import BytesIO, BufferedIOBase
 from typing import TYPE_CHECKING, Any
 
@@ -231,7 +232,7 @@ async def send(
         :param request: The request including destination URI, fields, payload.
         :param request_config: Configuration specific to this request.
         """
-        crt_request = await self._marshal_request(request)
+        crt_request, crt_body = await self._marshal_request(request)
         connection = await self._get_connection(request.destination)
         response_body = CRTResponseBody()
         response_factory = CRTResponseFactory(response_body)
@@ -242,8 +243,17 @@ async def send(
         )
         response_factory.set_done_callback(crt_stream)
         response_body.set_stream(crt_stream)
+        crt_stream.completion_future.add_done_callback(
+            partial(self._close_input_body, body=crt_body)
+        )
         return await response_factory.await_response()
 
+    def _close_input_body(
+        self, future: ConcurrentFuture[int], *, body: "BufferableByteStream | BytesIO"
+    ) -> None:
+        if future.exception(timeout=0):
+            body.close()
+
     async def _create_connection(
         self, url: core_interfaces.URI
     ) -> "crt_http.HttpClientConnection":
@@ -314,7 +324,7 @@ def _render_path(self, url: core_interfaces.URI) -> str:
 
     async def _marshal_request(
         self, request: http_aio_interfaces.HTTPRequest
-    ) -> "crt_http.HttpRequest":
+    ) -> tuple["crt_http.HttpRequest", "BufferableByteStream | BytesIO"]:
         """Create :py:class:`awscrt.http.HttpRequest` from
         :py:class:`smithy_http.aio.HTTPRequest`"""
         headers_list = []
@@ -343,13 +353,11 @@ async def _marshal_request(
             crt_body = BytesIO(body)
         else:
             # If the body is async, or potentially very large, start up a task to read
-            # it into the BytesIO object that CRT needs. By using asyncio.create_task
-            # we'll start the coroutine without having to explicitly await it.
+            # it into the intermediate object that CRT needs. By using
+            # asyncio.create_task we'll start the coroutine without having to
+            # explicitly await it.
             crt_body = BufferableByteStream()
-            if not isinstance(body, AsyncIterable):
-                # If the body isn't already an async iterable, wrap it in one. Objects
-                # with read methods will be read in chunks so as not to exhaust memory.
-                body = AsyncBytesReader(body)
+            body = AsyncBytesReader(body)
 
             # Start the read task in the background.
             read_task = asyncio.create_task(self._consume_body_async(body, crt_body))
@@ -365,13 +373,19 @@ async def _marshal_request(
             headers=headers,
             body_stream=crt_body,
         )
-        return crt_request
+        return crt_request, crt_body
 
     async def _consume_body_async(
-        self, source: AsyncIterable[bytes], dest: "BufferableByteStream"
+        self, source: AsyncBytesReader, dest: "BufferableByteStream"
     ) -> None:
-        async for chunk in source:
-            dest.write(chunk)
+        try:
+            async for chunk in source:
+                dest.write(chunk)
+        except Exception:
+            dest.close()
+            raise
+        finally:
+            await source.close()
         dest.end_stream()
 
     def __deepcopy__(self, memo: Any) -> "AWSCRTHTTPClient":

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

@@ -26,7 +26,7 @@ async def test_client_marshal_request() -> None:
         body=BytesIO(),
         fields=Fields(),
     )
-    crt_request = await client._marshal_request(request)  # type: ignore
+    crt_request, _ = await client._marshal_request(request)  # type: ignore
     assert crt_request.headers.get("host") == "example.com"  # type: ignore
     assert crt_request.headers.get("accept") == "*/*"  # type: ignore
     assert crt_request.method == "GET"  # type: ignore