diff --git a/changelog/fixed/somersaulting-sloth.md b/changelog/fixed/somersaulting-sloth.md new file mode 100644 index 0000000000..28d82fb6b1 --- /dev/null +++ b/changelog/fixed/somersaulting-sloth.md @@ -0,0 +1 @@ +`WorkflowStreamClient.subscribe` now propagates task cancellation instead of ending the subscription normally. diff --git a/temporalio/contrib/workflow_streams/_client.py b/temporalio/contrib/workflow_streams/_client.py index 202558fd38..8ed6cb5b51 100644 --- a/temporalio/contrib/workflow_streams/_client.py +++ b/temporalio/contrib/workflow_streams/_client.py @@ -585,7 +585,7 @@ async def subscribe( self._polled_run_id = handle.workflow_run_id result: PollResult = await handle.result() except asyncio.CancelledError: - return + raise except WorkflowUpdateFailedError as e: cause_type = getattr(e.cause, "type", None) if cause_type == TRUNCATED_OFFSET_ERROR_TYPE: @@ -611,7 +611,9 @@ async def subscribe( continue return raise - except WorkflowUpdateRPCTimeoutOrCancelledError: + except WorkflowUpdateRPCTimeoutOrCancelledError as e: + if isinstance(e.__cause__, asyncio.CancelledError): + raise e.__cause__ if await self._follow_continue_as_new(): continue return diff --git a/tests/contrib/workflow_streams/test_workflow_streams.py b/tests/contrib/workflow_streams/test_workflow_streams.py index a98d467c46..919bce6bd6 100644 --- a/tests/contrib/workflow_streams/test_workflow_streams.py +++ b/tests/contrib/workflow_streams/test_workflow_streams.py @@ -50,7 +50,7 @@ ) from temporalio.contrib.workflow_streams._types import _encode_payload from temporalio.converter import DataConverter, PayloadCodec -from temporalio.exceptions import ApplicationError +from temporalio.exceptions import ActivityError, ApplicationError from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation from temporalio.service import RPCError, RPCStatusCode from temporalio.testing import WorkflowEnvironment @@ -103,6 +103,37 @@ async def run(self) -> None: await workflow.wait_condition(lambda: self._closed) +@workflow.defn +class CancelSubscriptionWorkflow: + @workflow.init + def __init__(self) -> None: + self.stream = WorkflowStream() + self._cancel_requested = False + + @workflow.signal + def cancel_subscription(self) -> None: + self._cancel_requested = True + + @workflow.run + async def run(self) -> str: + self.stream.topic("events", type=bytes).publish(b"seed") + handle = workflow.start_activity( + "subscribe_until_cancelled", + start_to_close_timeout=timedelta(seconds=30), + heartbeat_timeout=timedelta(seconds=1), + cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED, + ) + await workflow.wait_condition(lambda: self._cancel_requested) + handle.cancel() + try: + result = await handle + except ActivityError as err: + result = type(err.cause).__name__ + self.stream.detach_pollers() + await workflow.wait_condition(workflow.all_handlers_finished) + return result + + @workflow.defn class ActivityPublishWorkflow: @workflow.init @@ -375,6 +406,31 @@ async def publish_items(count: int) -> None: client.topic("events", type=bytes).publish(f"item-{i}".encode()) +class CancellableSubscriber: + def __init__(self) -> None: + self.started = asyncio.Event() + + @activity.defn(name="subscribe_until_cancelled") + async def subscribe(self) -> str: + async def heartbeat() -> None: + while True: + activity.heartbeat() + await asyncio.sleep(0.1) + + heartbeat_task = asyncio.create_task(heartbeat()) + try: + stream = WorkflowStreamClient.from_within_activity() + try: + async for _ in stream.subscribe(result_type=bytes): + self.started.set() + except asyncio.CancelledError: + return "subscription-cancelled" + return "subscription-ended" + finally: + heartbeat_task.cancel() + await asyncio.gather(heartbeat_task, return_exceptions=True) + + @activity.defn(name="publish_multi_topic") async def publish_multi_topic(count: int) -> None: topics = ["a", "b", "c"] @@ -1091,7 +1147,7 @@ async def test_priority_flush(client: Client) -> None: @pytest.mark.asyncio async def test_iterator_cancellation(client: Client) -> None: """Cancelling a subscription iterator after it has yielded an item - completes cleanly.""" + propagates cancellation.""" async with new_worker( client, BasicWorkflowStreamWorkflow, @@ -1127,10 +1183,8 @@ async def subscribe_and_collect() -> None: async with _async_timeout(5): await first_item.wait() task.cancel() - try: + with pytest.raises(asyncio.CancelledError): await task - except asyncio.CancelledError: - pass assert len(items) == 1 assert items[0].data == b"seed" @@ -1138,6 +1192,25 @@ async def subscribe_and_collect() -> None: await handle.signal(BasicWorkflowStreamWorkflow.close) +@pytest.mark.asyncio +async def test_activity_subscription_propagates_cancellation(client: Client) -> None: + subscriber = CancellableSubscriber() + async with new_worker( + client, + CancelSubscriptionWorkflow, + activities=[subscriber.subscribe], + ) as worker: + handle = await client.start_workflow( + CancelSubscriptionWorkflow.run, + id=f"workflow-stream-activity-cancel-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + async with _async_timeout(5): + await subscriber.started.wait() + await handle.signal(CancelSubscriptionWorkflow.cancel_subscription) + assert await handle.result() == "subscription-cancelled" + + @pytest.mark.asyncio async def test_context_manager_flushes_on_exit(client: Client) -> None: """Context manager exit flushes all buffered items."""