mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-07-30 05:32:30 +08:00
108 lines
2.5 KiB
Python
108 lines
2.5 KiB
Python
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
|