diff --git a/changelog/fixed/sleepy-otter-stream.md b/changelog/fixed/sleepy-otter-stream.md new file mode 100644 index 0000000000..2fdfc99d3d --- /dev/null +++ b/changelog/fixed/sleepy-otter-stream.md @@ -0,0 +1 @@ +Preserve retained Workflow Stream items after workflow completion, failure, cancellation, termination, or timeout by draining them through read-only queries before subscriptions finish. diff --git a/temporalio/contrib/workflow_streams/README.md b/temporalio/contrib/workflow_streams/README.md index 2fb4f94855..e8ca461e22 100644 --- a/temporalio/contrib/workflow_streams/README.md +++ b/temporalio/contrib/workflow_streams/README.md @@ -13,7 +13,8 @@ scales with durable batches rather than tokens. Under the hood the stream is built directly on Temporal's existing message-passing primitives: Signals carry publishes, Updates serve -long-poll subscriptions, and a Query exposes the current global offset. +long-poll subscriptions, and Queries expose retained items and the current +global offset. The library packages the boilerplate that turns those primitives into a usable stream: batching to amortize per-event overhead, deduplication for exactly-once delivery, topic filtering, and continue-as-new helpers @@ -31,3 +32,16 @@ snippets throughout. For runnable end-to-end examples, see the [Workflow Streams samples](https://github.com/temporalio/samples-python/tree/main/workflow-streams). + +Subscriptions to completed, failed, canceled, terminated, or timed-out runs +drain their remaining retained items through the +`__temporal_workflow_stream_poll` Query, then finish. Each Query targets the +confirmed terminal run, including when the Workflow ID is reused. Queries +return immediately, paginate within the same response size limit as Updates, +and start at the retained base when the requested offset has been truncated. + +Upgrade workers before enabling upgraded clients. Existing clients work with +upgraded workers; upgraded clients surface missing-query errors from older +workers. Closed-stream reads require retained history and an available +compatible worker. Query rejection policies and server/query errors remain +visible; truncated items and expired history cannot be recovered. diff --git a/temporalio/contrib/workflow_streams/_client.py b/temporalio/contrib/workflow_streams/_client.py index 8ed6cb5b51..7ad2a1dfe0 100644 --- a/temporalio/contrib/workflow_streams/_client.py +++ b/temporalio/contrib/workflow_streams/_client.py @@ -8,7 +8,7 @@ sync payload converter. The **codec chain** (e.g. encryption, compression) is **not** run per item — it runs once at the envelope level when Temporal's SDK encodes the ``__temporal_workflow_stream_publish`` -signal args and the ``__temporal_workflow_stream_poll`` update result. +signal args and the ``__temporal_workflow_stream_poll`` update/query result. Running the codec per item as well would double-encrypt / double-compress, because the envelope path covers the items again. The per-item ``Payload`` still carries the encoding metadata (``encoding: json/plain``, @@ -161,10 +161,6 @@ def __init__( self._pending_seq: int = 0 self._pending_since: float | None = None self._topic_types: dict[str, type[Any]] = {} - # Run id the most recent poll's update was admitted to. Captured before - # waiting for the outcome so a mid-poll continue-as-new can be detected by - # describing that specific run. None until the first poll is admitted. - self._polled_run_id: str | None = None @classmethod def create( @@ -529,7 +525,11 @@ async def subscribe( """Async iterator that polls for new items. Automatically follows continue-as-new chains when the client - was created via :py:meth:`create`. + was created via :py:meth:`create`. After the workflow completes, + fails, is canceled, terminated, or times out, queries deliver its + remaining retained items before the iterator finishes. These reads + require retained history and an available worker with the poll query + handler. Query failures, including configured query rejection, propagate. Args: topics: Topic filter. A single topic name, a list of topic @@ -570,65 +570,68 @@ async def subscribe( else: topic_filter = topics offset = from_offset + poll_handle = self._handle + terminal_handle: WorkflowHandle[Any, Any] | None = None while True: - try: - # Wait only for ACCEPTED so the handle (and the run id it was - # admitted to) is available before we block on the outcome; if - # the run continues-as-new mid-poll, result() fails but we still - # know which run to inspect. - handle = await self._handle.start_update( + if terminal_handle is not None: + result = await terminal_handle.query( "__temporal_workflow_stream_poll", PollInput(topics=topic_filter, from_offset=offset), - wait_for_stage=WorkflowUpdateStage.ACCEPTED, result_type=PollResult, ) - self._polled_run_id = handle.workflow_run_id - result: PollResult = await handle.result() - except asyncio.CancelledError: - raise - except WorkflowUpdateFailedError as e: - cause_type = getattr(e.cause, "type", None) - if cause_type == TRUNCATED_OFFSET_ERROR_TYPE: - # Subscriber fell behind truncation. Retry from - # offset 0 which the stream treats as "from the - # beginning of whatever exists" (i.e., from - # base_offset). - offset = 0 - continue - if cause_type == STREAM_DRAINING_ERROR_TYPE: - # Workflow is detaching for continue-as-new. Back off and - # retry; the poll lands on the successor run once the - # rollover completes. - cooldown_secs = poll_cooldown.total_seconds() - if cooldown_secs > 0: - await asyncio.sleep(cooldown_secs) - continue - if cause_type == "AcceptedUpdateCompletedWorkflow": - # Workflow returned (or continued-as-new) before - # this poll's update completed. Either follow the - # chain or exit cleanly. - if await self._follow_continue_as_new(): + else: + # Each attempt tracks its own accepted run, including when the + # same stream client is shared by concurrent subscriptions. + polled_run_id: str | None = None + try: + handle = await poll_handle.start_update( + "__temporal_workflow_stream_poll", + PollInput(topics=topic_filter, from_offset=offset), + wait_for_stage=WorkflowUpdateStage.ACCEPTED, + result_type=PollResult, + ) + polled_run_id = handle.workflow_run_id + result = await handle.result() + except asyncio.CancelledError: + raise + except ( + WorkflowUpdateFailedError, + WorkflowUpdateRPCTimeoutOrCancelledError, + RPCError, + ) as err: + if isinstance(err, WorkflowUpdateFailedError): + cause_type = getattr(err.cause, "type", None) + if cause_type == TRUNCATED_OFFSET_ERROR_TYPE: + offset = 0 + continue + if cause_type == STREAM_DRAINING_ERROR_TYPE: + cooldown_secs = poll_cooldown.total_seconds() + if cooldown_secs > 0: + await asyncio.sleep(cooldown_secs) + continue + if isinstance(err.__cause__, asyncio.CancelledError): + raise err.__cause__ + desc = await self._describe_polled_run(poll_handle, polled_run_id) + if ( + desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW + and self._client is not None + ): + poll_handle = self._client.get_workflow_handle( + self._workflow_id + ) + continue + if desc.run_id and desc.status in ( + WorkflowExecutionStatus.COMPLETED, + WorkflowExecutionStatus.FAILED, + WorkflowExecutionStatus.CANCELED, + WorkflowExecutionStatus.TERMINATED, + WorkflowExecutionStatus.TIMED_OUT, + ): + terminal_handle = poll_handle._client.get_workflow_handle( + self._workflow_id, run_id=desc.run_id + ) continue - return - raise - except WorkflowUpdateRPCTimeoutOrCancelledError as e: - if isinstance(e.__cause__, asyncio.CancelledError): - raise e.__cause__ - if await self._follow_continue_as_new(): - continue - return - except RPCError as e: - # Workflow may have completed between polls; subscribe - # exits cleanly on terminal status so callers don't - # have to wrap the iterator in error handling for the - # normal end-of-stream case. - if e.status != RPCStatusCode.NOT_FOUND: raise - if await self._follow_continue_as_new(): - continue - if await self._workflow_in_terminal_state(): - return - raise converter = self._payload_converter() for wire_item in result.items: payload = _decode_payload(wire_item.data) @@ -643,61 +646,23 @@ async def subscribe( offset=wire_item.offset, ) offset = result.next_offset - cooldown_secs = poll_cooldown.total_seconds() - if not result.more_ready and cooldown_secs > 0: - await asyncio.sleep(cooldown_secs) - - async def _describe_polled_run(self) -> WorkflowExecutionDescription: - """Describe the specific run the most recent poll was admitted to. - - Describing that run (rather than the latest) is what lets a - continue-as-new be detected: a rolled-over run is closed with status - CONTINUED_AS_NEW, whereas the latest run would report RUNNING. Falls - back to the latest run when no run id has been captured yet, or when no - client is available to target a specific run. - """ - if self._client is not None: - return await self._client.get_workflow_handle( - self._workflow_id, run_id=self._polled_run_id + if terminal_handle is not None: + if not result.more_ready: + return + elif not result.more_ready: + cooldown_secs = poll_cooldown.total_seconds() + if cooldown_secs > 0: + await asyncio.sleep(cooldown_secs) + + async def _describe_polled_run( + self, handle: WorkflowHandle[Any, Any], polled_run_id: str | None + ) -> WorkflowExecutionDescription: + """Describe the accepted run, or the targeted execution before acceptance.""" + if polled_run_id is not None: + return await handle._client.get_workflow_handle( + self._workflow_id, run_id=polled_run_id ).describe() - return await self._handle.describe() - - async def _follow_continue_as_new(self) -> bool: - """Check if the polled run continued-as-new and re-target the handle. - - Returns True if the handle was updated (caller should retry). The - successor run id is not needed — re-targeting to an unpinned handle - makes the next poll address the latest (successor) run. - """ - if self._client is None: - return False - try: - desc = await self._describe_polled_run() - except Exception: - return False - if desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW: - self._handle = self._client.get_workflow_handle(self._workflow_id) - return True - return False - - async def _workflow_in_terminal_state(self) -> bool: - """Return True if the polled run has reached a terminal state. - - Used by ``subscribe()`` to distinguish "workflow finished — - stream is done" from "wrong workflow id" when a poll RPC - returns NOT_FOUND. - """ - try: - desc = await self._describe_polled_run() - except Exception: - return False - return desc.status in ( - WorkflowExecutionStatus.COMPLETED, - WorkflowExecutionStatus.FAILED, - WorkflowExecutionStatus.CANCELED, - WorkflowExecutionStatus.TERMINATED, - WorkflowExecutionStatus.TIMED_OUT, - ) + return await handle.describe() async def get_offset(self) -> int: """Query the current global offset (base_offset + log length).""" diff --git a/temporalio/contrib/workflow_streams/_stream.py b/temporalio/contrib/workflow_streams/_stream.py index ae8608c3b9..e2da3a130b 100644 --- a/temporalio/contrib/workflow_streams/_stream.py +++ b/temporalio/contrib/workflow_streams/_stream.py @@ -52,6 +52,7 @@ _PUBLISH_SIGNAL = "__temporal_workflow_stream_publish" _POLL_UPDATE = "__temporal_workflow_stream_poll" +_POLL_QUERY = "__temporal_workflow_stream_poll" _OFFSET_QUERY = "__temporal_workflow_stream_offset" _MAX_POLL_RESPONSE_BYTES = 1_000_000 @@ -84,6 +85,7 @@ class WorkflowStream: - ``__temporal_workflow_stream_publish`` signal — external publish with dedup - ``__temporal_workflow_stream_poll`` update — long-poll subscription + - ``__temporal_workflow_stream_poll`` query — immediate retained-item read - ``__temporal_workflow_stream_offset`` query — current log length Note: @@ -158,6 +160,7 @@ def __init__(self, prior_state: WorkflowStreamState | None = None) -> None: workflow.set_update_handler( _POLL_UPDATE, self._on_poll, validator=self._validate_poll ) + workflow.set_query_handler(_POLL_QUERY, self._collect_batch) workflow.set_query_handler(_OFFSET_QUERY, self._on_offset) def _publish_to_topic(self, topic: str, value: Any) -> None: @@ -337,9 +340,9 @@ async def continue_as_new( def truncate(self, up_to_offset: int) -> None: """Discard log entries before ``up_to_offset``. - After truncation, polls requesting an offset before the new - base will receive an ApplicationError. All global offsets - remain monotonic. + After truncation, poll updates requesting a nonzero offset before + the new base receive an ApplicationError. Poll queries read from + the retained base instead. All global offsets remain monotonic. Raises ApplicationError (not ValueError) when ``up_to_offset`` is past the end of the log so that callers invoking this from @@ -407,21 +410,18 @@ async def _on_poll(self, payload: PollInput) -> PollResult: or self._detaching ), ) - log_offset = payload.from_offset - self._base_offset - if log_offset < 0: - if payload.from_offset == 0: - # "From the beginning" — start at whatever is available. - log_offset = 0 - else: - # Subscriber had a specific position that's been - # truncated. ApplicationError fails this update (client - # gets the error) without crashing the workflow task — - # avoids a poison pill during replay. - raise ApplicationError( - f"Requested offset {payload.from_offset} has been truncated. " - f"Current base offset is {self._base_offset}.", - type=TRUNCATED_OFFSET_ERROR_TYPE, - ) + if payload.from_offset != 0 and payload.from_offset < self._base_offset: + # ApplicationError fails this update without failing the workflow task. + raise ApplicationError( + f"Requested offset {payload.from_offset} has been truncated. " + f"Current base offset is {self._base_offset}.", + type=TRUNCATED_OFFSET_ERROR_TYPE, + ) + return self._collect_batch(payload) + + def _collect_batch(self, payload: PollInput) -> PollResult: + """Read retained items immediately without changing stream state.""" + log_offset = max(0, payload.from_offset - self._base_offset) all_new = self._log[log_offset:] if payload.topics: topic_set = set(payload.topics) diff --git a/temporalio/contrib/workflow_streams/_types.py b/temporalio/contrib/workflow_streams/_types.py index a58cf75aeb..9bfcbe467b 100644 --- a/temporalio/contrib/workflow_streams/_types.py +++ b/temporalio/contrib/workflow_streams/_types.py @@ -5,7 +5,7 @@ ``Payload`` by the payload converter at publish time, and the resulting bytes/metadata are preserved per item so subscribers can decode with ``subscribe(result_type=T)``. The codec chain (e.g. encryption, compression) -applies once at the outer signal/update envelope level — not separately to each +applies once at the outer signal/update/query envelope level — not separately to each embedded item — so codec behavior is symmetric between workflow-side and client-side publishing. @@ -106,7 +106,10 @@ class PublishInput: @dataclass class PollInput: - """Update payload: request to poll for new items. + """Update/query payload: request to poll for new items. + + Updates wait for items and reject truncated nonzero offsets. Queries + return immediately and clamp offsets to the retained base. .. warning:: This class is experimental and may change in future versions. @@ -127,7 +130,7 @@ class _WorkflowStreamWireItem: @dataclass class PollResult: - """Update response: items matching the poll request. + """Update/query response: items matching the poll request. .. warning:: This class is experimental and may change in future versions. diff --git a/tests/contrib/workflow_streams/test_closed_streams.py b/tests/contrib/workflow_streams/test_closed_streams.py new file mode 100644 index 0000000000..d05224c333 --- /dev/null +++ b/tests/contrib/workflow_streams/test_closed_streams.py @@ -0,0 +1,518 @@ +"""Closed-run stream reads and subscription run isolation.""" + +from __future__ import annotations + +import asyncio +import json +import uuid +from dataclasses import dataclass +from datetime import timedelta +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from temporalio import workflow +from temporalio.client import ( + Client, + WorkflowExecutionStatus, + WorkflowHandle, + WorkflowQueryFailedError, + WorkflowQueryRejectedError, + WorkflowUpdateFailedError, +) +from temporalio.common import QueryRejectCondition, RawValue +from temporalio.contrib.workflow_streams import ( + PollInput, + PollResult, + WorkflowStream, + WorkflowStreamClient, +) +from temporalio.contrib.workflow_streams._types import ( + _decode_payload, + _WorkflowStreamWireItem, +) +from temporalio.converter import DataConverter +from temporalio.exceptions import ApplicationError +from temporalio.service import RPCError, RPCStatusCode +from tests.helpers import assert_eq_eventually, new_worker + +POLL = "__temporal_workflow_stream_poll" +WIRE_ABC = "ChgKCGVuY29kaW5nEgxiaW5hcnkvcGxhaW4SA2FiYw==" +TERMINAL_STATUSES = [ + WorkflowExecutionStatus.COMPLETED, + WorkflowExecutionStatus.FAILED, + WorkflowExecutionStatus.CANCELED, + WorkflowExecutionStatus.TERMINATED, + WorkflowExecutionStatus.TIMED_OUT, +] + + +def mock_stream( + status: WorkflowExecutionStatus = WorkflowExecutionStatus.COMPLETED, +) -> tuple[WorkflowStreamClient, Any, Any, Any]: + client = MagicMock(spec=Client) + client.data_converter = DataConverter.default + poll = MagicMock(spec=WorkflowHandle) + poll.id = "workflow" + poll._client = client + poll.start_update = AsyncMock( + side_effect=RPCError("closed", RPCStatusCode.NOT_FOUND, b"") + ) + poll.describe = AsyncMock( + return_value=SimpleNamespace(status=status, run_id="closed-run") + ) + read = MagicMock(spec=WorkflowHandle) + read.describe = poll.describe + read.query = AsyncMock(return_value=PollResult()) + client.get_workflow_handle.return_value = read + return WorkflowStreamClient(poll, client=client), client, poll, read + + +@pytest.mark.parametrize("status", TERMINAL_STATUSES) +async def test_workflow_streams_terminal_pages(status: WorkflowExecutionStatus) -> None: + stream, client, poll, read = mock_stream(status) + read.query.side_effect = [ + PollResult([_WorkflowStreamWireItem("events", WIRE_ABC, 3)], 4, True), + PollResult([_WorkflowStreamWireItem("events", WIRE_ABC, 5)], 6, False), + ] + items = [item async for item in stream.subscribe("events", 3, result_type=bytes)] + assert [(item.offset, item.data) for item in items] == [(3, b"abc"), (5, b"abc")] + client.get_workflow_handle.assert_called_once_with("workflow", run_id="closed-run") + assert poll.start_update.await_count == 1 + assert [call.args for call in read.query.await_args_list] == [ + (POLL, PollInput(["events"], 3)), + (POLL, PollInput(["events"], 4)), + ] + + +async def test_workflow_streams_direct_handle_keeps_default_converter() -> None: + _, client, poll, read = mock_stream() + read.query.return_value = PollResult( + [_WorkflowStreamWireItem("events", WIRE_ABC, 0)], 1 + ) + stream = WorkflowStreamClient(poll) + items = [item async for item in stream.subscribe(result_type=RawValue)] + assert items[0].data.payload.data == b"abc" + client.get_workflow_handle.assert_called_once_with("workflow", run_id="closed-run") + + +@pytest.mark.parametrize( + "failure", ["running", "unknown", "missing_run", "describe", "query"] +) +async def test_workflow_streams_unconfirmed_closure_and_query_errors( + failure: str, +) -> None: + stream, _, poll, read = mock_stream() + expected: Exception + if failure in ("running", "unknown", "missing_run"): + if failure == "missing_run": + poll.describe.return_value.run_id = "" + else: + poll.describe.return_value.status = ( + WorkflowExecutionStatus.RUNNING if failure == "running" else None + ) + expected = poll.start_update.side_effect + elif failure == "describe": + expected = RPCError("missing", RPCStatusCode.NOT_FOUND, b"") + poll.describe.side_effect = expected + else: + expected = WorkflowQueryFailedError("missing handler or query rejected") + read.query.side_effect = expected + with pytest.raises(type(expected)) as err: + _ = [item async for item in stream.subscribe()] + assert err.value is expected + + +async def test_workflow_streams_run_tracking_is_local_to_each_attempt() -> None: + stream, client, poll, _ = mock_stream() + admitted: asyncio.Queue[str] = asyncio.Queue() + outcomes: dict[str, asyncio.Future[PollResult]] = {} + reads: dict[str, Any] = {} + + async def start_update(*_args: Any, **_kwargs: Any) -> Any: + run_id = f"run-{len(outcomes)}" + future: asyncio.Future[PollResult] = asyncio.get_running_loop().create_future() + outcomes[run_id] = future + admitted.put_nowait(run_id) + return SimpleNamespace(workflow_run_id=run_id, result=lambda: future) + + def pinned_handle(workflow_id: str, *, run_id: str | None = None) -> Any: + assert workflow_id == "workflow" + assert run_id is not None + if run_id not in reads: + read = MagicMock(spec=WorkflowHandle) + read.describe = AsyncMock( + return_value=SimpleNamespace( + status=WorkflowExecutionStatus.COMPLETED, run_id=run_id + ) + ) + read.query = AsyncMock( + return_value=PollResult( + [_WorkflowStreamWireItem(run_id, WIRE_ABC, 0)], 1 + ) + ) + reads[run_id] = read + return reads[run_id] + + poll.start_update.side_effect = start_update + client.get_workflow_handle.side_effect = pinned_handle + + async def consume() -> list[str]: + return [item.topic async for item in stream.subscribe()] + + first = asyncio.create_task(consume()) + assert await admitted.get() == "run-0" + second = asyncio.create_task(consume()) + assert await admitted.get() == "run-1" + for run_id in ["run-1", "run-0"]: + outcomes[run_id].set_exception( + WorkflowUpdateFailedError( + ApplicationError("closed", type="AcceptedUpdateCompletedWorkflow") + ) + ) + assert await first == ["run-0"] + assert await second == ["run-1"] + + +async def test_workflow_streams_continue_as_new_resets_accepted_run() -> None: + stream, client, poll, read = mock_stream() + accepted = SimpleNamespace( + workflow_run_id="old-run", + result=AsyncMock( + side_effect=WorkflowUpdateFailedError( + ApplicationError("closed", type="AcceptedUpdateCompletedWorkflow") + ) + ), + ) + poll.start_update.side_effect = None + poll.start_update.return_value = accepted + old = MagicMock(spec=WorkflowHandle) + old.describe = AsyncMock( + return_value=SimpleNamespace( + status=WorkflowExecutionStatus.CONTINUED_AS_NEW, run_id="old-run" + ) + ) + successor = MagicMock(spec=WorkflowHandle) + successor._client = client + successor.start_update = AsyncMock( + side_effect=RPCError("closed", RPCStatusCode.NOT_FOUND, b"") + ) + successor.describe = AsyncMock( + return_value=SimpleNamespace( + status=WorkflowExecutionStatus.COMPLETED, run_id="successor-run" + ) + ) + read.query.return_value = PollResult( + [_WorkflowStreamWireItem("events", WIRE_ABC, 0)], 1 + ) + + def get_handle(workflow_id: str, *, run_id: str | None = None) -> Any: + assert workflow_id == "workflow" + return old if run_id == "old-run" else successor if run_id is None else read + + client.get_workflow_handle.side_effect = get_handle + assert [item.data async for item in stream.subscribe()] == [b"abc"] + old.describe.assert_awaited_once() + successor.describe.assert_awaited_once() + assert client.get_workflow_handle.call_args_list[-1].kwargs == { + "run_id": "successor-run" + } + + +async def test_workflow_streams_shared_wire_vector() -> None: + converter = DataConverter.default + request = '{"topics":["events"],"from_offset":3}' + response = ( + '{"items":[{"topic":"events","data":"' + + WIRE_ABC + + '","offset":3}],"next_offset":4,"more_ready":false}' + ) + for value, expected in [ + (PollInput(["events"], 3), request), + (PollResult([_WorkflowStreamWireItem("events", WIRE_ABC, 3)], 4), response), + ]: + payloads = await converter.encode([value]) + assert json.loads(payloads[0].data) == json.loads(expected) + decoded = (await converter.decode(payloads, [type(value)]))[0] + assert decoded == value + assert _decode_payload(WIRE_ABC).data == b"abc" + assert _decode_payload(WIRE_ABC).metadata == {"encoding": b"binary/plain"} + + +@dataclass +class ClosedStreamInput: + count: int = 3 + size: int = 1 + truncate: int = 0 + fail: bool = False + tail: bool = False + + +@workflow.defn +class ClosedStreamWorkflow: + @workflow.init + def __init__(self, arg: ClosedStreamInput) -> None: + self.stream = WorkflowStream() + self.closed = False + + @workflow.signal + def close(self) -> None: + self.closed = True + + @workflow.signal + def detach(self) -> None: + self.stream.detach_pollers() + + @workflow.run + async def run(self, arg: ClosedStreamInput) -> None: + for i in range(arg.count): + self.stream.topic("events" if i % 2 == 0 else "other").publish( + bytes([i]) * arg.size + ) + self.stream.truncate(arg.truncate) + await workflow.wait_condition(lambda: self.closed) + if arg.tail: + self.stream.topic("events").publish(b"tail") + if arg.fail: + raise ApplicationError("closed stream failure") + + +@pytest.mark.parametrize("status", TERMINAL_STATUSES) +async def test_workflow_streams_read_after_terminal_cold_replay( + client: Client, status: WorkflowExecutionStatus +) -> None: + async with new_worker( + client, ClosedStreamWorkflow, max_cached_workflows=0 + ) as worker: + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(fail=status == WorkflowExecutionStatus.FAILED), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + run_timeout=timedelta(seconds=2) + if status == WorkflowExecutionStatus.TIMED_OUT + else None, + ) + stream = WorkflowStreamClient(handle) + await assert_eq_eventually(3, stream.get_offset) + if status == WorkflowExecutionStatus.CANCELED: + await handle.cancel() + elif status == WorkflowExecutionStatus.TERMINATED: + await handle.terminate() + elif status != WorkflowExecutionStatus.TIMED_OUT: + await handle.signal(ClosedStreamWorkflow.close) + + async def execution_status() -> WorkflowExecutionStatus: + status = (await handle.describe()).status + assert status is not None + return status + + await assert_eq_eventually(status, execution_status) + before = await handle.fetch_history() + items = [item async for item in stream.subscribe(result_type=bytes)] + assert [item.data for item in items] == [b"\x00", b"\x01", b"\x02"] + assert [item.offset for item in items] == [0, 1, 2] + assert (await handle.fetch_history()).events == before.events + + +@pytest.mark.parametrize( + "arg,topics,offset,expected", + [ + (ClosedStreamInput(count=0), None, 0, []), + (ClosedStreamInput(), ["missing"], 0, []), + (ClosedStreamInput(), None, 1, [1, 2]), + (ClosedStreamInput(), None, 3, []), + (ClosedStreamInput(), None, 10, []), + (ClosedStreamInput(count=6, truncate=3), None, 1, [3, 4, 5]), + (ClosedStreamInput(count=8, size=200_000), None, 0, list(range(8))), + (ClosedStreamInput(count=8, size=200_000), ["events"], 0, [0, 2, 4, 6]), + ], +) +async def test_workflow_streams_closed_retention_and_pagination( + client: Client, + arg: ClosedStreamInput, + topics: list[str] | None, + offset: int, + expected: list[int], +) -> None: + async with new_worker(client, ClosedStreamWorkflow) as worker: + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + arg, + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.signal(ClosedStreamWorkflow.close) + await handle.result() + stream = WorkflowStreamClient.create(client, handle.id) + items = [ + item async for item in stream.subscribe(topics, offset, result_type=bytes) + ] + assert [item.offset for item in items] == expected + assert [item.data for item in items] == [ + bytes([i]) * arg.size for i in expected + ] + + +@pytest.mark.filterwarnings( + "ignore::temporalio.workflow.UnfinishedUpdateHandlersWarning" +) +async def test_workflow_streams_accepted_poll_drains_unread_tail( + client: Client, +) -> None: + async with new_worker(client, ClosedStreamWorkflow) as worker: + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(tail=True), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + stream = WorkflowStreamClient(handle) + await assert_eq_eventually(3, stream.get_offset) + accepted = asyncio.Event() + original_start = handle.start_update + + async def start(*args: Any, **kwargs: Any) -> Any: + update = await original_start(*args, **kwargs) + accepted.set() + return update + + async def consume() -> list[bytes]: + return [ + item.data + async for item in stream.subscribe( + from_offset=3, result_type=bytes, poll_cooldown=timedelta(0) + ) + ] + + with patch.object(handle, "start_update", side_effect=start): + task = asyncio.create_task(consume()) + await accepted.wait() + await handle.signal(ClosedStreamWorkflow.close) + assert await task == [b"tail"] + + +async def test_workflow_streams_query_is_immediate_while_draining( + client: Client, +) -> None: + async with new_worker(client, ClosedStreamWorkflow) as worker: + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(count=6, truncate=3), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.signal(ClosedStreamWorkflow.detach) + result = await handle.query(POLL, PollInput([], 1), result_type=PollResult) + assert [item.offset for item in result.items] == [3, 4, 5] + empty = await handle.query( + POLL, PollInput(["missing"], 6), result_type=PollResult + ) + assert empty == PollResult([], 6) + await handle.signal(ClosedStreamWorkflow.close) + await handle.result() + + +async def test_workflow_streams_closed_codec_and_raw_conversion(client: Client) -> None: + from tests.test_converter import SimpleCodec + + config = client.config() + config["data_converter"] = DataConverter(payload_codec=SimpleCodec()) + codec_client = Client(**config) + async with new_worker( + codec_client, ClosedStreamWorkflow, max_cached_workflows=0 + ) as worker: + handle = await codec_client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.signal(ClosedStreamWorkflow.close) + await handle.result() + stream = WorkflowStreamClient.create(codec_client, handle.id) + typed = [item.data async for item in stream.subscribe(result_type=bytes)] + raw = [ + item.data.payload.data + async for item in stream.subscribe(result_type=RawValue) + ] + assert typed == raw == [b"\x00", b"\x01", b"\x02"] + + +async def test_workflow_streams_closed_explicit_run_after_workflow_id_reuse( + client: Client, +) -> None: + async with new_worker(client, ClosedStreamWorkflow) as worker: + workflow_id = f"closed-stream-{uuid.uuid4()}" + old = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(), + id=workflow_id, + task_queue=worker.task_queue, + ) + await old.signal(ClosedStreamWorkflow.close) + await old.result() + new = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(count=1), + id=workflow_id, + task_queue=worker.task_queue, + ) + pinned = WorkflowStreamClient( + client.get_workflow_handle(workflow_id, run_id=old.result_run_id) + ) + + async def consume() -> list[int]: + return [item.offset async for item in pinned.subscribe()] + + assert await asyncio.gather(consume(), consume()) == [[0, 1, 2], [0, 1, 2]] + await new.signal(ClosedStreamWorkflow.close) + await new.result() + latest = WorkflowStreamClient.create(client, workflow_id) + assert [item.offset async for item in latest.subscribe()] == [0] + + +@workflow.defn +class LegacyClosedStreamWorkflow: + @workflow.run + async def run(self) -> None: + pass + + +async def test_workflow_streams_missing_query_and_configured_query_rejection( + client: Client, +) -> None: + async with new_worker( + client, ClosedStreamWorkflow, LegacyClosedStreamWorkflow + ) as worker: + legacy = await client.start_workflow( + LegacyClosedStreamWorkflow.run, + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await legacy.result() + with pytest.raises(WorkflowQueryFailedError): + _ = [item async for item in WorkflowStreamClient(legacy).subscribe()] + + handle = await client.start_workflow( + ClosedStreamWorkflow.run, + ClosedStreamInput(), + id=f"closed-stream-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await handle.signal(ClosedStreamWorkflow.close) + await handle.result() + config = client.config() + config["default_workflow_query_reject_condition"] = ( + QueryRejectCondition.NOT_OPEN + ) + stream = WorkflowStreamClient.create(Client(**config), handle.id) + with pytest.raises(WorkflowQueryRejectedError): + _ = [item async for item in stream.subscribe()] + + missing = WorkflowStreamClient.create(client, f"missing-{uuid.uuid4()}") + with pytest.raises(RPCError) as err: + _ = [item async for item in missing.subscribe()] + assert err.value.status == RPCStatusCode.NOT_FOUND diff --git a/tests/contrib/workflow_streams/test_workflow_streams.py b/tests/contrib/workflow_streams/test_workflow_streams.py index 79d0fa5409..65dbc454fa 100644 --- a/tests/contrib/workflow_streams/test_workflow_streams.py +++ b/tests/contrib/workflow_streams/test_workflow_streams.py @@ -30,7 +30,6 @@ from temporalio import activity, nexus, workflow from temporalio.client import ( Client, - WorkflowExecutionStatus, WorkflowHandle, WorkflowUpdateFailedError, WorkflowUpdateStage, @@ -2185,114 +2184,6 @@ async def test_continue_as_new_helper(client: Client) -> None: await new_handle.signal(ContinueAsNewHelperWorkflow.close) -@pytest.mark.asyncio -async def test_follow_continue_as_new_describes_polled_run(client: Client) -> None: - """Regression test for continue-as-new detection. - - ``_follow_continue_as_new`` must describe the *specific run the poll was - admitted to* — a rolled-over run is closed with status CONTINUED_AS_NEW, - whereas the latest (successor) run reports RUNNING. The previous - implementation described the latest run, so the check never fired and a poll - failure during a rollover stopped the subscription instead of following it. - - Driving the exact poll-failure race deterministically is impractical (the - workflow drains in-flight polls before continuing-as-new), so this asserts - the helper's decision directly against a real post-rollover run. - """ - async with new_worker(client, ContinueAsNewHelperWorkflow) as worker: - handle = await client.start_workflow( - ContinueAsNewHelperWorkflow.run, - CANWorkflowInputTyped(), - id=f"workflow-stream-can-follow-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - await handle.signal( - "__temporal_workflow_stream_publish", - PublishInput( - items=[PublishEntry(topic="events", data=_wire_bytes(b"item-0"))], - publisher_id="pub", - sequence=1, - ), - ) - old_run_id = handle.result_run_id - - await handle.signal(ContinueAsNewHelperWorkflow.trigger_continue) - new_handle = client.get_workflow_handle(handle.id) - await assert_eq_eventually(True, lambda: _is_different_run(handle, new_handle)) - - # The fix's premise: the polled (old) run reports CONTINUED_AS_NEW; the - # latest run reports RUNNING. - old_desc = await client.get_workflow_handle( - handle.id, run_id=old_run_id - ).describe() - assert old_desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW - latest_desc = await client.get_workflow_handle(handle.id).describe() - assert latest_desc.status == WorkflowExecutionStatus.RUNNING - - # The client follows the rollover when it describes the polled run, but - # the previous latest-run behavior (polled_run_id unset) would not. - following = WorkflowStreamClient.create(client, handle.id) - following._polled_run_id = old_run_id - assert await following._follow_continue_as_new() is True - - latest_only = WorkflowStreamClient.create(client, handle.id) - latest_only._polled_run_id = None # describes the latest run, as the bug did - assert await latest_only._follow_continue_as_new() is False - - await new_handle.signal(ContinueAsNewHelperWorkflow.close) - - -@pytest.mark.asyncio -async def test_subscribe_captures_polled_run_id(client: Client) -> None: - """The ``subscribe`` loop must record the run each poll was admitted to. - - ``_follow_continue_as_new`` relies on ``_polled_run_id`` to describe the - *polled* run rather than the latest one (see - ``test_follow_continue_as_new_describes_polled_run``). This verifies the - real subscribe flow populates that field, rather than setting it by hand. - """ - async with new_worker(client, BasicWorkflowStreamWorkflow) as worker: - handle = await client.start_workflow( - BasicWorkflowStreamWorkflow.run, - id=f"workflow-stream-polled-run-{uuid.uuid4()}", - task_queue=worker.task_queue, - ) - await handle.signal( - "__temporal_workflow_stream_publish", - PublishInput( - items=[PublishEntry(topic="events", data=_wire_bytes(b"item-0"))], - publisher_id="pub", - sequence=1, - ), - ) - - stream = WorkflowStreamClient.create(client, handle.id) - received: list[WorkflowStreamItem] = [] - - async def consume() -> None: - async for item in stream.subscribe( - from_offset=0, poll_cooldown=timedelta(0), result_type=bytes - ): - received.append(item) - - async def received_count() -> int: - return len(received) - - task = asyncio.create_task(consume()) - try: - await assert_eq_eventually(1, received_count) - # The poll was admitted to the run we started; subscribe must have - # captured exactly that run id, not left it unset. - assert stream._polled_run_id == handle.result_run_id - finally: - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - await handle.signal(BasicWorkflowStreamWorkflow.close) - - @workflow.defn class DrainingGateWorkflow: """CAN workflow that detaches pollers and then *holds* in the draining state