diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index e79cdef77..7ff7d281f 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -3782,6 +3782,11 @@ def __init__( self._seq = seq self._input = input self._task = asyncio.Task(fn) + # Register with the instance like activity/child workflow tasks: on + # eviction or worker shutdown the task must be cancelled with the rest + # instead of being dropped pending (which logs "Task was destroyed but + # it is pending!" once per waiting workflow). + self._instance._register_task(self._task, name=f"nexus: {input.operation_name}") self._start_fut: asyncio.Future[str | None] = instance.create_future() self._result_fut: asyncio.Future[OutputT | None] = instance.create_future() self._payload_converter = payload_converter diff --git a/tests/nexus/test_nexus_worker_shutdown.py b/tests/nexus/test_nexus_worker_shutdown.py index 2a94027d5..6d0ebaa3e 100644 --- a/tests/nexus/test_nexus_worker_shutdown.py +++ b/tests/nexus/test_nexus_worker_shutdown.py @@ -2,6 +2,8 @@ import asyncio import concurrent.futures +import gc +import logging import threading import uuid from dataclasses import dataclass @@ -122,6 +124,74 @@ async def run(self, input: ShutdownTestCallerInput) -> str: ) +@workflow.defn(sandboxed=False) +class NexusPendingOperationCallerWorkflow: + """Caller that parks on a Nexus operation that never starts.""" + + @workflow.run + async def run(self, input: ShutdownTestCallerInput) -> str: + nexus_client = workflow.create_nexus_client( + service=ShutdownTestService, + endpoint=make_nexus_endpoint_name(input.task_queue), + ) + return await nexus_client.execute_operation( + ShutdownTestService.hang_until_cancelled, None + ) + + +async def test_nexus_operation_handle_task_not_destroyed_on_shutdown( + env: WorkflowEnvironment, caplog: pytest.LogCaptureFixture +): + """Dropping a workflow parked on a Nexus operation must not leak a pending task. + + The Nexus operation handle creates its asyncio task directly. Before it was + registered with the workflow instance, dropping the workflow on worker + shutdown (or eviction) left that task pending and asyncio logged "Task was + destroyed but it is pending!" once per waiting workflow -- unlike activity + and child workflow handles, which register their tasks with the instance. + """ + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with time-skipping server") + + caller_task_queue = str(uuid.uuid4()) + # No handler worker is started for this endpoint, so the Nexus operation + # never starts and the caller workflow stays parked on it. + await env.create_nexus_endpoint( + make_nexus_endpoint_name(caller_task_queue), caller_task_queue + ) + + with caplog.at_level(logging.ERROR, logger="asyncio"): + async with Worker( + env.client, + task_queue=caller_task_queue, + workflows=[NexusPendingOperationCallerWorkflow], + ): + handle = await env.client.start_workflow( + NexusPendingOperationCallerWorkflow.run, + ShutdownTestCallerInput( + operation="hang_until_cancelled", task_queue=caller_task_queue + ), + id=str(uuid.uuid4()), + task_queue=caller_task_queue, + ) + # Give the workflow time to schedule the Nexus operation before the + # worker block exits and drops it from the cache + await asyncio.sleep(1) + + # Give the loop a chance to GC the discarded workflow instance and its tasks + await asyncio.sleep(0.1) + gc.collect() + await asyncio.sleep(0.1) + + await handle.terminate("test cleanup") + + assert not [ + r + for r in caplog.records + if "Task was destroyed but it is pending" in r.getMessage() + ] + + async def test_nexus_worker_shutdown(env: WorkflowEnvironment): """Test that Nexus operations are cancelled when worker shuts down without graceful timeout.""" if env.supports_time_skipping: