fix test and lint

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

Author: timl3136

cadence/_internal/activity/_context.py

@@ -1,6 +1,5 @@
 import asyncio
 from concurrent.futures.thread import ThreadPoolExecutor
-from logging import getLogger
 from typing import Any
 
 from cadence import Client
@@ -9,8 +8,6 @@
 from cadence.activity import ActivityInfo, ActivityContext
 from cadence.api.v1.common_pb2 import Payload
 
-_logger = getLogger(__name__)
-
 
 class _Context(ActivityContext):
     def __init__(
@@ -42,10 +39,7 @@ def info(self) -> ActivityInfo:
         return self._info
 
     def heartbeat(self, *details: Any) -> None:
-        task = asyncio.ensure_future(
-            self._heartbeat_sender.send_heartbeat(*details)
-        )
-        task.add_done_callback(_on_heartbeat_done)
+        asyncio.ensure_future(self._heartbeat_sender.send_heartbeat(*details))
 
 
 class _SyncContext(_Context):
@@ -73,13 +67,6 @@ def client(self) -> Client:
         raise RuntimeError("client is only supported in async activities")
 
     def heartbeat(self, *details: Any) -> None:
-        future = asyncio.run_coroutine_threadsafe(
+        asyncio.run_coroutine_threadsafe(
             self._heartbeat_sender.send_heartbeat(*details), self._loop
         )
-        future.add_done_callback(_on_heartbeat_done)
-
-
-def _on_heartbeat_done(task: asyncio.Task | asyncio.Future) -> None:
-    exc = task.exception()
-    if exc is not None:
-        _logger.warning("Heartbeat failed: %s", exc)

cadence/_internal/activity/_heartbeat.py

@@ -1,9 +1,12 @@
+from logging import getLogger
 from typing import Any
 
 from cadence.api.v1.service_worker_pb2 import RecordActivityTaskHeartbeatRequest
 from cadence.api.v1.service_worker_pb2_grpc import WorkerAPIStub
 from cadence.data_converter import DataConverter
 
+_logger = getLogger(__name__)
+
 
 class _HeartbeatSender:
     def __init__(
@@ -19,11 +22,14 @@ def __init__(
         self._identity = identity
 
     async def send_heartbeat(self, *details: Any) -> None:
-        payload = self._data_converter.to_data(list(details))
-        await self._worker_stub.RecordActivityTaskHeartbeat(
-            RecordActivityTaskHeartbeatRequest(
-                task_token=self._task_token,
-                details=payload,
-                identity=self._identity,
+        try:
+            payload = self._data_converter.to_data(list(details))
+            await self._worker_stub.RecordActivityTaskHeartbeat(
+                RecordActivityTaskHeartbeatRequest(
+                    task_token=self._task_token,
+                    details=payload,
+                    identity=self._identity,
+                )
             )
-        )
+        except Exception:
+            _logger.warning("Heartbeat failed", exc_info=True)

tests/integration_tests/workflow/test_heartbeat.py

@@ -84,12 +84,14 @@ async def test_activity_without_heartbeat_times_out(helper: CadenceHelper):
             execution_start_to_close_timeout=timedelta(seconds=30),
         )
 
-        response: GetWorkflowExecutionHistoryResponse = await worker.client.workflow_stub.GetWorkflowExecutionHistory(
-            GetWorkflowExecutionHistoryRequest(
-                domain=DOMAIN_NAME,
-                workflow_execution=execution,
-                wait_for_new_event=True,
-                skip_archival=True,
+        response: GetWorkflowExecutionHistoryResponse = (
+            await worker.client.workflow_stub.GetWorkflowExecutionHistory(
+                GetWorkflowExecutionHistoryRequest(
+                    domain=DOMAIN_NAME,
+                    workflow_execution=execution,
+                    wait_for_new_event=True,
+                    skip_archival=True,
+                )
             )
         )