perf counter for Trio tasks

task_perf_counter.py

from collections import defaultdict
from time import perf_counter

import trio
from attr import attrs, attr


@attrs(slots=True)
class TimeInfo:
    deschedule_start = attr(type=float, default=0)
    elapsed_descheduled = attr(type=float, default=0)


class DescheduledTimeInstrument(trio.abc.Instrument):
    """Track elapsed descheduled time of selected tasks."""

    def __init__(self, time_fn=perf_counter):
        self._time_fn = time_fn
        self._info_by_task = defaultdict(TimeInfo)

    def after_task_step(self, task):
        info = self._info_by_task.get(task)
        if info:
            info.deschedule_start = self._time_fn()

    def before_task_step(self, task):
        info = self._info_by_task.get(task)
        if info:
            info.elapsed_descheduled += self._time_fn() - info.deschedule_start

    def task_exited(self, task):
        # unregister instrument if there are no more traced tasks
        if self._info_by_task.pop(task, None) and len(self._info_by_task) == 0:
            trio.hazmat.remove_instrument(self)

    def get_elapsed_descheduled_time(self, task):
        """
        Return elapsed descheduled time in seconds since the given task was
        first referenced by this method.  The initial reference always returns 0.
        """
        return self._info_by_task[task].elapsed_descheduled


_instrument = DescheduledTimeInstrument()


def task_perf_counter():
    """
    For the current Trio task, return the value (in fractional seconds) of a
    performance counter, i.e. a clock with the highest available resolution to
    measure a short duration.  It includes time elapsed during time.sleep,
    but not trio.sleep.  The reference point of the returned value is
    undefined, so that only the difference between the results of consecutive
    calls is valid.
    """
    trio.hazmat.add_instrument(_instrument)
    task = trio.hazmat.current_task()
    return perf_counter() - _instrument.get_elapsed_descheduled_time(task)


async def foo():
    s = ''
    for _ in range(1_000_000):
        s += 's'
    await trio.sleep(.5)
    return s


async def main():
    t0 = task_perf_counter()
    print('hi')
    await trio.sleep(.5)
    await foo()
    # async with trio.open_nursery() as nursery:
    #     nursery.start_soon(foo)
    print('...')
    await trio.sleep(.5)
    print(task_perf_counter() - t0)


if __name__ == '__main__':
    trio.run(main)