From 8d411e36614af1fd8b0d48e2388ba6a750fe11bc Mon Sep 17 00:00:00 2001
From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com>
Date: Thu, 3 Sep 2026 17:10:05 -0700
Subject: [PATCH] fix(extstore): report wall-clock duration for storage metrics

---
 CHANGELOG.md                      |  3 +++
 temporalio/converter/_extstore.py | 36 +++++++++++++++++++++++++------
 tests/test_extstore.py            | 36 ++++++++++++++++++++++++++++++-
 3 files changed, 68 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d0d80bc71..f17761a36 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -40,6 +40,9 @@ to include examples, links to docs, or any other relevant information.
 
 ### Fixed
 
+- **Experimental**: External storage metrics now report the wall-clock time storage was in flight.
+  Previously each batch's duration was summed, over-reporting the time whenever storage operations
+  ran concurrently.
 - `StrandsPlugin` now disables Botocore retries for its default Bedrock model so
   model request retries are handled exclusively by Temporal.
 - `temporalio.contrib.openai_agents` now honors the `retry-after-ms` and
diff --git a/temporalio/converter/_extstore.py b/temporalio/converter/_extstore.py
index ad6fa2e4b..c300448b9 100644
--- a/temporalio/converter/_extstore.py
+++ b/temporalio/converter/_extstore.py
@@ -40,19 +40,26 @@ class StorageOperationMetrics:
     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."""
-
     driver_names: set[str] = dataclasses.field(default_factory=set)
     """Names of the drivers that participated in the operations."""
 
+    _spans: list[tuple[float, float]] = dataclasses.field(default_factory=list)
+    """Monotonic-clock start and end of each recorded batch."""
+
+    @property
+    def total_duration(self) -> timedelta:
+        """Wall-clock time spent on external storage operations."""
+        # Batches may run concurrently, so summing each batch's duration would
+        # double-count operations that overlapped.
+        return timedelta(seconds=_union_seconds(self._spans))
+
     def record_batch(
-        self, count: int, size: int, duration: timedelta, driver_names: set[str]
+        self, count: int, size: int, start: float, end: float, driver_names: set[str]
     ) -> None:
         """Record metrics from a batch of storage operations."""
         self.payload_count += count
         self.total_size += size
-        self.total_duration += duration
+        self._spans.append((start, end))
         self.driver_names.update(driver_names)
 
     @contextlib.contextmanager
@@ -70,6 +77,22 @@ def track(self) -> Generator[Self, None, None]:
 )
 
 
+def _union_seconds(spans: list[tuple[float, float]]) -> float:
+    """Total length of the union of the given monotonic-clock spans, in seconds."""
+    ordered = sorted(spans)
+    if not ordered:
+        return 0.0
+    total = 0.0
+    span_start, span_end = ordered[0]
+    for start, end in ordered[1:]:
+        if start > span_end:
+            total += span_end - span_start
+            span_start, span_end = start, end
+        elif end > span_end:
+            span_end = end
+    return total + (span_end - span_start)
+
+
 async def _gather_cancel_on_error(
     coros: Sequence[Coroutine[Any, Any, _T]],
 ) -> list[_T]:
@@ -624,6 +647,7 @@ def _record_metrics(
             metrics.record_batch(
                 count,
                 size,
-                timedelta(seconds=time.monotonic() - start_time),
+                start_time,
+                time.monotonic(),
                 driver_names,
             )
diff --git a/tests/test_extstore.py b/tests/test_extstore.py
index 4a52c65c6..943fc7be2 100644
--- a/tests/test_extstore.py
+++ b/tests/test_extstore.py
@@ -2,6 +2,7 @@
 
 import asyncio
 from collections.abc import Sequence
+from datetime import timedelta
 
 import pytest
 
@@ -19,7 +20,11 @@
     StorageDriverStoreContext,
     StorageDriverWorkflowInfo,
 )
-from temporalio.converter._extstore import _REFERENCE_ENCODING, _StorageReference
+from temporalio.converter._extstore import (
+    _REFERENCE_ENCODING,
+    StorageOperationMetrics,
+    _StorageReference,
+)
 from temporalio.converter._payload_converter import JSONProtoPayloadConverter
 from temporalio.exceptions import ApplicationError
 
@@ -834,5 +839,34 @@ async def test_new_format_encode_round_trips(self):
         assert decoded[0] == value
 
 
+def test_storage_metrics_aggregates_batches() -> None:
+    metrics = StorageOperationMetrics()
+    assert metrics.total_duration == timedelta(0)
+
+    metrics.record_batch(2, 1024, 0.0, 10.0, {"s3"})
+    metrics.record_batch(3, 2048, 5.0, 15.0, {"gcs"})
+
+    assert metrics.payload_count == 5
+    assert metrics.total_size == 3072
+    assert metrics.driver_names == {"gcs", "s3"}
+    # Concurrent batches: summing their durations would report 20 seconds.
+    assert metrics.total_duration == timedelta(seconds=15)
+
+
+def test_storage_metrics_duration_sums_disjoint_batches() -> None:
+    metrics = StorageOperationMetrics()
+    metrics.record_batch(1, 1, 0.0, 10.0, {"s3"})
+    metrics.record_batch(1, 1, 20.0, 30.0, {"s3"})
+    assert metrics.total_duration == timedelta(seconds=20)
+
+
+def test_storage_metrics_duration_merges_adjacent_and_nested_batches() -> None:
+    metrics = StorageOperationMetrics()
+    metrics.record_batch(1, 1, 0.0, 10.0, {"s3"})
+    metrics.record_batch(1, 1, 10.0, 20.0, {"s3"})
+    metrics.record_batch(1, 1, 12.0, 18.0, {"s3"})
+    assert metrics.total_duration == timedelta(seconds=20)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])