diff --git a/temporalio/bridge/src/extension.rs b/temporalio/bridge/src/extension.rs index e5ba9cec7..d52b94ffb 100644 --- a/temporalio/bridge/src/extension.rs +++ b/temporalio/bridge/src/extension.rs @@ -82,6 +82,12 @@ pub fn new_worker_extension( temporalio_streams::RedisStreamExtension::from_config_bytes(&config, meter) .map_err(|err| PyValueError::new_err(err.to_string()))? } + #[cfg(feature = "streams")] + temporalio_streams::MEMORY_EXTENSION_NAME => { + let meter = runtime.core.telemetry().get_temporal_metric_meter(); + temporalio_streams::MemoryStreamExtension::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 index b1743fcac..06f35d7b8 100644 --- a/temporalio/contrib/streams/redis.py +++ b/temporalio/contrib/streams/redis.py @@ -45,11 +45,8 @@ _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. +class StreamPlugin(SimplePlugin): + """Base of the plugins that carry Stream calls to a temporalio-streams extension. .. warning:: This API is experimental. @@ -57,54 +54,18 @@ class RedisStreamPlugin(SimplePlugin): 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, + name: str, + extension_name: str, + config: ext_pb.RedisStreamConfig, + runtime: temporalio.runtime.Runtime | 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) + """Create the plugin over the bridge's extension with ``extension_name``.""" self._extension = temporalio.worker.WorkerExtension( - EXTENSION_NAME, config.SerializeToString(), runtime=runtime + 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(self)] - ) + super().__init__(name, interceptors=[_StreamInterceptor(self)]) def configure_worker( self, config: temporalio.worker.WorkerConfig @@ -190,6 +151,63 @@ async def on_workflow_activation_complete( self._temporal_stream_backend.call("BufferWorkflowCalls", request) +class RedisStreamPlugin(StreamPlugin): + """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) + super().__init__( + "temporalio.RedisStreamPlugin", EXTENSION_NAME, config, runtime + ) + + 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: @@ -285,7 +303,7 @@ def operation_token(self) -> str | None: class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor): - def __init__(self, plugin: RedisStreamPlugin) -> None: + def __init__(self, plugin: StreamPlugin) -> None: super().__init__() self._plugin = plugin @@ -515,7 +533,7 @@ class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor): def __init__( self, next: temporalio.worker.ActivityInboundInterceptor, - plugin: RedisStreamPlugin, + plugin: StreamPlugin, ) -> None: super().__init__(next) self._plugin = plugin @@ -554,7 +572,7 @@ async def _close_quietly(streams: _activity.ActivityStreams) -> None: def _activity_streams( - input: temporalio.worker.ExecuteActivityInput, plugin: RedisStreamPlugin + input: temporalio.worker.ExecuteActivityInput, plugin: StreamPlugin ) -> _activity.ActivityStreams: # A sync Activity has no client, so the plugin gives the backend and the converter. info = temporalio.activity.info() diff --git a/temporalio/contrib/streams/testing.py b/temporalio/contrib/streams/testing.py new file mode 100644 index 000000000..391fb1363 --- /dev/null +++ b/temporalio/contrib/streams/testing.py @@ -0,0 +1,50 @@ +"""A Stream plugin for tests, with no Redis. + +.. warning:: + This module is experimental. +""" + +from __future__ import annotations + +from datetime import timedelta + +import temporalio.runtime +from temporalio.contrib.streams.proto import v1 as ext_pb +from temporalio.contrib.streams.redis import StreamPlugin + +MEMORY_EXTENSION_NAME = "temporalio-streams-memory" + + +class MemoryStreamPlugin(StreamPlugin): + """Keeps Streams in this process, with the same rules as on Redis, for tests. + + Put it on the client, as with ``RedisStreamPlugin``. Every Worker and client that shares one + plugin shares its Streams, and they go when the process ends. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + *, + retention: timedelta | None = None, + safety_ttl: timedelta | None = None, + runtime: temporalio.runtime.Runtime | None = None, + ) -> None: + """Create the plugin. + + Args: + 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. + runtime: The runtime of the client. Defaults to the default runtime. + """ + config = ext_pb.RedisStreamConfig() + if retention is not None: + config.retention.FromTimedelta(retention) + if safety_ttl is not None: + config.safety_ttl.FromTimedelta(safety_ttl) + super().__init__( + "temporalio.MemoryStreamPlugin", MEMORY_EXTENSION_NAME, config, runtime + ) diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index 813436c32..4b351f30a 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -18,7 +18,12 @@ from temporalio.client import Client from temporalio.common import RawValue, RetryPolicy from temporalio.contrib import streams -from temporalio.contrib.streams.redis import RedisStreamPlugin, _WorkflowInbound +from temporalio.contrib.streams.redis import ( + RedisStreamPlugin, + StreamPlugin, + _WorkflowInbound, +) +from temporalio.contrib.streams.testing import MemoryStreamPlugin from temporalio.converter import DataConverter, PayloadCodec from temporalio.exceptions import ApplicationError from temporalio.worker import ( @@ -71,10 +76,20 @@ def redis_url() -> Iterator[str]: process.wait() -@pytest.fixture -async def stream_client(client: Client, redis_url: str) -> AsyncIterator[Client]: +@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"] = [RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}")] + config["plugins"] = [plugin] yield Client(**config) @@ -665,9 +680,7 @@ async def test_plugin_replays_a_history_made_without_it( replayer = Replayer( workflows=[RandomAroundActivity], plugins=[ - p - for p in stream_client.config()["plugins"] - if isinstance(p, RedisStreamPlugin) + p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin) ], ) await replayer.replay_workflow(history) @@ -697,7 +710,7 @@ async def test_worker_shutdown_drains_closes(stream_client: Client): async def test_plugin_shutdown_flushes_client_writers(stream_client: Client): [plugin] = [ - p for p in stream_client.config()["plugins"] if isinstance(p, RedisStreamPlugin) + p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin) ] stream = await streams.create_stream(stream_client, f"shutdown-{uuid.uuid4()}") stream.topic("t", type=str).publish("unflushed")