diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index 90ce6c205..d72d210a3 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -918,6 +918,7 @@ async def _handle_cache_eviction( # skips the instance's eviction job entirely, and a Run whose # watchers outlived it would keep a backend connection open forever. self._external_stream_runtimes.pop(act.run_id, None) + await self._abort_dead_external_output(run_id=act.run_id) self._pending_external_output_stages.pop(act.run_id, None) if self._external_stream_manager is not None: await self._external_stream_manager.evict_run(act.run_id) @@ -1306,6 +1307,56 @@ async def _promote_external_output( else: self._pending_external_output_stages.pop(run_id, None) + async def _abort_dead_external_output(self, *, run_id: str) -> None: + """Abort the staged batches of an evicted Run that History has rejected. + + A Run evicted because its completion was not accepted may have staged + a batch that no marker will ever name. History already says so when + the task was failed or timed out, so such a stage is aborted here + instead of waiting for a reader to find the barrier. A stage whose + marker is authoritative is never re-promoted from here: promotion is + attempted once after the completion, and a stage it left pending is + owed to a cold client's repair, as the contract tests pin down. + """ + backend = self._external_stream_backend + client = self._client + pending = self._pending_external_output_stages.get(run_id) + if backend is None or client is None or not pending: + return + + from temporalio.contrib.external_workflow_streams._output_client import ( + _apply_output_stage_decision, + _history_decision, + _HistoryDecision, + ) + + for manifest in pending: + workflow_id = manifest.stream_key.workflow_id + try: + decision = await _history_decision( + client=client, + workflow_id=workflow_id, + manifest=manifest, + ) + if decision is not _HistoryDecision.ABORT: + continue + await _apply_output_stage_decision( + backend=backend, + manifest=manifest, + decision=decision, + ) + except asyncio.CancelledError: + raise + except Exception as err: + self._stream_metrics.record(err) + logger.warning( + "Could not abort the dead external output stage %s for " + "Workflow %s at eviction; it remains pending for a client", + manifest.stage_token, + workflow_id, + exc_info=True, + ) + def _replay_stream_converters( self, plan: Any ) -> dict[int, temporalio.converter.DataConverter]: diff --git a/tests/streams/test_redis_replay.py b/tests/streams/test_redis_replay.py index 0493dbcf3..9d078a712 100644 --- a/tests/streams/test_redis_replay.py +++ b/tests/streams/test_redis_replay.py @@ -45,7 +45,7 @@ ) from temporalio.streams._wire import to_wire from temporalio.streams.providers import redis as redis_provider -from temporalio.streams.providers.redis import RedisStreams, _chain +from temporalio.streams.providers.redis import RedisStreams, _chain, _TopicLogBackend from temporalio.worker import Replayer, Worker from tests.streams.test_streams_conformance import StreamHost, take from tests.streams.test_streams_workflow import ( @@ -1208,6 +1208,59 @@ async def test_a_workflow_reader_at_end_skips_what_was_there_and_replays( await Replayer(workflows=[FromNowOnGo], plugins=[provider]).replay_workflow(history) +async def test_a_batch_staged_for_a_rejected_task_does_not_stay_in_the_log( + live_client: Client, monkeypatch: pytest.MonkeyPatch +): + # The first stage takes longer than the Workflow Task timeout, so the server + # times the task out and the worker's completion is rejected; the run is + # evicted and run again, and the second attempt stages the same batch under + # a new token. The first stage can never be named by a marker. It is + # settled against History when the run is evicted and its entries leave the + # log, so the log holds the batch once and nothing counts the dead one. + class SlowFirstStage(_TopicLogBackend): + delayed = False + + async def stage_output(self, manifest: Any, records: Any) -> Any: + if not SlowFirstStage.delayed: + SlowFirstStage.delayed = True + await asyncio.sleep(3) + return await super().stage_output(manifest, records) + + monkeypatch.setattr(redis_provider, "_TopicLogBackend", SlowFirstStage) + provider = RedisStreams( + url=redis_url(), key_prefix=f"streams-redis-{uuid.uuid4().hex}" + ) + workflow_id = f"streams-redis-rejected-{uuid.uuid4().hex}" + try: + async with Worker( + live_client, + task_queue=f"tq-{workflow_id}", + workflows=[OneLine], + plugins=[provider], + ): + handle = await live_client.start_workflow( + OneLine.run, + id=workflow_id, + task_queue=f"tq-{workflow_id}", + task_timeout=timedelta(seconds=1), + ) + await asyncio.wait_for(handle.result(), 60) + assert SlowFirstStage.delayed + # Counted before any reader could settle a barrier: the worker did. + assert ( + await _log_length(provider, live_client, workflow_id, DECISIONS.name) + == 2 + ) + stream = provider.get_stream_handle(live_client, workflow_id) + records = [r async for r in stream.read(topic=DECISIONS)] + assert [(r.kind, r.value) for r in records] == [ + (RecordKind.DATA, {"from": "workflow"}), + (RecordKind.FINISH, None), + ] + finally: + await provider.close() + + async def _channel_support(client: Client) -> Any: """What the server offers the readers: no channels, independent ones or linked ones.""" from tests.contrib.external_workflow_streams.conftest import server_channel_support