diff --git a/temporalio/converter/_extstore.py b/temporalio/converter/_extstore.py index a946b2c0f..1c6bfca5f 100644 --- a/temporalio/converter/_extstore.py +++ b/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) diff --git a/tests/test_extstore.py b/tests/test_extstore.py index 9a058c582..d6104cd05 100644 --- a/tests/test_extstore.py +++ b/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"])