aioinsights.py

"""
This library can help you debug your asyncio application, to understand who
blocks event loop.

Install parameters:

  - ``total_threshold`` (seconds, 0.1 by default) - we analyze callbacks which
    took longer than this time to execute
  - ``func_threshold`` (seconds, 0.01 by default) - we log functions which took
    longer that this time to execute (total time)
  - ``with_metrics`` (bool, False by default) - enables Prometheus metrics
    support: ``loop_callback_time``

Example:

.. code-block:: python

    import asyncio
    import aioinsights

    aioinsights.install()
    asyncio.run(main())

Logged:

.. code-block:: text

    WARNING:aioinsights:200.7 ms spent in {built-in method time.sleep}, called by aioinsights.py:107::b

Note: doesn't work with uvloop.
"""
import time
import asyncio
import logging
import cProfile
from typing import ContextManager
from pstats import Stats, SortKey
from contextlib import nullcontext


_logger = logging.getLogger(__name__)


def _func_repr(func_key):
    module, line, func = func_key
    if module == '~':
        # special case for built-in functions
        if func.startswith('<') and func.endswith('>'):
            return f'{{{func[1:-1]}}}'
        else:
            return func
    else:
        return f'{module}:{line}::{func}'


class _HandleMixin:
    _total_threshold: float
    _func_threshold: float
    _metric_ctx: ContextManager

    def _run(self) -> None:
        started = time.monotonic()
        pr = cProfile.Profile()
        with self._metric_ctx:
            pr.enable()
            super()._run()
            pr.disable()
        stopped = time.monotonic()
        if stopped - started > self._total_threshold:
            st = Stats(pr).strip_dirs().sort_stats(SortKey.TIME)
            for fn in st.fcn_list:
                cc, nc, tt, ct, callers = st.stats[fn]
                if tt > self._func_threshold:
                    fn_repr = _func_repr(fn)
                    if callers:
                        caller = next(callers.keys().__iter__())
                        caller_repr = _func_repr(caller)
                        if len(callers) > 1:
                            caller_repr += '+'
                    else:
                        caller_repr = 'N/A'
                    _logger.warning(
                        '%.1f ms spent in %s, called by %s',
                        tt * 1000,
                        fn_repr,
                        caller_repr,
                    )


def install(*, total_threshold=0.1, func_threshold=0.01, with_metrics=False):
    from asyncio import events

    if not isinstance(events.Handle, _HandleMixin):
        if with_metrics:
            import prometheus_client

            metric = prometheus_client.Histogram(
                'loop_callback_time',
                'Event-loop callback execution time (seconds)',
            )
            metric_ctx = metric.time()
        else:
            metric_ctx = nullcontext()

        params = dict(
            _total_threshold=total_threshold,
            _func_threshold=func_threshold,
            _metric_ctx=metric_ctx,
        )
        events.Handle = type(
            'Handle', (_HandleMixin, events.Handle), params,
        )
        events.TimerHandle = type(
            'TimerHandle', (_HandleMixin, events.TimerHandle), params,
        )


if __name__ == '__main__':
    from prometheus_client import start_http_server

    install(with_metrics=True)
    logging.basicConfig()
    start_http_server(9000)

    async def a():
        async def b():
            await asyncio.sleep(0.01)
            time.sleep(.2)
            await asyncio.sleep(0.01)

        await asyncio.sleep(0.01)
        await b()
        await asyncio.sleep(0.01)

    asyncio.run(a())
    try:
        time.sleep(1000)
    except KeyboardInterrupt:
        print('')