diff --git a/benchmarks/overhead.py b/benchmarks/overhead.py index 01acc00..a287754 100644 --- a/benchmarks/overhead.py +++ b/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)}') diff --git a/src/perf_timer/__init__.py b/src/perf_timer/__init__.py index ffbf549..84e2ab6 100644 --- a/src/perf_timer/__init__.py +++ b/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__ diff --git a/src/perf_timer/_trio.py b/src/perf_timer/_trio.py index 328dbe2..2d32f70 100644 --- a/src/perf_timer/_trio.py +++ b/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) diff --git a/tests/test_trio.py b/tests/test_trio.py index b7840bb..00a6c86 100644 --- a/tests/test_trio.py +++ b/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