From 868348d18e0797e7b88c8bb361d7ee0de4716aae Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 12:20:13 -0700 Subject: [PATCH] Added the Redis and memory stream store plugins. --- temporalio/contrib/streams/__init__.py | 19 +- temporalio/contrib/streams/_plugin.py | 201 +++++++++++++++++++++ temporalio/contrib/streams/memory.py | 24 +++ temporalio/contrib/streams/redis.py | 79 ++++++++ tests/contrib/streams/_support.py | 48 +++++ tests/contrib/streams/test_store_plugin.py | 198 ++++++++++++++++++++ 6 files changed, 566 insertions(+), 3 deletions(-) create mode 100644 temporalio/contrib/streams/_plugin.py create mode 100644 temporalio/contrib/streams/memory.py create mode 100644 temporalio/contrib/streams/redis.py create mode 100644 tests/contrib/streams/_support.py create mode 100644 tests/contrib/streams/test_store_plugin.py diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 4363d17636..d672e8a52e 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -10,19 +10,30 @@ The contract: -1. **A new attempt supersedes the old one.** When a reader sees the first +1. **Producers write with an identity.** A producer has an id, an attempt + and a sequence that starts at one for each attempt. A retry of the + newest batch with the same content is written once and returns the + original position. A retry with different content, or with a lower + sequence, is refused with :class:`StreamProducerError`. Content is + compared over the payloads the converter produced, before the codec, so + a codec with a fresh nonce per call still deduplicates a retry. +2. **A new attempt supersedes the old one.** When a reader sees the first record of a producer's newer attempt, it yields a :attr:`RecordKind.SUPERSEDED` record first. No store holds that record, so every store reports a retry the same way. -2. **A cursor belongs to one stream.** Hand it back to resume strictly after +3. **A cursor belongs to one stream.** Hand it back to resume strictly after the record it names. A cursor from another store or another stream is refused with :class:`StreamCursorError`, and one whose record retention dropped with :class:`StreamExpiredError`. A read with no cursor starts at :data:`BEGINNING` or :data:`END`. -3. **Topics are defined once.** :func:`topic` defines a topic with the type +4. **Topics are defined once.** :func:`topic` defines a topic with the type its records decode to, and every party shares that definition. A call that names no topic addresses :data:`DEFAULT_TOPIC`. +A store is configured as a plugin: +``Client.connect(..., plugins=[RedisStreams(...)])`` registers it on the +client and on every Worker built from that client. + The stored record is ``temporal.sdk.streams.v1.StreamRecord``, in :mod:`temporalio.bridge.proto.streams.v1`, with the user's value in ``body`` as an ordinary payload. @@ -42,6 +53,7 @@ StreamStorageError, StreamUnsupportedError, ) +from temporalio.contrib.streams._plugin import StreamStorePlugin from temporalio.contrib.streams._record import ( Cursor, RecordKind, @@ -70,6 +82,7 @@ "StreamRef", "StreamRefusedError", "StreamStorageError", + "StreamStorePlugin", "StreamTopic", "StreamUnsupportedError", "Supersession", diff --git a/temporalio/contrib/streams/_plugin.py b/temporalio/contrib/streams/_plugin.py new file mode 100644 index 0000000000..737f0de1c4 --- /dev/null +++ b/temporalio/contrib/streams/_plugin.py @@ -0,0 +1,201 @@ +"""The plugin a store configuration is, and how each context finds the store. + +A configuration class is registered once, on the client: +``Client.connect(plugins=[RedisStreams(...)])``. Connecting the client asks +Core for the store it names. Core serves that one store to the client's +handles and to the Activities of every Worker built from the client, so an +Activity's append and an outside read meet in the same place. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any, TypeVar, cast + +import temporalio.bridge.client +from temporalio.bridge.proto.streams import ( + DeleteOwnerRequest, + StreamOwnerKind, + StreamStoreConfig, +) +from temporalio.bridge.streams import StreamCallFailure, StreamStore +from temporalio.bridge.streams_generated import StreamService +from temporalio.contrib.streams._errors import error_from_failure +from temporalio.plugin import SimplePlugin + +if TYPE_CHECKING: + from temporalio.client import Client, ClientConfig + from temporalio.service import ConnectConfig, ServiceClient + from temporalio.worker import WorkerConfig + +__all__ = ["StreamStorePlugin"] + +T = TypeVar("T") + + +class StreamStorePlugin(SimplePlugin): + """The base class of the store configuration classes. + + Use :class:`temporalio.contrib.streams.redis.RedisStreams` or + :class:`temporalio.contrib.streams.memory.MemoryStreams`. One instance + is one store, on one client connection, since Core asks about each + stream's owner over the connection that connected the store. A second + connection, such as another namespace's endpoint, needs its own instance. + Registering a second, different store on the same client raises + ``ValueError``, because a handle could otherwise pick either. + """ + + def __init__(self, name: str, config: StreamStoreConfig) -> None: + """Name the plugin and the store Core connects to.""" + super().__init__(name) + self._config = config + self._service: StreamService | None = None + self._store: StreamStore | None = None + self._connecting: asyncio.Lock | None = None + self._service_client: ServiceClient | None = None + # Whether Client.connect connected the store, which a Worker needs. + self._connected_by_connect = False + + def configure_client(self, config: ClientConfig) -> ClientConfig: + """Register this store on the client. + + Raises: + ValueError: Another store is already registered on the client. + """ + for plugin in config.get("plugins", []): + if isinstance(plugin, StreamStorePlugin) and plugin is not self: + raise ValueError(_two_stores(plugin, self)) + return super().configure_client(config) + + async def connect_service_client( + self, + config: ConnectConfig, + next: Callable[[ConnectConfig], Awaitable[ServiceClient]], + ) -> ServiceClient: + """Connect the client, then the store, so Workers built from it have one. + + A lazy client connects here too, since Core needs its connection to + ask about a stream's owner. + + Raises: + temporalio.contrib.streams.StreamError: Core could not reach the + store. + """ + service_client = await next(config) + await self._connected(service_client) + self._connected_by_connect = True + return service_client + + def configure_worker(self, config: WorkerConfig) -> WorkerConfig: + """Check that the Worker's client carries this store. + + Raises: + ValueError: The plugin is not registered on the Worker's client, + or that client was not made with ``Client.connect``. A + client made otherwise is refused even after one of its handles + connected the store. + """ + config = super().configure_worker(config) + client = config["client"] # type:ignore[reportTypedDictNotRequiredAccess] + if ( + self not in client.config()["plugins"] + or not self._connected_by_connect + or client.service_client is not self._service_client + ): + raise ValueError( + f"stream store {self.name()!r} must be registered on the Worker's " + "client, as Client.connect(..., plugins=[store])" + ) + return config + + async def delete_workflow_streams(self, namespace: str, workflow_id: str) -> int: + """Delete every stream of ``workflow_id`` in ``namespace``. + + An admin helper for cleanup and compliance requests. Deleting a + Workflow does not delete its streams. Their retention bounds what is + left, and this removes it at once, across every run chain of the + Workflow id. It asks Temporal nothing, so it can run after the + Workflow is gone. On Redis it walks the keyspace, so run it rarely, + not on a hot path. + + Delete only after the chain closed and its producers stopped. A late + append writes a fresh stream with no dedupe state and no close flag. + + Returns: + How many store keys were deleted. + + Raises: + ValueError: No client connected this store yet. + temporalio.contrib.streams.StreamStorageError: The store failed; + some keys may be deleted. + """ + if self._service is None: + raise ValueError(f"no client has connected stream store {self.name()!r}") + request = DeleteOwnerRequest( + namespace=namespace, + owner_kind=StreamOwnerKind.STREAM_OWNER_KIND_WORKFLOW, + workflow_id=workflow_id, + ) + return (await call(self._service.delete_owner(request))).deleted + + async def _service_for(self, client: Client) -> StreamService: + """The stream service, connecting the store if no client did yet. + + Raises: + ValueError: Another client connection already connected the store. + """ + await self._connected(client.service_client) + assert self._service is not None + return self._service + + async def _connected(self, service_client: ServiceClient) -> None: + if self._service_client is service_client: + return + if self._connecting is None: + self._connecting = asyncio.Lock() + async with self._connecting: + if self._service_client is service_client: + return + if self._service_client is not None: + raise ValueError( + f"stream store {self.name()!r} serves one client connection, " + "and another one already connected it; give each connection " + "a store instance of its own" + ) + bridge: temporalio.bridge.client.Client = await cast( + Any, service_client + )._connected_client() + self._store = await call(StreamStore.connect(bridge, self._config)) + self._service = StreamService(self._store) + self._service_client = service_client + + +async def call(pending: Awaitable[T]) -> T: + """Await a stream call, raising Core's failure as the error a caller catches.""" + try: + return await pending + except StreamCallFailure as failure: + raise error_from_failure(failure.failure) from None + + +def _two_stores(held: StreamStorePlugin, new: StreamStorePlugin) -> str: + return ( + f"stream store {held.name()!r} is already registered, so {new.name()!r} " + "cannot be; register one store per client" + ) + + +def store_for_client(client: Client) -> StreamStorePlugin: + """The store registered on ``client``. + + Raises: + ValueError: No store is registered on the client. + """ + for plugin in client.config()["plugins"]: + if isinstance(plugin, StreamStorePlugin): + return plugin + raise ValueError( + "no stream store is registered on this client; pass one as a plugin, for " + "example Client.connect(..., plugins=[RedisStreams(...)])" + ) diff --git a/temporalio/contrib/streams/memory.py b/temporalio/contrib/streams/memory.py new file mode 100644 index 0000000000..198a518bfd --- /dev/null +++ b/temporalio/contrib/streams/memory.py @@ -0,0 +1,24 @@ +"""Streams in this process's memory, for tests and local development.""" + +from __future__ import annotations + +from temporalio.bridge.proto.streams import MemoryStoreConfig, StreamStoreConfig +from temporalio.contrib.streams._plugin import StreamStorePlugin + +__all__ = ["MemoryStreams"] + + +class MemoryStreams(StreamStorePlugin): + """A store in this process's memory. + + Core holds it, so the client's handles and every Worker built from the + client share it. Nothing survives the process, and another process sees + none of it. + """ + + def __init__(self) -> None: + """Create the store's configuration.""" + super().__init__( + "temporalio.contrib.streams.MemoryStreams", + StreamStoreConfig(memory=MemoryStoreConfig()), + ) diff --git a/temporalio/contrib/streams/redis.py b/temporalio/contrib/streams/redis.py new file mode 100644 index 0000000000..7d0be4409a --- /dev/null +++ b/temporalio/contrib/streams/redis.py @@ -0,0 +1,79 @@ +"""Streams in the application's Redis.""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import timedelta + +from temporalio.bridge.proto.streams import RedisStoreConfig, StreamStoreConfig +from temporalio.contrib.streams._plugin import StreamStorePlugin + +__all__ = ["RedisStreams"] + + +class RedisStreams(StreamStorePlugin): + """A store in a Redis or a Redis Cluster that Core connects to. + + One instance per Redis deployment. Pass it to the client as a plugin, + and open handles from the client with + :func:`temporalio.contrib.streams.get_stream_handle`. Redis 7.0 or later + is required, with ``maxmemory-policy noeviction``. + """ + + def __init__( + self, + url: str | Sequence[str], + *, + cluster: bool = False, + key_prefix: str = "temporal-streams", + retention: timedelta = timedelta(days=7), + blocking_reads_per_node: int | None = None, + response_timeout: timedelta | None = None, + ) -> None: + """Create the store's configuration. + + Args: + url: A ``redis://`` or ``rediss://`` URL, or the seed nodes of a + cluster. Credentials and TLS settings come from the URL. In + this release TLS verifies the server against the system's + root certificates only, with no custom CA and no client + certificate. + cluster: Whether ``url`` names a Redis Cluster. Every key of one + Workflow's streams shares a hash tag, so a cluster serves + them from one slot. + key_prefix: Prepended to every key, so streams can share a Redis + with other data and an ACL can scope them. + retention: How long a stream keeps a record, and how long after + its last write the stream itself lives. + blocking_reads_per_node: How many reads may wait on one node at + once, each on a connection of its own. Core's default when + unset. + response_timeout: How long a short call waits for its reply. + Core's default when unset. + + Raises: + ValueError: No URL was given, ``key_prefix`` is empty, or + ``retention`` is shorter than a millisecond. + """ + urls = [url] if isinstance(url, str) else list(url) + if not urls: + raise ValueError("RedisStreams needs a URL") + if not key_prefix: + raise ValueError("key_prefix must not be empty") + if retention < timedelta(milliseconds=1): + raise ValueError( + f"retention must be at least a millisecond, got {retention}" + ) + config = RedisStoreConfig( + urls=urls, + cluster=cluster, + key_prefix=key_prefix, + blocking_reads_per_node=blocking_reads_per_node or 0, + ) + config.retention.FromTimedelta(retention) + if response_timeout is not None: + config.response_timeout.FromTimedelta(response_timeout) + super().__init__( + "temporalio.contrib.streams.RedisStreams", + StreamStoreConfig(redis=config), + ) diff --git a/tests/contrib/streams/_support.py b/tests/contrib/streams/_support.py new file mode 100644 index 0000000000..de4da231a5 --- /dev/null +++ b/tests/contrib/streams/_support.py @@ -0,0 +1,48 @@ +"""What the stream tests share.""" + +from __future__ import annotations + +import asyncio +import uuid +from typing import Any + +from temporalio.client import Client +from temporalio.contrib.streams import StreamStorePlugin +from temporalio.converter import DataConverter + + +async def connect_with( + client: Client, + store: StreamStorePlugin | None, + converter: DataConverter | None = None, +) -> Client: + """A client of the test server that carries ``store``. + + It connects, since connecting is what hands the store to Workers built + from it. + """ + return await Client.connect( + client.service_client.config.target_host, + namespace=client.namespace, + data_converter=converter or client.data_converter, + plugins=[store] if store is not None else [], + ) + + +def new_workflow_id() -> str: + return f"streams-{uuid.uuid4().hex}" + + +async def read_all(records: Any, timeout: float = 10.0) -> list: + """Read until the stream ends, which is when its owner closes.""" + out = [] + + async def pull() -> None: + async for record in records: + out.append(record) + + try: + await asyncio.wait_for(pull(), timeout) + finally: + await records.aclose() + return out diff --git a/tests/contrib/streams/test_store_plugin.py b/tests/contrib/streams/test_store_plugin.py new file mode 100644 index 0000000000..039527c1c3 --- /dev/null +++ b/tests/contrib/streams/test_store_plugin.py @@ -0,0 +1,198 @@ +"""A store configuration registers as a client plugin, and Core connects to it.""" + +from __future__ import annotations + +import hashlib +import uuid +from datetime import timedelta + +import pytest + +from temporalio import activity +from temporalio.api.common.v1 import Payload +from temporalio.bridge.proto.streams import ( + AppendRecord, + AppendRequest, + LatestRequest, + NamedProducer, + StreamAddress, + StreamOwnerKind, +) +from temporalio.bridge.proto.streams.v1 import StreamRecordKind +from temporalio.client import Client +from temporalio.contrib.streams import StreamError, StreamNotFoundError +from temporalio.contrib.streams._plugin import call +from temporalio.contrib.streams.memory import MemoryStreams +from temporalio.contrib.streams.redis import RedisStreams +from temporalio.worker import Worker +from tests.contrib.streams._support import connect_with + + +async def test_one_store_per_client(client: Client): + first, second = MemoryStreams(), MemoryStreams() + config = client.config() + config["plugins"] = [first, second] + with pytest.raises(ValueError, match="is already registered"): + Client(**config) + # The same instance twice is one store, not two. + config["plugins"] = [first, first] + Client(**config) + + +async def test_a_worker_takes_the_store_of_its_client(client: Client): + store = MemoryStreams() + streams_client = await connect_with(client, store) + Worker(streams_client, task_queue=str(uuid.uuid4()), activities=[_noop]) + # The store on the client and the Worker configures the Worker twice. + with pytest.warns(UserWarning, match="same plugin type"): + Worker( + streams_client, + task_queue=str(uuid.uuid4()), + activities=[_noop], + plugins=[store], + ) + + +async def test_a_store_on_the_worker_alone_is_refused(client: Client): + with pytest.raises(ValueError, match="registered on the Worker's client"): + Worker(client, task_queue="tq", activities=[_noop], plugins=[MemoryStreams()]) + # A client made without connecting has not connected its store either. + config = client.config() + config["plugins"] = [MemoryStreams()] + with pytest.raises(ValueError, match="registered on the Worker's client"): + Worker(Client(**config), task_queue="tq", activities=[_noop]) + + +async def test_a_client_made_without_connecting_connects_its_store_on_use( + client: Client, +): + store = MemoryStreams() + config = client.config() + config["plugins"] = [store] + service = await store._service_for(Client(**config)) + # An owner that does not exist is Core's answer, so the call reached the store. + stream = StreamAddress( + namespace=client.namespace, + owner_kind=StreamOwnerKind.STREAM_OWNER_KIND_WORKFLOW, + workflow_id=f"no-such-owner-{uuid.uuid4()}", + topic="out", + ) + with pytest.raises(StreamNotFoundError): + await call(service.latest(LatestRequest(stream=stream))) + + +async def test_a_store_core_cannot_reach_fails_the_connect(client: Client): + with pytest.raises(StreamError): + await connect_with( + client, + RedisStreams("redis://127.0.0.1:1", response_timeout=timedelta(seconds=1)), + ) + + +def test_redis_settings_go_into_the_store_config(): + store = RedisStreams( + ["redis://a:6379", "redis://b:6379"], + cluster=True, + key_prefix="app", + retention=timedelta(hours=2), + blocking_reads_per_node=4, + response_timeout=timedelta(seconds=3), + ) + config = store._config.redis + assert list(config.urls) == ["redis://a:6379", "redis://b:6379"] + assert config.cluster + assert config.key_prefix == "app" + assert config.retention.ToTimedelta() == timedelta(hours=2) + assert config.blocking_reads_per_node == 4 + assert config.response_timeout.ToTimedelta() == timedelta(seconds=3) + # Unset settings leave Core's defaults. + defaults = RedisStreams("redis://a:6379")._config.redis + assert defaults.blocking_reads_per_node == 0 + assert not defaults.HasField("response_timeout") + + +def test_redis_setting_mistakes_are_value_errors(): + with pytest.raises(ValueError, match="needs a URL"): + RedisStreams([]) + with pytest.raises(ValueError, match="key_prefix"): + RedisStreams("redis://a", key_prefix="") + with pytest.raises(ValueError, match="retention"): + RedisStreams("redis://a", retention=timedelta(microseconds=10)) + + +async def test_deleting_removes_what_a_workflow_stream_holds(client: Client): + store = MemoryStreams() + with pytest.raises(ValueError, match="no client has connected"): + await store.delete_workflow_streams(client.namespace, "wf") + streams_client = await connect_with(client, store) + service = await store._service_for(streams_client) + workflow_id = f"stream-owner-{uuid.uuid4()}" + await streams_client.start_workflow( + "StreamOwner", id=workflow_id, task_queue=f"unpolled-{uuid.uuid4()}" + ) + stream = address(client, workflow_id) + await call( + service.append( + AppendRequest( + stream=stream, + named=NamedProducer(producer_id="p", attempt=1), + sequence=1, + records=[ + AppendRecord( + kind=StreamRecordKind.STREAM_RECORD_KIND_DATA, + body=Payload(data=b"a"), + content_hash=hashlib.sha256(b"a").digest(), + ) + ], + digest=hashlib.sha256(b"batch").digest(), + ) + ) + ) + assert (await call(service.latest(LatestRequest(stream=stream)))).cursor + + assert await store.delete_workflow_streams(client.namespace, workflow_id) > 0 + assert (await call(service.latest(LatestRequest(stream=stream)))).cursor == "" + + +async def test_one_store_serves_one_client_connection(client: Client): + store = MemoryStreams() + first = await connect_with(client, store) + # Core asks about owners over the connection that connected the store, so a + # second connection would get the first one's answers. + with pytest.raises(ValueError, match="one client connection"): + await connect_with(client, store) + config = client.config() + config["plugins"] = [store] + with pytest.raises(ValueError, match="one client connection"): + await store._service_for(Client(**config)) + # A client copied from the connected one shares its connection. + copied = first.config() + assert await store._service_for(Client(**copied)) + + +async def test_a_worker_needs_a_client_made_with_connect_after_a_lazy_connect( + client: Client, +): + store = MemoryStreams() + config = client.config() + config["plugins"] = [store] + lazy = Client(**config) + # A handle call connects the store, but the client still wasn't connected + # through the plugin, so Workers built from it are refused all the same. + await store._service_for(lazy) + with pytest.raises(ValueError, match="registered on the Worker's client"): + Worker(lazy, task_queue=str(uuid.uuid4()), activities=[_noop]) + + +@activity.defn +async def _noop() -> None: + pass + + +def address(client: Client, workflow_id: str) -> StreamAddress: + return StreamAddress( + namespace=client.namespace, + owner_kind=StreamOwnerKind.STREAM_OWNER_KIND_WORKFLOW, + workflow_id=workflow_id, + topic="out", + )