aiohttp 2.x+ client tracing support

ddtrace/context.py

@@ -181,9 +181,13 @@ def get(self):
                 # clean the current state
                 self._trace = []
                 self._finished_spans = 0
-                self._parent_trace_id = None
-                self._parent_span_id = None
-                self._sampling_priority = None
+
+                # Don't clear out the parent trace IDs as this may be a cloned
+                # context in a new thread/task without a outer span
+                if not self._parent_trace_id or (trace and trace[0].trace_id == self._parent_trace_id):
+                    self._parent_trace_id = None
+                    self._parent_span_id = None
+                    self._sampling_priority = None
                 return trace, sampled
 
             elif self._partial_flush_enabled:

ddtrace/contrib/aiohttp/middlewares.py

@@ -1,4 +1,5 @@
 import asyncio
+import functools
 
 from ..asyncio import context_provider
 from ...compat import stringify
@@ -8,14 +9,73 @@
 from ...settings import config
 
 
+try:
+    from aiohttp.web import middleware
+    AIOHTTP_2x = True
+except ImportError:
+    AIOHTTP_2x = False
+
+    def middleware(f):
+        return f
+
+
 CONFIG_KEY = 'datadog_trace'
 REQUEST_CONTEXT_KEY = 'datadog_context'
 REQUEST_CONFIG_KEY = '__datadog_trace_config'
 REQUEST_SPAN_KEY = '__datadog_request_span'
 
 
 @asyncio.coroutine
-def trace_middleware(app, handler):
+@middleware
+def trace_middleware_2x(request, handler, app=None):
+    # application configs
+    if app is None:
+        app = request.app
+
+    tracer = app[CONFIG_KEY]['tracer']
+    service = app[CONFIG_KEY]['service']
+    distributed_tracing = app[CONFIG_KEY]['distributed_tracing_enabled']
+
+    # Create a new context based on the propagated information.
+    if distributed_tracing:
+        propagator = HTTPPropagator()
+        context = propagator.extract(request.headers)
+        # Only need to active the new context if something was propagated
+        if context.trace_id:
+            tracer.context_provider.activate(context)
+
+    # trace the handler
+    request_span = tracer.trace(
+        'aiohttp.request',
+        service=service,
+            span_type=SpanTypes.WEB,
+    )
+        request_span.set_tag(SPAN_MEASURED_KEY)
+
+    # Configure trace search sample rate
+    # DEV: aiohttp is special case maintains separate configuration from config api
+    analytics_enabled = app[CONFIG_KEY]['analytics_enabled']
+    if (config.analytics_enabled and analytics_enabled is not False) or analytics_enabled is True:
+        request_span.set_tag(
+            ANALYTICS_SAMPLE_RATE_KEY,
+            app[CONFIG_KEY].get('analytics_sample_rate', True)
+        )
+
+    # 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
+    request[REQUEST_SPAN_KEY] = request_span
+        request[REQUEST_CONFIG_KEY] = app[CONFIG_KEY]
+    try:
+            response = yield from handler(request)
+        return response
+    except Exception:
+        request_span.set_traceback()
+        raise
+
+
+@asyncio.coroutine
+def trace_middleware_1x(app, handler):
     """
     ``aiohttp`` middleware that traces the handler execution.
     Because handlers are run in different tasks for each request, we attach the Context
@@ -24,50 +84,10 @@ def trace_middleware(app, handler):
     * the Task is used by the internal automatic instrumentation
     * the ``Context`` attached to the request can be freely used in the application code
     """
-    @asyncio.coroutine
-    def attach_context(request):
-        # application configs
-        tracer = app[CONFIG_KEY]['tracer']
-        service = app[CONFIG_KEY]['service']
-        distributed_tracing = app[CONFIG_KEY]['distributed_tracing_enabled']
-
-        # Create a new context based on the propagated information.
-        if distributed_tracing:
-            propagator = HTTPPropagator()
-            context = propagator.extract(request.headers)
-            # Only need to active the new context if something was propagated
-            if context.trace_id:
-                tracer.context_provider.activate(context)
-
-        # trace the handler
-        request_span = tracer.trace(
-            'aiohttp.request',
-            service=service,
-            span_type=SpanTypes.WEB,
-        )
-        request_span.set_tag(SPAN_MEASURED_KEY)
+    return functools.partial(trace_middleware_2x, handler=handler, app=app)
 
-        # Configure trace search sample rate
-        # DEV: aiohttp is special case maintains separate configuration from config api
-        analytics_enabled = app[CONFIG_KEY]['analytics_enabled']
-        if (config.analytics_enabled and analytics_enabled is not False) or analytics_enabled is True:
-            request_span.set_tag(
-                ANALYTICS_SAMPLE_RATE_KEY,
-                app[CONFIG_KEY].get('analytics_sample_rate', True)
-            )
-
-        # 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
-        request[REQUEST_SPAN_KEY] = request_span
-        request[REQUEST_CONFIG_KEY] = app[CONFIG_KEY]
-        try:
-            response = yield from handler(request)
-            return response
-        except Exception:
-            request_span.set_traceback()
-            raise
-    return attach_context
+
+trace_middleware = trace_middleware_2x if AIOHTTP_2x else trace_middleware_1x
 
 
 @asyncio.coroutine
@@ -114,14 +134,15 @@ def on_prepare(request, response):
     request_span.finish()
 
 
-def trace_app(app, tracer, service='aiohttp-web'):
+def trace_app(app, tracer, service='aiohttp-web', distributed_tracing=True):
     """
     Tracing function that patches the ``aiohttp`` application so that it will be
     traced using the given ``tracer``.
 
     :param app: aiohttp application to trace
     :param tracer: tracer instance to use
     :param service: service name of tracer
+    :param distributed_tracing: set to True to enable distributed tracing
     """
 
     # safe-guard: don't trace an application twice
@@ -133,7 +154,7 @@ def trace_app(app, tracer, service='aiohttp-web'):
     app[CONFIG_KEY] = {
         'tracer': tracer,
         'service': service,
-        'distributed_tracing_enabled': True,
+        'distributed_tracing_enabled': distributed_tracing,
         'analytics_enabled': None,
         'analytics_sample_rate': 1.0,
     }