[sampling] add distributed tracing capabilities

ddtrace/api.py

@@ -12,40 +12,67 @@
 
 TRACE_COUNT_HEADER = 'X-Datadog-Trace-Count'
 
+_VERSIONS = {'v0.4': {'traces': '/v0.4/traces',
+                      'services': '/v0.4/services',
+                      'compatibility_mode': False,
+                      'fallback': 'v0.3'},
+             'v0.3': {'traces': '/v0.3/traces',
+                     'services': '/v0.3/services',
+                      'compatibility_mode': False,
+                      'fallback': 'v0.2'},
+             'v0.2': {'traces': '/v0.2/traces',
+                     'services': '/v0.2/services',
+                      'compatibility_mode': True,
+                      'fallback': None}}
+
 class API(object):
     """
     Send data to the trace agent using the HTTP protocol and JSON format
     """
-    def __init__(self, hostname, port, headers=None, encoder=None):
+    def __init__(self, hostname, port, headers=None, encoder=None, priority_sampling=False):
         self.hostname = hostname
         self.port = port
-        self._traces = '/v0.3/traces'
-        self._services = '/v0.3/services'
-        self._compatibility_mode = False
-        self._encoder = encoder or get_encoder()
 
-        # overwrite the Content-type with the one chosen in the Encoder
         self._headers = headers or {}
+        self._version = None
+
+        if priority_sampling:
+            self._set_version('v0.4', encoder=encoder)
+        else:
+            self._set_version('v0.3', encoder=encoder)
+
         self._headers.update({
-            'Content-Type': self._encoder.content_type,
             'Datadog-Meta-Lang': 'python',
             'Datadog-Meta-Lang-Version': PYTHON_VERSION,
             'Datadog-Meta-Lang-Interpreter': PYTHON_INTERPRETER,
             'Datadog-Meta-Tracer-Version': ddtrace.__version__,
         })
 
+    def _set_version(self, version, encoder=None):
+        if version not in _VERSIONS:
+            version = 'v0.2'
+        if version == self._version:
+            return
+        self._version = version
+        self._traces = _VERSIONS[version]['traces']
+        self._services = _VERSIONS[version]['services']
+        self._fallback = _VERSIONS[version]['fallback']
+        self._compatibility_mode = _VERSIONS[version]['compatibility_mode']
+        if self._compatibility_mode:
+            self._encoder = JSONEncoder()
+        else:
+            self._encoder = encoder or get_encoder()
+        # overwrite the Content-type with the one chosen in the Encoder
+        self._headers.update({'Content-Type': self._encoder.content_type})
+
     def _downgrade(self):
         """
         Downgrades the used encoder and API level. This method must fallback to a safe
         encoder and API, so that it will success despite users' configurations. This action
         ensures that the compatibility mode is activated so that the downgrade will be
         executed only once.
         """
-        self._compatibility_mode = True
-        self._traces = '/v0.2/traces'
-        self._services = '/v0.2/services'
-        self._encoder = JSONEncoder()
-        self._headers.update({'Content-Type': self._encoder.content_type})
+        self._set_version(self._fallback)
 
     def send_traces(self, traces):
         if not traces:
@@ -55,8 +82,8 @@ def send_traces(self, traces):
         response = self._put(self._traces, data, len(traces))
 
         # the API endpoint is not available so we should downgrade the connection and re-try the call
-        if response.status in [404, 415] and self._compatibility_mode is False:
-            log.debug('calling the endpoint "%s" but received %s; downgrading the API', self._traces, response.status)
+        if response.status in [404, 415] and self._fallback:
+            log.debug('calling endpoint "%s" but received %s; downgrading API', self._traces, response.status)
             self._downgrade()
             return self.send_traces(traces)
 
@@ -73,8 +100,8 @@ def send_services(self, services):
         response = self._put(self._services, data)
 
         # the API endpoint is not available so we should downgrade the connection and re-try the call
-        if response.status in [404, 415] and self._compatibility_mode is False:
-            log.debug('calling the endpoint "%s" but received 404; downgrading the API', self._services)
+        if response.status in [404, 415] and self._fallback:
+            log.debug('calling endpoint "%s" but received %s; downgrading API', self._services, response.status)
             self._downgrade()
             return self.send_services(services)
 

ddtrace/constants.py

@@ -1 +1,2 @@
 FILTERS_KEY = 'FILTERS'
+SAMPLING_PRIORITY_KEY = '_sampling_priority_v1'

ddtrace/context.py

@@ -20,25 +20,31 @@ class Context(object):
 
     This data structure is thread-safe.
     """
-    def __init__(self, trace_id=None, span_id=None, sampled=True):
+    def __init__(self, trace_id=None, span_id=None, sampled=True, sampling_priority=None):
         """
         Initialize a new thread-safe ``Context``.
 
         :param int trace_id: trace_id of parent span
         :param int span_id: span_id of parent span
         """
         self._trace = []
-        self._sampled = sampled
         self._finished_spans = 0
         self._current_span = None
         self._lock = threading.Lock()
-        self._parent_span_id = span_id
+
         self._parent_trace_id = trace_id
+        self._parent_span_id = span_id
+        self._sampled = sampled
+        self._sampling_priority = sampling_priority
+
+    def get_context_attributes(self):
+        """
+        Return the context propagatable attributes.
 
-    def _get_parent_span_ids(self):
-        """ Returns tuple of base trace_id, span_id for distributed tracing."""
+        Useful to propagate context to an external element.
+        """
         with self._lock:
-            return self._parent_trace_id, self._parent_span_id
+            return self._parent_trace_id, self._parent_span_id, self._sampling_priority
 
     def get_current_span(self):
         """
@@ -50,13 +56,28 @@ def get_current_span(self):
         with self._lock:
             return self._current_span
 
+    def _set_current_span(self, span):
+        """
+        Set current span internally.
+
+        Non-safe if not used with a lock. For internal Context usage only.
+        """
+        self._current_span = span
+        if span:
+            self._parent_trace_id = span.trace_id
+            self._parent_span_id = span.span_id
+            self._sampled = span.sampled
+            self._sampling_priority = span.get_sampling_priority()
+        else:
+            self._parent_span_id = None
+
     def add_span(self, span):
         """
         Add a span to the context trace list, keeping it as the last active span.
         """
         with self._lock:
-            self._current_span = span
-            self._sampled = span.sampled
+            self._set_current_span(span)
+
             self._trace.append(span)
             span._context = self
 
@@ -67,7 +88,7 @@ def close_span(self, span):
         """
         with self._lock:
             self._finished_spans += 1
-            self._current_span = span._parent
+            self._set_current_span(span._parent)
 
             # notify if the trace is not closed properly; this check is executed only
             # if the tracer debug_logging is enabled and when the root span is closed
@@ -114,9 +135,11 @@ def get(self):
                 sampled = self._sampled
                 # clean the current state
                 self._trace = []
-                self._sampled = False
                 self._finished_spans = 0
-                self._current_span = None
+                self._parent_trace_id = None
+                self._parent_span_id = None
+                self._sampling_priority = None
+                self._sampled = True
                 return trace, sampled
             else:
                 return None, None
@@ -145,9 +168,7 @@ def set(self, ctx):
     def get(self):
         ctx = getattr(self._locals, 'context', None)
         if not ctx:
-            # create a new Context if it's not available; this action
-            # is done once because the Context has the reset() method
-            # to reuse the same instance
+            # create a new Context if it's not available
             ctx = Context()
             self._locals.context = ctx
 

ddtrace/contrib/aiohttp/middlewares.py

@@ -3,6 +3,7 @@
 from ..asyncio import context_provider
 from ...ext import AppTypes, http
 from ...compat import stringify
+from ...context import Context
 
 
 CONFIG_KEY = 'datadog_trace'
@@ -11,6 +12,7 @@
 
 PARENT_TRACE_HEADER_ID = 'x-datadog-trace-id'
 PARENT_SPAN_HEADER_ID = 'x-datadog-parent-id'
+SAMPLING_PRIORITY_HEADER_ID = 'x-datadog-sampling-priority'
 
 
 @asyncio.coroutine
@@ -29,24 +31,30 @@ def attach_context(request):
         service = app[CONFIG_KEY]['service']
         distributed_tracing = app[CONFIG_KEY]['distributed_tracing_enabled']
 
+        context = tracer.get_call_context()
+
+        # Create a new context based on the propagated information
+        # Do not fill context with distributed sampling if the tracer is disabled
+        # because the would call the callee to generate references to data which
+        # has never been sent to agent.
+        if tracer.enabled and distributed_tracing:
+            trace_id = int(request.headers.get(PARENT_TRACE_HEADER_ID, 0))
+            parent_span_id = int(request.headers.get(PARENT_SPAN_HEADER_ID, 0))
+            sampling_priority = request.headers.get(SAMPLING_PRIORITY_HEADER_ID)
+            # keep sampling priority as None if not propagated, to support older client versions on the parent side
+            if sampling_priority:
+                sampling_priority = int(sampling_priority)
+
+            context = Context(trace_id=trace_id, span_id=parent_span_id, sampling_priority=sampling_priority)
+
         # trace the handler
-        request_span = tracer.trace(
+        request_span = tracer.start_span(
             'aiohttp.request',
             service=service,
             span_type=http.TYPE,
+            child_of=context,
         )
 
-        if distributed_tracing:
-            # set parent trace/span IDs if present:
-            # http://pypi.datadoghq.com/trace/docs/#distributed-tracing
-            parent_trace_id = request.headers.get(PARENT_TRACE_HEADER_ID)
-            if parent_trace_id is not None:
-                request_span.trace_id = int(parent_trace_id)
-
-            parent_span_id = request.headers.get(PARENT_SPAN_HEADER_ID)
-            if parent_span_id is not None:
-                request_span.parent_id = int(parent_span_id)
-
         # attach the context and the root span to the request; the Context
         # may be freely used by the application code
         request[REQUEST_CONTEXT_KEY] = request_span.context

ddtrace/contrib/asyncio/helpers.py

@@ -18,6 +18,9 @@ def set_call_context(task, ctx):
     """
     Updates the ``Context`` for the given Task. Useful when you need to
     pass the context among different tasks.
+
+    This method is available for backward-compatibility. Use the
+    ``AsyncioContextProvider`` API to set the current active ``Context``.
     """
     setattr(task, CONTEXT_ATTR, ctx)
 
@@ -74,7 +77,7 @@ def _wrap_executor(fn, args, tracer, ctx):
     # the AsyncioContextProvider knows that this is a new thread
     # so it is legit to pass the Context in the thread-local storage;
     # fn() will be executed outside the asyncio loop as a synchronous code
-    tracer._context_provider._local.set(ctx)
+    tracer.context_provider.activate(ctx)
     return fn(*args)
 
 
@@ -104,16 +107,11 @@ def _wrapped_create_task(wrapped, instance, args, kwargs):
 
     ctx = getattr(current_task, CONTEXT_ATTR, None)
     span = ctx.get_current_span() if ctx else None
-    if span:
-        parent_trace_id, parent_span_id = span.trace_id, span.span_id
-    elif ctx:
-        parent_trace_id, parent_span_id = ctx._get_parent_span_ids()
-    else:
-        parent_trace_id = parent_span_id = None
-
-    if parent_trace_id and parent_span_id:
+    if ctx:
+        parent_trace_id, parent_span_id, sampling_priority = ctx.get_context_attributes()
+
         # current task has a context, so parent a new context to the base context
-        new_ctx = Context(trace_id=parent_trace_id, span_id=parent_span_id)
+        new_ctx = Context(trace_id=parent_trace_id, span_id=parent_span_id, sampling_priority=sampling_priority)
         set_call_context(new_task, new_ctx)
 
     return new_task

ddtrace/contrib/asyncio/provider.py

@@ -14,8 +14,28 @@ class AsyncioContextProvider(DefaultContextProvider):
     execution. It must be used in asynchronous programming that relies
     in the built-in ``asyncio`` library. Framework instrumentation that
     is built on top of the ``asyncio`` library, can use this provider.
+
+    This Context Provider inherits from ``DefaultContextProvider`` because
+    it uses a thread-local storage when the ``Context`` is propagated to
+    a different thread, than the one that is running the async loop.
     """
-    def __call__(self, loop=None):
+    def activate(self, context, loop=None):
+        """Sets the scoped ``Context`` for the current running ``Task``.
+        """
+        try:
+            loop = loop or asyncio.get_event_loop()
+        except RuntimeError:
+            # detects if a loop is available in the current thread;
+            # This happens when a new thread is created from the one that is running
+            # the async loop
+            return self._local.set(context)
+
+        # the current unit of work (if tasks are used)
+        task = asyncio.Task.current_task(loop=loop)
+        setattr(task, CONTEXT_ATTR, context)
+        return context
+
+    def active(self, loop=None):
         """
         Returns the scoped Context for this execution flow. The ``Context`` uses
         the current task as a carrier so if a single task is used for the entire application,

ddtrace/contrib/gevent/greenlet.py

@@ -29,7 +29,11 @@ def __init__(self, *args, **kwargs):
         # the context is always available made exception of the main greenlet
         if ctx:
             # create a new context that inherits the current active span
-            new_ctx = Context()
-            new_ctx._sampled = ctx._sampled
+            # TODO: a better API for Context, should get the tuple at once
+            new_ctx = Context(
+                trace_id=ctx._parent_trace_id,
+                span_id=ctx._parent_span_id,
+                sampled=ctx._sampled,
+            )
             new_ctx._current_span = ctx._current_span
             setattr(self, CONTEXT_ATTR, new_ctx)

ddtrace/contrib/gevent/provider.py

@@ -15,12 +15,19 @@ class GeventContextProvider(BaseContextProvider):
     in the ``gevent`` library. Framework instrumentation that uses the
     gevent WSGI server (or gevent in general), can use this provider.
     """
-    def __call__(self):
+    def activate(self, context):
+        """Sets the scoped ``Context`` for the current running ``Greenlet``.
+        """
+        current_g = gevent.getcurrent()
+        if current_g is not None:
+            setattr(current_g, CONTEXT_ATTR, context)
+            return context
+
+    def active(self):
         """
         Returns the scoped ``Context`` for this execution flow. The ``Context``
         uses the ``Greenlet`` class as a carrier, and everytime a greenlet
-        is created it receives the "parent" context. The main greenlet
-        will never have an attached ``Context``.
+        is created it receives the "parent" context.
         """
         current_g = gevent.getcurrent()
         ctx = getattr(current_g, CONTEXT_ATTR, None)
@@ -29,9 +36,14 @@ def __call__(self):
             return ctx
 
         # the Greenlet doesn't have a Context so it's created and attached
-        # unless it's the main greenlet; in that case we must be sure
-        # that no Context is generated
-        if current_g.parent:
+        # TODO: previous implementation avoided to add a Context to the main
+        # greenlet because it could have side-effects when switching back
+        # and forth between different executions. This results in issues such
+        # as: https://github.com/DataDog/dd-trace-py/issues/309
+        # and is required for Distributed Tracing when providing a new arbitrary
+        # Context. On the other hand, it's imperative to double check if there
+        # are side effects.
+        if current_g:
             ctx = Context()
             setattr(current_g, CONTEXT_ATTR, ctx)
             return ctx