Fix storage stats to report wall-clock duration instead of sum

temporalio/converter/_extstore.py

@@ -46,22 +46,22 @@ class StorageOperationMetrics:
     driver_names: set[str] = dataclasses.field(default_factory=set)
     """Names of the drivers that participated in the operations."""
 
-    def record_batch(
-        self, count: int, size: int, duration: timedelta, driver_names: set[str]
-    ) -> None:
+    def record_batch(self, count: int, size: int, 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.driver_names.update(driver_names)
 
     @contextlib.contextmanager
     def track(self) -> Generator[Self, None, None]:
-        """Set this instance as the current metrics context and reset on exit."""
+        """Set this instance as the current metrics context and measure
+        wall-clock duration of the enclosed block."""
         token = _current_storage_metrics.set(self)
+        start = time.monotonic()
         try:
             yield self
         finally:
+            self.total_duration = timedelta(seconds=time.monotonic() - start)
             _current_storage_metrics.reset(token)
 
 
@@ -357,8 +357,6 @@ def _with_store_context(self, ctx: StorageDriverStoreContext) -> ExternalStorage
         return result
 
     async def _store_payload(self, payload: Payload) -> Payload:
-        start_time = time.monotonic()
-
         driver = self._select_driver(self._store_context, payload)
         if driver is None:
             return payload
@@ -379,7 +377,7 @@ async def _store_payload(self, payload: Payload) -> Payload:
             )
         reference_payload.external_payloads.add().size_bytes = external_size
 
-        ExternalStorage._record_metrics(1, external_size, start_time, {driver.name()})
+        ExternalStorage._record_metrics(1, external_size, {driver.name()})
 
         return reference_payload
 
@@ -395,8 +393,6 @@ async def _store_payload_sequence(
         if len(payloads) == 1:
             return [await self._store_payload(payloads[0])]
 
-        start_time = time.monotonic()
-
         results = list(payloads)
 
         to_store: list[tuple[int, Payload, StorageDriver]] = []
@@ -448,9 +444,7 @@ async def _store_payload_sequence(
             external_count += len(claims)
             driver_names.add(driver.name())
 
-        ExternalStorage._record_metrics(
-            external_count, external_size, start_time, driver_names
-        )
+        ExternalStorage._record_metrics(external_count, external_size, driver_names)
 
         return results
 
@@ -480,7 +474,6 @@ async def _retrieve_payload(self, payload: Payload) -> Payload:
         if ref is None:
             return payload
 
-        start_time = time.monotonic()
         driver = self._get_driver_by_name(ref.driver_name)
         context = StorageDriverRetrieveContext()
         claim = StorageDriverClaim(claim_data=dict(ref.claim_data))
@@ -491,9 +484,7 @@ async def _retrieve_payload(self, payload: Payload) -> Payload:
 
         stored_payload = stored_payloads[0]
 
-        ExternalStorage._record_metrics(
-            1, stored_payload.ByteSize(), start_time, {driver.name()}
-        )
+        ExternalStorage._record_metrics(1, stored_payload.ByteSize(), {driver.name()})
 
         return stored_payload
 
@@ -509,8 +500,6 @@ async def _retrieve_payload_sequence(
         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]]] = {}
@@ -564,9 +553,7 @@ async def _retrieve_payload_sequence(
         for i, retrieved_payload in enumerate(stored_list):
             results[retrieve_indices[i]] = retrieved_payload
 
-        ExternalStorage._record_metrics(
-            external_count, external_size, start_time, driver_names
-        )
+        ExternalStorage._record_metrics(external_count, external_size, driver_names)
 
         return results
 
@@ -587,14 +574,7 @@ def _validate_payload_length(
             )
 
     @staticmethod
-    def _record_metrics(
-        count: int, size: int, start_time: float, driver_names: set[str]
-    ):
+    def _record_metrics(count: int, size: int, driver_names: set[str]):
         metrics = _current_storage_metrics.get()
         if metrics is not None:
-            metrics.record_batch(
-                count,
-                size,
-                timedelta(seconds=time.monotonic() - start_time),
-                driver_names,
-            )
+            metrics.record_batch(count, size, driver_names)

tests/test_extstore.py

@@ -1,7 +1,9 @@
 """Tests for external storage functionality."""
 
 import asyncio
+import time
 from collections.abc import Sequence
+from datetime import timedelta
 
 import pytest
 
@@ -17,7 +19,11 @@
     StorageDriverRetrieveContext,
     StorageDriverStoreContext,
 )
-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
 
@@ -801,5 +807,44 @@ async def test_new_format_encode_round_trips(self):
         assert decoded[0] == value
 
 
+class TestStorageOperationMetrics:
+    def test_track_records_wall_clock_duration(self):
+        """total_duration should reflect the wall-clock span of the track()
+        context, not the sum of individual record_batch calls."""
+        metrics = StorageOperationMetrics()
+        with metrics.track():
+            metrics.record_batch(2, 100, {"driver-a"})
+            metrics.record_batch(3, 200, {"driver-b"})
+
+        assert metrics.payload_count == 5
+        assert metrics.total_size == 300
+        assert metrics.driver_names == {"driver-a", "driver-b"}
+
+    @pytest.mark.asyncio
+    async def test_concurrent_operations_report_wall_clock(self):
+        """When operations run concurrently, total_duration should reflect
+        wall-clock time (~sleep_time), not the sum of all operations
+        (~sleep_time * concurrency)."""
+        metrics = StorageOperationMetrics()
+
+        async def simulate_driver(name: str, size: int):
+            await asyncio.sleep(0.05)
+            metrics.record_batch(1, size, {name})
+
+        start = time.monotonic()
+        with metrics.track():
+            await asyncio.gather(
+                simulate_driver("a", 100),
+                simulate_driver("b", 200),
+                simulate_driver("c", 300),
+            )
+        elapsed = timedelta(seconds=time.monotonic() - start)
+
+        assert metrics.payload_count == 3
+        assert metrics.total_size == 600
+        assert metrics.driver_names == {"a", "b", "c"}
+        assert metrics.total_duration <= elapsed
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])