diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index f171bfa5c..bdd12a239 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -13,13 +13,14 @@ TYPE_CHECKING, Any, Concatenate, + TypeVar, cast, overload, ) import nexusrpc from nexusrpc import OutputT -from typing_extensions import Required, Self, TypedDict +from typing_extensions import Required, Self, TypedDict, Unpack import temporalio.activity import temporalio.api.workflowservice.v1 @@ -110,11 +111,19 @@ ) if TYPE_CHECKING: - from temporalio.common import StreamBackend + from temporalio.common import ( + Creating, + StreamBackend, + StreamDefinition, + StreamHandle, + StreamOptions, + ) from ._interceptor import Interceptor from ._plugin import Plugin +_StreamPayloadT = TypeVar("_StreamPayloadT") + class Client: """Client for accessing Temporal. @@ -310,7 +319,7 @@ def __init__( # Only a client with Stream backends loads the Streams package. from temporalio.common._streams import _plugin as _stream_plugin - config = _stream_plugin.configure_client(config) + config = _stream_plugin.configure_client(config, self) self._init_from_config(config) @@ -911,6 +920,79 @@ def get_workflow_handle( result_type=result_type, ) + def create_stream( + self, + definition: type[StreamDefinition[_StreamPayloadT]], + *, + id: str, + backend: str | None = None, + **options: Unpack[StreamOptions], + ) -> Creating[StreamHandle[_StreamPayloadT]]: + """Create a Stream that the client owns, or get the one that exists. + + A create is create-or-get by Stream id. On a Stream that exists, it ignores ``options`` + and returns that Stream. A Stream must exist before its first publish. The client closes + a Stream that it owns, and a Workflow only writes to it. Await the call to get the + handle, or use it in ``async with``, which closes the Stream at the end of the block:: + + async with client.create_stream(Turns, id="chat-42") as stream: + await stream.publish(Turn("user", "hi")) + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class. + id: The Stream id, which the caller picks, as for a Workflow id. + backend: The name of the backend. Defaults to the registered backend with + ``default=True``. + options: The options of a new Stream, such as ``retention``. See + :py:class:`temporalio.common.StreamOptions`. + + Raises: + StreamConfigError: The client has no backend with that name, or no backend has + ``default=True`` and the call gives no backend. No call is made. + StreamClosedError: A Stream of that id was created and closed before. + StreamExpiredError: A Stream of that id expired. + """ + from temporalio.common._streams import _stream_client + + return _stream_client.create( + self, definition, id=id, backend=backend, options=options + ) + + def get_stream_handle( + self, + definition: type[StreamDefinition[_StreamPayloadT]], + *, + id: str, + backend: str | None = None, + ) -> StreamHandle[_StreamPayloadT]: + """Get a Stream by its id. It makes no call. + + To create a Stream that the client owns, use :py:meth:`create_stream`. A handle that + came as data, such as from a Workflow's query, needs no call of this method. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class. + id: The Stream id. + backend: The name of the backend. Defaults to the registered backend with + ``default=True``. + + Raises: + StreamConfigError: The client has no backend with that name, or no backend has + ``default=True`` and the call gives no backend. + ValueError: The id is empty. + """ + from temporalio.common._streams import _stream_client + + return _stream_client.get_stream_handle( + self, definition, id=id, backend=backend + ) + def get_workflow_handle_for( self, workflow: ( diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index c1104fce3..2d77f6c63 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -63,6 +63,7 @@ Superseded, ) from temporalio.common._streams._plugin import StreamBackend + from temporalio.common._streams._reader import Subscription __all__ = [ "ActivityIDConflictPolicy", @@ -119,6 +120,7 @@ "StreamStorageError", "StreamTooLargeError", "StreamUnsupportedError", + "Subscription", "Superseded", "TypedSearchAttributes", "VersioningBehavior", @@ -1607,6 +1609,7 @@ def wait_sync(self, timeout: float | None = None) -> None: "StreamStorageError", "StreamTooLargeError", "StreamUnsupportedError", + "Subscription", "Superseded", "defn", ] diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 299ee3ced..c803acdb7 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -25,6 +25,7 @@ Superseded, ) from temporalio.common._streams._plugin import StreamBackend +from temporalio.common._streams._reader import Subscription __all__ = [ "END", @@ -49,6 +50,7 @@ "StreamStorageError", "StreamTooLargeError", "StreamUnsupportedError", + "Subscription", "Superseded", "defn", ] diff --git a/temporalio/common/_streams/_access.py b/temporalio/common/_streams/_access.py new file mode 100644 index 000000000..a68ed45c9 --- /dev/null +++ b/temporalio/common/_streams/_access.py @@ -0,0 +1,137 @@ +"""Stream calls outside Workflow code, which Activity and client handles share.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator, Callable, Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +import temporalio.converter +from temporalio.common._streams import _convert, _reader +from temporalio.common._streams._backend import BackendCalls, CallOptions +from temporalio.common._streams._items import End, close_failure +from temporalio.common._streams._reader import Item, Position, Subscription +from temporalio.common._streams._writer import StreamWriter +from temporalio.common._streams.system import models + + +@dataclass +class Access: + """The calls of one Stream handle, bound to a backend.""" + + backend: BackendCalls + converter: temporalio.converter.DataConverter + namespace: str + stream_id: str + writer: Callable[[], StreamWriter] + + def publish(self, values: Sequence[Any]) -> asyncio.Future[int]: + """Start the write of the values as Records, in order, and return its future. + + The future gives the offset of the last Record. + """ + if not values: + raise TypeError("publish takes one or more values") + return self.writer().write(values) + + async def items( + self, read: ReadSpec, position: Position + ) -> AsyncGenerator[Item, None]: + """Read the Stream's items in order, and move ``position`` past each page.""" + items = _reader.read( + self.backend.service(self.namespace, _convert.raw(self.converter)), + self.converter, + self.stream_id, + position=position, + payload_type=read.payload_type, + from_offset=read.offset, + last=read.last, + wait_new_records=read.wait_new_records, + create_wait=read.create_wait, + ) + try: + async for item in items: + yield item + finally: + await items.aclose() + + async def get_offset(self) -> int: + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + + Offsets are sparse and above zero, so ``offset=0`` reads from the start. + """ + service = self.backend.service(self.namespace) + response = await service.describe_stream( + handle=models.StreamHandle(id=self.stream_id) + ) + return response.last_offset + 1 if response.last_offset else 0 + + async def close(self, result: Any | None, failure: BaseException | None) -> None: + """Write this Producer's Records, then close the Stream.""" + writer = self.writer() + await writer.close() + producer = models._ProducerTransferTypeConverter().from_transfer_type( + writer.producer, models.Producer + ) + service = self.backend.service(self.namespace, self.converter) + # The generated close takes its Producer from a source, so the handle sends the request + # with its own. + await service._stream_service_close_stream( + models.CloseStreamRequest( + handle=models.StreamHandle(id=self.stream_id), + producer=producer, + result=result, + failure=close_failure(failure), + ) + ) + + +@dataclass +class ReadSpec: + """What one read asks for.""" + + payload_type: Any + offset: int | End | None + last: int | None + wait_new_records: bool + create_wait: timedelta | None + + +def read_spec( + payload_type: Any, + *, + offset: int | End | None, + last: int | None, + wait_new_records: bool, + create_wait: timedelta | None, +) -> ReadSpec: + """The options of a read, and the type that it decodes each value with. + + Raises: + ValueError: ``last`` is not above zero, or the call also gives ``offset``. + """ + if last is not None and (last < 1 or offset is not None): + raise ValueError("last takes a count above zero and no offset") + return ReadSpec( + payload_type=payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + + +def subscription( + items: Callable[[Position], AsyncGenerator[Item, None]], + read: ReadSpec, + options: CallOptions | None = None, +) -> Subscription[Any]: + """A read over the items that ``items`` gives.""" + # A read from END or of the last Records has no resume offset until its first page. + if read.last is not None or read.offset is End.END: + position = Position(None) + else: + position = Position(read.offset or 0) + return Subscription(items(position), position, options) diff --git a/temporalio/common/_streams/_convert.py b/temporalio/common/_streams/_convert.py new file mode 100644 index 000000000..5d538bcbd --- /dev/null +++ b/temporalio/common/_streams/_convert.py @@ -0,0 +1,66 @@ +"""Value conversion for Stream Records outside Workflow code.""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Sequence +from typing import Any + +import temporalio.api.common.v1 +import temporalio.common +import temporalio.converter +import temporalio.exceptions + + +class _RawPayloadConverter(temporalio.converter.PayloadConverter): + # The generated client runs the codec on Record bodies and keeps them as RawValue, so the + # handle converts each body with the Stream's payload type. + def to_payloads( + self, values: Sequence[Any] + ) -> list[temporalio.api.common.v1.Payload]: + payloads = [] + for value in values: + if not isinstance(value, temporalio.common.RawValue): + raise TypeError(f"expected a RawValue Record body, got {type(value)}") + payloads.append(value.payload) + return payloads + + def from_payloads( + self, + payloads: Sequence[temporalio.api.common.v1.Payload], + type_hints: list[type] | None = None, + ) -> list[Any]: + return [temporalio.common.RawValue(payload) for payload in payloads] + + +def raw( + converter: temporalio.converter.DataConverter, +) -> temporalio.converter.DataConverter: + """The converter with Record bodies kept as RawValue. Its codec still runs on them.""" + return dataclasses.replace(converter, payload_converter_class=_RawPayloadConverter) + + +def failure( + converter: temporalio.converter.DataConverter, failure: BaseException +) -> BaseException: + """Convert the RawValue payloads of a read's close failure with the user's converter.""" + if ( + not isinstance(failure, temporalio.exceptions.FailureError) + or not failure.failure + ): + return failure + return converter.failure_converter.from_failure( + failure.failure, converter.payload_converter + ) + + +def decode( + converter: temporalio.converter.DataConverter, body: Any, type_hint: Any +) -> Any: + """Convert a Record body that a read gave as RawValue, after the codec.""" + if not isinstance(body, temporalio.common.RawValue): + return body + [value] = converter.payload_converter.from_payloads( + [body.payload], [type_hint] if type_hint else None + ) + return value diff --git a/temporalio/common/_streams/_handle.py b/temporalio/common/_streams/_handle.py index 1a12d1ec3..deb0f39a2 100644 --- a/temporalio/common/_streams/_handle.py +++ b/temporalio/common/_streams/_handle.py @@ -2,20 +2,27 @@ from __future__ import annotations +import abc import asyncio +import contextvars import typing -from collections.abc import Callable, Coroutine, Generator +from collections.abc import Callable, Coroutine, Generator, Mapping from datetime import timedelta from types import TracebackType -from typing import Any, Generic, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar import google.protobuf.json_format from typing_extensions import Self, TypedDict import temporalio.converter import temporalio.exceptions +from temporalio.common._streams import _errors from temporalio.common._streams.system import models +if TYPE_CHECKING: + from temporalio.common._streams._items import End + from temporalio.common._streams._reader import Subscription + T = TypeVar("T") @@ -171,12 +178,118 @@ def from_transfer_type( return _from_json(value, type_hint) +# The client whose data the converter decodes. A handle in the data takes it, so client code uses +# a Stream that a Workflow hands out with no other call. +decoding_client: contextvars.ContextVar[Any] = contextvars.ContextVar( + "__temporal_stream_decoding_client", default=None +) + + +class ContextCalls(abc.ABC): + """The Stream calls of a handle in one context, such as client code. + + The handle passes each call to the calls of the context where it runs. A call that the + context can't make raises :py:class:`StreamUnsupportedError`. + """ + + @abc.abstractmethod + def publish( + self, + handle: StreamHandle[Any], + values: tuple[Any, ...], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> asyncio.Future[int]: + """Start a write. See :py:meth:`StreamHandle.publish`.""" + + @abc.abstractmethod + def subscribe( + self, + handle: StreamHandle[Any], + *, + offset: int | End | None, + last: int | None, + wait_new_records: bool, + create_wait: timedelta | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read. See :py:meth:`StreamHandle.subscribe`.""" + + @abc.abstractmethod + async def get_offset( + self, + handle: StreamHandle[Any], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> int: + """Get the next offset. See :py:meth:`StreamHandle.get_offset`.""" + + @abc.abstractmethod + async def close( + self, + handle: StreamHandle[Any], + result: Any, + failure: BaseException | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> None: + """Close the Stream. See :py:meth:`StreamHandle.close`.""" + + @abc.abstractmethod + async def enter(self, handle: StreamHandle[Any]) -> None: + """Start an ``async with`` block of the handle.""" + + @abc.abstractmethod + async def exit(self, handle: StreamHandle[Any], exc: BaseException | None) -> None: + """End an ``async with`` block of the handle. See :py:meth:`StreamHandle.__aexit__`.""" + + def publish_nowait( + self, handle: StreamHandle[Any], values: tuple[Any, ...] + ) -> None: + """Write without an await. See :py:meth:`StreamHandle.publish_nowait`.""" + del handle, values + raise _errors.StreamUnsupportedError( + "publish_nowait is for sync Workflow code. Elsewhere, use publish." + ) + + def close_nowait( + self, handle: StreamHandle[Any], result: Any, failure: BaseException | None + ) -> None: + """Close without an await. See :py:meth:`StreamHandle.close_nowait`.""" + del handle, result, failure + raise _errors.StreamUnsupportedError( + "close_nowait is for sync Workflow code. Elsewhere, use close." + ) + + def producer( + self, handle: StreamHandle[Any], name: str, epoch: int + ) -> StreamHandle[Any]: + """Return a handle of a named Producer. See :py:meth:`StreamHandle.producer`.""" + del handle, name, epoch + raise _errors.StreamUnsupportedError( + "Only client code writes as a named Producer" + ) + + class StreamHandle(Generic[T]): - """A handle on one Stream in Workflow, Activity and client code. ``T`` is the payload type. + """A handle on one Stream. ``T`` is the payload type. + + One handle type works in Workflow, Activity and client code. The SDK picks the write path + and the Producer from where the call runs. A call that the context can't make raises + :py:class:`StreamUnsupportedError`. - A handle is data. It's passed as a Workflow or Activity argument in the proto JSON form of - the ``temporal.api.stream.v1.StreamHandle`` message. The parameter's type, such as - ``StreamHandle[Turn]``, gives the payload type. + A handle is data. It's passed as a Workflow or Activity argument, or returned from a query, + in the proto JSON form of the ``temporal.api.stream.v1.StreamHandle`` message. The type + hint, such as ``StreamHandle[Turn]``, gives the payload type. + + In client code, a handle uses the client that made it, such as with + :py:meth:`temporalio.client.Client.get_stream_handle`, or the client that decoded it, such + as from a query, an Update or a Workflow result. A handle that no client decoded, such as + one that user code builds from JSON, raises :py:class:`StreamConfigError` at its first call. + All handles of one client and Stream write as one Producer, and :py:meth:`producer` returns + a handle that writes as a named Producer. A write error that no caller got raises at the + next call of the handle. .. warning:: This API is experimental. @@ -198,6 +311,10 @@ def __init__( self._backend = backend # The type that a read decodes each value with. None decodes with no type hint. self._payload_type = payload_type + # Client code binds a handle to a client and keeps its writer here. Workflow and Activity + # calls keep no state in the handle. + self._client: Any = decoding_client.get() + self._client_calls: ContextCalls | None = None @property def id(self) -> str: @@ -219,6 +336,177 @@ def backend(self) -> str | None: """Name of the backend that stores the Stream. None means the default one.""" return self._backend + def _calls(self) -> ContextCalls: + # The context modules import on first use, so Workflow code never imports the client's. + from temporalio.common._streams import _stream_client + + return _stream_client.calls(self) + + def publish( + self, + *values: T, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> asyncio.Future[int]: + """Start the write of one Record for each value, in order, and return its future. + + The write has its place in the Stream's order when the call returns, also when no + caller awaits the future. The writes of one handle keep the order of the calls. A cancel + of the future doesn't stop the write. A list is one value. In client code, the next call + of the handle raises a write error that no caller got. + + Args: + values: One value for each Record. + rpc_metadata: Headers of each call. The Redis and memory backends make no RPC, so + they send none. + rpc_timeout: The deadline of each call. A call past it raises ``RPCError`` with + ``DEADLINE_EXCEEDED``. + + Returns: + The future of the write. It gives the offset of the last Record, or raises + ``StreamError`` when the backend refused the write, such as + ``StreamNotFoundError`` for a Stream that nobody created. + + Raises: + StreamClosedError: This handle closed the Stream. + StreamConfigError: No client decoded the handle, or the client has no backend of + that name. + """ + return self._calls().publish(self, values, rpc_metadata, rpc_timeout) + + def subscribe( + self, + *, + offset: int | End | None = None, + last: int | None = None, + wait_new_records: bool = True, + create_wait: timedelta | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> Subscription[T]: + """Read the Stream's items in order. + + The read ends after a ``Closed`` item, or at the newest Record when + ``wait_new_records`` is False. To stop after a quiet period, use your own timer, and + resume later with ``offset=`` set to the read's ``next_offset``. A ``Superseded`` item + withdraws earlier Records of a failed attempt. + + Args: + offset: Start at this offset, inclusive, or at ``common.END`` for new + Records only. Defaults to the first Record. + last: Start with the newest ``last`` Records, oldest first, such as for the first + view of a chat. Can't be used with ``offset``. A Record that a later + ``Superseded`` item withdraws doesn't count. At most 1000. + wait_new_records: Wait for new Records at the end of the Stream. + create_wait: How long to wait for a Stream that doesn't exist yet, such as the + Stream of a Workflow that just started. Defaults to no wait. + rpc_metadata: Headers of each call of the read. See :py:meth:`publish`. + rpc_timeout: The deadline of each call of the read, not of the whole read. A call + that waits for new Records returns before the deadline. + """ + return self._calls().subscribe( + self, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + rpc_metadata=rpc_metadata, + rpc_timeout=rpc_timeout, + ) + + async def get_offset( + self, + *, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> int: + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. + + Offsets are sparse and above zero, so ``offset=0`` reads from the start. + ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`publish`. + """ + return await self._calls().get_offset(self, rpc_metadata, rpc_timeout) + + async def close( + self, + *, + result: Any = None, + failure: BaseException | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> None: + """Write this handle's Records, then close the Stream. + + Readers get the Records before the close, then a ``Closed`` item with ``result`` or + ``failure``. A cancel failure, such as a ``CancelledError``, closes the Stream as + canceled. ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`publish`. + + Raises: + StreamNotOwnerError: The code doesn't own the Stream, such as a client with a + Workflow's Stream. Only the run chain closes a Workflow's Stream, and only a + client closes a client's Stream. + """ + await self._calls().close(self, result, failure, rpc_metadata, rpc_timeout) + + async def __aenter__(self) -> Self: + """Use the Stream in ``async with``, which closes it at the end. + + Raises: + StreamNotOwnerError: The code can't close the Stream. See :py:meth:`close`. + """ + await self._calls().enter(self) + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream, with the exception as its failure if there is one. + + A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as + ``KeyboardInterrupt``, leaves it open. If this handle already closed the Stream, this + does nothing. + """ + await self._calls().exit(self, exc) + + def publish_nowait(self, *values: T) -> None: + """Write one Record for each value, in order, without an await, in sync Workflow code. + + Raises: + StreamUnsupportedError: The code isn't Workflow code. + """ + self._calls().publish_nowait(self, values) + + def close_nowait( + self, *, result: Any = None, failure: BaseException | None = None + ) -> None: + """Close the Stream without an await, in sync Workflow code. + + Raises: + StreamUnsupportedError: The code isn't Workflow code. + """ + self._calls().close_nowait(self, result, failure) + + def producer(self, name: str, epoch: int) -> StreamHandle[T]: + """Return this Stream as a handle that writes as the Producer ``name`` at ``epoch``. + + Use it for a writer that must fence its earlier instances, such as a server that + restarts. Keep ``epoch`` in durable storage, and raise it at each start. After a write + at a higher epoch, a write at a lower one raises :py:class:`StreamFencedError`. The + Records of the earlier epochs aren't withdrawn, since they're real messages. Run only + one live writer for each name. + + Args: + name: The Producer name, the same across restarts. + epoch: The writer's epoch, 1 or more. + + Raises: + StreamUnsupportedError: The code isn't client code. + """ + return self._calls().producer(self, name, epoch) + def _fields(self) -> dict[str, str]: # The fields that name the Stream. They need no call, so a handle whose id is not known # yet still compares, hashes and prints. diff --git a/temporalio/common/_streams/_plugin.py b/temporalio/common/_streams/_plugin.py index 1763ea9b8..51b2f85ce 100644 --- a/temporalio/common/_streams/_plugin.py +++ b/temporalio/common/_streams/_plugin.py @@ -157,11 +157,18 @@ def default_name(backends: Sequence[StreamBackend]) -> str: def configure_client( - config: temporalio.client.ClientConfig, + config: temporalio.client.ClientConfig, client: temporalio.client.Client ) -> temporalio.client.ClientConfig: - """Check the client's backends, and apply the plugin configuration of each one.""" + """Check the client's backends, and apply the plugin configuration of each one. + + The decodes of ``client`` bind each Stream handle in the data to it. + """ + # The client Stream module imports this one. + from temporalio.common._streams._stream_client import bind_decodes + backends = check_backends(config.get("stream_backends", ()), "client") config["stream_backends"] = backends + config["data_converter"] = bind_decodes(config["data_converter"], client) for backend in backends: config = backend._configure_client(config) return config diff --git a/temporalio/common/_streams/_reader.py b/temporalio/common/_streams/_reader.py new file mode 100644 index 000000000..964bade2b --- /dev/null +++ b/temporalio/common/_streams/_reader.py @@ -0,0 +1,326 @@ +"""Reads of a Stream, for Activities and clients.""" + +from __future__ import annotations + +import array +import asyncio +import dataclasses +from collections.abc import AsyncGenerator +from datetime import timedelta +from typing import Any, Generic, TypeVar + +import temporalio.common +import temporalio.converter +from temporalio.api.enums.v1 import StreamProducerKind +from temporalio.common._streams import _backend, _convert, _errors +from temporalio.common._streams._items import ( + Closed, + CloseStatus, + End, + Record, + Superseded, + producer_key, + withdraws, +) +from temporalio.common._streams.system import client as service_client +from temporalio.common._streams.system import models + +T = TypeVar("T") + +Item = Record[Any] | Superseded | Closed + +# The service waits at most 60 s per read, so a longer wait is a loop of reads. +_MAX_WAIT = timedelta(seconds=60) + +# An empty page that comes back in less than half of its wait is followed by a pause, so a +# Store that ignores the wait isn't asked again at once. The pause doubles up to the last one. +_FIRST_PAUSE = 0.05 +_LAST_PAUSE = 2.0 + + +@dataclasses.dataclass +class Position: + """Where a read resumes from. The read moves it past each page that it finished.""" + + next_offset: int | None + + +class Subscription(Generic[T]): + """The items of one read of a Stream, in order. Iterate it with ``async for``. + + It yields :py:class:`Record` items, :py:class:`Superseded` items that withdraw earlier + Records, and a :py:class:`Closed` item last when the Stream is closed. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + items: AsyncGenerator[Item, None], + position: Position, + options: _backend.CallOptions | None = None, + ) -> None: + """Create a read. Use a handle's ``subscribe`` instead.""" + self._items = items + self._position = position + # A client read's rpc_metadata and rpc_timeout, for each call of each step. + self._options = options + # The Producer and the offsets of each Record that this read gave, without values, so a + # Superseded item can list the offsets that it withdraws. + self._shown: dict[ + tuple[Any, ...], tuple[models.Producer, array.array[int]] + ] = {} + + @property + def next_offset(self) -> int | None: + """The offset to resume from, after the last item or the last page that the read got. + + Pass it as ``offset`` to a later read. After a ``Superseded`` item, it's that item's own + offset, since the first Record of its batch has the same offset and comes next. A read + from there gives the ``Superseded`` item again, which withdraws nothing new. After a + Record that failed to decode, it's past that Record. It's None before the first page of + a read from ``END`` or of a read with ``last``. + """ + return self._position.next_offset + + def __aiter__(self) -> Subscription[T]: + """The read is its own iterator.""" + return self + + async def __anext__(self) -> Record[T] | Superseded | Closed: + """The next item. + + Raises: + StreamError: A Record's value failed to decode with the Stream's payload type. The + read ends, and :py:attr:`next_offset` is past that Record. + """ + with _backend.using(self._options): + item = await self._items.__anext__() + if isinstance(item, Superseded): + # The Record of the same entry has the same offset and comes next. + self._position.next_offset = item.offset + else: + self._position.next_offset = max( + self._position.next_offset or 0, item.offset + 1 + ) + if ( + isinstance(item, Record) + and item.producer is not None + # Nothing withdraws a client's Records, so a long read of them keeps no offsets. + and item.producer.kind != StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT + ): + key = producer_key(item.producer) + _, runs = self._shown.setdefault(key, (item.producer, array.array("Q"))) + # A batch has consecutive offsets, so a run of them costs one first and last pair. + if runs and runs[-1] + 1 == item.offset: + runs[-1] = item.offset + else: + runs.extend((item.offset, item.offset)) + elif isinstance(item, Superseded): + item = dataclasses.replace( + item, withdrawn_offsets=self._withdraw(item.facts) + ) + return item + + def _withdraw(self, facts: models.StreamSuperseded) -> tuple[int, ...]: + offsets: list[int] = [] + for key, (producer, runs) in list(self._shown.items()): + if withdraws(facts, producer): + for first, last in zip(runs[::2], runs[1::2]): + offsets.extend(range(first, last + 1)) + # A withdrawn Record stays withdrawn, so the read forgets it. + del self._shown[key] + return tuple(sorted(offsets)) + + async def aclose(self) -> None: + """End the read early, from the task that reads it. It makes no more calls. + + If another task is waiting in the read for new Records, this raises ``RuntimeError``. + Cancel that task instead. + """ + with _backend.using(self._options): + await self._items.aclose() + + +async def read( + service: service_client.StreamServiceClientExtension, + converter: temporalio.converter.DataConverter, + stream_id: str, + *, + position: Position, + payload_type: Any, + from_offset: int | End | None, + last: int | None = None, + wait_new_records: bool, + create_wait: timedelta | None, +) -> AsyncGenerator[Item, None]: + """Yield the items of a Stream in order, until a close. + + With ``wait_new_records`` False, stop at the newest Record. With ``last``, start with the + newest ``last`` Records. + """ + loop = asyncio.get_running_loop() + create_until = loop.time() + create_wait.total_seconds() if create_wait else None + + def request() -> models.ReadRecordsRequest: + create_left = ( + timedelta(seconds=max(create_until - loop.time(), 0)) + if create_until is not None + else None + ) + return models.ReadRecordsRequest( + handle=models.StreamHandle(id=stream_id), + # END applies only to the first page. Later pages start at the last page's next offset. + from_offset=from_offset if isinstance(from_offset, int) else None, + start=models.StreamReadStart.StreamReadStartEnd + if from_offset is End.END + else models.StreamReadStart.StreamReadStartBeginning + if from_offset is None and not last + else None, + last=last, + create_wait=min(create_left, _backend.longest_wait(_MAX_WAIT)) + if create_left + else None, + ) + + while True: + pages = _pages(service, request(), wait_new_records) + try: + async for page in pages: + for record in page.records or []: + try: + item = await _item(converter, record, payload_type) + if item is None: + continue + except Exception as err: + # Move next_offset past the Record, so a resume doesn't fail on it again. + # The error gives the offset, for a read of that Record. + offset = record.offset or 0 + position.next_offset = max( + position.next_offset or 0, offset + 1 + ) + raise _errors.StreamError( + f"Can't decode the Record at the offset {offset}: {err}" + ) from err + yield item + if isinstance(item, Closed): + return + if page.next_offset is not None: + # A resume starts where the page ended. + position.next_offset = max( + position.next_offset or 0, page.next_offset + ) + return + except _errors.StreamNotFoundError: + if create_until is None or loop.time() >= create_until: + raise + finally: + await pages.aclose() + + +async def _pages( + service: service_client.StreamServiceClientExtension, + request: models.ReadRecordsRequest, + wait_new_records: bool, +) -> AsyncGenerator[models.ReadRecordsResponse, None]: + # A read that waits long-polls each page. A read with no wait reads pages until a page has + # no Record and doesn't advance. + if request.last: + # The newest Records come in one page that never waits. A waiting read then continues + # after them. + page = await _read_page(service, request) + yield page + if page.done or not wait_new_records or page.next_offset is None: + return + request = dataclasses.replace(request, last=None, from_offset=page.next_offset) + while True: + if wait_new_records: + page = await _read_until_records( + service, request, _backend.longest_wait(_MAX_WAIT) + ) + else: + page = await _read_page(service, request) + yield page + if page.done: + return + if not wait_new_records and ( + page.next_offset is None + # An empty page that didn't advance is at the newest Record. + or (not page.records and page.next_offset <= (request.from_offset or 0)) + ): + return + if page.next_offset is not None: + request = dataclasses.replace(request, from_offset=page.next_offset) + + +async def _read_page( + service: service_client.StreamServiceClientExtension, + request: models.ReadRecordsRequest, +) -> models.ReadRecordsResponse: + # The bodies stay RawValue, so the read decodes each one with the Stream's payload type. + return await service._stream_service_read_records( + request, temporalio.common.RawValue + ) + + +async def _read_until_records( + service: service_client.StreamServiceClientExtension, + request: models.ReadRecordsRequest, + deadline: timedelta, +) -> models.ReadRecordsResponse: + # Each read waits for what is left of the deadline, and an empty page moves the next read to + # its next offset. At the deadline the last page comes back, so the caller keeps its place. + loop = asyncio.get_running_loop() + end = loop.time() + deadline.total_seconds() + pause = _FIRST_PAUSE + while True: + wait = min(max(end - loop.time(), 0.001), _MAX_WAIT.total_seconds()) + request = dataclasses.replace(request, wait=timedelta(seconds=wait)) + asked = loop.time() + page = await _read_page(service, request) + if page.records or page.done: + return page + if page.next_offset is not None: + request = dataclasses.replace(request, from_offset=page.next_offset) + left = end - loop.time() + if left <= 0: + return page + if loop.time() - asked < wait / 2: + await asyncio.sleep(min(pause, left)) + pause = min(pause * 2, _LAST_PAUSE) + else: + pause = _FIRST_PAUSE + + +async def _item( + converter: temporalio.converter.DataConverter, + record: models.StreamRecord[temporalio.common.RawValue], + payload_type: Any, +) -> Item | None: + offset = record.offset or 0 + match record.item: + case None: + # A newer backend's item kind that this SDK doesn't know is empty, so the read skips + # it. + return None + case models.StreamItemData(value=data): + value = _convert.decode(converter, data.body, payload_type) + return Record(data=value, offset=offset, producer=data.producer) + case models.StreamItemSuperseded(value=facts): + return Superseded(offset=offset, facts=facts) + case models.StreamItemClose(value=close): + result = _convert.decode(converter, close.result, None) + failure = close.failure + # A close with no status takes it from its failure. + status = close.status or ( + CloseStatus.FAILED if failure is not None else CloseStatus.COMPLETED + ) + return Closed( + offset=offset, + result=result, + failure=_convert.failure(converter, failure) + if failure is not None + else None, + status=CloseStatus(status), + ) diff --git a/temporalio/common/_streams/_stream_client.py b/temporalio/common/_streams/_stream_client.py new file mode 100644 index 000000000..7e95fd3f0 --- /dev/null +++ b/temporalio/common/_streams/_stream_client.py @@ -0,0 +1,382 @@ +"""Streams in client code.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import itertools +import os +import uuid +import weakref +from collections.abc import AsyncGenerator, Mapping, Sequence +from datetime import timedelta +from typing import Any, TypeVar + +from typing_extensions import Self + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.client +import temporalio.converter +from temporalio.api.common.v1 import Payload +from temporalio.api.enums.v1 import StreamProducerKind +from temporalio.common._streams import _backend, _errors, _ids +from temporalio.common._streams._access import Access, ReadSpec, read_spec, subscription +from temporalio.common._streams._backend import BackendCalls +from temporalio.common._streams._defn import StreamDefinition, definition_of +from temporalio.common._streams._handle import ( + ContextCalls, + Creating, + StreamHandle, + StreamOptions, + decoding_client, + exit_failure, +) +from temporalio.common._streams._items import End +from temporalio.common._streams._plugin import client_backend +from temporalio.common._streams._reader import Item, Position, Subscription +from temporalio.common._streams._writer import StreamWriter + +T = TypeVar("T") + +# The handles of one client and Stream write as one Producer. Each client gets its own key, +# since two clients can have two runtime objects, whose writers would send clashing sequences. +_process_key = str(uuid.uuid4()) +_clients = itertools.count(1) +_CLIENT_KEYS: weakref.WeakKeyDictionary[temporalio.client.Client, str] = ( + weakref.WeakKeyDictionary() +) + + +def _new_process_key() -> None: + # A forked child has the parent's key and count, so two children would give their clients + # the same Producer ids. + global _process_key, _clients + _process_key = str(uuid.uuid4()) + _clients = itertools.count(1) + _CLIENT_KEYS.clear() + + +# Windows has no fork. +_register_at_fork = getattr(os, "register_at_fork", None) +if _register_at_fork is not None: + _register_at_fork(after_in_child=_new_process_key) + + +def _client_producer( + client: temporalio.client.Client, named: tuple[str, int] | None +) -> stream_v1.Producer: + if named is not None: + name, epoch = named + return stream_v1.Producer( + kind=StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT, + id=_ids.named_client_producer_id(name), + epoch=stream_v1.ProducerEpoch(attempt=epoch), + ) + key = _CLIENT_KEYS.get(client) + if key is None: + key = f"{_process_key}-{next(_clients)}" + _CLIENT_KEYS[client] = key + return stream_v1.Producer( + kind=StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT, + id=_ids.client_producer_id(client.identity, key), + epoch=stream_v1.ProducerEpoch(attempt=1), + ) + + +def calls(handle: StreamHandle[Any]) -> ClientCalls: + """The client calls of a handle, which the handle keeps. + + Raises: + StreamConfigError: No client decoded the handle, or the client has no backend of it. + """ + bound = handle._client_calls + if bound is None: + client = handle._client + if client is None: + # A lookup of some client would pick the wrong one when a process has two. So only + # a decode of a client binds a handle to it. + raise _errors.StreamConfigError( + f"The handle of the Stream {handle.id} has no client, since no client decoded " + "it. Use client.get_stream_handle(definition, id=handle.id)." + ) + _, backend = client_backend(client, handle.backend) + bound = ClientCalls(client, backend) + handle._client_calls = bound + assert isinstance(bound, ClientCalls) + return bound + + +class ClientCalls(ContextCalls): + """The Stream calls of one handle in client code. + + Each publish starts its write at the call. All handles of one client and Stream write as one + Producer, unless a handle is of a named Producer. + """ + + def __init__( + self, + client: temporalio.client.Client, + backend: BackendCalls, + named: tuple[str, int] | None = None, + ) -> None: + """Create the calls of a handle of ``client``. ``named`` is a Producer's name and epoch.""" + self._client = client + self._backend = backend + self._named = named + self._writer: StreamWriter | None = None + self._closed = False + + def _access(self, handle: StreamHandle[Any]) -> Access: + client = self._client + backend = self._backend + + def writer() -> StreamWriter: + if self._writer is None: + self._writer = StreamWriter( + writers=backend.writers(client.namespace, client.data_converter), + stream_id=handle.id, + producer=_client_producer(client, self._named), + free=backend.free_writer, + ) + return self._writer + + return Access( + backend=backend, + converter=client.data_converter, + namespace=client.namespace, + stream_id=handle.id, + writer=writer, + ) + + def _raise_lost_write(self) -> None: + # Client code has no end, such as a run's, to raise a write error that no caller got. So + # the handle's next call raises it. + error = self._writer.unretrieved_error() if self._writer else None + if error is not None: + raise error + + def publish( + self, + handle: StreamHandle[Any], + values: tuple[Any, ...], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> asyncio.Future[int]: + """Start a write. The write's task takes the call options.""" + self._raise_lost_write() + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return self._access(handle).publish(values) + + def subscribe( + self, + handle: StreamHandle[Any], + *, + offset: int | End | None, + last: int | None, + wait_new_records: bool, + create_wait: timedelta | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read.""" + self._raise_lost_write() + read = read_spec( + handle._payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + return subscription( + lambda position: self._items(handle, read, position), + read, + _backend.options_of(rpc_metadata, rpc_timeout), + ) + + async def _items( + self, handle: StreamHandle[Any], read: ReadSpec, position: Position + ) -> AsyncGenerator[Item, None]: + items = self._access(handle).items(read, position) + try: + async for item in items: + yield item + finally: + await items.aclose() + + async def get_offset( + self, + handle: StreamHandle[Any], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> int: + """Get the next offset.""" + self._raise_lost_write() + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + return await self._access(handle).get_offset() + + def _check_close(self, handle: StreamHandle[Any]) -> None: + if handle.workflow_id: + # Only the run chain can close a Workflow's Stream. The Store refuses a client's + # close too, but this check makes no call. + raise _errors.StreamNotOwnerError( + f"Can't close the Stream {handle.id} from a client, since a Workflow owns it" + ) + + async def close( + self, + handle: StreamHandle[Any], + result: Any, + failure: BaseException | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> None: + """Write this handle's Records, then close the Stream.""" + self._check_close(handle) + self._closed = True + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + await self._access(handle).close(result, failure) + # The close runs first, so a lost write doesn't leave the Stream open. + self._raise_lost_write() + + async def enter(self, handle: StreamHandle[Any]) -> None: + """Check that the client can close the Stream at the end of the block.""" + # Fail early, since the close at the end would raise and hide the block's own error. + self._check_close(handle) + + async def exit(self, handle: StreamHandle[Any], exc: BaseException | None) -> None: + """Close the Stream, unless this handle closed it already.""" + if self._closed: + return + close, failure = exit_failure(exc) + if close: + await self.close(handle, None, failure, {}, None) + + def producer( + self, handle: StreamHandle[Any], name: str, epoch: int + ) -> StreamHandle[Any]: + """Return a handle of the same Stream and client that writes as a named Producer.""" + if epoch < 1: + raise ValueError("A Producer's epoch is 1 or more") + named: StreamHandle[Any] = StreamHandle( + id=handle.id, + workflow_id=handle.workflow_id, + first_execution_run_id=handle.first_execution_run_id, + backend=handle.backend, + payload_type=handle._payload_type, + ) + named._client = self._client + named._client_calls = ClientCalls(self._client, self._backend, (name, epoch)) + return named + + +def create( + client: temporalio.client.Client, + definition: type[StreamDefinition[T]], + *, + id: str, + backend: str | None, + options: StreamOptions, +) -> Creating[StreamHandle[T]]: + """Create a Stream that the client owns, or get the one that exists. + + See :py:meth:`temporalio.client.Client.create_stream`. + """ + + async def create() -> StreamHandle[T]: + info = definition_of(definition) + backend_name, calls = client_backend(client, backend) + # A closed, expired or another owner's Stream fails the create. + await calls.service(client.namespace).create_stream( + definition, id=id, retention=options.get("retention") + ) + return _bound( + client, calls, id=id, backend_name=backend_name, payload_type=info.type + ) + + return Creating(create) + + +def get_stream_handle( + client: temporalio.client.Client, + definition: type[StreamDefinition[T]], + *, + id: str, + backend: str | None, +) -> StreamHandle[T]: + """Get a Stream by its id. + + See :py:meth:`temporalio.client.Client.get_stream_handle`. + """ + info = definition_of(definition) + if not id: + # An empty id would fail only at the first call, with an unclear error. + raise ValueError("The Stream id is empty") + backend_name, calls = client_backend(client, backend) + return _bound( + client, calls, id=id, backend_name=backend_name, payload_type=info.type + ) + + +def _bound( + client: temporalio.client.Client, + backend: BackendCalls, + *, + id: str, + backend_name: str, + payload_type: Any, +) -> StreamHandle[Any]: + handle: StreamHandle[Any] = StreamHandle( + id=id, backend=backend_name, payload_type=payload_type + ) + handle._client = client + handle._client_calls = ClientCalls(client, backend) + return handle + + +class _BindingPayloadConverter( + temporalio.converter.PayloadConverter, temporalio.converter.WithSerializationContext +): + # Wraps the payload converter of a client with Stream backends. A query result, an Update or + # Workflow result, or any other decode of the client binds each handle in it to the client, + # so code uses a handle that came as data with no other call. + def __init__( + self, + inner: temporalio.converter.PayloadConverter, + client: weakref.ref[temporalio.client.Client], + ) -> None: + self.inner = inner + self._client = client + + def to_payloads(self, values: Sequence[Any]) -> list[Payload]: + return self.inner.to_payloads(values) + + def from_payloads( + self, payloads: Sequence[Payload], type_hints: list[type] | None = None + ) -> list[Any]: + token = decoding_client.set(self._client()) + try: + return self.inner.from_payloads(payloads, type_hints) + finally: + decoding_client.reset(token) + + def with_context(self, context: temporalio.converter.SerializationContext) -> Self: + inner = self.inner + if isinstance(inner, temporalio.converter.WithSerializationContext): + inner = inner.with_context(context) + return self if inner is self.inner else type(self)(inner, self._client) + + +def bind_decodes( + converter: temporalio.converter.DataConverter, client: temporalio.client.Client +) -> temporalio.converter.DataConverter: + """A copy of a client's data converter whose decodes bind each Stream handle to the client.""" + inner = converter.payload_converter + if isinstance(inner, _BindingPayloadConverter): + # A config that came from another client has that client's converter. + inner = inner.inner + bound = dataclasses.replace(converter) + # The client holds its converter, so the converter holds the client weakly. + object.__setattr__( + bound, "payload_converter", _BindingPayloadConverter(inner, weakref.ref(client)) + ) + return bound diff --git a/temporalio/common/_streams/_writer.py b/temporalio/common/_streams/_writer.py new file mode 100644 index 000000000..ce05e166a --- /dev/null +++ b/temporalio/common/_streams/_writer.py @@ -0,0 +1,210 @@ +"""The writer of an Activity or a client Producer.""" + +from __future__ import annotations + +import asyncio +import weakref +from collections.abc import Callable, Generator, Sequence +from typing import TYPE_CHECKING, Any + +import temporalio.api.stream.v1 as stream_v1 +from temporalio.common._streams import _errors + +if TYPE_CHECKING: + from temporalio.common._streams.writer_system import client as writer_client + from temporalio.common._streams.writer_system.client import Writer + + +class WriteFuture(asyncio.Future[int]): + """The future of a write. A caller that gets its result or its error sets ``retrieved``. + + A write error that no caller gets still raises, at the end of a Workflow run or an Activity + attempt, or at the next call of a client's handle. + """ + + retrieved = False + + def result(self) -> int: + """The last offset of the write.""" + self.retrieved = True + return super().result() + + def exception(self) -> BaseException | None: + """The error of the write, or None.""" + self.retrieved = True + return super().exception() + + def __await__(self) -> Generator[Any, None, int]: + """Wait for the write, and return its last offset.""" + # The await of asyncio's future doesn't call result(), so this one does. + if not self.done(): + self._asyncio_future_blocking = True + yield self + return self.result() + + __iter__ = __await__ + + +def unretrieved_error(future: asyncio.Future[int]) -> BaseException | None: + """The error of a failed write that no caller got, else None.""" + if not isinstance(future, WriteFuture) or future.retrieved or future.cancelled(): + return None + return asyncio.Future.exception(future) + + +class StreamWriter: + """Writes the Records of one Producer to one Stream, through the writer client.""" + + def __init__( + self, + *, + writers: writer_client.StreamWriterClientExtension, + stream_id: str, + producer: stream_v1.Producer, + free: Callable[[int], None] | None = None, + previous_completed_task: stream_v1.WorkflowTaskEpoch | None = None, + failed_tasks: Sequence[int] = (), + ) -> None: + """Create a writer. It opens on the first write. + + ``writers`` encodes the values with the user's converter and codec. ``free`` releases + the crate's writer if Python frees this one without a close. + """ + self._writers = writers + self._stream_id = stream_id + self._producer = producer + self._free = free + self._release: weakref.finalize[Any, StreamWriter] | None = None + self._previous_completed_task = previous_completed_task + self._failed_tasks = list(failed_tasks) + self._writer: Writer | None = None + self._open_lock = asyncio.Lock() + # The last write that started. Each write waits for the one before it, so the order is + # the order of the calls, even before a write's task first runs. + self._last: asyncio.Task[int] | None = None + # The writes that failed. The first error that no caller got raises later. + self._failed: list[WriteFuture] = [] + self._closed = False + + @property + def producer(self) -> stream_v1.Producer: + """The Producer that the writer writes as.""" + return self._producer + + async def withdraw_earlier_attempts(self) -> None: + """Open the writer of a new attempt, which withdraws the Records of earlier attempts. + + A Stream that doesn't exist has nothing to withdraw. A closed or expired Stream accepts + no write, so the attempt continues, and a later write raises the real error. + """ + try: + await self.open() + except ( + _errors.StreamNotFoundError, + _errors.StreamClosedError, + _errors.StreamExpiredError, + ): + pass + + async def open(self) -> Writer: + """Open the crate's writer once. + + The open of a new attempt withdraws the Records of earlier attempts, even if it writes + nothing. + """ + from temporalio.common._streams.writer_system import models + + async with self._open_lock: + if self._writer is None: + producer = models._ProducerTransferTypeConverter().from_transfer_type( + self._producer, models.Producer + ) + previous = self._previous_completed_task + self._writer = await self._writers.open_writer( + handle=models.StreamHandle(id=self._stream_id), + producer=producer, + previous_completed_task=models.WorkflowTaskEpoch( + started_event_id=previous.started_event_id + ) + if previous is not None + else None, + failed_tasks=self._failed_tasks, + ) + if self._free is not None: + # The crate keeps a writer for each open that nobody closed or freed. So + # free it when Python frees this writer. + self._release = weakref.finalize( + self, self._free, self._writer.writer_id + ) + return self._writer + + def write(self, values: Sequence[Any]) -> WriteFuture: + """Start the write of the values as Records, and return its future. + + Writes start in the order of the calls, also when nobody awaits a future. A cancel of + the future doesn't stop the write, since the write keeps its place in the order. The + future gives the last offset. + + Raises: + StreamClosedError: The writer is closed. + """ + if self._closed: + raise _errors.StreamClosedError("The writer is closed") + loop = asyncio.get_running_loop() + task = loop.create_task(self._write_after(self._last, values)) + self._last = task + future = WriteFuture(loop=loop) + + def done(task: asyncio.Task[int]) -> None: + if self._last is task: + self._last = None + error = None if task.cancelled() else task.exception() + if future.cancelled(): + return + if task.cancelled(): + future.cancel() + elif error is not None: + future.set_exception(error) + self._failed.append(future) + else: + future.set_result(task.result()) + + task.add_done_callback(done) + return future + + def unretrieved_error(self) -> BaseException | None: + """The first write error that no caller got, or None. + + The caller owns the error, so the later errors of the same writes don't raise again. + """ + failed, self._failed = self._failed, [] + error = next((e for e in map(unretrieved_error, failed) if e is not None), None) + for future in failed: + future.retrieved = True + return error + + async def _write_after( + self, previous: asyncio.Task[int] | None, values: Sequence[Any] + ) -> int: + from temporalio.common._streams.writer_system import models + + if previous is not None: + # The earlier write's error belongs to its own caller. + await asyncio.wait([previous]) + writer = await self.open() + answer = await writer.write_records( + [models.AppendRecord(body=value) for value in values] + ) + return answer.last_offset or 0 + + async def close(self) -> None: + """Wait for the writes that started, then close the writer. The Stream stays open.""" + if self._closed: + return + self._closed = True + if self._last is not None: + await asyncio.wait([self._last]) + if self._release is not None: + self._release.detach() + if self._writer is not None: + await self._writer.close() diff --git a/tests/contrib/streams/fake_extension.py b/tests/contrib/streams/fake_extension.py new file mode 100644 index 000000000..74a22d561 --- /dev/null +++ b/tests/contrib/streams/fake_extension.py @@ -0,0 +1,175 @@ +"""An in-memory fake of the temporalio-streams runtime, for the facade tests.""" + +from __future__ import annotations + +import asyncio + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.api.streamservice.v1 as stream_pb +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.bridge.streams import StreamCallError +from temporalio.client import Client +from temporalio.common import StreamBackend +from temporalio.common._streams._backend import _ExtensionCalls +from temporalio.common._streams.proto import v1 as ext_pb +from temporalio.service import RPCStatusCode + + +def _error(code: RPCStatusCode, reason: StreamErrorReason.ValueType) -> Exception: + details = stream_v1.StreamErrorDetails(reason=reason).SerializeToString() + return StreamCallError(int(code), StreamErrorReason.Name(reason), details) + + +class _Stream: + def __init__(self) -> None: + self.records: list[stream_v1.StreamRecord] = [] + self.closed = False + self.changed = asyncio.Event() + + +class FakeExtension: + def __init__(self) -> None: + self.streams: dict[str, _Stream] = {} + self.writers: dict[int, tuple[str, stream_v1.Producer]] = {} + self.freed: list[int] = [] + self.namespaces: set[str] = set() + self.next_offset = 1 + + def _append(self, stream_id: str, record: stream_v1.StreamRecord) -> int: + stream = self.streams.get(stream_id) + if stream is None: + raise _error( + RPCStatusCode.NOT_FOUND, StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND + ) + if stream.closed: + raise _error( + RPCStatusCode.FAILED_PRECONDITION, + StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + record.offset = self.next_offset + self.next_offset += 1 + stream.records.append(record) + stream.changed.set() + stream.changed = asyncio.Event() + return record.offset + + def free_writer(self, writer_id: int) -> None: + self.freed.append(writer_id) + + async def call(self, method: str, input: bytes) -> bytes: + if method == "CreateStream": + create = stream_pb.CreateStreamRequest.FromString(input) + self.namespaces.add(create.namespace) + stream = self.streams.setdefault(create.stream_id, _Stream()) + if stream.closed: + raise _error( + RPCStatusCode.FAILED_PRECONDITION, + StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + return stream_pb.CreateStreamResponse( + handle=stream_v1.StreamHandle(id=create.stream_id), created=True + ).SerializeToString() + if method == "OpenWriter": + open_writer = ext_pb.OpenWriterRequest.FromString(input) + writer_id = len(self.writers) + 1 + self.writers[writer_id] = (open_writer.handle.id, open_writer.producer) + return ext_pb.OpenWriterResponse(writer_id=writer_id).SerializeToString() + if method == "WriteRecords": + write = ext_pb.WriteRecordsRequest.FromString(input) + stream_id, producer = self.writers[write.writer_id] + offsets = [ + self._append( + stream_id, + stream_v1.StreamRecord( + data=stream_v1.StreamData(body=record.body, producer=producer) + ), + ) + for record in write.records + ] + return ext_pb.WriteRecordsResponse( + first_offset=offsets[0] if offsets else 0, + last_offset=offsets[-1] if offsets else 0, + ).SerializeToString() + if method == "CloseWriter": + return ext_pb.CloseWriterResponse( + last_offset=self.next_offset - 1 + ).SerializeToString() + if method == "CloseStream": + close = stream_pb.CloseStreamRequest.FromString(input) + self._append( + close.handle.id, + stream_v1.StreamRecord( + close=stream_v1.StreamClose( + producer=close.producer, + result=close.result if close.HasField("result") else None, + ) + ), + ) + self.streams[close.handle.id].closed = True + return b"" + if method == "DescribeStream": + describe = stream_pb.DescribeStreamRequest.FromString(input) + records = self.streams[describe.handle.id].records + return stream_pb.DescribeStreamResponse( + last_offset=records[-1].offset if records else 0 + ).SerializeToString() + if method == "ReadRecords": + return ( + await self._read(stream_pb.ReadRecordsRequest.FromString(input)) + ).SerializeToString() + raise _error( + RPCStatusCode.UNIMPLEMENTED, + StreamErrorReason.STREAM_ERROR_REASON_UNSUPPORTED, + ) + + async def _read( + self, request: stream_pb.ReadRecordsRequest + ) -> stream_pb.ReadRecordsResponse: + loop = asyncio.get_running_loop() + create_until = loop.time() + request.create_wait.ToTimedelta().total_seconds() + while request.handle.id not in self.streams: + if loop.time() >= create_until: + raise _error( + RPCStatusCode.NOT_FOUND, + StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND, + ) + await asyncio.sleep(0.01) + stream = self.streams[request.handle.id] + offset = request.from_offset or 1 + wait = request.wait.ToTimedelta().total_seconds() + until = loop.time() + wait + while True: + found = [record for record in stream.records if record.offset >= offset] + if found: + done = any(record.HasField("close") for record in found) + return stream_pb.ReadRecordsResponse( + records=found, next_offset=found[-1].offset + 1, done=done + ) + left = until - loop.time() + if left <= 0: + return stream_pb.ReadRecordsResponse(next_offset=offset) + try: + await asyncio.wait_for(stream.changed.wait(), left) + except asyncio.TimeoutError: + pass + + +class _FakeBackend(StreamBackend): + def __init__(self, calls: _ExtensionCalls) -> None: + super().__init__("fake-streams", backend_name="fake", default=True) + self._calls = calls + + def _backend_for( + self, + client: Client, # pyright: ignore[reportUnusedParameter] + ) -> _ExtensionCalls: + return self._calls + + +def fake_client(client: Client) -> tuple[Client, FakeExtension]: + """A copy of ``client`` whose one Stream backend, ``fake``, is a fake runtime.""" + extension = FakeExtension() + backend = _FakeBackend(_ExtensionCalls(extension)) + config = client.config() + config["stream_backends"] = [backend] + return Client(**config), extension diff --git a/tests/contrib/streams/test_facade.py b/tests/contrib/streams/test_facade.py new file mode 100644 index 000000000..40ef3b9ad --- /dev/null +++ b/tests/contrib/streams/test_facade.py @@ -0,0 +1,214 @@ +import asyncio +import gc +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta + +import pytest + +from temporalio import common +from temporalio.api.common.v1 import Payload +from temporalio.api.enums.v1 import StreamErrorReason, StreamProducerKind +from temporalio.api.stream.v1 import StreamErrorDetails +from temporalio.bridge.streams import StreamCallError +from temporalio.client import Client +from temporalio.common._streams import _errors, _stream_client +from temporalio.converter import DataConverter, PayloadCodec +from temporalio.exceptions import ApplicationError +from temporalio.service import RPCError, RPCStatusCode +from tests.contrib.streams.fake_extension import fake_client + + +@dataclass +class Turn: + kind: str + text: str = "" + + +def _details(reason: StreamErrorReason.ValueType) -> bytes: + return StreamErrorDetails(reason=reason).SerializeToString() + + +def test_runtime_errors_map_to_stream_errors(): + fenced = StreamErrorReason.STREAM_ERROR_REASON_FENCED + err = _errors.from_call_error( + StreamCallError( + int(RPCStatusCode.FAILED_PRECONDITION), "fenced", _details(fenced) + ) + ) + assert isinstance(err, common.StreamFencedError) + assert err.reason == fenced + # A reason with no class of its own gives the base class, with the reason. + sequence = StreamErrorReason.STREAM_ERROR_REASON_SEQUENCE + behind = _errors.from_call_error( + StreamCallError( + int(RPCStatusCode.FAILED_PRECONDITION), "behind", _details(sequence) + ) + ) + assert type(behind) is common.StreamError and behind.reason == sequence + other = _errors.from_call_error( + StreamCallError(int(RPCStatusCode.INTERNAL), "odd", b"") + ) + assert type(other) is common.StreamError + assert other.reason == StreamErrorReason.STREAM_ERROR_REASON_UNSPECIFIED + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +async def test_client_writes_reads_and_closes(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream(Turns, id="chat-42") + first = await stream.publish(Turn("model")) + second = await stream.publish(Turn("user", "hello")) + assert second == first + 1 + assert await stream.get_offset() == second + 1 + await stream.close(result="bye") + items = [item async for item in stream.subscribe()] + assert [type(item) for item in items] == [ + common.Record, + common.Record, + common.Closed, + ] + turn, token, closed = items + assert isinstance(turn, common.Record) and turn.data == Turn("model") + assert isinstance(token, common.Record) and token.data == Turn("user", "hello") + assert turn.producer is not None + assert turn.producer.kind == StreamProducerKind.STREAM_PRODUCER_KIND_CLIENT + assert isinstance(closed, common.Closed) and closed.result == "bye" + assert extension.namespaces == {client.namespace} + + resumed = [item async for item in stream.subscribe(offset=token.offset)] + assert [item.offset for item in resumed] == [token.offset, closed.offset] + + +async def test_read_waits_for_create(client: Client): + client, _ = fake_client(client) + stream = client.get_stream_handle(Tokens, id="wf-1:run-1:tokens") + with pytest.raises(common.StreamNotFoundError): + async for _ in stream.subscribe(): + pass + + async def create_later() -> None: + created = await client.create_stream(Tokens, id=stream.id) + await created.publish("five") + await created.close() + + task = asyncio.create_task(create_later()) + items = [item async for item in stream.subscribe(create_wait=timedelta(seconds=5))] + await task + assert isinstance(items[0], common.Record) and items[0].data == "five" + assert isinstance(items[-1], common.Closed) + + +def test_stream_errors_fail_the_workflow_not_the_task(): + closed = StreamErrorReason.STREAM_ERROR_REASON_CLOSED + err = common.StreamClosedError("closed") + assert isinstance(err, ApplicationError) + assert err.non_retryable and err.type == "StreamClosedError" + assert err.reason == closed + assert not common.StreamStorageError("down").non_retryable + mapped = _errors.from_call_error( + StreamCallError( + int(RPCStatusCode.FAILED_PRECONDITION), "closed", _details(closed) + ) + ) + assert isinstance(mapped, common.StreamClosedError) and mapped.reason == closed + # A Workflow await gets the error from a Local Activity, with the reason in its details. + returned = _errors.from_activity_error( + ApplicationError( + "closed", StreamErrorDetails(reason=closed), type="StreamClosedError" + ) + ) + assert isinstance(returned, common.StreamClosedError) + assert returned.reason == closed + + +async def test_handles_of_one_client_share_a_producer_and_free_it(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream(Tokens, id="two-writers") + other = client.get_stream_handle(Tokens, id="two-writers") + await stream.publish("1") + await other.publish("2") + items = [item async for item in stream.subscribe(wait_new_records=False)] + producers = { + item.producer.id + for item in items + if isinstance(item, common.Record) and item.producer + } + assert len(producers) == 1 + # A handle that Python frees without a close frees its writer. So a server that makes a + # handle per request doesn't keep a writer for each. + del other + gc.collect() + assert len(extension.freed) == 1 + + +async def test_rpc_timeout_bounds_each_call_and_not_the_read(client: Client): + client, extension = fake_client(client) + stream = await client.create_stream(Tokens, id="slow") + call = extension.call + + async def slow_describe(method: str, input: bytes) -> bytes: + if method == "DescribeStream": + await asyncio.sleep(1) + return await call(method, input) + + setattr(extension, "call", slow_describe) + with pytest.raises(RPCError) as err: + await stream.get_offset(rpc_timeout=timedelta(milliseconds=100)) + assert err.value.status == RPCStatusCode.DEADLINE_EXCEEDED + + async def publish_later() -> None: + await asyncio.sleep(1) + await stream.publish("late") + + # Each call of a waiting read returns before the deadline, so the read continues. + reader = stream.subscribe(rpc_timeout=timedelta(milliseconds=300)) + publisher = asyncio.create_task(publish_later()) + item = await anext(reader) + await reader.aclose() + await publisher + assert isinstance(item, common.Record) and item.data == "late" + + +class _CountingCodec(PayloadCodec): + def __init__(self) -> None: + self.calls: list[tuple[str, int]] = [] + + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + self.calls.append(("encode", len(payloads))) + return list(payloads) + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + self.calls.append(("decode", len(payloads))) + return list(payloads) + + +async def test_codec_runs_once_for_each_write_and_page(client: Client): + # A remote codec costs a round trip for each call, so one call takes all the payloads of + # a message. + codec = _CountingCodec() + config = client.config() + config["data_converter"] = DataConverter(payload_codec=codec) + client, _ = fake_client(Client(**config)) + stream = await client.create_stream(Tokens, id="coded") + await stream.publish(*(str(n) for n in range(5))) + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert [item.data for item in items if isinstance(item, common.Record)] == [ + str(n) for n in range(5) + ] + assert codec.calls == [("encode", 5), ("decode", 5)] + + +async def test_forked_process_gets_its_own_client_producers(client: Client): + # A forked child has the parent's key and count, so the fork hook replaces them. + before = _stream_client._client_producer(client, None).id + _stream_client._new_process_key() + after = _stream_client._client_producer(client, None).id + assert before != after diff --git a/tests/contrib/streams/test_reader.py b/tests/contrib/streams/test_reader.py new file mode 100644 index 000000000..2bbe63e81 --- /dev/null +++ b/tests/contrib/streams/test_reader.py @@ -0,0 +1,92 @@ +from collections.abc import AsyncGenerator, Mapping +from datetime import timedelta +from typing import Any, get_origin + +from typing_extensions import override + +from temporalio.common._streams import _reader +from temporalio.common._streams._items import Record, Superseded +from temporalio.common._streams.system import client, models + + +class _Pages(client.StreamServiceClientExtension): + """Answers each read with the next page, and keeps the requests.""" + + def __init__(self, *pages: models.ReadRecordsResponse) -> None: + self.pages = list(pages) + self.requests: list[models.ReadRecordsRequest] = [] + + @override + def _create_nexus_client(self, *, service: str, endpoint: str) -> Any: + return self + + async def start_operation( + self, + operation: str, + arg: Any, + *, + id: str, + result_type: type[Any] | None = None, + rpc_metadata: Mapping[str, str | bytes] = {}, + rpc_timeout: timedelta | None = None, + ) -> Any: + del id, rpc_metadata, rpc_timeout + assert operation == "ReadRecords" + assert get_origin(result_type) is models.ReadRecordsResponse + self.requests.append(arg) + page = self.pages.pop(0) + + class Done: + async def result(self) -> Any: + return page + + return Done() + + +async def test_a_waiting_read_carries_the_offset_until_records_come(): + record = models.StreamRecord( + item=models.StreamItemData(models.StreamData(body=None)), offset=9 + ) + service = _Pages( + models.ReadRecordsResponse(next_offset=8), + models.ReadRecordsResponse(records=[record], next_offset=10), + ) + page = await _reader._read_until_records( + service, + models.ReadRecordsRequest( + handle=models.StreamHandle(id="wf:run"), from_offset=1 + ), + timedelta(seconds=5), + ) + assert page.records and page.records[0].offset == 9 + first, second = service.requests + assert (first.from_offset, second.from_offset) == (1, 8) + assert second.wait is not None + assert timedelta() < second.wait <= timedelta(seconds=60) + + +async def test_superseded_lists_each_withdrawn_offset_once(): + kind = models.StreamProducerKind.StreamProducerKindActivity + failed = models.Producer( + id="act:run:1:1", kind=kind, epoch=models.ProducerEpoch(attempt=1) + ) + other = models.Producer( + id="act:run:2:2", kind=kind, epoch=models.ProducerEpoch(attempt=1) + ) + # A batch has consecutive offsets, and other Producers' Records come between batches. + written = [(10, failed), (11, failed), (12, failed), (13, other), (20, failed)] + facts = models.StreamSuperseded( + producer_id=failed.id, producer_kind=kind, epoch=models.ProducerEpoch(attempt=2) + ) + + async def items() -> AsyncGenerator[_reader.Item, None]: + for offset, producer in written: + yield Record(data=None, offset=offset, producer=producer) + yield Superseded(offset=30, facts=facts) + yield Superseded(offset=31, facts=facts) + + read: _reader.Subscription[None] = _reader.Subscription( + items(), _reader.Position(0) + ) + superseded = [item async for item in read if isinstance(item, Superseded)] + assert [item.withdrawn_offsets for item in superseded] == [(10, 11, 12, 20), ()] diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py new file mode 100644 index 000000000..e56210d49 --- /dev/null +++ b/tests/contrib/streams/test_stream_client.py @@ -0,0 +1,139 @@ +import asyncio +import uuid +from dataclasses import dataclass + +import pytest + +from temporalio import common, workflow +from temporalio.client import Client +from temporalio.converter import DataConverter +from tests.contrib.streams.fake_extension import fake_client +from tests.helpers import new_worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +async def test_unknown_backend_raises_before_any_call(client: Client): + client, extension = fake_client(client) + with pytest.raises(common.StreamConfigError, match="no Stream backend redis"): + await client.create_stream(Turns, id="chat-1", backend="redis") + with pytest.raises(common.StreamConfigError, match="It has fake"): + client.get_stream_handle(Turns, id="chat-1", backend="redis") + assert not extension.streams and not extension.namespaces + stream = await client.create_stream(Turns, id="chat-1") + assert stream.backend == "fake" + + +async def test_client_without_backend_has_no_transport(client: Client): + with pytest.raises(common.StreamConfigError, match="no Stream backend"): + client.get_stream_handle(Turns, id="chat-1") + with pytest.raises(common.StreamConfigError, match="no Stream backend"): + await client.create_stream(Turns, id="chat-1") + + +async def test_only_a_decode_of_the_client_binds_a_handle(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Turns, id="chat-2") + converter = DataConverter.default.payload_converter + unbound = converter.from_payload( + converter.to_payload(stream), common.StreamHandle[Turn] + ) + assert unbound == stream + with pytest.raises(common.StreamConfigError, match="get_stream_handle"): + await unbound.get_offset() + payloads = await client.data_converter.encode([stream]) + [bound] = await client.data_converter.decode(payloads, [common.StreamHandle[Turn]]) + assert await bound.get_offset() == 0 + # A client made from another client's config binds to itself. + other = Client(**client.config(active_config=True)) + [rebound] = await other.data_converter.decode(payloads, [common.StreamHandle[Turn]]) + assert rebound._client is other + + +async def test_publish_starts_the_write_at_the_call(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Tokens, id="ordered") + futures = [stream.publish(str(n)) for n in range(3)] + # A cancel of a future doesn't stop its write, which already has its place. + futures[1].cancel() + stream.publish("3") + await stream.close() + items = [item async for item in stream.subscribe()] + assert [item.data for item in items if isinstance(item, common.Record)] == [ + "0", + "1", + "2", + "3", + ] + assert [await futures[0], await futures[2]] == [items[0].offset, items[2].offset] + with pytest.raises(common.StreamClosedError): + stream.publish("4") + + +async def test_next_call_raises_a_write_error_that_no_caller_got(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Tokens, id="lost") + await client.get_stream_handle(Tokens, id="lost").close() + late = stream.publish("late") + await asyncio.wait([late]) + with pytest.raises(common.StreamClosedError): + await stream.get_offset() + # The call owns the error, so the next one makes its call. + assert await stream.get_offset() > 0 + + +async def test_sync_workflow_calls_raise_in_client_code(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Tokens, id="no-nowait") + with pytest.raises(common.StreamUnsupportedError, match="sync Workflow code"): + stream.publish_nowait("1") + with pytest.raises(common.StreamUnsupportedError, match="sync Workflow code"): + stream.close_nowait() + + +@workflow.defn +class HandsOutStream: + @workflow.run + async def run(self) -> dict[str, str]: + return {"id": "handed-out"} + + @workflow.query + def stream(self) -> dict[str, str]: + # The proto JSON of a handle, as a Workflow's Stream that a query returns. + return {"id": "handed-out"} + + +async def test_query_and_workflow_results_bind_their_client(client: Client): + client, _ = fake_client(client) + created = await client.create_stream(Tokens, id="handed-out") + async with new_worker(client, HandsOutStream) as worker: + # By name, since the method's own return type is the JSON. + handle = await client.start_workflow( + "HandsOutStream", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + result_type=common.StreamHandle[str], + ) + result = await handle.result() + queried = await handle.query("stream", result_type=common.StreamHandle[str]) + # The handle names no backend, so it uses the default one. + assert queried.id == "handed-out" and queried.backend is None + await queried.publish("from the query result") + await result.publish("from the Workflow result") + items = [item async for item in created.subscribe(wait_new_records=False)] + assert [item.data for item in items if isinstance(item, common.Record)] == [ + "from the query result", + "from the Workflow result", + ]