diff --git a/tests/test_dns.py b/tests/test_dns.py
index 3c19655c..4b333d99 100644
--- a/tests/test_dns.py
+++ b/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
diff --git a/uvloop/__init__.py b/uvloop/__init__.py
index a98cc5ac..e3d33d50 100644
--- a/uvloop/__init__.py
+++ b/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
diff --git a/uvloop/dns.pyx b/uvloop/dns.pyx
index b7d02f32..01cdd5b8 100644
--- a/uvloop/dns.pyx
+++ b/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,
diff --git a/uvloop/includes/compat.h b/uvloop/includes/compat.h
index 1321ddef..7c7105e5 100644
--- a/uvloop/includes/compat.h
+++ b/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
diff --git a/uvloop/includes/python.pxd b/uvloop/includes/python.pxd
index 8c2b01c1..e5aa27dd 100644
--- a/uvloop/includes/python.pxd
+++ b/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
diff --git a/uvloop/loop.pxd b/uvloop/loop.pxd
index 87062eba..104a6f56 100644
--- a/uvloop/loop.pxd
+++ b/uvloop/loop.pxd
@@ -230,3 +230,5 @@ include "request.pxd"
include "handles/udp.pxd"
include "server.pxd"
+
+include "tracing.pxd"
diff --git a/uvloop/loop.pyx b/uvloop/loop.pyx
index 78ff7188..caa64324 100644
--- a/uvloop/loop.pyx
+++ b/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
diff --git a/uvloop/tracing.pxd b/uvloop/tracing.pxd
new file mode 100644
index 00000000..e0ac2272
--- /dev/null
+++ b/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()
diff --git a/uvloop/tracing.py b/uvloop/tracing.py
new file mode 100644
index 00000000..eba028e5
--- /dev/null
+++ b/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."""
diff --git a/uvloop/tracing.pyx b/uvloop/tracing.pyx
new file mode 100644
index 00000000..ad6f7339
--- /dev/null
+++ b/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)