From 0bf0f0ab4ab49d8ad97e5ddbadfaf89bc939ebf8 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 16:57:39 -0700 Subject: [PATCH 01/20] Added a Workflow read of a Stream that waits for a predicate. One Local Activity reads pages until the Workflow's predicate is true, so History gets one marker for the whole read and no Record. The predicate and the Records pass through a registry in the Worker process, since a closure can't be serialized and Records must stay out of History. --- temporalio/client/_stream_read.py | 288 ++++++++++++++++++++++++ temporalio/client/_stream_store.py | 93 +++++++- temporalio/worker/_stream_read.py | 95 ++++++++ temporalio/worker/_stream_store.py | 66 +++--- temporalio/worker/_workflow.py | 2 + temporalio/worker/_workflow_instance.py | 44 +++- temporalio/workflow/__init__.py | 8 + temporalio/workflow/_context.py | 17 ++ temporalio/workflow/_streams.py | 116 ++++++++++ 9 files changed, 691 insertions(+), 38 deletions(-) create mode 100644 temporalio/client/_stream_read.py create mode 100644 temporalio/worker/_stream_read.py create mode 100644 temporalio/workflow/_streams.py diff --git a/temporalio/client/_stream_read.py b/temporalio/client/_stream_read.py new file mode 100644 index 0000000000..43e5564395 --- /dev/null +++ b/temporalio/client/_stream_read.py @@ -0,0 +1,288 @@ +"""The process side of a Stream read that Workflow code waits on. + +A Workflow's read runs in a Local Activity, and the Local Activity calls a predicate +of the Workflow. A predicate is a closure and can't be serialized, and the Records +must stay out of History. So both pass through the registry of this module, in the +memory of the Worker process. History keeps only a small marker with the range of +offsets of the read. + +.. warning:: + This API is an experimental prototype and subject to change. +""" + +from __future__ import annotations + +import base64 +import dataclasses +import hashlib +import json +import struct +import threading +import time +from collections.abc import Awaitable, Callable, Sequence +from datetime import timedelta +from typing import Any + +import temporalio.api.common.v1 +import temporalio.api.stream.v1 +import temporalio.api.streamservice.v1 +import temporalio.common +import temporalio.exceptions + +_MARKER_METADATA_KEY = "temporal-stream-read" +_MARKER_METADATA_VALUE = b"v1" +_MAX_PAGE_RECORDS = 1000 +# A Store waits at most this long for one page. +_MAX_PAGE_WAIT = timedelta(seconds=60) +_DIGEST_SIZE = 16 + + +def _encode_offset(offset: bytes | None) -> str | None: + return None if offset is None else base64.b64encode(offset).decode() + + +def _decode_offset(offset: str | None) -> bytes | None: + return None if offset is None else base64.b64decode(offset) + + +def _records_digest( + records: Sequence[temporalio.api.stream.v1.StreamRecord], +) -> bytes: + """Give the first bytes of a SHA-256 over the offsets and bodies of ``records``. + + A length before each part keeps two different splits of the same bytes apart. + """ + digest = hashlib.sha256() + for record in records: + body = record.body.SerializeToString(deterministic=True) + for part in (record.offset, body): + digest.update(struct.pack(">Q", len(part))) + digest.update(part) + return digest.digest()[:_DIGEST_SIZE] + + +@dataclasses.dataclass(frozen=True) +class _StreamReadMarker: + """The result of a read's Local Activity, and so the only part of it in History.""" + + read_id: str + store: str + stream_id: str + from_offset: bytes | None + last_offset: bytes | None + record_count: int + matched: bool + done: bool + digest: bytes + + @classmethod + def for_records( + cls, + *, + read_id: str, + store: str, + stream_id: str, + from_offset: bytes | None, + records: Sequence[temporalio.api.stream.v1.StreamRecord], + matched: bool, + done: bool, + ) -> _StreamReadMarker: + """Build the marker of a read that gave ``records``.""" + return cls( + read_id=read_id, + store=store, + stream_id=stream_id, + from_offset=from_offset, + last_offset=records[-1].offset if records else from_offset, + record_count=len(records), + matched=matched, + done=done, + digest=_records_digest(records), + ) + + def to_raw_value(self) -> temporalio.common.RawValue: + """Give the marker as a payload that a Worker can find in any activation.""" + data = { + "read_id": self.read_id, + "store": self.store, + "stream_id": self.stream_id, + "from_offset": _encode_offset(self.from_offset), + "last_offset": _encode_offset(self.last_offset), + "record_count": self.record_count, + "matched": self.matched, + "done": self.done, + "digest": self.digest.hex(), + } + return temporalio.common.RawValue( + temporalio.api.common.v1.Payload( + metadata={ + "encoding": b"json/plain", + _MARKER_METADATA_KEY: _MARKER_METADATA_VALUE, + }, + data=json.dumps(data, separators=(",", ":"), sort_keys=True).encode(), + ) + ) + + @classmethod + def from_payload( + cls, payload: temporalio.api.common.v1.Payload + ) -> _StreamReadMarker | None: + """Give the marker in ``payload``, or ``None`` when the payload isn't one.""" + if payload.metadata.get(_MARKER_METADATA_KEY) != _MARKER_METADATA_VALUE: + return None + data = json.loads(payload.data) + return cls( + read_id=data["read_id"], + store=data["store"], + stream_id=data["stream_id"], + from_offset=_decode_offset(data["from_offset"]), + last_offset=_decode_offset(data["last_offset"]), + record_count=data["record_count"], + matched=data["matched"], + done=data["done"], + digest=bytes.fromhex(data["digest"]), + ) + + def describe_range(self) -> str: + """Name the Stream and the range of the read, for an error message.""" + start = self.from_offset.hex() if self.from_offset else "the beginning" + end = self.last_offset.hex() if self.last_offset else "the beginning" + return ( + f"the Stream '{self.stream_id}' on the Store '{self.store}', " + f"{self.record_count} Records after {start} up to {end}" + ) + + +@dataclasses.dataclass(eq=False) +class _PendingStreamRead: + """One read of Workflow code, from the start of its Local Activity to its result.""" + + until: Callable[[Any], bool] | None + """The predicate of the Workflow, or ``None``.""" + + to_record: Callable[[temporalio.api.stream.v1.StreamRecord], Any] + """Turns a proto Record into the typed Record of the Workflow.""" + + _result: ( + tuple[_StreamReadMarker, list[temporalio.api.stream.v1.StreamRecord]] | None + ) = None + + def fill( + self, + marker: _StreamReadMarker, + records: list[temporalio.api.stream.v1.StreamRecord], + ) -> None: + """Keep the Records of ``marker``, whose payloads the codec already decoded.""" + # One assignment, so a reader on another thread sees both or neither. + self._result = (marker, records) + + def records_for( + self, marker: _StreamReadMarker + ) -> list[temporalio.api.stream.v1.StreamRecord] | None: + """Give the Records of ``marker``, or ``None`` when this read has none for it.""" + result = self._result + if result is None or result[0] != marker: + return None + return result[1] + + +class _PendingStreamReads: + """The pending reads of the process, by Workflow run id and read id. + + Workflow code adds and removes an entry on its own thread. The read Activity and the + Worker look an entry up on the event loop of the Worker. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._reads: dict[tuple[str, str], _PendingStreamRead] = {} + + def add(self, run_id: str, read_id: str, read: _PendingStreamRead) -> None: + with self._lock: + self._reads[(run_id, read_id)] = read + + def get(self, run_id: str, read_id: str) -> _PendingStreamRead | None: + with self._lock: + return self._reads.get((run_id, read_id)) + + def remove(self, run_id: str, read_id: str, read: _PendingStreamRead) -> 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._reads.get((run_id, read_id)) is read: + del self._reads[(run_id, read_id)] + + def drop_run(self, run_id: str) -> None: + with self._lock: + for key in [key for key in self._reads if key[0] == run_id]: + del self._reads[key] + + def __len__(self) -> int: + with self._lock: + return len(self._reads) + + +_PENDING_STREAM_READS = _PendingStreamReads() + +_ReadPage = Callable[ + [temporalio.api.streamservice.v1.ReadRecordsRequest], + Awaitable[temporalio.api.streamservice.v1.ReadRecordsResponse], +] + + +@dataclasses.dataclass(frozen=True) +class _ReadUntilResult: + records: list[temporalio.api.stream.v1.StreamRecord] + matched: bool + done: bool + + +async def _read_until( # pyright: ignore[reportUnusedFunction] + read_page: _ReadPage, + *, + ref: temporalio.api.stream.v1.StreamRef, + from_offset: bytes | None, + max_records: int, + timeout: timedelta, + read: _PendingStreamRead, +) -> _ReadUntilResult: + """Read pages until the predicate of ``read`` is true, or a limit stops the read.""" + deadline = time.monotonic() + timeout.total_seconds() + records: list[temporalio.api.stream.v1.StreamRecord] = [] + cursor = from_offset + while len(records) < max_records: + time_left = deadline - time.monotonic() + if time_left <= 0: + break + request = temporalio.api.streamservice.v1.ReadRecordsRequest( + ref=ref, + from_offset=cursor or b"", + max_records=min(max_records - len(records), _MAX_PAGE_RECORDS), + ) + request.wait.FromTimedelta(min(timedelta(seconds=time_left), _MAX_PAGE_WAIT)) + response = await read_page(request) + page = list(response.records)[: max_records - len(records)] + for index, record in enumerate(page): + records.append(record) + if read.until is not None and _call_until(read, record): + # A done page ends with the last Record, so a match on it is the end too. + last = index == len(response.records) - 1 + return _ReadUntilResult(records, True, response.done and last) + if response.done and len(page) == len(response.records): + return _ReadUntilResult(records, False, True) + if page: + cursor = page[-1].offset + return _ReadUntilResult(records, False, False) + + +def _call_until( + read: _PendingStreamRead, record: temporalio.api.stream.v1.StreamRecord +) -> bool: + assert read.until is not None + try: + return bool(read.until(read.to_record(record))) + except Exception as err: + # The Workflow gets this error from its Local Activity, in Workflow code. + raise temporalio.exceptions.ApplicationError( + str(err), type=type(err).__name__, non_retryable=True + ) from err diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index 74ed6961b8..81c5bd7b49 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -17,9 +17,15 @@ import temporalio.converter import temporalio.converter._payload_converter import temporalio.nexus.system._support +import temporalio.exceptions import temporalio.nexus.system.stream_service import temporalio.runtime import temporalio.workflow +from temporalio.client._stream_read import ( + _PENDING_STREAM_READS, + _read_until, + _StreamReadMarker, +) from temporalio.converter import StorageDriverActivityInfo, StorageDriverWorkflowInfo from temporalio.exceptions import StreamConfigError @@ -138,12 +144,15 @@ async def describe_stream( def worker_activities(self) -> Sequence[Callable[..., Any]]: """The Activities that a Worker registers for this Store when it runs Workflows. - 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. + The default is two Activities, named for this Store. One answers + :py:meth:`call_from_workflow`, which runs it as a Local Activity, or as a regular + Activity for a close. The other answers :py:func:`temporalio.workflow.read_stream`, + which always runs it by its name. A Store that overrides + :py:meth:`call_from_workflow` overrides this method to match, and a Store that + overrides this method keeps the read Activity in its list for Workflow code that + reads. """ - return [_store_call_activity(self)] + return [_store_call_activity(self), _store_read_activity(self)] async def call_from_workflow( self, @@ -394,6 +403,80 @@ def _with_local_activity_producer( request.sequence = 1 +_STORE_READ_ACTIVITY = "__temporal_stream_store_read" + + +def _store_read_activity_name(store_name: str) -> str: + return f"{_STORE_READ_ACTIVITY}:{store_name}" + + +def _store_read_activity( + store: StreamStore, +) -> Callable[[str, bytes, int, float], Awaitable[temporalio.common.RawValue]]: + @temporalio.activity.defn(name=_store_read_activity_name(store.name)) + async def read( + read_id: str, request: bytes, max_records: int, timeout_seconds: float + ) -> temporalio.common.RawValue: + info = temporalio.activity.info() + pending = _PENDING_STREAM_READS.get(info.workflow_run_id or "", read_id) + if pending is None: + # The predicate and the Records live in the Workflow's process, so only a + # Local Activity of that Workflow can run this read. + raise temporalio.exceptions.ApplicationError( + f"The Stream read '{read_id}' has no pending read in this process. " + "Only temporalio.workflow.read_stream runs this Activity.", + non_retryable=True, + ) + read_request = temporalio.api.streamservice.v1.ReadRecordsRequest.FromString( + request + ) + # The codec decodes the bodies in each page, so the predicate and the Workflow + # get the user's values. + page_reader = _store_page_reader( + store, temporalio.activity.client().data_converter + ) + from_offset = read_request.from_offset or None + result = await _read_until( + page_reader, + ref=read_request.ref, + from_offset=from_offset, + max_records=max_records, + timeout=timedelta(seconds=timeout_seconds), + read=pending, + ) + marker = _StreamReadMarker.for_records( + read_id=read_id, + store=store.name, + stream_id=read_request.ref.id, + from_offset=from_offset, + records=result.records, + matched=result.matched, + done=result.done, + ) + pending.fill(marker, result.records) + return marker.to_raw_value() + + return read + + +def _store_page_reader( + 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. diff --git a/temporalio/worker/_stream_read.py b/temporalio/worker/_stream_read.py new file mode 100644 index 0000000000..fdc95cfee0 --- /dev/null +++ b/temporalio/worker/_stream_read.py @@ -0,0 +1,95 @@ +"""The Worker side of :py:func:`temporalio.workflow.read_stream`.""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +import temporalio.api.stream.v1 +import temporalio.common +import temporalio.workflow +from temporalio.client._stream_read import ( + _PENDING_STREAM_READS, + _PendingStreamRead, + _StreamReadMarker, +) +from temporalio.client._stream_store import ( + _STREAM_SERVICE, + StreamConfig, + _store_read_activity_name, +) +from temporalio.nexus.system.stream_service.models import ( + ReadRecordsRequest, + StreamRecord, +) +from temporalio.nexus.system.stream_service.workflow._resources.stream import Stream + +from ._stream_store import _store_errors, _workflow_target + +if TYPE_CHECKING: + from ._workflow_instance import _WorkflowInstanceImpl + +# The read Activity gets this much time after the read's own timeout to give its result. +_READ_ACTIVITY_GRACE = timedelta(seconds=10) + + +async def _read_stream_in_workflow( # pyright: ignore[reportUnusedFunction] + instance: _WorkflowInstanceImpl, + config: StreamConfig, + stream: Stream[Any], + *, + until: Callable[[StreamRecord[Any]], bool] | None, + from_offset: bytes | None, + max_records: int, + timeout: timedelta, +) -> temporalio.workflow.StreamRead[Any]: + request = ReadRecordsRequest(ref=stream.ref, from_offset=from_offset) + payload_converter, _ = instance._system_nexus_converters_for( + _STREAM_SERVICE.name, "ReadRecords", request + ) + request_message = payload_converter._to_transfer_value(request) + # The choice runs in Workflow code, as it does for the other Stream calls. + store = config._store_for(stream.ref, _workflow_target()) + 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) + + run_id = temporalio.workflow.info().run_id + read_id = str(temporalio.workflow.uuid4()) + pending = _PendingStreamRead(until=until, to_record=to_record) + _PENDING_STREAM_READS.add(run_id, read_id, pending) + try: + with _store_errors(): + result = await temporalio.workflow.execute_local_activity( + _store_read_activity_name(store.name), + args=[ + read_id, + request_message.SerializeToString(), + max_records, + timeout.total_seconds(), + ], + result_type=temporalio.common.RawValue, + start_to_close_timeout=timeout + _READ_ACTIVITY_GRACE, + retry_policy=config.workflow_call_options.retry_policy, + ) + marker = _StreamReadMarker.from_payload(result.payload) + if marker is None or marker.read_id != read_id: + raise RuntimeError( + f"The read of the Stream '{stream.ref.id}' got a result that is not " + "its read marker" + ) + records = pending.records_for(marker) + if records is None: + raise RuntimeError( + f"The Worker has no Records for the read of {marker.describe_range()}" + ) + return temporalio.workflow.StreamRead( + records=[to_record(record) for record in records], + last_offset=marker.last_offset, + matched=marker.matched, + done=marker.done, + ) + finally: + _PENDING_STREAM_READS.remove(run_id, read_id, pending) diff --git a/temporalio/worker/_stream_store.py b/temporalio/worker/_stream_store.py index 947cc7aeba..1ba16abfa9 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 @@ -24,18 +25,48 @@ from ._workflow_instance import _WorkflowInstanceImpl +_NO_STREAM_STORE_MESSAGE = ( + "The client of this Worker has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...])." +) + + def _no_stream_config_error() -> temporalio.exceptions.ApplicationError: # pyright: ignore[reportUnusedFunction] # A `StreamConfigError` would fail the Workflow Task, and its retries would fail the # same way for good. A Worker with no config is a setup that a user has to see, so # the Workflow fails instead. return temporalio.exceptions.ApplicationError( - "The client of this Worker has no Stream Store. Give one in " - "stream_config=StreamConfig(stores=[...]).", - type="StreamConfigError", - non_retryable=True, + _NO_STREAM_STORE_MESSAGE, type="StreamConfigError", non_retryable=True ) +@contextlib.contextmanager +def _store_errors() -> Iterator[None]: + """Give Workflow code the Store's error of a Local 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 +111,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/_workflow.py b/temporalio/worker/_workflow.py index 3e20e6adbb..5acf27df39 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_READS from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo from temporalio.worker.workflow_sandbox._runner import SandboxedWorkflowRunner @@ -678,6 +679,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_READS.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 07f63a69dc..1176497689 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -61,6 +61,8 @@ import temporalio.workflow from temporalio.client._stream_store import _STREAM_SERVICE from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo +from temporalio.nexus.system.stream_service.models import StreamRecord +from temporalio.nexus.system.stream_service.workflow._resources.stream import Stream from temporalio.nexus.system.stream_service.workflow._system_nexus_interceptor import ( _start_system_nexus_operation as _start_stream_system_nexus_operation, ) @@ -95,7 +97,12 @@ WorkflowInboundInterceptor, WorkflowOutboundInterceptor, ) -from ._stream_store import _no_stream_config_error, _WorkflowStreamStore +from ._stream_read import _read_stream_in_workflow +from ._stream_store import ( + _NO_STREAM_STORE_MESSAGE, + _no_stream_config_error, + _WorkflowStreamStore, +) logger = logging.getLogger(__name__) @@ -1994,6 +2001,27 @@ def workflow_upsert_search_attributes( else [update.value] ) + async def workflow_read_stream( + self, + stream: Stream[Any], + *, + until: Callable[[StreamRecord[Any]], bool] | None, + from_offset: bytes | None, + max_records: int, + timeout: timedelta, + ) -> temporalio.workflow.StreamRead[Any]: + if self._stream_store is None: + raise temporalio.exceptions.StreamConfigError(_NO_STREAM_STORE_MESSAGE) + return await _read_stream_in_workflow( + self, + self._stream_store.config, + stream, + until=until, + from_offset=from_offset, + max_records=max_records, + timeout=timeout, + ) + async def workflow_sleep( self, duration: float, @@ -2368,11 +2396,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/__init__.py b/temporalio/workflow/__init__.py index 60b02cc543..748d5b2502 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -155,6 +155,11 @@ logger, unsafe, ) +from ._streams import ( + DEFAULT_READ_STREAM_MAX_RECORDS, + StreamRead, + read_stream, +) from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowConfig, @@ -262,6 +267,9 @@ "SandboxImportNotificationPolicy", "logger", "unsafe", + "DEFAULT_READ_STREAM_MAX_RECORDS", + "StreamRead", + "read_stream", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", diff --git a/temporalio/workflow/_context.py b/temporalio/workflow/_context.py index a33759c37a..1a051c6ab3 100644 --- a/temporalio/workflow/_context.py +++ b/temporalio/workflow/_context.py @@ -21,10 +21,16 @@ from ._exceptions import _NotInWorkflowEventLoopError if TYPE_CHECKING: + from temporalio.nexus.system.stream_service.models import StreamRecord + from temporalio.nexus.system.stream_service.workflow._resources.stream import ( + Stream, + ) + from ._activities import ActivityCancellationType, ActivityHandle from ._event_groups import EventGroup from ._exceptions import ContinueAsNewVersioningBehavior, VersioningIntent from ._nexus import NexusOperationCancellationType, NexusOperationHandle + from ._streams import StreamRead from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowHandle, @@ -399,6 +405,17 @@ def workflow_payload_converter(self) -> temporalio.converter.PayloadConverter: . @abstractmethod def workflow_random(self) -> Random: ... + @abstractmethod + async def workflow_read_stream( + self, + stream: Stream[Any], + *, + until: Callable[[StreamRecord[Any]], bool] | None, + from_offset: bytes | None, + max_records: int, + timeout: timedelta, + ) -> StreamRead[Any]: ... + @abstractmethod def workflow_set_query_handler( self, name: str | None, handler: Callable | None diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py new file mode 100644 index 0000000000..be241e9c75 --- /dev/null +++ b/temporalio/workflow/_streams.py @@ -0,0 +1,116 @@ +"""A Workflow read of a Stream that waits for a predicate. + +.. warning:: + This API is an experimental prototype and subject to change. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import TYPE_CHECKING, Generic, TypeVar + +from ._context import _Runtime + +if TYPE_CHECKING: + # The generated Stream modules import `temporalio.workflow`, so this module names + # their types for the type checker only. + from temporalio.nexus.system.stream_service.models import StreamRecord + from temporalio.nexus.system.stream_service.workflow._resources.stream import ( + Stream, + ) + +__all__ = ["StreamRead", "read_stream"] + +BodyT = TypeVar("BodyT") + +DEFAULT_READ_STREAM_MAX_RECORDS = 10_000 +"""The most Records that one :py:func:`read_stream` gives when ``max_records`` is ``None``.""" + + +@dataclass(frozen=True) +class StreamRead(Generic[BodyT]): + """What one :py:func:`read_stream` gave. + + .. warning:: + This API is an experimental prototype and subject to change. + """ + + records: Sequence[StreamRecord[BodyT]] + """The Records of the read in order, typed as a client's read gives them.""" + + last_offset: bytes | None + """The offset of the last Record, or the ``from_offset`` of a read with none. + + Give it as the ``from_offset`` of the next read. + """ + + matched: bool + """Whether ``until`` was true for the last Record.""" + + done: bool + """Whether the Stream is closed and its last Record is in :py:attr:`records`.""" + + +async def read_stream( + stream: Stream[BodyT], + *, + until: Callable[[StreamRecord[BodyT]], bool] | None = None, + from_offset: bytes | None = None, + max_records: int | None = None, + timeout: timedelta, +) -> StreamRead[BodyT]: + """Read a Stream until a predicate is true, in one Local Activity. + + The read starts after ``from_offset`` and stops at the first of these. ``until`` + is true for a Record, ``max_records`` Records were read, the Stream is closed and + its last Record was read, or ``timeout`` passed. A Store's Local Activity reads + the pages, and History gets one marker for the whole read. The marker has the + range of offsets and no Record. A replay reads that range from the Store again + before the Workflow sees the result, so the Records must still be in the Store. + + Args: + stream: The Workflow ``Stream`` to read, from + :py:func:`temporalio.workflow.get_stream` or + :py:func:`temporalio.workflow.create_stream`. + until: Called with each typed Record in order. The read stops after the first + Record for which it is true, and Records after it stay for the next read. + It runs in the Worker process, on the Worker's event loop, while the Local + Activity reads. It runs outside the Workflow's own thread. So it must not + change any state, and it can't call a function of ``temporalio.workflow``. + It can read the Workflow's state through a reference such as ``self``. + When it sees a change of that state is not defined. A replay does not call + it, since the replay takes the end of the read from History. ``None`` + reads until another limit stops the read. + from_offset: Read after this offset. ``None`` is the beginning. + max_records: The most Records to read. ``None`` is + :py:data:`DEFAULT_READ_STREAM_MAX_RECORDS`. + timeout: How long the read waits for a match. It must be positive. + + Returns: + The Records of the read, the offset to read after next, and why it stopped. + + Raises: + temporalio.exceptions.StreamError: The Store failed the read. + temporalio.exceptions.ApplicationError: ``until`` raised. The error has the + message of the exception. + temporalio.exceptions.StreamConfigError: The Worker's client has no Stream + config. This fails the Workflow Task. + + .. warning:: + This API is an experimental prototype and subject to change. + """ + if timeout <= timedelta(): + raise ValueError("timeout must be positive") + if max_records is None: + max_records = DEFAULT_READ_STREAM_MAX_RECORDS + elif max_records < 1: + raise ValueError("max_records must be at least 1") + return await _Runtime.current().workflow_read_stream( + stream, + until=until, + from_offset=from_offset, + max_records=max_records, + timeout=timeout, + ) From a5d9ad6c7e154d4c5fa5ac2b90307f8924a60a94 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 16:58:04 -0700 Subject: [PATCH 02/20] Filled in the Records of a Workflow read before a replay. History keeps only the marker of a read, so a replay reads its range from the Store again. The Worker checks the count, the last offset and a digest before the Workflow sees the Records. --- temporalio/client/_stream_read.py | 42 +++++++++++++++++++++++++++ temporalio/worker/_replayer.py | 5 +++- temporalio/worker/_stream_read.py | 48 ++++++++++++++++++++++++++++++- temporalio/worker/_workflow.py | 7 +++++ 4 files changed, 100 insertions(+), 2 deletions(-) diff --git a/temporalio/client/_stream_read.py b/temporalio/client/_stream_read.py index 43e5564395..b96aca6146 100644 --- a/temporalio/client/_stream_read.py +++ b/temporalio/client/_stream_read.py @@ -286,3 +286,45 @@ def _call_until( raise temporalio.exceptions.ApplicationError( str(err), type=type(err).__name__, non_retryable=True ) from err + + +async def _read_range( # pyright: ignore[reportUnusedFunction] + read_page: _ReadPage, marker: _StreamReadMarker +) -> list[temporalio.api.stream.v1.StreamRecord]: + """Read the Records of ``marker`` again, and check them against it. + + Raises: + RuntimeError: The Store failed, the Records are gone, or they don't match. + """ + records: list[temporalio.api.stream.v1.StreamRecord] = [] + cursor = marker.from_offset + ref = temporalio.api.stream.v1.StreamRef(id=marker.stream_id, store=marker.store) + while len(records) < marker.record_count: + left = marker.record_count - len(records) + request = temporalio.api.streamservice.v1.ReadRecordsRequest( + ref=ref, + from_offset=cursor or b"", + max_records=min(left, _MAX_PAGE_RECORDS), + ) + try: + response = await read_page(request) + except Exception as err: + raise RuntimeError( + f"The replay of a Workflow read can't read {marker.describe_range()}: " + f"{err}" + ) from err + if not response.records: + raise RuntimeError( + f"The replay of a Workflow read can't read {marker.describe_range()}: " + f"the Store gave {len(records)} Records, and the Records after them are " + "gone" + ) + records.extend(list(response.records)[:left]) + cursor = records[-1].offset + last_offset = records[-1].offset if records else marker.from_offset + if last_offset != marker.last_offset or _records_digest(records) != marker.digest: + raise RuntimeError( + f"The replay of a Workflow read read {marker.describe_range()}, and the " + "Records differ from the ones of the first run" + ) + return records diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 7ad6bdeaaf..13ced88af0 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -70,7 +70,10 @@ 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 or a describe. For each + :py:func:`temporalio.workflow.read_stream` in the history, it reads the + Records of that read from the Store again, so the replay needs the same + ``data_converter`` and Stores that still have them. 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 index fdc95cfee0..5440448715 100644 --- a/temporalio/worker/_stream_read.py +++ b/temporalio/worker/_stream_read.py @@ -1,4 +1,7 @@ -"""The Worker side of :py:func:`temporalio.workflow.read_stream`.""" +"""The Worker side of :py:func:`temporalio.workflow.read_stream`. + +Workflow code starts the read here, and the Worker gives a replayed read its Records. +""" from __future__ import annotations @@ -7,16 +10,21 @@ from typing import TYPE_CHECKING, Any import temporalio.api.stream.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_READS, _PendingStreamRead, + _read_range, _StreamReadMarker, ) from temporalio.client._stream_store import ( _STREAM_SERVICE, StreamConfig, + _store_page_reader, _store_read_activity_name, ) from temporalio.nexus.system.stream_service.models import ( @@ -93,3 +101,41 @@ def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: ) finally: _PENDING_STREAM_READS.remove(run_id, read_id, pending) + + +async def _fill_stream_reads( # pyright: ignore[reportUnusedFunction] + act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, + config: StreamConfig, + data_converter: temporalio.converter.DataConverter, +) -> None: + """Give each read that this activation resolves its Records, before the Workflow runs. + + History has only the marker of a read. On a replay, the Records come from the Store + again, through the same codec as the read Activity used. + + Raises: + RuntimeError: The Store failed, the Records are gone, or they don't match the + marker. The activation fails, and so does the Workflow Task. + """ + 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 + marker = _StreamReadMarker.from_payload(resolve.result.completed.result) + if marker is None: + continue + pending = _PENDING_STREAM_READS.get(act.run_id, marker.read_id) + # On a first run the read Activity already gave its Records. + if pending is None or pending.records_for(marker) is not None: + continue + try: + store = config._get_store_by_name(marker.store) + except temporalio.exceptions.StreamConfigError as err: + raise RuntimeError( + f"The replay of a Workflow read can't read {marker.describe_range()}: " + f"{err}" + ) from err + page_reader = _store_page_reader(store, data_converter) + pending.fill(marker, await _read_range(page_reader, marker)) diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index 5acf27df39..f7de6f791a 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -41,6 +41,7 @@ WorkflowInboundInterceptor, WorkflowInterceptorClassInput, ) +from ._stream_read import _fill_stream_reads from ._stream_store import _WorkflowStreamStore from ._workflow_instance import ( _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS, @@ -433,6 +434,12 @@ async def _handle_activation( await self._decode_nexus_previous_run_outcome( act.run_id, init_job, previous_run_outcome, download_metrics ) + if self._stream_store is not None: + # A replay has only the marker of a Stream read, so the Records come + # from the Store before the Workflow sees the result. + await _fill_stream_reads( + act, self._stream_store.config, self._data_converter + ) if not workflow: assert init_job workflow = _RunningWorkflow( From f216560a33f16345ddb206d31445d69afd572aae Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 16:59:16 -0700 Subject: [PATCH 03/20] Tested and measured the Workflow read. The tests cover a first run, replays on a Worker and with the Replayer, and the limits of a read. The measurements show the History events and the Workflow Tasks that one read costs. --- tests/helpers/streams.py | 12 + tests/nexus/test_stream_read_workflow.py | 760 +++++++++++++++++++++++ 2 files changed, 772 insertions(+) create mode 100644 tests/nexus/test_stream_read_workflow.py diff --git a/tests/helpers/streams.py b/tests/helpers/streams.py index beb4eaaa45..b9606345d1 100644 --- a/tests/helpers/streams.py +++ b/tests/helpers/streams.py @@ -184,3 +184,15 @@ 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 Workflow read 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 + fail_next_task = False + results: dict[str, tuple[object, ...]] = {} diff --git a/tests/nexus/test_stream_read_workflow.py b/tests/nexus/test_stream_read_workflow.py new file mode 100644 index 0000000000..45c1ca09bf --- /dev/null +++ b/tests/nexus/test_stream_read_workflow.py @@ -0,0 +1,760 @@ +"""Tests of a Workflow read of a Stream that waits for a predicate.""" + +from __future__ import annotations + +import asyncio +import collections +import time +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any + +import pytest + +from temporalio import workflow +from temporalio.api.enums.v1 import EventType +from temporalio.api.streamservice.v1 import ReadRecordsRequest, ReadRecordsResponse +from temporalio.client import ( + Client, + StreamConfig, + WorkflowFailureError, + WorkflowHandle, + WorkflowHistory, +) +from temporalio.converter import DataConverter +from temporalio.exceptions import ApplicationError, CancelledError +from temporalio.nexus.system.stream_service import 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, assert_task_fail_eventually, new_worker + from tests.helpers.streams import RecordingStreamStore, StreamReadState + from tests.test_converter import SimpleCodec + +# Every body carries this text, so a search of History shows whether a body leaks. +_BODY_MARK = "secret-body" + + +@dataclass +class Token: + kind: str + text: str + + +@dataclass +class ReadArgs: + stream_id: str + until_kind: str | None = "done" + timeout_seconds: float = 10.0 + max_records: int | None = None + + +@dataclass +class ReadResult: + texts: list[str] + matched: bool + done: bool + last_offset: str | None + typed: bool + + +def _result(read: workflow.StreamRead[Token]) -> ReadResult: + return ReadResult( + texts=[record.body.text for record in read.records], + matched=read.matched, + 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: ReadResult) -> ReadResult: + # A replay gives no result to the test, so the Workflow keeps it here. + StreamReadState.results[workflow.info().run_id] = ( + tuple(result.texts), + result.matched, + result.done, + result.last_offset, + ) + return result + + +def _stream(stream_id: str) -> WorkflowStream[Token]: + return workflow.get_stream(StreamRef(id=stream_id), body_type=Token) + + +@workflow.defn +class ReadUntilKindWorkflow: + @workflow.run + async def run(self, args: ReadArgs) -> ReadResult: + 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 + + read = await workflow.read_stream( + _stream(args.stream_id), + until=until if until_kind is not None else None, + max_records=args.max_records, + timeout=timedelta(seconds=args.timeout_seconds), + ) + return _keep(_result(read)) + + +@workflow.defn +class ReadUntilWantedWorkflow: + 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) -> ReadResult: + read = await workflow.read_stream( + _stream(stream_id), + until=lambda record: record.body.text == self.wanted, + timeout=timedelta(seconds=20), + ) + return _result(read) + + +@workflow.defn +class TwoReadsWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[str]: + stream = _stream(stream_id) + first = await workflow.read_stream( + stream, + until=lambda record: record.body.text.endswith("m1"), + timeout=timedelta(seconds=10), + ) + await workflow.sleep(timedelta(milliseconds=1)) + second = await workflow.read_stream( + stream, + until=lambda record: record.body.text.endswith("m2"), + from_offset=first.last_offset, + timeout=timedelta(seconds=10), + ) + return [record.body.text for record in [*first.records, *second.records]] + + +@workflow.defn +class FailTaskOnceAfterReadWorkflow: + @workflow.run + async def run(self, stream_id: str) -> ReadResult: + read = await workflow.read_stream( + _stream(stream_id), + until=lambda record: record.body.kind == "done", + timeout=timedelta(seconds=10), + ) + # The failure comes in a later Workflow Task, so the retry replays the read. + await workflow.sleep(timedelta(milliseconds=1)) + if StreamReadState.fail_next_task: + StreamReadState.fail_next_task = False + raise RuntimeError("Failed the Workflow Task once") + return _result(read) + + +@workflow.defn +class TimeoutThenReadWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[ReadResult]: + stream = _stream(stream_id) + first = await workflow.read_stream( + stream, + until=lambda record: record.body.kind == "done", + timeout=timedelta(seconds=1), + ) + second = await workflow.read_stream( + stream, + until=lambda record: record.body.kind == "done", + from_offset=first.last_offset, + timeout=timedelta(seconds=10), + ) + return [_result(first), _result(second)] + + +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 workflow.read_stream( + _stream(stream_id), + until=_PREDICATES[predicate], + timeout=timedelta(seconds=10), + ) + except ApplicationError as err: + return f"{err.type}: {err.message}" + return "no error" + + +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 + + +def _tokens(count: int, *, prefix: str = "t") -> list[Token]: + return [Token("text", f"{_BODY_MARK}-{prefix}{index}") for index in range(count)] + + +def _done(text: str = "done") -> Token: + return Token("done", f"{_BODY_MARK}-{text}") + + +async def _new_stream( + client: Client, + store: Any = None, + data_converter: DataConverter | None = None, +) -> tuple[Client, Any, str]: + store = store or MemoryStreamStore() + config = client.config() + config["stream_config"] = StreamConfig(stores=[store]) + if data_converter is not None: + config["data_converter"] = data_converter + stream_client = Client(**config) + 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 _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: list[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 _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())) + + +def _assert_no_bodies(history: WorkflowHistory) -> None: + # The Workflow returns the bodies, so only its completion event has them. + events = [ + event + for event in history.events + if event.event_type != EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED + ] + assert len(events) == len(history.events) - 1 + # The JSON has payload data in base64, so the raw events are searched too. + assert _BODY_MARK not in WorkflowHistory(history.workflow_id, events).to_json() + for event in events: + assert _BODY_MARK.encode() not in event.SerializeToString() + + +async def _start_read( + stream_client: Client, + args: ReadArgs, + task_queue: str, + **kwargs: Any, +) -> WorkflowHandle[ReadUntilKindWorkflow, ReadResult]: + return await stream_client.start_workflow( + ReadUntilKindWorkflow.run, + args, + id=f"workflow-{uuid.uuid4()}", + task_queue=task_queue, + **kwargs, + ) + + +_RUNNERS = pytest.mark.parametrize( + "runner", + [SandboxedWorkflowRunner(), UnsandboxedWorkflowRunner()], + ids=["sandboxed", "unsandboxed"], +) + + +async def _run_match( + client: Client, + runner: WorkflowRunner = SandboxedWorkflowRunner(), + data_converter: DataConverter | None = None, +) -> tuple[Client, Any, WorkflowHandle[Any, ReadResult], ReadResult, list[Token]]: + stream_client, store, stream_id = await _new_stream( + client, data_converter=data_converter + ) + tokens = [*_tokens(4), _done()] + await _append(stream_client, stream_id, *tokens) + async with new_worker( + stream_client, ReadUntilKindWorkflow, workflow_runner=runner + ) as worker: + handle = await _start_read( + stream_client, ReadArgs(stream_id), worker.task_queue + ) + return stream_client, store, handle, await handle.result(), tokens + + +@_RUNNERS +async def test_read_until_match(client: Client, runner: WorkflowRunner): + _, _, handle, result, tokens = await _run_match(client, runner) + assert result == ReadResult( + texts=_texts(tokens), + matched=True, + done=False, + last_offset=result.last_offset, + typed=True, + ) + assert result.last_offset + history = await handle.fetch_history() + assert _markers(history) == ["core_local_activity"] + _assert_no_bodies(history) + + +async def test_read_waits_for_records(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = [*_tokens(4), _done()] + async with new_worker(stream_client, ReadUntilKindWorkflow) as worker: + handle = await _start_read( + stream_client, ReadArgs(stream_id), worker.task_queue + ) + await _append_over(stream_client, stream_id, tokens, 1.0) + result = await handle.result() + assert result.texts == _texts(tokens) + assert result.matched + 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 +async def test_read_predicate_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, ReadUntilWantedWorkflow, workflow_runner=runner + ) as worker: + handle = await stream_client.start_workflow( + ReadUntilWantedWorkflow.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(ReadUntilWantedWorkflow.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.matched + + +async def test_read_replays_with_replayer(client: Client): + _, store, handle, result, _ = await _run_match(client) + run_id = handle.result_run_id + assert run_id + first_run = StreamReadState.results.pop(run_id) + replay_store = RecordingStreamStore(store) + StreamReadState.until_calls = 0 + await Replayer( + workflows=[ReadUntilKindWorkflow], + stream_config=StreamConfig(stores=[replay_store]), + ).replay_workflow(await handle.fetch_history()) + assert StreamReadState.results.pop(run_id) == first_run + assert first_run[0] == tuple(result.texts) + assert replay_store.operations == ["ReadRecords"] + assert StreamReadState.until_calls == 0 + + +async def test_read_replays_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) + async with new_worker( + stream_client, TwoReadsWorkflow, max_cached_workflows=0 + ) as worker: + handle = await stream_client.start_workflow( + TwoReadsWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == _texts(tokens[:4]) + history = await handle.fetch_history() + assert _markers(history) == ["core_local_activity", "core_local_activity"] + # Two reads, plus a replay of the first read before the Workflow Task of the second. + reads = store.operations.count("ReadRecords") + print(f"\nWorker without cache: {reads} ReadRecords calls") + assert reads >= 3 + + +async def test_read_replay_fails_when_records_are_gone(client: Client): + _, _, handle, _, _ = await _run_match(client) + history = await handle.fetch_history() + with pytest.raises(Exception) as err: + await Replayer( + workflows=[ReadUntilKindWorkflow], + stream_config=StreamConfig(stores=[MemoryStreamStore()]), + ).replay_workflow(history) + stream_id = history.events[0].workflow_execution_started_event_attributes + assert "can't read the Stream" in str(err.value) + assert f"'{_stream_id_of(stream_id)}'" in str(err.value) + + +def _stream_id_of(started: Any) -> str: + return DataConverter.default.payload_converter.from_payloads( + started.input.payloads, [ReadArgs] + )[0].stream_id + + +async def test_read_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 + async with new_worker(stream_client, FailTaskOnceAfterReadWorkflow) as worker: + handle = await stream_client.start_workflow( + FailTaskOnceAfterReadWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + result = await handle.result() + assert result.texts == _texts(tokens) + assert result.matched + counts = _event_counts(await handle.fetch_history()) + assert counts["WORKFLOW_TASK_FAILED"] == 1 + assert counts["MARKER_RECORDED"] == 1 + + +async def test_read_timeout_gives_records_so_far(client: Client): + stream_client, _, stream_id = await _new_stream(client) + first = _tokens(3) + rest = [*_tokens(2, prefix="rest"), _done()] + await _append(stream_client, stream_id, *first) + async with new_worker(stream_client, TimeoutThenReadWorkflow) as worker: + handle = await stream_client.start_workflow( + TimeoutThenReadWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await asyncio.sleep(1.5) + await _append(stream_client, stream_id, *rest) + timed_out, matched = await handle.result() + assert timed_out.texts == _texts(first) + assert not timed_out.matched and not timed_out.done + assert matched.texts == _texts(rest) + assert matched.matched + + +async def test_read_stops_at_max_records(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = _tokens(10) + await _append(stream_client, stream_id, *tokens) + async with new_worker(stream_client, ReadUntilKindWorkflow) as worker: + handle = await _start_read( + stream_client, + ReadArgs(stream_id, until_kind=None, max_records=3), + worker.task_queue, + ) + result = await handle.result() + assert result.texts == _texts(tokens[:3]) + assert not result.matched and not result.done + + +async def test_read_of_closed_stream_is_done(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = _tokens(4) + await _append(stream_client, stream_id, *tokens) + await stream_client.get_stream(StreamRef(id=stream_id), body_type=Token).close() + async with new_worker(stream_client, ReadUntilKindWorkflow) as worker: + handle = await _start_read( + stream_client, ReadArgs(stream_id), worker.task_queue + ) + result = await handle.result() + assert result.texts == _texts(tokens) + assert result.done and not result.matched + + +async def test_read_with_payload_codec(client: Client): + data_converter = DataConverter(payload_codec=SimpleCodec()) + StreamReadState.until_calls = 0 + _, store, handle, result, tokens = await _run_match( + client, data_converter=data_converter + ) + assert result.texts == _texts(tokens) + assert result.typed and result.matched + # The predicate saw each Record decoded, or its type check would have raised. + assert StreamReadState.until_calls == len(tokens) + run_id = handle.result_run_id + assert run_id + first_run = StreamReadState.results.pop(run_id) + history = await handle.fetch_history() + marker = next( + event.marker_recorded_event_attributes + for event in history.events + if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED + ) + assert marker.details["result"].payloads[0].metadata["simple-codec"] == b"true" + await Replayer( + workflows=[ReadUntilKindWorkflow], + data_converter=data_converter, + stream_config=StreamConfig(stores=[store]), + ).replay_workflow(history) + assert StreamReadState.results.pop(run_id) == first_run + + +async def _cancel_waiting_read(client: Client, task_timeout: timedelta | None) -> float: + """Cancel a Workflow while its read waits, and give the seconds until it ends.""" + stream_client, store, stream_id = await _new_stream(client, TrackingStore()) + async with new_worker(stream_client, ReadUntilKindWorkflow) as worker: + handle = await _start_read( + stream_client, + ReadArgs(stream_id, 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_read_cancel(client: Client): + assert await _cancel_waiting_read(client, None) < 5 + + +async def test_read_cancel_with_short_task_timeout(client: Client): + assert await _cancel_waiting_read(client, timedelta(seconds=2)) < 5 + + +async def test_read_predicate_that_raises(client: Client): + assert await _failing_predicate(client, "raise") == ( + f"ValueError: bad token {_BODY_MARK}-t0" + ) + + +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)) + async with new_worker( + stream_client, FailingPredicateWorkflow, workflow_runner=runner + ) as worker: + return await stream_client.execute_workflow( + FailingPredicateWorkflow.run, + args=[stream_id, predicate], + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + + +async def test_read_without_stream_config_fails_the_task(client: Client): + async with new_worker(client, ReadUntilKindWorkflow) as worker: + handle = await _start_read( + client, + ReadArgs(f"stream-{uuid.uuid4()}"), + worker.task_queue, + execution_timeout=timedelta(seconds=30), + ) + await assert_task_fail_eventually( + handle, + message_contains=( + "The client of this Worker has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...])." + ), + ) + await handle.terminate() + + +@_RUNNERS +async def test_read_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_read_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 _run_flowing_read( + client: Client, *, task_timeout: timedelta | None = None, 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, ReadUntilKindWorkflow) as worker: + handle = await _start_read( + stream_client, + ReadArgs(stream_id, timeout_seconds=30), + worker.task_queue, + task_timeout=task_timeout, + ) + 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_read(client: Client): + history, _ = await _run_flowing_read(client) + print(f"\nA. 40 Records over 2 s, History events: {_event_counts(history)}") + + +async def test_measure_idle_wait(client: Client): + stream_client, _, stream_id = await _new_stream(client) + async with new_worker(stream_client, ReadUntilKindWorkflow) as worker: + handle = await _start_read( + stream_client, + ReadArgs(stream_id, 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"\nB. Idle wait of 6 s, History events: {counts}") + print(f"B. Workflow Tasks: {counts['WORKFLOW_TASK_COMPLETED']} completed") + + +async def test_measure_match_latency(client: Client): + latencies = [] + for _ in range(5): + _, latency = await _run_flowing_read(client) + latencies.append(latency) + print( + "\nC. Append of the match to the end of the Workflow: " + + ", ".join(f"{latency * 1000:.0f} ms" for latency in latencies) + ) From 609c2928e33596f9bed0a243ad943542114efb6f Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 17:28:24 -0700 Subject: [PATCH 04/20] Gave the Workflow read a stop reason and added a fold over its Records. The API review asked for an idle stop, a byte cap and one reason for the end of each read. A fold keeps its value in the marker, so its replay needs no Store and calls no function of the Workflow. --- temporalio/client/_stream_read.py | 303 ++++++++----- temporalio/client/_stream_store.py | 59 ++- temporalio/worker/_replayer.py | 4 +- temporalio/worker/_stream_read.py | 90 +++- temporalio/worker/_workflow_instance.py | 31 +- temporalio/workflow/__init__.py | 8 + temporalio/workflow/_context.py | 18 +- temporalio/workflow/_streams.py | 248 +++++++++-- tests/helpers/streams.py | 1 + tests/nexus/test_stream_read_workflow.py | 527 ++++++++++++++++++++++- 10 files changed, 1092 insertions(+), 197 deletions(-) diff --git a/temporalio/client/_stream_read.py b/temporalio/client/_stream_read.py index b96aca6146..d0b36bdb6e 100644 --- a/temporalio/client/_stream_read.py +++ b/temporalio/client/_stream_read.py @@ -1,10 +1,10 @@ -"""The process side of a Stream read that Workflow code waits on. +"""The process side of a Stream read or fold that Workflow code waits on. -A Workflow's read runs in a Local Activity, and the Local Activity calls a predicate -of the Workflow. A predicate is a closure and can't be serialized, and the Records -must stay out of History. So both pass through the registry of this module, in the -memory of the Worker process. History keeps only a small marker with the range of -offsets of the read. +A Workflow's read runs in a Local Activity, and the Local Activity calls functions of +the Workflow. A function is a closure and can't be serialized, and the Records must +stay out of History. So both pass through the registry of this module, in the memory +of the Worker process. History keeps only a small marker with the range of offsets of +the read, and the value of a fold. .. warning:: This API is an experimental prototype and subject to change. @@ -28,37 +28,50 @@ import temporalio.api.streamservice.v1 import temporalio.common import temporalio.exceptions +from temporalio.workflow._streams import StreamReadStop _MARKER_METADATA_KEY = "temporal-stream-read" -_MARKER_METADATA_VALUE = b"v1" +_MARKER_METADATA_VALUE = b"v2" _MAX_PAGE_RECORDS = 1000 # A Store waits at most this long for one page. -_MAX_PAGE_WAIT = timedelta(seconds=60) +_MAX_PAGE_WAIT = 60.0 _DIGEST_SIZE = 16 -def _encode_offset(offset: bytes | None) -> str | None: - return None if offset is None else base64.b64encode(offset).decode() +def _encode_bytes(value: bytes | None) -> str | None: + return None if value is None else base64.b64encode(value).decode() -def _decode_offset(offset: str | None) -> bytes | None: - return None if offset is None else base64.b64decode(offset) +def _decode_bytes(value: str | None) -> bytes | None: + return None if value is None else base64.b64decode(value) -def _records_digest( - records: Sequence[temporalio.api.stream.v1.StreamRecord], -) -> bytes: - """Give the first bytes of a SHA-256 over the offsets and bodies of ``records``. +class _RecordDigest: + """A SHA-256 over the offsets and bodies of Records, in the order that they come. A length before each part keeps two different splits of the same bytes apart. """ - digest = hashlib.sha256() - for record in records: + + def __init__(self) -> None: + self._hash = hashlib.sha256() + + def update(self, record: temporalio.api.stream.v1.StreamRecord) -> None: body = record.body.SerializeToString(deterministic=True) for part in (record.offset, body): - digest.update(struct.pack(">Q", len(part))) - digest.update(part) - return digest.digest()[:_DIGEST_SIZE] + 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) @@ -71,34 +84,11 @@ class _StreamReadMarker: from_offset: bytes | None last_offset: bytes | None record_count: int - matched: bool + stop: StreamReadStop done: bool digest: bytes - - @classmethod - def for_records( - cls, - *, - read_id: str, - store: str, - stream_id: str, - from_offset: bytes | None, - records: Sequence[temporalio.api.stream.v1.StreamRecord], - matched: bool, - done: bool, - ) -> _StreamReadMarker: - """Build the marker of a read that gave ``records``.""" - return cls( - read_id=read_id, - store=store, - stream_id=stream_id, - from_offset=from_offset, - last_offset=records[-1].offset if records else from_offset, - record_count=len(records), - matched=matched, - done=done, - digest=_records_digest(records), - ) + value: temporalio.api.common.v1.Payload | None = None + """The value of a fold. A read of Records has none.""" def to_raw_value(self) -> temporalio.common.RawValue: """Give the marker as a payload that a Worker can find in any activation.""" @@ -106,12 +96,17 @@ def to_raw_value(self) -> temporalio.common.RawValue: "read_id": self.read_id, "store": self.store, "stream_id": self.stream_id, - "from_offset": _encode_offset(self.from_offset), - "last_offset": _encode_offset(self.last_offset), + "from_offset": _encode_bytes(self.from_offset), + "last_offset": _encode_bytes(self.last_offset), "record_count": self.record_count, - "matched": self.matched, + "stop": self.stop.value, "done": self.done, "digest": self.digest.hex(), + "value": _encode_bytes( + None + if self.value is None + else self.value.SerializeToString(deterministic=True) + ), } return temporalio.common.RawValue( temporalio.api.common.v1.Payload( @@ -131,16 +126,20 @@ def from_payload( if payload.metadata.get(_MARKER_METADATA_KEY) != _MARKER_METADATA_VALUE: return None data = json.loads(payload.data) + value = _decode_bytes(data["value"]) return cls( read_id=data["read_id"], store=data["store"], stream_id=data["stream_id"], - from_offset=_decode_offset(data["from_offset"]), - last_offset=_decode_offset(data["last_offset"]), + from_offset=_decode_bytes(data["from_offset"]), + last_offset=_decode_bytes(data["last_offset"]), record_count=data["record_count"], - matched=data["matched"], + stop=StreamReadStop(data["stop"]), done=data["done"], digest=bytes.fromhex(data["digest"]), + value=None + if value is None + else temporalio.api.common.v1.Payload.FromString(value), ) def describe_range(self) -> str: @@ -153,16 +152,30 @@ def describe_range(self) -> str: ) +@dataclasses.dataclass(frozen=True) +class _StreamFold: + """The function and the start value of a fold.""" + + function: Callable[[Any, Any], Any] + initial: Any + + @dataclasses.dataclass(eq=False) class _PendingStreamRead: """One read of Workflow code, from the start of its Local Activity to its result.""" - until: Callable[[Any], bool] | None - """The predicate of the Workflow, or ``None``.""" - to_record: Callable[[temporalio.api.stream.v1.StreamRecord], Any] """Turns a proto Record into the typed Record of the Workflow.""" + until: Callable[..., bool] | None + """The predicate of the Workflow, or ``None``. + + A read calls it with the typed Record, and a fold with the value and the Record. + """ + + fold: _StreamFold | None = None + """The fold, or ``None`` for a read that gives its Records.""" + _result: ( tuple[_StreamReadMarker, list[temporalio.api.stream.v1.StreamRecord]] | None ) = None @@ -217,9 +230,9 @@ def drop_run(self, run_id: str) -> None: for key in [key for key in self._reads if key[0] == run_id]: del self._reads[key] - def __len__(self) -> int: + def keys(self) -> list[tuple[str, str]]: with self._lock: - return len(self._reads) + return list(self._reads) _PENDING_STREAM_READS = _PendingStreamReads() @@ -231,10 +244,80 @@ def __len__(self) -> int: @dataclasses.dataclass(frozen=True) -class _ReadUntilResult: +class _ReadLimits: + """The stops of a read, as the read Activity gets them.""" + + max_records: int + max_bytes: int | None + idle_timeout: float | None + """Seconds with no new Record before the read ends, or ``None``.""" + timeout: float + """Seconds for the whole read.""" + + +@dataclasses.dataclass(frozen=True) +class _ReadOutcome: records: list[temporalio.api.stream.v1.StreamRecord] - matched: bool + """The Records of a read. A fold keeps none.""" + record_count: int + last_offset: bytes | None + digest: bytes + stop: StreamReadStop done: bool + value: Any + """The value of a fold, or ``None`` for a read.""" + + +class _ReadState: + """What a read got so far, for one attempt of its Activity.""" + + def __init__(self, read: _PendingStreamRead, from_offset: bytes | None) -> None: + self._read = read + self._digest = _RecordDigest() + self.records: list[temporalio.api.stream.v1.StreamRecord] = [] + self.count = 0 + self.bytes = 0 + self.last_offset = from_offset + self.value = read.fold.initial if read.fold is not None else None + + def take(self, record: temporalio.api.stream.v1.StreamRecord, size: int) -> bool: + """Add ``record`` to the read, and give whether ``until`` is true after it.""" + self.count += 1 + self.bytes += size + self.last_offset = record.offset + self._digest.update(record) + read = self._read + if read.fold is None: + self.records.append(record) + if read.until is None: + return False + return bool(_call_workflow_function(read.until, read.to_record(record))) + typed = read.to_record(record) + self.value = _call_workflow_function(read.fold.function, self.value, typed) + if read.until is None: + return False + return bool(_call_workflow_function(read.until, self.value, typed)) + + def end(self, stop: StreamReadStop, *, done: bool = False) -> _ReadOutcome: + return _ReadOutcome( + records=self.records, + record_count=self.count, + last_offset=self.last_offset, + digest=self._digest.digest(), + stop=stop, + done=done, + value=self.value, + ) + + +def _call_workflow_function(function: Callable[..., Any], *args: Any) -> Any: + try: + return function(*args) + except Exception as err: + # The Workflow gets this error from its Local Activity, in Workflow code. + raise temporalio.exceptions.ApplicationError( + str(err), type=type(err).__name__, non_retryable=True + ) from err async def _read_until( # pyright: ignore[reportUnusedFunction] @@ -242,50 +325,72 @@ async def _read_until( # pyright: ignore[reportUnusedFunction] *, ref: temporalio.api.stream.v1.StreamRef, from_offset: bytes | None, - max_records: int, - timeout: timedelta, + limits: _ReadLimits, read: _PendingStreamRead, -) -> _ReadUntilResult: - """Read pages until the predicate of ``read`` is true, or a limit stops the read.""" - deadline = time.monotonic() + timeout.total_seconds() - records: list[temporalio.api.stream.v1.StreamRecord] = [] - cursor = from_offset - while len(records) < max_records: - time_left = deadline - time.monotonic() - if time_left <= 0: - break +) -> _ReadOutcome: + """Read pages until a stop of :py:class:`StreamReadStop` holds.""" + start = time.monotonic() + deadline = start + limits.timeout + # The idle stop counts from the last page that gave a Record. + last_arrival = start + state = _ReadState(read, from_offset) + + def idle(now: float) -> bool: + return ( + limits.idle_timeout is not None + and now - last_arrival >= limits.idle_timeout + ) + + while True: + if state.count >= limits.max_records: + return state.end(StreamReadStop.MAX_RECORDS) + bytes_left = None + if limits.max_bytes is not None: + bytes_left = limits.max_bytes - state.bytes + # Every Record has some bytes, so none fits in a cap that is used up. + if bytes_left <= 0: + return state.end(StreamReadStop.MAX_BYTES) + now = time.monotonic() + if now >= deadline: + return state.end( + StreamReadStop.IDLE if idle(now) else StreamReadStop.TIMEOUT + ) + wait = min(_MAX_PAGE_WAIT, deadline - now) + if limits.idle_timeout is not None: + wait = min(wait, max(0.0, last_arrival + limits.idle_timeout - now)) request = temporalio.api.streamservice.v1.ReadRecordsRequest( ref=ref, - from_offset=cursor or b"", - max_records=min(max_records - len(records), _MAX_PAGE_RECORDS), + from_offset=state.last_offset or b"", + max_records=min(limits.max_records - state.count, _MAX_PAGE_RECORDS), + max_bytes=bytes_left or 0, ) - request.wait.FromTimedelta(min(timedelta(seconds=time_left), _MAX_PAGE_WAIT)) + request.wait.FromTimedelta(timedelta(seconds=wait)) response = await read_page(request) - page = list(response.records)[: max_records - len(records)] - for index, record in enumerate(page): - records.append(record) - if read.until is not None and _call_until(read, record): - # A done page ends with the last Record, so a match on it is the end too. - last = index == len(response.records) - 1 - return _ReadUntilResult(records, True, response.done and last) - if response.done and len(page) == len(response.records): - return _ReadUntilResult(records, False, True) - if page: - cursor = page[-1].offset - return _ReadUntilResult(records, False, False) - - -def _call_until( - read: _PendingStreamRead, record: temporalio.api.stream.v1.StreamRecord -) -> bool: - assert read.until is not None - try: - return bool(read.until(read.to_record(record))) - except Exception as err: - # The Workflow gets this error from its Local Activity, in Workflow code. - raise temporalio.exceptions.ApplicationError( - str(err), type=type(err).__name__, non_retryable=True - ) from err + if response.records: + last_arrival = time.monotonic() + for index, record in enumerate(response.records): + if state.count >= limits.max_records: + return state.end(StreamReadStop.MAX_RECORDS) + size = record.ByteSize() + # A Store gives the first Record of a page even when it passes the cap, and + # the read keeps it only when it is the first Record of the read. + if ( + limits.max_bytes is not None + and state.count > 0 + and state.bytes + size > limits.max_bytes + ): + return state.end(StreamReadStop.MAX_BYTES) + matched = state.take(record, size) + # A done page ends with the last Record of the Stream. + done = response.done and index == len(response.records) - 1 + if matched: + return state.end(StreamReadStop.MATCH, done=done) + if done: + return state.end(StreamReadStop.DONE, done=True) + if response.done and not response.records: + return state.end(StreamReadStop.DONE, done=True) + if not response.records and idle(time.monotonic()): + return state.end(StreamReadStop.IDLE) async def _read_range( # pyright: ignore[reportUnusedFunction] diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index 81c5bd7b49..d4381b1d75 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -24,6 +24,7 @@ from temporalio.client._stream_read import ( _PENDING_STREAM_READS, _read_until, + _ReadLimits, _StreamReadMarker, ) from temporalio.converter import StorageDriverActivityInfo, StorageDriverWorkflowInfo @@ -146,11 +147,11 @@ def worker_activities(self) -> Sequence[Callable[..., Any]]: The default is two Activities, named for this Store. One answers :py:meth:`call_from_workflow`, which runs it as a Local Activity, or as a regular - Activity for a close. The other answers :py:func:`temporalio.workflow.read_stream`, - which always runs it by its name. A Store that overrides - :py:meth:`call_from_workflow` overrides this method to match, and a Store that - overrides this method keeps the read Activity in its list for Workflow code that - reads. + Activity for a close. The other answers :py:func:`temporalio.workflow.read_stream` + and :py:func:`temporalio.workflow.reduce_stream`, which always run it by its name. + A Store that overrides :py:meth:`call_from_workflow` overrides this method to + match. A Store that overrides this method keeps the read Activity in its list for + Workflow code that reads. """ return [_store_call_activity(self), _store_read_activity(self)] @@ -412,48 +413,68 @@ def _store_read_activity_name(store_name: str) -> str: def _store_read_activity( store: StreamStore, -) -> Callable[[str, bytes, int, float], Awaitable[temporalio.common.RawValue]]: +) -> Callable[ + [str, bytes, int, int | None, float | None, float], + Awaitable[temporalio.common.RawValue], +]: @temporalio.activity.defn(name=_store_read_activity_name(store.name)) async def read( - read_id: str, request: bytes, max_records: int, timeout_seconds: float + read_id: str, + request: bytes, + max_records: int, + max_bytes: int | None, + idle_timeout_seconds: float | None, + timeout_seconds: float, ) -> temporalio.common.RawValue: info = temporalio.activity.info() pending = _PENDING_STREAM_READS.get(info.workflow_run_id or "", read_id) if pending is None: - # The predicate and the Records live in the Workflow's process, so only a - # Local Activity of that Workflow can run this read. + # The functions of the Workflow live in its process, so only a Local + # Activity of that Workflow can run this read. raise temporalio.exceptions.ApplicationError( f"The Stream read '{read_id}' has no pending read in this process. " - "Only temporalio.workflow.read_stream runs this Activity.", + "Only temporalio.workflow.read_stream and reduce_stream run this " + "Activity.", non_retryable=True, ) read_request = temporalio.api.streamservice.v1.ReadRecordsRequest.FromString( request ) - # The codec decodes the bodies in each page, so the predicate and the Workflow + # The codec decodes the bodies in each page, so the functions of the Workflow # get the user's values. page_reader = _store_page_reader( store, temporalio.activity.client().data_converter ) from_offset = read_request.from_offset or None - result = await _read_until( + outcome = await _read_until( page_reader, ref=read_request.ref, from_offset=from_offset, - max_records=max_records, - timeout=timedelta(seconds=timeout_seconds), + limits=_ReadLimits( + max_records=max_records, + max_bytes=max_bytes, + idle_timeout=idle_timeout_seconds, + timeout=timeout_seconds, + ), read=pending, ) - marker = _StreamReadMarker.for_records( + marker = _StreamReadMarker( read_id=read_id, store=store.name, stream_id=read_request.ref.id, from_offset=from_offset, - records=result.records, - matched=result.matched, - done=result.done, + last_offset=outcome.last_offset, + record_count=outcome.record_count, + stop=outcome.stop, + done=outcome.done, + digest=outcome.digest, + # History keeps the value of a fold, so a replay needs no Store. + value=temporalio.activity.payload_converter().to_payload(outcome.value) + if pending.fold is not None + else None, ) - pending.fill(marker, result.records) + if pending.fold is None: + pending.fill(marker, outcome.records) return marker.to_raw_value() return read diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 13ced88af0..701dc8c227 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -73,7 +73,9 @@ def __init__( replay calls no method of its Stores for a create or a describe. For each :py:func:`temporalio.workflow.read_stream` in the history, it reads the Records of that read from the Store again, so the replay needs the same - ``data_converter`` and Stores that still have them. + ``data_converter`` and Stores that still have them. A + :py:func:`temporalio.workflow.reduce_stream` takes its value from the + history and calls no Store. 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 index 5440448715..f30817b149 100644 --- a/temporalio/worker/_stream_read.py +++ b/temporalio/worker/_stream_read.py @@ -1,4 +1,4 @@ -"""The Worker side of :py:func:`temporalio.workflow.read_stream`. +"""The Worker side of :py:func:`temporalio.workflow.read_stream` and ``reduce_stream``. Workflow code starts the read here, and the Worker gives a replayed read its Records. """ @@ -19,6 +19,7 @@ _PENDING_STREAM_READS, _PendingStreamRead, _read_range, + _StreamFold, _StreamReadMarker, ) from temporalio.client._stream_store import ( @@ -42,17 +43,21 @@ _READ_ACTIVITY_GRACE = timedelta(seconds=10) -async def _read_stream_in_workflow( # pyright: ignore[reportUnusedFunction] +async def _run_read( instance: _WorkflowInstanceImpl, config: StreamConfig, stream: Stream[Any], + limits: temporalio.workflow._StreamReadLimits, *, - until: Callable[[StreamRecord[Any]], bool] | None, - from_offset: bytes | None, - max_records: int, - timeout: timedelta, -) -> temporalio.workflow.StreamRead[Any]: - request = ReadRecordsRequest(ref=stream.ref, from_offset=from_offset) + until: Callable[..., bool] | None, + fold: _StreamFold | None, +) -> tuple[ + _StreamReadMarker, + list[temporalio.api.stream.v1.StreamRecord], + Callable[[temporalio.api.stream.v1.StreamRecord], Any], +]: + """Run the read Activity, and give its marker, its proto Records and their typing.""" + request = ReadRecordsRequest(ref=stream.ref, from_offset=limits.from_offset) payload_converter, _ = instance._system_nexus_converters_for( _STREAM_SERVICE.name, "ReadRecords", request ) @@ -66,7 +71,7 @@ def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: run_id = temporalio.workflow.info().run_id read_id = str(temporalio.workflow.uuid4()) - pending = _PendingStreamRead(until=until, to_record=to_record) + pending = _PendingStreamRead(to_record=to_record, until=until, fold=fold) _PENDING_STREAM_READS.add(run_id, read_id, pending) try: with _store_errors(): @@ -75,11 +80,15 @@ def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: args=[ read_id, request_message.SerializeToString(), - max_records, - timeout.total_seconds(), + limits.max_records, + limits.max_bytes, + None + if limits.idle_timeout is None + else limits.idle_timeout.total_seconds(), + limits.timeout.total_seconds(), ], result_type=temporalio.common.RawValue, - start_to_close_timeout=timeout + _READ_ACTIVITY_GRACE, + start_to_close_timeout=limits.timeout + _READ_ACTIVITY_GRACE, retry_policy=config.workflow_call_options.retry_policy, ) marker = _StreamReadMarker.from_payload(result.payload) @@ -88,21 +97,63 @@ def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: f"The read of the Stream '{stream.ref.id}' got a result that is not " "its read marker" ) + if fold is not None: + return marker, [], to_record records = pending.records_for(marker) if records is None: raise RuntimeError( f"The Worker has no Records for the read of {marker.describe_range()}" ) - return temporalio.workflow.StreamRead( - records=[to_record(record) for record in records], - last_offset=marker.last_offset, - matched=marker.matched, - done=marker.done, - ) + return marker, records, to_record finally: _PENDING_STREAM_READS.remove(run_id, read_id, pending) +async def _read_stream_in_workflow( # pyright: ignore[reportUnusedFunction] + instance: _WorkflowInstanceImpl, + config: StreamConfig, + stream: Stream[Any], + until: Callable[[StreamRecord[Any]], bool] | None, + limits: temporalio.workflow._StreamReadLimits, +) -> temporalio.workflow.StreamRead[Any]: + marker, records, to_record = await _run_read( + instance, config, stream, limits, until=until, fold=None + ) + return temporalio.workflow.StreamRead( + records=[to_record(record) for record in records], + last_offset=marker.last_offset, + stop=marker.stop, + done=marker.done, + ) + + +async def _reduce_stream_in_workflow( # pyright: ignore[reportUnusedFunction] + instance: _WorkflowInstanceImpl, + config: StreamConfig, + stream: Stream[Any], + fold: _StreamFold, + until: Callable[[Any, StreamRecord[Any]], bool] | None, + value_type: Any, + limits: temporalio.workflow._StreamReadLimits, +) -> temporalio.workflow.StreamReduce[Any]: + marker, _, _ = await _run_read( + instance, config, stream, limits, until=until, fold=fold + ) + if marker.value is None: + raise RuntimeError( + f"The fold of {marker.describe_range()} got a marker with no value" + ) + # A first run and a replay both take the value from the marker. + value = instance.workflow_payload_converter().from_payload(marker.value, value_type) + return temporalio.workflow.StreamReduce( + value=value, + count=marker.record_count, + last_offset=marker.last_offset, + stop=marker.stop, + done=marker.done, + ) + + async def _fill_stream_reads( # pyright: ignore[reportUnusedFunction] act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, config: StreamConfig, @@ -124,7 +175,8 @@ async def _fill_stream_reads( # pyright: ignore[reportUnusedFunction] if not resolve.is_local or not resolve.result.HasField("completed"): continue marker = _StreamReadMarker.from_payload(resolve.result.completed.result) - if marker is None: + # History has the value of a fold, so its replay needs no Store. + if marker is None or marker.value is not None: continue pending = _PENDING_STREAM_READS.get(act.run_id, marker.read_id) # On a first run the read Activity already gave its Records. diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 1176497689..470399dad4 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -59,6 +59,7 @@ import temporalio.exceptions import temporalio.nexus.system import temporalio.workflow +from temporalio.client._stream_read import _StreamFold from temporalio.client._stream_store import _STREAM_SERVICE from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo from temporalio.nexus.system.stream_service.models import StreamRecord @@ -97,7 +98,7 @@ WorkflowInboundInterceptor, WorkflowOutboundInterceptor, ) -from ._stream_read import _read_stream_in_workflow +from ._stream_read import _read_stream_in_workflow, _reduce_stream_in_workflow from ._stream_store import ( _NO_STREAM_STORE_MESSAGE, _no_stream_config_error, @@ -2004,22 +2005,34 @@ def workflow_upsert_search_attributes( async def workflow_read_stream( self, stream: Stream[Any], - *, until: Callable[[StreamRecord[Any]], bool] | None, - from_offset: bytes | None, - max_records: int, - timeout: timedelta, + limits: temporalio.workflow._StreamReadLimits, ) -> temporalio.workflow.StreamRead[Any]: if self._stream_store is None: raise temporalio.exceptions.StreamConfigError(_NO_STREAM_STORE_MESSAGE) return await _read_stream_in_workflow( + self, self._stream_store.config, stream, until, limits + ) + + async def workflow_reduce_stream( + self, + stream: Stream[Any], + function: Callable[[Any, StreamRecord[Any]], Any], + initial: Any, + until: Callable[[Any, StreamRecord[Any]], bool] | None, + value_type: Any, + limits: temporalio.workflow._StreamReadLimits, + ) -> temporalio.workflow.StreamReduce[Any]: + if self._stream_store is None: + raise temporalio.exceptions.StreamConfigError(_NO_STREAM_STORE_MESSAGE) + return await _reduce_stream_in_workflow( self, self._stream_store.config, stream, - until=until, - from_offset=from_offset, - max_records=max_records, - timeout=timeout, + _StreamFold(function=function, initial=initial), + until, + value_type, + limits, ) async def workflow_sleep( diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index 748d5b2502..11512d7c53 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -158,7 +158,11 @@ from ._streams import ( DEFAULT_READ_STREAM_MAX_RECORDS, StreamRead, + StreamReadStop, + StreamReduce, + _StreamReadLimits, read_stream, + reduce_stream, ) from ._workflow_ops import ( ChildWorkflowCancellationType, @@ -269,7 +273,10 @@ "unsafe", "DEFAULT_READ_STREAM_MAX_RECORDS", "StreamRead", + "StreamReadStop", + "StreamReduce", "read_stream", + "reduce_stream", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", @@ -312,6 +319,7 @@ "_update_validator", "_UpdateDefinition", "_NexusClient", + "_StreamReadLimits", "_build_log_context", "_imports_passed_through", "_in_sandbox", diff --git a/temporalio/workflow/_context.py b/temporalio/workflow/_context.py index 1a051c6ab3..7d22ba4b2d 100644 --- a/temporalio/workflow/_context.py +++ b/temporalio/workflow/_context.py @@ -30,7 +30,7 @@ from ._event_groups import EventGroup from ._exceptions import ContinueAsNewVersioningBehavior, VersioningIntent from ._nexus import NexusOperationCancellationType, NexusOperationHandle - from ._streams import StreamRead + from ._streams import StreamRead, StreamReduce, _StreamReadLimits from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowHandle, @@ -409,13 +409,21 @@ def workflow_random(self) -> Random: ... async def workflow_read_stream( self, stream: Stream[Any], - *, until: Callable[[StreamRecord[Any]], bool] | None, - from_offset: bytes | None, - max_records: int, - timeout: timedelta, + limits: _StreamReadLimits, ) -> StreamRead[Any]: ... + @abstractmethod + async def workflow_reduce_stream( + self, + stream: Stream[Any], + function: Callable[[Any, StreamRecord[Any]], Any], + initial: Any, + until: Callable[[Any, StreamRecord[Any]], bool] | None, + value_type: Any, + limits: _StreamReadLimits, + ) -> StreamReduce[Any]: ... + @abstractmethod def workflow_set_query_handler( self, name: str | None, handler: Callable | None diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py index be241e9c75..cc5616d11e 100644 --- a/temporalio/workflow/_streams.py +++ b/temporalio/workflow/_streams.py @@ -1,4 +1,4 @@ -"""A Workflow read of a Stream that waits for a predicate. +"""A Workflow read of a Stream that waits for a predicate, and a fold over its Records. .. warning:: This API is an experimental prototype and subject to change. @@ -6,6 +6,7 @@ from __future__ import annotations +import enum from collections.abc import Callable, Sequence from dataclasses import dataclass from datetime import timedelta @@ -21,12 +22,49 @@ Stream, ) -__all__ = ["StreamRead", "read_stream"] +__all__ = [ + "DEFAULT_READ_STREAM_MAX_RECORDS", + "StreamRead", + "StreamReadStop", + "StreamReduce", + "read_stream", + "reduce_stream", +] BodyT = TypeVar("BodyT") +ValueT = TypeVar("ValueT") -DEFAULT_READ_STREAM_MAX_RECORDS = 10_000 -"""The most Records that one :py:func:`read_stream` gives when ``max_records`` is ``None``.""" +DEFAULT_READ_STREAM_MAX_RECORDS = 1_000 +"""The most Records of one read or fold when ``max_records`` is ``None``.""" + + +class StreamReadStop(enum.Enum): + """Why a read or a fold of a Stream ended. + + When several hold, the first in this order is the reason. ``MATCH``, ``DONE``, + ``MAX_RECORDS``, ``MAX_BYTES``, ``IDLE``, ``TIMEOUT``. + + .. warning:: + This API is an experimental prototype and subject to change. + """ + + MATCH = "MATCH" + """``until`` was true for the last Record.""" + + DONE = "DONE" + """The Stream is closed and its last Record was read.""" + + MAX_RECORDS = "MAX_RECORDS" + """The read got ``max_records`` Records.""" + + MAX_BYTES = "MAX_BYTES" + """The next Record would pass ``max_bytes``. The next read gets it.""" + + IDLE = "IDLE" + """No new Record arrived for ``idle_timeout``.""" + + TIMEOUT = "TIMEOUT" + """``timeout`` passed.""" @dataclass(frozen=True) @@ -46,11 +84,79 @@ class StreamRead(Generic[BodyT]): Give it as the ``from_offset`` of the next read. """ - matched: bool - """Whether ``until`` was true for the last Record.""" + stop: StreamReadStop + """Why the read ended.""" done: bool - """Whether the Stream is closed and its last Record is in :py:attr:`records`.""" + """Whether the Stream is closed and its last Record is in :py:attr:`records`. + + It is the state of the Stream, so it can be true when :py:attr:`stop` is ``MATCH``. + """ + + +@dataclass(frozen=True) +class StreamReduce(Generic[ValueT]): + """What one :py:func:`reduce_stream` gave. + + .. warning:: + This API is an experimental prototype and subject to change. + """ + + value: ValueT + """The value after the last Record, or ``initial`` when the fold got none.""" + + count: int + """The number of Records that the fold got.""" + + last_offset: bytes | None + """The offset of the last Record, or the ``from_offset`` of a fold with none. + + Give it as the ``from_offset`` of the next read or fold. + """ + + stop: StreamReadStop + """Why the fold ended.""" + + done: bool + """Whether the Stream is closed and the fold got its last Record.""" + + +@dataclass(frozen=True) +class _StreamReadLimits: + """Where a read starts and what stops it, after the checks of the public functions.""" + + from_offset: bytes | None + idle_timeout: timedelta | None + timeout: timedelta + max_records: int + max_bytes: int | None + + @staticmethod + def check( + *, + from_offset: bytes | None, + idle_timeout: timedelta | None, + timeout: timedelta, + max_records: int | None, + max_bytes: int | None, + ) -> _StreamReadLimits: + if timeout <= timedelta(): + raise ValueError("timeout must be positive") + if idle_timeout is not None and idle_timeout < timedelta(): + raise ValueError("idle_timeout must not be negative") + if max_records is None: + max_records = DEFAULT_READ_STREAM_MAX_RECORDS + elif max_records < 1: + raise ValueError("max_records must be at least 1") + if max_bytes is not None and max_bytes < 1: + raise ValueError("max_bytes must be at least 1") + return _StreamReadLimits( + from_offset=from_offset, + idle_timeout=idle_timeout, + timeout=timeout, + max_records=max_records, + max_bytes=max_bytes, + ) async def read_stream( @@ -58,17 +164,18 @@ async def read_stream( *, until: Callable[[StreamRecord[BodyT]], bool] | None = None, from_offset: bytes | None = None, - max_records: int | None = None, + idle_timeout: timedelta | None = None, timeout: timedelta, + max_records: int | None = None, + max_bytes: int | None = None, ) -> StreamRead[BodyT]: """Read a Stream until a predicate is true, in one Local Activity. - The read starts after ``from_offset`` and stops at the first of these. ``until`` - is true for a Record, ``max_records`` Records were read, the Stream is closed and - its last Record was read, or ``timeout`` passed. A Store's Local Activity reads - the pages, and History gets one marker for the whole read. The marker has the - range of offsets and no Record. A replay reads that range from the Store again - before the Workflow sees the result, so the Records must still be in the Store. + The read starts after ``from_offset`` and ends at the first stop of + :py:class:`StreamReadStop`. A Store's Local Activity reads the pages, and History + gets one marker for the whole read. The marker has the range of offsets and no + Record. A replay reads that range from the Store again before the Workflow sees the + result, so the Records must still be in the Store. Args: stream: The Workflow ``Stream`` to read, from @@ -76,20 +183,29 @@ async def read_stream( :py:func:`temporalio.workflow.create_stream`. until: Called with each typed Record in order. The read stops after the first Record for which it is true, and Records after it stay for the next read. - It runs in the Worker process, on the Worker's event loop, while the Local - Activity reads. It runs outside the Workflow's own thread. So it must not - change any state, and it can't call a function of ``temporalio.workflow``. - It can read the Workflow's state through a reference such as ``self``. - When it sees a change of that state is not defined. A replay does not call - it, since the replay takes the end of the read from History. ``None`` - reads until another limit stops the read. + It is called once for each Record, when the Record arrives. A later change + of the state that it reads does not check an earlier Record again. It runs + in the Worker process, on the Worker's event loop, while the Local Activity + reads. It runs outside the Workflow's own thread. So it must not change any + state, and it can't call a function of ``temporalio.workflow``. It can read + the Workflow's state through a reference such as ``self``. When it sees a + change of that state is not defined. A replay does not call it, since the + replay takes the end of the read from History. ``None`` reads until another + stop. from_offset: Read after this offset. ``None`` is the beginning. + idle_timeout: End the read when no new Record arrived for this long. It counts + from the last Record that the read got, or from the start of the read when + it got none. ``timedelta(0)`` gives what is in the Stream now and waits for + nothing. ``None`` is no idle stop. + timeout: The most time for the read. It must be positive. max_records: The most Records to read. ``None`` is :py:data:`DEFAULT_READ_STREAM_MAX_RECORDS`. - timeout: How long the read waits for a match. It must be positive. + max_bytes: A cap on the bytes of the Records, as the Store counts them for a + page of the client ``read``. A Record that would pass the cap ends the read + before it, and the next read gets it. The first Record is always returned. Returns: - The Records of the read, the offset to read after next, and why it stopped. + The Records of the read, the offset to read after next, and why it ended. Raises: temporalio.exceptions.StreamError: The Store failed the read. @@ -101,16 +217,84 @@ async def read_stream( .. warning:: This API is an experimental prototype and subject to change. """ - if timeout <= timedelta(): - raise ValueError("timeout must be positive") - if max_records is None: - max_records = DEFAULT_READ_STREAM_MAX_RECORDS - elif max_records < 1: - raise ValueError("max_records must be at least 1") - return await _Runtime.current().workflow_read_stream( - stream, - until=until, + limits = _StreamReadLimits.check( from_offset=from_offset, + idle_timeout=idle_timeout, + timeout=timeout, max_records=max_records, + max_bytes=max_bytes, + ) + return await _Runtime.current().workflow_read_stream(stream, until, limits) + + +async def reduce_stream( + stream: Stream[BodyT], + function: Callable[[ValueT, StreamRecord[BodyT]], ValueT], + initial: ValueT, + *, + until: Callable[[ValueT, StreamRecord[BodyT]], bool] | None = None, + value_type: type[ValueT] | None = None, + from_offset: bytes | None = None, + idle_timeout: timedelta | None = None, + timeout: timedelta, + max_records: int | None = None, + max_bytes: int | None = None, +) -> StreamReduce[ValueT]: + """Fold the Records of a Stream into one value, in one Local Activity. + + The Local Activity calls ``value = function(value, record)`` for each Record in + order, then ``until(value, record)``. The fold ends at the first stop of + :py:class:`StreamReadStop`, with the same rules as :py:func:`read_stream`. No + Record comes back to the Workflow. History keeps the value in the marker of the + Local Activity, so a replay takes the value from History and needs no Store. The + size of the value is the caller's choice, and History pays for it. + + Args: + stream: The Workflow ``Stream`` to fold. + function: Gives the next value from the value so far and a typed Record. It + changes no state other than the value that it returns, so it must not + change ``initial`` or a value in place. It runs in the Worker process, on + the Worker's event loop, outside the Workflow's own thread. It does no I/O, + and it can't call a function of ``temporalio.workflow``. A replay does not + call it. + initial: The value before the first Record. + until: Called with the value after each Record and that Record. The fold stops + after the first Record for which it is true. It is called once for each + Record, when the Record arrives. It has the rules of ``function``. + value_type: The type that the Workflow decodes the value as. ``None`` takes the + type of ``initial``. + from_offset: Fold the Records after this offset. ``None`` is the beginning. + idle_timeout: As for :py:func:`read_stream`. + timeout: As for :py:func:`read_stream`. + max_records: As for :py:func:`read_stream`. + max_bytes: As for :py:func:`read_stream`. + + Returns: + The value, the number of Records, the offset to read after next, and why the + fold ended. + + Raises: + temporalio.exceptions.StreamError: The Store failed the read. + temporalio.exceptions.ApplicationError: ``function`` or ``until`` raised. The + error has the message of the exception. + temporalio.exceptions.StreamConfigError: The Worker's client has no Stream + config. This fails the Workflow Task. + + .. warning:: + This API is an experimental prototype and subject to change. + """ + limits = _StreamReadLimits.check( + from_offset=from_offset, + idle_timeout=idle_timeout, timeout=timeout, + max_records=max_records, + max_bytes=max_bytes, + ) + return await _Runtime.current().workflow_reduce_stream( + stream, + function, + initial, + until, + value_type if value_type is not None else type(initial), + limits, ) diff --git a/tests/helpers/streams.py b/tests/helpers/streams.py index b9606345d1..5cbfe2c7e6 100644 --- a/tests/helpers/streams.py +++ b/tests/helpers/streams.py @@ -194,5 +194,6 @@ class StreamReadState: """ until_calls = 0 + function_calls = 0 fail_next_task = False results: dict[str, tuple[object, ...]] = {} diff --git a/tests/nexus/test_stream_read_workflow.py b/tests/nexus/test_stream_read_workflow.py index 45c1ca09bf..d51c5a5441 100644 --- a/tests/nexus/test_stream_read_workflow.py +++ b/tests/nexus/test_stream_read_workflow.py @@ -15,6 +15,7 @@ from temporalio import workflow from temporalio.api.enums.v1 import EventType +from temporalio.api.stream.v1 import StreamRef as StreamRefMessage from temporalio.api.streamservice.v1 import ReadRecordsRequest, ReadRecordsResponse from temporalio.client import ( Client, @@ -23,6 +24,7 @@ WorkflowHandle, WorkflowHistory, ) +from temporalio.client._stream_read import _PENDING_STREAM_READS from temporalio.converter import DataConverter from temporalio.exceptions import ApplicationError, CancelledError from temporalio.nexus.system.stream_service import StreamRecord, StreamRef @@ -56,12 +58,15 @@ class ReadArgs: until_kind: str | None = "done" timeout_seconds: float = 10.0 max_records: int | None = None + idle_timeout_seconds: float | None = None + max_bytes: int | None = None + from_offset: str | None = None @dataclass class ReadResult: texts: list[str] - matched: bool + stop: str done: bool last_offset: str | None typed: bool @@ -70,7 +75,7 @@ class ReadResult: def _result(read: workflow.StreamRead[Token]) -> ReadResult: return ReadResult( texts=[record.body.text for record in read.records], - matched=read.matched, + stop=read.stop.value, done=read.done, last_offset=read.last_offset.hex() if read.last_offset else None, typed=all( @@ -84,7 +89,7 @@ def _keep(result: ReadResult) -> ReadResult: # A replay gives no result to the test, so the Workflow keeps it here. StreamReadState.results[workflow.info().run_id] = ( tuple(result.texts), - result.matched, + result.stop, result.done, result.last_offset, ) @@ -109,8 +114,13 @@ def until(record: StreamRecord[Token]) -> bool: read = await workflow.read_stream( _stream(args.stream_id), until=until if until_kind is not None else None, - max_records=args.max_records, + from_offset=bytes.fromhex(args.from_offset) if args.from_offset else None, + idle_timeout=None + if args.idle_timeout_seconds is None + else timedelta(seconds=args.idle_timeout_seconds), timeout=timedelta(seconds=args.timeout_seconds), + max_records=args.max_records, + max_bytes=args.max_bytes, ) return _keep(_result(read)) @@ -226,6 +236,192 @@ async def run(self, stream_id: str, predicate: str) -> str: return "no error" +@workflow.defn +class MaxBytesThenNextWorkflow: + @workflow.run + async def run(self, stream_id: str, max_bytes: int) -> list[ReadResult]: + stream = _stream(stream_id) + first = await workflow.read_stream( + stream, + idle_timeout=timedelta(0), + timeout=timedelta(seconds=10), + max_bytes=max_bytes, + ) + second = await workflow.read_stream( + stream, + from_offset=first.last_offset, + timeout=timedelta(seconds=10), + max_records=1, + ) + return [_result(first), _result(second)] + + +@dataclass +class ReduceArgs: + stream_id: str + until: str = "done" + timeout_seconds: float = 10.0 + + +@dataclass +class ReduceResult: + value: str + count: int + stop: str + done: bool + last_offset: str | None + + +def _join(value: str, record: StreamRecord[Token]) -> str: + StreamReadState.function_calls += 1 + return value + record.body.text + + +def _until_done(value: str, record: StreamRecord[Token]) -> bool: + del value + StreamReadState.until_calls += 1 + return record.body.kind == "done" + + +def _until_tag(value: str, record: StreamRecord[Token]) -> bool: + del record + StreamReadState.until_calls += 1 + return value.endswith("tag!") + + +_FOLD_UNTIL: dict[str, Callable[[str, StreamRecord[Token]], bool]] = { + "done": _until_done, + "tag": _until_tag, +} + + +def _reduce_result(reduce: workflow.StreamReduce[str]) -> ReduceResult: + return ReduceResult( + value=reduce.value, + count=reduce.count, + stop=reduce.stop.value, + done=reduce.done, + last_offset=reduce.last_offset.hex() if reduce.last_offset else None, + ) + + +@workflow.defn +class JoinTextWorkflow: + @workflow.run + async def run(self, args: ReduceArgs) -> ReduceResult: + reduce = await workflow.reduce_stream( + _stream(args.stream_id), + _join, + "", + until=_FOLD_UNTIL[args.until], + timeout=timedelta(seconds=args.timeout_seconds), + ) + result = _reduce_result(reduce) + StreamReadState.results[workflow.info().run_id] = ( + result.value, + result.count, + result.stop, + result.done, + result.last_offset, + ) + return result + + +@workflow.defn +class TwoFoldsWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[ReduceResult]: + stream = _stream(stream_id) + first = await workflow.reduce_stream( + stream, + _join, + "", + until=lambda value, _: value.endswith("m1"), + timeout=timedelta(seconds=10), + ) + await workflow.sleep(timedelta(milliseconds=1)) + second = await workflow.reduce_stream( + stream, + _join, + "", + until=lambda value, _: value.endswith("m2"), + from_offset=first.last_offset, + timeout=timedelta(seconds=10), + ) + return [_reduce_result(first), _reduce_result(second)] + + +def _raise_in_function(value: str, record: StreamRecord[Token]) -> str: + raise ValueError(f"bad fold after {value!r} at {record.body.text}") + + +@workflow.defn +class FailingFunctionWorkflow: + @workflow.run + async def run(self, stream_id: str) -> str: + try: + await workflow.reduce_stream( + _stream(stream_id), + _raise_in_function, + "", + timeout=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, record: StreamRecord[Token]) -> Summary: + return Summary(value.count + 1, [*value.kinds, record.body.kind]) + + +def _collect(value: list[str], record: StreamRecord[Token]) -> list[str]: + return [*value, record.body.text] + + +@dataclass +class TypedFoldResult: + summary: Summary + summary_type: str + texts: list[str] + texts_type: str + + +@workflow.defn +class TypedFoldWorkflow: + @workflow.run + async def run(self, stream_id: str) -> TypedFoldResult: + stream = _stream(stream_id) + summary = await workflow.reduce_stream( + stream, + _summarize, + Summary(0, []), + idle_timeout=timedelta(0), + timeout=timedelta(seconds=10), + ) + no_texts: list[str] = [] + texts = await workflow.reduce_stream( + stream, + _collect, + no_texts, + value_type=list[str], + idle_timeout=timedelta(0), + timeout=timedelta(seconds=10), + ) + return TypedFoldResult( + summary=summary.value, + summary_type=type(summary.value).__name__, + texts=texts.value, + texts_type=type(texts.value).__name__, + ) + + class TrackingStore(RecordingStreamStore): """A memory Store that knows whether a read waits, and whether one was cancelled.""" @@ -366,7 +562,7 @@ async def test_read_until_match(client: Client, runner: WorkflowRunner): _, _, handle, result, tokens = await _run_match(client, runner) assert result == ReadResult( texts=_texts(tokens), - matched=True, + stop="MATCH", done=False, last_offset=result.last_offset, typed=True, @@ -387,7 +583,7 @@ async def test_read_waits_for_records(client: Client): await _append_over(stream_client, stream_id, tokens, 1.0) result = await handle.result() assert result.texts == _texts(tokens) - assert result.matched + assert result.stop == "MATCH" assert _markers(await handle.fetch_history()) == ["core_local_activity"] @@ -438,7 +634,7 @@ async def read_waits() -> None: await _append(stream_client, stream_id, wanted, *_tokens(2, prefix="after")) result = await handle.result() assert result.texts == _texts([*first, wanted]) - assert result.matched + assert result.stop == "MATCH" async def test_read_replays_with_replayer(client: Client): @@ -517,7 +713,7 @@ async def test_read_survives_a_failed_workflow_task(client: Client): ) result = await handle.result() assert result.texts == _texts(tokens) - assert result.matched + assert result.stop == "MATCH" counts = _event_counts(await handle.fetch_history()) assert counts["WORKFLOW_TASK_FAILED"] == 1 assert counts["MARKER_RECORDED"] == 1 @@ -539,9 +735,9 @@ async def test_read_timeout_gives_records_so_far(client: Client): await _append(stream_client, stream_id, *rest) timed_out, matched = await handle.result() assert timed_out.texts == _texts(first) - assert not timed_out.matched and not timed_out.done + assert timed_out.stop == "TIMEOUT" and not timed_out.done assert matched.texts == _texts(rest) - assert matched.matched + assert matched.stop == "MATCH" async def test_read_stops_at_max_records(client: Client): @@ -556,7 +752,7 @@ async def test_read_stops_at_max_records(client: Client): ) result = await handle.result() assert result.texts == _texts(tokens[:3]) - assert not result.matched and not result.done + assert result.stop == "MAX_RECORDS" and not result.done async def test_read_of_closed_stream_is_done(client: Client): @@ -570,7 +766,7 @@ async def test_read_of_closed_stream_is_done(client: Client): ) result = await handle.result() assert result.texts == _texts(tokens) - assert result.done and not result.matched + assert result.stop == "DONE" and result.done async def test_read_with_payload_codec(client: Client): @@ -580,7 +776,7 @@ async def test_read_with_payload_codec(client: Client): client, data_converter=data_converter ) assert result.texts == _texts(tokens) - assert result.typed and result.matched + assert result.typed and result.stop == "MATCH" # The predicate saw each Record decoded, or its type check would have raised. assert StreamReadState.until_calls == len(tokens) run_id = handle.result_run_id @@ -758,3 +954,308 @@ async def test_measure_match_latency(client: Client): "\nC. Append of the match to the end of the Workflow: " + ", ".join(f"{latency * 1000:.0f} ms" for latency in latencies) ) + + +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 test_read_idle_timeout(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = _tokens(3) + await _append(stream_client, stream_id, *tokens) + start = time.monotonic() + _, result = await _run_workflow( + stream_client, + ReadUntilKindWorkflow, + ReadArgs(stream_id, idle_timeout_seconds=1, timeout_seconds=30), + ) + seconds = time.monotonic() - start + assert result.texts == _texts(tokens) + assert result.stop == "IDLE" and not result.done + assert 1 <= seconds < 3 + + +async def test_read_idle_timeout_zero(client: Client): + stream_client, _, stream_id = await _new_stream(client) + empty_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=empty_id, body_type=Token) + tokens = _tokens(3) + await _append(stream_client, stream_id, *tokens) + for read_id, expected in [(stream_id, _texts(tokens)), (empty_id, [])]: + start = time.monotonic() + _, result = await _run_workflow( + stream_client, + ReadUntilKindWorkflow, + ReadArgs(read_id, idle_timeout_seconds=0, timeout_seconds=30), + ) + assert result.texts == expected + assert result.stop == "IDLE" + assert time.monotonic() - start < 1 + + +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] + + +@pytest.mark.parametrize( + "stop", + ["MATCH", "DONE", "IDLE", "TIMEOUT", "MAX_RECORDS", "MAX_BYTES"], +) +async def test_read_stop_reason(client: Client, stop: str): + stream_client, store, stream_id = await _new_stream(client) + tokens = _tokens(5) + if stop == "MATCH": + tokens[2] = _done() + await _append(stream_client, stream_id, *tokens) + if stop == "DONE": + await stream_client.get_stream(StreamRef(id=stream_id), body_type=Token).close() + if stop == "MAX_BYTES": + sizes = await _record_sizes(store, stream_id) + # The third Record passes the cap by one byte. + cap = sum(sizes[:3]) - 1 + _, (first, second) = await _run_workflow( + stream_client, MaxBytesThenNextWorkflow, stream_id, cap + ) + assert first.texts == _texts(tokens[:2]) + assert first.stop == "MAX_BYTES" + assert second.texts == _texts(tokens[2:3]) + # A cap below the first Record still gives that Record. + _, (alone, _) = await _run_workflow( + stream_client, MaxBytesThenNextWorkflow, stream_id, 1 + ) + assert alone.texts == _texts(tokens[:1]) + assert alone.stop == "MAX_BYTES" + return + args = { + "MATCH": ReadArgs(stream_id), + "DONE": ReadArgs(stream_id), + "IDLE": ReadArgs(stream_id, idle_timeout_seconds=0.5, timeout_seconds=30), + "TIMEOUT": ReadArgs(stream_id, timeout_seconds=1), + "MAX_RECORDS": ReadArgs(stream_id, max_records=3), + }[stop] + _, result = await _run_workflow(stream_client, ReadUntilKindWorkflow, args) + assert result.stop == stop + expected = {"MATCH": tokens[:3], "MAX_RECORDS": tokens[:3]}.get(stop, tokens) + assert result.texts == _texts(expected) + assert result.done == (stop == "DONE") + + +async def test_reduce_joins_text_and_replays_without_store(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = [*_tokens(3), _done(), *_tokens(2, prefix="after")] + await _append(stream_client, stream_id, *tokens) + handle, result = await _run_workflow( + stream_client, JoinTextWorkflow, ReduceArgs(stream_id) + ) + assert result == ReduceResult( + value="".join(_texts(tokens[:4])), + count=4, + stop="MATCH", + done=False, + last_offset=result.last_offset, + ) + history = await handle.fetch_history() + assert _markers(history) == ["core_local_activity"] + run_id = handle.result_run_id + assert run_id + first_run = StreamReadState.results.pop(run_id) + empty_store = RecordingStreamStore(MemoryStreamStore()) + StreamReadState.function_calls = 0 + StreamReadState.until_calls = 0 + await Replayer( + workflows=[JoinTextWorkflow], + stream_config=StreamConfig(stores=[empty_store]), + ).replay_workflow(history) + assert StreamReadState.results.pop(run_id) == first_run + assert empty_store.operations == [] + assert StreamReadState.function_calls == 0 + assert StreamReadState.until_calls == 0 + + +async def test_reduce_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, JoinTextWorkflow, ReduceArgs(stream_id, until="tag") + ) + assert result.value == "xtag!" + assert result.count == 3 + assert result.stop == "MATCH" + + +async def test_reduce_on_a_worker_without_cache(client: Client): + stream_client, store, stream_id = await _new_stream( + client, RecordingStreamStore(MemoryStreamStore()) + ) + texts = ["a", "m1", "b", "m2", "c"] + await _append(stream_client, stream_id, *[Token("text", text) for text in texts]) + handle, (first, second) = await _run_workflow( + stream_client, + TwoFoldsWorkflow, + stream_id, + worker_options={"max_cached_workflows": 0}, + ) + assert (first.value, first.count, first.stop) == ("am1", 2, "MATCH") + assert (second.value, second.count, second.stop) == ("bm2", 2, "MATCH") + assert _markers(await handle.fetch_history()) == ["core_local_activity"] * 2 + # A replay of the first fold takes its value from History. + assert store.operations.count("ReadRecords") == 2 + + +async def test_reduce_function_that_raises(client: Client): + stream_client, _, stream_id = await _new_stream(client) + await _append(stream_client, stream_id, *_tokens(1)) + _, result = await _run_workflow(stream_client, FailingFunctionWorkflow, stream_id) + assert result == f"ValueError: bad fold after '' at {_BODY_MARK}-t0" + + +async def test_reduce_typed_values(client: Client): + stream_client, _, stream_id = await _new_stream(client) + tokens = [*_tokens(2), _done()] + await _append(stream_client, stream_id, *tokens) + _, result = await _run_workflow(stream_client, TypedFoldWorkflow, stream_id) + assert result == TypedFoldResult( + summary=Summary(3, ["text", "text", "done"]), + summary_type="Summary", + texts=_texts(tokens), + texts_type="list", + ) + + +async def test_read_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, ReadUntilKindWorkflow, max_cached_workflows=0 + ) as worker: + handle = await _start_read( + stream_client, + ReadArgs(stream_id, timeout_seconds=30), + worker.task_queue, + task_timeout=timedelta(seconds=2), + ) + # The registry keeps the entry of the read 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_READS.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_READS.keys() + ) + + +async def _signal_wait(client: Client) -> float: + stream_client, _, stream_id = await _new_stream(client) + async with new_worker(stream_client, ReadUntilWantedWorkflow) as worker: + handle = await stream_client.start_workflow( + ReadUntilWantedWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await asyncio.sleep(1) + await handle.signal(ReadUntilWantedWorkflow.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( + "\nD. Signal to the next Workflow Task Started while a read 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, JoinTextWorkflow) as worker: + handle = await stream_client.start_workflow( + JoinTextWorkflow.run, + ReduceArgs(stream_id, timeout_seconds=30), + id=f"workflow-{uuid.uuid4()}", + task_queue=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() + marker = next( + event.marker_recorded_event_attributes + for event in history.events + if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED + ) + payload = marker.details["result"].payloads[0] + print(f"\nE. 40 Records folded over 2 s, History events: {_event_counts(history)}") + print( + f"E. Marker payload: {payload.ByteSize()} bytes, data {len(payload.data)} " + f"bytes, folded value {len(result.value)} characters" + ) From b2b9a5f0e3e628182b5781fc7dc162401840bf33 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 23:33:23 -0700 Subject: [PATCH 05/20] Fitted the Workflow read to the Workflow append and close. Workflow code now appends and closes, so a replay skips the Store for those calls too, and a close's Store error comes from a regular Activity. Also sorted an import of the merge and added the changelog entry. --- changelog/added/patient-reader-heron.md | 1 + temporalio/client/_stream_store.py | 2 +- temporalio/worker/_replayer.py | 3 ++- temporalio/worker/_stream_store.py | 2 +- 4 files changed, 5 insertions(+), 3 deletions(-) create mode 100644 changelog/added/patient-reader-heron.md diff --git a/changelog/added/patient-reader-heron.md b/changelog/added/patient-reader-heron.md new file mode 100644 index 0000000000..6b62372d2c --- /dev/null +++ b/changelog/added/patient-reader-heron.md @@ -0,0 +1 @@ +Add `workflow.read_stream` and `workflow.reduce_stream`, which read a Stream from Workflow code in one Local Activity that waits until a predicate matches, the Stream is done, or a limit is reached. diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index d4381b1d75..141c3ceaaa 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -16,8 +16,8 @@ import temporalio.common import temporalio.converter import temporalio.converter._payload_converter -import temporalio.nexus.system._support import temporalio.exceptions +import temporalio.nexus.system._support import temporalio.nexus.system.stream_service import temporalio.runtime import temporalio.workflow diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 701dc8c227..3465a09819 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -70,7 +70,8 @@ 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 for a create or a describe. For each + replay calls no method of its Stores for a create, an append, a close or a + describe. For each :py:func:`temporalio.workflow.read_stream` in the history, it reads the Records of that read from the Store again, so the replay needs the same ``data_converter`` and Stores that still have them. A diff --git a/temporalio/worker/_stream_store.py b/temporalio/worker/_stream_store.py index 1ba16abfa9..7a7245c691 100644 --- a/temporalio/worker/_stream_store.py +++ b/temporalio/worker/_stream_store.py @@ -42,7 +42,7 @@ def _no_stream_config_error() -> temporalio.exceptions.ApplicationError: # pyri @contextlib.contextmanager def _store_errors() -> Iterator[None]: - """Give Workflow code the Store's error of a Local Activity, not its wrapper.""" + """Give Workflow code the Store's error of an Activity, not its wrapper.""" try: yield except temporalio.exceptions.ApplicationError as err: From cdf7cc59ae05aeb088b5da1a8ef4e23ecbea8e6f Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 23:37:01 -0700 Subject: [PATCH 06/20] Pointed the refused Workflow read at `workflow.read_stream`. The refusal told Workflow code to read in an Activity. This branch has a read for Workflow code, so the message names it. --- temporalio/client/_stream_store.py | 5 +++-- tests/nexus/test_stream_store_workflow.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index 141c3ceaaa..1963a23377 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -337,8 +337,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 " + "temporalio.workflow.read_stream or reduce_stream, which read in one " + "Local Activity." ), ), "CloseStream": _StreamOperation( diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py index 07ae119841..003cf1598c 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 " + "temporalio.workflow.read_stream or reduce_stream, which read in one " + "Local Activity." ) From b4cac1d82da84fa0774eea810dba960ebf65b627 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 01:03:29 -0700 Subject: [PATCH 07/20] Failed the Workflow for a read with no Stream config. A read or a fold on a Worker with no Stream config failed the Workflow Task, which retried for good. They raise the same non-retryable `ApplicationError` as the other Workflow Stream calls. --- temporalio/worker/_stream_store.py | 11 +++--- temporalio/worker/_workflow_instance.py | 10 ++---- temporalio/workflow/_streams.py | 16 ++++----- tests/nexus/test_stream_read_workflow.py | 45 ++++++++++++++++-------- 4 files changed, 45 insertions(+), 37 deletions(-) diff --git a/temporalio/worker/_stream_store.py b/temporalio/worker/_stream_store.py index 7a7245c691..3c6720d850 100644 --- a/temporalio/worker/_stream_store.py +++ b/temporalio/worker/_stream_store.py @@ -25,18 +25,15 @@ from ._workflow_instance import _WorkflowInstanceImpl -_NO_STREAM_STORE_MESSAGE = ( - "The client of this Worker has no Stream Store. Give one in " - "stream_config=StreamConfig(stores=[...])." -) - - def _no_stream_config_error() -> temporalio.exceptions.ApplicationError: # pyright: ignore[reportUnusedFunction] # A `StreamConfigError` would fail the Workflow Task, and its retries would fail the # same way for good. A Worker with no config is a setup that a user has to see, so # the Workflow fails instead. return temporalio.exceptions.ApplicationError( - _NO_STREAM_STORE_MESSAGE, type="StreamConfigError", non_retryable=True + "The client of this Worker has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...]).", + type="StreamConfigError", + non_retryable=True, ) diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 470399dad4..8290f04e8a 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -99,11 +99,7 @@ WorkflowOutboundInterceptor, ) from ._stream_read import _read_stream_in_workflow, _reduce_stream_in_workflow -from ._stream_store import ( - _NO_STREAM_STORE_MESSAGE, - _no_stream_config_error, - _WorkflowStreamStore, -) +from ._stream_store import _no_stream_config_error, _WorkflowStreamStore logger = logging.getLogger(__name__) @@ -2009,7 +2005,7 @@ async def workflow_read_stream( limits: temporalio.workflow._StreamReadLimits, ) -> temporalio.workflow.StreamRead[Any]: if self._stream_store is None: - raise temporalio.exceptions.StreamConfigError(_NO_STREAM_STORE_MESSAGE) + raise _no_stream_config_error() return await _read_stream_in_workflow( self, self._stream_store.config, stream, until, limits ) @@ -2024,7 +2020,7 @@ async def workflow_reduce_stream( limits: temporalio.workflow._StreamReadLimits, ) -> temporalio.workflow.StreamReduce[Any]: if self._stream_store is None: - raise temporalio.exceptions.StreamConfigError(_NO_STREAM_STORE_MESSAGE) + raise _no_stream_config_error() return await _reduce_stream_in_workflow( self, self._stream_store.config, diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py index cc5616d11e..8193b6f72f 100644 --- a/temporalio/workflow/_streams.py +++ b/temporalio/workflow/_streams.py @@ -209,10 +209,10 @@ async def read_stream( Raises: temporalio.exceptions.StreamError: The Store failed the read. - temporalio.exceptions.ApplicationError: ``until`` raised. The error has the - message of the exception. - temporalio.exceptions.StreamConfigError: The Worker's client has no Stream - config. This fails the Workflow Task. + temporalio.exceptions.ApplicationError: ``until`` raised, and the error has + the message of the exception. Or the Worker's client has no Stream config, + and the error has the ``type`` ``StreamConfigError``. That one is not + retryable, so an uncaught one fails the Workflow. .. warning:: This API is an experimental prototype and subject to change. @@ -275,10 +275,10 @@ async def reduce_stream( Raises: temporalio.exceptions.StreamError: The Store failed the read. - temporalio.exceptions.ApplicationError: ``function`` or ``until`` raised. The - error has the message of the exception. - temporalio.exceptions.StreamConfigError: The Worker's client has no Stream - config. This fails the Workflow Task. + temporalio.exceptions.ApplicationError: ``function`` or ``until`` raised, and + the error has the message of the exception. Or the Worker's client has no + Stream config, and the error has the ``type`` ``StreamConfigError``. That + one is not retryable, so an uncaught one fails the Workflow. .. warning:: This API is an experimental prototype and subject to change. diff --git a/tests/nexus/test_stream_read_workflow.py b/tests/nexus/test_stream_read_workflow.py index d51c5a5441..c41a8a3270 100644 --- a/tests/nexus/test_stream_read_workflow.py +++ b/tests/nexus/test_stream_read_workflow.py @@ -38,7 +38,7 @@ # 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, assert_task_fail_eventually, new_worker + from tests.helpers import assert_eventually, new_worker from tests.helpers.streams import RecordingStreamStore, StreamReadState from tests.test_converter import SimpleCodec @@ -867,22 +867,37 @@ async def _failing_predicate( ) -async def test_read_without_stream_config_fails_the_task(client: Client): - async with new_worker(client, ReadUntilKindWorkflow) as worker: - handle = await _start_read( - client, - ReadArgs(f"stream-{uuid.uuid4()}"), - worker.task_queue, - execution_timeout=timedelta(seconds=30), - ) - await assert_task_fail_eventually( - handle, - message_contains=( +async def test_read_and_reduce_without_stream_config_fail_the_workflow( + client: Client, +): + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker(client, ReadUntilKindWorkflow, JoinTextWorkflow) as worker: + handles: list[WorkflowHandle[Any, Any]] = [ + await _start_read( + client, + ReadArgs(stream_id), + worker.task_queue, + execution_timeout=timedelta(seconds=30), + ), + await client.start_workflow( + JoinTextWorkflow.run, + ReduceArgs(stream_id), + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=30), + ), + ] + for handle in handles: + 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 + assert cause.message == ( "The client of this Worker has no Stream Store. Give one in " "stream_config=StreamConfig(stores=[...])." - ), - ) - await handle.terminate() + ) @_RUNNERS From c360a9adf5f59044dbfcf89f207e0bff0cc43b78 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 01:45:54 -0700 Subject: [PATCH 08/20] Replaced the Workflow read with take, filter and reduce calls. The three calls share one scan, which Workflow code runs in a Local Activity and other code runs against the Store. The result goes into History, so a replay calls no Store and needs no config. A counter numbers the calls, and a result over 1 MiB fails at the await. --- changelog/added/patient-reader-heron.md | 2 +- temporalio/client/__init__.py | 22 + temporalio/client/_stream_read.py | 770 ++++++++++++---------- temporalio/client/_stream_store.py | 101 +-- temporalio/worker/_replayer.py | 11 +- temporalio/worker/_stream_read.py | 253 +++---- temporalio/worker/_worker.py | 3 + temporalio/worker/_workflow.py | 11 +- temporalio/worker/_workflow_instance.py | 45 +- temporalio/workflow/__init__.py | 18 +- temporalio/workflow/_context.py | 25 +- temporalio/workflow/_streams.py | 540 +++++++++++---- tests/nexus/test_stream_store_workflow.py | 4 +- 13 files changed, 1000 insertions(+), 805 deletions(-) diff --git a/changelog/added/patient-reader-heron.md b/changelog/added/patient-reader-heron.md index 6b62372d2c..0508d5456f 100644 --- a/changelog/added/patient-reader-heron.md +++ b/changelog/added/patient-reader-heron.md @@ -1 +1 @@ -Add `workflow.read_stream` and `workflow.reduce_stream`, which read a Stream from Workflow code in one Local Activity that waits until a predicate matches, the Stream is done, or a limit is reached. +Add `take_stream`, `filter_stream` and `reduce_stream` to `temporalio.workflow` and `temporalio.client`, which client, Activity and Workflow code call to scan a Stream until a predicate matches, the Stream is done, or a limit is reached. diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index 6fa6dd710c..eeebf78685 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -27,6 +27,18 @@ ServiceClient, TLSConfig, ) +from temporalio.workflow._streams import ( + DEFAULT_STREAM_IDLE_TIMEOUT, + DEFAULT_STREAM_MAX_BYTES, + DEFAULT_STREAM_MAX_RECORDS, + MAX_STREAM_RESULT_BYTES, + StreamRead, + StreamReadStop, + StreamReduce, + filter_stream, + reduce_stream, + take_stream, +) from ..common import HeaderCodecBehavior from ..types import ( @@ -356,6 +368,16 @@ "TaskReachabilityType", "CloudOperationsClient", "Plugin", + "DEFAULT_STREAM_IDLE_TIMEOUT", + "DEFAULT_STREAM_MAX_BYTES", + "DEFAULT_STREAM_MAX_RECORDS", + "MAX_STREAM_RESULT_BYTES", + "StreamRead", + "StreamReadStop", + "StreamReduce", + "filter_stream", + "reduce_stream", + "take_stream", "StreamConfig", "StreamSelectContext", "StreamStore", diff --git a/temporalio/client/_stream_read.py b/temporalio/client/_stream_read.py index d0b36bdb6e..ea87c19ceb 100644 --- a/temporalio/client/_stream_read.py +++ b/temporalio/client/_stream_read.py @@ -1,10 +1,10 @@ -"""The process side of a Stream read or fold that Workflow code waits on. +"""The scan of :py:func:`temporalio.workflow.take_stream`, ``filter_stream`` and +``reduce_stream``, for client, Activity and Workflow code. -A Workflow's read runs in a Local Activity, and the Local Activity calls functions of -the Workflow. A function is a closure and can't be serialized, and the Records must -stay out of History. So both pass through the registry of this module, in the memory -of the Worker process. History keeps only a small marker with the range of offsets of -the read, and the value of a fold. +Client and Activity code run the scan straight against the Store. Workflow code runs it in a +Local Activity, and the result of that Local Activity is the result of the call. 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. .. warning:: This API is an experimental prototype and subject to change. @@ -14,422 +14,528 @@ import base64 import dataclasses -import hashlib +import inspect import json -import struct import threading import time -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable from datetime import timedelta from typing import Any +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 -from temporalio.workflow._streams import StreamReadStop - -_MARKER_METADATA_KEY = "temporal-stream-read" -_MARKER_METADATA_VALUE = b"v2" +import temporalio.nexus.system +import temporalio.nexus.system.stream_service.client +import temporalio.workflow +from temporalio.nexus.system.stream_service.models import ( + ReadRecordsRequest, + StreamRecord, +) +from temporalio.workflow._streams import ( + MAX_STREAM_RESULT_BYTES, + StreamRead, + StreamReadStop, + StreamReduce, + _StreamCall, + _StreamCallKind, +) + +from ._client import Client +from ._stream_store import ( + _STREAM_SERVICE, + StreamStore, + _activity_target, + _store_page_reader, +) + +_RESULT_METADATA_KEY = "temporal-stream-read" +_RESULT_METADATA_VALUE = b"v3" +_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 -_DIGEST_SIZE = 16 +_ReadPage = Callable[ + [temporalio.api.streamservice.v1.ReadRecordsRequest], + Awaitable[temporalio.api.streamservice.v1.ReadRecordsResponse], +] +_ToRecord = Callable[[temporalio.api.stream.v1.StreamRecord], Any] -def _encode_bytes(value: bytes | None) -> str | None: + +def _encode_offset(value: bytes | None) -> str | None: return None if value is None else base64.b64encode(value).decode() -def _decode_bytes(value: str | None) -> bytes | None: +def _decode_offset(value: str | None) -> bytes | None: return None if value is None else base64.b64decode(value) -class _RecordDigest: - """A SHA-256 over the offsets and bodies of Records, in the order that they come. +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 - A length before each part keeps two different splits of the same bytes apart. - """ - def __init__(self) -> None: - self._hash = hashlib.sha256() +@dataclasses.dataclass +class _ScanOutcome: + kept: list[temporalio.api.stream.v1.StreamRecord] + """The Records that the call gives back. A fold keeps none.""" + 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: + """What a scan got so far.""" + + def __init__( + self, call: _StreamCall, to_record: _ToRecord, *, in_local_activity: bool + ) -> None: + self._call = call + self._to_record = to_record + self._in_local_activity = in_local_activity + self.kept: list[temporalio.api.stream.v1.StreamRecord] = [] + 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: + return self._call.max_bytes - self.bytes + + def look( + self, record: temporalio.api.stream.v1.StreamRecord, *, 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._to_record(record) + keep = True + if call.predicate is not None: + keep = bool(self._call_function("predicate", call.predicate, typed)) + size = record.ByteSize() + # The first Record of a call is taken even when it passes the cap. + if keep and self.count > 0 and self.bytes + size > call.max_bytes: + return StreamReadStop.MaxBytes + self.last_offset = record.offset + 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) + 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: + return _ScanOutcome( + kept=self.kept, + count=self.count, + last_offset=self.last_offset, + stop=stop, + done=done, + value=self.value, + ) - def update(self, record: temporalio.api.stream.v1.StreamRecord) -> None: - body = record.body.SerializeToString(deterministic=True) - for part in (record.offset, body): - self._hash.update(struct.pack(">Q", len(part))) - self._hash.update(part) - def digest(self) -> bytes: - return self._hash.digest()[:_DIGEST_SIZE] +async def _scan( + read_page: _ReadPage, + *, + ref: temporalio.api.stream.v1.StreamRef, + call: _StreamCall, + to_record: _ToRecord, + in_local_activity: bool, +) -> _ScanOutcome: + """Read pages after ``call.from_offset`` until a stop of :py:class:`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, to_record, 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 -def _records_digest( - records: Sequence[temporalio.api.stream.v1.StreamRecord], -) -> bytes: - digest = _RecordDigest() - for record in records: - digest.update(record) - return digest.digest() + while True: + if page_cap and state.count > 0 and state.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) + ) + request = temporalio.api.streamservice.v1.ReadRecordsRequest( + ref=ref, + from_offset=state.last_offset or b"", + max_records=min(call.max_records - state.count, _MAX_PAGE_RECORDS) + if page_cap + else _MAX_PAGE_RECORDS, + max_bytes=max(state.bytes_left(), 1) if page_cap else 0, + ) + request.wait.FromTimedelta(timedelta(seconds=wait)) + response = await read_page(request) + if response.records: + last_arrival = time.monotonic() + for index, record in enumerate(response.records): + # A done page ends with the last Record of the Stream. + ends_stream = response.done and index == len(response.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 response.done and not response.records: + return state.end(StreamReadStop.Done, done=True) + if not response.records and idle(time.monotonic()): + return state.end(StreamReadStop.Idle) @dataclasses.dataclass(frozen=True) -class _StreamReadMarker: - """The result of a read's Local Activity, and so the only part of it in History.""" +class _StreamCallResult: + """The result of one call, as the payload that History keeps for a Workflow.""" - read_id: str + kind: _StreamCallKind store: str stream_id: str from_offset: bytes | None last_offset: bytes | None - record_count: int + count: int stop: StreamReadStop done: bool - digest: bytes - value: temporalio.api.common.v1.Payload | None = None - """The value of a fold. A read of Records has none.""" - - def to_raw_value(self) -> temporalio.common.RawValue: - """Give the marker as a payload that a Worker can find in any activation.""" - data = { - "read_id": self.read_id, + 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.""" + + def to_payload(self) -> temporalio.api.common.v1.Payload: + """Give the result as one payload. The bodies stay in their binary form.""" + header = { + "kind": self.kind.value, "store": self.store, "stream_id": self.stream_id, - "from_offset": _encode_bytes(self.from_offset), - "last_offset": _encode_bytes(self.last_offset), - "record_count": self.record_count, - "stop": self.stop.value, + "from_offset": _encode_offset(self.from_offset), + "last_offset": _encode_offset(self.last_offset), + "count": self.count, + "stop": self.stop.name, "done": self.done, - "digest": self.digest.hex(), - "value": _encode_bytes( - None - if self.value is None - else self.value.SerializeToString(deterministic=True) - ), } - return temporalio.common.RawValue( + parts = [ temporalio.api.common.v1.Payload( - metadata={ - "encoding": b"json/plain", - _MARKER_METADATA_KEY: _MARKER_METADATA_VALUE, - }, - data=json.dumps(data, separators=(",", ":"), sort_keys=True).encode(), + metadata={"encoding": b"json/plain"}, + data=json.dumps(header, separators=(",", ":"), sort_keys=True).encode(), + ) + ] + if self.value is not None: + parts.append(self.value) + else: + 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={ + "encoding": b"binary/plain", + _RESULT_METADATA_KEY: _RESULT_METADATA_VALUE, + }, + data=temporalio.api.common.v1.Payloads(payloads=parts).SerializeToString(), ) @classmethod def from_payload( cls, payload: temporalio.api.common.v1.Payload - ) -> _StreamReadMarker | None: - """Give the marker in ``payload``, or ``None`` when the payload isn't one.""" - if payload.metadata.get(_MARKER_METADATA_KEY) != _MARKER_METADATA_VALUE: + ) -> _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 - data = json.loads(payload.data) - value = _decode_bytes(data["value"]) + header_part, body = temporalio.api.common.v1.Payloads.FromString( + payload.data + ).payloads + header = json.loads(header_part.data) + kind = _StreamCallKind(header["kind"]) + records: list[temporalio.api.stream.v1.StreamRecord] = [] + value = None + if kind is _StreamCallKind.REDUCE: + value = body + else: + records = list( + temporalio.api.streamservice.v1.ReadRecordsResponse.FromString( + body.data + ).records + ) return cls( - read_id=data["read_id"], - store=data["store"], - stream_id=data["stream_id"], - from_offset=_decode_bytes(data["from_offset"]), - last_offset=_decode_bytes(data["last_offset"]), - record_count=data["record_count"], - stop=StreamReadStop(data["stop"]), - done=data["done"], - digest=bytes.fromhex(data["digest"]), - value=None - if value is None - else temporalio.api.common.v1.Payload.FromString(value), - ) - - def describe_range(self) -> str: - """Name the Stream and the range of the read, for an error message.""" - start = self.from_offset.hex() if self.from_offset else "the beginning" - end = self.last_offset.hex() if self.last_offset else "the beginning" - return ( - f"the Stream '{self.stream_id}' on the Store '{self.store}', " - f"{self.record_count} Records after {start} up to {end}" + 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, ) -@dataclasses.dataclass(frozen=True) -class _StreamFold: - """The function and the start value of a fold.""" - - function: Callable[[Any, Any], Any] - initial: Any +def _result_payload( + *, + call: _StreamCall, + store: str, + stream_id: str, + outcome: _ScanOutcome, + payload_converter: temporalio.converter.PayloadConverter, +) -> temporalio.api.common.v1.Payload: + """Give the result as a payload, or raise when it is over the limit.""" + payload = _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, + value=payload_converter.to_payload(outcome.value) + if call.kind is _StreamCallKind.REDUCE + else None, + ).to_payload() + size = payload.ByteSize() + if size > MAX_STREAM_RESULT_BYTES: + 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 {MAX_STREAM_RESULT_BYTES} bytes. Lower " + "max_bytes or max_records, or fold to a smaller value.", + type="StreamResultTooLarge", + non_retryable=True, + ) + return payload @dataclasses.dataclass(eq=False) -class _PendingStreamRead: - """One read of Workflow code, from the start of its Local Activity to its result.""" - - to_record: Callable[[temporalio.api.stream.v1.StreamRecord], Any] - """Turns a proto Record into the typed Record of the Workflow.""" - - until: Callable[..., bool] | None - """The predicate of the Workflow, or ``None``. +class _PendingStreamCall: + """One call of Workflow code, from the start of its Local Activity to its result. - A read calls it with the typed Record, and a fold with the value and the Record. + It holds only what can't go into the Local Activity's input. """ - fold: _StreamFold | None = None - """The fold, or ``None`` for a read that gives its Records.""" + call: _StreamCall + """The checked call, with the caller's functions.""" - _result: ( - tuple[_StreamReadMarker, list[temporalio.api.stream.v1.StreamRecord]] | None - ) = None + store: StreamStore + """The Store that the call reads.""" - def fill( - self, - marker: _StreamReadMarker, - records: list[temporalio.api.stream.v1.StreamRecord], - ) -> None: - """Keep the Records of ``marker``, whose payloads the codec already decoded.""" - # One assignment, so a reader on another thread sees both or neither. - self._result = (marker, records) - - def records_for( - self, marker: _StreamReadMarker - ) -> list[temporalio.api.stream.v1.StreamRecord] | None: - """Give the Records of ``marker``, or ``None`` when this read has none for it.""" - result = self._result - if result is None or result[0] != marker: - return None - return result[1] + ref: temporalio.api.stream.v1.StreamRef + to_record: _ToRecord + """Turns a proto Record into the typed Record of the Workflow.""" -class _PendingStreamReads: - """The pending reads of the process, by Workflow run id and read id. - Workflow code adds and removes an entry on its own thread. The read Activity and the - Worker look an entry up on the event loop of the Worker. +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 looks an + entry up on the event loop of the Worker. """ def __init__(self) -> None: self._lock = threading.Lock() - self._reads: dict[tuple[str, str], _PendingStreamRead] = {} + self._calls: dict[tuple[str, str], _PendingStreamCall] = {} - def add(self, run_id: str, read_id: str, read: _PendingStreamRead) -> None: + def add(self, run_id: str, call_id: str, call: _PendingStreamCall) -> None: with self._lock: - self._reads[(run_id, read_id)] = read + self._calls[(run_id, call_id)] = call - def get(self, run_id: str, read_id: str) -> _PendingStreamRead | None: + def get(self, run_id: str, call_id: str) -> _PendingStreamCall | None: with self._lock: - return self._reads.get((run_id, read_id)) + return self._calls.get((run_id, call_id)) - def remove(self, run_id: str, read_id: str, read: _PendingStreamRead) -> None: + 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._reads.get((run_id, read_id)) is read: - del self._reads[(run_id, read_id)] + 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._reads if key[0] == run_id]: - del self._reads[key] + 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._reads) - - -_PENDING_STREAM_READS = _PendingStreamReads() - -_ReadPage = Callable[ - [temporalio.api.streamservice.v1.ReadRecordsRequest], - Awaitable[temporalio.api.streamservice.v1.ReadRecordsResponse], -] - - -@dataclasses.dataclass(frozen=True) -class _ReadLimits: - """The stops of a read, as the read Activity gets them.""" - - max_records: int - max_bytes: int | None - idle_timeout: float | None - """Seconds with no new Record before the read ends, or ``None``.""" - timeout: float - """Seconds for the whole read.""" - - -@dataclasses.dataclass(frozen=True) -class _ReadOutcome: - records: list[temporalio.api.stream.v1.StreamRecord] - """The Records of a read. A fold keeps none.""" - record_count: int - last_offset: bytes | None - digest: bytes - stop: StreamReadStop - done: bool - value: Any - """The value of a fold, or ``None`` for a read.""" + return list(self._calls) -class _ReadState: - """What a read got so far, for one attempt of its Activity.""" +_PENDING_STREAM_CALLS = _PendingStreamCalls() - def __init__(self, read: _PendingStreamRead, from_offset: bytes | None) -> None: - self._read = read - self._digest = _RecordDigest() - self.records: list[temporalio.api.stream.v1.StreamRecord] = [] - self.count = 0 - self.bytes = 0 - self.last_offset = from_offset - self.value = read.fold.initial if read.fold is not None else None - - def take(self, record: temporalio.api.stream.v1.StreamRecord, size: int) -> bool: - """Add ``record`` to the read, and give whether ``until`` is true after it.""" - self.count += 1 - self.bytes += size - self.last_offset = record.offset - self._digest.update(record) - read = self._read - if read.fold is None: - self.records.append(record) - if read.until is None: - return False - return bool(_call_workflow_function(read.until, read.to_record(record))) - typed = read.to_record(record) - self.value = _call_workflow_function(read.fold.function, self.value, typed) - if read.until is None: - return False - return bool(_call_workflow_function(read.until, self.value, typed)) - - def end(self, stop: StreamReadStop, *, done: bool = False) -> _ReadOutcome: - return _ReadOutcome( - records=self.records, - record_count=self.count, - last_offset=self.last_offset, - digest=self._digest.digest(), - stop=stop, - done=done, - value=self.value, - ) +_STREAM_CALL_ACTIVITY = "__temporal_stream_call" -def _call_workflow_function(function: Callable[..., Any], *args: Any) -> Any: - try: - return function(*args) - except Exception as err: - # The Workflow gets this error from its Local Activity, in Workflow code. +@temporalio.activity.defn(name=_STREAM_CALL_ACTIVITY) +async def _stream_call_activity(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( - str(err), type=type(err).__name__, non_retryable=True - ) from err - - -async def _read_until( # pyright: ignore[reportUnusedFunction] - read_page: _ReadPage, - *, - ref: temporalio.api.stream.v1.StreamRef, - from_offset: bytes | None, - limits: _ReadLimits, - read: _PendingStreamRead, -) -> _ReadOutcome: - """Read pages until a stop of :py:class:`StreamReadStop` holds.""" - start = time.monotonic() - deadline = start + limits.timeout - # The idle stop counts from the last page that gave a Record. - last_arrival = start - state = _ReadState(read, from_offset) - - def idle(now: float) -> bool: - return ( - limits.idle_timeout is not None - and now - last_arrival >= limits.idle_timeout + f"The Stream call {call_id} has no pending call in this process. Only " + "take_stream, filter_stream and reduce_stream in Workflow code run this " + "Activity.", + non_retryable=True, ) - - while True: - if state.count >= limits.max_records: - return state.end(StreamReadStop.MAX_RECORDS) - bytes_left = None - if limits.max_bytes is not None: - bytes_left = limits.max_bytes - state.bytes - # Every Record has some bytes, so none fits in a cap that is used up. - if bytes_left <= 0: - return state.end(StreamReadStop.MAX_BYTES) - now = time.monotonic() - if now >= deadline: - return state.end( - StreamReadStop.IDLE if idle(now) else StreamReadStop.TIMEOUT - ) - wait = min(_MAX_PAGE_WAIT, deadline - now) - if limits.idle_timeout is not None: - wait = min(wait, max(0.0, last_arrival + limits.idle_timeout - now)) - request = temporalio.api.streamservice.v1.ReadRecordsRequest( - ref=ref, - from_offset=state.last_offset or b"", - max_records=min(limits.max_records - state.count, _MAX_PAGE_RECORDS), - max_bytes=bytes_left or 0, + # The codec decodes the bodies of each page, so the caller's functions get its values. + outcome = await _scan( + _store_page_reader(pending.store, temporalio.activity.client().data_converter), + ref=pending.ref, + call=pending.call, + to_record=pending.to_record, + in_local_activity=True, + ) + return temporalio.common.RawValue( + _result_payload( + call=pending.call, + store=pending.store.name, + stream_id=pending.ref.id, + outcome=outcome, + payload_converter=temporalio.activity.payload_converter(), ) - request.wait.FromTimedelta(timedelta(seconds=wait)) - response = await read_page(request) - if response.records: - last_arrival = time.monotonic() - for index, record in enumerate(response.records): - if state.count >= limits.max_records: - return state.end(StreamReadStop.MAX_RECORDS) - size = record.ByteSize() - # A Store gives the first Record of a page even when it passes the cap, and - # the read keeps it only when it is the first Record of the read. - if ( - limits.max_bytes is not None - and state.count > 0 - and state.bytes + size > limits.max_bytes - ): - return state.end(StreamReadStop.MAX_BYTES) - matched = state.take(record, size) - # A done page ends with the last Record of the Stream. - done = response.done and index == len(response.records) - 1 - if matched: - return state.end(StreamReadStop.MATCH, done=done) - if done: - return state.end(StreamReadStop.DONE, done=True) - if response.done and not response.records: - return state.end(StreamReadStop.DONE, done=True) - if not response.records and idle(time.monotonic()): - return state.end(StreamReadStop.IDLE) - - -async def _read_range( # pyright: ignore[reportUnusedFunction] - read_page: _ReadPage, marker: _StreamReadMarker -) -> list[temporalio.api.stream.v1.StreamRecord]: - """Read the Records of ``marker`` again, and check them against it. - - Raises: - RuntimeError: The Store failed, the Records are gone, or they don't match. - """ - records: list[temporalio.api.stream.v1.StreamRecord] = [] - cursor = marker.from_offset - ref = temporalio.api.stream.v1.StreamRef(id=marker.stream_id, store=marker.store) - while len(records) < marker.record_count: - left = marker.record_count - len(records) - request = temporalio.api.streamservice.v1.ReadRecordsRequest( - ref=ref, - from_offset=cursor or b"", - max_records=min(left, _MAX_PAGE_RECORDS), + ) + + +def _stream_client(stream: Any) -> Client: + if isinstance(stream, temporalio.nexus.system.stream_service.client.Stream): + client = getattr(getattr(stream, "_nexus_client", None), "_client", None) + if isinstance(client, Client): + return client + raise TypeError("The client Stream has no Temporal client") + if temporalio.activity.in_activity(): + # A Workflow Stream outside Workflow code is one that a function of a call or an + # Activity got, so the Activity's client reaches its Store. + return temporalio.activity.client() + 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 straight against the Store.""" + client = _stream_client(stream) + stream_config = client.config(active_config=True).get("stream_config") + if stream_config is None: + raise temporalio.exceptions.StreamConfigError( + "This client has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...])." ) - try: - response = await read_page(request) - except Exception as err: - raise RuntimeError( - f"The replay of a Workflow read can't read {marker.describe_range()}: " - f"{err}" - ) from err - if not response.records: - raise RuntimeError( - f"The replay of a Workflow read can't read {marker.describe_range()}: " - f"the Store gave {len(records)} Records, and the Records after them are " - "gone" - ) - records.extend(list(response.records)[:left]) - cursor = records[-1].offset - last_offset = records[-1].offset if records else marker.from_offset - if last_offset != marker.last_offset or _records_digest(records) != marker.digest: - raise RuntimeError( - f"The replay of a Workflow read read {marker.describe_range()}, and the " - "Records differ from the ones of the first run" + request = ReadRecordsRequest(ref=stream.ref, from_offset=call.from_offset) + data_converter = client.data_converter + serialization_context = temporalio.nexus.system._get_serialization_context( + _STREAM_SERVICE.name, "ReadRecords", request + ) + if serialization_context is not None: + data_converter = data_converter.with_context(serialization_context) + payload_converter = temporalio.nexus.system._get_payload_converter( + data_converter.payload_converter, data_converter.failure_converter + ) + message = payload_converter._to_transfer_value(request) + assert isinstance(message, temporalio.api.streamservice.v1.ReadRecordsRequest) + store = stream_config._store_for(stream.ref, _activity_target()) + 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) + + outcome = await _scan( + _store_page_reader(store, data_converter), + ref=message.ref, + call=call, + to_record=to_record, + in_local_activity=False, + ) + # The limit is the same as in Workflow code, so a call gives the same result anywhere. + _result_payload( + call=call, + store=store.name, + stream_id=message.ref.id, + outcome=outcome, + payload_converter=data_converter.payload_converter, + ) + 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 records + return StreamRead( + records=[to_record(record) for record in 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 1963a23377..f3bd345a55 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -16,17 +16,10 @@ import temporalio.common import temporalio.converter import temporalio.converter._payload_converter -import temporalio.exceptions import temporalio.nexus.system._support import temporalio.nexus.system.stream_service import temporalio.runtime import temporalio.workflow -from temporalio.client._stream_read import ( - _PENDING_STREAM_READS, - _read_until, - _ReadLimits, - _StreamReadMarker, -) from temporalio.converter import StorageDriverActivityInfo, StorageDriverWorkflowInfo from temporalio.exceptions import StreamConfigError @@ -145,15 +138,13 @@ async def describe_stream( def worker_activities(self) -> Sequence[Callable[..., Any]]: """The Activities that a Worker registers for this Store when it runs Workflows. - The default is two Activities, named for this Store. One answers - :py:meth:`call_from_workflow`, which runs it as a Local Activity, or as a regular - Activity for a close. The other answers :py:func:`temporalio.workflow.read_stream` - and :py:func:`temporalio.workflow.reduce_stream`, which always run it by its name. - A Store that overrides :py:meth:`call_from_workflow` overrides this method to - match. A Store that overrides this method keeps the read Activity in its list for - Workflow code that reads. + 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. 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), _store_read_activity(self)] + return [_store_call_activity(self)] async def call_from_workflow( self, @@ -338,8 +329,8 @@ def check_response( temporalio.api.streamservice.v1.ReadRecordsResponse, workflow_refusal=( "Workflow code can't run a ReadRecords operation on a Store. Use " - "temporalio.workflow.read_stream or reduce_stream, which read in one " - "Local Activity." + "temporalio.workflow.take_stream, filter_stream or reduce_stream, which " + "read in one Local Activity." ), ), "CloseStream": _StreamOperation( @@ -405,82 +396,6 @@ def _with_local_activity_producer( request.sequence = 1 -_STORE_READ_ACTIVITY = "__temporal_stream_store_read" - - -def _store_read_activity_name(store_name: str) -> str: - return f"{_STORE_READ_ACTIVITY}:{store_name}" - - -def _store_read_activity( - store: StreamStore, -) -> Callable[ - [str, bytes, int, int | None, float | None, float], - Awaitable[temporalio.common.RawValue], -]: - @temporalio.activity.defn(name=_store_read_activity_name(store.name)) - async def read( - read_id: str, - request: bytes, - max_records: int, - max_bytes: int | None, - idle_timeout_seconds: float | None, - timeout_seconds: float, - ) -> temporalio.common.RawValue: - info = temporalio.activity.info() - pending = _PENDING_STREAM_READS.get(info.workflow_run_id or "", read_id) - if pending is None: - # The functions of the Workflow live in its process, so only a Local - # Activity of that Workflow can run this read. - raise temporalio.exceptions.ApplicationError( - f"The Stream read '{read_id}' has no pending read in this process. " - "Only temporalio.workflow.read_stream and reduce_stream run this " - "Activity.", - non_retryable=True, - ) - read_request = temporalio.api.streamservice.v1.ReadRecordsRequest.FromString( - request - ) - # The codec decodes the bodies in each page, so the functions of the Workflow - # get the user's values. - page_reader = _store_page_reader( - store, temporalio.activity.client().data_converter - ) - from_offset = read_request.from_offset or None - outcome = await _read_until( - page_reader, - ref=read_request.ref, - from_offset=from_offset, - limits=_ReadLimits( - max_records=max_records, - max_bytes=max_bytes, - idle_timeout=idle_timeout_seconds, - timeout=timeout_seconds, - ), - read=pending, - ) - marker = _StreamReadMarker( - read_id=read_id, - store=store.name, - stream_id=read_request.ref.id, - from_offset=from_offset, - last_offset=outcome.last_offset, - record_count=outcome.record_count, - stop=outcome.stop, - done=outcome.done, - digest=outcome.digest, - # History keeps the value of a fold, so a replay needs no Store. - value=temporalio.activity.payload_converter().to_payload(outcome.value) - if pending.fold is not None - else None, - ) - if pending.fold is None: - pending.fill(marker, outcome.records) - return marker.to_raw_value() - - return read - - def _store_page_reader( store: StreamStore, data_converter: temporalio.converter.DataConverter ) -> Callable[ diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 3465a09819..e3b4b6d7fa 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -70,13 +70,10 @@ 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 for a create, an append, a close or a - describe. For each - :py:func:`temporalio.workflow.read_stream` in the history, it reads the - Records of that read from the Store again, so the replay needs the same - ``data_converter`` and Stores that still have them. A - :py:func:`temporalio.workflow.reduce_stream` takes its value from the - history and calls no Store. + replay calls no method of its Stores. A + :py:func:`temporalio.workflow.take_stream`, ``filter_stream`` or + ``reduce_stream`` takes its result from the history, so its replay needs no + Stream config. 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 index f30817b149..10a3d57f79 100644 --- a/temporalio/worker/_stream_read.py +++ b/temporalio/worker/_stream_read.py @@ -1,193 +1,140 @@ -"""The Worker side of :py:func:`temporalio.workflow.read_stream` and ``reduce_stream``. - -Workflow code starts the read here, and the Worker gives a replayed read its Records. -""" +"""The Workflow side of :py:func:`temporalio.workflow.take_stream`, ``filter_stream`` and +``reduce_stream``.""" from __future__ import annotations -from collections.abc import Callable from datetime import timedelta from typing import TYPE_CHECKING, Any import temporalio.api.stream.v1 -import temporalio.bridge.proto.workflow_activation +import temporalio.api.streamservice.v1 import temporalio.common -import temporalio.converter -import temporalio.exceptions import temporalio.workflow from temporalio.client._stream_read import ( - _PENDING_STREAM_READS, - _PendingStreamRead, - _read_range, - _StreamFold, - _StreamReadMarker, -) -from temporalio.client._stream_store import ( - _STREAM_SERVICE, - StreamConfig, - _store_page_reader, - _store_read_activity_name, + _PENDING_STREAM_CALLS, + _STREAM_CALL_ACTIVITY, + _PendingStreamCall, + _StreamCallResult, ) +from temporalio.client._stream_store import _STREAM_SERVICE, StreamConfig from temporalio.nexus.system.stream_service.models import ( ReadRecordsRequest, StreamRecord, ) from temporalio.nexus.system.stream_service.workflow._resources.stream import Stream +from temporalio.workflow._streams import ( + StreamRead, + StreamReduce, + _StreamCall, + _StreamCallKind, +) -from ._stream_store import _store_errors, _workflow_target +from ._stream_store import _no_stream_config_error, _store_errors, _workflow_target if TYPE_CHECKING: from ._workflow_instance import _WorkflowInstanceImpl -# The read Activity gets this much time after the read's own timeout to give its result. -_READ_ACTIVITY_GRACE = timedelta(seconds=10) +# 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 _run_read( + +async def _stream_call_in_workflow( # pyright: ignore[reportUnusedFunction] instance: _WorkflowInstanceImpl, - config: StreamConfig, - stream: Stream[Any], - limits: temporalio.workflow._StreamReadLimits, - *, - until: Callable[..., bool] | None, - fold: _StreamFold | None, -) -> tuple[ - _StreamReadMarker, - list[temporalio.api.stream.v1.StreamRecord], - Callable[[temporalio.api.stream.v1.StreamRecord], Any], -]: - """Run the read Activity, and give its marker, its proto Records and their typing.""" - request = ReadRecordsRequest(ref=stream.ref, from_offset=limits.from_offset) + 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 ) - request_message = payload_converter._to_transfer_value(request) - # The choice runs in Workflow code, as it does for the other Stream calls. - store = config._store_for(stream.ref, _workflow_target()) + 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 takes the result from History, so it 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 - read_id = str(temporalio.workflow.uuid4()) - pending = _PendingStreamRead(to_record=to_record, until=until, fold=fold) - _PENDING_STREAM_READS.add(run_id, read_id, pending) + pending = None + if store is not None: + pending = _PendingStreamCall( + call=call, store=store, ref=message.ref, to_record=to_record + ) + _PENDING_STREAM_CALLS.add(run_id, call_id, pending) try: with _store_errors(): - result = await temporalio.workflow.execute_local_activity( - _store_read_activity_name(store.name), - args=[ - read_id, - request_message.SerializeToString(), - limits.max_records, - limits.max_bytes, - None - if limits.idle_timeout is None - else limits.idle_timeout.total_seconds(), - limits.timeout.total_seconds(), - ], + raw = await temporalio.workflow.execute_local_activity( + _STREAM_CALL_ACTIVITY, + call_id, result_type=temporalio.common.RawValue, - start_to_close_timeout=limits.timeout + _READ_ACTIVITY_GRACE, - retry_policy=config.workflow_call_options.retry_policy, - ) - marker = _StreamReadMarker.from_payload(result.payload) - if marker is None or marker.read_id != read_id: - raise RuntimeError( - f"The read of the Stream '{stream.ref.id}' got a result that is not " - "its read marker" - ) - if fold is not None: - return marker, [], to_record - records = pending.records_for(marker) - if records is None: - raise RuntimeError( - f"The Worker has no Records for the read of {marker.describe_range()}" + start_to_close_timeout=call.timeout + _STREAM_CALL_GRACE, + retry_policy=config.workflow_call_options.retry_policy + if config is not None + else None, ) - return marker, records, to_record finally: - _PENDING_STREAM_READS.remove(run_id, read_id, pending) - - -async def _read_stream_in_workflow( # pyright: ignore[reportUnusedFunction] - instance: _WorkflowInstanceImpl, - config: StreamConfig, - stream: Stream[Any], - until: Callable[[StreamRecord[Any]], bool] | None, - limits: temporalio.workflow._StreamReadLimits, -) -> temporalio.workflow.StreamRead[Any]: - marker, records, to_record = await _run_read( - instance, config, stream, limits, until=until, fold=None - ) - return temporalio.workflow.StreamRead( - records=[to_record(record) for record in records], - last_offset=marker.last_offset, - stop=marker.stop, - done=marker.done, - ) - - -async def _reduce_stream_in_workflow( # pyright: ignore[reportUnusedFunction] - instance: _WorkflowInstanceImpl, - config: StreamConfig, - stream: Stream[Any], - fold: _StreamFold, - until: Callable[[Any, StreamRecord[Any]], bool] | None, - value_type: Any, - limits: temporalio.workflow._StreamReadLimits, -) -> temporalio.workflow.StreamReduce[Any]: - marker, _, _ = await _run_read( - instance, config, stream, limits, until=until, fold=fold - ) - if marker.value is None: - raise RuntimeError( - f"The fold of {marker.describe_range()} got a marker with no value" + 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 + ): + recorded = ( + "no Stream call" + if result is None + else _describe( + result.kind, result.store, result.stream_id, result.from_offset + ) ) - # A first run and a replay both take the value from the marker. - value = instance.workflow_payload_converter().from_payload(marker.value, value_type) - return temporalio.workflow.StreamReduce( - value=value, - count=marker.record_count, - last_offset=marker.last_offset, - stop=marker.stop, - done=marker.done, + 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, + ) + return StreamRead( + records=[to_record(record) for record in result.records], + last_offset=result.last_offset, + stop=result.stop, + done=result.done, ) - - -async def _fill_stream_reads( # pyright: ignore[reportUnusedFunction] - act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, - config: StreamConfig, - data_converter: temporalio.converter.DataConverter, -) -> None: - """Give each read that this activation resolves its Records, before the Workflow runs. - - History has only the marker of a read. On a replay, the Records come from the Store - again, through the same codec as the read Activity used. - - Raises: - RuntimeError: The Store failed, the Records are gone, or they don't match the - marker. The activation fails, and so does the Workflow Task. - """ - 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 - marker = _StreamReadMarker.from_payload(resolve.result.completed.result) - # History has the value of a fold, so its replay needs no Store. - if marker is None or marker.value is not None: - continue - pending = _PENDING_STREAM_READS.get(act.run_id, marker.read_id) - # On a first run the read Activity already gave its Records. - if pending is None or pending.records_for(marker) is not None: - continue - try: - store = config._get_store_by_name(marker.store) - except temporalio.exceptions.StreamConfigError as err: - raise RuntimeError( - f"The replay of a Workflow read can't read {marker.describe_range()}: " - f"{err}" - ) from err - page_reader = _store_page_reader(store, data_converter) - pending.fill(marker, await _read_range(page_reader, marker)) diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 4ecce29d15..95ce14304c 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 f7de6f791a..40294c256c 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -27,7 +27,7 @@ import temporalio.exceptions import temporalio.workflow from temporalio.bridge.worker import PollShutdownError -from temporalio.client._stream_read import _PENDING_STREAM_READS +from temporalio.client._stream_read import _PENDING_STREAM_CALLS from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo from temporalio.worker.workflow_sandbox._runner import SandboxedWorkflowRunner @@ -41,7 +41,6 @@ WorkflowInboundInterceptor, WorkflowInterceptorClassInput, ) -from ._stream_read import _fill_stream_reads from ._stream_store import _WorkflowStreamStore from ._workflow_instance import ( _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS, @@ -434,12 +433,6 @@ async def _handle_activation( await self._decode_nexus_previous_run_outcome( act.run_id, init_job, previous_run_outcome, download_metrics ) - if self._stream_store is not None: - # A replay has only the marker of a Stream read, so the Records come - # from the Store before the Workflow sees the result. - await _fill_stream_reads( - act, self._stream_store.config, self._data_converter - ) if not workflow: assert init_job workflow = _RunningWorkflow( @@ -686,7 +679,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_READS.drop_run(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 8290f04e8a..99d4a07820 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -59,11 +59,8 @@ import temporalio.exceptions import temporalio.nexus.system import temporalio.workflow -from temporalio.client._stream_read import _StreamFold from temporalio.client._stream_store import _STREAM_SERVICE from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo -from temporalio.nexus.system.stream_service.models import StreamRecord -from temporalio.nexus.system.stream_service.workflow._resources.stream import Stream from temporalio.nexus.system.stream_service.workflow._system_nexus_interceptor import ( _start_system_nexus_operation as _start_stream_system_nexus_operation, ) @@ -98,7 +95,7 @@ WorkflowInboundInterceptor, WorkflowOutboundInterceptor, ) -from ._stream_read import _read_stream_in_workflow, _reduce_stream_in_workflow +from ._stream_read import _stream_call_in_workflow from ._stream_store import _no_stream_config_error, _WorkflowStreamStore logger = logging.getLogger(__name__) @@ -318,6 +315,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 @@ -1998,37 +1996,18 @@ def workflow_upsert_search_attributes( else [update.value] ) - async def workflow_read_stream( - self, - stream: Stream[Any], - until: Callable[[StreamRecord[Any]], bool] | None, - limits: temporalio.workflow._StreamReadLimits, - ) -> temporalio.workflow.StreamRead[Any]: - if self._stream_store is None: - raise _no_stream_config_error() - return await _read_stream_in_workflow( - self, self._stream_store.config, stream, until, limits - ) - - async def workflow_reduce_stream( - self, - stream: Stream[Any], - function: Callable[[Any, StreamRecord[Any]], Any], - initial: Any, - until: Callable[[Any, StreamRecord[Any]], bool] | None, - value_type: Any, - limits: temporalio.workflow._StreamReadLimits, - ) -> temporalio.workflow.StreamReduce[Any]: - if self._stream_store is None: - raise _no_stream_config_error() - return await _reduce_stream_in_workflow( + async def workflow_stream_call( + self, stream: Any, call: temporalio.workflow._streams._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, + self._stream_store.config if self._stream_store is not None else None, stream, - _StreamFold(function=function, initial=initial), - until, - value_type, - limits, + call, + str(self._stream_call_count), ) async def workflow_sleep( diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index 11512d7c53..7f9ac3f6d7 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -156,13 +156,16 @@ unsafe, ) from ._streams import ( - DEFAULT_READ_STREAM_MAX_RECORDS, + DEFAULT_STREAM_IDLE_TIMEOUT, + DEFAULT_STREAM_MAX_BYTES, + DEFAULT_STREAM_MAX_RECORDS, + MAX_STREAM_RESULT_BYTES, StreamRead, StreamReadStop, StreamReduce, - _StreamReadLimits, - read_stream, + filter_stream, reduce_stream, + take_stream, ) from ._workflow_ops import ( ChildWorkflowCancellationType, @@ -271,12 +274,16 @@ "SandboxImportNotificationPolicy", "logger", "unsafe", - "DEFAULT_READ_STREAM_MAX_RECORDS", + "DEFAULT_STREAM_IDLE_TIMEOUT", + "DEFAULT_STREAM_MAX_BYTES", + "DEFAULT_STREAM_MAX_RECORDS", + "MAX_STREAM_RESULT_BYTES", "StreamRead", "StreamReadStop", "StreamReduce", - "read_stream", + "filter_stream", "reduce_stream", + "take_stream", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", @@ -319,7 +326,6 @@ "_update_validator", "_UpdateDefinition", "_NexusClient", - "_StreamReadLimits", "_build_log_context", "_imports_passed_through", "_in_sandbox", diff --git a/temporalio/workflow/_context.py b/temporalio/workflow/_context.py index 7d22ba4b2d..2580901004 100644 --- a/temporalio/workflow/_context.py +++ b/temporalio/workflow/_context.py @@ -21,16 +21,11 @@ from ._exceptions import _NotInWorkflowEventLoopError if TYPE_CHECKING: - from temporalio.nexus.system.stream_service.models import StreamRecord - from temporalio.nexus.system.stream_service.workflow._resources.stream import ( - Stream, - ) - from ._activities import ActivityCancellationType, ActivityHandle from ._event_groups import EventGroup from ._exceptions import ContinueAsNewVersioningBehavior, VersioningIntent from ._nexus import NexusOperationCancellationType, NexusOperationHandle - from ._streams import StreamRead, StreamReduce, _StreamReadLimits + from ._streams import _StreamCall from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowHandle, @@ -406,23 +401,7 @@ def workflow_payload_converter(self) -> temporalio.converter.PayloadConverter: . def workflow_random(self) -> Random: ... @abstractmethod - async def workflow_read_stream( - self, - stream: Stream[Any], - until: Callable[[StreamRecord[Any]], bool] | None, - limits: _StreamReadLimits, - ) -> StreamRead[Any]: ... - - @abstractmethod - async def workflow_reduce_stream( - self, - stream: Stream[Any], - function: Callable[[Any, StreamRecord[Any]], Any], - initial: Any, - until: Callable[[Any, StreamRecord[Any]], bool] | None, - value_type: Any, - limits: _StreamReadLimits, - ) -> StreamReduce[Any]: ... + async def workflow_stream_call(self, stream: Any, call: _StreamCall) -> Any: ... @abstractmethod def workflow_set_query_handler( diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py index 8193b6f72f..fce55e604e 100644 --- a/temporalio/workflow/_streams.py +++ b/temporalio/workflow/_streams.py @@ -1,4 +1,7 @@ -"""A Workflow read of a Stream that waits for a predicate, and a fold over its Records. +"""Calls that scan a Stream until a stop: take, filter and reduce. + +Client, Activity and Workflow code call the same three functions. A generated method of the +client ``Stream`` and of the Workflow ``Stream`` calls them too. .. warning:: This API is an experimental prototype and subject to change. @@ -6,97 +9,117 @@ from __future__ import annotations +import dataclasses import enum -from collections.abc import Callable, Sequence -from dataclasses import dataclass +import inspect +import typing +from collections.abc import Callable from datetime import timedelta -from typing import TYPE_CHECKING, Generic, TypeVar +from typing import TYPE_CHECKING, Any, Generic, TypeVar from ._context import _Runtime if TYPE_CHECKING: # The generated Stream modules import `temporalio.workflow`, so this module names # their types for the type checker only. + import temporalio.nexus.system.stream_service.client + import temporalio.nexus.system.stream_service.workflow._resources.stream from temporalio.nexus.system.stream_service.models import StreamRecord - from temporalio.nexus.system.stream_service.workflow._resources.stream import ( - Stream, - ) __all__ = [ - "DEFAULT_READ_STREAM_MAX_RECORDS", + "DEFAULT_STREAM_IDLE_TIMEOUT", + "DEFAULT_STREAM_MAX_BYTES", + "DEFAULT_STREAM_MAX_RECORDS", + "MAX_STREAM_RESULT_BYTES", "StreamRead", "StreamReadStop", "StreamReduce", - "read_stream", + "filter_stream", "reduce_stream", + "take_stream", ] BodyT = TypeVar("BodyT") ValueT = TypeVar("ValueT") -DEFAULT_READ_STREAM_MAX_RECORDS = 1_000 -"""The most Records of one read or fold when ``max_records`` is ``None``.""" +DEFAULT_STREAM_IDLE_TIMEOUT = timedelta(seconds=1) +"""How long a call waits for a new Record when ``idle_timeout`` is ``None``.""" + +DEFAULT_STREAM_MAX_RECORDS = 1_000 +"""The most Records of one call when ``max_records`` is ``None``.""" + +DEFAULT_STREAM_MAX_BYTES = 256 * 1024 +"""The most bytes of the Records of one call when ``max_bytes`` is ``None``.""" +MAX_STREAM_RESULT_BYTES = 1024 * 1024 +"""The most bytes of the encoded result of one call. + +In Workflow code the result 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 the same limit, so a call gives the same result everywhere. +""" -class StreamReadStop(enum.Enum): - """Why a read or a fold of a Stream ended. - When several hold, the first in this order is the reason. ``MATCH``, ``DONE``, - ``MAX_RECORDS``, ``MAX_BYTES``, ``IDLE``, ``TIMEOUT``. +class StreamReadStop(enum.IntEnum): + """Why a call on a Stream ended. + + When several hold, the first in the order of the values is the reason. .. warning:: This API is an experimental prototype and subject to change. """ - MATCH = "MATCH" - """``until`` was true for the last Record.""" + Match = 0 + """``until`` accepted a Record.""" - DONE = "DONE" - """The Stream is closed and its last Record was read.""" + Done = 1 + """The Stream is closed and the call looked at its last Record.""" - MAX_RECORDS = "MAX_RECORDS" - """The read got ``max_records`` Records.""" + MaxRecords = 2 + """The call got ``max_records`` Records.""" - MAX_BYTES = "MAX_BYTES" - """The next Record would pass ``max_bytes``. The next read gets it.""" + MaxBytes = 3 + """The next Record would pass ``max_bytes``. The next call gets it.""" - IDLE = "IDLE" + Idle = 4 """No new Record arrived for ``idle_timeout``.""" - TIMEOUT = "TIMEOUT" + Timeout = 5 """``timeout`` passed.""" -@dataclass(frozen=True) +@dataclasses.dataclass(slots=True, kw_only=True) class StreamRead(Generic[BodyT]): - """What one :py:func:`read_stream` gave. + """What :py:func:`take_stream` or :py:func:`filter_stream` gave. .. warning:: This API is an experimental prototype and subject to change. """ - records: Sequence[StreamRecord[BodyT]] - """The Records of the read in order, typed as a client's read gives them.""" + records: list[StreamRecord[BodyT]] + """The Records that the call gives back, in order.""" - last_offset: bytes | None - """The offset of the last Record, or the ``from_offset`` of a read with none. + last_offset: bytes | None = None + """Where the scan stopped. - Give it as the ``from_offset`` of the next read. + It is the offset of the last Record that the scan looked at, kept or not, or the + ``from_offset`` of a call that looked at none. Give it as the ``from_offset`` of the next + call, so a loop loses and repeats no Record. """ stop: StreamReadStop - """Why the read ended.""" + """Why the call ended.""" done: bool - """Whether the Stream is closed and its last Record is in :py:attr:`records`. + """Whether the Stream is closed and the scan looked at its last Record. - It is the state of the Stream, so it can be true when :py:attr:`stop` is ``MATCH``. + It is the state of the Stream, so it can be true when :py:attr:`stop` is ``Match``. """ -@dataclass(frozen=True) +@dataclasses.dataclass(slots=True, kw_only=True) class StreamReduce(Generic[ValueT]): - """What one :py:func:`reduce_stream` gave. + """What :py:func:`reduce_stream` gave. .. warning:: This API is an experimental prototype and subject to change. @@ -108,11 +131,8 @@ class StreamReduce(Generic[ValueT]): count: int """The number of Records that the fold got.""" - last_offset: bytes | None - """The offset of the last Record, or the ``from_offset`` of a fold with none. - - Give it as the ``from_offset`` of the next read or fold. - """ + last_offset: bytes | None = None + """Where the scan stopped, as in :py:attr:`StreamRead.last_offset`.""" stop: StreamReadStop """Why the fold ended.""" @@ -121,180 +141,408 @@ class StreamReduce(Generic[ValueT]): """Whether the Stream is closed and the fold got its last Record.""" -@dataclass(frozen=True) -class _StreamReadLimits: - """Where a read starts and what stops it, after the checks of the public functions.""" +class _StreamCallKind(enum.Enum): + TAKE = "take" + FILTER = "filter" + REDUCE = "reduce" + +@dataclasses.dataclass(frozen=True) +class _StreamCall: + """One call after the checks of the public functions. It holds the caller's functions.""" + + kind: _StreamCallKind from_offset: bytes | None - idle_timeout: timedelta | None timeout: timedelta + idle_timeout: timedelta max_records: int - max_bytes: int | None + max_bytes: int + until: Callable[..., object] | None = None + predicate: Callable[..., object] | None = None + function: Callable[..., object] | None = None + initial: Any = None + value_type: 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, - timeout: timedelta, max_records: int | None, max_bytes: int | None, - ) -> _StreamReadLimits: - if timeout <= timedelta(): - raise ValueError("timeout must be positive") - if idle_timeout is not None and idle_timeout < timedelta(): + until: Callable[..., object] | None = None, + predicate: Callable[..., object] | None = None, + function: Callable[..., object] | None = None, + initial: Any = None, + value_type: 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 = DEFAULT_STREAM_IDLE_TIMEOUT + elif idle_timeout < timedelta(): raise ValueError("idle_timeout must not be negative") if max_records is None: - max_records = DEFAULT_READ_STREAM_MAX_RECORDS + max_records = DEFAULT_STREAM_MAX_RECORDS elif max_records < 1: raise ValueError("max_records must be at least 1") - if max_bytes is not None and max_bytes < 1: + if max_bytes is None: + max_bytes = DEFAULT_STREAM_MAX_BYTES + elif max_bytes < 1: raise ValueError("max_bytes must be at least 1") - return _StreamReadLimits( + 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, - idle_timeout=idle_timeout, timeout=timeout, + idle_timeout=idle_timeout, max_records=max_records, max_bytes=max_bytes, + until=until, + predicate=predicate, + function=function, + initial=initial, + value_type=value_type, ) -async def read_stream( - stream: Stream[BodyT], - *, +def _check_plain_function( + label: str, value: Callable[..., object] | None, *, 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: Callable[..., object], initial: Any, value_type: Any +) -> Any: + if value_type is not None: + return value_type + 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: Any, call: _StreamCall) -> Any: + runtime = _Runtime.maybe_current() + if runtime is not None: + return await runtime.workflow_stream_call(stream, call) + # The client package imports this module, so its scan is imported at the first 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: Callable[[StreamRecord[BodyT]], bool] | None = None, from_offset: bytes | None = None, idle_timeout: timedelta | None = None, - timeout: timedelta, max_records: int | None = None, max_bytes: int | None = None, ) -> StreamRead[BodyT]: - """Read a Stream until a predicate is true, in one Local Activity. - - The read starts after ``from_offset`` and ends at the first stop of - :py:class:`StreamReadStop`. A Store's Local Activity reads the pages, and History - gets one marker for the whole read. The marker has the range of offsets and no - Record. A replay reads that range from the Store again before the Workflow sees the - result, so the Records must still be in the Store. + """Give every Record of a Stream up to a stop. + + ``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 of :py:class:`StreamReadStop` that holds, in the order of their + values. ``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 result are the same in each. In a + client or an Activity, the call reads the Store straight away. + + 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.client.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 Records and their bodies. 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 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 Workflow ``Stream`` to read, from - :py:func:`temporalio.workflow.get_stream` or - :py:func:`temporalio.workflow.create_stream`. - until: Called with each typed Record in order. The read stops after the first - Record for which it is true, and Records after it stay for the next read. - It is called once for each Record, when the Record arrives. A later change - of the state that it reads does not check an earlier Record again. It runs - in the Worker process, on the Worker's event loop, while the Local Activity - reads. It runs outside the Workflow's own thread. So it must not change any - state, and it can't call a function of ``temporalio.workflow``. It can read - the Workflow's state through a reference such as ``self``. When it sees a - change of that state is not defined. A replay does not call it, since the - replay takes the end of the read from History. ``None`` reads until another - stop. - from_offset: Read after this offset. ``None`` is the beginning. - idle_timeout: End the read when no new Record arrived for this long. It counts - from the last Record that the read got, or from the start of the read when - it got none. ``timedelta(0)`` gives what is in the Stream now and waits for - nothing. ``None`` is no idle stop. - timeout: The most time for the read. It must be positive. - max_records: The most Records to read. ``None`` is - :py:data:`DEFAULT_READ_STREAM_MAX_RECORDS`. - max_bytes: A cap on the bytes of the Records, as the Store counts them for a - page of the client ``read``. A Record that would pass the cap ends the read - before it, and the next read gets it. The first Record is always returned. + stream: The Stream to scan. + timeout: The most time for the call. It must be positive. + 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:`DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow code, a longer + wait keeps the Workflow Task open, so a Signal, an Update or a cancel waits for + the next Workflow Task heartbeat, and each heartbeat adds events to History. + max_records: The most Records to give. ``None`` is + :py:data:`DEFAULT_STREAM_MAX_RECORDS`, 1,000. + max_bytes: A cap on the bytes of the Records to give. A Record that would pass + it 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:`DEFAULT_STREAM_MAX_BYTES`, 256 KiB. Returns: - The Records of the read, the offset to read after next, and why it ended. + The Records, where the scan stopped, and why. Raises: - temporalio.exceptions.StreamError: The Store failed the read. - temporalio.exceptions.ApplicationError: ``until`` raised, and the error has - the message of the exception. Or the Worker's client has no Stream config, - and the error has the ``type`` ``StreamConfigError``. That one is not - retryable, so an uncaught one fails the Workflow. + 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.ApplicationError: A function raised, or returned an + awaitable. Or the encoded result is over :py:data:`MAX_STREAM_RESULT_BYTES`, + and the message names both sizes. In Workflow code it is also raised 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. """ - limits = _StreamReadLimits.check( + 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: 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]: + """Give the Records of a Stream that a predicate keeps, up to a stop. + + ``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 of :py:class:`StreamReadStop` that holds, in the order of their + values. ``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 result are the same in each. In a + client or an Activity, the call reads the Store straight away. + + 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.client.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 kept Records and their bodies. 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 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 scan. + predicate: Keeps a Record when it returns true. + timeout: The most time for the call. It must be positive. + 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:`DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow code, a longer + wait keeps the Workflow Task open, so a Signal, an Update or a cancel waits for + the next Workflow Task heartbeat, and each heartbeat adds events to History. + max_records: The most Records to keep. ``None`` is + :py:data:`DEFAULT_STREAM_MAX_RECORDS`, 1,000. + max_bytes: A cap on the bytes of the Records to keep. A Record that would pass + it 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:`DEFAULT_STREAM_MAX_BYTES`, 256 KiB. + + 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.ApplicationError: A function raised, or returned an + awaitable. Or the encoded result is over :py:data:`MAX_STREAM_RESULT_BYTES`, + and the message names both sizes. In Workflow code it is also raised 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 _Runtime.current().workflow_read_stream(stream, until, limits) + return await _run_stream_call(stream, call) async def reduce_stream( - stream: Stream[BodyT], + stream: temporalio.nexus.system.stream_service.client.Stream[BodyT] + | temporalio.nexus.system.stream_service.workflow._resources.stream.Stream[BodyT], function: Callable[[ValueT, StreamRecord[BodyT]], ValueT], initial: ValueT, - *, + timeout: timedelta, until: Callable[[ValueT, StreamRecord[BodyT]], bool] | None = None, - value_type: type[ValueT] | None = None, from_offset: bytes | None = None, idle_timeout: timedelta | None = None, - timeout: timedelta, max_records: int | None = None, max_bytes: int | None = None, + *, + value_type: type[ValueT] | None = None, ) -> StreamReduce[ValueT]: - """Fold the Records of a Stream into one value, in one Local Activity. - - The Local Activity calls ``value = function(value, record)`` for each Record in - order, then ``until(value, record)``. The fold ends at the first stop of - :py:class:`StreamReadStop`, with the same rules as :py:func:`read_stream`. No - Record comes back to the Workflow. History keeps the value in the marker of the - Local Activity, so a replay takes the value from History and needs no Store. The - size of the value is the caller's choice, and History pays for it. + """Fold the Records of a Stream into one value, up to a stop. + + 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 of :py:class:`StreamReadStop` that holds, in the order of their + values. ``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 result are the same in each. In a + client or an Activity, the call reads the Store straight away. + + 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.client.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. 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 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 Workflow ``Stream`` to fold. - function: Gives the next value from the value so far and a typed Record. It - changes no state other than the value that it returns, so it must not - change ``initial`` or a value in place. It runs in the Worker process, on - the Worker's event loop, outside the Workflow's own thread. It does no I/O, - and it can't call a function of ``temporalio.workflow``. A replay does not - call it. + 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. - until: Called with the value after each Record and that Record. The fold stops - after the first Record for which it is true. It is called once for each - Record, when the Record arrives. It has the rules of ``function``. - value_type: The type that the Workflow decodes the value as. ``None`` takes the - type of ``initial``. - from_offset: Fold the Records after this offset. ``None`` is the beginning. - idle_timeout: As for :py:func:`read_stream`. - timeout: As for :py:func:`read_stream`. - max_records: As for :py:func:`read_stream`. - max_bytes: As for :py:func:`read_stream`. + timeout: The most time for the call. It must be positive. + 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:`DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow code, a longer + wait keeps the Workflow Task open, so a Signal, an Update or a cancel waits for + the next Workflow Task heartbeat, and each heartbeat adds events to History. + max_records: The most Records to fold. ``None`` is + :py:data:`DEFAULT_STREAM_MAX_RECORDS`, 1,000. + max_bytes: A cap on the bytes of the Records to fold. A Record that would pass + it 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:`DEFAULT_STREAM_MAX_BYTES`, 256 KiB. + value_type: The type that Workflow code decodes the value as. ``None`` takes the + return annotation of ``function``, else the type of ``initial``, and no type + when ``initial`` is ``None``. Client and Activity code get the value itself. Returns: - The value, the number of Records, the offset to read after next, and why the - fold ended. + The value, the number of Records, where the scan stopped, and why. Raises: - temporalio.exceptions.StreamError: The Store failed the read. - temporalio.exceptions.ApplicationError: ``function`` or ``until`` raised, and - the error has the message of the exception. Or the Worker's client has no - Stream config, and the error has the ``type`` ``StreamConfigError``. That - one is not retryable, so an uncaught one fails the Workflow. + 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.ApplicationError: A function raised, or returned an + awaitable. Or the encoded result is over :py:data:`MAX_STREAM_RESULT_BYTES`, + and the message names both sizes. In Workflow code it is also raised 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. """ - limits = _StreamReadLimits.check( + call = _StreamCall.check( + _StreamCallKind.REDUCE, + timeout=timeout, from_offset=from_offset, idle_timeout=idle_timeout, - timeout=timeout, max_records=max_records, max_bytes=max_bytes, + until=until, + function=function, + initial=initial, + value_type=_fold_value_type(function, initial, value_type), ) - return await _Runtime.current().workflow_reduce_stream( - stream, - function, - initial, - until, - value_type if value_type is not None else type(initial), - limits, - ) + return await _run_stream_call(stream, call) diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py index 003cf1598c..9fda45007c 100644 --- a/tests/nexus/test_stream_store_workflow.py +++ b/tests/nexus/test_stream_store_workflow.py @@ -643,8 +643,8 @@ async def test_workflow_read_fails_with_store(client: Client): ) assert result == ( "Workflow code can't run a ReadRecords operation on a Store. Use " - "temporalio.workflow.read_stream or reduce_stream, which read in one " - "Local Activity." + "temporalio.workflow.take_stream, filter_stream or reduce_stream, which " + "read in one Local Activity." ) From 9b038517a153960ecf803e3f19185522a8f36726 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 01:47:58 -0700 Subject: [PATCH 09/20] Let BufferedAppender write inside the function of a Stream call. A fold can append each transformed Record to another Stream. The appender writes through the Activity's client there, and the call writes its Records before it returns, so History gets the result only after they are in the Store. --- temporalio/client/_stream_read.py | 37 +++++++++++++- temporalio/contrib/streams/_appender.py | 64 ++++++++++++++++++++++++- 2 files changed, 97 insertions(+), 4 deletions(-) diff --git a/temporalio/client/_stream_read.py b/temporalio/client/_stream_read.py index ea87c19ceb..148ac18a77 100644 --- a/temporalio/client/_stream_read.py +++ b/temporalio/client/_stream_read.py @@ -46,6 +46,7 @@ ) from ._client import Client +from ._stream_appender import _CALL_APPENDERS, _flush_call_appenders from ._stream_store import ( _STREAM_SERVICE, StreamStore, @@ -244,6 +245,38 @@ def idle(now: float) -> bool: return state.end(StreamReadStop.Idle) +async def _scan_and_flush( + read_page: _ReadPage, + *, + ref: temporalio.api.stream.v1.StreamRef, + call: _StreamCall, + to_record: _ToRecord, + in_local_activity: bool, +) -> _ScanOutcome: + """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. + """ + appenders: list[Any] = [] + token = _CALL_APPENDERS.set(appenders) + try: + outcome = await _scan( + read_page, + ref=ref, + call=call, + to_record=to_record, + in_local_activity=in_local_activity, + ) + except Exception: + # The error of the scan wins over an error of the write. + await _flush_call_appenders(appenders, raise_error=False) + raise + finally: + _CALL_APPENDERS.reset(token) + await _flush_call_appenders(appenders, raise_error=True) + return outcome + + @dataclasses.dataclass(frozen=True) class _StreamCallResult: """The result of one call, as the payload that History keeps for a Workflow.""" @@ -447,7 +480,7 @@ async def _stream_call_activity(call_id: str) -> temporalio.common.RawValue: non_retryable=True, ) # The codec decodes the bodies of each page, so the caller's functions get its values. - outcome = await _scan( + outcome = await _scan_and_flush( _store_page_reader(pending.store, temporalio.activity.client().data_converter), ref=pending.ref, call=pending.call, @@ -510,7 +543,7 @@ async def _stream_call_outside_workflow( # pyright: ignore[reportUnusedFunction def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: return payload_converter._from_transfer_value(record, record_type) - outcome = await _scan( + outcome = await _scan_and_flush( _store_page_reader(store, data_converter), ref=message.ref, call=call, diff --git a/temporalio/contrib/streams/_appender.py b/temporalio/contrib/streams/_appender.py index 7b194ddf09..ab2805c3d7 100644 --- a/temporalio/contrib/streams/_appender.py +++ b/temporalio/contrib/streams/_appender.py @@ -7,12 +7,15 @@ from __future__ import annotations import asyncio +import contextvars +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.nexus.system.stream_service.client import temporalio.workflow @@ -29,6 +32,12 @@ _DEFAULT_MAX_DELAY = timedelta(milliseconds=100) +# The appenders 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_APPENDERS: contextvars.ContextVar[list[BufferedAppender[Any]] | None] = ( + contextvars.ContextVar("temporal_stream_call_appenders", default=None) +) + class BufferedAppender(Generic[BodyT]): """Batches the appends of one Stream. @@ -48,6 +57,13 @@ 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 :py:func:`temporalio.workflow.take_stream`, ``filter_stream`` or + ``reduce_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 +119,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 +138,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() + appenders = _CALL_APPENDERS.get() + if appenders is not None and self not in appenders: + appenders.append(self) loop = asyncio.get_running_loop() if not self._buffer: self._first_at = loop.time() @@ -171,6 +196,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 +252,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. @@ -233,3 +281,15 @@ def _fail(self, error: BaseException, count: int) -> None: if self._timer is not None: self._timer.cancel() self._timer = None + + +async def _flush_call_appenders( # pyright: ignore[reportUnusedFunction] + appenders: list[BufferedAppender[Any]], *, raise_error: bool +) -> None: + """Write the Records of the appenders that the functions of a Stream call used.""" + for appender in appenders: + try: + await appender.flush() + except Exception: + if raise_error: + raise From 45aa145f2548db80ee78ac679e1ccf2c763496a7 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 01:49:57 -0700 Subject: [PATCH 10/20] Tested the Stream calls in Workflow, client and Activity code. The tests cover each stop of each call with its last offset, replays with no Stream config, a Query after the Stream is gone, and a fold that appends to another Stream. The measurement tests print the History events of each case. --- tests/helpers/streams.py | 3 +- tests/nexus/test_stream_read_client.py | 230 +++ tests/nexus/test_stream_read_workflow.py | 1747 ++++++++++++---------- 3 files changed, 1204 insertions(+), 776 deletions(-) create mode 100644 tests/nexus/test_stream_read_client.py diff --git a/tests/helpers/streams.py b/tests/helpers/streams.py index 5cbfe2c7e6..5d5b8bae30 100644 --- a/tests/helpers/streams.py +++ b/tests/helpers/streams.py @@ -187,7 +187,7 @@ async def describe_stream( class StreamReadState: - """State that a test of a Workflow read shares with its Workflows. + """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. @@ -196,4 +196,5 @@ class StreamReadState: 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 0000000000..f4ad1cd8ab --- /dev/null +++ b/tests/nexus/test_stream_read_client.py @@ -0,0 +1,230 @@ +"""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 ( + BufferedAppender, + Client, + StreamReadStop, + filter_stream, + reduce_stream, + take_stream, +) +from temporalio.exceptions import ApplicationError +from temporalio.nexus.system.stream_service import StreamRecord, StreamRef +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, + Token, + _append, + _call, + _check_stop, + _close, + _new_stream, + _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 filter_stream( + stream, + 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: + return True + + with pytest.raises(TypeError, match="until must be a plain function"): + await take_stream(stream, timedelta(seconds=5), coroutine_until) + with pytest.raises(ApplicationError) as err: + await filter_stream( + stream, lambda record: coroutine_until(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("bad fold") + + # In a client, the error of the caller's function reaches it as it is. + with pytest.raises(ValueError, match="bad fold"): + await reduce_stream(stream, raises, 0, timedelta(seconds=5)) + + +async def test_client_result_over_the_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) + stream = _client_stream(stream_client, stream_id) + with pytest.raises(ApplicationError) as err: + await take_stream( + stream, timedelta(seconds=5), idle_timeout=timedelta(0), max_bytes=1 << 22 + ) + assert err.value.type == "StreamResultTooLarge" + assert "over the limit of 1048576 bytes" in err.value.message + # The default cap of 256 KiB keeps a call under the limit. + read = await take_stream(stream, timedelta(seconds=5), idle_timeout=timedelta(0)) + assert read.stop is StreamReadStop.MaxBytes + + +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 reduce_stream( + _client_stream(stream_client, source_id), + 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_workflow.py b/tests/nexus/test_stream_read_workflow.py index c41a8a3270..1fc2d9d36e 100644 --- a/tests/nexus/test_stream_read_workflow.py +++ b/tests/nexus/test_stream_read_workflow.py @@ -1,4 +1,4 @@ -"""Tests of a Workflow read of a Stream that waits for a predicate.""" +"""Tests of the take, filter and reduce calls of Workflow code.""" from __future__ import annotations @@ -6,7 +6,7 @@ import collections import time import uuid -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from datetime import datetime, timedelta from typing import Any @@ -14,20 +14,26 @@ import pytest from temporalio import workflow -from temporalio.api.enums.v1 import EventType +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 ( + BufferedAppender, Client, StreamConfig, WorkflowFailureError, WorkflowHandle, WorkflowHistory, ) -from temporalio.client._stream_read import _PENDING_STREAM_READS -from temporalio.converter import DataConverter -from temporalio.exceptions import ApplicationError, CancelledError -from temporalio.nexus.system.stream_service import StreamRecord, StreamRef +from temporalio.client._stream_read import _PENDING_STREAM_CALLS +from temporalio.converter import DataConverter, PayloadCodec +from temporalio.exceptions import ApplicationError, CancelledError, StreamError +from temporalio.nexus.system.stream_service import ( + ProducerKindLocalActivity, + StreamRecord, + StreamRef, +) from temporalio.nexus.system.stream_service.workflow._resources.stream import ( Stream as WorkflowStream, ) @@ -39,10 +45,13 @@ # 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 RecordingStreamStore, StreamReadState - from tests.test_converter import SimpleCodec + from tests.helpers.streams import ( + DictStreamStore, + RecordingStreamStore, + StreamReadState, + ) -# Every body carries this text, so a search of History shows whether a body leaks. +# Every body carries this text, so a search of History shows whether a body is readable. _BODY_MARK = "secret-body" @@ -53,29 +62,50 @@ class Token: @dataclass -class ReadArgs: +class CallArgs: stream_id: str + call: str = "take" until_kind: str | None = "done" timeout_seconds: float = 10.0 - max_records: int | None = None idle_timeout_seconds: float | None = None + max_records: int | None = None max_bytes: int | None = None from_offset: str | None = None @dataclass -class ReadResult: +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 _result(read: workflow.StreamRead[Token]) -> ReadResult: - return ReadResult( +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: workflow.StreamRead[Token]) -> CallResult: + return CallResult( texts=[record.body.text for record in read.records], - stop=read.stop.value, + 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( @@ -85,10 +115,11 @@ def _result(read: workflow.StreamRead[Token]) -> ReadResult: ) -def _keep(result: ReadResult) -> ReadResult: +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, @@ -96,37 +127,86 @@ def _keep(result: ReadResult) -> ReadResult: return result -def _stream(stream_id: str) -> WorkflowStream[Token]: - return workflow.get_stream(StreamRef(id=stream_id), body_type=Token) +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: + 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 workflow.take_stream( + stream, + 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 workflow.filter_stream( + stream, + _not_skipped, + timeout, + until if until_kind else None, + from_offset, + idle_timeout, + args.max_records, + args.max_bytes, + ) + return _read_result(read) + reduce = await workflow.reduce_stream( + stream, + _collect, + [], + 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 ReadUntilKindWorkflow: +class CallWorkflow: + def __init__(self) -> None: + self.result: CallResult | None = None + @workflow.run - async def run(self, args: ReadArgs) -> ReadResult: - 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 - - read = await workflow.read_stream( - _stream(args.stream_id), - until=until if until_kind is not None else None, - from_offset=bytes.fromhex(args.from_offset) if args.from_offset else None, - idle_timeout=None - if args.idle_timeout_seconds is None - else timedelta(seconds=args.idle_timeout_seconds), - timeout=timedelta(seconds=args.timeout_seconds), - max_records=args.max_records, - max_bytes=args.max_bytes, - ) - return _keep(_result(read)) + 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 ReadUntilWantedWorkflow: +class TakeUntilWantedWorkflow: def __init__(self) -> None: self.wanted: str | None = None @@ -135,236 +215,214 @@ def want(self, text: str) -> None: self.wanted = text @workflow.run - async def run(self, stream_id: str) -> ReadResult: - read = await workflow.read_stream( + async def run(self, stream_id: str) -> CallResult: + read = await workflow.take_stream( _stream(stream_id), + timedelta(seconds=20), until=lambda record: record.body.text == self.wanted, - timeout=timedelta(seconds=20), + idle_timeout=timedelta(seconds=20), ) - return _result(read) + return _read_result(read) @workflow.defn -class TwoReadsWorkflow: +class TwoTakesWorkflow: @workflow.run async def run(self, stream_id: str) -> list[str]: stream = _stream(stream_id) - first = await workflow.read_stream( + first = await workflow.take_stream( stream, + timedelta(seconds=10), until=lambda record: record.body.text.endswith("m1"), - timeout=timedelta(seconds=10), ) await workflow.sleep(timedelta(milliseconds=1)) - second = await workflow.read_stream( + second = await workflow.take_stream( stream, + timedelta(seconds=10), until=lambda record: record.body.text.endswith("m2"), from_offset=first.last_offset, - timeout=timedelta(seconds=10), ) return [record.body.text for record in [*first.records, *second.records]] @workflow.defn -class FailTaskOnceAfterReadWorkflow: +class FailTaskOnceAfterTakeWorkflow: @workflow.run - async def run(self, stream_id: str) -> ReadResult: - read = await workflow.read_stream( + async def run(self, stream_id: str) -> CallResult: + read = await workflow.take_stream( _stream(stream_id), + timedelta(seconds=10), until=lambda record: record.body.kind == "done", - timeout=timedelta(seconds=10), ) - # The failure comes in a later Workflow Task, so the retry replays the read. + # 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 _result(read) + return _read_result(read) @workflow.defn -class TimeoutThenReadWorkflow: +class FilterLoopWorkflow: @workflow.run - async def run(self, stream_id: str) -> list[ReadResult]: + async def run(self, stream_id: str) -> tuple[list[str], int]: stream = _stream(stream_id) - first = await workflow.read_stream( - stream, - until=lambda record: record.body.kind == "done", - timeout=timedelta(seconds=1), - ) - second = await workflow.read_stream( - stream, - until=lambda record: record.body.kind == "done", - from_offset=first.last_offset, - timeout=timedelta(seconds=10), - ) - return [_result(first), _result(second)] - - -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 + 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 workflow.filter_stream( + stream, + 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 ( + workflow.StreamReadStop.Done, + workflow.StreamReadStop.Idle, + ): + return found, calls -_PREDICATES: dict[str, Callable[[StreamRecord[Token]], bool]] = { - "raise": _raise_on_record, - "workflow_now": _call_workflow_now, - "datetime_now": _call_datetime_now, -} +@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 workflow.take_stream( + stream, timedelta(seconds=10), from_offset=offset + ) + pages.append((len(read.records), read.stop.name)) + offset = read.last_offset + if read.stop is not workflow.StreamReadStop.MaxBytes: + return pages @workflow.defn -class FailingPredicateWorkflow: +class OverLimitWorkflow: @workflow.run - async def run(self, stream_id: str, predicate: str) -> str: + async def run(self, stream_id: str, call: str) -> str: + stream = _stream(stream_id) try: - await workflow.read_stream( - _stream(stream_id), - until=_PREDICATES[predicate], - timeout=timedelta(seconds=10), - ) + if call == "take": + await workflow.take_stream( + stream, timedelta(seconds=10), max_bytes=4 * 1024 * 1024 + ) + else: + await workflow.reduce_stream( + stream, + lambda value, record: value + "x" * (1024 * 1024), + "", + timedelta(seconds=10), + max_records=1, + ) except ApplicationError as err: - return f"{err.type}: {err.message}" + return f"{err.type} {err.non_retryable}: {err.message}" return "no error" @workflow.defn -class MaxBytesThenNextWorkflow: +class TwoStreamsAtOnceWorkflow: @workflow.run - async def run(self, stream_id: str, max_bytes: int) -> list[ReadResult]: - stream = _stream(stream_id) - first = await workflow.read_stream( - stream, + async def run(self, first_id: str, second_id: str) -> list[str]: + tasks = [ + asyncio.create_task( + workflow.filter_stream( + _stream(stream_id), + 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 workflow.take_stream( + _stream(stream_id), + timedelta(seconds=10), + from_offset=b"\x01" if variant == "other_offset" else None, idle_timeout=timedelta(0), - timeout=timedelta(seconds=10), - max_bytes=max_bytes, - ) - second = await workflow.read_stream( - stream, - from_offset=first.last_offset, - timeout=timedelta(seconds=10), - max_records=1, ) - return [_result(first), _result(second)] + return [record.body.text for record in read.records] -@dataclass -class ReduceArgs: - stream_id: str - until: str = "done" - timeout_seconds: float = 10.0 +async def _async_until(record: StreamRecord[Token]) -> bool: + return True -@dataclass -class ReduceResult: - value: str - count: int - stop: str - done: bool - last_offset: str | None +def _until_that_awaits(record: StreamRecord[Token]) -> Any: + return asyncio.sleep(0) -def _join(value: str, record: StreamRecord[Token]) -> str: - StreamReadState.function_calls += 1 - return value + record.body.text +@workflow.defn +class RefusedFunctionWorkflow: + @workflow.run + async def run(self, stream_id: str, which: str) -> str: + stream = _stream(stream_id) + try: + await workflow.take_stream( + stream, + timedelta(seconds=10), + until=_async_until if which == "coroutine" else _until_that_awaits, + ) + 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 _until_done(value: str, record: StreamRecord[Token]) -> bool: - del value - StreamReadState.until_calls += 1 - return record.body.kind == "done" +def _raise_on_record(record: StreamRecord[Token]) -> bool: + raise ValueError(f"bad token {record.body.text}") -def _until_tag(value: str, record: StreamRecord[Token]) -> bool: +def _call_workflow_now(record: StreamRecord[Token]) -> bool: del record - StreamReadState.until_calls += 1 - return value.endswith("tag!") - - -_FOLD_UNTIL: dict[str, Callable[[str, StreamRecord[Token]], bool]] = { - "done": _until_done, - "tag": _until_tag, -} - - -def _reduce_result(reduce: workflow.StreamReduce[str]) -> ReduceResult: - return ReduceResult( - value=reduce.value, - count=reduce.count, - stop=reduce.stop.value, - done=reduce.done, - last_offset=reduce.last_offset.hex() if reduce.last_offset else None, - ) - - -@workflow.defn -class JoinTextWorkflow: - @workflow.run - async def run(self, args: ReduceArgs) -> ReduceResult: - reduce = await workflow.reduce_stream( - _stream(args.stream_id), - _join, - "", - until=_FOLD_UNTIL[args.until], - timeout=timedelta(seconds=args.timeout_seconds), - ) - result = _reduce_result(reduce) - StreamReadState.results[workflow.info().run_id] = ( - result.value, - result.count, - result.stop, - result.done, - result.last_offset, - ) - return result + return workflow.now() > datetime.min -@workflow.defn -class TwoFoldsWorkflow: - @workflow.run - async def run(self, stream_id: str) -> list[ReduceResult]: - stream = _stream(stream_id) - first = await workflow.reduce_stream( - stream, - _join, - "", - until=lambda value, _: value.endswith("m1"), - timeout=timedelta(seconds=10), - ) - await workflow.sleep(timedelta(milliseconds=1)) - second = await workflow.reduce_stream( - stream, - _join, - "", - until=lambda value, _: value.endswith("m2"), - from_offset=first.last_offset, - timeout=timedelta(seconds=10), - ) - return [_reduce_result(first), _reduce_result(second)] +def _call_datetime_now(record: StreamRecord[Token]) -> bool: + del record + return datetime.now() > datetime.min -def _raise_in_function(value: str, record: StreamRecord[Token]) -> str: - raise ValueError(f"bad fold after {value!r} at {record.body.text}") +_PREDICATES: dict[str, Callable[[StreamRecord[Token]], bool]] = { + "raise": _raise_on_record, + "workflow_now": _call_workflow_now, + "datetime_now": _call_datetime_now, +} @workflow.defn -class FailingFunctionWorkflow: +class FailingPredicateWorkflow: @workflow.run - async def run(self, stream_id: str) -> str: + async def run(self, stream_id: str, predicate: str) -> str: try: - await workflow.reduce_stream( - _stream(stream_id), - _raise_in_function, - "", - timeout=timedelta(seconds=10), + await workflow.filter_stream( + _stream(stream_id), _PREDICATES[predicate], timedelta(seconds=10) ) except ApplicationError as err: return f"{err.type}: {err.message}" @@ -377,49 +435,89 @@ class Summary: kinds: list[str] -def _summarize(value: Summary, record: StreamRecord[Token]) -> Summary: +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 _collect(value: list[str], record: StreamRecord[Token]) -> list[str]: - return [*value, record.body.text] - - -@dataclass -class TypedFoldResult: - summary: Summary - summary_type: str - texts: list[str] - texts_type: str - - @workflow.defn -class TypedFoldWorkflow: +class FoldTypesWorkflow: @workflow.run - async def run(self, stream_id: str) -> TypedFoldResult: + async def run(self, stream_id: str) -> list[str]: stream = _stream(stream_id) - summary = await workflow.reduce_stream( + quick = timedelta(seconds=10) + now = timedelta(0) + # The return annotation wins over an `initial` of `None`. + annotated = await workflow.reduce_stream( + stream, _summarize, None, quick, idle_timeout=now + ) + # With no annotation, the type of `initial` decides. + of_initial = await workflow.reduce_stream( stream, - _summarize, + lambda value, record: Summary(value.count + 1, []), Summary(0, []), - idle_timeout=timedelta(0), - timeout=timedelta(seconds=10), + quick, + idle_timeout=now, + ) + # With no annotation and an `initial` of `None`, there is no type hint. + untyped = await workflow.reduce_stream( + stream, + lambda value, record: Summary(1, [record.body.kind]), + None, + quick, + idle_timeout=now, ) - no_texts: list[str] = [] - texts = await workflow.reduce_stream( + explicit = await workflow.reduce_stream( stream, - _collect, - no_texts, + lambda value, record: [*value, record.body.text], + (), + quick, + idle_timeout=now, value_type=list[str], - idle_timeout=timedelta(0), - timeout=timedelta(seconds=10), ) - return TypedFoldResult( - summary=summary.value, - summary_type=type(summary.value).__name__, - texts=texts.value, - texts_type=type(texts.value).__name__, + return [ + type(result.value).__name__ + for result in (annotated, of_initial, untyped, explicit) + ] + + +@workflow.defn +class FoldUntilValueWorkflow: + @workflow.run + async def run(self, stream_id: str) -> tuple[str, int, str]: + result = await workflow.reduce_stream( + _stream(stream_id), + 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 workflow.reduce_stream( + _stream(source_id), + 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): @@ -441,25 +539,51 @@ async def read_records(self, request: ReadRecordsRequest) -> ReadRecordsResponse self.active_reads -= 1 -def _tokens(count: int, *, prefix: str = "t") -> list[Token]: - return [Token("text", f"{_BODY_MARK}-{prefix}{index}") for index in range(count)] +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() - config = client.config() - config["stream_config"] = StreamConfig(stores=[store]) - if data_converter is not None: - config["data_converter"] = data_converter - stream_client = Client(**config) + 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 @@ -470,6 +594,15 @@ async def _append(stream_client: Client, stream_id: str, *tokens: Token) -> None 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: @@ -482,7 +615,7 @@ async def _append_over( return last -def _texts(tokens: list[Token]) -> list[str]: +def _texts(tokens: Sequence[Token]) -> list[str]: return [token.text for token in tokens] @@ -494,6 +627,15 @@ def _markers(history: WorkflowHistory) -> list[str]: ] +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_") @@ -502,28 +644,32 @@ def _event_counts(history: WorkflowHistory) -> dict[str, int]: return dict(sorted(counts.items())) -def _assert_no_bodies(history: WorkflowHistory) -> None: - # The Workflow returns the bodies, so only its completion event has them. - events = [ - event - for event in history.events - if event.event_type != EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED - ] - assert len(events) == len(history.events) - 1 - # The JSON has payload data in base64, so the raw events are searched too. - assert _BODY_MARK not in WorkflowHistory(history.workflow_id, events).to_json() - for event in events: - assert _BODY_MARK.encode() not in event.SerializeToString() +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_read( - stream_client: Client, - args: ReadArgs, - task_queue: str, - **kwargs: Any, -) -> WorkflowHandle[ReadUntilKindWorkflow, ReadResult]: +async def _start_call( + stream_client: Client, args: CallArgs, task_queue: str, **kwargs: Any +) -> WorkflowHandle[CallWorkflow, CallResult]: return await stream_client.start_workflow( - ReadUntilKindWorkflow.run, + CallWorkflow.run, args, id=f"workflow-{uuid.uuid4()}", task_queue=task_queue, @@ -531,59 +677,328 @@ async def _start_read( ) -_RUNNERS = pytest.mark.parametrize( - "runner", - [SandboxedWorkflowRunner(), UnsandboxedWorkflowRunner()], - ids=["sandboxed", "unsandboxed"], -) +_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 _run_match( - client: Client, - runner: WorkflowRunner = SandboxedWorkflowRunner(), - data_converter: DataConverter | None = None, -) -> tuple[Client, Any, WorkflowHandle[Any, ReadResult], ReadResult, list[Token]]: - stream_client, store, stream_id = await _new_stream( - client, data_converter=data_converter - ) - tokens = [*_tokens(4), _done()] +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) - async with new_worker( - stream_client, ReadUntilKindWorkflow, workflow_runner=runner - ) as worker: - handle = await _start_read( - stream_client, ReadArgs(stream_id), worker.task_queue - ) - return stream_client, store, handle, await handle.result(), 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") -@_RUNNERS -async def test_read_until_match(client: Client, runner: WorkflowRunner): - _, _, handle, result, tokens = await _run_match(client, runner) - assert result == ReadResult( - texts=_texts(tokens), - stop="MATCH", - done=False, - last_offset=result.last_offset, - typed=True, - ) - assert result.last_offset - history = await handle.fetch_history() - assert _markers(history) == ["core_local_activity"] - _assert_no_bodies(history) +@_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_read_waits_for_records(client: Client): +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, ReadUntilKindWorkflow) as worker: - handle = await _start_read( - stream_client, ReadArgs(stream_id), worker.task_queue + 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 result.stop == "Match" assert _markers(await handle.fetch_history()) == ["core_local_activity"] @@ -604,19 +1019,24 @@ async def delivered() -> None: await assert_eventually(delivered) +_RUNNERS = pytest.mark.parametrize( + "runner", + [SandboxedWorkflowRunner(), UnsandboxedWorkflowRunner()], + ids=["sandboxed", "unsandboxed"], +) + + @_RUNNERS -async def test_read_predicate_reads_workflow_state( - client: Client, runner: WorkflowRunner -): +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, ReadUntilWantedWorkflow, workflow_runner=runner + stream_client, TakeUntilWantedWorkflow, workflow_runner=runner ) as worker: handle = await stream_client.start_workflow( - ReadUntilWantedWorkflow.run, + TakeUntilWantedWorkflow.run, stream_id, id=f"workflow-{uuid.uuid4()}", task_queue=worker.task_queue, @@ -627,34 +1047,17 @@ async def read_waits() -> None: assert store.active_reads == 1 await assert_eventually(read_waits) - await handle.signal(ReadUntilWantedWorkflow.want, wanted.text) + 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" + assert result.stop == "Match" -async def test_read_replays_with_replayer(client: Client): - _, store, handle, result, _ = await _run_match(client) - run_id = handle.result_run_id - assert run_id - first_run = StreamReadState.results.pop(run_id) - replay_store = RecordingStreamStore(store) - StreamReadState.until_calls = 0 - await Replayer( - workflows=[ReadUntilKindWorkflow], - stream_config=StreamConfig(stores=[replay_store]), - ).replay_workflow(await handle.fetch_history()) - assert StreamReadState.results.pop(run_id) == first_run - assert first_run[0] == tuple(result.texts) - assert replay_store.operations == ["ReadRecords"] - assert StreamReadState.until_calls == 0 - - -async def test_read_replays_on_a_worker_without_cache(client: Client): +async def test_calls_on_a_worker_without_cache(client: Client): stream_client, store, stream_id = await _new_stream( client, RecordingStreamStore(MemoryStreamStore()) ) @@ -662,148 +1065,100 @@ async def test_read_replays_on_a_worker_without_cache(client: Client): Token("text", f"{_BODY_MARK}-{text}") for text in ["a", "m1", "b", "m2", "c"] ] await _append(stream_client, stream_id, *tokens) - async with new_worker( - stream_client, TwoReadsWorkflow, max_cached_workflows=0 - ) as worker: - handle = await stream_client.start_workflow( - TwoReadsWorkflow.run, - stream_id, - id=f"workflow-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - assert await handle.result() == _texts(tokens[:4]) - history = await handle.fetch_history() - assert _markers(history) == ["core_local_activity", "core_local_activity"] - # Two reads, plus a replay of the first read before the Workflow Task of the second. - reads = store.operations.count("ReadRecords") - print(f"\nWorker without cache: {reads} ReadRecords calls") - assert reads >= 3 - - -async def test_read_replay_fails_when_records_are_gone(client: Client): - _, _, handle, _, _ = await _run_match(client) - history = await handle.fetch_history() - with pytest.raises(Exception) as err: - await Replayer( - workflows=[ReadUntilKindWorkflow], - stream_config=StreamConfig(stores=[MemoryStreamStore()]), - ).replay_workflow(history) - stream_id = history.events[0].workflow_execution_started_event_attributes - assert "can't read the Stream" in str(err.value) - assert f"'{_stream_id_of(stream_id)}'" in str(err.value) - - -def _stream_id_of(started: Any) -> str: - return DataConverter.default.payload_converter.from_payloads( - started.input.payloads, [ReadArgs] - )[0].stream_id + 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_read_survives_a_failed_workflow_task(client: Client): +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 - async with new_worker(stream_client, FailTaskOnceAfterReadWorkflow) as worker: - handle = await stream_client.start_workflow( - FailTaskOnceAfterReadWorkflow.run, - stream_id, - id=f"workflow-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - result = await handle.result() + handle, result = await _run_workflow( + stream_client, FailTaskOnceAfterTakeWorkflow, stream_id + ) assert result.texts == _texts(tokens) - assert result.stop == "MATCH" + assert result.stop == "Match" counts = _event_counts(await handle.fetch_history()) assert counts["WORKFLOW_TASK_FAILED"] == 1 assert counts["MARKER_RECORDED"] == 1 -async def test_read_timeout_gives_records_so_far(client: Client): +async def _failing_predicate( + client: Client, + predicate: str, + runner: WorkflowRunner = SandboxedWorkflowRunner(), +) -> str: stream_client, _, stream_id = await _new_stream(client) - first = _tokens(3) - rest = [*_tokens(2, prefix="rest"), _done()] - await _append(stream_client, stream_id, *first) - async with new_worker(stream_client, TimeoutThenReadWorkflow) as worker: - handle = await stream_client.start_workflow( - TimeoutThenReadWorkflow.run, - stream_id, - id=f"workflow-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - await asyncio.sleep(1.5) - await _append(stream_client, stream_id, *rest) - timed_out, matched = await handle.result() - assert timed_out.texts == _texts(first) - assert timed_out.stop == "TIMEOUT" and not timed_out.done - assert matched.texts == _texts(rest) - assert matched.stop == "MATCH" + 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_read_stops_at_max_records(client: Client): - stream_client, _, stream_id = await _new_stream(client) - tokens = _tokens(10) - await _append(stream_client, stream_id, *tokens) - async with new_worker(stream_client, ReadUntilKindWorkflow) as worker: - handle = await _start_read( - stream_client, - ReadArgs(stream_id, until_kind=None, max_records=3), - worker.task_queue, - ) - result = await handle.result() - assert result.texts == _texts(tokens[:3]) - assert result.stop == "MAX_RECORDS" and not result.done +async def test_predicate_that_raises(client: Client): + assert await _failing_predicate(client, "raise") == ( + f"ValueError: bad token {_BODY_MARK}-t0" + ) -async def test_read_of_closed_stream_is_done(client: Client): - stream_client, _, stream_id = await _new_stream(client) - tokens = _tokens(4) - await _append(stream_client, stream_id, *tokens) - await stream_client.get_stream(StreamRef(id=stream_id), body_type=Token).close() - async with new_worker(stream_client, ReadUntilKindWorkflow) as worker: - handle = await _start_read( - stream_client, ReadArgs(stream_id), worker.task_queue - ) - result = await handle.result() - assert result.texts == _texts(tokens) - assert result.stop == "DONE" and result.done +@_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_read_with_payload_codec(client: Client): - data_converter = DataConverter(payload_codec=SimpleCodec()) - StreamReadState.until_calls = 0 - _, store, handle, result, tokens = await _run_match( - client, data_converter=data_converter - ) - assert result.texts == _texts(tokens) - assert result.typed and result.stop == "MATCH" - # The predicate saw each Record decoded, or its type check would have raised. - assert StreamReadState.until_calls == len(tokens) - run_id = handle.result_run_id - assert run_id - first_run = StreamReadState.results.pop(run_id) - history = await handle.fetch_history() - marker = next( - event.marker_recorded_event_attributes - for event in history.events - if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED +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 marker.details["result"].payloads[0].metadata["simple-codec"] == b"true" - await Replayer( - workflows=[ReadUntilKindWorkflow], - data_converter=data_converter, - stream_config=StreamConfig(stores=[store]), - ).replay_workflow(history) - assert StreamReadState.results.pop(run_id) == first_run + assert unsandboxed == "no error" -async def _cancel_waiting_read(client: Client, task_timeout: timedelta | None) -> float: - """Cancel a Workflow while its read waits, and give the seconds until it ends.""" +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, ReadUntilKindWorkflow) as worker: - handle = await _start_read( + async with new_worker(stream_client, CallWorkflow) as worker: + handle = await _start_call( stream_client, - ReadArgs(stream_id, timeout_seconds=30), + CallArgs(stream_id, timeout_seconds=30, idle_timeout_seconds=30), worker.task_queue, task_timeout=task_timeout, ) @@ -835,101 +1190,90 @@ async def read_stopped() -> None: "about 8 s with the default task timeout of 10 s." ), ) -async def test_read_cancel(client: Client): - assert await _cancel_waiting_read(client, None) < 5 - - -async def test_read_cancel_with_short_task_timeout(client: Client): - assert await _cancel_waiting_read(client, timedelta(seconds=2)) < 5 +async def test_cancel_of_a_long_wait(client: Client): + assert await _cancel_waiting_call(client, None) < 5 -async def test_read_predicate_that_raises(client: Client): - assert await _failing_predicate(client, "raise") == ( - f"ValueError: bad token {_BODY_MARK}-t0" - ) +async def test_cancel_with_short_task_timeout(client: Client): + assert await _cancel_waiting_call(client, timedelta(seconds=2)) < 5 -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)) +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, FailingPredicateWorkflow, workflow_runner=runner + stream_client, CallWorkflow, max_cached_workflows=0 ) as worker: - return await stream_client.execute_workflow( - FailingPredicateWorkflow.run, - args=[stream_id, predicate], - id=f"workflow-{uuid.uuid4()}", - task_queue=worker.task_queue, + handle = await _start_call( + stream_client, + CallArgs(stream_id, timeout_seconds=30, idle_timeout_seconds=30), + worker.task_queue, + task_timeout=timedelta(seconds=2), ) - - -async def test_read_and_reduce_without_stream_config_fail_the_workflow( - client: Client, -): - stream_id = f"stream-{uuid.uuid4()}" - async with new_worker(client, ReadUntilKindWorkflow, JoinTextWorkflow) as worker: - handles: list[WorkflowHandle[Any, Any]] = [ - await _start_read( - client, - ReadArgs(stream_id), - worker.task_queue, - execution_timeout=timedelta(seconds=30), - ), - await client.start_workflow( - JoinTextWorkflow.run, - ReduceArgs(stream_id), - id=f"workflow-{uuid.uuid4()}", - task_queue=worker.task_queue, - execution_timeout=timedelta(seconds=30), - ), - ] - for handle in handles: - 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 - assert cause.message == ( - "The client of this Worker has no Stream Store. Give one in " - "stream_config=StreamConfig(stores=[...])." + # 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() + ) -@_RUNNERS -async def test_read_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 _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_read_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() +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 unsandboxed == "no error" + assert result.texts == _texts(tokens[:2]) + assert result.stop == "MaxBytes" -async def _run_flowing_read( - client: Client, *, task_timeout: timedelta | None = None, seconds: float = 2.0 +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, ReadUntilKindWorkflow) as worker: - handle = await _start_read( + async with new_worker(stream_client, CallWorkflow) as worker: + handle = await _start_call( stream_client, - ReadArgs(stream_id, timeout_seconds=30), + CallArgs(stream_id, timeout_seconds=30, idle_timeout_seconds=5), worker.task_queue, - task_timeout=task_timeout, ) appended = await _append_over(stream_client, stream_id, tokens, seconds) result = await handle.result() @@ -938,282 +1282,92 @@ async def _run_flowing_read( return await handle.fetch_history(), latency -async def test_measure_flowing_read(client: Client): - history, _ = await _run_flowing_read(client) +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_idle_wait(client: Client): +async def test_measure_take_of_forty(client: Client): stream_client, _, stream_id = await _new_stream(client) - async with new_worker(stream_client, ReadUntilKindWorkflow) as worker: - handle = await _start_read( - stream_client, - ReadArgs(stream_id, 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"\nB. Idle wait of 6 s, History events: {counts}") - print(f"B. Workflow Tasks: {counts['WORKFLOW_TASK_COMPLETED']} completed") - - -async def test_measure_match_latency(client: Client): - latencies = [] - for _ in range(5): - _, latency = await _run_flowing_read(client) - latencies.append(latency) - print( - "\nC. Append of the match to the end of the Workflow: " - + ", ".join(f"{latency * 1000:.0f} ms" for latency in latencies) + 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}") -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, +@workflow.defn +class KeepHitsWorkflow: + @workflow.run + async def run(self, stream_id: str) -> int: + read = await workflow.filter_stream( + _stream(stream_id), + lambda record: record.body.kind == "hit", + timedelta(seconds=10), + idle_timeout=timedelta(0), ) - return handle, await handle.result() + return len(read.records) -async def test_read_idle_timeout(client: Client): +async def test_measure_filter_that_keeps_five(client: Client): stream_client, _, stream_id = await _new_stream(client) - tokens = _tokens(3) - await _append(stream_client, stream_id, *tokens) - start = time.monotonic() - _, result = await _run_workflow( - stream_client, - ReadUntilKindWorkflow, - ReadArgs(stream_id, idle_timeout_seconds=1, timeout_seconds=30), - ) - seconds = time.monotonic() - start - assert result.texts == _texts(tokens) - assert result.stop == "IDLE" and not result.done - assert 1 <= seconds < 3 - - -async def test_read_idle_timeout_zero(client: Client): - stream_client, _, stream_id = await _new_stream(client) - empty_id = f"stream-{uuid.uuid4()}" - await stream_client.create_stream(id=empty_id, body_type=Token) - tokens = _tokens(3) - await _append(stream_client, stream_id, *tokens) - for read_id, expected in [(stream_id, _texts(tokens)), (empty_id, [])]: - start = time.monotonic() - _, result = await _run_workflow( - stream_client, - ReadUntilKindWorkflow, - ReadArgs(read_id, idle_timeout_seconds=0, timeout_seconds=30), - ) - assert result.texts == expected - assert result.stop == "IDLE" - assert time.monotonic() - start < 1 - - -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] - - -@pytest.mark.parametrize( - "stop", - ["MATCH", "DONE", "IDLE", "TIMEOUT", "MAX_RECORDS", "MAX_BYTES"], -) -async def test_read_stop_reason(client: Client, stop: str): - stream_client, store, stream_id = await _new_stream(client) - tokens = _tokens(5) - if stop == "MATCH": - tokens[2] = _done() - await _append(stream_client, stream_id, *tokens) - if stop == "DONE": - await stream_client.get_stream(StreamRef(id=stream_id), body_type=Token).close() - if stop == "MAX_BYTES": - sizes = await _record_sizes(store, stream_id) - # The third Record passes the cap by one byte. - cap = sum(sizes[:3]) - 1 - _, (first, second) = await _run_workflow( - stream_client, MaxBytesThenNextWorkflow, stream_id, cap - ) - assert first.texts == _texts(tokens[:2]) - assert first.stop == "MAX_BYTES" - assert second.texts == _texts(tokens[2:3]) - # A cap below the first Record still gives that Record. - _, (alone, _) = await _run_workflow( - stream_client, MaxBytesThenNextWorkflow, stream_id, 1 - ) - assert alone.texts == _texts(tokens[:1]) - assert alone.stop == "MAX_BYTES" - return - args = { - "MATCH": ReadArgs(stream_id), - "DONE": ReadArgs(stream_id), - "IDLE": ReadArgs(stream_id, idle_timeout_seconds=0.5, timeout_seconds=30), - "TIMEOUT": ReadArgs(stream_id, timeout_seconds=1), - "MAX_RECORDS": ReadArgs(stream_id, max_records=3), - }[stop] - _, result = await _run_workflow(stream_client, ReadUntilKindWorkflow, args) - assert result.stop == stop - expected = {"MATCH": tokens[:3], "MAX_RECORDS": tokens[:3]}.get(stop, tokens) - assert result.texts == _texts(expected) - assert result.done == (stop == "DONE") - - -async def test_reduce_joins_text_and_replays_without_store(client: Client): - stream_client, _, stream_id = await _new_stream(client) - tokens = [*_tokens(3), _done(), *_tokens(2, prefix="after")] + 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) - handle, result = await _run_workflow( - stream_client, JoinTextWorkflow, ReduceArgs(stream_id) - ) - assert result == ReduceResult( - value="".join(_texts(tokens[:4])), - count=4, - stop="MATCH", - done=False, - last_offset=result.last_offset, - ) - history = await handle.fetch_history() - assert _markers(history) == ["core_local_activity"] - run_id = handle.result_run_id - assert run_id - first_run = StreamReadState.results.pop(run_id) - empty_store = RecordingStreamStore(MemoryStreamStore()) - StreamReadState.function_calls = 0 - StreamReadState.until_calls = 0 - await Replayer( - workflows=[JoinTextWorkflow], - stream_config=StreamConfig(stores=[empty_store]), - ).replay_workflow(history) - assert StreamReadState.results.pop(run_id) == first_run - assert empty_store.operations == [] - assert StreamReadState.function_calls == 0 - assert StreamReadState.until_calls == 0 + 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_reduce_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, JoinTextWorkflow, ReduceArgs(stream_id, until="tag") - ) - assert result.value == "xtag!" - assert result.count == 3 - assert result.stop == "MATCH" - - -async def test_reduce_on_a_worker_without_cache(client: Client): - stream_client, store, stream_id = await _new_stream( - client, RecordingStreamStore(MemoryStreamStore()) - ) - texts = ["a", "m1", "b", "m2", "c"] - await _append(stream_client, stream_id, *[Token("text", text) for text in texts]) - handle, (first, second) = await _run_workflow( - stream_client, - TwoFoldsWorkflow, - stream_id, - worker_options={"max_cached_workflows": 0}, - ) - assert (first.value, first.count, first.stop) == ("am1", 2, "MATCH") - assert (second.value, second.count, second.stop) == ("bm2", 2, "MATCH") - assert _markers(await handle.fetch_history()) == ["core_local_activity"] * 2 - # A replay of the first fold takes its value from History. - assert store.operations.count("ReadRecords") == 2 - - -async def test_reduce_function_that_raises(client: Client): - stream_client, _, stream_id = await _new_stream(client) - await _append(stream_client, stream_id, *_tokens(1)) - _, result = await _run_workflow(stream_client, FailingFunctionWorkflow, stream_id) - assert result == f"ValueError: bad fold after '' at {_BODY_MARK}-t0" - - -async def test_reduce_typed_values(client: Client): +async def test_measure_idle_wait(client: Client): stream_client, _, stream_id = await _new_stream(client) - tokens = [*_tokens(2), _done()] - await _append(stream_client, stream_id, *tokens) - _, result = await _run_workflow(stream_client, TypedFoldWorkflow, stream_id) - assert result == TypedFoldResult( - summary=Summary(3, ["text", "text", "done"]), - summary_type="Summary", - texts=_texts(tokens), - texts_type="list", - ) - - -async def test_read_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, ReadUntilKindWorkflow, max_cached_workflows=0 - ) as worker: - handle = await _start_read( + async with new_worker(stream_client, CallWorkflow) as worker: + handle = await _start_call( stream_client, - ReadArgs(stream_id, timeout_seconds=30), + 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 read 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_READS.keys() - ) - ) + await asyncio.sleep(6) await _append(stream_client, stream_id, _done()) - result = await handle.result() + await handle.result() counts = _event_counts(await handle.fetch_history()) - reads = store.operations.count("ReadRecords") + 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( - 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_READS.keys() + "\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, ReadUntilWantedWorkflow) as worker: + async with new_worker(stream_client, TakeUntilWantedWorkflow) as worker: handle = await stream_client.start_workflow( - ReadUntilWantedWorkflow.run, + TakeUntilWantedWorkflow.run, stream_id, id=f"workflow-{uuid.uuid4()}", task_queue=worker.task_queue, ) await asyncio.sleep(1) - await handle.signal(ReadUntilWantedWorkflow.want, "never") + await handle.signal(TakeUntilWantedWorkflow.want, "never") async def delivered() -> float: events = (await handle.fetch_history()).events @@ -1244,7 +1398,7 @@ async def delivered() -> float: async def test_measure_signal_wait(client: Client): waits = [await _signal_wait(client) for _ in range(3)] print( - "\nD. Signal to the next Workflow Task Started while a read waits: " + "\nSignal to the next Workflow Task Started while a call waits: " + ", ".join(f"{wait:.2f} s" for wait in waits) ) @@ -1252,25 +1406,68 @@ async def test_measure_signal_wait(client: Client): 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, JoinTextWorkflow) as worker: - handle = await stream_client.start_workflow( - JoinTextWorkflow.run, - ReduceArgs(stream_id, timeout_seconds=30), - id=f"workflow-{uuid.uuid4()}", - task_queue=worker.task_queue, + 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() - marker = next( - event.marker_recorded_event_attributes - for event in history.events - if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED - ) - payload = marker.details["result"].payloads[0] - print(f"\nE. 40 Records folded over 2 s, History events: {_event_counts(history)}") + payload = _marker_payload(history) print( - f"E. Marker payload: {payload.ByteSize()} bytes, data {len(payload.data)} " - f"bytes, folded value {len(result.value)} characters" + 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" From 19569c77019b8a05e4fae88a3a36665e60cf911f Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 01:54:52 -0700 Subject: [PATCH 11/20] Fixed the type and docstring findings of the Stream calls. The type checkers and pydocstyle flagged an unused import, private helpers that other modules use, and module docstrings. The tests that misuse a call on purpose now say so to the checker. --- temporalio/client/_stream_read.py | 8 ++++---- temporalio/client/_stream_store.py | 2 +- temporalio/worker/_stream_read.py | 3 +-- temporalio/workflow/_streams.py | 4 +--- tests/nexus/test_stream_read_client.py | 10 +++++++--- tests/nexus/test_stream_read_workflow.py | 20 +++++++++++--------- 6 files changed, 25 insertions(+), 22 deletions(-) diff --git a/temporalio/client/_stream_read.py b/temporalio/client/_stream_read.py index 148ac18a77..3f9971125c 100644 --- a/temporalio/client/_stream_read.py +++ b/temporalio/client/_stream_read.py @@ -1,5 +1,4 @@ -"""The scan of :py:func:`temporalio.workflow.take_stream`, ``filter_stream`` and -``reduce_stream``, for client, Activity and Workflow code. +"""The scan of the take, filter and reduce calls, for client, Activity and Workflow code. Client and Activity code run the scan straight against the Store. Workflow code runs it in a Local Activity, and the result of that Local Activity is the result of the call. The @@ -31,7 +30,6 @@ import temporalio.exceptions import temporalio.nexus.system import temporalio.nexus.system.stream_service.client -import temporalio.workflow from temporalio.nexus.system.stream_service.models import ( ReadRecordsRequest, StreamRecord, @@ -466,7 +464,9 @@ def keys(self) -> list[tuple[str, str]]: @temporalio.activity.defn(name=_STREAM_CALL_ACTIVITY) -async def _stream_call_activity(call_id: str) -> temporalio.common.RawValue: +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) diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index f3bd345a55..8315b37cc5 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -396,7 +396,7 @@ def _with_local_activity_producer( request.sequence = 1 -def _store_page_reader( +def _store_page_reader( # pyright: ignore[reportUnusedFunction] store: StreamStore, data_converter: temporalio.converter.DataConverter ) -> Callable[ [temporalio.api.streamservice.v1.ReadRecordsRequest], diff --git a/temporalio/worker/_stream_read.py b/temporalio/worker/_stream_read.py index 10a3d57f79..4864984d2f 100644 --- a/temporalio/worker/_stream_read.py +++ b/temporalio/worker/_stream_read.py @@ -1,5 +1,4 @@ -"""The Workflow side of :py:func:`temporalio.workflow.take_stream`, ``filter_stream`` and -``reduce_stream``.""" +"""The Workflow side of the take, filter and reduce calls of a Stream.""" from __future__ import annotations diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py index fce55e604e..0ed847a304 100644 --- a/temporalio/workflow/_streams.py +++ b/temporalio/workflow/_streams.py @@ -214,9 +214,7 @@ def check( ) -def _check_plain_function( - label: str, value: Callable[..., object] | None, *, required: bool -) -> None: +def _check_plain_function(label: str, value: object, *, required: bool) -> None: if value is None: if required: raise TypeError(f"{label} is required") diff --git a/tests/nexus/test_stream_read_client.py b/tests/nexus/test_stream_read_client.py index f4ad1cd8ab..90fc733892 100644 --- a/tests/nexus/test_stream_read_client.py +++ b/tests/nexus/test_stream_read_client.py @@ -164,19 +164,23 @@ async def test_client_functions_must_be_plain(client: Client): 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 take_stream(stream, timedelta(seconds=5), coroutine_until) + await take_stream(stream, timedelta(seconds=5), awaiting) with pytest.raises(ApplicationError) as err: await filter_stream( - stream, lambda record: coroutine_until(record), timedelta(seconds=5) + stream, 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("bad fold") + 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"): diff --git a/tests/nexus/test_stream_read_workflow.py b/tests/nexus/test_stream_read_workflow.py index 1fc2d9d36e..6defb7f271 100644 --- a/tests/nexus/test_stream_read_workflow.py +++ b/tests/nexus/test_stream_read_workflow.py @@ -9,7 +9,7 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass from datetime import datetime, timedelta -from typing import Any +from typing import Any, cast import pytest @@ -138,6 +138,7 @@ def until(record: StreamRecord[Token]) -> bool: 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 @@ -168,10 +169,11 @@ def fold_until(value: list[str], record: StreamRecord[Token]) -> bool: args.max_bytes, ) return _read_result(read) + no_texts: list[str] = [] reduce = await workflow.reduce_stream( stream, _collect, - [], + no_texts, timeout, fold_until if until_kind else None, from_offset, @@ -370,10 +372,12 @@ async def run(self, stream_id: str) -> list[str]: 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) @@ -382,12 +386,10 @@ 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 workflow.take_stream( - stream, - timedelta(seconds=10), - until=_async_until if which == "coroutine" else _until_that_awaits, - ) + await workflow.take_stream(stream, timedelta(seconds=10), until=until) except TypeError as err: return f"TypeError: {err}" except ApplicationError as err: @@ -464,14 +466,14 @@ async def run(self, stream_id: str) -> list[str]: untyped = await workflow.reduce_stream( stream, lambda value, record: Summary(1, [record.body.kind]), - None, + cast(Any, None), quick, idle_timeout=now, ) explicit = await workflow.reduce_stream( stream, lambda value, record: [*value, record.body.text], - (), + cast(Any, ()), quick, idle_timeout=now, value_type=list[str], From b1cf57d220a0f2afea079cead8ab9eea5b799841 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 02:21:05 -0700 Subject: [PATCH 12/20] Fitted the Stream calls to the move of BufferedAppender. The client must not import a contrib package. So the scan keeps the list of writers to flush, and the appender adds itself to it. --- temporalio/client/_stream_read.py | 36 +++++++++++++++++++----- temporalio/contrib/streams/_appender.py | 26 +++-------------- temporalio/workflow/_streams.py | 12 ++++---- tests/nexus/test_stream_read_client.py | 2 +- tests/nexus/test_stream_read_workflow.py | 2 +- 5 files changed, 41 insertions(+), 37 deletions(-) diff --git a/temporalio/client/_stream_read.py b/temporalio/client/_stream_read.py index 3f9971125c..390a0b32e4 100644 --- a/temporalio/client/_stream_read.py +++ b/temporalio/client/_stream_read.py @@ -12,6 +12,7 @@ from __future__ import annotations import base64 +import contextvars import dataclasses import inspect import json @@ -19,7 +20,7 @@ import time from collections.abc import Awaitable, Callable from datetime import timedelta -from typing import Any +from typing import Any, Protocol import temporalio.activity import temporalio.api.common.v1 @@ -44,7 +45,6 @@ ) from ._client import Client -from ._stream_appender import _CALL_APPENDERS, _flush_call_appenders from ._stream_store import ( _STREAM_SERVICE, StreamStore, @@ -59,6 +59,19 @@ # 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 +) + _ReadPage = Callable[ [temporalio.api.streamservice.v1.ReadRecordsRequest], Awaitable[temporalio.api.streamservice.v1.ReadRecordsResponse], @@ -255,8 +268,8 @@ async def _scan_and_flush( In Workflow code the result goes into History only after those Records are in the Store. """ - appenders: list[Any] = [] - token = _CALL_APPENDERS.set(appenders) + flushes: list[_CallFlush] = [] + token = _CALL_FLUSHES.set(flushes) try: outcome = await _scan( read_page, @@ -267,14 +280,23 @@ async def _scan_and_flush( ) except Exception: # The error of the scan wins over an error of the write. - await _flush_call_appenders(appenders, raise_error=False) + await _flush_all(flushes, raise_error=False) raise finally: - _CALL_APPENDERS.reset(token) - await _flush_call_appenders(appenders, raise_error=True) + _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 + + @dataclasses.dataclass(frozen=True) class _StreamCallResult: """The result of one call, as the payload that History keeps for a Workflow.""" diff --git a/temporalio/contrib/streams/_appender.py b/temporalio/contrib/streams/_appender.py index ab2805c3d7..20570da727 100644 --- a/temporalio/contrib/streams/_appender.py +++ b/temporalio/contrib/streams/_appender.py @@ -7,7 +7,6 @@ from __future__ import annotations import asyncio -import contextvars from collections.abc import Awaitable, Callable from datetime import timedelta from types import TracebackType @@ -16,6 +15,7 @@ from typing_extensions import Self import temporalio.activity +import temporalio.client._stream_read import temporalio.nexus.system.stream_service.client import temporalio.workflow @@ -32,12 +32,6 @@ _DEFAULT_MAX_DELAY = timedelta(milliseconds=100) -# The appenders 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_APPENDERS: contextvars.ContextVar[list[BufferedAppender[Any]] | None] = ( - contextvars.ContextVar("temporal_stream_call_appenders", default=None) -) - class BufferedAppender(Generic[BodyT]): """Batches the appends of one Stream. @@ -140,9 +134,9 @@ def append(self, *records: BodyT) -> None: self._check_delay() if not temporalio.workflow.in_workflow(): self._client_stream = self._outside_workflow_stream() - appenders = _CALL_APPENDERS.get() - if appenders is not None and self not in appenders: - appenders.append(self) + 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() @@ -281,15 +275,3 @@ def _fail(self, error: BaseException, count: int) -> None: if self._timer is not None: self._timer.cancel() self._timer = None - - -async def _flush_call_appenders( # pyright: ignore[reportUnusedFunction] - appenders: list[BufferedAppender[Any]], *, raise_error: bool -) -> None: - """Write the Records of the appenders that the functions of a Stream call used.""" - for appender in appenders: - try: - await appender.flush() - except Exception: - if raise_error: - raise diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py index 0ed847a304..1909d3216b 100644 --- a/temporalio/workflow/_streams.py +++ b/temporalio/workflow/_streams.py @@ -281,8 +281,8 @@ async def take_stream( 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.client.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.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. @@ -375,8 +375,8 @@ async def filter_stream( 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.client.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.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. @@ -473,8 +473,8 @@ async def reduce_stream( 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.client.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.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. diff --git a/tests/nexus/test_stream_read_client.py b/tests/nexus/test_stream_read_client.py index 90fc733892..56826355bd 100644 --- a/tests/nexus/test_stream_read_client.py +++ b/tests/nexus/test_stream_read_client.py @@ -10,13 +10,13 @@ from temporalio import activity, workflow from temporalio.client import ( - BufferedAppender, Client, StreamReadStop, filter_stream, reduce_stream, take_stream, ) +from temporalio.contrib.streams import BufferedAppender from temporalio.exceptions import ApplicationError from temporalio.nexus.system.stream_service import StreamRecord, StreamRef from temporalio.testing import MemoryStreamStore diff --git a/tests/nexus/test_stream_read_workflow.py b/tests/nexus/test_stream_read_workflow.py index 6defb7f271..c8fa125202 100644 --- a/tests/nexus/test_stream_read_workflow.py +++ b/tests/nexus/test_stream_read_workflow.py @@ -19,7 +19,6 @@ from temporalio.api.stream.v1 import StreamRef as StreamRefMessage from temporalio.api.streamservice.v1 import ReadRecordsRequest, ReadRecordsResponse from temporalio.client import ( - BufferedAppender, Client, StreamConfig, WorkflowFailureError, @@ -27,6 +26,7 @@ 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 ( From d7564a7809cfdbadef00b9ac19f2c2489f0bbf33 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 02:21:18 -0700 Subject: [PATCH 13/20] Moved Core to the Stream WIT with take, filter and reduce. The WIT now declares the three calls as methods of the Stream resource. Only that file differs from the old pin, so the bridge needs no rebuild. --- temporalio/bridge/sdk-core | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index 59f4b33f1d..81155733bf 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit 59f4b33f1d90bb5ac4582ea465b4e8597c507214 +Subproject commit 81155733bf7d88ff829780251953e6640e847825 From b1a61aaf223410a524cdd634ddcc12846316e1d2 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 02:21:35 -0700 Subject: [PATCH 14/20] Pinned the Stream generator to the nexgen build with helper methods. The Stream WIT marks take, filter and reduce with `@nexus.helper`. Only a fork of nexgen knows that tag until its `main` has it. --- .github/workflows/ci.yml | 4 ++-- scripts/gen_nexus_system_api.py | 8 +++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c60bb33161..55173da580 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/scripts/gen_nexus_system_api.py b/scripts/gen_nexus_system_api.py index 48998fc7c8..0b073adca9 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}": From 47167798b8eb821cd1ed00f1d43fbfa3f0ec7784 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 02:21:48 -0700 Subject: [PATCH 15/20] Regenerated the Streams API with take, filter and reduce. The client Stream and the Workflow Stream get the three calls as methods. Each body is one call of a support helper, and the result types are generated models. --- .../nexus/system/stream_service/__init__.py | 6 ++ .../client/_resources/stream.py | 70 ++++++++++++++++++ .../nexus/system/stream_service/models.py | 31 ++++++++ .../workflow/_resources/stream.py | 73 +++++++++++++++++++ 4 files changed, 180 insertions(+) diff --git a/temporalio/nexus/system/stream_service/__init__.py b/temporalio/nexus/system/stream_service/__init__.py index 35b198843c..b68279df6f 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 c3fa617b6d..afac0ac8b8 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 4848a2ea58..fd3a0865dd 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 df548d2528..a237c82dad 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( From f81fa77dd8d3660c70491d89aa5cc952c37ec263 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 02:41:30 -0700 Subject: [PATCH 16/20] Wired the generated take, filter and reduce to the support helpers. The methods of the Stream are the public API, so the module functions and the hand-written result types go. A client reads its pages with `Stream.read`, so its interceptors see each one, and only Workflow code limits the size of a result, since only there it goes into History. --- changelog/added/patient-reader-heron.md | 2 +- temporalio/client/__init__.py | 22 - temporalio/client/_impl.py | 4 + temporalio/client/_stream_read.py | 336 +++++++------ temporalio/client/_stream_store.py | 6 +- temporalio/common.py | 40 ++ temporalio/contrib/streams/_appender.py | 5 +- temporalio/nexus/system/_support.py | 463 ++++++++++++++++++ temporalio/worker/_replayer.py | 5 +- temporalio/worker/_stream_read.py | 9 +- temporalio/worker/_workflow_instance.py | 3 +- temporalio/workflow/__init__.py | 22 - temporalio/workflow/_context.py | 3 +- temporalio/workflow/_streams.py | 546 ---------------------- tests/nexus/test_stream_read_client.py | 225 ++++++++- tests/nexus/test_stream_read_workflow.py | 92 ++-- tests/nexus/test_stream_store_workflow.py | 6 +- 17 files changed, 967 insertions(+), 822 deletions(-) delete mode 100644 temporalio/workflow/_streams.py diff --git a/changelog/added/patient-reader-heron.md b/changelog/added/patient-reader-heron.md index 0508d5456f..ab1bf2f93b 100644 --- a/changelog/added/patient-reader-heron.md +++ b/changelog/added/patient-reader-heron.md @@ -1 +1 @@ -Add `take_stream`, `filter_stream` and `reduce_stream` to `temporalio.workflow` and `temporalio.client`, which client, Activity and Workflow code call to scan a Stream until a predicate matches, the Stream is done, or a limit is reached. +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. diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index eeebf78685..6fa6dd710c 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -27,18 +27,6 @@ ServiceClient, TLSConfig, ) -from temporalio.workflow._streams import ( - DEFAULT_STREAM_IDLE_TIMEOUT, - DEFAULT_STREAM_MAX_BYTES, - DEFAULT_STREAM_MAX_RECORDS, - MAX_STREAM_RESULT_BYTES, - StreamRead, - StreamReadStop, - StreamReduce, - filter_stream, - reduce_stream, - take_stream, -) from ..common import HeaderCodecBehavior from ..types import ( @@ -368,16 +356,6 @@ "TaskReachabilityType", "CloudOperationsClient", "Plugin", - "DEFAULT_STREAM_IDLE_TIMEOUT", - "DEFAULT_STREAM_MAX_BYTES", - "DEFAULT_STREAM_MAX_RECORDS", - "MAX_STREAM_RESULT_BYTES", - "StreamRead", - "StreamReadStop", - "StreamReduce", - "filter_stream", - "reduce_stream", - "take_stream", "StreamConfig", "StreamSelectContext", "StreamStore", diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index 4766126ae8..5b234ec886 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 index 390a0b32e4..aa57c64b0d 100644 --- a/temporalio/client/_stream_read.py +++ b/temporalio/client/_stream_read.py @@ -1,9 +1,9 @@ """The scan of the take, filter and reduce calls, for client, Activity and Workflow code. -Client and Activity code run the scan straight against the Store. Workflow code runs it in a -Local Activity, and the result of that Local Activity is the result of the call. 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. +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 result of the call. 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. .. warning:: This API is an experimental prototype and subject to change. @@ -18,9 +18,10 @@ import json import threading import time -from collections.abc import Awaitable, Callable +from abc import ABC, abstractmethod +from collections.abc import Callable, Sequence from datetime import timedelta -from typing import Any, Protocol +from typing import Any, Generic, Protocol, TypeVar import temporalio.activity import temporalio.api.common.v1 @@ -29,29 +30,20 @@ import temporalio.common import temporalio.converter import temporalio.exceptions -import temporalio.nexus.system import temporalio.nexus.system.stream_service.client +from temporalio.nexus.system._support import _StreamCall, _StreamCallKind from temporalio.nexus.system.stream_service.models import ( - ReadRecordsRequest, - StreamRecord, -) -from temporalio.workflow._streams import ( - MAX_STREAM_RESULT_BYTES, StreamRead, StreamReadStop, + StreamRecord, StreamReduce, - _StreamCall, - _StreamCallKind, ) - -from ._client import Client -from ._stream_store import ( - _STREAM_SERVICE, - StreamStore, - _activity_target, - _store_page_reader, +from temporalio.nexus.system.stream_service.workflow._resources.stream import ( + Stream as _WorkflowStream, ) +from ._stream_store import StreamStore, _store_page_reader + _RESULT_METADATA_KEY = "temporal-stream-read" _RESULT_METADATA_VALUE = b"v3" _RECORDS_MESSAGE_TYPE = b"temporal.api.streamservice.v1.ReadRecordsResponse" @@ -72,11 +64,137 @@ async def flush(self) -> object: ... "temporal_stream_call_flushes", default=None ) -_ReadPage = Callable[ - [temporalio.api.streamservice.v1.ReadRecordsRequest], - Awaitable[temporalio.api.streamservice.v1.ReadRecordsResponse], -] +# 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: @@ -114,8 +232,8 @@ def _call_function( @dataclasses.dataclass -class _ScanOutcome: - kept: list[temporalio.api.stream.v1.StreamRecord] +class _ScanOutcome(Generic[_RecordT]): + kept: list[_RecordT] """The Records that the call gives back. A fold keeps none.""" count: int """How many Records the call kept or folded.""" @@ -126,16 +244,16 @@ class _ScanOutcome: """The value of a fold, or ``None``.""" -class _ScanState: +class _ScanState(Generic[_RecordT]): """What a scan got so far.""" def __init__( - self, call: _StreamCall, to_record: _ToRecord, *, in_local_activity: bool + self, call: _StreamCall, pages: _Pages[_RecordT], *, in_local_activity: bool ) -> None: self._call = call - self._to_record = to_record + self._pages = pages self._in_local_activity = in_local_activity - self.kept: list[temporalio.api.stream.v1.StreamRecord] = [] + self.kept: list[_RecordT] = [] self.count = 0 self.bytes = 0 self.last_offset = call.from_offset @@ -151,22 +269,20 @@ def _call_function( def bytes_left(self) -> int: return self._call.max_bytes - self.bytes - def look( - self, record: temporalio.api.stream.v1.StreamRecord, *, ends_stream: bool - ) -> StreamReadStop | None: + 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._to_record(record) + typed = self._pages.typed(record) keep = True if call.predicate is not None: keep = bool(self._call_function("predicate", call.predicate, typed)) - size = record.ByteSize() + 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 self.bytes + size > call.max_bytes: return StreamReadStop.MaxBytes - self.last_offset = record.offset + self.last_offset = self._pages.offset(record) if keep: self.count += 1 self.bytes += size @@ -186,7 +302,9 @@ def look( return StreamReadStop.MaxRecords return None - def end(self, stop: StreamReadStop, *, done: bool = False) -> _ScanOutcome: + def end( + self, stop: StreamReadStop, *, done: bool = False + ) -> _ScanOutcome[_RecordT]: return _ScanOutcome( kept=self.kept, count=self.count, @@ -198,20 +316,15 @@ def end(self, stop: StreamReadStop, *, done: bool = False) -> _ScanOutcome: async def _scan( - read_page: _ReadPage, - *, - ref: temporalio.api.stream.v1.StreamRef, - call: _StreamCall, - to_record: _ToRecord, - in_local_activity: bool, -) -> _ScanOutcome: - """Read pages after ``call.from_offset`` until a stop of :py:class:`StreamReadStop`.""" + 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, to_record, in_local_activity=in_local_activity) + 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 @@ -230,40 +343,35 @@ def idle(now: float) -> bool: wait = min( _MAX_PAGE_WAIT, deadline - now, max(0.0, last_arrival + idle_timeout - now) ) - request = temporalio.api.streamservice.v1.ReadRecordsRequest( - ref=ref, - from_offset=state.last_offset or b"", - max_records=min(call.max_records - state.count, _MAX_PAGE_RECORDS) - if page_cap - else _MAX_PAGE_RECORDS, - max_bytes=max(state.bytes_left(), 1) if page_cap else 0, + 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(state.bytes_left(), 1) if page_cap else 0, + wait=timedelta(seconds=wait), + ) ) - request.wait.FromTimedelta(timedelta(seconds=wait)) - response = await read_page(request) - if response.records: + if records: last_arrival = time.monotonic() - for index, record in enumerate(response.records): + for index, record in enumerate(records): # A done page ends with the last Record of the Stream. - ends_stream = response.done and index == len(response.records) - 1 + 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 response.done and not response.records: + if done and not records: return state.end(StreamReadStop.Done, done=True) - if not response.records and idle(time.monotonic()): + if not records and idle(time.monotonic()): return state.end(StreamReadStop.Idle) async def _scan_and_flush( - read_page: _ReadPage, - *, - ref: temporalio.api.stream.v1.StreamRef, - call: _StreamCall, - to_record: _ToRecord, - in_local_activity: bool, -) -> _ScanOutcome: + 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. @@ -271,13 +379,7 @@ async def _scan_and_flush( flushes: list[_CallFlush] = [] token = _CALL_FLUSHES.set(flushes) try: - outcome = await _scan( - read_page, - ref=ref, - call=call, - to_record=to_record, - in_local_activity=in_local_activity, - ) + 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) @@ -395,10 +497,10 @@ def _result_payload( call: _StreamCall, store: str, stream_id: str, - outcome: _ScanOutcome, + outcome: _ScanOutcome[temporalio.api.stream.v1.StreamRecord], payload_converter: temporalio.converter.PayloadConverter, ) -> temporalio.api.common.v1.Payload: - """Give the result as a payload, or raise when it is over the limit.""" + """Give the result for History as a payload, or raise when it is over the limit.""" payload = _StreamCallResult( kind=call.kind, store=store, @@ -414,10 +516,11 @@ def _result_payload( else None, ).to_payload() size = payload.ByteSize() - if size > MAX_STREAM_RESULT_BYTES: + 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 {MAX_STREAM_RESULT_BYTES} bytes. Lower " + 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, @@ -496,17 +599,19 @@ async def _stream_call_activity( # pyright: ignore[reportUnusedFunction] # 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 " - "take_stream, filter_stream and reduce_stream in Workflow code run this " - "Activity.", + 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( - _store_page_reader(pending.store, temporalio.activity.client().data_converter), - ref=pending.ref, + _StorePages( + pending.store, + temporalio.activity.client().data_converter, + pending.ref, + pending.to_record, + ), call=pending.call, - to_record=pending.to_record, in_local_activity=True, ) return temporalio.common.RawValue( @@ -520,16 +625,17 @@ async def _stream_call_activity( # pyright: ignore[reportUnusedFunction] ) -def _stream_client(stream: Any) -> Client: +def _client_stream( + stream: Any, +) -> temporalio.nexus.system.stream_service.client.Stream[Any]: if isinstance(stream, temporalio.nexus.system.stream_service.client.Stream): - client = getattr(getattr(stream, "_nexus_client", None), "_client", None) - if isinstance(client, Client): - return client - raise TypeError("The client Stream has no Temporal client") - if temporalio.activity.in_activity(): + 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 the Activity's client reaches its Store. - return temporalio.activity.client() + # 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." @@ -539,46 +645,12 @@ def _stream_client(stream: Any) -> Client: 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 straight against the Store.""" - client = _stream_client(stream) - stream_config = client.config(active_config=True).get("stream_config") - if stream_config is None: - raise temporalio.exceptions.StreamConfigError( - "This client has no Stream Store. Give one in " - "stream_config=StreamConfig(stores=[...])." - ) - request = ReadRecordsRequest(ref=stream.ref, from_offset=call.from_offset) - data_converter = client.data_converter - serialization_context = temporalio.nexus.system._get_serialization_context( - _STREAM_SERVICE.name, "ReadRecords", request - ) - if serialization_context is not None: - data_converter = data_converter.with_context(serialization_context) - payload_converter = temporalio.nexus.system._get_payload_converter( - data_converter.payload_converter, data_converter.failure_converter - ) - message = payload_converter._to_transfer_value(request) - assert isinstance(message, temporalio.api.streamservice.v1.ReadRecordsRequest) - store = stream_config._store_for(stream.ref, _activity_target()) - 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) + """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( - _store_page_reader(store, data_converter), - ref=message.ref, - call=call, - to_record=to_record, - in_local_activity=False, - ) - # The limit is the same as in Workflow code, so a call gives the same result anywhere. - _result_payload( - call=call, - store=store.name, - stream_id=message.ref.id, - outcome=outcome, - payload_converter=data_converter.payload_converter, + _ClientPages(_client_stream(stream)), call=call, in_local_activity=False ) if call.kind is _StreamCallKind.REDUCE: return StreamReduce( @@ -589,7 +661,7 @@ def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: done=outcome.done, ) return StreamRead( - records=[to_record(record) for record in outcome.kept], + 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 8315b37cc5..9e01f460bb 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -328,9 +328,9 @@ def check_response( temporalio.api.streamservice.v1.ReadRecordsRequest, temporalio.api.streamservice.v1.ReadRecordsResponse, workflow_refusal=( - "Workflow code can't run a ReadRecords operation on a Store. Use " - "temporalio.workflow.take_stream, filter_stream or reduce_stream, which " - "read in one Local Activity." + "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( diff --git a/temporalio/common.py b/temporalio/common.py index 8bcc66a4ae..24996df680 100644 --- a/temporalio/common.py +++ b/temporalio/common.py @@ -1455,6 +1455,46 @@ 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``. + +.. 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. + +.. 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/_appender.py b/temporalio/contrib/streams/_appender.py index 20570da727..af524cfe40 100644 --- a/temporalio/contrib/streams/_appender.py +++ b/temporalio/contrib/streams/_appender.py @@ -51,9 +51,8 @@ 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 :py:func:`temporalio.workflow.take_stream`, ``filter_stream`` or - ``reduce_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, + * 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 diff --git a/temporalio/nexus/system/_support.py b/temporalio/nexus/system/_support.py index 55f53b59a1..be5e5016bb 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,457 @@ 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 + 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") + 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, + 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. + The marker holds the result, with the Records and their bodies, so the encoded result + must fit in :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES`. 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 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 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. + + 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 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. + The marker holds the result, with the kept Records and their bodies, so the encoded + result must fit in :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES`. 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 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 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. + + 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 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 + 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/worker/_replayer.py b/temporalio/worker/_replayer.py index e3b4b6d7fa..6a0085833d 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -70,9 +70,8 @@ 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. A - :py:func:`temporalio.workflow.take_stream`, ``filter_stream`` or - ``reduce_stream`` takes its result from the history, so its replay needs no + replay calls no method of its Stores. The ``take``, ``filter`` or ``reduce`` + of a Workflow Stream takes its result from the history, so its replay needs no Stream config. Note, unlike the worker, for the replayer the workflow_task_executor diff --git a/temporalio/worker/_stream_read.py b/temporalio/worker/_stream_read.py index 4864984d2f..410ccdc17c 100644 --- a/temporalio/worker/_stream_read.py +++ b/temporalio/worker/_stream_read.py @@ -16,17 +16,14 @@ _StreamCallResult, ) from temporalio.client._stream_store import _STREAM_SERVICE, StreamConfig +from temporalio.nexus.system._support import _StreamCall, _StreamCallKind from temporalio.nexus.system.stream_service.models import ( ReadRecordsRequest, - StreamRecord, -) -from temporalio.nexus.system.stream_service.workflow._resources.stream import Stream -from temporalio.workflow._streams import ( StreamRead, + StreamRecord, StreamReduce, - _StreamCall, - _StreamCallKind, ) +from temporalio.nexus.system.stream_service.workflow._resources.stream import Stream from ._stream_store import _no_stream_config_error, _store_errors, _workflow_target diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 99d4a07820..3d2311b112 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 @@ -1997,7 +1998,7 @@ def workflow_upsert_search_attributes( ) async def workflow_stream_call( - self, stream: Any, call: temporalio.workflow._streams._StreamCall + 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. diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index 7f9ac3f6d7..60b02cc543 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -155,18 +155,6 @@ logger, unsafe, ) -from ._streams import ( - DEFAULT_STREAM_IDLE_TIMEOUT, - DEFAULT_STREAM_MAX_BYTES, - DEFAULT_STREAM_MAX_RECORDS, - MAX_STREAM_RESULT_BYTES, - StreamRead, - StreamReadStop, - StreamReduce, - filter_stream, - reduce_stream, - take_stream, -) from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowConfig, @@ -274,16 +262,6 @@ "SandboxImportNotificationPolicy", "logger", "unsafe", - "DEFAULT_STREAM_IDLE_TIMEOUT", - "DEFAULT_STREAM_MAX_BYTES", - "DEFAULT_STREAM_MAX_RECORDS", - "MAX_STREAM_RESULT_BYTES", - "StreamRead", - "StreamReadStop", - "StreamReduce", - "filter_stream", - "reduce_stream", - "take_stream", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", diff --git a/temporalio/workflow/_context.py b/temporalio/workflow/_context.py index 2580901004..4f7b8b82f4 100644 --- a/temporalio/workflow/_context.py +++ b/temporalio/workflow/_context.py @@ -21,11 +21,12 @@ 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 from ._nexus import NexusOperationCancellationType, NexusOperationHandle - from ._streams import _StreamCall from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowHandle, diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py deleted file mode 100644 index 1909d3216b..0000000000 --- a/temporalio/workflow/_streams.py +++ /dev/null @@ -1,546 +0,0 @@ -"""Calls that scan a Stream until a stop: take, filter and reduce. - -Client, Activity and Workflow code call the same three functions. A generated method of the -client ``Stream`` and of the Workflow ``Stream`` calls them too. - -.. warning:: - This API is an experimental prototype and subject to change. -""" - -from __future__ import annotations - -import dataclasses -import enum -import inspect -import typing -from collections.abc import Callable -from datetime import timedelta -from typing import TYPE_CHECKING, Any, Generic, TypeVar - -from ._context import _Runtime - -if TYPE_CHECKING: - # The generated Stream modules import `temporalio.workflow`, so this module names - # their types for the type checker only. - import temporalio.nexus.system.stream_service.client - import temporalio.nexus.system.stream_service.workflow._resources.stream - from temporalio.nexus.system.stream_service.models import StreamRecord - -__all__ = [ - "DEFAULT_STREAM_IDLE_TIMEOUT", - "DEFAULT_STREAM_MAX_BYTES", - "DEFAULT_STREAM_MAX_RECORDS", - "MAX_STREAM_RESULT_BYTES", - "StreamRead", - "StreamReadStop", - "StreamReduce", - "filter_stream", - "reduce_stream", - "take_stream", -] - -BodyT = TypeVar("BodyT") -ValueT = TypeVar("ValueT") - -DEFAULT_STREAM_IDLE_TIMEOUT = timedelta(seconds=1) -"""How long a call waits for a new Record when ``idle_timeout`` is ``None``.""" - -DEFAULT_STREAM_MAX_RECORDS = 1_000 -"""The most Records of one call when ``max_records`` is ``None``.""" - -DEFAULT_STREAM_MAX_BYTES = 256 * 1024 -"""The most bytes of the Records of one call when ``max_bytes`` is ``None``.""" - -MAX_STREAM_RESULT_BYTES = 1024 * 1024 -"""The most bytes of the encoded result of one call. - -In Workflow code the result 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 the same limit, so a call gives the same result everywhere. -""" - - -class StreamReadStop(enum.IntEnum): - """Why a call on a Stream ended. - - When several hold, the first in the order of the values is the reason. - - .. warning:: - This API is an experimental prototype and subject to change. - """ - - Match = 0 - """``until`` accepted a Record.""" - - Done = 1 - """The Stream is closed and the call looked at its last Record.""" - - MaxRecords = 2 - """The call got ``max_records`` Records.""" - - MaxBytes = 3 - """The next Record would pass ``max_bytes``. The next call gets it.""" - - Idle = 4 - """No new Record arrived for ``idle_timeout``.""" - - Timeout = 5 - """``timeout`` passed.""" - - -@dataclasses.dataclass(slots=True, kw_only=True) -class StreamRead(Generic[BodyT]): - """What :py:func:`take_stream` or :py:func:`filter_stream` gave. - - .. warning:: - This API is an experimental prototype and subject to change. - """ - - records: list[StreamRecord[BodyT]] - """The Records that the call gives back, in order.""" - - last_offset: bytes | None = None - """Where the scan stopped. - - It is the offset of the last Record that the scan looked at, kept or not, or the - ``from_offset`` of a call that looked at none. Give it as the ``from_offset`` of the next - call, so a loop loses and repeats no Record. - """ - - stop: StreamReadStop - """Why the call ended.""" - - done: bool - """Whether the Stream is closed and the scan looked at its last Record. - - It is the state of the Stream, so it can be true when :py:attr:`stop` is ``Match``. - """ - - -@dataclasses.dataclass(slots=True, kw_only=True) -class StreamReduce(Generic[ValueT]): - """What :py:func:`reduce_stream` gave. - - .. warning:: - This API is an experimental prototype and subject to change. - """ - - value: ValueT - """The value after the last Record, or ``initial`` when the fold got none.""" - - count: int - """The number of Records that the fold got.""" - - last_offset: bytes | None = None - """Where the scan stopped, as in :py:attr:`StreamRead.last_offset`.""" - - stop: StreamReadStop - """Why the fold ended.""" - - done: bool - """Whether the Stream is closed and the fold got its last Record.""" - - -class _StreamCallKind(enum.Enum): - TAKE = "take" - FILTER = "filter" - REDUCE = "reduce" - - -@dataclasses.dataclass(frozen=True) -class _StreamCall: - """One call after the checks of the public functions. It holds the caller's functions.""" - - kind: _StreamCallKind - from_offset: bytes | None - timeout: timedelta - idle_timeout: timedelta - max_records: int - max_bytes: int - until: Callable[..., object] | None = None - predicate: Callable[..., object] | None = None - function: Callable[..., object] | None = None - initial: Any = None - value_type: 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: Callable[..., object] | None = None, - predicate: Callable[..., object] | None = None, - function: Callable[..., object] | None = None, - initial: Any = None, - value_type: 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 = DEFAULT_STREAM_IDLE_TIMEOUT - elif idle_timeout < timedelta(): - raise ValueError("idle_timeout must not be negative") - if max_records is None: - max_records = DEFAULT_STREAM_MAX_RECORDS - elif max_records < 1: - raise ValueError("max_records must be at least 1") - if max_bytes is None: - max_bytes = 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, - until=until, - predicate=predicate, - function=function, - initial=initial, - value_type=value_type, - ) - - -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: Callable[..., object], initial: Any, value_type: Any -) -> Any: - if value_type is not None: - return value_type - 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: Any, call: _StreamCall) -> Any: - runtime = _Runtime.maybe_current() - if runtime is not None: - return await runtime.workflow_stream_call(stream, call) - # The client package imports this module, so its scan is imported at the first 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: 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]: - """Give every Record of a Stream up to a stop. - - ``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 of :py:class:`StreamReadStop` that holds, in the order of their - values. ``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 result are the same in each. In a - client or an Activity, the call reads the Store straight away. - - 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 Records and their bodies. 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 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 scan. - timeout: The most time for the call. It must be positive. - 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:`DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow code, a longer - wait keeps the Workflow Task open, so a Signal, an Update or a cancel waits for - the next Workflow Task heartbeat, and each heartbeat adds events to History. - max_records: The most Records to give. ``None`` is - :py:data:`DEFAULT_STREAM_MAX_RECORDS`, 1,000. - max_bytes: A cap on the bytes of the Records to give. A Record that would pass - it 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:`DEFAULT_STREAM_MAX_BYTES`, 256 KiB. - - 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.ApplicationError: A function raised, or returned an - awaitable. Or the encoded result is over :py:data:`MAX_STREAM_RESULT_BYTES`, - and the message names both sizes. In Workflow code it is also raised 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: 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]: - """Give the Records of a Stream that a predicate keeps, up to a stop. - - ``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 of :py:class:`StreamReadStop` that holds, in the order of their - values. ``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 result are the same in each. In a - client or an Activity, the call reads the Store straight away. - - 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 kept Records and their bodies. 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 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 scan. - predicate: Keeps a Record when it returns true. - timeout: The most time for the call. It must be positive. - 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:`DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow code, a longer - wait keeps the Workflow Task open, so a Signal, an Update or a cancel waits for - the next Workflow Task heartbeat, and each heartbeat adds events to History. - max_records: The most Records to keep. ``None`` is - :py:data:`DEFAULT_STREAM_MAX_RECORDS`, 1,000. - max_bytes: A cap on the bytes of the Records to keep. A Record that would pass - it 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:`DEFAULT_STREAM_MAX_BYTES`, 256 KiB. - - 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.ApplicationError: A function raised, or returned an - awaitable. Or the encoded result is over :py:data:`MAX_STREAM_RESULT_BYTES`, - and the message names both sizes. In Workflow code it is also raised 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: 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, - *, - value_type: type[ValueT] | None = None, -) -> StreamReduce[ValueT]: - """Fold the Records of a Stream into one value, up to a stop. - - 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 of :py:class:`StreamReadStop` that holds, in the order of their - values. ``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 result are the same in each. In a - client or an Activity, the call reads the Store straight away. - - 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. 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 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. - 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:`DEFAULT_STREAM_IDLE_TIMEOUT`, 1 second. In Workflow code, a longer - wait keeps the Workflow Task open, so a Signal, an Update or a cancel waits for - the next Workflow Task heartbeat, and each heartbeat adds events to History. - max_records: The most Records to fold. ``None`` is - :py:data:`DEFAULT_STREAM_MAX_RECORDS`, 1,000. - max_bytes: A cap on the bytes of the Records to fold. A Record that would pass - it 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:`DEFAULT_STREAM_MAX_BYTES`, 256 KiB. - value_type: The type that Workflow code decodes the value as. ``None`` takes the - return annotation of ``function``, else the type of ``initial``, and no type - when ``initial`` is ``None``. Client and Activity code get the value itself. - - 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.ApplicationError: A function raised, or returned an - awaitable. Or the encoded result is over :py:data:`MAX_STREAM_RESULT_BYTES`, - and the message names both sizes. In Workflow code it is also raised 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, - value_type=_fold_value_type(function, initial, value_type), - ) - return await _run_stream_call(stream, call) diff --git a/tests/nexus/test_stream_read_client.py b/tests/nexus/test_stream_read_client.py index 56826355bd..ae7547e124 100644 --- a/tests/nexus/test_stream_read_client.py +++ b/tests/nexus/test_stream_read_client.py @@ -11,14 +11,24 @@ from temporalio import activity, workflow from temporalio.client import ( Client, - StreamReadStop, - filter_stream, - reduce_stream, - take_stream, + Interceptor, + NexusOperationHandle, + OutboundInterceptor, + StartNexusOperationInput, ) from temporalio.contrib.streams import BufferedAppender from temporalio.exceptions import ApplicationError -from temporalio.nexus.system.stream_service import StreamRecord, StreamRef +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 @@ -26,12 +36,17 @@ _STOPS, CallArgs, CallResult, + CallWorkflow, + OverLimitWorkflow, Token, _append, _call, _check_stop, _close, _new_stream, + _offsets, + _record_sizes, + _run_workflow, _tokens, _with_store, ) @@ -57,8 +72,7 @@ async def _filter_loop(stream: Any) -> list[str]: found: list[str] = [] offset: bytes | None = None while True: - read = await filter_stream( - stream, + read = await stream.filter( lambda record: record.body.kind == "hit", timedelta(seconds=10), from_offset=offset, @@ -171,11 +185,9 @@ async def coroutine_until(record: StreamRecord[Token]) -> bool: awaiting: Any = coroutine_until with pytest.raises(TypeError, match="until must be a plain function"): - await take_stream(stream, timedelta(seconds=5), awaiting) + await stream.take(timedelta(seconds=5), awaiting) with pytest.raises(ApplicationError) as err: - await filter_stream( - stream, lambda record: awaiting(record), timedelta(seconds=5) - ) + 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") @@ -184,23 +196,191 @@ def raises(value: int, record: StreamRecord[Token]) -> int: # In a client, the error of the caller's function reaches it as it is. with pytest.raises(ValueError, match="bad fold"): - await reduce_stream(stream, raises, 0, timedelta(seconds=5)) + await stream.reduce(raises, 0, timedelta(seconds=5)) -async def test_client_result_over_the_limit(client: Client): +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) - with pytest.raises(ApplicationError) as err: - await take_stream( - stream, timedelta(seconds=5), idle_timeout=timedelta(0), max_bytes=1 << 22 + # 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, ) - assert err.value.type == "StreamResultTooLarge" - assert "over the limit of 1048576 bytes" in err.value.message - # The default cap of 256 KiB keeps a call under the limit. - read = await take_stream(stream, timedelta(seconds=5), idle_timeout=timedelta(0)) - assert read.stop is StreamReadStop.MaxBytes + # 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( @@ -219,8 +399,7 @@ def append_upper(value: int, record: StreamRecord[Token]) -> int: appender.append(record.body.text.upper()) return value + 1 - result = await reduce_stream( - _client_stream(stream_client, source_id), + result = await _client_stream(stream_client, source_id).reduce( append_upper, 0, timedelta(seconds=5), diff --git a/tests/nexus/test_stream_read_workflow.py b/tests/nexus/test_stream_read_workflow.py index c8fa125202..ff535235da 100644 --- a/tests/nexus/test_stream_read_workflow.py +++ b/tests/nexus/test_stream_read_workflow.py @@ -31,6 +31,8 @@ from temporalio.exceptions import ApplicationError, CancelledError, StreamError from temporalio.nexus.system.stream_service import ( ProducerKindLocalActivity, + StreamRead, + StreamReadStop, StreamRecord, StreamRef, ) @@ -101,7 +103,7 @@ def _collect(value: list[str], record: StreamRecord[Token]) -> list[str]: return [*value, record.body.text] -def _read_result(read: workflow.StreamRead[Token]) -> CallResult: +def _read_result(read: StreamRead[Token]) -> CallResult: return CallResult( texts=[record.body.text for record in read.records], count=len(read.records), @@ -147,8 +149,7 @@ def fold_until(value: list[str], record: StreamRecord[Token]) -> bool: 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 workflow.take_stream( - stream, + read = await stream.take( timeout, until if until_kind else None, from_offset, @@ -158,8 +159,7 @@ def fold_until(value: list[str], record: StreamRecord[Token]) -> bool: ) return _read_result(read) if args.call == "filter": - read = await workflow.filter_stream( - stream, + read = await stream.filter( _not_skipped, timeout, until if until_kind else None, @@ -170,8 +170,7 @@ def fold_until(value: list[str], record: StreamRecord[Token]) -> bool: ) return _read_result(read) no_texts: list[str] = [] - reduce = await workflow.reduce_stream( - stream, + reduce = await stream.reduce( _collect, no_texts, timeout, @@ -218,8 +217,7 @@ def want(self, text: str) -> None: @workflow.run async def run(self, stream_id: str) -> CallResult: - read = await workflow.take_stream( - _stream(stream_id), + read = await _stream(stream_id).take( timedelta(seconds=20), until=lambda record: record.body.text == self.wanted, idle_timeout=timedelta(seconds=20), @@ -232,14 +230,12 @@ class TwoTakesWorkflow: @workflow.run async def run(self, stream_id: str) -> list[str]: stream = _stream(stream_id) - first = await workflow.take_stream( - stream, + first = await stream.take( timedelta(seconds=10), until=lambda record: record.body.text.endswith("m1"), ) await workflow.sleep(timedelta(milliseconds=1)) - second = await workflow.take_stream( - stream, + second = await stream.take( timedelta(seconds=10), until=lambda record: record.body.text.endswith("m2"), from_offset=first.last_offset, @@ -251,8 +247,7 @@ async def run(self, stream_id: str) -> list[str]: class FailTaskOnceAfterTakeWorkflow: @workflow.run async def run(self, stream_id: str) -> CallResult: - read = await workflow.take_stream( - _stream(stream_id), + read = await _stream(stream_id).take( timedelta(seconds=10), until=lambda record: record.body.kind == "done", ) @@ -275,8 +270,7 @@ async def run(self, stream_id: str) -> tuple[list[str], int]: while True: calls += 1 # A filter with `max_records=1` gives the next Record that it keeps. - read = await workflow.filter_stream( - stream, + read = await stream.filter( lambda record: record.body.kind == "hit", timedelta(seconds=10), from_offset=offset, @@ -286,8 +280,8 @@ async def run(self, stream_id: str) -> tuple[list[str], int]: found.extend(record.body.text for record in read.records) offset = read.last_offset if read.stop in ( - workflow.StreamReadStop.Done, - workflow.StreamReadStop.Idle, + StreamReadStop.Done, + StreamReadStop.Idle, ): return found, calls @@ -300,12 +294,10 @@ async def run(self, stream_id: str) -> list[tuple[int, str]]: pages: list[tuple[int, str]] = [] offset: bytes | None = None while True: - read = await workflow.take_stream( - stream, timedelta(seconds=10), from_offset=offset - ) + 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 workflow.StreamReadStop.MaxBytes: + if read.stop is not StreamReadStop.MaxBytes: return pages @@ -316,12 +308,9 @@ async def run(self, stream_id: str, call: str) -> str: stream = _stream(stream_id) try: if call == "take": - await workflow.take_stream( - stream, timedelta(seconds=10), max_bytes=4 * 1024 * 1024 - ) + await stream.take(timedelta(seconds=10), max_bytes=4 * 1024 * 1024) else: - await workflow.reduce_stream( - stream, + await stream.reduce( lambda value, record: value + "x" * (1024 * 1024), "", timedelta(seconds=10), @@ -338,8 +327,7 @@ class TwoStreamsAtOnceWorkflow: async def run(self, first_id: str, second_id: str) -> list[str]: tasks = [ asyncio.create_task( - workflow.filter_stream( - _stream(stream_id), + _stream(stream_id).filter( lambda record: record.body.kind == "hit", timedelta(seconds=10), max_records=1, @@ -362,8 +350,7 @@ async def run(self, stream_id: str) -> list[str]: workflow.uuid4() if variant == "other_stream": stream_id += "-other" - read = await workflow.take_stream( - _stream(stream_id), + read = await _stream(stream_id).take( timedelta(seconds=10), from_offset=b"\x01" if variant == "other_offset" else None, idle_timeout=timedelta(0), @@ -389,7 +376,7 @@ async def run(self, stream_id: str, which: str) -> str: # 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 workflow.take_stream(stream, timedelta(seconds=10), until=until) + await stream.take(timedelta(seconds=10), until=until) except TypeError as err: return f"TypeError: {err}" except ApplicationError as err: @@ -423,8 +410,8 @@ class FailingPredicateWorkflow: @workflow.run async def run(self, stream_id: str, predicate: str) -> str: try: - await workflow.filter_stream( - _stream(stream_id), _PREDICATES[predicate], timedelta(seconds=10) + await _stream(stream_id).filter( + _PREDICATES[predicate], timedelta(seconds=10) ) except ApplicationError as err: return f"{err.type}: {err.message}" @@ -443,6 +430,10 @@ def _summarize(value: Summary | None, record: StreamRecord[Token]) -> Summary: 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 @@ -451,36 +442,28 @@ async def run(self, stream_id: str) -> list[str]: quick = timedelta(seconds=10) now = timedelta(0) # The return annotation wins over an `initial` of `None`. - annotated = await workflow.reduce_stream( - stream, _summarize, None, quick, idle_timeout=now - ) + annotated = await stream.reduce(_summarize, None, quick, idle_timeout=now) # With no annotation, the type of `initial` decides. - of_initial = await workflow.reduce_stream( - stream, + 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 workflow.reduce_stream( - stream, + untyped = await stream.reduce( lambda value, record: Summary(1, [record.body.kind]), cast(Any, None), quick, idle_timeout=now, ) - explicit = await workflow.reduce_stream( - stream, - lambda value, record: [*value, record.body.text], - cast(Any, ()), - quick, - idle_timeout=now, - value_type=list[str], + # 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, explicit) + for result in (annotated, of_initial, untyped, over_initial) ] @@ -488,8 +471,7 @@ async def run(self, stream_id: str) -> list[str]: class FoldUntilValueWorkflow: @workflow.run async def run(self, stream_id: str) -> tuple[str, int, str]: - result = await workflow.reduce_stream( - _stream(stream_id), + result = await _stream(stream_id).reduce( lambda value, record: value + record.body.text, "", timedelta(seconds=10), @@ -510,8 +492,7 @@ def append_upper(value: int, record: StreamRecord[Token]) -> int: return value + 1 try: - result = await workflow.reduce_stream( - _stream(source_id), + result = await _stream(source_id).reduce( append_upper, 0, timedelta(seconds=10), @@ -1309,8 +1290,7 @@ async def test_measure_take_of_forty(client: Client): class KeepHitsWorkflow: @workflow.run async def run(self, stream_id: str) -> int: - read = await workflow.filter_stream( - _stream(stream_id), + read = await _stream(stream_id).filter( lambda record: record.body.kind == "hit", timedelta(seconds=10), idle_timeout=timedelta(0), diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py index 9fda45007c..4255f890ad 100644 --- a/tests/nexus/test_stream_store_workflow.py +++ b/tests/nexus/test_stream_store_workflow.py @@ -642,9 +642,9 @@ async def test_workflow_read_fails_with_store(client: Client): task_queue=worker.task_queue, ) assert result == ( - "Workflow code can't run a ReadRecords operation on a Store. Use " - "temporalio.workflow.take_stream, filter_stream or reduce_stream, which " - "read in one Local Activity." + "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." ) From 116f765e64d245dd129e9dbe8ce93b49d3e29de8 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 03:02:26 -0700 Subject: [PATCH 17/20] Added a pipeline that runs a chain of Stream steps as one call. The take, filter and reduce calls don't chain. The pipeline compiles a chain into one of them, so in Workflow code it stays one Local Activity and one marker. --- changelog/added/chained-pipeline-wren.md | 1 + temporalio/contrib/streams/README.md | 44 +- temporalio/contrib/streams/__init__.py | 8 +- temporalio/contrib/streams/_pipeline.py | 513 +++++++++++++++++++++++ tests/contrib/streams/test_pipeline.py | 387 +++++++++++++++++ 5 files changed, 949 insertions(+), 4 deletions(-) create mode 100644 changelog/added/chained-pipeline-wren.md create mode 100644 temporalio/contrib/streams/_pipeline.py create mode 100644 tests/contrib/streams/test_pipeline.py diff --git a/changelog/added/chained-pipeline-wren.md b/changelog/added/chained-pipeline-wren.md new file mode 100644 index 0000000000..5a2e3d47f1 --- /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/temporalio/contrib/streams/README.md b/temporalio/contrib/streams/README.md index f64fa2d009..77cb5abd29 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,42 @@ 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. diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index d1ee83436d..d9c4d119cd 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/_pipeline.py b/temporalio/contrib/streams/_pipeline.py new file mode 100644 index 0000000000..5d8be87081 --- /dev/null +++ b/temporalio/contrib/streams/_pipeline.py @@ -0,0 +1,513 @@ +"""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. 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 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. + + 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. 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 goes into History, 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. + + 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. + + 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/tests/contrib/streams/test_pipeline.py b/tests/contrib/streams/test_pipeline.py new file mode 100644 index 0000000000..59c8ee4644 --- /dev/null +++ b/tests/contrib/streams/test_pipeline.py @@ -0,0 +1,387 @@ +"""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, 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.helpers import new_worker +from tests.helpers.streams import RecordingStreamStore +from tests.nexus.test_stream_read_workflow import ( + Token, + _append, + _close, + _markers, + _offsets, + _with_store, +) + +_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]), + ] + + +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 From e68591a0c7c26e71140b19184a5330156c4325da Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 03:50:18 -0700 Subject: [PATCH 18/20] Kept the Records of take and filter out of History when a Worker asks. A large take put every body into History. With `StreamReplaySource.STORE` the marker holds only offsets and a digest, and the Worker reads the Records from the Store before a replay. A replay follows the marker, and a failed read fails the Workflow Task with `StreamReplayError`. --- changelog/added/patient-reader-heron.md | 1 + temporalio/client/__init__.py | 2 + temporalio/client/_stream_read.py | 343 +++++++++++++++++++++--- temporalio/client/_stream_store.py | 49 ++++ temporalio/common.py | 8 +- temporalio/contrib/streams/README.md | 8 + temporalio/contrib/streams/_pipeline.py | 26 +- temporalio/exceptions.py | 25 ++ temporalio/nexus/system/_support.py | 77 ++++-- temporalio/worker/_replayer.py | 12 +- temporalio/worker/_stream_read.py | 96 ++++++- temporalio/worker/_workflow.py | 8 + 12 files changed, 572 insertions(+), 83 deletions(-) diff --git a/changelog/added/patient-reader-heron.md b/changelog/added/patient-reader-heron.md index ab1bf2f93b..8d95c851cf 100644 --- a/changelog/added/patient-reader-heron.md +++ b/changelog/added/patient-reader-heron.md @@ -1 +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/temporalio/client/__init__.py b/temporalio/client/__init__.py index 6fa6dd710c..8f567ea21a 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/_stream_read.py b/temporalio/client/_stream_read.py index aa57c64b0d..4ab84910ed 100644 --- a/temporalio/client/_stream_read.py +++ b/temporalio/client/_stream_read.py @@ -2,8 +2,13 @@ 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 result of the call. 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. +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. @@ -14,12 +19,14 @@ 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 Callable, Sequence +from collections.abc import Awaitable, Callable, Sequence from datetime import timedelta from typing import Any, Generic, Protocol, TypeVar @@ -42,10 +49,15 @@ Stream as _WorkflowStream, ) -from ._stream_store import StreamStore, _store_page_reader +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. @@ -235,6 +247,8 @@ def _call_function( 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 @@ -254,6 +268,7 @@ def __init__( 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 @@ -266,7 +281,9 @@ def _call_function( label, function, *args, in_local_activity=self._in_local_activity ) - def bytes_left(self) -> int: + 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: @@ -280,8 +297,14 @@ def look(self, record: _RecordT, *, ends_stream: bool) -> StreamReadStop | 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 self.bytes + size > call.max_bytes: + 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 @@ -292,6 +315,7 @@ def look(self, record: _RecordT, *, ends_stream: bool) -> StreamReadStop | None: ) 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): @@ -307,6 +331,7 @@ def end( ) -> _ScanOutcome[_RecordT]: return _ScanOutcome( kept=self.kept, + kept_after=self.kept_after, count=self.count, last_offset=self.last_offset, stop=stop, @@ -332,7 +357,8 @@ def idle(now: float) -> bool: return now - last_arrival >= idle_timeout while True: - if page_cap and state.count > 0 and state.bytes_left() <= 0: + 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() @@ -349,7 +375,9 @@ def idle(now: float) -> bool: max_records=min(call.max_records - state.count, _MAX_PAGE_RECORDS) if page_cap else _MAX_PAGE_RECORDS, - max_bytes=max(state.bytes_left(), 1) if page_cap else 0, + max_bytes=max(bytes_left, 1) + if page_cap and bytes_left is not None + else 0, wait=timedelta(seconds=wait), ) ) @@ -399,6 +427,52 @@ async def _flush_all(flushes: list[_CallFlush], *, raise_error: bool) -> None: 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.""" @@ -415,10 +489,11 @@ class _StreamCallResult: """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 to_payload(self) -> temporalio.api.common.v1.Payload: - """Give the result as one payload. The bodies stay in their binary form.""" - header = { + def _header(self) -> dict[str, Any]: + return { "kind": self.kind.value, "store": self.store, "stream_id": self.stream_id, @@ -428,6 +503,26 @@ def to_payload(self) -> temporalio.api.common.v1.Payload: "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"}, @@ -436,7 +531,7 @@ def to_payload(self) -> temporalio.api.common.v1.Payload: ] if self.value is not None: parts.append(self.value) - else: + elif replay is None: parts.append( temporalio.api.common.v1.Payload( metadata={ @@ -449,10 +544,7 @@ def to_payload(self) -> temporalio.api.common.v1.Payload: ) ) return temporalio.api.common.v1.Payload( - metadata={ - "encoding": b"binary/plain", - _RESULT_METADATA_KEY: _RESULT_METADATA_VALUE, - }, + metadata=metadata, data=temporalio.api.common.v1.Payloads(payloads=parts).SerializeToString(), ) @@ -463,19 +555,28 @@ def from_payload( """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 - header_part, body = temporalio.api.common.v1.Payloads.FromString( - payload.data - ).payloads - header = json.loads(header_part.data) + 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 = body + 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( - body.data + parts[1].data ).records ) return cls( @@ -489,19 +590,57 @@ def from_payload( 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, -) -> temporalio.api.common.v1.Payload: - """Give the result for History as a payload, or raise when it is over the limit.""" - payload = _StreamCallResult( +) -> 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, @@ -510,11 +649,16 @@ def _result_payload( count=outcome.count, stop=outcome.stop, done=outcome.done, - records=outcome.kept, + records=outcome.kept if replay is None else [], value=payload_converter.to_payload(outcome.value) if call.kind is _StreamCallKind.REDUCE else None, - ).to_payload() + 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: @@ -525,7 +669,7 @@ def _result_payload( type="StreamResultTooLarge", non_retryable=True, ) - return payload + return result, payload @dataclasses.dataclass(eq=False) @@ -546,12 +690,37 @@ class _PendingStreamCall: 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 looks an - entry up on the event loop of the Worker. + 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: @@ -614,15 +783,113 @@ async def _stream_call_activity( # pyright: ignore[reportUnusedFunction] call=pending.call, in_local_activity=True, ) - return temporalio.common.RawValue( - _result_payload( - call=pending.call, - store=pending.store.name, - stream_id=pending.ref.id, - outcome=outcome, - payload_converter=temporalio.activity.payload_converter(), - ) + 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( diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index 9e01f460bb..853104f362 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. @@ -499,6 +534,20 @@ class StreamConfig: Activity for a close. """ + workflow_read_replay_from: StreamReplaySource = 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 24996df680..8aaea19dbb 100644 --- a/temporalio/common.py +++ b/temporalio/common.py @@ -1476,7 +1476,9 @@ class HeaderCodecBehavior(IntEnum): 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``. +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. @@ -1488,7 +1490,9 @@ class HeaderCodecBehavior(IntEnum): 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. +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. diff --git a/temporalio/contrib/streams/README.md b/temporalio/contrib/streams/README.md index 77cb5abd29..2dbd2378de 100644 --- a/temporalio/contrib/streams/README.md +++ b/temporalio/contrib/streams/README.md @@ -61,3 +61,11 @@ scan stopped, so a loop that gives it as the next `from_offset` loses and repeat `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/_pipeline.py b/temporalio/contrib/streams/_pipeline.py index 5d8be87081..f4a6069a91 100644 --- a/temporalio/contrib/streams/_pipeline.py +++ b/temporalio/contrib/streams/_pipeline.py @@ -104,8 +104,11 @@ class Pipeline(Generic[ItemT]): 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. Client, Activity and Workflow code get the same - result. + 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. @@ -173,8 +176,12 @@ def records( 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 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. + 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. @@ -258,7 +265,8 @@ async def reduce( 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. Workflow code decodes the value as the return + 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: @@ -299,9 +307,10 @@ async def collect( :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 goes into History, so its encoded size must fit in 1 MiB. + 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. + go into History as :py:meth:`records` puts them there, so they follow the setting. Args: timeout: As in the call of the Stream. @@ -365,7 +374,8 @@ async def count( 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. + 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. diff --git a/temporalio/exceptions.py b/temporalio/exceptions.py index 658012ee4f..52cb089564 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 be5e5016bb..f4a791007c 100644 --- a/temporalio/nexus/system/_support.py +++ b/temporalio/nexus/system/_support.py @@ -479,7 +479,10 @@ class _StreamCall: timeout: timedelta idle_timeout: timedelta max_records: int - max_bytes: 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 @@ -511,6 +514,7 @@ def check( 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: @@ -528,6 +532,7 @@ def check( idle_timeout=idle_timeout, max_records=max_records, max_bytes=max_bytes, + default_max_bytes=default_max_bytes, until=until, predicate=predicate, function=function, @@ -619,12 +624,18 @@ async def take_stream( :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 Records and their bodies, so the encoded result - must fit in :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES`. 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 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 + 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 @@ -652,7 +663,9 @@ async def take_stream( 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. + :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. @@ -665,10 +678,10 @@ async def take_stream( 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 + 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. @@ -728,12 +741,19 @@ async def filter_stream( :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 kept Records and their bodies, so the encoded - result must fit in :py:data:`temporalio.common.MAX_STREAM_RESULT_BYTES`. 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 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 + 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 @@ -762,7 +782,9 @@ async def filter_stream( 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. + :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. @@ -775,10 +797,10 @@ async def filter_stream( 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 + 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. @@ -843,9 +865,12 @@ async def reduce_stream( :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 - 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 + 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 diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 6a0085833d..6361b3a0f0 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -70,9 +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. The ``take``, ``filter`` or ``reduce`` - of a Workflow Stream takes its result from the history, so its replay needs no - Stream config. + 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 index 410ccdc17c..c8ce965940 100644 --- a/temporalio/worker/_stream_read.py +++ b/temporalio/worker/_stream_read.py @@ -1,21 +1,36 @@ -"""The Workflow side of the take, filter and reduce calls of a Stream.""" +"""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 +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, @@ -68,8 +83,8 @@ async def _stream_call_in_workflow( # pyright: ignore[reportUnusedFunction] def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: return payload_converter._from_transfer_value(record, record_type) - # A replay takes the result from History, so it needs no Store and no config. The - # choice of a Store runs in Workflow code, as it does for the other Stream calls. + # 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 @@ -79,8 +94,22 @@ def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: 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 + 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: @@ -104,6 +133,7 @@ def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: 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" @@ -128,9 +158,63 @@ def to_record(record: temporalio.api.stream.v1.StreamRecord) -> Any: 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 result.records], + 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/_workflow.py b/temporalio/worker/_workflow.py index 40294c256c..3c26d2d40a 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -41,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, @@ -433,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( From 66019c8dc93eb5b3f17380d6191cc94666c0c863 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 03:50:18 -0700 Subject: [PATCH 19/20] Tested where a replay of a take or a filter gets its Records. The tests cover History without bodies, the reads of a replay, its failures, and a change of the setting while a Workflow runs. --- tests/contrib/streams/test_pipeline.py | 66 ++- tests/nexus/test_stream_read_replay_from.py | 442 ++++++++++++++++++++ 2 files changed, 505 insertions(+), 3 deletions(-) create mode 100644 tests/nexus/test_stream_read_replay_from.py diff --git a/tests/contrib/streams/test_pipeline.py b/tests/contrib/streams/test_pipeline.py index 59c8ee4644..00604e1d04 100644 --- a/tests/contrib/streams/test_pipeline.py +++ b/tests/contrib/streams/test_pipeline.py @@ -12,7 +12,12 @@ import temporalio.api.streamservice.v1 from temporalio import workflow -from temporalio.client import Client, StreamConfig, StreamStoreCallOptions +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 ( @@ -24,17 +29,22 @@ ) from temporalio.testing import MemoryStreamStore from temporalio.worker import Replayer -from tests.helpers import new_worker -from tests.helpers.streams import RecordingStreamStore 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) @@ -237,6 +247,24 @@ async def run(self, stream_id: str) -> list[ChainResult]: ] +@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.""" @@ -385,3 +413,35 @@ async def test_typed_ends(client: Client): 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/nexus/test_stream_read_replay_from.py b/tests/nexus/test_stream_read_replay_from.py new file mode 100644 index 0000000000..a847367a92 --- /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 From 2de2f0a6b3aa3f7a7c771f2258ef015b9a51bfc8 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 04:47:49 -0700 Subject: [PATCH 20/20] Fixed the API docs build for the replay setting. The docs tool read the bare enum default of the field as an alias, and failed on every link to the field. A field default is not an alias. --- temporalio/client/_stream_store.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index 853104f362..b13751ede4 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -534,7 +534,9 @@ class StreamConfig: Activity for a close. """ - workflow_read_replay_from: StreamReplaySource = StreamReplaySource.HISTORY + 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