fix: contextvar usage (#1936)

* refactor trace building

Instead of storing a mutable Context object in the contextvar, the
active span/context is stored instead. This removes the need to have
to clone the context which means that ddtrace will now work out of the
box with async frameworks that use contextvars.

* fix OT with the new context management

* Fix sampling

* remove deprecation for context_provider

* add current_execution_id function

* use execution id to group traces

this matches the previous functionality although we'll probably
want to change this in the near future

* Revert "use execution id to group traces"

This reverts commit 1d53a5c.

* various updates

- formatting
- added lock to _Trace
- updated tests
- docstrings

* support setting sampling_priority through context

* fix gevent tests

* fix asyncio

* enable partial flushing by default

Since multiple execution traces are no longer handled individually it
makes sense to enable partial flushing by default for long running
cross-execution traces.

* use get

* fix aiohttp tests

* store traces per tracer instance

* formatting

* test fixes

* fix futures test

* fix parenting issue for non active spans

* fix flask, asyncio and tornado tests

* clean up

- remove unrelated formatting
- remove new public apis (can be added later)
- add back context.sampling_priority

* sampling cleanup, backwards compatibility

* more cleanup

- revert unnecessary changes

* revert active() returning None

* handle trace-level tags

Trace-level tags like origin and sampling priority belong to a trace
and it's cumbersome to try to keep on the root span.

Context seems like the reasonable place for these things although
if feels a bit icky to overload context with it for a couple of
reasons:
  - the context tags have to be shared between spans (could be easy
    to forget)
  - it overloads the context to support being a "non-local span" and
    storing trace level state.

Ideally we could split this functionality into two structures...

maybe TraceContext and NonLocalSpan...?

* fix opentracing + gevent tests

* fix span.context active span

* rename methods

* get span from context if possible

* remove useless properties

* thread-safe context

* Add context test cases

* Remove backwards compatibility

This is done for several reasons:

- The existing functionality is undefined in a lot of cases so it
  doesn't make sense to support.
- Backwards support is not entirely possible since we are no longer in
  control of the context.

  One example of this: previously the active span of a context could be
  accessed from outside the execution that the span was created in via
  the span.context. This is no longer possible since the execution
  context is abstracted away.
- Adding backwards support for some of the features led to complexity
  which is largely unnecessary with the above points in mind.

* update grpc test

sampling priority is actually unnecessary in these traces

* Clean up docs

* use lock instead of copying

* Use attr for context

* Add update tags api

* Revert context fix, fix tests

The context fix isn't required with the new definition of context

* Apply suggestions from code review

thank u jd and gab

Co-authored-by: Gabriele N. Tornetta <[email protected]>
Co-authored-by: Julien Danjou <[email protected]>

* Doc updates

* Fix typehint

* get_call_context -> current_trace_context

The old API naming no longer made sense as Context is not tied to
a call stack. I think that trace context is more clear as to what
the context is used for.

* update tests

* Reinstate get_call_context but deprecate it

* Expand on context documentation

* Implement local root as span attribute

This makes local root lookup constant time and a bit more defined
at the cost of another span attribute.

* Fix tests

* Address code review comments

- Doc fixes/updates
- Direct return

* Handle possibly None context

* Inline type hints

* delete init context attrs, fix eq, add unit tests

* Fix get_call_context and add back tests

It should always return a context

* Fix context attrs

* Move updating active span to context provider

Previously the active span updating was being done in the
tracer._on_span_finish callback which may not be called from the same
execution in which the span was created (eg. when a span is finished in
a future's callback). This could result in the wrong span being
activated (in the execution in which the span finished).

This has a disadvantage however. There is a possibility of a memory leak
when the following happens: there are executions A and B. A span is
created in A and finished in B. When this happens then a reference to
the span in kept in the contextvar for A. It will be removed/updated
when the context provider is queried for an active span or if a span is
activated. If this doesn't occur then the span (and any spans referenced
from the span) will be kept in memory so long as the contextvars context
lives.

The worst-case here is that a really large trace is created in an
execution and the spans are finished in another and no subsequent
tracing is done for the rest of the application lifecycle. Then this
trace will remain in memory for so long as the executor exists.

* Fix docs

* Add docs for cross-execution tracing

* Move self.enabled shortcut to after warning logic

* Doc cleanup

* Add comment for why the active span is replaced

* Store newly created contexts on a span

* Fixup the context documentation

Co-authored-by: Gabriele N. Tornetta <[email protected]>
Co-authored-by: Julien Danjou <[email protected]>
Co-authored-by: Brett Langdon <[email protected]>

Author: 4 people

ddtrace/context.py

@@ -1,98 +1,102 @@
 import threading
 from typing import Optional
 from typing import TYPE_CHECKING
+from typing import Text
+
+import attr
 
 from .constants import ORIGIN_KEY
 from .constants import SAMPLING_PRIORITY_KEY
+from .internal.compat import NumericType
 from .internal.logger import get_logger
 from .utils.deprecation import deprecated
-from .utils.formats import asbool
-from .utils.formats import get_env
 
 
 if TYPE_CHECKING:
-    from ddtrace import Span
+    from .span import Span
+    from .span import _MetaDictType
+    from .span import _MetricDictType
 
 log = get_logger(__name__)
 
 
+@attr.s(eq=False, slots=True)
 class Context(object):
+    """Represents the state required to propagate a trace across execution
+    boundaries.
     """
-    Context is used to keep track of a hierarchy of spans for the current
-    execution flow. During each logical execution, the same ``Context`` is
-    used to represent a single logical trace, even if the trace is built
-    asynchronously.
-
-    A single code execution may use multiple ``Context`` if part of the execution
-    must not be related to the current tracing. As example, a delayed job may
-    compose a standalone trace instead of being related to the same trace that
-    generates the job itself. On the other hand, if it's part of the same
-    ``Context``, it will be related to the original trace.
-
-    This data structure is thread-safe.
-    """
-
-    _partial_flush_enabled = asbool(get_env("tracer", "partial_flush_enabled", default=False))
-    _partial_flush_min_spans = int(get_env("tracer", "partial_flush_min_spans", default=500))  # type: ignore[arg-type]
-
-    def __init__(
-        self,
-        trace_id=None,  # type: Optional[int]
-        span_id=None,  # type: Optional[int]
-        sampling_priority=None,  # type: Optional[int]
-        dd_origin=None,  # type: Optional[str]
-    ):
-        # type: (...) -> 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._current_span = None  # type: Optional[Span]
-        self._lock = threading.RLock()
 
-        self._parent_trace_id = trace_id
-        self._parent_span_id = span_id
-        self._sampling_priority = sampling_priority
-        self._local_root_span = None  # type: Optional[Span]
-        self.dd_origin = dd_origin
+    trace_id = attr.ib(default=None, type=Optional[int])
+    span_id = attr.ib(default=None, type=Optional[int])
+    _dd_origin = attr.ib(default=None, type=Optional[str], repr=False)
+    _sampling_priority = attr.ib(default=None, type=Optional[NumericType], repr=False)
+    _lock = attr.ib(factory=threading.RLock, type=threading.RLock, repr=False)
+    _meta = attr.ib(factory=dict)  # type: _MetaDictType
+    _metrics = attr.ib(factory=dict)  # type: _MetricDictType
+
+    def __attrs_post_init__(self):
+        self.dd_origin = self._dd_origin
+        self.sampling_priority = self._sampling_priority
+        del self._dd_origin
+        del self._sampling_priority
+
+    def __eq__(self, other):
+        with self._lock:
+            return (
+                self.trace_id == other.trace_id
+                and self.span_id == other.span_id
+                and self._meta == other._meta
+                and self._metrics == other._metrics
+            )
 
-    @property
-    def trace_id(self):
-        """Return current context trace_id."""
+    def _with_span(self, span):
+        # type: (Span) -> Context
+        """Return a shallow copy of the context with the given span."""
         with self._lock:
-            return self._parent_trace_id
+            ctx = self.__class__(trace_id=span.trace_id, span_id=span.span_id)
+            ctx._lock = self._lock
+            ctx._meta = self._meta
+            ctx._metrics = self._metrics
+            return ctx
 
-    @property
-    def span_id(self):
-        """Return current context span_id."""
+    def _update_tags(self, span):
+        # type: (Span) -> None
         with self._lock:
-            return self._parent_span_id
+            span.meta.update(self._meta)
+            span.metrics.update(self._metrics)
 
     @property
     def sampling_priority(self):
-        """Return current context sampling priority."""
-        with self._lock:
-            return self._sampling_priority
+        # type: () -> Optional[NumericType]
+        """Return the context sampling priority for the trace."""
+        return self._metrics.get(SAMPLING_PRIORITY_KEY)
 
     @sampling_priority.setter
     def sampling_priority(self, value):
-        # type: (int) -> None
-        """Set sampling priority."""
+        # type: (Optional[NumericType]) -> None
         with self._lock:
-            self._sampling_priority = value
+            if value is None:
+                if SAMPLING_PRIORITY_KEY in self._metrics:
+                    del self._metrics[SAMPLING_PRIORITY_KEY]
+                return
+            self._metrics[SAMPLING_PRIORITY_KEY] = value
 
-    def _clone(self):
-        # type: () -> Context
+    @property
+    def dd_origin(self):
+        # type: () -> Optional[Text]
+        """Get the origin of the trace."""
+        return self._meta.get(ORIGIN_KEY)
+
+    @dd_origin.setter
+    def dd_origin(self, value):
+        # type: (Optional[Text]) -> None
+        """Set the origin of the trace."""
         with self._lock:
-            new_ctx = Context(
-                trace_id=self._parent_trace_id,
-                span_id=self._parent_span_id,
-                sampling_priority=self._sampling_priority,
-            )
-            new_ctx._current_span = self._current_span
-            return new_ctx
+            if value is None:
+                if ORIGIN_KEY in self._meta:
+                    del self._meta[ORIGIN_KEY]
+                return
+            self._meta[ORIGIN_KEY] = value
 
     @deprecated("Cloning contexts will no longer be required in 0.50", version="0.50")
     def clone(self):
@@ -101,106 +105,4 @@ def clone(self):
         Partially clones the current context.
         It copies everything EXCEPT the registered and finished spans.
         """
-        return self._clone()
-
-    def _get_current_root_span(self):
-        # type: () -> Optional[Span]
-        with self._lock:
-            return self._local_root_span
-
-    @deprecated("Please use tracer.current_root_span() instead", version="0.50")
-    def get_current_root_span(self):
-        # type: () -> Optional[Span]
-        """
-        Return the root span of the context or None if it does not exist.
-        """
-        return self._get_current_root_span()
-
-    def _get_current_span(self):
-        # type: () -> Optional[Span]
-        with self._lock:
-            return self._current_span
-
-    @deprecated("Please use tracer.current_span() instead", version="0.50")
-    def get_current_span(self):
-        # type: () -> Optional[Span]
-        """
-        Return the last active span that corresponds to the last inserted
-        item in the trace list. This cannot be considered as the current active
-        span in asynchronous environments, because some spans can be closed
-        earlier while child spans still need to finish their traced execution.
-        """
-        return self._get_current_span()
-
-    def _set_current_span(self, span):
-        # type: (Optional[Span]) -> None
-        """
-        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
-        else:
-            self._parent_span_id = None
-
-    def _add_span(self, span):
-        # type: (Span) -> None
-        with self._lock:
-            # Assume the first span added to the context is the local root
-            if self._local_root_span is None:
-                self._local_root_span = span
-            self._set_current_span(span)
-            span._context = self
-
-    @deprecated("Context will no longer support active span management in a later version.", version="0.50")
-    def add_span(self, span):
-        # type: (Span) -> None
-        """Activates span in the context."""
-        return self._add_span(span)
-
-    def _close_span(self, span):
-        # type: (Span) -> None
-        with self._lock:
-            if span == self._local_root_span:
-                if self.dd_origin is not None:
-                    span.meta[ORIGIN_KEY] = self.dd_origin
-                if self._sampling_priority is not None:
-                    span.metrics[SAMPLING_PRIORITY_KEY] = self._sampling_priority
-
-            # If a parent exists to the closing span and it is unfinished, then
-            # activate it next.
-            if span._parent and not span._parent.finished:
-                self._set_current_span(span._parent)
-            # Else if the span is the local root of this context, then clear the
-            # context so the next trace can be started.
-            elif span == self._local_root_span:
-                self._set_current_span(span._parent)
-                self._local_root_span = None
-                self._parent_trace_id = None
-                self._parent_span_id = None
-                self._sampling_priority = None
-            # Else the span that is closing is closing after its parent.
-            # This is most likely an error. To ensure future traces are not
-            # affected clear out the context and set the current span to
-            # ``None``.
-            else:
-                log.debug(
-                    "span %r closing after its parent %r, this is an error when not using async", span, span._parent
-                )
-                self._set_current_span(None)
-                self._local_root_span = None
-                self._parent_trace_id = None
-                self._parent_span_id = None
-                self._sampling_priority = None
-
-    @deprecated(message="Context will no longer support active span management in a later version.", version="0.50")
-    def close_span(self, span):
-        # type: (Span) -> None
-        """Updates the context after a span has finished.
-
-        The next active span becomes `span`'s parent.
-        """
-        return self._close_span(span)
+        return self

ddtrace/contrib/asyncio/helpers.py

@@ -7,7 +7,6 @@
 
 import ddtrace
 
-from ...context import Context
 from .provider import CONTEXT_ATTR
 from .wrappers import wrapped_create_task
 
@@ -29,7 +28,7 @@ def ensure_future(coro_or_future, *, loop=None, tracer=None):
     If the current task already has a Context, it will be attached to the new Task so the Trace list will be preserved.
     """
     tracer = tracer or ddtrace.tracer
-    current_ctx = tracer.get_call_context()
+    current_ctx = tracer.current_trace_context()
     task = asyncio.ensure_future(coro_or_future, loop=loop)
     set_call_context(task, current_ctx)
     return task
@@ -53,12 +52,10 @@ def run_in_executor(loop, executor, func, *args, tracer=None):
 
     """
     tracer = tracer or ddtrace.tracer
-    ctx = Context()
-    current_ctx = tracer.get_call_context()
-    ctx._current_span = current_ctx._current_span
+    current_ctx = tracer.current_trace_context()
 
     # prepare the future using an executor wrapper
-    future = loop.run_in_executor(executor, _wrap_executor, func, args, tracer, ctx)
+    future = loop.run_in_executor(executor, _wrap_executor, func, args, tracer, current_ctx)
     return future
 
 

ddtrace/contrib/asyncio/patch.py

@@ -1,12 +1,11 @@
 import asyncio
+import sys
 
 from ddtrace.vendor.wrapt import ObjectProxy
 from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
 
-from ...internal.compat import CONTEXTVARS_IS_AVAILABLE
 from ...utils.wrappers import unwrap as _u
 from .wrappers import wrapped_create_task
-from .wrappers import wrapped_create_task_contextvars
 
 
 def patch():
@@ -17,14 +16,13 @@ def patch():
         return
     setattr(asyncio, "_datadog_patch", True)
 
-    wrapper = wrapped_create_task_contextvars if CONTEXTVARS_IS_AVAILABLE else wrapped_create_task
+    if sys.version_info < (3, 7, 0):
+        _w(asyncio.BaseEventLoop, "create_task", wrapped_create_task)
 
-    _w(asyncio.BaseEventLoop, "create_task", wrapper)
-
-    # also patch event loop if not inheriting the wrapped create_task from BaseEventLoop
-    loop = asyncio.get_event_loop()
-    if not isinstance(loop.create_task, ObjectProxy):
-        _w(loop, "create_task", wrapper)
+        # also patch event loop if not inheriting the wrapped create_task from BaseEventLoop
+        loop = asyncio.get_event_loop()
+        if not isinstance(loop.create_task, ObjectProxy):
+            _w(loop, "create_task", wrapped_create_task)
 
 
 def unpatch():
@@ -33,9 +31,10 @@ def unpatch():
     if getattr(asyncio, "_datadog_patch", False):
         setattr(asyncio, "_datadog_patch", False)
 
-    _u(asyncio.BaseEventLoop, "create_task")
+    if sys.version_info < (3, 7, 0):
+        _u(asyncio.BaseEventLoop, "create_task")
 
-    # also unpatch event loop if not inheriting the already unwrapped create_task from BaseEventLoop
-    loop = asyncio.get_event_loop()
-    if isinstance(loop.create_task, ObjectProxy):
-        _u(loop, "create_task")
+        # also unpatch event loop if not inheriting the already unwrapped create_task from BaseEventLoop
+        loop = asyncio.get_event_loop()
+        if isinstance(loop.create_task, ObjectProxy):
+            _u(loop, "create_task")