[tornado] TracerStackContext is now an instance and not a class with classmethods

Author: Emanuele Palazzetti

ddtrace/contrib/tornado/__init__.py

@@ -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/patch.py

@@ -3,7 +3,7 @@
 
 from wrapt import wrap_function_wrapper as _w
 
-from . import handlers, application, decorators, template, futures, compat, TracerStackContext
+from . import handlers, application, decorators, template, futures, compat, context_provider
 from ...util import unwrap as _u
 
 
@@ -40,7 +40,7 @@ def patch():
 
     # configure the global tracer
     ddtrace.tracer.configure(
-        context_provider=TracerStackContext,
+        context_provider=context_provider,
         wrap_executor=decorators.wrap_executor,
     )
 

ddtrace/contrib/tornado/stack_context.py

@@ -2,9 +2,10 @@
 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
@@ -19,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):
         """
@@ -53,10 +55,9 @@ 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.
@@ -68,15 +69,14 @@ def active(cls):
             # 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)
+            return self._local.get()
         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
+            for stack in reversed(_state.contexts[0]):
+                if isinstance(stack, self.__class__) and stack._active:
+                    return stack._context
 
-    @classmethod
-    def activate(cls, ctx):
+    def activate(self, ctx):
         """
         Set the active ``Context`` for this async execution. If a ``TracerStackContext``
         is not found, the context is discarded.
@@ -87,12 +87,13 @@ def activate(cls, ctx):
         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)
+            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, cls) and stack_ctx.active:
-                    stack_ctx.context = ctx
+                if isinstance(stack_ctx, self.__class__) and stack_ctx._active:
+                    stack_ctx._context = ctx
+        return ctx
 
 
 def run_with_trace_context(func, *args, **kwargs):