From 5bd5c52620811d5f72bbd3b8cf992f17909e4405 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 21:02:19 -0700 Subject: [PATCH 1/4] Added RedisStreamPlugin over the linked temporalio-streams crate. The bridge links the crate behind a streams feature that is on by default, so the plugin can build the Redis worker extension by name. The plugin buffers each Workflow Activation's Stream calls and hands them to the crate at the end of the Activation, and it gives Activities and clients their writers. --- temporalio/bridge/Cargo.lock | 131 ++++++++ temporalio/bridge/Cargo.toml | 23 +- temporalio/bridge/src/extension.rs | 13 +- temporalio/contrib/streams/redis.py | 461 ++++++++++++++++++++++++++++ tests/contrib/streams/test_redis.py | 269 ++++++++++++++++ 5 files changed, 893 insertions(+), 4 deletions(-) create mode 100644 temporalio/contrib/streams/redis.py create mode 100644 tests/contrib/streams/test_redis.py diff --git a/temporalio/bridge/Cargo.lock b/temporalio/bridge/Cargo.lock index 48e4a9416b..be25623919 100644 --- a/temporalio/bridge/Cargo.lock +++ b/temporalio/bridge/Cargo.lock @@ -29,6 +29,32 @@ version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" +[[package]] +name = "arc-swap" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "380962b3d93b99642e50b9febf72c071e97e8b32d86eae9cf4ef582bb960ae6e" +dependencies = [ + "rustversion", +] + +[[package]] +name = "arcstr" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" + +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-trait" version = "0.1.92" @@ -307,7 +333,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", ] [[package]] @@ -344,6 +374,12 @@ dependencies = [ "libc", ] +[[package]] +name = "crc16" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "338089f42c427b86394a5ee60ff321da23a5c89c9d89514c829687b26359fcff" + [[package]] name = "crc32fast" version = "1.5.2" @@ -540,6 +576,26 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "event-listener" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" +dependencies = [ + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "fastrand" version = "2.5.0" @@ -1493,6 +1549,12 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + [[package]] name = "parking_lot" version = "0.12.5" @@ -2033,6 +2095,39 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "redis" +version = "1.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5b98ae99461e02895c3a96a25dbe4d852a3d5e5b815d31d5d9e4b5344d99e9" +dependencies = [ + "arc-swap", + "arcstr", + "async-lock", + "backon", + "bytes", + "cfg-if", + "combine", + "crc16", + "futures-channel", + "futures-util", + "itoa", + "log", + "percent-encoding", + "pin-project-lite", + "rand 0.10.3", + "rustls", + "rustls-native-certs", + "ryu", + "sha1_smol", + "socket2", + "tokio", + "tokio-rustls", + "tokio-util", + "url", + "xxhash-rust", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2382,6 +2477,12 @@ dependencies = [ "serde", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sharded-slab" version = "0.1.7" @@ -2591,6 +2692,7 @@ dependencies = [ "temporalio-client", "temporalio-common", "temporalio-sdk-core", + "temporalio-streams", "tokio", "tokio-rustls", "tokio-stream", @@ -2774,6 +2876,28 @@ dependencies = [ "zip", ] +[[package]] +name = "temporalio-streams" +version = "0.1.0" +source = "git+https://github.com/moetemp/sdk-streams?rev=5b366595abc08adf26358dd1a54afdf3fb3d4097#5b366595abc08adf26358dd1a54afdf3fb3d4097" +dependencies = [ + "async-trait", + "aws-lc-rs", + "crc16", + "parking_lot", + "prost", + "prost-wkt-types", + "redis", + "temporalio-common", + "temporalio-sdk-core", + "thiserror", + "tokio", + "tokio-util", + "tonic", + "tonic-prost-build", + "tracing", +] + [[package]] name = "termtree" version = "0.5.1" @@ -2883,6 +3007,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "libc", "pin-project-lite", "tokio", @@ -3620,6 +3745,12 @@ dependencies = [ "rustix", ] +[[package]] +name = "xxhash-rust" +version = "0.8.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "550a2b930b62486a393c52d5c3b84bff264b28aa437ed64694d31e93b1757af7" + [[package]] name = "yoke" version = "0.8.3" diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml index 4e45b00ba3..3ce419b03d 100644 --- a/temporalio/bridge/Cargo.toml +++ b/temporalio/bridge/Cargo.toml @@ -25,18 +25,23 @@ crate-type = ["cdylib"] # The extra `tokio-rustls/*` entry is sdk-python-specific: the bridge has its own # direct tokio-rustls dep (client.rs custom verifier) that Ruby lacked, so its # provider must track the build too — else `cargo tree -i ring` would still find ring. -default = ["tls-ring"] +default = ["tls-ring", "streams"] tls-ring = [ "temporalio-sdk-core/tls-ring", "temporalio-client/tls-ring", + "temporalio-streams?/tls-ring", "tokio-rustls/ring", ] tls-aws-lc = [ "temporalio-sdk-core/tls-aws-lc", "temporalio-client/tls-aws-lc", + "temporalio-streams?/tls-aws-lc", "tokio-rustls/aws_lc_rs", ] -fips = ["tls-aws-lc", "dep:aws-lc-rs"] +fips = ["tls-aws-lc", "dep:aws-lc-rs", "streams", "temporalio-streams?/fips"] +# Links the Stream runtime as the `temporalio-streams-redis` worker extension. No stream code +# runs unless a plugin builds the extension. +streams = ["dep:temporalio-streams"] [dependencies] anyhow = "1.0" @@ -84,6 +89,20 @@ tonic = "0.14" tracing = "0.1" url = "2.2" +# The rev moves with the sdk-core pin, and the patch below makes the crate use the submodule's +# crates, so the bridge sees one type of each. +[dependencies.temporalio-streams] +git = "https://github.com/moetemp/sdk-streams" +rev = "5b366595abc08adf26358dd1a54afdf3fb3d4097" +default-features = false +optional = true + +# The key must match the sdk-rust URL in the sdk-streams manifest byte for byte, or Cargo +# ignores the patch with only a warning. +[patch."https://github.com/moetemp/sdk-rust"] +temporalio-common = { path = "./sdk-core/crates/common" } +temporalio-sdk-core = { path = "./sdk-core/crates/sdk-core" } + [profile.release] opt-level = 3 debug = false diff --git a/temporalio/bridge/src/extension.rs b/temporalio/bridge/src/extension.rs index 4958765503..e5ba9cec77 100644 --- a/temporalio/bridge/src/extension.rs +++ b/temporalio/bridge/src/extension.rs @@ -61,8 +61,11 @@ fn call_error(err: ExtensionError) -> PyErr { /// Builds the extension with `name` from its encoded config. Only the extensions that this build /// links are known. -// A release build links no extension yet, so it never reads the config. -#[cfg_attr(not(debug_assertions), allow(unused_variables))] +// A release build with no extension feature never reads the config. +#[cfg_attr( + not(any(debug_assertions, feature = "streams")), + allow(unused_variables) +)] pub fn new_worker_extension( runtime_ref: &runtime::RuntimeRef, name: &str, @@ -73,6 +76,12 @@ pub fn new_worker_extension( // An extension may start tasks on the runtime that it lives on. let _guard = runtime.core.tokio_handle().enter(); let extension: Arc = match name { + #[cfg(feature = "streams")] + temporalio_streams::EXTENSION_NAME => { + let meter = runtime.core.telemetry().get_temporal_metric_meter(); + temporalio_streams::RedisStreamExtension::from_config_bytes(&config, meter) + .map_err(|err| PyValueError::new_err(err.to_string()))? + } #[cfg(debug_assertions)] test_extension::NAME => Arc::new(test_extension::TestExtension::new(&config)?), _ => { diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py new file mode 100644 index 0000000000..d0d6b70177 --- /dev/null +++ b/temporalio/contrib/streams/redis.py @@ -0,0 +1,461 @@ +"""The Redis Stream plugin. + +.. warning:: + This module is experimental. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +from collections.abc import Generator, Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +import temporalio.activity +import temporalio.api.common.v1 +import temporalio.api.streamservice.v1 as stream_pb +import temporalio.client +import temporalio.converter +import temporalio.exceptions +import temporalio.runtime +import temporalio.worker +import temporalio.workflow +from temporalio.contrib.streams import _activity, _errors, _ids, _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 +from temporalio.contrib.streams.system import models +from temporalio.plugin import SimplePlugin + +EXTENSION_NAME = "temporalio-streams-redis" + +_SERVICE = "temporal.api.streamservice.v1.StreamService" +# The plugin puts these into the headers of each Activity and Local Activity that a Workflow +# schedules. The key keeps the Producer id the same across retries. +_REF_HEADER = "__temporal_stream_ref" +_KEY_HEADER = "__temporal_stream_key" + + +class RedisStreamPlugin(SimplePlugin): + """Carries Stream calls to Redis, through the temporalio-streams runtime in the bridge. + + Put it on the client. A Worker takes its client's plugins, so one plugin serves Workflow + code, Activities and the client. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + url: str | Sequence[str], + *, + cluster: bool = False, + key_prefix: str | None = None, + retention: timedelta | None = None, + safety_ttl: timedelta | None = None, + durable_writes: int = 0, + durable_timeout: timedelta | None = None, + response_timeout: timedelta | None = None, + runtime: temporalio.runtime.Runtime | None = None, + ) -> None: + """Create the plugin. It connects to Redis on the first Stream call. + + Args: + url: A Redis URL, or the seed URLs of a cluster. Credentials and TLS come from the + URL. + cluster: Whether the URLs are of a Redis cluster. + key_prefix: The prefix of every key. Defaults to ``temporal-streams``. + retention: The default Retention of a new Stream. Defaults to 7 days. + safety_ttl: How long an open Workflow-owned Stream stays after its last write. + Defaults to 90 days. + durable_writes: When above zero, each Workflow flush waits for this many replicas. + 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. + """ + config = ext_pb.RedisStreamConfig( + urls=[url] if isinstance(url, str) else list(url), + cluster=cluster, + key_prefix=key_prefix or "", + durable_writes=durable_writes, + ) + for name, 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) + self._extension = temporalio.worker.WorkerExtension( + EXTENSION_NAME, config.SerializeToString(), runtime=runtime + ) + self._temporal_stream_backend = StreamBackend(self._extension) + super().__init__( + "temporalio.RedisStreamPlugin", interceptors=[_StreamInterceptor()] + ) + + def configure_worker( + self, config: temporalio.worker.WorkerConfig + ) -> temporalio.worker.WorkerConfig: + """Add the plugin's interceptors and give Core the extension.""" + config = super().configure_worker(config) + if config.get("extension") is not None: + raise ValueError("The Worker has another extension") + config["extension"] = self._extension + return config + + async def on_workflow_activation_complete( + self, input: temporalio.worker.WorkflowActivationCompleteInput + ) -> None: + """Hand the Activation's Stream calls to the runtime, encoded.""" + calls = [item for item in input.items if isinstance(item, _WorkflowCall)] + if not calls: + return + info = input.info + request = ext_pb.BufferWorkflowCallsRequest( + run_id=input.run_id, + withdraw_run=_withdraw_run(info), + chain_continues=any(call.chain_continues for call in calls), + ) + for call in _merge_appends(calls): + request.calls.append( + await _encode_call(call, input.data_converter, info.namespace) + ) + self._temporal_stream_backend.call("BufferWorkflowCalls", request) + + +def _withdraw_run(info: temporalio.workflow.Info) -> str: + # A retry run withdraws the Records of the run it retries. + if info.attempt > 1 and info.continued_run_id: + return info.continued_run_id + return "" + + +@dataclass(frozen=True) +class _WorkflowCall: + operation: str + request: Any + when_run_closes: bool = False + chain_continues: bool = False + + +def _merge_appends(calls: list[_WorkflowCall]) -> list[_WorkflowCall]: + # Each append is one batch with its own sequence, so adjacent appends to one Stream go + # as one. + merged: list[_WorkflowCall] = [] + for call in calls: + last = merged[-1] if merged else None + if ( + last is not None + and last.operation == call.operation == "AppendRecords" + and last.request.stream == call.request.stream + ): + records = [*(last.request.records or []), *(call.request.records or [])] + merged[-1] = dataclasses.replace( + last, request=dataclasses.replace(last.request, records=records) + ) + else: + merged.append(call) + return merged + + +async def _encode_call( + call: _WorkflowCall, + converter: temporalio.converter.DataConverter, + namespace: str, +) -> ext_pb.WorkflowCall: + request = call.request + if call.operation == "CreateStream": + create = stream_pb.CreateStreamRequest( + namespace=namespace, stream_id=request.stream_id + ) + if request.retention is not None: + create.retention.FromTimedelta(request.retention) + return ext_pb.WorkflowCall(create=create) + stream = stream_pb.StreamRef(id=request.stream.id if request.stream else "") + if call.operation == "AppendRecords": + records = request.records or [] + bodies = await converter.encode([record.body for record in records]) + return ext_pb.WorkflowCall( + append=stream_pb.AppendRecordsRequest( + namespace=namespace, + stream=stream, + records=[ + stream_pb.AppendRecord(topic=record.topic, body=body) + for record, body in zip(records, bodies) + ], + ) + ) + close = stream_pb.CloseStreamRequest(namespace=namespace, stream=stream) + if request.result is not None: + [result] = await converter.encode([request.result]) + close.result.CopyFrom(result) + if request.failure is not None: + await converter.encode_failure(request.failure, close.failure) + return ext_pb.WorkflowCall(close=close, when_run_closes=call.when_run_closes) + + +class _DoneHandle(temporalio.workflow.NexusOperationHandle[Any]): + # The call is in the Activation's list, so the handle is done at once. + def __init__(self, result: Any) -> None: + self._result = result + + def cancel(self) -> bool: + return False + + def __await__(self) -> Generator[Any, Any, Any]: + return self._value().__await__() + + async def _value(self) -> Any: + return self._result + + @property + def operation_token(self) -> str | None: + return None + + +class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor): + def workflow_interceptor_class( + self, input: temporalio.worker.WorkflowInterceptorClassInput + ) -> type[temporalio.worker.WorkflowInboundInterceptor] | None: + return _WorkflowInbound + + def intercept_activity( + self, next: temporalio.worker.ActivityInboundInterceptor + ) -> temporalio.worker.ActivityInboundInterceptor: + return _ActivityInbound(next) + + +class _WorkflowInbound(temporalio.worker.WorkflowInboundInterceptor): + def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: + super().init(_WorkflowOutbound(outbound)) + + async def execute_workflow( + self, input: temporalio.worker.ExecuteWorkflowInput + ) -> Any: + try: + result = await super().execute_workflow(input) + except temporalio.workflow.ContinueAsNewError: + # The next run continues the run chain's Streams. + raise + except asyncio.CancelledError: + _close_with_run( + failure=temporalio.exceptions.CancelledError("Workflow cancelled") + ) + raise + except BaseException as err: + if _fails_run(err): + _close_with_run(failure=err, chain_continues=_chain_continues(err)) + raise + # A cron run that completes still goes on with the chain. + _close_with_run( + result=result, + chain_continues=bool(temporalio.workflow.info().cron_schedule), + ) + return result + + +def _fails_run(err: BaseException) -> bool: + # Any other exception fails the Workflow Task, and the run goes on. + return temporalio.workflow._Runtime.current().workflow_is_failure_exception(err) + + +def _chain_continues(err: BaseException) -> bool: + # The server starts another run of the chain, which goes on writing the default Stream, + # so the runtime skips the run's closes. + info = temporalio.workflow.info() + if info.cron_schedule: + return True + policy = info.retry_policy + if policy is None or ( + policy.maximum_attempts != 0 and info.attempt >= policy.maximum_attempts + ): + return False + return not _non_retryable(err, policy.non_retryable_error_types) + + +def _non_retryable(err: BaseException, types: Sequence[str] | None) -> bool: + if isinstance(err, temporalio.exceptions.ApplicationError): + return err.non_retryable or (err.type or "") in (types or []) + return type(err).__name__ in (types or []) + + +def _close_with_run( + *, + result: Any | None = None, + failure: BaseException | None = None, + chain_continues: bool = False, +) -> None: + for stream in _workflow._state().streams.values(): + if not stream.auto_close or not stream._used or stream._closed: + continue + stream._closed = True + temporalio.workflow.send_to_plugins( + _WorkflowCall( + "CloseStream", + models.CloseStreamRequest( + stream=models.StreamRef(id=stream.ref.id), + result=result, + failure=failure, + ), + when_run_closes=True, + chain_continues=chain_continues, + ), + # A close is idempotent, and its run may have no task left to send it on. + during_replay=True, + ) + + +class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): + _temporal_stream_transport = True + + async def start_system_nexus_operation( + self, input: temporalio.worker.StartNexusOperationInput[Any, Any] + ) -> temporalio.workflow.NexusOperationHandle[Any]: + if input.service != _SERVICE: + return await super().start_system_nexus_operation(input) + operation = input.operation_name + if operation == "CreateStream": + result: Any = models.CreateStreamResponse( + stream=models.StreamRef(id=input.input.stream_id) + ) + elif operation == "AppendRecords": + result = models.AppendRecordsResponse() + elif operation == "CloseStream": + result = models.CloseStreamResponse() + else: + raise _errors.StreamError(f"Workflow code can't call {operation}") + temporalio.workflow.send_to_plugins( + _WorkflowCall(operation, input.input), + # The runtime already has the Records of a replayed task. A create or a close is + # idempotent, so it goes again, for a run whose buffer an eviction dropped. + during_replay=operation != "AppendRecords", + ) + return _DoneHandle(result) + + def start_activity( + self, input: temporalio.worker.StartActivityInput + ) -> temporalio.workflow.ActivityHandle[Any]: + input.headers = {**input.headers, **_activity_headers()} + return super().start_activity(input) + + def start_local_activity( + self, input: temporalio.worker.StartLocalActivityInput + ) -> temporalio.workflow.ActivityHandle[Any]: + input.headers = {**input.headers, **_activity_headers()} + return super().start_local_activity(input) + + +def _activity_headers() -> dict[str, temporalio.api.common.v1.Payload]: + info = temporalio.workflow.info() + # Both ends use the default converter, so a custom user converter can't break the headers. + converter = temporalio.converter.PayloadConverter.default + ref = _ids.default_stream_id(info.workflow_id, info.first_execution_run_id) + return { + _REF_HEADER: converter.to_payload(ref), + _KEY_HEADER: converter.to_payload(str(temporalio.workflow.uuid4())), + } + + +class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor): + async def execute_activity( + self, input: temporalio.worker.ExecuteActivityInput + ) -> Any: + streams = _activity_streams(input) + if streams is None: + return await super().execute_activity(input) + token = _activity.current.set(streams) + try: + if ( + streams.default_ref is not None + and temporalio.activity.info().attempt > 1 + ): + # The open of a new attempt withdraws the Records of the earlier ones, also + # when this attempt publishes nothing. + await streams.writer(streams.default_ref.id).open() + result = await super().execute_activity(input) + except BaseException: + await _close_quietly(streams) + raise + else: + await streams.close() + return result + finally: + _activity.current.reset(token) + + +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 _plugin_backend(client: temporalio.client.Client) -> StreamBackend | None: + for plugin in client.config()["plugins"]: + if isinstance(plugin, RedisStreamPlugin): + return plugin._temporal_stream_backend + return None + + +def _activity_streams( + input: temporalio.worker.ExecuteActivityInput, +) -> _activity.ActivityStreams | None: + try: + client = temporalio.activity.client() + except RuntimeError: + return None + backend = _plugin_backend(client) + if backend is None: + return None + info = temporalio.activity.info() + converter = temporalio.converter.PayloadConverter.default + ref = input.headers.get(_REF_HEADER) + key = input.headers.get(_KEY_HEADER) + key_text = converter.from_payload(key, str) if key is not None else "" + epoch = stream_pb.ProducerEpoch(attempt=info.attempt) + if info.is_local: + task = info.workflow_task + if task is not None: + epoch.task_started_event_id = task.started_event_id + epoch.task_attempt = task.attempt + if task.started_time is not None: + epoch.task_started_time.FromDatetime(task.started_time) + run_id = info.workflow_run_id or "" + producer = stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_LOCAL_ACTIVITY, + id=_ids.local_activity_producer_id(run_id, key_text), + run_id=run_id, + ) + else: + epoch.attempt_scheduled_time.FromDatetime(info.current_attempt_scheduled_time) + if info.in_workflow: + run_id = info.workflow_run_id or "" + producer = stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_ACTIVITY, + id=_ids.activity_producer_id(run_id, info.activity_id, key_text), + run_id=run_id, + ) + else: + producer = stream_pb.Producer( + kind=stream_pb.ProducerKind.PRODUCER_KIND_STANDALONE_ACTIVITY, + id=_ids.standalone_activity_producer_id( + info.activity_id, info.activity_run_id or "" + ), + ) + producer.epoch.CopyFrom(epoch) + return _activity.ActivityStreams( + backend=backend, + converter=client.data_converter, + namespace=info.namespace, + producer=producer, + default_ref=StreamRef(converter.from_payload(ref, str)) + if ref is not None + else None, + ) diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py new file mode 100644 index 0000000000..0924b1eb47 --- /dev/null +++ b/tests/contrib/streams/test_redis.py @@ -0,0 +1,269 @@ +import asyncio +import shutil +import socket +import subprocess +import uuid +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass +from datetime import timedelta + +import pytest + +from temporalio import activity, workflow +from temporalio.client import Client +from temporalio.common import RetryPolicy +from temporalio.contrib import streams +from temporalio.contrib.streams.redis import RedisStreamPlugin +from temporalio.worker import UnsandboxedWorkflowRunner +from tests.helpers import 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 +async def stream_client(client: Client, redis_url: str) -> AsyncIterator[Client]: + config = client.config() + config["plugins"] = [RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}")] + yield Client(**config) + + +@dataclass +class Turn: + kind: str + text: str = "" + + +@activity.defn +async def call_model(prompt: str) -> str: + tokens = streams.activity_stream().topic("tokens", type=str) + for word in prompt.split(): + tokens.publish(word) + return prompt.upper() + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + turns = streams.workflow_stream().topic("turns", type=Turn) + turns.publish(Turn("model")) + reply = await workflow.execute_activity( + call_model, prompt, start_to_close_timeout=timedelta(seconds=30) + ) + turns.publish(Turn("done", reply)) + return reply + + +async def read_all(stream: streams.ClientStream) -> list[streams.Record | object]: + return [item async for item in stream.subscribe(idle_timeout=timedelta(seconds=20))] + + +async def test_workflow_and_activity_write_the_default_stream(stream_client: Client): + async with new_worker(stream_client, Agent, activities=[call_model]) as worker: + handle = await stream_client.start_workflow( + Agent.run, "hi there", id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + stream = streams.for_handle(handle, wait_for_create=timedelta(seconds=20)) + turns = stream.topic("turns", type=Turn) + stream.topic("tokens", type=str) + items = await read_all(stream) + assert await handle.result() == "HI THERE" + data = [ + (item.topic, item.data) for item in items if isinstance(item, streams.Record) + ] + assert data == [ + ("turns", Turn("model")), + ("tokens", "hi"), + ("tokens", "there"), + ("turns", Turn("done", "HI THERE")), + ] + closed = items[-1] + assert isinstance(closed, streams.Closed) and closed.result == "HI THERE" + only_turns = [item async for item in turns.subscribe()] + assert [type(item) for item in only_turns] == [ + streams.Record, + streams.Record, + streams.Closed, + ] + + +@workflow.defn +class Chat: + @workflow.run + async def run(self, turns: int) -> None: + index = streams.workflow_stream().topic("turns", type=int) + for n in range(turns): + with streams.create_workflow_stream(f"turn-{n}") as turn: + index.publish(n) + turn.topic("tokens", type=str).publish(f"token-{n}") + await workflow.sleep(timedelta(milliseconds=10)) + + +async def test_scoped_streams_close_at_the_end_of_their_block(stream_client: Client): + async with new_worker(stream_client, Chat) as worker: + handle = await stream_client.start_workflow( + Chat.run, 2, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + turn = streams.for_handle(handle, name="turn-1") + items = [item async for item in turn.topic("tokens", type=str).subscribe()] + assert isinstance(items[0], streams.Record) and items[0].data == "token-1" + assert isinstance(items[-1], streams.Closed) + + +@workflow.defn +class WriteToRef: + @workflow.run + async def run(self, ref: streams.StreamRef) -> None: + streams.workflow_stream(ref).topic("t", type=str).publish("from workflow") + + +async def test_client_stream_takes_workflow_writes(stream_client: Client): + stream = await streams.create_stream(stream_client, f"chat-{uuid.uuid4()}") + stream.topic("t", type=str).publish("from client") + await stream.flush() + async with new_worker(stream_client, WriteToRef) as worker: + await stream_client.execute_workflow( + WriteToRef.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 == ["from client", "from workflow"] + assert isinstance(items[-1], streams.Closed) + + +attempts: dict[str, int] = {} + + +@activity.defn +async def flaky(key: str) -> None: + attempts[key] = attempts.get(key, 0) + 1 + tokens = streams.activity_stream().topic("tokens", type=str) + tokens.publish(f"attempt-{attempts[key]}") + await streams.activity_stream().flush() + if attempts[key] == 1: + raise RuntimeError("first attempt fails") + + +@workflow.defn +class Retrying: + @workflow.run + async def run(self, key: str) -> None: + streams.workflow_stream().topic("turns", type=str).publish("start") + await workflow.execute_activity( + flaky, + key, + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def test_activity_retry_withdraws_the_failed_attempt(stream_client: Client): + key = str(uuid.uuid4()) + async with new_worker(stream_client, Retrying, activities=[flaky]) as worker: + handle = await stream_client.start_workflow( + Retrying.run, key, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + items = await read_all(streams.for_handle(handle)) + records = [item for item in items if isinstance(item, streams.Record)] + superseded = [item for item in items if isinstance(item, streams.Superseded)] + assert [record.data for record in records] == ["start", "attempt-1", "attempt-2"] + assert len(superseded) == 1 + withdrawn = [record.data for record in records if superseded[0].withdraws(record)] + assert withdrawn == ["attempt-1"] + + +@workflow.defn +class PublishThenFailTask: + @workflow.run + async def run(self, key: str) -> None: + turns = streams.workflow_stream().topic("turns", type=str) + turns.publish("before") + await workflow.sleep(timedelta(milliseconds=10)) + turns.publish("after") + if attempts.setdefault(key, 0) == 0: + attempts[key] = 1 + raise RuntimeError("fail this Workflow Task once") + + +async def test_failed_workflow_task_is_withdrawn(stream_client: Client): + key = str(uuid.uuid4()) + # The test's counter must survive the retry of the task, so no sandbox. + async with new_worker( + stream_client, + PublishThenFailTask, + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + handle = await stream_client.start_workflow( + PublishThenFailTask.run, + key, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + items = await read_all(streams.for_handle(handle)) + records = [item for item in items if isinstance(item, streams.Record)] + # The failed task wrote nothing, so the reader gets each Record once. + assert [record.data for record in records] == ["before", "after"] + assert not any(isinstance(item, streams.Superseded) for item in items) + assert isinstance(items[-1], streams.Closed) + + +@workflow.defn +class Reading: + @workflow.run + async def run(self) -> str: + try: + streams.workflow_stream().topic("t", type=str).subscribe() + except TypeError as err: + return str(err) + return "read" + + +async def test_workflow_code_can_not_read(stream_client: Client): + async with new_worker(stream_client, Reading) as worker: + result = await stream_client.execute_workflow( + Reading.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + assert "can't read" in result + await asyncio.sleep(0) From 7ac12f4312c7b352d94a5daa02efcd029160c91f Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 21:28:18 -0700 Subject: [PATCH 2/4] Fixed the Redis plugin's review findings. Activity keys come from a per-run counter, so the plugin leaves the Workflow's random values alone. A retried attempt withdraws without a create, the last run always closes the default Stream, and chain_continues covers retries, cron and the execution timeout. The streams feature is now off by default and on in the official builds. --- pyproject.toml | 12 +- temporalio/bridge/Cargo.toml | 5 +- temporalio/contrib/streams/redis.py | 232 ++++++++++++++++++++-------- tests/contrib/streams/test_chain.py | 68 ++++++++ tests/contrib/streams/test_redis.py | 190 +++++++++++++++++++++++ 5 files changed, 435 insertions(+), 72 deletions(-) create mode 100644 tests/contrib/streams/test_chain.py diff --git a/pyproject.toml b/pyproject.toml index 396a11bc59..a9492c59ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,8 +97,8 @@ dev = [ ] [tool.poe.tasks] -build-develop = "uv run maturin develop --uv" -build-develop-with-release = { cmd = "uv run maturin develop --release --uv" } +build-develop = "uv run maturin develop --uv --features streams" +build-develop-with-release = { cmd = "uv run maturin develop --release --uv --features streams" } # FIPS builds: swap the rustls stack onto aws-lc-rs FIPS mode (aws-lc-fips-sys), eliminating `ring`. build-develop-fips = { cmd = "uv run maturin develop --uv --no-default-features --features fips", env = { TEMPORALIO_FIPS = "1" } } build-wheel-fips = { cmd = "uv run maturin build --release --no-default-features --features fips", env = { TEMPORALIO_FIPS = "1" } } @@ -164,18 +164,20 @@ filterwarnings = [ before-all = "pip install protoc-wheel-0" build = "cp310-win_amd64 cp310-manylinux_x86_64 cp310-manylinux_aarch64 cp310-musllinux_x86_64 cp310-musllinux_aarch64 cp310-macosx_x86_64 cp310-macosx_arm64" build-verbosity = 1 +# The official wheels link the Stream runtime. +environment = { MATURIN_PEP517_ARGS = "--features streams" } [tool.cibuildwheel.macos] -environment = { MACOSX_DEPLOYMENT_TARGET = "10.12" } +environment = { MACOSX_DEPLOYMENT_TARGET = "10.12", MATURIN_PEP517_ARGS = "--features streams" } [tool.cibuildwheel.linux] before-all = "sh scripts/cibuildwheel_before_all_linux.sh" before-build = "sh scripts/cibuildwheel_before_build_linux.sh" -environment = { PATH = "$PATH:$HOME/.cargo/bin", CARGO_NET_GIT_FETCH_WITH_CLI = "true" } +environment = { PATH = "$PATH:$HOME/.cargo/bin", CARGO_NET_GIT_FETCH_WITH_CLI = "true", MATURIN_PEP517_ARGS = "--features streams" } [[tool.cibuildwheel.overrides]] select = "*musllinux*" -environment = { PATH = "$PATH:$HOME/.cargo/bin", CARGO_NET_GIT_FETCH_WITH_CLI = "true", RUSTFLAGS = "-C target-feature=-crt-static" } +environment = { PATH = "$PATH:$HOME/.cargo/bin", CARGO_NET_GIT_FETCH_WITH_CLI = "true", RUSTFLAGS = "-C target-feature=-crt-static", MATURIN_PEP517_ARGS = "--features streams" } [tool.mypy] ignore_missing_imports = true diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml index 3ce419b03d..8a10eb2940 100644 --- a/temporalio/bridge/Cargo.toml +++ b/temporalio/bridge/Cargo.toml @@ -25,7 +25,7 @@ crate-type = ["cdylib"] # The extra `tokio-rustls/*` entry is sdk-python-specific: the bridge has its own # direct tokio-rustls dep (client.rs custom verifier) that Ruby lacked, so its # provider must track the build too — else `cargo tree -i ring` would still find ring. -default = ["tls-ring", "streams"] +default = ["tls-ring"] tls-ring = [ "temporalio-sdk-core/tls-ring", "temporalio-client/tls-ring", @@ -40,7 +40,8 @@ tls-aws-lc = [ ] fips = ["tls-aws-lc", "dep:aws-lc-rs", "streams", "temporalio-streams?/fips"] # Links the Stream runtime as the `temporalio-streams-redis` worker extension. No stream code -# runs unless a plugin builds the extension. +# runs unless a plugin builds the extension. The official builds turn it on. It is off by +# default, since its git dependencies would break offline builds of the sdist. streams = ["dep:temporalio-streams"] [dependencies] diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index d0d6b70177..447c0dcc5d 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -10,13 +10,14 @@ import dataclasses from collections.abc import Generator, Sequence from dataclasses import dataclass -from datetime import timedelta +from datetime import datetime, timedelta from typing import Any import temporalio.activity import temporalio.api.common.v1 import temporalio.api.streamservice.v1 as stream_pb import temporalio.client +import temporalio.common import temporalio.converter import temporalio.exceptions import temporalio.runtime @@ -28,6 +29,7 @@ from temporalio.contrib.streams.proto import v1 as ext_pb from temporalio.contrib.streams.system import models from temporalio.plugin import SimplePlugin +from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner EXTENSION_NAME = "temporalio-streams-redis" @@ -94,8 +96,9 @@ def __init__( EXTENSION_NAME, config.SerializeToString(), runtime=runtime ) self._temporal_stream_backend = StreamBackend(self._extension) + self._converter = temporalio.converter.DataConverter.default super().__init__( - "temporalio.RedisStreamPlugin", interceptors=[_StreamInterceptor()] + "temporalio.RedisStreamPlugin", interceptors=[_StreamInterceptor(self)] ) def configure_worker( @@ -106,13 +109,32 @@ def configure_worker( if config.get("extension") is not None: raise ValueError("The Worker has another extension") config["extension"] = self._extension + # Records get no serialization context, so Workflows, Activities and clients encode + # them alike. + client = config.get("client") + if client is not None: + self._converter = client.data_converter + runner = config.get("workflow_runner") + if isinstance(runner, SandboxedWorkflowRunner): + # The interceptor and the facade must share the run's Stream state, also when the + # restrictions don't pass temporalio through. + config["workflow_runner"] = dataclasses.replace( + runner, + restrictions=runner.restrictions.with_passthrough_modules( + "temporalio.contrib.streams" + ), + ) return config async def on_workflow_activation_complete( self, input: temporalio.worker.WorkflowActivationCompleteInput ) -> None: """Hand the Activation's Stream calls to the runtime, encoded.""" - calls = [item for item in input.items if isinstance(item, _WorkflowCall)] + calls = [ + item + for item in input.items + if getattr(item, "_temporal_stream_call", False) + ] if not calls: return info = input.info @@ -123,8 +145,10 @@ async def on_workflow_activation_complete( ) for call in _merge_appends(calls): request.calls.append( - await _encode_call(call, input.data_converter, info.namespace) + await _encode_call(call, self._converter, info.namespace) ) + # The call only adds to the run's buffer under a short lock, so it doesn't block the + # event loop. self._temporal_stream_backend.call("BufferWorkflowCalls", request) @@ -137,6 +161,10 @@ def _withdraw_run(info: temporalio.workflow.Info) -> str: @dataclass(frozen=True) class _WorkflowCall: + # Marks the item by attribute, which holds even if a sandbox imports a second copy of the + # class. + _temporal_stream_call = True + operation: str request: Any when_run_closes: bool = False @@ -219,6 +247,10 @@ def operation_token(self) -> str | None: class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor): + def __init__(self, plugin: RedisStreamPlugin) -> None: + super().__init__() + self._plugin = plugin + def workflow_interceptor_class( self, input: temporalio.worker.WorkflowInterceptorClassInput ) -> type[temporalio.worker.WorkflowInboundInterceptor] | None: @@ -227,7 +259,7 @@ def workflow_interceptor_class( def intercept_activity( self, next: temporalio.worker.ActivityInboundInterceptor ) -> temporalio.worker.ActivityInboundInterceptor: - return _ActivityInbound(next) + return _ActivityInbound(next, self._plugin) class _WorkflowInbound(temporalio.worker.WorkflowInboundInterceptor): @@ -242,19 +274,29 @@ async def execute_workflow( except temporalio.workflow.ContinueAsNewError: # The next run continues the run chain's Streams. raise - except asyncio.CancelledError: - _close_with_run( - failure=temporalio.exceptions.CancelledError("Workflow cancelled") - ) - raise except BaseException as err: - if _fails_run(err): - _close_with_run(failure=err, chain_continues=_chain_continues(err)) + if ( + isinstance(err, asyncio.CancelledError) + and temporalio.workflow.cancellation_reason() is not None + ): + # The run ends cancelled. + _close_with_run( + failure=temporalio.exceptions.CancelledError("Workflow cancelled") + ) + elif _fails_run(err): + _close_with_run( + failure=err, + chain_continues=chain_continues( + temporalio.workflow.info(), temporalio.workflow.now(), err + ), + ) raise # A cron run that completes still goes on with the chain. _close_with_run( result=result, - chain_continues=bool(temporalio.workflow.info().cron_schedule), + chain_continues=chain_continues( + temporalio.workflow.info(), temporalio.workflow.now(), None + ), ) return result @@ -264,18 +306,34 @@ def _fails_run(err: BaseException) -> bool: return temporalio.workflow._Runtime.current().workflow_is_failure_exception(err) -def _chain_continues(err: BaseException) -> bool: - # The server starts another run of the chain, which goes on writing the default Stream, - # so the runtime skips the run's closes. - info = temporalio.workflow.info() +def chain_continues( + info: temporalio.workflow.Info, now: datetime, failure: BaseException | None +) -> bool: + """Whether the server starts another run of the chain after this run ends. + + Then the next run goes on with the default Stream, so the runtime skips this run's closes. + A cron run always continues. A failed run continues when its retry policy has attempts and + time left for the failure. + """ if info.cron_schedule: return True policy = info.retry_policy - if policy is None or ( - policy.maximum_attempts != 0 and info.attempt >= policy.maximum_attempts - ): + if failure is None or policy is None: + return False + if policy.maximum_attempts != 0 and info.attempt >= policy.maximum_attempts: + return False + if _non_retryable(failure, policy.non_retryable_error_types): return False - return not _non_retryable(err, policy.non_retryable_error_types) + if info.execution_timeout is not None: + backoff = policy.initial_interval * ( + policy.backoff_coefficient ** (info.attempt - 1) + ) + if policy.maximum_interval is not None: + backoff = min(backoff, policy.maximum_interval) + expiration = info.workflow_start_time + info.execution_timeout + if now + backoff >= expiration: + return False + return True def _non_retryable(err: BaseException, types: Sequence[str] | None) -> bool: @@ -290,17 +348,35 @@ def _close_with_run( failure: BaseException | None = None, chain_continues: bool = False, ) -> None: - for stream in _workflow._state().streams.values(): - if not stream.auto_close or not stream._used or stream._closed: + info = temporalio.workflow.info() + default_id = _ids.default_stream_id(info.workflow_id, info.first_execution_run_id) + state = _workflow._state() + # The default Stream closes with the run even when this run wrote nothing, since an + # earlier run of the chain may have written it. A close of a Stream with no key does + # nothing. + ids = ( + [] + if (default := state.streams.get(default_id)) and default._closed + else [default_id] + ) + for stream in state.streams.values(): + if stream.ref.id == default_id or stream._closed: continue - stream._closed = True + if stream.auto_close and stream._used: + ids.append(stream.ref.id) + for stream in state.streams.values(): + if stream.ref.id in ids: + stream._closed = True + for stream_id in ids: temporalio.workflow.send_to_plugins( _WorkflowCall( "CloseStream", - models.CloseStreamRequest( - stream=models.StreamRef(id=stream.ref.id), - result=result, - failure=failure, + _convert_in_sandbox( + models.CloseStreamRequest( + stream=models.StreamRef(id=stream_id), + result=result, + failure=failure, + ) ), when_run_closes=True, chain_continues=chain_continues, @@ -310,9 +386,52 @@ def _close_with_run( ) +def _convert_in_sandbox(request: Any) -> Any: + # A value from Workflow code may be of a class that only the sandbox knows, so the + # Workflow's payload converter turns it into a Payload here. The hook outside runs only + # the codec. + converter = temporalio.workflow.payload_converter() + if isinstance(request, models.AppendRecordsRequest): + return dataclasses.replace( + request, + records=[ + dataclasses.replace( + record, + body=temporalio.common.RawValue(converter.to_payload(record.body)), + ) + for record in request.records or [] + ], + ) + if isinstance(request, models.CloseStreamRequest) and request.result is not None: + return dataclasses.replace( + request, + result=temporalio.common.RawValue(converter.to_payload(request.result)), + ) + return request + + class _WorkflowOutbound(temporalio.worker.WorkflowOutboundInterceptor): _temporal_stream_transport = True + def __init__(self, next: temporalio.worker.WorkflowOutboundInterceptor) -> None: + super().__init__(next) + # 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. + self._activity_calls = 0 + + def _activity_headers(self) -> dict[str, temporalio.api.common.v1.Payload]: + self._activity_calls += 1 + info = temporalio.workflow.info() + # Both ends use the default converter, so a custom user converter can't break the + # headers. + converter = temporalio.converter.PayloadConverter.default + ref = _ids.default_stream_id(info.workflow_id, info.first_execution_run_id) + return { + _REF_HEADER: converter.to_payload(ref), + _KEY_HEADER: converter.to_payload(str(self._activity_calls)), + } + async def start_system_nexus_operation( self, input: temporalio.worker.StartNexusOperationInput[Any, Any] ) -> temporalio.workflow.NexusOperationHandle[Any]: @@ -330,7 +449,7 @@ async def start_system_nexus_operation( else: raise _errors.StreamError(f"Workflow code can't call {operation}") temporalio.workflow.send_to_plugins( - _WorkflowCall(operation, input.input), + _WorkflowCall(operation, _convert_in_sandbox(input.input)), # The runtime already has the Records of a replayed task. A create or a close is # idempotent, so it goes again, for a run whose buffer an eviction dropped. during_replay=operation != "AppendRecords", @@ -340,34 +459,29 @@ async def start_system_nexus_operation( def start_activity( self, input: temporalio.worker.StartActivityInput ) -> temporalio.workflow.ActivityHandle[Any]: - input.headers = {**input.headers, **_activity_headers()} + input.headers = {**input.headers, **self._activity_headers()} return super().start_activity(input) def start_local_activity( self, input: temporalio.worker.StartLocalActivityInput ) -> temporalio.workflow.ActivityHandle[Any]: - input.headers = {**input.headers, **_activity_headers()} + input.headers = {**input.headers, **self._activity_headers()} return super().start_local_activity(input) -def _activity_headers() -> dict[str, temporalio.api.common.v1.Payload]: - info = temporalio.workflow.info() - # Both ends use the default converter, so a custom user converter can't break the headers. - converter = temporalio.converter.PayloadConverter.default - ref = _ids.default_stream_id(info.workflow_id, info.first_execution_run_id) - return { - _REF_HEADER: converter.to_payload(ref), - _KEY_HEADER: converter.to_payload(str(temporalio.workflow.uuid4())), - } - - class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor): + def __init__( + self, + next: temporalio.worker.ActivityInboundInterceptor, + plugin: RedisStreamPlugin, + ) -> None: + super().__init__(next) + self._plugin = plugin + async def execute_activity( self, input: temporalio.worker.ExecuteActivityInput ) -> Any: - streams = _activity_streams(input) - if streams is None: - return await super().execute_activity(input) + streams = _activity_streams(input, self._plugin) token = _activity.current.set(streams) try: if ( @@ -375,8 +489,8 @@ async def execute_activity( and temporalio.activity.info().attempt > 1 ): # The open of a new attempt withdraws the Records of the earlier ones, also - # when this attempt publishes nothing. - await streams.writer(streams.default_ref.id).open() + # when this attempt publishes nothing. It creates no Stream. + await streams.writer(streams.default_ref.id).withdraw_earlier_attempts() result = await super().execute_activity(input) except BaseException: await _close_quietly(streams) @@ -397,23 +511,10 @@ async def _close_quietly(streams: _activity.ActivityStreams) -> None: ) -def _plugin_backend(client: temporalio.client.Client) -> StreamBackend | None: - for plugin in client.config()["plugins"]: - if isinstance(plugin, RedisStreamPlugin): - return plugin._temporal_stream_backend - return None - - def _activity_streams( - input: temporalio.worker.ExecuteActivityInput, -) -> _activity.ActivityStreams | None: - try: - client = temporalio.activity.client() - except RuntimeError: - return None - backend = _plugin_backend(client) - if backend is None: - return None + input: temporalio.worker.ExecuteActivityInput, plugin: RedisStreamPlugin +) -> _activity.ActivityStreams: + # A sync Activity has no client, so the plugin gives the backend and the converter. info = temporalio.activity.info() converter = temporalio.converter.PayloadConverter.default ref = input.headers.get(_REF_HEADER) @@ -451,11 +552,12 @@ def _activity_streams( ) producer.epoch.CopyFrom(epoch) return _activity.ActivityStreams( - backend=backend, - converter=client.data_converter, + backend=plugin._temporal_stream_backend, + converter=plugin._converter, namespace=info.namespace, producer=producer, default_ref=StreamRef(converter.from_payload(ref, str)) if ref is not None else None, + loop=asyncio.get_running_loop(), ) diff --git a/tests/contrib/streams/test_chain.py b/tests/contrib/streams/test_chain.py new file mode 100644 index 0000000000..7e98918071 --- /dev/null +++ b/tests/contrib/streams/test_chain.py @@ -0,0 +1,68 @@ +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from typing import Any + +import pytest + +from temporalio.common import RetryPolicy +from temporalio.contrib.streams.redis import chain_continues +from temporalio.exceptions import ApplicationError + +START = datetime(2026, 10, 21, tzinfo=timezone.utc) + + +def info(**fields: Any) -> Any: + base: dict[str, Any] = { + "cron_schedule": None, + "retry_policy": None, + "attempt": 1, + "execution_timeout": None, + "workflow_start_time": START, + } + return SimpleNamespace(**{**base, **fields}) + + +RETRY = RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3) + + +@pytest.mark.parametrize( + "case, run, failure, expected", + [ + ("success", info(), None, False), + ("cron success", info(cron_schedule="* * * * *"), None, True), + ("cron failure", info(cron_schedule="* * * * *"), RuntimeError("x"), True), + ("no retry policy", info(), RuntimeError("x"), False), + ("attempts left", info(retry_policy=RETRY), RuntimeError("x"), True), + ("last attempt", info(retry_policy=RETRY, attempt=3), RuntimeError("x"), False), + ( + "unlimited attempts", + info(retry_policy=RetryPolicy(maximum_attempts=0), attempt=50), + RuntimeError("x"), + True, + ), + ( + "non-retryable error", + info(retry_policy=RETRY), + ApplicationError("x", non_retryable=True), + False, + ), + ( + "non-retryable type", + info( + retry_policy=RetryPolicy(non_retryable_error_types=["Bad"]), + ), + ApplicationError("x", type="Bad"), + False, + ), + ( + "no time left", + info(retry_policy=RETRY, execution_timeout=timedelta(seconds=10)), + RuntimeError("x"), + False, + ), + ], +) +def test_chain_continues(case: str, run: Any, failure: Any, expected: bool): + # "no time left" runs 9.5 s into a 10 s chain, so a 1 s backoff passes the expiration. + now = START + timedelta(seconds=9.5) + assert chain_continues(run, now, failure) == expected, case diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index 0924b1eb47..77a6646c2a 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -1,4 +1,6 @@ import asyncio +import concurrent.futures +import dataclasses import shutil import socket import subprocess @@ -15,6 +17,10 @@ from temporalio.contrib import streams from temporalio.contrib.streams.redis import RedisStreamPlugin from temporalio.worker import UnsandboxedWorkflowRunner +from temporalio.worker.workflow_sandbox import ( + SandboxedWorkflowRunner, + SandboxRestrictions, +) from tests.helpers import new_worker # Ports that other local services own. @@ -267,3 +273,187 @@ async def test_workflow_code_can_not_read(stream_client: Client): ) assert "can't read" in result await asyncio.sleep(0) + + +@activity.defn +async def noop() -> None: + pass + + +@workflow.defn +class RandomAroundActivity: + @workflow.run + async def run(self) -> bool: + before = workflow.random().getstate() + handle = workflow.start_activity( + noop, start_to_close_timeout=timedelta(seconds=30) + ) + same = workflow.random().getstate() == before + await handle + return same + + +async def test_plugin_leaves_the_workflow_random_alone(stream_client: Client): + async with new_worker( + stream_client, RandomAroundActivity, activities=[noop] + ) as worker: + assert await stream_client.execute_workflow( + RandomAroundActivity.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + + +retries: dict[str, int] = {} + + +@activity.defn +async def fails_once(key: str) -> None: + retries[key] = retries.get(key, 0) + 1 + if retries[key] == 1: + raise RuntimeError("first attempt fails") + + +@workflow.defn +class NoStreams: + @workflow.run + async def run(self, key: str) -> None: + await workflow.execute_activity( + fails_once, + key, + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=10)), + ) + + +async def test_workflow_without_streams_creates_none(stream_client: Client): + async with new_worker(stream_client, NoStreams, activities=[fails_once]) as w: + handle = await stream_client.start_workflow( + NoStreams.run, + str(uuid.uuid4()), + id=f"wf-{uuid.uuid4()}", + task_queue=w.task_queue, + ) + await handle.result() + # The retry withdrew nothing and the run's close found no Stream, so none exists. + with pytest.raises(streams.StreamNotFoundError): + await streams.for_handle(handle).get_offset() + + +@workflow.defn +class WritesThenContinues: + @workflow.run + async def run(self, write: bool) -> str: + if write: + streams.workflow_stream().topic("t", type=str).publish("first run") + workflow.continue_as_new(False) + return "second run" + + +async def test_last_run_closes_a_stream_it_did_not_write(stream_client: Client): + async with new_worker(stream_client, WritesThenContinues) as worker: + handle = await stream_client.start_workflow( + WritesThenContinues.run, + True, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == "second run" + stream = streams.for_handle(handle) + stream.topic("t", type=str) + items = await read_all(stream) + assert [i.data for i in items if isinstance(i, streams.Record)] == ["first run"] + closed = items[-1] + assert isinstance(closed, streams.Closed) and closed.result == "second run" + + +@workflow.defn +class CancelledBlock: + @workflow.run + async def run(self) -> None: + with streams.create_workflow_stream("turn") as turn: + turn.topic("t", type=str).publish("started") + await workflow.wait_condition(lambda: False) + + +async def test_cancelled_block_closes_with_a_failure(stream_client: Client): + async with new_worker(stream_client, CancelledBlock) as worker: + handle = await stream_client.start_workflow( + CancelledBlock.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + turn = streams.for_handle( + handle, name="turn", wait_for_create=timedelta(seconds=20) + ) + turn.topic("t", type=str) + first = await anext(aiter(turn.subscribe())) + assert isinstance(first, streams.Record) + await handle.cancel() + items = await read_all(turn) + closed = items[-1] + assert isinstance(closed, streams.Closed) + assert closed.failure is not None + + +@activity.defn +def sync_tokens(count: int) -> None: + tokens = streams.activity_stream().topic("tokens", type=int) + for n in range(count): + tokens.publish(n) + + +@workflow.defn +class SyncWriter: + @workflow.run + async def run(self) -> None: + await workflow.execute_activity( + sync_tokens, 3, start_to_close_timeout=timedelta(seconds=30) + ) + + +async def test_sync_activity_publishes(stream_client: Client): + async with new_worker( + stream_client, + SyncWriter, + activities=[sync_tokens], + activity_executor=concurrent.futures.ThreadPoolExecutor(2), + ) as worker: + handle = await stream_client.start_workflow( + SyncWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + stream = streams.for_handle(handle) + stream.topic("tokens", type=int) + items = await read_all(stream) + assert [i.data for i in items if isinstance(i, streams.Record)] == [0, 1, 2] + + +@workflow.defn +class PlainWriter: + @workflow.run + async def run(self) -> str: + streams.workflow_stream().topic("t", type=str).publish("plain") + with streams.create_workflow_stream("scoped") as scoped: + scoped.topic("t", type=str).publish("scoped") + return "done" + + +async def test_temporal_only_passthrough_keeps_records(stream_client: Client): + # Only the minimum modules and temporalio pass through. Values are plain, since a + # dataclass from a sandbox with its own dataclasses module can't convert outside it. + runner = SandboxedWorkflowRunner( + restrictions=dataclasses.replace( + SandboxRestrictions.default, + passthrough_modules=SandboxRestrictions.passthrough_modules_with_temporal, + ) + ) + async with new_worker(stream_client, PlainWriter, workflow_runner=runner) as worker: + handle = await stream_client.start_workflow( + PlainWriter.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + items = await read_all(streams.for_handle(handle)) + scoped = await read_all(streams.for_handle(handle, name="scoped")) + assert [i.data for i in items if isinstance(i, streams.Record)] == ["plain"] + assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" + assert [i.data for i in scoped if isinstance(i, streams.Record)] == ["scoped"] + assert isinstance(scoped[-1], streams.Closed) From 6a510f7c5946cdb1f509b69a58fd2fb79578a408 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 21:35:36 -0700 Subject: [PATCH 3/4] Moved the crate pin to the s5 close-tracker fix. The crate hardened how it tracks closes in flight, which the Worker shutdown drain depends on. --- temporalio/bridge/Cargo.lock | 2 +- temporalio/bridge/Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/temporalio/bridge/Cargo.lock b/temporalio/bridge/Cargo.lock index be25623919..f39ddca904 100644 --- a/temporalio/bridge/Cargo.lock +++ b/temporalio/bridge/Cargo.lock @@ -2879,7 +2879,7 @@ dependencies = [ [[package]] name = "temporalio-streams" version = "0.1.0" -source = "git+https://github.com/moetemp/sdk-streams?rev=5b366595abc08adf26358dd1a54afdf3fb3d4097#5b366595abc08adf26358dd1a54afdf3fb3d4097" +source = "git+https://github.com/moetemp/sdk-streams?rev=0981f62b1454e6501e3173d2115d94cb25eb9036#0981f62b1454e6501e3173d2115d94cb25eb9036" dependencies = [ "async-trait", "aws-lc-rs", diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml index 8a10eb2940..29392c3f42 100644 --- a/temporalio/bridge/Cargo.toml +++ b/temporalio/bridge/Cargo.toml @@ -94,7 +94,7 @@ url = "2.2" # crates, so the bridge sees one type of each. [dependencies.temporalio-streams] git = "https://github.com/moetemp/sdk-streams" -rev = "5b366595abc08adf26358dd1a54afdf3fb3d4097" +rev = "0981f62b1454e6501e3173d2115d94cb25eb9036" default-features = false optional = true From 63de9a4da128b9c3c1ba85b4c9e4330e9b561caf Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 22:16:40 -0700 Subject: [PATCH 4/4] Left GeneratorExit alone in the plugin's Workflow interceptor. Python closes an evicted run's coroutine outside the Workflow's loop. The interceptor then called the runtime, and each eviction printed an Exception ignored traceback. --- temporalio/contrib/streams/redis.py | 4 ++ tests/contrib/streams/test_redis.py | 87 ++++++++++++++++++++++++++++- 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py index 447c0dcc5d..4e4acddda6 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -275,6 +275,10 @@ async def execute_workflow( # The next run continues the run chain's Streams. raise except BaseException as err: + if not isinstance(err, (Exception, asyncio.CancelledError)): + # A GeneratorExit or another eviction exception closes the coroutine outside + # the Workflow's loop. The run did not end, so there is nothing to close. + raise if ( isinstance(err, asyncio.CancelledError) and temporalio.workflow.cancellation_reason() is not None diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index 77a6646c2a..953c22f1eb 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -1,13 +1,15 @@ import asyncio import concurrent.futures import dataclasses +import gc import shutil import socket import subprocess import uuid -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Generator, Iterator from dataclasses import dataclass from datetime import timedelta +from typing import Any, cast import pytest @@ -15,8 +17,12 @@ from temporalio.client import Client from temporalio.common import RetryPolicy from temporalio.contrib import streams -from temporalio.contrib.streams.redis import RedisStreamPlugin -from temporalio.worker import UnsandboxedWorkflowRunner +from temporalio.contrib.streams.redis import RedisStreamPlugin, _WorkflowInbound +from temporalio.worker import ( + ExecuteWorkflowInput, + UnsandboxedWorkflowRunner, + WorkflowInboundInterceptor, +) from temporalio.worker.workflow_sandbox import ( SandboxedWorkflowRunner, SandboxRestrictions, @@ -457,3 +463,78 @@ async def test_temporal_only_passthrough_keeps_records(stream_client: Client): assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" assert [i.data for i in scoped if isinstance(i, streams.Record)] == ["scoped"] assert isinstance(scoped[-1], streams.Closed) + + +@workflow.defn +class WaitsForever: + @workflow.run + async def run(self) -> None: + streams.workflow_stream().topic("t", type=str).publish("waiting") + await workflow.wait_condition(lambda: False) + + +class _Suspends: + def __await__(self) -> Generator[None, None, None]: + yield + + +class _NeverEnds(WorkflowInboundInterceptor): + def __init__(self) -> None: # pyright: ignore[reportMissingSuperCall] + pass + + async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any: + await _Suspends() + + +def test_closing_an_evicted_run_touches_no_runtime(): + # Python closes the coroutine of an evicted run with GeneratorExit, outside the + # 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.send(None) + run.close() + + +@pytest.mark.parametrize("cache, safe", [(0, True), (10, True), (10, False)]) +async def test_eviction_prints_nothing( + stream_client: Client, capfd: pytest.CaptureFixture[str], cache: int, safe: bool +): + # With no cache, each Workflow Task evicts the run. With a cache, the Worker's shutdown + # evicts it. Either way the run's coroutine closes. + async with new_worker( + stream_client, + WaitsForever, + max_cached_workflows=cache, + disable_safe_workflow_eviction=not safe, + ) as w: + handle = await stream_client.start_workflow( + WaitsForever.run, id=f"wf-{uuid.uuid4()}", task_queue=w.task_queue + ) + stream = streams.for_handle(handle, wait_for_create=timedelta(seconds=20)) + assert isinstance(await anext(aiter(stream.subscribe())), streams.Record) + await handle.terminate() + gc.collect() + captured = capfd.readouterr() + assert "Exception ignored" not in captured.err + captured.out + + +@workflow.defn +class CreatesEmpty: + @workflow.run + async def run(self) -> None: + streams.create_workflow_stream("empty") + await workflow.sleep(timedelta(milliseconds=10)) + + +async def test_get_offset_of_a_created_empty_stream(stream_client: Client): + async with new_worker(stream_client, CreatesEmpty) as worker: + handle = await stream_client.start_workflow( + CreatesEmpty.run, id=f"wf-{uuid.uuid4()}", task_queue=worker.task_queue + ) + await handle.result() + # A created Stream with no Record gives the start offset, so a reader can begin there. + assert await streams.for_handle(handle, name="empty").get_offset() == 0 + # A Stream that nobody created is still an error. + with pytest.raises(streams.StreamNotFoundError): + await streams.for_handle(handle, name="missing").get_offset()