Don't hang worker shutdown when a Nexus task is cancelled while encoding its failure

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.

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:

tests/nexus/test_nexus_worker_shutdown.py

@@ -4,19 +4,28 @@
 import concurrent.futures
 import threading
 import uuid
+from collections.abc import Sequence
 from dataclasses import dataclass
 from datetime import timedelta
 from typing import Literal
 
 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()
+        # 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]:
+        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