Search code examples
pythonperformanceasynchronouspython-asynciometrics

How to measure time spent in blocking code while using asyncio in Python?


I'm currently migrating some Python code that used to be blocking to use asyncio with async/await. It is a lot of code to migrate at once so I would prefer to do it gradually and have metrics. With that thing in mind I want to create a decorator to wrap some functions and know how long they are blocking the event loop. For example:

def measure_blocking_code(f):
    def wrapper(*args, **kwargs):
        # ?????
        # It should measure JUST 1 second
        # not 5 which is what the whole async function takes
    return wrapper

@measure_blocking_code
async def my_function():
    my_blocking_function()  # Takes 1 seconds
    await my_async_function()  # Takes 2 seconds
    await my_async_function_2()  # Takes 2 seconds

I know the event loop has a debug function that already report this, but I need to get that information for specific functions.


Solution

  • TLDR;

    This decorator does the job:

    def measure_blocking_code(f):
        async def wrapper(*args, **kwargs):
            t = 0
            coro = f()
            try:
                while True:
                    t0 = time.perf_counter()
                    future = coro.send(None)
                    t1 = time.perf_counter()
                    t += t1 - t0
                    while not future.done():
                        await asyncio.sleep(0)
                    future.result()  # raises exceptions if any
            except StopIteration as e:
                print(f'Function took {t:.2e} sec')
                return e.value
        return wrapper
    

    Explanation

    This workaround exploits the conventions used in asyncio implementation in cPython. These conventions are a superset of PEP-492. In other words:

    1. You can generally use async/await without knowing these details.
    2. This might not work with other async libraries like trio.

    An asyncio coro object (coro) can be executed by calling .send() member. This will only run the blocking code, until an async call yields a Future object. By only measuring the time spent in .send(), the duration of the blocking code can be determined.