From a912bfb1fd489685d783460649726f65a086a7ba Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Thu, 5 Nov 2026 01:58:35 -0700 Subject: [PATCH] Added Workflow Streams, which write through Local Activities. --- temporalio/contrib/stream/__init__.py | 15 +- temporalio/contrib/stream/testing.py | 4 +- temporalio/nexus/system/_support.py | 69 +- temporalio/streams/_activity.py | 60 + temporalio/streams/_extension.py | 423 +++++- temporalio/streams/_ids.py | 20 +- temporalio/streams/_reader.py | 4 +- temporalio/streams/_workflow.py | 762 ++++++++++ temporalio/streams/_writes.py | 265 ++++ temporalio/workflow/__init__.py | 14 + temporalio/workflow/_streams.py | 348 +++++ tests/contrib/streams/_support.py | 17 +- tests/contrib/streams/test_redis.py | 69 + tests/contrib/streams/test_stream_client.py | 100 ++ tests/contrib/streams/test_workflow_module.py | 1320 +++++++++++++++++ 15 files changed, 3459 insertions(+), 31 deletions(-) create mode 100644 temporalio/streams/_activity.py create mode 100644 temporalio/streams/_workflow.py create mode 100644 temporalio/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/contrib/stream/__init__.py b/temporalio/contrib/stream/__init__.py index f61b3f5e80..1e8dc6667f 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:attr:`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:data:`temporalio.workflow.create_stream` and :py:attr:`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 on that queue +that polls Activities must run the close, and must register the Redis backend. A Worker with +``no_remote_activities=True`` logs a warning, since the close waits for another Worker. + .. warning:: This package is experimental. """ @@ -34,6 +38,11 @@ class RedisStreamBackend(ExtensionStreamBackend): each Stream ref that it creates keeps 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 on that queue that polls Activities must have this + backend. A Worker with ``no_remote_activities=True`` logs a warning, since the close waits + for another Worker. + .. warning:: This API is experimental. """ diff --git a/temporalio/contrib/stream/testing.py b/temporalio/contrib/stream/testing.py index cc9edb9e44..aa7b4571c3 100644 --- a/temporalio/contrib/stream/testing.py +++ b/temporalio/contrib/stream/testing.py @@ -19,7 +19,9 @@ 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 its ``name``, ``memory`` by default. It has Redis's rules, so tests can swap the - Store without changing Workflow code. + Store without changing Workflow code. As on Redis, a Workflow's close is a regular Activity + on the Workflow's task queue, so a Worker on that queue that polls Activities must have the + backend. .. warning:: This API is experimental. diff --git a/temporalio/nexus/system/_support.py b/temporalio/nexus/system/_support.py index 5b8eed4b39..d984f708eb 100644 --- a/temporalio/nexus/system/_support.py +++ b/temporalio/nexus/system/_support.py @@ -490,51 +490,74 @@ def wake(_: object, done: asyncio.Future[None] = done) -> None: def request_id() -> str: """The id of one create call, which its retries send again. - In a Workflow it comes from ``workflow.uuid4()``, so a replay sends the same id. + Outside a Workflow it's a random UUID. In a Workflow it's empty, since the Workflow create + sets an id that is the same in each run of the chain. """ - from temporalio.workflow import in_workflow, uuid4 + from temporalio.workflow import in_workflow - return str(uuid4() if in_workflow() else uuid.uuid4()) + # A workflow.uuid4() here would move the user's workflow.random() sequence, and the + # Workflow create replaces the id anyway. + return "" if in_workflow() else str(uuid.uuid4()) # 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. The Streams API imports this module, so the sources import +# the API only when they run. def workflow_producer() -> typing.Any: - """The Producer of the Workflow run's writes.""" - raise _no_workflow_streams("workflow_producer") + """The Workflow Producer of the run.""" + from temporalio.streams import _ids + from temporalio.workflow import info + + return _ids.workflow_producer(info().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 Streams send their Producer, and a call without one fails. - raise RuntimeError("A Stream write outside a Workflow needs its Producer") + from temporalio.streams import _activity + from temporalio.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 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 + ) 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.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.streams import _activity + from temporalio.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/temporalio/streams/_activity.py b/temporalio/streams/_activity.py new file mode 100644 index 0000000000..c1a9b9bff5 --- /dev/null +++ b/temporalio/streams/_activity.py @@ -0,0 +1,60 @@ +"""The Stream state of one Activity attempt.""" + +from __future__ import annotations + +import asyncio +import contextvars +from dataclasses import dataclass, field + +import temporalio.api.stream.v1 as stream_v1 +import temporalio.converter +from temporalio.streams._backend import BackendCalls +from temporalio.streams._writer import StreamWriter + + +@dataclass +class ActivityStreams: + """What a Stream backend gives the Streams of one Activity attempt.""" + + backend: BackendCalls + converter: temporalio.converter.DataConverter + namespace: str + producer: stream_v1.Producer + loop: asyncio.AbstractEventLoop + # The attempt's Workflow ID, which the close Activity checks against a bound 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) + + 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: + """Close every writer of the attempt.""" + errors: list[BaseException] = [] + 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/streams/_extension.py b/temporalio/streams/_extension.py index a0986b3dee..05df874b9e 100644 --- a/temporalio/streams/_extension.py +++ b/temporalio/streams/_extension.py @@ -2,24 +2,53 @@ from __future__ import annotations +import asyncio +import dataclasses import logging import weakref from collections.abc import Sequence -from typing import Any +from datetime import timedelta +from typing import Any, TypeVar +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.streams import _activity, _ids, _writes +from temporalio.streams import _workflow as streams_workflow from temporalio.streams._backend import BackendCalls, _ExtensionCalls from temporalio.streams._errors import StreamConfigError from temporalio.streams._plugin import StreamBackend +from temporalio.streams._ref import StreamRef from temporalio.streams.proto import v1 as ext_pb +from temporalio.streams.system import models +from temporalio.streams.system.workflow import ( + _system_nexus_interceptor as system_nexus, +) +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" +_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 +87,7 @@ def __init__( name, backend_name=backend_name, default=default, + activities=_writes.ACTIVITIES, ) def _new_backend( @@ -107,6 +137,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.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 +249,318 @@ 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: + _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) + + +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 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) + + 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__( + 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 = [ref.id for ref 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[StreamRef[Any]]: + # Find the Stream refs in call arguments. A ref can be in a field, a slot, a list or a dict. + refs: list[StreamRef[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, StreamRef): + refs.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 refs + + +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 "" + previous: stream_v1.WorkflowTaskEpoch | None = None + failed: tuple[int, ...] = () + if info.is_local: + local = stream_v1.LocalActivityEpoch(attempt=info.attempt) + 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) + # The dispatching task orders the epoch first, so a later task's write withdraws a + # failed task's Records. + local.task.started_event_id = task.started_event_id + local.task.attempt = task.attempt + if task.started_time is not None: + local.task.started_time.FromDatetime(task.started_time) + run_id = info.workflow_run_id or "" + producer = stream_v1.Producer( + id=_ids.local_activity_producer_id(run_id, key_text), + epoch=stream_v1.ProducerEpoch(local_activity=local), + run_id=run_id, + ) + else: + attempt = stream_v1.ActivityEpoch(attempt=info.attempt) + attempt.attempt_scheduled_time.FromDatetime(info.current_attempt_scheduled_time) + if info.in_workflow: + run_id = info.workflow_run_id or "" + producer = stream_v1.Producer( + id=_ids.activity_producer_id(run_id, info.activity_id, key_text), + epoch=stream_v1.ProducerEpoch(activity=attempt), + run_id=run_id, + ) + else: + producer = stream_v1.Producer( + id=_ids.standalone_activity_producer_id( + info.activity_id, info.activity_run_id or "" + ), + epoch=stream_v1.ProducerEpoch(standalone_activity=attempt), + ) + return _activity.ActivityStreams( + backend=interceptor.backend, + converter=interceptor.converter, + namespace=info.namespace, + producer=producer, + loop=asyncio.get_running_loop(), + workflow_id=info.workflow_id, + previous_completed_task=previous, + failed_tasks=failed, + ) diff --git a/temporalio/streams/_ids.py b/temporalio/streams/_ids.py index 1924e9b4c2..bf200c2a97 100644 --- a/temporalio/streams/_ids.py +++ b/temporalio/streams/_ids.py @@ -1,9 +1,11 @@ -"""Producer ids. They must match the temporalio-streams crate byte for byte.""" +"""Producers and their ids. The ids must match the temporalio-streams crate byte for byte.""" from __future__ import annotations import urllib.parse +from temporalio.streams.system import models + def encode_part(text: str) -> str: """Percent-encode every byte but the unreserved ones.""" @@ -19,6 +21,22 @@ def workflow_producer_id(run_id: str) -> str: return f"wf:{_join(run_id)}" +def workflow_producer(run_id: str) -> models.Producer: + """The Workflow Producer of a run. + + The run writes through Local Activities, which send their own Producers, so its Workflow + Producer only closes. Its epoch is the Workflow case with no task, which fences nothing, + since that Producer writes no Record. + """ + return models.Producer( + id=workflow_producer_id(run_id), + run_id=run_id, + epoch=models.ProducerEpoch( + kind=models.ProducerEpochKindWorkflow(models.WorkflowTaskEpoch()) + ), + ) + + def local_activity_producer_id(run_id: str, call_key: str) -> str: """Producer id of a Local Activity call.""" return f"la:{_join(run_id, call_key)}" diff --git a/temporalio/streams/_reader.py b/temporalio/streams/_reader.py index f0f26d4ad5..b101476f50 100644 --- a/temporalio/streams/_reader.py +++ b/temporalio/streams/_reader.py @@ -65,7 +65,9 @@ def __init__( # A client read's rpc_metadata and rpc_timeout, for each call of each step. self._options = options # The Producer and the offsets of each Record that this read gave, without values, so a - # Superseded item can list the offsets that it withdraws. + # Superseded item can list the offsets that it withdraws. Each write Local Activity of a + # run is its own Producer, so a read of a Workflow's writes keeps one entry for each + # write, not for each Record, and a Superseded item scans those entries. self._shown: dict[ tuple[Any, ...], tuple[models.Producer, array.array[int]] ] = {} diff --git a/temporalio/streams/_workflow.py b/temporalio/streams/_workflow.py new file mode 100644 index 0000000000..cbf36eef2f --- /dev/null +++ b/temporalio/streams/_workflow.py @@ -0,0 +1,762 @@ +"""Streams in Workflow code. + +Workflow code creates a Stream with :py:data:`temporalio.workflow.create_stream`. A ref that the +Workflow receives as data works the same through :py:func:`temporalio.workflow.to_stream`, such +as a ref 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 Callable, Generator +from dataclasses import dataclass, field +from datetime import timedelta +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.exceptions +import temporalio.nexus.system._support as _support +import temporalio.workflow +from temporalio.streams import _errors +from temporalio.streams._ref import StreamOptions, StreamRef +from temporalio.streams._writer import WriteFuture, unretrieved_error +from temporalio.streams.system import models +from temporalio.streams.system.workflow._resources import ( + Stream as GeneratedStream, +) +from temporalio.streams.system.workflow._resources import ( + stream as generated_stream, +) +from temporalio.streams.system.workflow.operations import ( + create_stream as generated_create, +) +from temporalio.workflow._streams import WorkflowStream + +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. +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 refs can identify it, so the state is keyed by + # id. + ref: models.StreamRef + closed: bool = False + + +@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.ref != request.ref + ): + 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 append 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. + deferred_closes: list[StreamRef[Any]] = 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) + # 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) + 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[T](id, 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 + + +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 + if isinstance(request, stream_pb.AppendRecordsRequest) and ( + stream_id not in state.withdrawn + ): + previous_run = _support.withdraw_run() + if previous_run: + # A retry run's first append to a Stream withdraws the Records of the run it + # retries. The Store refuses a run withdrawal on a create with no owner, and a + # Workflow's Stream has none, so only an append sends it. + request.withdraw_run = previous_run + 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(), *state.calls) 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 _payload(value: Any) -> Any: + # 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").to_payload(value) + + +def _model_ref(ref: StreamRef[Any]) -> models.StreamRef: + return models.StreamRef(id=ref.id, owner=ref.owner) + + +def _proto_ref(ref: StreamRef[Any]) -> stream_v1.StreamRef: + return models._StreamRefTransferTypeConverter().to_transfer_type(_model_ref(ref)) + + +def _generated_ref(ref: StreamRef[Any]) -> models.StreamRef: + return models.StreamRef(id=ref.id, owner=ref.owner, backend=ref.backend or None) + + +def _check( + ref: StreamRef[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_ref(ref)) + if run_stream.closed: + raise _errors.StreamClosedError( + f"Can't {action} the Stream {ref.id}, since this run closed it" + ) + if not backend_checked: + _check_backend(ref.backend, action) + return run_stream + + +def _start_append(ref: StreamRef[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(ref, "write to") + call = generated_stream.start_append_records( + ref=_generated_ref(ref), + records=[models.AppendRecord(body=value) for value in values], + ) + return _call_future(call, lambda response: response.last_offset or 0) + + +def _append(ref: StreamRef[Any], values: tuple[Any, ...]) -> asyncio.Future[int]: + if not values: + raise TypeError("publish takes one or more values") + _check(ref, "write to") + records = [stream_v1.AppendRecord(body=_payload(v)) for v in values] + return _write( + ref.id, + APPEND_ACTIVITY, + stream_pb.AppendRecordsRequest(ref=_proto_ref(ref), records=records), + ) + + +def _start_close(ref: StreamRef[Any]) -> asyncio.Future[int]: + _closing(ref) + call = generated_stream.start_close_stream(ref=_generated_ref(ref)) + return _call_future(call, lambda _: 0) + + +def check_owner(ref: StreamRef[Any]) -> None: + """Check that this run chain may close the Stream. + + A Workflow closes a standalone Stream. Only its own run chain closes a bound Stream. + """ + owner = ref.owner + if owner is None: + return + info = temporalio.workflow.info() + if (owner.business_id, owner.run_id) != ( + info.workflow_id, + info.first_execution_run_id, + ): + # 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 {ref.id}, since it is bound to another run chain" + ) + + +def _closing(ref: StreamRef[Any]) -> _RunStream: + check_owner(ref) + return _check(ref, "close") + + +def _close(ref: StreamRef[Any]) -> asyncio.Future[int]: + _closing(ref).closed = True + return _write( + ref.id, CLOSE_ACTIVITY, stream_pb.CloseStreamRequest(ref=_proto_ref(ref)) + ) + + +def publish(ref: StreamRef[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(ref, 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", + ref.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 + + +async def close(ref: StreamRef[Any]) -> 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 + # close write. + await GeneratedStream(_generated_ref(ref), None).close() + + +async def exit_block(ref: StreamRef[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_ref(ref)).closed: + return + # A run that continues as new has no retry run, so its block end closes the Stream, though + # ContinueAsNewError isn't an Exception. + continued = isinstance(exc, temporalio.workflow.ContinueAsNewError) + if not (continued or _support.exit_close(ref, exc)): + return + if temporalio.workflow.info().cron_schedule: + return + if continued or _closes_at_block_end(exc): + await close(ref) + else: + # A refused close raises here, as it does for a close at the block end. + _closing(ref) + state.deferred_closes.append(ref) + + +def publish_nowait(ref: StreamRef[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(ref, values)) + + +def close_nowait(ref: StreamRef[Any]) -> None: + """Start the close. The end of the run waits for it.""" + _write_nowait(_start_close(ref)) + + +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 ref in state.deferred_closes: + if not state.stream(_model_ref(ref)).closed: + _write_nowait(_start_close(ref)) + + +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 _conflict_policy(or_get: bool) -> models.StreamIdConflictPolicy: + return ( + models.StreamIdConflictPolicy.StreamIdConflictPolicyUseExisting + if or_get + else models.StreamIdConflictPolicy.StreamIdConflictPolicyFail + ) + + +async def create_stream( + payload_type: type[T] | None, + id: str, + *, + backend: str | None, + options: StreamOptions, + or_get: bool, +) -> WorkflowStream[T]: + """Create a standalone Stream, or with ``or_get`` get the one that exists. + + See :py:data:`temporalio.workflow.create_stream`. + """ + ref: StreamRef[T] = _creating(id, backend) + await generated_create.create_stream( + payload_type or object, + id=id, + backend=ref.backend, + retention=options.get("retention"), + id_conflict_policy=_conflict_policy(or_get), + ) + return WorkflowStream(ref) + + +def create_stream_nowait( + payload_type: type[T] | None, + id: str, + *, + backend: str | None, + options: StreamOptions, + or_get: bool, +) -> WorkflowStream[T]: + """Create a standalone Stream, or get it, without an await. + + See :py:data:`temporalio.workflow.create_stream_nowait`. + """ + ref: StreamRef[T] = _creating(id, backend) + call = generated_create.start_create_stream( + payload_type or object, + id=id, + backend=ref.backend, + retention=options.get("retention"), + id_conflict_policy=_conflict_policy(or_get), + ) + _write_nowait(_call_future(call, lambda _: 0)) + return WorkflowStream(ref) + + +def _creating(id: str, backend: str | None) -> StreamRef[Any]: + # The checks raise before the generated call. A Workflow creates only standalone Streams, so + # the ref has no owner. + if not id: + raise ValueError("The Stream id is empty") + ref: StreamRef[Any] = StreamRef(id=id, backend=_check_backend(backend, "create")) + _check(ref, "create", backend_checked=True) + return ref + + +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( + ref: StreamRef[Any], request: models.CreateStreamRequest +) -> asyncio.Future[int]: + _check(ref, "create", backend_checked=True) + proto = models._CreateStreamRequestTransferTypeConverter().to_transfer_type(request) + proto.stream_ref.CopyFrom(_proto_ref(ref)) + # The Store refuses a run withdrawal on a create with no owner, and a Workflow's Stream has + # none. The run's first append withdraws the previous run instead. + proto.ClearField("withdraw_run") + # A later run of the chain repeats the same create, so the id is the same in each run. The + # Store then gives a retry run or a cron run the Stream that the first run made. + info = temporalio.workflow.info() + proto.request_id = f"{info.first_execution_run_id}:{ref.id}" + return _write(ref.id, CREATE_ACTIVITY, proto) + + +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_model(ref: models.StreamRef, backend: str | None) -> StreamRef[Any]: + return StreamRef(id=ref.id, owner=ref.owner, backend=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 Streams, 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.""" + model_ref = request.stream_ref + stream = _from_model(model_ref, _check_backend(model_ref.backend, "create")) + write = _start_create(stream, request) + return _WriteHandle( + write, lambda _: models.CreateStreamResult(ref=_generated_ref(stream)) + ) + + +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_model(request.ref, request.ref.backend) + write = _append(stream, tuple(record.body for record in request.records)) + 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.""" + stream = _from_model(request.ref, request.ref.backend) + write = _close(stream) + 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 to_stream(ref: StreamRef[T]) -> WorkflowStream[T]: + """Get the Stream of a ref 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(ref) diff --git a/temporalio/streams/_writes.py b/temporalio/streams/_writes.py new file mode 100644 index 0000000000..15931b86f6 --- /dev/null +++ b/temporalio/streams/_writes.py @@ -0,0 +1,265 @@ +"""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.streams import _activity, _backend, _errors, _ids +from temporalio.streams._workflow import ( + APPEND_ACTIVITY, + CLOSE_ACTIVITY, + CREATE_ACTIVITY, +) +from temporalio.streams.system import models + + +def _context() -> tuple[_activity.ActivityStreams, _backend.StreamServiceCalls]: + 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 + local = producer.epoch.local_activity + started = local.task.started_time if local.task.HasField("started_time") else None + task = models.WorkflowTaskEpoch( + started_event_id=local.task.started_event_id or None, + attempt=local.task.attempt or None, + started_time=models.Timestamp(seconds=started.seconds, nanos=started.nanos) + if started is not None + else None, + ) + return models.Producer( + id=producer.id, + run_id=producer.run_id or None, + epoch=models.ProducerEpoch( + kind=models.ProducerEpochKindLocalActivity( + models.LocalActivityEpoch(task=task, 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 _model_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 standalone Stream, or with ``USE_EXISTING`` get it.""" + _, service = _context() + # The request has the Workflow's request id, so a retry of this Activity isn't a conflict. A + # closed, expired or another owner's Stream fails the create. + await service._stream_service_create_stream( + models._CreateStreamRequestTransferTypeConverter().from_transfer_type( + request, models.CreateStreamRequest + ), + temporalio.common.RawValue, + ) + 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.ref.id, producer_pb.id) + task = producer_pb.epoch.local_activity.task + epoch = ( + task.started_event_id, + task.attempt, + task.started_time.seconds, + 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 service.send_append_records( + models.AppendRecordsRequest( + ref=_model_ref(request.ref), + 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 not isinstance(err, _errors.StreamTooLargeError): + if err.non_retryable: + _PENDING.pop(key, epoch) + raise + 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. + 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 Stream as the Workflow run that scheduled this Activity.""" + streams, service = _context() + owner = request.ref.owner if request.ref.HasField("owner") else None + if owner is not None and owner.business_id != streams.workflow_id: + # The Store refuses an Activity Producer's close, so the Activity closes as the run. It + # refuses a Stream bound to another Workflow, and Workflow code checks the run chain. + raise _errors.StreamNotOwnerError( + f"Can't close the Stream {request.ref.id} for another Workflow" + ) + await service.send_close_stream( + models.CloseStreamRequest( + ref=_model_ref(request.ref), + producer=_ids.workflow_producer(streams.producer.run_id), + ) + ) + return 0 + + +ACTIVITIES = [create_stream, append_records, close_stream] diff --git a/temporalio/workflow/__init__.py b/temporalio/workflow/__init__.py index c631fd88e8..d55d65996a 100644 --- a/temporalio/workflow/__init__.py +++ b/temporalio/workflow/__init__.py @@ -151,6 +151,14 @@ logger, unsafe, ) +from ._streams import ( + WorkflowStream, + create_or_get_stream, + create_or_get_stream_nowait, + create_stream, + create_stream_nowait, + to_stream, +) from ._workflow_ops import ( ChildWorkflowCancellationType, ChildWorkflowConfig, @@ -258,6 +266,12 @@ "SandboxImportNotificationPolicy", "logger", "unsafe", + "WorkflowStream", + "create_or_get_stream", + "create_or_get_stream_nowait", + "create_stream", + "create_stream_nowait", + "to_stream", "ChildWorkflowCancellationType", "ChildWorkflowConfig", "ChildWorkflowHandle", diff --git a/temporalio/workflow/_streams.py b/temporalio/workflow/_streams.py new file mode 100644 index 0000000000..f89358fc52 --- /dev/null +++ b/temporalio/workflow/_streams.py @@ -0,0 +1,348 @@ +"""Streams in Workflow code. + +These calls import the Streams API of :py:mod:`temporalio.streams` on first use, so a Workflow +with no Stream never imports it. +""" + +from __future__ import annotations + +import asyncio +from types import TracebackType +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +from typing_extensions import Self, Unpack + +if TYPE_CHECKING: + import temporalio.api.stream.v1 + from temporalio.streams import StreamOptions, StreamRef + +T = TypeVar("T") +_U = TypeVar("_U") + + +class WorkflowStream(Generic[T]): + """A Stream in Workflow code, with its calls. ``T`` is the payload type. + + Get one with :py:data:`create_stream`, :py:data:`create_or_get_stream`, or + :py:func:`to_stream` for a ref that came as data. Its :py:attr:`ref` 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, ref: StreamRef[T]) -> None: + """Create a Stream. Use ``create_stream`` or ``to_stream`` instead.""" + self._ref = ref + + @property + def ref(self) -> StreamRef[T]: + """The data that identifies the Stream, to pass on.""" + return self._ref + + @property + def proto_ref(self) -> temporalio.api.stream.v1.StreamRef: + """The ``StreamRef`` wire message of the Stream, for a field of a proto message.""" + from temporalio.streams import _ref + + return _ref.to_proto(self._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. 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.streams import _workflow + + return _workflow.publish(self._ref, 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.streams import _workflow + + _workflow.publish_nowait(self._ref, values) + + async def close(self) -> None: + """Close the Stream after the writes before it, once the Workflow Task commits. + + A close carries no value, so code with a last value publishes it first. Readers get the + Records before the close, then their read ends. Later writes in the run raise + :py:class:`temporalio.streams.StreamClosedError`. + + 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's Workflow Producer. 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. + + Raises: + StreamNotOwnerError: The Stream is bound to another run chain, which alone closes + it. It raises with no call. + """ + from temporalio.streams import _workflow + + await _workflow.close(self._ref) + + def close_nowait(self) -> 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.streams import _workflow + + _workflow.close_nowait(self._ref) + + 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.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._ref) + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream at the end of the block. + + If this run already closed the Stream, this does nothing. A cancel of the block or of + the run closes the Stream. 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 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. 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.streams import _workflow + + await _workflow.exit_block(self._ref, exc) + + def __repr__(self) -> str: + """Show the ref.""" + return f"{type(self).__name__}({self._ref!r})" + + +# Python has no type arguments on a call, so a call that takes the payload type is an object +# whose subscript binds the type, as in workflow.create_stream[Turn]("session-7-turns"). Type +# checkers then see WorkflowStream[Turn]. With no subscript, T is Any. + + +class CreateStream(Generic[T]): + """The call of :py:data:`create_stream` and :py:data:`create_or_get_stream`. + + Subscript it with the payload type, as in ``workflow.create_stream[Turn]("turns-7")``. + + .. warning:: + This API is experimental. + """ + + def __init__(self, payload_type: type[T] | None, *, or_get: bool) -> None: + """Create the call. Use the module's ``create_stream`` instead.""" + self._payload_type = payload_type + self._or_get = or_get + + def __getitem__(self, payload_type: type[_U]) -> CreateStream[_U]: + """The call for Streams of ``payload_type``. + + It takes a class, such as ``Turn`` or ``list[str]``. Type checkers refuse a union, such + as ``str | None``, so use ``Any`` or a class that holds the value instead. + """ + return CreateStream(payload_type, or_get=self._or_get) + + async def __call__( + self, + id: str, + *, + backend: str | None = None, + **options: Unpack[StreamOptions], + ) -> WorkflowStream[T]: + """Create the Stream, and return it. + + Args: + id: The Stream id, which the Workflow picks, such as a session id. + backend: The name of the backend. Defaults to the registered backend with + ``default=True``. A replay must create the Stream on the same backend, so keep + the Worker's default, or gate a new one with ``workflow.patched``. 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.streams.StreamOptions`. + + Raises: + StreamAlreadyExistsError: ``create_stream`` of an id that exists, which no earlier + create of this run chain made. + 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.streams import _workflow + + return await _workflow.create_stream( + self._payload_type, + id, + backend=backend, + options=options, + or_get=self._or_get, + ) + + +class CreateStreamNowait(Generic[T]): + """The call of :py:data:`create_stream_nowait` and :py:data:`create_or_get_stream_nowait`. + + Subscript it with the payload type, as in ``workflow.create_stream_nowait[Turn]("turns-7")``. + + .. warning:: + This API is experimental. + """ + + def __init__(self, payload_type: type[T] | None, *, or_get: bool) -> None: + """Create the call. Use the module's ``create_stream_nowait`` instead.""" + self._payload_type = payload_type + self._or_get = or_get + + def __getitem__(self, payload_type: type[_U]) -> CreateStreamNowait[_U]: + """The call for Streams of ``payload_type``. It takes a class, as for ``create_stream``.""" + return CreateStreamNowait(payload_type, or_get=self._or_get) + + def __call__( + self, + id: str, + *, + backend: str | None = None, + **options: Unpack[StreamOptions], + ) -> WorkflowStream[T]: + """Start the create, and return the Stream at once. + + It takes the arguments of :py:meth:`CreateStream.__call__`. A backend error raises in + the run as an error that the run doesn't catch, so a ``StreamError`` fails the run. + """ + from temporalio.streams import _workflow + + return _workflow.create_stream_nowait( + self._payload_type, + id, + backend=backend, + options=options, + or_get=self._or_get, + ) + + +create_stream: CreateStream[Any] = CreateStream(None, or_get=False) +"""Create a Stream, which fails if a Stream of the id exists. + +Give the payload type as a type argument:: + + turns = await workflow.create_stream[Turn](f"{session_id}-turns") + +A Stream must exist before its first publish. A Workflow creates only standalone Streams, under an +id of its choice, and the Stream doesn't end with the run. Each run of a chain sends the same create +for an id, so a retry run or a cron run that creates the Stream of an earlier run gets that Stream, +and the options don't change. A run after continue-as-new uses the Stream through the ref that the +earlier run forwards, with :py:func:`to_stream`. The create and each publish are Local Activities, +and the create's error is raised here. 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. See +:py:meth:`CreateStream.__call__`. + +.. warning:: + This API is experimental. +""" + +create_or_get_stream: CreateStream[Any] = CreateStream(None, or_get=True) +"""Create a Stream, or get the one of the id that exists. + +It takes the arguments of :py:data:`create_stream`. On a Stream that exists, it ignores the +options. It still fails on a closed or expired Stream. + +.. warning:: + This API is experimental. +""" + +create_stream_nowait: CreateStreamNowait[Any] = CreateStreamNowait(None, or_get=False) +"""Create a Stream without an await, which fails the run where :py:data:`create_stream` raises. + +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:data:`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. +""" + +create_or_get_stream_nowait: CreateStreamNowait[Any] = CreateStreamNowait( + None, or_get=True +) +"""Create a Stream, or get the one that exists, without an await. + +It is to :py:data:`create_or_get_stream` what :py:data:`create_stream_nowait` is to +:py:data:`create_stream`. + +.. warning:: + This API is experimental. +""" + + +def to_stream(ref: StreamRef[T]) -> WorkflowStream[T]: + """Get the Stream of a ref 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.streams import _workflow + + return _workflow.to_stream(ref) diff --git a/tests/contrib/streams/_support.py b/tests/contrib/streams/_support.py index 9f1a05f106..476ba7e0cf 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 describes a Stream in tests.""" from __future__ import annotations @@ -9,6 +9,10 @@ import pytest +from temporalio.client import Client +from temporalio.streams._plugin import client_backend +from temporalio.streams.system import models + _BRIDGE = Path(__file__).parents[3] / "temporalio" / "bridge" @@ -39,3 +43,14 @@ 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")) + + +async def stream_status( + client: Client, stream_id: str, backend: str | None = None +) -> models.StreamStatus: + """The status that a describe gives, such as whether the Stream is closed.""" + _, calls = client_backend(client, backend) + response = await calls.service(client.namespace).send_describe_stream( + models.DescribeStreamRequest(ref=models.StreamRef(id=stream_id)) + ) + return response.status or models.StreamStatus.StreamStatusUnspecified diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py new file mode 100644 index 0000000000..f046e6e8f8 --- /dev/null +++ b/tests/contrib/streams/test_redis.py @@ -0,0 +1,69 @@ +import uuid +from datetime import timedelta +from typing import Any + +from temporalio import streams, workflow +from temporalio.client import Client, Stream +from temporalio.streams.system import models +from temporalio.worker import ( + UnsandboxedWorkflowRunner, +) +from tests.contrib.streams._support import stream_status +from tests.helpers import new_worker + + +async def read_all(stream: Stream[Any]) -> list[Any]: + return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] + + +def records(items: list[Any]) -> list[streams.Record[Any]]: + return [item for item in items if isinstance(item, streams.Record)] + + +attempts: dict[str, int] = {} + + +@workflow.defn +class PublishThenFailTask: + @workflow.run + async def run(self, key: str) -> None: + stream = await workflow.create_stream[str]( + f"{workflow.info().workflow_id}-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() + stream_id = f"{handle.id}-tokens" + items = await read_all(stream_client.get_stream[str](stream_id)) + assert ( + await stream_status(stream_client, stream_id) + == models.StreamStatus.StreamStatusClosed + ) + # 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, streams.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]] diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index dfda0625ec..26ce750cad 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 from temporalio.streams.system import models +from temporalio.worker import Replayer from tests.contrib.streams.fake_extension import fake_client from tests.helpers import new_worker @@ -425,3 +426,102 @@ async def test_stream_takes_the_named_backend_or_the_default(client: Client): await two.to_stream(streams.StreamRef(id="moving")).get_offset() with pytest.raises(ValueError): two.get_stream[str]("") + + +@workflow.defn +class WriteToStream: + @workflow.run + async def run(self, ref: streams.StreamRef[Turn]) -> None: + await workflow.to_stream(ref).publish(Turn("workflow")) + + +async def test_client_stream_takes_workflow_writes(stream_client: Client): + stream = await stream_client.create_stream[Turn](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, + stream.ref, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await stream.close() + items = [item async for item in stream.subscribe()] + data = [item.data for item in items if isinstance(item, streams.Record)] + assert data == [Turn("client"), Turn("a"), Turn("b"), Turn("workflow")] + + # Another Stream of the same id reads it from an offset. + again = stream_client.get_stream[Turn](stream.ref.id) + resumed = [ + item.data + async for item in again.subscribe(offset=first + 1) + if isinstance(item, streams.Record) + ] + assert resumed == [Turn("a"), Turn("b"), Turn("workflow")] + + +@workflow.defn +class Agent: + def __init__(self) -> None: + self._tokens: workflow.WorkflowStream[str] | None = None + + @workflow.run + async def run(self, prompt: str) -> str: + await workflow.sleep(timedelta(milliseconds=200)) + self._tokens = stream = await workflow.create_stream[str]( + f"{workflow.info().workflow_id}-tokens" + ) + await stream.publish(*prompt.split()) + await stream.close() + return prompt.upper() + + @workflow.query + def stream(self) -> streams.StreamRef[str] | None: + # A query has no ref before the run creates the Stream. + return self._tokens.ref if self._tokens else None + + +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 + ) + ref = await handle.query(Agent.stream) + while ref is None: + await asyncio.sleep(0.05) + ref = await handle.query(Agent.stream) + reader = stream_client.to_stream(ref) + items = [item async for item in reader.subscribe()] + await handle.result() + assert [item.data for item in items if isinstance(item, streams.Record)] == [ + "hi", + "there", + ] + # A Workflow creates only standalone Streams, so the ref has no owner. + assert reader.ref.owner is None and reader.ref.backend is not None + + +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[Turn]("") + assert not extension.namespaces + + +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[str]("d-1")).ref.backend == "b" + assert two.get_stream[str]("d-1").ref.backend == "b" + # A Worker serves only one backend, and so does a Replayer. + with pytest.raises(streams.StreamConfigError, match="one backend"): + new_worker(Client(**config), Agent) + with pytest.raises(streams.StreamConfigError, match="one backend"): + Replayer(workflows=[Agent], stream_backends=config["stream_backends"]) diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py new file mode 100644 index 0000000000..2cee910b08 --- /dev/null +++ b/tests/contrib/streams/test_workflow_module.py @@ -0,0 +1,1320 @@ +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, streams, 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.contrib.stream.testing import MemoryStreamBackend +from temporalio.converter import ( + DataConverter, + DefaultPayloadConverter, + PayloadCodec, + SerializationContext, + WithSerializationContext, + WorkflowSerializationContext, +) +from temporalio.exceptions import ApplicationError, CancelledError +from temporalio.streams import _workflow, _writes +from temporalio.streams._backend import _ExtensionCalls +from temporalio.streams.system import models +from temporalio.streams.system.workflow._resources.stream import ( + stream as generated_stream, +) +from temporalio.worker import Replayer, UnsandboxedWorkflowRunner +from tests.contrib.streams._support import stream_status +from tests.helpers import assert_eventually, assert_task_fail_eventually, new_worker + +OPEN = models.StreamStatus.StreamStatusOpen +CLOSED = models.StreamStatus.StreamStatusClosed + + +@dataclass +class Turn: + kind: str + text: str = "" + + +def own_id(name: str = "tokens") -> str: + # A Workflow picks the id of each Stream that it creates. These tests use its Workflow ID. + return f"{workflow.info().workflow_id}-{name}" + + +def stream_id(handle: WorkflowHandle[Any, Any], name: str = "tokens") -> str: + return f"{handle.id}-{name}" + + +async def read( + handle: WorkflowHandle[Any, Any], + name: str = "tokens", + payload_type: type[Any] = str, + backend: str | None = None, +) -> list[object]: + stream = handle._client.get_stream[payload_type]( + stream_id(handle, name), backend=backend + ) + return [item async for item in stream.subscribe(wait_new_records=False)] + + +async def status(handle: WorkflowHandle[Any, Any], name: str = "tokens") -> object: + return await stream_status(handle._client, stream_id(handle, name)) + + +def data(items: Sequence[object]) -> list[object]: + return [item.data for item in items if isinstance(item, streams.Record)] + + +def reason(err: streams.StreamError) -> str: + return StreamErrorReason.Name(StreamErrorReason.ValueType(err.reason)) + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + # Two channels are two Streams. + turns = await workflow.create_stream[Turn](own_id("turns")) + tokens = await workflow.create_stream[str](own_id("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() + 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", Turn) + tokens = await read(handle) + assert await status(handle, "turns") == await status(handle) == CLOSED + assert data(turns) == [Turn("user", "hi there"), Turn("model")] + assert data(tokens) == ["hi", "there"] + + +@workflow.defn +class CreatesTwice: + @workflow.run + async def run(self, client_id: str) -> list[str]: + first = await workflow.create_stream[str](own_id()) + # The second create is the same create as the first, so it gets the same Stream. + again = await workflow.create_stream[str](own_id()) + results = [str(again.ref == first.ref)] + await first.close() + try: + await workflow.create_stream[str](client_id) + except streams.StreamAlreadyExistsError as err: + results.append(reason(err)) + # The create-or-get gets the Stream that exists, and ignores its options. + theirs = await workflow.create_or_get_stream[str]( + client_id, retention=timedelta(hours=1) + ) + await theirs.publish("one") + return results + + +async def test_create_fails_on_an_id_that_another_caller_made(stream_client: Client): + theirs = await stream_client.create_stream[str](f"client-{uuid.uuid4()}") + async with new_worker(stream_client, CreatesTwice) as worker: + result = await stream_client.execute_workflow( + CreatesTwice.run, + theirs.ref.id, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == ["True", "STREAM_ERROR_REASON_ALREADY_EXISTS"] + items = [item async for item in theirs.subscribe(wait_new_records=False)] + assert data(items) == ["one"] + + +@workflow.defn +class CreatesBetweenDraws: + @workflow.run + async def run(self) -> bool: + before = workflow.random().getstate() + stream = await workflow.create_stream[str](own_id()) + workflow.create_or_get_stream_nowait[str](own_id()) + drawn = workflow.random().getstate() != before + await stream.close() + return drawn + + +async def test_create_leaves_the_workflow_random_values_alone(stream_client: Client): + async with new_worker(stream_client, CreatesBetweenDraws) as worker: + drawn = await stream_client.execute_workflow( + CreatesBetweenDraws.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert not drawn + + +@workflow.defn +class GeneratedCalls: + @workflow.run + async def run(self) -> list[str]: + tokens = await workflow.create_stream[str](own_id()) + generated = generated_stream( + ref=_workflow._generated_ref(tokens.ref), body_type=str + ) + offset = await generated.publish("a", "b") + refused = [] + try: + await generated.read() + except streams.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"] + # Each write is a marker, as with the Streams, 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 + + +@workflow.defn +class RefusedCalls: + @workflow.run + async def run(self) -> list[str]: + tokens = await workflow.create_stream[str](own_id()) + refused = [] + # A Stream bound to another run chain fails before the block runs, and before any call. + owner = streams.Execution( + type=streams.ExecutionType.ExecutionTypeWorkflow, + business_id="other", + run_id="run-1", + ) + other = workflow.to_stream(streams.StreamRef[str](id="other", owner=owner)) + try: + async with other: + refused.append("entered") + except streams.StreamNotOwnerError as err: + refused.append(type(err).__name__) + try: + await other.close() + except streams.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 == ["StreamNotOwnerError", "StreamNotOwnerError"] + + +@workflow.defn +class Blocks: + @workflow.run + async def run(self) -> None: + async with await workflow.create_stream[str](own_id("first")) as first: + await first.publish("token-0") + try: + async with await workflow.create_stream[str](own_id("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() + assert data(await read(handle, "first")) == ["token-0"] + assert data(await read(handle, "second")) == ["token-1"] + assert await status(handle, "first") == await status(handle, "second") == CLOSED + + +@workflow.defn +class Continuing: + @workflow.run + async def run(self, ref: streams.StreamRef[str] | None) -> str: + if ref is None: + created = await workflow.create_stream[str](own_id()) + await created.publish("run 1") + # The next run gets the Stream only through the ref that this run forwards. + workflow.continue_as_new(created.ref) + stream = workflow.to_stream(ref) + await stream.publish("run 2") + await stream.close() + return "done" + + +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" + assert data(await read(handle)) == ["run 1", "run 2"] + assert await status(handle) == CLOSED + + +@workflow.defn +class CreatesInEachRun: + @workflow.run + async def run(self, first: bool) -> str: + # The next run sends the create of the first run, so it gets the same Stream. + stream = await workflow.create_stream[str](own_id()) + if first: + await stream.publish("run 1") + workflow.continue_as_new(False) + await stream.publish("run 2") + await stream.close() + return "done" + + +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, + ) + assert await handle.result() == "done" + assert data(await read(handle)) == ["run 1", "run 2"] + + +@workflow.defn +class Mistakes: + def __init__(self) -> None: + self._stream: workflow.WorkflowStream[str] | None = None + + @workflow.run + async def run(self) -> list[str]: + errors = [] + stream = await workflow.create_stream[str](own_id()) + self._stream = stream + await stream.publish("last") + await stream.close() + try: + await stream.publish("late") + except streams.StreamClosedError as err: + errors.append(reason(err)) + # A closed Stream fails a create with CLOSED, which comes before ALREADY_EXISTS. + for create in (workflow.create_stream, workflow.create_or_get_stream): + try: + await create[str](own_id()) + except streams.StreamClosedError as err: + errors.append(reason(err)) + return errors + + @workflow.query + def publish_in_query(self) -> str: + assert self._stream is not None + try: + self._stream.publish("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"] * 3 + assert await handle.query(Mistakes.publish_in_query) == "ReadOnlyContextError" + assert data(await read(handle)) == ["last"] + + +@workflow.defn +class ClosesAReceivedStream: + @workflow.run + async def run(self, ref: streams.StreamRef[str]) -> str: + stream = workflow.to_stream(ref) + await stream.publish("from workflow") + await stream.close() + return "closed" + + +async def test_workflow_closes_a_standalone_stream_that_it_received( + stream_client: Client, +): + # A standalone Stream has no owner, so a Workflow closes it as its run's Producer. + stream = await stream_client.create_stream[str](f"received-{uuid.uuid4()}") + async with new_worker(stream_client, ClosesAReceivedStream) as worker: + assert ( + await stream_client.execute_workflow( + ClosesAReceivedStream.run, + stream.ref, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + == "closed" + ) + items = [item async for item in stream.subscribe()] + assert data(items) == ["from workflow"] + assert await stream_status(stream_client, stream.ref.id) == CLOSED + + +@workflow.defn +class Failing: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream[str](own_id()) + 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() + assert data(await read(handle)) == ["before"] + # The Stream stays open, so a later reader still waits for a close. + assert await status(handle) == OPEN + + +@workflow.defn +class RetriesItsCreate: + @workflow.run + async def run(self) -> str: + attempt = workflow.info().attempt + stream = await workflow.create_stream[str](own_id()) + await stream.publish(f"run-{attempt}") + if attempt == 1: + raise ApplicationError("the first run fails") + await stream.close() + return f"created on attempt {attempt}" + + +async def test_retry_run_creates_the_stream_of_its_first_run(stream_client: Client): + # The retry run sends the create of the first run, so the create gets that Stream. + async with new_worker(stream_client, RetriesItsCreate) as worker: + handle = await stream_client.start_workflow( + RetriesItsCreate.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + retry_policy=RetryPolicy( + initial_interval=timedelta(milliseconds=100), maximum_attempts=2 + ), + ) + assert await handle.result() == "created on attempt 2" + assert data(await read(handle)) == ["run-1", "run-2"] + assert await status(handle) == CLOSED + + +@workflow.defn +class FailsInItsBlockOnce: + @workflow.run + async def run(self) -> str: + attempt = workflow.info().attempt + async with await workflow.create_stream[str](own_id()) 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) + # Only the block end of the run that succeeds closes the Stream. + assert await status(handle) == CLOSED + written = [item for item in items if isinstance(item, streams.Record)] + superseded = [item for item in items if isinstance(item, streams.Superseded)] + assert [r.data for r in written] == ["run-1", "run-2"] + # The retry run's first write withdraws the Records of the run that it retries. + assert [r.data for r in written if superseded[0].withdraws(r)] == ["run-1"] + + +@workflow.defn +class EndsItsBlock: + @workflow.run + async def run(self, mode: str) -> str: + try: + async with await workflow.create_stream[str](own_id()) 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. + 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() + assert data(await read(handle)) == [mode] + assert await status(handle) == CLOSED, mode + # 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 await workflow.create_stream[str](own_id()) 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_the_stream(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() + assert data(await read(handle)) == ["started"] + assert await status(handle) == CLOSED, f"local={local}" + + +@workflow.defn +class CronBlock: + @workflow.run + async def run(self) -> str: + # Each cron run sends the create of the first run, so it gets the same Stream. + async with await workflow.create_stream[str](own_id()) 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 = stream_client.get_stream[str](stream_id(handle)).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 await status(handle) == OPEN + assert isinstance(first, streams.Record) and isinstance(second, streams.Record) + + +failed_close_tasks: set[str] = set() + + +@workflow.defn +class SchedulesACloseThenFailsItsTask: + @workflow.run + async def run(self) -> None: + stream = await workflow.create_stream[str](own_id()) + await stream.publish("token") + stream.close_nowait() + 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. The retried task creates again with the same request id, so the create + # succeeds. + 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) + assert await status(handle) == CLOSED + written = [item for item in items if isinstance(item, streams.Record)] + superseded = [item for item in items if isinstance(item, streams.Superseded)] + # The retried task writes the token again, and withdraws the failed task's copy. + assert [r.data for r in written] == ["token", "token"] + assert [r for r in written if any(s.withdraws(r) for s in superseded)] == [ + written[0] + ] + 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[str](own_id(), 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[str](own_id()) + 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() + try: + self._stream.publish_nowait("late") + except streams.StreamClosedError as err: + return reason(err) + 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() + assert data(await read(handle)) == ["a", "p"] + assert await status(handle) == CLOSED + + +@workflow.defn +class CreatesInInit: + @workflow.init + def __init__(self, prompt: str) -> None: + self._stream = workflow.create_or_get_stream_nowait[str](own_id()) + self._stream.publish_nowait(prompt) + + @workflow.run + async def run(self, prompt: str) -> str: + await self._stream.publish(prompt.upper()) + await self._stream.close() + 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" + assert data(await read(handle)) == ["hi", "HI"] + assert await status(handle) == CLOSED + + +@workflow.defn +class PublishesWithNoAwait: + @workflow.run + async def run(self) -> list[int]: + tokens = await workflow.create_stream[str](own_id()) + 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, streams.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 = workflow.to_stream(streams.StreamRef[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[str](own_id()) + 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, streams.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.to_stream(streams.StreamRef[str](id=own_id())) + 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() + missing = stream_id(handle) + assert not [key for key in _writes._PENDING._entries if key[1] == missing] + + +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[str](own_id()) + 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, streams.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[str](own_id()) + try: + await tokens.publish(*["x" * 1_400_000] * 4) + except streams.StreamTooLargeError as err: + return reason(err) + 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[str](own_id()) + + +@workflow.defn +class CreatesOnANamedBackend: + @workflow.run + async def run(self, backend: str) -> str: + stream = await workflow.create_stream[str](own_id(), backend=backend) + await stream.publish(f"on {backend}") + await workflow.to_stream(stream.ref).close() + return stream.ref.backend or "" + + +@workflow.defn +class CreatesAfterAnIdleTime: + @workflow.run + async def run(self) -> str: + await workflow.create_stream[str]( + own_id(), retention=timedelta(milliseconds=500) + ) + await workflow.sleep(timedelta(seconds=1.5)) + try: + await workflow.create_or_get_stream[str](own_id()) + except streams.StreamExpiredError as err: + return reason(err) + return "created" + + +async def test_create_of_an_expired_workflow_stream_raises(stream_client: Client): + # A Workflow's standalone Stream expires one Retention after its last write. + async with new_worker(stream_client, CreatesAfterAnIdleTime) as worker: + assert ( + await stream_client.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 gives 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" + assert data(await read(named, backend="memory")) == ["on memory"] + assert await stream_status(memory_client, stream_id(named), "memory") == 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[str](own_id()) + await stream.publish("a") + stream.close_nowait() + 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 await status(handle) == CLOSED + assert isinstance(failed.value.cause, CancelledError) + types = await _event_types(handle) + assert EventType.EVENT_TYPE_ACTIVITY_TASK_CANCEL_REQUESTED not in types + assert data(await read(handle)) == ["a"] + + +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: workflow.WorkflowStream[str] | None = None + + @workflow.run + async def run(self) -> str: + self.late = await workflow.create_stream[str](own_id("late")) + tokens = await workflow.create_stream[str](own_id()) + 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[str](own_id()) + 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"] + + +@workflow.defn +class ContinuesInItsBlock: + @workflow.run + async def run(self, ref: streams.StreamRef[str] | None) -> str: + if ref is None: + async with await workflow.create_stream[str](own_id()) as stream: + await stream.publish("run 1") + workflow.continue_as_new(stream.ref) + try: + await workflow.to_stream(ref).publish("run 2") + except streams.StreamClosedError as err: + return reason(err) + return "wrote" + + +async def test_continue_as_new_in_a_block_closes_the_stream(stream_client: Client): + # The run that continues as new has no retry run, so its block end closes the Stream. + async with new_worker(stream_client, ContinuesInItsBlock) as worker: + handle = await stream_client.start_workflow( + ContinuesInItsBlock.run, + None, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == "STREAM_ERROR_REASON_CLOSED" + assert data(await read(handle)) == ["run 1"] + assert await status(handle) == CLOSED