From a2ee29d38693e78a54d0cc45c770ff29034f231b Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 1 Oct 2026 00:57:26 -0500 Subject: [PATCH 1/3] Report an activity's outcome even when a cancel lands during encoding _RunningActivity.cancel() already skips cancelling the task once the activity is done, but done was only set after the result or failure had been encoded. A cancel from Core (workflow cancel, heartbeat failure, or worker shutdown with the default zero grace period) that arrived while encode() or encode_failure() was awaiting raised CancelledError past the handler's except Exception, so complete_activity_task was never called. Core kept the activity outstanding, wait_all_completed never returned and worker.shutdown() hung. Mark the activity done as soon as its code returns or raises, before encoding, and make wait_all_completed collect task exceptions so it keeps its no-raise contract. Add a regression test with a codec that yields during encode; it fails on main with a 20s shutdown timeout. --- CHANGELOG.md | 4 ++ temporalio/worker/_activity.py | 27 ++++++++---- tests/worker/test_activity.py | 79 +++++++++++++++++++++++++++++++++- 3 files changed, 101 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a8995696..61bb8db6d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,10 @@ to include examples, links to docs, or any other relevant information. ### Fixed +- A cancel that arrives while an activity's result or failure is still being encoded no longer + interrupts the reporting, so the completion is still sent. Previously the activity stayed open + on the server side and worker shutdown waited on it forever. + ### Security ## [1.34.0] - 2026-09-30 diff --git a/temporalio/worker/_activity.py b/temporalio/worker/_activity.py index f72df9cf4..98f79b8dd 100644 --- a/temporalio/worker/_activity.py +++ b/temporalio/worker/_activity.py @@ -197,7 +197,14 @@ async def drain_poll_queue(self) -> None: async def wait_all_completed(self) -> None: running_tasks = [v.task for v in self._running_activities.values() if v.task] if running_tasks: - await asyncio.gather(*running_tasks, return_exceptions=False) + # Never let a task exception escape and stall shutdown + for result in await asyncio.gather(*running_tasks, return_exceptions=True): + if isinstance(result, BaseException) and not isinstance( + result, asyncio.CancelledError + ): + logger.warning( + "Activity task raised during worker shutdown", exc_info=result + ) def _handle_cancel_activity_task( self, @@ -223,7 +230,7 @@ def _heartbeat(self, task_token: bytes, *details: Any) -> None: # converter is async, we have to schedule it. If the activity is done, # we do not schedule any more. Technically this should be impossible to # call if the activity is done because this sync call can only be called - # inside the activity and done is set to False when the activity + # inside the activity and done is set to True when the activity # returns. logger = temporalio.activity.logger activity = self._running_activities.get(task_token) @@ -348,9 +355,15 @@ async def _handle_start_activity_task( context, StorageDriverStoreContext(target=store_target) ) try: - result = await self._execute_activity( - start, running_activity, task_token, data_converter - ) + try: + result = await self._execute_activity( + start, running_activity, task_token, data_converter + ) + finally: + # The activity code has finished, so a cancel from here on must + # not interrupt reporting its outcome; done stops cancel() from + # cancelling this task + running_activity.done = True [payload] = await data_converter.encode([result]) completion.result.completed.result.CopyFrom(payload) except BaseException as err: @@ -472,9 +485,7 @@ async def _handle_start_activity_task( # Do final completion try: - # We mark the activity as done and let the currently running - # heartbeat task finish - running_activity.done = True + # Let the currently running heartbeat task finish if running_activity.last_heartbeat_task: try: await running_activity.last_heartbeat_task diff --git a/tests/worker/test_activity.py b/tests/worker/test_activity.py index c2add4769..f33e83d79 100644 --- a/tests/worker/test_activity.py +++ b/tests/worker/test_activity.py @@ -1,5 +1,6 @@ import asyncio import concurrent.futures +import dataclasses import logging import logging.handlers import os @@ -30,7 +31,11 @@ WorkflowHandle, ) from temporalio.common import RawValue, RetryPolicy -from temporalio.converter import DefaultPayloadConverter +from temporalio.converter import ( + DataConverter, + DefaultPayloadConverter, + PayloadCodec, +) from temporalio.exceptions import ( ActivityError, ApplicationError, @@ -1131,6 +1136,78 @@ async def wait_on_event() -> str: assert "Worker graceful shutdown" == await handle.result() +class _SlowEncodeCodec(PayloadCodec): + """Yields while encoding, like a codec that calls a remote KMS.""" + + def __init__(self) -> None: + self.encoding_started = asyncio.Event() + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + self.encoding_started.set() + await asyncio.sleep(3) + return list(payloads) + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return list(payloads) + + +async def test_activity_cancel_during_result_encode_still_completes( + client: Client, worker: ExternalWorker, caplog: pytest.LogCaptureFixture +): + @activity.defn + async def fail_with_details() -> NoReturn: + # Details send the failure through the codec + raise ApplicationError("boom", {"detail": "x"}) + + codec = _SlowEncodeCodec() + config = client.config() + config["data_converter"] = dataclasses.replace( + DataConverter.default, payload_codec=codec + ) + act_task_queue = str(uuid.uuid4()) + act_worker = Worker( + Client(**config), task_queue=act_task_queue, activities=[fail_with_details] + ) + run_task = asyncio.create_task(act_worker.run()) + handle = await client.start_workflow( + "kitchen_sink", + KSWorkflowParams( + actions=[ + KSAction( + execute_activity=KSExecuteActivityAction( + name="fail_with_details", + task_queue=act_task_queue, + retry_max_attempts=1, + ) + ) + ] + ), + id=str(uuid.uuid4()), + task_queue=worker.task_queue, + ) + # Shut down while the failure is still being encoded. With no grace period + # the worker cancels the activity at once, which must not keep its outcome + # from being reported, or shutdown would wait on it forever. + await codec.encoding_started.wait() + with caplog.at_level(logging.DEBUG, logger="temporalio.worker._activity"): + await asyncio.wait_for(act_worker.shutdown(), 20) + await run_task + # The cancel must have reached the activity while it was still encoding, + # otherwise this test proves nothing + assert any( + rec.getMessage().startswith("Cancelling activity") for rec in caplog.records + ) + with pytest.raises(WorkflowFailureError) as err: + await handle.result() + assert isinstance(err.value.cause, ActivityError) + assert isinstance(err.value.cause.cause, ApplicationError) + assert err.value.cause.cause.message == "boom" + + @activity.defn def picklable_wait_on_event() -> str: activity.wait_for_worker_shutdown_sync(20) From 7a4c129000744fd004ed40ee4d653ee1a9b24dbd Mon Sep 17 00:00:00 2001 From: David Hyde Date: Thu, 1 Oct 2026 01:00:39 -0500 Subject: [PATCH 2/3] clean up changelog --- CHANGELOG.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 61bb8db6d..51baf99b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,8 +29,7 @@ to include examples, links to docs, or any other relevant information. ### Fixed - A cancel that arrives while an activity's result or failure is still being encoded no longer - interrupts the reporting, so the completion is still sent. Previously the activity stayed open - on the server side and worker shutdown waited on it forever. + interrupts the reporting, so the completion is still sent. ### Security From abb57512372a99c635c2deb73d83d4090f72de86 Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 1 Oct 2026 01:25:07 -0500 Subject: [PATCH 3/3] Release the slow codec only after the cancel has been delivered A fixed sleep in encode let a slow enough cancel arrive after encoding had finished, in which case an unfixed worker would also pass. The codec now blocks until the test has seen the "Cancelling activity" log, which the cancel handler emits in the same synchronous block as cancel(). --- tests/worker/test_activity.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/tests/worker/test_activity.py b/tests/worker/test_activity.py index f33e83d79..9d303f9dd 100644 --- a/tests/worker/test_activity.py +++ b/tests/worker/test_activity.py @@ -1136,17 +1136,18 @@ async def wait_on_event() -> str: assert "Worker graceful shutdown" == await handle.result() -class _SlowEncodeCodec(PayloadCodec): - """Yields while encoding, like a codec that calls a remote KMS.""" +class _BlockingEncodeCodec(PayloadCodec): + """Blocks in encode until released, like a codec waiting on a remote KMS.""" def __init__(self) -> None: self.encoding_started = asyncio.Event() + self.release = asyncio.Event() async def encode( self, payloads: Sequence[temporalio.api.common.v1.Payload] ) -> list[temporalio.api.common.v1.Payload]: self.encoding_started.set() - await asyncio.sleep(3) + await self.release.wait() return list(payloads) async def decode( @@ -1163,7 +1164,7 @@ async def fail_with_details() -> NoReturn: # Details send the failure through the codec raise ApplicationError("boom", {"detail": "x"}) - codec = _SlowEncodeCodec() + codec = _BlockingEncodeCodec() config = client.config() config["data_converter"] = dataclasses.replace( DataConverter.default, payload_codec=codec @@ -1194,13 +1195,21 @@ async def fail_with_details() -> NoReturn: # from being reported, or shutdown would wait on it forever. await codec.encoding_started.wait() with caplog.at_level(logging.DEBUG, logger="temporalio.worker._activity"): - await asyncio.wait_for(act_worker.shutdown(), 20) + shutdown_task = asyncio.create_task(act_worker.shutdown()) + + # Only let encoding finish once the cancel has reached the activity, + # so it cannot complete first and let an unfixed worker pass + async def cancel_logged() -> None: + while not any( + rec.getMessage().startswith("Cancelling activity") + for rec in caplog.records + ): + await asyncio.sleep(0.01) + + await asyncio.wait_for(cancel_logged(), 20) + codec.release.set() + await asyncio.wait_for(shutdown_task, 20) await run_task - # The cancel must have reached the activity while it was still encoding, - # otherwise this test proves nothing - assert any( - rec.getMessage().startswith("Cancelling activity") for rec in caplog.records - ) with pytest.raises(WorkflowFailureError) as err: await handle.result() assert isinstance(err.value.cause, ActivityError)