Emit workflow task duration information via logging (#1386)

Author: jmaeagle99

temporalio/bridge/worker.py

@@ -20,6 +20,7 @@
 import temporalio.bridge.runtime
 import temporalio.bridge.temporal_sdk_bridge
 import temporalio.converter
+import temporalio.converter._extstore
 from temporalio.api.common.v1.message_pb2 import Payload
 from temporalio.bridge._visitor import VisitorFunctions
 from temporalio.bridge.temporal_sdk_bridge import (
@@ -302,19 +303,33 @@ async def decode_activation(
     activation: temporalio.bridge.proto.workflow_activation.WorkflowActivation,
     data_converter: temporalio.converter.DataConverter,
     decode_headers: bool,
-) -> None:
-    """Decode all payloads in the activation."""
-    await CommandAwarePayloadVisitor(
-        skip_search_attributes=True, skip_headers=not decode_headers
-    ).visit(_Visitor(data_converter._decode_payload_sequence), activation)
+) -> temporalio.converter._extstore.StorageOperationMetrics:
+    """Decode all payloads in the activation.
+
+    Returns:
+        Metrics from any external storage retrieval operations that occurred.
+    """
+    metrics = temporalio.converter._extstore.StorageOperationMetrics()
+    with metrics.track():
+        await CommandAwarePayloadVisitor(
+            skip_search_attributes=True, skip_headers=not decode_headers
+        ).visit(_Visitor(data_converter._decode_payload_sequence), activation)
+    return metrics
 
 
 async def encode_completion(
     completion: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion,
     data_converter: temporalio.converter.DataConverter,
     encode_headers: bool,
-) -> None:
-    """Encode all payloads in the completion."""
-    await CommandAwarePayloadVisitor(
-        skip_search_attributes=True, skip_headers=not encode_headers
-    ).visit(_Visitor(data_converter._encode_payload_sequence), completion)
+) -> temporalio.converter._extstore.StorageOperationMetrics:
+    """Encode all payloads in the completion.
+
+    Returns:
+        Metrics from any external storage store operations that occurred.
+    """
+    metrics = temporalio.converter._extstore.StorageOperationMetrics()
+    with metrics.track():
+        await CommandAwarePayloadVisitor(
+            skip_search_attributes=True, skip_headers=not encode_headers
+        ).visit(_Visitor(data_converter._encode_payload_sequence), completion)
+    return metrics

temporalio/converter/_extstore.py

@@ -5,10 +5,14 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
+import contextvars
 import dataclasses
+import time
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Coroutine, Mapping, Sequence
+from collections.abc import Callable, Coroutine, Generator, Mapping, Sequence
 from dataclasses import dataclass
+from datetime import timedelta
 from typing import Any, ClassVar, TypeVar
 
 from typing_extensions import Self
@@ -25,6 +29,40 @@
 _REFERENCE_ENCODING = b"json/external-storage-reference"
 
 
+@dataclass
+class StorageOperationMetrics:
+    """Accumulates metrics from external storage operations."""
+
+    payload_count: int = 0
+    """Number of payloads stored or retrieved externally."""
+
+    total_size: int = 0
+    """Total size in bytes of externally stored/retrieved payloads."""
+
+    total_duration: timedelta = dataclasses.field(default_factory=timedelta)
+    """Wall-clock time spent on external storage operations."""
+
+    def record_batch(self, count: int, size: int, duration: timedelta) -> None:
+        """Record metrics from a batch of storage operations."""
+        self.payload_count += count
+        self.total_size += size
+        self.total_duration += duration
+
+    @contextlib.contextmanager
+    def track(self) -> Generator[Self, None, None]:
+        """Set this instance as the current metrics context and reset on exit."""
+        token = _current_storage_metrics.set(self)
+        try:
+            yield self
+        finally:
+            _current_storage_metrics.reset(token)
+
+
+_current_storage_metrics: contextvars.ContextVar[StorageOperationMetrics | None] = (
+    contextvars.ContextVar("_current_storage_metrics", default=None)
+)
+
+
 async def _gather_cancel_on_error(
     coros: Sequence[Coroutine[Any, Any, _T]],
 ) -> list[_T]:
@@ -255,6 +293,7 @@ def _get_driver_by_name(self, name: str) -> StorageDriver:
         return driver
 
     async def _store_payload(self, payload: Payload) -> Payload:
+        start_time = time.monotonic()
         context = StorageDriverStoreContext(serialization_context=self._context)
 
         driver = self._select_driver(context, payload)
@@ -265,6 +304,7 @@ async def _store_payload(self, payload: Payload) -> Payload:
 
         self._validate_claim_length(claims, expected=1, driver=driver)
 
+        external_size = payload.ByteSize()
         reference = _StorageReference(
             driver_name=driver.name(),
             driver_claim=claims[0],
@@ -274,7 +314,10 @@ async def _store_payload(self, payload: Payload) -> Payload:
             raise ValueError(
                 f"Failed to serialize storage reference for driver '{driver.name()}'"
             )
-        reference_payload.external_payloads.add().size_bytes = payload.ByteSize()
+        reference_payload.external_payloads.add().size_bytes = external_size
+
+        ExternalStorage._record_metrics(1, external_size, start_time)
+
         return reference_payload
 
     async def _store_payloads(self, payloads: Payloads):
@@ -289,6 +332,8 @@ async def _store_payload_sequence(
         if len(payloads) == 1:
             return [await self._store_payload(payloads[0])]
 
+        start_time = time.monotonic()
+
         results = list(payloads)
         context = StorageDriverStoreContext(serialization_context=self._context)
 
@@ -315,6 +360,8 @@ async def _store_payload_sequence(
             ]
         )
 
+        external_count = 0
+        external_size = 0
         for (driver, indexed_payloads), claims in zip(driver_group_list, all_claims):
             indices = [idx for idx, _ in indexed_payloads]
             sizes = [p.ByteSize() for _, p in indexed_payloads]
@@ -333,13 +380,20 @@ async def _store_payload_sequence(
                     )
                 reference_payload.external_payloads.add().size_bytes = sizes[i]
                 results[indices[i]] = reference_payload
+                external_size += sizes[i]
+
+            external_count += len(claims)
+
+        ExternalStorage._record_metrics(external_count, external_size, start_time)
 
         return results
 
     async def _retrieve_payload(self, payload: Payload) -> Payload:
         if len(payload.external_payloads) == 0:
             return payload
 
+        start_time = time.monotonic()
+
         reference = self._claim_converter.from_payload(payload, _StorageReference)
         if not isinstance(reference, _StorageReference):
             return payload
@@ -351,7 +405,11 @@ async def _retrieve_payload(self, payload: Payload) -> Payload:
 
         self._validate_payload_length(stored_payloads, expected=1, driver=driver)
 
-        return stored_payloads[0]
+        stored_payload = stored_payloads[0]
+
+        ExternalStorage._record_metrics(1, stored_payload.ByteSize(), start_time)
+
+        return stored_payload
 
     async def _retrieve_payloads(self, payloads: Payloads):
         stored_payloads = await self._retrieve_payload_sequence(payloads.payloads)
@@ -362,11 +420,13 @@ async def _retrieve_payload_sequence(
         self,
         payloads: Sequence[Payload],
     ) -> list[Payload]:
-        results = list(payloads)
-
         if len(payloads) == 1:
             return [await self._retrieve_payload(payloads[0])]
 
+        start_time = time.monotonic()
+
+        results = list(payloads)
+
         driver_claims: dict[StorageDriver, list[tuple[int, StorageDriverClaim]]] = {}
         for index, payload in enumerate(payloads):
             if len(payload.external_payloads) == 0:
@@ -394,6 +454,8 @@ async def _retrieve_payload_sequence(
             ]
         )
 
+        external_count = 0
+        external_size = 0
         for (driver, indexed_claims), stored_payloads in zip(
             driver_claim_list, all_stored
         ):
@@ -407,13 +469,18 @@ async def _retrieve_payload_sequence(
 
             for idx, stored_payload in zip(indices, stored_payloads):
                 stored_by_index[idx] = stored_payload
+                external_size += stored_payload.ByteSize()
+
+            external_count += len(stored_payloads)
 
         retrieve_indices = sorted(stored_by_index.keys())
         stored_list = [stored_by_index[idx] for idx in retrieve_indices]
 
         for i, retrieved_payload in enumerate(stored_list):
             results[retrieve_indices[i]] = retrieved_payload
 
+        ExternalStorage._record_metrics(external_count, external_size, start_time)
+
         return results
 
     def _validate_claim_length(
@@ -431,3 +498,11 @@ def _validate_payload_length(
             raise ValueError(
                 f"Driver '{driver.name()}' returned {len(payloads)} payloads, expected {expected}",
             )
+
+    @staticmethod
+    def _record_metrics(count: int, size: int, start_time: float):
+        metrics = _current_storage_metrics.get()
+        if metrics is not None:
+            metrics.record_batch(
+                count, size, timedelta(seconds=time.monotonic() - start_time)
+            )

temporalio/worker/_workflow.py

@@ -9,9 +9,10 @@
 import os
 import sys
 import threading
+import time
 from collections.abc import Awaitable, Callable, MutableMapping, Sequence
 from dataclasses import dataclass
-from datetime import timezone
+from datetime import timedelta, timezone
 from types import TracebackType
 
 import temporalio.api.common.v1
@@ -21,6 +22,7 @@
 import temporalio.bridge.worker
 import temporalio.common
 import temporalio.converter
+import temporalio.converter._extstore
 import temporalio.converter._payload_limits
 import temporalio.exceptions
 import temporalio.workflow
@@ -255,6 +257,8 @@ async def _handle_activation(
         completion.successful.SetInParent()
         workflow = None
         data_converter = self._data_converter
+        task_start_time = time.monotonic()
+        download_metrics = temporalio.converter._extstore.StorageOperationMetrics()
         try:
             if LOG_PROTOS:
                 logger.debug("Received workflow activation:\n%s", act)
@@ -291,7 +295,7 @@ async def _handle_activation(
                             workflow_context=workflow_context,
                         ),
                     )
-            await temporalio.bridge.worker.decode_activation(
+            download_metrics = await temporalio.bridge.worker.decode_activation(
                 act,
                 data_converter,
                 decode_headers=self._encode_headers,
@@ -399,9 +403,10 @@ async def _handle_activation(
                 ),
             )
 
+        upload_metrics = temporalio.converter._extstore.StorageOperationMetrics()
         try:
             try:
-                await temporalio.bridge.worker.encode_completion(
+                upload_metrics = await temporalio.bridge.worker.encode_completion(
                     completion,
                     data_converter,
                     encode_headers=self._encode_headers,
@@ -429,6 +434,71 @@ async def _handle_activation(
                 "Failed completing activation on workflow with run ID %s", act.run_id
             )
 
+        # Log workflow task duration with external storage metrics
+        self._log_workflow_task_duration(
+            act, task_start_time, download_metrics, upload_metrics
+        )
+
+    @staticmethod
+    def _log_workflow_task_duration(
+        act: temporalio.bridge.proto.workflow_activation.WorkflowActivation,
+        task_start_time: float,
+        download_metrics: temporalio.converter._extstore.StorageOperationMetrics,
+        upload_metrics: temporalio.converter._extstore.StorageOperationMetrics,
+    ) -> None:
+        task_duration = timedelta(seconds=time.monotonic() - task_start_time)
+
+        def _fmt_duration(td: timedelta) -> str:
+            secs = td.total_seconds()
+            if secs >= 1:
+                return f"{secs:.3f}s"
+            return f"{secs * 1000:.3f}ms"
+
+        msg_details: dict[str, object] = {
+            "event_id": act.history_length,
+            "workflow_task_duration": _fmt_duration(task_duration),
+        }
+        extra: dict[str, object] = {
+            "event_id": act.history_length,
+            "workflow_task_duration": task_duration,
+        }
+        if download_metrics.payload_count > 0:
+            msg_details["payload_download_count"] = download_metrics.payload_count
+            msg_details["payload_download_size"] = download_metrics.total_size
+            msg_details["payload_download_duration"] = _fmt_duration(
+                download_metrics.total_duration
+            )
+            extra["payload_download_count"] = download_metrics.payload_count
+            extra["payload_download_size"] = download_metrics.total_size
+            extra["payload_download_duration"] = download_metrics.total_duration
+        if upload_metrics.payload_count > 0:
+            msg_details["payload_upload_count"] = upload_metrics.payload_count
+            msg_details["payload_upload_size"] = upload_metrics.total_size
+            msg_details["payload_upload_duration"] = _fmt_duration(
+                upload_metrics.total_duration
+            )
+            extra["payload_upload_count"] = upload_metrics.payload_count
+            extra["payload_upload_size"] = upload_metrics.total_size
+            extra["payload_upload_duration"] = upload_metrics.total_duration
+        if task_duration.total_seconds() > 10:
+            logger.warning(
+                "[TMPRL1104] Workflow task exceeded 10 seconds (%s)",
+                msg_details,
+                extra=extra,
+            )
+        elif task_duration.total_seconds() > 5:
+            logger.info(
+                "[TMPRL1104] Workflow task exceeded 5 seconds (%s)",
+                msg_details,
+                extra=extra,
+            )
+        else:
+            logger.debug(
+                "[TMPRL1104] Workflow task duration information (%s)",
+                msg_details,
+                extra=extra,
+            )
+
     async def _handle_cache_eviction(
         self,
         act: temporalio.bridge.proto.workflow_activation.WorkflowActivation,

tests/helpers/__init__.py

@@ -422,12 +422,12 @@ def __init__(self) -> None:
         self.log_queue: queue.Queue[logging.LogRecord] = queue.Queue()
 
     @contextmanager
-    def logs_captured(self, *loggers: logging.Logger):
+    def logs_captured(self, *loggers: logging.Logger, level: int = logging.INFO):
         handler = logging.handlers.QueueHandler(self.log_queue)
 
         prev_levels = [l.level for l in loggers]
         for l in loggers:
-            l.setLevel(logging.INFO)
+            l.setLevel(level)
             l.addHandler(handler)
         try:
             yield self
@@ -447,6 +447,15 @@ def find(
                 return record
         return None
 
+    def find_all(
+        self, pred: Callable[[logging.LogRecord], bool]
+    ) -> list[logging.LogRecord]:
+        return [
+            record
+            for record in cast(list[logging.LogRecord], self.log_queue.queue)
+            if pred(record)
+        ]
+
 
 class LogHandler:
     @staticmethod