diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c60bb3316..55173da58 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -228,8 +228,8 @@ jobs: - run: uv tool install poethepoet - run: uv sync --all-extras - run: cargo install --locked nexgen --version 0.2.8 --features advanced --force - # The Stream package needs the `main` of nexgen, which no release has yet. - - run: cargo install --locked --git https://github.com/temporalio/nexgen --rev 9653a0f5a845df0c5059554f1b5813db6a74cd1a --features advanced --root "$RUNNER_TEMP/nexgen-stream" nexgen + # The Stream package needs `@nexus.helper`, which only a fork of nexgen has yet. + - run: cargo install --locked --git https://github.com/moetemp/nexgen --rev 4eab992027a449192a18ee3b3ab190a23e5ae174 --features advanced --root "$RUNNER_TEMP/nexgen-stream" nexgen - run: echo "NEX_GEN_STREAM_BIN=$RUNNER_TEMP/nexgen-stream/bin/nexgen" >> "$GITHUB_ENV" - run: poe build-develop - run: poe gen-protos diff --git a/changelog/added/chained-pipeline-wren.md b/changelog/added/chained-pipeline-wren.md new file mode 100644 index 000000000..5a2e3d47f --- /dev/null +++ b/changelog/added/chained-pipeline-wren.md @@ -0,0 +1 @@ +Add `temporalio.contrib.streams.pipeline`, which chains `filter`, `map`, `take_until` and `limit` steps over a Stream and runs the whole chain as one `take`, `filter` or `reduce` call. In Workflow code a chain is one Local Activity and one marker in History. diff --git a/changelog/added/patient-reader-heron.md b/changelog/added/patient-reader-heron.md new file mode 100644 index 000000000..8d95c851c --- /dev/null +++ b/changelog/added/patient-reader-heron.md @@ -0,0 +1,2 @@ +Add the `take`, `filter` and `reduce` methods to the client `Stream` and the Workflow `Stream`. Client, Activity and Workflow code call them to scan a Stream until a predicate matches, the Stream is done, or a limit is reached. +Add `StreamConfig.workflow_read_replay_from`. With `StreamReplaySource.STORE`, a `take` or a `filter` of Workflow code keeps only the offsets and a digest of its Records in History, and a replay reads the Records from the Store again. diff --git a/scripts/gen_nexus_system_api.py b/scripts/gen_nexus_system_api.py index 48998fc7c..0b073adca 100644 --- a/scripts/gen_nexus_system_api.py +++ b/scripts/gen_nexus_system_api.py @@ -28,8 +28,9 @@ workflow_init_path = base_dir / "temporalio" / "workflow" / "__init__.py" NEX_GEN_VERSION = "0.2.8" # The Stream package needs a newer nexgen, which rejects the other WITs. One pin is enough -# again once the api WITs and a nexgen release agree. No release has nexgen commit 9653a0f5, -# the `main` of nexgen, yet, so NEX_GEN_STREAM_BIN points at a build of it. +# again once the api WITs and a nexgen release agree. The Stream WIT uses `@nexus.helper`, +# which only the fork https://github.com/moetemp/nexgen has, at commit 4eab9920. So +# NEX_GEN_STREAM_BIN points at a build of that commit. NEX_GEN_STREAM_VERSION = "0.2.12" proto_files = [ gen_protos.api_proto_dir @@ -87,7 +88,8 @@ def nex_gen_stream_command() -> list[str]: bin_path = os.environ.get("NEX_GEN_STREAM_BIN") if not bin_path: raise RuntimeError( - "NEX_GEN_STREAM_BIN must point at a nexgen build of commit 9653a0f5" + "NEX_GEN_STREAM_BIN must point at a nexgen build of commit 4eab9920 " + "of https://github.com/moetemp/nexgen" ) version = subprocess.check_output([bin_path, "--version"], text=True).strip() if version != f"nexgen {NEX_GEN_STREAM_VERSION}": diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index 59f4b33f1..81155733b 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit 59f4b33f1d90bb5ac4582ea465b4e8597c507214 +Subproject commit 81155733bf7d88ff829780251953e6640e847825 diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index 6fa6dd710..8f567ea21 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -187,6 +187,7 @@ ) from ._stream_store import ( StreamConfig, + StreamReplaySource, StreamSelectContext, StreamStore, StreamStoreCallOptions, @@ -357,6 +358,7 @@ "CloudOperationsClient", "Plugin", "StreamConfig", + "StreamReplaySource", "StreamSelectContext", "StreamStore", "StreamStoreCallOptions", diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index 4766126ae..5b234ec88 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -24,6 +24,7 @@ import temporalio.api.errordetails.v1 import temporalio.api.failure.v1 import temporalio.api.schedule.v1 +import temporalio.api.streamservice.v1 import temporalio.api.taskqueue.v1 import temporalio.api.update.v1 import temporalio.api.workflowservice.v1 @@ -126,6 +127,7 @@ ScheduleUpdate, ScheduleUpdateInput, ) +from ._stream_read import _note_page_sizes from ._stream_store import ( _STREAM_SERVICE, StreamConfig, @@ -1686,6 +1688,8 @@ async def _start_stream_store_operation( response = await _call_stream_store( store, data_converter, input.operation, request ) + if isinstance(response, temporalio.api.streamservice.v1.ReadRecordsResponse): + _note_page_sizes(response) handle: NexusOperationHandle[Any] = NexusOperationHandle( self._client, input.id, diff --git a/temporalio/client/_stream_read.py b/temporalio/client/_stream_read.py new file mode 100644 index 000000000..4ab84910e --- /dev/null +++ b/temporalio/client/_stream_read.py @@ -0,0 +1,935 @@ +"""The scan of the take, filter and reduce calls, for client, Activity and Workflow code. + +Client and Activity code run the scan with the pages of the client ``Stream.read``. Workflow +code runs it in a Local Activity that reads the Store, and the result of that Local Activity is +the marker of the call in History. The functions of the caller can't be serialized, so they +reach the Local Activity through the registry of this module, in the memory of the Worker +process. + +A marker of a ``take`` or a ``filter`` holds the kept Records, or only their offsets and a +digest when the Worker's ``StreamConfig.workflow_read_replay_from`` is ``STORE``. A replay +of the second kind reads the Records from the Store again with :py:func:`_read_kept_records`. + +.. warning:: + This API is an experimental prototype and subject to change. +""" + +from __future__ import annotations + +import base64 +import contextvars +import dataclasses +import hashlib +import inspect +import json +import struct +import threading +import time +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable, Sequence +from datetime import timedelta +from typing import Any, Generic, Protocol, TypeVar + +import temporalio.activity +import temporalio.api.common.v1 +import temporalio.api.stream.v1 +import temporalio.api.streamservice.v1 +import temporalio.common +import temporalio.converter +import temporalio.exceptions +import temporalio.nexus.system.stream_service.client +from temporalio.nexus.system._support import _StreamCall, _StreamCallKind +from temporalio.nexus.system.stream_service.models import ( + StreamRead, + StreamReadStop, + StreamRecord, + StreamReduce, +) +from temporalio.nexus.system.stream_service.workflow._resources.stream import ( + Stream as _WorkflowStream, +) + +from ._stream_store import StreamReplaySource, StreamStore, _store_page_reader + +_RESULT_METADATA_KEY = "temporal-stream-read" +_RESULT_METADATA_VALUE = b"v3" +# A marker with this key holds the offsets of its Records, and not the Records. +_REPLAY_FROM_METADATA_KEY = "temporal-stream-replay-from" +_REPLAY_FROM_STORE = b"store" +_DIGEST_SIZE = 16 +_SHOWN_OFFSETS = 5 +_RECORDS_MESSAGE_TYPE = b"temporal.api.streamservice.v1.ReadRecordsResponse" +_MAX_PAGE_RECORDS = 1000 +# A Store waits at most this long for one page. +_MAX_PAGE_WAIT = 60.0 + + +class _CallFlush(Protocol): + """A writer that a function of a call used, such as a ``BufferedAppender``.""" + + async def flush(self) -> object: ... + + +# The writers that the functions of one take, filter or reduce call used, so the call writes +# their Records before it returns. It is unset outside such a function. +_CALL_FLUSHES: contextvars.ContextVar[list[_CallFlush] | None] = contextvars.ContextVar( + "temporal_stream_call_flushes", default=None +) + +# The size of each Record of one page that the scan of a client or an Activity reads, by its +# offset. The scan sets it around one `Stream.read`, and the client's Store call fills it. Only +# that call has the proto Records, and `max_bytes` counts their encoded size. +_PAGE_SIZES: contextvars.ContextVar[dict[bytes, int] | None] = contextvars.ContextVar( + "temporal_stream_page_sizes", default=None +) + + +def _note_page_sizes( # pyright: ignore[reportUnusedFunction] + response: temporalio.api.streamservice.v1.ReadRecordsResponse, +) -> None: + """Give the scan that waits for this page the size of each of its Records.""" + sizes = _PAGE_SIZES.get() + if sizes is None: + return + for record in response.records: + sizes[record.offset] = record.ByteSize() + + +_ToRecord = Callable[[temporalio.api.stream.v1.StreamRecord], Any] +_RecordT = TypeVar("_RecordT") + + +@dataclasses.dataclass(frozen=True) +class _PageRequest: + """One page that the scan asks for.""" + + from_offset: bytes | None + max_records: int + max_bytes: int + """The cap on the bytes of the page, or 0 for no cap.""" + wait: timedelta + + +class _Pages(ABC, Generic[_RecordT]): + """Reads the pages of one scan, and tells what the scan needs of each Record.""" + + @abstractmethod + async def read(self, request: _PageRequest) -> tuple[Sequence[_RecordT], bool]: + """Give the Records of the next page, and whether it ends a closed Stream.""" + + @abstractmethod + def offset(self, record: _RecordT) -> bytes: + """Give the offset of a Record of the last page.""" + + @abstractmethod + def size(self, record: _RecordT) -> int: + """Give the encoded size of a Record of the last page, after the codec.""" + + @abstractmethod + def typed(self, record: _RecordT) -> Any: + """Give the Record that the caller's functions get.""" + + +class _StorePages(_Pages[temporalio.api.stream.v1.StreamRecord]): + """The pages of the Local Activity of Workflow code, read from its Store. + + The Records stay proto Records, so the result in History keeps the bodies as the Store + gave them. + """ + + def __init__( + self, + store: StreamStore, + data_converter: temporalio.converter.DataConverter, + ref: temporalio.api.stream.v1.StreamRef, + to_record: _ToRecord, + ) -> None: + self._read_page = _store_page_reader(store, data_converter) + self._ref = ref + self._to_record = to_record + + async def read( + self, request: _PageRequest + ) -> tuple[Sequence[temporalio.api.stream.v1.StreamRecord], bool]: + message = temporalio.api.streamservice.v1.ReadRecordsRequest( + ref=self._ref, + from_offset=request.from_offset or b"", + max_records=request.max_records, + max_bytes=request.max_bytes, + ) + message.wait.FromTimedelta(request.wait) + response = await self._read_page(message) + return response.records, response.done + + def offset(self, record: temporalio.api.stream.v1.StreamRecord) -> bytes: + return record.offset + + def size(self, record: temporalio.api.stream.v1.StreamRecord) -> int: + return record.ByteSize() + + def typed(self, record: temporalio.api.stream.v1.StreamRecord) -> Any: + return self._to_record(record) + + +class _ClientPages(_Pages[StreamRecord[Any]]): + """The pages of client and Activity code, each one ``read`` of a client ``Stream``.""" + + def __init__( + self, stream: temporalio.nexus.system.stream_service.client.Stream[Any] + ) -> None: + self._stream = stream + self._sizes: dict[bytes, int] = {} + + async def read( + self, request: _PageRequest + ) -> tuple[Sequence[StreamRecord[Any]], bool]: + sizes: dict[bytes, int] = {} + token = _PAGE_SIZES.set(sizes) + try: + response = await self._stream.read( + from_offset=request.from_offset, + max_records=request.max_records, + wait=request.wait, + max_bytes=request.max_bytes or None, + ) + finally: + _PAGE_SIZES.reset(token) + self._sizes = sizes + return response.records, response.done + + def offset(self, record: StreamRecord[Any]) -> bytes: + return record.offset + + def size(self, record: StreamRecord[Any]) -> int: + # Only a Store call measures a Record. An interceptor that answers a read itself + # gives Records that none measured, so they count no bytes. + return self._sizes.get(record.offset, 0) + + def typed(self, record: StreamRecord[Any]) -> Any: + return record + + +def _encode_offset(value: bytes | None) -> str | None: + return None if value is None else base64.b64encode(value).decode() + + +def _decode_offset(value: str | None) -> bytes | None: + return None if value is None else base64.b64decode(value) + + +def _call_function( + label: str, function: Callable[..., Any], *args: Any, in_local_activity: bool +) -> Any: + try: + result = function(*args) + except temporalio.exceptions.ApplicationError: + raise + except Exception as err: + if not in_local_activity: + raise + # Workflow code gets this error from the Local Activity, at the `await` of its call. + raise temporalio.exceptions.ApplicationError( + str(err), type=type(err).__name__, non_retryable=True + ) from err + if inspect.isawaitable(result): + if inspect.iscoroutine(result): + result.close() + raise temporalio.exceptions.ApplicationError( + f"The {label} of a Stream call returned an awaitable. It must be a plain " + "function that returns its answer.", + type="TypeError", + non_retryable=True, + ) + return result + + +@dataclasses.dataclass +class _ScanOutcome(Generic[_RecordT]): + kept: list[_RecordT] + """The Records that the call gives back. A fold keeps none.""" + kept_after: list[bytes | None] + """For each kept Record, the offset of the Record that the scan looked at before it.""" + count: int + """How many Records the call kept or folded.""" + last_offset: bytes | None + stop: StreamReadStop + done: bool + value: Any + """The value of a fold, or ``None``.""" + + +class _ScanState(Generic[_RecordT]): + """What a scan got so far.""" + + def __init__( + self, call: _StreamCall, pages: _Pages[_RecordT], *, in_local_activity: bool + ) -> None: + self._call = call + self._pages = pages + self._in_local_activity = in_local_activity + self.kept: list[_RecordT] = [] + self.kept_after: list[bytes | None] = [] + self.count = 0 + self.bytes = 0 + self.last_offset = call.from_offset + self.value = call.initial if call.kind is _StreamCallKind.REDUCE else None + + def _call_function( + self, label: str, function: Callable[..., Any], *args: Any + ) -> Any: + return _call_function( + label, function, *args, in_local_activity=self._in_local_activity + ) + + def bytes_left(self) -> int | None: + if self._call.max_bytes is None: + return None + return self._call.max_bytes - self.bytes + + def look(self, record: _RecordT, *, ends_stream: bool) -> StreamReadStop | None: + """Look at the next Record, and give the stop that holds after it, if any.""" + call = self._call + typed: Any = None + if call.kind is not _StreamCallKind.TAKE or call.until is not None: + typed = self._pages.typed(record) + keep = True + if call.predicate is not None: + keep = bool(self._call_function("predicate", call.predicate, typed)) + size = self._pages.size(record) + # The first Record of a call is taken even when it passes the cap. + if ( + keep + and self.count > 0 + and call.max_bytes is not None + and self.bytes + size > call.max_bytes + ): + return StreamReadStop.MaxBytes + before = self.last_offset + self.last_offset = self._pages.offset(record) + if keep: + self.count += 1 + self.bytes += size + if call.function is not None: + self.value = self._call_function( + "function", call.function, self.value, typed + ) + else: + self.kept.append(record) + self.kept_after.append(before) + if call.until is not None: + args = (typed,) if call.function is None else (self.value, typed) + if self._call_function("until", call.until, *args): + return StreamReadStop.Match + if ends_stream: + return StreamReadStop.Done + if self.count >= call.max_records: + return StreamReadStop.MaxRecords + return None + + def end( + self, stop: StreamReadStop, *, done: bool = False + ) -> _ScanOutcome[_RecordT]: + return _ScanOutcome( + kept=self.kept, + kept_after=self.kept_after, + count=self.count, + last_offset=self.last_offset, + stop=stop, + done=done, + value=self.value, + ) + + +async def _scan( + pages: _Pages[_RecordT], *, call: _StreamCall, in_local_activity: bool +) -> _ScanOutcome[_RecordT]: + """Read pages after ``call.from_offset`` until a stop of ``StreamReadStop``.""" + idle_timeout = call.idle_timeout.total_seconds() + start = time.monotonic() + deadline = start + call.timeout.total_seconds() + # The idle stop counts from the last page that gave a Record. + last_arrival = start + state = _ScanState(call, pages, in_local_activity=in_local_activity) + # A filter counts only its kept Records, so the Store's cap of a page can't serve it. + page_cap = call.kind is not _StreamCallKind.FILTER + + def idle(now: float) -> bool: + return now - last_arrival >= idle_timeout + + while True: + bytes_left = state.bytes_left() + if page_cap and state.count > 0 and bytes_left is not None and bytes_left <= 0: + # Every Record has some bytes, so none fits in a cap that is used up. + return state.end(StreamReadStop.MaxBytes) + now = time.monotonic() + if now >= deadline: + return state.end( + StreamReadStop.Idle if idle(now) else StreamReadStop.Timeout + ) + wait = min( + _MAX_PAGE_WAIT, deadline - now, max(0.0, last_arrival + idle_timeout - now) + ) + records, done = await pages.read( + _PageRequest( + from_offset=state.last_offset, + max_records=min(call.max_records - state.count, _MAX_PAGE_RECORDS) + if page_cap + else _MAX_PAGE_RECORDS, + max_bytes=max(bytes_left, 1) + if page_cap and bytes_left is not None + else 0, + wait=timedelta(seconds=wait), + ) + ) + if records: + last_arrival = time.monotonic() + for index, record in enumerate(records): + # A done page ends with the last Record of the Stream. + ends_stream = done and index == len(records) - 1 + stop = state.look(record, ends_stream=ends_stream) + if stop is not None: + return state.end( + stop, done=ends_stream and stop is not StreamReadStop.MaxBytes + ) + if done and not records: + return state.end(StreamReadStop.Done, done=True) + if not records and idle(time.monotonic()): + return state.end(StreamReadStop.Idle) + + +async def _scan_and_flush( + pages: _Pages[_RecordT], *, call: _StreamCall, in_local_activity: bool +) -> _ScanOutcome[_RecordT]: + """Run :py:func:`_scan`, then write the Records that the caller's functions appended. + + In Workflow code the result goes into History only after those Records are in the Store. + """ + flushes: list[_CallFlush] = [] + token = _CALL_FLUSHES.set(flushes) + try: + outcome = await _scan(pages, call=call, in_local_activity=in_local_activity) + except Exception: + # The error of the scan wins over an error of the write. + await _flush_all(flushes, raise_error=False) + raise + finally: + _CALL_FLUSHES.reset(token) + await _flush_all(flushes, raise_error=True) + return outcome + + +async def _flush_all(flushes: list[_CallFlush], *, raise_error: bool) -> None: + for flush in flushes: + try: + await flush.flush() + except Exception: + if raise_error: + raise + + +class _RecordDigest: + """A SHA-256 over the offset, the body and the Producer of Records, in their order. + + A length before each part keeps two different splits of the same bytes apart. + """ + + def __init__(self) -> None: + self._hash = hashlib.sha256() + + def update(self, record: temporalio.api.stream.v1.StreamRecord) -> None: + for part in ( + record.offset, + record.body.SerializeToString(deterministic=True), + record.producer.SerializeToString(deterministic=True), + ): + self._hash.update(struct.pack(">Q", len(part))) + self._hash.update(part) + + def digest(self) -> bytes: + return self._hash.digest()[:_DIGEST_SIZE] + + +def _records_digest( + records: Sequence[temporalio.api.stream.v1.StreamRecord], +) -> bytes: + digest = _RecordDigest() + for record in records: + digest.update(record) + return digest.digest() + + +@dataclasses.dataclass(frozen=True) +class _StoreReplay: + """What a marker without Records holds, so a replay can read its Records again.""" + + call_id: str + """The number of the call in its Workflow run.""" + digest: bytes + """The digest of the kept Records, as the codec decoded them.""" + through: bytes | None = None + """For a take, the offset of the last kept Record. The Records are the ones after + ``from_offset`` up to this one.""" + kept: tuple[tuple[bytes | None, bytes], ...] = () + """For a filter, each kept Record as the offset before it and its own offset.""" + + +@dataclasses.dataclass(frozen=True) +class _StreamCallResult: + """The result of one call, as the payload that History keeps for a Workflow.""" + + kind: _StreamCallKind + store: str + stream_id: str + from_offset: bytes | None + last_offset: bytes | None + count: int + stop: StreamReadStop + done: bool + records: list[temporalio.api.stream.v1.StreamRecord] + """The kept Records of a take or a filter, with their bodies.""" + value: temporalio.api.common.v1.Payload | None + """The value of a fold.""" + replay: _StoreReplay | None = None + """What a replay needs to read the Records from the Store, for a marker without them.""" + + def _header(self) -> dict[str, Any]: + return { + "kind": self.kind.value, + "store": self.store, + "stream_id": self.stream_id, + "from_offset": _encode_offset(self.from_offset), + "last_offset": _encode_offset(self.last_offset), + "count": self.count, + "stop": self.stop.name, + "done": self.done, + } + + def to_payload(self) -> temporalio.api.common.v1.Payload: + """Give the result as one payload. The bodies stay in their binary form.""" + header = self._header() + metadata = { + "encoding": b"binary/plain", + _RESULT_METADATA_KEY: _RESULT_METADATA_VALUE, + } + replay = self.replay + if replay is not None: + metadata[_REPLAY_FROM_METADATA_KEY] = _REPLAY_FROM_STORE + header["call_id"] = replay.call_id + header["digest"] = replay.digest.hex() + if self.kind is _StreamCallKind.TAKE: + header["through"] = _encode_offset(replay.through) + else: + header["kept"] = [ + [_encode_offset(before), _encode_offset(offset)] + for before, offset in replay.kept + ] + parts = [ + temporalio.api.common.v1.Payload( + metadata={"encoding": b"json/plain"}, + data=json.dumps(header, separators=(",", ":"), sort_keys=True).encode(), + ) + ] + if self.value is not None: + parts.append(self.value) + elif replay is None: + parts.append( + temporalio.api.common.v1.Payload( + metadata={ + "encoding": b"binary/protobuf", + "messageType": _RECORDS_MESSAGE_TYPE, + }, + data=temporalio.api.streamservice.v1.ReadRecordsResponse( + records=self.records + ).SerializeToString(), + ) + ) + return temporalio.api.common.v1.Payload( + metadata=metadata, + data=temporalio.api.common.v1.Payloads(payloads=parts).SerializeToString(), + ) + + @classmethod + def from_payload( + cls, payload: temporalio.api.common.v1.Payload + ) -> _StreamCallResult | None: + """Give the result in ``payload``, or ``None`` when the payload isn't one.""" + if payload.metadata.get(_RESULT_METADATA_KEY) != _RESULT_METADATA_VALUE: + return None + parts = temporalio.api.common.v1.Payloads.FromString(payload.data).payloads + header = json.loads(parts[0].data) + kind = _StreamCallKind(header["kind"]) + records: list[temporalio.api.stream.v1.StreamRecord] = [] + value = None + replay = None + if kind is _StreamCallKind.REDUCE: + value = parts[1] + elif payload.metadata.get(_REPLAY_FROM_METADATA_KEY) == _REPLAY_FROM_STORE: + replay = _StoreReplay( + call_id=header["call_id"], + digest=bytes.fromhex(header["digest"]), + through=_decode_offset(header.get("through")), + kept=tuple( + (_decode_offset(before), _decode_offset(offset) or b"") + for before, offset in header.get("kept", []) + ), + ) + else: + records = list( + temporalio.api.streamservice.v1.ReadRecordsResponse.FromString( + parts[1].data + ).records + ) + return cls( + kind=kind, + store=header["store"], + stream_id=header["stream_id"], + from_offset=_decode_offset(header["from_offset"]), + last_offset=_decode_offset(header["last_offset"]), + count=header["count"], + stop=StreamReadStop[header["stop"]], + done=header["done"], + records=records, + value=value, + replay=replay, + ) + + def describe_records(self) -> str: + """Name the Stream and the offsets of the kept Records, for an error message.""" + start = self.from_offset.hex() if self.from_offset else "the beginning" + where = f"the Stream '{self.stream_id}' on the Store '{self.store}'" + replay = self.replay + if replay is None or self.kind is _StreamCallKind.TAKE: + through = replay.through if replay is not None else self.last_offset + end = through.hex() if through else start + return f"{where}, {self.count} Records after {start} up to {end}" + offsets = [offset.hex() for _, offset in replay.kept] + shown = ", ".join(offsets[:_SHOWN_OFFSETS]) + if len(offsets) > _SHOWN_OFFSETS: + shown += f" and {len(offsets) - _SHOWN_OFFSETS} more" + return f"{where}, {self.count} Records at the offsets {shown}" + + +def _result_payload( + *, + call: _StreamCall, + call_id: str, + replay_from: StreamReplaySource, + store: str, + stream_id: str, + outcome: _ScanOutcome[temporalio.api.stream.v1.StreamRecord], + payload_converter: temporalio.converter.PayloadConverter, +) -> tuple[_StreamCallResult, temporalio.api.common.v1.Payload]: + """Give the result for History, as itself and as a payload. + + Raises when a result with Records or a value is over the limit. + """ + replay = None + if replay_from is StreamReplaySource.STORE and call.kind is not ( + _StreamCallKind.REDUCE + ): + replay = _StoreReplay( + call_id=call_id, + digest=_records_digest(outcome.kept), + through=outcome.kept[-1].offset + if call.kind is _StreamCallKind.TAKE and outcome.kept + else None, + kept=tuple( + (before, record.offset) + for before, record in zip(outcome.kept_after, outcome.kept) + ) + if call.kind is _StreamCallKind.FILTER + else (), + ) + result = _StreamCallResult( + kind=call.kind, + store=store, + stream_id=stream_id, + from_offset=call.from_offset, + last_offset=outcome.last_offset, + count=outcome.count, + stop=outcome.stop, + done=outcome.done, + records=outcome.kept if replay is None else [], + value=payload_converter.to_payload(outcome.value) + if call.kind is _StreamCallKind.REDUCE + else None, + replay=replay, + ) + payload = result.to_payload() + if replay is not None: + # The marker holds no Record, so its size doesn't grow with their bodies. + return result, payload + size = payload.ByteSize() + limit = temporalio.common.MAX_STREAM_RESULT_BYTES + if size > limit: + raise temporalio.exceptions.ApplicationError( + f"The result of the {call.kind.value} of the Stream '{stream_id}' is {size} " + f"bytes encoded, over the limit of {limit} bytes. Lower " + "max_bytes or max_records, or fold to a smaller value.", + type="StreamResultTooLarge", + non_retryable=True, + ) + return result, payload + + +@dataclasses.dataclass(eq=False) +class _PendingStreamCall: + """One call of Workflow code, from the start of its Local Activity to its result. + + It holds only what can't go into the Local Activity's input. + """ + + call: _StreamCall + """The checked call, with the caller's functions.""" + + store: StreamStore + """The Store that the call reads.""" + + ref: temporalio.api.stream.v1.StreamRef + + to_record: _ToRecord + """Turns a proto Record into the typed Record of the Workflow.""" + + replay_from: StreamReplaySource = StreamReplaySource.HISTORY + """What the marker of a new call holds, from the Worker's Stream config.""" + + _filled: ( + tuple[_StreamCallResult, list[temporalio.api.stream.v1.StreamRecord]] | None + ) = None + + def fill( + self, + result: _StreamCallResult, + records: list[temporalio.api.stream.v1.StreamRecord], + ) -> None: + """Keep the Records of a marker without them, as the codec decoded them.""" + # One assignment, so a reader on another thread sees both or neither. + self._filled = (result, records) + + def records_for( + self, result: _StreamCallResult + ) -> list[temporalio.api.stream.v1.StreamRecord] | None: + """Give the Records of ``result``, or ``None`` when this call has none for it.""" + filled = self._filled + if filled is None or filled[0] != result: + return None + return filled[1] + + +class _PendingStreamCalls: + """The pending calls of the process, by Workflow run id and the call's number. + + Workflow code adds and removes an entry on its own thread. The Local Activity and the + read before a replay look an entry up on the event loop of the Worker. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._calls: dict[tuple[str, str], _PendingStreamCall] = {} + + def add(self, run_id: str, call_id: str, call: _PendingStreamCall) -> None: + with self._lock: + self._calls[(run_id, call_id)] = call + + def get(self, run_id: str, call_id: str) -> _PendingStreamCall | None: + with self._lock: + return self._calls.get((run_id, call_id)) + + def remove(self, run_id: str, call_id: str, call: _PendingStreamCall) -> None: + # A replay of the same run can add a new entry under the same key before the + # code of an evicted run gets to its own removal. + with self._lock: + if self._calls.get((run_id, call_id)) is call: + del self._calls[(run_id, call_id)] + + def drop_run(self, run_id: str) -> None: + with self._lock: + for key in [key for key in self._calls if key[0] == run_id]: + del self._calls[key] + + def keys(self) -> list[tuple[str, str]]: + with self._lock: + return list(self._calls) + + +_PENDING_STREAM_CALLS = _PendingStreamCalls() + +_STREAM_CALL_ACTIVITY = "__temporal_stream_call" + + +@temporalio.activity.defn(name=_STREAM_CALL_ACTIVITY) +async def _stream_call_activity( # pyright: ignore[reportUnusedFunction] + call_id: str, +) -> temporalio.common.RawValue: + """Run the scan of one call of Workflow code, and give its result.""" + info = temporalio.activity.info() + pending = _PENDING_STREAM_CALLS.get(info.workflow_run_id or "", call_id) + if pending is None: + # The functions of the caller live in its process, so only a Local Activity of + # that Workflow can run the scan. + raise temporalio.exceptions.ApplicationError( + f"The Stream call {call_id} has no pending call in this process. Only the " + "take, filter and reduce of a Workflow Stream run this Activity.", + non_retryable=True, + ) + # The codec decodes the bodies of each page, so the caller's functions get its values. + outcome = await _scan_and_flush( + _StorePages( + pending.store, + temporalio.activity.client().data_converter, + pending.ref, + pending.to_record, + ), + call=pending.call, + in_local_activity=True, + ) + result, payload = _result_payload( + call=pending.call, + call_id=call_id, + replay_from=pending.replay_from, + store=pending.store.name, + stream_id=pending.ref.id, + outcome=outcome, + payload_converter=temporalio.activity.payload_converter(), + ) + if result.replay is not None: + # The marker has no Record, so the first run gives the Workflow the Records of + # this scan and reads none again. The Workflow compares the marker that it gets + # with the parsed one. + parsed = _StreamCallResult.from_payload(payload) + assert parsed is not None + pending.fill(parsed, outcome.kept) + return temporalio.common.RawValue(payload) + + +def _replay_error( + result: _StreamCallResult, cause: str +) -> temporalio.exceptions.StreamReplayError: + return temporalio.exceptions.StreamReplayError( + f"The replay of a {result.kind.value} can't read " + f"{result.describe_records()} from the Store: {cause}", + store=result.store, + stream_id=result.stream_id, + ) + + +async def _read_kept_records( # pyright: ignore[reportUnusedFunction] + read_page: Callable[ + [temporalio.api.streamservice.v1.ReadRecordsRequest], + Awaitable[temporalio.api.streamservice.v1.ReadRecordsResponse], + ], + result: _StreamCallResult, +) -> list[temporalio.api.stream.v1.StreamRecord]: + """Read the Records of a marker without them from the Store, and check them. + + A take reads its range of offsets in pages. A filter reads each kept Record with one + read after the offset before it, so it reads none of the Records that it dropped. + + Raises: + temporalio.exceptions.StreamReplayError: The Store failed, the Records are gone, + or they don't match the count or the digest of the marker. + """ + replay = result.replay + assert replay is not None + + ref = temporalio.api.stream.v1.StreamRef(id=result.stream_id, store=result.store) + + async def read( + after: bytes | None, max_records: int + ) -> Sequence[temporalio.api.stream.v1.StreamRecord]: + try: + response = await read_page( + temporalio.api.streamservice.v1.ReadRecordsRequest( + ref=ref, from_offset=after or b"", max_records=max_records + ) + ) + except Exception as err: + # An application error, such as a `StreamError`, already names its type. + cause = ( + str(err) + if isinstance(err, temporalio.exceptions.ApplicationError) + else f"{type(err).__name__}: {err}" + ) + raise _replay_error(result, cause) from err + return response.records + + records: list[temporalio.api.stream.v1.StreamRecord] = [] + if result.kind is _StreamCallKind.TAKE: + cursor = result.from_offset + while len(records) < result.count: + left = result.count - len(records) + page = await read(cursor, min(left, _MAX_PAGE_RECORDS)) + if not page: + raise _replay_error( + result, + f"the Store gave {len(records)} Records, and the Records after " + "them are gone", + ) + records.extend(list(page)[:left]) + cursor = records[-1].offset + last = records[-1].offset if records else None + if last != replay.through: + raise _replay_error( + result, + "the last Record that the Store gave is at " + f"{last.hex() if last else 'no offset'}, not at the end of the range", + ) + else: + if len(replay.kept) != result.count: + raise _replay_error( + result, + f"the marker has {len(replay.kept)} offsets for {result.count} Records", + ) + for before, offset in replay.kept: + page = await read(before, 1) + if not page or page[0].offset != offset: + raise _replay_error(result, f"the Record at {offset.hex()} is gone") + records.append(page[0]) + if _records_digest(records) != replay.digest: + raise _replay_error( + result, "the Records differ from the digest of the first run" + ) + return records + + +def _client_stream( + stream: Any, +) -> temporalio.nexus.system.stream_service.client.Stream[Any]: + if isinstance(stream, temporalio.nexus.system.stream_service.client.Stream): + return stream + if isinstance(stream, _WorkflowStream) and temporalio.activity.in_activity(): + # A Workflow Stream outside Workflow code is one that a function of a call or an + # Activity got, so it reads through the Activity's client. + return temporalio.activity.client().get_stream( + stream.ref, body_type=stream.body_type + ) + raise TypeError( + "A Workflow Stream can be read only in Workflow code or an Activity. Use " + "client.get_stream for a client Stream." + ) + + +async def _stream_call_outside_workflow( # pyright: ignore[reportUnusedFunction] + stream: Any, call: _StreamCall +) -> StreamRead[Any] | StreamReduce[Any]: + """Run the scan of a call of client or Activity code with the pages of ``Stream.read``. + + The result doesn't go into History, so it has no size limit and isn't encoded. + """ + outcome = await _scan_and_flush( + _ClientPages(_client_stream(stream)), call=call, in_local_activity=False + ) + if call.kind is _StreamCallKind.REDUCE: + return StreamReduce( + value=outcome.value, + count=outcome.count, + last_offset=outcome.last_offset, + stop=outcome.stop, + done=outcome.done, + ) + return StreamRead( + records=outcome.kept, + last_offset=outcome.last_offset, + stop=outcome.stop, + done=outcome.done, + ) diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index 74ed6961b..b13751ede 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -1,6 +1,7 @@ from __future__ import annotations import dataclasses +import enum from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable, Sequence from datetime import timedelta @@ -27,6 +28,40 @@ import temporalio.bridge.streams +class StreamReplaySource(enum.Enum): + """Where a replay of a ``take`` or a ``filter`` of Workflow code gets its Records. + + :py:attr:`StreamConfig.workflow_read_replay_from` holds it. It decides what a new call + records in History. A replay follows what History holds, so a change of the setting + doesn't break a Workflow that is running. A ``reduce`` always keeps its value in + History, since a replay can't run the function again. Client and Activity code have no + History, so the setting doesn't touch them. + + .. warning:: + This API is experimental and subject to change. + """ + + HISTORY = "history" + """The marker in History holds the kept Records with their bodies. + + A replay takes the Records from History, so it calls no Store and needs no Stream + config. The encoded result must fit in + :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES`. This is the default. + """ + + STORE = "store" + """The marker in History holds no Record and no body. + + It holds the offsets of the kept Records, their count and a digest of them. A replay + reads the Records from the Store again before the Workflow sees them, and checks the + count and the digest. So the replay needs the Stream config, the same payload codec, + and the Records. The Records of a call must stay in the Store, and Retention must keep + them, for as long as the Workflow can replay. When they can't be read, the Workflow + Task fails with :py:class:`temporalio.exceptions.StreamReplayError`. The result has no + size limit, and ``max_bytes=None`` means no cap on the bytes of a call. + """ + + @dataclasses.dataclass(frozen=True) class StreamStoreCallOptions: """Options of the Activity that calls a Stream Store for Workflow code. @@ -141,7 +176,8 @@ def worker_activities(self) -> Sequence[Callable[..., Any]]: The default is one Activity, named for this Store, that :py:meth:`call_from_workflow` runs as a Local Activity, or as a regular Activity for a close. A Store that overrides one of the two methods overrides the other to - match. + match. The take, filter and reduce calls of Workflow code run an Activity of the + SDK that reads the Store through its methods, so this list doesn't need it. """ return [_store_call_activity(self)] @@ -327,8 +363,9 @@ def check_response( temporalio.api.streamservice.v1.ReadRecordsRequest, temporalio.api.streamservice.v1.ReadRecordsResponse, workflow_refusal=( - "Workflow code can't read a Stream on a Store in this release. " - "Call an Activity that reads it through its client." + "Workflow code can't run a ReadRecords operation on a Store. Use the " + "take, filter or reduce method of the Workflow Stream, which reads in one " + "Local Activity." ), ), "CloseStream": _StreamOperation( @@ -394,6 +431,24 @@ def _with_local_activity_producer( request.sequence = 1 +def _store_page_reader( # pyright: ignore[reportUnusedFunction] + store: StreamStore, data_converter: temporalio.converter.DataConverter +) -> Callable[ + [temporalio.api.streamservice.v1.ReadRecordsRequest], + Awaitable[temporalio.api.streamservice.v1.ReadRecordsResponse], +]: + async def read_page( + request: temporalio.api.streamservice.v1.ReadRecordsRequest, + ) -> temporalio.api.streamservice.v1.ReadRecordsResponse: + response = await _call_stream_store( + store, data_converter, "ReadRecords", request + ) + assert isinstance(response, temporalio.api.streamservice.v1.ReadRecordsResponse) + return response + + return read_page + + @dataclasses.dataclass(frozen=True) class StreamSelectContext: """Context passed to :py:attr:`StreamConfig.store_selector` calls. @@ -479,6 +534,22 @@ class StreamConfig: Activity for a close. """ + workflow_read_replay_from: StreamReplaySource = dataclasses.field( + default=StreamReplaySource.HISTORY + ) + """Where a replay of a ``take`` or a ``filter`` of Workflow code gets its Records. + + It applies to every ``take`` and ``filter`` of Workflow code on a Worker whose client + has this config, and decides what a new call records in History. + :py:attr:`StreamReplaySource.HISTORY`, the default, puts the kept Records with their + bodies into History, so a replay needs no Store. :py:attr:`StreamReplaySource.STORE` + puts only their offsets, their count and a digest into History, so a replay reads + them from the Store again. With ``STORE`` the Records of a call must stay in the + Store for as long as the Workflow can replay. A replay follows the marker in History + and not this setting, so a Workflow that is running replays after a change of it. + ``reduce`` always keeps its value in History. + """ + _store_map: dict[str, StreamStore] = dataclasses.field( init=False, repr=False, compare=False ) diff --git a/temporalio/common.py b/temporalio/common.py index 8bcc66a4a..8aaea19db 100644 --- a/temporalio/common.py +++ b/temporalio/common.py @@ -1455,6 +1455,50 @@ class HeaderCodecBehavior(IntEnum): """Only automatically encode and decode headers in workflow activation encoding and decoding.""" +DEFAULT_STREAM_IDLE_TIMEOUT = timedelta(seconds=1) +"""How long a ``take``, ``filter`` or ``reduce`` of a Stream waits for a new Record. + +It applies when the call's ``idle_timeout`` is ``None``. + +.. warning:: + This API is experimental and subject to change. +""" + +DEFAULT_STREAM_MAX_RECORDS = 1_000 +"""The most Records of one ``take``, ``filter`` or ``reduce`` of a Stream. + +It applies when the call's ``max_records`` is ``None``. + +.. warning:: + This API is experimental and subject to change. +""" + +DEFAULT_STREAM_MAX_BYTES = 256 * 1024 +"""The most bytes of the Records of one ``take``, ``filter`` or ``reduce`` of a Stream. + +It applies when the call's ``max_bytes`` is ``None``. A ``take`` or a ``filter`` of Workflow +code on a Worker whose ``StreamConfig.workflow_read_replay_from`` is ``STORE`` has no default +cap, since no body goes into History. + +.. warning:: + This API is experimental and subject to change. +""" + +MAX_STREAM_RESULT_BYTES = 1024 * 1024 +"""The most bytes of the encoded result of one Stream call in Workflow code. + +In Workflow code the result of a ``take``, ``filter`` or ``reduce`` goes into History, so this +bounds what one call adds to it. A call over it raises +:py:class:`temporalio.exceptions.ApplicationError`. Client and Activity code have no such +limit, since their result doesn't go into History and isn't encoded. A ``take`` or a +``filter`` on a Worker whose ``StreamConfig.workflow_read_replay_from`` is ``STORE`` has no +such limit either, since its marker holds offsets and no Record. + +.. warning:: + This API is experimental and subject to change. +""" + + @dataclass class _CompositeEvent: # pyright: ignore[reportUnusedClass] # This should always be present, but is sometimes lazily set internally diff --git a/temporalio/contrib/streams/README.md b/temporalio/contrib/streams/README.md index f64fa2d00..2dbd2378d 100644 --- a/temporalio/contrib/streams/README.md +++ b/temporalio/contrib/streams/README.md @@ -2,8 +2,9 @@ > This package is experimental. Its API can change in a later release. -`temporalio.contrib.streams` holds helpers that sit on top of the Streams API of the SDK. A -helper takes a client `Stream` or a Workflow `Stream`. +`temporalio.contrib.streams` holds two helpers that sit on top of the Streams API of the SDK. +Each one takes a client `Stream` or a Workflow `Stream`, and works the same in client, Activity +and Workflow code. ## `BufferedAppender` @@ -21,3 +22,50 @@ async with BufferedAppender(tokens) as appender: ``` The end of the `async with` block writes the Records that are still in the buffer. + +## `pipeline` + +`pipeline` chains steps over a Stream: `filter`, `map`, `take_until` and `limit`. A step runs +nothing. The end of the chain runs the whole chain as one `take`, `filter` or `reduce` call of +the Stream, so in Workflow code the chain is one Local Activity and one marker in History. + +```python +from datetime import timedelta + +from temporalio.contrib.streams import pipeline + +text = await ( + pipeline(tokens) + .filter(lambda record: record.body.kind == "text") + .map(lambda record: record.body.text) + .take_until(lambda text: text.endswith("\n")) + .reduce(lambda joined, text: joined + text, "", timeout=timedelta(minutes=1)) +) +``` + +An item is a `StreamRecord` until the first `map`, and the value of that map after it. The +functions of the steps must keep no state, since a Local Activity can run them again. + +Each end takes `timeout`, `from_offset`, `idle_timeout`, `max_records` and `max_bytes`, as the +three calls do. Each result has `last_offset`, `stop` and `done`. `last_offset` is where the +scan stopped, so a loop that gives it as the next `from_offset` loses and repeats no Record. + +| End | Call of the Stream | Gives back | Goes into History in Workflow code | +| --- | --- | --- | --- | +| `records(...)` | one `filter`, or one `take` when the chain keeps every Record | `StreamRead[BodyT]` with the kept Records | the Records | +| `reduce(function, initial, ...)` | one `reduce` | `StreamReduce[ValueT]` with the folded value | the value | +| `collect(...)` | one `reduce`, or as `records` when the chain has no `map` | `StreamReduce[list[T]]` with the items | the list | +| `count(...)` | one `reduce` | `StreamReduce[int]` with the number of items | the count | + +`records` runs only a chain of `filter` steps, with at most one `take_until` and at most one +`limit` as its last step. For any other chain it raises `TypeError`, and `collect` or +`reduce` runs it. In Workflow code a result must fit in 1 MiB once encoded, since it goes +into History. + +The table shows the default. A Worker whose `StreamConfig` has +`workflow_read_replay_from=StreamReplaySource.STORE` keeps the Records of `records`, and of +`collect` with no `map`, out of History. Their marker then holds the offsets of the Records, a +count and a digest, with no size limit and no default `max_bytes`. A replay reads the Records +from the Store again, so they must stay there for as long as the Workflow can replay. The +folding ends, `reduce`, `count` and `collect` with a `map`, keep their value in History under +both settings. diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index d1ee83436..d9c4d119c 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -1,14 +1,18 @@ """Helpers that sit on top of the Streams API. -:py:class:`BufferedAppender` batches the appends of one Stream. It takes a client Stream or a -Workflow Stream. +:py:class:`BufferedAppender` batches the appends of one Stream. :py:func:`pipeline` chains +steps over a Stream and runs the whole chain as one ``take``, ``filter`` or ``reduce`` call. +Both take a client Stream or a Workflow Stream. .. warning:: This package is experimental and may change in future versions. """ from temporalio.contrib.streams._appender import BufferedAppender +from temporalio.contrib.streams._pipeline import Pipeline, pipeline __all__ = [ "BufferedAppender", + "Pipeline", + "pipeline", ] diff --git a/temporalio/contrib/streams/_appender.py b/temporalio/contrib/streams/_appender.py index 7b194ddf0..af524cfe4 100644 --- a/temporalio/contrib/streams/_appender.py +++ b/temporalio/contrib/streams/_appender.py @@ -7,12 +7,15 @@ from __future__ import annotations import asyncio +from collections.abc import Awaitable, Callable from datetime import timedelta from types import TracebackType -from typing import TYPE_CHECKING, Generic, TypeVar +from typing import TYPE_CHECKING, Any, Generic, TypeVar from typing_extensions import Self +import temporalio.activity +import temporalio.client._stream_read import temporalio.nexus.system.stream_service.client import temporalio.workflow @@ -48,6 +51,12 @@ class BufferedAppender(Generic[BodyT]): one write. Each write is the Workflow ``append`` of the Stream, so it is one Local Activity and one marker in History. A replay takes the results from History and writes nothing. + * In a function of the ``take``, ``filter`` or ``reduce`` of a Stream, the appender + writes as in an Activity, also when it was made with a Workflow Stream. In Workflow code that function runs in the call's Local Activity, + so the Records get the Producer and epoch of that Local Activity. The call writes + the buffered Records before it returns, and a failed write fails the call. A replay + doesn't call the function, so it writes nothing. A Workflow Task that fails and runs + again calls it again, and the appender writes again with a higher epoch. Use it in ``async with``, which writes the remaining Records at the end of the block, or await :py:meth:`flush`. Without ``async with`` or a :py:meth:`flush`, the Records still in @@ -103,6 +112,10 @@ def __init__( self._error: BaseException | None = None self._last_offset: bytes | None = None self._unwritten = 0 + # A Workflow Stream can't append outside Workflow code, so it writes through this. + self._client_stream: ( + temporalio.nexus.system.stream_service.client.Stream[BodyT] | None + ) = None def append(self, *records: BodyT) -> None: """Add one Record for each value to the buffer, and return at once. @@ -118,6 +131,11 @@ def append(self, *records: BodyT) -> None: if not records: raise TypeError("append takes one or more values") self._check_delay() + if not temporalio.workflow.in_workflow(): + self._client_stream = self._outside_workflow_stream() + flushes = temporalio.client._stream_read._CALL_FLUSHES.get() + if flushes is not None and self not in flushes: + flushes.append(self) loop = asyncio.get_running_loop() if not self._buffer: self._first_at = loop.time() @@ -171,6 +189,24 @@ def _raise_error(self) -> None: if self._error is not None: raise self._error + def _outside_workflow_stream( + self, + ) -> temporalio.nexus.system.stream_service.client.Stream[BodyT]: + stream = self._stream + if isinstance(stream, temporalio.nexus.system.stream_service.client.Stream): + return stream + if self._client_stream is not None: + return self._client_stream + if not temporalio.activity.in_activity(): + raise TypeError( + "A BufferedAppender of a Workflow Stream appends only in Workflow code, " + "in an Activity, or in a function of a Stream call" + ) + # The Activity's client reaches the Store, as a client Stream of the Activity does. + return temporalio.activity.client().get_stream( + stream.ref, body_type=stream.body_type + ) + def _check_delay(self) -> None: if self._max_delay is not None and temporalio.workflow.in_workflow(): raise ValueError( @@ -209,7 +245,12 @@ async def _write(self, batch: list[BodyT]) -> None: # The Store numbers each batch, since the appender sends no `sequence` and never # sends a batch twice. A `sequence` counts for a Producer on a Stream, and two # appenders of one Activity attempt would each start at 1. - response = await self._stream.append(*batch) + append: Callable[..., Awaitable[Any]] = ( + self._stream.append + if temporalio.workflow.in_workflow() or self._client_stream is None + else self._client_stream.append + ) + response = await append(*batch) self._last_offset = response.last_offset except asyncio.CancelledError: # A later call must not report a cancel that its caller didn't ask for. diff --git a/temporalio/contrib/streams/_pipeline.py b/temporalio/contrib/streams/_pipeline.py new file mode 100644 index 000000000..f4a6069a9 --- /dev/null +++ b/temporalio/contrib/streams/_pipeline.py @@ -0,0 +1,523 @@ +"""Steps that chain over a Stream and run as one call. + +.. warning:: + This module is experimental and subject to change. +""" + +from __future__ import annotations + +import dataclasses +import enum +import types +import typing +from collections.abc import Awaitable, Callable +from datetime import timedelta +from typing import Any, Generic, Protocol, TypeVar + +import temporalio.common +from temporalio.nexus.system.stream_service import ( + StreamRead, + StreamRecord, + StreamReduce, +) + +BodyT = TypeVar("BodyT") +ItemT = TypeVar("ItemT") +NewItemT = TypeVar("NewItemT") +ValueT = TypeVar("ValueT") + +# The value that the fold of a chain carries from Record to Record. It holds the caller's +# value, the count of each `limit` step, whether the chain ended, and how many items reached +# the end. A Local Activity can run its fold again, so no step keeps a count elsewhere. +_State = tuple[Any, tuple[int, ...], bool, int] + + +class _Stream(Protocol[BodyT]): + """The three calls of a client ``Stream`` and of a Workflow ``Stream``.""" + + async def take( + self, + timeout: timedelta, + until: Callable[[StreamRecord[BodyT]], bool] | None = None, + from_offset: bytes | None = None, + idle_timeout: timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamRead[BodyT]: ... + + async def filter( + self, + predicate: Callable[[StreamRecord[BodyT]], bool], + timeout: timedelta, + until: Callable[[StreamRecord[BodyT]], bool] | None = None, + from_offset: bytes | None = None, + idle_timeout: timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamRead[BodyT]: ... + + async def reduce( + self, + function: Callable[[ValueT, StreamRecord[BodyT]], ValueT], + initial: ValueT, + timeout: timedelta, + until: Callable[[ValueT, StreamRecord[BodyT]], bool] | None = None, + from_offset: bytes | None = None, + idle_timeout: timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamReduce[ValueT]: ... + + +class _StepKind(enum.Enum): + FILTER = "filter" + MAP = "map" + TAKE_UNTIL = "take_until" + LIMIT = "limit" + + +@dataclasses.dataclass(frozen=True) +class _Step: + kind: _StepKind + function: Callable[[Any], Any] | None = None + count: int = 0 + + +def pipeline(stream: _Stream[BodyT]) -> Pipeline[StreamRecord[BodyT]]: + """Start a chain of steps over a Stream. + + Args: + stream: A client ``Stream`` or a Workflow ``Stream``. + + Returns: + A chain with no steps, whose items are the Records of the Stream. + + .. warning:: + This API is experimental and subject to change. + """ + return Pipeline(stream, ()) + + +class Pipeline(Generic[ItemT]): + """A chain of steps over a Stream. Its end runs the whole chain as one call. + + Each step gives a new chain and runs nothing. An item is a ``StreamRecord`` until the + first :py:meth:`map`, and the value of that map after it. An end makes one ``take``, + ``filter`` or ``reduce`` call of the Stream, so in Workflow code the chain is one Local + Activity and one marker in History. The marker of :py:meth:`records` follows the + Worker's :py:attr:`temporalio.client.StreamConfig.workflow_read_replay_from`, as a + ``take`` or a ``filter`` does. The folding ends, :py:meth:`reduce`, :py:meth:`collect` + with a ``map``, and :py:meth:`count`, keep their value in History under both values. + Client, Activity and Workflow code get the same result. + + The functions of the steps must be plain functions that keep no state. In Workflow code + a Local Activity can run them again after a failure, with the same function objects. + + ``last_offset`` of each result is where the scan stopped, also when a + :py:meth:`take_until` or a :py:meth:`limit` ended the chain. Give it as the + ``from_offset`` of the next end, so a loop loses and repeats no Record. + + .. warning:: + This API is experimental and subject to change. + """ + + def __init__(self, stream: _Stream[Any], steps: tuple[_Step, ...]) -> None: + """Make a chain. Use :py:func:`pipeline` to start one.""" + self._stream = stream + self._steps = steps + + def _then(self, step: _Step) -> Pipeline[Any]: + return Pipeline(self._stream, (*self._steps, step)) + + def filter(self, predicate: Callable[[ItemT], bool]) -> Pipeline[ItemT]: + """Keep the items that ``predicate`` accepts.""" + return self._then(_Step(_StepKind.FILTER, predicate)) + + def map(self, function: Callable[[ItemT], NewItemT]) -> Pipeline[NewItemT]: + """Turn each item into the value of ``function``. + + In Workflow code, :py:meth:`collect` decodes the items as the return annotation of + the last ``map``. Give one when the items are not plain JSON values. + """ + return self._then(_Step(_StepKind.MAP, function)) + + def take_until(self, predicate: Callable[[ItemT], bool]) -> Pipeline[ItemT]: + """End the chain after the first item that ``predicate`` accepts. + + That item still goes on to the next step. + """ + return self._then(_Step(_StepKind.TAKE_UNTIL, predicate)) + + def limit(self, count: int) -> Pipeline[ItemT]: + """End the chain after ``count`` items passed this step. + + Raises: + ValueError: ``count`` is less than 1. + """ + if count < 1: + raise ValueError("count must be at least 1") + return self._then(_Step(_StepKind.LIMIT, count=count)) + + def records( + self: Pipeline[StreamRecord[BodyT]], + timeout: timedelta, + from_offset: bytes | None = None, + idle_timeout: timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> Awaitable[StreamRead[BodyT]]: + """Give the Records that reach the end of the chain. + + It makes one ``filter`` call of the Stream, or one ``take`` call when the chain + keeps every Record. The filters become the call's ``predicate``, the + :py:meth:`take_until` its ``until``, and the :py:meth:`limit` its ``max_records``, + with no count of their own. A limit over ``max_records``, or over 1,000 when + ``max_records`` is ``None``, ends at that count. ``max_records`` and ``max_bytes`` + count the Records that the chain keeps. A limit ends the call with the stop + ``MaxRecords``, and a take_until with ``Match``. + + In Workflow code the marker in History follows the Worker's + :py:attr:`temporalio.client.StreamConfig.workflow_read_replay_from`. With + ``HISTORY``, the default, the Records of the result go into History, so their + encoded size must fit in 1 MiB. ``max_bytes``, 256 KiB by default, keeps them under + it. With ``STORE``, only their offsets and a digest go into History, there is no + size limit and no default cap, and a replay reads the Records from the Store again. + + The chain may hold any number of :py:meth:`filter` steps, at most one + :py:meth:`take_until`, and at most one :py:meth:`limit` as its last step. + + Args: + timeout: As in ``take`` and ``filter`` of the Stream. + from_offset: As in ``take`` and ``filter`` of the Stream. + idle_timeout: As in ``take`` and ``filter`` of the Stream. + max_records: As in ``take`` and ``filter`` of the Stream. + max_bytes: As in ``take`` and ``filter`` of the Stream. + + Returns: + The kept Records, where the scan stopped, and why. + + Raises: + TypeError: The chain has another step, such as a ``map``. The message names + that step. Such a chain ends with :py:meth:`collect` or :py:meth:`reduce`. + """ + filters: list[Callable[[Any], Any]] = [] + until: Callable[[StreamRecord[BodyT]], bool] | None = None + limit: int | None = None + for index, step in enumerate(self._steps): + name = _misfit(step, index, len(self._steps), until is not None) + if name is not None: + raise TypeError( + f"records() can't run this chain, since its step {name} doesn't fit. " + "It runs only filter steps, one take_until and a limit as the last " + "step. End the chain with collect() or reduce()." + ) + if step.kind is _StepKind.FILTER: + assert step.function is not None + filters.append(step.function) + elif step.kind is _StepKind.TAKE_UNTIL: + assert step.function is not None + until = _until_after(tuple(filters), step.function) + else: + limit = step.count + if limit is not None: + ceiling = ( + max_records + if max_records is not None + else temporalio.common.DEFAULT_STREAM_MAX_RECORDS + ) + max_records = min(limit, ceiling) + stream: _Stream[BodyT] = self._stream + if not filters: + return stream.take( + timeout, + until=until, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + return stream.filter( + _all_of(tuple(filters)), + timeout, + until=until, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + + async def reduce( + self, + function: Callable[[ValueT, ItemT], ValueT], + initial: ValueT, + timeout: timedelta, + from_offset: bytes | None = None, + idle_timeout: timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamReduce[ValueT]: + """Fold the items that reach the end of the chain into one value. + + It makes one ``reduce`` call of the Stream. ``value = function(value, item)`` runs + for each item that reaches the end. ``function`` must not change ``initial`` or a + value in place. ``count`` of the result is the number of those items. + ``max_records`` and ``max_bytes`` count every Record that the scan looks at. A + take_until or a limit that ends the chain ends the call with the stop ``Match``. + + In Workflow code the value goes into History with the counts of the chain, so its + encoded size must fit in 1 MiB. This holds under both values of + ``StreamConfig.workflow_read_replay_from``, since a replay can't run the function. Workflow code decodes the value as the return + annotation of ``function``, else as the type of ``initial``. + + Args: + function: Gives the next value from the value so far and an item. + initial: The value before the first item. + timeout: As in ``reduce`` of the Stream. + from_offset: As in ``reduce`` of the Stream. + idle_timeout: As in ``reduce`` of the Stream. + max_records: As in ``reduce`` of the Stream. + max_bytes: As in ``reduce`` of the Stream. + + Returns: + The value, the number of items that reached the end, where the scan stopped, + and why. + """ + return await self._fold( + function, + initial, + _value_type(function, initial), + timeout=timeout, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + + async def collect( + self, + timeout: timedelta, + from_offset: bytes | None = None, + idle_timeout: timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamReduce[list[ItemT]]: + """Give the items that reach the end of the chain as a list. + + A chain with a :py:meth:`map` makes one ``reduce`` call of the Stream, as + :py:meth:`reduce` does. A Record can't be the value of a fold, so a chain with no + ``map`` runs as :py:meth:`records`, with one ``take`` or ``filter`` call. + + In Workflow code the list of a ``map`` goes into History under both values of + ``StreamConfig.workflow_read_replay_from``, so its encoded size must fit in 1 MiB. + The items of a map are decoded as the return annotation of its function. Records + go into History as :py:meth:`records` puts them there, so they follow the setting. + + Args: + timeout: As in the call of the Stream. + from_offset: As in the call of the Stream. + idle_timeout: As in the call of the Stream. + max_records: As in the call of the Stream. + max_bytes: As in the call of the Stream. + + Returns: + The items, their number, where the scan stopped, and why. + + Raises: + TypeError: The chain has no ``map`` and :py:meth:`records` can't run it. + """ + maps = [step for step in self._steps if step.kind is _StepKind.MAP] + if not maps: + try: + running = typing.cast("Pipeline[StreamRecord[Any]]", self).records( + timeout, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + except TypeError as err: + raise TypeError( + "collect() of Records runs as records(), since a Record can't be the " + f"value of a fold. Add a map before collect(). {err}" + ) from None + read = await running + return StreamReduce( + value=typing.cast("list[ItemT]", list(read.records)), + count=len(read.records), + last_offset=read.last_offset, + stop=read.stop, + done=read.done, + ) + assert maps[-1].function is not None + item_type = _return_type(maps[-1].function) + empty: list[ItemT] = [] + return await self._fold( + _append_item, + empty, + types.GenericAlias(list, (Any if item_type is None else item_type,)), + timeout=timeout, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + + async def count( + self, + timeout: timedelta, + from_offset: bytes | None = None, + idle_timeout: timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamReduce[int]: + """Count the items that reach the end of the chain. + + It makes one ``reduce`` call of the Stream, as :py:meth:`reduce` does. In Workflow + code the count goes into History with the counts of the chain, a few bytes that fit + well in the 1 MiB of a result. This holds under both values of + ``StreamConfig.workflow_read_replay_from``. + + Args: + timeout: As in ``reduce`` of the Stream. + from_offset: As in ``reduce`` of the Stream. + idle_timeout: As in ``reduce`` of the Stream. + max_records: As in ``reduce`` of the Stream. + max_bytes: As in ``reduce`` of the Stream. + + Returns: + The number of items, where the scan stopped, and why. + """ + return await self._fold( + _add_one, + 0, + int, + timeout=timeout, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + + async def _fold( + self, + function: Callable[[ValueT, Any], ValueT], + initial: ValueT, + value_type: Any, + *, + timeout: timedelta, + from_offset: bytes | None, + idle_timeout: timedelta | None, + max_records: int | None, + max_bytes: int | None, + ) -> StreamReduce[ValueT]: + steps = self._steps + limits = sum(1 for step in steps if step.kind is _StepKind.LIMIT) + + def fold(state: _State, record: StreamRecord[Any]) -> _State: + value, counts, ended, reached = state + item: Any = record + next_counts = list(counts) + limit_index = 0 + for step in steps: + if step.kind is _StepKind.FILTER: + assert step.function is not None + if not step.function(item): + return value, tuple(next_counts), ended, reached + elif step.kind is _StepKind.MAP: + assert step.function is not None + item = step.function(item) + elif step.kind is _StepKind.TAKE_UNTIL: + assert step.function is not None + if step.function(item): + ended = True + else: + next_counts[limit_index] += 1 + if next_counts[limit_index] >= step.count: + ended = True + limit_index += 1 + return function(value, item), tuple(next_counts), ended, reached + 1 + + # The call decodes the value of History as the return annotation of `fold`, so the + # caller's value keeps its type in Workflow code. + fold.__annotations__ = { + "return": types.GenericAlias( + tuple, (value_type, tuple[int, ...], bool, int) + ) + } + start: _State = (initial, (0,) * limits, False, 0) + result = await self._stream.reduce( + fold, + start, + timeout, + until=_has_ended, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + value, _, _, reached = result.value + return StreamReduce( + value=value, + count=reached, + last_offset=result.last_offset, + stop=result.stop, + done=result.done, + ) + + +def _misfit(step: _Step, index: int, steps: int, has_until: bool) -> str | None: + """Name the step when ``records`` can't run it, else give ``None``.""" + if step.kind is _StepKind.MAP: + return "map" + if step.kind is _StepKind.TAKE_UNTIL and has_until: + return "take_until, the second one" + if step.kind is _StepKind.LIMIT and index != steps - 1: + return "limit, which isn't the last step" + return None + + +def _all_of( + predicates: tuple[Callable[[Any], Any], ...], +) -> Callable[[StreamRecord[Any]], bool]: + def predicate(record: StreamRecord[Any]) -> bool: + return all(check(record) for check in predicates) + + return predicate + + +def _until_after( + filters: tuple[Callable[[Any], Any], ...], until: Callable[[Any], Any] +) -> Callable[[StreamRecord[Any]], bool]: + # The take_until sees only the Records that pass the filters before it. + def predicate(record: StreamRecord[Any]) -> bool: + return all(check(record) for check in filters) and bool(until(record)) + + return predicate + + +def _has_ended(state: _State, record: StreamRecord[Any]) -> bool: + del record + return state[2] + + +def _append_item(value: list[Any], item: Any) -> list[Any]: + return [*value, item] + + +def _add_one(value: int, item: Any) -> int: + del item + return value + 1 + + +def _return_type(function: Callable[..., Any]) -> Any: + try: + return typing.get_type_hints(function).get("return") + except Exception: + return None + + +def _value_type(function: Callable[..., Any], initial: Any) -> Any: + found = _return_type(function) + if found is not None: + return found + return type(initial) if initial is not None else Any diff --git a/temporalio/exceptions.py b/temporalio/exceptions.py index 658012ee4..52cb08956 100644 --- a/temporalio/exceptions.py +++ b/temporalio/exceptions.py @@ -315,6 +315,31 @@ class StreamConfigError(RuntimeError): """ +class StreamReplayError(RuntimeError): + """A replay can't read the Records of a ``take`` or a ``filter`` from the Store again. + + A call of Workflow code on a Worker whose + :py:attr:`temporalio.client.StreamConfig.workflow_read_replay_from` is ``STORE`` keeps + only the offsets, the count and a digest of its Records in History. Before a replay + gives the call its result, the Worker reads those Records from the Store. When the + Records expired, the Stream is gone, the Store fails, the Worker or the ``Replayer`` + has no Stream config, or the Records differ from the digest, the Worker raises this + error. It fails the Workflow Task and not the Workflow, so the task runs again later. + The message names the Stream, the offsets and the cause. + + .. warning:: + This API is experimental and subject to change. + """ + + def __init__(self, message: str, *, store: str, stream_id: str) -> None: + """Initialize the error with the Store and the Stream that the replay reads.""" + super().__init__(message) + self.store = store + """The name of the Store of the call.""" + self.stream_id = stream_id + """The id of the Stream of the call.""" + + class CancelledError(FailureError): """Error raised on workflow/activity cancellation.""" diff --git a/temporalio/nexus/system/_support.py b/temporalio/nexus/system/_support.py index 55f53b59a..f4a791007 100644 --- a/temporalio/nexus/system/_support.py +++ b/temporalio/nexus/system/_support.py @@ -1,6 +1,10 @@ +from __future__ import annotations + import collections.abc import dataclasses +import enum import functools +import inspect import typing import urllib.parse import uuid @@ -20,14 +24,19 @@ import temporalio.nexus.system if typing.TYPE_CHECKING: + import temporalio.nexus.system.stream_service.client + import temporalio.nexus.system.stream_service.workflow._resources.stream from temporalio.nexus.system.stream_service.models import ( Producer, ReadRecordsRequest, ReadRecordsResponse, + StreamRead, StreamRecord, + StreamReduce, ) _BodyT = typing.TypeVar("_BodyT") +_ValueT = typing.TypeVar("_ValueT") class SignalWithStartWorkflowModelRequest(typing.Protocol): @@ -449,3 +458,482 @@ def versioning_override_from_proto( if proto.auto_upgrade: return temporalio.common.AutoUpgradeVersioningOverride() raise ValueError("unknown versioning override proto shape") + + +# The take, filter and reduce calls of a Stream. The generated `take`, `filter` and `reduce` +# methods of the client Stream and of the Workflow Stream call the three functions below. + + +class _StreamCallKind(enum.Enum): + TAKE = "take" + FILTER = "filter" + REDUCE = "reduce" + + +@dataclasses.dataclass(frozen=True) +class _StreamCall: + """One call after the checks of the helpers. It holds the caller's functions.""" + + kind: _StreamCallKind + from_offset: bytes | None + timeout: timedelta + idle_timeout: timedelta + max_records: int + max_bytes: int | None + """The cap on the bytes of the Records, or ``None`` for no cap.""" + default_max_bytes: bool = False + """Whether the caller gave no ``max_bytes``, so the cap is the default one.""" + until: collections.abc.Callable[..., object] | None = None + predicate: collections.abc.Callable[..., object] | None = None + function: collections.abc.Callable[..., object] | None = None + initial: typing.Any = None + value_type: typing.Any = None + """The type that Workflow code decodes the value of a fold as, or ``None`` for no hint.""" + + @staticmethod + def check( + kind: _StreamCallKind, + *, + timeout: timedelta, + from_offset: bytes | None, + idle_timeout: timedelta | None, + max_records: int | None, + max_bytes: int | None, + until: collections.abc.Callable[..., object] | None = None, + predicate: collections.abc.Callable[..., object] | None = None, + function: collections.abc.Callable[..., object] | None = None, + initial: typing.Any = None, + ) -> _StreamCall: + if not isinstance(timeout, timedelta) or timeout <= timedelta(): + raise ValueError("timeout must be a positive timedelta") + if idle_timeout is None: + idle_timeout = temporalio.common.DEFAULT_STREAM_IDLE_TIMEOUT + elif idle_timeout < timedelta(): + raise ValueError("idle_timeout must not be negative") + if max_records is None: + max_records = temporalio.common.DEFAULT_STREAM_MAX_RECORDS + elif max_records < 1: + raise ValueError("max_records must be at least 1") + default_max_bytes = max_bytes is None + if max_bytes is None: + max_bytes = temporalio.common.DEFAULT_STREAM_MAX_BYTES + elif max_bytes < 1: + raise ValueError("max_bytes must be at least 1") + for label, value, required in ( + ("until", until, False), + ("predicate", predicate, kind is _StreamCallKind.FILTER), + ("function", function, kind is _StreamCallKind.REDUCE), + ): + _check_plain_function(label, value, required=required) + return _StreamCall( + kind=kind, + from_offset=from_offset, + timeout=timeout, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + default_max_bytes=default_max_bytes, + until=until, + predicate=predicate, + function=function, + initial=initial, + value_type=_fold_value_type(function, initial) + if function is not None + else None, + ) + + +def _check_plain_function(label: str, value: object, *, required: bool) -> None: + if value is None: + if required: + raise TypeError(f"{label} is required") + return + if not callable(value): + raise TypeError(f"{label} must be a function") + # The scan calls the function and uses its answer at once, so it can't await one. + if inspect.iscoroutinefunction(value) or inspect.iscoroutinefunction( + getattr(value, "__call__", None) + ): + raise TypeError( + f"{label} must be a plain function, not a coroutine function. " + "The scan calls it for each Record and can't await it." + ) + + +def _fold_value_type( + function: collections.abc.Callable[..., object], initial: typing.Any +) -> typing.Any: + # The generated `reduce` takes no type for the value, so Workflow code decodes the value + # of History as the return annotation of the function, else as the type of `initial`. + try: + hints = typing.get_type_hints(function) + except Exception: + hints = {} + if "return" in hints: + return hints["return"] + return type(initial) if initial is not None else None + + +async def _run_stream_call(stream: typing.Any, call: _StreamCall) -> typing.Any: + # The generated modules import this module, and both packages below import them. So + # the two imports wait for the first call. + from temporalio.workflow._context import _Runtime + + runtime = _Runtime.maybe_current() + if runtime is not None: + return await runtime.workflow_stream_call(stream, call) + from temporalio.client._stream_read import _stream_call_outside_workflow + + return await _stream_call_outside_workflow(stream, call) + + +async def take_stream( + stream: temporalio.nexus.system.stream_service.client.Stream[_BodyT] + | temporalio.nexus.system.stream_service.workflow._resources.stream.Stream[_BodyT], + /, + *, + timeout: timedelta, + until: collections.abc.Callable[[StreamRecord[_BodyT]], bool] | None, + from_offset: bytes | None, + idle_timeout: timedelta | None, + max_records: int | None, + max_bytes: int | None, +) -> StreamRead[_BodyT]: + """Give every Record of a Stream up to a stop. The ``take`` method of a Stream calls it. + + ``until(record)`` is called for each Record, and the first Record that it accepts is the + last one given back. + + The scan starts after ``from_offset`` and reads pages with the Store's long poll. It ends + at the first stop that holds, in the order of ``StreamReadStop``. ``Match`` is a Record + that ``until`` accepted. ``Done`` is the last Record of a closed Stream. ``MaxRecords`` + and ``MaxBytes`` are the limits, and ``Idle`` and ``Timeout`` are the two waits. + ``last_offset`` in the result is where the scan stopped. Give it as the ``from_offset`` of + the next call, so a loop loses and repeats no Record. + + Client, Activity and Workflow code call it with a client ``Stream`` or a Workflow + ``Stream``. The scan, its stops, its defaults and its count of bytes are the same in + each. In a client or an Activity, each page is one ``read`` of the client ``Stream``, so + the client's interceptors see it. A Workflow ``Stream`` in an Activity reads through the + Activity's client. + + The functions of the caller must be plain functions that change no state. The scan calls + each one for a Record and uses its answer at once. A + :py:class:`temporalio.contrib.streams.BufferedAppender` can append in them, and the call + writes the appender's Records before it returns. A failed write fails the call with the + :py:class:`temporalio.exceptions.StreamError`. + + In Workflow code only, the scan runs in one Local Activity and History gets one marker. + What the marker holds follows the Worker's + :py:attr:`temporalio.client.StreamConfig.workflow_read_replay_from`. With ``HISTORY``, + the default, it holds the Records and their bodies, so the encoded result must fit in + :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES`. A replay then takes the result + from History, so it calls no Store and needs no Stream config. With ``STORE``, it holds + the range of offsets, the count and a digest of the Records, and no body, with no size + limit. A replay then reads the Records from the Store again, so they must stay there + for as long as the Workflow can replay. When it can't read them, the Workflow Task + fails with :py:class:`temporalio.exceptions.StreamReplayError`. A replay follows the + marker, not the setting. Records that the scan looked at and didn't keep stay out of + History. The Worker's data converter, with its payload codec, encodes the marker. A + replay doesn't call the functions, so an appender in them writes nothing. A Workflow Task that fails and runs again calls them + again, and an appender writes again with a higher epoch. The functions run on the + Worker's event loop, not in the Workflow's thread. They can read Workflow state through + a reference such as ``self``, but can't call ``temporalio.workflow``. Each sees a Record + once, so a later change of that state doesn't check an earlier Record again. + + Args: + stream: The Stream to scan. + timeout: The most time for the call. It must be positive. In Workflow code a call + keeps its Workflow Task open until it returns, so a Signal, an Update or a + cancel reaches the Workflow only then. The default ``idle_timeout`` bounds that + on a quiet Stream, and a short ``timeout``, such as 2 seconds, bounds it while + Records still flow. Only a call that runs longer than about 80% of the Workflow + Task timeout makes them wait for a Workflow Task heartbeat. + until: Ends the call after the first Record that it accepts. ``None`` scans + until another stop. + from_offset: Scan the Records after this offset. ``None`` is the beginning. + idle_timeout: End the call with ``Idle`` when no new Record arrived for this long. + It counts from the last Record that the scan got, or from the start of the + call. ``timedelta(0)`` gives what is in the Stream now. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow + code each Workflow Task heartbeat of a long wait adds events to History. + max_records: The most Records to give. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_MAX_RECORDS`, 1,000. + max_bytes: A cap on the bytes of the Records to give. A Record counts its encoded + size as the Store gave it, after the payload codec decoded it. A Record that + would pass the cap ends the call before that Record, and the next call gets it. + The first Record of a call is always taken, also when it is larger. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_MAX_BYTES`, 256 KiB. In Workflow + code on a Worker whose ``workflow_read_replay_from`` is ``STORE``, ``None`` is + no cap. + + Returns: + The Records, where the scan stopped, and why. + + Raises: + TypeError: A function is a coroutine function, or the Stream is of the wrong kind. + ValueError: A limit is out of range. + temporalio.exceptions.StreamError: The Store failed the call, or an appender in a + function failed its write. + temporalio.exceptions.StreamConfigError: In client or Activity code, the client + has no Stream Store. + temporalio.exceptions.ApplicationError: A function returned an awaitable. In + Workflow code also when a function raised, when the encoded result with the + Records is over :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES` with the + ``type`` ``StreamResultTooLarge``, or when the Worker's client has no Stream + config with the ``type`` ``StreamConfigError``. These are not retryable, so an + uncaught one fails the Workflow. + temporalio.workflow.NondeterminismError: In Workflow code, a replay of a call + whose Store, Stream id or ``from_offset`` differs from the one in History. + + .. warning:: + This API is an experimental prototype and subject to change. + """ + call = _StreamCall.check( + _StreamCallKind.TAKE, + timeout=timeout, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + until=until, + ) + return await _run_stream_call(stream, call) + + +async def filter_stream( + stream: temporalio.nexus.system.stream_service.client.Stream[_BodyT] + | temporalio.nexus.system.stream_service.workflow._resources.stream.Stream[_BodyT], + /, + *, + predicate: collections.abc.Callable[[StreamRecord[_BodyT]], bool], + timeout: timedelta, + until: collections.abc.Callable[[StreamRecord[_BodyT]], bool] | None, + from_offset: bytes | None, + idle_timeout: timedelta | None, + max_records: int | None, + max_bytes: int | None, +) -> StreamRead[_BodyT]: + """Give the Records that a predicate keeps, up to a stop. ``Stream.filter`` calls it. + + ``predicate(record)`` decides whether a Record is given back. ``until(record)`` is + called for every Record that the scan looks at, kept or not. ``max_records`` and + ``max_bytes`` count the kept Records. ``max_records=1`` gives the first Record that + ``predicate`` keeps. + + The scan starts after ``from_offset`` and reads pages with the Store's long poll. It ends + at the first stop that holds, in the order of ``StreamReadStop``. ``Match`` is a Record + that ``until`` accepted. ``Done`` is the last Record of a closed Stream. ``MaxRecords`` + and ``MaxBytes`` are the limits, and ``Idle`` and ``Timeout`` are the two waits. + ``last_offset`` in the result is where the scan stopped. Give it as the ``from_offset`` of + the next call, so a loop loses and repeats no Record. + + Client, Activity and Workflow code call it with a client ``Stream`` or a Workflow + ``Stream``. The scan, its stops, its defaults and its count of bytes are the same in + each. In a client or an Activity, each page is one ``read`` of the client ``Stream``, so + the client's interceptors see it. A Workflow ``Stream`` in an Activity reads through the + Activity's client. + + The functions of the caller must be plain functions that change no state. The scan calls + each one for a Record and uses its answer at once. A + :py:class:`temporalio.contrib.streams.BufferedAppender` can append in them, and the call + writes the appender's Records before it returns. A failed write fails the call with the + :py:class:`temporalio.exceptions.StreamError`. + + In Workflow code only, the scan runs in one Local Activity and History gets one marker. + What the marker holds follows the Worker's + :py:attr:`temporalio.client.StreamConfig.workflow_read_replay_from`. With ``HISTORY``, + the default, it holds the kept Records and their bodies, so the encoded result must fit + in :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES`. A replay then takes the result + from History, so it calls no Store and needs no Stream config. With ``STORE``, it holds + the offset of each kept Record and of the Record before it, the count and a digest, + and no body, with no size limit. A replay then reads only the kept Records from the + Store again, so they must stay there for as long as the Workflow can replay. When it + can't read them, the Workflow Task fails with + :py:class:`temporalio.exceptions.StreamReplayError`. A replay follows the marker, not + the setting. Records that the scan looked at and didn't keep stay out of History. The + Worker's data converter, with its payload codec, encodes the marker. A replay doesn't + call the functions, so an appender in them writes nothing. A Workflow Task that fails and runs again + calls them again, and an appender writes again with a higher epoch. The functions run + on the Worker's event loop, not in the Workflow's thread. They can read Workflow state + through a reference such as ``self``, but can't call ``temporalio.workflow``. Each sees + a Record once, so a later change of that state doesn't check an earlier Record again. + + Args: + stream: The Stream to scan. + predicate: Keeps a Record when it returns true. + timeout: The most time for the call. It must be positive. In Workflow code a call + keeps its Workflow Task open until it returns, so a Signal, an Update or a + cancel reaches the Workflow only then. The default ``idle_timeout`` bounds that + on a quiet Stream, and a short ``timeout``, such as 2 seconds, bounds it while + Records still flow. Only a call that runs longer than about 80% of the Workflow + Task timeout makes them wait for a Workflow Task heartbeat. + until: Ends the call after the first Record that it accepts, kept or not. + ``None`` scans until another stop. + from_offset: Scan the Records after this offset. ``None`` is the beginning. + idle_timeout: End the call with ``Idle`` when no new Record arrived for this long. + It counts from the last Record that the scan got, or from the start of the + call. ``timedelta(0)`` gives what is in the Stream now. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow + code each Workflow Task heartbeat of a long wait adds events to History. + max_records: The most Records to keep. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_MAX_RECORDS`, 1,000. + max_bytes: A cap on the bytes of the Records to keep. A Record counts its encoded + size as the Store gave it, after the payload codec decoded it. A Record that + would pass the cap ends the call before that Record, and the next call gets it. + The first Record of a call is always taken, also when it is larger. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_MAX_BYTES`, 256 KiB. In Workflow + code on a Worker whose ``workflow_read_replay_from`` is ``STORE``, ``None`` is + no cap. + + Returns: + The kept Records, where the scan stopped, and why. + + Raises: + TypeError: A function is a coroutine function, or the Stream is of the wrong kind. + ValueError: A limit is out of range. + temporalio.exceptions.StreamError: The Store failed the call, or an appender in a + function failed its write. + temporalio.exceptions.StreamConfigError: In client or Activity code, the client + has no Stream Store. + temporalio.exceptions.ApplicationError: A function returned an awaitable. In + Workflow code also when a function raised, when the encoded result with the + Records is over :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES` with the + ``type`` ``StreamResultTooLarge``, or when the Worker's client has no Stream + config with the ``type`` ``StreamConfigError``. These are not retryable, so an + uncaught one fails the Workflow. + temporalio.workflow.NondeterminismError: In Workflow code, a replay of a call + whose Store, Stream id or ``from_offset`` differs from the one in History. + + .. warning:: + This API is an experimental prototype and subject to change. + """ + call = _StreamCall.check( + _StreamCallKind.FILTER, + timeout=timeout, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + until=until, + predicate=predicate, + ) + return await _run_stream_call(stream, call) + + +async def reduce_stream( + stream: temporalio.nexus.system.stream_service.client.Stream[_BodyT] + | temporalio.nexus.system.stream_service.workflow._resources.stream.Stream[_BodyT], + /, + *, + function: collections.abc.Callable[[_ValueT, StreamRecord[_BodyT]], _ValueT], + initial: _ValueT, + timeout: timedelta, + until: collections.abc.Callable[[_ValueT, StreamRecord[_BodyT]], bool] | None, + from_offset: bytes | None, + idle_timeout: timedelta | None, + max_records: int | None, + max_bytes: int | None, +) -> StreamReduce[_ValueT]: + """Fold the Records of a Stream into one value, up to a stop. ``Stream.reduce`` calls it. + + The scan calls ``value = function(value, record)`` for each Record in order, then + ``until(value, record)``. No Record comes back, only the value. ``function`` must not + change ``initial`` or a value in place. + + The scan starts after ``from_offset`` and reads pages with the Store's long poll. It ends + at the first stop that holds, in the order of ``StreamReadStop``. ``Match`` is a Record + that ``until`` accepted. ``Done`` is the last Record of a closed Stream. ``MaxRecords`` + and ``MaxBytes`` are the limits, and ``Idle`` and ``Timeout`` are the two waits. + ``last_offset`` in the result is where the scan stopped. Give it as the ``from_offset`` of + the next call, so a loop loses and repeats no Record. + + Client, Activity and Workflow code call it with a client ``Stream`` or a Workflow + ``Stream``. The scan, its stops, its defaults and its count of bytes are the same in + each. In a client or an Activity, each page is one ``read`` of the client ``Stream``, so + the client's interceptors see it. A Workflow ``Stream`` in an Activity reads through the + Activity's client. + + The functions of the caller must be plain functions that change no state. The scan calls + each one for a Record and uses its answer at once. A + :py:class:`temporalio.contrib.streams.BufferedAppender` can append in them, and the call + writes the appender's Records before it returns. A failed write fails the call with the + :py:class:`temporalio.exceptions.StreamError`. + + In Workflow code only, the scan runs in one Local Activity and History gets one marker. + The marker holds the result, with the value, so the encoded result must fit in + :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES`. Workflow code decodes the value as + the return annotation of ``function``, else as the type of ``initial``, and with no type + when ``initial`` is ``None``. Records that the scan looked at stay out of History. The + value goes into History under both values of the Worker's + :py:attr:`temporalio.client.StreamConfig.workflow_read_replay_from`, since a replay + can't run the function again. The Worker's data converter, with its payload codec, + encodes the marker. A replay takes the result from History, so it calls no Store and + needs no Stream config. A replay doesn't call the functions either, so an appender in + them writes nothing. A Workflow Task that + fails and runs again calls them again, and an appender writes again with a higher epoch. + The functions run on the Worker's event loop, not in the Workflow's thread. They can + read Workflow state through a reference such as ``self``, but can't call + ``temporalio.workflow``. Each sees a Record once, so a later change of that state + doesn't check an earlier Record again. + + Args: + stream: The Stream to fold. + function: Gives the next value from the value so far and a Record. + initial: The value before the first Record. + timeout: The most time for the call. It must be positive. In Workflow code a call + keeps its Workflow Task open until it returns, so a Signal, an Update or a + cancel reaches the Workflow only then. The default ``idle_timeout`` bounds that + on a quiet Stream, and a short ``timeout``, such as 2 seconds, bounds it while + Records still flow. Only a call that runs longer than about 80% of the Workflow + Task timeout makes them wait for a Workflow Task heartbeat. + until: Called with the value after each Record and that Record. Ends the + fold after the first Record that it accepts. + from_offset: Scan the Records after this offset. ``None`` is the beginning. + idle_timeout: End the call with ``Idle`` when no new Record arrived for this long. + It counts from the last Record that the scan got, or from the start of the + call. ``timedelta(0)`` gives what is in the Stream now. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow + code each Workflow Task heartbeat of a long wait adds events to History. + max_records: The most Records to fold. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_MAX_RECORDS`, 1,000. + max_bytes: A cap on the bytes of the Records to fold. A Record counts its encoded + size as the Store gave it, after the payload codec decoded it. A Record that + would pass the cap ends the call before that Record, and the next call gets it. + The first Record of a call is always taken, also when it is larger. ``None`` is + :py:data:`temporalio.common.DEFAULT_STREAM_MAX_BYTES`, 256 KiB. + + Returns: + The value, the number of Records, where the scan stopped, and why. + + Raises: + TypeError: A function is a coroutine function, or the Stream is of the wrong kind. + ValueError: A limit is out of range. + temporalio.exceptions.StreamError: The Store failed the call, or an appender in a + function failed its write. + temporalio.exceptions.StreamConfigError: In client or Activity code, the client + has no Stream Store. + temporalio.exceptions.ApplicationError: A function returned an awaitable. In + Workflow code also when a function raised, when the encoded result is over + :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES` with the ``type`` + ``StreamResultTooLarge``, or when the Worker's client has no Stream config + with the ``type`` ``StreamConfigError``. These are not retryable, so an + uncaught one fails the Workflow. + temporalio.workflow.NondeterminismError: In Workflow code, a replay of a call + whose Store, Stream id or ``from_offset`` differs from the one in History. + + .. warning:: + This API is an experimental prototype and subject to change. + """ + call = _StreamCall.check( + _StreamCallKind.REDUCE, + timeout=timeout, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + until=until, + function=function, + initial=initial, + ) + return await _run_stream_call(stream, call) diff --git a/temporalio/nexus/system/stream_service/__init__.py b/temporalio/nexus/system/stream_service/__init__.py index 35b198843..b68279df6 100644 --- a/temporalio/nexus/system/stream_service/__init__.py +++ b/temporalio/nexus/system/stream_service/__init__.py @@ -27,7 +27,10 @@ ReadRecordsRequest, ReadRecordsResponse, StreamIdConflictPolicy, + StreamRead, + StreamReadStop, StreamRecord, + StreamReduce, StreamRef, StreamStatus, Timestamp, @@ -61,7 +64,10 @@ "ReadRecordsRequest", "ReadRecordsResponse", "StreamIdConflictPolicy", + "StreamRead", + "StreamReadStop", "StreamRecord", + "StreamReduce", "StreamRef", "StreamStatus", "Timestamp", diff --git a/temporalio/nexus/system/stream_service/client/_resources/stream.py b/temporalio/nexus/system/stream_service/client/_resources/stream.py index c3fa617b6..afac0ac8b 100644 --- a/temporalio/nexus/system/stream_service/client/_resources/stream.py +++ b/temporalio/nexus/system/stream_service/client/_resources/stream.py @@ -25,13 +25,16 @@ Producer, ReadRecordsRequest, ReadRecordsResponse, + StreamRead, StreamRecord, + StreamReduce, StreamRef, ) from ...services import StreamService from .._poll import PollIterator BodyT = typing.TypeVar("BodyT") +ValueT = typing.TypeVar("ValueT") class Stream(typing.Generic[BodyT]): @@ -257,5 +260,72 @@ async def describe( ) return await handle.result() + async def take( + self, + timeout: datetime.timedelta, + until: collections.abc.Callable[[StreamRecord[BodyT]], bool] | None = None, + from_offset: bytes | None = None, + idle_timeout: datetime.timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamRead[BodyT]: + """Read every Record up to a stop.""" + return await _support.take_stream( + self, + timeout=timeout, + until=until, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + + async def filter( + self, + predicate: collections.abc.Callable[[StreamRecord[BodyT]], bool], + timeout: datetime.timedelta, + until: collections.abc.Callable[[StreamRecord[BodyT]], bool] | None = None, + from_offset: bytes | None = None, + idle_timeout: datetime.timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamRead[BodyT]: + """Read the Records that the predicate keeps, up to a stop.""" + return await _support.filter_stream( + self, + predicate=predicate, + timeout=timeout, + until=until, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + + async def reduce( + self, + function: collections.abc.Callable[[ValueT, StreamRecord[BodyT]], ValueT], + initial: ValueT, + timeout: datetime.timedelta, + until: collections.abc.Callable[[ValueT, StreamRecord[BodyT]], bool] + | None = None, + from_offset: bytes | None = None, + idle_timeout: datetime.timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamReduce[ValueT]: + """Fold the Records into one value, up to a stop.""" + return await _support.reduce_stream( + self, + function=function, + initial=initial, + timeout=timeout, + until=until, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + __all__ = ["Stream"] diff --git a/temporalio/nexus/system/stream_service/models.py b/temporalio/nexus/system/stream_service/models.py index 4848a2ea5..fd3a0865d 100644 --- a/temporalio/nexus/system/stream_service/models.py +++ b/temporalio/nexus/system/stream_service/models.py @@ -38,7 +38,17 @@ def _stream_status_from_proto(value: int) -> StreamStatus: return StreamStatus.StreamStatusUnspecified +class StreamReadStop(enum.IntEnum): + Match = 0 + Done = 1 + MaxRecords = 2 + MaxBytes = 3 + Idle = 4 + Timeout = 5 + + BodyT = typing.TypeVar("BodyT") +ValueT = typing.TypeVar("ValueT") _BodyT = typing_extensions.TypeVar("_BodyT", default=object) @@ -1336,6 +1346,27 @@ def to_transfer_type( )(DescribeStreamResponse) # pyright: ignore[reportUnusedCallResult] +@dataclasses.dataclass(slots=True, kw_only=True) +class StreamRead(typing.Generic[BodyT]): + """What one read of many pages gave.""" + + records: list[StreamRecord[BodyT]] + last_offset: bytes | None = None + stop: StreamReadStop + done: bool + + +@dataclasses.dataclass(slots=True, kw_only=True) +class StreamReduce(typing.Generic[ValueT]): + """What one fold of many pages gave.""" + + value: ValueT + count: int + last_offset: bytes | None = None + stop: StreamReadStop + done: bool + + @dataclasses.dataclass(slots=True) class ProducerKindWorkflow: value: WorkflowProducer diff --git a/temporalio/nexus/system/stream_service/workflow/_resources/stream.py b/temporalio/nexus/system/stream_service/workflow/_resources/stream.py index df548d252..a237c82da 100644 --- a/temporalio/nexus/system/stream_service/workflow/_resources/stream.py +++ b/temporalio/nexus/system/stream_service/workflow/_resources/stream.py @@ -2,6 +2,8 @@ from __future__ import annotations +import collections.abc +import datetime import types import typing @@ -20,10 +22,14 @@ DescribeStreamRequest, DescribeStreamResponse, Producer, + StreamRead, + StreamRecord, + StreamReduce, StreamRef, ) BodyT = typing.TypeVar("BodyT") +ValueT = typing.TypeVar("ValueT") class Stream(typing.Generic[BodyT]): @@ -96,6 +102,73 @@ async def describe( handle = await _describe_stream(request) return await handle + async def take( + self, + timeout: datetime.timedelta, + until: collections.abc.Callable[[StreamRecord[BodyT]], bool] | None = None, + from_offset: bytes | None = None, + idle_timeout: datetime.timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamRead[BodyT]: + """Read every Record up to a stop.""" + return await _support.take_stream( + self, + timeout=timeout, + until=until, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + + async def filter( + self, + predicate: collections.abc.Callable[[StreamRecord[BodyT]], bool], + timeout: datetime.timedelta, + until: collections.abc.Callable[[StreamRecord[BodyT]], bool] | None = None, + from_offset: bytes | None = None, + idle_timeout: datetime.timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamRead[BodyT]: + """Read the Records that the predicate keeps, up to a stop.""" + return await _support.filter_stream( + self, + predicate=predicate, + timeout=timeout, + until=until, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + + async def reduce( + self, + function: collections.abc.Callable[[ValueT, StreamRecord[BodyT]], ValueT], + initial: ValueT, + timeout: datetime.timedelta, + until: collections.abc.Callable[[ValueT, StreamRecord[BodyT]], bool] + | None = None, + from_offset: bytes | None = None, + idle_timeout: datetime.timedelta | None = None, + max_records: int | None = None, + max_bytes: int | None = None, + ) -> StreamReduce[ValueT]: + """Fold the Records into one value, up to a stop.""" + return await _support.reduce_stream( + self, + function=function, + initial=initial, + timeout=timeout, + until=until, + from_offset=from_offset, + idle_timeout=idle_timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + @typing.overload def get_stream( diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 7ad6bdeaa..6361b3a0f 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -70,7 +70,15 @@ def __init__( worker had, so a replay makes the same commands. A history that records a Stream call needs it. With no Stream config, Workflow code gets an error in place of that call, and the replay fails with a nondeterminism error. A - replay calls no method of its Stores. + replay calls no method of its Stores for a create, an append, a close or a + describe. A ``reduce`` of a Workflow Stream, and a ``take`` or a ``filter`` + whose marker holds its Records, take the result from the history and need no + Stream config. A ``take`` or a ``filter`` that ran under + ``workflow_read_replay_from=StreamReplaySource.STORE`` has only the offsets of + its Records in the history. Its replay reads them from the Store again, so it + needs this config, the same ``data_converter``, and Stores that still have the + Records. Without them the replay fails with + :py:class:`temporalio.exceptions.StreamReplayError`. Note, unlike the worker, for the replayer the workflow_task_executor will default to a new thread pool executor with no max_workers set that diff --git a/temporalio/worker/_stream_read.py b/temporalio/worker/_stream_read.py new file mode 100644 index 000000000..c8ce96594 --- /dev/null +++ b/temporalio/worker/_stream_read.py @@ -0,0 +1,220 @@ +"""The Workflow side of the take, filter and reduce calls of a Stream. + +Workflow code runs a call here, and the Worker reads the Records of a replayed call from the +Store when its marker in History has only their offsets. +""" + +from __future__ import annotations + +import dataclasses +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +import temporalio.api.stream.v1 +import temporalio.api.streamservice.v1 +import temporalio.bridge.proto.workflow_activation +import temporalio.common +import temporalio.converter +import temporalio.exceptions +import temporalio.workflow +from temporalio.client._stream_read import ( + _PENDING_STREAM_CALLS, + _STREAM_CALL_ACTIVITY, + _PendingStreamCall, + _read_kept_records, + _replay_error, + _StreamCallResult, +) +from temporalio.client._stream_store import ( + _STREAM_SERVICE, + StreamConfig, + StreamReplaySource, + _store_page_reader, +) +from temporalio.nexus.system._support import _StreamCall, _StreamCallKind +from temporalio.nexus.system.stream_service.models import ( + ReadRecordsRequest, + StreamRead, + StreamRecord, + StreamReduce, +) +from temporalio.nexus.system.stream_service.workflow._resources.stream import Stream + +from ._stream_store import _no_stream_config_error, _store_errors, _workflow_target + +if TYPE_CHECKING: + from ._workflow_instance import _WorkflowInstanceImpl + +# The Local Activity gets this much time after the call's own timeout to give its result. +_STREAM_CALL_GRACE = timedelta(seconds=10) + + +def _describe( + kind: _StreamCallKind, store: str | None, stream_id: str, from_offset: bytes | None +) -> str: + start = from_offset.hex() if from_offset else "the beginning" + on_store = f" on the Store '{store}'" if store else "" + return f"a {kind.value} of the Stream '{stream_id}'{on_store} after {start}" + + +async def _stream_call_in_workflow( # pyright: ignore[reportUnusedFunction] + instance: _WorkflowInstanceImpl, + config: StreamConfig | None, + stream: Any, + call: _StreamCall, + call_id: str, +) -> StreamRead[Any] | StreamReduce[Any]: + """Run a call of Workflow code in one Local Activity, or take its result from History.""" + if not isinstance(stream, Stream): + raise TypeError( + "Workflow code takes a Workflow Stream, from temporalio.workflow.get_stream " + "or create_stream" + ) + if config is None and not instance._is_replaying: + raise _no_stream_config_error() + request = ReadRecordsRequest(ref=stream.ref, from_offset=call.from_offset) + payload_converter, _ = instance._system_nexus_converters_for( + _STREAM_SERVICE.name, "ReadRecords", request + ) + message = payload_converter._to_transfer_value(request) + assert isinstance(message, temporalio.api.streamservice.v1.ReadRecordsRequest) + record_type: Any = StreamRecord[stream.body_type] # type: ignore[name-defined] + + def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: + return payload_converter._from_transfer_value(record, record_type) + + # A replay of a marker with Records needs no Store and no config. The choice of a + # Store runs in Workflow code, as it does for the other Stream calls. + store = ( + config._store_for(stream.ref, _workflow_target()) + if config is not None + else None + ) + store_name = store.name if store is not None else stream.ref.store + run_id = temporalio.workflow.info().run_id + pending = None + if store is not None: + assert config is not None + replay_from = config.workflow_read_replay_from + if ( + replay_from is StreamReplaySource.STORE + and call.kind is not _StreamCallKind.REDUCE + and call.default_max_bytes + ): + # No body goes into History, so the default cap that keeps a marker small + # has no use here. + call = dataclasses.replace(call, max_bytes=None) + pending = _PendingStreamCall( + call=call, + store=store, + ref=message.ref, + to_record=to_record, + replay_from=replay_from, + ) + _PENDING_STREAM_CALLS.add(run_id, call_id, pending) + try: + with _store_errors(): + raw = await temporalio.workflow.execute_local_activity( + _STREAM_CALL_ACTIVITY, + call_id, + result_type=temporalio.common.RawValue, + start_to_close_timeout=call.timeout + _STREAM_CALL_GRACE, + retry_policy=config.workflow_call_options.retry_policy + if config is not None + else None, + ) + finally: + if pending is not None: + _PENDING_STREAM_CALLS.remove(run_id, call_id, pending) + result = _StreamCallResult.from_payload(raw.payload) + if ( + result is None + or result.kind is not call.kind + or result.stream_id != message.ref.id + or (store_name and result.store != store_name) + or result.from_offset != call.from_offset + or (result.replay is not None and result.replay.call_id != call_id) + ): + recorded = ( + "no Stream call" + if result is None + else _describe( + result.kind, result.store, result.stream_id, result.from_offset + ) + ) + raise temporalio.workflow.NondeterminismError( + f"Stream call {call_id} of the Workflow is " + f"{_describe(call.kind, store_name, message.ref.id, call.from_offset)}, " + f"and History has {recorded}" + ) + if call.kind is _StreamCallKind.REDUCE: + assert result.value is not None + return StreamReduce( + value=instance.workflow_payload_converter().from_payload( + result.value, call.value_type + ), + count=result.count, + last_offset=result.last_offset, + stop=result.stop, + done=result.done, + ) + records = result.records + if result.replay is not None: + # The Worker read these Records before the activation, or the Local Activity of + # the first run kept them. + filled = pending.records_for(result) if pending is not None else None + if filled is None: + raise _replay_error( + result, + "the Worker has no Stream config" + if config is None + else "the Worker read no Records for this call", + ) + records = filled + return StreamRead( + records=[to_record(record) for record in records], + last_offset=result.last_offset, + stop=result.stop, + done=result.done, + ) + + +async def _read_store_records( # pyright: ignore[reportUnusedFunction] + act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, + config: StreamConfig | None, + data_converter: temporalio.converter.DataConverter, +) -> None: + """Read the Records of each marker without Records in this activation. + + It runs before the Workflow sees the activation. The read goes through the same codec + as the Local Activity of the first run. It runs for every such marker, whatever the + Worker's ``workflow_read_replay_from`` is. + + Raises: + temporalio.exceptions.StreamReplayError: The Worker has no Stream config, the + Store failed, or the Records are gone or differ. The Workflow Task fails. + """ + for job in act.jobs: + if not job.HasField("resolve_activity"): + continue + resolve = job.resolve_activity + if not resolve.is_local or not resolve.result.HasField("completed"): + continue + result = _StreamCallResult.from_payload(resolve.result.completed.result) + if result is None or result.replay is None: + continue + if config is None: + raise _replay_error(result, "the Worker has no Stream config") + pending = _PENDING_STREAM_CALLS.get(act.run_id, result.replay.call_id) + # Without a pending call, the Workflow code makes no such call, and its own check + # reports that. A first run already has the Records of its Local Activity. + if pending is None or pending.records_for(result) is not None: + continue + try: + store = config._get_store_by_name(result.store) + except temporalio.exceptions.StreamConfigError as err: + raise _replay_error(result, str(err)) from err + records = await _read_kept_records( + _store_page_reader(store, data_converter), result + ) + pending.fill(result, records) diff --git a/temporalio/worker/_stream_store.py b/temporalio/worker/_stream_store.py index 947cc7aeb..3c6720d85 100644 --- a/temporalio/worker/_stream_store.py +++ b/temporalio/worker/_stream_store.py @@ -3,8 +3,9 @@ from __future__ import annotations import asyncio +import contextlib import dataclasses -from collections.abc import Generator +from collections.abc import Generator, Iterator from typing import TYPE_CHECKING, Any from nexusrpc import OutputT @@ -36,6 +37,33 @@ def _no_stream_config_error() -> temporalio.exceptions.ApplicationError: # pyri ) +@contextlib.contextmanager +def _store_errors() -> Iterator[None]: + """Give Workflow code the Store's error of an Activity, not its wrapper.""" + try: + yield + except temporalio.exceptions.ApplicationError as err: + # Core hands back a failed Local Activity's error without a wrapper. + stream_error = temporalio.exceptions._stream_error_from_failure(err) + if stream_error is not None: + raise stream_error from None + raise + except temporalio.exceptions.ActivityError as err: + if err.cause is not None: + stream_error = temporalio.exceptions._stream_error_from_failure(err.cause) + if stream_error is not None: + raise stream_error from None + if isinstance( + err.cause, + ( + temporalio.exceptions.ApplicationError, + temporalio.exceptions.TimeoutError, + ), + ): + raise err.cause + raise + + def _workflow_target() -> StorageDriverWorkflowInfo: # External storage builds the same value for a Workflow. info = temporalio.workflow.info() @@ -80,33 +108,10 @@ def start_operation( options = self.config.workflow_call_options async def result() -> OutputT: - try: + with _store_errors(): response = await store.call_from_workflow( input.operation_name, request, options ) - except temporalio.exceptions.ApplicationError as err: - # Core hands back a failed Local Activity's error without a wrapper. - stream_error = temporalio.exceptions._stream_error_from_failure(err) - if stream_error is not None: - raise stream_error from None - raise - except temporalio.exceptions.ActivityError as err: - # Workflow code gets the Store's error, not the Local Activity's wrapper. - if err.cause is not None: - stream_error = temporalio.exceptions._stream_error_from_failure( - err.cause - ) - if stream_error is not None: - raise stream_error from None - if isinstance( - err.cause, - ( - temporalio.exceptions.ApplicationError, - temporalio.exceptions.TimeoutError, - ), - ): - raise err.cause - raise return payload_converter._from_transfer_value( stream_operation.check_response(store, response), input.output_type ) diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 4ecce29d1..95ce14304 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -24,6 +24,7 @@ import temporalio.common import temporalio.runtime import temporalio.service +from temporalio.client._stream_read import _stream_call_activity from temporalio.common import ( HeaderCodecBehavior, VersioningBehavior, @@ -504,6 +505,8 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf if stream_config is not None and config.get("workflows"): for store in stream_config.stores: activities.extend(store.worker_activities()) + # The take, filter and reduce calls of Workflow code run this Local Activity. + activities.append(_stream_call_activity) if activities: _warn_if_activity_executor_max_workers_is_inconsistent(config) self._activity_worker = _ActivityWorker( diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index 3e20e6adb..3c26d2d40 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -27,6 +27,7 @@ import temporalio.exceptions import temporalio.workflow from temporalio.bridge.worker import PollShutdownError +from temporalio.client._stream_read import _PENDING_STREAM_CALLS from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo from temporalio.worker.workflow_sandbox._runner import SandboxedWorkflowRunner @@ -40,6 +41,7 @@ WorkflowInboundInterceptor, WorkflowInterceptorClassInput, ) +from ._stream_read import _read_store_records from ._stream_store import _WorkflowStreamStore from ._workflow_instance import ( _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS, @@ -432,6 +434,13 @@ async def _handle_activation( await self._decode_nexus_previous_run_outcome( act.run_id, init_job, previous_run_outcome, download_metrics ) + # A marker of a take or a filter can hold only the offsets of its Records, so + # they come from the Store before the Workflow sees the result. + await _read_store_records( + act, + self._stream_store.config if self._stream_store is not None else None, + self._data_converter, + ) if not workflow: assert init_job workflow = _RunningWorkflow( @@ -678,6 +687,7 @@ async def _handle_cache_eviction( # Remove from map and send completion if act.run_id in self._running_workflows: del self._running_workflows[act.run_id] + _PENDING_STREAM_CALLS.drop_run(act.run_id) try: await self._bridge_worker().complete_workflow_activation( temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion( diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 07f63a69d..3d2311b11 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -58,6 +58,7 @@ import temporalio.converter._serialization_context import temporalio.exceptions import temporalio.nexus.system +import temporalio.nexus.system._support import temporalio.workflow from temporalio.client._stream_store import _STREAM_SERVICE from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo @@ -95,6 +96,7 @@ WorkflowInboundInterceptor, WorkflowOutboundInterceptor, ) +from ._stream_read import _stream_call_in_workflow from ._stream_store import _no_stream_config_error, _WorkflowStreamStore logger = logging.getLogger(__name__) @@ -314,6 +316,7 @@ def __init__(self, det: WorkflowInstanceDetails) -> None: # reads the previous run's result and failure with the same context. self._nexus_serialization_context = det.nexus_serialization_context self._stream_store = det.stream_store + self._stream_call_count = 0 if self._nexus_serialization_context: self._result_payload_converter = self._payload_converter_with_context( self._nexus_serialization_context @@ -1994,6 +1997,20 @@ def workflow_upsert_search_attributes( else [update.value] ) + async def workflow_stream_call( + self, stream: Any, call: temporalio.nexus.system._support._StreamCall + ) -> Any: + # A counter of the instance, and not a random value, numbers the calls. A change of + # the code's random calls then leaves the numbers of a replay alone. + self._stream_call_count += 1 + return await _stream_call_in_workflow( + self, + self._stream_store.config if self._stream_store is not None else None, + stream, + call, + str(self._stream_call_count), + ) + async def workflow_sleep( self, duration: float, @@ -2368,11 +2385,19 @@ def _system_nexus_converters( ) -> tuple[ temporalio.nexus.system._SystemNexusPayloadConverter, temporalio.converter.FailureConverter, + ]: + return self._system_nexus_converters_for( + input.service, input.operation_name, input.input + ) + + def _system_nexus_converters_for( + self, service: str, operation: str, request: Any + ) -> tuple[ + temporalio.nexus.system._SystemNexusPayloadConverter, + temporalio.converter.FailureConverter, ]: serialization_context = temporalio.nexus.system._get_serialization_context( - input.service, - input.operation_name, - input.input, + service, operation, request ) user_payload_converter = self._context_free_payload_converter user_failure_converter = self._context_free_failure_converter diff --git a/temporalio/workflow/_context.py b/temporalio/workflow/_context.py index a33759c37..4f7b8b82f 100644 --- a/temporalio/workflow/_context.py +++ b/temporalio/workflow/_context.py @@ -21,6 +21,8 @@ from ._exceptions import _NotInWorkflowEventLoopError if TYPE_CHECKING: + from temporalio.nexus.system._support import _StreamCall + from ._activities import ActivityCancellationType, ActivityHandle from ._event_groups import EventGroup from ._exceptions import ContinueAsNewVersioningBehavior, VersioningIntent @@ -399,6 +401,9 @@ def workflow_payload_converter(self) -> temporalio.converter.PayloadConverter: . @abstractmethod def workflow_random(self) -> Random: ... + @abstractmethod + async def workflow_stream_call(self, stream: Any, call: _StreamCall) -> Any: ... + @abstractmethod def workflow_set_query_handler( self, name: str | None, handler: Callable | None diff --git a/tests/contrib/streams/test_pipeline.py b/tests/contrib/streams/test_pipeline.py new file mode 100644 index 000000000..00604e1d0 --- /dev/null +++ b/tests/contrib/streams/test_pipeline.py @@ -0,0 +1,447 @@ +"""Tests of the pipeline of `temporalio.contrib.streams`.""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +import pytest +from google.protobuf.message import Message + +import temporalio.api.streamservice.v1 +from temporalio import workflow +from temporalio.client import ( + Client, + StreamConfig, + StreamReplaySource, + StreamStoreCallOptions, +) +from temporalio.common import RetryPolicy +from temporalio.contrib.streams import Pipeline, pipeline +from temporalio.nexus.system.stream_service import ( + StreamRead, + StreamReadStop, + StreamRecord, + StreamReduce, + StreamRef, +) +from temporalio.testing import MemoryStreamStore +from temporalio.worker import Replayer +from tests.nexus.test_stream_read_workflow import ( + _BODY_MARK, + Token, + _append, + _close, + _marker_payload, + _markers, + _offsets, + _with_store, +) + +# The Workflows and the test share the values of `StreamReadState`. +with workflow.unsafe.imports_passed_through(): + from tests.helpers import new_worker + from tests.helpers.streams import RecordingStreamStore, StreamReadState + +_QUICK = timedelta(seconds=10) +_NOW = timedelta(0) + +# Two lines of text, with Records of another kind between them. +_TOKENS = [ + Token("text", "a"), + Token("skip", "x"), + Token("text", "b\n"), + Token("text", "c"), + Token("skip", "y"), + Token("text", "d\n"), + Token("text", "e"), +] + + +def _is_text(record: StreamRecord[Token]) -> bool: + return record.body.kind == "text" + + +def _text(record: StreamRecord[Token]) -> str: + return record.body.text + + +def _ends_line(text: str) -> bool: + return text.endswith("\n") + + +def _join(joined: str, text: str) -> str: + return joined + text + + +def _record_ends_line(record: StreamRecord[Token]) -> bool: + return record.body.text.endswith("\n") + + +@dataclass +class ChainResult: + name: str + values: list[str] + count: int + stop: str + last_offset: str | None + done: bool + + +def _of_read(name: str, read: StreamRead[Token]) -> ChainResult: + return ChainResult( + name, + [record.body.text for record in read.records], + len(read.records), + read.stop.name, + read.last_offset.hex() if read.last_offset else None, + read.done, + ) + + +def _of_reduce(name: str, result: StreamReduce[Any], values: list[str]) -> ChainResult: + return ChainResult( + name, + values, + result.count, + result.stop.name, + result.last_offset.hex() if result.last_offset else None, + result.done, + ) + + +async def _run_chains(stream: Any) -> list[ChainResult]: + """Run each step and each end once over the Stream of `_TOKENS`.""" + start = pipeline(stream) + texts = start.filter(_is_text) + results = [ + _of_read("filter records", await texts.records(_QUICK, idle_timeout=_NOW)), + _of_read( + "take_until records", + await texts.take_until(_record_ends_line).records( + _QUICK, idle_timeout=_NOW + ), + ), + _of_read( + "limit records", await texts.limit(3).records(_QUICK, idle_timeout=_NOW) + ), + _of_read("all records", await start.records(_QUICK, idle_timeout=_NOW)), + ] + line = ( + await texts.map(_text) + .take_until(_ends_line) + .reduce(_join, "", _QUICK, idle_timeout=_NOW) + ) + results.append(_of_reduce("reduce", line, [line.value])) + collected = await texts.map(_text).limit(4).collect(_QUICK, idle_timeout=_NOW) + results.append(_of_reduce("collect", collected, collected.value)) + records = await texts.limit(2).collect(_QUICK, idle_timeout=_NOW) + results.append( + _of_reduce( + "collect records", records, [record.body.text for record in records.value] + ) + ) + count = await texts.count(_QUICK, idle_timeout=_NOW) + results.append(_of_reduce("count", count, [str(count.value)])) + early = await start.limit(3).filter(_is_text).count(_QUICK, idle_timeout=_NOW) + results.append(_of_reduce("limit then filter", early, [str(early.value)])) + return results + + +def _expected(offsets: list[str]) -> list[ChainResult]: + return [ + ChainResult("filter records", ["a", "b\n", "c", "d\n", "e"], 5, "Done", offsets[6], True), + ChainResult("take_until records", ["a", "b\n"], 2, "Match", offsets[2], False), + ChainResult("limit records", ["a", "b\n", "c"], 3, "MaxRecords", offsets[3], False), + ChainResult("all records", [token.text for token in _TOKENS], 7, "Done", offsets[6], True), + ChainResult("reduce", ["ab\n"], 2, "Match", offsets[2], False), + ChainResult("collect", ["a", "b\n", "c", "d\n"], 4, "Match", offsets[5], False), + ChainResult("collect records", ["a", "b\n"], 2, "MaxRecords", offsets[2], False), + ChainResult("count", ["5"], 5, "Done", offsets[6], True), + # The limit counts every Record, and the filter after it drops the skipped one. + ChainResult("limit then filter", ["2"], 2, "Match", offsets[2], False), + ] # fmt: skip + + +def _workflow_stream(stream_id: str) -> Any: + return workflow.get_stream(StreamRef(id=stream_id), body_type=Token) + + +@workflow.defn +class ChainsWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[ChainResult]: + return await _run_chains(_workflow_stream(stream_id)) + + +@workflow.defn +class OneChainWorkflow: + @workflow.run + async def run(self, stream_id: str) -> str: + result = ( + await pipeline(_workflow_stream(stream_id)) + .filter(_is_text) + .map(_text) + .take_until(_ends_line) + .reduce(_join, "", _QUICK, idle_timeout=_NOW) + ) + return result.value + + +async def _hits_one_by_one(stream: Any) -> list[str]: + found: list[str] = [] + offset: bytes | None = None + while True: + read = ( + await pipeline(stream) + .filter(_is_text) + .limit(1) + .records(_QUICK, from_offset=offset, idle_timeout=_NOW) + ) + found.extend(record.body.text for record in read.records) + offset = read.last_offset + if read.stop in (StreamReadStop.Done, StreamReadStop.Idle): + return found + + +async def _lines(stream: Any) -> list[str]: + lines: list[str] = [] + offset: bytes | None = None + while True: + line = ( + await pipeline(stream) + .filter(_is_text) + .map(_text) + .take_until(_ends_line) + .reduce(_join, "", _QUICK, from_offset=offset, idle_timeout=_NOW) + ) + if line.count: + lines.append(line.value) + offset = line.last_offset + if line.stop in (StreamReadStop.Done, StreamReadStop.Idle): + return lines + + +@workflow.defn +class LoopsWorkflow: + @workflow.run + async def run(self, stream_id: str) -> tuple[list[str], list[str]]: + stream = _workflow_stream(stream_id) + return await _hits_one_by_one(stream), await _lines(stream) + + +@workflow.defn +class LimitAndUntilWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[ChainResult]: + texts = pipeline(_workflow_stream(stream_id)).filter(_is_text).map(_text) + limited = await texts.limit(4).collect(_QUICK, idle_timeout=_NOW) + line = await texts.take_until(_ends_line).reduce( + _join, "", _QUICK, idle_timeout=_NOW + ) + return [ + _of_reduce("limit", limited, limited.value), + _of_reduce("take_until", line, [line.value]), + ] + + +@workflow.defn +class RecordsEndsWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[int]: + stream = _workflow_stream(stream_id) + line = ( + await pipeline(stream) + .filter(_is_text) + .take_until(_record_ends_line) + .records(_QUICK, idle_timeout=_NOW) + ) + first = await pipeline(stream).limit(3).records(_QUICK, idle_timeout=_NOW) + StreamReadState.results[workflow.info().run_id] = tuple( + tuple(record.body.text for record in read.records) for read in (line, first) + ) + return [len(line.records), len(first.records)] + + +class SmallPagesStore(RecordingStreamStore): + """Gives two Records a page, and fails one read of the scan when it is armed.""" + + def __init__(self, *, fail_read: int | None) -> None: + super().__init__(MemoryStreamStore()) + self.fail_read = fail_read + self.armed = False + self.reads = 0 + + async def before_call(self, operation: str, request: Message) -> None: + if not isinstance(request, temporalio.api.streamservice.v1.ReadRecordsRequest): + return + request.max_records = min(request.max_records or 2, 2) + if not self.armed: + return + self.reads += 1 + if self.reads == self.fail_read: + raise RuntimeError("The Store failed one read") + + +def _retrying_client(client: Client, store: SmallPagesStore) -> Client: + config = client.config() + config["stream_config"] = StreamConfig( + stores=[store], + workflow_call_options=StreamStoreCallOptions( + retry_policy=RetryPolicy( + maximum_attempts=3, initial_interval=timedelta(milliseconds=10) + ) + ), + ) + return Client(**config) + + +async def _stream_of_tokens( + client: Client, stream_client: Client | None = None +) -> tuple[Client, str, list[str]]: + stream_client = stream_client or _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=stream_id, body_type=Token) + await _append(stream_client, stream_id, *_TOKENS) + await _close(stream_client, stream_id) + return stream_client, stream_id, await _offsets(stream_client, stream_id) + + +async def _execute(stream_client: Client, workflow_class: type, *args: Any) -> Any: + async with new_worker(stream_client, workflow_class) as worker: + handle = await stream_client.start_workflow( + getattr(workflow_class, "run"), + args=list(args), + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + return handle, await handle.result() + + +async def test_each_step_and_end_in_workflow_code(client: Client): + stream_client, stream_id, offsets = await _stream_of_tokens(client) + _, results = await _execute(stream_client, ChainsWorkflow, stream_id) + assert results == _expected(offsets) + + +async def test_each_step_and_end_from_a_client(client: Client): + stream_client, stream_id, offsets = await _stream_of_tokens(client) + stream = stream_client.get_stream(StreamRef(id=stream_id), body_type=Token) + assert await _run_chains(stream) == _expected(offsets) + + +async def test_a_chain_is_one_marker(client: Client): + stream_client, stream_id, _ = await _stream_of_tokens(client) + handle, line = await _execute(stream_client, OneChainWorkflow, stream_id) + assert line == "ab\n" + assert len(_markers(await handle.fetch_history())) == 1 + + +async def test_records_refuses_a_chain_that_does_not_fit(client: Client): + stream = client.get_stream(StreamRef(id="unused"), body_type=Token) + texts = pipeline(stream).filter(_is_text) + cases: list[tuple[Pipeline[Any], str]] = [ + (texts.map(_text), "map"), + (texts.take_until(_is_text).take_until(_is_text), "take_until, the second one"), + (texts.limit(1).filter(_is_text), "limit, which isn't the last step"), + ] + for chain, step in cases: + with pytest.raises(TypeError) as err: + chain.records(_QUICK) + assert str(err.value) == ( + f"records() can't run this chain, since its step {step} doesn't fit. It " + "runs only filter steps, one take_until and a limit as the last step. End " + "the chain with collect() or reduce()." + ) + + +async def test_loops_read_every_record_once(client: Client): + stream_client, stream_id, _ = await _stream_of_tokens(client) + _, (hits, lines) = await _execute(stream_client, LoopsWorkflow, stream_id) + assert hits == ["a", "b\n", "c", "d\n", "e"] + assert lines == ["ab\n", "cd\n", "e"] + stream = stream_client.get_stream(StreamRef(id=stream_id), body_type=Token) + assert await _hits_one_by_one(stream) == hits + assert await _lines(stream) == lines + + +async def test_failed_local_activity_attempt_gives_the_clean_result(client: Client): + clean_store = SmallPagesStore(fail_read=None) + clean_client, clean_id, clean_offsets = await _stream_of_tokens( + client, _retrying_client(client, clean_store) + ) + _, clean = await _execute(clean_client, LimitAndUntilWorkflow, clean_id) + + failing_store = SmallPagesStore(fail_read=2) + failing_client, failing_id, failing_offsets = await _stream_of_tokens( + client, _retrying_client(client, failing_store) + ) + failing_store.armed = True + _, after_failure = await _execute(failing_client, LimitAndUntilWorkflow, failing_id) + + # The second read failed after the first page went through the steps. + assert failing_store.reads > 2 + # The two Streams are different, so the results name a Record by its place. + for results, offsets in ((clean, clean_offsets), (after_failure, failing_offsets)): + for result in results: + result.last_offset = str(offsets.index(result.last_offset or "")) + assert after_failure == clean + assert clean[0].values == ["a", "b\n", "c", "d\n"] + assert clean[1].values == ["ab\n"] + + +async def test_replayer_replays_each_end(client: Client): + stream_client, stream_id, offsets = await _stream_of_tokens(client) + handle, results = await _execute(stream_client, ChainsWorkflow, stream_id) + assert results == _expected(offsets) + await Replayer(workflows=[ChainsWorkflow]).replay_workflow( + await handle.fetch_history() + ) + + +async def test_typed_ends(client: Client): + stream_client, stream_id, _ = await _stream_of_tokens(client) + stream = stream_client.get_stream(StreamRef(id=stream_id), body_type=Token) + texts: Pipeline[str] = pipeline(stream).filter(_is_text).map(_text) + read: StreamRead[Token] = await pipeline(stream).limit(1).records(_QUICK) + joined: StreamReduce[str] = await texts.reduce(_join, "", _QUICK, idle_timeout=_NOW) + items: StreamReduce[list[str]] = await texts.collect(_QUICK, idle_timeout=_NOW) + count: StreamReduce[int] = await texts.count(_QUICK, idle_timeout=_NOW) + assert read.records[0].body == _TOKENS[0] + assert joined.value == "ab\ncd\ne" + assert items.value == ["a", "b\n", "c", "d\n", "e"] + assert count.value == 5 + + +async def test_records_end_follows_the_replay_setting(client: Client): + store = MemoryStreamStore() + config = client.config() + config["stream_config"] = StreamConfig( + stores=[store], workflow_read_replay_from=StreamReplaySource.STORE + ) + stream_client = Client(**config) + tokens = [Token(token.kind, f"{_BODY_MARK}-{token.text}") for token in _TOKENS] + stream_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=stream_id, body_type=Token) + await _append(stream_client, stream_id, *tokens) + handle, counts = await _execute(stream_client, RecordsEndsWorkflow, stream_id) + # A filter and a take, each with no body in its marker. + assert counts == [2, 3] + history = await handle.fetch_history() + assert _markers(history) == ["core_local_activity"] * 2 + for index in (0, 1): + marker = _marker_payload(history, index) + assert marker.metadata["temporal-stream-replay-from"] == b"store" + assert _BODY_MARK not in history.to_json() + run_id = handle.result_run_id or "" + first_run = StreamReadState.results.pop(run_id) + assert first_run == ( + tuple(token.text for token in tokens[:3] if token.kind == "text"), + tuple(token.text for token in tokens[:3]), + ) + await Replayer( + workflows=[RecordsEndsWorkflow], stream_config=StreamConfig(stores=[store]) + ).replay_workflow(history) + assert StreamReadState.results.pop(run_id) == first_run diff --git a/tests/helpers/streams.py b/tests/helpers/streams.py index beb4eaaa4..5d5b8bae3 100644 --- a/tests/helpers/streams.py +++ b/tests/helpers/streams.py @@ -184,3 +184,17 @@ async def describe_stream( else StreamStatus.STREAM_STATUS_OPEN, last_offset=stream.records[-1].offset if stream.records else b"", ) + + +class StreamReadState: + """State that a test of a Stream call shares with its Workflows. + + A Workflow file imports this module with ``imports_passed_through``, so the sandbox + and the test see one copy of these values. + """ + + until_calls = 0 + function_calls = 0 + fail_next_task = False + variant = "" + results: dict[str, tuple[object, ...]] = {} diff --git a/tests/nexus/test_stream_read_client.py b/tests/nexus/test_stream_read_client.py new file mode 100644 index 000000000..ae7547e12 --- /dev/null +++ b/tests/nexus/test_stream_read_client.py @@ -0,0 +1,413 @@ +"""Tests of the take, filter and reduce calls of client and Activity code.""" + +from __future__ import annotations + +import uuid +from datetime import timedelta +from typing import Any + +import pytest + +from temporalio import activity, workflow +from temporalio.client import ( + Client, + Interceptor, + NexusOperationHandle, + OutboundInterceptor, + StartNexusOperationInput, +) +from temporalio.contrib.streams import BufferedAppender +from temporalio.exceptions import ApplicationError +from temporalio.nexus.system.stream_service import ( + ReadRecordsRequest, + ReadRecordsResponse, + StreamReadStop, + StreamRecord, + StreamRef, +) +from temporalio.nexus.system.stream_service.client import Stream as ClientStream +from temporalio.nexus.system.stream_service.workflow._resources.stream import ( + Stream as WorkflowStream, +) +from temporalio.testing import MemoryStreamStore +from tests.helpers import new_worker +from tests.helpers.streams import RecordingStreamStore +from tests.nexus.test_stream_read_workflow import ( + _STOPS, + CallArgs, + CallResult, + CallWorkflow, + OverLimitWorkflow, + Token, + _append, + _call, + _check_stop, + _close, + _new_stream, + _offsets, + _record_sizes, + _run_workflow, + _tokens, + _with_store, +) + + +def _client_stream(stream_client: Client, stream_id: str) -> Any: + return stream_client.get_stream(StreamRef(id=stream_id), body_type=Token) + + +@activity.defn +async def call_in_activity(args: CallArgs) -> CallResult: + stream = activity.client().get_stream(StreamRef(id=args.stream_id), body_type=Token) + return await _call(args, stream) + + +@activity.defn +async def filter_loop_in_activity(stream_id: str) -> list[str]: + stream = activity.client().get_stream(StreamRef(id=stream_id), body_type=Token) + return await _filter_loop(stream) + + +async def _filter_loop(stream: Any) -> list[str]: + found: list[str] = [] + offset: bytes | None = None + while True: + read = await stream.filter( + lambda record: record.body.kind == "hit", + timedelta(seconds=10), + from_offset=offset, + idle_timeout=timedelta(0), + max_records=1, + ) + found.extend(record.body.text for record in read.records) + offset = read.last_offset + if read.stop in (StreamReadStop.Done, StreamReadStop.Idle): + return found + + +@workflow.defn +class ActivityCallWorkflow: + @workflow.run + async def run(self, args: CallArgs) -> CallResult: + return await workflow.execute_activity( + call_in_activity, args, start_to_close_timeout=timedelta(seconds=30) + ) + + +@workflow.defn +class ActivityLoopWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[str]: + return await workflow.execute_activity( + filter_loop_in_activity, + stream_id, + start_to_close_timeout=timedelta(seconds=30), + ) + + +async def _loop_stream(stream_client: Client) -> tuple[str, list[str]]: + stream_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=stream_id, body_type=Token) + tokens = [ + Token("hit" if index % 4 == 1 else "text", f"record-{index}") + for index in range(20) + ] + await _append(stream_client, stream_id, *tokens) + await _close(stream_client, stream_id) + return stream_id, [token.text for token in tokens if token.kind == "hit"] + + +@pytest.mark.parametrize("call", ["take", "filter", "reduce"]) +@pytest.mark.parametrize("stop", _STOPS) +async def test_client_call_stop_and_last_offset(client: Client, call: str, stop: str): + stream_client = _with_store(client, MemoryStreamStore()) + + async def run(stream_id: str, options: dict[str, Any]) -> CallResult: + args = CallArgs(stream_id, call=call, **options) + result = await _call(args, _client_stream(stream_client, stream_id)) + assert result.typed + return result + + await _check_stop(stream_client, call, stop, run) + + +async def test_client_filter_loop_reads_to_the_end(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream_id, hits = await _loop_stream(stream_client) + assert await _filter_loop(_client_stream(stream_client, stream_id)) == hits + + +@pytest.mark.parametrize("stop", _STOPS) +async def test_activity_call_stop_and_last_offset(client: Client, stop: str): + stream_client = _with_store(client, MemoryStreamStore()) + # Each stop once, and the three calls in turn. + call = ["take", "filter", "reduce"][_STOPS.index(stop) % 3] + async with new_worker( + stream_client, ActivityCallWorkflow, activities=[call_in_activity] + ) as worker: + + async def run(stream_id: str, options: dict[str, Any]) -> CallResult: + return await stream_client.execute_workflow( + ActivityCallWorkflow.run, + CallArgs(stream_id, call=call, **options), + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + + await _check_stop(stream_client, call, stop, run) + + +async def test_activity_filter_loop_reads_to_the_end(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream_id, hits = await _loop_stream(stream_client) + async with new_worker( + stream_client, ActivityLoopWorkflow, activities=[filter_loop_in_activity] + ) as worker: + found = await stream_client.execute_workflow( + ActivityLoopWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert found == hits + + +async def test_client_functions_must_be_plain(client: Client): + stream_client, _, stream_id = await _new_stream(client) + await _append(stream_client, stream_id, *_tokens(1)) + stream = _client_stream(stream_client, stream_id) + + async def coroutine_until(record: StreamRecord[Token]) -> bool: + del record + return True + + # The type checker refuses both functions, and the call refuses them at run time. + awaiting: Any = coroutine_until + + with pytest.raises(TypeError, match="until must be a plain function"): + await stream.take(timedelta(seconds=5), awaiting) + with pytest.raises(ApplicationError) as err: + await stream.filter(lambda record: awaiting(record), timedelta(seconds=5)) + assert err.value.non_retryable + assert err.value.message.startswith("The predicate of a Stream call returned an") + + def raises(value: int, record: StreamRecord[Token]) -> int: + raise ValueError(f"bad fold after {value} at {record.offset.hex()}") + + # In a client, the error of the caller's function reaches it as it is. + with pytest.raises(ValueError, match="bad fold"): + await stream.reduce(raises, 0, timedelta(seconds=5)) + + +async def test_client_result_has_no_limit(client: Client): + stream_client, _, stream_id = await _new_stream(client) + for _ in range(6): + await _append(stream_client, stream_id, *[Token("text", "x" * 20_000)] * 10) + await _close(stream_client, stream_id) + stream = _client_stream(stream_client, stream_id) + # Over 1 MiB of Records, which a client gets since its result doesn't go into History. + read = await stream.take( + timedelta(seconds=5), idle_timeout=timedelta(0), max_bytes=1 << 22 + ) + assert len(read.records) == 60 + assert read.stop is StreamReadStop.Done + _, message = await _run_workflow( + stream_client, OverLimitWorkflow, stream_id, "take" + ) + assert message.startswith("StreamResultTooLarge True: ") + + +def test_stream_methods(): + for stream_class in (ClientStream, WorkflowStream): + for name in ("take", "filter", "reduce"): + assert callable(getattr(stream_class, name, None)), (stream_class, name) + assert hasattr(ClientStream, "read") + assert not hasattr(WorkflowStream, "read") + + +class _PageCounter(Interceptor): + """Counts the ReadRecords operations that the client starts.""" + + def __init__(self) -> None: + super().__init__() + self.reads = 0 + + def intercept_client(self, next: OutboundInterceptor) -> OutboundInterceptor: + return _PageCounterOutbound(self, next) + + +class _PageCounterOutbound(OutboundInterceptor): + def __init__(self, counter: _PageCounter, next: OutboundInterceptor) -> None: + super().__init__(next) + self._counter = counter + + async def start_nexus_operation( + self, input: StartNexusOperationInput + ) -> NexusOperationHandle[Any]: + if input.operation == "ReadRecords": + self._counter.reads += 1 + return await super().start_nexus_operation(input) + + +def _with_interceptor(stream_client: Client, interceptor: Interceptor) -> Client: + config = stream_client.config() + config["interceptors"] = [*config["interceptors"], interceptor] + return Client(**config) + + +async def test_client_interceptor_sees_each_page(client: Client): + store = RecordingStreamStore(MemoryStreamStore()) + stream_client, _, stream_id = await _new_stream(client, store) + for start in range(0, 2500, 1000): + await _append(stream_client, stream_id, *_tokens(1000, prefix=f"p{start}-")) + counter = _PageCounter() + read = await _client_stream( + _with_interceptor(stream_client, counter), stream_id + ).take( + timedelta(seconds=10), + idle_timeout=timedelta(0), + max_records=2500, + max_bytes=1 << 30, + ) + assert len(read.records) == 2500 + assert read.stop is StreamReadStop.MaxRecords + # A page holds at most 1,000 Records, so the take reads three pages. + assert counter.reads == store.operations.count("ReadRecords") == 3 + + +@pytest.mark.parametrize("call", ["take", "filter", "reduce"]) +async def test_max_bytes_stops_at_the_same_offset_everywhere(client: Client, call: str): + store = MemoryStreamStore() + stream_client, _, stream_id = await _new_stream(client, store) + # Every other Record is one that the filter drops, so it counts only the kept ones. + tokens = [ + Token("skip" if index % 2 else "text", f"record-{index}") for index in range(12) + ] + await _append(stream_client, stream_id, *tokens) + sizes = await _record_sizes(store, stream_id) + args = CallArgs( + stream_id, + call=call, + until_kind=None, + idle_timeout_seconds=0, + max_bytes=sum(sizes[:3]) - 1, + ) + in_client = await _call(args, _client_stream(stream_client, stream_id)) + async with new_worker( + stream_client, + CallWorkflow, + ActivityCallWorkflow, + activities=[call_in_activity], + ) as worker: + in_activity = await stream_client.execute_workflow( + ActivityCallWorkflow.run, + args, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + in_workflow = await stream_client.execute_workflow( + CallWorkflow.run, + args, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert in_client.stop == "MaxBytes" + assert in_client == in_activity == in_workflow + assert in_client.count == 2 + # The dropped Records count no bytes, so the filter scans past two of them. + offsets = await _offsets(stream_client, stream_id) + assert in_client.last_offset == offsets[3 if call == "filter" else 1] + + +class _ReadAnswerer(Interceptor): + """Answers each ReadRecords operation itself, with no call of the Store.""" + + def __init__(self, records: list[StreamRecord[Token]]) -> None: + super().__init__() + self.records = records + + def intercept_client(self, next: OutboundInterceptor) -> OutboundInterceptor: + return _ReadAnswererOutbound(self, next) + + +class _AnsweredHandle: + def __init__(self, response: ReadRecordsResponse[Token]) -> None: + self._response = response + + async def result(self) -> ReadRecordsResponse[Token]: + return self._response + + +class _ReadAnswererOutbound(OutboundInterceptor): + def __init__(self, answerer: _ReadAnswerer, next: OutboundInterceptor) -> None: + super().__init__(next) + self._answerer = answerer + + async def start_nexus_operation( + self, input: StartNexusOperationInput + ) -> NexusOperationHandle[Any]: + if input.operation != "ReadRecords": + return await super().start_nexus_operation(input) + request = input.arg + assert isinstance(request, ReadRecordsRequest) + records = [ + record + for record in self._answerer.records + if request.from_offset is None or record.offset > request.from_offset + ] + response = ReadRecordsResponse( + records=records, + last_offset=records[-1].offset if records else request.from_offset, + done=True, + ) + # The generated `read` only awaits `result()` of the handle. + handle: Any = _AnsweredHandle(response) + return handle + + +async def test_interceptor_that_answers_a_read_counts_no_bytes(client: Client): + stream_client, _, stream_id = await _new_stream(client) + records = [ + StreamRecord(body=Token("text", "x" * 1000), offset=bytes([index])) + for index in range(1, 6) + ] + answered = _with_interceptor(stream_client, _ReadAnswerer(records)) + # No Store call measured these Records, so a cap of one byte doesn't stop the take. + read = await _client_stream(answered, stream_id).take( + timedelta(seconds=5), idle_timeout=timedelta(0), max_bytes=1 + ) + assert [record.offset for record in read.records] == [ + record.offset for record in records + ] + assert read.stop is StreamReadStop.Done + assert read.last_offset == bytes([5]) + + +async def test_client_appender_in_a_function_writes_before_the_call_returns( + client: Client, +): + store = RecordingStreamStore(MemoryStreamStore()) + stream_client, _, source_id = await _new_stream(client, store) + target = await stream_client.create_stream( + id=f"stream-{uuid.uuid4()}", body_type=str + ) + await _append(stream_client, source_id, *_tokens(5)) + # A long delay, so only the end of the call writes the buffered Records. + appender = BufferedAppender(target, max_delay=timedelta(hours=1)) + + def append_upper(value: int, record: StreamRecord[Token]) -> int: + appender.append(record.body.text.upper()) + return value + 1 + + result = await _client_stream(stream_client, source_id).reduce( + append_upper, + 0, + timedelta(seconds=5), + idle_timeout=timedelta(0), + ) + assert result.count == 5 + assert store.operations.count("AppendRecords") == 2 + page = await target.read() + assert [record.body for record in page.records] == [ + token.text.upper() for token in _tokens(5) + ] diff --git a/tests/nexus/test_stream_read_replay_from.py b/tests/nexus/test_stream_read_replay_from.py new file mode 100644 index 000000000..a847367a9 --- /dev/null +++ b/tests/nexus/test_stream_read_replay_from.py @@ -0,0 +1,442 @@ +"""Tests of `StreamConfig.workflow_read_replay_from` for the take and filter of Workflow code.""" + +from __future__ import annotations + +import dataclasses +import uuid +from typing import Any + +import pytest +from google.protobuf.message import Message + +from temporalio import workflow +from temporalio.api.common.v1 import Payload, Payloads +from temporalio.api.stream.v1 import StreamRef as StreamRefMessage +from temporalio.api.streamservice.v1 import ReadRecordsRequest +from temporalio.client import ( + Client, + StreamConfig, + StreamReplaySource, + WorkflowHistory, +) +from temporalio.converter import DataConverter +from temporalio.testing import MemoryStreamStore +from temporalio.worker import Replayer +from tests.nexus.test_stream_read_workflow import ( + _BODY_MARK, + CallArgs, + CallResult, + CallWorkflow, + KeepHitsWorkflow, + Token, + XorCodec, + _append, + _call, + _close, + _done, + _keep, + _marker_payload, + _markers, + _texts, + _tokens, +) + +with workflow.unsafe.imports_passed_through(): + from tests.helpers import assert_eventually, new_worker + from tests.helpers.streams import ( + DictStreamStore, + RecordingStreamStore, + StreamReadState, + ) + +_HISTORY = StreamReplaySource.HISTORY +_STORE = StreamReplaySource.STORE +_READ_CALLS = pytest.mark.parametrize("call", ["take", "filter"]) + +# The Records of a call. The filter drops the two of the kind "skip". +_CALL_TOKENS = [ + Token("text", f"{_BODY_MARK}-t0"), + Token("skip", f"{_BODY_MARK}-s1"), + Token("text", f"{_BODY_MARK}-t2"), + Token("skip", f"{_BODY_MARK}-s3"), + Token("text", f"{_BODY_MARK}-t4"), + _done(), +] + + +class RequestsStore(RecordingStreamStore): + """A Store that keeps each read request, to show what a replay reads.""" + + def __init__(self, inner: Any) -> None: + super().__init__(inner) + self.reads: list[ReadRecordsRequest] = [] + + async def before_call(self, operation: str, request: Message) -> None: + if isinstance(request, ReadRecordsRequest): + self.reads.append(request) + + +@workflow.defn +class QuietCallWorkflow: + """Runs the call of ``args`` and gives back no text, so no body is in its result.""" + + @workflow.run + async def run(self, args: CallArgs) -> int: + return _keep(await _call(args)).count + + +def _config(store: Any, replay_from: StreamReplaySource) -> StreamConfig: + return StreamConfig(stores=[store], workflow_read_replay_from=replay_from) + + +def _client( + client: Client, + store: Any, + replay_from: StreamReplaySource, + data_converter: DataConverter | None = None, +) -> Client: + config = client.config() + config["stream_config"] = _config(store, replay_from) + if data_converter is not None: + config["data_converter"] = data_converter + return Client(**config) + + +async def _stream_of(client: Client, *tokens: Token, closed: bool = False) -> str: + stream_id = f"stream-{uuid.uuid4()}" + await client.create_stream(id=stream_id, body_type=Token) + for start in range(0, len(tokens), 10): + await _append(client, stream_id, *tokens[start : start + 10]) + if closed: + await _close(client, stream_id) + return stream_id + + +async def _run( + client: Client, workflow_class: type, *args: Any +) -> tuple[WorkflowHistory, Any]: + async with new_worker(client, workflow_class) as worker: + handle = await client.start_workflow( + getattr(workflow_class, "run"), + args=list(args), + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + result = await handle.result() + history = await handle.fetch_history() + return history, result + + +async def _plain_marker( + history: WorkflowHistory, + data_converter: DataConverter | None = None, + index: int = 0, +) -> Payload: + payload = _marker_payload(history, index) + if data_converter is not None and data_converter.payload_codec is not None: + [payload] = await data_converter.payload_codec.decode([payload]) + return payload + + +def _replays_from_store(payload: Payload) -> bool: + return payload.metadata.get("temporal-stream-replay-from") == b"store" + + +def _no_body_in(history: WorkflowHistory) -> bool: + return _BODY_MARK not in history.to_json() and not any( + _BODY_MARK.encode() in event.SerializeToString() for event in history.events + ) + + +async def _store_records(store: Any, stream_id: str) -> list[Any]: + response = await store.inner.read_records( + ReadRecordsRequest(ref=StreamRefMessage(id=stream_id, store=store.name)) + ) + return list(response.records) + + +def _run_result(history: WorkflowHistory) -> tuple[object, ...]: + return StreamReadState.results.pop( + history.events[ + 0 + ].workflow_execution_started_event_attributes.original_execution_run_id + ) + + +async def _replay_failure( + history: WorkflowHistory, workflow_class: type, **options: Any +) -> str: + result = await Replayer(workflows=[workflow_class], **options).replay_workflow( + history, raise_on_replay_failure=False + ) + assert result.replay_failure is not None + return str(result.replay_failure) + + +@_READ_CALLS +@pytest.mark.parametrize("codec", [False, True], ids=["no_codec", "codec"]) +async def test_store_keeps_every_body_out_of_history_and_replays( + client: Client, call: str, codec: bool +): + data_converter = DataConverter(payload_codec=XorCodec()) if codec else None + store = RequestsStore(MemoryStreamStore()) + history_client = _client(client, store, _HISTORY, data_converter) + store_client = _client(client, store, _STORE, data_converter) + stream_id = await _stream_of(store_client, *_CALL_TOKENS) + args = CallArgs(stream_id, call=call) + history_run, _ = await _run(history_client, QuietCallWorkflow, args) + by_history = _run_result(history_run) + history, count = await _run(store_client, QuietCallWorkflow, args) + first_run = _run_result(history) + # A first run gives Workflow code the same result under both settings. + assert first_run == by_history + kept = [0, 2, 4, 5] if call == "filter" else [0, 1, 2, 3, 4, 5] + assert first_run[0] == tuple(_CALL_TOKENS[index].text for index in kept) + assert count == len(kept) + # A codec hides the bodies under HISTORY too, and without one they show there. + assert _no_body_in(history_run) == codec + assert _no_body_in(history) + marker = await _plain_marker(history, data_converter) + assert _replays_from_store(marker) + assert _markers(history) == ["core_local_activity"] + + replayer_options: dict[str, Any] = {} + if data_converter is not None: + replayer_options["data_converter"] = data_converter + store.reads.clear() + await Replayer( + workflows=[QuietCallWorkflow], + stream_config=_config(store, _HISTORY), + **replayer_options, + ).replay_workflow(history) + assert _run_result(history) == first_run + if call == "filter": + # One read of one Record for each kept Record, after the offset before it. The + # two dropped Records are never read. + assert [request.max_records for request in store.reads] == [1] * len(kept) + offsets = [record.offset for record in await _store_records(store, stream_id)] + assert [request.from_offset for request in store.reads] == [ + offsets[index - 1] if index else b"" for index in kept + ] + else: + assert len(store.reads) == 1 + + +@pytest.mark.parametrize("call", ["take", "filter"]) +async def test_measure_store_marker(client: Client, call: str): + sizes: dict[StreamReplaySource, int] = {} + for replay_from in (_HISTORY, _STORE): + store_client = _client(client, MemoryStreamStore(), replay_from) + if call == "take": + stream_id = await _stream_of(store_client, *_tokens(40), closed=True) + history, count = await _run( + store_client, QuietCallWorkflow, CallArgs(stream_id, until_kind=None) + ) + assert count == 40 + else: + tokens = [ + Token("hit" if index % 8 == 0 else "text", f"{_BODY_MARK}-{index}") + for index in range(40) + ] + stream_id = await _stream_of(store_client, *tokens, closed=True) + history, kept = await _run(store_client, KeepHitsWorkflow, stream_id) + assert kept == 5 + StreamReadState.results.clear() + payload = await _plain_marker(history) + assert _replays_from_store(payload) == (replay_from is _STORE) + assert _no_body_in(history) == (replay_from is _STORE) + sizes[replay_from] = payload.ByteSize() + label = ( + "take of 40 small Records" if call == "take" else "filter that keeps 5 of 40" + ) + print( + f"\nMarker of a {label}: HISTORY {sizes[_HISTORY]} bytes, " + f"STORE {sizes[_STORE]} bytes" + ) + assert sizes[_STORE] < 600 + assert sizes[_STORE] < sizes[_HISTORY] + + +async def _store_history( + client: Client, call: str +) -> tuple[DictStreamStore, str, WorkflowHistory]: + store = DictStreamStore() + store_client = _client(client, store, _STORE) + stream_id = await _stream_of(store_client, *_CALL_TOKENS) + history, _ = await _run(store_client, CallWorkflow, CallArgs(stream_id, call=call)) + StreamReadState.results.clear() + return store, stream_id, history + + +def _describe(call: str, stream_id: str) -> str: + where = f"the Stream '{stream_id}' on the Store 'dict'" + if call == "take": + return f"{where}, 6 Records after the beginning up to 0000000000000006" + return ( + f"{where}, 4 Records at the offsets 0000000000000001, 0000000000000003, " + "0000000000000005, 0000000000000006" + ) + + +@_READ_CALLS +async def test_replay_without_the_records_fails_the_task(client: Client, call: str): + store, stream_id, history = await _store_history(client, call) + prefix = f"The replay of a {call} can't read {_describe(call, stream_id)} from the Store: " + + no_config = await _replay_failure(history, CallWorkflow) + assert 'r#type: "StreamReplayError"' in no_config + assert prefix + "the Worker has no Stream config" in no_config + + del store.streams[stream_id] + deleted = await _replay_failure( + history, CallWorkflow, stream_config=_config(store, _STORE) + ) + assert 'r#type: "StreamReplayError"' in deleted + assert prefix + f"StreamError: No Stream '{stream_id}'" in deleted + + +@_READ_CALLS +async def test_changed_body_fails_the_replay_through_the_digest( + client: Client, call: str +): + store, stream_id, history = await _store_history(client, call) + record = store.streams[stream_id].records[2] + record.body.data = record.body.data.replace(b"t2", b"t9") + failure = await _replay_failure( + history, CallWorkflow, stream_config=_config(store, _HISTORY) + ) + assert 'r#type: "StreamReplayError"' in failure + assert ( + f"The replay of a {call} can't read {_describe(call, stream_id)} from the " + "Store: the Records differ from the digest of the first run" + ) in failure + + +@workflow.defn +class TwoCallsWorkflow: + def __init__(self) -> None: + self.go = False + + @workflow.signal + def proceed(self) -> None: + self.go = True + + @workflow.run + async def run(self, args: CallArgs) -> list[CallResult]: + first = await _call(args) + await workflow.wait_condition(lambda: self.go) + second = await _call(dataclasses.replace(args, from_offset=first.last_offset)) + StreamReadState.results[workflow.info().run_id] = ( + tuple(first.texts), + tuple(second.texts), + ) + return [first, second] + + +@_READ_CALLS +@pytest.mark.parametrize( + "order", + [(_HISTORY, _STORE), (_STORE, _HISTORY)], + ids=["history_first", "store_first"], +) +async def test_change_of_the_setting_keeps_a_running_workflow( + client: Client, call: str, order: tuple[StreamReplaySource, StreamReplaySource] +): + store = MemoryStreamStore() + first_client = _client(client, store, order[0]) + second_client = _client(client, store, order[1]) + stream_id = await _stream_of(first_client, *_CALL_TOKENS, _done("again")) + task_queue = f"tq-{uuid.uuid4()}" + async with new_worker(first_client, TwoCallsWorkflow, task_queue=task_queue): + handle = await first_client.start_workflow( + TwoCallsWorkflow.run, + CallArgs(stream_id, call=call), + id=f"workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + + async def first_call_recorded() -> None: + assert _markers(await handle.fetch_history()) == ["core_local_activity"] + + await assert_eventually(first_call_recorded) + # A new Worker with the other setting replays the first call and runs the second. + async with new_worker(second_client, TwoCallsWorkflow, task_queue=task_queue): + await handle.signal(TwoCallsWorkflow.proceed) + first, second = await handle.result() + kept = [0, 2, 4, 5] if call == "filter" else [0, 1, 2, 3, 4, 5] + assert first.texts == [_CALL_TOKENS[index].text for index in kept] + assert second.texts == [_done("again").text] + history = await handle.fetch_history() + markers = [await _plain_marker(history, index=index) for index in (0, 1)] + assert [_replays_from_store(marker) for marker in markers] == [ + replay_from is _STORE for replay_from in order + ] + run_result = StreamReadState.results.pop(handle.result_run_id or "") + for replay_from in (_HISTORY, _STORE): + await Replayer( + workflows=[TwoCallsWorkflow], stream_config=_config(store, replay_from) + ).replay_workflow(history) + assert StreamReadState.results.pop(handle.result_run_id or "") == run_result + + +async def test_reduce_is_the_same_under_both_settings(client: Client): + store = MemoryStreamStore() + stream_id = await _stream_of(_client(client, store, _HISTORY), *_CALL_TOKENS) + markers: list[Payload] = [] + results: list[CallResult] = [] + for replay_from in (_HISTORY, _STORE): + history, result = await _run( + _client(client, store, replay_from), + CallWorkflow, + CallArgs(stream_id, call="reduce"), + ) + StreamReadState.results.clear() + markers.append(await _plain_marker(history)) + results.append(result) + # The value of a fold is in History, so a replay needs no Store. + await Replayer(workflows=[CallWorkflow]).replay_workflow(history) + StreamReadState.results.clear() + assert results[0] == results[1] + assert results[0].texts == _texts(_CALL_TOKENS) + assert markers[0] == markers[1] + assert not _replays_from_store(markers[1]) + + +async def test_take_over_one_mib_with_store(client: Client): + store = MemoryStreamStore() + store_client = _client(client, store, _STORE) + tokens = [Token("text", f"{index}-" + "x" * 20_000) for index in range(60)] + stream_id = await _stream_of(store_client, *tokens, closed=True) + history, result = await _run( + store_client, CallWorkflow, CallArgs(stream_id, until_kind=None) + ) + # No default cap and no limit of 1 MiB, so one take gives every Record. + assert (result.count, result.stop, result.done) == (60, "Done", True) + assert result.texts == _texts(tokens) + first_run = _run_result(history) + marker = await _plain_marker(history) + assert _replays_from_store(marker) + assert marker.ByteSize() < 1000 + header = Payloads.FromString(marker.data).payloads[0].data.decode() + assert '"count":60' in header + await Replayer( + workflows=[CallWorkflow], stream_config=_config(store, _STORE) + ).replay_workflow(history) + assert _run_result(history) == first_run + # A take with a cap of its own keeps it. + history, capped = await _run( + store_client, + CallWorkflow, + CallArgs(stream_id, until_kind=None, max_bytes=100_000), + ) + StreamReadState.results.clear() + assert (capped.count, capped.stop) == (4, "MaxBytes") + # Under HISTORY the default cap of 256 KiB holds. + _, by_history = await _run( + _client(client, store, _HISTORY), + CallWorkflow, + CallArgs(stream_id, until_kind=None), + ) + StreamReadState.results.clear() + assert by_history.stop == "MaxBytes" + assert by_history.count < 60 diff --git a/tests/nexus/test_stream_read_workflow.py b/tests/nexus/test_stream_read_workflow.py new file mode 100644 index 000000000..ff535235d --- /dev/null +++ b/tests/nexus/test_stream_read_workflow.py @@ -0,0 +1,1455 @@ +"""Tests of the take, filter and reduce calls of Workflow code.""" + +from __future__ import annotations + +import asyncio +import collections +import time +import uuid +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any, cast + +import pytest + +from temporalio import workflow +from temporalio.api.common.v1 import Payload, Payloads +from temporalio.api.enums.v1 import EventType, StreamErrorReason +from temporalio.api.stream.v1 import StreamRef as StreamRefMessage +from temporalio.api.streamservice.v1 import ReadRecordsRequest, ReadRecordsResponse +from temporalio.client import ( + Client, + StreamConfig, + WorkflowFailureError, + WorkflowHandle, + WorkflowHistory, +) +from temporalio.client._stream_read import _PENDING_STREAM_CALLS +from temporalio.contrib.streams import BufferedAppender +from temporalio.converter import DataConverter, PayloadCodec +from temporalio.exceptions import ApplicationError, CancelledError, StreamError +from temporalio.nexus.system.stream_service import ( + ProducerKindLocalActivity, + StreamRead, + StreamReadStop, + StreamRecord, + StreamRef, +) +from temporalio.nexus.system.stream_service.workflow._resources.stream import ( + Stream as WorkflowStream, +) +from temporalio.testing import MemoryStreamStore +from temporalio.worker import Replayer, UnsandboxedWorkflowRunner, WorkflowRunner +from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner + +# The Workflows and the test share the values of `StreamReadState`, so the sandbox +# must not load its own copy of the module. +with workflow.unsafe.imports_passed_through(): + from tests.helpers import assert_eventually, new_worker + from tests.helpers.streams import ( + DictStreamStore, + RecordingStreamStore, + StreamReadState, + ) + +# Every body carries this text, so a search of History shows whether a body is readable. +_BODY_MARK = "secret-body" + + +@dataclass +class Token: + kind: str + text: str + + +@dataclass +class CallArgs: + stream_id: str + call: str = "take" + until_kind: str | None = "done" + timeout_seconds: float = 10.0 + idle_timeout_seconds: float | None = None + max_records: int | None = None + max_bytes: int | None = None + from_offset: str | None = None + + +@dataclass +class CallResult: + texts: list[str] + """The texts of the Records of a take or a filter, or the value of a fold.""" + count: int + stop: str + done: bool + last_offset: str | None + typed: bool + + +def _seconds(value: float | None) -> timedelta | None: + return None if value is None else timedelta(seconds=value) + + +def _stream(stream_id: str) -> WorkflowStream[Token]: + return workflow.get_stream(StreamRef(id=stream_id), body_type=Token) + + +def _not_skipped(record: StreamRecord[Token]) -> bool: + return record.body.kind != "skip" + + +def _collect(value: list[str], record: StreamRecord[Token]) -> list[str]: + StreamReadState.function_calls += 1 + return [*value, record.body.text] + + +def _read_result(read: StreamRead[Token]) -> CallResult: + return CallResult( + texts=[record.body.text for record in read.records], + count=len(read.records), + stop=read.stop.name, + done=read.done, + last_offset=read.last_offset.hex() if read.last_offset else None, + typed=all( + isinstance(record, StreamRecord) and isinstance(record.body, Token) + for record in read.records + ), + ) + + +def _keep(result: CallResult) -> CallResult: + # A replay gives no result to the test, so the Workflow keeps it here. + StreamReadState.results[workflow.info().run_id] = ( + tuple(result.texts), + result.count, + result.stop, + result.done, + result.last_offset, + ) + return result + + +async def _call(args: CallArgs, stream: Any = None) -> CallResult: + """Run the call of ``args``. Client and Activity code give a client Stream.""" + stream = stream or _stream(args.stream_id) + until_kind = args.until_kind + + def until(record: StreamRecord[Token]) -> bool: + StreamReadState.until_calls += 1 + assert isinstance(record.body, Token) + return record.body.kind == until_kind + + def fold_until(value: list[str], record: StreamRecord[Token]) -> bool: + del value + StreamReadState.until_calls += 1 + return record.body.kind == until_kind + + timeout = timedelta(seconds=args.timeout_seconds) + from_offset = bytes.fromhex(args.from_offset) if args.from_offset else None + idle_timeout = _seconds(args.idle_timeout_seconds) + # The arguments go by position, in the order that the generated methods use. + if args.call == "take": + read = await stream.take( + timeout, + until if until_kind else None, + from_offset, + idle_timeout, + args.max_records, + args.max_bytes, + ) + return _read_result(read) + if args.call == "filter": + read = await stream.filter( + _not_skipped, + timeout, + until if until_kind else None, + from_offset, + idle_timeout, + args.max_records, + args.max_bytes, + ) + return _read_result(read) + no_texts: list[str] = [] + reduce = await stream.reduce( + _collect, + no_texts, + timeout, + fold_until if until_kind else None, + from_offset, + idle_timeout, + args.max_records, + args.max_bytes, + ) + return CallResult( + texts=reduce.value, + count=reduce.count, + stop=reduce.stop.name, + done=reduce.done, + last_offset=reduce.last_offset.hex() if reduce.last_offset else None, + typed=isinstance(reduce.value, list), + ) + + +@workflow.defn +class CallWorkflow: + def __init__(self) -> None: + self.result: CallResult | None = None + + @workflow.run + async def run(self, args: CallArgs) -> CallResult: + self.result = _keep(await _call(args)) + return self.result + + @workflow.query + def texts(self) -> list[str]: + assert self.result is not None + return self.result.texts + + +@workflow.defn +class TakeUntilWantedWorkflow: + def __init__(self) -> None: + self.wanted: str | None = None + + @workflow.signal + def want(self, text: str) -> None: + self.wanted = text + + @workflow.run + async def run(self, stream_id: str) -> CallResult: + read = await _stream(stream_id).take( + timedelta(seconds=20), + until=lambda record: record.body.text == self.wanted, + idle_timeout=timedelta(seconds=20), + ) + return _read_result(read) + + +@workflow.defn +class TwoTakesWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[str]: + stream = _stream(stream_id) + first = await stream.take( + timedelta(seconds=10), + until=lambda record: record.body.text.endswith("m1"), + ) + await workflow.sleep(timedelta(milliseconds=1)) + second = await stream.take( + timedelta(seconds=10), + until=lambda record: record.body.text.endswith("m2"), + from_offset=first.last_offset, + ) + return [record.body.text for record in [*first.records, *second.records]] + + +@workflow.defn +class FailTaskOnceAfterTakeWorkflow: + @workflow.run + async def run(self, stream_id: str) -> CallResult: + read = await _stream(stream_id).take( + timedelta(seconds=10), + until=lambda record: record.body.kind == "done", + ) + # The failure comes in a later Workflow Task, so the retry replays the call. + await workflow.sleep(timedelta(milliseconds=1)) + if StreamReadState.fail_next_task: + StreamReadState.fail_next_task = False + raise RuntimeError("Failed the Workflow Task once") + return _read_result(read) + + +@workflow.defn +class FilterLoopWorkflow: + @workflow.run + async def run(self, stream_id: str) -> tuple[list[str], int]: + stream = _stream(stream_id) + found: list[str] = [] + offset: bytes | None = None + calls = 0 + while True: + calls += 1 + # A filter with `max_records=1` gives the next Record that it keeps. + read = await stream.filter( + lambda record: record.body.kind == "hit", + timedelta(seconds=10), + from_offset=offset, + idle_timeout=timedelta(0), + max_records=1, + ) + found.extend(record.body.text for record in read.records) + offset = read.last_offset + if read.stop in ( + StreamReadStop.Done, + StreamReadStop.Idle, + ): + return found, calls + + +@workflow.defn +class TakeLoopWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[tuple[int, str]]: + stream = _stream(stream_id) + pages: list[tuple[int, str]] = [] + offset: bytes | None = None + while True: + read = await stream.take(timedelta(seconds=10), from_offset=offset) + pages.append((len(read.records), read.stop.name)) + offset = read.last_offset + if read.stop is not StreamReadStop.MaxBytes: + return pages + + +@workflow.defn +class OverLimitWorkflow: + @workflow.run + async def run(self, stream_id: str, call: str) -> str: + stream = _stream(stream_id) + try: + if call == "take": + await stream.take(timedelta(seconds=10), max_bytes=4 * 1024 * 1024) + else: + await stream.reduce( + lambda value, record: value + "x" * (1024 * 1024), + "", + timedelta(seconds=10), + max_records=1, + ) + except ApplicationError as err: + return f"{err.type} {err.non_retryable}: {err.message}" + return "no error" + + +@workflow.defn +class TwoStreamsAtOnceWorkflow: + @workflow.run + async def run(self, first_id: str, second_id: str) -> list[str]: + tasks = [ + asyncio.create_task( + _stream(stream_id).filter( + lambda record: record.body.kind == "hit", + timedelta(seconds=10), + max_records=1, + ) + ) + for stream_id in (first_id, second_id) + ] + # `workflow.wait` is the deterministic form of `asyncio.wait`. + await workflow.wait(tasks) + return [record.body.text for task in tasks for record in task.result().records] + + +@workflow.defn +class VariantWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[str]: + # A replay sets `variant` to change the code after the first run. + variant = StreamReadState.variant + if variant == "uuid": + workflow.uuid4() + if variant == "other_stream": + stream_id += "-other" + read = await _stream(stream_id).take( + timedelta(seconds=10), + from_offset=b"\x01" if variant == "other_offset" else None, + idle_timeout=timedelta(0), + ) + return [record.body.text for record in read.records] + + +async def _async_until(record: StreamRecord[Token]) -> bool: + del record + return True + + +def _until_that_awaits(record: StreamRecord[Token]) -> Any: + del record + return asyncio.sleep(0) + + +@workflow.defn +class RefusedFunctionWorkflow: + @workflow.run + async def run(self, stream_id: str, which: str) -> str: + stream = _stream(stream_id) + # The type checker refuses both functions, and the call refuses them at run time. + until: Any = _async_until if which == "coroutine" else _until_that_awaits + try: + await stream.take(timedelta(seconds=10), until=until) + except TypeError as err: + return f"TypeError: {err}" + except ApplicationError as err: + return f"ApplicationError {err.type} {err.non_retryable}: {err.message}" + return "no error" + + +def _raise_on_record(record: StreamRecord[Token]) -> bool: + raise ValueError(f"bad token {record.body.text}") + + +def _call_workflow_now(record: StreamRecord[Token]) -> bool: + del record + return workflow.now() > datetime.min + + +def _call_datetime_now(record: StreamRecord[Token]) -> bool: + del record + return datetime.now() > datetime.min + + +_PREDICATES: dict[str, Callable[[StreamRecord[Token]], bool]] = { + "raise": _raise_on_record, + "workflow_now": _call_workflow_now, + "datetime_now": _call_datetime_now, +} + + +@workflow.defn +class FailingPredicateWorkflow: + @workflow.run + async def run(self, stream_id: str, predicate: str) -> str: + try: + await _stream(stream_id).filter( + _PREDICATES[predicate], timedelta(seconds=10) + ) + except ApplicationError as err: + return f"{err.type}: {err.message}" + return "no error" + + +@dataclass +class Summary: + count: int + kinds: list[str] + + +def _summarize(value: Summary | None, record: StreamRecord[Token]) -> Summary: + if value is None: + value = Summary(0, []) + return Summary(value.count + 1, [*value.kinds, record.body.kind]) + + +def _texts_so_far(value: Sequence[str], record: StreamRecord[Token]) -> list[str]: + return [*value, record.body.text] + + +@workflow.defn +class FoldTypesWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[str]: + stream = _stream(stream_id) + quick = timedelta(seconds=10) + now = timedelta(0) + # The return annotation wins over an `initial` of `None`. + annotated = await stream.reduce(_summarize, None, quick, idle_timeout=now) + # With no annotation, the type of `initial` decides. + of_initial = await stream.reduce( + lambda value, record: Summary(value.count + 1, []), + Summary(0, []), + quick, + idle_timeout=now, + ) + # With no annotation and an `initial` of `None`, there is no type hint. + untyped = await stream.reduce( + lambda value, record: Summary(1, [record.body.kind]), + cast(Any, None), + quick, + idle_timeout=now, + ) + # The return annotation wins over the type of `initial` too. + over_initial = await stream.reduce( + _texts_so_far, cast(Any, ()), quick, idle_timeout=now + ) + return [ + type(result.value).__name__ + for result in (annotated, of_initial, untyped, over_initial) + ] + + +@workflow.defn +class FoldUntilValueWorkflow: + @workflow.run + async def run(self, stream_id: str) -> tuple[str, int, str]: + result = await _stream(stream_id).reduce( + lambda value, record: value + record.body.text, + "", + timedelta(seconds=10), + until=lambda value, record: value.endswith("tag!"), + ) + return result.value, result.count, result.stop.name + + +@workflow.defn +class FoldThatAppendsWorkflow: + @workflow.run + async def run(self, source_id: str, target_id: str) -> str: + target = workflow.get_stream(StreamRef(id=target_id), body_type=str) + appender = BufferedAppender(target) + + def append_upper(value: int, record: StreamRecord[Token]) -> int: + appender.append(record.body.text.upper()) + return value + 1 + + try: + result = await _stream(source_id).reduce( + append_upper, + 0, + timedelta(seconds=10), + idle_timeout=timedelta(0), + ) + except StreamError as err: + return StreamErrorReason.Name(err.reason) + return f"{result.count} {result.stop.name}" + + +class TrackingStore(RecordingStreamStore): + """A memory Store that knows whether a read waits, and whether one was cancelled.""" + + def __init__(self) -> None: + super().__init__(MemoryStreamStore()) + self.active_reads = 0 + self.cancelled_reads = 0 + + async def read_records(self, request: ReadRecordsRequest) -> ReadRecordsResponse: + self.active_reads += 1 + try: + return await super().read_records(request) + except asyncio.CancelledError: + self.cancelled_reads += 1 + raise + finally: + self.active_reads -= 1 + + +class XorCodec(PayloadCodec): + """A codec that leaves no text of a payload readable.""" + + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + return [ + Payload(metadata={"xor-codec": b"true"}, data=_xor(payload)) + for payload in payloads + ] + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + return [ + Payload.FromString(bytes(byte ^ 0x5A for byte in payload.data)) + for payload in payloads + ] + + +def _xor(payload: Payload) -> bytes: + return bytes(byte ^ 0x5A for byte in payload.SerializeToString()) + + +def _tokens(count: int, *, prefix: str = "t", kind: str = "text") -> list[Token]: + return [Token(kind, f"{_BODY_MARK}-{prefix}{index}") for index in range(count)] + + +def _done(text: str = "done") -> Token: + return Token("done", f"{_BODY_MARK}-{text}") + + +def _with_store( + client: Client, store: Any, data_converter: DataConverter | None = None +) -> Client: + config = client.config() + config["stream_config"] = StreamConfig(stores=[store]) + if data_converter is not None: + config["data_converter"] = data_converter + return Client(**config) + + +async def _new_stream( + client: Client, + store: Any = None, + data_converter: DataConverter | None = None, +) -> tuple[Client, Any, str]: + store = store or MemoryStreamStore() + stream_client = _with_store(client, store, data_converter) + stream_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=stream_id, body_type=Token) + return stream_client, store, stream_id + + +async def _append(stream_client: Client, stream_id: str, *tokens: Token) -> None: + stream = stream_client.get_stream(StreamRef(id=stream_id), body_type=Token) + await stream.append(*tokens) + + +async def _close(stream_client: Client, stream_id: str) -> None: + await stream_client.get_stream(StreamRef(id=stream_id), body_type=Token).close() + + +async def _offsets(stream_client: Client, stream_id: str) -> list[str]: + stream = stream_client.get_stream(StreamRef(id=stream_id), body_type=Token) + return [record.offset.hex() async for record in stream.iter_read_items()] + + +async def _append_over( + stream_client: Client, stream_id: str, tokens: list[Token], seconds: float +) -> float: + """Append ``tokens`` one by one over ``seconds``, and give the time of the last.""" + last = time.monotonic() + for token in tokens: + await asyncio.sleep(seconds / len(tokens)) + last = time.monotonic() + await _append(stream_client, stream_id, token) + return last + + +def _texts(tokens: Sequence[Token]) -> list[str]: + return [token.text for token in tokens] + + +def _markers(history: WorkflowHistory) -> list[str]: + return [ + event.marker_recorded_event_attributes.marker_name + for event in history.events + if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED + ] + + +def _marker_payload(history: WorkflowHistory, index: int = 0) -> Payload: + markers = [ + event.marker_recorded_event_attributes + for event in history.events + if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED + ] + return markers[index].details["result"].payloads[0] + + +def _event_counts(history: WorkflowHistory) -> dict[str, int]: + counts = collections.Counter( + EventType.Name(event.event_type).removeprefix("EVENT_TYPE_") + for event in history.events + ) + return dict(sorted(counts.items())) + + +async def _run_workflow( + stream_client: Client, + workflow_class: type, + *args: Any, + worker_options: dict[str, Any] | None = None, + **start_options: Any, +) -> tuple[WorkflowHandle[Any, Any], Any]: + run = getattr(workflow_class, "run") + async with new_worker( + stream_client, workflow_class, **(worker_options or {}) + ) as worker: + handle = await stream_client.start_workflow( + run, + args=list(args), + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + **start_options, + ) + return handle, await handle.result() + + +async def _start_call( + stream_client: Client, args: CallArgs, task_queue: str, **kwargs: Any +) -> WorkflowHandle[CallWorkflow, CallResult]: + return await stream_client.start_workflow( + CallWorkflow.run, + args, + id=f"workflow-{uuid.uuid4()}", + task_queue=task_queue, + **kwargs, + ) + + +_CALLS = pytest.mark.parametrize("call", ["take", "filter", "reduce"]) +_STOPS = ["Match", "Done", "MaxRecords", "MaxBytes", "Idle", "Timeout"] + +# The five Records of a stop test. The filter drops the Record of the kind "skip". +_STOP_TOKENS = [ + Token("text", f"{_BODY_MARK}-t0"), + Token("skip", f"{_BODY_MARK}-s1"), + Token("text", f"{_BODY_MARK}-t2"), + Token("text", f"{_BODY_MARK}-t3"), + Token("text", f"{_BODY_MARK}-t4"), +] + + +def _stop_case( + call: str, stop: str +) -> tuple[list[Token], dict[str, Any], list[int], int]: + """Give the Records, the arguments, the indexes of the result and of ``last_offset``.""" + tokens = list(_STOP_TOKENS) + every = [0, 1, 2, 3, 4] + kept = every if call != "filter" else [0, 2, 3, 4] + if stop == "Match": + tokens[3] = _done() + return tokens, {}, [index for index in kept if index <= 3], 3 + if stop == "Done": + return tokens, {"until_kind": None}, kept, 4 + if stop == "MaxRecords": + return tokens, {"max_records": 2}, kept[:2], kept[1] + if stop == "MaxBytes": + # The first Record passes the cap and comes back, and the next kept one stops the + # call. A filter looks at the dropped Record in between. + return tokens, {"max_bytes": 1}, kept[:1], 1 if call == "filter" else 0 + if stop == "Idle": + return tokens, {"idle_timeout_seconds": 0.3, "until_kind": None}, kept, 4 + return ( + tokens, + {"timeout_seconds": 1, "idle_timeout_seconds": 30, "until_kind": None}, + kept, + 4, + ) + + +async def _check_stop( + stream_client: Client, + call: str, + stop: str, + run: Callable[[str, dict[str, Any]], Any], +) -> None: + """Append the Records of a stop case, run the call, and check its result.""" + stream_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=stream_id, body_type=Token) + tokens, options, indexes, last = _stop_case(call, stop) + await _append(stream_client, stream_id, *tokens) + if stop == "Done": + await _close(stream_client, stream_id) + offsets = await _offsets(stream_client, stream_id) + result = await run(stream_id, options) + texts = [tokens[index].text for index in indexes] + if call == "reduce": + # A fold gets every Record that it looks at. + texts = [tokens[index].text for index in range(last + 1)] + assert (result.stop, result.texts, result.last_offset) == ( + stop, + texts, + offsets[last], + ) + assert result.count == len(texts) + assert result.done == (stop == "Done") + + +@_CALLS +@pytest.mark.parametrize("stop", _STOPS) +async def test_call_stop_and_last_offset(client: Client, call: str, stop: str): + stream_client = _with_store(client, MemoryStreamStore()) + + async def run(stream_id: str, options: dict[str, Any]) -> CallResult: + _, result = await _run_workflow( + stream_client, CallWorkflow, CallArgs(stream_id, call=call, **options) + ) + assert result.typed + return result + + await _check_stop(stream_client, call, stop, run) + + +async def test_filter_loop_reads_to_the_end(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens: list[Token] = [] + for index in range(20): + kind = "hit" if index % 4 == 1 else "text" + tokens.append(Token(kind, f"{_BODY_MARK}-{kind}{index}")) + await _append(stream_client, stream_id, *tokens) + await _close(stream_client, stream_id) + handle, (found, calls) = await _run_workflow( + stream_client, FilterLoopWorkflow, stream_id + ) + hits = [token.text for token in tokens if token.kind == "hit"] + assert found == hits + # One call for each of the 5 matches, and one that reaches the end. + assert calls == 6 + history = await handle.fetch_history() + assert len(_markers(history)) == 6 + print( + f"\nC. A filter loop over 5 matches, History events: {_event_counts(history)}" + ) + + +async def test_filter_until_accepts_a_dropped_record(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = [ + Token("text", f"{_BODY_MARK}-t0"), + Token("skip", f"{_BODY_MARK}-s1"), + Token("text", f"{_BODY_MARK}-t2"), + ] + await _append(stream_client, stream_id, *tokens) + offsets = await _offsets(stream_client, stream_id) + _, result = await _run_workflow( + stream_client, + CallWorkflow, + CallArgs(stream_id, call="filter", until_kind="skip"), + ) + assert result.texts == _texts(tokens[:1]) + assert result.stop == "Match" + assert result.last_offset == offsets[1] + + +@_CALLS +async def test_replay_needs_no_stream_config_and_calls_no_store( + client: Client, call: str +): + store = RecordingStreamStore(MemoryStreamStore()) + stream_client, _, stream_id = await _new_stream(client, store) + await _append(stream_client, stream_id, *_STOP_TOKENS[:3], _done()) + handle, _ = await _run_workflow( + stream_client, CallWorkflow, CallArgs(stream_id, call=call) + ) + run_id = handle.result_run_id + assert run_id + first_run = StreamReadState.results.pop(run_id) + history = await handle.fetch_history() + assert _markers(history) == ["core_local_activity"] + StreamReadState.until_calls = 0 + StreamReadState.function_calls = 0 + operations = len(store.operations) + await Replayer(workflows=[CallWorkflow]).replay_workflow(history) + assert StreamReadState.results.pop(run_id) == first_run + await Replayer( + workflows=[CallWorkflow], stream_config=StreamConfig(stores=[store]) + ).replay_workflow(history) + assert StreamReadState.results.pop(run_id) == first_run + assert len(store.operations) == operations + assert StreamReadState.until_calls == 0 + assert StreamReadState.function_calls == 0 + + +async def test_query_after_the_stream_is_deleted(client: Client): + store = DictStreamStore() + stream_client, _, stream_id = await _new_stream(client, store) + tokens = [*_tokens(3), _done()] + await _append(stream_client, stream_id, *tokens) + task_queue = f"tq-{uuid.uuid4()}" + async with new_worker(stream_client, CallWorkflow, task_queue=task_queue): + handle = await _start_call( + stream_client, CallArgs(stream_id, idle_timeout_seconds=0), task_queue + ) + assert (await handle.result()).texts == _texts(tokens) + del store.streams[stream_id] + # A new Worker has no cache, so the Query replays the whole History. + async with new_worker(stream_client, CallWorkflow, task_queue=task_queue): + assert await handle.query(CallWorkflow.texts) == _texts(tokens) + assert stream_id not in store.streams + + +async def test_take_loop_over_max_bytes(client: Client): + stream_client, _, stream_id = await _new_stream(client) + # 40 Records of about 10 KiB pass the default cap of 256 KiB. + tokens = [ + Token("text", f"{_BODY_MARK}-{index}-" + "x" * 10_000) for index in range(40) + ] + for start in range(0, len(tokens), 10): + await _append(stream_client, stream_id, *tokens[start : start + 10]) + await _close(stream_client, stream_id) + _, pages = await _run_workflow(stream_client, TakeLoopWorkflow, stream_id) + assert [stop for _, stop in pages] == ["MaxBytes", "Done"] + assert sum(count for count, _ in pages) == len(tokens) + print(f"\nTake loop over 400 KiB: {pages}") + + +@pytest.mark.parametrize("call", ["take", "reduce"]) +async def test_result_over_the_limit_raises_at_the_await(client: Client, call: str): + stream_client, _, stream_id = await _new_stream(client) + tokens = [Token("text", "x" * 20_000) for _ in range(60)] + for start in range(0, len(tokens), 10): + await _append(stream_client, stream_id, *tokens[start : start + 10]) + await _close(stream_client, stream_id) + _, message = await _run_workflow(stream_client, OverLimitWorkflow, stream_id, call) + prefix = f"StreamResultTooLarge True: The result of the {call} of the Stream " + assert message.startswith(prefix + f"'{stream_id}' is ") + assert "bytes encoded, over the limit of 1048576 bytes." in message + + +@pytest.mark.parametrize( + "variant", ["other_stream", "other_offset", "uuid"], ids=lambda value: value +) +async def test_replay_checks_the_stream_and_offset_of_a_call( + client: Client, variant: str +): + stream_client, _, stream_id = await _new_stream(client) + await _append(stream_client, stream_id, *_tokens(2)) + StreamReadState.variant = "" + handle, _ = await _run_workflow(stream_client, VariantWorkflow, stream_id) + history = await handle.fetch_history() + StreamReadState.variant = variant + try: + result = await Replayer(workflows=[VariantWorkflow]).replay_workflow( + history, raise_on_replay_failure=False + ) + finally: + StreamReadState.variant = "" + if variant == "uuid": + assert result.replay_failure is None + return + failure = str(result.replay_failure) + # The Workflow code raised the SDK's non-determinism error, which failed the task. + assert 'r#type: "NondeterminismError"' in failure + assert "Stream call 1 of the Workflow is a take of the Stream" in failure + assert f"History has a take of the Stream '{stream_id}' on the Store" in failure + + +async def test_payload_codec_hides_every_body_in_history(client: Client): + data_converter = DataConverter(payload_codec=XorCodec()) + stream_client, _, stream_id = await _new_stream( + client, data_converter=data_converter + ) + await _append(stream_client, stream_id, *_tokens(3), _done()) + for call in ("take", "filter", "reduce"): + handle, result = await _run_workflow( + stream_client, CallWorkflow, CallArgs(stream_id, call=call) + ) + assert result.count == 4 and result.stop == "Match" + history = await handle.fetch_history() + payload = _marker_payload(history) + assert payload.metadata["xor-codec"] == b"true" + assert _BODY_MARK not in history.to_json() + for event in history.events: + assert _BODY_MARK.encode() not in event.SerializeToString() + await Replayer( + workflows=[CallWorkflow], data_converter=data_converter + ).replay_workflow(history) + + +@pytest.mark.parametrize("which", ["coroutine", "awaitable"]) +async def test_function_must_be_plain(client: Client, which: str): + stream_client, _, stream_id = await _new_stream(client) + await _append(stream_client, stream_id, *_tokens(1)) + _, message = await _run_workflow( + stream_client, RefusedFunctionWorkflow, stream_id, which + ) + if which == "coroutine": + assert message.startswith("TypeError: until must be a plain function") + else: + assert message == ( + "ApplicationError TypeError True: The until of a Stream call returned an " + "awaitable. It must be a plain function that returns its answer." + ) + + +async def test_fold_value_types(client: Client): + stream_client, _, stream_id = await _new_stream(client) + await _append(stream_client, stream_id, *_tokens(2)) + _, names = await _run_workflow(stream_client, FoldTypesWorkflow, stream_id) + assert names == ["Summary", "Summary", "dict", "list"] + + +async def test_fold_until_sees_the_value(client: Client): + stream_client, _, stream_id = await _new_stream(client) + texts = ["x", "ta", "g!", "y"] + await _append(stream_client, stream_id, *[Token("text", text) for text in texts]) + _, result = await _run_workflow(stream_client, FoldUntilValueWorkflow, stream_id) + assert list(result) == ["xtag!", 3, "Match"] + + +async def test_two_calls_on_two_streams_at_once(client: Client): + stream_client, _, first_id = await _new_stream(client) + second_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=second_id, body_type=Token) + await _append(stream_client, first_id, *_tokens(2), Token("hit", "first-hit")) + await _append(stream_client, second_id, Token("hit", "second-hit")) + handle, texts = await _run_workflow( + stream_client, TwoStreamsAtOnceWorkflow, first_id, second_id + ) + assert texts == ["first-hit", "second-hit"] + assert _markers(await handle.fetch_history()) == ["core_local_activity"] * 2 + + +async def test_default_idle_stop_on_a_quiet_stream(client: Client): + stream_client, _, stream_id = await _new_stream(client) + start = time.monotonic() + handle, result = await _run_workflow( + stream_client, CallWorkflow, CallArgs(stream_id, timeout_seconds=30) + ) + seconds = time.monotonic() - start + assert result.stop == "Idle" and result.texts == [] + assert 1 <= seconds < 3 + counts = _event_counts(await handle.fetch_history()) + # The one Workflow Task held the call, so no heartbeat added a task. + assert counts["WORKFLOW_TASK_COMPLETED"] == 1 + print(f"\nB. A quiet Stream with the default idle stop, History events: {counts}") + + +async def test_take_waits_for_records(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = [*_tokens(4), _done()] + async with new_worker(stream_client, CallWorkflow) as worker: + handle = await _start_call( + stream_client, + CallArgs(stream_id, idle_timeout_seconds=5), + worker.task_queue, + ) + await _append_over(stream_client, stream_id, tokens, 1.0) + result = await handle.result() + assert result.texts == _texts(tokens) + assert result.stop == "Match" + assert _markers(await handle.fetch_history()) == ["core_local_activity"] + + +async def _wait_for_signal_delivery(handle: WorkflowHandle[Any, Any]) -> None: + async def delivered() -> None: + events = (await handle.fetch_history()).events + signaled = [ + index + for index, event in enumerate(events) + if event.event_type == EventType.EVENT_TYPE_WORKFLOW_EXECUTION_SIGNALED + ] + assert signaled + assert any( + event.event_type == EventType.EVENT_TYPE_WORKFLOW_TASK_STARTED + for event in events[signaled[0] :] + ) + + await assert_eventually(delivered) + + +_RUNNERS = pytest.mark.parametrize( + "runner", + [SandboxedWorkflowRunner(), UnsandboxedWorkflowRunner()], + ids=["sandboxed", "unsandboxed"], +) + + +@_RUNNERS +async def test_until_reads_workflow_state(client: Client, runner: WorkflowRunner): + stream_client, store, stream_id = await _new_stream(client, TrackingStore()) + first = _tokens(2) + wanted = Token("text", f"{_BODY_MARK}-wanted") + await _append(stream_client, stream_id, *first) + async with new_worker( + stream_client, TakeUntilWantedWorkflow, workflow_runner=runner + ) as worker: + handle = await stream_client.start_workflow( + TakeUntilWantedWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + task_timeout=timedelta(seconds=2), + ) + + async def read_waits() -> None: + assert store.active_reads == 1 + + await assert_eventually(read_waits) + await handle.signal(TakeUntilWantedWorkflow.want, wanted.text) + # A Workflow Task heartbeat brings the Signal while the Local Activity reads. + await _wait_for_signal_delivery(handle) + await asyncio.sleep(0.5) + await _append(stream_client, stream_id, wanted, *_tokens(2, prefix="after")) + result = await handle.result() + assert result.texts == _texts([*first, wanted]) + assert result.stop == "Match" + + +async def test_calls_on_a_worker_without_cache(client: Client): + stream_client, store, stream_id = await _new_stream( + client, RecordingStreamStore(MemoryStreamStore()) + ) + tokens = [ + Token("text", f"{_BODY_MARK}-{text}") for text in ["a", "m1", "b", "m2", "c"] + ] + await _append(stream_client, stream_id, *tokens) + handle, texts = await _run_workflow( + stream_client, + TwoTakesWorkflow, + stream_id, + worker_options={"max_cached_workflows": 0}, + ) + assert texts == _texts(tokens[:4]) + assert _markers(await handle.fetch_history()) == ["core_local_activity"] * 2 + # The Workflow Task of the second call replays the first from History. + assert store.operations.count("ReadRecords") == 2 + + +async def test_call_survives_a_failed_workflow_task(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = [*_tokens(3), _done()] + await _append(stream_client, stream_id, *tokens) + StreamReadState.fail_next_task = True + handle, result = await _run_workflow( + stream_client, FailTaskOnceAfterTakeWorkflow, stream_id + ) + assert result.texts == _texts(tokens) + assert result.stop == "Match" + counts = _event_counts(await handle.fetch_history()) + assert counts["WORKFLOW_TASK_FAILED"] == 1 + assert counts["MARKER_RECORDED"] == 1 + + +async def _failing_predicate( + client: Client, + predicate: str, + runner: WorkflowRunner = SandboxedWorkflowRunner(), +) -> str: + stream_client, _, stream_id = await _new_stream(client) + await _append(stream_client, stream_id, *_tokens(1)) + _, result = await _run_workflow( + stream_client, + FailingPredicateWorkflow, + stream_id, + predicate, + worker_options={"workflow_runner": runner}, + ) + return result + + +async def test_predicate_that_raises(client: Client): + assert await _failing_predicate(client, "raise") == ( + f"ValueError: bad token {_BODY_MARK}-t0" + ) + + +@_RUNNERS +async def test_predicate_cannot_call_workflow_functions( + client: Client, runner: WorkflowRunner +): + # The predicate runs on the Worker's event loop, where no Workflow runs. + assert await _failing_predicate(client, "workflow_now", runner) == ( + "_NotInWorkflowEventLoopError: Not in workflow event loop" + ) + + +async def test_predicate_in_sandbox_keeps_its_restrictions(client: Client): + sandboxed = await _failing_predicate(client, "datetime_now") + assert sandboxed.startswith("RestrictedWorkflowAccessError: ") + unsandboxed = await _failing_predicate( + client, "datetime_now", UnsandboxedWorkflowRunner() + ) + assert unsandboxed == "no error" + + +async def test_calls_without_stream_config_fail_the_workflow(client: Client): + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker(client, CallWorkflow) as worker: + for call in ("take", "filter", "reduce"): + handle = await _start_call( + client, + CallArgs(stream_id, call=call), + worker.task_queue, + execution_timeout=timedelta(seconds=30), + ) + with pytest.raises(WorkflowFailureError) as err: + await asyncio.wait_for(handle.result(), timeout=10) + cause = err.value.cause + assert isinstance(cause, ApplicationError) + assert cause.type == "StreamConfigError" + assert cause.non_retryable + + +async def _cancel_waiting_call(client: Client, task_timeout: timedelta | None) -> float: + """Cancel a Workflow while its call waits, and give the seconds until it ends.""" + stream_client, store, stream_id = await _new_stream(client, TrackingStore()) + async with new_worker(stream_client, CallWorkflow) as worker: + handle = await _start_call( + stream_client, + CallArgs(stream_id, timeout_seconds=30, idle_timeout_seconds=30), + worker.task_queue, + task_timeout=task_timeout, + ) + + async def read_waits() -> None: + assert store.active_reads == 1 + + await assert_eventually(read_waits) + start = time.monotonic() + await handle.cancel() + with pytest.raises(WorkflowFailureError) as err: + await handle.result() + seconds = time.monotonic() - start + assert isinstance(err.value.cause, CancelledError) + + async def read_stopped() -> None: + assert store.active_reads == 0 + assert store.cancelled_reads == 1 + + await assert_eventually(read_stopped) + print(f"\nCancel with task_timeout={task_timeout}: ended after {seconds:.1f} s") + return seconds + + +@pytest.mark.xfail( + strict=True, + reason=( + "The cancel reaches the Workflow only at the next Workflow Task heartbeat, " + "about 8 s with the default task timeout of 10 s." + ), +) +async def test_cancel_of_a_long_wait(client: Client): + assert await _cancel_waiting_call(client, None) < 5 + + +async def test_cancel_with_short_task_timeout(client: Client): + assert await _cancel_waiting_call(client, timedelta(seconds=2)) < 5 + + +async def test_call_spans_heartbeats_on_a_worker_without_cache(client: Client): + stream_client, store, stream_id = await _new_stream(client, TrackingStore()) + first = _tokens(2) + await _append(stream_client, stream_id, *first) + async with new_worker( + stream_client, CallWorkflow, max_cached_workflows=0 + ) as worker: + handle = await _start_call( + stream_client, + CallArgs(stream_id, timeout_seconds=30, idle_timeout_seconds=30), + worker.task_queue, + task_timeout=timedelta(seconds=2), + ) + # The registry keeps the entry of the call at every sample while it waits. + samples = [] + for _ in range(16): + await asyncio.sleep(0.25) + samples.append( + any( + run_id == handle.result_run_id + for run_id, _ in _PENDING_STREAM_CALLS.keys() + ) + ) + await _append(stream_client, stream_id, _done()) + result = await handle.result() + counts = _event_counts(await handle.fetch_history()) + reads = store.operations.count("ReadRecords") + print( + f"\nHeartbeats without cache: {counts}, ReadRecords calls {reads}, " + f"cancelled {store.cancelled_reads}, entry in {sum(samples)} of " + f"{len(samples)} samples" + ) + assert result.texts == _texts([*first, _done()]) + assert result.stop == "Match" + # Heartbeats happened, and the one Local Activity ran through them uncancelled. + assert counts["WORKFLOW_TASK_COMPLETED"] >= 3 + assert reads == 2 + assert store.cancelled_reads == 0 + assert all(samples) + assert not any( + run_id == handle.result_run_id for run_id, _ in _PENDING_STREAM_CALLS.keys() + ) + + +async def _record_sizes(store: Any, stream_id: str) -> list[int]: + response = await store.read_records( + ReadRecordsRequest(ref=StreamRefMessage(id=stream_id, store=store.name)) + ) + return [record.ByteSize() for record in response.records] + + +async def test_take_max_bytes_boundary(client: Client): + stream_client, store, stream_id = await _new_stream(client) + tokens = _tokens(5) + await _append(stream_client, stream_id, *tokens) + sizes = await _record_sizes(store, stream_id) + # The third Record passes the cap by one byte. + _, result = await _run_workflow( + stream_client, + CallWorkflow, + CallArgs(stream_id, until_kind=None, max_bytes=sum(sizes[:3]) - 1), + ) + assert result.texts == _texts(tokens[:2]) + assert result.stop == "MaxBytes" + + +async def _run_flowing_take( + client: Client, *, seconds: float = 2.0 +) -> tuple[WorkflowHistory, float]: + stream_client, _, stream_id = await _new_stream(client) + tokens = [*_tokens(39), _done()] + async with new_worker(stream_client, CallWorkflow) as worker: + handle = await _start_call( + stream_client, + CallArgs(stream_id, timeout_seconds=30, idle_timeout_seconds=5), + worker.task_queue, + ) + appended = await _append_over(stream_client, stream_id, tokens, seconds) + result = await handle.result() + latency = time.monotonic() - appended + assert result.texts == _texts(tokens) + return await handle.fetch_history(), latency + + +async def test_measure_flowing_take(client: Client): + history, _ = await _run_flowing_take(client) + payload = _marker_payload(history) + print(f"\nA. 40 Records over 2 s, History events: {_event_counts(history)}") + print(f"A. Result payload: {payload.ByteSize()} bytes") + + +async def test_measure_take_of_forty(client: Client): + stream_client, _, stream_id = await _new_stream(client) + await _append(stream_client, stream_id, *_tokens(40)) + await _close(stream_client, stream_id) + handle, result = await _run_workflow( + stream_client, CallWorkflow, CallArgs(stream_id, until_kind=None) + ) + assert (result.count, result.stop) == (40, "Done") + payload = _marker_payload(await handle.fetch_history()) + header = Payloads.FromString(payload.data).payloads[0].data.decode() + print(f"\nResult of a take of 40 small Records: {payload.ByteSize()} bytes") + print(f"Its header: {header}") + + +@workflow.defn +class KeepHitsWorkflow: + @workflow.run + async def run(self, stream_id: str) -> int: + read = await _stream(stream_id).filter( + lambda record: record.body.kind == "hit", + timedelta(seconds=10), + idle_timeout=timedelta(0), + ) + return len(read.records) + + +async def test_measure_filter_that_keeps_five(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = [ + Token("hit" if index % 8 == 0 else "text", f"{_BODY_MARK}-{index}") + for index in range(40) + ] + await _append(stream_client, stream_id, *tokens) + await _close(stream_client, stream_id) + handle, kept = await _run_workflow(stream_client, KeepHitsWorkflow, stream_id) + assert kept == 5 + payload = _marker_payload(await handle.fetch_history()) + print(f"\nResult of a filter that keeps 5 of 40: {payload.ByteSize()} bytes") + + +async def test_measure_idle_wait(client: Client): + stream_client, _, stream_id = await _new_stream(client) + async with new_worker(stream_client, CallWorkflow) as worker: + handle = await _start_call( + stream_client, + CallArgs(stream_id, timeout_seconds=30, idle_timeout_seconds=30), + worker.task_queue, + task_timeout=timedelta(seconds=2), + ) + await asyncio.sleep(6) + await _append(stream_client, stream_id, _done()) + await handle.result() + counts = _event_counts(await handle.fetch_history()) + print(f"\nA wait of 6 s with idle_timeout=30 s, History events: {counts}") + + +async def test_measure_match_latency(client: Client): + latencies = [] + for _ in range(5): + _, latency = await _run_flowing_take(client) + latencies.append(latency) + print( + "\nAppend of the match to the end of the Workflow: " + + ", ".join(f"{latency * 1000:.0f} ms" for latency in latencies) + ) + + +async def _signal_wait(client: Client) -> float: + stream_client, _, stream_id = await _new_stream(client) + async with new_worker(stream_client, TakeUntilWantedWorkflow) as worker: + handle = await stream_client.start_workflow( + TakeUntilWantedWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await asyncio.sleep(1) + await handle.signal(TakeUntilWantedWorkflow.want, "never") + + async def delivered() -> float: + events = (await handle.fetch_history()).events + signals = [ + event + for event in events + if event.event_type == EventType.EVENT_TYPE_WORKFLOW_EXECUTION_SIGNALED + ] + assert signals + signaled = signals[0] + started = [ + event + for event in events + if event.event_type == EventType.EVENT_TYPE_WORKFLOW_TASK_STARTED + and event.event_id > signaled.event_id + ] + assert started + return ( + started[0].event_time.ToDatetime() - signaled.event_time.ToDatetime() + ).total_seconds() + + seconds = await assert_eventually(delivered, timeout=timedelta(seconds=15)) + await handle.terminate() + return seconds + + +@pytest.mark.timeout(120) +async def test_measure_signal_wait(client: Client): + waits = [await _signal_wait(client) for _ in range(3)] + print( + "\nSignal to the next Workflow Task Started while a call waits: " + + ", ".join(f"{wait:.2f} s" for wait in waits) + ) + + +async def test_measure_flowing_fold(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = [*_tokens(39), _done()] + async with new_worker(stream_client, CallWorkflow) as worker: + handle = await _start_call( + stream_client, + CallArgs( + stream_id, call="reduce", timeout_seconds=30, idle_timeout_seconds=5 + ), + worker.task_queue, + ) + await _append_over(stream_client, stream_id, tokens, 2.0) + result = await handle.result() + assert result.count == 40 + history = await handle.fetch_history() + payload = _marker_payload(history) + print( + f"\nD. A fold over 40 Records over 2 s, History events: {_event_counts(history)}" + ) + print(f"D. Result payload: {payload.ByteSize()} bytes") + + +async def test_fold_appends_each_record_to_another_stream(client: Client): + store = RecordingStreamStore(MemoryStreamStore()) + stream_client, _, source_id = await _new_stream(client, store) + target_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=target_id, body_type=str) + tokens = _tokens(40) + await _append(stream_client, source_id, *tokens) + await _close(stream_client, source_id) + handle, result = await _run_workflow( + stream_client, FoldThatAppendsWorkflow, source_id, target_id + ) + assert result == "40 Done" + target = stream_client.get_stream(StreamRef(id=target_id), body_type=str) + records = [record async for record in target.iter_read_items()] + assert [record.body for record in records] == [ + text.upper() for text in _texts(tokens) + ] + run_id = handle.result_run_id + producer = records[0].producer + assert producer is not None and isinstance(producer.kind, ProducerKindLocalActivity) + assert producer.kind.value.run_id == run_id + history = await handle.fetch_history() + counts = _event_counts(history) + # One marker for the fold. The appends went to the Store from its Local Activity. + assert _markers(history) == ["core_local_activity"] + assert "TIMER_STARTED" not in counts + print(f"\nE. A fold that appends 40 Records to another Stream: {counts}") + operations = len(store.operations) + await Replayer(workflows=[FoldThatAppendsWorkflow]).replay_workflow(history) + await Replayer( + workflows=[FoldThatAppendsWorkflow], stream_config=StreamConfig(stores=[store]) + ).replay_workflow(history) + assert len(store.operations) == operations + assert len([record async for record in target.iter_read_items()]) == 40 + + +async def test_failed_append_in_a_function_fails_the_call(client: Client): + stream_client, _, source_id = await _new_stream(client) + await _append(stream_client, source_id, *_tokens(2)) + _, result = await _run_workflow( + stream_client, + FoldThatAppendsWorkflow, + source_id, + f"stream-{uuid.uuid4()}", + ) + assert result == "STREAM_ERROR_REASON_NOT_FOUND" diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py index 07ae11984..4255f890a 100644 --- a/tests/nexus/test_stream_store_workflow.py +++ b/tests/nexus/test_stream_store_workflow.py @@ -642,8 +642,9 @@ async def test_workflow_read_fails_with_store(client: Client): task_queue=worker.task_queue, ) assert result == ( - "Workflow code can't read a Stream on a Store in this release. " - "Call an Activity that reads it through its client." + "Workflow code can't run a ReadRecords operation on a Store. Use the " + "take, filter or reduce method of the Workflow Stream, which reads in one " + "Local Activity." )