From f23845a4599b19f45c549d8df48699ea593c916c Mon Sep 17 00:00:00 2001 From: parbhatkapila4 Date: Sun, 25 Oct 2026 12:40:49 +0530 Subject: [PATCH 1/2] Don't hang worker shutdown when a Nexus task is cancelled while encoding its failure --- changelog/fixed/stubborn-nexus-heron.md | 1 + temporalio/worker/_nexus.py | 18 +++- tests/nexus/test_nexus_worker_shutdown.py | 106 ++++++++++++++++++++++ 3 files changed, 123 insertions(+), 2 deletions(-) create mode 100644 changelog/fixed/stubborn-nexus-heron.md diff --git a/changelog/fixed/stubborn-nexus-heron.md b/changelog/fixed/stubborn-nexus-heron.md new file mode 100644 index 0000000000..175aebfc84 --- /dev/null +++ b/changelog/fixed/stubborn-nexus-heron.md @@ -0,0 +1 @@ +Fix worker shutdown hanging when a Nexus task is cancelled while its failure is still being encoded. diff --git a/temporalio/worker/_nexus.py b/temporalio/worker/_nexus.py index a2f4b8ca75..37131e333b 100644 --- a/temporalio/worker/_nexus.py +++ b/temporalio/worker/_nexus.py @@ -326,7 +326,14 @@ async def _handle_cancel_operation_task( data_converter.payload_converter, completion.failure, ) - await self._encode_completion(completion, data_converter) + try: + await self._encode_completion(completion, data_converter) + except asyncio.CancelledError: + # Core still needs a completion, or worker shutdown hangs + completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( + task_token=task_token, + ack_cancel=task_cancellation.is_cancelled(), + ) await self._complete_task(completion) except Exception: logger.exception("Failed to send Nexus task completion") @@ -391,7 +398,14 @@ async def _handle_start_operation_task( if isinstance(err, concurrent.futures.BrokenExecutor): self._fail_worker_exception_queue.put_nowait(err) - await self._encode_completion(completion, data_converter) + try: + await self._encode_completion(completion, data_converter) + except asyncio.CancelledError: + # Core still needs a completion, or worker shutdown hangs + completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( + task_token=task_token, + ack_cancel=task_cancellation.is_cancelled(), + ) await self._complete_task(completion) except Exception: diff --git a/tests/nexus/test_nexus_worker_shutdown.py b/tests/nexus/test_nexus_worker_shutdown.py index 6ae14987d5..49c899b446 100644 --- a/tests/nexus/test_nexus_worker_shutdown.py +++ b/tests/nexus/test_nexus_worker_shutdown.py @@ -4,6 +4,7 @@ import concurrent.futures import threading import uuid +from collections.abc import Sequence from dataclasses import dataclass from datetime import timedelta from typing import Literal @@ -11,12 +12,20 @@ import nexusrpc import pytest from nexusrpc.handler import ( + CancelOperationContext, + OperationHandler, StartOperationContext, + StartOperationResultAsync, + operation_handler, service_handler, sync_operation, ) from temporalio import nexus, workflow +from temporalio.api.common.v1 import Payload +from temporalio.client import Client +from temporalio.converter import DataConverter, PayloadCodec +from temporalio.exceptions import ApplicationError from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker from tests.helpers.nexus import ( @@ -34,6 +43,18 @@ class ShutdownTestService: hang_until_cancelled: nexusrpc.Operation[None, str] wait_for_shutdown_sync: nexusrpc.Operation[None, str] check_shutdown: nexusrpc.Operation[None, str] + fail_with_details: nexusrpc.Operation[None, str] + fail_cancel_with_details: nexusrpc.Operation[None, str] + + +class FailCancelWithDetailsHandler(OperationHandler[None, str]): + async def start( + self, ctx: StartOperationContext, input: None + ) -> StartOperationResultAsync: + return StartOperationResultAsync(token="token") + + async def cancel(self, ctx: CancelOperationContext, token: str) -> None: + raise ApplicationError("Cancel failed", "some details") @service_handler(service=ShutdownTestService) @@ -83,6 +104,15 @@ async def check_shutdown(self, _ctx: StartOperationContext, _input: None) -> str self.shutdown_check_after = nexus.is_worker_shutdown() return "done" + @sync_operation + async def fail_with_details(self, _ctx: StartOperationContext, _input: None) -> str: + # The details give the failure a payload for the codec to encode + raise ApplicationError("Operation failed", "some details") + + @operation_handler + def fail_cancel_with_details(self) -> OperationHandler[None, str]: + return FailCancelWithDetailsHandler() + @dataclass class ShutdownTestCallerInput: @@ -91,6 +121,8 @@ class ShutdownTestCallerInput: "hang_until_cancelled", "wait_for_shutdown_sync", "check_shutdown", + "fail_with_details", + "fail_cancel_with_details", ] task_queue: str @@ -120,6 +152,16 @@ async def run(self, input: ShutdownTestCallerInput) -> str: return await nexus_client.execute_operation( ShutdownTestService.check_shutdown, None ) + case "fail_with_details": + return await nexus_client.execute_operation( + ShutdownTestService.fail_with_details, None + ) + case "fail_cancel_with_details": + handle = await nexus_client.start_operation( + ShutdownTestService.fail_cancel_with_details, None + ) + handle.cancel() + return await handle async def test_nexus_worker_shutdown(env: WorkflowEnvironment): @@ -348,3 +390,67 @@ async def test_is_worker_shutdown(env: WorkflowEnvironment): assert handler.shutdown_check_before is False assert handler.shutdown_check_after is True assert result == "done" + + +class SlowEncodePayloadCodec(PayloadCodec): + def __init__(self) -> None: + self.encode_started = asyncio.Event() + + async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: + self.encode_started.set() + # Long enough for the worker shutdown to cancel the task mid-encode + await asyncio.sleep(3) + return list(payloads) + + async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: + return list(payloads) + + +@pytest.mark.parametrize("operation", ["fail_with_details", "fail_cancel_with_details"]) +async def test_nexus_worker_shutdown_while_encoding_failure( + env: WorkflowEnvironment, + operation: Literal["fail_with_details", "fail_cancel_with_details"], +): + """Test that shutdown doesn't hang when it cancels a task mid-encode.""" + if env.supports_time_skipping_v1: + pytest.skip("Nexus tests don't work with time-skipping server") + + handler_task_queue = str(uuid.uuid4()) + caller_task_queue = str(uuid.uuid4()) + await env.create_nexus_endpoint( + make_nexus_endpoint_name(handler_task_queue), handler_task_queue + ) + + codec = SlowEncodePayloadCodec() + handler_client_config = env.client.config() + handler_client_config["data_converter"] = DataConverter(payload_codec=codec) + + with concurrent.futures.ThreadPoolExecutor() as executor: + async with Worker( + env.client, + task_queue=caller_task_queue, + workflows=[ShutdownTestCallerWorkflow], + ): + handler_worker = Worker( + Client(**handler_client_config), + task_queue=handler_task_queue, + nexus_service_handlers=[ShutdownTestServiceHandler()], + nexus_task_executor=executor, + ) + + handler_worker_task = asyncio.create_task(handler_worker.run()) + + await env.client.start_workflow( + ShutdownTestCallerWorkflow.run, + ShutdownTestCallerInput( + operation=operation, + task_queue=handler_task_queue, + ), + id=str(uuid.uuid4()), + task_queue=caller_task_queue, + ) + + await codec.encode_started.wait() + + await asyncio.wait_for(handler_worker.shutdown(), 10) + await handler_worker_task From a1ec746f070190ac753c3d4999f8a31fa9879ac2 Mon Sep 17 00:00:00 2001 From: parbhatkapila4 Date: Mon, 26 Oct 2026 22:58:54 +0530 Subject: [PATCH 2/2] Wait on the cancel instead of sleeping in the Nexus shutdown test --- tests/nexus/test_nexus_worker_shutdown.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/nexus/test_nexus_worker_shutdown.py b/tests/nexus/test_nexus_worker_shutdown.py index 49c899b446..ca58a0f83c 100644 --- a/tests/nexus/test_nexus_worker_shutdown.py +++ b/tests/nexus/test_nexus_worker_shutdown.py @@ -398,8 +398,8 @@ def __init__(self) -> None: async def encode(self, payloads: Sequence[Payload]) -> list[Payload]: self.encode_started.set() - # Long enough for the worker shutdown to cancel the task mid-encode - await asyncio.sleep(3) + # Block until the worker shutdown cancels the task mid-encode + await asyncio.Future() return list(payloads) async def decode(self, payloads: Sequence[Payload]) -> list[Payload]: