[tornado] patch concurrent.futures to propagate Tornado context from the main thread

Author: Emanuele Palazzetti

ddtrace/contrib/tornado/futures.py

@@ -0,0 +1,33 @@
+from ddtrace import tracer
+from ddtrace.context import Context
+
+
+def _wrap_submit(func, instance, args, kwargs):
+    """
+    Wrap `Executor` method used to submit a work executed in another
+    thread. This wrapper ensures that a new `Context` is created and
+    properly propagated using an intermediate function.
+    """
+    # create a new Context with the right active Span
+    # TODO: the current implementation doesn't provide the GlobalTracer
+    # singleton, so we should rely in our top-level import
+    ctx = Context()
+    current_ctx = tracer.get_call_context()
+    ctx._current_span = current_ctx._current_span
+
+    # extract the target function that must be executed in
+    # a new thread and the `target` arguments
+    fn = args[0]
+    fn_args = args[1:]
+    return func(_wrap_execution, ctx, fn, fn_args, kwargs)
+
+def _wrap_execution(ctx, fn, args, kwargs):
+    """
+    Intermediate target function that is executed in a new thread;
+    it receives the original function with arguments and keyword
+    arguments, including our tracing `Context`. The current context
+    provider sets the Active context in a thread local storage
+    variable because it's outside the asynchronous loop.
+    """
+    tracer.context_provider.activate(ctx)
+    return fn(*args, **kwargs)

ddtrace/contrib/tornado/patch.py

@@ -1,9 +1,10 @@
 import ddtrace
 import tornado
+import concurrent
 
 from wrapt import wrap_function_wrapper as _w
 
-from . import handlers, application, decorators, template, TracerStackContext
+from . import handlers, application, decorators, template, futures, TracerStackContext
 from ...util import unwrap as _u
 
 
@@ -31,6 +32,12 @@ def patch():
     # patch Template system
     _w('tornado.template', 'Template.generate', template.generate)
 
+    # patch Python Futures when an Executor pool is used
+    # TODO: this may be a generic module and should be moved
+    # in a separate contributions when we want to support multi-threading
+    # context propagation
+    _w('concurrent.futures', 'ThreadPoolExecutor.submit', futures._wrap_submit)
+
     # configure the global tracer
     ddtrace.tracer.configure(
         context_provider=TracerStackContext,
@@ -53,3 +60,4 @@ def unpatch():
     _u(tornado.web.Application, '__init__')
     _u(tornado.concurrent, 'run_on_executor')
     _u(tornado.template.Template, 'generate')
+    _u(concurrent.futures.ThreadPoolExecutor, 'submit')

ddtrace/contrib/tornado/stack_context.py

@@ -1,3 +1,4 @@
+from tornado.ioloop import IOLoop
 from tornado.stack_context import StackContextInconsistentError, _state
 
 from ...context import Context
@@ -59,20 +60,39 @@ def active(cls):
         """
         Return the ``Context`` from the current execution flow. This method can be
         used inside a Tornado coroutine to retrieve and use the current tracing context.
+        If used in a separated Thread, the `_state` thread-local storage is used to
+        propagate the current Active context from the `MainThread`.
         """
-        for ctx in reversed(_state.contexts[0]):
-            if isinstance(ctx, cls) and ctx.active:
-                return ctx.context
+        io_loop = getattr(IOLoop._current, 'instance', None)
+        if io_loop is None:
+            # if a Tornado loop is not available, it means that this method
+            # has been called from a synchronous code, so we can rely in a
+            # thread-local storage
+            return getattr(_state, '__datadog_context', None)
+        else:
+            # we're inside a Tornado loop so the TracerStackContext is used
+            for ctx in reversed(_state.contexts[0]):
+                if isinstance(ctx, cls) and ctx.active:
+                    return ctx.context
 
     @classmethod
     def activate(cls, ctx):
         """
         Set the active ``Context`` for this async execution. If a ``TracerStackContext``
         is not found, the context is discarded.
+        If used in a separated Thread, the `_state` thread-local storage is used to
+        propagate the current Active context from the `MainThread`.
         """
-        for stack_ctx in reversed(_state.contexts[0]):
-            if isinstance(stack_ctx, cls) and stack_ctx.active:
-                stack_ctx.context = ctx
+        io_loop = getattr(IOLoop._current, 'instance', None)
+        if io_loop is None:
+            # because we're outside of an asynchronous execution, we store
+            # the current context in a thread-local storage
+            setattr(_state, '__datadog_context', ctx)
+        else:
+            # we're inside a Tornado loop so the TracerStackContext is used
+            for stack_ctx in reversed(_state.contexts[0]):
+                if isinstance(stack_ctx, cls) and stack_ctx.active:
+                    stack_ctx.context = ctx
 
 
 def run_with_trace_context(func, *args, **kwargs):

tests/contrib/tornado/test_executor_decorator.py

@@ -43,6 +43,36 @@ def test_on_executor_handler(self):
         eq_(0, executor_span.error)
         ok_(executor_span.duration >= 0.05)
 
+    def test_on_executor_submit(self):
+        # it should propagate the context when a handler uses directly the `executor.submit()`
+        response = self.fetch('/executor_submit_handler/')
+        eq_(200, response.code)
+
+        traces = self.tracer.writer.pop_traces()
+        eq_(2, len(traces))
+        eq_(1, len(traces[0]))
+        eq_(1, len(traces[1]))
+
+        # this trace yields the execution of the thread
+        request_span = traces[1][0]
+        eq_('tornado-web', request_span.service)
+        eq_('tornado.request', request_span.name)
+        eq_('http', request_span.span_type)
+        eq_('tests.contrib.tornado.web.app.ExecutorSubmitHandler', request_span.resource)
+        eq_('GET', request_span.get_tag('http.method'))
+        eq_('200', request_span.get_tag('http.status_code'))
+        eq_('/executor_submit_handler/', request_span.get_tag('http.url'))
+        eq_(0, request_span.error)
+        ok_(request_span.duration >= 0.05)
+
+        # this trace is executed in a different thread
+        executor_span = traces[0][0]
+        eq_('tornado-web', executor_span.service)
+        eq_('tornado.executor.query', executor_span.name)
+        eq_(executor_span.parent_id, request_span.span_id)
+        eq_(0, executor_span.error)
+        ok_(executor_span.duration >= 0.05)
+
     def test_on_delayed_executor_handler(self):
         # it should trace a handler that uses @run_on_executor but that doesn't
         # wait for its termination

tests/contrib/tornado/utils.py

@@ -2,8 +2,7 @@
 
 from ddtrace.contrib.tornado import patch, unpatch
 
-from .web import app
-from .web.compat import reload_module
+from .web import app, compat
 from ...test_tracer import get_dummy_tracer
 
 
@@ -16,7 +15,8 @@ class TornadoTestCase(AsyncHTTPTestCase):
     def get_app(self):
         # patch Tornado and reload module app
         patch()
-        reload_module(app)
+        compat.reload_module(compat)
+        compat.reload_module(app)
         # create a dummy tracer and a Tornado web application
         self.tracer = get_dummy_tracer()
         settings = self.get_settings()
@@ -34,4 +34,5 @@ def tearDown(self):
         super(TornadoTestCase, self).tearDown()
         # unpatch Tornado
         unpatch()
-        reload_module(app)
+        compat.reload_module(compat)
+        compat.reload_module(app)

tests/contrib/tornado/web/app.py

@@ -159,6 +159,22 @@ def get(self):
         self.write('OK')
 
 
+class ExecutorSubmitHandler(tornado.web.RequestHandler):
+    executor = ThreadPoolExecutor(max_workers=3)
+
+    def query(self):
+        tracer = self.settings['datadog_trace']['tracer']
+        with tracer.trace('tornado.executor.query'):
+            time.sleep(0.05)
+
+    @tornado.gen.coroutine
+    def get(self):
+        # run the query in another Executor, without using
+        # Tornado decorators
+        yield self.executor.submit(self.query)
+        self.write('OK')
+
+
 class ExecutorDelayedHandler(tornado.web.RequestHandler):
     # used automatically by the @run_on_executor decorator
     executor = ThreadPoolExecutor(max_workers=3)
@@ -296,6 +312,7 @@ def make_app(settings={}):
         (r'/template_exception/', TemplateExceptionHandler),
         # handlers that spawn new threads
         (r'/executor_handler/', ExecutorHandler),
+        (r'/executor_submit_handler/', ExecutorSubmitHandler),
         (r'/executor_delayed_handler/', ExecutorDelayedHandler),
         (r'/executor_custom_handler/', ExecutorCustomHandler),
         (r'/executor_custom_args_handler/', ExecutorCustomArgsHandler),