context.py

import base64
import re
from typing import Any
from typing import Optional
from typing import Text

from ddtrace._trace._span_link import SpanLink
from ddtrace.constants import _ORIGIN_KEY
from ddtrace.constants import _SAMPLING_PRIORITY_KEY
from ddtrace.constants import _USER_ID_KEY
from ddtrace.internal.compat import NumericType
from ddtrace.internal.constants import MAX_UINT_64BITS as _MAX_UINT_64BITS
from ddtrace.internal.constants import W3C_TRACEPARENT_KEY
from ddtrace.internal.constants import W3C_TRACESTATE_KEY
from ddtrace.internal.logger import get_logger
from ddtrace.internal.native._native import ContextData
from ddtrace.internal.threads import RLock
from ddtrace.internal.utils.http import w3c_get_dd_list_member as _w3c_get_dd_list_member


_ContextState = tuple[
    Optional[int],  # trace_id
    Optional[int],  # span_id
    dict[str, str],  # _meta
    dict[str, NumericType],  # _metrics
    list[SpanLink],  # span_links
    dict[str, Any],  # baggage
    bool,  # is_remote
    bool,  # _reactivate
]


_DD_ORIGIN_INVALID_CHARS_REGEX = re.compile(r"[^\x20-\x7E]+")

log = get_logger(__name__)


class Context(ContextData):
    """Represents the state required to propagate a trace across execution
    boundaries.
    """

    __slots__ = [
        "_lock",
        "__weakref__",
    ]

    def __init__(
        self,
        trace_id: Optional[int] = None,
        span_id: Optional[int] = None,
        dd_origin: Optional[str] = None,
        sampling_priority: Optional[float] = None,
        meta: Optional[dict[str, str]] = None,
        metrics: Optional[dict[str, NumericType]] = None,
        lock: Optional[RLock] = None,
        span_links: Optional[list[SpanLink]] = None,
        baggage: Optional[dict[str, Any]] = None,
        is_remote: bool = True,
    ):
        # ContextData.__new__ already populated trace_id/span_id/_meta/_metrics/
        # _baggage/_span_links/_is_remote/_reactivate.
        if dd_origin is not None and _DD_ORIGIN_INVALID_CHARS_REGEX.search(dd_origin) is None:
            self._meta[_ORIGIN_KEY] = dd_origin
        if sampling_priority is not None:
            self._metrics[_SAMPLING_PRIORITY_KEY] = sampling_priority

        if lock is not None:
            self._lock = lock
        else:
            # DEV: A `forksafe.RLock` is not necessary here since Contexts
            # are recreated by the tracer after fork
            # https://github.com/DataDog/dd-trace-py/blob/a1932e8ddb704d259ea8a3188d30bf542f59fd8d/ddtrace/tracer.py#L489-L508
            self._lock = RLock()

    def __getstate__(self) -> _ContextState:
        return (
            self.trace_id,
            self.span_id,
            self._meta,
            self._metrics,
            self._span_links,
            self._baggage,
            self._is_remote,
            self._reactivate,
            # Note: self._lock is not serializable
        )

    def __setstate__(self, state: _ContextState) -> None:
        (
            self.trace_id,
            self.span_id,
            self._meta,
            self._metrics,
            self._span_links,
            self._baggage,
            self._is_remote,
            self._reactivate,
        ) = state
        # We cannot serialize and lock, so we must recreate it unless we already have one
        self._lock = RLock()

    def __enter__(self) -> "Context":
        self._lock.acquire()
        return self

    def __exit__(self, *args: Any) -> None:
        self._lock.release()

    @property
    def sampling_priority(self) -> 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: Optional[NumericType]) -> None:
        with self._lock:
            if value is None:
                if _SAMPLING_PRIORITY_KEY in self._metrics:
                    del self._metrics[_SAMPLING_PRIORITY_KEY]
                return
            self._metrics[_SAMPLING_PRIORITY_KEY] = value

    @property
    def _traceparent(self) -> str:
        tp = self._meta.get(W3C_TRACEPARENT_KEY)
        if self.span_id is None or self.trace_id is None:
            # if we only have a traceparent then we'll forward it
            # if we don't have a span id or trace id value we can't build a valid traceparent
            return tp or ""

        # determine the trace_id value
        if tp:
            # grab the original traceparent trace id, not the converted value
            trace_id = tp.split("-")[1]
        else:
            trace_id = f"{self.trace_id:032x}"

        return f"00-{trace_id}-{self.span_id:016x}-{self._traceflags}"

    @property
    def _traceflags(self) -> str:
        return "01" if self.sampling_priority and self.sampling_priority > 0 else "00"

    @property
    def _tracestate(self) -> str:
        dd_list_member = _w3c_get_dd_list_member(self)

        # if there's a preexisting tracestate we need to update it to preserve other vendor data
        ts = self._meta.get(W3C_TRACESTATE_KEY, "")
        if ts and dd_list_member:
            # cut out the original dd list member from tracestate so we can replace it with the new one we created
            ts_w_out_dd = re.sub("dd=(.+?)(?:,|$)", "", ts)
            if ts_w_out_dd:
                ts = f"dd={dd_list_member},{ts_w_out_dd}"
            else:
                ts = f"dd={dd_list_member}"
        # if there is no original tracestate value then tracestate is just the dd list member we created
        elif dd_list_member:
            ts = f"dd={dd_list_member}"
        return ts

    @property
    def dd_origin(self) -> Optional[Text]:
        """Get the origin of the trace."""
        return self._meta.get(_ORIGIN_KEY)

    @dd_origin.setter
    def dd_origin(self, value: Optional[Text]) -> None:
        """Set the origin of the trace."""
        with self._lock:
            if value is None:
                if _ORIGIN_KEY in self._meta:
                    del self._meta[_ORIGIN_KEY]
                return
            self._meta[_ORIGIN_KEY] = value

    @property
    def dd_user_id(self) -> Optional[Text]:
        """Get the user ID of the trace."""
        user_id = self._meta.get(_USER_ID_KEY)
        if user_id:
            return str(base64.b64decode(user_id), encoding="utf-8")
        return None

    @dd_user_id.setter
    def dd_user_id(self, value: Optional[Text]) -> None:
        """Set the user ID of the trace."""
        with self._lock:
            if value is None:
                if _USER_ID_KEY in self._meta:
                    del self._meta[_USER_ID_KEY]
                return
            self._meta[_USER_ID_KEY] = str(base64.b64encode(bytes(value, encoding="utf-8")), encoding="utf-8")

    @property
    def _trace_id_64bits(self) -> Optional[int]:
        """Return the trace ID as a 64-bit value."""
        if self.trace_id is None:
            return None
        else:
            return _MAX_UINT_64BITS & self.trace_id

    def set_baggage_item(self, key: str, value: Any) -> None:
        """Sets a baggage item in this span context.
        Note that this operation mutates the baggage of this span context
        """
        with self._lock:
            self._baggage[key] = value

    def copy(self, trace_id: int, span_id: int) -> "Context":
        """Return a shallow copy of the context with the given correlation IDs."""
        # PERF: copy() is run once per child span. Construct via ContextData.__new__
        # directly, passing the shared _meta/_metrics/_baggage references straight into
        # native construction, to avoid the overhead of __init__'s dd_origin regex and
        # other processing. This optimization holds true if we trust that this data has
        # been validated already.
        with self._lock:
            ctx = Context.__new__(
                Context,
                trace_id=trace_id,
                span_id=span_id,
                meta=self._meta,
                metrics=self._metrics,
                baggage=self._baggage,
                is_remote=False,
            )
        ctx._lock = self._lock
        return ctx

    def _with_baggage_item(self, key: str, value: Any) -> "Context":
        """Returns a copy of this span with a new baggage item.
        Useful for instantiating new child span contexts.
        """
        new_baggage = dict(self._baggage)
        new_baggage[key] = value

        ctx = self.__class__(trace_id=self.trace_id, span_id=self.span_id)
        ctx._meta = self._meta
        ctx._metrics = self._metrics
        ctx._baggage = new_baggage
        return ctx

    def get_baggage_item(self, key: str) -> Optional[Any]:
        """Gets a baggage item in this span context."""
        return self._baggage.get(key, None)

    def get_all_baggage_items(self) -> dict[str, Any]:
        """Returns all baggage items in this span context."""
        return self._baggage

    def remove_baggage_item(self, key: str) -> None:
        """Remove a baggage item from this span context."""
        with self._lock:
            if key in self._baggage:
                del self._baggage[key]

    def remove_all_baggage_items(self) -> None:
        """Removes all baggage items from this span context."""
        with self._lock:
            self._baggage.clear()

    def __eq__(self, other: Any) -> bool:
        # NOTE: span_id/_reactivate are deliberately excluded. A Context compares
        # equal to any per-span copy() of itself (differing only in span_id) --
        # e.g. Span.context builds one such copy per child span -- so equality
        # here means "same trace-level state", not "same span".
        if isinstance(other, Context):
            with self._lock:
                return (
                    self.trace_id == other.trace_id
                    and self._meta == other._meta
                    and self._metrics == other._metrics
                    and self._span_links == other._span_links
                    and self._baggage == other._baggage
                    and self._is_remote == other._is_remote
                )
        return False

    def __repr__(self) -> str:
        return (
            f"Context(trace_id={self.trace_id}, span_id={self.span_id}, _meta={self._meta}, "
            f"_metrics={self._metrics}, _span_links={self._span_links}, _baggage={self._baggage}, "
            f"_is_remote={self._is_remote})"
        )

    def __hash__(self) -> int:
        return hash(self.trace_id)