UVloop tracing test dns tracing

test_dns_tracing.py

import asyncio
import contextvars
import uvloop
from types import SimpleNamespace

asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())

trace_context = contextvars.ContextVar('trace_context')


class MyCollector(uvloop.TracingCollector):
    def __init__(self, loop=None):
        self._loop = loop if loop else asyncio.get_event_loop()
        uvloop.TracingCollector.__init__(self)

    @property
    def loop(self):
        return self._loop

    ## Collector methods
    def dns_request_begin(self, host, port, family, type, proto, flags):
        ctx = SimpleNamespace(dns_request_begin = self.loop.time())
        trace_context.set(ctx)

    def dns_request_end(self, data):
        delta = loop.time() - trace_context.get().dns_request_begin
        print(f"DNS request took {delta}")


async def main(loop):
    with uvloop.tracing(MyCollector(loop)):
        for i in range(100):
            host = await loop.getaddrinfo('google.com', 80)

loop = asyncio.get_event_loop()
loop.run_until_complete(main(loop))