Add an asynchronous bytes provider

This adds a class that allows bytes to be asynchronously exchanged.
This primarily serves the purpose of enabling event streaming. Event
streams will create a provider under the hood that will be written
to. That provider will then be passed as the trasnport request body,
which will read it as an async bytes iterable.

Author: JordonPhillips

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

@@ -1,10 +1,12 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
+import asyncio
 from asyncio import iscoroutinefunction
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
 from io import BytesIO
-from typing import Self, cast
+from typing import Any, Self, cast
 
+from ..exceptions import SmithyException
 from ..interfaces import BytesReader
 from .interfaces import AsyncByteStream, StreamingBlob
 
@@ -271,3 +273,153 @@ async def __anext__(self) -> bytes:
         if data:
             return data
         raise StopAsyncIteration
+
+
+class AsyncBytesProvider:
+    """A buffer that allows chunks of bytes to be exchanged asynchronously.
+
+    Bytes are written in chunks to an internal buffer, that is then drained via an async
+    iterator.
+    """
+
+    def __init__(
+        self, intial_data: bytes | None = None, max_buffered_chunks: int = 16
+    ) -> None:
+        """Initialize the AsyncBytesProvider.
+
+        :param initial_data: An initial chunk of bytes to make available.
+        :param max_buffered_chunks: The maximum number of chunks of data to buffer.
+            Calls to ``write`` will block until the number of chunks is less than this
+            number. Default is 16.
+        """
+        if intial_data is not None:
+            self._data = [intial_data]
+        else:
+            self._data = []
+
+        if max_buffered_chunks < 1:
+            raise ValueError(
+                "The maximum number of buffered chunks must be greater than 0."
+            )
+
+        self._closed = False
+        self._closing = False
+        self._flushing = False
+        self._max_buffered_chunks = max_buffered_chunks
+
+        # Create a Condition to synchronize access to the data chunk pool.
+        self._data_condition = asyncio.Condition()
+
+    async def write(self, data: bytes) -> None:
+        if self._closed:
+            raise SmithyException("Attempted to write to a closed provider.")
+
+        # Acquire a lock on the data buffer, releasing it automatically when the
+        # block exits.
+        async with self._data_condition:
+
+            # Wait for the number of chunks in the buffer to be less than the
+            # specified maximum. This also releases the lock until the condition
+            # is met.
+            await self._data_condition.wait_for(self._can_write)
+
+            # The provider could have been closed while waiting to write, so an
+            # additional check is done here for safety.
+            if self._closed or self._closing:
+                # Notify to allow other coroutines to check their conditions.
+                self._data_condition.notify()
+                raise SmithyException(
+                    "Attempted to write to a closed or closing provider."
+                )
+
+            # Add a new chunk of data to the buffer and notify the next waiting
+            # coroutine.
+            self._data.append(data)
+            self._data_condition.notify()
+
+    def _can_write(self) -> bool:
+        return (
+            self._closed
+            or self._closing
+            or (len(self._data) < self._max_buffered_chunks and not self._flushing)
+        )
+
+    @property
+    def closed(self) -> bool:
+        """Returns whether the provider is closed."""
+        return self._closed
+
+    async def flush(self) -> None:
+        """Waits for all buffered data to be consumed."""
+        if self._closed:
+            return
+
+        # Acquire a lock on the data buffer, releasing it automatically when the
+        # block exits.
+        async with self._data_condition:
+            # Block writes
+            self._flushing = True
+
+            # Wait for the stream to be closed or for the data buffer to be empty.
+            await self._data_condition.wait_for(lambda: len(self._data) == 0)
+
+            # Unblock writes
+            self._flushing = False
+
+    async def close(self, flush: bool = False) -> None:
+        """Closes the provider.
+
+        Pending writing tasks queued after this will fail, so such tasks should be
+        awaited before this. Write tasks queued before this may succeed, however.
+
+        :param flush: Whether to flush buffered data before closing. If false, all
+            buffered data will be lost. Default is False.
+        """
+        if self._closed:
+            return
+
+        # Acquire a lock on the data buffer, releasing it automatically when the
+        # block exits. Notably this will not wait on a condition to move forward.
+        async with self._data_condition:
+            self._closing = True
+            if flush:
+                await self._data_condition.wait_for(lambda: len(self._data) == 0)
+            else:
+                # Clear out any pending data, freeing up memory.
+                self._data.clear()
+
+            self._closed = True
+            self._closing = False
+
+            # Notify all waiting coroutines that the provider has closed.
+            self._data_condition.notify_all()
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> bytes:
+        # Acquire a lock on the data buffer, releasing it automatically when the
+        # block exits.
+        async with self._data_condition:
+
+            # Wait for the stream to be closed or for the data buffer to be non-empty.
+            # This also releases the lock until the condition is met.
+            await self._data_condition.wait_for(
+                lambda: self._closed or len(self._data) > 0
+            )
+
+            # If the provider is closed, end the iteration.
+            if self._closed:
+                raise StopAsyncIteration
+
+            # Pop the next chunk of data from the buffer, then notify any waiting
+            # coroutines, returning immediately after.
+            result = self._data.pop()
+            self._data_condition.notify()
+            return result
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        await self.close()

python-packages/smithy-core/tests/unit/aio/test_types.py

@@ -1,11 +1,17 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
+import asyncio
 from io import BytesIO
 from typing import Self
 
 import pytest
 
-from smithy_core.aio.types import AsyncBytesReader, SeekableAsyncBytesReader
+from smithy_core.aio.types import (
+    AsyncBytesProvider,
+    AsyncBytesReader,
+    SeekableAsyncBytesReader,
+)
+from smithy_core.exceptions import SmithyException
 
 
 class _AsyncIteratorWrapper:
@@ -270,3 +276,166 @@ async def test_seekable_iter_chunks() -> None:
 
     assert reader.tell() == 3
     assert result == b"foo"
+
+
+async def test_provider_requires_positive_max_chunks() -> None:
+    with pytest.raises(ValueError):
+        AsyncBytesProvider(max_buffered_chunks=-1)
+
+
+async def drain_provider(provider: AsyncBytesProvider, dest: list[bytes]) -> None:
+    async for chunk in provider:
+        dest.append(chunk)
+
+
+async def test_provider_reads_initial_data() -> None:
+    provider = AsyncBytesProvider(intial_data=b"foo")
+    result: list[bytes] = []
+
+    # Start the read task in the background.
+    read_task = asyncio.create_task(drain_provider(provider, result))
+
+    # Wait for the buffer to drain. At that point all the data should
+    # be read, but the read dask won't actually be complete yet
+    # because it's still waiting on future data.
+    await provider.flush()
+    assert result == [b"foo"]
+    assert not read_task.done()
+
+    # Now actually close the provider, which will let the read task
+    # complete.
+    await provider.close()
+    await read_task
+
+    # The result should not have changed
+    assert result == [b"foo"]
+
+
+async def test_provider_reads_written_data() -> None:
+    provider = AsyncBytesProvider()
+    result: list[bytes] = []
+
+    # Start the read task in the background.
+    read_task = asyncio.create_task(drain_provider(provider, result))
+    await provider.write(b"foo")
+
+    # Wait for the buffer to drain. At that point all the data should
+    # be read, but the read dask won't actually be complete yet
+    # because it's still waiting on future data.
+    await provider.flush()
+    assert result == [b"foo"]
+    assert not read_task.done()
+
+    # Now actually close the provider, which will let the read task
+    # complete.
+    await provider.close()
+    await read_task
+
+    # The result should not have changed
+    assert result == [b"foo"]
+
+
+async def test_close_stops_writes() -> None:
+    provider = AsyncBytesProvider()
+    await provider.close()
+    with pytest.raises(SmithyException):
+        await provider.write(b"foo")
+
+
+async def test_close_deletes_buffered_data() -> None:
+    provider = AsyncBytesProvider(b"foo")
+    await provider.close()
+    result: list[bytes] = []
+    await drain_provider(provider, result)
+    assert result == []
+
+    # We weren't able to read data, which is what we want. But here we dig into
+    # the internals to be sure that the buffer is clear and no data is haning
+    # around.
+    assert provider._data == []  # type: ignore
+
+
+async def test_only_max_chunks_buffered() -> None:
+    # Initialize the provider with a max buffer of one and immediately have it
+    # filled with an initial chunk.
+    provider = AsyncBytesProvider(b"foo", max_buffered_chunks=1)
+
+    # Schedule a write task. Using create_task immediately enqueues it, though it
+    # won't start executing until its turn in the loop.
+    write_task = asyncio.create_task(provider.write(b"bar"))
+
+    # Suspend the current coroutine so the write task can take over. It shouldn't
+    # complete because it should be waiting on the buffer to drain. One tenth of
+    # a second is way more than enough time for it to complete under normal
+    # circumstances.
+    await asyncio.sleep(0.1)
+    assert not write_task.done()
+
+    # Now begin the read task in the background. Since it's draining the buffer, the
+    # write task will become unblocked.
+    result: list[bytes] = []
+    read_task = asyncio.create_task(drain_provider(provider, result))
+
+    # The read task won't be done until we close the provider, but the write task
+    # should be able to complete now.
+    await write_task
+
+    # The write task and read task don't necessarily complete at the same time,
+    # so we wait until the buffer is empty here.
+    await provider.flush()
+    assert result == [b"foo", b"bar"]
+
+    # Now we can close the provider and wait for the read task to end.
+    await provider.close()
+    await read_task
+
+
+async def test_close_stops_queued_writes() -> None:
+    # Initialize the provider with a max buffer of one and immediately have it
+    # filled with an initial chunk.
+    provider = AsyncBytesProvider(b"foo", max_buffered_chunks=1)
+
+    # Schedule a write task. Using create_task immediately enqueues it, though it
+    # can't complete until the buffer is free.
+    write_task = asyncio.create_task(provider.write(b"bar"))
+
+    # Now close the provider. The write task will raise an error.
+    await provider.close()
+
+    with pytest.raises(SmithyException):
+        await write_task
+
+
+async def test_close_with_flush() -> None:
+    # Initialize the provider with a max buffer of one and immediately have it
+    # filled with an initial chunk.
+    provider = AsyncBytesProvider(b"foo", max_buffered_chunks=1)
+
+    # Schedule a write task. Using create_task immediately enqueues it, though it
+    # can't complete until the buffer is free.
+    write_task = asyncio.create_task(provider.write(b"bar"))
+
+    # Now flush the provider and close it. The read task will be able to read the
+    # alredy buffered data, but the write task will fail.
+    close_task = asyncio.create_task(provider.close(flush=True))
+
+    # There is a timing issue to when a write will fail. If they're in the queue
+    # before the close task, they may still make it through. Here the current
+    # coroutine is suspended so that both the write task and close task have a
+    # chance to check their conditions and set necessary state.
+    await asyncio.sleep(0.1)
+
+    # Now we can start the read task. We can immediately await it because the close
+    # task will complete in the background, which will then stop the iteration.
+    result: list[bytes] = []
+    await drain_provider(provider, result)
+
+    # Ensure that the close task is complete.
+    await close_task
+
+    # The write will have been blocked by the close task, so the read task will
+    # only see the initial data. The write task will raise an exception as the
+    # provider closed before it could write its data.
+    assert result == [b"foo"]
+    with pytest.raises(SmithyException):
+        await write_task