add cancellation of heartbeat when activity completes

Signed-off-by: Tim Li <ltim@uber.com>

Author: timl3136

cadence/_internal/activity/_activity_executor.py

@@ -1,7 +1,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from logging import getLogger
 from traceback import format_exception
-from typing import Any, Callable, cast
+from typing import Any, Callable, Union, cast
 from google.protobuf.duration import to_timedelta
 from google.protobuf.timestamp import to_datetime
 
@@ -47,7 +47,9 @@ async def execute(self, task: PollForActivityTaskResponse):
             _logger.exception("Activity failed")
             await self._report_failure(task, e)
 
-    def _create_context(self, task: PollForActivityTaskResponse) -> _Context:
+    def _create_context(
+        self, task: PollForActivityTaskResponse
+    ) -> Union[_Context, _SyncContext]:
         activity_type = task.activity_type.name
         try:
             activity_def = cast(BaseDefinition, self._registry(activity_type))

cadence/_internal/activity/_context.py

@@ -25,8 +25,19 @@ def __init__(
 
     async def execute(self, payload: Payload) -> Any:
         params = self._to_params(payload)
-        with self._activate():
-            return await self._activity_def.impl_fn(*params)
+        try:
+            with self._activate():
+                return await self._activity_def.impl_fn(*params)
+        finally:
+            await self._cancel_pending_heartbeats()
+
+    async def _cancel_pending_heartbeats(self) -> None:
+        if not self._heartbeat_tasks:
+            return
+        tasks = list(self._heartbeat_tasks)
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
 
     def _to_params(self, payload: Payload) -> list[Any]:
         return self._activity_def.signature.params_from_payload(
@@ -40,9 +51,11 @@ def info(self) -> ActivityInfo:
         return self._info
 
     def heartbeat(self, *details: Any) -> None:
-        task = asyncio.create_task(self._heartbeat_sender.send_heartbeat(*details))
-        self._heartbeat_tasks.add(task)
-        task.add_done_callback(self._heartbeat_tasks.discard)
+        heartbeat_task = asyncio.create_task(
+            self._heartbeat_sender.send_heartbeat(*details)
+        )
+        self._heartbeat_tasks.add(heartbeat_task)
+        heartbeat_task.add_done_callback(self._heartbeat_tasks.discard)
 
 
 class _SyncContext(_Context):

tests/cadence/_internal/activity/test_activity_executor.py

@@ -326,8 +326,9 @@ async def activity_fn():
     executor = ActivityExecutor(client, "task_list", "identity", 1, reg.get_activity)
 
     await executor.execute(fake_task("activity_type", ""))
-    # Allow the event loop to process the fire-and-forget heartbeat task
-    await asyncio.sleep(0)
+    # ensure_future schedules the heartbeat as a fire-and-forget task;
+    # it needs multiple event-loop iterations to complete.
+    await asyncio.sleep(0.1)
 
     worker_stub.RespondActivityTaskCompleted.assert_called_once()
     worker_stub.RecordActivityTaskHeartbeat.assert_called_once_with(