Swallowed cancellation (#879)

* Add SDK Flag for ensure cancellation is respected when an activity is completing

* Linting

* Revert "Linting"

This reverts commit 898ed1d.

* Linting

* Remove log statements

* Remove test until enabled outside of the replayer

* Merge

* Adding additional replay test, and a global override for a non-replay test

* Address some test comments, fix sdk-core commit change

Author: tconley1428

temporalio/worker/_workflow_instance.py

@@ -17,6 +17,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import timedelta
+from enum import IntEnum
 from typing import (
     Any,
     Awaitable,
@@ -367,8 +368,8 @@ def activate(
         self._continue_as_new_suggested = act.continue_as_new_suggested
         self._time_ns = act.timestamp.ToNanoseconds()
         self._is_replaying = act.is_replaying
-
         self._current_thread_id = threading.get_ident()
+        self._current_internal_flags = act.available_internal_flags
         activation_err: Optional[Exception] = None
         try:
             # Split into job sets with patches, then signals + updates, then
@@ -1613,6 +1614,20 @@ async def run_activity() -> Any:
                     # dict with its new seq
                     self._pending_activities[handle._seq] = handle
                 except asyncio.CancelledError:
+                    # If an activity future completes at the same time as a cancellation is being processed, the cancellation would be swallowed
+                    # _WorkflowLogicFlag.RAISE_ON_CANCELLING_COMPLETED_ACTIVITY will correctly reraise the exception
+                    if handle._result_fut.done():
+                        # TODO in next release, check sdk flag when not replaying instead of global override, remove the override, and set flag use
+                        if (
+                            (
+                                not self._is_replaying
+                                and _raise_on_cancelling_completed_activity_override
+                            )
+                            or _WorkflowLogicFlag.RAISE_ON_CANCELLING_COMPLETED_ACTIVITY
+                            in self._current_internal_flags
+                        ):
+                            # self._current_completion.successful.used_internal_flags.append(WorkflowLogicFlag.RAISE_ON_CANCELLING_COMPLETED_ACTIVITY)
+                            raise
                     # Send a cancel request to the activity
                     handle._apply_cancel_command(self._add_command())
 
@@ -3112,3 +3127,13 @@ def _make_unfinished_signal_handler_message(
             [{"name": name, "count": count} for name, count in names.most_common()]
         )
     )
+
+
+class _WorkflowLogicFlag(IntEnum):
+    """Flags that may be set on task/activation completion to differentiate new from old workflow behavior."""
+
+    RAISE_ON_CANCELLING_COMPLETED_ACTIVITY = 1
+
+
+# Used by tests to validate behavior prior to SDK flag becoming default
+_raise_on_cancelling_completed_activity_override = False

tests/worker/test_replayer.py

@@ -1,4 +1,6 @@
 import asyncio
+import random
+import re
 import sys
 import uuid
 from dataclasses import dataclass
@@ -18,6 +20,7 @@
     WorkflowInboundInterceptor,
     WorkflowInterceptorClassInput,
 )
+from temporalio.workflow import NondeterminismError
 from tests.helpers import assert_eq_eventually
 from tests.worker.test_workflow import (
     ActivityAndSignalsWhileWorkflowDown,
@@ -503,3 +506,55 @@ async def test_replayer_alternate_async_ordering() -> None:
         "counter-2",
         "act-done",
     ]
+
+
+@activity.defn
+async def short_activity_async():
+    delay = random.uniform(0.05, 0.15)  # 50~150ms delay
+    await asyncio.sleep(delay)
+    return 1
+
+
+@workflow.defn
+class QuickActivityWorkflow:
+    @workflow.run
+    async def run(self, total_seconds: float = 10.0):
+        workflow.logger.info("Duration: %f", total_seconds)
+        end = workflow.now() + timedelta(seconds=total_seconds)
+        while True:
+            workflow.logger.info("Stage 1")
+            res = await workflow.execute_activity(
+                short_activity_async, schedule_to_close_timeout=timedelta(seconds=10)
+            )
+            workflow.logger.info("Stage 2, %s", res)
+
+            if workflow.now() > end:
+                break
+
+
+async def test_swallowed_activity_cancellation() -> None:
+    with (
+        Path(__file__)
+        .with_name("test_replayer_swallowed_activity_cancellation.json")
+        .open() as f
+    ):
+        history = f.read()
+    with pytest.raises(NondeterminismError):
+        await Replayer(
+            workflows=[QuickActivityWorkflow],
+            interceptors=[WorkerWorkflowResultInterceptor()],
+        ).replay_workflow(WorkflowHistory.from_json("fake", history))
+
+
+async def test_swallowed_activity_cancellation_no_flag() -> None:
+    with (
+        Path(__file__)
+        .with_name("test_replayer_swallowed_activity_cancellation.json")
+        .open() as f
+    ):
+        history = f.read()
+        history = re.sub(r'"langUsedFlags": \[\s*1\s*]', "", history)
+    await Replayer(
+        workflows=[QuickActivityWorkflow],
+        interceptors=[WorkerWorkflowResultInterceptor()],
+    ).replay_workflow(WorkflowHistory.from_json("fake", history))