fix await pending heartbeats instead of cancelling them

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

Author: timl3136

cadence/_internal/activity/_context.py

@@ -21,22 +21,20 @@ def __init__(
         self._info = info
         self._activity_def = activity_def
         self._heartbeat_sender = heartbeat_sender
-        self._heartbeat_tasks: set[asyncio.Task[None]] = set()
+        self._heartbeat_tasks: set[asyncio.Future[None]] = set()
 
     async def execute(self, payload: Payload) -> Any:
         params = self._to_params(payload)
         try:
             with self._activate():
                 return await self._activity_def.impl_fn(*params)
         finally:
-            await self._cancel_pending_heartbeats()
+            await self._wait_pending_heartbeats()
 
-    async def _cancel_pending_heartbeats(self) -> None:
+    async def _wait_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]:
@@ -73,7 +71,10 @@ def __init__(
     async def execute(self, payload: Payload) -> Any:
         params = self._to_params(payload)
         self._loop = asyncio.get_running_loop()
-        return await self._loop.run_in_executor(self._executor, self._run, params)
+        try:
+            return await self._loop.run_in_executor(self._executor, self._run, params)
+        finally:
+            await self._wait_pending_heartbeats()
 
     def _run(self, args: list[Any]) -> Any:
         with self._activate():
@@ -83,6 +84,9 @@ def client(self) -> Client:
         raise RuntimeError("client is only supported in async activities")
 
     def heartbeat(self, *details: Any) -> None:
-        asyncio.run_coroutine_threadsafe(
+        future = asyncio.run_coroutine_threadsafe(
             self._heartbeat_sender.send_heartbeat(*details), self._loop
         )
+        wrapped = asyncio.wrap_future(future, loop=self._loop)
+        self._heartbeat_tasks.add(wrapped)
+        wrapped.add_done_callback(self._heartbeat_tasks.discard)

tests/cadence/_internal/activity/test_activity_executor.py

@@ -326,9 +326,6 @@ async def activity_fn():
     executor = ActivityExecutor(client, "task_list", "identity", 1, reg.get_activity)
 
     await executor.execute(fake_task("activity_type", ""))
-    # 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(
@@ -359,9 +356,6 @@ def activity_fn():
     executor = ActivityExecutor(client, "task_list", "identity", 1, reg.get_activity)
 
     await executor.execute(fake_task("activity_type", ""))
-    # run_coroutine_threadsafe schedules the heartbeat from a thread;
-    # the coroutine 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(