From 1ff36562bc7a643b129f781157129262a7d059f6 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 17:50:24 -0700 Subject: [PATCH 1/9] Added the Stream Stores and the hooks of the generated Streams API. A client takes its Stream Stores in a `StreamConfig`, and each call of the generated Streams API goes to one of them. The Stores run Core's Stream runtime through the bridge. --- .github/actions/install-redis/action.yml | 25 + .github/workflows/ci.yml | 5 + changelog/added/chattering-ledger-otter.md | 1 + pyproject.toml | 17 +- scripts/gen_nexus_system_api.py | 80 +- scripts/gen_payload_visitor.py | 34 +- temporalio/bridge/Cargo.lock | 212 +++++- temporalio/bridge/Cargo.toml | 20 +- temporalio/bridge/sdk-core | 2 +- temporalio/bridge/src/lib.rs | 67 ++ temporalio/bridge/src/streams.rs | 130 ++++ temporalio/bridge/streams.py | 104 +++ temporalio/client/__init__.py | 12 + temporalio/client/_client.py | 24 +- temporalio/client/_impl.py | 51 ++ temporalio/client/_stream_appender.py | 193 +++++ temporalio/client/_stream_store.py | 554 ++++++++++++++ temporalio/contrib/redis/__init__.py | 1 + temporalio/contrib/redis/streams.py | 116 +++ temporalio/exceptions.py | 109 +++ temporalio/nexus/system/__init__.py | 38 +- temporalio/nexus/system/_support.py | 153 +++- temporalio/testing/__init__.py | 2 + temporalio/testing/_stream_store.py | 57 ++ temporalio/testing/_workflow.py | 12 + temporalio/worker/_interceptor.py | 11 +- temporalio/worker/_replayer.py | 12 + temporalio/worker/_stream_store.py | 121 ++++ temporalio/worker/_worker.py | 27 +- temporalio/worker/_workflow.py | 4 + temporalio/worker/_workflow_instance.py | 113 ++- tests/contrib/redis/__init__.py | 0 tests/contrib/redis/test_streams.py | 243 +++++++ tests/helpers/streams.py | 186 +++++ tests/nexus/test_stream_appender.py | 159 ++++ tests/nexus/test_stream_store.py | 630 ++++++++++++++++ tests/nexus/test_stream_store_workflow.py | 806 +++++++++++++++++++++ 37 files changed, 4247 insertions(+), 84 deletions(-) create mode 100644 .github/actions/install-redis/action.yml create mode 100644 changelog/added/chattering-ledger-otter.md create mode 100644 temporalio/bridge/src/streams.rs create mode 100644 temporalio/bridge/streams.py create mode 100644 temporalio/client/_stream_appender.py create mode 100644 temporalio/client/_stream_store.py create mode 100644 temporalio/contrib/redis/__init__.py create mode 100644 temporalio/contrib/redis/streams.py create mode 100644 temporalio/testing/_stream_store.py create mode 100644 temporalio/worker/_stream_store.py create mode 100644 tests/contrib/redis/__init__.py create mode 100644 tests/contrib/redis/test_streams.py create mode 100644 tests/helpers/streams.py create mode 100644 tests/nexus/test_stream_appender.py create mode 100644 tests/nexus/test_stream_store.py create mode 100644 tests/nexus/test_stream_store_workflow.py diff --git a/.github/actions/install-redis/action.yml b/.github/actions/install-redis/action.yml new file mode 100644 index 0000000000..de0a2d6390 --- /dev/null +++ b/.github/actions/install-redis/action.yml @@ -0,0 +1,25 @@ +name: Install Redis +description: Install and start a Redis server for the tests of the Redis Stream Store. +runs: + using: composite + steps: + - if: runner.os == 'Linux' + shell: bash + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends redis-server + - if: runner.os == 'macOS' + shell: bash + run: brew install redis + # Redis has no official Windows build, so the Redis Store tests skip there. The port is not the + # default one, since the Linux package starts its own server on 6379. + - if: runner.os != 'Windows' + shell: bash + run: | + redis-server --port 6380 --save "" --appendonly no --daemonize yes + for _ in $(seq 1 50); do + redis-cli -p 6380 ping && break + sleep 0.1 + done + redis-cli -p 6380 ping + echo "STREAMS_REDIS_URL=redis://127.0.0.1:6380" >> "$GITHUB_ENV" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9e3f5e2daa..c60bb33161 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -58,6 +58,8 @@ jobs: if: ${{ matrix.clippyLinter }} - run: poe build-develop - run: poe lint + - name: Install and start Redis for the Redis Stream Store tests + uses: ./.github/actions/install-redis - run: mkdir junit-xml - run: poe test ${{matrix.pytestExtraArgs}} -s --junit-xml=junit-xml/${{ matrix.python }}--${{ matrix.os }}.xml timeout-minutes: 20 @@ -226,6 +228,9 @@ jobs: - run: uv tool install poethepoet - run: uv sync --all-extras - run: cargo install --locked nexgen --version 0.2.8 --features advanced --force + # The Stream package needs the `main` of nexgen, which no release has yet. + - run: cargo install --locked --git https://github.com/temporalio/nexgen --rev 9653a0f5a845df0c5059554f1b5813db6a74cd1a --features advanced --root "$RUNNER_TEMP/nexgen-stream" nexgen + - run: echo "NEX_GEN_STREAM_BIN=$RUNNER_TEMP/nexgen-stream/bin/nexgen" >> "$GITHUB_ENV" - run: poe build-develop - run: poe gen-protos - name: Check generation unchanged diff --git a/changelog/added/chattering-ledger-otter.md b/changelog/added/chattering-ledger-otter.md new file mode 100644 index 0000000000..f15fdc478b --- /dev/null +++ b/changelog/added/chattering-ledger-otter.md @@ -0,0 +1 @@ +Add the experimental Streams API. Client and Activity code create a Stream, append Records to it, read them and close it, and Workflow code creates and describes one. A client takes its Stream Stores in `stream_config=StreamConfig(stores=[...])`, such as `temporalio.contrib.redis.streams.RedisStreamStore`, or `temporalio.testing.MemoryStreamStore` in tests. diff --git a/pyproject.toml b/pyproject.toml index 2e25bb1fc9..8d427a07a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,8 +93,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-redis" +build-develop-with-release = { cmd = "uv run maturin develop --release --uv --features streams-redis" } # 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" } } @@ -153,18 +153,20 @@ filterwarnings = [ before-all = "pip install protoc-wheel-0" build = "cp311-win_amd64 cp311-manylinux_x86_64 cp311-manylinux_aarch64 cp311-musllinux_x86_64 cp311-musllinux_aarch64 cp311-macosx_x86_64 cp311-macosx_arm64" build-verbosity = 1 +# The official wheels link the Stream runtime. +environment = { MATURIN_PEP517_ARGS = "--features streams-redis" } [tool.cibuildwheel.macos] -environment = { MACOSX_DEPLOYMENT_TARGET = "10.12" } +environment = { MACOSX_DEPLOYMENT_TARGET = "10.12", MATURIN_PEP517_ARGS = "--features streams-redis" } [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-redis" } [[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-redis" } [tool.mypy] ignore_missing_imports = true @@ -173,12 +175,17 @@ exclude = [ 'temporalio/api', 'temporalio/bridge/proto', 'temporalio/nexus/system/workflow_service', + 'temporalio/nexus/system/stream_service', ] [[tool.mypy.overrides]] module = "temporalio.nexus.system.workflow_service.*" ignore_errors = true +[[tool.mypy.overrides]] +module = "temporalio.nexus.system.stream_service.*" +ignore_errors = true + [tool.pydocstyle] convention = "google" # https://github.com/PyCQA/pydocstyle/issues/363#issuecomment-625563088 diff --git a/scripts/gen_nexus_system_api.py b/scripts/gen_nexus_system_api.py index 3b16ce1bb3..48998fc7c8 100644 --- a/scripts/gen_nexus_system_api.py +++ b/scripts/gen_nexus_system_api.py @@ -23,9 +23,14 @@ ) wit_deps_dir = wit_input_dir / "deps" workflow_output_dir = base_dir / "temporalio" / "nexus" / "system" / "workflow_service" +stream_output_dir = base_dir / "temporalio" / "nexus" / "system" / "stream_service" notification_output_dir = base_dir / "temporalio" / "nexus" / "notifications" workflow_init_path = base_dir / "temporalio" / "workflow" / "__init__.py" NEX_GEN_VERSION = "0.2.8" +# The Stream package needs a newer nexgen, which rejects the other WITs. One pin is enough +# again once the api WITs and a nexgen release agree. No release has nexgen commit 9653a0f5, +# the `main` of nexgen, yet, so NEX_GEN_STREAM_BIN points at a build of it. +NEX_GEN_STREAM_VERSION = "0.2.12" proto_files = [ gen_protos.api_proto_dir / "temporal" @@ -36,10 +41,22 @@ gen_protos.api_proto_dir / "temporal" / "api" + / "streamservice" + / "v1" + / "request_response.proto", + gen_protos.api_proto_dir + / "temporal" + / "api" / "notificationservice" / "v1" / "request_response.proto", ] +# The Workflow functions of these generated packages are exported by temporalio.workflow. +# They are in import order, so that sorting the imports keeps them between the markers. +workflow_export_packages = [ + ("temporalio.nexus.system.stream_service.workflow", stream_output_dir / "workflow"), + ("temporalio.nexus.system.workflow_service", workflow_output_dir), +] def nex_gen_command() -> list[str]: @@ -66,6 +83,20 @@ def nex_gen_command() -> list[str]: return ["nexgen"] +def nex_gen_stream_command() -> list[str]: + bin_path = os.environ.get("NEX_GEN_STREAM_BIN") + if not bin_path: + raise RuntimeError( + "NEX_GEN_STREAM_BIN must point at a nexgen build of commit 9653a0f5" + ) + version = subprocess.check_output([bin_path, "--version"], text=True).strip() + if version != f"nexgen {NEX_GEN_STREAM_VERSION}": + raise RuntimeError( + f"NEX_GEN_STREAM_BIN is {version}, not nexgen {NEX_GEN_STREAM_VERSION}" + ) + return [bin_path] + + def build_descriptor_set(descriptor_path: Path) -> None: subprocess.check_call( [ @@ -85,14 +116,15 @@ def generate_package( output_dir: Path, descriptor_path: Path, *, - native_api: bool, + flags: list[str], ) -> None: - args = [*command, "python", str(wit_input_dir / wit_name), str(wit_deps_dir)] - if native_api: - args.append("--native-api") subprocess.check_call( [ - *args, + *command, + "python", + str(wit_input_dir / wit_name), + str(wit_deps_dir), + *flags, "--system-nexus", "--support-package", "temporalio.nexus.system._support", @@ -132,17 +164,26 @@ def replace_marker_block( def generate_workflow_exports() -> None: - exports = generated_exports(workflow_output_dir / "__init__.py") + package_exports = [ + (module, generated_exports(package_dir / "__init__.py")) + for module, package_dir in workflow_export_packages + ] import_block = [ "# BEGIN GENERATED NEXUS SYSTEM EXPORTS\n", - "from temporalio.nexus.system.workflow_service import (\n", - *[f" {export},\n" for export in exports], - ")\n", + *[ + line + for module, exports in package_exports + for line in [ + f"from {module} import (\n", + *[f" {export},\n" for export in exports], + ")\n", + ] + ], "# END GENERATED NEXUS SYSTEM EXPORTS\n", ] all_block = [ " # BEGIN GENERATED NEXUS SYSTEM __ALL__\n", - *[f' "{export}",\n' for export in exports], + *[f' "{export}",\n' for _, exports in package_exports for export in exports], " # END GENERATED NEXUS SYSTEM __ALL__\n", ] content = workflow_init_path.read_text() @@ -162,6 +203,7 @@ def generate_workflow_exports() -> None: def generate_nexus_system_api() -> None: required_paths = [ wit_input_dir / "workflow-service.wit", + wit_input_dir / "stream-service.wit", wit_input_dir / "notification-service.wit", wit_deps_dir, *proto_files, @@ -173,23 +215,31 @@ def generate_nexus_system_api() -> None: descriptor_path = Path(temp_dir) / "temporal_api.bin" build_descriptor_set(descriptor_path) command = nex_gen_command() - for wit_name, output_dir, native_api in ( - ("workflow-service.wit", workflow_output_dir, True), - ("notification-service.wit", notification_output_dir, False), + stream_command = nex_gen_stream_command() + for wit_name, output_dir, package_command, flags in ( + ("workflow-service.wit", workflow_output_dir, command, ["--native-api"]), + ( + "stream-service.wit", + stream_output_dir, + stream_command, + ["--targets", "workflow,client"], + ), + ("notification-service.wit", notification_output_dir, command, []), ): shutil.rmtree(output_dir, ignore_errors=True) output_dir.parent.mkdir(parents=True, exist_ok=True) generate_package( - command, + package_command, wit_name, output_dir, descriptor_path, - native_api=native_api, + flags=flags, ) workflow_output_dir.parent.joinpath("__init__.py").touch() generate_workflow_exports() format_paths = [ str(workflow_output_dir), + str(stream_output_dir), str(notification_output_dir), str(workflow_init_path), ] diff --git a/scripts/gen_payload_visitor.py b/scripts/gen_payload_visitor.py index ba5ddc112b..eb41d72864 100644 --- a/scripts/gen_payload_visitor.py +++ b/scripts/gen_payload_visitor.py @@ -24,20 +24,30 @@ def discover_system_nexus_roots() -> list[Descriptor]: - module_path = base_dir / "temporalio" / "nexus" / "system" / "workflow_service" - spec = spec_from_file_location( - "temporalio_nexus_system_workflow_service", - module_path / "__init__.py", - submodule_search_locations=[str(module_path)], - ) - if spec is None or spec.loader is None: - raise RuntimeError(f"Cannot load generated system service from {module_path}") - module = module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) + services: list[type] = [] + system_path = base_dir / "temporalio" / "nexus" / "system" + for module_path in sorted(p.parent for p in system_path.glob("*/services.py")): + spec = spec_from_file_location( + f"temporalio_nexus_system_{module_path.name}", + module_path / "__init__.py", + submodule_search_locations=[str(module_path)], + ) + if spec is None or spec.loader is None: + raise RuntimeError( + f"Cannot load generated system service from {module_path}" + ) + module = module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + services.extend( + value + for value in vars(module.services).values() + if isinstance(value, type) and nexusrpc.get_service_definition(value) + ) + annotations = [a for service in services for a in get_type_hints(service).values()] roots: list[Descriptor] = [] - for annotation in get_type_hints(module._services.WorkflowService).values(): + for annotation in annotations: if get_origin(annotation) is not nexusrpc.Operation: continue for operation_type in get_args(annotation): diff --git a/temporalio/bridge/Cargo.lock b/temporalio/bridge/Cargo.lock index 4e2c0a2f60..0a7f33d260 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.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84a22b5d51209f6de83d37a0b33d320acb8da88823e42a94c075140708f032a5" +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.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0876700b0d007ad25d5508865a4259c3af9af6d296bf663a0d586d0bc3369f6f" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-trait" version = "0.1.92" @@ -181,6 +207,15 @@ version = "2.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "bon" version = "3.10.3" @@ -265,7 +300,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.3.1", "rand_core 0.10.1", ] @@ -306,7 +341,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4356eba6976435498d05a311ee9fc486da3bc81c200d3f0d7dd9703b41913f92" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", ] [[package]] @@ -343,6 +382,15 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "cpufeatures" version = "0.3.1" @@ -352,6 +400,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.5" @@ -376,6 +430,16 @@ version = "0.8.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + [[package]] name = "darling" version = "0.24.1" @@ -433,6 +497,16 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "dirs" version = "7.0.0" @@ -548,6 +622,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" @@ -728,6 +822,16 @@ dependencies = [ "slab", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "gethostname" version = "1.1.0" @@ -1501,6 +1605,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" @@ -2041,6 +2151,39 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "redis" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db7f3d75b5a3301746bff39a767bf4c5c42506b08513527308b3875bfb64a031" +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" @@ -2390,6 +2533,23 @@ dependencies = [ "serde", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -2593,6 +2753,8 @@ dependencies = [ "temporalio-client", "temporalio-common", "temporalio-sdk-core", + "temporalio-streams", + "temporalio-streams-redis", "tokio", "tokio-rustls", "tokio-stream", @@ -2776,6 +2938,42 @@ dependencies = [ "zip", ] +[[package]] +name = "temporalio-streams" +version = "0.1.0" +dependencies = [ + "async-trait", + "parking_lot", + "prost", + "prost-wkt-types", + "sha2", + "temporalio-common", + "thiserror", + "tokio", + "tonic", +] + +[[package]] +name = "temporalio-streams-redis" +version = "0.1.0" +source = "git+https://github.com/temporal-community/sdk-streams?rev=63799c4fd9431a9f12bed78fdff66f144f1e0ac4#63799c4fd9431a9f12bed78fdff66f144f1e0ac4" +dependencies = [ + "async-trait", + "aws-lc-rs", + "crc16", + "parking_lot", + "prost", + "prost-wkt-types", + "redis", + "temporalio-common", + "temporalio-streams", + "thiserror", + "tokio", + "tonic-prost-build", + "tracing", + "url", +] + [[package]] name = "termtree" version = "0.5.1" @@ -3141,6 +3339,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c" +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "typetag" version = "0.2.23" @@ -3631,6 +3835,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.4" diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml index 4b23ae6ce9..a49caf5eee 100644 --- a/temporalio/bridge/Cargo.toml +++ b/temporalio/bridge/Cargo.toml @@ -29,14 +29,21 @@ default = ["tls-ring"] tls-ring = [ "temporalio-sdk-core/tls-ring", "temporalio-client/tls-ring", + "temporalio-streams-redis?/tls-ring", "tokio-rustls/ring", ] tls-aws-lc = [ "temporalio-sdk-core/tls-aws-lc", "temporalio-client/tls-aws-lc", + "temporalio-streams-redis?/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-redis", "temporalio-streams-redis/fips"] +# Links the Stream runtime and its memory Store. Off by default. +streams = ["dep:temporalio-streams"] +# Adds the Redis Store and turns on `streams`. The official builds turn it on. Cargo fetches the +# git dependency below for every build, with this feature off too. +streams-redis = ["streams", "dep:temporalio-streams-redis"] [dependencies] anyhow = "1.0" @@ -72,6 +79,10 @@ temporalio-sdk-core = { version = "0.11.0", path = "./sdk-core/crates/sdk-core", "envconfig", "prometheus", ] } +# The Stream runtime and its memory Store, from the submodule. +temporalio-streams = { path = "./sdk-core/crates/streams", optional = true } +# The Redis Store, which lives in its own repo. +temporalio-streams-redis = { git = "https://github.com/temporal-community/sdk-streams", rev = "63799c4fd9431a9f12bed78fdff66f144f1e0ac4", default-features = false, optional = true } tokio = "1.26" # Matches the sdk-core client crate's rustls stack; used to build the custom # server certificate verifier for ClientTlsConfig.verification_server_name. @@ -84,6 +95,13 @@ tonic = "0.14" tracing = "0.1" url = "2.2" +# The Redis crate takes temporalio-common and temporalio-streams from crates.io. This makes it use +# the submodule's crates, so the bridge sees one copy of each type. The versions must be semver +# compatible. +[patch.crates-io] +temporalio-common = { path = "./sdk-core/crates/common" } +temporalio-streams = { path = "./sdk-core/crates/streams" } + [profile.release] opt-level = 3 debug = false diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index c675c32469..9a110b3cda 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit c675c32469af707e1c979a7b808c70b4a0608b82 +Subproject commit 9a110b3cda949199b7792348ff5d265b1d1c414a diff --git a/temporalio/bridge/src/lib.rs b/temporalio/bridge/src/lib.rs index a4fda26db6..1eb7b0be62 100644 --- a/temporalio/bridge/src/lib.rs +++ b/temporalio/bridge/src/lib.rs @@ -6,6 +6,7 @@ mod client_rpc_generated; mod envconfig; mod metric; mod runtime; +mod streams; mod testing; mod worker; @@ -48,6 +49,16 @@ fn temporal_sdk_bridge(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(init_runtime, m)?)?; m.add_function(wrap_pyfunction!(raise_in_thread, m)?)?; + // Stream stuff + m.add("StreamCallError", py.get_type::())?; + #[cfg(feature = "streams")] + { + m.add_class::()?; + m.add_function(wrap_pyfunction!(new_memory_stream_runtime, m)?)?; + #[cfg(feature = "streams-redis")] + m.add_function(wrap_pyfunction!(new_redis_stream_runtime, m)?)?; + } + // Testing stuff m.add_class::()?; m.add_function(wrap_pyfunction!(start_dev_server, m)?)?; @@ -104,6 +115,62 @@ fn raise_in_thread(py: Python, thread_id: std::os::raw::c_long, exc: &Bound<'_, runtime::raise_in_thread(py, thread_id, exc) } +#[cfg(feature = "streams")] +#[pyfunction] +#[pyo3(signature = (runtime_ref, store_name, retention=None, last_write_ttl=None))] +fn new_memory_stream_runtime( + runtime_ref: &runtime::RuntimeRef, + store_name: String, + retention: Option, + last_write_ttl: Option, +) -> PyResult { + streams::new_memory_stream_runtime(runtime_ref, store_name, retention, last_write_ttl) +} + +#[cfg(feature = "streams-redis")] +#[pyfunction] +#[pyo3(signature = ( + runtime_ref, + urls, + *, + cluster, + key_prefix, + retention, + last_write_ttl, + response_timeout, + store_name, + wait_replicas, + wait_timeout, +))] +#[allow(clippy::too_many_arguments)] +fn new_redis_stream_runtime( + runtime_ref: &runtime::RuntimeRef, + urls: Vec, + cluster: bool, + key_prefix: Option, + retention: Option, + last_write_ttl: Option, + response_timeout: Option, + store_name: String, + wait_replicas: i32, + wait_timeout: Option, +) -> PyResult { + streams::new_redis_stream_runtime( + runtime_ref, + streams::RedisStreamOptions { + urls, + cluster, + key_prefix, + retention, + last_write_ttl, + response_timeout, + store_name, + wait_replicas, + wait_timeout, + }, + ) +} + #[pyfunction] fn start_dev_server<'a>( py: Python<'a>, diff --git a/temporalio/bridge/src/streams.rs b/temporalio/bridge/src/streams.rs new file mode 100644 index 0000000000..0f9218523c --- /dev/null +++ b/temporalio/bridge/src/streams.rs @@ -0,0 +1,130 @@ +//! Bindings to the Stream runtime of the `temporalio-streams` crate, on the memory Store or, with +//! the `streams-redis` feature, on the Redis Store of `temporalio-streams-redis`. Python calls the +//! five Stream service operations by name, with the encoded request, and gets the encoded response. + +use pyo3::exceptions::PyException; +#[cfg(feature = "streams")] +use pyo3::prelude::*; + +#[cfg(feature = "streams")] +use crate::runtime; + +// The exception args are the gRPC code, the message and the encoded details, as in `RPCError`. +pyo3::create_exception!(temporal_sdk_bridge, StreamCallError, PyException); + +/// The crate's runtime object. Each Store on the bridge holds one. +#[cfg(feature = "streams")] +#[pyclass(frozen)] +pub struct StreamRuntimeRef { + streams: std::sync::Arc, + runtime: runtime::Runtime, +} + +#[cfg(feature = "streams")] +#[pymethods] +impl StreamRuntimeRef { + /// Runs one Stream operation on an encoded request. The future gives the encoded response. + fn call<'p>( + &self, + py: Python<'p>, + operation: String, + request: Vec, + ) -> PyResult> { + self.runtime + .assert_same_process("call the Stream runtime")?; + let streams = self.streams.clone(); + self.runtime.future_into_py(py, async move { + streams.call(&operation, &request).await.map_err(call_error) + }) + } +} + +#[cfg(feature = "streams")] +fn call_error(err: temporalio_streams::StreamError) -> PyErr { + StreamCallError::new_err((err.code() as u32, err.message.clone(), err.details())) +} + +/// Builds the runtime on a new memory Store. A `None` duration takes the crate's default. +#[cfg(feature = "streams")] +pub fn new_memory_stream_runtime( + runtime_ref: &runtime::RuntimeRef, + store_name: String, + retention: Option, + last_write_ttl: Option, +) -> PyResult { + let runtime = runtime_ref.runtime.clone(); + runtime.assert_same_process("create a Stream runtime")?; + // The runtime can spawn Tokio tasks, which need an entered Tokio runtime. + let _guard = runtime.core.tokio_handle().enter(); + let meter = runtime.core.telemetry().get_temporal_metric_meter(); + let streams = temporalio_streams::MemoryStreamRuntime::new( + store_name, + temporalio_streams::MemoryStreamOptions { + retention, + last_write_ttl, + }, + meter, + ); + Ok(StreamRuntimeRef { streams, runtime }) +} + +/// The options of the Redis Store, as Python passes them. A `None` takes the crate's default. +#[cfg(feature = "streams-redis")] +pub struct RedisStreamOptions { + pub urls: Vec, + pub cluster: bool, + pub key_prefix: Option, + pub retention: Option, + pub last_write_ttl: Option, + pub response_timeout: Option, + pub store_name: String, + pub wait_replicas: i32, + pub wait_timeout: Option, +} + +/// Builds the runtime on Redis. It connects at its first call, so a server that is down fails that +/// call and not this one. A bad option raises `ValueError` with the crate's message. +#[cfg(feature = "streams-redis")] +pub fn new_redis_stream_runtime( + runtime_ref: &runtime::RuntimeRef, + options: RedisStreamOptions, +) -> PyResult { + use pyo3::exceptions::PyValueError; + use temporalio_streams_redis::{proto::RedisStreamConfig, RedisStreamRuntime, RuntimeOptions}; + + let runtime = runtime_ref.runtime.clone(); + runtime.assert_same_process("create a Stream runtime")?; + // The crate reads an empty or zero field as its default, so `None` maps to that. + let config = RedisStreamConfig { + urls: options.urls, + cluster: options.cluster, + key_prefix: options.key_prefix.unwrap_or_default(), + retention: proto_duration("retention", options.retention)?, + last_write_ttl: proto_duration("last_write_ttl", options.last_write_ttl)?, + response_timeout: proto_duration("response_timeout", options.response_timeout)?, + store_name: options.store_name, + wait_replicas: options.wait_replicas, + wait_timeout: proto_duration("wait_timeout", options.wait_timeout)?, + ..Default::default() + }; + let options = + RuntimeOptions::try_from(config).map_err(|err| PyValueError::new_err(err.to_string()))?; + // The runtime can spawn Tokio tasks, which need an entered Tokio runtime. + let _guard = runtime.core.tokio_handle().enter(); + let meter = runtime.core.telemetry().get_temporal_metric_meter(); + let streams = RedisStreamRuntime::create(options, meter); + Ok(StreamRuntimeRef { streams, runtime }) +} + +// The config's durations are the `prost-wkt-types` type that the Redis crate's protos use. The +// bridge doesn't depend on that crate, so the field's type picks the conversion. +#[cfg(feature = "streams-redis")] +fn proto_duration(field: &str, value: Option) -> PyResult> +where + T: TryFrom, + T::Error: std::fmt::Display, +{ + value.map(T::try_from).transpose().map_err(|err| { + pyo3::exceptions::PyValueError::new_err(format!("The {field} is too long: {err}")) + }) +} diff --git a/temporalio/bridge/streams.py b/temporalio/bridge/streams.py new file mode 100644 index 0000000000..b189517c25 --- /dev/null +++ b/temporalio/bridge/streams.py @@ -0,0 +1,104 @@ +"""The Stream runtime that this bridge build links. (unstable) + +Nothing in this module should be considered stable. The API may change. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import timedelta +from typing import Any + +import temporalio.bridge.runtime +import temporalio.bridge.temporal_sdk_bridge +import temporalio.exceptions + +StreamCallError: type[Exception] = temporalio.bridge.temporal_sdk_bridge.StreamCallError + + +class StreamRuntime: + """The runtime of the ``temporalio-streams`` crate on one Store.""" + + @staticmethod + def new_memory( + runtime: temporalio.bridge.runtime.Runtime, + *, + name: str, + retention: timedelta | None, + last_write_ttl: timedelta | None, + ) -> StreamRuntime: + """Build the runtime on a new memory Store. + + Raises: + temporalio.exceptions.StreamConfigError: The bridge build has no Stream + runtime. + """ + new = getattr( + temporalio.bridge.temporal_sdk_bridge, "new_memory_stream_runtime", None + ) + if new is None: + raise temporalio.exceptions.StreamConfigError( + "This build of temporalio has no Stream runtime" + ) + return StreamRuntime(new(runtime._ref, name, retention, last_write_ttl)) + + @staticmethod + def new_redis( + runtime: temporalio.bridge.runtime.Runtime, + *, + urls: Sequence[str], + cluster: bool, + key_prefix: str | None, + retention: timedelta | None, + last_write_ttl: timedelta | None, + response_timeout: timedelta | None, + name: str, + wait_replicas: int, + wait_timeout: timedelta | None, + ) -> StreamRuntime: + """Build the runtime on Redis. It connects at its first call. + + Raises: + temporalio.exceptions.StreamConfigError: The bridge build has no Redis + Stream Store. + ValueError: An option is bad. + """ + new = getattr( + temporalio.bridge.temporal_sdk_bridge, "new_redis_stream_runtime", None + ) + if new is None: + raise temporalio.exceptions.StreamConfigError( + "This build of temporalio has no Redis Stream Store" + ) + return StreamRuntime( + new( + runtime._ref, + list(urls), + cluster=cluster, + key_prefix=key_prefix, + retention=retention, + last_write_ttl=last_write_ttl, + response_timeout=response_timeout, + store_name=name, + wait_replicas=wait_replicas, + wait_timeout=wait_timeout, + ) + ) + + def __init__(self, ref: Any) -> None: + """Wrap the bridge's runtime object. Use ``new_memory`` or ``new_redis`` instead.""" + self._ref = ref + + async def call(self, operation: str, request: bytes) -> bytes: + """Run one Stream operation on an encoded request and return the encoded response. + + Raises: + temporalio.exceptions.StreamError: The call failed. + """ + try: + return await self._ref.call(operation, request) + except StreamCallError as err: + _, message, details = err.args + raise temporalio.exceptions._stream_error_from_call( + message, details + ) from None diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index 2eef41a39d..7305b4be59 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -185,6 +185,13 @@ ScheduleUpdate, ScheduleUpdateInput, ) +from ._stream_appender import BufferedAppender +from ._stream_store import ( + StreamConfig, + StreamSelectContext, + StreamStore, + StreamStoreCallOptions, +) from ._worker_versioning import ( BuildIdOp, BuildIdOpAddNewCompatible, @@ -350,6 +357,11 @@ "TaskReachabilityType", "CloudOperationsClient", "Plugin", + "BufferedAppender", + "StreamConfig", + "StreamSelectContext", + "StreamStore", + "StreamStoreCallOptions", "Callback", "_ClientImpl", "_apply_headers", diff --git a/temporalio/client/_client.py b/temporalio/client/_client.py index efef5cb528..d42d7d9b0e 100644 --- a/temporalio/client/_client.py +++ b/temporalio/client/_client.py @@ -25,6 +25,7 @@ import temporalio.api.workflowservice.v1 import temporalio.common import temporalio.converter +import temporalio.nexus.system.stream_service.client import temporalio.runtime import temporalio.service import temporalio.workflow @@ -94,6 +95,7 @@ ScheduleBackfill, ScheduleHandle, ) +from ._stream_store import StreamConfig, _refuse_store_plugins from ._worker_versioning import ( BuildIdOp, TaskReachabilityType, @@ -114,7 +116,9 @@ from ._plugin import Plugin -class Client: +class Client( + temporalio.nexus.system.stream_service.client.StreamServiceClientExtension +): """Client for accessing Temporal. Most users will use :py:meth:`connect` to create a client. The @@ -156,6 +160,7 @@ async def connect( grpc_compression: GrpcCompression = GrpcCompression.GZIP, payload_limits: PayloadLimitsConfig = PayloadLimitsConfig(), header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_config: StreamConfig | None = None, ) -> Self: """Connect to a Temporal server. @@ -225,7 +230,18 @@ async def connect( payload_limits: Warning thresholds for outbound payload/memo sizes. Over-threshold fields are logged but still sent. Set a threshold to 0 to disable it. header_codec_behavior: Encoding behavior for headers sent by the client. + stream_config: The Stream Stores of this client and the selector that + picks one for a new Stream, as + :py:class:`temporalio.converter.ExternalStorage` holds drivers and a + selector. A client with it sends each Stream operation to one of its + Stores. This release has no Streams on the server, so with ``None`` a + Stream operation raises + :py:class:`temporalio.exceptions.StreamConfigError`. + + .. warning:: + This API is experimental and subject to change. """ + _refuse_store_plugins(plugins) connect_config = temporalio.service.ConnectConfig( target_host=target_host, api_key=api_key, @@ -262,6 +278,7 @@ def make_lambda( default_workflow_query_reject_condition=default_workflow_query_reject_condition, header_codec_behavior=header_codec_behavior, plugins=plugins, + stream_config=stream_config, ) def __init__( @@ -275,11 +292,13 @@ def __init__( default_workflow_query_reject_condition: None | (temporalio.common.QueryRejectCondition) = None, header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_config: StreamConfig | None = None, ): """Create a Temporal client from a service client. See :py:meth:`connect` for details on the parameters. """ + _refuse_store_plugins(plugins) # Store the config for tracking config = ClientConfig( service_client=service_client, @@ -289,6 +308,7 @@ def __init__( interceptors=interceptors, default_workflow_query_reject_condition=default_workflow_query_reject_condition, header_codec_behavior=header_codec_behavior, + stream_config=stream_config, ) self._initial_config = config.copy() @@ -3029,6 +3049,7 @@ class ClientConnectConfig(TypedDict, total=False): grpc_compression: GrpcCompression payload_limits: PayloadLimitsConfig header_codec_behavior: HeaderCodecBehavior + stream_config: StreamConfig | None class ClientConfig(TypedDict, total=False): @@ -3043,3 +3064,4 @@ class ClientConfig(TypedDict, total=False): temporalio.common.QueryRejectCondition | None ] header_codec_behavior: Required[HeaderCodecBehavior] + stream_config: Required[StreamConfig | None] diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index 72986ffa70..4766126ae8 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -33,6 +33,7 @@ import temporalio.exceptions import temporalio.nexus import temporalio.nexus._operation_context +import temporalio.nexus.system from temporalio.activity import ActivityCancellationDetails from temporalio.converter import ( ActivitySerializationContext, @@ -125,6 +126,13 @@ ScheduleUpdate, ScheduleUpdateInput, ) +from ._stream_store import ( + _STREAM_SERVICE, + StreamConfig, + _activity_target, + _call_stream_store, + _request_ref, +) from ._worker_versioning import WorkerBuildIdVersionSets, WorkerTaskReachability from ._workflow import ( WorkflowExecutionAsyncIterator, @@ -1565,6 +1573,18 @@ async def start_nexus_operation( self, input: StartNexusOperationInput ) -> NexusOperationHandle[Any]: """Start a nexus operation and return a handle to it.""" + if ( + temporalio.nexus.system.is_system_endpoint(input.endpoint) + and input.service == _STREAM_SERVICE.name + ): + # This release has no Streams on the server, so a Stream call needs a Store. + stream_config = self._client._config.get("stream_config") + if stream_config is None: + raise temporalio.exceptions.StreamConfigError( + "This client has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...])." + ) + return await self._start_stream_store_operation(stream_config, input) nexus_context = temporalio.converter.NexusSerializationContext( endpoint=input.endpoint, service=input.service, @@ -1648,6 +1668,37 @@ async def start_nexus_operation( operation=input.operation, ) + async def _start_stream_store_operation( + self, stream_config: StreamConfig, input: StartNexusOperationInput + ) -> NexusOperationHandle[Any]: + data_converter = self._client.data_converter + serialization_context = temporalio.nexus.system._get_serialization_context( + input.service, input.operation, input.arg + ) + if serialization_context is not None: + data_converter = data_converter.with_context(serialization_context) + payload_converter = temporalio.nexus.system._get_payload_converter( + data_converter.payload_converter, data_converter.failure_converter + ) + # The generated model converts to its proto message with no payload in between. + request = payload_converter._to_transfer_value(input.arg) + store = stream_config._store_for(_request_ref(input.arg), _activity_target()) + response = await _call_stream_store( + store, data_converter, input.operation, request + ) + handle: NexusOperationHandle[Any] = NexusOperationHandle( + self._client, + input.id, + result_type=input.result_type, + endpoint=input.endpoint, + service=input.service, + operation=input.operation, + ) + handle._known_outcome = payload_converter._from_transfer_value( + response, input.result_type + ) + return handle + async def describe_nexus_operation( self, input: DescribeNexusOperationInput ) -> NexusOperationExecutionDescription: diff --git a/temporalio/client/_stream_appender.py b/temporalio/client/_stream_appender.py new file mode 100644 index 0000000000..ec8d3d49d3 --- /dev/null +++ b/temporalio/client/_stream_appender.py @@ -0,0 +1,193 @@ +"""A helper that batches the appends of one Stream.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta +from types import TracebackType +from typing import Generic, TypeVar + +from typing_extensions import Self + +import temporalio.nexus.system.stream_service.client + +from ._stream_store import _check_not_negative + +BodyT = TypeVar("BodyT") + +# The Store refuses an append of more Records than this. Many small Records fit under its byte +# limit, but they keep the Store busy for long. +_MAX_BATCH_RECORDS = 1000 + + +class BufferedAppender(Generic[BodyT]): + """Batches the appends of one Stream. + + Use it in an Activity or a client that appends often, such as one that streams model + tokens. :py:meth:`append` returns at once and does no I/O. The appender writes the + buffered Records as one batch at most ``max_delay`` after the first of them, or at once + when they fill a batch of 1,000 Records. One write is in flight at a time. Records + appended during a write wait for the next batch, and they keep their order. + + Use it in ``async with``, which writes the remaining Records at the end of the block, or + await :py:meth:`flush`. Without ``async with`` or a :py:meth:`flush`, the Records still in + the buffer are lost. The end of an Activity attempt doesn't write them. + + ``await flush()`` is the backpressure point. A caller that appends faster than the Store + accepts Records must await it from time to time, so the buffer doesn't grow without bound. + + After a batch fails, the appender writes nothing more, so a reader sees no gap. Each later + :py:meth:`append` and :py:meth:`flush` raises the error of that batch. + + The appender doesn't count bytes, so a batch over the size limit of the Store fails with + the Store's error. + + It takes a client Stream, so Workflow code has no use for it. + + .. warning:: + This API is experimental and subject to change. + """ + + def __init__( + self, + stream: temporalio.nexus.system.stream_service.client.Stream[BodyT], + *, + max_delay: timedelta = timedelta(milliseconds=100), + ) -> None: + """Create an appender. It starts no write until the first append. + + Args: + stream: The Stream to append to. + max_delay: The longest time that a Record waits in the buffer before its batch + starts to write. It must not be negative. + """ + _check_not_negative(max_delay=max_delay) + self._stream = stream + self._max_delay = max_delay + self._buffer: list[BodyT] = [] + # When the oldest buffered Record arrived, so its batch starts at most max_delay later. + self._first_at = 0.0 + self._timer: asyncio.TimerHandle | None = None + self._writing: asyncio.Task[None] | None = None + self._appended = 0 + # Count of the Records that a write took or that a failed batch left, written or not. + self._settled = 0 + self._error: BaseException | None = None + self._last_offset: bytes | None = None + self._unwritten = 0 + + def append(self, *records: BodyT) -> None: + """Add one Record for each value to the buffer, and return at once. + + A list is one value. + + Raises: + TypeError: No value was given. + temporalio.exceptions.StreamError: A batch failed, so the appender takes no more + Records. + """ + self._raise_error() + if not records: + raise TypeError("append takes one or more values") + loop = asyncio.get_running_loop() + if not self._buffer: + self._first_at = loop.time() + self._buffer.extend(records) + self._appended += len(records) + self._schedule() + + async def flush(self) -> bytes | None: + """Write every Record appended before the call. + + Returns: + The ``last_offset`` of the last written batch, or ``None`` if the appender has + written nothing. + + Raises: + temporalio.exceptions.StreamError: A batch failed. + """ + target = self._appended + while self._settled < target: + if self._writing is None: + self._start() + assert self._writing is not None + # The write goes on when the caller is cancelled, so no batch is left half done. + await asyncio.shield(self._writing) + self._raise_error() + return self._last_offset + + @property + def unwritten(self) -> int: + """How many Records the failed batch and the ones after it left unwritten.""" + return self._unwritten + + async def __aenter__(self) -> Self: + """Use the appender in ``async with``.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Write the buffered Records. A write error doesn't hide an error of the block.""" + try: + await self.flush() + except Exception: + if exc is None: + raise + + def _raise_error(self) -> None: + if self._error is not None: + raise self._error + + def _schedule(self) -> None: + if self._writing is not None or not self._buffer: + return + if len(self._buffer) >= _MAX_BATCH_RECORDS: + self._start() + elif self._timer is None: + self._timer = asyncio.get_running_loop().call_at( + self._first_at + self._max_delay.total_seconds(), self._start + ) + + def _start(self) -> None: + if self._timer is not None: + self._timer.cancel() + self._timer = None + if self._writing is not None or not self._buffer: + return + batch = self._buffer[:_MAX_BATCH_RECORDS] + del self._buffer[:_MAX_BATCH_RECORDS] + self._writing = asyncio.get_running_loop().create_task(self._write(batch)) + + async def _write(self, batch: list[BodyT]) -> None: + try: + # The Store numbers each batch, since the appender sends no `sequence` and never + # sends a batch twice. A `sequence` counts for a Producer on a Stream, and two + # appenders of one Activity attempt would each start at 1. + response = await self._stream.append(*batch) + self._last_offset = response.last_offset + except asyncio.CancelledError: + # A later call must not report a cancel that its caller didn't ask for. + self._fail(RuntimeError("The write of a batch was cancelled"), len(batch)) + raise + except Exception as err: + self._fail(err, len(batch)) + finally: + self._settled += len(batch) + self._writing = None + # Schedule the Records appended during the write, for when their delay ends or + # they fill a batch. + self._schedule() + + def _fail(self, error: BaseException, count: int) -> None: + # The Records after a failed batch are never written, so a reader sees no gap. + self._error = error + self._unwritten += count + len(self._buffer) + self._settled += len(self._buffer) + self._buffer.clear() + if self._timer is not None: + self._timer.cancel() + self._timer = None diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py new file mode 100644 index 0000000000..622f00bf18 --- /dev/null +++ b/temporalio/client/_stream_store.py @@ -0,0 +1,554 @@ +from __future__ import annotations + +import dataclasses +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable, Sequence +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +import google.protobuf.message +import nexusrpc + +import temporalio.activity +import temporalio.api.streamservice.v1 +import temporalio.bridge._visitor +import temporalio.bridge.runtime +import temporalio.common +import temporalio.converter +import temporalio.nexus.system.stream_service +import temporalio.runtime +import temporalio.workflow +from temporalio.converter import StorageDriverActivityInfo, StorageDriverWorkflowInfo +from temporalio.exceptions import StreamConfigError + +if TYPE_CHECKING: + import temporalio.bridge.streams + + +@dataclasses.dataclass(frozen=True) +class StreamStoreCallOptions: + """Options of the Local Activity that calls a Stream Store for Workflow code. + + Workflow code can't do I/O, so the default + :py:meth:`StreamStore.call_from_workflow` calls its Store from a Local Activity. + + A :py:class:`StreamConfig` holds them in + :py:attr:`StreamConfig.workflow_call_options`, and every Store of that config + gets the same options. + + .. warning:: + This API is experimental and subject to change. + """ + + start_to_close_timeout: timedelta | None = timedelta(seconds=60) + """Most time for one attempt.""" + + schedule_to_close_timeout: timedelta | None = None + """Most time for all attempts.""" + + retry_policy: temporalio.common.RetryPolicy = dataclasses.field( + default_factory=lambda: temporalio.common.RetryPolicy(maximum_attempts=1) + ) + """Retries of a failed call. The default is one attempt.""" + + def __post_init__(self) -> None: + """Check that a timeout is set.""" + if not self.start_to_close_timeout and not self.schedule_to_close_timeout: + raise ValueError( + "start_to_close_timeout or schedule_to_close_timeout must be set" + ) + + +class StreamStore(ABC): + """A Store that answers the Stream operations of a client. + + A client with a Store sends each Stream operation to it. This release has no + Streams on the server, so a client needs a Store for any Stream operation. Use + :py:class:`temporalio.contrib.redis.streams.RedisStreamStore`, or + :py:class:`temporalio.testing.MemoryStreamStore` in tests. + + To keep Streams in another system, write a Store that implements :py:attr:`name` + and the five operations. Each operation takes and gives a message of + ``temporalio.api.streamservice.v1``. The bodies of the Records arrive as payloads + that the SDK already encoded with the data converter and the payload codec, so a + Store keeps payloads and not the user's types. A failed call raises + :py:class:`temporalio.exceptions.StreamError` with a reason. + + .. warning:: + This API is experimental and subject to change. + """ + + @property + @abstractmethod + def name(self) -> str: + """Name of the Store, which the refs of its Streams carry.""" + ... + + @abstractmethod + async def create_stream( + self, request: temporalio.api.streamservice.v1.CreateStreamRequest + ) -> temporalio.api.streamservice.v1.CreateStreamResponse: + """Create a Stream, or get the open one when the conflict policy allows it. + + The ref of the response carries the name of this Store. + """ + ... + + @abstractmethod + async def append_records( + self, request: temporalio.api.streamservice.v1.AppendRecordsRequest + ) -> temporalio.api.streamservice.v1.AppendRecordsResponse: + """Append one batch of Records as one Producer. + + A batch that is sent again with the same Producer, epoch and sequence is + written once. + """ + ... + + @abstractmethod + async def read_records( + self, request: temporalio.api.streamservice.v1.ReadRecordsRequest + ) -> temporalio.api.streamservice.v1.ReadRecordsResponse: + """Read one page of Records after an offset. + + At the end of the Stream it waits up to the request's ``wait`` for a new Record. + """ + ... + + @abstractmethod + async def close_stream( + self, request: temporalio.api.streamservice.v1.CloseStreamRequest + ) -> temporalio.api.streamservice.v1.CloseStreamResponse: + """Close a Stream. Readers drain what is left, and later appends fail.""" + ... + + @abstractmethod + async def describe_stream( + self, request: temporalio.api.streamservice.v1.DescribeStreamRequest + ) -> temporalio.api.streamservice.v1.DescribeStreamResponse: + """Give the status and the newest offset of a Stream.""" + ... + + def worker_activities(self) -> Sequence[Callable[..., Any]]: + """The Activities that a Worker registers for this Store when it runs Workflows. + + The default is one Activity, named for this Store, that + :py:meth:`call_from_workflow` runs as a Local Activity. A Store that overrides + one of the two methods overrides the other to match. + """ + return [_store_call_activity(self)] + + async def call_from_workflow( + self, + operation: str, + request: google.protobuf.message.Message, + options: StreamStoreCallOptions, + ) -> google.protobuf.message.Message: + """Answer one Stream operation for Workflow code. + + It runs in Workflow code, so it does no I/O itself. The default runs the + Activity of :py:meth:`worker_activities` as a Local Activity, with the timeouts + and the retry policy of ``options``. A replay takes the result from History and + calls no method of this Store. + + Args: + operation: The Nexus operation name, ``CreateStream`` or ``DescribeStream`` + in this release. + request: The request message of ``temporalio.api.streamservice.v1``. + options: The :py:attr:`StreamConfig.workflow_call_options` of the config. + + Returns: + The response message of the operation. + """ + stream_operation = _STREAM_OPERATIONS[operation] + # Workflow code can start only a create and a describe in this release, and their + # messages carry no payload. So the Local Activity has no payload codec step. + response = await temporalio.workflow.execute_local_activity( + _store_call_activity_name(self.name), + args=[operation, request.SerializeToString()], + result_type=bytes, + schedule_to_close_timeout=options.schedule_to_close_timeout, + start_to_close_timeout=options.start_to_close_timeout, + retry_policy=options.retry_policy, + ) + return stream_operation.response_type.FromString(response) + + +class _BridgeStreamStore(StreamStore): # pyright: ignore[reportUnusedClass] + """A Store that runs on the Stream runtime of the bridge. + + It builds the runtime at its first call, so the Store can be built before the + Runtime and its Tokio threads exist. + """ + + def __init__( + self, + *, + name: str, + runtime: temporalio.runtime.Runtime | None, + ) -> None: + self._name = name + self._runtime = runtime + self._stream_runtime: temporalio.bridge.streams.StreamRuntime | None = None + + @property + def name(self) -> str: + """Name of the Store.""" + return self._name + + @abstractmethod + def _new_stream_runtime( + self, runtime: temporalio.bridge.runtime.Runtime + ) -> temporalio.bridge.streams.StreamRuntime: ... + + async def create_stream( + self, request: temporalio.api.streamservice.v1.CreateStreamRequest + ) -> temporalio.api.streamservice.v1.CreateStreamResponse: + """See base class.""" + return await self._call_bridge("CreateStream", request) + + async def append_records( + self, request: temporalio.api.streamservice.v1.AppendRecordsRequest + ) -> temporalio.api.streamservice.v1.AppendRecordsResponse: + """See base class.""" + return await self._call_bridge("AppendRecords", request) + + async def read_records( + self, request: temporalio.api.streamservice.v1.ReadRecordsRequest + ) -> temporalio.api.streamservice.v1.ReadRecordsResponse: + """See base class.""" + return await self._call_bridge("ReadRecords", request) + + async def close_stream( + self, request: temporalio.api.streamservice.v1.CloseStreamRequest + ) -> temporalio.api.streamservice.v1.CloseStreamResponse: + """See base class.""" + return await self._call_bridge("CloseStream", request) + + async def describe_stream( + self, request: temporalio.api.streamservice.v1.DescribeStreamRequest + ) -> temporalio.api.streamservice.v1.DescribeStreamResponse: + """See base class.""" + return await self._call_bridge("DescribeStream", request) + + async def _call_bridge( + self, operation: str, request: google.protobuf.message.Message + ) -> Any: + """Answer one Stream operation with the Stream runtime of the bridge. + + Raises: + temporalio.exceptions.StreamError: The call failed. + temporalio.exceptions.StreamConfigError: This build of temporalio has no + Stream runtime, or not this Store. + """ + if self._stream_runtime is None: + runtime = self._runtime or temporalio.runtime.Runtime.default() + self._stream_runtime = self._new_stream_runtime(runtime._core_runtime) + # The bridge takes and gives the messages in their wire form. + response = await self._stream_runtime.call( + operation, request.SerializeToString() + ) + return _STREAM_OPERATIONS[operation].response_type.FromString(response) + + +@dataclasses.dataclass(frozen=True) +class _StreamOperation: + """How a Store answers one operation of the Stream service.""" + + store_method: str + """The name of the :py:class:`StreamStore` method that answers it.""" + + request_type: type[google.protobuf.message.Message] + + response_type: type[google.protobuf.message.Message] + + workflow_refusal: str | None = None + """Why Workflow code can't start it in this release, or ``None`` when it can.""" + + async def call( + self, store: StreamStore, request: google.protobuf.message.Message + ) -> google.protobuf.message.Message: + """Call the method of ``store`` for this operation and check its answer.""" + return self.check_response( + store, await getattr(store, self.store_method)(request) + ) + + def check_response( + self, store: StreamStore, response: object + ) -> google.protobuf.message.Message: + """Give ``response`` back when it has the response type of this operation.""" + # A Store that someone wrote can give another type, and the error names it here + # rather than in the conversion to the user's type. + if not isinstance(response, self.response_type): + raise TypeError( + f"The Stream Store '{store.name}' answered {self.store_method} with " + f"{type(response).__name__}, not {self.response_type.__name__}" + ) + return response + + +_STREAM_OPERATIONS: dict[str, _StreamOperation] = { + "CreateStream": _StreamOperation( + "create_stream", + temporalio.api.streamservice.v1.CreateStreamRequest, + temporalio.api.streamservice.v1.CreateStreamResponse, + ), + "AppendRecords": _StreamOperation( + "append_records", + temporalio.api.streamservice.v1.AppendRecordsRequest, + temporalio.api.streamservice.v1.AppendRecordsResponse, + workflow_refusal=( + "Workflow code can't append to a Stream on a Store in this release. " + "Call a Local Activity that appends through its client." + ), + ), + "ReadRecords": _StreamOperation( + "read_records", + temporalio.api.streamservice.v1.ReadRecordsRequest, + temporalio.api.streamservice.v1.ReadRecordsResponse, + workflow_refusal=( + "Workflow code can't read a Stream on a Store in this release. " + "Call an Activity that reads it through its client." + ), + ), + "CloseStream": _StreamOperation( + "close_stream", + temporalio.api.streamservice.v1.CloseStreamRequest, + temporalio.api.streamservice.v1.CloseStreamResponse, + workflow_refusal=( + "Workflow code can't close a Stream on a Store in this release. " + "Call an Activity that closes it through its client." + ), + ), + "DescribeStream": _StreamOperation( + "describe_stream", + temporalio.api.streamservice.v1.DescribeStreamRequest, + temporalio.api.streamservice.v1.DescribeStreamResponse, + ), +} +"""The operations of the Stream service, by their Nexus operation name.""" + +_STORE_CALL_ACTIVITY = "__temporal_stream_store_call" + + +def _store_call_activity_name(store_name: str) -> str: + # The name carries the Store's name, so two Stores on one Worker don't clash. + return f"{_STORE_CALL_ACTIVITY}:{store_name}" + + +def _store_call_activity( + store: StreamStore, +) -> Callable[[str, bytes], Awaitable[bytes]]: + @temporalio.activity.defn(name=_store_call_activity_name(store.name)) + async def call(operation: str, request: bytes) -> bytes: + stream_operation = _STREAM_OPERATIONS[operation] + response = await stream_operation.call( + store, stream_operation.request_type.FromString(request) + ) + return response.SerializeToString() + + return call + + +@dataclasses.dataclass(frozen=True) +class StreamSelectContext: + """Context passed to :py:attr:`StreamConfig.store_selector` calls. + + .. warning:: + This API is experimental and subject to change. + """ + + id: str + """The id of the Stream.""" + + target: StorageDriverActivityInfo | StorageDriverWorkflowInfo | None = None + """The Workflow or Activity that makes the call, or ``None`` for a plain client. + + A call in an Activity has that Activity here. A Workflow's call runs the selector + in Workflow code, and has that Workflow here. + """ + + +@dataclasses.dataclass(frozen=True) +class StreamConfig: + """The Stream Stores of a client, and the rule that picks one for a new Stream. + + A client with this config sends each Stream operation to one of its Stores. + This release has no Streams on the server, so a client without this config + raises :py:class:`temporalio.exceptions.StreamConfigError` for a Stream + operation. It holds Stores and a selector, as + :py:class:`temporalio.converter.ExternalStorage` holds drivers and a selector. + + A ref carries the name of its Store, so a call on a known Stream goes to the + Store of that name. A call whose ref names no Store, such as a create, goes to + the one Store, or to the Store that :py:attr:`store_selector` returns. + + .. warning:: + This API is experimental and subject to change. + """ + + stores: Sequence[StreamStore] + """At least one Store. Each Store has a unique name. + + A ref carries the :py:attr:`StreamStore.name` of its Store, and a call looks + the Store up by that name. With more than one Store, :py:attr:`store_selector` + must also be set. A Worker serves every Store of its client. + """ + + store_selector: Callable[[StreamSelectContext], StreamStore] | None = None + """Picks the Store of a new Stream. Required with more than one Store. + + It gets a :py:class:`StreamSelectContext` for each call whose ref names no + Store, and returns one of the instances in :py:attr:`stores`. When it is + ``None`` and there is one Store, every new Stream goes to that Store. + + For a call of Workflow code, the selector runs in Workflow code, also on a + replay. So it must give the same Store for the same context. The Local Activity + of the chosen Store is in History, and another answer on a replay is a + non-determinism error. A Stream that code builds from an id only goes through the + selector on each call, so the selector must also give the same Store for the same + id. + """ + + workflow_call_options: StreamStoreCallOptions = dataclasses.field( + default_factory=StreamStoreCallOptions + ) + """Options of a Store's call for Workflow code, for every Store of this config. + + :py:meth:`StreamStore.call_from_workflow` gets them. The default runs its Local + Activity with their timeouts and retry policy. + """ + + _store_map: dict[str, StreamStore] = dataclasses.field( + init=False, repr=False, compare=False + ) + """Name-keyed index of :py:attr:`stores`, built at construction time.""" + + def __post_init__(self) -> None: + """Check the Stores and build the name-keyed map of them. + + Raises :py:class:`temporalio.exceptions.StreamConfigError` if there is no + Store, if two Stores share a name, or if more than one Store is given + without a :py:attr:`store_selector`. + """ + if not self.stores: + raise StreamConfigError("StreamConfig.stores must have at least one Store") + store_map: dict[str, StreamStore] = {} + for store in self.stores: + if store.name in store_map: + raise StreamConfigError( + f"StreamConfig.stores has more than one Store named '{store.name}'. " + "Each Store needs a unique name." + ) + store_map[store.name] = store + if len(self.stores) > 1 and self.store_selector is None: + raise StreamConfigError( + "StreamConfig.store_selector must be set when there is more than one Store" + ) + object.__setattr__(self, "_store_map", store_map) + + def _select_store(self, context: StreamSelectContext) -> StreamStore: + selector = self.store_selector + if selector is None: + return self.stores[0] + store = selector(context) + if ( + not isinstance(store, StreamStore) + or self._store_map.get(store.name) is not store + ): + raise StreamConfigError( + f"The store_selector returned {store!r}, which isn't one of " + "StreamConfig.stores" + ) + return store + + def _get_store_by_name(self, name: str) -> StreamStore: + store = self._store_map.get(name) + if store is None: + raise StreamConfigError( + f"The ref names the Stream Store '{name}', which this client doesn't have" + ) + return store + + def _store_for( + self, + ref: temporalio.nexus.system.stream_service.StreamRef[Any], + target: StorageDriverActivityInfo | StorageDriverWorkflowInfo | None, + ) -> StreamStore: + if ref.store: + return self._get_store_by_name(ref.store) + return self._select_store(StreamSelectContext(id=ref.id, target=target)) + + +def _request_ref( # pyright: ignore[reportUnusedFunction] + request: Any, +) -> temporalio.nexus.system.stream_service.StreamRef[Any]: + # A create calls its ref `stream_ref` and every other request calls it `ref`, so + # the field is found by its type. + if dataclasses.is_dataclass(request) and not isinstance(request, type): + for field in dataclasses.fields(request): + value = getattr(request, field.name) + if isinstance(value, temporalio.nexus.system.stream_service.StreamRef): + return value + raise TypeError(f"The request {type(request).__name__} has no StreamRef") + + +def _activity_target() -> StorageDriverActivityInfo | None: # pyright: ignore[reportUnusedFunction] + # External storage builds the same value from the running Activity. + if not temporalio.activity.in_activity(): + return None + info = temporalio.activity.info() + return StorageDriverActivityInfo( + id=info.activity_id, + type=info.activity_type, + run_id=info.activity_run_id, + namespace=info.namespace, + ) + + +def _refuse_store_plugins( # pyright: ignore[reportUnusedFunction] + plugins: Sequence[object], +) -> None: + for plugin in plugins: + if isinstance(plugin, StreamStore): + raise StreamConfigError( + f"The Stream Store '{plugin.name}' is in plugins. Give it in " + "stream_config=StreamConfig(stores=[...]) instead." + ) + + +def _check_not_negative(**durations: timedelta | None) -> None: # pyright: ignore[reportUnusedFunction] + for label, value in durations.items(): + if value is not None and value < timedelta(): + raise ValueError(f"{label} must not be negative") + + +def _stream_service() -> nexusrpc.ServiceDefinition: + service = nexusrpc.get_service_definition( + temporalio.nexus.system.stream_service.StreamService + ) + if service is None: + raise RuntimeError("The generated StreamService has no service definition") + return service + + +_STREAM_SERVICE = _stream_service() + + +async def _call_stream_store( # pyright: ignore[reportUnusedFunction] + store: StreamStore, + data_converter: temporalio.converter.DataConverter, + operation: str, + request: google.protobuf.message.Message, +) -> google.protobuf.message.Message: + stream_operation = _STREAM_OPERATIONS[operation] + if not data_converter.payload_codec: + return await stream_operation.call(store, request) + # The worker package imports the client, so a module level import would cycle. + from temporalio.bridge.worker import _Visitor + + visitor = temporalio.bridge._visitor.PayloadVisitor(skip_search_attributes=True) + await visitor.visit(_Visitor(data_converter._encode_payload_sequence), request) + response = await stream_operation.call(store, request) + await visitor.visit(_Visitor(data_converter._decode_payload_sequence), response) + return response diff --git a/temporalio/contrib/redis/__init__.py b/temporalio/contrib/redis/__init__.py new file mode 100644 index 0000000000..34dad945bf --- /dev/null +++ b/temporalio/contrib/redis/__init__.py @@ -0,0 +1 @@ +"""Redis integrations for the Temporal SDK.""" diff --git a/temporalio/contrib/redis/streams.py b/temporalio/contrib/redis/streams.py new file mode 100644 index 0000000000..389601406e --- /dev/null +++ b/temporalio/contrib/redis/streams.py @@ -0,0 +1,116 @@ +"""The Redis Stream Store. + +A client takes it in ``stream_config``, and its Workers use the same Store:: + + from temporalio.client import Client, StreamConfig + from temporalio.contrib.redis.streams import RedisStreamStore + + client = await Client.connect( + "localhost:7233", + stream_config=StreamConfig(stores=[RedisStreamStore("redis://localhost:6379")]), + ) + +.. warning:: + This module is experimental and subject to change. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from datetime import timedelta + +import temporalio.bridge.runtime +import temporalio.bridge.streams +import temporalio.client._stream_store +import temporalio.runtime + + +class RedisStreamStore(temporalio.client._stream_store._BridgeStreamStore): + """A Stream Store that keeps its Streams in Redis 7.0 or later. + + It runs the Stream runtime of the bridge on the Redis Store of the + ``temporalio-streams-redis`` crate. It connects at its first call, so building + it needs no server, and a server that is down fails that call. A failed call + raises :py:class:`temporalio.exceptions.StreamError`. + + A key has no namespace, so a Stream id is unique in one Redis and key prefix. + Give each namespace its own key prefix to keep their Streams apart. + + .. warning:: + This API is experimental and subject to change. + """ + + def __init__( + self, + url: str | Sequence[str], + *, + cluster: bool = False, + key_prefix: str | None = None, + retention: timedelta | None = None, + last_write_ttl: timedelta | None = None, + response_timeout: timedelta | None = None, + wait_replicas: int = 0, + wait_timeout: timedelta | None = None, + name: str = "redis", + runtime: temporalio.runtime.Runtime | None = None, + ) -> None: + """Create a Redis Store. It builds its Stream runtime and connects at its first call. + + A bad option raises ``ValueError`` at that first call, with the message of the + Redis crate. + + Args: + url: A ``redis://`` or ``rediss://`` 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. ``None`` takes ``temporal-streams``. + A Stream id is unique in one Redis and key prefix. + retention: The default Retention of a new Stream, which is how long the + Stream stays readable after its close. ``None`` takes 7 days. + last_write_ttl: How long an open Stream stays after its last write. + ``None`` keeps an open Stream and every Record until its close. + response_timeout: The longest time that one Redis call can take. ``None`` + takes 5 seconds. + wait_replicas: How many replicas must have each write before it answers, + with the Redis ``WAIT`` command. Zero sends no ``WAIT``. + wait_timeout: How long a write waits for its replicas, at most half of + ``response_timeout``. ``None`` takes 1 second. + name: Name of the Store, which the refs of its Streams carry, so that every + process can route a ref that it gets. + runtime: Runtime whose Tokio threads run the Store. ``None`` takes + :py:meth:`temporalio.runtime.Runtime.default` at the first call. + """ + temporalio.client._stream_store._check_not_negative( + retention=retention, + last_write_ttl=last_write_ttl, + response_timeout=response_timeout, + wait_timeout=wait_timeout, + ) + super().__init__(name=name, runtime=runtime) + self._urls = [url] if isinstance(url, str) else list(url) + self._cluster = cluster + self._key_prefix = key_prefix + self._retention = retention + self._last_write_ttl = last_write_ttl + self._response_timeout = response_timeout + self._wait_replicas = wait_replicas + self._wait_timeout = wait_timeout + + def _new_stream_runtime( + self, runtime: temporalio.bridge.runtime.Runtime + ) -> temporalio.bridge.streams.StreamRuntime: + return temporalio.bridge.streams.StreamRuntime.new_redis( + runtime, + urls=self._urls, + cluster=self._cluster, + key_prefix=self._key_prefix, + retention=self._retention, + last_write_ttl=self._last_write_ttl, + response_timeout=self._response_timeout, + name=self._name, + wait_replicas=self._wait_replicas, + wait_timeout=self._wait_timeout, + ) + + +__all__ = ["RedisStreamStore"] diff --git a/temporalio/exceptions.py b/temporalio/exceptions.py index 03a623c21e..63135cd9fc 100644 --- a/temporalio/exceptions.py +++ b/temporalio/exceptions.py @@ -8,6 +8,7 @@ import temporalio.api.enums.v1 import temporalio.api.failure.v1 +import temporalio.api.stream.v1 class TemporalError(Exception): @@ -198,6 +199,114 @@ def category(self) -> ApplicationErrorCategory: return self._category +_STREAM_ERROR_REASON_UNSPECIFIED = ( + temporalio.api.enums.v1.StreamErrorReason.STREAM_ERROR_REASON_UNSPECIFIED +) + +# The one reason that a retry can fix. +_STREAM_RETRYABLE_REASONS: frozenset[int] = frozenset( + {temporalio.api.enums.v1.StreamErrorReason.STREAM_ERROR_REASON_STORAGE_UNAVAILABLE} +) + + +class StreamError(ApplicationError): + """A failed Stream call. ``reason`` tells the cases apart. + + One class covers every failed call in client, Activity and Workflow code, so match + on ``reason`` and not on the type. The reason is a ``StreamErrorReason``, such as + ``STREAM_ERROR_REASON_NOT_FOUND`` (nobody created the Stream), + ``STREAM_ERROR_REASON_CLOSED`` (the Stream is closed), or + ``STREAM_ERROR_REASON_ALREADY_EXISTS`` (a strict create of a Stream that exists). + A Store outage, ``STREAM_ERROR_REASON_STORAGE_UNAVAILABLE``, is the one retryable + reason. + + .. warning:: + This API is experimental and subject to change. + """ + + def __init__( + self, + message: str, + *, + reason: temporalio.api.enums.v1.StreamErrorReason.ValueType | None = None, + ) -> None: + """Initialize a Stream error. ``reason`` defaults to the unspecified reason.""" + if reason is None: + reason = _STREAM_ERROR_REASON_UNSPECIFIED + # The details carry the reason, so the error keeps it when it crosses an + # Activity or a Nexus operation. + super().__init__( + message, + temporalio.api.stream.v1.StreamErrorDetails(reason=reason), + type="StreamError", + non_retryable=reason not in _STREAM_RETRYABLE_REASONS, + ) + self._reason = reason + + @property + def reason(self) -> temporalio.api.enums.v1.StreamErrorReason.ValueType: + """Why the Stream call failed.""" + return self._reason + + +def _stream_error_from_reason(reason: int, message: str) -> StreamError: + if reason not in temporalio.api.enums.v1.StreamErrorReason.values(): + # A reason that this SDK doesn't know gives the unspecified reason. + reason = _STREAM_ERROR_REASON_UNSPECIFIED + return StreamError( + message, reason=temporalio.api.enums.v1.StreamErrorReason.ValueType(reason) + ) + + +def _stream_error_from_call( # pyright: ignore[reportUnusedFunction] + message: str, details: bytes +) -> StreamError: + # A Store on the bridge fails with the encoded details, as a gRPC status has them. + proto = temporalio.api.stream.v1.StreamErrorDetails.FromString(details) + return _stream_error_from_reason(proto.reason, message) + + +def _stream_error_from_failure( # pyright: ignore[reportUnusedFunction] + err: BaseException, +) -> StreamError | None: + # A Local Activity or a Nexus operation hands back a plain application error, so the + # error is rebuilt from its type and details. + if isinstance(err, StreamError): + return err + if not isinstance(err, ApplicationError) or err.type != "StreamError": + return None + reason: int = _STREAM_ERROR_REASON_UNSPECIFIED + for detail in err.details: + if isinstance(detail, temporalio.api.stream.v1.StreamErrorDetails): + reason = detail.reason + break + return _stream_error_from_reason(reason, err.message) + + +class StreamConfigError(RuntimeError): + """A Stream call can't find its Store, because of how the Stores are set up. + + It is not a :py:class:`StreamError`. A Store didn't fail. The setup of the Stores + is wrong, and a retry won't help. These are its cases. + + * A Stream call is made on a client with no Stream config. This release has no + Streams on the server, so every Stream call needs a Store. + * A Workflow's Stream call is made on a Worker whose client has no Stream config. + This fails the Workflow Task. + * A :py:class:`temporalio.client.StreamConfig` has no Store, has two Stores with + one name, or has more than one Store and no selector. + * A ref names a Store that the config of the client doesn't have. + * The selector returns a Store that isn't one of the config's Stores. + * A :py:class:`temporalio.client.StreamStore` is given in ``plugins``. + * A Workflow's Stream call has a ref that names a Store that the config of the + Worker's client doesn't have. This fails the Workflow Task. + * This build of temporalio has no Stream runtime, or no Redis Stream Store. + + .. warning:: + This API is experimental and subject to change. + """ + + class CancelledError(FailureError): """Error raised on workflow/activity cancellation.""" diff --git a/temporalio/nexus/system/__init__.py b/temporalio/nexus/system/__init__.py index 1f8f7c6d29..d3096b1a00 100644 --- a/temporalio/nexus/system/__init__.py +++ b/temporalio/nexus/system/__init__.py @@ -19,6 +19,7 @@ from temporalio.bridge._visitor_functions import VisitorFunctions from temporalio.converter import BinaryProtoPayloadConverter, CompositePayloadConverter from temporalio.converter._payload_converter import ( + _get_transfer_type_converter, _TemporalTransferTypePayloadConverter, ) @@ -123,6 +124,27 @@ def from_payloads( with _user_converter_context(self._user_converters): return self._outer_payload_converter.from_payloads(payloads, type_hints) + def _to_transfer_value(self, value: Any) -> Any: + """Convert a system Nexus model to its proto message, with no payload in between.""" + converter = _get_transfer_type_converter(type(value)) + if converter is None: + raise TypeError(f"{type(value).__name__} has no transfer type") + with _user_converter_context(self._user_converters): + return converter.to_transfer_type(value) + + def _from_transfer_value(self, value: Any, type_hint: type | None) -> Any: + """Convert a proto message to the system Nexus model of ``type_hint``. + + With no ``type_hint``, or one that has no transfer type, it gives the message. + """ + converter = ( + _get_transfer_type_converter(type_hint) if type_hint is not None else None + ) + if converter is None or type_hint is None: + return value + with _user_converter_context(self._user_converters): + return converter.from_transfer_type(value, type_hint) + def is_system_endpoint(endpoint: str) -> bool: """Return whether a Nexus endpoint is the Temporal system endpoint. @@ -171,7 +193,7 @@ async def maybe_visit_payload( def _get_payload_converter( # pyright: ignore[reportUnusedFunction] user_payload_converter: temporalio.converter.PayloadConverter, user_failure_converter: temporalio.converter.FailureConverter, -) -> temporalio.converter.PayloadConverter: +) -> _SystemNexusPayloadConverter: """Return the fixed payload converter for system Nexus outer envelopes.""" return _SystemNexusPayloadConverter(user_payload_converter, user_failure_converter) @@ -182,9 +204,14 @@ def _get_serialization_context( # pyright: ignore[reportUnusedFunction] request: Any, ) -> temporalio.converter.SerializationContext | None: """Return the serialization context for a system Nexus operation.""" + from .stream_service.workflow import ( + __nexus_operation_registry__ as stream_operation_registry, + ) from .workflow_service import __nexus_operation_registry__ - operation_info = __nexus_operation_registry__.get((service, operation)) + operation_info = __nexus_operation_registry__.get( + (service, operation) + ) or stream_operation_registry.get((service, operation)) if operation_info is None or operation_info.serialization_context is None: return None return operation_info.serialization_context(request) @@ -194,3 +221,10 @@ def _get_serialization_context( # pyright: ignore[reportUnusedFunction] "TEMPORAL_SYSTEM_ENDPOINT", "is_system_endpoint", ] + +# The generated Workflow modules of `stream_service` import `temporalio.workflow` at their top, +# and `temporalio.workflow` imports `create_stream` and `get_stream` from them. So this package +# loads `temporalio.workflow` as its last statement, when every name above exists. Then +# `temporalio.workflow` is complete before a generated Workflow module starts, whichever module +# is the first import. +import temporalio.workflow # noqa: E402 # pyright: ignore[reportUnusedImport] diff --git a/temporalio/nexus/system/_support.py b/temporalio/nexus/system/_support.py index 15a8e80102..55f53b59a1 100644 --- a/temporalio/nexus/system/_support.py +++ b/temporalio/nexus/system/_support.py @@ -1,9 +1,15 @@ import collections.abc +import dataclasses +import functools import typing -from datetime import timedelta +import urllib.parse +import uuid +from datetime import datetime, timedelta import google.protobuf.duration_pb2 +import google.protobuf.timestamp_pb2 +import temporalio.activity import temporalio.api.common.v1.message_pb2 as common_pb2 import temporalio.api.enums.v1.workflow_pb2 as workflow_enums_pb2 import temporalio.api.failure.v1.message_pb2 as failure_pb2 @@ -13,12 +19,157 @@ import temporalio.converter import temporalio.nexus.system +if typing.TYPE_CHECKING: + from temporalio.nexus.system.stream_service.models import ( + Producer, + ReadRecordsRequest, + ReadRecordsResponse, + StreamRecord, + ) + +_BodyT = typing.TypeVar("_BodyT") + class SignalWithStartWorkflowModelRequest(typing.Protocol): namespace: str id: str +def request_id() -> str: + import temporalio.workflow + + # Each run of a chain sends the same create, and the Store answers a create with the + # request id that made the Stream as a success. One id for the chain also replays. + if temporalio.workflow.in_workflow(): + return temporalio.workflow.info().first_execution_run_id + return str(uuid.uuid4()) + + +def _producer_id(*parts: str) -> str: + # Quoting keeps a ':' inside an id from shifting the parts. + return ":".join(urllib.parse.quote(part, safe="") for part in parts) + + +# One id for the process, because nothing fences a client write. +@functools.cache +def _client_producer_id() -> str: + return _producer_id("client", str(uuid.uuid4())) + + +def _timestamp(value: datetime) -> google.protobuf.timestamp_pb2.Timestamp: + proto = google.protobuf.timestamp_pb2.Timestamp() + proto.FromDatetime(value) + return proto + + +def default_producer() -> "Producer": + from temporalio.nexus.system.stream_service.models import ( + ActivityEpoch, + ActivityProducer, + ClientProducer, + LocalActivityEpoch, + LocalActivityProducer, + Producer, + ProducerKindClient, + ProducerKindLocalActivity, + ProducerKindStandaloneActivity, + ProducerKindWorkflowActivity, + Timestamp, + WorkflowTaskEpoch, + ) + + if not temporalio.activity.in_activity(): + return Producer( + id=_client_producer_id(), kind=ProducerKindClient(ClientProducer()) + ) + info = temporalio.activity.info() + if info.is_local: + task = info.workflow_task + if task is None or info.workflow_run_id is None: + raise RuntimeError( + "A Stream write from a Local Activity needs its Workflow Task, and the " + "Worker's Core gave none" + ) + started_time = None + if task.started_time is not None: + proto_time = _timestamp(task.started_time) + started_time = Timestamp(seconds=proto_time.seconds, nanos=proto_time.nanos) + # The activity id of a Local Activity is its sequence number in the run, so a + # retried Workflow Task writes as the same Producer with a higher epoch. + return Producer( + id=_producer_id("la", info.workflow_run_id, info.activity_id), + kind=ProducerKindLocalActivity( + LocalActivityProducer( + run_id=info.workflow_run_id, + epoch=LocalActivityEpoch( + task=WorkflowTaskEpoch( + started_event_id=task.started_event_id, + attempt=task.attempt, + started_time=started_time, + ), + attempt=info.attempt, + ), + ) + ), + ) + attempt_time = _timestamp(info.current_attempt_scheduled_time) + epoch = ActivityEpoch( + attempt_scheduled_time=Timestamp( + seconds=attempt_time.seconds, nanos=attempt_time.nanos + ), + attempt=info.attempt, + ) + if info.workflow_run_id is not None: + # The scheduled time stays the same across attempts, and it tells apart two uses + # of one activity id in a run. + schedule_key = str(_timestamp(info.scheduled_time).ToNanoseconds()) + return Producer( + id=_producer_id( + "act", info.workflow_run_id, info.activity_id, schedule_key + ), + kind=ProducerKindWorkflowActivity( + ActivityProducer(run_id=info.workflow_run_id, epoch=epoch) + ), + ) + if info.activity_run_id is None: + raise RuntimeError("A Standalone Activity has no activity run id") + return Producer( + id=_producer_id("sa", info.activity_id, info.activity_run_id), + kind=ProducerKindStandaloneActivity(ActivityProducer(epoch=epoch)), + ) + + +def workflow_producer() -> "Producer": + import temporalio.workflow + from temporalio.nexus.system.stream_service.models import ( + Producer, + ProducerKindWorkflow, + WorkflowProducer, + ) + + run_id = temporalio.workflow.info().run_id + return Producer( + id=_producer_id("wf", run_id), + kind=ProducerKindWorkflow(WorkflowProducer(run_id=run_id)), + ) + + +def next_read_request( + request: "ReadRecordsRequest[_BodyT]", response: "ReadRecordsResponse[_BodyT]" +) -> "ReadRecordsRequest[_BodyT] | None": + if response.done or (request.wait is None and not response.records): + return None + if response.last_offset is None: + return request + return dataclasses.replace(request, from_offset=response.last_offset) + + +def resume_read_request( + request: "ReadRecordsRequest[_BodyT]", item: "StreamRecord[_BodyT]" +) -> "ReadRecordsRequest[_BodyT]": + return dataclasses.replace(request, from_offset=item.offset) + + def retry_policy_from_proto( proto: common_pb2.RetryPolicy, ) -> temporalio.common.RetryPolicy: diff --git a/temporalio/testing/__init__.py b/temporalio/testing/__init__.py index 0cd652d86b..5574a0a39b 100644 --- a/temporalio/testing/__init__.py +++ b/temporalio/testing/__init__.py @@ -1,11 +1,13 @@ """Test framework for workflows and activities.""" from ._activity import ActivityEnvironment +from ._stream_store import MemoryStreamStore from ._timeskipping import TimeSkipper, TimeSkippingConfig from ._workflow import WorkflowEnvironment __all__ = [ "ActivityEnvironment", + "MemoryStreamStore", "TimeSkipper", "WorkflowEnvironment", "TimeSkippingConfig", diff --git a/temporalio/testing/_stream_store.py b/temporalio/testing/_stream_store.py new file mode 100644 index 0000000000..86ba0ac16d --- /dev/null +++ b/temporalio/testing/_stream_store.py @@ -0,0 +1,57 @@ +"""A Stream Store in the memory of the process, for tests.""" + +from __future__ import annotations + +from datetime import timedelta + +import temporalio.bridge.runtime +import temporalio.bridge.streams +import temporalio.client._stream_store +import temporalio.runtime + + +class MemoryStreamStore(temporalio.client._stream_store._BridgeStreamStore): + """A Stream Store that keeps its Streams in the memory of this process, for tests. + + It runs the same Stream runtime and rules as every other Store on the bridge. A + failed call raises :py:class:`temporalio.exceptions.StreamError`. + + .. warning:: + This API is experimental and subject to change. + """ + + def __init__( + self, + *, + name: str = "memory", + retention: timedelta | None = None, + last_write_ttl: timedelta | None = None, + runtime: temporalio.runtime.Runtime | None = None, + ) -> None: + """Create a memory Store. It builds its Stream runtime at its first call. + + Args: + name: Name of the Store, which the refs of its Streams carry. + retention: How long a closed Stream stays. ``None`` takes the default of + the Stream runtime. + last_write_ttl: How long an open Stream stays after its last write. + ``None`` keeps an open Stream until its close. + runtime: Runtime whose Tokio threads run the Store. ``None`` takes + :py:meth:`temporalio.runtime.Runtime.default` at the first call. + """ + temporalio.client._stream_store._check_not_negative( + retention=retention, last_write_ttl=last_write_ttl + ) + super().__init__(name=name, runtime=runtime) + self._retention = retention + self._last_write_ttl = last_write_ttl + + def _new_stream_runtime( + self, runtime: temporalio.bridge.runtime.Runtime + ) -> temporalio.bridge.streams.StreamRuntime: + return temporalio.bridge.streams.StreamRuntime.new_memory( + runtime, + name=self._name, + retention=self._retention, + last_write_ttl=self._last_write_ttl, + ) diff --git a/temporalio/testing/_workflow.py b/temporalio/testing/_workflow.py index 6c8354162c..54add4e363 100644 --- a/temporalio/testing/_workflow.py +++ b/temporalio/testing/_workflow.py @@ -102,6 +102,7 @@ async def start_local( dev_server_download_ttl: timedelta | None = None, dev_server_startup_timeout: timedelta | None = None, ui_port: int | None = None, + stream_config: temporalio.client.StreamConfig | None = None, ) -> WorkflowEnvironment: """Start a full Temporal server locally, downloading if necessary. @@ -165,6 +166,8 @@ async def start_local( become ready after spawning. Defaults to 5 seconds and does not include download time. ui_port: UI port to use if UI is enabled. + stream_config: See parameter of the same name on + :py:meth:`temporalio.client.Client.connect`. Returns: The started CLI dev server workflow environment. @@ -194,6 +197,7 @@ async def start_local( dev_server_download_ttl=dev_server_download_ttl, dev_server_startup_timeout=dev_server_startup_timeout, ui_port=ui_port, + stream_config=stream_config, ) @classmethod @@ -216,6 +220,7 @@ async def start_time_skipping( test_server_extra_args: Sequence[str] = [], test_server_download_ttl: timedelta | None = None, test_server_startup_timeout: timedelta | None = None, + stream_config: temporalio.client.StreamConfig | None = None, ) -> WorkflowEnvironment: """Start a V1 time skipping workflow environment. @@ -276,6 +281,8 @@ async def start_time_skipping( test_server_startup_timeout: Maximum time to wait for the test server to become ready after spawning. Defaults to 5 seconds and does not include download time. + stream_config: See parameter of the same name on + :py:meth:`temporalio.client.Client.connect`. Returns: The started workflow environment with time skipping. @@ -296,6 +303,7 @@ async def start_time_skipping( test_server_extra_args=test_server_extra_args, test_server_download_ttl=test_server_download_ttl, test_server_startup_timeout=test_server_startup_timeout, + stream_config=stream_config, ) @classmethod @@ -582,6 +590,7 @@ async def _bootstrap_dev_server( dev_server_download_ttl: timedelta | None = None, dev_server_startup_timeout: timedelta | None = None, ui_port: int | None = None, + stream_config: temporalio.client.StreamConfig | None = None, ) -> tuple[ temporalio.bridge.testing.EphemeralServer, temporalio.client.Client, @@ -649,6 +658,7 @@ async def _bootstrap_dev_server( rpc_metadata=rpc_metadata, identity=identity, runtime=runtime, + stream_config=stream_config, ) except: try: @@ -719,6 +729,7 @@ async def _create( test_server_extra_args: Sequence[str], test_server_download_ttl: timedelta | None, test_server_startup_timeout: timedelta | None, + stream_config: temporalio.client.StreamConfig | None, ) -> _V1WorkflowEnvironment: runtime = runtime or temporalio.runtime.Runtime.default() download_ttl_ms = None @@ -753,6 +764,7 @@ async def _create( rpc_metadata=rpc_metadata, identity=identity, runtime=runtime, + stream_config=stream_config, ) except: try: diff --git a/temporalio/worker/_interceptor.py b/temporalio/worker/_interceptor.py index 6196f5efe1..9cfab54022 100644 --- a/temporalio/worker/_interceptor.py +++ b/temporalio/worker/_interceptor.py @@ -21,6 +21,9 @@ import temporalio.nexus import temporalio.nexus._util import temporalio.workflow +from temporalio.nexus.system.stream_service.workflow._system_nexus_interceptor import ( + _SystemNexusWorkflowOutboundInterceptorBase as _StreamSystemNexusInterceptorBase, +) from temporalio.nexus.system.workflow_service._system_nexus_interceptor import ( _SystemNexusWorkflowOutboundInterceptorBase, ) @@ -455,7 +458,9 @@ async def handle_update_handler(self, input: HandleUpdateInput) -> Any: return await self.next.handle_update_handler(input) -class WorkflowOutboundInterceptor(_SystemNexusWorkflowOutboundInterceptorBase): +class WorkflowOutboundInterceptor( + _SystemNexusWorkflowOutboundInterceptorBase, _StreamSystemNexusInterceptorBase +): """Outbound interceptor to wrap calls made from within workflows. This should be extended by any workflow outbound interceptors. @@ -472,9 +477,7 @@ def __init__(self, next: WorkflowOutboundInterceptor) -> None: """ self.next = next - def _next_system_nexus_interceptor( - self, - ) -> _SystemNexusWorkflowOutboundInterceptorBase: + def _next_system_nexus_interceptor(self) -> WorkflowOutboundInterceptor: return self.next def continue_as_new(self, input: ContinueAsNewInput) -> NoReturn: diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 62a27b29cf..540e9d4052 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -22,6 +22,7 @@ from ..common import HeaderCodecBehavior from ._interceptor import Interceptor +from ._stream_store import _WorkflowStreamStore from ._worker import load_default_build_id from ._workflow import _WorkflowWorker from ._workflow_instance import ( @@ -56,6 +57,7 @@ def __init__( runtime: temporalio.runtime.Runtime | None = None, disable_safe_workflow_eviction: bool = False, header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + stream_config: temporalio.client.StreamConfig | None = None, ) -> None: """Create a replayer to replay workflows from history. @@ -64,6 +66,11 @@ def __init__( the replayer that were passed to the worker when the workflow originally ran. + ``stream_config`` takes the Stream config of the client that the + worker had, so a replay makes the same commands. A history with a Stream + call needs it, since a Stream call with no Stream config fails the + Workflow Task. A replay calls no method of its Stores. + Note, unlike the worker, for the replayer the workflow_task_executor will default to a new thread pool executor with no max_workers set that will be shared across all replay calls and never explicitly shut down. @@ -86,6 +93,7 @@ def __init__( runtime=runtime, disable_safe_workflow_eviction=disable_safe_workflow_eviction, header_codec_behavior=header_codec_behavior, + stream_config=stream_config, ) self._initial_config = self._config.copy() self._default_workflow_logic_flags = set(_DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS) @@ -287,6 +295,9 @@ def on_eviction_hook( default_workflow_logic_flags=frozenset( self._default_workflow_logic_flags ), + stream_store=_WorkflowStreamStore.from_config( + self._config.get("stream_config") + ), ) external_storage = data_converter.external_storage storage_driver_types = ( @@ -435,6 +446,7 @@ class ReplayerConfig(TypedDict, total=False): runtime: temporalio.runtime.Runtime | None disable_safe_workflow_eviction: bool header_codec_behavior: HeaderCodecBehavior + stream_config: temporalio.client.StreamConfig | None @dataclass(frozen=True) diff --git a/temporalio/worker/_stream_store.py b/temporalio/worker/_stream_store.py new file mode 100644 index 0000000000..da9faa689c --- /dev/null +++ b/temporalio/worker/_stream_store.py @@ -0,0 +1,121 @@ +"""How Workflow code reaches the Stream Stores of its Worker's client.""" + +from __future__ import annotations + +import asyncio +import dataclasses +from collections.abc import Generator +from typing import TYPE_CHECKING, Any + +from nexusrpc import OutputT + +import temporalio.exceptions +import temporalio.workflow +from temporalio.client._stream_store import ( + _STREAM_OPERATIONS, + StreamConfig, + _request_ref, +) +from temporalio.converter import StorageDriverWorkflowInfo + +from ._interceptor import StartNexusOperationInput + +if TYPE_CHECKING: + from ._workflow_instance import _WorkflowInstanceImpl + + +def _workflow_target() -> StorageDriverWorkflowInfo: + # External storage builds the same value for a Workflow. + info = temporalio.workflow.info() + return StorageDriverWorkflowInfo( + id=info.workflow_id, + type=info.workflow_type, + run_id=info.run_id, + namespace=info.namespace, + ) + + +@dataclasses.dataclass(frozen=True) +class _WorkflowStreamStore: + """The Worker's Stream Stores as Workflow code reaches them.""" + + config: StreamConfig + """The Stream config of the Worker's client.""" + + @classmethod + def from_config(cls, config: StreamConfig | None) -> _WorkflowStreamStore | None: + """Build it from the Stream config of a client, or give ``None`` when it has none.""" + return cls(config=config) if config is not None else None + + def start_operation( + self, + instance: _WorkflowInstanceImpl, + input: StartNexusOperationInput[Any, OutputT], + ) -> temporalio.workflow.NexusOperationHandle[OutputT]: + """Answer a Stream call of Workflow code with the Store that the call goes to.""" + stream_operation = _STREAM_OPERATIONS[input.operation_name] + # The Workflow `Stream` has no `append`, `close` or `read`, but the generated + # Workflow service still has their operations, so Workflow code can start one + # through a Nexus client. In this release, Workflow code writes and reads a Stream + # on a Store through its own Activities. + if stream_operation.workflow_refusal is not None: + raise temporalio.exceptions.ApplicationError( + stream_operation.workflow_refusal, non_retryable=True + ) + payload_converter, _ = instance._system_nexus_converters(input) + request = payload_converter._to_transfer_value(input.input) + # The choice runs in Workflow code, so a ref of a Store that the config doesn't + # have raises `StreamConfigError` here, with no I/O, and fails the Workflow Task. + store = self.config._store_for(_request_ref(input.input), _workflow_target()) + options = self.config.workflow_call_options + + async def result() -> OutputT: + try: + response = await store.call_from_workflow( + input.operation_name, request, options + ) + except temporalio.exceptions.ApplicationError as err: + # Core hands back a failed Local Activity's error without a wrapper. + stream_error = temporalio.exceptions._stream_error_from_failure(err) + if stream_error is not None: + raise stream_error from None + raise + except temporalio.exceptions.ActivityError as err: + # Workflow code gets the Store's error, not the Local Activity's wrapper. + if err.cause is not None: + stream_error = temporalio.exceptions._stream_error_from_failure( + err.cause + ) + if stream_error is not None: + raise stream_error from None + if isinstance( + err.cause, + ( + temporalio.exceptions.ApplicationError, + temporalio.exceptions.TimeoutError, + ), + ): + raise err.cause + raise + return payload_converter._from_transfer_value( + stream_operation.check_response(store, response), input.output_type + ) + + return _StreamStoreOperationHandle( + instance.create_task(result(), name=f"stream store: {input.operation_name}") + ) + + +class _StreamStoreOperationHandle(temporalio.workflow.NexusOperationHandle[OutputT]): + def __init__(self, task: asyncio.Task[OutputT]) -> None: + self._task = task + + def cancel(self) -> bool: + return self._task.cancel() + + def __await__(self) -> Generator[Any, Any, OutputT]: + return self._task.__await__() + + @property + def operation_token(self) -> str | None: + return None diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index e06bb99a4b..bbf68e2bf4 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -34,6 +34,7 @@ from ._interceptor import Interceptor from ._nexus import _NexusWorker from ._plugin import Plugin +from ._stream_store import _WorkflowStreamStore from ._tuning import WorkerTuner from ._workflow import ( _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, @@ -498,7 +499,11 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf or bridge_client.config.runtime or temporalio.runtime.Runtime.default() ) - activities = config.get("activities") + activities = list(config.get("activities") or []) + stream_config = client_config.get("stream_config") + if stream_config is not None and config.get("workflows"): + for store in stream_config.stores: + activities.extend(store.worker_activities()) if activities: _warn_if_activity_executor_max_workers_is_inconsistent(config) self._activity_worker = _ActivityWorker( @@ -578,6 +583,7 @@ def check_activity(activity: str): encode_headers=client_config["header_codec_behavior"] != HeaderCodecBehavior.NO_CODEC, max_workflow_task_external_storage_concurrency=max_workflow_task_external_storage_concurrency, + stream_store=_WorkflowStreamStore.from_config(stream_config), ) tuner = config.get("tuner") @@ -667,7 +673,9 @@ def check_activity(activity: str): enable_workflows=self._workflow_worker is not None, enable_local_activities=self._activity_worker is not None and self._workflow_worker is not None, - enable_remote_activities=self._activity_worker is not None + # The Activities of the Stream Stores run as Local Activities, so + # alone they don't make this worker poll for Activity Tasks. + enable_remote_activities=bool(config.get("activities")) and not config.get("no_remote_activities"), enable_nexus=self._nexus_worker is not None, ), @@ -751,13 +759,19 @@ def client(self, value: temporalio.client.Client) -> None: Changing the client will make sure the worker starts using it for the next calls it makes. However, outstanding client calls will still complete with the existing client. The new client cannot be "lazy" and - must be using the same runtime as the current client. + must be using the same runtime and Stream Stores as the current client. """ bridge_client = _extract_bridge_client_for_worker(value) if self._runtime is not bridge_client._runtime: raise ValueError( "New client is not on the same runtime as the existing client" ) + # Running Workflows route their Stream calls by the Store that the worker + # started with, so a new client can't change it. + if _stream_config_stores(value) != _stream_config_stores(self.client): + raise ValueError( + "New client doesn't have the Stream Stores of the existing client" + ) assert bridge_client._bridge_client self._bridge_worker.replace_client(bridge_client._bridge_client) self._config["client"] = value @@ -1013,6 +1027,13 @@ async def __aexit__(self, exc_type: type[BaseException] | None, *args: Any) -> N raise self._async_context_run_exception +def _stream_config_stores( + client: temporalio.client.Client, +) -> list[temporalio.client.StreamStore]: + config = client.config(active_config=True).get("stream_config") + return list(config.stores) if config is not None else [] + + class WorkerConfig(TypedDict, total=False): """TypedDict of config originally passed to :py:class:`Worker`.""" diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index d8380a6b90..3e20e6adbb 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -40,6 +40,7 @@ WorkflowInboundInterceptor, WorkflowInterceptorClassInput, ) +from ._stream_store import _WorkflowStreamStore from ._workflow_instance import ( _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS, PatchActivationInput, @@ -128,6 +129,7 @@ def __init__( encode_headers: bool, max_workflow_task_external_storage_concurrency: int, default_workflow_logic_flags: frozenset[_WorkflowLogicFlag] | None = None, + stream_store: _WorkflowStreamStore | None = None, ) -> None: # Debug mode is enabled if specified or if the TEMPORAL_DEBUG env var is truthy debug_mode = debug_mode or bool(os.environ.get("TEMPORAL_DEBUG")) @@ -197,6 +199,7 @@ def __init__( self._max_workflow_task_external_storage_concurrency = ( max_workflow_task_external_storage_concurrency ) + self._stream_store = stream_store self._throw_after_activation: Exception | None = None # If debug mode is enabled, disable deadlock detection @@ -819,6 +822,7 @@ def _create_workflow_instance( nexus_serialization_context=temporalio.converter._serialization_context._propagated_nexus_serialization_context( init ), + stream_store=self._stream_store, ) if defn.sandboxed: return self._workflow_runner.create_instance(det) diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 03463e2c03..09f1ae53fc 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -59,7 +59,14 @@ import temporalio.exceptions import temporalio.nexus.system import temporalio.workflow +from temporalio.client._stream_store import _STREAM_SERVICE from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo +from temporalio.nexus.system.stream_service.workflow._system_nexus_interceptor import ( + _start_system_nexus_operation as _start_stream_system_nexus_operation, +) +from temporalio.nexus.system.stream_service.workflow._system_nexus_interceptor import ( + _SystemNexusWorkflowOutboundInterceptorTerminal as _StreamSystemNexusInterceptorTerminal, +) from temporalio.nexus.system.workflow_service._system_nexus_interceptor import ( _start_system_nexus_operation, _SystemNexusWorkflowOutboundInterceptorTerminal, @@ -88,6 +95,7 @@ WorkflowInboundInterceptor, WorkflowOutboundInterceptor, ) +from ._stream_store import _WorkflowStreamStore logger = logging.getLogger(__name__) @@ -185,6 +193,10 @@ class WorkflowInstanceDetails: temporalio.converter.NexusSerializationContext | None ) = None """Serialization context of the Nexus operation handler that started the workflow, if any.""" + stream_store: _WorkflowStreamStore | None = None + """The Worker's Stream Stores as Workflow code reaches them, or ``None`` when the + Worker's client has none. This release has no Streams on the server, so with ``None`` + a Stream call fails the Workflow Task.""" class WorkflowInstance(ABC): @@ -301,6 +313,7 @@ def __init__(self, det: WorkflowInstanceDetails) -> None: # the Nexus caller, which decodes them with the handler's Nexus context. A retried or cron run # reads the previous run's result and failure with the same context. self._nexus_serialization_context = det.nexus_serialization_context + self._stream_store = det.stream_store if self._nexus_serialization_context: self._result_payload_converter = self._payload_converter_with_context( self._nexus_serialization_context @@ -1830,23 +1843,25 @@ async def workflow_start_nexus_operation( event_groups: Sequence[temporalio.workflow.EventGroup] | None = None, ) -> temporalio.workflow.NexusOperationHandle[OutputT]: if temporalio.nexus.system.is_system_endpoint(endpoint): - return await _start_system_nexus_operation( - self._outbound, - StartNexusOperationInput( - endpoint=temporalio.nexus.system.TEMPORAL_SYSTEM_ENDPOINT, - service=service, - operation=operation, - input=input, - output_type=output_type, - schedule_to_close_timeout=schedule_to_close_timeout, - schedule_to_start_timeout=schedule_to_start_timeout, - start_to_close_timeout=start_to_close_timeout, - cancellation_type=cancellation_type, - headers=None, - summary=summary, - event_groups=event_groups, - ), + system_input = StartNexusOperationInput( + endpoint=temporalio.nexus.system.TEMPORAL_SYSTEM_ENDPOINT, + service=service, + operation=operation, + input=input, + output_type=output_type, + schedule_to_close_timeout=schedule_to_close_timeout, + schedule_to_start_timeout=schedule_to_start_timeout, + start_to_close_timeout=start_to_close_timeout, + cancellation_type=cancellation_type, + headers=None, + summary=summary, + event_groups=event_groups, ) + if service == _STREAM_SERVICE.name: + return await _start_stream_system_nexus_operation( + self._outbound, system_input + ) + return await _start_system_nexus_operation(self._outbound, system_input) return await self._outbound.start_nexus_operation( StartNexusOperationInput( endpoint=endpoint, @@ -2288,6 +2303,7 @@ async def _outbound_start_nexus_operation( # examples of the evolution of the resulting handle state machine in the sync and async # Nexus response cases. handle: _NexusOperationHandle[OutputT] + payload_converter: temporalio.converter.PayloadConverter async def operation_handle_fn() -> OutputT: return cast( @@ -2299,25 +2315,7 @@ async def operation_handle_fn() -> OutputT: ) if temporalio.nexus.system.is_system_endpoint(input.endpoint): - serialization_context = temporalio.nexus.system._get_serialization_context( - input.service, - input.operation_name, - input.input, - ) - user_payload_converter = self._context_free_payload_converter - user_failure_converter = self._context_free_failure_converter - if serialization_context is not None: - user_payload_converter = self._payload_converter_with_context( - serialization_context - ) - user_failure_converter = self._failure_converter_with_context( - serialization_context - ) - payload_converter = temporalio.nexus.system._get_payload_converter( - user_payload_converter, - user_failure_converter, - ) - failure_converter = user_failure_converter + payload_converter, failure_converter = self._system_nexus_converters(input) else: serialization_context = temporalio.converter.NexusSerializationContext( endpoint=input.endpoint, @@ -2355,9 +2353,46 @@ async def _intercept_system_nexus_operation( async def _schedule_system_nexus_operation( self, input: StartNexusOperationInput[Any, OutputT] - ) -> _NexusOperationHandle[OutputT]: + ) -> temporalio.workflow.NexusOperationHandle[OutputT]: + # This release has no Streams on the server, so a Worker answers a Stream call with + # a Store of its client. On a Worker whose client has none, the `StreamConfigError` + # fails the Workflow Task. Every other system Nexus operation goes to the server. + if input.service == _STREAM_SERVICE.name: + store = self._stream_store + if store is None: + raise temporalio.exceptions.StreamConfigError( + "The client of this Worker has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...])." + ) + return store.start_operation(self, input) return await self._outbound_start_nexus_operation(input) + def _system_nexus_converters( + self, input: StartNexusOperationInput[Any, Any] + ) -> tuple[ + temporalio.nexus.system._SystemNexusPayloadConverter, + temporalio.converter.FailureConverter, + ]: + serialization_context = temporalio.nexus.system._get_serialization_context( + input.service, + input.operation_name, + input.input, + ) + user_payload_converter = self._context_free_payload_converter + user_failure_converter = self._context_free_failure_converter + if serialization_context is not None: + user_payload_converter = self._payload_converter_with_context( + serialization_context + ) + user_failure_converter = self._failure_converter_with_context( + serialization_context + ) + payload_converter = temporalio.nexus.system._get_payload_converter( + user_payload_converter, + user_failure_converter, + ) + return payload_converter, user_failure_converter + #### Miscellaneous helpers #### # These are in alphabetical order. @@ -3360,7 +3395,9 @@ async def handle_update_handler(self, input: HandleUpdateInput) -> Any: class _WorkflowOutboundImpl( - _SystemNexusWorkflowOutboundInterceptorTerminal, WorkflowOutboundInterceptor + _SystemNexusWorkflowOutboundInterceptorTerminal, + _StreamSystemNexusInterceptorTerminal, + WorkflowOutboundInterceptor, ): def __init__(self, instance: _WorkflowInstanceImpl) -> None: # type: ignore # We are intentionally not calling the base class's __init__ here @@ -3407,7 +3444,7 @@ async def _intercept_system_nexus_operation( async def start_system_nexus_operation( self, input: StartNexusOperationInput[Any, OutputT] - ) -> _NexusOperationHandle[OutputT]: + ) -> temporalio.workflow.NexusOperationHandle[OutputT]: return await self._instance._schedule_system_nexus_operation(input) def start_local_activity( diff --git a/tests/contrib/redis/__init__.py b/tests/contrib/redis/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/contrib/redis/test_streams.py b/tests/contrib/redis/test_streams.py new file mode 100644 index 0000000000..0bae0921f1 --- /dev/null +++ b/tests/contrib/redis/test_streams.py @@ -0,0 +1,243 @@ +from __future__ import annotations + +import asyncio +import os +import socket +import uuid +from collections.abc import Awaitable +from datetime import timedelta + +import pytest + +from temporalio import activity, workflow +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.client import Client, StreamConfig, StreamStore +from temporalio.contrib.redis.streams import RedisStreamStore +from temporalio.exceptions import StreamError +from temporalio.nexus.system.stream_service import ( + ClientProducer, + Producer, + ProducerKindClient, + ProducerKindLocalActivity, + ProducerKindWorkflowActivity, + StreamRef, + StreamStatus, +) +from temporalio.nexus.system.stream_service.client import Stream +from temporalio.nexus.system.stream_service.workflow import create_stream +from tests.helpers import new_worker + +_REDIS_URL = os.environ.get("STREAMS_REDIS_URL") + +pytestmark = pytest.mark.skipif( + not _REDIS_URL, reason="STREAMS_REDIS_URL is not set to a Redis server" +) + + +def _store() -> RedisStreamStore: + assert _REDIS_URL + # Each test has its own key prefix, so that runs on one Redis don't share Streams. + return RedisStreamStore(_REDIS_URL, key_prefix=f"test-{uuid.uuid4()}") + + +def _with_store(client: Client, store: StreamStore) -> Client: + config = client.config() + config["stream_config"] = StreamConfig(stores=[store]) + return Client(**config) + + +async def _create_stream(client: Client) -> Stream[str]: + return await client.create_stream(id=f"stream-{uuid.uuid4()}", body_type=str) + + +async def _stream_error(call: Awaitable[object]) -> StreamError: + with pytest.raises(StreamError) as err: + await call + return err.value + + +@activity.defn +async def append_values(ref: StreamRef[str], values: list[str]) -> None: + await activity.client().get_stream(ref, body_type=str).append(*values) + + +@workflow.defn +class AppendWorkflow: + @workflow.run + async def run(self, stream_id: str) -> StreamRef[str]: + stream = await create_stream(id=stream_id, body_type=str) + await workflow.execute_local_activity( + append_values, + args=[stream.ref, ["a", "b"]], + start_to_close_timeout=timedelta(seconds=10), + ) + return stream.ref + + +@activity.defn +async def append_in_activity(stream_id: str) -> None: + stream = activity.client().get_stream(StreamRef(id=stream_id), body_type=str) + await stream.append("from-activity") + + +@workflow.defn +class RunActivityWorkflow: + @workflow.run + async def run(self, stream_id: str) -> None: + await workflow.execute_activity( + append_in_activity, + stream_id, + start_to_close_timeout=timedelta(seconds=10), + ) + + +async def test_client_creates_appends_reads_and_closes(client: Client): + stream_client = _with_store(client, _store()) + stream = await _create_stream(stream_client) + assert stream.created + assert stream.ref.store == "redis" + + appended = await stream.append("a", "b") + assert appended.first_offset is not None + assert appended.last_offset is not None + records = [record async for record in stream.iter_read_items()] + assert [(record.body, record.offset) for record in records] == [ + ("a", appended.first_offset), + ("b", appended.last_offset), + ] + + await stream.close() + described = await stream.describe() + assert described.status == StreamStatus.StreamStatusClosed + assert described.last_offset == appended.last_offset + page = await stream.read(from_offset=appended.last_offset) + assert page.done + assert page.records == [] + + +async def test_workflow_local_activity_appends_with_local_activity_producer( + client: Client, +): + stream_client = _with_store(client, _store()) + async with new_worker( + stream_client, AppendWorkflow, activities=[append_values] + ) as worker: + handle = await stream_client.start_workflow( + AppendWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + ref = await handle.result() + + stream = stream_client.get_stream(ref, body_type=str) + records = [record async for record in stream.iter_read_items()] + assert [record.body for record in records] == ["a", "b"] + run_id = handle.result_run_id + for record in records: + assert record.producer is not None + assert isinstance(record.producer.kind, ProducerKindLocalActivity) + assert record.producer.kind.value.run_id == run_id + + +async def test_activity_appends_with_activity_producer(client: Client): + stream_client = _with_store(client, _store()) + stream = await _create_stream(stream_client) + async with new_worker( + stream_client, RunActivityWorkflow, activities=[append_in_activity] + ) as worker: + handle = await stream_client.start_workflow( + RunActivityWorkflow.run, + stream.ref.id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + + [record] = [record async for record in stream.iter_read_items()] + assert record.body == "from-activity" + assert record.producer is not None + assert isinstance(record.producer.kind, ProducerKindWorkflowActivity) + assert record.producer.kind.value.run_id == handle.result_run_id + + +async def test_client_append_with_sequence_writes_once(client: Client): + stream = await _create_stream(_with_store(client, _store())) + producer = Producer(id="mine", kind=ProducerKindClient(ClientProducer())) + first = await stream.append("a", "b", producer=producer, sequence=5) + second = await stream.append("a", "b", producer=producer, sequence=5) + assert second == first + records = [record async for record in stream.iter_read_items()] + assert [record.body for record in records] == ["a", "b"] + assert [record.offset for record in records] == [ + first.first_offset, + first.last_offset, + ] + + +async def test_failed_calls_raise_stream_error_with_reason(client: Client): + stream_client = _with_store(client, _store()) + stream = await _create_stream(stream_client) + + err = await _stream_error( + stream_client.create_stream(id=stream.ref.id, body_type=str) + ) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_ALREADY_EXISTS + + unknown = stream_client.get_stream( + StreamRef(id=f"stream-{uuid.uuid4()}"), body_type=str + ) + err = await _stream_error(unknown.describe()) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND + + await stream.close() + err = await _stream_error(stream.append("a")) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_CLOSED + + +async def test_read_with_wait_gets_a_later_record(client: Client): + stream_client = _with_store(client, _store()) + stream = await _create_stream(stream_client) + + async def append_later() -> None: + await asyncio.sleep(0.5) + await stream.append("late") + + append = asyncio.create_task(append_later()) + page = await stream.read(wait=timedelta(seconds=10)) + await append + assert [record.body for record in page.records] == ["late"] + + +async def test_key_prefixes_keep_streams_apart(client: Client): + first = _with_store(client, _store()) + second = _with_store(client, _store()) + stream = await _create_stream(first) + await stream.append("a") + + other = second.get_stream(StreamRef(id=stream.ref.id), body_type=str) + err = await _stream_error(other.describe()) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND + created = await second.create_stream(id=stream.ref.id, body_type=str) + assert created.created + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +async def test_store_without_server_fails_at_first_call(client: Client): + store = RedisStreamStore(f"redis://127.0.0.1:{_free_port()}") + stream_client = _with_store(client, store) + err = await _stream_error(_create_stream(stream_client)) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_STORAGE_UNAVAILABLE + assert not err.non_retryable + + +async def test_bad_option_fails_with_the_crate_message(client: Client): + assert _REDIS_URL + store = RedisStreamStore([_REDIS_URL, _REDIS_URL]) + with pytest.raises(ValueError, match="cluster"): + await _create_stream(_with_store(client, store)) diff --git a/tests/helpers/streams.py b/tests/helpers/streams.py new file mode 100644 index 0000000000..beb4eaaa45 --- /dev/null +++ b/tests/helpers/streams.py @@ -0,0 +1,186 @@ +"""Stream Stores that the Stream tests share.""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Awaitable, Callable +from typing import TypeVar + +from google.protobuf.message import Message + +from temporalio.api.enums.v1 import ( + StreamErrorReason, + StreamIdConflictPolicy, + StreamStatus, +) +from temporalio.api.stream.v1 import StreamRecord, StreamRef +from temporalio.api.streamservice.v1 import ( + AppendRecordsRequest, + AppendRecordsResponse, + CloseStreamRequest, + CloseStreamResponse, + CreateStreamRequest, + CreateStreamResponse, + DescribeStreamRequest, + DescribeStreamResponse, + ReadRecordsRequest, + ReadRecordsResponse, +) +from temporalio.client import StreamStore +from temporalio.exceptions import StreamError + +_RequestT = TypeVar("_RequestT", bound=Message) +_ResponseT = TypeVar("_ResponseT", bound=Message) + + +class RecordingStreamStore(StreamStore): + """A Store that passes each call to another Store and keeps the operation names. + + It has the name of the Store that it wraps, since that Store writes its own name + in the ref of a create. A test overrides :py:meth:`before_call` to fail a call or + to hold it. + """ + + def __init__(self, inner: StreamStore) -> None: + self.inner = inner + self.operations: list[str] = [] + + @property + def name(self) -> str: + return self.inner.name + + async def before_call(self, operation: str, request: Message) -> None: + del operation, request + + async def _forward( + self, + operation: str, + request: _RequestT, + call: Callable[[_RequestT], Awaitable[_ResponseT]], + ) -> _ResponseT: + self.operations.append(operation) + await self.before_call(operation, request) + return await call(request) + + async def create_stream(self, request: CreateStreamRequest) -> CreateStreamResponse: + return await self._forward("CreateStream", request, self.inner.create_stream) + + async def append_records( + self, request: AppendRecordsRequest + ) -> AppendRecordsResponse: + return await self._forward("AppendRecords", request, self.inner.append_records) + + async def read_records(self, request: ReadRecordsRequest) -> ReadRecordsResponse: + return await self._forward("ReadRecords", request, self.inner.read_records) + + async def close_stream(self, request: CloseStreamRequest) -> CloseStreamResponse: + return await self._forward("CloseStream", request, self.inner.close_stream) + + async def describe_stream( + self, request: DescribeStreamRequest + ) -> DescribeStreamResponse: + return await self._forward( + "DescribeStream", request, self.inner.describe_stream + ) + + +@dataclasses.dataclass +class _DictStream: + request_id: str + records: list[StreamRecord] = dataclasses.field(default_factory=list) + closed: bool = False + + +def _offset(index: int) -> bytes: + return index.to_bytes(8, "big") + + +class DictStreamStore(StreamStore): + """A Store written from scratch, with no bridge, that keeps its Streams in a dict. + + It has enough of each operation for the tests, and no dedupe, no wait and no + Retention. + """ + + def __init__(self, name: str = "dict") -> None: + self._name = name + self.streams: dict[str, _DictStream] = {} + + @property + def name(self) -> str: + return self._name + + def _stream(self, ref: StreamRef) -> _DictStream: + stream = self.streams.get(ref.id) + if stream is None: + raise StreamError( + f"No Stream '{ref.id}'", + reason=StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND, + ) + return stream + + async def create_stream(self, request: CreateStreamRequest) -> CreateStreamResponse: + ref = StreamRef(id=request.stream_ref.id, store=self.name) + stream = self.streams.get(ref.id) + if stream is None: + self.streams[ref.id] = _DictStream(request_id=request.request_id) + return CreateStreamResponse(ref=ref, created=True) + if stream.request_id == request.request_id: + return CreateStreamResponse(ref=ref, created=True) + if ( + request.id_conflict_policy + == StreamIdConflictPolicy.STREAM_ID_CONFLICT_POLICY_USE_EXISTING + ): + return CreateStreamResponse(ref=ref, created=False) + raise StreamError( + f"The Stream '{ref.id}' exists", + reason=StreamErrorReason.STREAM_ERROR_REASON_ALREADY_EXISTS, + ) + + async def append_records( + self, request: AppendRecordsRequest + ) -> AppendRecordsResponse: + stream = self._stream(request.ref) + if stream.closed: + raise StreamError( + f"The Stream '{request.ref.id}' is closed", + reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED, + ) + first = len(stream.records) + 1 + for record in request.records: + stream.records.append( + StreamRecord( + producer=request.producer, + body=record.body, + offset=_offset(len(stream.records) + 1), + ) + ) + return AppendRecordsResponse( + first_offset=_offset(first), last_offset=_offset(len(stream.records)) + ) + + async def read_records(self, request: ReadRecordsRequest) -> ReadRecordsResponse: + stream = self._stream(request.ref) + start = int.from_bytes(request.from_offset, "big") if request.from_offset else 0 + page = stream.records[start : start + (request.max_records or 1000)] + last_offset = page[-1].offset if page else request.from_offset + return ReadRecordsResponse( + records=page, + last_offset=last_offset, + done=stream.closed and start + len(page) == len(stream.records), + ) + + async def close_stream(self, request: CloseStreamRequest) -> CloseStreamResponse: + self._stream(request.ref).closed = True + return CloseStreamResponse() + + async def describe_stream( + self, request: DescribeStreamRequest + ) -> DescribeStreamResponse: + stream = self._stream(request.ref) + return DescribeStreamResponse( + status=StreamStatus.STREAM_STATUS_CLOSED + if stream.closed + else StreamStatus.STREAM_STATUS_OPEN, + last_offset=stream.records[-1].offset if stream.records else b"", + ) diff --git a/tests/nexus/test_stream_appender.py b/tests/nexus/test_stream_appender.py new file mode 100644 index 0000000000..fd320ff474 --- /dev/null +++ b/tests/nexus/test_stream_appender.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta + +import pytest +from google.protobuf.message import Message + +import temporalio.api.streamservice.v1 +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.client import BufferedAppender, Client +from temporalio.nexus.system.stream_service import StreamRef +from temporalio.nexus.system.stream_service.client import Stream +from temporalio.testing import MemoryStreamStore +from tests.helpers.streams import RecordingStreamStore +from tests.nexus.test_stream_store import _create_stream, _stream_error, _with_store + +# Long enough that only a flush, a full batch or the end of `async with` starts a write. +_NEVER = timedelta(hours=1) + + +class RecordingStore(RecordingStreamStore): + def __init__(self) -> None: + super().__init__(MemoryStreamStore()) + self.batches: list[int] = [] + self.append_started = asyncio.Event() + self.release: asyncio.Event | None = None + + async def before_call(self, operation: str, request: Message) -> None: + if isinstance(request, temporalio.api.streamservice.v1.AppendRecordsRequest): + self.batches.append(len(request.records)) + self.append_started.set() + if self.release is not None: + await self.release.wait() + + +async def _bodies(stream: Stream[str]) -> list[str]: + return [record.body async for record in stream.iter_read_items()] + + +async def test_appends_inside_one_delay_go_as_one_batch(client: Client): + store = RecordingStore() + stream = await _create_stream(client, store) + appender = BufferedAppender(stream, max_delay=_NEVER) + appender.append("a") + appender.append("b", "c") + appender.append("d") + await appender.flush() + assert store.batches == [4] + assert await _bodies(stream) == ["a", "b", "c", "d"] + + +async def test_flush_returns_the_last_offset(client: Client): + stream = await _create_stream(client, MemoryStreamStore()) + appender = BufferedAppender(stream) + assert await appender.flush() is None + appender.append("a", "b") + offset = await appender.flush() + assert offset is not None + assert offset == (await stream.describe()).last_offset + + +async def test_appends_during_a_write_go_in_the_next_batch(client: Client): + store = RecordingStore() + stream = await _create_stream(client, store) + store.release = asyncio.Event() + appender = BufferedAppender(stream, max_delay=_NEVER) + appender.append("a") + first = asyncio.create_task(appender.flush()) + await store.append_started.wait() + appender.append("b") + appender.append("c") + assert store.batches == [1] + store.release.set() + await first + await appender.flush() + assert store.batches == [1, 2] + assert await _bodies(stream) == ["a", "b", "c"] + + +async def test_a_large_buffer_goes_as_batches_of_at_most_1000(client: Client): + store = RecordingStore() + stream = await _create_stream(client, store) + appender = BufferedAppender(stream, max_delay=_NEVER) + values = [str(index) for index in range(2500)] + appender.append(*values) + # A full batch doesn't wait for the delay or a flush. + await store.append_started.wait() + await appender.flush() + assert store.batches == [1000, 1000, 500] + assert await _bodies(stream) == values + + +async def test_the_timer_alone_writes_the_batch(client: Client): + store = RecordingStore() + stream = await _create_stream(client, store) + appender = BufferedAppender(stream, max_delay=timedelta(milliseconds=10)) + appender.append("a", "b") + page = await stream.read(wait=timedelta(seconds=10)) + assert [record.body for record in page.records] == ["a", "b"] + assert store.batches == [2] + + +async def test_a_failed_batch_stops_the_appender(client: Client): + store = RecordingStore() + stream = await _create_stream(client, store) + await stream.close() + store.release = asyncio.Event() + appender = BufferedAppender(stream, max_delay=_NEVER) + appender.append("a") + first = asyncio.create_task(appender.flush()) + await store.append_started.wait() + appender.append("b", "c") + store.release.set() + err = await _stream_error(first) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_CLOSED + assert await _stream_error(appender.flush()) is err + with pytest.raises(type(err)) as later: + appender.append("d") + assert later.value is err + assert appender.unwritten == 3 + assert store.batches == [1] + + +async def test_async_with_flushes_at_its_end(client: Client): + store = RecordingStore() + stream = await _create_stream(client, store) + async with BufferedAppender(stream, max_delay=_NEVER) as appender: + appender.append("a", "b") + assert store.batches == [2] + assert await _bodies(stream) == ["a", "b"] + + +async def test_async_with_keeps_the_error_of_its_block(client: Client): + store = RecordingStore() + stream = await _create_stream(client, store) + await stream.close() + appender = BufferedAppender(stream, max_delay=_NEVER) + with pytest.raises(RuntimeError, match="the block failed"): + async with appender: + appender.append("a") + raise RuntimeError("the block failed") + # The flush at the end of the block ran and failed. + assert store.batches == [1] + assert appender.unwritten == 1 + err = await _stream_error(appender.flush()) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_CLOSED + + +async def test_append_without_a_value_fails(client: Client): + stream = _with_store(client, MemoryStreamStore()).get_stream( + StreamRef(id=f"stream-{uuid.uuid4()}"), body_type=str + ) + appender = BufferedAppender(stream) + with pytest.raises(TypeError, match="one or more values"): + appender.append() + with pytest.raises(ValueError, match="max_delay must not be negative"): + BufferedAppender(stream, max_delay=timedelta(seconds=-1)) diff --git a/tests/nexus/test_stream_store.py b/tests/nexus/test_stream_store.py new file mode 100644 index 0000000000..db37459f9b --- /dev/null +++ b/tests/nexus/test_stream_store.py @@ -0,0 +1,630 @@ +from __future__ import annotations + +import ast +import dataclasses +import subprocess +import sys +import uuid +from collections.abc import Awaitable, Callable +from datetime import datetime, timedelta, timezone +from pathlib import Path +from unittest import mock + +import google.protobuf.timestamp_pb2 +import pytest + +import temporalio +import temporalio.activity +import temporalio.api.stream.v1 +import temporalio.api.workflowservice.v1 +import temporalio.client._stream_store +import temporalio.common +import temporalio.converter._payload_converter +import temporalio.exceptions +import temporalio.nexus.system +import temporalio.nexus.system._support +import temporalio.workflow +from temporalio.api.enums.v1 import StreamErrorReason +from temporalio.client import ( + Client, + StreamConfig, + StreamSelectContext, + StreamStore, +) +from temporalio.client._stream_store import _STREAM_OPERATIONS, _request_ref +from temporalio.converter import DataConverter, StorageDriverActivityInfo +from temporalio.exceptions import ApplicationError, StreamConfigError, StreamError +from temporalio.nexus.system.stream_service import ( + ClientProducer, + DescribeStreamRequest, + Producer, + ProducerKindClient, + ProducerKindLocalActivity, + ProducerKindStandaloneActivity, + ProducerKindWorkflowActivity, + StreamIdConflictPolicy, + StreamRef, + StreamService, + StreamStatus, + Timestamp, +) +from temporalio.nexus.system.stream_service.client import ( + Stream, + StreamServiceClientExtension, +) +from temporalio.testing import ( + ActivityEnvironment, + MemoryStreamStore, + WorkflowEnvironment, +) +from tests import DEV_SERVER_DOWNLOAD_VERSION +from tests.helpers.streams import DictStreamStore +from tests.test_converter import SimpleCodec + + +def _with_store(client: Client, store: StreamStore) -> Client: + return _with_config(client, StreamConfig(stores=[store])) + + +def _with_config(client: Client, stream_config: StreamConfig) -> Client: + config = client.config() + config["stream_config"] = stream_config + return Client(**config) + + +def _two_stores( + selector: Callable[[StreamSelectContext], StreamStore], +) -> tuple[MemoryStreamStore, MemoryStreamStore, StreamConfig]: + first = MemoryStreamStore(name="first") + second = MemoryStreamStore(name="second") + return first, second, StreamConfig(stores=[first, second], store_selector=selector) + + +async def _create_stream(client: Client, store: StreamStore) -> Stream[str]: + return await _with_store(client, store).create_stream( + id=f"stream-{uuid.uuid4()}", body_type=str + ) + + +async def _only_producer(stream: Stream[str]) -> Producer: + [record] = [record async for record in stream.iter_read_items()] + assert record.producer is not None + return record.producer + + +async def _append_in_activity( + client: Client, info: temporalio.activity.Info +) -> Producer: + stream = await _create_stream(client, MemoryStreamStore()) + env = ActivityEnvironment() + env.info = info + await env.run(stream.append, "a") + return await _only_producer(stream) + + +def _timestamp(value: datetime) -> Timestamp: + proto = google.protobuf.timestamp_pb2.Timestamp() + proto.FromDatetime(value) + return Timestamp(seconds=proto.seconds, nanos=proto.nanos) + + +async def _stream_error(call: Awaitable[object]) -> StreamError: + with pytest.raises(StreamError) as err: + await call + return err.value + + +async def test_client_stream_uses_store(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + + stream = await stream_client.create_stream( + id=f"stream-{uuid.uuid4()}", body_type=str + ) + assert stream.created + assert stream.ref.store == "memory" + appended = await stream.append("a", "b") + assert appended.first_offset is not None + assert appended.last_offset is not None + assert appended.first_offset < appended.last_offset + records = [record async for record in stream.iter_read_items()] + assert [(record.body, record.offset) for record in records] == [ + ("a", appended.first_offset), + ("b", appended.last_offset), + ] + + other = stream_client.get_stream(stream.ref, body_type=str) + assert [record.body async for record in other.iter_read_items()] == ["a", "b"] + + described = await stream.describe() + assert described.status == StreamStatus.StreamStatusOpen + assert described.last_offset == appended.last_offset + + await stream.close() + described = await stream.describe() + assert described.status == StreamStatus.StreamStatusClosed + assert described.last_offset == appended.last_offset + page = await stream.read(from_offset=records[-1].offset) + assert page.done + assert page.records == [] + pages = [page async for page in stream.iter_read()] + assert pages[-1].done + assert [record.body for page in pages for record in page.records] == ["a", "b"] + + +async def test_client_stream_uses_a_store_written_from_scratch(client: Client): + store = DictStreamStore() + config = client.config() + config["data_converter"] = DataConverter(payload_codec=SimpleCodec()) + config["stream_config"] = StreamConfig(stores=[store]) + stream_client = Client(**config) + + stream_id = f"stream-{uuid.uuid4()}" + stream = await stream_client.create_stream(id=stream_id, body_type=str) + assert stream.created + assert stream.ref.store == "dict" + appended = await stream.append("a", "b") + assert [record.body async for record in stream.iter_read_items()] == ["a", "b"] + # The Store gets payloads that the client already encoded with the codec. + stored = store.streams[stream_id].records + assert [record.body.metadata["simple-codec"] for record in stored] == [ + b"true", + b"true", + ] + described = await stream.describe() + assert described.status == StreamStatus.StreamStatusOpen + assert described.last_offset == appended.last_offset + err = await _stream_error(stream_client.create_stream(id=stream_id, body_type=str)) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_ALREADY_EXISTS + + await stream.close() + err = await _stream_error(stream.append("c")) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_CLOSED + + +def test_stream_operations_match_the_stream_service(): + service = temporalio.client._stream_store._STREAM_SERVICE + assert set(_STREAM_OPERATIONS) == set(service.operation_definitions) + for name, operation in _STREAM_OPERATIONS.items(): + definition = service.operation_definitions[name] + for model, message in ( + (definition.input_type, operation.request_type), + (definition.output_type, operation.response_type), + ): + converter = ( + temporalio.converter._payload_converter._get_transfer_type_converter( + model + ) + ) + assert converter is not None + assert converter.transfer_type is message + assert callable(getattr(StreamStore, operation.store_method)) + + +async def test_client_strict_create_of_existing_stream_fails(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + await stream_client.create_stream(id=stream_id, body_type=str) + err = await _stream_error(stream_client.create_stream(id=stream_id, body_type=str)) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_ALREADY_EXISTS + assert err.non_retryable + existing = await stream_client.create_stream( + id=stream_id, + body_type=str, + id_conflict_policy=StreamIdConflictPolicy.StreamIdConflictPolicyUseExisting, + ) + # The response's `created` is false, which the generated model reads as unset. + assert not existing.created + + +async def test_client_read_or_describe_of_unknown_stream_fails(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream = stream_client.get_stream( + StreamRef(id=f"stream-{uuid.uuid4()}"), body_type=str + ) + for call in (stream.read(), stream.describe()): + err = await _stream_error(call) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND + + +async def test_client_append_after_close_fails(client: Client): + stream = await _create_stream(client, MemoryStreamStore()) + await stream.append("a") + await stream.close() + err = await _stream_error(stream.append("b")) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_CLOSED + assert err.type == "StreamError" + + +async def test_client_stream_store_uses_payload_codec(client: Client): + store = MemoryStreamStore() + config = client.config() + config["data_converter"] = DataConverter(payload_codec=SimpleCodec()) + config["stream_config"] = StreamConfig(stores=[store]) + stream_client = Client(**config) + + stream = await stream_client.create_stream( + id=f"stream-{uuid.uuid4()}", body_type=str + ) + await stream.append("a", "b") + # A client without the codec reads the bodies as the Store keeps them. + raw = _with_store(client, store).get_stream( + StreamRef(id=stream.ref.id), body_type=temporalio.common.RawValue + ) + stored = [record.body.payload async for record in raw.iter_read_items()] + assert [body.metadata["simple-codec"] for body in stored] == [b"true", b"true"] + assert [record.body async for record in stream.iter_read_items()] == ["a", "b"] + + +async def test_client_append_has_client_producer(client: Client): + stream = await _create_stream(client, MemoryStreamStore()) + await stream.append("a") + producer = await _only_producer(stream) + assert producer.id.startswith("client:") + assert isinstance(producer.kind, ProducerKindClient) + + +async def test_workflow_activity_append_has_activity_producer(client: Client): + scheduled = datetime(2026, 11, 12, 10, 0, 0, 123456, tzinfo=timezone.utc) + attempt_scheduled = scheduled + timedelta(seconds=30) + producer = await _append_in_activity( + client, + dataclasses.replace( + ActivityEnvironment.default_info(), + activity_id="step:1", + attempt=3, + workflow_run_id="run-1", + scheduled_time=scheduled, + current_attempt_scheduled_time=attempt_scheduled, + ), + ) + epoch_nanos = ( + (scheduled - datetime(1970, 1, 1, tzinfo=timezone.utc)) + // timedelta(microseconds=1) + * 1000 + ) + assert producer.id == f"act:run-1:step%3A1:{epoch_nanos}" + assert isinstance(producer.kind, ProducerKindWorkflowActivity) + assert producer.kind.value.run_id == "run-1" + assert producer.kind.value.epoch.attempt == 3 + assert producer.kind.value.epoch.attempt_scheduled_time == _timestamp( + attempt_scheduled + ) + + +async def test_standalone_activity_append_has_activity_producer(client: Client): + attempt_scheduled = datetime(2026, 11, 12, 10, 0, 30, tzinfo=timezone.utc) + producer = await _append_in_activity( + client, + dataclasses.replace( + ActivityEnvironment.default_info(), + activity_id="job", + activity_run_id="job-run", + attempt=2, + workflow_id=None, + workflow_namespace=None, + workflow_run_id=None, + workflow_type=None, + current_attempt_scheduled_time=attempt_scheduled, + ), + ) + assert producer.id == "sa:job:job-run" + assert isinstance(producer.kind, ProducerKindStandaloneActivity) + assert not producer.kind.value.run_id + assert producer.kind.value.epoch.attempt == 2 + assert producer.kind.value.epoch.attempt_scheduled_time == _timestamp( + attempt_scheduled + ) + + +async def test_local_activity_append_has_local_activity_producer(client: Client): + task_started = datetime(2026, 11, 12, 10, 0, 0, 123456, tzinfo=timezone.utc) + producer = await _append_in_activity( + client, + dataclasses.replace( + ActivityEnvironment.default_info(), + activity_id="3", + attempt=2, + is_local=True, + workflow_run_id="run:1", + workflow_task=temporalio.activity.WorkflowTaskAttempt( + started_event_id=7, attempt=1, started_time=task_started + ), + ), + ) + assert producer.id == "la:run%3A1:3" + assert isinstance(producer.kind, ProducerKindLocalActivity) + assert producer.kind.value.run_id == "run:1" + assert producer.kind.value.epoch.task.started_event_id == 7 + assert producer.kind.value.epoch.task.attempt == 1 + assert producer.kind.value.epoch.task.started_time == _timestamp(task_started) + assert producer.kind.value.epoch.attempt == 2 + + +async def test_local_activity_append_without_workflow_task_fails(client: Client): + with pytest.raises(RuntimeError, match="Workflow Task"): + await _append_in_activity( + client, + dataclasses.replace(ActivityEnvironment.default_info(), is_local=True), + ) + + +async def test_explicit_producer_is_sent(client: Client): + stream = await _create_stream(client, MemoryStreamStore()) + with mock.patch.object( + temporalio.nexus.system._support, "default_producer" + ) as default_producer: + await stream.append( + "a", producer=Producer(id="mine", kind=ProducerKindClient(ClientProducer())) + ) + default_producer.assert_not_called() + producer = await _only_producer(stream) + assert producer.id == "mine" + assert isinstance(producer.kind, ProducerKindClient) + + +async def test_client_append_with_sequence_writes_once(client: Client): + stream = await _create_stream(client, MemoryStreamStore()) + producer = Producer(id="mine", kind=ProducerKindClient(ClientProducer())) + first = await stream.append("a", "b", producer=producer, sequence=5) + second = await stream.append("a", "b", producer=producer, sequence=5) + assert second == first + records = [record async for record in stream.iter_read_items()] + assert [record.body for record in records] == ["a", "b"] + assert [record.offset for record in records] == [ + first.first_offset, + first.last_offset, + ] + + +def test_client_stream_methods_are_generated(): + assert "create_stream" not in vars(Client) + assert "get_stream" not in vars(Client) + assert Client.create_stream is StreamServiceClientExtension.create_stream + assert Client.get_stream is StreamServiceClientExtension.get_stream + + +def test_no_hand_written_stream_classes(): + root = Path(temporalio.__file__).parent + generated = root / "nexus" / "system" / "stream_service" + names = {"Stream", "WorkflowStream", "StreamCalls", "StreamClient"} + found: list[tuple[str, str]] = [] + for source in sorted(root.rglob("*.py")): + if generated in source.parents: + continue + for node in ast.walk(ast.parse(source.read_text(encoding="utf8"))): + if isinstance(node, ast.ClassDef) and node.name in names: + found.append((source.relative_to(root).as_posix(), node.name)) + # The Workflow Streams contrib library is older than the generated Stream API and + # doesn't use it. + found.remove(("contrib/workflow_streams/_stream.py", "WorkflowStream")) + assert found == [] + for name in ("Stream", "WorkflowStream"): + value = getattr(temporalio.workflow, name, None) + assert value is None or value.__module__.startswith( + "temporalio.nexus.system.stream_service." + ) + + +def test_each_stream_module_can_be_the_first_import(): + modules = [ + "temporalio.nexus.system.stream_service.workflow", + "temporalio.nexus.system.stream_service.client", + "temporalio.nexus.system.stream_service", + "temporalio.nexus.system", + "temporalio.workflow", + "temporalio.client", + "temporalio.worker", + "temporalio.contrib.redis.streams", + ] + failures: dict[str, str] = {} + for module in modules: + # Only a new interpreter has no module loaded, so the import order is real. + result = subprocess.run( + [sys.executable, "-c", f"import {module}"], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + failures[module] = result.stderr.strip().splitlines()[-1] + assert failures == {} + + +async def test_client_without_stream_store_raises_and_sends_nothing(client: Client): + requests: list[ + temporalio.api.workflowservice.v1.StartNexusOperationExecutionRequest + ] = [] + + async def start_nexus_operation_execution( + req: temporalio.api.workflowservice.v1.StartNexusOperationExecutionRequest, + **_: object, + ) -> temporalio.api.workflowservice.v1.StartNexusOperationExecutionResponse: + requests.append(req) + return temporalio.api.workflowservice.v1.StartNexusOperationExecutionResponse( + run_id="run" + ) + + nexus_client = client.create_nexus_client( + StreamService, temporalio.nexus.system.TEMPORAL_SYSTEM_ENDPOINT + ) + with mock.patch.object( + client.workflow_service, + "start_nexus_operation_execution", + start_nexus_operation_execution, + ): + with pytest.raises(StreamConfigError) as err: + await nexus_client.start_operation( + StreamService.describe_stream, + DescribeStreamRequest(ref=StreamRef(id="stream")), + id="operation", + ) + # The generated client method goes the same way. + with pytest.raises(StreamConfigError): + await client.create_stream(id="stream", body_type=str) + assert str(err.value) == ( + "This client has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...])." + ) + assert requests == [] + + +def test_stream_config_refuses_a_bad_set_of_stores(): + with pytest.raises(StreamConfigError, match="at least one Store"): + StreamConfig(stores=[]) + same_name = [MemoryStreamStore(name="a"), MemoryStreamStore(name="a")] + with pytest.raises(StreamConfigError, match="more than one Store named 'a'"): + StreamConfig(stores=same_name, store_selector=lambda _: same_name[0]) + with pytest.raises(StreamConfigError, match="store_selector must be set"): + StreamConfig(stores=[MemoryStreamStore(name="a"), MemoryStreamStore(name="b")]) + assert not issubclass(StreamConfigError, StreamError) + + +def test_client_refuses_a_store_in_plugins(client: Client): + config = client.config() + config["plugins"] = [MemoryStreamStore()] # type: ignore[list-item] + with pytest.raises(StreamConfigError, match="plugins"): + Client(**config) + + +async def test_selector_picks_the_store_of_a_new_stream(client: Client): + contexts: list[StreamSelectContext] = [] + + def select(context: StreamSelectContext) -> StreamStore: + contexts.append(context) + return second + + first, second, stream_config = _two_stores(select) + stream_client = _with_config(client, stream_config) + stream_id = f"stream-{uuid.uuid4()}" + stream = await stream_client.create_stream(id=stream_id, body_type=str) + assert stream.ref.store == "second" + assert contexts == [StreamSelectContext(id=stream_id, target=None)] + + await stream_client.get_stream(stream.ref, body_type=str).append("a") + # A call through a ref that names its Store doesn't ask the selector. + assert len(contexts) == 1 + in_second = _with_store(client, second).get_stream(stream.ref, body_type=str) + assert [record.body async for record in in_second.iter_read_items()] == ["a"] + in_first = _with_store(client, first).get_stream( + StreamRef(id=stream_id), body_type=str + ) + err = await _stream_error(in_first.describe()) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND + + +async def test_selector_gets_the_activity_that_makes_the_call(client: Client): + contexts: list[StreamSelectContext] = [] + + def select(context: StreamSelectContext) -> StreamStore: + contexts.append(context) + return first + + first, _, stream_config = _two_stores(select) + stream_client = _with_config(client, stream_config) + stream_id = f"stream-{uuid.uuid4()}" + env = ActivityEnvironment() + env.info = dataclasses.replace( + ActivityEnvironment.default_info(), + activity_id="job", + activity_type="make_stream", + activity_run_id="job-run", + namespace="streams", + ) + + async def create() -> Stream[str]: + return await stream_client.create_stream(id=stream_id, body_type=str) + + stream = await env.run(create) + assert stream.ref.store == "first" + assert contexts == [ + StreamSelectContext( + id=stream_id, + target=StorageDriverActivityInfo( + id="job", type="make_stream", run_id="job-run", namespace="streams" + ), + ) + ] + + +async def test_selector_that_returns_another_store_fails(client: Client): + _, _, stream_config = _two_stores(lambda _: MemoryStreamStore(name="first")) + stream_client = _with_config(client, stream_config) + with pytest.raises(StreamConfigError, match="isn't one of"): + await stream_client.create_stream(id=f"stream-{uuid.uuid4()}", body_type=str) + + +async def test_ref_of_a_store_that_the_client_lacks_fails(client: Client): + stream_client = _with_store(client, MemoryStreamStore(name="first")) + stream = stream_client.get_stream( + StreamRef(id=f"stream-{uuid.uuid4()}", store="other"), body_type=str + ) + with pytest.raises(StreamConfigError, match="'other'"): + await stream.describe() + + +def test_request_without_a_stream_ref_fails(): + request = Producer(id="mine", kind=ProducerKindClient(ClientProducer())) + with pytest.raises(TypeError, match="Producer has no StreamRef"): + _request_ref(request) + + +async def test_start_local_gives_its_client_the_stream_config(): + stream_config = StreamConfig(stores=[MemoryStreamStore()]) + async with await WorkflowEnvironment.start_local( + stream_config=stream_config, + dev_server_download_version=DEV_SERVER_DOWNLOAD_VERSION, + ) as env: + assert env.client.config()["stream_config"] is stream_config + stream = await env.client.create_stream( + id=f"stream-{uuid.uuid4()}", body_type=str + ) + assert stream.ref.store == "memory" + + +def test_stream_error_carries_its_reason(): + err = StreamError( + "the Stream is closed", reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED + ) + assert err.type == "StreamError" + assert err.message == "the Stream is closed" + assert err.non_retryable + assert list(err.details) == [ + temporalio.api.stream.v1.StreamErrorDetails( + reason=StreamErrorReason.STREAM_ERROR_REASON_CLOSED + ) + ] + assert not StreamError( + "down", reason=StreamErrorReason.STREAM_ERROR_REASON_STORAGE_UNAVAILABLE + ).non_retryable + assert StreamError("x").reason == ( + StreamErrorReason.STREAM_ERROR_REASON_UNSPECIFIED + ) + + +def test_stream_error_from_application_error(): + from_failure = temporalio.exceptions._stream_error_from_failure + rebuilt = from_failure( + ApplicationError( + "the Stream exists", + temporalio.api.stream.v1.StreamErrorDetails( + reason=StreamErrorReason.STREAM_ERROR_REASON_ALREADY_EXISTS + ), + type="StreamError", + ) + ) + assert isinstance(rebuilt, StreamError) + assert rebuilt.reason == StreamErrorReason.STREAM_ERROR_REASON_ALREADY_EXISTS + assert rebuilt.message == "the Stream exists" + unknown = from_failure( + ApplicationError( + "x", + temporalio.api.stream.v1.StreamErrorDetails( + reason=StreamErrorReason.ValueType(999) + ), + type="StreamError", + ) + ) + assert unknown is not None + assert unknown.reason == StreamErrorReason.STREAM_ERROR_REASON_UNSPECIFIED + assert from_failure(ApplicationError("x", type="RuntimeError")) is None + assert from_failure(RuntimeError("x")) is None diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py new file mode 100644 index 0000000000..7913a60498 --- /dev/null +++ b/tests/nexus/test_stream_store_workflow.py @@ -0,0 +1,806 @@ +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Callable, Sequence +from datetime import timedelta +from typing import Any + +import pytest +from google.protobuf.message import Message + +import temporalio.common +import temporalio.exceptions +from temporalio import activity, workflow +from temporalio.api.enums.v1 import EventType, StreamErrorReason +from temporalio.api.streamservice.v1 import ( + CreateStreamRequest as CreateStreamRequestMessage, +) +from temporalio.api.streamservice.v1 import ( + CreateStreamResponse as CreateStreamResponseMessage, +) +from temporalio.api.streamservice.v1 import ( + DescribeStreamRequest as DescribeStreamRequestMessage, +) +from temporalio.api.streamservice.v1 import ( + DescribeStreamResponse as DescribeStreamResponseMessage, +) +from temporalio.client import ( + Client, + StreamConfig, + StreamSelectContext, + StreamStore, + StreamStoreCallOptions, + WorkflowHandle, +) +from temporalio.common import RetryPolicy +from temporalio.converter import DataConverter, StorageDriverWorkflowInfo +from temporalio.exceptions import ApplicationError, StreamError +from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT +from temporalio.nexus.system.stream_service import ( + AppendRecord, + AppendRecordsRequest, + ClientProducer, + CloseStreamRequest, + Producer, + ProducerKindClient, + ProducerKindLocalActivity, + ReadRecordsRequest, + StreamIdConflictPolicy, + StreamRecord, + StreamRef, + StreamService, +) +from temporalio.nexus.system.stream_service.workflow import create_stream, get_stream +from temporalio.nexus.system.stream_service.workflow._resources.stream import ( + Stream as WorkflowStream, +) +from temporalio.testing import MemoryStreamStore +from temporalio.worker import ( + ActivityInboundInterceptor, + ExecuteActivityInput, + Interceptor, + Replayer, +) +from tests.helpers import assert_task_fail_eventually, new_worker +from tests.helpers.streams import DictStreamStore, RecordingStreamStore +from tests.test_converter import SimpleCodec + +# Workflow code writes these values, so the History shows only whether a Record body +# leaks into it. +_BODIES = ["first-record-body", "second-record-body"] + + +@activity.defn +async def append_values(ref: StreamRef[str], values: list[str]) -> None: + await activity.client().get_stream(ref, body_type=str).append(*values) + + +async def _append_in_local_activity(ref: StreamRef[str], values: list[str]) -> None: + await workflow.execute_local_activity( + append_values, args=[ref, values], start_to_close_timeout=timedelta(seconds=10) + ) + + +@workflow.defn +class AppendWorkflow: + @workflow.run + async def run(self, stream_id: str) -> StreamRef[str]: + stream = await create_stream(id=stream_id, body_type=str) + await _append_in_local_activity(stream.ref, _BODIES) + return stream.ref + + +@workflow.defn +class RefusedOperationWorkflow: + @workflow.run + async def run(self, operation: str) -> str: + # The Workflow `Stream` has no method for these operations, so the test starts + # them through a Nexus client, as Workflow code still can. + ref = StreamRef[str](id="stream") + producer = Producer(id="workflow", kind=ProducerKindClient(ClientProducer())) + requests: dict[str, Any] = { + "AppendRecords": AppendRecordsRequest( + ref=ref, records=[AppendRecord(body="a")], producer=producer + ), + "CloseStream": CloseStreamRequest(ref=ref, producer=producer), + "ReadRecords": ReadRecordsRequest(ref=ref), + } + nexus_client = workflow.create_nexus_client( + service=StreamService, endpoint=TEMPORAL_SYSTEM_ENDPOINT + ) + try: + await nexus_client.start_operation(operation, requests[operation]) + except ApplicationError as err: + assert err.non_retryable + return err.message + return "started" + + +@workflow.defn +class DescribeResultWorkflow: + @workflow.run + async def run(self, stream_id: str) -> str: + stream = await create_stream(id=stream_id, body_type=str) + try: + await stream.describe() + except StreamError as err: + return f"StreamError {StreamErrorReason.Name(err.reason)}: {err.message}" + except ApplicationError as err: + return f"ApplicationError {err.type}: {err.message}" + except temporalio.exceptions.TimeoutError as err: + return f"TimeoutError {err.type and err.type.name}" + return "described" + + +@workflow.defn +class DescribeStreamWorkflow: + @workflow.run + async def run(self, stream_id: str) -> str: + stream = get_stream(StreamRef(id=stream_id), body_type=str) + try: + await stream.describe() + except StreamError as err: + assert isinstance(err, ApplicationError) + return StreamErrorReason.Name(err.reason) + return "described" + + +@workflow.defn +class CreateThenContinueAsNewWorkflow: + @workflow.run + async def run(self, stream_id: str, run: int) -> list[bool]: + stream = await create_stream(id=stream_id, body_type=str) + await _append_in_local_activity(stream.ref, [f"run-{run}"]) + if run == 1: + workflow.continue_as_new(args=[stream_id, 2]) + return [bool(stream.created)] + + +@workflow.defn +class DescribeOtherStoreWorkflow: + @workflow.run + async def run(self, stream_id: str) -> None: + await get_stream( + StreamRef(id=stream_id, store="other"), body_type=str + ).describe() + + +@workflow.defn +class CreateAndDescribeWorkflow: + @workflow.run + async def run(self, stream_id: str) -> StreamRef[str]: + stream = await create_stream(id=stream_id, body_type=str) + await stream.describe() + return stream.ref + + +class CountingStore(RecordingStreamStore): + def __init__(self, name: str) -> None: + super().__init__(MemoryStreamStore(name=name)) + + +@activity.defn +async def user_activity() -> None: + pass + + +@workflow.defn +class CreateAndActivityWorkflow: + @workflow.run + async def run(self, stream_id: str) -> None: + await create_stream(id=stream_id, body_type=str) + await workflow.execute_local_activity( + user_activity, start_to_close_timeout=timedelta(seconds=10) + ) + + +_fail_next_task = True + + +@workflow.defn(sandboxed=False) +class AppendThenFailTaskOnceWorkflow: + @workflow.run + async def run(self, stream_id: str) -> None: + stream = await create_stream( + id=stream_id, + body_type=str, + id_conflict_policy=StreamIdConflictPolicy.StreamIdConflictPolicyUseExisting, + ) + await _append_in_local_activity(stream.ref, _BODIES) + global _fail_next_task + if _fail_next_task: + _fail_next_task = False + raise RuntimeError("Failed the Workflow Task once") + + +@activity.defn +async def append_from_local_activity(stream_id: str) -> None: + stream = activity.client().get_stream(StreamRef(id=stream_id), body_type=str) + await stream.append("a") + + +@workflow.defn +class UserLocalActivityAppendWorkflow: + @workflow.run + async def run(self, stream_id: str) -> None: + await workflow.execute_local_activity( + append_from_local_activity, + stream_id, + start_to_close_timeout=timedelta(seconds=10), + ) + + +class FlakyDescribeStore(RecordingStreamStore): + def __init__(self, *, failures: int = 0, delay: timedelta = timedelta()) -> None: + super().__init__(MemoryStreamStore()) + self._failures = failures + self._delay = delay + self.describes = 0 + + async def before_call(self, operation: str, request: Message) -> None: + if operation == "DescribeStream": + self.describes += 1 + await asyncio.sleep(self._delay.total_seconds()) + if self.describes <= self._failures: + raise RuntimeError("describe failed") + + +class FlakyCreateStore(CountingStore): + def __init__(self, name: str) -> None: + super().__init__(name) + self.failed_creates = 0 + + async def before_call(self, operation: str, request: Message) -> None: + if operation == "CreateStream" and not self.failed_creates: + self.failed_creates += 1 + raise RuntimeError("create failed") + + +class RecordingInterceptor(Interceptor): + def __init__(self) -> None: + self.activity_types: list[str] = [] + + def intercept_activity( + self, next: ActivityInboundInterceptor + ) -> ActivityInboundInterceptor: + return _RecordingActivityInbound(next, self.activity_types) + + +class _RecordingActivityInbound(ActivityInboundInterceptor): + def __init__( + self, next: ActivityInboundInterceptor, activity_types: list[str] + ) -> None: + super().__init__(next) + self._activity_types = activity_types + + async def execute_activity(self, input: ExecuteActivityInput) -> Any: + self._activity_types.append(activity.info().activity_type) + return await super().execute_activity(input) + + +def _with_store(client: Client, store: StreamStore) -> Client: + return _with_config(client, StreamConfig(stores=[store])) + + +def _with_config(client: Client, stream_config: StreamConfig) -> Client: + config = client.config() + config["stream_config"] = stream_config + return Client(**config) + + +def _two_counting_stores( + contexts: list[StreamSelectContext], +) -> tuple[CountingStore, CountingStore, StreamConfig]: + first = CountingStore("first") + second = CountingStore("second") + + def select(context: StreamSelectContext) -> StreamStore: + contexts.append(context) + return second + + return first, second, StreamConfig(stores=[first, second], store_selector=select) + + +async def _run_create_and_describe( + client: Client, stream_id: str +) -> tuple[WorkflowHandle[CreateAndDescribeWorkflow, StreamRef[str]], StreamRef[str]]: + async with new_worker(client, CreateAndDescribeWorkflow) as worker: + handle = await client.start_workflow( + CreateAndDescribeWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + return handle, await handle.result() + + +async def _read_all(client: Client, stream_id: str) -> list[StreamRecord[str]]: + stream = client.get_stream(StreamRef(id=stream_id), body_type=str) + return [record async for record in stream.iter_read_items()] + + +def _producers(records: list[StreamRecord[str]]) -> list[Producer]: + producers = [record.producer for record in records] + assert all(producer is not None for producer in producers) + return [producer for producer in producers if producer is not None] + + +async def _run_append_workflow( + client: Client, +) -> tuple[WorkflowHandle[AppendWorkflow, StreamRef[str]], StreamRef[str]]: + async with new_worker(client, AppendWorkflow, activities=[append_values]) as worker: + handle = await client.start_workflow( + AppendWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + return handle, await handle.result() + + +async def _describe_result( + client: Client, + store: StreamStore, + options: StreamStoreCallOptions | None = None, +) -> str: + stream_config = StreamConfig( + stores=[store], workflow_call_options=options or StreamStoreCallOptions() + ) + stream_client = _with_config(client, stream_config) + async with new_worker(stream_client, DescribeResultWorkflow) as worker: + return await stream_client.execute_workflow( + DescribeResultWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + + +async def test_workflow_stream_uses_store(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + handle, ref = await _run_append_workflow(stream_client) + + stream = stream_client.get_stream(ref, body_type=str) + records = [record async for record in stream.iter_read_items()] + assert [record.body for record in records] == _BODIES + run_id = handle.result_run_id + assert run_id + first, second = _producers(records) + assert first == second + # The create is the run's first Local Activity, so the append is its second. + assert first.id == f"la:{run_id}:2" + assert isinstance(first.kind, ProducerKindLocalActivity) + assert first.kind.value.run_id == run_id + assert first.kind.value.epoch.task.started_event_id > 0 + assert first.kind.value.epoch.attempt == 1 + + +async def test_workflow_stream_uses_a_store_written_from_scratch(client: Client): + store = DictStreamStore() + stream_client = _with_store(client, store) + stream_id = f"stream-{uuid.uuid4()}" + _, ref = await _run_create_and_describe(stream_client, stream_id) + assert ref == StreamRef(id=stream_id, store="dict") + assert list(store.streams) == [stream_id] + + stream = stream_client.get_stream(ref, body_type=str) + await stream.append("a") + assert [record.body async for record in stream.iter_read_items()] == ["a"] + + +async def test_workflow_stream_store_uses_payload_codec(client: Client): + store = MemoryStreamStore() + config = client.config() + config["data_converter"] = DataConverter(payload_codec=SimpleCodec()) + config["stream_config"] = StreamConfig(stores=[store]) + stream_client = Client(**config) + _, ref = await _run_append_workflow(stream_client) + + # A client without the codec reads the bodies as the Store keeps them. + raw = _with_store(client, store).get_stream( + StreamRef(id=ref.id), body_type=temporalio.common.RawValue + ) + stored = [record.body.payload async for record in raw.iter_read_items()] + assert [body.metadata["simple-codec"] for body in stored] == [b"true", b"true"] + stream = stream_client.get_stream(ref, body_type=str) + assert [record.body async for record in stream.iter_read_items()] == _BODIES + + +async def test_workflow_stream_history_has_no_bodies(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + handle, _ = await _run_append_workflow(stream_client) + events = (await handle.fetch_history()).events + markers = [ + event.marker_recorded_event_attributes.marker_name + for event in events + if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED + ] + # One marker for the create and one for the Local Activity that appends. + assert markers == ["core_local_activity", "core_local_activity"] + for event in events: + for body in _BODIES: + assert body.encode() not in event.SerializeToString() + + +async def test_workflow_stream_replay_does_not_call_store(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + handle, _ = await _run_append_workflow(stream_client) + replay_store = CountingStore("memory") + await Replayer( + workflows=[AppendWorkflow], + stream_config=StreamConfig(stores=[replay_store]), + ).replay_workflow(await handle.fetch_history()) + assert replay_store.operations == [] + + +async def test_workflow_create_without_store_fails_the_task(client: Client): + async with new_worker(client, CreateAndDescribeWorkflow) as worker: + handle = await client.start_workflow( + CreateAndDescribeWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=30), + ) + await assert_task_fail_eventually( + handle, + message_contains=( + "The client of this Worker has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...])." + ), + ) + await handle.terminate() + + +@pytest.mark.parametrize( + "operation, message", + [ + ( + "AppendRecords", + "Workflow code can't append to a Stream on a Store in this release. " + "Call a Local Activity that appends through its client.", + ), + ( + "CloseStream", + "Workflow code can't close a Stream on a Store in this release. " + "Call an Activity that closes it through its client.", + ), + ( + "ReadRecords", + "Workflow code can't read a Stream on a Store in this release. " + "Call an Activity that reads it through its client.", + ), + ], +) +async def test_workflow_stream_operation_fails_with_store( + client: Client, operation: str, message: str +): + stream_client = _with_store(client, MemoryStreamStore()) + async with new_worker(stream_client, RefusedOperationWorkflow) as worker: + result = await stream_client.execute_workflow( + RefusedOperationWorkflow.run, + operation, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == message + + +def test_workflow_stream_has_no_append_close_or_read(): + for name in ("append", "close", "read"): + assert not hasattr(WorkflowStream, name) + assert hasattr(WorkflowStream, "describe") + + +async def test_workflow_stream_store_error_reaches_workflow_code(client: Client): + store = FlakyDescribeStore(failures=1) + result = await _describe_result(client, store) + assert result == "ApplicationError RuntimeError: describe failed" + assert store.describes == 1 + + +async def test_workflow_describe_of_unknown_stream_raises_stream_error(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + async with new_worker(stream_client, DescribeStreamWorkflow) as worker: + reason = await stream_client.execute_workflow( + DescribeStreamWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert reason == "STREAM_ERROR_REASON_NOT_FOUND" + + +async def test_strict_create_in_each_run_of_a_chain_succeeds(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker( + stream_client, CreateThenContinueAsNewWorkflow, activities=[append_values] + ) as worker: + handle = await stream_client.start_workflow( + CreateThenContinueAsNewWorkflow.run, + args=[stream_id, 1], + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == [True] + first_run_id = handle.first_execution_run_id + last_run = await stream_client.get_workflow_handle(handle.id).describe() + assert first_run_id + assert last_run.run_id != first_run_id + records = await _read_all(stream_client, stream_id) + assert [record.body for record in records] == ["run-1", "run-2"] + first, second = _producers(records) + assert isinstance(first.kind, ProducerKindLocalActivity) + assert isinstance(second.kind, ProducerKindLocalActivity) + assert first.kind.value.run_id == first_run_id + assert second.kind.value.run_id == last_run.run_id + + +async def test_client_append_after_local_activity_append_keeps_offset_order( + client: Client, +): + stream_client = _with_store(client, MemoryStreamStore()) + _, ref = await _run_append_workflow(stream_client) + stream = stream_client.get_stream(ref, body_type=str) + appended = await stream.append("from-client") + records = await _read_all(stream_client, ref.id) + assert [record.body for record in records] == [*_BODIES, "from-client"] + offsets = [record.offset for record in records] + assert offsets == sorted(offsets) + assert len(set(offsets)) == len(offsets) + assert offsets[-1] == appended.last_offset + kinds = [type(producer.kind) for producer in _producers(records)] + assert kinds == [ + ProducerKindLocalActivity, + ProducerKindLocalActivity, + ProducerKindClient, + ] + + +async def test_workflow_stream_store_call_options_retry(client: Client): + options = StreamStoreCallOptions( + retry_policy=RetryPolicy( + initial_interval=timedelta(milliseconds=10), maximum_attempts=2 + ) + ) + store = FlakyDescribeStore(failures=1) + assert await _describe_result(client, store, options) == "described" + assert store.describes == 2 + + +async def test_workflow_stream_store_call_options_timeout(client: Client): + options = StreamStoreCallOptions(start_to_close_timeout=timedelta(milliseconds=200)) + store = FlakyDescribeStore(delay=timedelta(seconds=5)) + result = await _describe_result(client, store, options) + assert result == "TimeoutError START_TO_CLOSE" + + +async def test_worker_keeps_its_stream_stores(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + async with new_worker(stream_client, AppendWorkflow) as worker: + with pytest.raises(ValueError, match="Stream Stores"): + worker.client = client + worker.client = stream_client + + +async def test_activity_interceptors_see_stream_store_activity(client: Client): + interceptor = RecordingInterceptor() + stream_client = _with_store(client, MemoryStreamStore()) + async with new_worker( + stream_client, + CreateAndActivityWorkflow, + activities=[user_activity], + interceptors=[interceptor], + ) as worker: + await stream_client.execute_workflow( + CreateAndActivityWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + # The create runs as the SDK's Local Activity, before the Workflow's own Activity. + assert interceptor.activity_types == [ + "__temporal_stream_store_call:memory", + "user_activity", + ] + + +async def test_workflow_task_retry_appends_with_higher_epoch(client: Client): + global _fail_next_task + _fail_next_task = True + stream_client = _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker( + stream_client, AppendThenFailTaskOnceWorkflow, activities=[append_values] + ) as worker: + await stream_client.execute_workflow( + AppendThenFailTaskOnceWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + records = await _read_all(stream_client, stream_id) + # The failed Workflow Task records no marker, so its retry runs the Local Activity + # again and writes a second copy. + assert [record.body for record in records] == [*_BODIES, *_BODIES] + producers = _producers(records) + assert len({producer.id for producer in producers}) == 1 + + def epoch(producer: Producer) -> tuple[int, int, int]: + assert isinstance(producer.kind, ProducerKindLocalActivity) + local = producer.kind.value.epoch + return (local.task.started_event_id, local.task.attempt, local.attempt) + + first, second, third, fourth = (epoch(producer) for producer in producers) + assert first == second + assert third == fourth + assert third > first + + +async def test_user_local_activity_append_has_local_activity_producer(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream = await stream_client.create_stream( + id=f"stream-{uuid.uuid4()}", body_type=str + ) + async with new_worker( + stream_client, + UserLocalActivityAppendWorkflow, + activities=[append_from_local_activity], + ) as worker: + handle = await stream_client.start_workflow( + UserLocalActivityAppendWorkflow.run, + stream.ref.id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.result() + [producer] = _producers(await _read_all(stream_client, stream.ref.id)) + run_id = handle.result_run_id + assert producer.id == f"la:{run_id}:1" + assert isinstance(producer.kind, ProducerKindLocalActivity) + assert producer.kind.value.run_id == run_id + assert producer.kind.value.epoch.task.started_event_id > 0 + assert producer.kind.value.epoch.attempt == 1 + + +async def test_workflow_ref_of_another_store_fails_the_task(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + async with new_worker(stream_client, DescribeOtherStoreWorkflow) as worker: + handle = await stream_client.start_workflow( + DescribeOtherStoreWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=30), + ) + await assert_task_fail_eventually( + handle, message_contains="names the Stream Store 'other'" + ) + await handle.terminate() + + +async def test_workflow_create_goes_to_the_selected_store(client: Client): + contexts: list[StreamSelectContext] = [] + first, second, stream_config = _two_counting_stores(contexts) + stream_id = f"stream-{uuid.uuid4()}" + handle, ref = await _run_create_and_describe( + _with_config(client, stream_config), stream_id + ) + assert ref.store == "second" + assert first.operations == [] + assert second.operations == ["CreateStream", "DescribeStream"] + # The describe goes by the Store that the ref names, so the selector runs once. + assert contexts == [ + StreamSelectContext( + id=stream_id, + target=StorageDriverWorkflowInfo( + id=handle.id, + type="CreateAndDescribeWorkflow", + run_id=handle.result_run_id, + namespace=client.namespace, + ), + ) + ] + + +async def test_workflow_replay_calls_no_store(client: Client): + _, _, stream_config = _two_counting_stores([]) + stream_id = f"stream-{uuid.uuid4()}" + handle, _ = await _run_create_and_describe( + _with_config(client, stream_config), stream_id + ) + contexts: list[StreamSelectContext] = [] + first, second, replay_config = _two_counting_stores(contexts) + await Replayer( + workflows=[CreateAndDescribeWorkflow], stream_config=replay_config + ).replay_workflow(await handle.fetch_history()) + # The selector runs in Workflow code, so a replay runs it again with the Workflow. + assert [context.id for context in contexts] == [stream_id] + target = contexts[0].target + assert isinstance(target, StorageDriverWorkflowInfo) + assert (target.id, target.run_id) == (handle.id, handle.result_run_id) + assert first.operations == [] + assert second.operations == [] + + +async def test_workflow_create_with_several_stores_takes_the_config_options( + client: Client, +): + first = CountingStore("first") + second = FlakyCreateStore("second") + stream_config = StreamConfig( + stores=[first, second], + store_selector=lambda _: second, + workflow_call_options=StreamStoreCallOptions( + retry_policy=RetryPolicy( + initial_interval=timedelta(milliseconds=10), maximum_attempts=2 + ) + ), + ) + # The default options make one attempt, so only the config's retry policy lets + # the create pass after its first attempt fails. + _, ref = await _run_create_and_describe( + _with_config(client, stream_config), f"stream-{uuid.uuid4()}" + ) + assert ref.store == "second" + assert first.operations == [] + assert second.operations == ["CreateStream", "CreateStream", "DescribeStream"] + + +class OwnActivityStore(CountingStore): + """A Store that brings its own Activity for the calls of Workflow code.""" + + def __init__(self) -> None: + super().__init__("own") + self.workflow_calls: list[str] = [] + + def worker_activities(self) -> Sequence[Callable[..., Any]]: + @activity.defn(name="own_store_call") + async def own_store_call(operation: str, request: bytes) -> bytes: + if operation == "CreateStream": + created = await self.create_stream( + CreateStreamRequestMessage.FromString(request) + ) + return created.SerializeToString() + described = await self.describe_stream( + DescribeStreamRequestMessage.FromString(request) + ) + return described.SerializeToString() + + return [own_store_call] + + async def call_from_workflow( + self, operation: str, request: Message, options: StreamStoreCallOptions + ) -> Message: + self.workflow_calls.append(operation) + response = await workflow.execute_local_activity( + "own_store_call", + args=[operation, request.SerializeToString()], + result_type=bytes, + start_to_close_timeout=options.start_to_close_timeout, + ) + if operation == "CreateStream": + return CreateStreamResponseMessage.FromString(response) + return DescribeStreamResponseMessage.FromString(response) + + +async def test_workflow_call_goes_through_the_store_own_activity(client: Client): + store = OwnActivityStore() + interceptor = RecordingInterceptor() + stream_client = _with_store(client, store) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker( + stream_client, CreateAndDescribeWorkflow, interceptors=[interceptor] + ) as worker: + ref = await stream_client.execute_workflow( + CreateAndDescribeWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert ref == StreamRef(id=stream_id, store="own") + assert store.workflow_calls == ["CreateStream", "DescribeStream"] + assert store.operations == ["CreateStream", "DescribeStream"] + assert interceptor.activity_types == ["own_store_call", "own_store_call"] From 5e4c62bd7d67a9b419e477b35f61659361a24c3d Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 18:55:01 -0700 Subject: [PATCH 2/9] Moved Core to the Stream WIT with a Workflow append and close. The Stream WIT now gives the Workflow target `append` and `close`, and keeps `read` out of it. --- temporalio/bridge/sdk-core | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index 9a110b3cda..59f4b33f1d 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit 9a110b3cda949199b7792348ff5d265b1d1c414a +Subproject commit 59f4b33f1d90bb5ac4582ea465b4e8597c507214 From 7841bf0923675c26bba0d00af89418e06a60e3f5 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 18:59:53 -0700 Subject: [PATCH 3/9] Ran a Workflow's Stream append in the Store's Local Activity. The Workflow `Stream` has `append` again, and Workflow code can't do I/O. The Local Activity writes as its own Producer with a fixed attempt, so a retry in one Workflow Task writes once. --- temporalio/client/_stream_store.py | 52 +++-- temporalio/worker/_stream_store.py | 6 +- tests/nexus/test_stream_store_workflow.py | 239 +++++++++++++++++++--- 3 files changed, 250 insertions(+), 47 deletions(-) diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index 622f00bf18..21183b750e 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -15,6 +15,8 @@ import temporalio.bridge.runtime import temporalio.common import temporalio.converter +import temporalio.converter._payload_converter +import temporalio.nexus.system._support import temporalio.nexus.system.stream_service import temporalio.runtime import temporalio.workflow @@ -49,7 +51,9 @@ class StreamStoreCallOptions: retry_policy: temporalio.common.RetryPolicy = dataclasses.field( default_factory=lambda: temporalio.common.RetryPolicy(maximum_attempts=1) ) - """Retries of a failed call. The default is one attempt.""" + """Retries of a failed call. The default is one attempt. In one Workflow Task, each + attempt of a Workflow's append without an explicit ``producer`` sends the same + Producer, epoch and sequence, so the Store writes the batch once.""" def __post_init__(self) -> None: """Check that a timeout is set.""" @@ -152,8 +156,8 @@ async def call_from_workflow( calls no method of this Store. Args: - operation: The Nexus operation name, ``CreateStream`` or ``DescribeStream`` - in this release. + operation: The Nexus operation name, ``CreateStream``, ``AppendRecords`` + or ``DescribeStream`` in this release. request: The request message of ``temporalio.api.streamservice.v1``. options: The :py:attr:`StreamConfig.workflow_call_options` of the config. @@ -161,8 +165,8 @@ async def call_from_workflow( The response message of the operation. """ stream_operation = _STREAM_OPERATIONS[operation] - # Workflow code can start only a create and a describe in this release, and their - # messages carry no payload. So the Local Activity has no payload codec step. + # Workflow code runs no payload codec, so the Activity encodes the bodies of an + # append before the Store gets them. response = await temporalio.workflow.execute_local_activity( _store_call_activity_name(self.name), args=[operation, request.SerializeToString()], @@ -297,10 +301,6 @@ def check_response( "append_records", temporalio.api.streamservice.v1.AppendRecordsRequest, temporalio.api.streamservice.v1.AppendRecordsResponse, - workflow_refusal=( - "Workflow code can't append to a Stream on a Store in this release. " - "Call a Local Activity that appends through its client." - ), ), "ReadRecords": _StreamOperation( "read_records", @@ -341,15 +341,43 @@ def _store_call_activity( ) -> Callable[[str, bytes], Awaitable[bytes]]: @temporalio.activity.defn(name=_store_call_activity_name(store.name)) async def call(operation: str, request: bytes) -> bytes: - stream_operation = _STREAM_OPERATIONS[operation] - response = await stream_operation.call( - store, stream_operation.request_type.FromString(request) + message = _STREAM_OPERATIONS[operation].request_type.FromString(request) + if isinstance(message, temporalio.api.streamservice.v1.AppendRecordsRequest): + _with_local_activity_producer(message) + response = await _call_stream_store( + store, temporalio.activity.client().data_converter, operation, message ) return response.SerializeToString() return call +def _with_local_activity_producer( + request: temporalio.api.streamservice.v1.AppendRecordsRequest, +) -> None: + # On a Store, a Workflow appends as the Local Activity that runs the append, so the + # Store can fence the write of a failed Workflow Task. A Producer that Workflow code + # gave of another kind reaches the Store as it is. + if ( + request.producer.WhichOneof("kind") != "workflow" + or not temporalio.activity.info().is_local + ): + return + converter = temporalio.converter._payload_converter._get_transfer_type_converter( + temporalio.nexus.system.stream_service.Producer + ) + if converter is None: + raise TypeError("The generated Producer has no proto transfer type") + request.producer.CopyFrom( + converter.to_transfer_type(temporalio.nexus.system._support.default_producer()) + ) + # Each retry of this Local Activity in one Workflow Task sends the same Producer, epoch + # and sequence, so the Store writes the batch once and gives back its first offsets. + request.producer.local_activity.epoch.attempt = 1 + if not request.sequence: + request.sequence = 1 + + @dataclasses.dataclass(frozen=True) class StreamSelectContext: """Context passed to :py:attr:`StreamConfig.store_selector` calls. diff --git a/temporalio/worker/_stream_store.py b/temporalio/worker/_stream_store.py index da9faa689c..bd134bf4b1 100644 --- a/temporalio/worker/_stream_store.py +++ b/temporalio/worker/_stream_store.py @@ -54,10 +54,8 @@ def start_operation( ) -> temporalio.workflow.NexusOperationHandle[OutputT]: """Answer a Stream call of Workflow code with the Store that the call goes to.""" stream_operation = _STREAM_OPERATIONS[input.operation_name] - # The Workflow `Stream` has no `append`, `close` or `read`, but the generated - # Workflow service still has their operations, so Workflow code can start one - # through a Nexus client. In this release, Workflow code writes and reads a Stream - # on a Store through its own Activities. + # Workflow code can start any operation of the generated service through a Nexus + # client, also one that the Workflow `Stream` leaves out. if stream_operation.workflow_refusal is not None: raise temporalio.exceptions.ApplicationError( stream_operation.workflow_refusal, non_retryable=True diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py index 7913a60498..f6c520e1a3 100644 --- a/tests/nexus/test_stream_store_workflow.py +++ b/tests/nexus/test_stream_store_workflow.py @@ -9,6 +9,7 @@ import pytest from google.protobuf.message import Message +import temporalio.api.streamservice.v1 import temporalio.common import temporalio.exceptions from temporalio import activity, workflow @@ -40,6 +41,7 @@ from temporalio.nexus.system.stream_service import ( AppendRecord, AppendRecordsRequest, + AppendRecordsResponse, ClientProducer, CloseStreamRequest, Producer, @@ -84,6 +86,15 @@ async def _append_in_local_activity(ref: StreamRef[str], values: list[str]) -> N @workflow.defn class AppendWorkflow: + @workflow.run + async def run(self, stream_id: str) -> StreamRef[str]: + stream = await create_stream(id=stream_id, body_type=str) + await stream.append(*_BODIES) + return stream.ref + + +@workflow.defn +class LocalActivityAppendWorkflow: @workflow.run async def run(self, stream_id: str) -> StreamRef[str]: stream = await create_stream(id=stream_id, body_type=str) @@ -91,6 +102,44 @@ async def run(self, stream_id: str) -> StreamRef[str]: return stream.ref +# Workflow code appends either with the Workflow `Stream` or with a Local Activity of +# its own, and both write as the Local Activity that runs the append. +_APPEND_WORKFLOWS = pytest.mark.parametrize( + "append_workflow", [AppendWorkflow, LocalActivityAppendWorkflow] +) + + +@workflow.defn +class AppendOffsetsWorkflow: + @workflow.run + async def run(self, stream_id: str) -> AppendRecordsResponse: + stream = await create_stream(id=stream_id, body_type=str) + return await stream.append(*_BODIES) + + +@workflow.defn +class AppendToStreamWorkflow: + @workflow.run + async def run(self, stream_id: str) -> str: + stream = get_stream(StreamRef(id=stream_id), body_type=str) + try: + await stream.append("a") + except StreamError as err: + assert isinstance(err, ApplicationError) + return StreamErrorReason.Name(err.reason) + return "appended" + + +@workflow.defn +class ExplicitProducerWorkflow: + @workflow.run + async def run(self, stream_id: str) -> None: + stream = await create_stream(id=stream_id, body_type=str) + await stream.append( + "a", producer=Producer(id="mine", kind=ProducerKindClient(ClientProducer())) + ) + + @workflow.defn class RefusedOperationWorkflow: @workflow.run @@ -201,13 +250,16 @@ async def run(self, stream_id: str) -> None: @workflow.defn(sandboxed=False) class AppendThenFailTaskOnceWorkflow: @workflow.run - async def run(self, stream_id: str) -> None: + async def run(self, stream_id: str, in_local_activity: bool) -> None: stream = await create_stream( id=stream_id, body_type=str, id_conflict_policy=StreamIdConflictPolicy.StreamIdConflictPolicyUseExisting, ) - await _append_in_local_activity(stream.ref, _BODIES) + if in_local_activity: + await _append_in_local_activity(stream.ref, _BODIES) + else: + await stream.append(*_BODIES) global _fail_next_task if _fail_next_task: _fail_next_task = False @@ -257,6 +309,34 @@ async def before_call(self, operation: str, request: Message) -> None: raise RuntimeError("create failed") +class LostAnswerStore(RecordingStreamStore): + """Keeps each append request, and can lose the answer of the first append.""" + + def __init__(self, *, lose_first_answer: bool = True) -> None: + super().__init__(MemoryStreamStore()) + self._lose_first_answer = lose_first_answer + self.appends: list[temporalio.api.streamservice.v1.AppendRecordsRequest] = [] + self.first_response: ( + temporalio.api.streamservice.v1.AppendRecordsResponse | None + ) = None + + async def append_records( + self, request: temporalio.api.streamservice.v1.AppendRecordsRequest + ) -> temporalio.api.streamservice.v1.AppendRecordsResponse: + kept = temporalio.api.streamservice.v1.AppendRecordsRequest() + kept.CopyFrom(request) + self.appends.append(kept) + response = await super().append_records(request) + if self.first_response is None: + self.first_response = response + if self._lose_first_answer: + raise StreamError( + "Lost the answer of a write", + reason=StreamErrorReason.STREAM_ERROR_REASON_STORAGE_UNAVAILABLE, + ) + return response + + class RecordingInterceptor(Interceptor): def __init__(self) -> None: self.activity_types: list[str] = [] @@ -327,11 +407,13 @@ def _producers(records: list[StreamRecord[str]]) -> list[Producer]: async def _run_append_workflow( - client: Client, -) -> tuple[WorkflowHandle[AppendWorkflow, StreamRef[str]], StreamRef[str]]: - async with new_worker(client, AppendWorkflow, activities=[append_values]) as worker: - handle = await client.start_workflow( - AppendWorkflow.run, + client: Client, append_workflow: type[Any] = AppendWorkflow +) -> tuple[WorkflowHandle[Any, StreamRef[str]], StreamRef[str]]: + async with new_worker( + client, append_workflow, activities=[append_values] + ) as worker: + handle: WorkflowHandle[Any, StreamRef[str]] = await client.start_workflow( + append_workflow.run, f"stream-{uuid.uuid4()}", id=f"workflow-{uuid.uuid4()}", task_queue=worker.task_queue, @@ -357,9 +439,10 @@ async def _describe_result( ) -async def test_workflow_stream_uses_store(client: Client): +@_APPEND_WORKFLOWS +async def test_workflow_stream_uses_store(client: Client, append_workflow: type[Any]): stream_client = _with_store(client, MemoryStreamStore()) - handle, ref = await _run_append_workflow(stream_client) + handle, ref = await _run_append_workflow(stream_client, append_workflow) stream = stream_client.get_stream(ref, body_type=str) records = [record async for record in stream.iter_read_items()] @@ -389,13 +472,16 @@ async def test_workflow_stream_uses_a_store_written_from_scratch(client: Client) assert [record.body async for record in stream.iter_read_items()] == ["a"] -async def test_workflow_stream_store_uses_payload_codec(client: Client): +@_APPEND_WORKFLOWS +async def test_workflow_stream_store_uses_payload_codec( + client: Client, append_workflow: type[Any] +): store = MemoryStreamStore() config = client.config() config["data_converter"] = DataConverter(payload_codec=SimpleCodec()) config["stream_config"] = StreamConfig(stores=[store]) stream_client = Client(**config) - _, ref = await _run_append_workflow(stream_client) + _, ref = await _run_append_workflow(stream_client, append_workflow) # A client without the codec reads the bodies as the Store keeps them. raw = _with_store(client, store).get_stream( @@ -407,28 +493,34 @@ async def test_workflow_stream_store_uses_payload_codec(client: Client): assert [record.body async for record in stream.iter_read_items()] == _BODIES -async def test_workflow_stream_history_has_no_bodies(client: Client): +@_APPEND_WORKFLOWS +async def test_workflow_stream_history_has_no_bodies( + client: Client, append_workflow: type[Any] +): stream_client = _with_store(client, MemoryStreamStore()) - handle, _ = await _run_append_workflow(stream_client) + handle, _ = await _run_append_workflow(stream_client, append_workflow) events = (await handle.fetch_history()).events markers = [ event.marker_recorded_event_attributes.marker_name for event in events if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED ] - # One marker for the create and one for the Local Activity that appends. + # One marker for the create and one for the append. assert markers == ["core_local_activity", "core_local_activity"] for event in events: for body in _BODIES: assert body.encode() not in event.SerializeToString() -async def test_workflow_stream_replay_does_not_call_store(client: Client): +@_APPEND_WORKFLOWS +async def test_workflow_stream_replay_does_not_call_store( + client: Client, append_workflow: type[Any] +): stream_client = _with_store(client, MemoryStreamStore()) - handle, _ = await _run_append_workflow(stream_client) + handle, _ = await _run_append_workflow(stream_client, append_workflow) replay_store = CountingStore("memory") await Replayer( - workflows=[AppendWorkflow], + workflows=[append_workflow], stream_config=StreamConfig(stores=[replay_store]), ).replay_workflow(await handle.fetch_history()) assert replay_store.operations == [] @@ -456,11 +548,6 @@ async def test_workflow_create_without_store_fails_the_task(client: Client): @pytest.mark.parametrize( "operation, message", [ - ( - "AppendRecords", - "Workflow code can't append to a Stream on a Store in this release. " - "Call a Local Activity that appends through its client.", - ), ( "CloseStream", "Workflow code can't close a Stream on a Store in this release. " @@ -487,10 +574,10 @@ async def test_workflow_stream_operation_fails_with_store( assert result == message -def test_workflow_stream_has_no_append_close_or_read(): - for name in ("append", "close", "read"): - assert not hasattr(WorkflowStream, name) - assert hasattr(WorkflowStream, "describe") +def test_workflow_stream_has_append_and_close_and_no_read(): + for name in ("append", "close", "describe", "__aenter__", "__aexit__"): + assert hasattr(WorkflowStream, name) + assert not hasattr(WorkflowStream, "read") async def test_workflow_stream_store_error_reaches_workflow_code(client: Client): @@ -538,11 +625,12 @@ async def test_strict_create_in_each_run_of_a_chain_succeeds(client: Client): assert second.kind.value.run_id == last_run.run_id -async def test_client_append_after_local_activity_append_keeps_offset_order( - client: Client, +@_APPEND_WORKFLOWS +async def test_client_append_after_workflow_append_keeps_offset_order( + client: Client, append_workflow: type[Any] ): stream_client = _with_store(client, MemoryStreamStore()) - _, ref = await _run_append_workflow(stream_client) + _, ref = await _run_append_workflow(stream_client, append_workflow) stream = stream_client.get_stream(ref, body_type=str) appended = await stream.append("from-client") records = await _read_all(stream_client, ref.id) @@ -607,7 +695,10 @@ async def test_activity_interceptors_see_stream_store_activity(client: Client): ] -async def test_workflow_task_retry_appends_with_higher_epoch(client: Client): +@pytest.mark.parametrize("in_local_activity", [False, True]) +async def test_workflow_task_retry_appends_with_higher_epoch( + client: Client, in_local_activity: bool +): global _fail_next_task _fail_next_task = True stream_client = _with_store(client, MemoryStreamStore()) @@ -617,7 +708,7 @@ async def test_workflow_task_retry_appends_with_higher_epoch(client: Client): ) as worker: await stream_client.execute_workflow( AppendThenFailTaskOnceWorkflow.run, - stream_id, + args=[stream_id, in_local_activity], id=f"workflow-{uuid.uuid4()}", task_queue=worker.task_queue, ) @@ -804,3 +895,89 @@ async def test_workflow_call_goes_through_the_store_own_activity(client: Client) assert store.workflow_calls == ["CreateStream", "DescribeStream"] assert store.operations == ["CreateStream", "DescribeStream"] assert interceptor.activity_types == ["own_store_call", "own_store_call"] + + +async def test_workflow_append_returns_offsets_of_records_in_order(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker(stream_client, AppendOffsetsWorkflow) as worker: + appended = await stream_client.execute_workflow( + AppendOffsetsWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + records = await _read_all(stream_client, stream_id) + assert [record.body for record in records] == _BODIES + assert [appended.first_offset, appended.last_offset] == [ + record.offset for record in records + ] + + +async def test_workflow_append_to_closed_stream_raises_stream_error(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream = await stream_client.create_stream( + id=f"stream-{uuid.uuid4()}", body_type=str + ) + await stream.close() + async with new_worker(stream_client, AppendToStreamWorkflow) as worker: + reason = await stream_client.execute_workflow( + AppendToStreamWorkflow.run, + stream.ref.id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert reason == "STREAM_ERROR_REASON_CLOSED" + + +async def test_workflow_append_retry_in_one_task_writes_once(client: Client): + store = LostAnswerStore() + stream_config = StreamConfig( + stores=[store], + workflow_call_options=StreamStoreCallOptions( + retry_policy=RetryPolicy( + initial_interval=timedelta(milliseconds=10), maximum_attempts=2 + ) + ), + ) + stream_client = _with_config(client, stream_config) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker(stream_client, AppendOffsetsWorkflow) as worker: + appended = await stream_client.execute_workflow( + AppendOffsetsWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + first, second = store.appends + assert first.producer == second.producer + assert first.producer.local_activity.epoch.attempt == 1 + assert first.sequence == second.sequence == 1 + records = await _read_all(stream_client, stream_id) + assert [record.body for record in records] == _BODIES + assert store.first_response is not None + assert appended.first_offset == store.first_response.first_offset + assert appended.last_offset == store.first_response.last_offset + assert [appended.first_offset, appended.last_offset] == [ + record.offset for record in records + ] + + +async def test_workflow_explicit_producer_reaches_store(client: Client): + store = LostAnswerStore(lose_first_answer=False) + stream_client = _with_store(client, store) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker(stream_client, ExplicitProducerWorkflow) as worker: + await stream_client.execute_workflow( + ExplicitProducerWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + [producer] = _producers(await _read_all(stream_client, stream_id)) + assert producer.id == "mine" + assert isinstance(producer.kind, ProducerKindClient) + [request] = store.appends + assert request.producer.id == "mine" + assert request.producer.WhichOneof("kind") == "client" + assert request.sequence == 0 From fc6b8d3b3faffdcd044a6ce59056653adecf0f95 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Fri, 13 Nov 2026 19:03:03 -0700 Subject: [PATCH 4/9] Ran a Workflow's Stream close in an Activity after its Workflow Task. A Local Activity reaches the Store before its Workflow Task completes, so a failed task would leave a closed Stream. A Worker with Stream Stores now polls for Activity Tasks, so it can run the close. --- changelog/added/chattering-ledger-otter.md | 2 +- temporalio/client/_stream_store.py | 71 ++++-- temporalio/worker/_worker.py | 9 +- tests/contrib/redis/test_streams.py | 35 +++ tests/nexus/test_stream_store_workflow.py | 251 ++++++++++++++++++--- 5 files changed, 306 insertions(+), 62 deletions(-) diff --git a/changelog/added/chattering-ledger-otter.md b/changelog/added/chattering-ledger-otter.md index f15fdc478b..b628c9c7f8 100644 --- a/changelog/added/chattering-ledger-otter.md +++ b/changelog/added/chattering-ledger-otter.md @@ -1 +1 @@ -Add the experimental Streams API. Client and Activity code create a Stream, append Records to it, read them and close it, and Workflow code creates and describes one. A client takes its Stream Stores in `stream_config=StreamConfig(stores=[...])`, such as `temporalio.contrib.redis.streams.RedisStreamStore`, or `temporalio.testing.MemoryStreamStore` in tests. +Add the experimental Streams API. Client, Activity and Workflow code create a Stream, append Records to it and close it, and client and Activity code read it. A client takes its Stream Stores in `stream_config=StreamConfig(stores=[...])`, such as `temporalio.contrib.redis.streams.RedisStreamStore`, or `temporalio.testing.MemoryStreamStore` in tests. diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index 21183b750e..bf68c129b1 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -29,10 +29,12 @@ @dataclasses.dataclass(frozen=True) class StreamStoreCallOptions: - """Options of the Local Activity that calls a Stream Store for Workflow code. + """Options of the Activity that calls a Stream Store for Workflow code. Workflow code can't do I/O, so the default - :py:meth:`StreamStore.call_from_workflow` calls its Store from a Local Activity. + :py:meth:`StreamStore.call_from_workflow` calls its Store from an Activity. A close + runs it as a regular Activity, and every other operation runs it as a Local + Activity. Both get these options. A :py:class:`StreamConfig` holds them in :py:attr:`StreamConfig.workflow_call_options`, and every Store of that config @@ -137,8 +139,9 @@ def worker_activities(self) -> Sequence[Callable[..., Any]]: """The Activities that a Worker registers for this Store when it runs Workflows. The default is one Activity, named for this Store, that - :py:meth:`call_from_workflow` runs as a Local Activity. A Store that overrides - one of the two methods overrides the other to match. + :py:meth:`call_from_workflow` runs as a Local Activity, or as a regular Activity + for a close. A Store that overrides one of the two methods overrides the other to + match. """ return [_store_call_activity(self)] @@ -151,13 +154,15 @@ async def call_from_workflow( """Answer one Stream operation for Workflow code. It runs in Workflow code, so it does no I/O itself. The default runs the - Activity of :py:meth:`worker_activities` as a Local Activity, with the timeouts - and the retry policy of ``options``. A replay takes the result from History and - calls no method of this Store. + Activity of :py:meth:`worker_activities` with the timeouts and the retry policy + of ``options``. A close runs it as a regular Activity on the Workflow's Task + Queue, after the Workflow Task completes. Every other operation runs it as a + Local Activity. A replay takes the result from History and calls no method of + this Store. Args: - operation: The Nexus operation name, ``CreateStream``, ``AppendRecords`` - or ``DescribeStream`` in this release. + operation: The Nexus operation name, ``CreateStream``, ``AppendRecords``, + ``CloseStream`` or ``DescribeStream`` in this release. request: The request message of ``temporalio.api.streamservice.v1``. options: The :py:attr:`StreamConfig.workflow_call_options` of the config. @@ -167,14 +172,29 @@ async def call_from_workflow( stream_operation = _STREAM_OPERATIONS[operation] # Workflow code runs no payload codec, so the Activity encodes the bodies of an # append before the Store gets them. - response = await temporalio.workflow.execute_local_activity( - _store_call_activity_name(self.name), - args=[operation, request.SerializeToString()], - result_type=bytes, - schedule_to_close_timeout=options.schedule_to_close_timeout, - start_to_close_timeout=options.start_to_close_timeout, - retry_policy=options.retry_policy, - ) + if operation == "CloseStream": + # A Local Activity reaches the Store before its Workflow Task completes. If + # that task then failed, the Stream would stay closed and the retry of the + # task could not append. The server schedules a regular Activity only when + # the task completes. + response = await temporalio.workflow.execute_activity( + _store_call_activity_name(self.name), + args=[operation, request.SerializeToString()], + result_type=bytes, + task_queue=temporalio.workflow.info().task_queue, + schedule_to_close_timeout=options.schedule_to_close_timeout, + start_to_close_timeout=options.start_to_close_timeout, + retry_policy=options.retry_policy, + ) + else: + response = await temporalio.workflow.execute_local_activity( + _store_call_activity_name(self.name), + args=[operation, request.SerializeToString()], + result_type=bytes, + schedule_to_close_timeout=options.schedule_to_close_timeout, + start_to_close_timeout=options.start_to_close_timeout, + retry_policy=options.retry_policy, + ) return stream_operation.response_type.FromString(response) @@ -315,10 +335,6 @@ def check_response( "close_stream", temporalio.api.streamservice.v1.CloseStreamRequest, temporalio.api.streamservice.v1.CloseStreamResponse, - workflow_refusal=( - "Workflow code can't close a Stream on a Store in this release. " - "Call an Activity that closes it through its client." - ), ), "DescribeStream": _StreamOperation( "describe_stream", @@ -411,6 +427,10 @@ class StreamConfig: Store of that name. A call whose ref names no Store, such as a create, goes to the one Store, or to the Store that :py:attr:`store_selector` returns. + A Worker with Workflows and this config polls for Activity Tasks, since a Workflow's + close runs the Store's Activity as a regular Activity. With ``no_remote_activities`` + on that Worker, another Worker of the Task Queue must run the close. + .. warning:: This API is experimental and subject to change. """ @@ -431,8 +451,8 @@ class StreamConfig: ``None`` and there is one Store, every new Stream goes to that Store. For a call of Workflow code, the selector runs in Workflow code, also on a - replay. So it must give the same Store for the same context. The Local Activity - of the chosen Store is in History, and another answer on a replay is a + replay. So it must give the same Store for the same context. The Activity of the + chosen Store is in History, and another answer on a replay is a non-determinism error. A Stream that code builds from an id only goes through the selector on each call, so the selector must also give the same Store for the same id. @@ -443,8 +463,9 @@ class StreamConfig: ) """Options of a Store's call for Workflow code, for every Store of this config. - :py:meth:`StreamStore.call_from_workflow` gets them. The default runs its Local - Activity with their timeouts and retry policy. + :py:meth:`StreamStore.call_from_workflow` gets them. The default runs the Store's + Activity with their timeouts and retry policy, as a Local Activity, or as a regular + Activity for a close. """ _store_map: dict[str, StreamStore] = dataclasses.field( diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index bbf68e2bf4..4ecce29d15 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -668,14 +668,15 @@ def check_activity(activity: str): # are not running an activity worker at all. Otherwise shutdown # will not proceed properly. no_remote_activities=config.get("no_remote_activities") - or not config.get("activities"), + or self._activity_worker is None, task_types=temporalio.bridge.worker.WorkerTaskTypes( enable_workflows=self._workflow_worker is not None, enable_local_activities=self._activity_worker is not None and self._workflow_worker is not None, - # The Activities of the Stream Stores run as Local Activities, so - # alone they don't make this worker poll for Activity Tasks. - enable_remote_activities=bool(config.get("activities")) + # A Workflow's Stream close runs a Store's Activity as a regular + # Activity, so the Activities of the Stream Stores alone make this + # worker poll for Activity Tasks. + enable_remote_activities=self._activity_worker is not None and not config.get("no_remote_activities"), enable_nexus=self._nexus_worker is not None, ), diff --git a/tests/contrib/redis/test_streams.py b/tests/contrib/redis/test_streams.py index 0bae0921f1..95a5b3a37e 100644 --- a/tests/contrib/redis/test_streams.py +++ b/tests/contrib/redis/test_streams.py @@ -74,6 +74,16 @@ async def run(self, stream_id: str) -> StreamRef[str]: return stream.ref +@workflow.defn +class AppendAndCloseWorkflow: + @workflow.run + async def run(self, stream_id: str) -> StreamRef[str]: + stream = await create_stream(id=stream_id, body_type=str) + await stream.append("a", "b") + await stream.close() + return stream.ref + + @activity.defn async def append_in_activity(stream_id: str) -> None: stream = activity.client().get_stream(StreamRef(id=stream_id), body_type=str) @@ -140,6 +150,31 @@ async def test_workflow_local_activity_appends_with_local_activity_producer( assert record.producer.kind.value.run_id == run_id +async def test_workflow_appends_and_closes(client: Client): + stream_client = _with_store(client, _store()) + async with new_worker(stream_client, AppendAndCloseWorkflow) as worker: + handle = await stream_client.start_workflow( + AppendAndCloseWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + ref = await handle.result() + + stream = stream_client.get_stream(ref, body_type=str) + records = [record async for record in stream.iter_read_items()] + assert [record.body for record in records] == ["a", "b"] + run_id = handle.result_run_id + for record in records: + assert record.producer is not None + assert isinstance(record.producer.kind, ProducerKindLocalActivity) + assert record.producer.kind.value.run_id == run_id + described = await stream.describe() + assert described.status == StreamStatus.StreamStatusClosed + err = await _stream_error(stream.append("c")) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_CLOSED + + async def test_activity_appends_with_activity_producer(client: Client): stream_client = _with_store(client, _store()) stream = await _create_stream(stream_client) diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py index f6c520e1a3..480a931118 100644 --- a/tests/nexus/test_stream_store_workflow.py +++ b/tests/nexus/test_stream_store_workflow.py @@ -39,11 +39,8 @@ from temporalio.exceptions import ApplicationError, StreamError from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT from temporalio.nexus.system.stream_service import ( - AppendRecord, - AppendRecordsRequest, AppendRecordsResponse, ClientProducer, - CloseStreamRequest, Producer, ProducerKindClient, ProducerKindLocalActivity, @@ -52,6 +49,7 @@ StreamRecord, StreamRef, StreamService, + StreamStatus, ) from temporalio.nexus.system.stream_service.workflow import create_stream, get_stream from temporalio.nexus.system.stream_service.workflow._resources.stream import ( @@ -63,6 +61,11 @@ ExecuteActivityInput, Interceptor, Replayer, + StartActivityInput, + Worker, + WorkflowInboundInterceptor, + WorkflowInterceptorClassInput, + WorkflowOutboundInterceptor, ) from tests.helpers import assert_task_fail_eventually, new_worker from tests.helpers.streams import DictStreamStore, RecordingStreamStore @@ -144,16 +147,10 @@ async def run(self, stream_id: str) -> None: class RefusedOperationWorkflow: @workflow.run async def run(self, operation: str) -> str: - # The Workflow `Stream` has no method for these operations, so the test starts - # them through a Nexus client, as Workflow code still can. - ref = StreamRef[str](id="stream") - producer = Producer(id="workflow", kind=ProducerKindClient(ClientProducer())) + # The Workflow `Stream` has no `read`, so the test starts its operation through + # a Nexus client, as Workflow code still can. requests: dict[str, Any] = { - "AppendRecords": AppendRecordsRequest( - ref=ref, records=[AppendRecord(body="a")], producer=producer - ), - "CloseStream": CloseStreamRequest(ref=ref, producer=producer), - "ReadRecords": ReadRecordsRequest(ref=ref), + "ReadRecords": ReadRecordsRequest(ref=StreamRef[str](id="stream")), } nexus_client = workflow.create_nexus_client( service=StreamService, endpoint=TEMPORAL_SYSTEM_ENDPOINT @@ -337,6 +334,86 @@ async def append_records( return response +class CloseRequestStore(RecordingStreamStore): + """Keeps each close request.""" + + def __init__(self) -> None: + super().__init__(MemoryStreamStore()) + self.closes: list[temporalio.api.streamservice.v1.CloseStreamRequest] = [] + + async def close_stream( + self, request: temporalio.api.streamservice.v1.CloseStreamRequest + ) -> temporalio.api.streamservice.v1.CloseStreamResponse: + kept = temporalio.api.streamservice.v1.CloseStreamRequest() + kept.CopyFrom(request) + self.closes.append(kept) + return await super().close_stream(request) + + +@workflow.defn +class CloseWorkflow: + @workflow.run + async def run(self, stream_id: str) -> StreamRef[str]: + stream = await create_stream(id=stream_id, body_type=str) + await stream.append("a") + await stream.close() + return stream.ref + + +@workflow.defn +class AsyncWithCloseWorkflow: + @workflow.run + async def run(self, stream_id: str) -> str: + async with create_stream(id=stream_id, body_type=str) as stream: + await stream.append("a") + described = await stream.describe() + assert described.status is not None + return described.status.name + + +_fail_close_task = True +_started_activities: list[str] = [] + + +@workflow.defn(sandboxed=False) +class CloseThenFailTaskOnceWorkflow: + @workflow.run + async def run(self, stream_id: str) -> None: + stream = await create_stream( + id=stream_id, + body_type=str, + id_conflict_policy=StreamIdConflictPolicy.StreamIdConflictPolicyUseExisting, + ) + await stream.append("a") + close = asyncio.create_task(stream.close()) + # A few turns of the loop let the close schedule its Activity in this task. + for _ in range(10): + await asyncio.sleep(0) + global _fail_close_task + if _fail_close_task: + _fail_close_task = False + raise RuntimeError("Failed the Workflow Task once") + await close + + +class StartedActivityInterceptor(Interceptor): + def workflow_interceptor_class( + self, input: WorkflowInterceptorClassInput + ) -> type[WorkflowInboundInterceptor]: + return _StartedActivityInbound + + +class _StartedActivityInbound(WorkflowInboundInterceptor): + def init(self, outbound: WorkflowOutboundInterceptor) -> None: + super().init(_StartedActivityOutbound(outbound)) + + +class _StartedActivityOutbound(WorkflowOutboundInterceptor): + def start_activity(self, input: StartActivityInput) -> workflow.ActivityHandle[Any]: + _started_activities.append(str(input.activity)) + return super().start_activity(input) + + class RecordingInterceptor(Interceptor): def __init__(self) -> None: self.activity_types: list[str] = [] @@ -545,33 +622,19 @@ async def test_workflow_create_without_store_fails_the_task(client: Client): await handle.terminate() -@pytest.mark.parametrize( - "operation, message", - [ - ( - "CloseStream", - "Workflow code can't close a Stream on a Store in this release. " - "Call an Activity that closes it through its client.", - ), - ( - "ReadRecords", - "Workflow code can't read a Stream on a Store in this release. " - "Call an Activity that reads it through its client.", - ), - ], -) -async def test_workflow_stream_operation_fails_with_store( - client: Client, operation: str, message: str -): +async def test_workflow_read_fails_with_store(client: Client): stream_client = _with_store(client, MemoryStreamStore()) async with new_worker(stream_client, RefusedOperationWorkflow) as worker: result = await stream_client.execute_workflow( RefusedOperationWorkflow.run, - operation, + "ReadRecords", id=f"workflow-{uuid.uuid4()}", task_queue=worker.task_queue, ) - assert result == message + assert result == ( + "Workflow code can't read a Stream on a Store in this release. " + "Call an Activity that reads it through its client." + ) def test_workflow_stream_has_append_and_close_and_no_read(): @@ -981,3 +1044,127 @@ async def test_workflow_explicit_producer_reaches_store(client: Client): assert request.producer.id == "mine" assert request.producer.WhichOneof("kind") == "client" assert request.sequence == 0 + + +async def _describe_status(client: Client, stream_id: str) -> StreamStatus | None: + stream = client.get_stream(StreamRef(id=stream_id), body_type=str) + return (await stream.describe()).status + + +async def test_workflow_close_runs_in_activity_after_its_task(client: Client): + store = CloseRequestStore() + stream_client = _with_store(client, store) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker(stream_client, CloseWorkflow) as worker: + handle = await stream_client.start_workflow( + CloseWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + ref = await handle.result() + assert await _describe_status(stream_client, stream_id) == ( + StreamStatus.StreamStatusClosed + ) + with pytest.raises(StreamError) as err: + await stream_client.get_stream(ref, body_type=str).append("b") + assert err.value.reason == StreamErrorReason.STREAM_ERROR_REASON_CLOSED + # The close keeps the Workflow Producer of the request. + [close] = store.closes + assert close.producer.WhichOneof("kind") == "workflow" + assert close.producer.workflow.run_id == handle.result_run_id + events = (await handle.fetch_history()).events + scheduled = [ + event.activity_task_scheduled_event_attributes + for event in events + if event.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED + ] + assert [attributes.activity_type.name for attributes in scheduled] == [ + "__temporal_stream_store_call:memory" + ] + assert scheduled[0].task_queue.name == worker.task_queue + markers = [ + event.marker_recorded_event_attributes.marker_name + for event in events + if event.event_type == EventType.EVENT_TYPE_MARKER_RECORDED + ] + # One marker for the create and one for the append, and none for the close. + assert markers == ["core_local_activity", "core_local_activity"] + + +async def test_workflow_async_with_closes_the_stream(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker(stream_client, AsyncWithCloseWorkflow) as worker: + status = await stream_client.execute_workflow( + AsyncWithCloseWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert status == StreamStatus.StreamStatusOpen.name + assert await _describe_status(stream_client, stream_id) == ( + StreamStatus.StreamStatusClosed + ) + records = await _read_all(stream_client, stream_id) + assert [record.body for record in records] == ["a"] + + +async def test_worker_with_workflows_only_runs_the_close(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + async with Worker( + stream_client, task_queue=str(uuid.uuid4()), workflows=[CloseWorkflow] + ) as worker: + await stream_client.execute_workflow( + CloseWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await _describe_status(stream_client, stream_id) == ( + StreamStatus.StreamStatusClosed + ) + + +async def test_failed_workflow_task_leaves_the_stream_open(client: Client): + global _fail_close_task + _fail_close_task = True + _started_activities.clear() + stream_client = _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker( + stream_client, + CloseThenFailTaskOnceWorkflow, + interceptors=[StartedActivityInterceptor()], + ) as worker: + await stream_client.execute_workflow( + CloseThenFailTaskOnceWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + # The failed task started the close too, but the server never scheduled it. + assert _started_activities == ["__temporal_stream_store_call:memory"] * 2 + records = await _read_all(stream_client, stream_id) + # The retried task appends again to the open Stream, and then closes it. + assert [record.body for record in records] == ["a", "a"] + assert await _describe_status(stream_client, stream_id) == ( + StreamStatus.StreamStatusClosed + ) + + +async def test_activity_interceptors_see_the_close_activity(client: Client): + interceptor = RecordingInterceptor() + stream_client = _with_store(client, MemoryStreamStore()) + async with new_worker( + stream_client, CloseWorkflow, interceptors=[interceptor] + ) as worker: + await stream_client.execute_workflow( + CloseWorkflow.run, + f"stream-{uuid.uuid4()}", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + # The create and the append run as Local Activities, and the close as an Activity. + assert interceptor.activity_types == ["__temporal_stream_store_call:memory"] * 3 From 5d78e9f8813dc53c0f59cbe7cf28606239c2b1a1 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 00:43:50 -0700 Subject: [PATCH 5/9] Let a Stream config hold several Stores with no selector. A create can name its Store with `store=`, so a config with several Stores needs no selector. A call that names no Store on such a config raises `StreamConfigError` at the call. --- temporalio/client/_stream_store.py | 33 ++++++++++------- temporalio/exceptions.py | 5 ++- tests/nexus/test_stream_store.py | 45 ++++++++++++++++++++++- tests/nexus/test_stream_store_workflow.py | 27 ++++++++++++++ 4 files changed, 92 insertions(+), 18 deletions(-) diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index bf68c129b1..d8b6f85656 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -424,8 +424,10 @@ class StreamConfig: :py:class:`temporalio.converter.ExternalStorage` holds drivers and a selector. A ref carries the name of its Store, so a call on a known Stream goes to the - Store of that name. A call whose ref names no Store, such as a create, goes to - the one Store, or to the Store that :py:attr:`store_selector` returns. + Store of that name. A create can name its Store with ``store=``. A call whose ref + names no Store goes to the Store that :py:attr:`store_selector` returns, or to the + one Store when there is no selector. With more than one Store and no selector, + such a call raises :py:class:`temporalio.exceptions.StreamConfigError`. A Worker with Workflows and this config polls for Activity Tasks, since a Workflow's close runs the Store's Activity as a regular Activity. With ``no_remote_activities`` @@ -439,16 +441,18 @@ class StreamConfig: """At least one Store. Each Store has a unique name. A ref carries the :py:attr:`StreamStore.name` of its Store, and a call looks - the Store up by that name. With more than one Store, :py:attr:`store_selector` - must also be set. A Worker serves every Store of its client. + the Store up by that name. With more than one Store and no + :py:attr:`store_selector`, each call must name its Store. A Worker serves every + Store of its client. """ store_selector: Callable[[StreamSelectContext], StreamStore] | None = None - """Picks the Store of a new Stream. Required with more than one Store. + """Picks the Store of a call whose ref names no Store. - It gets a :py:class:`StreamSelectContext` for each call whose ref names no - Store, and returns one of the instances in :py:attr:`stores`. When it is - ``None`` and there is one Store, every new Stream goes to that Store. + It gets a :py:class:`StreamSelectContext` for each such call, and returns one of + the instances in :py:attr:`stores`. When it is ``None`` and there is one Store, + such a call goes to that Store. When it is ``None`` and there is more than one + Store, such a call raises :py:class:`temporalio.exceptions.StreamConfigError`. For a call of Workflow code, the selector runs in Workflow code, also on a replay. So it must give the same Store for the same context. The Activity of the @@ -477,8 +481,7 @@ def __post_init__(self) -> None: """Check the Stores and build the name-keyed map of them. Raises :py:class:`temporalio.exceptions.StreamConfigError` if there is no - Store, if two Stores share a name, or if more than one Store is given - without a :py:attr:`store_selector`. + Store, or if two Stores share a name. """ if not self.stores: raise StreamConfigError("StreamConfig.stores must have at least one Store") @@ -490,15 +493,17 @@ def __post_init__(self) -> None: "Each Store needs a unique name." ) store_map[store.name] = store - if len(self.stores) > 1 and self.store_selector is None: - raise StreamConfigError( - "StreamConfig.store_selector must be set when there is more than one Store" - ) object.__setattr__(self, "_store_map", store_map) def _select_store(self, context: StreamSelectContext) -> StreamStore: selector = self.store_selector if selector is None: + if len(self.stores) > 1: + raise StreamConfigError( + "StreamConfig has more than one Store and no store_selector, and " + f"the call on the Stream '{context.id}' names no Store. Name one " + "with store=, or set StreamConfig.store_selector." + ) return self.stores[0] store = selector(context) if ( diff --git a/temporalio/exceptions.py b/temporalio/exceptions.py index 63135cd9fc..702f5885bb 100644 --- a/temporalio/exceptions.py +++ b/temporalio/exceptions.py @@ -293,8 +293,9 @@ class StreamConfigError(RuntimeError): Streams on the server, so every Stream call needs a Store. * A Workflow's Stream call is made on a Worker whose client has no Stream config. This fails the Workflow Task. - * A :py:class:`temporalio.client.StreamConfig` has no Store, has two Stores with - one name, or has more than one Store and no selector. + * A :py:class:`temporalio.client.StreamConfig` has no Store, or has two Stores + with one name. + * A call names no Store on a config with more than one Store and no selector. * A ref names a Store that the config of the client doesn't have. * The selector returns a Store that isn't one of the config's Stores. * A :py:class:`temporalio.client.StreamStore` is given in ``plugins``. diff --git a/tests/nexus/test_stream_store.py b/tests/nexus/test_stream_store.py index db37459f9b..0b10411c5f 100644 --- a/tests/nexus/test_stream_store.py +++ b/tests/nexus/test_stream_store.py @@ -474,11 +474,52 @@ def test_stream_config_refuses_a_bad_set_of_stores(): same_name = [MemoryStreamStore(name="a"), MemoryStreamStore(name="a")] with pytest.raises(StreamConfigError, match="more than one Store named 'a'"): StreamConfig(stores=same_name, store_selector=lambda _: same_name[0]) - with pytest.raises(StreamConfigError, match="store_selector must be set"): - StreamConfig(stores=[MemoryStreamStore(name="a"), MemoryStreamStore(name="b")]) assert not issubclass(StreamConfigError, StreamError) +def test_stream_config_takes_several_stores_with_no_selector(): + first, second = MemoryStreamStore(name="a"), MemoryStreamStore(name="b") + stream_config = StreamConfig(stores=[first, second]) + assert stream_config.store_selector is None + assert stream_config._store_for(StreamRef(id="stream", store="b"), None) is second + with pytest.raises(StreamConfigError, match="more than one Store"): + stream_config._store_for(StreamRef(id="stream"), None) + + +async def test_client_create_with_store_on_several_stores_needs_no_selector( + client: Client, +): + first = MemoryStreamStore(name="first") + second = MemoryStreamStore(name="second") + stream_client = _with_config(client, StreamConfig(stores=[first, second])) + stream = await stream_client.create_stream( + id=f"stream-{uuid.uuid4()}", store="second", body_type=str + ) + assert stream.ref.store == "second" + await stream.append("a") + in_second = _with_store(client, second).get_stream(stream.ref, body_type=str) + assert [record.body async for record in in_second.iter_read_items()] == ["a"] + in_first = _with_store(client, first).get_stream( + StreamRef(id=stream.ref.id), body_type=str + ) + err = await _stream_error(in_first.describe()) + assert err.reason == StreamErrorReason.STREAM_ERROR_REASON_NOT_FOUND + + +async def test_client_create_without_store_on_several_stores_fails(client: Client): + stream_config = StreamConfig( + stores=[MemoryStreamStore(name="first"), MemoryStreamStore(name="second")] + ) + stream_client = _with_config(client, stream_config) + with pytest.raises(StreamConfigError) as err: + await stream_client.create_stream(id="stream", body_type=str) + assert str(err.value) == ( + "StreamConfig has more than one Store and no store_selector, and the call " + "on the Stream 'stream' names no Store. Name one with store=, or set " + "StreamConfig.store_selector." + ) + + def test_client_refuses_a_store_in_plugins(client: Client): config = client.config() config["plugins"] = [MemoryStreamStore()] # type: ignore[list-item] diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py index 480a931118..1403994790 100644 --- a/tests/nexus/test_stream_store_workflow.py +++ b/tests/nexus/test_stream_store_workflow.py @@ -835,6 +835,33 @@ async def test_workflow_ref_of_another_store_fails_the_task(client: Client): await handle.terminate() +async def test_workflow_create_without_store_on_several_stores_fails_the_task( + client: Client, +): + stream_config = StreamConfig( + stores=[CountingStore("first"), CountingStore("second")] + ) + stream_client = _with_config(client, stream_config) + stream_id = f"stream-{uuid.uuid4()}" + async with new_worker(stream_client, CreateAndDescribeWorkflow) as worker: + handle = await stream_client.start_workflow( + CreateAndDescribeWorkflow.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=30), + ) + await assert_task_fail_eventually( + handle, + message_contains=( + "StreamConfig has more than one Store and no store_selector, and the " + f"call on the Stream '{stream_id}' names no Store. Name one with " + "store=, or set StreamConfig.store_selector." + ), + ) + await handle.terminate() + + async def test_workflow_create_goes_to_the_selected_store(client: Client): contexts: list[StreamSelectContext] = [] first, second, stream_config = _two_counting_stores(contexts) From 208d1b0ae00fc38a475217e0636b249df21620f6 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 00:53:42 -0700 Subject: [PATCH 6/9] Failed the Workflow for a Stream call with no Stream config. A failed Workflow Task retries for good, and a Worker with no Stream config never recovers. Workflow code now gets a non-retryable `ApplicationError`, so a user sees the failed Workflow. --- temporalio/client/_stream_store.py | 7 ++++++ temporalio/exceptions.py | 11 +++++++-- temporalio/worker/_replayer.py | 7 +++--- temporalio/worker/_stream_store.py | 12 ++++++++++ temporalio/worker/_workflow_instance.py | 12 ++++------ tests/nexus/test_stream_store_workflow.py | 28 +++++++++++++++-------- 6 files changed, 55 insertions(+), 22 deletions(-) diff --git a/temporalio/client/_stream_store.py b/temporalio/client/_stream_store.py index d8b6f85656..74ed6961b8 100644 --- a/temporalio/client/_stream_store.py +++ b/temporalio/client/_stream_store.py @@ -423,6 +423,13 @@ class StreamConfig: operation. It holds Stores and a selector, as :py:class:`temporalio.converter.ExternalStorage` holds drivers and a selector. + On a Worker whose client has no config, a Stream call of Workflow code raises + :py:class:`temporalio.exceptions.ApplicationError` with the ``type`` + ``StreamConfigError``, and an uncaught one fails the Workflow. Workflow code can + catch this error. A Workflow that catches it must not go on to make the same call + on a Worker that has a config, since the History then has no record of the first + call. + A ref carries the name of its Store, so a call on a known Stream goes to the Store of that name. A create can name its Store with ``store=``. A call whose ref names no Store goes to the Store that :py:attr:`store_selector` returns, or to the diff --git a/temporalio/exceptions.py b/temporalio/exceptions.py index 702f5885bb..658012ee4f 100644 --- a/temporalio/exceptions.py +++ b/temporalio/exceptions.py @@ -291,8 +291,8 @@ class StreamConfigError(RuntimeError): * A Stream call is made on a client with no Stream config. This release has no Streams on the server, so every Stream call needs a Store. - * A Workflow's Stream call is made on a Worker whose client has no Stream config. - This fails the Workflow Task. + * A Workflow's Stream call has a ref that names no Store, on a Worker whose + client has more than one Store and no selector. This fails the Workflow Task. * A :py:class:`temporalio.client.StreamConfig` has no Store, or has two Stores with one name. * A call names no Store on a config with more than one Store and no selector. @@ -303,6 +303,13 @@ class StreamConfigError(RuntimeError): Worker's client doesn't have. This fails the Workflow Task. * This build of temporalio has no Stream runtime, or no Redis Stream Store. + A Workflow's Stream call on a Worker whose client has no Stream config doesn't + raise this class. Workflow code gets an :py:class:`ApplicationError` with the + ``type`` ``StreamConfigError`` and ``non_retryable`` set, so an uncaught one fails + the Workflow. Workflow code can catch this error. A Workflow that catches it must + not go on to make the same call on a Worker that has a config, since the History + then has no record of the first call. + .. warning:: This API is experimental and subject to change. """ diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 540e9d4052..7ad6bdeaaf 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -67,9 +67,10 @@ def __init__( ran. ``stream_config`` takes the Stream config of the client that the - worker had, so a replay makes the same commands. A history with a Stream - call needs it, since a Stream call with no Stream config fails the - Workflow Task. A replay calls no method of its Stores. + worker had, so a replay makes the same commands. A history that records a + Stream call needs it. With no Stream config, Workflow code gets an error in + place of that call, and the replay fails with a nondeterminism error. A + replay calls no method of its Stores. Note, unlike the worker, for the replayer the workflow_task_executor will default to a new thread pool executor with no max_workers set that diff --git a/temporalio/worker/_stream_store.py b/temporalio/worker/_stream_store.py index bd134bf4b1..947cc7aeba 100644 --- a/temporalio/worker/_stream_store.py +++ b/temporalio/worker/_stream_store.py @@ -24,6 +24,18 @@ from ._workflow_instance import _WorkflowInstanceImpl +def _no_stream_config_error() -> temporalio.exceptions.ApplicationError: # pyright: ignore[reportUnusedFunction] + # A `StreamConfigError` would fail the Workflow Task, and its retries would fail the + # same way for good. A Worker with no config is a setup that a user has to see, so + # the Workflow fails instead. + return temporalio.exceptions.ApplicationError( + "The client of this Worker has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...]).", + type="StreamConfigError", + non_retryable=True, + ) + + def _workflow_target() -> StorageDriverWorkflowInfo: # External storage builds the same value for a Workflow. info = temporalio.workflow.info() diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 09f1ae53fc..07f63a69dc 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -95,7 +95,7 @@ WorkflowInboundInterceptor, WorkflowOutboundInterceptor, ) -from ._stream_store import _WorkflowStreamStore +from ._stream_store import _no_stream_config_error, _WorkflowStreamStore logger = logging.getLogger(__name__) @@ -196,7 +196,7 @@ class WorkflowInstanceDetails: stream_store: _WorkflowStreamStore | None = None """The Worker's Stream Stores as Workflow code reaches them, or ``None`` when the Worker's client has none. This release has no Streams on the server, so with ``None`` - a Stream call fails the Workflow Task.""" + a Stream call fails the Workflow.""" class WorkflowInstance(ABC): @@ -2355,15 +2355,11 @@ async def _schedule_system_nexus_operation( self, input: StartNexusOperationInput[Any, OutputT] ) -> temporalio.workflow.NexusOperationHandle[OutputT]: # This release has no Streams on the server, so a Worker answers a Stream call with - # a Store of its client. On a Worker whose client has none, the `StreamConfigError` - # fails the Workflow Task. Every other system Nexus operation goes to the server. + # a Store of its client. Every other system Nexus operation goes to the server. if input.service == _STREAM_SERVICE.name: store = self._stream_store if store is None: - raise temporalio.exceptions.StreamConfigError( - "The client of this Worker has no Stream Store. Give one in " - "stream_config=StreamConfig(stores=[...])." - ) + raise _no_stream_config_error() return store.start_operation(self, input) return await self._outbound_start_nexus_operation(input) diff --git a/tests/nexus/test_stream_store_workflow.py b/tests/nexus/test_stream_store_workflow.py index 1403994790..07ae119841 100644 --- a/tests/nexus/test_stream_store_workflow.py +++ b/tests/nexus/test_stream_store_workflow.py @@ -32,6 +32,7 @@ StreamSelectContext, StreamStore, StreamStoreCallOptions, + WorkflowFailureError, WorkflowHandle, ) from temporalio.common import RetryPolicy @@ -603,7 +604,9 @@ async def test_workflow_stream_replay_does_not_call_store( assert replay_store.operations == [] -async def test_workflow_create_without_store_fails_the_task(client: Client): +async def test_workflow_create_without_stream_config_fails_the_workflow( + client: Client, +): async with new_worker(client, CreateAndDescribeWorkflow) as worker: handle = await client.start_workflow( CreateAndDescribeWorkflow.run, @@ -612,14 +615,21 @@ async def test_workflow_create_without_store_fails_the_task(client: Client): task_queue=worker.task_queue, execution_timeout=timedelta(seconds=30), ) - await assert_task_fail_eventually( - handle, - message_contains=( - "The client of this Worker has no Stream Store. Give one in " - "stream_config=StreamConfig(stores=[...])." - ), - ) - await handle.terminate() + with pytest.raises(WorkflowFailureError) as err: + await asyncio.wait_for(handle.result(), timeout=10) + cause = err.value.cause + assert isinstance(cause, ApplicationError) + assert not isinstance(cause, StreamError) + assert cause.type == "StreamConfigError" + assert cause.non_retryable + assert cause.message == ( + "The client of this Worker has no Stream Store. Give one in " + "stream_config=StreamConfig(stores=[...])." + ) + # A Replayer with no Stream config fails the Workflow the same way. + await Replayer(workflows=[CreateAndDescribeWorkflow]).replay_workflow( + await handle.fetch_history() + ) async def test_workflow_read_fails_with_store(client: Client): From 0c0ccd56f8c316f9b5ecaae13af4a3c5af08da67 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 01:24:18 -0700 Subject: [PATCH 7/9] Let BufferedAppender batch the appends of Workflow code. A timer of the Workflow event loop goes into History, so in Workflow code the appender starts its write at the next yield instead. The appends of one activation become one Local Activity. --- temporalio/client/_stream_appender.py | 71 +++++-- tests/nexus/test_stream_appender_workflow.py | 197 +++++++++++++++++++ 2 files changed, 252 insertions(+), 16 deletions(-) create mode 100644 tests/nexus/test_stream_appender_workflow.py diff --git a/temporalio/client/_stream_appender.py b/temporalio/client/_stream_appender.py index ec8d3d49d3..1fda0cc776 100644 --- a/temporalio/client/_stream_appender.py +++ b/temporalio/client/_stream_appender.py @@ -5,33 +5,51 @@ import asyncio from datetime import timedelta from types import TracebackType -from typing import Generic, TypeVar +from typing import TYPE_CHECKING, Generic, TypeVar from typing_extensions import Self import temporalio.nexus.system.stream_service.client +import temporalio.workflow from ._stream_store import _check_not_negative +if TYPE_CHECKING: + # The generated Workflow modules import `temporalio.workflow`, which a runtime import + # here would pull into the client package. + import temporalio.nexus.system.stream_service.workflow._resources.stream + BodyT = TypeVar("BodyT") # The Store refuses an append of more Records than this. Many small Records fit under its byte # limit, but they keep the Store busy for long. _MAX_BATCH_RECORDS = 1000 +_DEFAULT_MAX_DELAY = timedelta(milliseconds=100) + class BufferedAppender(Generic[BodyT]): """Batches the appends of one Stream. - Use it in an Activity or a client that appends often, such as one that streams model - tokens. :py:meth:`append` returns at once and does no I/O. The appender writes the - buffered Records as one batch at most ``max_delay`` after the first of them, or at once - when they fill a batch of 1,000 Records. One write is in flight at a time. Records - appended during a write wait for the next batch, and they keep their order. + Use it in code that appends often, such as an Activity that streams model tokens. + :py:meth:`append` returns at once and does no I/O. One write is in flight at a time. + Records appended during a write wait for the next batch, and they keep their order. A batch + holds at most 1,000 Records. + + It takes a client Stream or a Workflow Stream. When a write starts depends on where + :py:meth:`append` runs, not on the type of the Stream: + + * In a client or an Activity, the appender writes the buffered Records at most + ``max_delay`` after the first of them, or at once when they fill a batch. + * In Workflow code, a timer would go into History, so the appender uses none. The first + write starts when the calling code next yields, so the appends of one activation go as + one write. Each write is the Workflow ``append`` of the Stream, so it is one Local + Activity and one marker in History. A replay takes the results from History and writes + nothing. Use it in ``async with``, which writes the remaining Records at the end of the block, or await :py:meth:`flush`. Without ``async with`` or a :py:meth:`flush`, the Records still in - the buffer are lost. The end of an Activity attempt doesn't write them. + the buffer are lost. The end of an Activity attempt or of a Workflow doesn't write them. ``await flush()`` is the backpressure point. A caller that appends faster than the Store accepts Records must await it from time to time, so the buffer doesn't grow without bound. @@ -42,32 +60,39 @@ class BufferedAppender(Generic[BodyT]): The appender doesn't count bytes, so a batch over the size limit of the Store fails with the Store's error. - It takes a client Stream, so Workflow code has no use for it. - .. warning:: This API is experimental and subject to change. """ def __init__( self, - stream: temporalio.nexus.system.stream_service.client.Stream[BodyT], + stream: temporalio.nexus.system.stream_service.client.Stream[BodyT] + | temporalio.nexus.system.stream_service.workflow._resources.stream.Stream[ + BodyT + ], *, - max_delay: timedelta = timedelta(milliseconds=100), + max_delay: timedelta | None = None, ) -> None: """Create an appender. It starts no write until the first append. Args: - stream: The Stream to append to. + stream: The Stream to append to, a client Stream or a Workflow Stream. max_delay: The longest time that a Record waits in the buffer before its batch - starts to write. It must not be negative. + starts to write, in a client or an Activity. It must not be negative. + ``None`` means 100 milliseconds. Workflow code must leave it ``None``. + + Raises: + ValueError: ``max_delay`` is negative, or it is set in Workflow code. """ _check_not_negative(max_delay=max_delay) self._stream = stream self._max_delay = max_delay + self._check_delay() self._buffer: list[BodyT] = [] # When the oldest buffered Record arrived, so its batch starts at most max_delay later. self._first_at = 0.0 - self._timer: asyncio.TimerHandle | None = None + # The timer of a client or an Activity, or the start at the next yield of Workflow code. + self._timer: asyncio.Handle | None = None self._writing: asyncio.Task[None] | None = None self._appended = 0 # Count of the Records that a write took or that a failed batch left, written or not. @@ -89,6 +114,7 @@ def append(self, *records: BodyT) -> None: self._raise_error() if not records: raise TypeError("append takes one or more values") + self._check_delay() loop = asyncio.get_running_loop() if not self._buffer: self._first_at = loop.time() @@ -142,14 +168,27 @@ def _raise_error(self) -> None: if self._error is not None: raise self._error + def _check_delay(self) -> None: + if self._max_delay is not None and temporalio.workflow.in_workflow(): + raise ValueError( + "BufferedAppender takes no max_delay in Workflow code, since a delay " + "would put a timer into History. Leave max_delay as None." + ) + def _schedule(self) -> None: if self._writing is not None or not self._buffer: return - if len(self._buffer) >= _MAX_BATCH_RECORDS: + if temporalio.workflow.in_workflow(): + # A timer of the Workflow event loop is durable. A start at the next yield + # takes every Record of this activation into one write, and adds no event. + if self._timer is None: + self._timer = asyncio.get_running_loop().call_soon(self._start) + elif len(self._buffer) >= _MAX_BATCH_RECORDS: self._start() elif self._timer is None: + max_delay = self._max_delay or _DEFAULT_MAX_DELAY self._timer = asyncio.get_running_loop().call_at( - self._first_at + self._max_delay.total_seconds(), self._start + self._first_at + max_delay.total_seconds(), self._start ) def _start(self) -> None: diff --git a/tests/nexus/test_stream_appender_workflow.py b/tests/nexus/test_stream_appender_workflow.py new file mode 100644 index 0000000000..0bb9cf2216 --- /dev/null +++ b/tests/nexus/test_stream_appender_workflow.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta +from typing import Any + +from google.protobuf.message import Message + +import temporalio.api.streamservice.v1 +from temporalio import workflow +from temporalio.api.enums.v1 import EventType, StreamErrorReason +from temporalio.api.history.v1 import HistoryEvent +from temporalio.client import BufferedAppender, Client, WorkflowHandle +from temporalio.exceptions import StreamError +from temporalio.nexus.system.stream_service import StreamRef +from temporalio.nexus.system.stream_service.workflow import create_stream, get_stream +from temporalio.testing import MemoryStreamStore +from temporalio.worker import Replayer +from tests.helpers import new_worker +from tests.helpers.streams import RecordingStreamStore +from tests.nexus.test_stream_store import _with_store + +_BURST = [f"record-{index}" for index in range(40)] + + +class BatchStore(RecordingStreamStore): + """Keeps the number of Records in each append.""" + + def __init__(self) -> None: + super().__init__(MemoryStreamStore()) + self.batches: list[int] = [] + + async def before_call(self, operation: str, request: Message) -> None: + if isinstance(request, temporalio.api.streamservice.v1.AppendRecordsRequest): + self.batches.append(len(request.records)) + + +async def _yield_a_few_times() -> None: + # A turn of the loop with no timer, so the write starts and waits for its Local Activity. + for _ in range(5): + await asyncio.sleep(0) + + +@workflow.defn +class BurstWorkflow: + @workflow.run + async def run(self, stream_id: str) -> str: + stream = await create_stream(id=stream_id, body_type=str) + appender = BufferedAppender(stream) + for value in _BURST: + appender.append(value) + last_offset = await appender.flush() + assert last_offset is not None + return last_offset.hex() + + +@workflow.defn +class InFlightWorkflow: + @workflow.run + async def run(self, stream_id: str) -> None: + stream = await create_stream(id=stream_id, body_type=str) + async with BufferedAppender(stream) as appender: + appender.append("a") + await _yield_a_few_times() + # The write of "a" waits for its Local Activity, so these go into the next one. + appender.append("b") + appender.append("c", "d") + + +@workflow.defn +class FailedAppendWorkflow: + @workflow.run + async def run(self, stream_id: str) -> list[Any]: + appender = BufferedAppender(get_stream(StreamRef(id=stream_id), body_type=str)) + appender.append("a") + await _yield_a_few_times() + appender.append("b", "c") + try: + await appender.flush() + except StreamError as err: + first = err + else: + return ["flushed"] + try: + appender.append("d") + except StreamError as err: + later = err + else: + return ["appended"] + return [ + StreamErrorReason.Name(first.reason), + later is first, + appender.unwritten, + ] + + +@workflow.defn +class MaxDelayWorkflow: + @workflow.run + async def run(self, stream_id: str) -> str: + stream = get_stream(StreamRef(id=stream_id), body_type=str) + try: + BufferedAppender(stream, max_delay=timedelta(milliseconds=10)) + except ValueError as err: + return str(err) + return "built" + + +async def _run( + client: Client, workflow_type: type[Any], stream_id: str +) -> tuple[WorkflowHandle[Any, Any], Any]: + async with new_worker(client, workflow_type) as worker: + handle: WorkflowHandle[Any, Any] = await client.start_workflow( + workflow_type.run, + stream_id, + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=30), + ) + return handle, await handle.result() + + +def _of_type(events: list[HistoryEvent], event_type: EventType.ValueType) -> int: + return sum(1 for event in events if event.event_type == event_type) + + +async def _bodies(client: Client, stream_id: str) -> list[str]: + stream = client.get_stream(StreamRef(id=stream_id), body_type=str) + return [record.body async for record in stream.iter_read_items()] + + +async def test_appends_of_one_activation_are_one_append(client: Client): + store = BatchStore() + stream_client = _with_store(client, store) + stream_id = f"stream-{uuid.uuid4()}" + handle, last_offset = await _run(stream_client, BurstWorkflow, stream_id) + + assert store.batches == [40] + assert await _bodies(stream_client, stream_id) == _BURST + described = await stream_client.get_stream( + StreamRef(id=stream_id), body_type=str + ).describe() + assert described.last_offset is not None + assert last_offset == described.last_offset.hex() + events = list((await handle.fetch_history()).events) + # One marker for the create and one for the append. + assert _of_type(events, EventType.EVENT_TYPE_MARKER_RECORDED) == 2 + assert _of_type(events, EventType.EVENT_TYPE_TIMER_STARTED) == 0 + + +async def test_appends_during_a_workflow_append_join_the_next_one(client: Client): + store = BatchStore() + stream_client = _with_store(client, store) + stream_id = f"stream-{uuid.uuid4()}" + handle, _ = await _run(stream_client, InFlightWorkflow, stream_id) + + assert store.batches == [1, 3] + assert await _bodies(stream_client, stream_id) == ["a", "b", "c", "d"] + events = list((await handle.fetch_history()).events) + assert _of_type(events, EventType.EVENT_TYPE_MARKER_RECORDED) == 3 + assert _of_type(events, EventType.EVENT_TYPE_TIMER_STARTED) == 0 + + +async def test_a_failed_workflow_append_stops_the_appender(client: Client): + store = BatchStore() + _, result = await _run( + _with_store(client, store), FailedAppendWorkflow, f"stream-{uuid.uuid4()}" + ) + assert result == ["STREAM_ERROR_REASON_NOT_FOUND", True, 3] + assert store.batches == [1] + + +async def test_replay_of_workflow_appender_writes_nothing(client: Client): + stream_client = _with_store(client, MemoryStreamStore()) + stream_id = f"stream-{uuid.uuid4()}" + handle, _ = await _run(stream_client, InFlightWorkflow, stream_id) + + replay_store = BatchStore() + await Replayer( + workflows=[InFlightWorkflow], + stream_config=_with_store(client, replay_store).config()["stream_config"], + ).replay_workflow(await handle.fetch_history()) + assert replay_store.operations == [] + assert await _bodies(stream_client, stream_id) == ["a", "b", "c", "d"] + + +async def test_workflow_appender_refuses_a_max_delay(client: Client): + _, message = await _run( + _with_store(client, MemoryStreamStore()), + MaxDelayWorkflow, + f"stream-{uuid.uuid4()}", + ) + assert message == ( + "BufferedAppender takes no max_delay in Workflow code, since a delay would put " + "a timer into History. Leave max_delay as None." + ) From 17f1a04a7c47dffb0d78e8702ad3e00604280b04 Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 02:16:58 -0700 Subject: [PATCH 8/9] Moved BufferedAppender to temporalio.contrib.streams. The public Streams API is the generated code. A helper on top of it is a library, so it belongs in contrib. --- temporalio/client/__init__.py | 2 -- temporalio/contrib/streams/__init__.py | 14 ++++++++++++++ .../streams/_appender.py} | 15 +++++++++------ tests/contrib/streams/__init__.py | 0 .../streams/test_appender.py} | 3 ++- .../streams/test_appender_workflow.py} | 3 ++- 6 files changed, 27 insertions(+), 10 deletions(-) create mode 100644 temporalio/contrib/streams/__init__.py rename temporalio/{client/_stream_appender.py => contrib/streams/_appender.py} (95%) create mode 100644 tests/contrib/streams/__init__.py rename tests/{nexus/test_stream_appender.py => contrib/streams/test_appender.py} (98%) rename tests/{nexus/test_stream_appender_workflow.py => contrib/streams/test_appender_workflow.py} (98%) diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index 7305b4be59..6fa6dd710c 100644 --- a/temporalio/client/__init__.py +++ b/temporalio/client/__init__.py @@ -185,7 +185,6 @@ ScheduleUpdate, ScheduleUpdateInput, ) -from ._stream_appender import BufferedAppender from ._stream_store import ( StreamConfig, StreamSelectContext, @@ -357,7 +356,6 @@ "TaskReachabilityType", "CloudOperationsClient", "Plugin", - "BufferedAppender", "StreamConfig", "StreamSelectContext", "StreamStore", diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py new file mode 100644 index 0000000000..d1ee83436d --- /dev/null +++ b/temporalio/contrib/streams/__init__.py @@ -0,0 +1,14 @@ +"""Helpers that sit on top of the Streams API. + +:py:class:`BufferedAppender` batches the appends of one Stream. It takes a client Stream or a +Workflow Stream. + +.. warning:: + This package is experimental and may change in future versions. +""" + +from temporalio.contrib.streams._appender import BufferedAppender + +__all__ = [ + "BufferedAppender", +] diff --git a/temporalio/client/_stream_appender.py b/temporalio/contrib/streams/_appender.py similarity index 95% rename from temporalio/client/_stream_appender.py rename to temporalio/contrib/streams/_appender.py index 1fda0cc776..7b194ddf09 100644 --- a/temporalio/client/_stream_appender.py +++ b/temporalio/contrib/streams/_appender.py @@ -1,4 +1,8 @@ -"""A helper that batches the appends of one Stream.""" +"""A helper that batches the appends of one Stream. + +.. warning:: + This module is experimental and subject to change. +""" from __future__ import annotations @@ -12,11 +16,9 @@ import temporalio.nexus.system.stream_service.client import temporalio.workflow -from ._stream_store import _check_not_negative - if TYPE_CHECKING: - # The generated Workflow modules import `temporalio.workflow`, which a runtime import - # here would pull into the client package. + # The Workflow Stream class is in a private module of the generated code. Only its type + # is needed, since the appender calls the same `append` on both Streams. import temporalio.nexus.system.stream_service.workflow._resources.stream BodyT = TypeVar("BodyT") @@ -84,7 +86,8 @@ def __init__( Raises: ValueError: ``max_delay`` is negative, or it is set in Workflow code. """ - _check_not_negative(max_delay=max_delay) + if max_delay is not None and max_delay < timedelta(): + raise ValueError("max_delay must not be negative") self._stream = stream self._max_delay = max_delay self._check_delay() diff --git a/tests/contrib/streams/__init__.py b/tests/contrib/streams/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/nexus/test_stream_appender.py b/tests/contrib/streams/test_appender.py similarity index 98% rename from tests/nexus/test_stream_appender.py rename to tests/contrib/streams/test_appender.py index fd320ff474..a864335387 100644 --- a/tests/nexus/test_stream_appender.py +++ b/tests/contrib/streams/test_appender.py @@ -9,7 +9,8 @@ import temporalio.api.streamservice.v1 from temporalio.api.enums.v1 import StreamErrorReason -from temporalio.client import BufferedAppender, Client +from temporalio.client import Client +from temporalio.contrib.streams import BufferedAppender from temporalio.nexus.system.stream_service import StreamRef from temporalio.nexus.system.stream_service.client import Stream from temporalio.testing import MemoryStreamStore diff --git a/tests/nexus/test_stream_appender_workflow.py b/tests/contrib/streams/test_appender_workflow.py similarity index 98% rename from tests/nexus/test_stream_appender_workflow.py rename to tests/contrib/streams/test_appender_workflow.py index 0bb9cf2216..998149b2e9 100644 --- a/tests/nexus/test_stream_appender_workflow.py +++ b/tests/contrib/streams/test_appender_workflow.py @@ -11,7 +11,8 @@ from temporalio import workflow from temporalio.api.enums.v1 import EventType, StreamErrorReason from temporalio.api.history.v1 import HistoryEvent -from temporalio.client import BufferedAppender, Client, WorkflowHandle +from temporalio.client import Client, WorkflowHandle +from temporalio.contrib.streams import BufferedAppender from temporalio.exceptions import StreamError from temporalio.nexus.system.stream_service import StreamRef from temporalio.nexus.system.stream_service.workflow import create_stream, get_stream From 4c5003d08a36cbf2415eef92b1279418808aaf0f Mon Sep 17 00:00:00 2001 From: Moe Dashti Date: Sat, 14 Nov 2026 03:16:38 -0700 Subject: [PATCH 9/9] Added a README to temporalio.contrib.streams. The other contrib packages have one, and the appender's rules for Workflow code need a place. --- temporalio/contrib/streams/README.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 temporalio/contrib/streams/README.md diff --git a/temporalio/contrib/streams/README.md b/temporalio/contrib/streams/README.md new file mode 100644 index 0000000000..f64fa2d009 --- /dev/null +++ b/temporalio/contrib/streams/README.md @@ -0,0 +1,23 @@ +# Streams helpers + +> This package is experimental. Its API can change in a later release. + +`temporalio.contrib.streams` holds helpers that sit on top of the Streams API of the SDK. A +helper takes a client `Stream` or a Workflow `Stream`. + +## `BufferedAppender` + +`BufferedAppender` batches the appends of one Stream. `append` returns at once, and the +appender writes the buffered Records as one batch. In a client or an Activity a batch waits +at most `max_delay`. In Workflow code a batch starts when the code next yields, so it adds no +timer to History. + +```python +from temporalio.contrib.streams import BufferedAppender + +async with BufferedAppender(tokens) as appender: + for chunk in chunks: + appender.append(chunk) +``` + +The end of the `async with` block writes the Records that are still in the buffer.