diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index 1bcd2f008..ba622ecb5 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -65,6 +65,7 @@ from temporalio.common._streams._plugin import StreamBackend from temporalio.common._streams._reader import Subscription from temporalio.common._streams._stream_client import ClientStream + from temporalio.common._streams._workflow import WorkflowStream __all__ = [ "ActivityIDConflictPolicy", @@ -130,6 +131,7 @@ "WorkerDeploymentVersion", "WorkflowIDConflictPolicy", "WorkflowIDReusePolicy", + "WorkflowStream", "defn", ] @@ -1614,6 +1616,7 @@ def wait_sync(self, timeout: float | None = None) -> None: "StreamUnsupportedError", "Subscription", "Superseded", + "WorkflowStream", "defn", ] ) diff --git a/temporalio/common/_streams/__init__.py b/temporalio/common/_streams/__init__.py index 8e4785be2..4c828426f 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -27,6 +27,7 @@ from temporalio.common._streams._plugin import StreamBackend from temporalio.common._streams._reader import Subscription from temporalio.common._streams._stream_client import ClientStream +from temporalio.common._streams._workflow import WorkflowStream __all__ = [ "END", @@ -54,5 +55,6 @@ "StreamUnsupportedError", "Subscription", "Superseded", + "WorkflowStream", "defn", ] diff --git a/temporalio/common/_streams/_activity.py b/temporalio/common/_streams/_activity.py new file mode 100644 index 000000000..d6e87b288 --- /dev/null +++ b/temporalio/common/_streams/_activity.py @@ -0,0 +1,71 @@ +"""The Stream state of one Activity attempt.""" + +from __future__ import annotations + +import asyncio +import contextvars +from dataclasses import dataclass, field +from typing import Any + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.converter +from temporalio.common._streams._backend import BackendCalls +from temporalio.common._streams._writer import StreamWriter + + +@dataclass +class ActivityStreams: + """What a Stream backend gives the Activity handles for one Activity attempt.""" + + backend: BackendCalls + converter: temporalio.converter.DataConverter + namespace: str + producer: stream_v1.Producer + loop: asyncio.AbstractEventLoop + backend_name: str = "" + # Whether the Worker's backend is the default, which a handle with no backend name uses. + default: bool = False + # The attempt's Workflow id, which a close sends so the Store can check the Stream's owner. + workflow_id: str | None = None + # A Local Activity's writes use these to withdraw the run's orphaned Records. + previous_completed_task: stream_v1.WorkflowTaskEpoch | None = None + failed_tasks: tuple[int, ...] = () + writers: dict[str, StreamWriter] = field(default_factory=dict) + # The attempt's BufferedPublishers, which the end of the attempt flushes. + publishers: list[Any] = field(default_factory=list) + + def writer(self, stream_id: str) -> StreamWriter: + """The attempt's writer for a Stream, one per Stream.""" + writer = self.writers.get(stream_id) + if writer is None: + writer = StreamWriter( + writers=self.backend.writers(self.namespace, self.converter), + stream_id=stream_id, + producer=self.producer, + free=self.backend.free_writer, + previous_completed_task=self.previous_completed_task, + failed_tasks=self.failed_tasks, + ) + self.writers[stream_id] = writer + return writer + + async def close(self) -> None: + """Flush the attempt's publishers, then close every writer of the attempt.""" + errors = [] + for publisher in self.publishers: + try: + await publisher._flush_at_attempt_end() + except Exception as err: + errors.append(err) + for writer in self.writers.values(): + try: + await writer.close() + except Exception as err: + errors.append(err) + if errors: + raise errors[0] + + +current: contextvars.ContextVar[ActivityStreams | None] = contextvars.ContextVar( + "__temporal_activity_streams", default=None +) diff --git a/temporalio/common/_streams/_extension.py b/temporalio/common/_streams/_extension.py index dbb68c82d..d7b71dfd6 100644 --- a/temporalio/common/_streams/_extension.py +++ b/temporalio/common/_streams/_extension.py @@ -2,24 +2,50 @@ from __future__ import annotations +import asyncio +import dataclasses import logging import weakref -from collections.abc import Sequence -from typing import Any +from collections.abc import Mapping, Sequence +from datetime import timedelta +from typing import Any, NoReturn +import temporalio.activity +import temporalio.api.common.v1 +import temporalio.api.stream.v1 as stream_v1 +import temporalio.api.streamservice.v1 as stream_pb import temporalio.bridge.streams import temporalio.client +import temporalio.common import temporalio.converter import temporalio.runtime import temporalio.service import temporalio.worker +import temporalio.workflow +from temporalio.api.enums.v1 import StreamProducerKind +from temporalio.common._streams import _activity, _ids, _writes +from temporalio.common._streams import _workflow as streams_workflow from temporalio.common._streams._backend import BackendCalls, _ExtensionCalls from temporalio.common._streams._errors import StreamConfigError +from temporalio.common._streams._handle import StreamHandle, workflow_owner from temporalio.common._streams._plugin import StreamBackend from temporalio.common._streams.proto import v1 as ext_pb +from temporalio.common._streams.system import models +from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner logger = logging.getLogger(__name__) +# The backend puts this key into the headers of each Local Activity that a Workflow schedules, +# and of each Activity of a run that uses Streams. It keeps the Producer id the same on retries. +_KEY_HEADER = "__temporal_stream_key" +# Continue-as-new passes the run chain's open Streams to the next run in this header. +_OWNED_STREAMS_HEADER = "__temporal_owned_streams" +_CLOSE_TIMEOUT = timedelta(seconds=30) +# Only a Store error is retryable. Each Stream error sets its own non_retryable flag. +_CLOSE_RETRY = temporalio.common.RetryPolicy( + initial_interval=timedelta(seconds=1), maximum_interval=timedelta(seconds=30) +) + class ExtensionStreamBackend(StreamBackend): """Base class of the Redis and memory backends, which call the temporalio-streams runtime. @@ -58,6 +84,7 @@ def __init__( name, backend_name=backend_name, default=default, + activities=_writes.ACTIVITIES, ) def _new_backend( @@ -107,6 +134,82 @@ def _interceptor_of(self, interceptors: Sequence[Any]) -> _StreamInterceptor | N return interceptor return None + def _configure_worker( + self, config: temporalio.worker.WorkerConfig + ) -> temporalio.worker.WorkerConfig: + # The Worker gets the Stream interceptor from its client. + config = super()._configure_worker(config) + client = config["client"] # type: ignore[reportTypedDictNotRequiredAccess] + interceptors = [ + *client.config(active_config=True).get("interceptors", []), + *config.get("interceptors", []), + ] + if any( + isinstance(other, _StreamInterceptor) and other.plugin is not self + for other in interceptors + ): + raise StreamConfigError( + "The Worker's client has more than one Stream backend. A Worker serves only " + "one backend, so give the Worker a client with one backend." + ) + mine = self._interceptor_of(interceptors) + if mine is None: + raise StreamConfigError( + f"The Worker's client has no {type(self).__name__}. Register it with " + "stream_backends=." + ) + if config.get("no_remote_activities"): + # Another Worker on the task queue can run the close Activity, so it isn't an error. + logger.warning( + "The Worker on the task queue %s polls no Activities, so the close of a " + "Workflow's Stream waits for another Worker on that queue with the backend", + config.get("task_queue"), + ) + # Records get no serialization context, so Workflows, Activities and clients encode + # them alike. + mine.converter = client.data_converter + runner = config.get("workflow_runner") + if isinstance(runner, SandboxedWorkflowRunner): + # The interceptor and temporalio.workflow must share the run's Stream state, even + # when the restrictions don't pass temporalio through. + config["workflow_runner"] = dataclasses.replace( + runner, + restrictions=runner.restrictions.with_passthrough_modules( + "temporalio.common._streams" + ), + ) + return config + + def _configure_replayer( + self, config: temporalio.worker.ReplayerConfig + ) -> temporalio.worker.ReplayerConfig: + # A replay runs no Activity, so it needs the Workflow interceptor and no runtime. + config = super()._configure_replayer(config) + if any( + isinstance(other, _StreamInterceptor) and other.plugin is not self + for other in config.get("interceptors", []) + ): + raise StreamConfigError( + "The Replayer has another Stream backend. A Replayer serves only one " + "backend, as a Worker does." + ) + config["interceptors"] = [ + *config.get("interceptors", []), + _StreamInterceptor(self, None, None), + ] + return config + + async def _close_in_workflow(self, request: stream_pb.CloseStreamRequest) -> int: + # An Activity runs only after the Workflow Task commits, so a failed task leaves the + # Stream open for its retry. The Activity closes the Stream as the run chain that owns it. + return await temporalio.workflow.execute_activity( + streams_workflow.CLOSE_ACTIVITY, + request, + result_type=int, + start_to_close_timeout=_CLOSE_TIMEOUT, + retry_policy=_CLOSE_RETRY, + ) + async def shutdown(self) -> None: """Close every open writer, including the client's. @@ -143,3 +246,306 @@ def __init__( self.plugin = plugin self.backend = backend self.converter = converter or temporalio.converter.DataConverter.default + + def workflow_interceptor_class( + self, input: temporalio.worker.WorkflowInterceptorClassInput + ) -> type[temporalio.worker.WorkflowInboundInterceptor] | None: + # Workflow code reads the Worker's backend from the interceptor class. + return type( + "_WorkflowInbound", + (_WorkflowInbound,), + { + "_backend": self.plugin, + "_backend_name": self.plugin.backend_name, + "_default": self.plugin.default, + }, + ) + + def intercept_activity( + self, next: temporalio.worker.ActivityInboundInterceptor + ) -> temporalio.worker.ActivityInboundInterceptor: + return _ActivityInbound(next, self) + + +class _WorkflowInbound(temporalio.worker.WorkflowInboundInterceptor): + _backend: ExtensionStreamBackend | None = None + _backend_name = "" + _default = False + + def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: + assert self._backend is not None + super().init( + _WorkflowOutbound( + outbound, self._backend, self._backend_name, self._default + ) + ) + + async def execute_workflow( + self, input: temporalio.worker.ExecuteWorkflowInput + ) -> Any: + _take_owned_streams(input.headers) + _receive(input.args) + try: + result = await super().execute_workflow(input) + except BaseException as err: + if _ends_run(err): + streams_workflow.close_at_run_end(err) + await streams_workflow.drain() + raise + streams_workflow.close_at_run_end(None) + await streams_workflow.drain() + return result + + async def handle_signal(self, input: temporalio.worker.HandleSignalInput) -> None: + _receive(input.args) + await super().handle_signal(input) + + async def handle_update_handler( + self, input: temporalio.worker.HandleUpdateInput + ) -> Any: + _receive(input.args) + return await super().handle_update_handler(input) + + +def _receive(args: Sequence[Any]) -> None: + # A run that receives a Stream uses Streams. The inputs, Signals and Updates are in + # History, so a replay sets the same flag. + if _stream_arguments(args): + streams_workflow._state().uses_streams = True + + +def _ends_run(err: BaseException) -> bool: + # True if err ends the run, which must send its unawaited writes first. Any other exception + # fails the Workflow Task, and an eviction closes the coroutine outside the Workflow's loop. + if isinstance(err, temporalio.workflow.ContinueAsNewError): + return True + if isinstance(err, asyncio.CancelledError): + return temporalio.workflow.cancellation_reason() is not None + return isinstance( + err, Exception + ) and temporalio.workflow._Runtime.current().workflow_is_failure_exception(err) + + +def _owned_ref(stream_id: str) -> models.StreamRef: + # The run chain owns each Stream that it creates. + info = temporalio.workflow.info() + return models.StreamRef( + id=stream_id, + owner=workflow_owner(info.workflow_id, info.first_execution_run_id), + ) + + +def _take_owned_streams( + headers: Mapping[str, temporalio.api.common.v1.Payload], +) -> None: + payload = headers.get(_OWNED_STREAMS_HEADER) + if payload is None: + return + state = streams_workflow._state() + state.uses_streams = True + converter = temporalio.converter.PayloadConverter.default + for stream_id, backend, *_ in converter.from_payload(payload, list[Any]): + state.stream(_owned_ref(stream_id)).backend = backend + + +class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): + _temporal_stream_transport = True + + def __init__( + self, + next: temporalio.worker.WorkflowOutboundInterceptor, + backend: ExtensionStreamBackend, + backend_name: str, + default: bool, + ) -> None: + super().__init__(next) + self._temporal_stream_backend = backend + self._temporal_stream_backend_name = backend_name + self._temporal_stream_default = default + # Counts the run's Activity calls, so a key is the same on replay and on a retry. It + # doesn't call workflow.uuid4(), so the user's random values don't change. + self._activity_calls = 0 + + def _activity_headers(self) -> dict[str, temporalio.api.common.v1.Payload]: + self._activity_calls += 1 + # Both ends use the default converter, so a custom user converter can't break the + # header. + converter = temporalio.converter.PayloadConverter.default + return {_KEY_HEADER: converter.to_payload(str(self._activity_calls))} + + def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: + # The next run needs each owned Stream's backend for stream_handle. + owned = [ + [stream_id, run_stream.backend] + for stream_id, run_stream in streams_workflow._state().streams.items() + if run_stream.backend is not None and not run_stream.closed + ] + if owned: + converter = temporalio.converter.PayloadConverter.default + input.headers = { + **input.headers, + _OWNED_STREAMS_HEADER: converter.to_payload(owned), + } + super().continue_as_new(input) + + def start_activity( + self, input: temporalio.worker.StartActivityInput + ) -> temporalio.workflow.ActivityHandle[Any]: + # The key is in History, so add it only for a run that uses Streams or passes one to + # this Activity. Count every call, so the keys are the same on replay. + headers = self._activity_headers() + state = streams_workflow._state() + if not state.uses_streams and _stream_arguments(input.args): + state.uses_streams = True + if state.uses_streams: + input.headers = {**input.headers, **headers} + return super().start_activity(input) + + def start_local_activity( + self, input: temporalio.worker.StartLocalActivityInput + ) -> temporalio.workflow.ActivityHandle[Any]: + # A Local Activity's headers aren't in History, so every Local Activity gets the key. + input.headers = {**input.headers, **self._activity_headers()} + return super().start_local_activity(input) + + +class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor): + def __init__( + self, + next: temporalio.worker.ActivityInboundInterceptor, + interceptor: _StreamInterceptor, + ) -> None: + super().__init__(next) + self._interceptor = interceptor + + async def execute_activity( + self, input: temporalio.worker.ExecuteActivityInput + ) -> Any: + streams = _activity_streams(input, self._interceptor) + token = _activity.current.set(streams) + try: + if temporalio.activity.info().attempt > 1: + # Withdraw the Records of earlier attempts, even if this attempt publishes + # nothing. This creates no Stream. + stream_ids = [handle.id for handle in _stream_arguments(input.args)] + for stream_id in dict.fromkeys(stream_ids): + await streams.writer(stream_id).withdraw_earlier_attempts() + result = await super().execute_activity(input) + except BaseException: + await _close_quietly(streams) + raise + else: + await streams.close() + return result + finally: + _activity.current.reset(token) + + +def _stream_arguments(args: Sequence[Any]) -> list[StreamHandle[Any]]: + # Find the Stream handles in call arguments. A handle can be in a field, a slot, a list or + # a dict. + handles: list[StreamHandle[Any]] = [] + seen: set[int] = set() + + def walk(value: Any, depth: int) -> None: + if depth > 4 or id(value) in seen: + return + seen.add(id(value)) + if isinstance(value, StreamHandle): + handles.append(value) + elif isinstance(value, (str, bytes, int, float, bool)) or value is None: + return + elif isinstance(value, dict): + for item in value.values(): + walk(item, depth + 1) + elif isinstance(value, (list, tuple, set, frozenset)): + for item in value: + walk(item, depth + 1) + else: + for item in _fields_of(value): + walk(item, depth + 1) + + for arg in args: + walk(arg, 0) + return handles + + +def _fields_of(value: Any) -> list[Any]: + fields = list(getattr(value, "__dict__", {}).values()) + for klass in type(value).__mro__: + slots = getattr(klass, "__slots__", ()) + for slot in (slots,) if isinstance(slots, str) else slots: + if hasattr(value, slot): + fields.append(getattr(value, slot)) + return fields + + +async def _close_quietly(streams: _activity.ActivityStreams) -> None: + try: + await asyncio.shield(streams.close()) + except Exception: + temporalio.activity.logger.warning( + "Failed flushing Stream Records of a failed Activity", exc_info=True + ) + + +def _activity_streams( + input: temporalio.worker.ExecuteActivityInput, interceptor: _StreamInterceptor +) -> _activity.ActivityStreams: + # A sync Activity has no client, so the interceptor gives the backend and the converter. + assert interceptor.backend is not None + info = temporalio.activity.info() + converter = temporalio.converter.PayloadConverter.default + key = input.headers.get(_KEY_HEADER) + key_text = converter.from_payload(key, str) if key is not None else "" + epoch = stream_v1.ProducerEpoch(attempt=info.attempt) + previous: stream_v1.WorkflowTaskEpoch | None = None + failed: tuple[int, ...] = () + if info.is_local: + task = info.workflow_task + if task is not None: + # Always set it, with 0 if no task completed. If it's unset, the Store doesn't + # withdraw the run's orphaned Local Activity Records. + previous = stream_v1.WorkflowTaskEpoch( + started_event_id=task.previous_started_event_id + ) + failed = tuple(task.failed_tasks) + epoch.task_started_event_id = task.started_event_id + epoch.task_attempt = task.attempt + if task.started_time is not None: + epoch.task_started_time.FromDatetime(task.started_time) + run_id = info.workflow_run_id or "" + producer = stream_v1.Producer( + kind=StreamProducerKind.STREAM_PRODUCER_KIND_LOCAL_ACTIVITY, + id=_ids.local_activity_producer_id(run_id, key_text), + run_id=run_id, + ) + else: + epoch.attempt_scheduled_time.FromDatetime(info.current_attempt_scheduled_time) + if info.in_workflow: + run_id = info.workflow_run_id or "" + producer = stream_v1.Producer( + kind=StreamProducerKind.STREAM_PRODUCER_KIND_ACTIVITY, + id=_ids.activity_producer_id(run_id, info.activity_id, key_text), + run_id=run_id, + ) + else: + producer = stream_v1.Producer( + kind=StreamProducerKind.STREAM_PRODUCER_KIND_STANDALONE_ACTIVITY, + id=_ids.standalone_activity_producer_id( + info.activity_id, info.activity_run_id or "" + ), + ) + producer.epoch.CopyFrom(epoch) + return _activity.ActivityStreams( + backend=interceptor.backend, + converter=interceptor.converter, + namespace=info.namespace, + producer=producer, + loop=asyncio.get_running_loop(), + backend_name=interceptor.plugin.backend_name, + default=interceptor.plugin.default, + workflow_id=info.workflow_id, + previous_completed_task=previous, + failed_tasks=failed, + ) diff --git a/temporalio/common/_streams/_workflow.py b/temporalio/common/_streams/_workflow.py new file mode 100644 index 000000000..08634e1fa --- /dev/null +++ b/temporalio/common/_streams/_workflow.py @@ -0,0 +1,755 @@ +"""Streams in Workflow code. + +Workflow code creates a Stream with :py:func:`temporalio.workflow.create_stream` and gets it +again with :py:func:`temporalio.workflow.stream_handle`. A Stream that the Workflow receives as +an argument works the same, such as a :py:class:`WorkflowStream` that a run passes on through +continue-as-new. + +Each create and publish is a Local Activity. Each close is a command that the backend runs after +the Workflow Task commits, and on Redis it's a regular Activity. So each one has a place in the +run's commands. To add, remove or move one in a Workflow that has open runs, gate the change +with ``workflow.patched``, or the replay of those runs fails as nondeterministic. Publishes that +wait behind a write to the same Stream are sent as one Local Activity. If a task fails or the +Worker crashes before the close commits, the Stream stays open for the task's retry. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Generator +from dataclasses import dataclass, field +from datetime import timedelta +from types import TracebackType +from typing import Any, NamedTuple, TypeVar + +import google.protobuf.message + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.api.streamservice.v1 as stream_pb +import temporalio.converter +import temporalio.exceptions +import temporalio.workflow +from temporalio.api.enums.v1 import StreamErrorReason, StreamProducerKind +from temporalio.common._streams import _errors, _ids +from temporalio.common._streams._defn import StreamDefinition, definition_of +from temporalio.common._streams._handle import ( + Creating, + StreamHandle, + StreamOptions, + exit_failure, + workflow_owner, +) +from temporalio.common._streams._items import close_status +from temporalio.common._streams.system import models + +T = TypeVar("T") + +# The create and the append are Local Activities, so History has only their results. The Redis +# close is a regular Activity on the Workflow's task queue. +CREATE_ACTIVITY = "__temporal_streams_create" +APPEND_ACTIVITY = "__temporal_streams_append" +CLOSE_ACTIVITY = "__temporal_streams_close" +_WRITE_TIMEOUT = timedelta(seconds=30) +# Limits of one merged append before the codec. Half of the Store's 4 MiB batch limit leaves room +# for most codecs. The Local Activity splits a merge by call if the codec grows it past the limit. +_MERGE_RECORDS = 1000 +_MERGE_BYTES = 2 * 1024 * 1024 + + +@dataclass +class _RunStream: + # What this run knows about one Stream. Many handles can name it, so the state is keyed by id. + ref: models.StreamRef + closed: bool = False + # Set for a Stream that the run chain created, which continue-as-new passes on. + backend: str | None = None + + +@dataclass +class _Write: + # One Local Activity that writes to a Stream, and the calls that it answers. Each call has the + # Record count up to its own last Record, and its future. + activity: str + request: google.protobuf.message.Message + calls: list[tuple[int, asyncio.Future[int]]] + size: int = 0 + started: bool = False + + def take(self, request: google.protobuf.message.Message) -> bool: + # An append that waits behind the previous write merges the appends that come after it. + # So a burst of awaits costs two markers, not one for each await. + if self.started or not isinstance(request, stream_pb.AppendRecordsRequest): + return False + mine = self.request + if ( + not isinstance(mine, stream_pb.AppendRecordsRequest) + or mine.stream != request.stream + ): + return False + size = sum(record.ByteSize() for record in request.records) + count = len(mine.records) + len(request.records) + if count > _MERGE_RECORDS or self.size + size > _MERGE_BYTES: + return False + mine.records.extend(request.records) + self.size += size + return True + + +@dataclass +class _RunState: + streams: dict[str, _RunStream] = field(default_factory=dict) + # The last write of each Stream. Each write waits for it, so writes keep their order. + tails: dict[str, asyncio.Task[None]] = field(default_factory=dict) + # The last write of each Stream that hasn't started, which later appends can join. + waiting: dict[str, _Write] = field(default_factory=dict) + # The Streams on which a retry run's first write withdrew the previous run's Records. + withdrawn: set[str] = field(default_factory=set) + # Whether the run created, bound or received a Stream. Only then do its Activities get a + # Producer key, so a run with no Stream adds nothing to History. + uses_streams: bool = False + # The async with blocks whose close waits for the run's outcome, with each block's failure. + deferred_closes: list[tuple[WorkflowStream[Any], BaseException | None]] = field( + default_factory=list + ) + # The publishes that failed before a caller got the error. A caller can still get it, so + # the run's end checks them again. + failed_publishes: list[asyncio.Future[int]] = field(default_factory=list) + + def stream(self, ref: models.StreamRef) -> _RunStream: + known = self.streams.get(ref.id) + if known is None: + known = _RunStream(ref) + self.streams[ref.id] = known + return known + + +def _state() -> _RunState: + # Keep the state on the run, since its tasks refer to the run's loop and would keep a weak + # map entry alive. + runtime = temporalio.workflow._Runtime.current() + state = getattr(runtime, "_temporal_streams_state", None) + if state is None: + state = _RunState() + setattr(runtime, "_temporal_streams_state", state) + return state + + +class WorkerBackend(NamedTuple): + """The Stream backend of the Worker that runs the Workflow.""" + + name: str + default: bool + + +def _transport() -> Any: + # The run's outbound interceptor that the Worker's backend marks, or None. + outbound = getattr(temporalio.workflow._Runtime.current(), "_outbound", None) + while outbound is not None: + if getattr(outbound, "_temporal_stream_transport", False): + return outbound + outbound = getattr(outbound, "next", None) + return None + + +def worker_backend() -> WorkerBackend | None: + """The Worker's backend, which marks the run's outbound interceptor, or None.""" + outbound = _transport() + if outbound is None: + return None + return WorkerBackend( + getattr(outbound, "_temporal_stream_backend_name", ""), + getattr(outbound, "_temporal_stream_default", False), + ) + + +async def _close_in_workflow(request: stream_pb.CloseStreamRequest) -> int: + # The Worker's backend runs the close after the Workflow Task commits. + outbound = _transport() + if outbound is None: + raise no_worker_backend() + return await outbound._temporal_stream_backend._close_in_workflow(request) + + +def no_worker_backend() -> _errors.StreamConfigError: + """The error of Workflow code on a Worker with no Stream backend.""" + return _errors.StreamConfigError( + "No Stream backend is on the Worker. Register one on its client with " + "stream_backends=." + ) + + +def _check_backend(backend: str | None, action: str) -> str: + # These errors come from the Worker's config, not from History. So they fail the Workflow + # Task, and a Worker with the right config retries it. + worker = worker_backend() + if worker is None: + raise no_worker_backend() + if not backend: + if not worker.default: + raise _errors.no_default_backend( + "Worker", + [worker.name], + "workflow.create_stream(definition, backend=...)", + ) + return worker.name + if backend != worker.name: + raise _errors.StreamConfigError( + f"Can't {action} a Stream of the backend {backend}. The Worker has only " + f"{worker.name}." + ) + return worker.name + + +def _records(request: google.protobuf.message.Message) -> int: + if isinstance(request, stream_pb.AppendRecordsRequest): + return len(request.records) + return 0 + + +class _WriteFuture(asyncio.Future[int]): + # A caller that gets the result or the error sets retrieved. The run's end raises only the + # errors that no caller got. + retrieved = False + + def result(self) -> int: + self.retrieved = True + return super().result() + + def exception(self) -> BaseException | None: + self.retrieved = True + return super().exception() + + def __await__(self) -> Generator[Any, None, int]: + # The await of asyncio's future doesn't call result(), so this one does. + if not self.done(): + self._asyncio_future_blocking = True + yield self + return self.result() + + __iter__ = __await__ + + +def _write( + stream_id: str, activity: str, request: google.protobuf.message.Message +) -> asyncio.Future[int]: + state = _state() + loop = asyncio.get_running_loop() + future = _WriteFuture(loop=loop) + waiting = state.waiting.get(stream_id) + if waiting is not None and waiting.take(request): + waiting.calls.append((_records(waiting.request), future)) + return future + info = temporalio.workflow.info() + if ( + info.attempt > 1 + and info.continued_run_id + and stream_id not in state.withdrawn + and isinstance( + request, (stream_pb.CreateStreamRequest, stream_pb.AppendRecordsRequest) + ) + ): + # A retry run's first write to a Stream withdraws the Records of the run it retries. + setattr(request, "withdraw_run", info.continued_run_id) + state.withdrawn.add(stream_id) + write = _Write(activity, request, [(_records(request), future)]) + if isinstance(request, stream_pb.AppendRecordsRequest): + write.size = sum(record.ByteSize() for record in request.records) + previous = state.tails.get(stream_id) + + async def run() -> None: + if previous is not None: + try: + await previous + except BaseException: + pass + write.started = True + if state.waiting.get(stream_id) is write: + del state.waiting[stream_id] + try: + if isinstance(write.request, stream_pb.AppendRecordsRequest): + # The Local Activity returns each merged call's last offset or error. + counts = [ + count - before + for (before, _), (count, _) in zip( + [(0, None), *write.calls], write.calls + ) + ] + answers = await temporalio.workflow.execute_local_activity( + activity, + args=[write.request, counts], + result_type=list[Any], + start_to_close_timeout=_WRITE_TIMEOUT, + ) + elif activity == CLOSE_ACTIVITY: + assert isinstance(write.request, stream_pb.CloseStreamRequest) + answers = [await _close_in_workflow(write.request)] + else: + answers = [ + await temporalio.workflow.execute_local_activity( + activity, + write.request, + result_type=int, + start_to_close_timeout=_WRITE_TIMEOUT, + ) + ] + except ( + temporalio.exceptions.ActivityError, + temporalio.exceptions.ApplicationError, + ) as err: + for _, call in write.calls: + if not call.done(): + call.set_exception(_errors.from_activity_error(err)) + return + except BaseException: + for _, call in write.calls: + call.cancel() + raise + for (_, call), answer in zip(write.calls, answers): + if call.done(): + continue + if isinstance(answer, dict): + call.set_exception( + _errors.from_reason(int(answer["reason"]), str(answer["message"])) + ) + else: + call.set_result(int(answer)) + + task = loop.create_task(run()) + state.tails[stream_id] = task + state.waiting[stream_id] = write + return future + + +def _write_nowait(future: asyncio.Future[int]) -> None: + runtime = temporalio.workflow._Runtime.current() + + def done(future: asyncio.Future[int]) -> None: + # No caller awaits the write, so the run raises its error after this callback, as an + # error that the run doesn't catch. + if not future.cancelled() and future.exception() is not None: + setattr(runtime, "_current_activation_error", future.exception()) + + future.add_done_callback(done) + + +async def drain() -> None: + """Wait for the run's writes and closes that nobody awaits, before the run ends. + + A publish error that no caller got raises here, as an error that the run doesn't catch. A + cancel of the run doesn't stop the writes. The run waits for them, then ends as canceled. + """ + state = _state() + cancel: asyncio.CancelledError | None = None + # A Signal or Update handler can start a write while the run waits here. + while pending := [task for task in state.tails.values() if not task.done()]: + for task in pending: + try: + await asyncio.shield(task) + except asyncio.CancelledError as err: + if not task.done(): + cancel = err + except BaseException: + # A failed write already gave its error to its callers. + pass + if cancel is not None: + raise cancel + for future in state.failed_publishes: + error = _unretrieved_error(future) + if error is not None: + raise error + + +def _unretrieved_error(future: asyncio.Future[int]) -> BaseException | None: + # The error of a failed write that no caller got, else None. + if not isinstance(future, _WriteFuture) or future.retrieved or future.cancelled(): + return None + return asyncio.Future.exception(future) + + +def _converters() -> tuple[ + temporalio.converter.PayloadConverter, temporalio.converter.FailureConverter +]: + # Records get no serialization context, so Workflows, Activities and clients encode them + # alike. The backend runs only the codec. + runtime = temporalio.workflow._Runtime.current() + return ( + getattr(runtime, "_context_free_payload_converter"), + getattr(runtime, "_context_free_failure_converter"), + ) + + +def _payload(value: Any) -> Any: + return _converters()[0].to_payload(value) + + +class WorkflowStream(StreamHandle[T]): + """A Stream in Workflow code. Workflow code writes and closes it, and can't read it. + + Pass it to an Activity, a child Workflow or the next run of continue-as-new as an + argument. Use it in ``async with`` to close it at the end of the block. The end of a run + closes only the Streams whose ``async with`` block deferred the close to the run end. A + Stream that the run chain never closes stays open until its ``safety_ttl``. + + .. warning:: + This API is experimental. + """ + + def _model_ref(self) -> models.StreamRef: + return models.StreamRef( + id=self.id, + owner=workflow_owner(self.workflow_id, self.first_execution_run_id), + ) + + def _ref(self) -> stream_v1.StreamRef: + return models._StreamRefTransferTypeConverter().to_transfer_type( + self._model_ref() + ) + + def publish(self, *values: T) -> asyncio.Future[int]: + """Start the write of one Record for each value, in order, and return its future. + + The write has its place in the Stream's order when the call returns, also when no + caller awaits the future. Writes to one Stream keep the order of the calls. The write + is a Local Activity, so History has its result and not the values. Publishes that wait + behind a write to the same Stream are sent as one Local Activity, and each call gets the + offset of its own last Record. A list is one value. + + A caller that awaits the future, or calls its ``result()`` or ``exception()``, gets the + write's error and owns it. An error that no caller gets raises at the end of the run, as + an error of :py:meth:`publish_nowait` does. So a ``StreamError`` fails the run. The + Workflow's logger logs the error when the write fails, if no caller waits for it then. + + Returns: + The future of the write. It gives the offset of the last Record, or raises + ``StreamError`` when the backend refused the write, such as + ``StreamNotFoundError`` for a Stream that nobody created, or + ``StreamClosedError`` for a Stream that another caller closed. + + Raises: + StreamClosedError: This run closed the Stream. + temporalio.workflow.ReadOnlyContextError: The code must not change state, such as + a query handler. + StreamConfigError: The Worker has no Stream backend, or not the Stream's backend. + It fails the Workflow Task, not the run. + """ + state = _state() + future = self._append(values) + loop = asyncio.get_running_loop() + + def warn(future: asyncio.Future[int]) -> None: + error = _unretrieved_error(future) + if error is not None: + temporalio.workflow.logger.warning( + "A publish to the Stream %s failed, and no caller has its error. If no " + "caller gets the error, it fails the run when the run ends: %s", + self.id, + error, + ) + + def failed(future: asyncio.Future[int]) -> None: + if _unretrieved_error(future) is not None: + state.failed_publishes.append(future) + # A caller that already waits for the future gets the error before warn runs. + loop.call_soon(warn, future) + + future.add_done_callback(failed) + return future + + def publish_nowait(self, *values: T) -> None: + """Write one Record for each value, in order, without an await. + + Use it in sync Workflow code, such as ``@workflow.init`` and sync signal and update + handlers. It's like ``asyncio.Queue.put_nowait`` for code that can't await ``put``. A + backend error raises in the run as an error that the run doesn't catch, so a + ``StreamError`` fails the run. + """ + _write_nowait(self._append(values)) + + def _check(self, action: str, *, backend_checked: bool = False) -> _RunStream: + if temporalio.workflow.unsafe.is_read_only(): + raise temporalio.workflow.ReadOnlyContextError( + f"Can't {action} a Stream in read-only code" + ) + state = _state() + state.uses_streams = True + run_stream = state.stream(self._model_ref()) + if run_stream.closed: + raise _errors.StreamClosedError( + f"Can't {action} the Stream {self.id}, since this run closed it", + reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + if not backend_checked: + _check_backend(self.backend, action) + return run_stream + + def _append(self, values: tuple[Any, ...]) -> asyncio.Future[int]: + if not values: + raise TypeError("publish takes one or more values") + self._check("write to") + records = [stream_v1.AppendRecord(body=_payload(v)) for v in values] + return _write( + self.id, + APPEND_ACTIVITY, + stream_pb.AppendRecordsRequest(stream=self._ref(), records=records), + ) + + async def close( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Close the Stream after the writes before it, once the Workflow Task commits. + + If the task fails first, or the Worker crashes, the Stream stays open, and the task's + retry writes to it again. On Redis the close is a regular Activity on the Workflow's + task queue, which closes the Stream as the run chain that owns it. So every Worker on + that queue must register the Redis backend. A Worker that polls Activities runs the + close, and a Worker with ``no_remote_activities=True`` doesn't. History keeps + ``result`` and ``failure`` as the close Activity's input, so the payload size limit + applies to them. + + Readers get the Records before the close, then a ``Closed`` item with ``result`` or + ``failure``. A cancel failure, such as a ``CancelledError`` or an ``ActivityError`` + that a cancel caused, closes the Stream as canceled. Later writes in this run raise + :py:class:`StreamClosedError`. + + Raises: + StreamNotOwnerError: This run chain didn't create the Stream. + """ + await self._close(result, failure) + + def close_nowait( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Close the Stream without an await, for sync Workflow code. + + It is to :py:meth:`close` what ``asyncio.Queue.put_nowait`` is to ``put``. It starts the + close, and the end of the run waits for it. A backend error raises in the run as an + error that the run doesn't catch, so a ``StreamError`` fails the run. + """ + _write_nowait(self._close(result, failure)) + + def _closing(self) -> _RunStream: + info = temporalio.workflow.info() + if (self.workflow_id, self.first_execution_run_id) != ( + info.workflow_id, + info.first_execution_run_id, + ): + # Only the run chain that created a Stream can close it. The backend refuses the + # close too, but this check fails before the Workflow schedules a command. + raise _errors.StreamNotOwnerError( + f"Can't close the Stream {self.id}, since this run chain didn't create it", + reason=StreamErrorReason.STREAM_ERROR_REASON_NOT_OWNER, + ) + return self._check("close") + + def _close(self, result: Any, failure: BaseException | None) -> asyncio.Future[int]: + self._closing().closed = True + request = stream_pb.CloseStreamRequest(stream=self._ref()) + if result is not None: + request.result.CopyFrom(_payload(result)) + if failure is not None: + payload_converter, failure_converter = _converters() + failure_converter.to_failure(failure, payload_converter, request.failure) + status = close_status(failure) + if status is not None: + request.status = status + return _write(self.id, CLOSE_ACTIVITY, request) + + async def __aenter__(self) -> WorkflowStream[T]: + """Use the Stream in ``async with``, which closes it at the end.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream, with the exception as its failure if there is one. + + A cancel of the block or of the run closes the Stream as canceled. An exception that + fails the run with no retry run closes it too. Examples are an exception in a run with + no retry policy, a non-retryable ``ApplicationError``, and an exception on the last + attempt. + + For any other exception, the close waits for the run's end. The Stream then closes with + that exception if the run returns, continues as new, or fails with no retry run. It + stays open if a retry run comes next, since that run writes to the same Stream. It also + stays open if the exception fails only the Workflow Task, since the task's retry runs + the block again. + + In a cron run, the block end leaves the Stream open, since the Stream spans the cron + chain and the next run creates it again. An eviction cancels the block, and the Worker + sends no commands for an evicted run. An explicit :py:meth:`close` always closes the + Stream. + """ + state = _state() + if state.stream(self._model_ref()).closed: + return + close, failure = exit_failure(exc) + if not close or temporalio.workflow.info().cron_schedule: + return + if _closes_at_block_end(exc): + await self.close(failure=failure) + else: + # A refused close raises here, as it does for a close at the block end. + self._closing() + state.deferred_closes.append((self, failure)) + + +def _closes_at_block_end(exc: BaseException | None) -> bool: + # For other exceptions, only the run's end shows whether a retry run will write to the Stream. + if exc is None or isinstance(exc, asyncio.CancelledError) or _cancelled(exc): + return True + runtime = temporalio.workflow._Runtime.current() + return runtime.workflow_is_failure_exception(exc) and not _retried(exc) + + +def close_at_run_end(err: BaseException | None) -> None: + """Close the Streams whose block end deferred the close, unless a retry run comes next. + + ``err`` is the exception that ends the run, or None if the run returns. + """ + state = _state() + runtime = temporalio.workflow._Runtime.current() + if ( + isinstance(err, Exception) + and not _cancelled(err) + and runtime.workflow_is_failure_exception(err) + and _retried(err) + ): + return + for stream, failure in state.deferred_closes: + if not state.stream(stream._model_ref()).closed: + _write_nowait(stream._close(None, failure)) + + +def _retried(exc: BaseException) -> bool: + # Whether the server starts a retry run after the run fails with exc, by the server's rule. + info = temporalio.workflow.info() + policy = info.retry_policy + if policy is None or 0 < policy.maximum_attempts <= info.attempt: + return False + failure_type: str | None + if isinstance(exc, temporalio.exceptions.ApplicationError): + if exc.non_retryable: + return False + failure_type = exc.type + elif isinstance( + exc, + (temporalio.exceptions.CancelledError, temporalio.exceptions.TerminatedError), + ): + return False + elif isinstance(exc, temporalio.exceptions.TimeoutError): + return exc.type in ( + temporalio.exceptions.TimeoutType.START_TO_CLOSE, + temporalio.exceptions.TimeoutType.HEARTBEAT, + ) + elif isinstance(exc, temporalio.exceptions.FailureError): + return True + else: + # For any other exception, the run's failure has the class name as its type. + failure_type = type(exc).__name__ + return failure_type not in (policy.non_retryable_error_types or ()) + + +def _cancelled(exc: BaseException) -> bool: + # A cancelled run ends as canceled with no retry, even when the cancel reached the block as + # the error of an Activity, a Local Activity or a child. + return ( + temporalio.workflow.cancellation_reason() is not None + and temporalio.exceptions.is_cancelled_exception(exc) + ) + + +def _own_handle( + definition: type[StreamDefinition[T]], backend: str | None +) -> WorkflowStream[T]: + info = temporalio.workflow.info() + known = definition_of(definition) + return WorkflowStream( + id=_ids.scoped_stream_id( + info.workflow_id, info.first_execution_run_id, known.name + ), + workflow_id=info.workflow_id, + first_execution_run_id=info.first_execution_run_id, + backend=backend, + payload_type=known.type, + ) + + +def create_stream( + definition: type[StreamDefinition[T]], + *, + backend: str | None, + options: StreamOptions, +) -> Creating[WorkflowStream[T]]: + """Create a Stream that this run chain owns, or get the one that exists. + + See :py:func:`temporalio.workflow.create_stream`. + """ + + async def create() -> WorkflowStream[T]: + stream, task = _create(definition, options, backend) + await task + return stream + + return Creating(create) + + +def create_stream_nowait( + definition: type[StreamDefinition[T]], + *, + backend: str | None, + options: StreamOptions, +) -> WorkflowStream[T]: + """Create a Stream that this run chain owns, or get it, without an await. + + See :py:func:`temporalio.workflow.create_stream_nowait`. + """ + stream, task = _create(definition, options, backend) + _write_nowait(task) + return stream + + +def _create( + definition: type[StreamDefinition[T]], + options: StreamOptions, + backend: str | None, +) -> tuple[WorkflowStream[T], asyncio.Future[int]]: + stream = _own_handle(definition, _check_backend(backend, "create")) + run_stream = stream._check("create", backend_checked=True) + run_id = temporalio.workflow.info().run_id + # A retry run's withdrawal names the owner's run id. The Store keeps the ref's run chain, and + # checks each close against it. + request = stream_pb.CreateStreamRequest( + stream_id=stream.id, + owner=stream_v1.Producer( + kind=StreamProducerKind.STREAM_PRODUCER_KIND_WORKFLOW, + id=_ids.workflow_producer_id(run_id), + run_id=run_id, + ), + ref=stream._ref(), + ) + retention = options.get("retention") + if retention is not None: + request.retention.FromTimedelta(retention) + run_stream.backend = stream.backend + return stream, _write(stream.id, CREATE_ACTIVITY, request) + + +def stream_handle( + definition: type[StreamDefinition[T]], *, backend: str | None = None +) -> WorkflowStream[T]: + """Get a Stream that this run chain created, by its definition. + + See :py:func:`temporalio.workflow.stream_handle`. + """ + stream = _own_handle(definition, backend) + if not temporalio.workflow.unsafe.is_read_only(): + _state().uses_streams = True + if backend is None: + known = _state().streams.get(stream.id) + stream._backend = (known.backend if known else None) or _check_backend( + None, "get" + ) + return stream diff --git a/temporalio/common/_streams/_writes.py b/temporalio/common/_streams/_writes.py new file mode 100644 index 000000000..4eb9efa7d --- /dev/null +++ b/temporalio/common/_streams/_writes.py @@ -0,0 +1,298 @@ +"""The Activities that write for Workflow code. + +Each Workflow ``create_stream`` is a Local Activity. The publishes that wait behind a write to +the same Stream share one Local Activity. History has the result and not the values. Each +Workflow ``close`` is a regular Activity, so it runs only after its Workflow Task completes. +""" + +from __future__ import annotations + +import dataclasses +from collections import OrderedDict +from typing import Any + +import temporalio.activity +import temporalio.api.stream.v1 as stream_v1 +import temporalio.api.streamservice.v1 as stream_pb +import temporalio.common +from temporalio.api.enums.v1 import StreamErrorReason, StreamStatus +from temporalio.common._streams import _activity, _errors, _ids +from temporalio.common._streams._workflow import ( + APPEND_ACTIVITY, + CLOSE_ACTIVITY, + CREATE_ACTIVITY, +) +from temporalio.common._streams.system import client as service_client +from temporalio.common._streams.system import models + + +def _context() -> tuple[ + _activity.ActivityStreams, service_client.StreamServiceClientExtension +]: + streams = _activity.current.get() + if streams is None: + raise _errors.StreamConfigError("No Stream backend runs this Local Activity") + return streams, streams.backend.service(streams.namespace, streams.converter) + + +def _producer(streams: _activity.ActivityStreams) -> models.Producer: + # All attempts of one Workflow Task write with one epoch and sequence, so the Store dedupes a + # retry. A later task's higher epoch withdraws a failed task's Records for the same call key. + producer = streams.producer + epoch = producer.epoch + started = epoch.task_started_time if epoch.HasField("task_started_time") else None + return models.Producer( + kind=models.StreamProducerKind(producer.kind), + id=producer.id, + run_id=producer.run_id or None, + epoch=models.ProducerEpoch( + task_started_event_id=epoch.task_started_event_id or None, + task_attempt=epoch.task_attempt or None, + task_started_time=models.Timestamp( + seconds=started.seconds, nanos=started.nanos + ) + if started is not None + else None, + attempt=1, + ), + ) + + +def _history(streams: _activity.ActivityStreams) -> dict[str, Any]: + # The run's task history, so the write withdraws the run's orphaned Local Activity Records. + previous = streams.previous_completed_task + return { + "previous_completed_task": models.WorkflowTaskEpoch( + started_event_id=previous.started_event_id + ) + if previous is not None + else None, + "failed_tasks": list(streams.failed_tasks), + } + + +def _ref(ref: stream_v1.StreamRef) -> models.StreamRef: + return models._StreamRefTransferTypeConverter().from_transfer_type( + ref, models.StreamRef + ) + + +@temporalio.activity.defn(name=CREATE_ACTIVITY) +async def create_stream(request: stream_pb.CreateStreamRequest) -> int: + """Create a Workflow's Stream, or get the one that exists.""" + _, service = _context() + created = await service.create_stream( + stream_id=request.stream_id, + retention=request.retention.ToTimedelta() + if request.HasField("retention") + else None, + withdraw_run=request.withdraw_run or None, + owner=models.Producer( + kind=models.StreamProducerKind(request.owner.kind), + id=request.owner.id, + run_id=request.owner.run_id or None, + ), + ref=_ref(request.ref) if request.HasField("ref") else None, + ) + if created.status == StreamStatus.STREAM_STATUS_CLOSED: + raise _errors.StreamClosedError( + f"The Stream {request.stream_id} was created and closed before", + reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + if created.status == StreamStatus.STREAM_STATUS_EXPIRED: + raise _errors.StreamExpiredError( + f"The Stream {request.stream_id} expired", + reason=StreamErrorReason.STREAM_ERROR_REASON_EXPIRED, + ) + return 0 + + +# The Store's limits on the Records that it keeps, as the crate counts them. +_MAX_RECORD_BYTES = 1536 * 1024 +_MAX_BATCH_BYTES = 4 * 1024 * 1024 + + +@dataclasses.dataclass +class _Encoded: + # The batches of one append, with each batch's calls and encoded bodies, and the result of + # each batch that the Store accepted. + batches: list[tuple[list[int], list[stream_v1.AppendRecord]]] + size: int + answers: dict[int, Any] = dataclasses.field(default_factory=dict) + + +_Key = tuple[str, str, str] +_Epoch = tuple[int, int, int, int] + + +class _Pending: + # The appends that haven't succeeded, one for each Producer. A retry in this process resends + # the same bytes, so the Store's dedupe matches them, even with a random codec. + + def __init__(self, limit: int) -> None: + self._limit = limit + self._size = 0 + self._entries: OrderedDict[_Key, tuple[_Epoch, _Encoded]] = OrderedDict() + + def get(self, key: _Key, epoch: _Epoch) -> _Encoded | None: + found = self._entries.get(key) + return found[1] if found is not None and found[0] == epoch else None + + def put(self, key: _Key, epoch: _Epoch, encoded: _Encoded) -> None: + found = self._entries.get(key) + if found is not None: + if found[0] > epoch: + # The Store fences this write, since the Producer has a later epoch. + return + self.pop(key, found[0]) + self._entries[key] = (epoch, encoded) + self._size += encoded.size + while self._size > self._limit and len(self._entries) > 1: + _, (_, oldest) = self._entries.popitem(last=False) + self._size -= oldest.size + + def pop(self, key: _Key, epoch: _Epoch) -> None: + found = self._entries.get(key) + if found is not None and found[0] == epoch: + del self._entries[key] + self._size -= found[1].size + + +# A run that ends while its append retries leaves its entry, so evict the oldest past this size. +_PENDING = _Pending(64 * 1024 * 1024) + + +async def _encode( + streams: _activity.ActivityStreams, + request: stream_pb.AppendRecordsRequest, + counts: list[int], + producer: stream_v1.Producer, +) -> _Encoded: + # The Workflow converted the values, so only the codec runs on them here. + bodies = [record.body for record in request.records] + codec = streams.converter.payload_codec + if codec is not None: + bodies = await codec.encode(bodies) + records = [stream_v1.AppendRecord(body=body) for body in bodies] + sizes = [ + stream_v1.StreamRecord( + data=stream_v1.StreamData(producer=producer, body=record.body) + ).ByteSize() + for record in records + ] + if sum(sizes) <= _MAX_BATCH_BYTES and max(sizes, default=0) <= _MAX_RECORD_BYTES: + return _Encoded([(list(range(len(counts))), records)], sum(sizes)) + # The codec grew the merge past the Store's limit, so send each call as its own batch, as + # with no merge. A call that's too large by itself fails by itself. + batches: list[tuple[list[int], list[stream_v1.AppendRecord]]] = [] + start = 0 + for call, count in enumerate(counts): + batches.append(([call], records[start : start + count])) + start += count + return _Encoded(batches, sum(sizes)) + + +@temporalio.activity.defn(name=APPEND_ACTIVITY) +async def append_records( + request: stream_pb.AppendRecordsRequest, counts: list[int] +) -> list[Any]: + """Append the Records of a Workflow's merged publishes. + + ``counts`` has the Record count of each publish. Returns each publish's last offset, or the + error of a publish that the Store refused as too large. + """ + streams, service = _context() + producer = _producer(streams) + producer_pb = models._ProducerTransferTypeConverter().to_transfer_type(producer) + key = (streams.namespace, request.stream.id, producer_pb.id) + task = producer_pb.epoch + epoch = ( + task.task_started_event_id, + task.task_attempt, + task.task_started_time.seconds, + task.task_started_time.nanos, + ) + encoded = _PENDING.get(key, epoch) + if encoded is None: + encoded = await _encode(streams, request, counts, producer_pb) + _PENDING.put(key, epoch, encoded) + # The codec already ran on the bodies. + service = streams.backend.service( + streams.namespace, dataclasses.replace(streams.converter, payload_codec=None) + ) + answers: list[Any] = [0] * len(counts) + for index, (calls, records) in enumerate(encoded.batches): + answer = encoded.answers.get(index) + if answer is None: + try: + response = await ( + await service.append_records( + stream=_ref(request.stream), + producer=producer, + sequence=index + 1, + records=[ + models.AppendRecord( + body=temporalio.common.RawValue(record.body) + ) + for record in records + ], + withdraw_run=(request.withdraw_run or None) + if index == 0 + else None, + **_history(streams), + ) + ) + answer = response.last_offset or 0 + except _errors.StreamError as err: + if err.reason != StreamErrorReason.STREAM_ERROR_REASON_TOO_LARGE: + if err.non_retryable: + _PENDING.pop(key, epoch) + raise + answer = {"reason": err.reason, "message": err.message} + encoded.answers[index] = answer + # A Store gives one batch consecutive offsets, so each call's last Record is the + # batch's last one less the Records after it. + after = sum(counts[call] for call in calls) + for call in calls: + after -= counts[call] + answers[call] = answer - after if isinstance(answer, int) else answer + _PENDING.pop(key, epoch) + return answers + + +@temporalio.activity.defn(name=CLOSE_ACTIVITY) +async def close_stream(request: stream_pb.CloseStreamRequest) -> int: + """Close a Workflow's Stream, as the run chain that owns it.""" + streams, service = _context() + # The Store refuses an Activity Producer's close of a Workflow's Stream, so the Activity + # writes the close as the run that scheduled it, and only for that run's chain. + if request.stream.owner.business_id != streams.workflow_id: + raise _errors.StreamNotOwnerError( + f"Can't close the Stream {request.stream.id} for another Workflow", + reason=StreamErrorReason.STREAM_ERROR_REASON_NOT_OWNER, + ) + run_id = streams.producer.run_id + failure: Any = None + if request.HasField("failure"): + failure = streams.converter.failure_converter.from_failure( + request.failure, streams.converter.payload_converter + ) + await ( + await service.close_stream( + stream=_ref(request.stream), + producer=models.Producer( + kind=models.StreamProducerKind.StreamProducerKindWorkflow, + id=_ids.workflow_producer_id(run_id), + run_id=run_id, + ), + result=temporalio.common.RawValue(request.result) + if request.HasField("result") + else None, + failure=failure, + status=models.StreamCloseStatus(request.status) if request.status else None, + ) + ) + return 0 + + +ACTIVITIES = [create_stream, append_records, close_stream] diff --git a/temporalio/contrib/stream/__init__.py b/temporalio/contrib/stream/__init__.py index 4e8c7fcd2..19e24062e 100644 --- a/temporalio/contrib/stream/__init__.py +++ b/temporalio/contrib/stream/__init__.py @@ -1,8 +1,8 @@ """The Redis Stream backend. -Register it on the client. A client creates a Stream with -:py:meth:`temporalio.client.Client.create_stream`. ``backend=`` picks a backend that isn't the -default:: +Register it on the client. Workflow code and clients create Streams with +:py:func:`temporalio.workflow.create_stream` and :py:meth:`temporalio.client.Client.create_stream`. +``backend=`` picks a backend that isn't the default:: from temporalio.contrib.stream import RedisStreamBackend @@ -11,6 +11,10 @@ stream_backends=[RedisStreamBackend("redis://localhost", default=True)], ) +A Workflow's close is a regular Activity on the Workflow's task queue. So a Worker with a Redis +backend must poll Activities on its task queue for the close Activity. All Workers that share a +task queue must register the Redis backend. + .. warning:: This package is experimental. """ @@ -34,6 +38,10 @@ class RedisStreamBackend(ExtensionStreamBackend): Stream handle that it creates stores that name. A Worker serves only one backend. A client can have more than one, and ``backend=`` picks one. + A Workflow's close is a regular Activity on the Workflow's task queue, which runs after the + Workflow Task commits. So a Worker with this backend must poll Activities on its task queue, + and all Workers that share a task queue must register the backend. + .. warning:: This API is experimental. """ diff --git a/temporalio/contrib/stream/testing.py b/temporalio/contrib/stream/testing.py index 8d393ae59..79a1beb7b 100644 --- a/temporalio/contrib/stream/testing.py +++ b/temporalio/contrib/stream/testing.py @@ -19,7 +19,8 @@ class MemoryStreamBackend(ExtensionStreamBackend): Register it on the client, as with ``RedisStreamBackend``. Every Worker and client that shares one backend shares its Streams. The Streams are lost when the process ends. It registers the backend name ``memory``. It has Redis's rules, so tests can swap the Store - without changing Workflow code. + without changing Workflow code. As on Redis, a Workflow's close is a regular Activity, so + all Workers that share a task queue must register the backend. .. warning:: This API is experimental. diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index c631fd88e..ef1c20a37 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -151,6 +151,11 @@ logger, unsafe, ) +from ._streams import ( + create_stream, + create_stream_nowait, + stream_handle, +) from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowConfig, @@ -258,6 +263,9 @@ "SandboxImportNotificationPolicy", "logger", "unsafe", + "create_stream", + "create_stream_nowait", + "stream_handle", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py new file mode 100644 index 000000000..c220c869a --- /dev/null +++ b/temporalio/workflow/_streams.py @@ -0,0 +1,122 @@ +"""Streams in Workflow code. + +These functions import the Streams API of :py:mod:`temporalio.common` on first use, so a +Workflow with no Stream never imports it. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, TypeVar + +from typing_extensions import Unpack + +if TYPE_CHECKING: + from temporalio.common import ( + Creating, + StreamDefinition, + StreamOptions, + WorkflowStream, + ) + +T = TypeVar("T") + + +def create_stream( + definition: type[StreamDefinition[T]], + *, + backend: str | None = None, + **options: Unpack[StreamOptions], +) -> Creating[WorkflowStream[T]]: + """Create a Stream that this run chain owns, or get the one that exists. + + A Stream must exist before its first publish. The Stream id is + ``::``, with each part percent-encoded, so a create is + create-or-get by that id. On a Stream that exists, it ignores ``options`` and returns that + Stream. So a Workflow Task retry, a retry run and the run after continue-as-new all get the + same Stream. The create and each publish are Local Activities, and the create's error is + raised here. Only the run chain can close the Stream. A close runs after its Workflow Task + commits, as a regular Activity on Redis, so a task that fails first leaves the Stream open + for its retry. On Redis, the end of the run chain doesn't close a Stream outside an + ``async with`` block. So a Stream that nobody closes stays open until its ``safety_ttl``. + + Await the call to get the handle, or use it in ``async with``, which closes the Stream at + the end of the block:: + + async with workflow.create_stream(Tokens) as tokens: + await tokens.publish("hi") + + An exception in the block closes the Stream with that exception as the failure. If a retry + run can come next, the close waits for the run's end. The Stream then stays open only for a + retry run, or for the retry of a failed Workflow Task. A cron run leaves the Stream open. + See :py:meth:`temporalio.common.WorkflowStream.__aexit__`. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class. Its name is the Stream name. + backend: The name of the backend. Defaults to the registered backend with + ``default=True``. The default is part of the Workflow's code, since a replay must + create the Stream on the same backend. A Worker serves only one backend, so the + name must be the Worker's backend name. + options: The options of a new Stream, such as ``retention``. See + :py:class:`temporalio.common.StreamOptions`. + + Raises: + StreamError: The backend refused the create. It's ``StreamClosedError`` or + ``StreamExpiredError`` if a Stream with that id closed or expired. + ReadOnlyContextError: The code must not change state, such as a query handler. + StreamConfigError: The Worker has no Stream backend or none with that name, or no + backend has ``default=True`` and the call gives no backend. It fails the Workflow + Task, not the run. + """ + from temporalio.common._streams import _workflow + + return _workflow.create_stream(definition, backend=backend, options=options) + + +def create_stream_nowait( + definition: type[StreamDefinition[T]], + *, + backend: str | None = None, + **options: Unpack[StreamOptions], +) -> WorkflowStream[T]: + """Create a Stream that this run chain owns, or get it, without an await. + + Use it in sync Workflow code, such as ``@workflow.init``. It's like + ``asyncio.Queue.put_nowait`` for code that can't await ``put``. It takes the arguments of + :py:func:`create_stream`. A backend error raises in the run as an error that the run + doesn't catch, so a ``StreamError`` fails the run. + + .. warning:: + This API is experimental. + """ + from temporalio.common._streams import _workflow + + return _workflow.create_stream_nowait(definition, backend=backend, options=options) + + +def stream_handle( + definition: type[StreamDefinition[T]], *, backend: str | None = None +) -> WorkflowStream[T]: + """Get a Stream that this run chain created, by its definition. It makes no call. + + A later run of the chain, such as one after continue-as-new, uses it to write to or close + the Stream. In ``@workflow.init`` of a run that continued as new, the run hasn't read the + backends of the chain's Streams yet. So pass ``backend=`` there, or call it in the run + method. + + .. warning:: + This API is experimental. + + Args: + definition: The Stream definition class, with the same name as in the create. + backend: The backend of the Stream. Defaults to the backend that the run chain created + it on, or else to the registered backend with ``default=True``. + + Raises: + StreamConfigError: No backend is given, known or marked as the default. + """ + from temporalio.common._streams import _workflow + + return _workflow.stream_handle(definition, backend=backend) diff --git a/tests/contrib/streams/_support.py b/tests/contrib/streams/_support.py index 9f1a05f10..27e0a9597 100644 --- a/tests/contrib/streams/_support.py +++ b/tests/contrib/streams/_support.py @@ -1,4 +1,4 @@ -"""Finds the Stream crates that the bridge links, for the shared vectors and protos.""" +"""Finds the Stream crates that the bridge links, and binds a Workflow's Stream in tests.""" from __future__ import annotations @@ -6,9 +6,17 @@ import json import subprocess from pathlib import Path +from typing import Any, TypeVar import pytest +from temporalio import common +from temporalio.client import WorkflowHandle +from temporalio.common._streams import _ids +from temporalio.common._streams._defn import definition_of + +T = TypeVar("T") + _BRIDGE = Path(__file__).parents[3] / "temporalio" / "bridge" @@ -39,3 +47,16 @@ def crate_file(relative: str, crate: str = "temporalio-streams") -> Path: def vectors(name: str) -> object: return json.loads(crate_file(f"testdata/{name}").read_text(encoding="utf-8")) + + +def workflow_stream( + handle: WorkflowHandle[Any, Any], + definition: type[common.StreamDefinition[T]], + backend: str | None = None, +) -> common.ClientStream[T]: + """The Stream of ``definition`` that the Workflow of ``handle`` created.""" + # A Workflow's Stream id is ::. + stream_id = _ids.scoped_stream_id( + handle.id, handle.first_execution_run_id or "", definition_of(definition).name + ) + return handle._client.stream(definition, id=stream_id, backend=backend) diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py new file mode 100644 index 000000000..b1ec38ae2 --- /dev/null +++ b/tests/contrib/streams/test_redis.py @@ -0,0 +1,66 @@ +import uuid +from datetime import timedelta +from typing import Any + +from temporalio import common, workflow +from temporalio.client import Client +from temporalio.worker import ( + UnsandboxedWorkflowRunner, +) +from tests.contrib.streams._support import workflow_stream +from tests.helpers import new_worker + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +async def read_all(stream: common.ClientStream[Any]) -> list[Any]: + return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] + + +def records(items: list[Any]) -> list[common.Record[Any]]: + return [item for item in items if isinstance(item, common.Record)] + + +attempts: dict[str, int] = {} + + +@workflow.defn +class PublishThenFailTask: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream(Tokens) + await stream.publish("before") + await workflow.sleep(timedelta(milliseconds=10)) + await stream.publish("after") + if attempts.setdefault(key, 0) == 0: + attempts[key] = 1 + raise RuntimeError("fail this Workflow Task once") + await stream.close() + + +async def test_failed_workflow_task_is_withdrawn(stream_client: Client): + key = str(uuid.uuid4()) + # The test's counter must keep its value across the task's retry, so the Workflow has no + # sandbox. + async with new_worker( + stream_client, + PublishThenFailTask, + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await stream_client.start_workflow( + PublishThenFailTask.run, + key, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = await read_all(workflow_stream(handle, Tokens)) + # The retried task writes "after" again, and withdraws the copy of the failed task. + written = records(items) + superseded = [item for item in items if isinstance(item, common.Superseded)] + assert [record.data for record in written] == ["before", "after", "after"] + withdrawn = [r for r in written if any(i.withdraws(r) for i in superseded)] + assert withdrawn == [written[1]] + assert isinstance(items[-1], common.Closed) diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index ce1d1bf3e..7d3125174 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -13,6 +13,7 @@ from temporalio.contrib.stream.testing import MemoryStreamBackend from temporalio.converter import DataConverter, DefaultFailureConverter from temporalio.exceptions import ApplicationError +from temporalio.worker import Replayer from tests.contrib.streams.fake_extension import fake_client from tests.helpers import new_worker @@ -91,6 +92,80 @@ async def test_query_result_binds_its_client_stream(client: Client): ] +@workflow.defn +class WriteToStream: + @workflow.run + async def run(self, stream: common.WorkflowStream[Turn]) -> None: + await stream.publish(Turn("workflow")) + + +async def test_client_stream_takes_workflow_writes(stream_client: Client): + stream = await stream_client.create_stream(Turns, id=f"chat-{uuid.uuid4()}") + assert await stream.get_offset() == 0 + first = await stream.publish(Turn("client")) + last = await stream.publish(Turn("a"), Turn("b")) + assert last > first + async with new_worker(stream_client, WriteToStream) as worker: + await stream_client.execute_workflow( + WriteToStream.run, + common.WorkflowStream.from_handle(stream), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await stream.close(result="bye") + items = [item async for item in stream.subscribe()] + data = [item.data for item in items if isinstance(item, common.Record)] + assert data == [Turn("client"), Turn("a"), Turn("b"), Turn("workflow")] + assert isinstance(items[-1], common.Closed) and items[-1].result == "bye" + + # Another handle of the same Stream reads it from an offset. + again = stream_client.stream(Turns, id=stream.id) + resumed = [ + item.data + async for item in again.subscribe(offset=first + 1) + if isinstance(item, common.Record) + ] + assert resumed == [Turn("a"), Turn("b"), Turn("workflow")] + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + await workflow.sleep(timedelta(milliseconds=200)) + stream = await workflow.create_stream(Tokens) + await stream.publish(*prompt.split()) + await stream.close(result=prompt.upper()) + return prompt.upper() + + @workflow.query + def stream(self) -> common.WorkflowStream[str]: + # The handle needs no call, so a reader gets it before the create. + return workflow.stream_handle(Tokens) + + +async def test_client_reads_a_workflow_stream(stream_client: Client): + async with new_worker(stream_client, Agent) as worker: + handle = await stream_client.start_workflow( + Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + stream = await handle.query(Agent.stream) + reader = await handle.query("stream", result_type=common.ClientStream[str]) + # The Workflow creates its Stream later, so the read waits for the create. + items = [ + item async for item in reader.subscribe(create_wait=timedelta(seconds=20)) + ] + assert stream.id == reader.id + assert [item.data for item in items if isinstance(item, common.Record)] == [ + "hi", + "there", + ] + assert isinstance(items[-1], common.Closed) and items[-1].result == "HI THERE" + assert reader.workflow_id == handle.id + with pytest.raises(common.StreamNotOwnerError): + await reader.close() + + async def test_async_with_closes_a_client_stream(stream_client: Client): async with stream_client.create_stream(Tokens, id=f"block-{uuid.uuid4()}") as done: await done.publish("a") @@ -292,6 +367,15 @@ async def test_read_keeps_no_offsets_of_client_records(stream_client: Client): assert read._shown == {} +async def test_empty_stream_id_raises_with_no_call(client: Client): + client, extension = fake_client(client) + with pytest.raises(ValueError, match="id is empty"): + client.stream(Turns, id="") + with pytest.raises(ValueError, match="id is empty"): + client.stream(Turns, stream=common.WorkflowStream(id="")) + assert not extension.namespaces + + async def test_create_of_an_expired_stream_raises(stream_client: Client): # An open client Stream expires one Retention after its last write. stream_id = f"expired-{uuid.uuid4()}" @@ -336,6 +420,22 @@ async def test_clients_of_one_runtime_share_the_backend_calls(stream_client: Cli ] +async def test_registration_rules(client: Client): + config = client.config() + config["stream_backends"] = [ + MemoryStreamBackend(), + MemoryStreamBackend(name="b", default=True), + ] + two = Client(**config) + assert (await two.create_stream(Tokens, id="d-1")).backend == "b" + assert two.stream(Tokens, id="d-1").backend == "b" + # A Worker serves only one backend, and so does a Replayer. + with pytest.raises(common.StreamConfigError, match="one backend"): + new_worker(Client(**config), Agent) + with pytest.raises(common.StreamConfigError, match="one backend"): + Replayer(workflows=[Agent], stream_backends=config["stream_backends"]) + + async def test_default_is_native_with_no_default_backend(client: Client): config = client.config() config["stream_backends"] = [MemoryStreamBackend()] diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py new file mode 100644 index 000000000..ca599fdb0 --- /dev/null +++ b/tests/contrib/streams/test_workflow_module.py @@ -0,0 +1,1154 @@ +import asyncio +import dataclasses +import logging +import os +import uuid +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +import pytest + +from temporalio import activity, common, workflow +from temporalio.api.common.v1 import Payload +from temporalio.api.enums.v1 import EventType, StreamErrorReason +from temporalio.client import ( + Client, + WorkflowFailureError, + WorkflowHandle, +) +from temporalio.common import RetryPolicy +from temporalio.common._streams import _writes +from temporalio.common._streams._backend import _ExtensionCalls +from temporalio.contrib.stream.testing import MemoryStreamBackend +from temporalio.converter import ( + DataConverter, + DefaultPayloadConverter, + PayloadCodec, + SerializationContext, + WithSerializationContext, + WorkflowSerializationContext, +) +from temporalio.exceptions import ApplicationError, CancelledError +from temporalio.worker import Replayer, UnsandboxedWorkflowRunner +from tests.contrib.streams._support import workflow_stream +from tests.helpers import assert_eventually, assert_task_fail_eventually, new_worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@common.defn(name="turns") +class Turns(common.StreamDefinition[Turn]): ... + + +@common.defn(name="tokens") +class Tokens(common.StreamDefinition[str]): ... + + +async def read( + handle: WorkflowHandle[Any, Any], + definition: type[common.StreamDefinition[Any]] = Tokens, +) -> list[object]: + stream = workflow_stream(handle, definition) + return [item async for item in stream.subscribe(wait_new_records=False)] + + +def data(items: Sequence[object]) -> list[object]: + return [item.data for item in items if isinstance(item, common.Record)] + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + # Two channels are two Streams. + turns = await workflow.create_stream(Turns) + tokens = await workflow.create_stream(Tokens) + await turns.publish(Turn("user", prompt), Turn("model")) + await tokens.publish(*prompt.split()) + # A run's end doesn't close these Streams, so the run closes them before it returns. + await turns.close(result=prompt.upper()) + await tokens.close() + return prompt.upper() + + +async def test_create_publish_and_close(stream_client: Client): + async with new_worker(stream_client, Agent) as worker: + handle = await stream_client.start_workflow( + Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == "HI THERE" + turns = await read(handle, Turns) + tokens = await read(handle, Tokens) + assert data(turns) == [Turn("user", "hi there"), Turn("model")] + assert data(tokens) == ["hi", "there"] + closed = turns[-1] + assert isinstance(closed, common.Closed) and closed.result == "HI THERE" + assert isinstance(tokens[-1], common.Closed) + + +@common.defn(name="first") +class First(common.StreamDefinition[str]): ... + + +@common.defn(name="second") +class Second(common.StreamDefinition[str]): ... + + +@workflow.defn +class Blocks: + @workflow.run + async def run(self) -> None: + async with workflow.create_stream(First) as first: + await first.publish("token-0") + try: + async with workflow.create_stream(Second) as second: + await second.publish("token-1") + raise ApplicationError("turn failed") + except ApplicationError: + pass + + +async def test_async_with_closes_the_stream(stream_client: Client): + async with new_worker(stream_client, Blocks) as worker: + handle = await stream_client.start_workflow( + Blocks.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + first = await read(handle, First) + second = await read(handle, Second) + assert data(first) == ["token-0"] + assert isinstance(first[-1], common.Closed) and first[-1].failure is None + assert data(second) == ["token-1"] + assert isinstance(second[-1], common.Closed) + assert str(second[-1].failure) == "turn failed" + + +@workflow.defn +class Continuing: + @workflow.run + async def run(self, stream: common.WorkflowStream[str] | None) -> str: + if stream is None: + stream = await workflow.create_stream(Tokens) + await stream.publish("run 1") + workflow.continue_as_new(stream) + assert stream._payload_type is str + assert stream == workflow.stream_handle(Tokens) + await stream.publish("run 2") + await stream.close(result="done") + return "done" + + +@workflow.defn +class CreatesInEachRun: + @workflow.run + async def run(self, first: bool) -> str: + # The id has the first run id, so the run after continue-as-new gets the same Stream. + # The Stream exists then, so its create ignores the other Retention. + retention = timedelta(days=1 if first else 2) + stream = await workflow.create_stream(Tokens, retention=retention) + await stream.publish("run 1" if first else "run 2") + if first: + workflow.continue_as_new(False) + await stream.close() + return stream.id + + +async def test_create_in_the_next_run_gets_the_same_stream(stream_client: Client): + async with new_worker(stream_client, CreatesInEachRun) as worker: + handle = await stream_client.start_workflow( + CreatesInEachRun.run, + True, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + stream_id = await handle.result() + items = await read(handle) + assert stream_id == workflow_stream(handle, Tokens).id + assert data(items) == ["run 1", "run 2"] + assert isinstance(items[-1], common.Closed) + + +async def test_continue_as_new_passes_the_stream(stream_client: Client): + async with new_worker(stream_client, Continuing) as worker: + handle = await stream_client.start_workflow( + Continuing.run, None, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == "done" + items = await read(handle) + assert data(items) == ["run 1", "run 2"] + closed = [item for item in items if isinstance(item, common.Closed)] + assert len(closed) == 1 and closed[0].result == "done" + + +@workflow.defn +class Mistakes: + def __init__(self) -> None: + self._stream: common.WorkflowStream[str] | None = None + + @workflow.run + async def run(self) -> list[str]: + errors = [] + stream = await workflow.create_stream(Tokens) + self._stream = stream + await stream.close(result="early") + try: + await stream.publish("late") + except common.StreamClosedError as err: + errors.append(StreamErrorReason.Name(err.reason)) + try: + await workflow.create_stream(Tokens) + except common.StreamClosedError as err: + errors.append(StreamErrorReason.Name(err.reason)) + return errors + + @workflow.query + def publish_in_query(self) -> str: + assert self._stream is not None + try: + self._stream._append(("x",)) + except workflow.ReadOnlyContextError as err: + return type(err).__name__ + return "wrote" + + +async def test_replay_reproducible_errors_raise_in_workflow_code( + stream_client: Client, +): + async with new_worker(stream_client, Mistakes) as worker: + handle = await stream_client.start_workflow( + Mistakes.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == [ + "STREAM_ERROR_REASON_CLOSED", + "STREAM_ERROR_REASON_CLOSED", + ] + assert await handle.query(Mistakes.publish_in_query) == "ReadOnlyContextError" + items = await read(handle) + assert isinstance(items[-1], common.Closed) and items[-1].result == "early" + + +@workflow.defn +class ClosesAReceivedStream: + @workflow.run + async def run(self, stream: common.WorkflowStream[str]) -> str: + await stream.publish("from workflow") + try: + await stream.close(result="ended by a writer") + except common.StreamNotOwnerError as err: + return StreamErrorReason.Name(err.reason) + return "closed" + + +async def test_workflow_closes_only_its_own_streams(stream_client: Client): + stream = await stream_client.create_stream(Tokens, id=f"owned-{uuid.uuid4()}") + async with new_worker(stream_client, ClosesAReceivedStream) as worker: + assert ( + await stream_client.execute_workflow( + ClosesAReceivedStream.run, + common.WorkflowStream.from_handle(stream), + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + == "STREAM_ERROR_REASON_NOT_OWNER" + ) + # The Workflow that received the Stream can't close it, so it stays open. + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert data(items) == ["from workflow"] + + +@workflow.defn +class Failing: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream(Tokens) + await stream.publish("before") + raise ApplicationError("run failed", non_retryable=True) + + +async def test_run_end_closes_nothing(stream_client: Client): + async with new_worker(stream_client, Failing) as worker: + handle = await stream_client.start_workflow( + Failing.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + with pytest.raises(WorkflowFailureError): + await handle.result() + items = await read(handle) + # The Stream stays open, so a later reader still waits for a close. + assert data(items) == ["before"] + assert not any(isinstance(item, common.Closed) for item in items) + + +@workflow.defn +class FailsInItsBlockOnce: + @workflow.run + async def run(self) -> str: + attempt = workflow.info().attempt + async with workflow.create_stream(Tokens) as stream: + await stream.publish(f"run-{attempt}") + if attempt == 1: + raise ApplicationError("the first run fails") + return f"done on attempt {attempt}" + + +async def test_block_end_leaves_the_stream_open_for_a_retry_run(stream_client: Client): + async with new_worker(stream_client, FailsInItsBlockOnce) as worker: + handle = await stream_client.start_workflow( + FailsInItsBlockOnce.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + retry_policy=RetryPolicy( + initial_interval=timedelta(milliseconds=100), maximum_attempts=3 + ), + ) + assert await handle.result() == "done on attempt 2" + items = await read(handle) + written = [item for item in items if isinstance(item, common.Record)] + superseded = [item for item in items if isinstance(item, common.Superseded)] + assert [r.data for r in written] == ["run-1", "run-2"] + assert [r.data for r in written if superseded[0].withdraws(r)] == ["run-1"] + # Only the block end of the run that succeeds closes the Stream. + closes = [item for item in items if isinstance(item, common.Closed)] + assert len(closes) == 1 and items[-1] is closes[0] + assert closes[0].failure is None + + +@workflow.defn +class EndsItsBlock: + @workflow.run + async def run(self, mode: str) -> str: + try: + async with workflow.create_stream(Tokens) as stream: + await stream.publish(mode) + if mode == "non_retryable": + raise ApplicationError("block failed", non_retryable=True) + if mode == "listed_type": + raise ApplicationError("block failed", type="Listed") + if mode == "caught_task_failure": + raise ValueError("block failed") + raise ApplicationError("block failed") + except (ApplicationError, ValueError) as err: + if mode.startswith("caught"): + return "caught" + if mode == "then_non_retryable": + raise ApplicationError("run failed", non_retryable=True) from err + raise + + +async def test_block_end_closes_when_no_retry_run_follows(stream_client: Client): + # Each run ends with no retry run, so each block end closes its Stream, at once or at the + # run's end, with the block's exception. + retries = RetryPolicy( + initial_interval=timedelta(milliseconds=100), maximum_attempts=3 + ) + cases = { + "non_retryable": retries, + "last_attempt": RetryPolicy(maximum_attempts=1), + "listed_type": dataclasses.replace( + retries, non_retryable_error_types=["Listed"] + ), + "caught": retries, + "caught_task_failure": None, + "then_non_retryable": retries, + } + async with new_worker(stream_client, EndsItsBlock) as worker: + handles = { + mode: await stream_client.start_workflow( + EndsItsBlock.run, + mode, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + retry_policy=policy, + ) + for mode, policy in cases.items() + } + for mode, handle in handles.items(): + if mode.startswith("caught"): + assert await handle.result() == "caught" + else: + with pytest.raises(WorkflowFailureError): + await handle.result() + items = await read(handle) + assert data(items) == [mode] + closed = items[-1] + assert isinstance(closed, common.Closed), mode + assert closed.status == common.CloseStatus.FAILED + assert isinstance(closed.failure, ApplicationError) + assert closed.failure.message == "block failed" + # A close at the run's end is a command after the run's code, even on replay. + replayer = Replayer( + workflows=[EndsItsBlock], + stream_backends=stream_client.config()["stream_backends"], + ) + for mode in ("caught", "then_non_retryable"): + await replayer.replay_workflow(await handles[mode].fetch_history()) + + +# The Activity of a test's run sets this when it starts, so the test cancels the run then. +waiting: dict[str, asyncio.Event] = {} + + +@activity.defn +async def waits_for_cancel() -> None: + waiting[activity.info().workflow_id or ""].set() + while True: + activity.heartbeat() + await asyncio.sleep(0.1) + + +@workflow.defn +class CancelledInItsBlock: + @workflow.run + async def run(self, local: bool) -> None: + async with workflow.create_stream(Tokens) as stream: + await stream.publish("started") + if local: + await workflow.execute_local_activity( + waits_for_cancel, start_to_close_timeout=timedelta(minutes=1) + ) + else: + await workflow.execute_activity( + waits_for_cancel, + start_to_close_timeout=timedelta(minutes=1), + heartbeat_timeout=timedelta(seconds=1), + ) + + +async def test_cancel_during_an_activity_closes_as_canceled(stream_client: Client): + # The cancel reaches the block as an ActivityError, and a cancelled run gets no retry run. + async with new_worker( + stream_client, CancelledInItsBlock, activities=[waits_for_cancel] + ) as worker: + for local, policy in ((False, RetryPolicy(maximum_attempts=3)), (True, None)): + workflow_id = f"wf-{uuid.uuid4()}" + waiting[workflow_id] = asyncio.Event() + handle = await stream_client.start_workflow( + CancelledInItsBlock.run, + local, + id=workflow_id, + task_queue=worker.task_queue, + retry_policy=policy, + ) + await asyncio.wait_for(waiting[workflow_id].wait(), 20) + await handle.cancel() + with pytest.raises(WorkflowFailureError): + await handle.result() + items = await read(handle) + assert data(items) == ["started"] + closed = items[-1] + assert isinstance(closed, common.Closed), f"local={local}" + assert closed.status == common.CloseStatus.CANCELED + + +@workflow.defn +class CronBlock: + @workflow.run + async def run(self) -> str: + async with workflow.create_stream(Tokens) as stream: + await stream.publish("turn") + return "done" + + +async def test_block_end_leaves_a_cron_stream_open(stream_client: Client): + # The Stream spans the cron chain, so a close would fail the next run's create. + async with new_worker(stream_client, CronBlock) as worker: + handle = await stream_client.start_workflow( + CronBlock.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + cron_schedule="@every 1s", + ) + reader = workflow_stream(handle, Tokens).subscribe( + create_wait=timedelta(seconds=20) + ) + first = await asyncio.wait_for(anext(reader), 20) + second = await asyncio.wait_for(anext(reader), 20) + await reader.aclose() + await handle.terminate() + assert isinstance(first, common.Record) and isinstance(second, common.Record) + items = await read(handle) + assert not any(isinstance(item, common.Closed) for item in items) + + +failed_close_tasks: set[str] = set() + + +@workflow.defn +class SchedulesACloseThenFailsItsTask: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream(Tokens) + await stream.publish("token") + stream.close_nowait(result="closed") + workflow_id = workflow.info().workflow_id + if workflow_id not in failed_close_tasks: + failed_close_tasks.add(workflow_id) + raise RuntimeError("fail the Workflow Task that schedules the close") + + +async def test_failed_task_writes_no_close(stream_client: Client): + # The test's set must keep its value across the task's retry, so the Workflow has no + # sandbox. + async with new_worker( + stream_client, + SchedulesACloseThenFailsItsTask, + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await stream_client.start_workflow( + SchedulesACloseThenFailsItsTask.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = await read(handle) + closes = [item for item in items if isinstance(item, common.Closed)] + assert len(closes) == 1 and closes[0].result == "closed" + history = await handle.fetch_history() + # The failed task's close was never sent, so History has one close Activity. + scheduled = [ + event + for event in history.events + if event.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED + ] + assert [ + event.activity_task_scheduled_event_attributes.activity_type.name + for event in scheduled + ] == ["__temporal_streams_close"] + + +@workflow.defn +class CreatesOn: + @workflow.run + async def run(self, backend: str | None) -> None: + await workflow.create_stream(Tokens, backend=backend) + + +async def test_worker_config_errors_fail_the_task( + client: Client, stream_client: Client +): + # A Worker with no plugin, or with another backend, fails the Workflow Task, so a Worker + # with the right config can continue the run. + for task_client, backend, message in ( + (client, None, "No Stream backend"), + (stream_client, "elsewhere", "backend elsewhere"), + ): + async with new_worker(task_client, CreatesOn) as worker: + handle = await task_client.start_workflow( + CreatesOn.run, + backend, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await assert_task_fail_eventually(handle, message_contains=message) + await handle.terminate() + + +@workflow.defn +class SyncHandlers: + @workflow.init + def __init__(self, prompt: str) -> None: + # The create runs first, so the Stream exists for the signal's writes. + self._stream = workflow.create_stream_nowait(Tokens) + self._prompt = prompt + self._done = False + + @workflow.run + async def run(self, prompt: str) -> None: + assert prompt == self._prompt + await workflow.wait_condition(lambda: self._done) + + @workflow.signal + def token(self, text: str) -> None: + self._stream.publish_nowait(text, self._prompt) + + @workflow.update + def finish(self) -> str: + self._stream.close_nowait(result="finished") + try: + self._stream.publish_nowait("late") + except common.StreamClosedError as err: + return StreamErrorReason.Name(err.reason) + finally: + self._done = True + return "wrote" + + +async def test_sync_handlers_publish_and_close_with_no_await(stream_client: Client): + async with new_worker(stream_client, SyncHandlers) as worker: + handle = await stream_client.start_workflow( + SyncHandlers.run, "p", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.signal(SyncHandlers.token, "a") + assert ( + await handle.execute_update(SyncHandlers.finish) + == "STREAM_ERROR_REASON_CLOSED" + ) + await handle.result() + items = await read(handle) + assert data(items) == ["a", "p"] + assert isinstance(items[-1], common.Closed) and items[-1].result == "finished" + + +@workflow.defn +class CreatesInInit: + @workflow.init + def __init__(self, prompt: str) -> None: + self._stream = workflow.create_stream_nowait(Tokens) + self._stream.publish_nowait(prompt) + + @workflow.run + async def run(self, prompt: str) -> str: + await self._stream.publish(prompt.upper()) + await self._stream.close(result="done") + return "done" + + +async def test_init_creates_and_publishes_with_no_await(stream_client: Client): + async with new_worker(stream_client, CreatesInInit) as worker: + handle = await stream_client.start_workflow( + CreatesInInit.run, + "hi", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == "done" + items = await read(handle) + assert data(items) == ["hi", "HI"] + assert isinstance(items[-1], common.Closed) and items[-1].result == "done" + + +@workflow.defn +class PublishesWithNoAwait: + @workflow.run + async def run(self) -> list[int]: + tokens = await workflow.create_stream(Tokens) + first = tokens.publish("a") + second = tokens.publish("b") + last = await tokens.publish("c") + # Nobody awaits this write, and the run's end still waits for it. + tokens.publish("d") + return [first.result(), second.result(), last] + + +async def test_publish_starts_the_write_and_returns_its_future(stream_client: Client): + async with new_worker(stream_client, PublishesWithNoAwait) as worker: + handle = await stream_client.start_workflow( + PublishesWithNoAwait.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + offsets = await handle.result() + items = await read(handle) + written = [item for item in items if isinstance(item, common.Record)] + assert [r.data for r in written] == ["a", "b", "c", "d"] + assert offsets == [r.offset for r in written[:3]] + + +@workflow.defn +class LeavesAFailedPublish: + @workflow.run + async def run(self, stream_id: str, retrieve: bool) -> str: + # Nobody created the Stream, so the backend refuses the write. + missing = common.WorkflowStream[str](id=stream_id) + future = missing.publish("lost") + errors: list[BaseException | None] = [] + if retrieve: + future.add_done_callback(lambda done: errors.append(done.exception())) + await workflow.wait_condition(lambda: bool(errors)) + return "done" + + +async def test_publish_error_that_no_caller_gets_raises_at_the_run_end( + stream_client: Client, caplog: pytest.LogCaptureFixture +): + # The error raises as if the run didn't catch it, as with publish_nowait. + lost_id = f"missing-{uuid.uuid4()}" + with caplog.at_level(logging.WARNING, logger="temporalio.workflow"): + async with new_worker(stream_client, LeavesAFailedPublish) as worker: + lost = await stream_client.start_workflow( + LeavesAFailedPublish.run, + args=[lost_id, False], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + with pytest.raises(WorkflowFailureError) as failed: + await lost.result() + assert isinstance(failed.value.cause, ApplicationError) + assert failed.value.cause.type == "StreamNotFoundError" + # A done callback that reads the error owns it, so the run completes. + owned = await stream_client.start_workflow( + LeavesAFailedPublish.run, + args=[f"missing-{uuid.uuid4()}", True], + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await owned.result() == "done" + # The Workflow's logger logs the error that no caller got when the write fails. + warnings = [ + r.getMessage() for r in caplog.records if "no caller has" in r.getMessage() + ] + assert len(warnings) == 1 and lost_id in warnings[0] + + +@workflow.defn +class Gathers: + @workflow.run + async def run(self) -> list[int]: + tokens = await workflow.create_stream(Tokens) + tokens.publish_nowait("0") + # The first write starts, so the next ones wait behind it. + await asyncio.sleep(0) + return list( + await asyncio.gather(*(tokens.publish(str(n)) for n in range(1, 5))) + ) + + +async def test_concurrent_publishes_keep_their_order(stream_client: Client): + async with new_worker(stream_client, Gathers) as worker: + handle = await stream_client.start_workflow( + Gathers.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + offsets = await handle.result() + items = await read(handle) + assert data(items) == ["0", "1", "2", "3", "4"] + # The four writes that waited were sent as one Local Activity, and each call got the offset + # of its own Record. + assert ( + offsets + == [item.offset for item in items if isinstance(item, common.Record)][1:] + ) + history = await handle.fetch_history() + markers = [ + e + for e in history.events + if e.event_type == EventType.EVENT_TYPE_MARKER_RECORDED + ] + # The create, the first write and the four that waited. + assert len(markers) == 3 + # Replay uses the markers and writes nothing. + await Replayer( + workflows=[Gathers], stream_backends=stream_client.config()["stream_backends"] + ).replay_workflow(history) + assert data(await read(handle)) == ["0", "1", "2", "3", "4"] + + +class _RandomCodec(PayloadCodec): + # Adds random bytes to each payload, as the random IV of an encryption codec does. + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + return [ + Payload( + metadata={"encoding": b"binary/random"}, + data=os.urandom(8) + p.SerializeToString(), + ) + for p in payloads + ] + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + return [ + _inner(p.data[8:]) if _encoded(p, b"binary/random") else p for p in payloads + ] + + +class _GrowingCodec(PayloadCodec): + # Triples each payload, so a merge under the Workflow's cap exceeds the Store's batch limit + # after the codec. + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + return [ + Payload( + metadata={ + "encoding": b"binary/grown", + "size": str(len(p.SerializeToString())).encode(), + }, + data=p.SerializeToString() * 3, + ) + for p in payloads + ] + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + return [ + _inner(p.data[: int(p.metadata["size"])]) + if _encoded(p, b"binary/grown") + else p + for p in payloads + ] + + +def _encoded(payload: Payload, encoding: bytes) -> bool: + return payload.metadata.get("encoding") == encoding + + +def _inner(data: bytes) -> Payload: + inner = Payload() + inner.ParseFromString(data) + return inner + + +def with_codec(client: Client, codec: PayloadCodec) -> Client: + config = client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, payload_codec=codec + ) + return Client(**config) + + +async def test_retried_write_goes_in_once( + stream_client: Client, monkeypatch: pytest.MonkeyPatch +): + original = _ExtensionCalls.call_async + sent: list[list[bytes]] = [] + + async def loses_the_first_answer( + self: Any, method: str, request: Any, response: Any + ) -> Any: + answer = await original(self, method, request, response) + if method == "AppendRecords": + sent.append([record.body.SerializeToString() for record in request.records]) + if len(sent) == 1: + raise RuntimeError("the answer is lost") + return answer + + monkeypatch.setattr(_ExtensionCalls, "call_async", loses_the_first_answer) + client = with_codec(stream_client, _RandomCodec()) + async with new_worker(client, Gathers) as worker: + handle = await client.start_workflow( + Gathers.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + items = await read(handle) + assert data(items) == ["0", "1", "2", "3", "4"] + # The retry sends the bytes of the first attempt, so the Store's dedupe matches it as the + # same batch, even though the codec adds random bytes. + assert sent[0] == sent[1] + + +@workflow.defn +class PublishesToAMissingStream: + @workflow.run + async def run(self) -> None: + # No create, so the Store refuses the append, and no retry sends it again. + stream = workflow.stream_handle(Tokens) + await stream.publish("x" * 100_000) + + +async def test_refused_write_keeps_no_pending_bytes(stream_client: Client): + async with new_worker(stream_client, PublishesToAMissingStream) as worker: + handle = await stream_client.start_workflow( + PublishesToAMissingStream.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + with pytest.raises(WorkflowFailureError): + await handle.result() + stream_id = workflow_stream(handle, Tokens).id + assert not [key for key in _writes._PENDING._entries if key[1] == stream_id] + + +def _entry(size: int) -> Any: + return _writes._Encoded([], size) + + +def test_pending_appends_keep_the_newest_epoch_and_a_bounded_size(): + pending = _writes._Pending(100) + one, two = ("ns", "s", "la:1"), ("ns", "s", "la:2") + first, retried = _entry(10), _entry(10) + pending.put(one, (5, 1, 0, 0), first) + # A rerun under a later epoch replaces the earlier one. + pending.put(one, (9, 1, 0, 0), retried) + assert pending.get(one, (5, 1, 0, 0)) is None + assert pending.get(one, (9, 1, 0, 0)) is retried + # A write of an earlier epoch is fenced, so its put and its pop don't change the later entry. + pending.put(one, (5, 1, 0, 0), first) + pending.pop(one, (5, 1, 0, 0)) + assert pending.get(one, (9, 1, 0, 0)) is retried + # Past the size limit, the oldest entry is evicted. + pending.put(two, (9, 1, 0, 0), _entry(95)) + assert pending.get(one, (9, 1, 0, 0)) is None + assert pending.get(two, (9, 1, 0, 0)) is not None + + +@workflow.defn +class PublishesLargeValues: + @workflow.run + async def run(self, size: int) -> list[int]: + tokens = await workflow.create_stream(Tokens) + return list( + await asyncio.gather( + tokens.publish(*["a" * size] * 3), tokens.publish(*["b" * size] * 3) + ) + ) + + +async def test_merge_that_the_codec_grows_past_the_batch_limit(stream_client: Client): + # The two publishes merge under the Workflow's cap. The codec makes the merge larger than the + # batch limit, so each publish is sent as its own batch. + client = with_codec(stream_client, _GrowingCodec()) + async with new_worker(client, PublishesLargeValues) as worker: + handle = await client.start_workflow( + PublishesLargeValues.run, + 300_000, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + offsets = await handle.result() + items = await read(handle) + written = [item for item in items if isinstance(item, common.Record)] + assert [len(r.data) for r in written] == [300_000] * 6 + assert offsets == [written[2].offset, written[5].offset] + + +@workflow.defn +class PublishesTooMuch: + @workflow.run + async def run(self) -> str: + tokens = await workflow.create_stream(Tokens) + try: + await tokens.publish(*["x" * 1_400_000] * 4) + except common.StreamTooLargeError as err: + return StreamErrorReason.Name(err.reason) + return "written" + + +async def test_batch_over_the_store_limit_raises_too_large(stream_client: Client): + # Each Record is under the Store's Record limit, and the publish's batch is over 4 MiB. + async with new_worker(stream_client, PublishesTooMuch) as worker: + handle = await stream_client.start_workflow( + PublishesTooMuch.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == "STREAM_ERROR_REASON_TOO_LARGE" + assert data(await read(handle)) == [] + + +@workflow.defn +class CreatesOnTheDefault: + @workflow.run + async def run(self) -> None: + await workflow.create_stream(Tokens) + + +@workflow.defn +class CreatesOnANamedBackend: + @workflow.run + async def run(self, backend: str) -> str: + stream = await workflow.create_stream(Tokens, backend=backend) + await stream.publish(f"on {backend}") + await workflow.stream_handle(Tokens).close(result="closed") + return stream.backend or "" + + +@workflow.defn +class CreatesAfterAnIdleTime: + @workflow.run + async def run(self) -> str: + await workflow.create_stream(Tokens) + await workflow.sleep(timedelta(seconds=1.5)) + try: + await workflow.create_stream(Tokens) + except common.StreamExpiredError as err: + return StreamErrorReason.Name(err.reason) + return "created" + + +async def test_create_of_an_expired_workflow_stream_raises(client: Client): + # An open Workflow Stream expires one safety TTL after its last write. + config = client.config() + config["stream_backends"] = [ + MemoryStreamBackend(default=True, safety_ttl=timedelta(milliseconds=500)) + ] + expiring = Client(**config) + async with new_worker(expiring, CreatesAfterAnIdleTime) as worker: + assert ( + await expiring.execute_workflow( + CreatesAfterAnIdleTime.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + == "STREAM_ERROR_REASON_EXPIRED" + ) + + +async def test_default_is_native_with_no_default_backend(client: Client): + # With no backend marked default=True, the default is native, which this SDK lacks. A create + # that names the registered backend still works. + config = client.config() + config["stream_backends"] = [MemoryStreamBackend()] + memory_client = Client(**config) + async with new_worker( + memory_client, CreatesOnTheDefault, CreatesOnANamedBackend + ) as worker: + handle = await memory_client.start_workflow( + CreatesOnTheDefault.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await assert_task_fail_eventually(handle, message_contains="default is native") + await handle.terminate() + named = await memory_client.start_workflow( + CreatesOnANamedBackend.run, + "memory", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await named.result() == "memory" + stream = workflow_stream(named, Tokens, backend="memory") + items = [item async for item in stream.subscribe()] + assert data(items) == ["on memory"] + assert isinstance(items[-1], common.Closed) and items[-1].result == "closed" + + +async def _event_types(handle: WorkflowHandle[Any, Any]) -> list[int]: + return [event.event_type async for event in handle.fetch_history_events()] + + +async def _close_is_scheduled(handle: WorkflowHandle[Any, Any]) -> None: + assert EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED in await _event_types(handle) + + +async def _cancel_is_handled(handle: WorkflowHandle[Any, Any]) -> None: + types = await _event_types(handle) + requested = EventType.EVENT_TYPE_WORKFLOW_EXECUTION_CANCEL_REQUESTED + assert requested in types + assert ( + EventType.EVENT_TYPE_WORKFLOW_TASK_COMPLETED in types[types.index(requested) :] + ) + + +@workflow.defn +class ClosesThenReturns: + @workflow.run + async def run(self) -> str: + stream = await workflow.create_stream(Tokens) + await stream.publish("a") + stream.close_nowait(result="bye") + return "done" + + +async def test_cancel_during_the_run_end_waits_for_the_close(stream_client: Client): + # The first Worker polls no Activities, so the run waits at its end for the close Activity. + async with new_worker( + stream_client, ClosesThenReturns, no_remote_activities=True + ) as worker: + handle = await stream_client.start_workflow( + ClosesThenReturns.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await assert_eventually(lambda: _close_is_scheduled(handle)) + await handle.cancel() + await assert_eventually(lambda: _cancel_is_handled(handle)) + async with new_worker(stream_client, task_queue=worker.task_queue): + with pytest.raises(WorkflowFailureError) as failed: + await handle.result() + assert isinstance(failed.value.cause, CancelledError) + types = await _event_types(handle) + assert EventType.EVENT_TYPE_ACTIVITY_TASK_CANCEL_REQUESTED not in types + items = await read(handle) + assert data(items) == ["a"] + assert isinstance(items[-1], common.Closed) and items[-1].result == "bye" + + +@common.defn(name="late") +class Late(common.StreamDefinition[str]): ... + + +class _SlowCodec(PayloadCodec): + # Delays the write of one value past the Workflow Task's heartbeat. + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + if any(b"slow" in payload.data for payload in payloads): + await asyncio.sleep(3) + return list(payloads) + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + return list(payloads) + + +@workflow.defn +class WritesFromASignalAtTheEnd: + def __init__(self) -> None: + self.late: common.WorkflowStream[str] | None = None + + @workflow.run + async def run(self) -> str: + self.late = await workflow.create_stream(Late) + tokens = await workflow.create_stream(Tokens) + tokens.close_nowait() + return "done" + + @workflow.signal + def write(self) -> None: + assert self.late is not None + self.late.publish_nowait("slow") + + +async def _signal_is_handled(handle: WorkflowHandle[Any, Any]) -> None: + types = await _event_types(handle) + signaled = EventType.EVENT_TYPE_WORKFLOW_EXECUTION_SIGNALED + assert signaled in types + assert EventType.EVENT_TYPE_WORKFLOW_TASK_STARTED in types[types.index(signaled) :] + + +async def test_run_end_waits_for_a_write_that_starts_while_it_waits( + stream_client: Client, +): + config = stream_client.config() + config["data_converter"] = DataConverter(payload_codec=_SlowCodec()) + slow = Client(**config) + async with new_worker( + slow, WritesFromASignalAtTheEnd, no_remote_activities=True + ) as worker: + # A short task timeout makes Core complete the task before the slow write ends. + handle = await slow.start_workflow( + WritesFromASignalAtTheEnd.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + task_timeout=timedelta(seconds=1), + ) + await assert_eventually(lambda: _close_is_scheduled(handle)) + await handle.signal(WritesFromASignalAtTheEnd.write) + await assert_eventually(lambda: _signal_is_handled(handle)) + async with new_worker(slow, task_queue=worker.task_queue): + assert await handle.result() == "done" + assert data(await read(handle, Late)) == ["slow"] + + +class _TagsWithContext(DefaultPayloadConverter, WithSerializationContext): + # Puts the Workflow id of the serialization context in front of each str value. + def __init__(self, workflow_id: str = "") -> None: + super().__init__() + self._workflow_id = workflow_id + + def with_context(self, context: SerializationContext) -> "_TagsWithContext": + if isinstance(context, WorkflowSerializationContext): + return _TagsWithContext(context.workflow_id) + return self + + def to_payloads(self, values: Sequence[Any]) -> list[Payload]: + if self._workflow_id: + values = [ + f"{self._workflow_id}:{v}" if isinstance(v, str) else v for v in values + ] + return super().to_payloads(values) + + +@workflow.defn +class PublishesOneValue: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream(Tokens) + await stream.publish("plain") + await stream.close() + + +async def test_workflow_records_get_no_serialization_context(stream_client: Client): + # Activities and clients convert Records with no context, so Workflows do too. + config = stream_client.config() + config["data_converter"] = DataConverter(payload_converter_class=_TagsWithContext) + tagging = Client(**config) + async with new_worker(tagging, PublishesOneValue) as worker: + handle = await tagging.start_workflow( + PublishesOneValue.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + assert data(await read(handle)) == ["plain"]