Keep activity heartbeats in order when encoding is slow

changelog/fixed/orderly-otter.md

@@ -0,0 +1 @@
+Activity heartbeats are now recorded in the order they were sent. With a payload codec that is slow to encode, an older heartbeat could be recorded after a newer one, so a retried activity resumed from stale heartbeat details.

temporalio/worker/_activity.py

@@ -165,9 +165,8 @@ async def raise_from_exception_queue() -> NoReturn:
 
                 if task.HasField("start"):
                     # Cancelled event and sync field will be updated inside
-                    # _run_activity when the activity function is obtained. Max
-                    # size of 1000 should be plenty for the heartbeat queue.
-                    activity = _RunningActivity(pending_heartbeats=asyncio.Queue(1000))
+                    # _run_activity when the activity function is obtained.
+                    activity = _RunningActivity()
                     activity.task = asyncio.create_task(
                         self._handle_start_activity_task(
                             task.task_token, task.start, activity
@@ -246,9 +245,9 @@ def _heartbeat(self, task_token: bytes, *details: Any) -> None:
         logger = temporalio.activity.logger
         activity = self._running_activities.get(task_token)
         if activity and not activity.done:
-            # Put on queue and schedule a task. We will let the queue-full error
-            # be thrown here.
-            activity.pending_heartbeats.put_nowait(details)
+            # Only the newest heartbeat gets sent, so it replaces one that hasn't
+            # been picked up yet
+            activity.pending_heartbeat = details
             activity.last_heartbeat_task = asyncio.create_task(
                 self._heartbeat_async(logger, activity, task_token)
             )
@@ -259,10 +258,20 @@ async def _heartbeat_async(
         activity: _RunningActivity,
         task_token: bytes,
     ) -> None:
-        # Drain the queue, only taking the last value to actually heartbeat
-        details: Sequence[Any] | None = None
-        while not activity.pending_heartbeats.empty():
-            details = activity.pending_heartbeats.get_nowait()
+        # One heartbeat at a time, so a slow encode can't let an older heartbeat
+        # reach core after a newer one
+        async with activity.heartbeat_lock:
+            await self._send_latest_heartbeat(logger, activity, task_token)
+
+    async def _send_latest_heartbeat(
+        self,
+        logger: logging.LoggerAdapter,
+        activity: _RunningActivity,
+        task_token: bytes,
+    ) -> None:
+        # Nothing to send if an earlier task already sent the newest heartbeat
+        details = activity.pending_heartbeat
+        activity.pending_heartbeat = None
         if details is None:
             return
 
@@ -716,12 +725,13 @@ def assert_activity_valid(self, activity: str) -> None:
 
 @dataclass
 class _RunningActivity:
-    pending_heartbeats: asyncio.Queue[Sequence[Any]]
     # Most of these optional values are set before use
     info: temporalio.activity.Info | None = None
     task: asyncio.Task | None = None
     cancelled_event: temporalio.common._CompositeEvent | None = None
     last_heartbeat_task: asyncio.Task | None = None
+    pending_heartbeat: Sequence[Any] | None = None
+    heartbeat_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
     cancel_thread_raiser: _ThreadExceptionRaiser | None = None
     sync: bool = False
     done: bool = False

tests/worker/test_activity.py

@@ -879,6 +879,56 @@ async def some_activity() -> str:
     assert list(timeout.last_heartbeat_details) == ["some details!"]
 
 
+class _SlowFirstHeartbeatCodec(PayloadCodec):
+    """Encodes the first heartbeat slower than later ones, like a codec waiting on a remote KMS."""
+
+    async def encode(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        if any(p.data == b'"first"' for p in payloads):
+            await asyncio.sleep(0.5)
+        return list(payloads)
+
+    async def decode(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        return list(payloads)
+
+
+async def test_activity_heartbeat_details_keep_order_with_slow_codec(
+    client: Client,
+    worker: ExternalWorker,
+    env: WorkflowEnvironment,
+    shared_state_manager: SharedStateManager,
+):
+    if env.supports_time_skipping:
+        pytest.skip("https://github.com/temporalio/sdk-java/issues/2459")
+
+    @activity.defn
+    async def some_activity() -> str:
+        info = activity.info()
+        if info.heartbeat_details:
+            return f"resumed from: {next(iter(info.heartbeat_details))}"
+        activity.heartbeat("first")
+        await asyncio.sleep(0.1)
+        activity.heartbeat("second")
+        await asyncio.sleep(1)
+        raise RuntimeError("Try again!")
+
+    config = client.config()
+    config["data_converter"] = dataclasses.replace(
+        client.data_converter, payload_codec=_SlowFirstHeartbeatCodec()
+    )
+    result = await _execute_workflow_with_activity(
+        Client(**config),
+        worker,
+        some_activity,
+        retry_max_attempts=2,
+        shared_state_manager=shared_state_manager,
+    )
+    assert result.result == "resumed from: second"
+
+
 @activity.defn
 def picklable_heartbeat_details_activity() -> str:
     info = activity.info()