From daf565f11c4d17dcd70bc98d7630e0daff274145 Mon Sep 17 00:00:00 2001
From: Dan Albert <dan@gingerhq.net>
Date: Sat, 10 Sep 2022 15:58:27 -0700
Subject: [PATCH] Add additional information to multi-event tracing.

Show the number of times the event was logged as well as an average run-
time.
---
 game/profiling.py | 71 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 60 insertions(+), 11 deletions(-)

diff --git a/game/profiling.py b/game/profiling.py
index 219453d9..a0f3f169 100644
--- a/game/profiling.py
+++ b/game/profiling.py
@@ -4,6 +4,7 @@ import logging
 import timeit
 from collections import defaultdict
 from contextlib import contextmanager
+from dataclasses import dataclass
 from datetime import timedelta
 from types import TracebackType
 from typing import Iterator, Optional, Type
@@ -11,15 +12,29 @@ from typing import Iterator, Optional, Type
 
 @contextmanager
 def logged_duration(event: str) -> Iterator[None]:
-    start = timeit.default_timer()
-    yield
-    end = timeit.default_timer()
-    logging.debug("%s took %s", event, timedelta(seconds=end - start))
+    timer = Timer()
+    with timer:
+        yield
+    logging.debug("%s took %s", event, timer.duration)
+
+
+@dataclass
+class CountedEvent:
+    count: int = 0
+    duration: timedelta = timedelta()
+
+    def increment(self, duration: timedelta) -> None:
+        self.count += 1
+        self.duration += duration
+
+    @property
+    def average(self) -> timedelta:
+        return self.duration / self.count
 
 
 class MultiEventTracer:
     def __init__(self) -> None:
-        self.events: dict[str, timedelta] = defaultdict(timedelta)
+        self.events: dict[str, CountedEvent] = defaultdict(CountedEvent)
 
     def __enter__(self) -> MultiEventTracer:
         return self
@@ -30,12 +45,46 @@ class MultiEventTracer:
         exc_val: Optional[BaseException],
         exc_tb: Optional[TracebackType],
     ) -> None:
-        for event, duration in self.events.items():
-            logging.debug("%s took %s", event, duration)
+        for event, counter in self.events.items():
+            logging.debug(
+                "%s %d times took %s (%s average)",
+                event,
+                counter.count,
+                counter.duration,
+                counter.average,
+            )
 
     @contextmanager
     def trace(self, event: str) -> Iterator[None]:
-        start = timeit.default_timer()
-        yield
-        end = timeit.default_timer()
-        self.events[event] += timedelta(seconds=end - start)
+        timer = Timer()
+        with timer:
+            yield
+        self.events[event].increment(timer.duration)
+
+
+class Timer:
+    def __init__(self) -> None:
+        self._start_time: float | None = None
+        self._end_time: float | None = None
+        self._duration: timedelta | None = None
+
+    @property
+    def duration(self) -> timedelta:
+        if self._duration is None:
+            raise RuntimeError(
+                "Cannot query the duration of the timer before it has elapsed"
+            )
+        return self._duration
+
+    def __enter__(self) -> None:
+        self._start_time = timeit.default_timer()
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        assert self._start_time is not None
+        self._end_time = timeit.default_timer()
+        self._duration = timedelta(seconds=self._end_time - self._start_time)