diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index 2eef41a39..dafe658fc 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -185,6 +185,7 @@ ScheduleUpdate, ScheduleUpdateInput, ) +from ._streams import Stream from ._worker_versioning import ( BuildIdOp, BuildIdOpAddNewCompatible, @@ -397,4 +398,5 @@ "ParamType", "ReturnType", "SelfType", + "Stream", ] diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index f171bfa5c..9e8ccbb56 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,10 +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 + from ._streams import Stream + +_StreamPayloadT = TypeVar("_StreamPayloadT") class Client: @@ -911,6 +921,102 @@ 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[Stream[_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 + Stream, 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( + self, + definition: type[StreamDefinition[_StreamPayloadT]], + *, + id: str, + backend: str | None = None, + ) -> Stream[_StreamPayloadT]: + """Get a Stream by its id. It makes no call. + + To create a Stream that the client owns, use :py:meth:`create_stream`. For a handle + that came as data, such as from a Workflow's query, use :py:meth:`to_stream`. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class. Reads decode each Record with its payload + type. + 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(self, definition, id=id, backend=backend) + + def to_stream( + self, + handle: StreamHandle[_StreamPayloadT], + *, + payload_type: type[_StreamPayloadT] | None = None, + ) -> Stream[_StreamPayloadT]: + """Get the Stream of a handle that came as data, such as from a query. It makes no call. + + .. warning:: + This API is experimental. + + Args: + handle: The handle of the Stream. + payload_type: The type that reads decode each Record with. Defaults to no type, so + a read decodes each Record as a JSON value. + + Raises: + StreamConfigError: The client has no backend of the handle, or the handle has no + backend and no backend has ``default=True``. + """ + from temporalio.common._streams import _stream_client + + return _stream_client.to_stream(self, handle, payload_type=payload_type) + def get_workflow_handle_for( self, workflow: ( diff --git a/temporalio/client/_streams.py b/temporalio/client/_streams.py new file mode 100644 index 000000000..e06c5d39c --- /dev/null +++ b/temporalio/client/_streams.py @@ -0,0 +1,202 @@ +"""Streams outside Workflow code. + +This module imports the Streams API of :py:mod:`temporalio.common` on first use, so a client +with no Stream never imports it. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping +from datetime import timedelta +from types import TracebackType +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +from typing_extensions import Self + +if TYPE_CHECKING: + from temporalio.common import End, StreamHandle, Subscription + from temporalio.common._streams._stream_client import StreamCalls + +T = TypeVar("T") + + +class Stream(Generic[T]): + """A Stream outside Workflow code, with its calls. ``T`` is the payload type. + + Get one with :py:meth:`temporalio.client.Client.create_stream`, + :py:meth:`temporalio.client.Client.get_stream`, or + :py:meth:`temporalio.client.Client.to_stream` for a handle that came as data. Its + :py:attr:`handle` is the data to pass on, such as to a Workflow. + + All Streams of one client and Stream id write as one Producer, and :py:meth:`producer` + returns a Stream that writes as a named Producer. A write error that no caller got raises at + the next call of the Stream. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, handle: StreamHandle[T], payload_type: Any, calls: StreamCalls + ) -> None: + """Create a Stream. Use a client's ``create_stream``, ``get_stream`` or ``to_stream``.""" + self._handle = handle + # The type that a read decodes each value with. None decodes with no type hint. + self._payload_type = payload_type + self._calls = calls + + @property + def handle(self) -> StreamHandle[T]: + """The data that names the Stream, to pass on.""" + return self._handle + + 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 Stream keep the order of the calls. A cancel + of the future doesn't stop the write. A list is one value. The next call of the Stream + 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 object closed the Stream. + """ + return self._calls.publish(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. + """ + from temporalio.common._streams._access import read_spec + + read = read_spec( + self._payload_type, + offset=offset, + last=last, + wait_new_records=wait_new_records, + create_wait=create_wait, + ) + return self._calls.subscribe(read, rpc_metadata, 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(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 the Records that this object published, 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. A client with a Workflow's Stream raises with + no call. + """ + await self._calls.close(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() + 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 object already closed the Stream, this + does nothing. + """ + await self._calls.exit(exc) + + def producer(self, name: str, epoch: int) -> Stream[T]: + """Return this Stream as a Stream 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:`temporalio.common.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. + """ + return Stream( + self._handle, self._payload_type, self._calls.producer(name, epoch) + ) + + def __repr__(self) -> str: + """Show the handle.""" + return f"{type(self).__name__}({self._handle!r})" diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index f50662821..cce40ae49 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 from temporalio.common._streams.system.models import ( Execution, ExecutionType, @@ -136,6 +137,7 @@ "StreamSuperseded", "StreamTooLargeError", "StreamUnsupportedError", + "Subscription", "Superseded", "TypedSearchAttributes", "VersioningBehavior", @@ -1632,6 +1634,7 @@ def wait_sync(self, timeout: float | None = None) -> None: "StreamSuperseded", "StreamTooLargeError", "StreamUnsupportedError", + "Subscription", "Superseded", "WorkflowTaskEpoch", "defn", diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index eb57006c7..f07d611df 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 # nexgen generates these wire models. Handles, Records, Superseded items and errors carry them. from temporalio.common._streams.system.models import ( @@ -68,6 +69,7 @@ "StreamSuperseded", "StreamTooLargeError", "StreamUnsupportedError", + "Subscription", "Superseded", "WorkflowTaskEpoch", "defn", diff --git a/temporalio/common/_streams/_access.py b/temporalio/common/_streams/_access.py new file mode 100644 index 000000000..5e106b64b --- /dev/null +++ b/temporalio/common/_streams/_access.py @@ -0,0 +1,137 @@ +"""Stream calls outside Workflow code, which Activity and client Streams 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, 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.send_describe_stream( + models.DescribeStreamRequest(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 Stream sends the request + # with its own. + await service.send_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/_reader.py b/temporalio/common/_streams/_reader.py new file mode 100644 index 000000000..da7a435b8 --- /dev/null +++ b/temporalio/common/_streams/_reader.py @@ -0,0 +1,323 @@ +"""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 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 Stream'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: _backend.StreamServiceCalls, + 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: _backend.StreamServiceCalls, + 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: _backend.StreamServiceCalls, + request: models.ReadRecordsRequest, +) -> models.ReadRecordsResponse: + # The bodies stay RawValue, so the read decodes each one with the Stream's payload type. + return await service.send_read_records(request, temporalio.common.RawValue) + + +async def _read_until_records( + service: _backend.StreamServiceCalls, + 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..44c3a473c --- /dev/null +++ b/temporalio/common/_streams/_stream_client.py @@ -0,0 +1,337 @@ +"""Streams outside Workflow code.""" + +from __future__ import annotations + +import abc +import asyncio +import itertools +import os +import uuid +import weakref +from collections.abc import AsyncGenerator, Mapping +from datetime import timedelta +from typing import Any, TypeVar + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.client +from temporalio.api.enums.v1 import StreamProducerKind +from temporalio.client._streams import Stream +from temporalio.common._streams import _backend, _errors, _ids +from temporalio.common._streams._access import Access, ReadSpec, subscription +from temporalio.common._streams._backend import BackendCalls +from temporalio.common._streams._defn import StreamDefinition, definition_of +from temporalio.common._streams._handle import ( + Creating, + StreamHandle, + StreamOptions, + exit_failure, +) +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 Streams of one client and Stream id 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), + ) + + +class StreamCalls(abc.ABC): + """The calls of one :py:class:`Stream`, such as in client code.""" + + @abc.abstractmethod + def publish( + self, + values: tuple[Any, ...], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> asyncio.Future[int]: + """Start a write. See :py:meth:`Stream.publish`.""" + + @abc.abstractmethod + def subscribe( + self, + read: ReadSpec, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read. See :py:meth:`Stream.subscribe`.""" + + @abc.abstractmethod + async def get_offset( + self, rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None + ) -> int: + """Get the next offset. See :py:meth:`Stream.get_offset`.""" + + @abc.abstractmethod + async def close( + self, + result: Any, + failure: BaseException | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> None: + """Close the Stream. See :py:meth:`Stream.close`.""" + + @abc.abstractmethod + async def enter(self) -> None: + """Start an ``async with`` block of the Stream.""" + + @abc.abstractmethod + async def exit(self, exc: BaseException | None) -> None: + """End an ``async with`` block of the Stream. See :py:meth:`Stream.__aexit__`.""" + + @abc.abstractmethod + def producer(self, name: str, epoch: int) -> StreamCalls: + """The calls of a named Producer. See :py:meth:`Stream.producer`.""" + + +class ClientCalls(StreamCalls): + """The calls of one Stream in client code. + + Each publish starts its write at the call. All Streams of one client and Stream id write as + one Producer, unless a Stream is of a named Producer. + """ + + def __init__( + self, + client: temporalio.client.Client, + backend: BackendCalls, + handle: StreamHandle[Any], + named: tuple[str, int] | None = None, + ) -> None: + """Create the calls of ``handle``. ``named`` is a Producer's name and epoch.""" + self._client = client + self._backend = backend + self._handle = handle + self._named = named + self._writer: StreamWriter | None = None + self._closed = False + + def _access(self) -> Access: + client = self._client + backend = self._backend + stream_id = self._handle.id + + def writer() -> StreamWriter: + if self._writer is None: + self._writer = StreamWriter( + writers=backend.writers(client.namespace, client.data_converter), + stream_id=stream_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=stream_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 Stream'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, + 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().publish(values) + + def subscribe( + self, + read: ReadSpec, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Start a read.""" + self._raise_lost_write() + return subscription( + lambda position: self._items(read, position), + read, + _backend.options_of(rpc_metadata, rpc_timeout), + ) + + async def _items( + self, read: ReadSpec, position: Position + ) -> AsyncGenerator[Item, None]: + items = self._access().items(read, position) + try: + async for item in items: + yield item + finally: + await items.aclose() + + async def get_offset( + self, 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().get_offset() + + def _check_close(self) -> None: + if self._handle.owner is not None: + # 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 {self._handle.id} from a client, since a Workflow " + "owns it" + ) + + async def close( + self, + result: Any, + failure: BaseException | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> None: + """Write the Records of these calls, then close the Stream.""" + self._check_close() + self._closed = True + with _backend.using(_backend.options_of(rpc_metadata, rpc_timeout)): + await self._access().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) -> 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() + + async def exit(self, exc: BaseException | None) -> None: + """Close the Stream, unless these calls closed it already.""" + if self._closed: + return + close, failure = exit_failure(exc) + if close: + await self.close(None, failure, {}, None) + + def producer(self, name: str, epoch: int) -> StreamCalls: + """The calls of the same Stream and client that write as a named Producer.""" + if epoch < 1: + raise ValueError("A Producer's epoch is 1 or more") + return ClientCalls(self._client, self._backend, self._handle, (name, epoch)) + + +def create( + client: temporalio.client.Client, + definition: type[StreamDefinition[T]], + *, + id: str, + backend: str | None, + options: StreamOptions, +) -> Creating[Stream[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() -> Stream[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") + ) + handle: StreamHandle[T] = StreamHandle(id=id, backend=backend_name) + return _stream(client, calls, handle, info.type) + + return Creating(create) + + +def get_stream( + client: temporalio.client.Client, + definition: type[StreamDefinition[T]], + *, + id: str, + backend: str | None, +) -> Stream[T]: + """Get a Stream by its id. + + See :py:meth:`temporalio.client.Client.get_stream`. + """ + 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) + handle: StreamHandle[T] = StreamHandle(id=id, backend=backend_name) + return _stream(client, calls, handle, info.type) + + +def to_stream( + client: temporalio.client.Client, + handle: StreamHandle[T], + *, + payload_type: type[T] | None, +) -> Stream[T]: + """The Stream of a handle that came as data. + + See :py:meth:`temporalio.client.Client.to_stream`. + """ + _, calls = client_backend( + client, handle.backend, hint="a handle that has its backend name" + ) + return _stream(client, calls, handle, payload_type) + + +def _stream( + client: temporalio.client.Client, + backend: BackendCalls, + handle: StreamHandle[T], + payload_type: Any, +) -> Stream[T]: + return Stream(handle, payload_type, ClientCalls(client, backend, handle)) diff --git a/temporalio/common/_streams/_writer.py b/temporalio/common/_streams/_writer.py new file mode 100644 index 000000000..dddf2b352 --- /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 Stream. + """ + + 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..eabc1bb6e --- /dev/null +++ b/tests/contrib/streams/test_facade.py @@ -0,0 +1,232 @@ +import asyncio +import gc +import subprocess +import sys +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta + +import pytest + +import temporalio.client +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(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.handle.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_streams_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(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 Stream that Python frees without a close frees its writer. So a server that makes a + # Stream 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 + + +async def test_client_streams_are_of_the_exported_class(client: Client): + client, _ = fake_client(client) + stream = client.get_stream(Tokens, id="exported") + assert isinstance(stream, temporalio.client.Stream) + assert "Stream" in temporalio.client.__all__ + + +def test_import_of_temporalio_client_loads_no_stream_code(): + code = "import sys, temporalio.client; print('temporalio.common._streams' in sys.modules)" + loaded = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=True + ).stdout + assert loaded.strip() == "False" diff --git a/tests/contrib/streams/test_reader.py b/tests/contrib/streams/test_reader.py new file mode 100644 index 000000000..468096a11 --- /dev/null +++ b/tests/contrib/streams/test_reader.py @@ -0,0 +1,76 @@ +from collections.abc import AsyncGenerator +from datetime import timedelta +from typing import Any, cast + +from typing_extensions import override + +import temporalio.common +from temporalio.common._streams import _backend, _reader +from temporalio.common._streams._items import Record, Superseded +from temporalio.common._streams.system import models + + +class _Pages(_backend.StreamServiceCalls): + """Answers each read with the next page, and keeps the requests.""" + + def __init__(self, *pages: models.ReadRecordsResponse) -> None: + super().__init__(cast(Any, None)) + self.pages = list(pages) + self.requests: list[models.ReadRecordsRequest] = [] + + @override + async def send_read_records( + self, request: models.ReadRecordsRequest, body_type: type[Any] + ) -> models.ReadRecordsResponse[Any]: + assert body_type is temporalio.common.RawValue + self.requests.append(request) + return self.pages.pop(0) + + +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..24c308d40 --- /dev/null +++ b/tests/contrib/streams/test_stream_client.py @@ -0,0 +1,135 @@ +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(Turns, id="chat-1", backend="redis") + with pytest.raises(common.StreamConfigError, match="It has fake"): + client.to_stream(common.StreamHandle(id="chat-1", backend="redis")) + assert not extension.streams and not extension.namespaces + stream = await client.create_stream(Turns, id="chat-1") + assert stream.handle.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(Turns, id="chat-1") + with pytest.raises(common.StreamConfigError, match="no Stream backend"): + await client.create_stream(Turns, id="chat-1") + + +async def test_to_stream_gives_the_stream_of_a_handle(client: Client): + client, _ = fake_client(client) + stream = await client.create_stream(Turns, id="chat-2") + converter = DataConverter.default.payload_converter + handle = converter.from_payload( + converter.to_payload(stream.handle), common.StreamHandle[Turn] + ) + assert handle == stream.handle + typed = client.to_stream(handle, payload_type=Turn) + assert typed.handle is handle and await typed.get_offset() == 0 + await typed.publish(Turn("model", "hi")) + [record] = [r async for r in typed.subscribe(wait_new_records=False)] + assert isinstance(record, common.Record) and record.data == Turn("model", "hi") + # With no payload type, a read decodes each Record as a JSON value. + untyped = client.to_stream(handle) + [raw] = [r async for r in untyped.subscribe(wait_new_records=False)] + assert isinstance(raw, common.Record) and raw.data == { + "kind": "model", + "text": "hi", + } + + +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(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 + + +@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_give_handles(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 its Stream uses the default one. + assert queried.id == "handed-out" and queried.backend is None + await client.to_stream(queried, payload_type=str).publish("from the query result") + await client.to_stream(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", + ]