context_manager.py

import abc
import threading
from ddtrace.vendor import six

from .logger import get_logger
from ..context import Context

log = get_logger(__name__)

try:
    from contextvars import ContextVar

    _DD_CONTEXTVAR = ContextVar("datadog_contextvar", default=None)
    CONTEXTVARS_IS_AVAILABLE = True
except ImportError:
    CONTEXTVARS_IS_AVAILABLE = False


class BaseContextManager(six.with_metaclass(abc.ABCMeta)):
    def __init__(self, reset=True):
        if reset:
            self.reset()

    @abc.abstractmethod
    def _has_active_context(self):
        pass

    @abc.abstractmethod
    def set(self, ctx):
        pass

    @abc.abstractmethod
    def get(self):
        pass

    def reset(self):
        pass


class ThreadLocalContext(BaseContextManager):
    """
    ThreadLocalContext can be used as a tracer global reference to create
    a different ``Context`` for each thread. In synchronous tracer, this
    is required to prevent multiple threads sharing the same ``Context``
    in different executions.
    """

    def __init__(self, reset=True):
        # always initialize a new thread-local context holder
        super(ThreadLocalContext, self).__init__(reset=True)

    def _has_active_context(self):
        """
        Determine whether we have a currently active context for this thread

        :returns: Whether an active context exists
        :rtype: bool
        """
        ctx = getattr(self._locals, "context", None)
        return ctx is not None

    def set(self, ctx):
        setattr(self._locals, "context", ctx)

    def get(self):
        ctx = getattr(self._locals, "context", None)
        if not ctx:
            # create a new Context if it's not available
            ctx = Context()
            self._locals.context = ctx

        return ctx

    def reset(self):
        self._locals = threading.local()


class ContextVarContextManager(BaseContextManager):
    """
    _ContextVarContext can be used in place of the ThreadLocalContext for Python
    3.7 and above to manage different ``Context`` objects for each thread and
    async task.
    """

    def _has_active_context(self):
        ctx = _DD_CONTEXTVAR.get()
        return ctx is not None

    def set(self, ctx):
        _DD_CONTEXTVAR.set(ctx)

    def get(self):
        ctx = _DD_CONTEXTVAR.get()
        if not ctx:
            ctx = Context()
            self.set(ctx)

        return ctx

    def reset(self):
        _DD_CONTEXTVAR.set(None)


if CONTEXTVARS_IS_AVAILABLE:
    DefaultContextManager = ContextVarContextManager
else:
    DefaultContextManager = ThreadLocalContext