Merge pull request #362 from palazzem/python-futures

[tornado] patch concurrent.futures if available

Author: Emanuele Palazzetti

ddtrace/context.py

@@ -51,6 +51,12 @@ def span_id(self):
         with self._lock:
             return self._parent_span_id
 
+    @property
+    def sampled(self):
+        """Return current context sampled flag."""
+        with self._lock:
+            return self._sampled
+
     @property
     def sampling_priority(self):
         """Return current context sampling priority."""

ddtrace/contrib/tornado/__init__.py

@@ -3,7 +3,7 @@
 Auto instrumentation is available using the ``patch`` function that **must be called before**
 importing the tornado library. The following is an example::
 
-    # patch before importing tornado
+    # patch before importing tornado and concurrent.futures
     from ddtrace import tracer, patch
     patch(tornado=True)
 
@@ -83,10 +83,10 @@ def notify(self):
 with require_modules(required_modules) as missing_modules:
     if not missing_modules:
         from .stack_context import run_with_trace_context, TracerStackContext
-        from .patch import patch, unpatch
 
-        # alias for API compatibility
-        context_provider = TracerStackContext
+        context_provider = TracerStackContext()
+
+        from .patch import patch, unpatch
 
         __all__ = [
             'patch',

ddtrace/contrib/tornado/application.py

@@ -2,7 +2,7 @@
 
 from tornado import template
 
-from . import decorators, TracerStackContext
+from . import decorators, context_provider
 from .constants import CONFIG_KEY
 
 from ...ext import AppTypes
@@ -37,7 +37,7 @@ def tracer_config(__init__, app, args, kwargs):
     # global tracer while here we can have a different instance (even if
     # this is not usual).
     tracer.configure(
-        context_provider=TracerStackContext,
+        context_provider=context_provider,
         wrap_executor=decorators.wrap_executor,
         enabled=settings.get('enabled', None),
         hostname=settings.get('agent_hostname', None),

ddtrace/contrib/tornado/compat.py

@@ -0,0 +1,9 @@
+from ..util import require_modules
+
+
+optional_modules = ['concurrent.futures']
+
+with require_modules(optional_modules) as missing_modules:
+    # detect if concurrent.futures is available as a Python
+    # stdlib or Python 2.7 backport
+    futures_available = len(missing_modules) == 0

ddtrace/contrib/tornado/futures.py

@@ -0,0 +1,28 @@
+import ddtrace
+
+
+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.
+    """
+    # propagate the same Context in the new thread
+    current_ctx = ddtrace.tracer.context_provider.active()
+
+    # 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, current_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.
+    """
+    ddtrace.tracer.context_provider.activate(ctx)
+    return fn(*args, **kwargs)

ddtrace/contrib/tornado/patch.py

@@ -3,7 +3,7 @@
 
 from wrapt import wrap_function_wrapper as _w
 
-from . import handlers, application, decorators, template, TracerStackContext
+from . import handlers, application, decorators, template, futures, compat, context_provider
 from ...util import unwrap as _u
 
 
@@ -25,15 +25,19 @@ def patch():
     _w('tornado.web', 'RequestHandler.on_finish', handlers.on_finish)
     _w('tornado.web', 'RequestHandler.log_exception', handlers.log_exception)
 
-    # patch Tornado decorators
-    _w('tornado.concurrent', 'run_on_executor', decorators._run_on_executor)
-
     # 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
+    if compat.futures_available:
+        _w('concurrent.futures', 'ThreadPoolExecutor.submit', futures._wrap_submit)
+
     # configure the global tracer
     ddtrace.tracer.configure(
-        context_provider=TracerStackContext,
+        context_provider=context_provider,
         wrap_executor=decorators.wrap_executor,
     )
 
@@ -53,3 +57,6 @@ def unpatch():
     _u(tornado.web.Application, '__init__')
     _u(tornado.concurrent, 'run_on_executor')
     _u(tornado.template.Template, 'generate')
+
+    if compat.futures_available:
+        _u('concurrent.futures.ThreadPoolExecutor', 'submit')

ddtrace/contrib/tornado/stack_context.py

@@ -1,9 +1,11 @@
+from tornado.ioloop import IOLoop
 from tornado.stack_context import StackContextInconsistentError, _state
 
 from ...context import Context
+from ...provider import DefaultContextProvider
 
 
-class TracerStackContext(object):
+class TracerStackContext(DefaultContextProvider):
     """
     A context manager that manages ``Context`` instances in a thread-local state.
     It must be used everytime a Tornado's handler or coroutine is used within a
@@ -18,8 +20,9 @@ class TracerStackContext(object):
     https://github.com/tornadoweb/tornado/issues/1063
     """
     def __init__(self):
-        self.active = True
-        self.context = Context()
+        super(TracerStackContext, self).__init__()
+        self._active = True
+        self._context = Context()
 
     def enter(self):
         """
@@ -52,27 +55,45 @@ def __exit__(self, type, value, traceback):
         self.new_contexts = None
 
     def deactivate(self):
-        self.active = False
+        self._active = False
 
-    @classmethod
-    def active(cls):
+    def active(self):
         """
         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
-
-    @classmethod
-    def activate(cls, ctx):
+        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 self._local.get()
+        else:
+            # we're inside a Tornado loop so the TracerStackContext is used
+            for stack in reversed(_state.contexts[0]):
+                if isinstance(stack, self.__class__) and stack._active:
+                    return stack._context
+
+    def activate(self, 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
+            self._local.set(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, self.__class__) and stack_ctx._active:
+                    stack_ctx._context = ctx
+        return ctx
 
 
 def run_with_trace_context(func, *args, **kwargs):

tests/contrib/tornado/test_executor_decorator.py

@@ -2,6 +2,7 @@
 import unittest
 
 from nose.tools import eq_, ok_
+from ddtrace.contrib.tornado.compat import futures_available
 
 from tornado import version_info
 
@@ -19,12 +20,11 @@ def test_on_executor_handler(self):
         eq_(200, response.code)
 
         traces = self.tracer.writer.pop_traces()
-        eq_(2, len(traces))
-        eq_(1, len(traces[0]))
-        eq_(1, len(traces[1]))
+        eq_(1, len(traces))
+        eq_(2, len(traces[0]))
 
         # this trace yields the execution of the thread
-        request_span = traces[1][0]
+        request_span = traces[0][0]
         eq_('tornado-web', request_span.service)
         eq_('tornado.request', request_span.name)
         eq_('http', request_span.span_type)
@@ -36,46 +36,39 @@ def test_on_executor_handler(self):
         ok_(request_span.duration >= 0.05)
 
         # this trace is executed in a different thread
-        executor_span = traces[0][0]
+        executor_span = traces[0][1]
         eq_('tornado-web', executor_span.service)
         eq_('tornado.executor.with', 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
-        response = self.fetch('/executor_delayed_handler/')
+    @unittest.skipUnless(futures_available, 'Futures must be available to test direct submit')
+    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)
 
-        # timeout for the background thread execution
-        time.sleep(0.1)
-
         traces = self.tracer.writer.pop_traces()
-        eq_(2, len(traces))
-        eq_(1, len(traces[0]))
-        eq_(1, len(traces[1]))
-
-        # order the `traces` list to have deterministic results
-        # (required only for this special use case)
-        traces.sort(key=lambda x: x[0].name, reverse=True)
+        eq_(1, len(traces))
+        eq_(2, len(traces[0]))
 
         # this trace yields the execution of the thread
         request_span = traces[0][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.ExecutorDelayedHandler', request_span.resource)
+        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_delayed_handler/', request_span.get_tag('http.url'))
+        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[1][0]
+        executor_span = traces[0][1]
         eq_('tornado-web', executor_span.service)
-        eq_('tornado.executor.with', executor_span.name)
+        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)
@@ -86,12 +79,11 @@ def test_on_executor_exception_handler(self):
         eq_(500, response.code)
 
         traces = self.tracer.writer.pop_traces()
-        eq_(2, len(traces))
-        eq_(1, len(traces[0]))
-        eq_(1, len(traces[1]))
+        eq_(1, len(traces))
+        eq_(2, len(traces[0]))
 
         # this trace yields the execution of the thread
-        request_span = traces[1][0]
+        request_span = traces[0][0]
         eq_('tornado-web', request_span.service)
         eq_('tornado.request', request_span.name)
         eq_('http', request_span.span_type)
@@ -104,7 +96,7 @@ def test_on_executor_exception_handler(self):
         ok_('Exception: Ouch!' in request_span.get_tag('error.stack'))
 
         # this trace is executed in a different thread
-        executor_span = traces[0][0]
+        executor_span = traces[0][1]
         eq_('tornado-web', executor_span.service)
         eq_('tornado.executor.with', executor_span.name)
         eq_(executor_span.parent_id, request_span.span_id)
@@ -123,12 +115,11 @@ def test_on_executor_custom_kwarg(self):
         eq_(200, response.code)
 
         traces = self.tracer.writer.pop_traces()
-        eq_(2, len(traces))
-        eq_(1, len(traces[0]))
-        eq_(1, len(traces[1]))
+        eq_(1, len(traces))
+        eq_(2, len(traces[0]))
 
         # this trace yields the execution of the thread
-        request_span = traces[1][0]
+        request_span = traces[0][0]
         eq_('tornado-web', request_span.service)
         eq_('tornado.request', request_span.name)
         eq_('http', request_span.span_type)
@@ -140,7 +131,7 @@ def test_on_executor_custom_kwarg(self):
         ok_(request_span.duration >= 0.05)
 
         # this trace is executed in a different thread
-        executor_span = traces[0][0]
+        executor_span = traces[0][1]
         eq_('tornado-web', executor_span.service)
         eq_('tornado.executor.with', executor_span.name)
         eq_(executor_span.parent_id, request_span.span_id)

tests/contrib/tornado/test_safety.py

@@ -48,6 +48,7 @@ class TestAppSafety(TornadoTestCase):
     """
     Ensure that the application patch has the proper safety guards.
     """
+
     def test_trace_unpatch(self):
         # the application must not be traced if unpatch() is called
         patch()
@@ -108,6 +109,27 @@ def test_arbitrary_resource_404(self):
         eq_('tornado.web.ErrorHandler', request_span.resource)
         eq_('/does_not_exist/', request_span.get_tag('http.url'))
 
+    @gen_test
+    def test_futures_without_context(self):
+        # ensures that if futures propagation is available, an empty
+        # context doesn't crash the system
+        from .web.compat import ThreadPoolExecutor
+
+        def job():
+            with self.tracer.trace('job'):
+                return 42
+
+        executor = ThreadPoolExecutor(max_workers=3)
+        yield executor.submit(job)
+
+        traces = self.tracer.writer.pop_traces()
+        eq_(1, len(traces))
+        eq_(1, len(traces[0]))
+
+        # this trace yields the execution of the thread
+        span = traces[0][0]
+        eq_('job', span.name)
+
 
 class TestCustomAppSafety(TornadoTestCase):
     """