POC Tracing support

tests/test_dns.py

@@ -1,10 +1,14 @@
 import asyncio
 import socket
+import sys
 import unittest
 
 from uvloop import _testbase as tb
 
 
+PY37 = sys.version_info >= (3, 7, 0)
+
+
 class BaseTestDNS:
 
     def _test_getaddrinfo(self, *args, **kwargs):
@@ -177,6 +181,57 @@ async def run():
         finally:
             self.loop.close()
 
+    @unittest.skipUnless(PY37, 'requires Python 3.7')
+    def test_getaddrinfo_tracing(self):
+        from time import monotonic
+        from uvloop import start_tracing, stop_tracing
+        from uvloop.tracing import Tracer, Span
+
+        class DummySpan(Span):
+            def __init__(self, name, parent=None):
+                self.name = name
+                self.parent = parent
+                self.start_time = monotonic()
+                self.finish_time = None
+                self.children = []
+                self.tags = {}
+
+            def set_tag(self, key, value):
+                self.tags[key] = value
+
+            def finish(self, finish_time=None):
+                self.finish_time = finish_time or monotonic()
+
+            @property
+            def is_finished(self):
+                return self.finish_time is not None
+
+
+        class DummyTracer(Tracer):
+            def start_span(self, name, parent_span):
+                span = DummySpan(name, parent_span)
+                parent_span.children.append(span)
+                return span
+
+        root_span = DummySpan('root')
+        start_tracing(DummyTracer(), root_span)
+        self.loop.run_until_complete(
+            self.loop.getaddrinfo('example.com', 80)
+        )
+        root_span.finish()
+        assert root_span.children
+        assert root_span.children[0].name == 'getaddrinfo'
+        assert root_span.children[0].tags['host'] == b'example.com'
+        assert root_span.children[0].tags['port'] == b'80'
+        assert root_span.children[0].is_finished
+        assert root_span.children[0].start_time < root_span.children[0].finish_time
+
+        stop_tracing()
+        self.loop.run_until_complete(
+            self.loop.getaddrinfo('example.com', 80)
+        )
+        assert len(root_span.children) == 1
+
 
 class Test_AIO_DNS(BaseTestDNS, tb.AIOTestCase):
     pass

uvloop/__init__.py

@@ -1,15 +1,21 @@
 import asyncio
 
 from asyncio.events import BaseDefaultEventLoopPolicy as __BasePolicy
+from sys import version_info
 
 from . import includes as __includes  # NOQA
 from . import _patch  # NOQA
 from .loop import Loop as __BaseLoop  # NOQA
 
+PY37 = version_info >= (3, 7, 0)
 
-__version__ = '0.11.0.dev0'
-__all__ = ('new_event_loop', 'EventLoopPolicy')
+if PY37:
+    from .loop import start_tracing, stop_tracing
+    __all__ = ('new_event_loop', 'EventLoopPolicy', 'start_tracing', 'stop_tracing')
+else:
+    __all__ = ('new_event_loop', 'EventLoopPolicy')
 
+__version__ = '0.11.0.dev0'
 
 class Loop(__BaseLoop, asyncio.AbstractEventLoop):
     pass

uvloop/dns.pyx

@@ -249,6 +249,7 @@ cdef class AddrInfoRequest(UVRequest):
     cdef:
         system.addrinfo hints
         object callback
+        object context
         uv.uv_getaddrinfo_t _req_data
 
     def __cinit__(self, Loop loop,
@@ -278,6 +279,11 @@ cdef class AddrInfoRequest(UVRequest):
             callback(ex)
             return
 
+        if PY37:
+            self.context = <object>PyContext_CopyCurrent()
+        else:
+            self.context = None
+
         memset(&self.hints, 0, sizeof(system.addrinfo))
         self.hints.ai_flags = flags
         self.hints.ai_family = family
@@ -336,6 +342,9 @@ cdef void __on_addrinfo_resolved(uv.uv_getaddrinfo_t *resolver,
         object callback = request.callback
         AddrInfo ai
 
+    if PY37:
+        PyContext_Enter(<PyContext*>request.context)
+
     try:
         if status < 0:
             callback(convert_error(status))
@@ -347,6 +356,8 @@ cdef void __on_addrinfo_resolved(uv.uv_getaddrinfo_t *resolver,
         loop._handle_exception(ex)
     finally:
         request.on_done()
+        if PY37:
+            PyContext_Exit(<PyContext*>request.context)
 
 
 cdef void __on_nameinfo_resolved(uv.uv_getnameinfo_t* req,

uvloop/includes/compat.h

@@ -35,6 +35,10 @@ typedef struct {
     PyObject_HEAD
 } PyContext;
 
+typedef struct {
+    PyObject_HEAD
+} PyContextVar;
+
 PyContext * PyContext_CopyCurrent(void) {
     abort();
     return NULL;
@@ -49,4 +53,8 @@ int PyContext_Exit(PyContext *ctx) {
     abort();
     return -1;
 }
+
+int PyContextVar_Get(PyContextVar *var, PyObject *default_value, PyObject **value) {
+    return -1;
+}
 #endif

uvloop/includes/python.pxd

@@ -17,6 +17,10 @@ cdef extern from "Python.h":
 
 cdef extern from "includes/compat.h":
     ctypedef struct PyContext
+    ctypedef struct PyContextVar
+    ctypedef struct PyObject
     PyContext* PyContext_CopyCurrent() except NULL
     int PyContext_Enter(PyContext *) except -1
     int PyContext_Exit(PyContext *) except -1
+    int PyContextVar_Get(
+        PyContextVar *var, object default_value, PyObject **value) except -1

uvloop/loop.pxd

@@ -230,3 +230,5 @@ include "request.pxd"
 include "handles/udp.pxd"
 
 include "server.pxd"
+
+include "tracing.pxd"

uvloop/loop.pyx

@@ -18,7 +18,9 @@ from .includes.python cimport PY_VERSION_HEX, \
                               PyContext, \
                               PyContext_CopyCurrent, \
                               PyContext_Enter, \
-                              PyContext_Exit
+                              PyContext_Exit, \
+			      PyContextVar, \
+			      PyContextVar_Get
 
 from libc.stdint cimport uint64_t
 from libc.string cimport memset, strerror, memcpy
@@ -821,13 +823,26 @@ cdef class Loop:
                 except Exception as ex:
                     if not fut.cancelled():
                         fut.set_exception(ex)
+
                 else:
                     if not fut.cancelled():
                         fut.set_result(data)
+
             else:
                 if not fut.cancelled():
                     fut.set_exception(result)
 
+            traced_context = __traced_context()
+            if traced_context:
+                traced_context.current_span().finish()
+
+        traced_context = __traced_context()
+        if traced_context:
+            traced_context.start_span(
+                "getaddrinfo",
+                tags={'host': host, 'port': port}
+            )
+
         AddrInfoRequest(self, host, port, family, type, proto, flags, callback)
         return fut
 
@@ -2976,6 +2991,8 @@ include "handles/udp.pyx"
 
 include "server.pyx"
 
+include "tracing.pyx"
+
 
 # Used in UVProcess
 cdef vint __atfork_installed = 0

uvloop/tracing.pxd

@@ -0,0 +1,10 @@
+cdef class TracedContext:
+    cdef:
+        object _tracer
+        object _span
+        object _root_span
+
+    cdef object start_span(self, name, tags=?)
+    cdef object current_span(self)
+
+cdef TracedContext __traced_context()

uvloop/tracing.py

@@ -0,0 +1,24 @@
+import abc
+
+class Span(abc.ABC):
+
+    @abc.abstractmethod
+    def set_tag(self, key, value):
+        """Tag the span with an arbitrary key and value."""
+
+    @abc.abstractmethod
+    def finish(self, finish_time=None):
+        """Indicate that the work represented by this span
+        has been completed or terminated."""
+
+    @abc.abstractproperty
+    def is_finished(self):
+        """Return True if the current span is already finished."""
+
+
+class Tracer(abc.ABC):
+
+    @abc.abstractmethod
+    def start_span(self, name, parent_span):
+        """Start a new Span with a specific name. The parent of the span
+        will be also passed as a paramter."""

uvloop/tracing.pyx

@@ -0,0 +1,71 @@
+from contextlib import contextmanager
+
+if PY37:
+    import contextvars
+    __traced_ctx = contextvars.ContextVar('__traced_ctx', default=None)
+else:
+    __traced_ctx = None
+
+
+cdef class TracedContext:
+    def __cinit__(self, tracer, root_span):
+        self._tracer = tracer
+        self._root_span = root_span
+        self._span = None
+
+    cdef object start_span(self, name, tags=None):
+        parent_span = self._span if self._span else self._root_span
+        span = self._tracer.start_span(name, parent_span)
+
+        if tags:
+            for key, value in tags.items():
+                span.set_tag(key, value)
+
+        self._span = span
+        return self._span
+
+    cdef object current_span(self):
+        return self._span
+
+
+cdef inline TracedContext __traced_context():
+    cdef:
+        PyObject* traced_context = NULL
+
+    if not PY37:
+        return
+
+    PyContextVar_Get(<PyContextVar*> __traced_ctx, None, &traced_context)
+
+    if <object>traced_context is None:
+        return
+    return <TracedContext>traced_context
+
+
+def start_tracing(tracer, root_span):
+    if not PY37:
+        raise RuntimeError(
+            "tracing only supported by Python 3.7 or newer versions")
+
+    traced_context = __traced_ctx.get(None)
+    if traced_context is not None:
+        raise RuntimeError("Tracing already started")
+
+    traced_context = TracedContext(tracer, root_span)
+    __traced_ctx.set(traced_context)
+
+
+def stop_tracing():
+    if not PY37:
+        raise RuntimeError(
+            "tracing only supported by Python 3.7 or newer versions")
+
+    traced_context = __traced_context()
+    if traced_context is None:
+        return
+
+    span = traced_context.current_span()
+    if span and not span.is_finished:
+        span.finish()
+
+    __traced_ctx.set(None)