For review

api/src/opentrons/util/performance_helpers.py

@@ -1,5 +1,6 @@
 """Performance helpers for tracking robot context."""
 
+import asyncio
 import functools
 from pathlib import Path
 from opentrons_shared_data.performance.dev_types import (
@@ -19,14 +20,12 @@
 _should_track = ff.enable_performance_metrics(
     RobotTypeEnum.robot_literal_to_enum(robot_configs.load().model)
 )
-STORE_EACH = _should_track
-
 
 class StubbedTracker(SupportsTracking):
     """A stubbed tracker that does nothing."""
 
     def __init__(
-        self, storage_location: Path, should_track: bool, store_each: bool
+        self, storage_location: Path, should_track: bool
     ) -> None:
         """Initialize the stubbed tracker."""
         pass
@@ -67,7 +66,7 @@ def _get_robot_context_tracker() -> SupportsTracking:
     global _robot_context_tracker
     if _robot_context_tracker is None:
         _robot_context_tracker = package_to_use(
-            get_performance_metrics_data_dir(), _should_track, STORE_EACH
+            get_performance_metrics_data_dir(), _should_track
         )
     return _robot_context_tracker
 
@@ -79,12 +78,12 @@ def track_analysis(func: F) -> F:
 
     # Typing a decorator that wraps a decorator with args, nope
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):  # type: ignore # noqa: ANN002, ANN003, ANN201
-        tracker = _get_robot_context_tracker()
-        tracked_func = tracker.track(RobotContextState.ANALYZING_PROTOCOL)(func)
-
-        result = tracked_func(*args, **kwargs)
+    async def wrapper(*args, **kwargs):  # type: ignore # noqa: ANN002, ANN003, ANN201
+        tracker: SupportsTracking = _get_robot_context_tracker()
 
-        return result
+        try:
+            return await tracker.track(func, RobotContextState.ANALYZING_PROTOCOL, *args, **kwargs)
+        finally:
+            tracker.store()
 
     return wrapper  # type: ignore

api/tests/opentrons/cli/test_cli.py

@@ -39,12 +39,6 @@ def override_data_store(tmp_path: Path) -> Iterator[None]:
     context_tracker._store = old_store  # type: ignore[attr-defined]
 
 
-@pytest.fixture
-def monkeypatch_set_store_each_to_true(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Override the STORE_EACH flag for the RobotContextTracker."""
-    monkeypatch.setattr(context_tracker, "_store_each", True)
-
-
 def verify_metrics_store_file(file_path: Path, expected_length: int) -> None:
     """Verify that the metrics store file contains the expected number of lines."""
     with open(file_path, "r") as f:
@@ -290,9 +284,9 @@ def test_python_error_line_numbers(
     assert error["detail"] == expected_detail
 
 
-@pytest.mark.usefixtures("override_data_store", "monkeypatch_set_store_each_to_true")
-def test_track_analysis_with_store_each_run(tmp_path: Path) -> None:
-    """Test that the RobotContextTracker tracks analysis and stores each run."""
+@pytest.mark.usefixtures("override_data_store")
+def test_track_analysis_with_store(tmp_path: Path) -> None:
+    """Test that the RobotContextTracker tracks analysis and stores each execution."""
     protocol_source = textwrap.dedent(
         """
         requirements = {"apiLevel": "2.15"}

api/tests/opentrons/util/test_performance_helpers.py

@@ -10,7 +10,7 @@
 
 def test_return_function_unchanged() -> None:
     """Test that the function is returned unchanged when using StubbedTracker."""
-    tracker = StubbedTracker(Path("/path/to/storage"), True, store_each=True)
+    tracker = StubbedTracker(Path("/path/to/storage"), True)
 
     def func_to_track() -> None:
         pass

performance-metrics/src/performance_metrics/robot_context_tracker.py

@@ -49,9 +49,7 @@ class RobotContextTracker(SupportsTracking):
 
     METADATA_NAME: Final[Literal["robot_context_data"]] = "robot_context_data"
 
-    def __init__(
-        self, storage_location: Path, should_track: bool, store_each: bool
-    ) -> None:
+    def __init__(self, storage_location: Path, should_track: bool) -> None:
         """Initializes the RobotContextTracker with an empty storage list."""
         self._store = MetricsStore[RawContextData](
             MetricsMetadata(
@@ -61,12 +59,11 @@ def __init__(
             )
         )
         self._should_track = should_track
-        self._store_each = store_each
 
         if self._should_track:
             self._store.setup()
 
-    def track(self, state: RobotContextState) -> Callable:  # type: ignore
+    async def track(self, func_to_track: Callable, state: RobotContextState, *args, **kwargs) -> Callable:  # type: ignore
         """Decorator factory for tracking the execution duration and state of robot operations.
 
         Args:
@@ -76,60 +73,35 @@ def track(self, state: RobotContextState) -> Callable:  # type: ignore
             Callable: A decorator that wraps a function to track its execution duration and state.
         """
 
-        def inner_decorator(func: Callable[P, R]) -> Callable[P, R]:
-            if not self._should_track:
-                return func
-
-            if inspect.iscoroutinefunction(func):
-
-                @wraps(func)
-                async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-                    function_start_time = timing_function()
-                    duration_start_time = perf_counter_ns()
-                    try:
-                        result = await func(*args, **kwargs)
-                    finally:
-                        duration_end_time = perf_counter_ns()
-
-                        self._store.add(
-                            RawContextData(
-                                func_start=function_start_time,
-                                duration=duration_end_time - duration_start_time,
-                                state=state,
-                            )
-                        )
-
-                        if self._store_each:
-                            self.store()
-
-                    return result  # type: ignore
-
-            else:
-
-                @wraps(func)
-                def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-                    function_start_time = timing_function()
-                    duration_start_time = perf_counter_ns()
-                    try:
-                        result = func(*args, **kwargs)
-                    finally:
-                        duration_end_time = perf_counter_ns()
-                        self._store.add(
-                            RawContextData(
-                                func_start=function_start_time,
-                                duration=duration_end_time - duration_start_time,
-                                state=state,
-                            )
-                        )
-
-                        if self._store_each:
-                            self.store()
-
-                    return result
-
-            return wrapper  # type: ignore
-
-        return inner_decorator
+
+
+        if not self._should_track:
+            return func_to_track
+
+        function_start_time = timing_function()
+        duration_start_time = perf_counter_ns()
+
+        if inspect.iscoroutinefunction(func_to_track):
+            try:
+                result = await func_to_track(*args, **kwargs)
+            finally:
+                 duration_end_time = perf_counter_ns()
+        else:
+            try:
+                result = func_to_track(*args, **kwargs)
+            finally:
+                duration_end_time = perf_counter_ns()
+
+        self._store.add(
+            RawContextData(
+                func_start=function_start_time,
+                duration=duration_end_time - duration_start_time,
+                state=state,
+            )
+        )
+        return result
+
+
 
     def store(self) -> None:
         """Returns the stored context data and clears the storage list."""

performance-metrics/tests/performance_metrics/test_metrics_store.py

@@ -17,9 +17,7 @@
 
 async def test_storing_to_file(tmp_path: Path) -> None:
     """Tests storing the tracked data to a file."""
-    robot_context_tracker = RobotContextTracker(
-        tmp_path, should_track=True, store_each=False
-    )
+    robot_context_tracker = RobotContextTracker(tmp_path, should_track=True)
 
     @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
     def starting_robot() -> None:

performance-metrics/tests/performance_metrics/test_robot_context_tracker.py

@@ -7,7 +7,6 @@
 from opentrons_shared_data.performance.dev_types import RobotContextState
 from time import sleep, time_ns
 from unittest.mock import patch
-from conftest import verify_metrics_store_file
 
 # Corrected times in seconds
 STARTING_TIME = 0.001
@@ -20,9 +19,7 @@
 @pytest.fixture
 def robot_context_tracker(tmp_path: Path) -> RobotContextTracker:
     """Fixture to provide a fresh instance of RobotContextTracker for each test."""
-    return RobotContextTracker(
-        storage_location=tmp_path, should_track=True, store_each=False
-    )
+    return RobotContextTracker(storage_location=tmp_path, should_track=True)
 
 
 def test_robot_context_tracker(robot_context_tracker: RobotContextTracker) -> None:
@@ -231,9 +228,7 @@ async def second_async_calibrating() -> None:
 
 def test_no_tracking(tmp_path: Path) -> None:
     """Tests that operations are not tracked when tracking is disabled."""
-    robot_context_tracker = RobotContextTracker(
-        tmp_path, should_track=False, store_each=False
-    )
+    robot_context_tracker = RobotContextTracker(tmp_path, should_track=False)
 
     @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
     def operation_without_tracking() -> None:
@@ -253,9 +248,7 @@ def operation_without_tracking() -> None:
 def test_using_non_linux_time_functions(tmp_path: Path) -> None:
     """Tests tracking operations using non-Linux time functions."""
     file_path = tmp_path / "test_file.csv"
-    robot_context_tracker = RobotContextTracker(
-        file_path, should_track=True, store_each=False
-    )
+    robot_context_tracker = RobotContextTracker(file_path, should_track=True)
 
     @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
     def starting_robot() -> None:
@@ -276,39 +269,3 @@ def calibrating_robot() -> None:
         data.duration > 0 for data in storage
     ), "All duration times should be greater than 0."
     assert len(storage) == 2, "Both operations should be tracked."
-
-
-def test_store_each(tmp_path: Path) -> None:
-    """Tests storing data after each operation when store_each is enabled."""
-    robot_context_tracker = RobotContextTracker(
-        tmp_path, should_track=True, store_each=True
-    )
-
-    @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
-    def starting_robot() -> None:
-        sleep(STARTING_TIME)
-
-    @robot_context_tracker.track(state=RobotContextState.CALIBRATING)
-    def calibrating_robot() -> None:
-        sleep(CALIBRATING_TIME)
-
-    starting_robot()
-
-    assert len(robot_context_tracker._store._data) == 0, (
-        "Operation should automatically stored when store_each is enabled."
-        "So _data should be cleared after every track"
-    )
-    verify_metrics_store_file(
-        robot_context_tracker._store.metadata.data_file_location, 1
-    )
-
-    calibrating_robot()
-
-    assert len(robot_context_tracker._store._data) == 0, (
-        "Operation should automatically stored when store_each is enabled."
-        "So _data should be cleared after every track"
-    )
-
-    verify_metrics_store_file(
-        robot_context_tracker._store.metadata.data_file_location, 2
-    )

shared-data/python/opentrons_shared_data/performance/dev_types.py

@@ -11,7 +11,7 @@ class SupportsTracking(Protocol):
     """Protocol for classes that support tracking of robot context."""
 
     def __init__(
-        self, storage_location: Path, should_track: bool, store_each: bool
+        self, storage_location: Path, should_track: bool
     ) -> None:
         """Initialize the tracker."""
         ...