aiohttp tracing

aiohttp_trace.py

import aiohttp


def request_tracer(results_collector):
    """
    Provides request tracing to aiohttp client sessions.
    :param results_collector: a dict to which the tracing results will be added.
    :return: an aiohttp.TraceConfig object.

    :example:

    >>> import asyncio
    >>> import aiohttp
    >>> from aiohttp_trace import request_tracer
    >>>
    >>>
    >>> async def func():
    >>>     trace = {}
    >>>     async with aiohttp.ClientSession(trace_configs=[request_tracer(trace)]) as client:
    >>>         async with client.get('https://github.com') as response:
    >>>             print(trace)
    >>>
    >>> asyncio.get_event_loop().run_until_complete(func())
    {'dns_lookup_and_dial': 43.3, 'connect': 334.29, 'transfer': 148.48, 'total': 526.08, 'is_redirect': False}
    """

    async def on_request_start(session, context, params):
        context.on_request_start = session.loop.time()
        context.is_redirect = False

    async def on_connection_create_start(session, context, params):
        since_start = session.loop.time() - context.on_request_start
        context.on_connection_create_start = since_start

    async def on_request_redirect(session, context, params):
        since_start = session.loop.time() - context.on_request_start
        context.on_request_redirect = since_start
        context.is_redirect = True

    async def on_dns_resolvehost_start(session, context, params):
        since_start = session.loop.time() - context.on_request_start
        context.on_dns_resolvehost_start = since_start

    async def on_dns_resolvehost_end(session, context, params):
        since_start = session.loop.time() - context.on_request_start
        context.on_dns_resolvehost_end = since_start

    async def on_connection_create_end(session, context, params):
        since_start = session.loop.time() - context.on_request_start
        context.on_connection_create_end = since_start
    
    async def on_request_chunk_sent(session, context, params):
        since_start = session.loop.time() - context.on_request_start
        context.on_request_chunk_sent = since_start

    async def on_request_end(session, context, params):
        total = session.loop.time() - context.on_request_start
        context.on_request_end = total

        dns_lookup_and_dial = context.on_dns_resolvehost_end - context.on_dns_resolvehost_start
        connect = context.on_connection_create_end - dns_lookup_and_dial
        transfer = total - context.on_connection_create_end
        is_redirect = context.is_redirect

        results_collector['dns_lookup_and_dial'] = round(dns_lookup_and_dial * 1000, 2)
        results_collector['connect'] = round(connect * 1000, 2)
        results_collector['transfer'] = round(transfer * 1000, 2)
        results_collector['total'] = round(total * 1000, 2)
        results_collector['is_redirect'] = is_redirect

    trace_config = aiohttp.TraceConfig()

    trace_config.on_request_start.append(on_request_start)
    trace_config.on_request_redirect.append(on_request_redirect)
    trace_config.on_dns_resolvehost_start.append(on_dns_resolvehost_start)
    trace_config.on_dns_resolvehost_end.append(on_dns_resolvehost_end)
    trace_config.on_connection_create_start.append(on_connection_create_start)
    trace_config.on_connection_create_end.append(on_connection_create_end)
    trace_config.on_request_end.append(on_request_end)
    trace_config.on_request_chunk_sent.append(on_request_chunk_sent)

    return trace_config

example.py

import aiohttp
import asyncio
from aiohttp.client_exceptions import ClientError
from aiohttp_trace import request_tracer


async def fetch(url):
    trace = {}
    async with aiohttp.ClientSession(trace_configs=[request_tracer(trace)]) as client:
        async with client.get(url) as response:
            print(trace)


urls = [
    'https://github.com'
]

loop = asyncio.get_event_loop()
loop.run_until_complete(asyncio.gather(*[fetch(url) for url in urls]))

@anthonynsimon, your request tracer is great! However, exceptions occur after the first iteration when used within a loop. This is because the aiohttp request relies on an existing connection, which means that the on_dns_xxx, on_connection_create_start, and on_connection_create_end methods are not called again. I solved this by adding try/except blocks like the following:

total = asyncio.get_event_loop().time() - context.on_request_start
context.on_request_end = total

try:
    # raise exception if context.on_dns_resolvehost_start doesn't exist
    dns_lookup_and_dial = context.on_dns_resolvehost_end - context.on_dns_resolvehost_start
    trace_data['dns_lookup_and_dial'] = round(dns_lookup_and_dial * 1000)

except AttributeError as e:
    dns_lookup_and_dial = 0

try:
    # raise exception if context.on_connection_create_end doesn't exist
    connect = context.on_connection_create_end - dns_lookup_and_dial
    trace_data['connect'] = round(connect * 1000)
    chunk_sent = context.on_request_chunk_sent - context.on_connection_create_end
    transfer = total - context.on_connection_create_end

except AttributeError as e:
    chunk_sent = context.on_request_chunk_sent
    transfer = total - dns_lookup_and_dial

trace_data['chunk_sent'] = round(chunk_sent * 1000)
trace_data['transfer'] = round(transfer * 1000)
trace_data['total'] = round(total * 1000)
trace_data['is_redirect'] = context.is_redirect

This works, but I always get something like this output in iterations after the first one:

trace: {'chunk_sent': 1, 'transfer': 201, 'total': 201, 'is_redirect': False, 'response_chunk_received': 0}

chunk_sent has a small value like 1 ms and transfer is equal to the total because dns_lookup_and_dial was zero in the iteration. So what part of the request took 200 ms? I only found an extra signal named on_request_headers_sent which you didn't implement in your code, but it also took a small time like 1 ms and doesn't have an effect on the request time.

I also suspect that the signal handlers are calling at the entry point of each aiohttp corresponding method or at their exit point. So should I change my exception blocks?