including child tasks in Trio measurement

benchmarks/overhead.py

@@ -21,7 +21,7 @@
 
 from perf_timer import (PerfTimer, ThreadPerfTimer, TrioPerfTimer,
                         AverageObserver, StdDevObserver, HistogramObserver,
-                        measure_overhead)
+                        measure_overhead, TrioHierarchyPerfTimer)
 from perf_timer._impl import _format_duration
 
 
@@ -39,7 +39,7 @@ async def main():
     print()
     print('compare types:')
     observer = default_observer
-    for timer_type in (PerfTimer, ThreadPerfTimer, TrioPerfTimer):
+    for timer_type in (PerfTimer, ThreadPerfTimer, TrioPerfTimer, TrioHierarchyPerfTimer):
         duration = measure_overhead(partial(timer_type, observer=observer))
         item = f'{timer_type.__name__}(observer={observer.__name__}):'
         print(f'    {item:45s}{_format(duration)}')

src/perf_timer/__init__.py

@@ -1,7 +1,8 @@
 from ._impl import (PerfTimer, ThreadPerfTimer, AverageObserver,
                     StdDevObserver, HistogramObserver, measure_overhead)
 try:
-    from ._trio import trio_perf_counter, TrioPerfTimer
+    from ._trio import (trio_perf_counter, trio_hierarchy_perf_counter,
+                        TrioPerfTimer, TrioHierarchyPerfTimer)
 except ImportError:
     pass
 from ._version import __version__

src/perf_timer/_trio.py

@@ -70,6 +70,57 @@ def trio_perf_counter():
     return perf_counter() - _instrument.get_elapsed_descheduled_time(task)
 
 
+class _HierarchyDescheduledTimeInstrument(trio.abc.Instrument):
+    """Trio instrument tracking elapsed descheduled time of given task and children
+
+    The implementation is similar to _DescheduledTimeInstrument, except on task
+    steps we walk up the task tree to find tracked tasks.  This probably has
+    significant overhead since it must be done even for untracked hierarchies.
+    """
+
+    def __init__(self, time_fn=perf_counter):
+        self._time_fn = time_fn
+        self._info_by_root_task = defaultdict(_TimeInfo)
+
+    def _parents_info(self, task: trio_lowlevel.Task, time_fn):
+        t = 0
+        while True:
+            info = self._info_by_root_task.get(task)
+            if info:
+                t = t or time_fn()
+                yield info, t
+            parent_nursery = task.parent_nursery
+            if parent_nursery:
+                task = parent_nursery.parent_task
+            else:
+                break
+
+    def task_exited(self, task: trio_lowlevel.Task):
+        if self._info_by_root_task.pop(task, None) and not self._info_by_root_task:
+            trio_lowlevel.remove_instrument(self)
+
+    def after_task_step(self, task: trio_lowlevel.Task):
+        # TODO: maintain global "tracked tasks" to provide a fast path?
+        for info, t in self._parents_info(task, self._time_fn):
+            info.deschedule_start = t
+
+    def before_task_step(self, task: trio_lowlevel.Task):
+        for info, t in self._parents_info(task, self._time_fn):
+            info.elapsed_descheduled += t - info.deschedule_start
+
+    def get_elapsed_descheduled_time(self, task: trio_lowlevel.Task):
+        return self._info_by_root_task[task].elapsed_descheduled
+
+
+_hierarchy_instrument = _HierarchyDescheduledTimeInstrument()
+
+
+def trio_hierarchy_perf_counter():
+    trio_lowlevel.add_instrument(_hierarchy_instrument)
+    task = trio_lowlevel.current_task()
+    return perf_counter() - _hierarchy_instrument.get_elapsed_descheduled_time(task)
+
+
 class TrioPerfTimer(PerfTimer):
     """Variant of PerfTimer which measures Trio task time
 
@@ -97,3 +148,10 @@ async def foo():
 
     def __init__(self, name, time_fn=trio_perf_counter, **kwargs):
         super().__init__(name, time_fn=time_fn, **kwargs)
+
+
+# TODO: The shorter trio_perf_counter and TrioPerfTimer names should include
+#  hierarchy, and non-hierarchy variants relegated to longer names.
+class TrioHierarchyPerfTimer(TrioPerfTimer):
+    def __init__(self, name, time_fn=trio_hierarchy_perf_counter, **kwargs):
+        super().__init__(name, time_fn=time_fn, **kwargs)

tests/test_trio.py

@@ -8,7 +8,14 @@
 except ImportError:
     import trio.hazmat as trio_lowlevel
 
-from perf_timer import trio_perf_counter, _trio, TrioPerfTimer, AverageObserver
+from perf_timer import (trio_perf_counter, trio_hierarchy_perf_counter, _trio,
+                        TrioPerfTimer, AverageObserver)
+
+
+async def _work(duration, count=1):
+    for _ in range(count):
+        time.sleep(duration)
+        await trio.sleep(0)
 
 
 async def test_descheduled_time_instrument():
@@ -62,16 +69,65 @@ async def _untracked_child():
 async def test_trio_perf_counter_time_sleep():
     # NOTE: subject to false pass due to reliance on wall time
     t0 = trio_perf_counter()
-    time.sleep(.01)
+    await _work(.1, 5)
     dt = trio_perf_counter() - t0
-    assert dt > .008
+    assert dt == pytest.approx(.5, rel=.15)
+
+
+async def test_trio_hierarchy_perf_counter():
+    t0 = trio_hierarchy_perf_counter()
+    async with trio.open_nursery() as nursery:
+        await _work(.05, 3)
+
+        @nursery.start_soon
+        async def _child():
+            async with trio.open_nursery() as nursery2:
+                @nursery2.start_soon
+                async def _child_child():
+                    await _work(.05, 5)
+
+                await _work(.05, 5)
+
+        @nursery.start_soon
+        async def _child_2():
+            await _work(.05, 10)
+
+    dt = trio_hierarchy_perf_counter() - t0
+    assert dt == pytest.approx(1.15, rel=.15)
 
 
-async def test_trio_perf_counter_unregister():
+async def test_trio_hierarchy_perf_counter_active_children():
+    # trio_hierarchy_perf_counter() should work even if the task already has children
+    async with trio.open_nursery() as nursery:
+        @nursery.start_soon
+        async def _child():
+            async with trio.open_nursery() as nursery2:
+                @nursery2.start_soon
+                async def _child_child():
+                    await _work(.1, 5)
+
+                await _work(.1, 5)
+
+        await trio.sleep(.5)
+        t0 = trio_hierarchy_perf_counter()
+
+        @nursery.start_soon
+        async def _child_2():
+            await _work(.1, 10)
+
+    dt = trio_hierarchy_perf_counter() - t0
+    assert dt == pytest.approx(1.5, rel=.15)
+
+
+@pytest.mark.parametrize('counter_fn, instrument', [
+    (trio_perf_counter, _trio._instrument),
+    (trio_hierarchy_perf_counter, _trio._hierarchy_instrument),
+])
+async def test_trio_perf_counter_unregister(counter_fn, instrument: trio.abc.Instrument):
     async def perf_counter_with_trio_sleep():
-        trio_perf_counter()
+        counter_fn()
         await trio.sleep(0)
-        trio_perf_counter()
+        counter_fn()
 
     async with trio.open_nursery() as nursery:
         nursery.start_soon(perf_counter_with_trio_sleep)
@@ -81,7 +137,7 @@ async def perf_counter_with_trio_sleep():
     # the Trio instrumentation to no longer be active (so remove call
     # will fail).
     with pytest.raises(KeyError):
-        trio_lowlevel.remove_instrument(_trio._instrument)
+        trio_lowlevel.remove_instrument(instrument)
 
 
 async def test_trio_perf_timer(autojump_clock):
@@ -116,3 +172,19 @@ async def foo():
     assert timer._sum == 15
     assert timer._max == 10
     del timer
+
+
+# TODO: test nested perf timers on the same task
+# async def _foo():
+#     with TrioPerfTimer('timer1'):
+#         async with trio.open_nursery() as nursery:
+#             @nursery.start_soon
+#             async def _child_1():
+#                 pass  # do work
+#
+#             # do work
+#
+#             with TrioPerfTimer('timer2'):
+#                 @nursery.start_soon
+#                 async def _child_2():
+#                     pass  # do work