diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 9801b8d20..3960442e8 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -9,6 +9,7 @@ This package is experimental. """ +from temporalio.contrib.streams import workflow from temporalio.contrib.streams._activity import activity_stream from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams._client import ( @@ -77,5 +78,6 @@ "for_handle", "for_ref", "topic", + "workflow", "workflow_stream", ] diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 1782a19de..28ee6bd01 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -9,10 +9,10 @@ import asyncio import dataclasses import logging -from collections.abc import Awaitable, Callable, Generator, Sequence +from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta -from typing import Any +from typing import Any, NoReturn import temporalio.activity import temporalio.api.common.v1 @@ -25,6 +25,7 @@ import temporalio.worker import temporalio.workflow from temporalio.contrib.streams import _activity, _errors, _ids, _workflow +from temporalio.contrib.streams import workflow as streams_workflow from temporalio.contrib.streams._backend import StreamBackend from temporalio.contrib.streams._stream import StreamRef from temporalio.contrib.streams.proto import v1 as ext_pb @@ -43,6 +44,9 @@ # schedules. The key keeps the Producer id the same across retries. _REF_HEADER = "__temporal_stream_ref" _KEY_HEADER = "__temporal_stream_key" +# Continue-as-new passes the ids of the open Streams that close with the run chain to the next +# run, in this header. +_CLOSE_WITH_RUN_HEADER = "__temporal_stream_close_with_run" class StreamPlugin(SimplePlugin): @@ -58,8 +62,17 @@ def __init__( extension_name: str, config: ext_pb.RedisStreamConfig, runtime: temporalio.runtime.Runtime | None, + backend_name: str, + default: bool, ) -> None: - """Create the plugin over the bridge's extension with ``extension_name``.""" + """Create the plugin over the bridge's extension with ``extension_name``. + + ``backend_name`` is the name that Stream handles carry for this backend. ``default`` + makes it the client's default backend when the client has more than one. + """ + self._backend_name = backend_name + self._default = default + config.backend = backend_name self._extension = temporalio.worker.WorkerExtension( extension_name, config.SerializeToString(), runtime=runtime ) @@ -67,6 +80,16 @@ def __init__( self._converter = temporalio.converter.DataConverter.default super().__init__(name, interceptors=[_StreamInterceptor(self)]) + @property + def backend_name(self) -> str: + """The name that Stream handles carry for this backend.""" + return self._backend_name + + @property + def default(self) -> bool: + """Whether a client with more than one Stream plugin uses this one by default.""" + return self._default + def configure_worker( self, config: temporalio.worker.WorkerConfig ) -> temporalio.worker.WorkerConfig: @@ -173,6 +196,8 @@ def __init__( durable_timeout: timedelta | None = None, response_timeout: timedelta | None = None, runtime: temporalio.runtime.Runtime | None = None, + name: str = "redis", + default: bool = False, ) -> None: """Create the plugin. It connects to Redis on the first Stream call. @@ -188,6 +213,10 @@ def __init__( durable_timeout: How long that wait takes at most. Defaults to 1 s. response_timeout: How long one Redis call takes at most. Defaults to 5 s. runtime: The runtime of the client. Defaults to the default runtime. + name: The backend name that Stream handles carry, so every process routes a + handle it receives. + default: Make this the client's default backend when the client has more than + one Stream plugin. With none marked, the first one is the default. """ config = ext_pb.RedisStreamConfig( urls=[url] if isinstance(url, str) else list(url), @@ -195,16 +224,21 @@ def __init__( key_prefix=key_prefix or "", durable_writes=durable_writes, ) - for name, value in ( + for field, value in ( ("retention", retention), ("safety_ttl", safety_ttl), ("durable_timeout", durable_timeout), ("response_timeout", response_timeout), ): if value is not None: - getattr(config, name).FromTimedelta(value) + getattr(config, field).FromTimedelta(value) super().__init__( - "temporalio.RedisStreamPlugin", EXTENSION_NAME, config, runtime + "temporalio.RedisStreamPlugin", + EXTENSION_NAME, + config, + runtime, + name, + default, ) @@ -259,6 +293,8 @@ async def _encode_call( ) if request.retention is not None: create.retention.FromTimedelta(request.retention) + if request.auto_close is not None: + create.auto_close.FromTimedelta(request.auto_close) return ext_pb.WorkflowCall(create=create) ref = request.stream or models.StreamRef(id="") stream = stream_pb.StreamRef( @@ -311,11 +347,17 @@ class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interc def __init__(self, plugin: StreamPlugin) -> None: super().__init__() self._plugin = plugin + # Workflow code learns the Worker's backend name from the interceptor class. + self._inbound_class = type( + "_WorkflowInbound", + (_WorkflowInbound,), + {"_backend_name": plugin.backend_name}, + ) def workflow_interceptor_class( self, input: temporalio.worker.WorkflowInterceptorClassInput ) -> type[temporalio.worker.WorkflowInboundInterceptor] | None: - return _WorkflowInbound + return self._inbound_class def intercept_activity( self, next: temporalio.worker.ActivityInboundInterceptor @@ -324,12 +366,15 @@ def intercept_activity( class _WorkflowInbound(temporalio.worker.WorkflowInboundInterceptor): + _backend_name = "" + def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: - super().init(_WorkflowOutbound(outbound)) + super().init(_WorkflowOutbound(outbound, self._backend_name)) async def execute_workflow( self, input: temporalio.worker.ExecuteWorkflowInput ) -> Any: + _take_close_with_run(input.headers) try: result = await super().execute_workflow(input) except temporalio.workflow.ContinueAsNewError: @@ -395,7 +440,10 @@ def chain_continues( ) if policy.maximum_interval is not None: backoff = min(backoff, policy.maximum_interval) - expiration = info.workflow_start_time + info.execution_timeout + # The execution timeout covers the whole chain, so this run's start is not its start. + expiration = info.execution_expiration_time or ( + info.workflow_start_time + info.execution_timeout + ) if now + backoff >= expiration: return False return True @@ -432,16 +480,24 @@ def _close_with_run( for stream in state.streams.values(): if stream.ref.id in ids: stream._closed = True - for stream_id in ids: + # A received ref keeps its own owner. + refs = { + stream_id: state.streams[stream_id]._model_ref() + if stream_id in state.streams + else _owned_ref(stream_id) + for stream_id in ids + } + for run_stream in streams_workflow._state().streams.values(): + if run_stream.close_with_run and not run_stream.closed: + run_stream.closed = True + refs[run_stream.ref.id] = run_stream.ref + for ref in refs.values(): temporalio.workflow.send_to_plugins( _WorkflowCall( "CloseStream", _convert_in_sandbox( models.CloseStreamRequest( - # A received ref keeps its own owner. - stream=state.streams[stream_id]._model_ref() - if stream_id in state.streams - else _owned_ref(stream_id), + stream=ref, result=result, failure=failure, ) @@ -488,11 +544,26 @@ def _owned_ref(stream_id: str) -> models.StreamRef: ) +def _take_close_with_run( + headers: Mapping[str, temporalio.api.common.v1.Payload], +) -> None: + payload = headers.get(_CLOSE_WITH_RUN_HEADER) + if payload is None: + return + state = streams_workflow._state() + converter = temporalio.converter.PayloadConverter.default + for stream_id in converter.from_payload(payload, list[str]): + state.stream(_owned_ref(stream_id)).close_with_run = True + + class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): _temporal_stream_transport = True - def __init__(self, next: temporalio.worker.WorkflowOutboundInterceptor) -> None: + def __init__( + self, next: temporalio.worker.WorkflowOutboundInterceptor, backend_name: str + ) -> None: super().__init__(next) + self._temporal_stream_backend_name = backend_name # Numbers each Activity call of the run in order, so the key is the same on replay and # a retry keeps it. A counter, not workflow.uuid4(), leaves the user's random values # alone. @@ -534,6 +605,20 @@ async def start_system_nexus_operation( ) return _DoneHandle(result) + def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: + open_ids = [ + stream_id + for stream_id, run_stream in streams_workflow._state().streams.items() + if run_stream.close_with_run and not run_stream.closed + ] + if open_ids: + converter = temporalio.converter.PayloadConverter.default + input.headers = { + **input.headers, + _CLOSE_WITH_RUN_HEADER: converter.to_payload(open_ids), + } + super().continue_as_new(input) + def start_activity( self, input: temporalio.worker.StartActivityInput ) -> temporalio.workflow.ActivityHandle[Any]: diff --git a/temporalio/contrib/streams/testing.py b/temporalio/contrib/streams/testing.py index 4929dcb42..8b98eee4c 100644 --- a/temporalio/contrib/streams/testing.py +++ b/temporalio/contrib/streams/testing.py @@ -32,6 +32,8 @@ def __init__( safety_ttl: timedelta | None = None, durable_writes: int = 0, runtime: temporalio.runtime.Runtime | None = None, + name: str = "memory", + default: bool = False, ) -> None: """Create the plugin. @@ -43,6 +45,9 @@ def __init__( replicas. The memory Store has none, so each Workflow flush fails with the missing ``replica_wait`` capability. It is for tests of that failure. runtime: The runtime of the client. Defaults to the default runtime. + name: The backend name that Stream handles carry. + default: Make this the client's default backend when the client has more than + one Stream plugin. """ config = ext_pb.RedisStreamConfig(durable_writes=durable_writes) if retention is not None: @@ -50,5 +55,10 @@ def __init__( if safety_ttl is not None: config.safety_ttl.FromTimedelta(safety_ttl) super().__init__( - "temporalio.MemoryStreamPlugin", MEMORY_EXTENSION_NAME, config, runtime + "temporalio.MemoryStreamPlugin", + MEMORY_EXTENSION_NAME, + config, + runtime, + name, + default, ) diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/contrib/streams/workflow.py new file mode 100644 index 000000000..9526c334a --- /dev/null +++ b/temporalio/contrib/streams/workflow.py @@ -0,0 +1,393 @@ +"""Streams in Workflow code. + +Workflow code creates a Stream with :py:func:`create_stream` and gets one again with +:py:func:`stream_handle`. A Stream that the Workflow receives as an argument, such as a +:py:class:`WorkflowStream` that a run passes on through continue-as-new, works the same. + +.. warning:: + This module is experimental. +""" + +from __future__ import annotations + +import asyncio +import weakref +from collections.abc import Coroutine +from dataclasses import dataclass, field +from datetime import timedelta +from types import TracebackType +from typing import Any, Generic, TypeVar, overload + +import temporalio.exceptions +import temporalio.workflow +from temporalio.contrib.streams import _errors, _ids +from temporalio.contrib.streams._defn import TopicDefinition, definition_of +from temporalio.contrib.streams._handle import StreamHandle +from temporalio.contrib.streams.system import models +from temporalio.contrib.streams.system.operations import ( + append_records as _append_records, +) +from temporalio.contrib.streams.system.operations import close_stream as _close_stream +from temporalio.contrib.streams.system.operations import ( + create_stream as _create_stream, +) + +E = TypeVar("E") +T = TypeVar("T") +_R = TypeVar("_R") + +__all__ = [ + "WorkflowStream", + "WorkflowTopic", + "create_stream", + "create_stream_nowait", + "stream_handle", +] + + +@dataclass +class _RunStream: + # What this run knows about one Stream. Many handles may name it, so it is kept by id. + ref: models.StreamRef + close_with_run: bool = False + closed: bool = False + + +@dataclass +class _RunState: + streams: dict[str, _RunStream] = field(default_factory=dict) + + 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 + + +# One state per Workflow instance. The plugin passes temporalio.contrib.streams through the +# sandbox, so the plugin's interceptor sees the same state, and it goes when the run leaves +# the cache. +_states: weakref.WeakKeyDictionary[Any, _RunState] = weakref.WeakKeyDictionary() + + +def _state() -> _RunState: + runtime = temporalio.workflow._Runtime.current() + state = _states.get(runtime) + if state is None: + state = _RunState() + _states[runtime] = state + return state + + +def _backend_name() -> str | None: + # A Stream plugin marks its outbound interceptor with its backend name. With none, the + # call would go to the server as a System Nexus command, which only native Streams serve. + outbound = getattr(temporalio.workflow._Runtime.current(), "_outbound", None) + while outbound is not None: + if getattr(outbound, "_temporal_stream_transport", False): + return getattr(outbound, "_temporal_stream_backend_name", "") + outbound = getattr(outbound, "next", None) + return None + + +def _check_backend(backend: str | None, action: str) -> str: + # The Worker's config gives these errors, not the history, so they fail the Workflow Task + # and a Worker with the right config retries it. An ApplicationError would fail the run. + worker_backend = _backend_name() + if worker_backend is None: + raise RuntimeError( + "No Stream plugin is on the Worker, such as RedisStreamPlugin" + ) + if backend and backend != worker_backend: + raise RuntimeError( + f"Can't {action} a Stream of the backend {backend}. The Worker has only " + f"{worker_backend}." + ) + return worker_backend + + +def _now(call: Coroutine[Any, Any, _R]) -> _R: + # The plugin's interceptor answers at once, so the call ends in its first step. + try: + call.send(None) + except StopIteration as done: + return done.value + call.close() + raise RuntimeError("No Stream plugin took the call, such as RedisStreamPlugin") + + +class WorkflowTopic(Generic[T]): + """One Topic of a :py:class:`WorkflowStream`. All Topics of a Stream share one order. + + .. warning:: + This API is experimental. + """ + + def __init__(self, stream: WorkflowStream[Any], name: str, type: Any) -> None: + """Create a Topic. Use :py:meth:`WorkflowStream.topic` instead.""" + self._stream = stream + self._name = name + self._type = type + + @property + def name(self) -> str: + """Name of the Topic.""" + return self._name + + @property + def type(self) -> Any: + """Type of the Topic's values.""" + return self._type + + async def publish(self, *values: T) -> None: + """Add one Record for each value, in order. + + It returns when the Records are in the run's buffer. The backend writes them before + the Workflow Task completes, and before a Local Activity that the run starts later. A + list is one value. + + Raises: + StreamClosedError: This run closed the Stream. + ReadOnlyContextError: The code must not change state, such as a query handler. + RuntimeError: No Stream plugin is on the Worker, or the Stream is of a backend + that the Worker does not have. It fails the Workflow Task, not the run. + """ + self.publish_nowait(*values) + + def publish_nowait(self, *values: T) -> None: + """Add one Record for each value, in order, with no await. + + It is for sync Workflow code, such as ``@workflow.init`` and sync signal and update + handlers, as ``asyncio.Queue.put_nowait`` is for code that can't await ``put``. It has + the buffer and the errors of :py:meth:`publish`. + """ + self._stream._append(self._name, values) + + +class WorkflowStream(StreamHandle[E]): + """A Stream in Workflow code. Workflow code writes and closes it, and can't read it. + + Pass it to an Activity, a child Workflow or the next run of continue-as-new as an + argument. Use it in ``async with`` to close it at the end of the block. + + .. warning:: + This API is experimental. + """ + + def _model_ref(self) -> models.StreamRef: + return models.StreamRef( + id=self.id, + workflow_id=self.workflow_id, + first_execution_run_id=self.first_execution_run_id, + ) + + @overload + def topic(self, topic: TopicDefinition[T]) -> WorkflowTopic[T]: ... + + @overload + def topic(self, topic: str) -> WorkflowTopic[Any]: ... + + @overload + def topic(self, topic: str, *, type: type[T]) -> WorkflowTopic[T]: ... + + def topic( + self, topic: str | TopicDefinition[Any], *, type: Any = None + ) -> WorkflowTopic[Any]: + """A Topic of this Stream. + + Args: + topic: A Topic of the Stream's definition, such as ``Events.turns``, or a name. + A name of the definition takes its type. + type: The type of a Topic outside the definition, or of a Stream with no + definition. + + Raises: + StreamError: The Topic is not in the Stream's definition, with the reason + ``unknown_topic``. + """ + name, topic_type = self._topic_spec(topic, type) + return WorkflowTopic(self, name, topic_type) + + def _check(self, action: str, *, backend_checked: bool = False) -> _RunStream: + if temporalio.workflow.unsafe.is_read_only(): + raise _errors.ReadOnlyContextError( + f"Can't {action} a Stream in read-only code" + ) + run_stream = _state().stream(self._model_ref()) + if run_stream.closed: + raise _errors.StreamClosedError( + f"Can't {action} the Stream {self.id}, since this run closed it", + reason="stream_closed", + ) + if not backend_checked: + _check_backend(self.backend, action) + return run_stream + + def _append(self, topic: str, values: tuple[Any, ...]) -> None: + if not values: + raise TypeError("publish takes one or more values") + self._check("write to") + # The backend's answer has no offsets that replay could give again, so it is dropped. + _now( + _append_records.append_records( + stream=self._model_ref(), + records=[ + models.AppendRecord(topic=topic, body=value) for value in values + ], + ) + ) + + async def close( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Close the Stream. It goes in the run's buffer, as a publish does. + + Readers get the Records before it, then a ``Closed`` item with ``result`` or + ``failure``. Later writes in this run raise :py:class:`StreamClosedError`. + """ + self.close_nowait(result=result, failure=failure) + + def close_nowait( + self, *, result: Any | None = None, failure: BaseException | None = None + ) -> None: + """Close the Stream with no await, for sync Workflow code. + + It is to :py:meth:`close` what ``asyncio.Queue.put_nowait`` is to ``put``, with the same + buffer and errors. + """ + run_stream = self._check("close") + run_stream.closed = True + _now( + _close_stream.close_stream( + stream=self._model_ref(), result=result, failure=failure + ) + ) + + async def __aenter__(self) -> WorkflowStream[E]: + """Use the Stream in ``async with``, which closes it at the end.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the Stream, with the exception as its failure if there is one. + + A cancel closes it with a ``CancelledError`` failure. Another ``BaseException``, such as + an eviction, leaves it open, since the block did not end in Workflow code. + """ + if _state().stream(self._model_ref()).closed: + return + if isinstance(exc, asyncio.CancelledError): + await self.close( + failure=temporalio.exceptions.CancelledError(str(exc) or "Cancelled") + ) + elif exc is None or isinstance(exc, Exception): + await self.close(failure=exc) + + +def _own_handle( + definition: type[E], name: str | None, backend: str | None +) -> WorkflowStream[E]: + info = temporalio.workflow.info() + stream_name = name or definition_of(definition).name + return WorkflowStream( + id=_ids.scoped_stream_id( + info.workflow_id, info.first_execution_run_id, stream_name + ), + workflow_id=info.workflow_id, + first_execution_run_id=info.first_execution_run_id, + backend=backend, + definition=definition, + ) + + +async def create_stream( + definition: type[E], + *, + name: str | None = None, + retention: timedelta | None = None, + auto_close: timedelta | None = None, + close_with_run: bool = True, + backend: str | None = None, +) -> WorkflowStream[E]: + """Create a Stream that this run chain owns. + + Its id is ``::``, so a client finds it from the Workflow + handle with no call. The create goes in the run's buffer, as a publish does. A bad setting + fails the Workflow Task when the buffer goes out. + + Args: + definition: The Stream definition class. + name: The Stream name. Defaults to the definition's name. Give one Stream per turn, + for example, its own name. + retention: How long the Records stay after the close. Defaults to the backend's. + auto_close: Close the Stream with no result when no Record was appended for this + long, so a Stream that nobody closes does not stay open. Defaults to never. + close_with_run: Close the Stream when the run chain closes, with the result or the + failure of its last run. A run that continues as new passes it on. + backend: The name of the backend. Defaults to the Worker's. + + Raises: + StreamClosedError: This run closed the Stream. + ReadOnlyContextError: The code must not change state, such as a query handler. + RuntimeError: No Stream plugin is on the Worker, or it has no backend of that name. + It fails the Workflow Task, not the run. + ValueError: ``close_with_run=False`` for a Stream that already closes with the run + chain. + """ + return create_stream_nowait( + definition, + name=name, + retention=retention, + auto_close=auto_close, + close_with_run=close_with_run, + backend=backend, + ) + + +def create_stream_nowait( + definition: type[E], + *, + name: str | None = None, + retention: timedelta | None = None, + auto_close: timedelta | None = None, + close_with_run: bool = True, + backend: str | None = None, +) -> WorkflowStream[E]: + """Create a Stream that this run chain owns, with no await. + + It is for sync Workflow code, such as ``@workflow.init``, as ``asyncio.Queue.put_nowait`` + is for code that can't await ``put``. It takes the options of :py:func:`create_stream` + and has the same buffer and errors. + """ + stream = _own_handle(definition, name, _check_backend(backend, "create")) + run_stream = stream._check("create", backend_checked=True) + if run_stream.close_with_run and not close_with_run: + raise ValueError( + f"The Stream {stream.id} already closes with the run chain, and " + "close_with_run=False can't turn that off" + ) + _now( + _create_stream.create_stream( + stream_id=stream.id, retention=retention, auto_close=auto_close + ) + ) + run_stream.close_with_run = close_with_run + return stream + + +def stream_handle(definition: type[E], *, name: str | None = None) -> WorkflowStream[E]: + """A Stream that this run chain created, by its definition and name. It makes no call. + + A later run of the chain, such as one after continue-as-new, uses it to write the Stream + again. + + Args: + definition: The Stream definition class. + name: The Stream name. Defaults to the definition's name. + """ + return _own_handle(definition, name, _backend_name() or None) diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index 8f6e1d066..6ef1dae87 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -757,6 +757,11 @@ def _create_workflow_instance( workflow_id=init.workflow_id, workflow_type=init.workflow_type, priority=temporalio.common.Priority._from_proto(init.priority), + execution_expiration_time=init.workflow_execution_expiration_time.ToDatetime().replace( + tzinfo=timezone.utc + ) + if init.HasField("workflow_execution_expiration_time") + else None, ) last_failure = ( diff --git a/temporalio/workflow/_context.py b/temporalio/workflow/_context.py index 481942b0b..dfa9ddb18 100644 --- a/temporalio/workflow/_context.py +++ b/temporalio/workflow/_context.py @@ -136,6 +136,10 @@ class Info: workflow_type: str + execution_expiration_time: datetime | None = None + """When the execution timeout of the run chain passes, across retries and continue-as-new. + None when the chain has no execution timeout.""" + def _logger_details(self) -> Mapping[str, Any]: return { # TODO(cretz): worker ID? diff --git a/tests/contrib/streams/conftest.py b/tests/contrib/streams/conftest.py new file mode 100644 index 000000000..7a29e0b64 --- /dev/null +++ b/tests/contrib/streams/conftest.py @@ -0,0 +1,65 @@ +import shutil +import socket +import subprocess +import time +import uuid +from collections.abc import AsyncIterator, Iterator + +import pytest + +from temporalio.client import Client +from temporalio.contrib.streams.redis import RedisStreamPlugin, StreamPlugin +from temporalio.contrib.streams.testing import MemoryStreamPlugin + +# Ports that other local services own. +_RESERVED = {7333, 7861, 7931, 6431, 7101, 7102, 7103} + + +def _free_port() -> int: + while True: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + if port not in _RESERVED: + return port + + +# One scratch server serves every test module, and each test takes its own key prefix. +@pytest.fixture(scope="session") +def redis_url() -> Iterator[str]: + server = shutil.which("redis-server") + if server is None: + pytest.skip("redis-server is not installed") + port = _free_port() + process = subprocess.Popen( + [server, "--port", str(port), "--save", "", "--appendonly", "no"], + stdout=subprocess.DEVNULL, + ) + try: + for _ in range(100): + try: + socket.create_connection(("127.0.0.1", port), timeout=0.1).close() + break + except OSError: + time.sleep(0.05) + yield f"redis://127.0.0.1:{port}" + finally: + process.terminate() + process.wait() + + +@pytest.fixture(params=["memory", "redis"]) +async def stream_client( + client: Client, request: pytest.FixtureRequest +) -> AsyncIterator[Client]: + # Every live test runs on the memory Store and on Redis, since both must keep one set of + # rules. + plugin: StreamPlugin + if request.param == "memory": + plugin = MemoryStreamPlugin() + else: + redis_url = request.getfixturevalue("redis_url") + plugin = RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}") + config = client.config() + config["plugins"] = [plugin] + yield Client(**config) diff --git a/tests/contrib/streams/test_chain.py b/tests/contrib/streams/test_chain.py index 7e9891807..06b439edd 100644 --- a/tests/contrib/streams/test_chain.py +++ b/tests/contrib/streams/test_chain.py @@ -18,6 +18,7 @@ def info(**fields: Any) -> Any: "attempt": 1, "execution_timeout": None, "workflow_start_time": START, + "execution_expiration_time": None, } return SimpleNamespace(**{**base, **fields}) @@ -60,6 +61,17 @@ def info(**fields: Any) -> Any: RuntimeError("x"), False, ), + ( + "no time left in a later run of the chain", + info( + retry_policy=RETRY, + execution_timeout=timedelta(seconds=10), + workflow_start_time=START + timedelta(seconds=9), + execution_expiration_time=START + timedelta(seconds=10), + ), + RuntimeError("x"), + False, + ), ], ) def test_chain_continues(case: str, run: Any, failure: Any, expected: bool): diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index 17bb8bae6..1e7c5450a 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -2,13 +2,11 @@ import concurrent.futures import dataclasses import gc -import shutil -import socket -import subprocess import uuid -from collections.abc import AsyncIterator, Generator, Iterator, Sequence +from collections.abc import Generator, Sequence from dataclasses import dataclass from datetime import timedelta +from types import SimpleNamespace from typing import Any, cast import pytest @@ -20,7 +18,6 @@ from temporalio.contrib import streams from temporalio.contrib.streams.proto import v1 as ext_pb from temporalio.contrib.streams.redis import ( - RedisStreamPlugin, StreamPlugin, _WorkflowInbound, ) @@ -39,60 +36,6 @@ ) from tests.helpers import assert_task_fail_eventually, new_worker -# Ports that other local services own. -_RESERVED = {7333, 7861, 7931, 6431, 7101, 7102, 7103} - - -def _free_port() -> int: - while True: - with socket.socket() as sock: - sock.bind(("127.0.0.1", 0)) - port = sock.getsockname()[1] - if port not in _RESERVED: - return port - - -@pytest.fixture(scope="module") -def redis_url() -> Iterator[str]: - server = shutil.which("redis-server") - if server is None: - pytest.skip("redis-server is not installed") - port = _free_port() - process = subprocess.Popen( - [server, "--port", str(port), "--save", "", "--appendonly", "no"], - stdout=subprocess.DEVNULL, - ) - try: - for _ in range(100): - try: - socket.create_connection(("127.0.0.1", port), timeout=0.1).close() - break - except OSError: - import time - - time.sleep(0.05) - yield f"redis://127.0.0.1:{port}" - finally: - process.terminate() - process.wait() - - -@pytest.fixture(params=["memory", "redis"]) -async def stream_client( - client: Client, request: pytest.FixtureRequest -) -> AsyncIterator[Client]: - # Every live test runs on the memory Store and on Redis, since both must keep one set of - # rules. - plugin: StreamPlugin - if request.param == "memory": - plugin = MemoryStreamPlugin() - else: - redis_url = request.getfixturevalue("redis_url") - plugin = RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}") - config = client.config() - config["plugins"] = [plugin] - yield Client(**config) - @dataclass class Turn: @@ -729,7 +672,9 @@ def test_closing_an_evicted_run_touches_no_runtime(): # Workflow's loop. The interceptor must let it pass with no call to the runtime, or # Python prints "Exception ignored". inbound = _WorkflowInbound(_NeverEnds()) - run = inbound.execute_workflow(cast(ExecuteWorkflowInput, object())) + run = inbound.execute_workflow( + cast(ExecuteWorkflowInput, cast(object, SimpleNamespace(headers={}))) + ) run.send(None) run.close() diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py new file mode 100644 index 000000000..1ea1def5a --- /dev/null +++ b/tests/contrib/streams/test_workflow_module.py @@ -0,0 +1,368 @@ +import asyncio +import uuid +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +import pytest + +from temporalio import workflow +from temporalio.client import ( + Client, + WorkflowFailureError, + WorkflowHandle, +) +from temporalio.contrib import streams +from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.contrib.streams.redis import RedisStreamPlugin +from temporalio.contrib.streams.testing import MemoryStreamPlugin +from temporalio.exceptions import ApplicationError +from tests.helpers import assert_task_fail_eventually, new_worker + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@streams.defn(name="events") +class Events: + turns = streams.topic(Turn) + tokens = streams.topic(str) + + +async def read(handle: WorkflowHandle[Any, Any], name: str = "events") -> list[object]: + stream = streams.for_handle( + handle, name=name, wait_for_create=timedelta(seconds=20) + ) + stream.topic("turns", type=Turn) + stream.topic("tokens", type=str) + return [item async for item in stream.subscribe()] + + +def data(items: Sequence[object]) -> list[object]: + return [item.data for item in items if isinstance(item, streams.Record)] + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + stream = await stream_workflow.create_stream(Events) + turns = stream.topic(Events.turns) + await turns.publish(Turn("user", prompt), Turn("model")) + await stream.topic("tokens").publish(*prompt.split()) + return prompt.upper() + + +async def test_create_publish_and_close_with_run(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" + items = await read(handle) + assert data(items) == [Turn("user", "hi there"), Turn("model"), "hi", "there"] + closed = items[-1] + assert isinstance(closed, streams.Closed) and closed.result == "HI THERE" + + +@workflow.defn +class Turns: + @workflow.run + async def run(self, fail_at: int) -> None: + for n in range(2): + try: + async with await stream_workflow.create_stream( + Events, name=f"turn-{n}", close_with_run=False + ) as turn: + await turn.topic(Events.tokens).publish(f"token-{n}") + if n == fail_at: + raise ApplicationError("turn failed") + except ApplicationError: + pass + await workflow.sleep(timedelta(milliseconds=10)) + + +async def test_async_with_closes_the_stream(stream_client: Client): + async with new_worker(stream_client, Turns) as worker: + handle = await stream_client.start_workflow( + Turns.run, 1, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + first = await read(handle, "turn-0") + second = await read(handle, "turn-1") + assert data(first) == ["token-0"] + assert isinstance(first[-1], streams.Closed) and first[-1].failure is None + assert data(second) == ["token-1"] + assert isinstance(second[-1], streams.Closed) + assert str(second[-1].failure) == "turn failed" + + +@workflow.defn +class Continuing: + @workflow.run + async def run(self, stream: stream_workflow.WorkflowStream[Events] | None) -> str: + if stream is None: + stream = await stream_workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("run 1") + workflow.continue_as_new(stream) + assert stream.definition is Events + assert stream == stream_workflow.stream_handle(Events) + await stream.topic(Events.tokens).publish("run 2") + 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" + items = await read(handle) + assert data(items) == ["run 1", "run 2"] + closed = [item for item in items if isinstance(item, streams.Closed)] + assert len(closed) == 1 and closed[0].result == "done" + + +@workflow.defn +class Mistakes: + def __init__(self) -> None: + self._stream: stream_workflow.WorkflowStream[Events] | None = None + + @workflow.run + async def run(self) -> list[str]: + errors = [] + stream = await stream_workflow.create_stream(Events, close_with_run=False) + self._stream = stream + try: + stream.topic("unknown") + except streams.StreamError as err: + errors.append(err.reason) + await stream.close(result="early") + try: + await stream.topic(Events.tokens).publish("late") + except streams.StreamClosedError as err: + errors.append(err.reason) + try: + await stream_workflow.create_stream(Events) + except streams.StreamClosedError as err: + errors.append(err.reason) + return errors + + @workflow.query + def publish_in_query(self) -> str: + assert self._stream is not None + try: + self._stream.topic(Events.tokens)._stream._append("tokens", ("x",)) + except streams.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() == [ + "unknown_topic", + "stream_closed", + "stream_closed", + ] + assert await handle.query(Mistakes.publish_in_query) == "ReadOnlyContextError" + items = await read(handle) + assert isinstance(items[-1], streams.Closed) and items[-1].result == "early" + + +@workflow.defn +class Failing: + @workflow.run + async def run(self) -> None: + stream = await stream_workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("before") + raise ApplicationError("run failed", non_retryable=True) + + +async def test_close_with_run_takes_the_failure(stream_client: Client): + async with new_worker(stream_client, Failing) as worker: + handle = await stream_client.start_workflow( + Failing.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + with pytest.raises(WorkflowFailureError): + await handle.result() + items = await read(handle) + assert data(items) == ["before"] + assert isinstance(items[-1], streams.Closed) + assert str(items[-1].failure) == "run failed" + + +@workflow.defn +class LeavesOpen: + @workflow.run + async def run(self) -> None: + stream = await stream_workflow.create_stream( + Events, auto_close=timedelta(seconds=1), close_with_run=False + ) + await stream.topic(Events.tokens).publish("only") + + +async def test_auto_close_closes_a_quiet_stream(stream_client: Client): + async with new_worker(stream_client, LeavesOpen) as worker: + handle = await stream_client.start_workflow( + LeavesOpen.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + await asyncio.sleep(1.5) + stream = streams.for_handle(handle, name="events") + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert data(items) == ["only"] + assert isinstance(items[-1], streams.Closed) + assert items[-1].result is None and items[-1].failure is None + + +@workflow.defn +class CreatesOn: + @workflow.run + async def run(self, backend: str | None, again_without_close: bool) -> None: + await stream_workflow.create_stream(Events, backend=backend) + if again_without_close: + await stream_workflow.create_stream(Events, close_with_run=False) + + +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 take the run on. + for task_client, args, message in ( + (client, [None, False], "No Stream plugin"), + (stream_client, ["elsewhere", False], "backend elsewhere"), + (stream_client, [None, True], "close_with_run=False"), + ): + async with new_worker(task_client, CreatesOn) as worker: + handle = await task_client.start_workflow( + CreatesOn.run, + args=args, + 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 CronTurn: + @workflow.run + async def run(self) -> str: + stream = await stream_workflow.create_stream(Events) + await stream.topic(Events.tokens).publish("turn") + return "done" + + +async def test_cron_runs_share_an_open_stream(stream_client: Client): + # A cron chain goes on after each run, so no run closes the chain's Stream, and each run + # writes to it again. + async with new_worker(stream_client, CronTurn) as worker: + handle = await stream_client.start_workflow( + CronTurn.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + cron_schedule="@every 1s", + ) + stream = streams.for_handle( + handle, name="events", wait_for_create=timedelta(seconds=20) + ) + reader = stream.subscribe() + first = await reader.__anext__() + second = await reader.__anext__() + await reader.aclose() + await handle.terminate() + assert isinstance(first, streams.Record) and first.data == "turn" + assert isinstance(second, streams.Record) and second.data == "turn" + items = [item async for item in stream.subscribe(wait_new_records=False)] + assert not any(isinstance(item, streams.Closed) for item in items) + + +def test_plugins_take_their_backend_name(): + redis = RedisStreamPlugin( + "redis://127.0.0.1:1", response_timeout=timedelta(seconds=1), default=True + ) + assert (redis.backend_name, redis.default) == ("redis", True) + assert RedisStreamPlugin("redis://127.0.0.1:1", name="old").backend_name == "old" + assert MemoryStreamPlugin().backend_name == "memory" + + +@workflow.defn +class SyncHandlers: + @workflow.init + def __init__(self, prompt: str) -> None: + self._stream = stream_workflow.stream_handle(Events) + self._prompt = prompt + self._done = False + + @workflow.run + async def run(self, prompt: str) -> None: + # The init made the handle before the create, since a handle makes no call. + assert prompt == self._prompt + await stream_workflow.create_stream(Events, close_with_run=False) + await workflow.wait_condition(lambda: self._done) + + @workflow.signal + def token(self, text: str) -> None: + self._stream.topic(Events.tokens).publish_nowait(text, self._prompt) + + @workflow.update + def finish(self) -> str: + self._stream.close_nowait(result="finished") + try: + self._stream.topic(Events.tokens).publish_nowait("late") + except streams.StreamClosedError as err: + return err.reason + finally: + self._done = True + return "wrote" + + +async def test_sync_handlers_publish_and_close_with_no_await(stream_client: Client): + async with new_worker(stream_client, SyncHandlers) as worker: + handle = await stream_client.start_workflow( + SyncHandlers.run, "p", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.signal(SyncHandlers.token, "a") + assert await handle.execute_update(SyncHandlers.finish) == "stream_closed" + await handle.result() + items = await read(handle) + assert data(items) == ["a", "p"] + assert isinstance(items[-1], streams.Closed) and items[-1].result == "finished" + + +@workflow.defn +class CreatesInInit: + @workflow.init + def __init__(self, prompt: str) -> None: + self._stream = stream_workflow.create_stream_nowait(Events) + self._stream.topic(Events.tokens).publish_nowait(prompt) + + @workflow.run + async def run(self, prompt: str) -> str: + await self._stream.topic(Events.tokens).publish(prompt.upper()) + return "done" + + +async def test_init_creates_and_publishes_with_no_await(stream_client: Client): + async with new_worker(stream_client, CreatesInInit) as worker: + handle = await stream_client.start_workflow( + CreatesInInit.run, + "hi", + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == "done" + items = await read(handle) + assert data(items) == ["hi", "HI"] + assert isinstance(items[-1], streams.Closed) and items[-1].result == "done"