feat: implement heartbeat progress recovery through heartbeat details

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

Author: timl3136

cadence/_internal/activity/_activity_executor.py

@@ -62,6 +62,7 @@ def _create_context(
             self._data_converter,
             task.task_token,
             self._identity,
+            task.heartbeat_details,
         )
 
         if activity_def.strategy == ExecutionStrategy.ASYNC:

cadence/_internal/activity/_context.py

@@ -1,6 +1,6 @@
 import asyncio
 from concurrent.futures.thread import ThreadPoolExecutor
-from typing import Any
+from typing import Any, Type
 
 from cadence import Client
 from cadence._internal.activity._definition import BaseDefinition
@@ -55,6 +55,9 @@ def heartbeat(self, *details: Any) -> None:
         self._heartbeat_tasks.add(heartbeat_task)
         heartbeat_task.add_done_callback(self._heartbeat_tasks.discard)
 
+    def heartbeat_details(self, *types: Type) -> list[Any]:
+        return self._heartbeat_sender.get_details(*types)
+
 
 class _SyncContext(_Context):
     def __init__(

cadence/_internal/activity/_heartbeat.py

@@ -1,6 +1,7 @@
 from logging import getLogger
-from typing import Any
+from typing import Any, Type
 
+from cadence.api.v1.common_pb2 import Payload
 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
@@ -15,11 +16,16 @@ def __init__(
         data_converter: DataConverter,
         task_token: bytes,
         identity: str,
+        previous_details: Payload,
     ):
         self._worker_stub = worker_stub
         self._data_converter = data_converter
         self._task_token = task_token
         self._identity = identity
+        self._previous_details = previous_details
+
+    def get_details(self, *types: Type) -> list[Any]:
+        return self._data_converter.from_data(self._previous_details, list(types))
 
     async def send_heartbeat(self, *details: Any) -> None:
         try:

cadence/activity.py

@@ -68,6 +68,17 @@ def heartbeat(*details: Any) -> None:
     ActivityContext.get().heartbeat(*details)
 
 
+def heartbeat_details(*types: Type) -> list[Any]:
+    """Return heartbeat details from the previous attempt.
+
+    Pass type hints to decode the values into specific Python types:
+        step, total = activity.heartbeat_details(int, int)
+
+    Without type hints, returns raw JSON-decoded values.
+    """
+    return ActivityContext.get().heartbeat_details(*types)
+
+
 class ActivityContext(ABC):
     _var: ContextVar["ActivityContext"] = ContextVar("activity")
 
@@ -80,6 +91,9 @@ def client(self) -> Client: ...
     @abstractmethod
     def heartbeat(self, *details: Any) -> None: ...
 
+    @abstractmethod
+    def heartbeat_details(self, *types: Type) -> list[Any]: ...
+
     @contextmanager
     def _activate(self) -> Iterator[None]:
         token = ActivityContext._var.set(self)

cadence/data_converter.py

@@ -39,7 +39,7 @@ def _decode_whitespace_delimited(
     ) -> List[Any]:
         results: List[Any] = []
         start, end = 0, len(payload)
-        while start < end and len(results) < len(type_hints):
+        while start < end and (not type_hints or len(results) < len(type_hints)):
             remaining = payload[start:end]
             (value, value_end) = self._decoder.raw_decode(remaining)
             start += value_end + 1
@@ -51,6 +51,8 @@ def _decode_whitespace_delimited(
     def _convert_into(
         values: List[Any], type_hints: Sequence[Type | None]
     ) -> List[Any]:
+        if not type_hints:
+            return list(values)
         results: List[Any] = []
         for i, type_hint in enumerate(type_hints):
             if not type_hint or type_hint is Any:

tests/cadence/_internal/activity/test_activity_executor.py

@@ -281,7 +281,11 @@ def fake_info(activity_type: str) -> ActivityInfo:
     )
 
 
-def fake_task(activity_type: str, input_json: str) -> PollForActivityTaskResponse:
+def fake_task(
+    activity_type: str,
+    input_json: str,
+    heartbeat_details: str = "",
+) -> PollForActivityTaskResponse:
     return PollForActivityTaskResponse(
         task_token=b"task_token",
         workflow_domain="workflow_domain",
@@ -298,6 +302,9 @@ def fake_task(activity_type: str, input_json: str) -> PollForActivityTaskRespons
         scheduled_time=from_datetime(datetime(2020, 1, 2, 3)),
         started_time=from_datetime(datetime(2020, 1, 2, 4)),
         start_to_close_timeout=from_timedelta(timedelta(seconds=2)),
+        heartbeat_details=Payload(data=heartbeat_details.encode())
+        if heartbeat_details
+        else Payload(),
     )
 
 
@@ -365,3 +372,82 @@ def activity_fn():
             identity="identity",
         )
     )
+
+
+async def test_heartbeat_details_recovery_async(client):
+    worker_stub = client.worker_stub
+    worker_stub.RespondActivityTaskCompleted = AsyncMock(
+        return_value=RespondActivityTaskCompletedResponse()
+    )
+
+    reg = Registry()
+
+    @reg.activity(name="activity_type")
+    async def activity_fn():
+        return activity.heartbeat_details(str, int)
+
+    executor = ActivityExecutor(client, "task_list", "identity", 1, reg.get_activity)
+
+    await executor.execute(
+        fake_task("activity_type", "", heartbeat_details='"progress" 42')
+    )
+
+    worker_stub.RespondActivityTaskCompleted.assert_called_once_with(
+        RespondActivityTaskCompletedRequest(
+            task_token=b"task_token",
+            result=Payload(data=b'["progress",42]'),
+            identity="identity",
+        )
+    )
+
+
+async def test_heartbeat_details_recovery_sync(client):
+    worker_stub = client.worker_stub
+    worker_stub.RespondActivityTaskCompleted = AsyncMock(
+        return_value=RespondActivityTaskCompletedResponse()
+    )
+
+    reg = Registry()
+
+    @reg.activity(name="activity_type")
+    def activity_fn():
+        return activity.heartbeat_details(str, int)
+
+    executor = ActivityExecutor(client, "task_list", "identity", 1, reg.get_activity)
+
+    await executor.execute(
+        fake_task("activity_type", "", heartbeat_details='"progress" 42')
+    )
+
+    worker_stub.RespondActivityTaskCompleted.assert_called_once_with(
+        RespondActivityTaskCompletedRequest(
+            task_token=b"task_token",
+            result=Payload(data=b'["progress",42]'),
+            identity="identity",
+        )
+    )
+
+
+async def test_heartbeat_details_empty_when_no_previous_heartbeat(client):
+    worker_stub = client.worker_stub
+    worker_stub.RespondActivityTaskCompleted = AsyncMock(
+        return_value=RespondActivityTaskCompletedResponse()
+    )
+
+    reg = Registry()
+
+    @reg.activity(name="activity_type")
+    async def activity_fn():
+        return activity.heartbeat_details()
+
+    executor = ActivityExecutor(client, "task_list", "identity", 1, reg.get_activity)
+
+    await executor.execute(fake_task("activity_type", ""))
+
+    worker_stub.RespondActivityTaskCompleted.assert_called_once_with(
+        RespondActivityTaskCompletedRequest(
+            task_token=b"task_token",
+            result=Payload(data=b"[]"),
+            identity="identity",
+        )
+    )

tests/cadence/_internal/activity/test_heartbeat.py

@@ -32,6 +32,7 @@ def sender(worker_stub, data_converter) -> _HeartbeatSender:
         data_converter=data_converter,
         task_token=b"task_token",
         identity="test-identity",
+        previous_details=Payload(),
     )