From e7710c080588c90d82d38dd51f61d71668a1047d Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 25 Oct 2026 15:01:50 -0700 Subject: [PATCH 1/4] Added Workflow Streams, which write through Local Activities. Each write has the epoch of the Workflow Task that runs it, so the Store withdraws the writes of a failed task. A close is a regular Activity, which runs only after its Workflow Task commits. --- temporalio/common/__init__.py | 3 + temporalio/common/_streams/__init__.py | 2 + temporalio/common/_streams/_activity.py | 71 + temporalio/common/_streams/_extension.py | 410 +++++- temporalio/common/_streams/_workflow.py | 755 +++++++++++ temporalio/common/_streams/_writes.py | 298 +++++ temporalio/contrib/stream/__init__.py | 14 +- temporalio/contrib/stream/testing.py | 3 +- temporalio/workflow/__init__.py | 8 + temporalio/workflow/_streams.py | 122 ++ tests/contrib/streams/_support.py | 23 +- tests/contrib/streams/test_redis.py | 66 + tests/contrib/streams/test_stream_client.py | 100 ++ tests/contrib/streams/test_workflow_module.py | 1154 +++++++++++++++++ 14 files changed, 3022 insertions(+), 7 deletions(-) create mode 100644 temporalio/common/_streams/_activity.py create mode 100644 temporalio/common/_streams/_workflow.py create mode 100644 temporalio/common/_streams/_writes.py create mode 100644 temporalio/workflow/_streams.py create mode 100644 tests/contrib/streams/test_redis.py create mode 100644 tests/contrib/streams/test_workflow_module.py diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index 1bcd2f0080..ba622ecb55 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 8e4785be29..4c828426fd 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 0000000000..d6e87b288c --- /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 dbb68c82d1..d7b71dfd61 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 0000000000..08634e1fa4 --- /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 0000000000..4eb9efa7dd --- /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 4e8c7fcd20..19e24062e3 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 8d393ae59c..79a1beb7b2 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 c631fd88e8..ef1c20a37c 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 0000000000..c220c869a5 --- /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 9f1a05f106..27e0a95975 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 0000000000..b1ec38ae29 --- /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 ce1d1bf3e9..7d3125174b 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 0000000000..ca599fdb0b --- /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"] From 26d16f3e741ce9ab36c4d6a0a1f8b02de65d2879 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sat, 31 Oct 2026 03:31:28 -0700 Subject: [PATCH 2/4] Moved the Workflow writes to the new bindings, and renamed the handles. The handles make their writes with the generated Workflow API, and the sync calls use its start forms, so each write keeps its place at the call. The Worker's interceptor runs those calls as the same write Local Activities, so History doesn't change. A closed or expired Stream fails the create, and a close sends only its failure, which sets the close status. --- temporalio/common/_streams/_extension.py | 63 ++++- temporalio/common/_streams/_workflow.py | 259 ++++++++++++++---- temporalio/common/_streams/_writes.py | 68 ++--- temporalio/nexus/system/_support.py | 86 ++++-- tests/contrib/streams/test_stream_client.py | 2 +- tests/contrib/streams/test_workflow_module.py | 50 +++- 6 files changed, 401 insertions(+), 127 deletions(-) diff --git a/temporalio/common/_streams/_extension.py b/temporalio/common/_streams/_extension.py index d7b71dfd61..82a14a0a77 100644 --- a/temporalio/common/_streams/_extension.py +++ b/temporalio/common/_streams/_extension.py @@ -8,7 +8,7 @@ import weakref from collections.abc import Mapping, Sequence from datetime import timedelta -from typing import Any, NoReturn +from typing import Any, NoReturn, TypeVar import temporalio.activity import temporalio.api.common.v1 @@ -30,11 +30,15 @@ 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 _system_nexus_interceptor as system_nexus from temporalio.common._streams.system import models from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner logger = logging.getLogger(__name__) +_InputT = TypeVar("_InputT") +_OutputT = TypeVar("_OutputT") + # 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" @@ -326,10 +330,10 @@ def _ends_run(err: BaseException) -> bool: ) and temporalio.workflow._Runtime.current().workflow_is_failure_exception(err) -def _owned_ref(stream_id: str) -> models.StreamRef: +def _owned_handle(stream_id: str) -> models.StreamHandle: # The run chain owns each Stream that it creates. info = temporalio.workflow.info() - return models.StreamRef( + return models.StreamHandle( id=stream_id, owner=workflow_owner(info.workflow_id, info.first_execution_run_id), ) @@ -345,7 +349,7 @@ def _take_owned_streams( 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 + state.stream(_owned_handle(stream_id)).backend = backend class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): @@ -408,6 +412,57 @@ def start_local_activity( input.headers = {**input.headers, **self._activity_headers()} return super().start_local_activity(input) + async def start_system_nexus_operation( + self, input: temporalio.worker.StartNexusOperationInput[_InputT, _OutputT] + ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: + # The Stream service calls of the generated Workflow API run as the write Local + # Activities, so History has a marker for each write and no Nexus event. + return await system_nexus._start_system_nexus_operation( + _StreamServiceCalls(self.next), input + ) + + +class _StreamServiceCalls(system_nexus._SystemNexusWorkflowOutboundInterceptorBase): + # Gives each Stream service call of the generated dispatch to the Workflow's write path, and + # any other System Nexus call to the next interceptor. + def __init__(self, next: temporalio.worker.WorkflowOutboundInterceptor) -> None: + self._next = next + + def _next_system_nexus_interceptor( + self, + ) -> system_nexus._SystemNexusWorkflowOutboundInterceptorBase: + return self + + async def start_system_nexus_operation( + self, input: temporalio.worker.StartNexusOperationInput[_InputT, _OutputT] + ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: + return await self._next.start_system_nexus_operation(input) + + async def start_create_stream( + self, request: models.CreateStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: + return streams_workflow.run_create_stream(request) + + async def start_append_records( + self, request: models.AppendRecordsRequest[Any] + ) -> temporalio.workflow.NexusOperationHandle[models.AppendRecordsResponse]: + return streams_workflow.run_append_records(request) + + async def start_close_stream( + self, request: models.CloseStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.CloseStreamResponse]: + return streams_workflow.run_close_stream(request) + + async def start_read_records( + self, request: models.ReadRecordsRequest + ) -> temporalio.workflow.NexusOperationHandle[models.ReadRecordsResponse[Any]]: + raise streams_workflow.no_workflow_reads() + + async def start_describe_stream( + self, request: models.DescribeStreamRequest + ) -> temporalio.workflow.NexusOperationHandle[models.DescribeStreamResponse]: + raise streams_workflow.no_workflow_reads() + class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor): def __init__( diff --git a/temporalio/common/_streams/_workflow.py b/temporalio/common/_streams/_workflow.py index 08634e1fa4..1e2d3ba912 100644 --- a/temporalio/common/_streams/_workflow.py +++ b/temporalio/common/_streams/_workflow.py @@ -16,7 +16,7 @@ from __future__ import annotations import asyncio -from collections.abc import Generator +from collections.abc import Callable, Generator from dataclasses import dataclass, field from datetime import timedelta from types import TracebackType @@ -29,7 +29,6 @@ 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 ( @@ -39,10 +38,16 @@ exit_failure, workflow_owner, ) -from temporalio.common._streams._items import close_status +from temporalio.common._streams._items import close_failure from temporalio.common._streams.system import models +from temporalio.common._streams.system._resources import Stream as GeneratedStream +from temporalio.common._streams.system._resources import stream as generated_stream +from temporalio.common._streams.system.operations import ( + create_stream as generated_create, +) T = TypeVar("T") +OutputT = TypeVar("OutputT") # 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. @@ -59,7 +64,7 @@ @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 + handle: models.StreamHandle closed: bool = False # Set for a Stream that the run chain created, which continue-as-new passes on. backend: str | None = None @@ -83,7 +88,7 @@ def take(self, request: google.protobuf.message.Message) -> bool: mine = self.request if ( not isinstance(mine, stream_pb.AppendRecordsRequest) - or mine.stream != request.stream + or mine.handle != request.handle ): return False size = sum(record.ByteSize() for record in request.records) @@ -114,12 +119,15 @@ class _RunState: # 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) + # The calls of the generated Workflow API in flight. A call ends after its write, so the + # run's end waits for both. + calls: set[asyncio.Future[Any]] = field(default_factory=set) - def stream(self, ref: models.StreamRef) -> _RunStream: - known = self.streams.get(ref.id) + def stream(self, handle: models.StreamHandle) -> _RunStream: + known = self.streams.get(handle.id) if known is None: - known = _RunStream(ref) - self.streams[ref.id] = known + known = _RunStream(handle) + self.streams[handle.id] = known return known @@ -341,7 +349,9 @@ async def drain() -> None: 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()]: + while pending := [ + task for task in (*state.tails.values(), *state.calls) if not task.done() + ]: for task in pending: try: await asyncio.shield(task) @@ -394,15 +404,15 @@ class WorkflowStream(StreamHandle[T]): This API is experimental. """ - def _model_ref(self) -> models.StreamRef: - return models.StreamRef( + def _model_handle(self) -> models.StreamHandle: + return models.StreamHandle( 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 _proto_handle(self) -> stream_v1.StreamHandle: + return models._StreamHandleTransferTypeConverter().to_transfer_type( + self._model_handle() ) def publish(self, *values: T) -> asyncio.Future[int]: @@ -433,7 +443,7 @@ def publish(self, *values: T) -> asyncio.Future[int]: It fails the Workflow Task, not the run. """ state = _state() - future = self._append(values) + future = self._start_append(values) loop = asyncio.get_running_loop() def warn(future: asyncio.Future[int]) -> None: @@ -463,7 +473,18 @@ def publish_nowait(self, *values: T) -> None: 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)) + _write_nowait(self._start_append(values)) + + def _start_append(self, values: tuple[Any, ...]) -> asyncio.Future[int]: + # The checks raise at the call, as the write's place in the order does. + if not values: + raise TypeError("publish takes one or more values") + self._check("write to") + call = generated_stream.start_append_records( + handle=self._generated_handle(), + records=[models.AppendRecord(body=value) for value in values], + ) + return _call_future(call, lambda response: response.last_offset or 0) def _check(self, action: str, *, backend_checked: bool = False) -> _RunStream: if temporalio.workflow.unsafe.is_read_only(): @@ -472,11 +493,10 @@ def _check(self, action: str, *, backend_checked: bool = False) -> _RunStream: ) state = _state() state.uses_streams = True - run_stream = state.stream(self._model_ref()) + run_stream = state.stream(self._model_handle()) 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, + f"Can't {action} the Stream {self.id}, since this run closed it" ) if not backend_checked: _check_backend(self.backend, action) @@ -490,7 +510,9 @@ def _append(self, values: tuple[Any, ...]) -> asyncio.Future[int]: return _write( self.id, APPEND_ACTIVITY, - stream_pb.AppendRecordsRequest(stream=self._ref(), records=records), + stream_pb.AppendRecordsRequest( + handle=self._proto_handle(), records=records + ), ) async def close( @@ -514,7 +536,19 @@ async def close( Raises: StreamNotOwnerError: This run chain didn't create the Stream. """ - await self._close(result, failure) + await self._generated().close(result=result, failure=failure) + + def _generated(self) -> GeneratedStream[Any]: + # The generated Workflow API makes the call, and the Worker's interceptor runs it as the + # write Local Activity. + return GeneratedStream(self._generated_handle()) + + def _generated_handle(self) -> models.StreamHandle: + return models.StreamHandle( + id=self.id, + owner=workflow_owner(self.workflow_id, self.first_execution_run_id), + backend=self.backend or None, + ) def close_nowait( self, *, result: Any | None = None, failure: BaseException | None = None @@ -525,7 +559,16 @@ def close_nowait( 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)) + _write_nowait(self._start_close(result, failure)) + + def _start_close( + self, result: Any, failure: BaseException | None + ) -> asyncio.Future[int]: + self._closing() + call = generated_stream.start_close_stream( + handle=self._generated_handle(), result=result, failure=failure + ) + return _call_future(call, lambda _: 0) def _closing(self) -> _RunStream: info = temporalio.workflow.info() @@ -536,22 +579,19 @@ def _closing(self) -> _RunStream: # 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, + f"Can't close the Stream {self.id}, since this run chain didn't create it" ) 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()) + request = stream_pb.CloseStreamRequest(handle=self._proto_handle()) if result is not None: request.result.CopyFrom(_payload(result)) + failure = close_failure(failure) 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]: @@ -583,7 +623,7 @@ async def __aexit__( Stream. """ state = _state() - if state.stream(self._model_ref()).closed: + if state.stream(self._model_handle()).closed: return close, failure = exit_failure(exc) if not close or temporalio.workflow.info().cron_schedule: @@ -619,8 +659,8 @@ def close_at_run_end(err: BaseException | None) -> None: ): return for stream, failure in state.deferred_closes: - if not state.stream(stream._model_ref()).closed: - _write_nowait(stream._close(None, failure)) + if not state.stream(stream._model_handle()).closed: + _write_nowait(stream._start_close(None, failure)) def _retried(exc: BaseException) -> bool: @@ -689,8 +729,10 @@ def create_stream( """ async def create() -> WorkflowStream[T]: - stream, task = _create(definition, options, backend) - await task + stream = _creating(definition, backend) + await generated_create.create_stream( + definition, retention=options.get("retention"), backend=stream.backend + ) return stream return Creating(create) @@ -706,35 +748,146 @@ def create_stream_nowait( See :py:func:`temporalio.workflow.create_stream_nowait`. """ - stream, task = _create(definition, options, backend) - _write_nowait(task) + stream = _creating(definition, backend) + call = generated_create.start_create_stream( + definition, retention=options.get("retention"), backend=stream.backend + ) + _write_nowait(_call_future(call, lambda _: 0)) return stream -def _create( - definition: type[StreamDefinition[T]], - options: StreamOptions, - backend: str | None, -) -> tuple[WorkflowStream[T], asyncio.Future[int]]: +def _creating( + definition: type[StreamDefinition[T]], backend: str | None +) -> WorkflowStream[T]: + # The checks raise before the generated call, as they did before the write. stream = _own_handle(definition, _check_backend(backend, "create")) + stream._check("create", backend_checked=True) + return stream + + +def _call_future( + call: asyncio.Task[Any], result: Callable[[Any], int] +) -> asyncio.Future[int]: + # A call of the generated Workflow API starts its operation in a task, and the operation's + # handle gives the write's result. The caller's future follows both, and the run's end + # waits for them. + loop = asyncio.get_running_loop() + future = _WriteFuture(loop=loop) + + async def follow() -> None: + # The task sets the future before it ends, so the run's end sees a failed publish. + try: + started = await call + if isinstance(started, temporalio.workflow.NexusOperationHandle): + started = await started + except asyncio.CancelledError: + future.cancel() + raise + except BaseException as err: + future.set_exception(err) + else: + future.set_result(result(started)) + + task = loop.create_task(follow()) + calls = _state().calls + calls.add(task) + task.add_done_callback(calls.discard) + return future + + +def _start_create( + stream: WorkflowStream[Any], retention: timedelta | None, owner: stream_v1.Producer +) -> asyncio.Future[int]: 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(), + stream_id=stream.id, owner=owner, stream_handle=stream._proto_handle() ) - 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) + return _write(stream.id, CREATE_ACTIVITY, request) + + +class _WriteHandle(temporalio.workflow.NexusOperationHandle[OutputT]): + # The interceptor runs a call of the generated Workflow API as a write Local Activity, so + # the call's handle waits for the write. + def __init__( + self, write: asyncio.Future[int], result: Callable[[int], OutputT] + ) -> None: + self._write = write + self._result = result + + def cancel(self) -> bool: + # A write keeps its place in the Stream's order, so it can't be cancelled. + return False + + def __await__(self) -> Generator[Any, Any, OutputT]: + return self._wait().__await__() + + async def _wait(self) -> OutputT: + return self._result(await self._write) + + @property + def operation_token(self) -> str | None: + return None + + +def _from_handle(handle: models.StreamHandle) -> WorkflowStream[Any]: + owner = handle.owner + return WorkflowStream( + id=handle.id, + workflow_id=owner.business_id if owner else None, + first_execution_run_id=owner.run_id if owner else None, + backend=handle.backend, + ) + + +# The Worker's interceptor runs the Stream calls of the generated Workflow API with these. They +# build the same Local Activity requests as the handles, so History doesn't change. A retry run's +# withdrawal comes from _write, once for each Stream, and the write Activity sends its own +# Producer. + + +def run_create_stream( + request: models.CreateStreamRequest, +) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: + """Run a create of the generated Workflow API as the create Local Activity.""" + handle = request.stream_handle or models.StreamHandle(id=request.id) + stream = _from_handle(handle) + stream._backend = _check_backend(handle.backend, "create") + owner = ( + models._ProducerTransferTypeConverter().to_transfer_type(request.owner) + if request.owner is not None + else stream_v1.Producer() + ) + write = _start_create(stream, request.retention, owner) + return _WriteHandle(write, lambda _: models.CreateStreamResult(handle=handle)) + + +def run_append_records( + request: models.AppendRecordsRequest[Any], +) -> temporalio.workflow.NexusOperationHandle[models.AppendRecordsResponse]: + """Run an append of the generated Workflow API as the append Local Activity.""" + stream = _from_handle(request.handle) + write = stream._append(tuple(record.body for record in request.records or [])) + return _WriteHandle( + write, lambda offset: models.AppendRecordsResponse(last_offset=offset) + ) + + +def run_close_stream( + request: models.CloseStreamRequest, +) -> temporalio.workflow.NexusOperationHandle[models.CloseStreamResponse]: + """Run a close of the generated Workflow API as the close write.""" + write = _from_handle(request.handle)._close(request.result, request.failure) + return _WriteHandle(write, lambda _: models.CloseStreamResponse()) + + +def no_workflow_reads() -> _errors.StreamUnsupportedError: + """The error of a read or a describe in Workflow code.""" + return _errors.StreamUnsupportedError( + "Workflow code can't read a Stream. Read it from an Activity." + ) def stream_handle( diff --git a/temporalio/common/_streams/_writes.py b/temporalio/common/_streams/_writes.py index 4eb9efa7dd..3e3f4b725e 100644 --- a/temporalio/common/_streams/_writes.py +++ b/temporalio/common/_streams/_writes.py @@ -15,7 +15,6 @@ 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, @@ -71,9 +70,9 @@ def _history(streams: _activity.ActivityStreams) -> dict[str, Any]: } -def _ref(ref: stream_v1.StreamRef) -> models.StreamRef: - return models._StreamRefTransferTypeConverter().from_transfer_type( - ref, models.StreamRef +def _model_handle(handle: stream_v1.StreamHandle) -> models.StreamHandle: + return models._StreamHandleTransferTypeConverter().from_transfer_type( + handle, models.StreamHandle ) @@ -81,29 +80,24 @@ def _ref(ref: stream_v1.StreamRef) -> models.StreamRef: 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, + # The generated create takes its owner from a source, so the Activity sends the Workflow's + # request as it is. A closed, expired or another owner's Stream fails the create. + await service._stream_service_create_stream( + models.CreateStreamRequest( + 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, + ), + stream_handle=_model_handle(request.stream_handle), ), - ref=_ref(request.ref) if request.HasField("ref") else None, + temporalio.common.RawValue, ) - 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 @@ -204,7 +198,7 @@ async def append_records( streams, service = _context() producer = _producer(streams) producer_pb = models._ProducerTransferTypeConverter().to_transfer_type(producer) - key = (streams.namespace, request.stream.id, producer_pb.id) + key = (streams.namespace, request.handle.id, producer_pb.id) task = producer_pb.epoch epoch = ( task.task_started_event_id, @@ -225,9 +219,9 @@ async def append_records( answer = encoded.answers.get(index) if answer is None: try: - response = await ( - await service.append_records( - stream=_ref(request.stream), + response = await service._stream_service_append_records( + models.AppendRecordsRequest( + handle=_model_handle(request.handle), producer=producer, sequence=index + 1, records=[ @@ -244,11 +238,11 @@ async def append_records( ) answer = response.last_offset or 0 except _errors.StreamError as err: - if err.reason != StreamErrorReason.STREAM_ERROR_REASON_TOO_LARGE: + if not isinstance(err, _errors.StreamTooLargeError): if err.non_retryable: _PENDING.pop(key, epoch) raise - answer = {"reason": err.reason, "message": err.message} + answer = {"reason": int(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. @@ -266,10 +260,9 @@ async def close_stream(request: stream_pb.CloseStreamRequest) -> int: 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: + if request.handle.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, + f"Can't close the Stream {request.handle.id} for another Workflow" ) run_id = streams.producer.run_id failure: Any = None @@ -277,9 +270,9 @@ async def close_stream(request: stream_pb.CloseStreamRequest) -> int: failure = streams.converter.failure_converter.from_failure( request.failure, streams.converter.payload_converter ) - await ( - await service.close_stream( - stream=_ref(request.stream), + await service._stream_service_close_stream( + models.CloseStreamRequest( + handle=_model_handle(request.handle), producer=models.Producer( kind=models.StreamProducerKind.StreamProducerKindWorkflow, id=_ids.workflow_producer_id(run_id), @@ -289,7 +282,6 @@ async def close_stream(request: stream_pb.CloseStreamRequest) -> int: if request.HasField("result") else None, failure=failure, - status=models.StreamCloseStatus(request.status) if request.status else None, ) ) return 0 diff --git a/temporalio/nexus/system/_support.py b/temporalio/nexus/system/_support.py index a6b621ae01..e7b7d0cbb2 100644 --- a/temporalio/nexus/system/_support.py +++ b/temporalio/nexus/system/_support.py @@ -493,59 +493,95 @@ def wake(_: object, done: asyncio.Future[None] = done) -> None: # The Stream sources. Each one checks where it runs and gives the value of that context, or -# raises where it has none to give. - - -def _no_workflow_streams(source: str) -> RuntimeError: - # This SDK doesn't run Workflow Stream calls yet, so a generated Workflow call fails here. - return RuntimeError(f"Workflow code can't call the Stream service yet ({source})") +# raises where it has none to give. They use the Streams API's id rules, so they import it when +# they run. def workflow_stream_id(definition: type[typing.Any]) -> str: """The id of the Workflow's own Stream of ``definition``.""" - del definition - raise _no_workflow_streams("workflow_stream_id") + from temporalio.common._streams import _ids + from temporalio.common._streams._defn import definition_of + from temporalio.workflow import info + + run = info() + return _ids.scoped_stream_id( + run.workflow_id, run.first_execution_run_id, definition_of(definition).name + ) def workflow_handle_owner() -> typing.Any: - """The owner of the Workflow's own Stream in its handle.""" - raise _no_workflow_streams("workflow_handle_owner") + """The owner of the Workflow's own Stream in its handle: the run chain.""" + from temporalio.common._streams._handle import workflow_owner as handle_owner + from temporalio.workflow import info + + run = info() + return handle_owner(run.workflow_id, run.first_execution_run_id) def workflow_owner() -> typing.Any: """The Producer of the Workflow run that owns a Stream that it creates.""" - raise _no_workflow_streams("workflow_owner") + return workflow_producer() def workflow_producer() -> typing.Any: - """The Producer of the Workflow run's writes.""" - raise _no_workflow_streams("workflow_producer") + """The Producer of the Workflow run.""" + from temporalio.common._streams import _ids + from temporalio.common._streams.system import models + from temporalio.workflow import info + + run_id = info().run_id + return models.Producer( + kind=models.StreamProducerKind.StreamProducerKindWorkflow, + id=_ids.workflow_producer_id(run_id), + run_id=run_id, + ) def withdraw_run() -> str | None: """The previous run whose Records a retry run's first write withdraws.""" - from temporalio.workflow import in_workflow + from temporalio.workflow import in_workflow, info - if in_workflow(): - raise _no_workflow_streams("withdraw_run") - # Code outside a Workflow isn't a retry run. - return None + if not in_workflow(): + # Code outside a Workflow isn't a retry run. + return None + run = info() + return run.continued_run_id if run.attempt > 1 and run.continued_run_id else None def client_producer() -> typing.Any: """The Producer of a write outside a Workflow.""" - # Each client has its own Producer, and a generated call doesn't say which client makes it. - # So the SDK's handles send their Producer, and a call without one fails. - raise RuntimeError("A Stream write outside a Workflow needs its Producer") + from temporalio.common._streams import _activity + from temporalio.common._streams.system import models + + streams = _activity.current.get() + if streams is None: + # Each client has its own Producer, and a generated call doesn't say which client makes + # it. So the SDK's handles send their Producer, and a call without one fails. + raise RuntimeError("A Stream write outside a Workflow needs its Producer") + return models._ProducerTransferTypeConverter().from_transfer_type( + streams.producer, models.Producer + ) def local_activity_failed_tasks() -> list[int]: """The run's failed Workflow Tasks, whose Records a Local Activity's write withdraws.""" - # Only a Local Activity has them, and this SDK's Local Activities don't write here yet. - return [] + from temporalio.common._streams import _activity + + streams = _activity.current.get() + # Only a Local Activity has them. + return list(streams.failed_tasks) if streams is not None else [] def local_activity_previous_completed_task() -> typing.Any: """The run's last completed Workflow Task, for a Local Activity's write.""" - # Only a Local Activity has it, and this SDK's Local Activities don't write here yet. - return None + from temporalio.common._streams import _activity + from temporalio.common._streams.system import models + + streams = _activity.current.get() + previous = streams.previous_completed_task if streams is not None else None + # Only a Local Activity has it. + if previous is None: + return None + return models._WorkflowTaskEpochTransferTypeConverter().from_transfer_type( + previous, models.WorkflowTaskEpoch + ) diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 8f2166497a..17b034a691 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -428,7 +428,7 @@ async def test_empty_stream_id_raises_with_no_call(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="")) + client.stream(Turns, handle=common.WorkflowStream(id="")) assert not extension.namespaces diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py index ca599fdb0b..36e033845b 100644 --- a/tests/contrib/streams/test_workflow_module.py +++ b/tests/contrib/streams/test_workflow_module.py @@ -21,6 +21,9 @@ from temporalio.common import RetryPolicy from temporalio.common._streams import _writes from temporalio.common._streams._backend import _ExtensionCalls +from temporalio.common._streams.system._resources.stream import ( + stream as generated_stream, +) from temporalio.contrib.stream.testing import MemoryStreamBackend from temporalio.converter import ( DataConverter, @@ -92,6 +95,37 @@ async def test_create_publish_and_close(stream_client: Client): assert isinstance(tokens[-1], common.Closed) +@workflow.defn +class GeneratedCalls: + @workflow.run + async def run(self) -> list[str]: + tokens = await workflow.create_stream(Tokens) + generated = generated_stream(Tokens, handle=tokens._generated().handle) + offset = await generated.publish("a", "b") + refused = [] + try: + await generated.read() + except common.StreamUnsupportedError as err: + refused.append(type(err).__name__) + await tokens.close() + return [str(offset > 0), *refused] + + +async def test_generated_workflow_calls_run_as_local_activities(stream_client: Client): + async with new_worker(stream_client, GeneratedCalls) as worker: + handle = await stream_client.start_workflow( + GeneratedCalls.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert await handle.result() == ["True", "StreamUnsupportedError"] + tokens = await read(handle) + events = await _event_types(handle) + assert data(tokens) == ["a", "b"] + assert isinstance(tokens[-1], common.Closed) + # Each write is a marker, as with the handles, and History has no Nexus event. + assert events.count(EventType.EVENT_TYPE_MARKER_RECORDED) >= 2 + assert EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED not in events + + @common.defn(name="first") class First(common.StreamDefinition[str]): ... @@ -200,11 +234,15 @@ async def run(self) -> list[str]: try: await stream.publish("late") except common.StreamClosedError as err: - errors.append(StreamErrorReason.Name(err.reason)) + errors.append( + StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) + ) try: await workflow.create_stream(Tokens) except common.StreamClosedError as err: - errors.append(StreamErrorReason.Name(err.reason)) + errors.append( + StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) + ) return errors @workflow.query @@ -241,7 +279,7 @@ async def run(self, stream: common.WorkflowStream[str]) -> str: try: await stream.close(result="ended by a writer") except common.StreamNotOwnerError as err: - return StreamErrorReason.Name(err.reason) + return StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) return "closed" @@ -572,7 +610,7 @@ def finish(self) -> str: try: self._stream.publish_nowait("late") except common.StreamClosedError as err: - return StreamErrorReason.Name(err.reason) + return StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) finally: self._done = True return "wrote" @@ -909,7 +947,7 @@ async def run(self) -> str: try: await tokens.publish(*["x" * 1_400_000] * 4) except common.StreamTooLargeError as err: - return StreamErrorReason.Name(err.reason) + return StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) return "written" @@ -949,7 +987,7 @@ async def run(self) -> str: try: await workflow.create_stream(Tokens) except common.StreamExpiredError as err: - return StreamErrorReason.Name(err.reason) + return StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) return "created" From bfd9b40b7ecc25f57281551d1d0811f3aa3091ea Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 1 Nov 2026 01:56:09 -0700 Subject: [PATCH 3/4] Gave `StreamHandle` the Workflow calls, and dropped `WorkflowStream`. In Workflow code a handle passes each call to the Workflow calls, which keep the Local Activity writes and the closes at the run end. `workflow.get_stream_handle` replaces `workflow.stream_handle`. A read, `producer` and the RPC options raise in Workflow code, and `async with` on another run chain's Stream fails before the block. --- temporalio/common/__init__.py | 3 - temporalio/common/_streams/__init__.py | 2 - temporalio/common/_streams/_extension.py | 2 +- temporalio/common/_streams/_handle.py | 74 ++- temporalio/common/_streams/_workflow.py | 490 ++++++++---------- temporalio/workflow/__init__.py | 4 +- temporalio/workflow/_streams.py | 14 +- tests/contrib/streams/_support.py | 4 +- tests/contrib/streams/test_redis.py | 2 +- tests/contrib/streams/test_stream_client.py | 20 +- tests/contrib/streams/test_workflow_module.py | 67 ++- 11 files changed, 369 insertions(+), 313 deletions(-) diff --git a/temporalio/common/__init__.py b/temporalio/common/__init__.py index 37470ad879..2d77f6c63f 100644 --- a/temporalio/common/__init__.py +++ b/temporalio/common/__init__.py @@ -64,7 +64,6 @@ ) from temporalio.common._streams._plugin import StreamBackend from temporalio.common._streams._reader import Subscription - from temporalio.common._streams._workflow import WorkflowStream __all__ = [ "ActivityIDConflictPolicy", @@ -129,7 +128,6 @@ "WorkerDeploymentVersion", "WorkflowIDConflictPolicy", "WorkflowIDReusePolicy", - "WorkflowStream", "defn", ] @@ -1613,7 +1611,6 @@ 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 be870617d7..c803acdb7b 100644 --- a/temporalio/common/_streams/__init__.py +++ b/temporalio/common/_streams/__init__.py @@ -26,7 +26,6 @@ ) from temporalio.common._streams._plugin import StreamBackend from temporalio.common._streams._reader import Subscription -from temporalio.common._streams._workflow import WorkflowStream __all__ = [ "END", @@ -53,6 +52,5 @@ "StreamUnsupportedError", "Subscription", "Superseded", - "WorkflowStream", "defn", ] diff --git a/temporalio/common/_streams/_extension.py b/temporalio/common/_streams/_extension.py index 82a14a0a77..48c0f71665 100644 --- a/temporalio/common/_streams/_extension.py +++ b/temporalio/common/_streams/_extension.py @@ -378,7 +378,7 @@ def _activity_headers(self) -> dict[str, temporalio.api.common.v1.Payload]: 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. + # The next run needs each owned Stream's backend for get_stream_handle. owned = [ [stream_id, run_stream.backend] for stream_id, run_stream in streams_workflow._state().streams.items() diff --git a/temporalio/common/_streams/_handle.py b/temporalio/common/_streams/_handle.py index deb0f39a24..839b64b34f 100644 --- a/temporalio/common/_streams/_handle.py +++ b/temporalio/common/_streams/_handle.py @@ -283,6 +283,11 @@ class StreamHandle(Generic[T]): in the proto JSON form of the ``temporal.api.stream.v1.StreamHandle`` message. The type hint, such as ``StreamHandle[Turn]``, gives the payload type. + Workflow code writes to and closes a Stream, and can't read it. Each publish is a Local + Activity, and each close a command that the backend runs after the Workflow Task commits. A + Workflow passes a handle to an Activity, a child Workflow or the next run of continue-as-new + as an argument, or hands it out through a query. + In client code, a handle uses the client that made it, such as with :py:meth:`temporalio.client.Client.get_stream_handle`, or the client that decoded it, such as from a query, an Update or a Workflow result. A handle that no client decoded, such as @@ -338,6 +343,12 @@ def backend(self) -> str | None: def _calls(self) -> ContextCalls: # The context modules import on first use, so Workflow code never imports the client's. + import temporalio.workflow + + if temporalio.workflow.in_workflow(): + from temporalio.common._streams import _workflow + + return _workflow.WorkflowCalls() from temporalio.common._streams import _stream_client return _stream_client.calls(self) @@ -355,12 +366,21 @@ def publish( of the future doesn't stop the write. A list is one value. In client code, the next call of the handle raises a write error that no caller got. + In Workflow code, the 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 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. + Args: values: One value for each Record. - rpc_metadata: Headers of each call. The Redis and memory backends make no RPC, so - they send none. - rpc_timeout: The deadline of each call. A call past it raises ``RPCError`` with - ``DEADLINE_EXCEEDED``. + rpc_metadata: Headers of each call, outside Workflow code. The Redis and memory + backends make no RPC, so they send none. + rpc_timeout: The deadline of each call, outside Workflow code. A call past it + raises ``RPCError`` with ``DEADLINE_EXCEEDED``. Returns: The future of the write. It gives the offset of the last Record, or raises @@ -368,9 +388,14 @@ def publish( ``StreamNotFoundError`` for a Stream that nobody created. Raises: - StreamClosedError: This handle closed the Stream. + StreamClosedError: This handle closed the Stream, or in Workflow code, this run + closed it. StreamConfigError: No client decoded the handle, or the client has no backend of - that name. + that name. In Workflow code, the Worker has no Stream backend or not the + Stream's backend, which fails the Workflow Task and not the run. + temporalio.workflow.ReadOnlyContextError: The Workflow code must not change state, + such as a query handler. + StreamUnsupportedError: Workflow code gave ``rpc_metadata`` or ``rpc_timeout``. """ return self._calls().publish(self, values, rpc_metadata, rpc_timeout) @@ -403,6 +428,9 @@ def subscribe( rpc_metadata: Headers of each call of the read. See :py:meth:`publish`. rpc_timeout: The deadline of each call of the read, not of the whole read. A call that waits for new Records returns before the deadline. + + Raises: + StreamUnsupportedError: The code is Workflow code, which can't read a Stream. """ return self._calls().subscribe( self, @@ -424,6 +452,9 @@ async def get_offset( Offsets are sparse and above zero, so ``offset=0`` reads from the start. ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`publish`. + + Raises: + StreamUnsupportedError: The code is Workflow code, which can't read a Stream. """ return await self._calls().get_offset(self, rpc_metadata, rpc_timeout) @@ -441,6 +472,16 @@ async def close( ``failure``. A cancel failure, such as a ``CancelledError``, closes the Stream as canceled. ``rpc_metadata`` and ``rpc_timeout`` are as for :py:meth:`publish`. + In Workflow code, the close comes 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. A cancel failure includes an ``ActivityError`` that a cancel caused. + Later writes in the run raise :py:class:`StreamClosedError`. + Raises: StreamNotOwnerError: The code doesn't own the Stream, such as a client with a Workflow's Stream. Only the run chain closes a Workflow's Stream, and only a @@ -468,12 +509,29 @@ async def __aexit__( A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as ``KeyboardInterrupt``, leaves it open. If this handle already closed the Stream, this does nothing. + + In Workflow code, 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. """ await self._calls().exit(self, exc) def publish_nowait(self, *values: T) -> None: """Write one Record for each value, in order, without an await, in sync Workflow code. + 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. + Raises: StreamUnsupportedError: The code isn't Workflow code. """ @@ -484,6 +542,10 @@ def close_nowait( ) -> None: """Close the Stream without an await, in 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. + Raises: StreamUnsupportedError: The code isn't Workflow code. """ diff --git a/temporalio/common/_streams/_workflow.py b/temporalio/common/_streams/_workflow.py index 1e2d3ba912..8c562752ca 100644 --- a/temporalio/common/_streams/_workflow.py +++ b/temporalio/common/_streams/_workflow.py @@ -1,9 +1,8 @@ """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. +again with :py:func:`temporalio.workflow.get_stream_handle`. A Stream that the Workflow receives +as an argument works the same, such as a handle 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 @@ -16,11 +15,10 @@ from __future__ import annotations import asyncio -from collections.abc import Callable, Generator +from collections.abc import Callable, Generator, Mapping from dataclasses import dataclass, field from datetime import timedelta -from types import TracebackType -from typing import Any, NamedTuple, TypeVar +from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar import google.protobuf.message @@ -32,13 +30,15 @@ from temporalio.common._streams import _errors, _ids from temporalio.common._streams._defn import StreamDefinition, definition_of from temporalio.common._streams._handle import ( + ContextCalls, Creating, StreamHandle, StreamOptions, exit_failure, workflow_owner, ) -from temporalio.common._streams._items import close_failure +from temporalio.common._streams._items import End, close_failure +from temporalio.common._streams._writer import WriteFuture, unretrieved_error from temporalio.common._streams.system import models from temporalio.common._streams.system._resources import Stream as GeneratedStream from temporalio.common._streams.system._resources import stream as generated_stream @@ -46,6 +46,9 @@ create_stream as generated_create, ) +if TYPE_CHECKING: + from temporalio.common._streams._reader import Subscription + T = TypeVar("T") OutputT = TypeVar("OutputT") @@ -113,7 +116,7 @@ class _RunState: # 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( + deferred_closes: list[tuple[StreamHandle[Any], BaseException | None]] = field( default_factory=list ) # The publishes that failed before a caller got the error. A caller can still get it, so @@ -214,35 +217,12 @@ def _records(request: google.protobuf.message.Message) -> int: 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) + 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)) @@ -364,18 +344,11 @@ async def drain() -> None: if cancel is not None: raise cancel for future in state.failed_publishes: - error = _unretrieved_error(future) + 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 ]: @@ -392,72 +365,157 @@ 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. +def _model_handle(handle: StreamHandle[Any]) -> models.StreamHandle: + return models.StreamHandle( + id=handle.id, + owner=workflow_owner(handle.workflow_id, handle.first_execution_run_id), + ) - 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 _proto_handle(handle: StreamHandle[Any]) -> stream_v1.StreamHandle: + return models._StreamHandleTransferTypeConverter().to_transfer_type( + _model_handle(handle) + ) + + +def _generated_handle(handle: StreamHandle[Any]) -> models.StreamHandle: + return models.StreamHandle( + id=handle.id, + owner=workflow_owner(handle.workflow_id, handle.first_execution_run_id), + backend=handle.backend or None, + ) - def _model_handle(self) -> models.StreamHandle: - return models.StreamHandle( - id=self.id, - owner=workflow_owner(self.workflow_id, self.first_execution_run_id), + +def _check( + handle: StreamHandle[Any], 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(_model_handle(handle)) + if run_stream.closed: + raise _errors.StreamClosedError( + f"Can't {action} the Stream {handle.id}, since this run closed it" ) + if not backend_checked: + _check_backend(handle.backend, action) + return run_stream + + +def _start_append( + handle: StreamHandle[Any], values: tuple[Any, ...] +) -> asyncio.Future[int]: + # The checks raise at the call, as the write's place in the order does. + if not values: + raise TypeError("publish takes one or more values") + _check(handle, "write to") + call = generated_stream.start_append_records( + handle=_generated_handle(handle), + records=[models.AppendRecord(body=value) for value in values], + ) + return _call_future(call, lambda response: response.last_offset or 0) + + +def _append(handle: StreamHandle[Any], values: tuple[Any, ...]) -> asyncio.Future[int]: + if not values: + raise TypeError("publish takes one or more values") + _check(handle, "write to") + records = [stream_v1.AppendRecord(body=_payload(v)) for v in values] + return _write( + handle.id, + APPEND_ACTIVITY, + stream_pb.AppendRecordsRequest(handle=_proto_handle(handle), records=records), + ) + - def _proto_handle(self) -> stream_v1.StreamHandle: - return models._StreamHandleTransferTypeConverter().to_transfer_type( - self._model_handle() +def _start_close( + handle: StreamHandle[Any], result: Any, failure: BaseException | None +) -> asyncio.Future[int]: + _closing(handle) + call = generated_stream.start_close_stream( + handle=_generated_handle(handle), result=result, failure=failure + ) + return _call_future(call, lambda _: 0) + + +def _check_owner(handle: StreamHandle[Any]) -> None: + info = temporalio.workflow.info() + if (handle.workflow_id, handle.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 {handle.id}, since this run chain didn't create it" ) - 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. - """ + +def _closing(handle: StreamHandle[Any]) -> _RunStream: + _check_owner(handle) + return _check(handle, "close") + + +def _close( + handle: StreamHandle[Any], result: Any, failure: BaseException | None +) -> asyncio.Future[int]: + _closing(handle).closed = True + request = stream_pb.CloseStreamRequest(handle=_proto_handle(handle)) + if result is not None: + request.result.CopyFrom(_payload(result)) + failure = close_failure(failure) + if failure is not None: + payload_converter, failure_converter = _converters() + failure_converter.to_failure(failure, payload_converter, request.failure) + return _write(handle.id, CLOSE_ACTIVITY, request) + + +def _no_rpc_options( + rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None +) -> None: + if rpc_metadata or rpc_timeout is not None: + # The same error as the other calls that Workflow code can't make. + raise _errors.StreamUnsupportedError( + "A Stream call in Workflow code is a Local Activity, so it takes no " + "rpc_metadata or rpc_timeout" + ) + + +class WorkflowCalls(ContextCalls): + """The Stream calls in Workflow code. Workflow code writes and closes a Stream. + + Each create and publish is a Local Activity, and each close is a command that the backend + runs after the Workflow Task commits. + """ + + def publish( + self, + handle: StreamHandle[Any], + values: tuple[Any, ...], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> asyncio.Future[int]: + """Start the write Local Activity. An error that no caller gets fails the run.""" + _no_rpc_options(rpc_metadata, rpc_timeout) state = _state() - future = self._start_append(values) + future = _start_append(handle, values) loop = asyncio.get_running_loop() def warn(future: asyncio.Future[int]) -> None: - error = _unretrieved_error(future) + 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, + handle.id, error, ) def failed(future: asyncio.Future[int]) -> None: - if _unretrieved_error(future) is not 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) @@ -465,175 +523,77 @@ def failed(future: asyncio.Future[int]) -> None: 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._start_append(values)) - - def _start_append(self, values: tuple[Any, ...]) -> asyncio.Future[int]: - # The checks raise at the call, as the write's place in the order does. - if not values: - raise TypeError("publish takes one or more values") - self._check("write to") - call = generated_stream.start_append_records( - handle=self._generated_handle(), - records=[models.AppendRecord(body=value) for value in values], - ) - return _call_future(call, lambda response: response.last_offset or 0) - - 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_handle()) - if run_stream.closed: - raise _errors.StreamClosedError( - f"Can't {action} the Stream {self.id}, since this run closed it" - ) - 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( - handle=self._proto_handle(), records=records - ), - ) + def subscribe( + self, + handle: StreamHandle[Any], + *, + offset: int | End | None, + last: int | None, + wait_new_records: bool, + create_wait: timedelta | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> Subscription[Any]: + """Raise, since Workflow code can't read a Stream.""" + raise no_workflow_reads() + + async def get_offset( + self, + handle: StreamHandle[Any], + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | None, + ) -> int: + """Raise, since Workflow code can't describe a Stream.""" + raise no_workflow_reads() async def close( - self, *, result: Any | None = None, failure: BaseException | None = None + self, + handle: StreamHandle[Any], + result: Any, + failure: BaseException | None, + rpc_metadata: Mapping[str, str | bytes], + rpc_timeout: timedelta | 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._generated().close(result=result, failure=failure) - - def _generated(self) -> GeneratedStream[Any]: + """Close the Stream after the writes before it, once the Workflow Task commits.""" + _no_rpc_options(rpc_metadata, rpc_timeout) # The generated Workflow API makes the call, and the Worker's interceptor runs it as the # write Local Activity. - return GeneratedStream(self._generated_handle()) - - def _generated_handle(self) -> models.StreamHandle: - return models.StreamHandle( - id=self.id, - owner=workflow_owner(self.workflow_id, self.first_execution_run_id), - backend=self.backend or None, + await GeneratedStream(_generated_handle(handle)).close( + result=result, failure=failure ) - def close_nowait( - self, *, result: Any | None = None, failure: BaseException | None = None - ) -> None: - """Close the Stream without an await, for sync Workflow code. + async def enter(self, handle: StreamHandle[Any]) -> None: + """Check that this run chain can close the Stream at the end of the block.""" + # Fail before the block, since the close at its end would raise and hide the block's + # own error. + _check_owner(handle) - 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._start_close(result, failure)) - - def _start_close( - self, result: Any, failure: BaseException | None - ) -> asyncio.Future[int]: - self._closing() - call = generated_stream.start_close_stream( - handle=self._generated_handle(), result=result, failure=failure - ) - return _call_future(call, lambda _: 0) - - 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" - ) - return self._check("close") - - def _close(self, result: Any, failure: BaseException | None) -> asyncio.Future[int]: - self._closing().closed = True - request = stream_pb.CloseStreamRequest(handle=self._proto_handle()) - if result is not None: - request.result.CopyFrom(_payload(result)) - failure = close_failure(failure) - if failure is not None: - payload_converter, failure_converter = _converters() - failure_converter.to_failure(failure, payload_converter, request.failure) - 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. - """ + async def exit(self, handle: StreamHandle[Any], exc: BaseException | None) -> None: + """Close the Stream, or defer the close to the run's end.""" state = _state() - if state.stream(self._model_handle()).closed: + if state.stream(_model_handle(handle)).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) + await self.close(handle, None, failure, {}, None) else: # A refused close raises here, as it does for a close at the block end. - self._closing() - state.deferred_closes.append((self, failure)) + _closing(handle) + state.deferred_closes.append((handle, failure)) + + def publish_nowait( + self, handle: StreamHandle[Any], values: tuple[Any, ...] + ) -> None: + """Start the write. Its error raises in the run as an error that the run doesn't catch.""" + _write_nowait(_start_append(handle, values)) + + def close_nowait( + self, handle: StreamHandle[Any], result: Any, failure: BaseException | None + ) -> None: + """Start the close. The end of the run waits for it.""" + _write_nowait(_start_close(handle, result, failure)) def _closes_at_block_end(exc: BaseException | None) -> bool: @@ -659,8 +619,8 @@ def close_at_run_end(err: BaseException | None) -> None: ): return for stream, failure in state.deferred_closes: - if not state.stream(stream._model_handle()).closed: - _write_nowait(stream._start_close(None, failure)) + if not state.stream(_model_handle(stream)).closed: + _write_nowait(_start_close(stream, None, failure)) def _retried(exc: BaseException) -> bool: @@ -703,10 +663,10 @@ def _cancelled(exc: BaseException) -> bool: def _own_handle( definition: type[StreamDefinition[T]], backend: str | None -) -> WorkflowStream[T]: +) -> StreamHandle[T]: info = temporalio.workflow.info() known = definition_of(definition) - return WorkflowStream( + return StreamHandle( id=_ids.scoped_stream_id( info.workflow_id, info.first_execution_run_id, known.name ), @@ -722,13 +682,13 @@ def create_stream( *, backend: str | None, options: StreamOptions, -) -> Creating[WorkflowStream[T]]: +) -> Creating[StreamHandle[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]: + async def create() -> StreamHandle[T]: stream = _creating(definition, backend) await generated_create.create_stream( definition, retention=options.get("retention"), backend=stream.backend @@ -743,7 +703,7 @@ def create_stream_nowait( *, backend: str | None, options: StreamOptions, -) -> WorkflowStream[T]: +) -> StreamHandle[T]: """Create a Stream that this run chain owns, or get it, without an await. See :py:func:`temporalio.workflow.create_stream_nowait`. @@ -758,10 +718,10 @@ def create_stream_nowait( def _creating( definition: type[StreamDefinition[T]], backend: str | None -) -> WorkflowStream[T]: - # The checks raise before the generated call, as they did before the write. +) -> StreamHandle[T]: + # The checks raise before the generated call. stream = _own_handle(definition, _check_backend(backend, "create")) - stream._check("create", backend_checked=True) + _check(stream, "create", backend_checked=True) return stream @@ -772,7 +732,7 @@ def _call_future( # handle gives the write's result. The caller's future follows both, and the run's end # waits for them. loop = asyncio.get_running_loop() - future = _WriteFuture(loop=loop) + future = WriteFuture(loop=loop) async def follow() -> None: # The task sets the future before it ends, so the run's end sees a failed publish. @@ -796,11 +756,11 @@ async def follow() -> None: def _start_create( - stream: WorkflowStream[Any], retention: timedelta | None, owner: stream_v1.Producer + stream: StreamHandle[Any], retention: timedelta | None, owner: stream_v1.Producer ) -> asyncio.Future[int]: - run_stream = stream._check("create", backend_checked=True) + run_stream = _check(stream, "create", backend_checked=True) request = stream_pb.CreateStreamRequest( - stream_id=stream.id, owner=owner, stream_handle=stream._proto_handle() + stream_id=stream.id, owner=owner, stream_handle=_proto_handle(stream) ) if retention is not None: request.retention.FromTimedelta(retention) @@ -832,13 +792,13 @@ def operation_token(self) -> str | None: return None -def _from_handle(handle: models.StreamHandle) -> WorkflowStream[Any]: +def _from_model(handle: models.StreamHandle, backend: str | None) -> StreamHandle[Any]: owner = handle.owner - return WorkflowStream( + return StreamHandle( id=handle.id, workflow_id=owner.business_id if owner else None, first_execution_run_id=owner.run_id if owner else None, - backend=handle.backend, + backend=backend, ) @@ -853,8 +813,7 @@ def run_create_stream( ) -> temporalio.workflow.NexusOperationHandle[models.CreateStreamResult]: """Run a create of the generated Workflow API as the create Local Activity.""" handle = request.stream_handle or models.StreamHandle(id=request.id) - stream = _from_handle(handle) - stream._backend = _check_backend(handle.backend, "create") + stream = _from_model(handle, _check_backend(handle.backend, "create")) owner = ( models._ProducerTransferTypeConverter().to_transfer_type(request.owner) if request.owner is not None @@ -868,8 +827,8 @@ def run_append_records( request: models.AppendRecordsRequest[Any], ) -> temporalio.workflow.NexusOperationHandle[models.AppendRecordsResponse]: """Run an append of the generated Workflow API as the append Local Activity.""" - stream = _from_handle(request.handle) - write = stream._append(tuple(record.body for record in request.records or [])) + stream = _from_model(request.handle, request.handle.backend) + write = _append(stream, tuple(record.body for record in request.records or [])) return _WriteHandle( write, lambda offset: models.AppendRecordsResponse(last_offset=offset) ) @@ -879,7 +838,8 @@ def run_close_stream( request: models.CloseStreamRequest, ) -> temporalio.workflow.NexusOperationHandle[models.CloseStreamResponse]: """Run a close of the generated Workflow API as the close write.""" - write = _from_handle(request.handle)._close(request.result, request.failure) + stream = _from_model(request.handle, request.handle.backend) + write = _close(stream, request.result, request.failure) return _WriteHandle(write, lambda _: models.CloseStreamResponse()) @@ -890,12 +850,12 @@ def no_workflow_reads() -> _errors.StreamUnsupportedError: ) -def stream_handle( +def get_stream_handle( definition: type[StreamDefinition[T]], *, backend: str | None = None -) -> WorkflowStream[T]: +) -> StreamHandle[T]: """Get a Stream that this run chain created, by its definition. - See :py:func:`temporalio.workflow.stream_handle`. + See :py:func:`temporalio.workflow.get_stream_handle`. """ stream = _own_handle(definition, backend) if not temporalio.workflow.unsafe.is_read_only(): diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index ef1c20a37c..5efa835679 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -154,7 +154,7 @@ from ._streams import ( create_stream, create_stream_nowait, - stream_handle, + get_stream_handle, ) from ._workflow_ops import ( ChildWorkflowCancellationType, @@ -265,7 +265,7 @@ "unsafe", "create_stream", "create_stream_nowait", - "stream_handle", + "get_stream_handle", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py index c220c869a5..4388528eff 100644 --- a/temporalio/workflow/_streams.py +++ b/temporalio/workflow/_streams.py @@ -14,8 +14,8 @@ from temporalio.common import ( Creating, StreamDefinition, + StreamHandle, StreamOptions, - WorkflowStream, ) T = TypeVar("T") @@ -26,7 +26,7 @@ def create_stream( *, backend: str | None = None, **options: Unpack[StreamOptions], -) -> Creating[WorkflowStream[T]]: +) -> Creating[StreamHandle[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 @@ -48,7 +48,7 @@ def create_stream( 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__`. + See :py:meth:`temporalio.common.StreamHandle.__aexit__`. .. warning:: This API is experimental. @@ -80,7 +80,7 @@ def create_stream_nowait( *, backend: str | None = None, **options: Unpack[StreamOptions], -) -> WorkflowStream[T]: +) -> StreamHandle[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 @@ -96,9 +96,9 @@ def create_stream_nowait( return _workflow.create_stream_nowait(definition, backend=backend, options=options) -def stream_handle( +def get_stream_handle( definition: type[StreamDefinition[T]], *, backend: str | None = None -) -> WorkflowStream[T]: +) -> StreamHandle[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 @@ -119,4 +119,4 @@ def stream_handle( """ from temporalio.common._streams import _workflow - return _workflow.stream_handle(definition, backend=backend) + return _workflow.get_stream_handle(definition, backend=backend) diff --git a/tests/contrib/streams/_support.py b/tests/contrib/streams/_support.py index 27e0a95975..2faeeed8aa 100644 --- a/tests/contrib/streams/_support.py +++ b/tests/contrib/streams/_support.py @@ -53,10 +53,10 @@ def workflow_stream( handle: WorkflowHandle[Any, Any], definition: type[common.StreamDefinition[T]], backend: str | None = None, -) -> common.ClientStream[T]: +) -> common.StreamHandle[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) + return handle._client.get_stream_handle(definition, id=stream_id, backend=backend) diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index b1ec38ae29..deb38c9d55 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -15,7 +15,7 @@ class Tokens(common.StreamDefinition[str]): ... -async def read_all(stream: common.ClientStream[Any]) -> list[Any]: +async def read_all(stream: common.StreamHandle[Any]) -> list[Any]: return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 17b034a691..777eb87eb8 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -149,7 +149,7 @@ async def test_query_and_workflow_results_bind_their_client(client: Client): @workflow.defn class WriteToStream: @workflow.run - async def run(self, stream: common.WorkflowStream[Turn]) -> None: + async def run(self, stream: common.StreamHandle[Turn]) -> None: await stream.publish(Turn("workflow")) @@ -162,7 +162,7 @@ async def test_client_stream_takes_workflow_writes(stream_client: Client): async with new_worker(stream_client, WriteToStream) as worker: await stream_client.execute_workflow( WriteToStream.run, - common.WorkflowStream.from_handle(stream), + stream, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -173,7 +173,7 @@ async def test_client_stream_takes_workflow_writes(stream_client: Client): 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) + again = stream_client.get_stream_handle(Turns, id=stream.id) resumed = [ item.data async for item in again.subscribe(offset=first + 1) @@ -193,9 +193,9 @@ async def run(self, prompt: str) -> str: return prompt.upper() @workflow.query - def stream(self) -> common.WorkflowStream[str]: + def stream(self) -> common.StreamHandle[str]: # The handle needs no call, so a reader gets it before the create. - return workflow.stream_handle(Tokens) + return workflow.get_stream_handle(Tokens) async def test_client_reads_a_workflow_stream(stream_client: Client): @@ -203,13 +203,11 @@ async def test_client_reads_a_workflow_stream(stream_client: Client): 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]) + reader = await handle.query(Agent.stream) # 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", @@ -426,9 +424,7 @@ async def test_read_keeps_no_offsets_of_client_records(stream_client: Client): 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, handle=common.WorkflowStream(id="")) + client.get_stream_handle(Turns, id="") assert not extension.namespaces @@ -484,7 +480,7 @@ async def test_registration_rules(client: Client): ] two = Client(**config) assert (await two.create_stream(Tokens, id="d-1")).backend == "b" - assert two.stream(Tokens, id="d-1").backend == "b" + assert two.get_stream_handle(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) diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py index 36e033845b..0dee7e23d3 100644 --- a/tests/contrib/streams/test_workflow_module.py +++ b/tests/contrib/streams/test_workflow_module.py @@ -19,7 +19,7 @@ WorkflowHandle, ) from temporalio.common import RetryPolicy -from temporalio.common._streams import _writes +from temporalio.common._streams import _workflow, _writes from temporalio.common._streams._backend import _ExtensionCalls from temporalio.common._streams.system._resources.stream import ( stream as generated_stream, @@ -100,7 +100,7 @@ class GeneratedCalls: @workflow.run async def run(self) -> list[str]: tokens = await workflow.create_stream(Tokens) - generated = generated_stream(Tokens, handle=tokens._generated().handle) + generated = generated_stream(Tokens, handle=_workflow._generated_handle(tokens)) offset = await generated.publish("a", "b") refused = [] try: @@ -126,6 +126,49 @@ async def test_generated_workflow_calls_run_as_local_activities(stream_client: C assert EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED not in events +@workflow.defn +class RefusedCalls: + @workflow.run + async def run(self) -> list[str]: + tokens = await workflow.create_stream(Tokens) + refused = [] + for call in ( + lambda: tokens.subscribe(), + lambda: tokens.producer("ingest", 1), + lambda: tokens.publish("a", rpc_timeout=timedelta(seconds=1)), + ): + try: + call() + except common.StreamUnsupportedError as err: + refused.append(type(err).__name__) + try: + await tokens.get_offset() + except common.StreamUnsupportedError as err: + refused.append(type(err).__name__) + # Another run chain's Stream fails before the block runs. + other: common.StreamHandle[str] = common.StreamHandle( + id="other:run-1:tokens", workflow_id="other", first_execution_run_id="run-1" + ) + try: + async with other: + refused.append("entered") + except common.StreamNotOwnerError as err: + refused.append(type(err).__name__) + await tokens.close() + return refused + + +async def test_calls_that_workflow_code_cant_make_raise(stream_client: Client): + async with new_worker(stream_client, RefusedCalls) as worker: + result = await stream_client.execute_workflow( + RefusedCalls.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert result == [ + *["StreamUnsupportedError"] * 4, + "StreamNotOwnerError", + ] + + @common.defn(name="first") class First(common.StreamDefinition[str]): ... @@ -166,13 +209,13 @@ async def test_async_with_closes_the_stream(stream_client: Client): @workflow.defn class Continuing: @workflow.run - async def run(self, stream: common.WorkflowStream[str] | None) -> str: + async def run(self, stream: common.StreamHandle[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) + assert stream == workflow.get_stream_handle(Tokens) await stream.publish("run 2") await stream.close(result="done") return "done" @@ -223,7 +266,7 @@ async def test_continue_as_new_passes_the_stream(stream_client: Client): @workflow.defn class Mistakes: def __init__(self) -> None: - self._stream: common.WorkflowStream[str] | None = None + self._stream: common.StreamHandle[str] | None = None @workflow.run async def run(self) -> list[str]: @@ -249,7 +292,7 @@ async def run(self) -> list[str]: def publish_in_query(self) -> str: assert self._stream is not None try: - self._stream._append(("x",)) + self._stream.publish("x") except workflow.ReadOnlyContextError as err: return type(err).__name__ return "wrote" @@ -274,7 +317,7 @@ async def test_replay_reproducible_errors_raise_in_workflow_code( @workflow.defn class ClosesAReceivedStream: @workflow.run - async def run(self, stream: common.WorkflowStream[str]) -> str: + async def run(self, stream: common.StreamHandle[str]) -> str: await stream.publish("from workflow") try: await stream.close(result="ended by a writer") @@ -289,7 +332,7 @@ async def test_workflow_closes_only_its_own_streams(stream_client: Client): assert ( await stream_client.execute_workflow( ClosesAReceivedStream.run, - common.WorkflowStream.from_handle(stream), + stream, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -692,7 +735,7 @@ 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) + missing = common.StreamHandle[str](id=stream_id) future = missing.publish("lost") errors: list[BaseException | None] = [] if retrieve: @@ -869,7 +912,7 @@ 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) + stream = workflow.get_stream_handle(Tokens) await stream.publish("x" * 100_000) @@ -974,7 +1017,7 @@ class CreatesOnANamedBackend: 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") + await workflow.get_stream_handle(Tokens).close(result="closed") return stream.backend or "" @@ -1105,7 +1148,7 @@ async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: @workflow.defn class WritesFromASignalAtTheEnd: def __init__(self) -> None: - self.late: common.WorkflowStream[str] | None = None + self.late: common.StreamHandle[str] | None = None @workflow.run async def run(self) -> str: From f56d0d7413ecdde9434243f9e4142d0bda64e92c Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 3 Nov 2026 00:53:57 -0700 Subject: [PATCH 4/4] Moved the Workflow calls to `WorkflowStream`, with `to_stream`. --- temporalio/common/_streams/_activity.py | 2 +- temporalio/common/_streams/_extension.py | 2 +- temporalio/common/_streams/_workflow.py | 267 +++++++----------- temporalio/nexus/system/_support.py | 2 +- temporalio/workflow/__init__.py | 8 +- temporalio/workflow/_streams.py | 201 +++++++++++-- tests/contrib/streams/_support.py | 8 +- tests/contrib/streams/test_redis.py | 4 +- tests/contrib/streams/test_stream_client.py | 23 +- tests/contrib/streams/test_workflow_module.py | 72 +++-- 10 files changed, 350 insertions(+), 239 deletions(-) diff --git a/temporalio/common/_streams/_activity.py b/temporalio/common/_streams/_activity.py index d6e87b288c..36779e3cf2 100644 --- a/temporalio/common/_streams/_activity.py +++ b/temporalio/common/_streams/_activity.py @@ -15,7 +15,7 @@ @dataclass class ActivityStreams: - """What a Stream backend gives the Activity handles for one Activity attempt.""" + """What a Stream backend gives the Streams of one Activity attempt.""" backend: BackendCalls converter: temporalio.converter.DataConverter diff --git a/temporalio/common/_streams/_extension.py b/temporalio/common/_streams/_extension.py index 48c0f71665..0fb2fe7344 100644 --- a/temporalio/common/_streams/_extension.py +++ b/temporalio/common/_streams/_extension.py @@ -378,7 +378,7 @@ def _activity_headers(self) -> dict[str, temporalio.api.common.v1.Payload]: 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 get_stream_handle. + # The next run needs each owned Stream's backend for get_stream. owned = [ [stream_id, run_stream.backend] for stream_id, run_stream in streams_workflow._state().streams.items() diff --git a/temporalio/common/_streams/_workflow.py b/temporalio/common/_streams/_workflow.py index 8c562752ca..1de5900689 100644 --- a/temporalio/common/_streams/_workflow.py +++ b/temporalio/common/_streams/_workflow.py @@ -1,8 +1,9 @@ """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.get_stream_handle`. A Stream that the Workflow receives -as an argument works the same, such as a handle that a run passes on through continue-as-new. +again with :py:func:`temporalio.workflow.get_stream`. A handle that the Workflow receives as data +works the same through :py:func:`temporalio.workflow.to_stream`, such as a handle 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 @@ -15,10 +16,10 @@ from __future__ import annotations import asyncio -from collections.abc import Callable, Generator, Mapping +from collections.abc import Callable, Generator from dataclasses import dataclass, field from datetime import timedelta -from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar +from typing import Any, NamedTuple, TypeVar import google.protobuf.message @@ -30,14 +31,13 @@ from temporalio.common._streams import _errors, _ids from temporalio.common._streams._defn import StreamDefinition, definition_of from temporalio.common._streams._handle import ( - ContextCalls, Creating, StreamHandle, StreamOptions, exit_failure, workflow_owner, ) -from temporalio.common._streams._items import End, close_failure +from temporalio.common._streams._items import close_failure from temporalio.common._streams._writer import WriteFuture, unretrieved_error from temporalio.common._streams.system import models from temporalio.common._streams.system._resources import Stream as GeneratedStream @@ -45,9 +45,7 @@ from temporalio.common._streams.system.operations import ( create_stream as generated_create, ) - -if TYPE_CHECKING: - from temporalio.common._streams._reader import Subscription +from temporalio.workflow._streams import WorkflowStream T = TypeVar("T") OutputT = TypeVar("OutputT") @@ -366,10 +364,7 @@ def _payload(value: Any) -> Any: def _model_handle(handle: StreamHandle[Any]) -> models.StreamHandle: - return models.StreamHandle( - id=handle.id, - owner=workflow_owner(handle.workflow_id, handle.first_execution_run_id), - ) + return models.StreamHandle(id=handle.id, owner=handle.owner) def _proto_handle(handle: StreamHandle[Any]) -> stream_v1.StreamHandle: @@ -380,9 +375,7 @@ def _proto_handle(handle: StreamHandle[Any]) -> stream_v1.StreamHandle: def _generated_handle(handle: StreamHandle[Any]) -> models.StreamHandle: return models.StreamHandle( - id=handle.id, - owner=workflow_owner(handle.workflow_id, handle.first_execution_run_id), - backend=handle.backend or None, + id=handle.id, owner=handle.owner, backend=handle.backend or None ) @@ -441,9 +434,11 @@ def _start_close( return _call_future(call, lambda _: 0) -def _check_owner(handle: StreamHandle[Any]) -> None: +def check_owner(handle: StreamHandle[Any]) -> None: + """Check that this run chain owns the Stream, which only the owner closes.""" info = temporalio.workflow.info() - if (handle.workflow_id, handle.first_execution_run_id) != ( + owner = handle.owner + if owner is None or (owner.business_id, owner.run_id) != ( info.workflow_id, info.first_execution_run_id, ): @@ -455,7 +450,7 @@ def _check_owner(handle: StreamHandle[Any]) -> None: def _closing(handle: StreamHandle[Any]) -> _RunStream: - _check_owner(handle) + check_owner(handle) return _check(handle, "close") @@ -473,127 +468,72 @@ def _close( return _write(handle.id, CLOSE_ACTIVITY, request) -def _no_rpc_options( - rpc_metadata: Mapping[str, str | bytes], rpc_timeout: timedelta | None -) -> None: - if rpc_metadata or rpc_timeout is not None: - # The same error as the other calls that Workflow code can't make. - raise _errors.StreamUnsupportedError( - "A Stream call in Workflow code is a Local Activity, so it takes no " - "rpc_metadata or rpc_timeout" - ) +def publish(handle: StreamHandle[Any], values: tuple[Any, ...]) -> asyncio.Future[int]: + """Start the write Local Activity. An error that no caller gets fails the run. + See :py:meth:`temporalio.workflow.WorkflowStream.publish`. + """ + state = _state() + future = _start_append(handle, values) + loop = asyncio.get_running_loop() -class WorkflowCalls(ContextCalls): - """The Stream calls in Workflow code. Workflow code writes and closes a Stream. + 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", + handle.id, + error, + ) - Each create and publish is a Local Activity, and each close is a command that the backend - runs after the Workflow Task commits. - """ + 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) - def publish( - self, - handle: StreamHandle[Any], - values: tuple[Any, ...], - rpc_metadata: Mapping[str, str | bytes], - rpc_timeout: timedelta | None, - ) -> asyncio.Future[int]: - """Start the write Local Activity. An error that no caller gets fails the run.""" - _no_rpc_options(rpc_metadata, rpc_timeout) - state = _state() - future = _start_append(handle, 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", - handle.id, - error, - ) + future.add_done_callback(failed) + return future - 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 +async def close( + handle: StreamHandle[Any], result: Any, failure: BaseException | None +) -> None: + """Close the Stream after the writes before it, once the Workflow Task commits.""" + # The generated Workflow API makes the call, and the Worker's interceptor runs it as the + # write Local Activity. + await GeneratedStream(_generated_handle(handle)).close( + result=result, failure=failure + ) - def subscribe( - self, - handle: StreamHandle[Any], - *, - offset: int | End | None, - last: int | None, - wait_new_records: bool, - create_wait: timedelta | None, - rpc_metadata: Mapping[str, str | bytes], - rpc_timeout: timedelta | None, - ) -> Subscription[Any]: - """Raise, since Workflow code can't read a Stream.""" - raise no_workflow_reads() - - async def get_offset( - self, - handle: StreamHandle[Any], - rpc_metadata: Mapping[str, str | bytes], - rpc_timeout: timedelta | None, - ) -> int: - """Raise, since Workflow code can't describe a Stream.""" - raise no_workflow_reads() - - async def close( - self, - handle: StreamHandle[Any], - result: Any, - failure: BaseException | None, - rpc_metadata: Mapping[str, str | bytes], - rpc_timeout: timedelta | None, - ) -> None: - """Close the Stream after the writes before it, once the Workflow Task commits.""" - _no_rpc_options(rpc_metadata, rpc_timeout) - # The generated Workflow API makes the call, and the Worker's interceptor runs it as the - # write Local Activity. - await GeneratedStream(_generated_handle(handle)).close( - result=result, failure=failure - ) - async def enter(self, handle: StreamHandle[Any]) -> None: - """Check that this run chain can close the Stream at the end of the block.""" - # Fail before the block, since the close at its end would raise and hide the block's - # own error. - _check_owner(handle) +async def exit_block(handle: StreamHandle[Any], exc: BaseException | None) -> None: + """Close the Stream at the end of an ``async with`` block, or defer the close.""" + state = _state() + if state.stream(_model_handle(handle)).closed: + return + close_block, failure = exit_failure(exc) + if not close_block or temporalio.workflow.info().cron_schedule: + return + if _closes_at_block_end(exc): + await close(handle, None, failure) + else: + # A refused close raises here, as it does for a close at the block end. + _closing(handle) + state.deferred_closes.append((handle, failure)) - async def exit(self, handle: StreamHandle[Any], exc: BaseException | None) -> None: - """Close the Stream, or defer the close to the run's end.""" - state = _state() - if state.stream(_model_handle(handle)).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(handle, None, failure, {}, None) - else: - # A refused close raises here, as it does for a close at the block end. - _closing(handle) - state.deferred_closes.append((handle, failure)) - def publish_nowait( - self, handle: StreamHandle[Any], values: tuple[Any, ...] - ) -> None: - """Start the write. Its error raises in the run as an error that the run doesn't catch.""" - _write_nowait(_start_append(handle, values)) +def publish_nowait(handle: StreamHandle[Any], values: tuple[Any, ...]) -> None: + """Start the write. Its error raises in the run as an error that the run doesn't catch.""" + _write_nowait(_start_append(handle, values)) - def close_nowait( - self, handle: StreamHandle[Any], result: Any, failure: BaseException | None - ) -> None: - """Start the close. The end of the run waits for it.""" - _write_nowait(_start_close(handle, result, failure)) + +def close_nowait( + handle: StreamHandle[Any], result: Any, failure: BaseException | None +) -> None: + """Start the close. The end of the run waits for it.""" + _write_nowait(_start_close(handle, result, failure)) def _closes_at_block_end(exc: BaseException | None) -> bool: @@ -661,19 +601,21 @@ def _cancelled(exc: BaseException) -> bool: ) +def _own_id(definition: type[StreamDefinition[Any]]) -> str: + info = temporalio.workflow.info() + return _ids.scoped_stream_id( + info.workflow_id, info.first_execution_run_id, definition_of(definition).name + ) + + def _own_handle( definition: type[StreamDefinition[T]], backend: str | None ) -> StreamHandle[T]: info = temporalio.workflow.info() - known = definition_of(definition) return StreamHandle( - 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, + id=_own_id(definition), + owner=workflow_owner(info.workflow_id, info.first_execution_run_id), backend=backend, - payload_type=known.type, ) @@ -682,18 +624,18 @@ def create_stream( *, backend: str | None, options: StreamOptions, -) -> Creating[StreamHandle[T]]: +) -> 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() -> StreamHandle[T]: - stream = _creating(definition, backend) + async def create() -> WorkflowStream[T]: + handle = _creating(definition, backend) await generated_create.create_stream( - definition, retention=options.get("retention"), backend=stream.backend + definition, retention=options.get("retention"), backend=handle.backend ) - return stream + return WorkflowStream(handle) return Creating(create) @@ -703,17 +645,17 @@ def create_stream_nowait( *, backend: str | None, options: StreamOptions, -) -> StreamHandle[T]: +) -> 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 = _creating(definition, backend) + handle = _creating(definition, backend) call = generated_create.start_create_stream( - definition, retention=options.get("retention"), backend=stream.backend + definition, retention=options.get("retention"), backend=handle.backend ) _write_nowait(_call_future(call, lambda _: 0)) - return stream + return WorkflowStream(handle) def _creating( @@ -793,13 +735,7 @@ def operation_token(self) -> str | None: def _from_model(handle: models.StreamHandle, backend: str | None) -> StreamHandle[Any]: - owner = handle.owner - return StreamHandle( - id=handle.id, - workflow_id=owner.business_id if owner else None, - first_execution_run_id=owner.run_id if owner else None, - backend=backend, - ) + return StreamHandle(id=handle.id, owner=handle.owner, backend=backend) # The Worker's interceptor runs the Stream calls of the generated Workflow API with these. They @@ -850,19 +786,26 @@ def no_workflow_reads() -> _errors.StreamUnsupportedError: ) -def get_stream_handle( +def get_stream( definition: type[StreamDefinition[T]], *, backend: str | None = None -) -> StreamHandle[T]: - """Get a Stream that this run chain created, by its definition. +) -> WorkflowStream[T]: + """Get the run chain's own Stream of a definition. - See :py:func:`temporalio.workflow.get_stream_handle`. + See :py:func:`temporalio.workflow.get_stream`. """ - 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 + known = _state().streams.get(_own_id(definition)) + backend = (known.backend if known else None) or _check_backend(None, "get") + return WorkflowStream(_own_handle(definition, backend)) + + +def to_stream(handle: StreamHandle[T]) -> WorkflowStream[T]: + """Get the Stream of a handle that came as data. + + See :py:func:`temporalio.workflow.to_stream`. + """ + if not temporalio.workflow.unsafe.is_read_only(): + _state().uses_streams = True + return WorkflowStream(handle) diff --git a/temporalio/nexus/system/_support.py b/temporalio/nexus/system/_support.py index e7b7d0cbb2..5673ce84b9 100644 --- a/temporalio/nexus/system/_support.py +++ b/temporalio/nexus/system/_support.py @@ -556,7 +556,7 @@ def client_producer() -> typing.Any: streams = _activity.current.get() if streams is None: # Each client has its own Producer, and a generated call doesn't say which client makes - # it. So the SDK's handles send their Producer, and a call without one fails. + # it. So the SDK's Streams send their Producer, and a call without one fails. raise RuntimeError("A Stream write outside a Workflow needs its Producer") return models._ProducerTransferTypeConverter().from_transfer_type( streams.producer, models.Producer diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index 5efa835679..6c25ba4b84 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -152,9 +152,11 @@ unsafe, ) from ._streams import ( + WorkflowStream, create_stream, create_stream_nowait, - get_stream_handle, + get_stream, + to_stream, ) from ._workflow_ops import ( ChildWorkflowCancellationType, @@ -263,9 +265,11 @@ "SandboxImportNotificationPolicy", "logger", "unsafe", + "WorkflowStream", "create_stream", "create_stream_nowait", - "get_stream_handle", + "get_stream", + "to_stream", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py index 4388528eff..072041e45c 100644 --- a/temporalio/workflow/_streams.py +++ b/temporalio/workflow/_streams.py @@ -1,14 +1,16 @@ """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. +These functions and calls 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 +import asyncio +from types import TracebackType +from typing import TYPE_CHECKING, Any, Generic, TypeVar -from typing_extensions import Unpack +from typing_extensions import Self, Unpack if TYPE_CHECKING: from temporalio.common import ( @@ -21,12 +23,164 @@ T = TypeVar("T") +class WorkflowStream(Generic[T]): + """A Stream in Workflow code, with its calls. ``T`` is the payload type. + + Get one with :py:func:`create_stream`, :py:func:`get_stream`, or :py:func:`to_stream` for a + handle that came as data. Its :py:attr:`handle` is the data to pass on, such as to an + Activity, a child Workflow or the next run of continue-as-new, or through a query. + + Workflow code writes to and closes a Stream, and can't read it. Each publish is a Local + Activity, and each close a command that the backend runs after the Workflow Task commits. + + .. warning:: + This API is experimental. + """ + + def __init__(self, handle: StreamHandle[T]) -> None: + """Create a Stream. Use ``create_stream``, ``get_stream`` or ``to_stream`` instead.""" + self._handle = handle + + @property + def handle(self) -> StreamHandle[T]: + """The data that names the Stream, to pass on.""" + return self._handle + + 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. The writes to one Stream keep the order of the calls. A cancel + of the future doesn't stop the write. A list is one value. + + 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 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. + + Args: + values: One value for each Record. + + Returns: + The future of the write. It gives the offset of the last Record, or raises + ``StreamError`` when the backend refused the write, such as + ``StreamNotFoundError`` for a Stream that nobody created. + + Raises: + StreamClosedError: This run closed the Stream. + StreamConfigError: The Worker has no Stream backend or not the Stream's backend, + which fails the Workflow Task and not the run. + temporalio.workflow.ReadOnlyContextError: The Workflow code must not change state, + such as a query handler. + """ + from temporalio.common._streams import _workflow + + return _workflow.publish(self._handle, values) + + 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. + """ + from temporalio.common._streams import _workflow + + _workflow.publish_nowait(self._handle, values) + + async def close( + self, *, result: Any = None, failure: BaseException | None = None + ) -> None: + """Close the Stream after the writes before it, once the Workflow Task commits. + + Readers get the Records before the close, then a ``Closed`` item with ``result`` or + ``failure``. A cancel failure, such as a ``CancelledError``, closes the Stream as + canceled. + + 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. A cancel failure includes an ``ActivityError`` that a cancel caused. Later writes + in the run raise :py:class:`temporalio.common.StreamClosedError`. + + Raises: + StreamNotOwnerError: This run chain didn't create the Stream. Only the run chain + closes a Workflow's Stream, and only a client closes a client's Stream. It + raises with no call. + """ + from temporalio.common._streams import _workflow + + await _workflow.close(self._handle, result, failure) + + def close_nowait( + self, *, result: Any = None, failure: BaseException | None = None + ) -> None: + """Close the Stream without an await, in 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. + """ + from temporalio.common._streams import _workflow + + _workflow.close_nowait(self._handle, result, failure) + + async def __aenter__(self) -> Self: + """Use the Stream in ``async with``, which closes it at the end. + + Raises: + StreamNotOwnerError: This run chain can't close the Stream. See :py:meth:`close`. + """ + from temporalio.common._streams import _workflow + + # Fail before the block, since the close at its end would raise and hide the block's + # own error. + _workflow.check_owner(self._handle) + 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. + + If this run already closed the Stream, this does nothing. A cancel of the block or of + the run closes the Stream as canceled. Another ``BaseException``, such as + ``KeyboardInterrupt``, leaves it open. 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. + """ + from temporalio.common._streams import _workflow + + await _workflow.exit_block(self._handle, exc) + + def __repr__(self) -> str: + """Show the handle.""" + return f"{type(self).__name__}({self._handle!r})" + + def create_stream( definition: type[StreamDefinition[T]], *, backend: str | None = None, **options: Unpack[StreamOptions], -) -> Creating[StreamHandle[T]]: +) -> 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 @@ -39,7 +193,7 @@ def create_stream( 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 + Await the call to get the Stream, or use it in ``async with``, which closes the Stream at the end of the block:: async with workflow.create_stream(Tokens) as tokens: @@ -48,7 +202,7 @@ def create_stream( 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.StreamHandle.__aexit__`. + See :py:meth:`WorkflowStream.__aexit__`. .. warning:: This API is experimental. @@ -80,7 +234,7 @@ def create_stream_nowait( *, backend: str | None = None, **options: Unpack[StreamOptions], -) -> StreamHandle[T]: +) -> 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 @@ -96,15 +250,16 @@ def create_stream_nowait( return _workflow.create_stream_nowait(definition, backend=backend, options=options) -def get_stream_handle( +def get_stream( definition: type[StreamDefinition[T]], *, backend: str | None = None -) -> StreamHandle[T]: - """Get a Stream that this run chain created, by its definition. It makes no call. +) -> WorkflowStream[T]: + """Get the run chain's own Stream of a 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. + It's the Stream that :py:func:`create_stream` creates, so a query handler can hand out its + handle before the create. 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. @@ -119,4 +274,18 @@ def get_stream_handle( """ from temporalio.common._streams import _workflow - return _workflow.get_stream_handle(definition, backend=backend) + return _workflow.get_stream(definition, backend=backend) + + +def to_stream(handle: StreamHandle[T]) -> WorkflowStream[T]: + """Get the Stream of a handle that came as data, such as an argument. It makes no call. + + The Stream's calls check its backend. A Worker that doesn't have it fails the Workflow Task + at the first call. + + .. warning:: + This API is experimental. + """ + from temporalio.common._streams import _workflow + + return _workflow.to_stream(handle) diff --git a/tests/contrib/streams/_support.py b/tests/contrib/streams/_support.py index 2faeeed8aa..65d479e8e6 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, and binds a Workflow's Stream in tests.""" +"""Finds the Stream crates that the bridge links, and gets a Workflow's Stream in tests.""" from __future__ import annotations @@ -11,7 +11,7 @@ import pytest from temporalio import common -from temporalio.client import WorkflowHandle +from temporalio.client import Stream, WorkflowHandle from temporalio.common._streams import _ids from temporalio.common._streams._defn import definition_of @@ -53,10 +53,10 @@ def workflow_stream( handle: WorkflowHandle[Any, Any], definition: type[common.StreamDefinition[T]], backend: str | None = None, -) -> common.StreamHandle[T]: +) -> Stream[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.get_stream_handle(definition, id=stream_id, backend=backend) + return handle._client.get_stream(definition, id=stream_id, backend=backend) diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index deb38c9d55..71d438cc0c 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -3,7 +3,7 @@ from typing import Any from temporalio import common, workflow -from temporalio.client import Client +from temporalio.client import Client, Stream from temporalio.worker import ( UnsandboxedWorkflowRunner, ) @@ -15,7 +15,7 @@ class Tokens(common.StreamDefinition[str]): ... -async def read_all(stream: common.StreamHandle[Any]) -> list[Any]: +async def read_all(stream: Stream[Any]) -> list[Any]: return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 790b38e3b3..20815c438b 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -146,8 +146,8 @@ async def test_query_and_workflow_results_give_handles(client: Client): @workflow.defn class WriteToStream: @workflow.run - async def run(self, stream: common.StreamHandle[Turn]) -> None: - await stream.publish(Turn("workflow")) + async def run(self, handle: common.StreamHandle[Turn]) -> None: + await workflow.to_stream(handle).publish(Turn("workflow")) async def test_client_stream_takes_workflow_writes(stream_client: Client): @@ -159,7 +159,7 @@ async def test_client_stream_takes_workflow_writes(stream_client: Client): async with new_worker(stream_client, WriteToStream) as worker: await stream_client.execute_workflow( WriteToStream.run, - stream, + stream.handle, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -169,8 +169,8 @@ async def test_client_stream_takes_workflow_writes(stream_client: Client): 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.get_stream_handle(Turns, id=stream.id) + # Another Stream of the same id reads it from an offset. + again = stream_client.get_stream(Turns, id=stream.handle.id) resumed = [ item.data async for item in again.subscribe(offset=first + 1) @@ -192,7 +192,7 @@ async def run(self, prompt: str) -> str: @workflow.query def stream(self) -> common.StreamHandle[str]: # The handle needs no call, so a reader gets it before the create. - return workflow.get_stream_handle(Tokens) + return workflow.get_stream(Tokens).handle async def test_client_reads_a_workflow_stream(stream_client: Client): @@ -200,7 +200,7 @@ async def test_client_reads_a_workflow_stream(stream_client: Client): handle = await stream_client.start_workflow( Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - reader = await handle.query(Agent.stream) + reader = stream_client.to_stream(await handle.query(Agent.stream)) # 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)) @@ -210,7 +210,8 @@ async def test_client_reads_a_workflow_stream(stream_client: Client): "there", ] assert isinstance(items[-1], common.Closed) and items[-1].result == "HI THERE" - assert reader.workflow_id == handle.id + owner = reader.handle.owner + assert owner is not None and owner.business_id == handle.id with pytest.raises(common.StreamNotOwnerError): await reader.close() @@ -420,7 +421,7 @@ async def test_read_keeps_no_offsets_of_client_records(stream_client: Client): 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.get_stream_handle(Turns, id="") + client.get_stream(Turns, id="") assert not extension.namespaces @@ -475,8 +476,8 @@ async def test_registration_rules(client: Client): MemoryStreamBackend(name="b", default=True), ] two = Client(**config) - assert (await two.create_stream(Tokens, id="d-1")).backend == "b" - assert two.get_stream_handle(Tokens, id="d-1").backend == "b" + assert (await two.create_stream(Tokens, id="d-1")).handle.backend == "b" + assert two.get_stream(Tokens, id="d-1").handle.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) diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py index 0dee7e23d3..b4d64d2e14 100644 --- a/tests/contrib/streams/test_workflow_module.py +++ b/tests/contrib/streams/test_workflow_module.py @@ -21,6 +21,7 @@ from temporalio.common import RetryPolicy from temporalio.common._streams import _workflow, _writes from temporalio.common._streams._backend import _ExtensionCalls +from temporalio.common._streams._handle import workflow_owner from temporalio.common._streams.system._resources.stream import ( stream as generated_stream, ) @@ -100,7 +101,9 @@ class GeneratedCalls: @workflow.run async def run(self) -> list[str]: tokens = await workflow.create_stream(Tokens) - generated = generated_stream(Tokens, handle=_workflow._generated_handle(tokens)) + generated = generated_stream( + Tokens, handle=_workflow._generated_handle(tokens.handle) + ) offset = await generated.publish("a", "b") refused = [] try: @@ -132,28 +135,21 @@ class RefusedCalls: async def run(self) -> list[str]: tokens = await workflow.create_stream(Tokens) refused = [] - for call in ( - lambda: tokens.subscribe(), - lambda: tokens.producer("ingest", 1), - lambda: tokens.publish("a", rpc_timeout=timedelta(seconds=1)), - ): - try: - call() - except common.StreamUnsupportedError as err: - refused.append(type(err).__name__) - try: - await tokens.get_offset() - except common.StreamUnsupportedError as err: - refused.append(type(err).__name__) - # Another run chain's Stream fails before the block runs. - other: common.StreamHandle[str] = common.StreamHandle( - id="other:run-1:tokens", workflow_id="other", first_execution_run_id="run-1" + # Another run chain's Stream fails before the block runs, and before any call. + other = workflow.to_stream( + common.StreamHandle[str]( + id="other:run-1:tokens", owner=workflow_owner("other", "run-1") + ) ) try: async with other: refused.append("entered") except common.StreamNotOwnerError as err: refused.append(type(err).__name__) + try: + await other.close() + except common.StreamNotOwnerError as err: + refused.append(type(err).__name__) await tokens.close() return refused @@ -163,10 +159,7 @@ async def test_calls_that_workflow_code_cant_make_raise(stream_client: Client): result = await stream_client.execute_workflow( RefusedCalls.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue ) - assert result == [ - *["StreamUnsupportedError"] * 4, - "StreamNotOwnerError", - ] + assert result == ["StreamNotOwnerError", "StreamNotOwnerError"] @common.defn(name="first") @@ -209,13 +202,13 @@ async def test_async_with_closes_the_stream(stream_client: Client): @workflow.defn class Continuing: @workflow.run - async def run(self, stream: common.StreamHandle[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.get_stream_handle(Tokens) + async def run(self, handle: common.StreamHandle[str] | None) -> str: + if handle is None: + created = await workflow.create_stream(Tokens) + await created.publish("run 1") + workflow.continue_as_new(created.handle) + assert handle == workflow.get_stream(Tokens).handle + stream = workflow.to_stream(handle) await stream.publish("run 2") await stream.close(result="done") return "done" @@ -233,7 +226,7 @@ async def run(self, first: bool) -> str: if first: workflow.continue_as_new(False) await stream.close() - return stream.id + return stream.handle.id async def test_create_in_the_next_run_gets_the_same_stream(stream_client: Client): @@ -246,7 +239,7 @@ async def test_create_in_the_next_run_gets_the_same_stream(stream_client: Client ) stream_id = await handle.result() items = await read(handle) - assert stream_id == workflow_stream(handle, Tokens).id + assert stream_id == workflow_stream(handle, Tokens).handle.id assert data(items) == ["run 1", "run 2"] assert isinstance(items[-1], common.Closed) @@ -266,7 +259,7 @@ async def test_continue_as_new_passes_the_stream(stream_client: Client): @workflow.defn class Mistakes: def __init__(self) -> None: - self._stream: common.StreamHandle[str] | None = None + self._stream: workflow.WorkflowStream[str] | None = None @workflow.run async def run(self) -> list[str]: @@ -317,7 +310,8 @@ async def test_replay_reproducible_errors_raise_in_workflow_code( @workflow.defn class ClosesAReceivedStream: @workflow.run - async def run(self, stream: common.StreamHandle[str]) -> str: + async def run(self, handle: common.StreamHandle[str]) -> str: + stream = workflow.to_stream(handle) await stream.publish("from workflow") try: await stream.close(result="ended by a writer") @@ -332,7 +326,7 @@ async def test_workflow_closes_only_its_own_streams(stream_client: Client): assert ( await stream_client.execute_workflow( ClosesAReceivedStream.run, - stream, + stream.handle, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -735,7 +729,7 @@ 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.StreamHandle[str](id=stream_id) + missing = workflow.to_stream(common.StreamHandle[str](id=stream_id)) future = missing.publish("lost") errors: list[BaseException | None] = [] if retrieve: @@ -912,7 +906,7 @@ 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.get_stream_handle(Tokens) + stream = workflow.get_stream(Tokens) await stream.publish("x" * 100_000) @@ -925,7 +919,7 @@ async def test_refused_write_keeps_no_pending_bytes(stream_client: Client): ) with pytest.raises(WorkflowFailureError): await handle.result() - stream_id = workflow_stream(handle, Tokens).id + stream_id = workflow_stream(handle, Tokens).handle.id assert not [key for key in _writes._PENDING._entries if key[1] == stream_id] @@ -1017,8 +1011,8 @@ class CreatesOnANamedBackend: async def run(self, backend: str) -> str: stream = await workflow.create_stream(Tokens, backend=backend) await stream.publish(f"on {backend}") - await workflow.get_stream_handle(Tokens).close(result="closed") - return stream.backend or "" + await workflow.get_stream(Tokens).close(result="closed") + return stream.handle.backend or "" @workflow.defn @@ -1148,7 +1142,7 @@ async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: @workflow.defn class WritesFromASignalAtTheEnd: def __init__(self) -> None: - self.late: common.StreamHandle[str] | None = None + self.late: workflow.WorkflowStream[str] | None = None @workflow.run async def run(self) -> str: