From 331d45bb5edabf14152f28cb49c4abc6f974c985 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 08:04:07 -0700 Subject: [PATCH 01/29] reexport moved names to simplify diff --- ddtrace/internal/_exceptions.py | 38 +- ddtrace/internal/compat.py | 126 +-- ddtrace/internal/constants.py | 160 +--- ddtrace/internal/logger.py | 247 +----- ddtrace/internal/module.py | 776 +----------------- ddtrace/internal/schema.py | 1 + ddtrace/internal/threads.py | 221 +---- ddtrace/internal/utils/_exceptions.py | 37 + ddtrace/internal/utils/compat.py | 125 +++ ddtrace/internal/utils/constants.py | 159 ++++ ddtrace/internal/utils/logger.py | 246 ++++++ ddtrace/internal/utils/module.py | 771 +++++++++++++++++ .../internal/{ => utils}/schema/__init__.py | 0 .../internal/{ => utils}/schema/processor.py | 0 .../schema/span_attribute_schema.py | 2 +- ddtrace/internal/utils/threads.py | 218 +++++ .../internal/{ => utils}/wrapping/__init__.py | 147 ++-- .../internal/{ => utils}/wrapping/asyncs.py | 0 .../internal/{ => utils}/wrapping/context.py | 592 ++++++------- .../{ => utils}/wrapping/generators.py | 0 ddtrace/internal/wrapping.py | 1 + 21 files changed, 1891 insertions(+), 1976 deletions(-) create mode 100644 ddtrace/internal/schema.py create mode 100644 ddtrace/internal/utils/_exceptions.py create mode 100644 ddtrace/internal/utils/compat.py create mode 100644 ddtrace/internal/utils/constants.py create mode 100644 ddtrace/internal/utils/logger.py create mode 100644 ddtrace/internal/utils/module.py rename ddtrace/internal/{ => utils}/schema/__init__.py (100%) rename ddtrace/internal/{ => utils}/schema/processor.py (100%) rename ddtrace/internal/{ => utils}/schema/span_attribute_schema.py (98%) create mode 100644 ddtrace/internal/utils/threads.py rename ddtrace/internal/{ => utils}/wrapping/__init__.py (76%) rename ddtrace/internal/{ => utils}/wrapping/asyncs.py (100%) rename ddtrace/internal/{ => utils}/wrapping/context.py (50%) rename ddtrace/internal/{ => utils}/wrapping/generators.py (100%) create mode 100644 ddtrace/internal/wrapping.py diff --git a/ddtrace/internal/_exceptions.py b/ddtrace/internal/_exceptions.py index e2836e036be..73de4a62a71 100644 --- a/ddtrace/internal/_exceptions.py +++ b/ddtrace/internal/_exceptions.py @@ -1,37 +1 @@ -from typing import Optional -from typing import TypeVar - - -class DDBlockException(BaseException): - """ - Base class for any in-tree decision to abort the current operation - (web request blocking, AI Guard policy abort, future product blocks). - Inherits from BaseException so a generic ``except Exception:`` handler in - user code does not accidentally swallow a blocking decision. - """ - - -class BlockingException(DDBlockException): - """ - Exception raised when a request is blocked by ASM - It derives from BaseException (via DDBlockException) to avoid being caught by the general Exception handler - """ - - -E = TypeVar("E", bound=BaseException) - - -def find_exception( - exc: BaseException, - exception_type: type[E], -) -> Optional[E]: - """Traverse an exception and its children to find the first occurrence of a specific exception type.""" - if isinstance(exc, exception_type): - return exc - # The check matches both native Python3.11+ and `exceptiongroup` compatibility package versions of ExceptionGroup - if exc.__class__.__name__ in ("BaseExceptionGroup", "ExceptionGroup") and hasattr(exc, "exceptions"): - for sub_exc in exc.exceptions: - found = find_exception(sub_exc, exception_type) - if found: - return found - return None +from ddtrace.internal.utils._exceptions import * # noqa diff --git a/ddtrace/internal/compat.py b/ddtrace/internal/compat.py index a0a122bd980..8463b28f23e 100644 --- a/ddtrace/internal/compat.py +++ b/ddtrace/internal/compat.py @@ -1,125 +1 @@ -import ipaddress -import sys -from types import TracebackType -from typing import Any -from typing import Optional # noqa:F401 -from typing import Text # noqa:F401 -from typing import Union # noqa:F401 - -import wrapt - - -__all__ = [ - "maybe_stringify", -] - -PYTHON_VERSION_INFO = sys.version_info - - -def ensure_text(s, encoding="utf-8", errors="ignore") -> str: - if isinstance(s, str): - return s - if isinstance(s, bytes): - return s.decode(encoding, errors) - raise TypeError("Expected str or bytes but received %r" % (s.__class__)) - - -def ensure_binary(s, encoding="utf-8", errors="ignore") -> bytes: - if isinstance(s, bytes): - return s - if not isinstance(s, str): - raise TypeError("Expected str or bytes but received %r" % (s.__class__)) - return s.encode(encoding, errors) - - -NumericType = Union[int, float] - - -def is_integer(obj: Any) -> bool: - """Helper to determine if the provided ``obj`` is an integer type or not""" - # DEV: We have to make sure it is an integer and not a boolean - # >>> type(True) - # - # >>> isinstance(True, int) - # True - return isinstance(obj, int) and not isinstance(obj, bool) - - -def maybe_stringify(obj: Any) -> Optional[str]: - if obj is not None: - return str(obj) - return None - - -ExcInfoType = Union[tuple[type[BaseException], BaseException, Optional[TracebackType]], tuple[None, None, None]] - -# Sentinel value to represent "no exception" -NO_EXCEPTION: ExcInfoType = (None, None, None) - - -def is_valid_ip(ip: str) -> bool: - try: - # try parsing the IP address - ipaddress.ip_address(str(ip)) - return True - except Exception: - return False - - -def ip_is_global(ip: str) -> bool: - """ - is_global is Python 3+ only. This could raise a ValueError if the IP is not valid. - """ - parsed_ip = ipaddress.ip_address(str(ip)) - - return parsed_ip.is_global - - -# This fix was implemented in 3.9.8 -# https://github.com/python/cpython/issues/83860 -if PYTHON_VERSION_INFO >= (3, 9, 8): - from functools import singledispatchmethod -else: - from functools import singledispatchmethod - - def _register(self, cls, method=None): - if hasattr(cls, "__func__"): - setattr(cls, "__annotations__", cls.__func__.__annotations__) - return self.dispatcher.register(cls, func=method) - - singledispatchmethod.register = _register # type: ignore[method-assign] - - -def get_mp_context(): - import multiprocessing - - return multiprocessing.get_context("fork" if sys.platform != "win32" else "spawn") - - -def __getattr__(name: str) -> Any: - # These attributes are expensive to pre-compute, so we make them lazy - if name == "PYTHON_VERSION": - from platform import python_version - - globals()[name] = python_version() - - elif name == "PYTHON_INTERPRETER": - from platform import python_implementation - - globals()[name] = python_implementation() - - try: - return globals()[name] - except KeyError: - raise AttributeError(f"module '{__name__}' has no attribute '{name}'") - - -if hasattr(wrapt, "BaseObjectProxy"): - # This must be used for wrapt version >= 2.0.0 - wrapt_class: type = wrapt.BaseObjectProxy -else: - wrapt_class = wrapt.ObjectProxy - - -def is_wrapted(o: object) -> bool: - return isinstance(o, wrapt_class) +from ddtrace.internal.utils.compat import * # noqa diff --git a/ddtrace/internal/constants.py b/ddtrace/internal/constants.py index 1167f75f9e6..14a047f6b24 100644 --- a/ddtrace/internal/constants.py +++ b/ddtrace/internal/constants.py @@ -1,159 +1 @@ -from ddtrace.constants import AUTO_KEEP -from ddtrace.constants import AUTO_REJECT -from ddtrace.constants import USER_KEEP -from ddtrace.constants import USER_REJECT - - -PROPAGATION_STYLE_DATADOG = "datadog" -PROPAGATION_STYLE_B3_MULTI = "b3multi" -PROPAGATION_STYLE_B3_SINGLE = "b3" -_PROPAGATION_STYLE_W3C_TRACECONTEXT = "tracecontext" -_PROPAGATION_STYLE_NONE = "none" -_PROPAGATION_STYLE_DEFAULT = "datadog,tracecontext,baggage" -_PROPAGATION_STYLE_BAGGAGE = "baggage" -PROPAGATION_STYLE_ALL = ( - _PROPAGATION_STYLE_W3C_TRACECONTEXT, - PROPAGATION_STYLE_DATADOG, - PROPAGATION_STYLE_B3_MULTI, - PROPAGATION_STYLE_B3_SINGLE, - _PROPAGATION_STYLE_NONE, - _PROPAGATION_STYLE_BAGGAGE, -) -_PROPAGATION_BEHAVIOR_CONTINUE = "continue" -_PROPAGATION_BEHAVIOR_IGNORE = "ignore" -_PROPAGATION_BEHAVIOR_RESTART = "restart" -_PROPAGATION_BEHAVIOR_DEFAULT = _PROPAGATION_BEHAVIOR_CONTINUE -W3C_TRACESTATE_KEY = "tracestate" -W3C_TRACEPARENT_KEY = "traceparent" -W3C_TRACESTATE_PARENT_ID_KEY = "p" -W3C_TRACESTATE_ORIGIN_KEY = "o" -W3C_TRACESTATE_SAMPLING_PRIORITY_KEY = "s" -DEFAULT_SAMPLING_RATE_LIMIT = 100 -SAMPLING_HASH_MODULO = 1 << 64 -# Big prime number to make hashing better distributed, it has to be the same factor as the Agent -# and other tracers to allow chained sampling -SAMPLING_KNUTH_FACTOR = 1111111111111111111 -SAMPLING_DECISION_TRACE_TAG_KEY = "_dd.p.dm" -LAST_DD_PARENT_ID_KEY = "_dd.parent_id" -DEFAULT_SERVICE_NAME = "unnamed-python-service" -# Used to set the name of an integration on a span -COMPONENT = "component" -HIGHER_ORDER_TRACE_ID_BITS = "_dd.p.tid" -MAX_UINT_64BITS = (1 << 64) - 1 -MIN_INT_64BITS = -(2**63) -MAX_INT_64BITS = 2**63 - 1 -SAMPLING_DECISION_MAKER_INHERITED = "_dd.dm.inherited" -SAMPLING_DECISION_MAKER_SERVICE = "_dd.dm.service" -SAMPLING_DECISION_MAKER_RESOURCE = "_dd.dm.resource" -SPAN_LINK_KIND = "dd.kind" -SPAN_LINKS_KEY = "_dd.span_links" -SPAN_EVENTS_KEY = "events" -SPAN_API_DATADOG = "datadog" -SPAN_API_OTEL = "otel" -SPAN_API_OPENTRACING = "opentracing" -DEFAULT_BUFFER_SIZE = 20 << 20 # 20 MB -DEFAULT_MAX_PAYLOAD_SIZE = 20 << 20 # 20 MB -DEFAULT_PROCESSING_INTERVAL = 1.0 -DEFAULT_REUSE_CONNECTIONS = False -BLOCKED_RESPONSE_HTML = """You've been blocked

Sorry, you cannot access this page. Please contact the customer service team.

Security Response ID: [security_response_id]

""" # noqa: E501 -BLOCKED_RESPONSE_JSON = """{"errors":[{"title":"You've been blocked","detail":"Sorry, you cannot access this page. Please contact the customer service team. Security provided by Datadog."}],"security_response_id":"[security_response_id]"}""" # noqa: E501 -HTTP_REQUEST_BLOCKED = "http.request.blocked" -RESPONSE_HEADERS = "http.response.headers" -HTTP_REQUEST_QUERY = "http.request.query" -HTTP_REQUEST_COOKIE_VALUE = "http.request.cookie.value" -HTTP_REQUEST_COOKIE_NAME = "http.request.cookie.name" -HTTP_REQUEST_PATH = "http.request.path" -HTTP_REQUEST_HEADER_NAME = "http.request.header.name" -HTTP_REQUEST_HEADER = "http.request.header" -HTTP_REQUEST_PARAMETER = "http.request.parameter" -HTTP_REQUEST_BODY = "http.request.body" -HTTP_REQUEST_UPGRADED = "http.upgraded" -HTTP_REQUEST_PATH_PARAMETER = "http.request.path.parameter" -REQUEST_PATH_PARAMS = "http.request.path_params" -STATUS_403_TYPE_AUTO = {"status_code": 403, "type": "auto"} -PROCESS_TAGS = "_dd.tags.process" -PROPAGATED_HASH = "_dd.propagated_hash" -_SERVICE_SOURCE = "_dd.svc_src" - -CONTAINER_ID_HEADER_NAME = "Datadog-Container-Id" -CONTAINER_TAGS_HASH = "Datadog-Container-Tags-Hash" - -ENTITY_ID_HEADER_NAME = "Datadog-Entity-ID" - -EXTERNAL_ENV_HEADER_NAME = "Datadog-External-Env" -EXTERNAL_ENV_ENVIRONMENT_VARIABLE = "DD_EXTERNAL_ENV" - -MESSAGING_BATCH_COUNT = "messaging.batch_count" -MESSAGING_DESTINATION_NAME = "messaging.destination.name" -MESSAGING_MESSAGE_ID = "messaging.message_id" -MESSAGING_OPERATION = "messaging.operation" -MESSAGING_SYSTEM = "messaging.system" - -USER_AGENT_HEADER = "user-agent" -FLASK_ENDPOINT = "flask.endpoint" -FLASK_VIEW_ARGS = "flask.view_args" -FLASK_URL_RULE = "flask.url_rule" -FLASK_RESOURCE_FULL = "flask.resource.full" - -_HTTPLIB_NO_TRACE_REQUEST = "_dd_no_trace" -DEFAULT_TIMEOUT = 2.0 - -# baggage -DD_TRACE_BAGGAGE_MAX_ITEMS = 64 -DD_TRACE_BAGGAGE_MAX_BYTES = 8192 -BAGGAGE_TAG_PREFIX = "baggage." - -# W3C Trace Context tracestate (https://www.w3.org/TR/trace-context/): -# max 32 list-members; vendors SHOULD propagate at most 512 characters (we cap parsing to that size). -DD_TRACE_TRACESTATE_MAX_ITEMS = 32 -DD_TRACE_TRACESTATE_MAX_BYTES = 512 -# Per W3C Trace Context, oversized list-members are preferred targets when truncating by size. -DD_TRACE_TRACESTATE_ITEM_MAX_CHARS = 128 - -SPAN_EVENTS_HAS_EXCEPTION = "_dd.span_events.has_exception" -COLLECTOR_MAX_SIZE_PER_SPAN = 100 - -LOG_ATTR_TRACE_ID = "dd.trace_id" -LOG_ATTR_SPAN_ID = "dd.span_id" -LOG_ATTR_ENV = "dd.env" -LOG_ATTR_VERSION = "dd.version" -LOG_ATTR_SERVICE = "dd.service" -LOG_ATTR_VALUE_ZERO = "0" -LOG_ATTR_VALUE_EMPTY = "" - - -class SamplingMechanism(object): - DEFAULT = 0 - AGENT_RATE_BY_SERVICE = 1 - REMOTE_RATE = 2 # not used, this mechanism is deprecated - LOCAL_USER_TRACE_SAMPLING_RULE = 3 - MANUAL = 4 - APPSEC = 5 - REMOTE_RATE_USER = 6 # not used, this mechanism is deprecated - REMOTE_RATE_DATADOG = 7 # not used, this mechanism is deprecated - SPAN_SAMPLING_RULE = 8 - OTLP_INGEST_PROBABILISTIC_SAMPLING = 9 # not used in ddtrace - DATA_JOBS_MONITORING = 10 # not used in ddtrace - REMOTE_USER_TRACE_SAMPLING_RULE = 11 - REMOTE_DYNAMIC_TRACE_SAMPLING_RULE = 12 - AI_GUARD = 13 - - -SAMPLING_MECHANISM_TO_PRIORITIES = { - # TODO(munir): Update mapping to include single span sampling and appsec sampling mechanisms - SamplingMechanism.AGENT_RATE_BY_SERVICE: (AUTO_KEEP, AUTO_REJECT), - SamplingMechanism.DEFAULT: (AUTO_KEEP, AUTO_REJECT), - SamplingMechanism.MANUAL: (USER_KEEP, USER_REJECT), - SamplingMechanism.APPSEC: (AUTO_KEEP, AUTO_REJECT), - SamplingMechanism.LOCAL_USER_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), - SamplingMechanism.REMOTE_USER_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), - SamplingMechanism.REMOTE_DYNAMIC_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), -} -_KEEP_PRIORITY_INDEX = 0 -_REJECT_PRIORITY_INDEX = 1 - - -# List of support values in DD_TRACE_EXPERIMENTAL_FEATURES_ENABLED -class EXPERIMENTAL_FEATURES: - # Enables submitting runtime metrics as gauges (instead of distributions) - RUNTIME_METRICS = "DD_RUNTIME_METRICS_ENABLED" +from ddtrace.internal.utils.constants import * # noqa diff --git a/ddtrace/internal/logger.py b/ddtrace/internal/logger.py index c57a527650f..18e296e4630 100644 --- a/ddtrace/internal/logger.py +++ b/ddtrace/internal/logger.py @@ -1,246 +1 @@ -""" -Logging utilities for internal use. -Usage: - import ddtrace.internal.logger as logger - ddlog = logger.get_logger(__name__) - - # Otherwise default is set to 1 minute or DD_TRACE_LOGGING_RATE - logger.set_tag_rate_limit("waf::init", logger.HOUR) - - # "product" is required, but other keys are optional as well as kwargs exc_info and stack_info - # supported keys are - # product: product or integration name. Required - # more_info : more information to be logged after the main tag. Default is empty string - # stack_limit: limit the stack trace to this depth for stack_info. Default is 0 (0 is no limit) - # exec_limit: limit the stack trace to this depth for exec_info. Default is 1, only the top level (0 is no limit) - - info = # format the info string - ddlog.debug('waf::init', extra={"product": "appsec", "stack_limit": 4, "more_info": info}, stack_info=True) - # This will log the message only once per hour, counting the number of skipped messages, using "waf::init" as the - # tag to keep track of the rate limit - # Different log levels can be used, the rate limit is shared between all invocations and all levels for the same tag - - # example result - DEBUG appsec::waf::init[some more info] 1 additional messages skipped - (followed by the 4 first levels of the stack trace) - - Legacy support: - if extra is not used or product is absent, the log will be treated as legacy and will be logged as is using - filename and line number of the log call - -""" - -import collections -from dataclasses import dataclass -from dataclasses import field -import logging -import time -import traceback -from typing import DefaultDict -from typing import Optional -from typing import Tuple -from typing import Union - -from ddtrace.internal.settings import env - - -SECOND = 1 -MINUTE = 60 * SECOND -HOUR = 60 * MINUTE -DAY = 24 * HOUR - - -@dataclass -class LoggerPrefix: - prefix: str - level: Optional[int] = None - children: dict[str, "LoggerPrefix"] = field(default_factory=dict) - - def lookup(self, name: str) -> Optional[int]: - """ - Lookup the log level for a given logger name in the trie. - - The name is split by '.' and each part is used to traverse the trie. - If a part is not found, it returns the level of the closest parent node. - """ - parts = name.replace("_", ".").lower().split(".") - parts.pop(0) # remove the ddtrace prefix - current = self - while parts: - if (part := parts.pop(0)) not in current.children: - return current.level - current = current.children[part] - - return current.level - - @classmethod - def build_trie(cls): - trie = cls(prefix="ddtrace", level=None, children={}) - - for logger_name, level in ((k, v) for k, v in env.items() if k.startswith("_DD_") and k.endswith("_LOG_LEVEL")): - # Remove the _DD_ prefix and _LOG_LEVEL suffix - logger_name = logger_name[4:-10] - parts = logger_name.lower().split("_") - current = trie.children - while parts: - if (part := parts.pop(0)) not in current: - current[part] = cls(prefix=part, level=getattr(logging, level, None) if not parts else None) - current = current[part].children - - return trie - - -LOG_LEVEL_TRIE = LoggerPrefix.build_trie() - - -def get_logger(name: str) -> logging.Logger: - """ - Retrieve or create a ``Logger`` instance with consistent behavior for internal use. - - Configure all loggers with a rate limiter filter to prevent excessive logging. - - """ - logger = logging.getLogger(name) - # addFilter will only add the filter if it is not already present - logger.addFilter(log_filter) - - # Set the log level from the environment variable of the closest parent - # logger. - if name.startswith("ddtrace."): # for the whole of ddtrace we have DD_TRACE_DEBUG - if (level := LOG_LEVEL_TRIE.lookup(name)) is not None: - logger.setLevel(level) - - return logger - - -_RATE_LIMITS = {} - - -def set_tag_rate_limit(tag: str, rate: int) -> None: - """ - Set the rate limit for a specific tag. - - """ - _RATE_LIMITS[tag] = rate - - -# Class used for keeping track of a log lines current time bucket and the number of log lines skipped -class LoggingBucket: - def __init__(self, bucket: float, skipped: int): - self.bucket = bucket - self.skipped = skipped - - def __repr__(self): - return f"LoggingBucket({self.bucket}, {self.skipped})" - - def is_sampled(self, record: logging.LogRecord, rate: float) -> bool: - """ - Determine if the log line should be sampled based on the rate limit. - """ - current = time.monotonic() - if current - self.bucket >= rate: - self.bucket = current - record.skipped = self.skipped - self.skipped = 0 - return True - self.skipped += 1 - return False - - -# dict to keep track of the current time bucket per name/level/pathname/lineno - -_MINF = float("-inf") - -# IMPORTANT: Do not change typing types to built-ins until minimum Python version is 3.11+ -# Module-level tuple[...] and defaultdict[...] in Python 3.10 affect import timing. See packages.py for details. -key_type = Union[Tuple[str, int, str, int], str] # noqa: UP006 -_buckets: DefaultDict[key_type, LoggingBucket] = collections.defaultdict(lambda: LoggingBucket(_MINF, 0)) # noqa: UP006 - -# Allow 1 log record per name/level/pathname/lineno every 60 seconds by default -# Allow configuring via `DD_TRACE_LOGGING_RATE` -# DEV: `DD_TRACE_LOGGING_RATE=0` means to disable all rate limiting -_rate_limit = int(env.get("DD_TRACE_LOGGING_RATE", default=60)) - - -def log_filter(record: logging.LogRecord) -> bool: - """ - Function used to determine if a log record should be outputted or not (True = output, False = skip). - - This function will: - - Rate limit log records based on the logger name, record level, filename, and line number - """ - logger = logging.getLogger(record.name) - rate_limit = _RATE_LIMITS.get(record.msg, _rate_limit) - # If rate limiting has been disabled (`DD_TRACE_LOGGING_RATE=0`) then apply no rate limit - # If the logger is set to debug, then do not apply any limits to any log - if not rate_limit or logger.getEffectiveLevel() == logging.DEBUG: - must_be_propagated = True - else: - # Allow 1 log record by pathname/lineno every X seconds or message/levelno for product logs - # This way each unique log message can get logged at least once per time period - if hasattr(record, "product"): - key: key_type = record.msg - else: - key = (record.name, record.levelno, record.pathname, record.lineno) - # If rate limiting has been disabled (`DD_TRACE_LOGGING_RATE=0`) then apply no rate limit - # If the logger is set to debug, then do not apply any limits to any log - # Only log this message if the time bucket allows it - must_be_propagated = _buckets[key].is_sampled(record, rate_limit) - if must_be_propagated: - skipped = record.__dict__.pop("skipped", 0) - if skipped: - skip_str = f" [{skipped} skipped]" - else: - skip_str = "" - product = record.__dict__.pop("product", None) - # new syntax - if product: - more_info = record.__dict__.pop("more_info", "") - stack_limit = record.__dict__.pop("stack_limit", 0) - exec_limit = record.__dict__.pop("exec_limit", 1) - # format the stacks if they are present with the right depth - if stack_limit and record.stack_info: - record.stack_info = format_stack(record.stack_info, stack_limit) - string_buffer = [f"{product}::{record.msg}{more_info}{skip_str}"] - if record.stack_info: - string_buffer.append(record.stack_info) - if record.exc_info: - string_buffer.extend(traceback.format_exception(record.exc_info[1], limit=exec_limit or None)) - record.msg = "\n".join(string_buffer) - # clean the record for any subsequent handlers - record.stack_info = None - record.exc_info = None - else: - record.msg = f"{record.msg}{skip_str}" - return must_be_propagated - - -def format_stack(stack_info, limit) -> str: - stack = stack_info.split("\n") - if len(stack) <= limit * 2 + 1: - return stack_info - stack_str = "\n".join(stack[-2 * limit :]) - return f"{stack[0]}\n{stack_str}" - - -class LogInjectionState(object): - # Log injection is disabled - DISABLED = "false" - # Log injection is enabled, but not yet configured - ENABLED = "true" - # Log injection is enabled and configured for structured logging - # This value is deprecated, but kept for backwards compatibility - STRUCTURED = "structured" - - -def get_log_injection_state(raw_config: Optional[str]) -> bool: - if raw_config: - normalized = raw_config.lower().strip() - if normalized == LogInjectionState.STRUCTURED or normalized in ("true", "1"): - return True - elif normalized not in ("false", "0"): - logging.warning( - "Invalid log injection state '%s'. Expected 'true', 'false', or 'structured'. Defaulting to 'false'.", - normalized, - ) - return False +from ddtrace.internal.utils.logger import * # noqa diff --git a/ddtrace/internal/module.py b/ddtrace/internal/module.py index 576987aa975..b45ec5e1175 100644 --- a/ddtrace/internal/module.py +++ b/ddtrace/internal/module.py @@ -1,775 +1 @@ -import abc -from collections import defaultdict -from importlib._bootstrap import _init_module_attrs -from importlib.machinery import ModuleSpec -from importlib.util import find_spec -from importlib.util import spec_from_loader -from pathlib import Path -import sys -from types import CodeType -from types import FunctionType -from types import ModuleType -import typing as t -from weakref import WeakValueDictionary as wvdict - -from ddtrace.internal.logger import get_logger -from ddtrace.internal.utils import get_argument_value -from ddtrace.internal.wrapping.context import WrappingContext - - -if t.TYPE_CHECKING: - from importlib.abc import Loader - - -ModuleHookType = t.Callable[[ModuleType], None] -TransformerType = t.Callable[[CodeType, ModuleType], CodeType] -PreExecHookType = t.Callable[[t.Any, ModuleType], None] -PreExecHookCond = t.Union[str, t.Callable[[str], bool]] - -ImportExceptionHookType = t.Callable[[t.Any, ModuleType], None] -ImportExceptionHookCond = t.Union[str, t.Callable[[str], bool]] - - -log = get_logger(__name__) - - -_run_code = None -_run_module_transformers: list[TransformerType] = [] -_post_run_module_hooks: list[ModuleHookType] = [] - - -def _wrapped_run_code(*args: t.Any, **kwargs: t.Any) -> dict[str, t.Any]: - # DEV: If we are calling this wrapper then _run_code must have been set to - # the original runpy._run_code. - assert _run_code is not None - - code = t.cast(CodeType, get_argument_value(args, kwargs, 0, "code")) - mod_name = t.cast(str, get_argument_value(args, kwargs, 3, "mod_name")) - - module = sys.modules[mod_name] - - for transformer in _run_module_transformers: - code = transformer(code, module) - - kwargs.pop("code", None) - new_args = (code, *args[1:]) - - try: - return _run_code(*new_args, **kwargs) - finally: - module = sys.modules[mod_name] - for hook in _post_run_module_hooks: - hook(module) - - -def _patch_run_code() -> None: - global _run_code - - if _run_code is None: - import runpy - - _run_code = runpy._run_code # type: ignore[attr-defined] - runpy._run_code = _wrapped_run_code # type: ignore[attr-defined] - - -def register_run_module_transformer(transformer: TransformerType) -> None: - """Register a run module transformer.""" - _run_module_transformers.append(transformer) - - -def unregister_run_module_transformer(transformer: TransformerType) -> None: - """Unregister a run module transformer. - - If the transformer was not registered, a ``ValueError`` exception is raised. - """ - _run_module_transformers.remove(transformer) - - -def register_post_run_module_hook(hook: ModuleHookType) -> None: - """Register a post run module hook. - - The hooks gets called after the module is loaded. For this to work, the - hook needs to be registered during the interpreter initialization, e.g. as - part of a sitecustomize.py script. - """ - global _run_code, _post_run_module_hooks - - _patch_run_code() - - _post_run_module_hooks.append(hook) - - -def unregister_post_run_module_hook(hook: ModuleHookType) -> None: - """Unregister a post run module hook. - - If the hook was not registered, a ``ValueError`` exception is raised. - """ - global _post_run_module_hooks - - _post_run_module_hooks.remove(hook) - - -def origin(module: ModuleType) -> t.Optional[Path]: - """Get the origin source file of the module.""" - try: - # Do not access __dd_origin__ directly to avoid force-loading lazy - # modules. - return object.__getattribute__(module, "__dd_origin__") - except AttributeError: - try: - # DEV: Use object.__getattribute__ to avoid potential side-effects. - orig = Path(object.__getattribute__(module, "__file__")).resolve() - except (AttributeError, TypeError): - # Module is probably only partially initialised, so we look at its - # spec instead - try: - # DEV: Use object.__getattribute__ to avoid potential side-effects. - orig = Path(object.__getattribute__(module, "__spec__").origin).resolve() - except (AttributeError, ValueError, TypeError): - orig = None - - if orig is not None and orig.suffix == "pyc": - orig = orig.with_suffix(".py") - - if orig is not None: - # If we failed to find a valid origin we don't cache the value and - # try again the next time. - try: - module.__dd_origin__ = orig # type: ignore[attr-defined] - except AttributeError: - pass - - return orig - - -def _resolve(path: Path) -> t.Optional[Path]: - """Resolve a (relative) path with respect to sys.path.""" - for base in (Path(_) for _ in sys.path): - if base.is_dir(): - resolved_path = (base / path.expanduser()).resolve() - if resolved_path.is_file(): - return resolved_path - return None - - -# Borrowed from the wrapt module -# https://github.com/GrahamDumpleton/wrapt/blob/df0e62c2740143cceb6cafea4c306dae1c559ef8/src/wrapt/importer.py - - -def find_loader(fullname: str) -> t.Optional["Loader"]: - return getattr(find_spec(fullname), "loader", None) - - -def is_module_installed(module_name): - return find_loader(module_name) is not None - - -def is_namespace_spec(spec: ModuleSpec) -> bool: - return spec.origin is None and spec.submodule_search_locations is not None - - -class _ImportHookChainedLoader: - def __init__(self, loader: t.Optional["Loader"], spec: t.Optional[ModuleSpec] = None) -> None: - self.loader = loader - self.spec = spec - - self.callbacks: dict[t.Any, t.Callable[[ModuleType], None]] = {} - self.import_exception_callbacks: dict[t.Any, t.Callable[[ModuleType], None]] = {} - - self.transformers: dict[t.Any, TransformerType] = {} - - # A missing loader is generally an indication of a namespace package. - if loader is None or hasattr(loader, "create_module"): - self.create_module = self._create_module - if loader is None or hasattr(loader, "exec_module"): - self.exec_module = self._exec_module - - def __getattr__(self, name): - # Proxy any other attribute access to the underlying loader. - return getattr(self.loader, name) - - def namespace_module(self, spec: ModuleSpec) -> ModuleType: - module = ModuleType(spec.name) - # Pretend that we do not have a loader (this would be self), to - # allow _init_module_attrs to create the appropriate NamespaceLoader - # for the namespace module. - spec.loader = None - - _init_module_attrs(spec, module, override=True) - - # Chain the loaders - self.loader = spec.loader - module.__loader__ = spec.loader = self # type: ignore[assignment] - - return module - - def add_callback(self, key: t.Any, callback: t.Callable[[ModuleType], None]) -> None: - self.callbacks[key] = callback - - def add_import_exception_callback(self, key: t.Any, callback: t.Callable[[ModuleType], None]) -> None: - self.import_exception_callbacks[key] = callback - - def add_transformer(self, key: t.Any, transformer: TransformerType) -> None: - self.transformers[key] = transformer - - def call_back(self, module: ModuleType) -> None: - # Restore the original loader, if possible. Some specs might be native - # and won't support attribute assignment. - try: - module.__loader__ = self.loader - except (AttributeError, TypeError): - pass - try: - object.__getattribute__(module, "spec").loader = self.loader - except (AttributeError, TypeError): - pass - - if module.__name__ == "pkg_resources": - # DEV: pkg_resources support to prevent errors such as - # NotImplementedError: Can't perform this operation for unregistered - # loader type - module.register_loader_type(_ImportHookChainedLoader, module.DefaultProvider) - - for callback in self.callbacks.values(): - callback(module) - - def load_module(self, fullname: str) -> t.Optional[ModuleType]: - if self.loader is None: - if self.spec is None: - return None - sys.modules[self.spec.name] = module = self.namespace_module(self.spec) - else: - module = self.loader.load_module(fullname) - - try: - self.call_back(module) - except Exception: - log.exception("Failed to call back on module %s", module) - - return module - - def _create_module(self, spec): - if self.loader is not None: - return self.loader.create_module(spec) - - if is_namespace_spec(spec): - return self.namespace_module(spec) - - return None - - def _find_first_hook( - self, module: ModuleType, hooks_attr: str - ) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: - for _ in sys.meta_path: - if isinstance(_, ModuleWatchdog): - try: - for ( - cond, - hook, - ) in getattr(_, hooks_attr, []): - if (isinstance(cond, str) and cond == module.__name__) or ( - callable(cond) and cond(module.__name__) - ): - return hook - except Exception: - log.debug("Exception happened while processing %s", hooks_attr, exc_info=True) - return None - - def _find_first_exception_hook(self, module: ModuleType) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: - return self._find_first_hook(module, "_import_exception_hooks") - - def _find_first_pre_exec_hook(self, module: ModuleType) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: - return self._find_first_hook(module, "_pre_exec_module_hooks") - - def _exec_module(self, module: ModuleType) -> None: - # Collect and run only the first hook that matches the module. - - _get_code = getattr(self.loader, "get_code", None) - # DEV: avoid re-wrapping the loader's get_code method (eg: in case of repeated importlib.reload() calls) - if _get_code is not None and not getattr(_get_code, "_dd_get_code", False): - - def get_code(_loader, fullname): - code = _get_code(fullname) - - for callback in self.transformers.values(): - code = callback(code, module) - - return code - - setattr(get_code, "_dd_get_code", True) - - self.loader.get_code = get_code.__get__(self.loader, type(self.loader)) # type: ignore[union-attr] - - pre_exec_hook = self._find_first_pre_exec_hook(module) - - if pre_exec_hook: - pre_exec_hook(self, module) - else: - if self.loader is None: - spec = getattr(module, "__spec__", None) - if spec is not None and is_namespace_spec(spec): - sys.modules[spec.name] = module - else: - try: - self.loader.exec_module(module) - except Exception: - exception_hook = self._find_first_exception_hook(module) - if exception_hook is not None: - exception_hook(self, module) - - # Hide the chained loader method from the traceback - _, e, tb = sys.exc_info() - if e is not None and tb is not None: - e.__traceback__ = tb.tb_next - - raise - - try: - self.call_back(module) - except Exception: - log.exception("Failed to call back on module %s", module) - - -class BaseModuleWatchdog(abc.ABC): - """Base module watchdog. - - Invokes ``after_import`` every time a new module is imported. - """ - - _instance: t.Optional["BaseModuleWatchdog"] = None - - def __init__(self) -> None: - self._finding: set[str] = set() - - # DEV: pkg_resources support to prevent errors such as - # NotImplementedError: Can't perform this operation for unregistered - pkg_resources = sys.modules.get("pkg_resources") - if pkg_resources is not None: - for _ in range(5): - # The package might be in the process of being imported by - # another thread by the time this module watchdog is created. If - # we fail to find the method we try again after a short wait. - # This is likely to happen when a module watchdog is installed - # lazily after boot, e.g. as a response to a RC enablement - # record. This means that we are unlikely to add delays to the - # main or any other use threads. - try: - pkg_resources.register_loader_type(_ImportHookChainedLoader, pkg_resources.DefaultProvider) - break - except AttributeError: - from time import sleep - - sleep(0.1) - else: - log.warning("Cannot ensure correct support with pkg_resources") - - def _add_to_meta_path(self) -> None: - sys.meta_path.insert(0, self) # type: ignore[arg-type] - - @classmethod - def _find_in_meta_path(cls) -> t.Optional[int]: - for i, meta_path in enumerate(sys.meta_path): - if type(meta_path) is cls: - return i - return None - - @classmethod - def _remove_from_meta_path(cls) -> None: - i = cls._find_in_meta_path() - - if i is None: - log.warning("%s is not installed", cls.__name__) - return - - sys.meta_path.pop(i) - - def after_import(self, module: ModuleType) -> None: - raise NotImplementedError() - - def transform(self, code: CodeType, _module: ModuleType) -> CodeType: - return code - - def find_module(self, fullname: str, path: t.Optional[str] = None) -> t.Optional["Loader"]: - if fullname in self._finding: - return None - - self._finding.add(fullname) - - try: - original_loader = find_loader(fullname) - if original_loader is not None: - loader = ( - _ImportHookChainedLoader(original_loader) - if not isinstance(original_loader, _ImportHookChainedLoader) - else original_loader - ) - - loader.add_callback(type(self), self.after_import) - loader.add_transformer(type(self), self.transform) - - return t.cast("Loader", loader) - - finally: - self._finding.remove(fullname) - - return None - - def find_spec( - self, fullname: str, path: t.Optional[str] = None, target: t.Optional[ModuleType] = None - ) -> t.Optional[ModuleSpec]: - if fullname in self._finding: - return None - - self._finding.add(fullname) - - try: - # Iterate sys.meta_path directly to avoid re-entering the import machinery - # and acquiring per-module locks (A-B deadlock risk when a background thread - # calls importlib.files() during module __init__). - spec = None - for finder in sys.meta_path: - if finder is self: - continue - _find_spec = getattr(finder, "find_spec", None) - if _find_spec is not None: - spec = _find_spec(fullname, path, target) - else: - # Fallback for legacy finders that only implement find_module() - # (deprecated in 3.4, removed in 3.12) — mirrors CPython _find_spec(). - _find_module = getattr(finder, "find_module", None) - if _find_module is None: - continue - loader = _find_module(fullname, path) - spec = spec_from_loader(fullname, loader) if loader is not None else None - if spec is not None: - break - - if spec is None: - return None - - loader = getattr(spec, "loader", None) - - if not isinstance(loader, _ImportHookChainedLoader): - spec.loader = t.cast("Loader", _ImportHookChainedLoader(loader, spec)) - - t.cast(_ImportHookChainedLoader, spec.loader).add_callback(type(self), self.after_import) - t.cast(_ImportHookChainedLoader, spec.loader).add_transformer(type(self), self.transform) - - return spec - - finally: - self._finding.remove(fullname) - - @classmethod - def install(cls) -> None: - """Install the module watchdog.""" - if cls.is_installed(): - return - cls._instance = cls() - cls._instance._add_to_meta_path() - log.debug("%s installed", cls) - - @classmethod - def is_installed(cls) -> bool: - """Check whether this module watchdog class is installed.""" - return cls._instance is not None and type(cls._instance) is cls - - @classmethod - def uninstall(cls) -> None: - """Uninstall the module watchdog. - - This will uninstall only the most recently installed instance of this - class. - """ - if not cls.is_installed(): - return - - cls._remove_from_meta_path() - - cls._instance = None - - log.debug("%s uninstalled", cls) - - -class ModuleWatchdog(BaseModuleWatchdog): - """Module watchdog. - - Hooks into the import machinery to detect when modules are loaded/unloaded. - This is also responsible for triggering any registered import hooks. - - Subclasses might customize the default behavior by overriding the - ``after_import`` method, which is triggered on every module import, once - the subclass is installed. - """ - - def __init__(self) -> None: - super().__init__() - - self._hook_map: defaultdict[str, list[ModuleHookType]] = defaultdict(list) - # DEV: It would make more sense to make this a mapping of Path to ModuleType - # but the WeakValueDictionary causes an ignored exception on shutdown - # because the pathlib module is being garbage collected. - self._om: t.Optional[t.MutableMapping[str, ModuleType]] = None - # _pre_exec_module_hooks is a set of tuples (condition, hook) instead - # of a list to ensure that no hook is duplicated - self._pre_exec_module_hooks: set[tuple[PreExecHookCond, PreExecHookType]] = set() - self._import_exception_hooks: set[tuple[ImportExceptionHookCond, ImportExceptionHookType]] = set() - - @property - def _origin_map(self) -> t.MutableMapping[str, ModuleType]: - def modules_with_origin(modules: t.Iterable[ModuleType]) -> t.MutableMapping[str, ModuleType]: - result: t.MutableMapping[str, ModuleType] = wvdict() - - for m in modules: - module_origin = origin(m) - if module_origin is None: - continue - - try: - result[str(module_origin)] = m - except TypeError: - # This can happen if the module is a special object that - # does not allow for weak references. Quite likely this is - # an object created by a native extension. We make the - # assumption that this module does not contain valuable - # information that can be used at the Python runtime level. - pass - - return result - - if self._om is None: - try: - self._om = modules_with_origin(sys.modules.values()) - except RuntimeError: - # The state of sys.modules might have been mutated by another - # thread. We try to build the full mapping at the next occasion. - # For now we take the more expensive route of building a list of - # the current values, which might be incomplete. - return modules_with_origin(list(sys.modules.values())) - return self._om - - def after_import(self, module: ModuleType) -> None: - module_path = origin(module) - path = str(module_path) if module_path is not None else None - if path is not None: - self._origin_map[path] = module - - # Collect all hooks by module origin and name - hooks = [] - if path is not None and path in self._hook_map: - hooks.extend(self._hook_map[path]) - if module.__name__ in self._hook_map: - hooks.extend(self._hook_map[module.__name__]) - - if hooks: - log.debug("Calling %d registered hooks on import of module '%s'", len(hooks), module.__name__) - for hook in hooks: - hook(module) - - @classmethod - def get_by_origin(cls, _origin: Path) -> t.Optional[ModuleType]: - """Lookup a module by its origin.""" - if not cls.is_installed(): - return None - - instance = t.cast(ModuleWatchdog, cls._instance) - - resolved_path = _resolve(_origin) - if resolved_path is not None: - path = str(resolved_path) - module = instance._origin_map.get(path) - if module is not None: - return module - - # Check if this is the __main__ module - main_module = sys.modules.get("__main__") - if main_module is not None and origin(main_module) == resolved_path: - # Register for future lookups - instance._origin_map[path] = main_module - - return main_module - - return None - - @classmethod - def register_origin_hook(cls, origin: Path, hook: ModuleHookType) -> None: - """Register a hook to be called when the module with the given origin is - imported. - - The hook will be called with the module object as argument. - """ - cls.install() - - # DEV: Under the hypothesis that this is only ever called by the probe - # poller thread, there are no further actions to take. Should this ever - # change, then thread-safety might become a concern. - resolved_path = _resolve(origin) - if resolved_path is None: - log.warning("Cannot resolve module origin %s", origin) - return - - path = str(resolved_path) - - log.debug("Registering hook '%r' on path '%s'", hook, path) - instance = t.cast(ModuleWatchdog, cls._instance) - instance._hook_map[path].append(hook) - try: - module = instance.get_by_origin(resolved_path) - if module is None: - # The path was resolved but we still haven't seen any module - # that has it as origin. Nothing more we can do for now. - return - # Sanity check: the module might have been removed from sys.modules - # but not yet garbage collected. - try: - sys.modules[module.__name__] - except KeyError: - del instance._origin_map[path] - raise - except KeyError: - # The module is not loaded yet. Nothing more we can do. - return - - # The module was already imported so we invoke the hook straight-away - log.debug("Calling hook '%r' on already imported module '%s'", hook, module.__name__) - hook(module) - - @classmethod - def unregister_origin_hook(cls, origin: Path, hook: ModuleHookType) -> None: - """Unregister the hook registered with the given module origin and - argument. - """ - if not cls.is_installed(): - return - - resolved_path = _resolve(origin) - if resolved_path is None: - log.warning("Module origin %s cannot be resolved", origin) - return - - path = str(resolved_path) - - instance = t.cast(ModuleWatchdog, cls._instance) - if path not in instance._hook_map: - log.warning("No hooks registered for origin %s", origin) - return - - try: - if path in instance._hook_map: - hooks = instance._hook_map[path] - hooks.remove(hook) - if not hooks: - del instance._hook_map[path] - except ValueError: - log.warning("Hook %r not registered for origin %s", hook, origin) - return - - @classmethod - def register_module_hook(cls, module: str, hook: ModuleHookType) -> None: - """Register a hook to be called when the module with the given name is - imported. - - The hook will be called with the module object as argument. - """ - cls.install() - - log.debug("Registering hook '%r' on module '%s'", hook, module) - instance = t.cast(ModuleWatchdog, cls._instance) - instance._hook_map[module].append(hook) - try: - module_object = sys.modules[module] - except KeyError: - # The module is not loaded yet. Nothing more we can do. - return - - # The module was already imported so we invoke the hook straight-away - log.debug("Calling hook '%r' on already imported module '%s'", hook, module) - hook(module_object) - - @classmethod - def unregister_module_hook(cls, module: str, hook: ModuleHookType) -> None: - """Unregister the hook registered with the given module name and - argument. - """ - if not cls.is_installed(): - return - - instance = t.cast(ModuleWatchdog, cls._instance) - if module not in instance._hook_map: - log.warning("No hooks registered for module %s", module) - return - - try: - if module in instance._hook_map: - hooks = instance._hook_map[module] - hooks.remove(hook) - if not hooks: - del instance._hook_map[module] - except ValueError: - log.warning("Hook %r not registered for module %r", hook, module) - return - - @classmethod - def after_module_imported(cls, module: str) -> t.Callable[[ModuleHookType], None]: - def _(hook: ModuleHookType) -> None: - cls.register_module_hook(module, hook) - - return _ - - @classmethod - def register_pre_exec_module_hook( - cls: type["ModuleWatchdog"], cond: PreExecHookCond, hook: PreExecHookType - ) -> None: - """Register a hook to execute before/instead of exec_module. - - The pre exec_module hook is executed before the module is executed - to allow for changed modules to be executed as needed. To ensure - that the hook is applied only to the modules that are required, - the condition is evaluated against the module name. - """ - cls.install() - - log.debug("Registering pre_exec module hook '%r' on condition '%s'", hook, cond) - instance = t.cast(ModuleWatchdog, cls._instance) - instance._pre_exec_module_hooks.add((cond, hook)) - - @classmethod - def remove_pre_exec_module_hook(cls: type["ModuleWatchdog"], cond: PreExecHookCond, hook: PreExecHookType) -> None: - """Register a hook to execute before/instead of exec_module. Only for testing proposes""" - instance = t.cast(ModuleWatchdog, cls._instance) - instance._pre_exec_module_hooks.remove((cond, hook)) - - @classmethod - def register_import_exception_hook( - cls: type["ModuleWatchdog"], cond: ImportExceptionHookCond, hook: ImportExceptionHookType - ): - cls.install() - - instance = t.cast(ModuleWatchdog, cls._instance) - instance._import_exception_hooks.add((cond, hook)) - - -class LazyWrappingContext(WrappingContext): - def __return__(self, value: t.Any) -> t.Any: - # Update the global (i.e. the module) scope with the local scope of the - # wrapped function. - self.__frame__.f_globals.update(self.__frame__.f_locals) - - return super().__return__(value) - - -def lazy(f: t.Callable[[], None]) -> None: - LazyWrappingContext(t.cast(FunctionType, f)).wrap() - - _globals = sys._getframe(1).f_globals - - def __getattr__(name: str) -> t.Any: - f() - try: - return _globals[name] - except KeyError: - h = AttributeError(f"module {_globals['__name__']!r} has no attribute {name!r}") - h.__suppress_context__ = True - raise h - - _globals["__getattr__"] = __getattr__ +from ddtrace.internal.utils.module import * # noqa diff --git a/ddtrace/internal/schema.py b/ddtrace/internal/schema.py new file mode 100644 index 00000000000..9ba7325082f --- /dev/null +++ b/ddtrace/internal/schema.py @@ -0,0 +1 @@ +from ddtrace.internal.utils.schema import * # noqa diff --git a/ddtrace/internal/threads.py b/ddtrace/internal/threads.py index 84da0cabfa9..a6666307f10 100644 --- a/ddtrace/internal/threads.py +++ b/ddtrace/internal/threads.py @@ -1,220 +1 @@ -from time import monotonic_ns -import typing as t - -from ddtrace.internal import _threads as _threads_mod -from ddtrace.internal import forksafe -from ddtrace.internal._threads import PeriodicThread as _PeriodicThread -from ddtrace.internal._threads import periodic_threads -from ddtrace.internal.logger import get_logger - - -log = get_logger(__name__) - -# We try to import the stdlib locks from the _thread module, where they are -# implemented in C for CPython for most platforms. If that fails, we fall back -# to the threading module, which provides a pure Python implementation that -# should work on all platforms. We also make sure to grab a reference to the -# original lock classes, in case they get patched by monkey-patching libraries -# like gevent. -try: - from _thread import allocate_lock as Lock -except ImportError: - from threading import Lock - -try: - from _thread import RLock -except ImportError: - from threading import RLock - - -__all__ = [ - "Lock", - "PeriodicThread", - "RLock", -] - - -# Forking state management. This is a barrier to either prevent new threads -# from being started while forking, or to allow a thread to be started -# completely if a fork comes in the middle of it. -_forking = False -_forking_lock = Lock() - - -class BoundMethod(t.Protocol): - __self__: t.Any - - def __call__(self) -> None: ... - - -# List of threads that have requested to be started while forking. These will -# be started after the fork is complete. -_threads_to_start_after_fork: list[BoundMethod] = [] - - -def _safe_restart(start: t.Callable[[], None], name: t.Optional[str] = None) -> None: - """Invoke a post-fork thread-start callable, logging resource errors instead of raising. - - The native layer translates pthread_create failures (EAGAIN, ENOMEM) into - OSError. Post-fork restart is triggered automatically by forksafe hooks — - there is no explicit caller that can handle the error, so losing a - periodic thread to resource exhaustion must not crash the host. - Explicit start() calls let OSError propagate so the caller can react. - """ - try: - start() - except Exception as e: - log.error("failed to start periodic thread %s: %s", name, e) - - -class PeriodicThread(_PeriodicThread): - """A fork-safe periodic thread.""" - - __autorestart__ = True - - def start(self) -> None: - with _forking_lock: - # We cannot start a new thread while we are forking, because we are - # trying to stop them all. In that case, we take note of the thread - # and start it after the fork. - if not _forking: - super().start() - else: - _threads_to_start_after_fork.append(t.cast(BoundMethod, super().start)) - - -# Set of running periodic threads that need to be restarted after a fork. -_threads_to_restart_after_fork: set[_PeriodicThread] = set() - - -# A typical scenario is that of forking worker threads in a loop. For the -# parent process, this would mean having to stop and restart the threads in -# between forks, which is not ideal. Instead, we can use a timer to restart -# the threads after a certain amount of time has passed since the last fork. -# This way, we can avoid stopping and restarting the threads in between forks. -class ThreadRestartTimer(PeriodicThread): - __timeout__ = int(1e8) # nanoseconds - - _instance: t.Optional["ThreadRestartTimer"] = None - _timestamp = 0 - - def __init__(self): - super().__init__(self.__timeout__ / 1e9, self._restart_threads, name=f"{__name__}:{self.__class__.__name__}") - - def _restart_threads(self) -> None: - # Restart the threads after we have stopped calling fork for a while. - with _forking_lock: - # If we are forking, we will try again later. - if _forking: - return - - # If we haven't have calls to fork for a while, we can restart the - # threads. This way we avoid stopping and restarting the threads - # in between forks. - if monotonic_ns() >= self._timestamp: # 100ms - for thread in _threads_to_restart_after_fork.copy(): - if isinstance(thread, ThreadRestartTimer): - # Skip any ThreadRestartTimer instance, - # to avoid restarting orphaned timer instances that were - # caught in periodic_threads during a fork. - continue - log.debug("Restarting thread %s after fork", thread.name) - try: - thread._after_fork(force=True) - except Exception as e: - log.error("failed to restart periodic thread %s after fork: %s", thread.name, e) - _threads_to_restart_after_fork.clear() - - for thread_start in _threads_to_start_after_fork: - log.debug("Starting thread %s after fork", thread_start.__self__.name) - _safe_restart(thread_start, thread_start.__self__.name) - _threads_to_start_after_fork.clear() - - # We no longer need this thread so we clear it. - self.clear() - - @classmethod - def clear(cls): - """Clear the timer and stop it if it is running.""" - if cls._instance is not None: - cls._instance.stop() - cls._instance = None - - @classmethod - def touch(cls): - """Set the new expiration time for the timer.""" - cls._timestamp = monotonic_ns() + cls.__timeout__ - - @classmethod - def set(cls): - """Set the timer to restart the threads after a fork.""" - if cls._instance is None: - cls._instance = cls() - cls._instance.start() - else: - # We have already created the timer, so we let the forksafe logic - # handle the restart instead of creating a new instance. - cls._instance._after_fork() - - -@forksafe.register -def _after_fork_child(): - global _forking - - _forking = False - - # Keep child at-fork work minimal: thread restarts happen asynchronously in - # the child so application code can resume immediately after fork. Parent - # process threads are still restarted in _after_fork_parent() below. - for thread in _threads_to_restart_after_fork.copy(): - log.debug("Restarting thread %s after fork in child", thread.name) - try: - thread._after_fork(force=False) - except Exception as e: - log.error("failed to restart periodic thread %s after fork in child: %s", thread.name, e) - _threads_to_restart_after_fork.clear() - - for thread_start in _threads_to_start_after_fork.copy(): - log.debug("Starting thread %s after fork in child", thread_start.__self__.name) - _safe_restart(thread_start, thread_start.__self__.name) - _threads_to_start_after_fork.clear() - - -@forksafe.register_after_parent -def _after_fork_parent() -> None: - global _forking - - _forking = False - - if _threads_to_restart_after_fork or _threads_to_start_after_fork: - ThreadRestartTimer.set() - - -@forksafe.register_before_fork -def _before_fork() -> None: - global _threads_to_restart_after_fork, _forking_lock, _forking - - ThreadRestartTimer.touch() - - with _forking_lock: - _forking = True - - # Snapshot pending restarts first so a worker moving pending -> active - # concurrently cannot be missed between the two snapshots. - pending_threads = getattr(_threads_mod, "_pending_threads", lambda: ())() - _threads_to_restart_after_fork.update(pending_threads) - # Take note of all the periodic threads that are running and will need to be - # restarted. - _threads_to_restart_after_fork.update(periodic_threads.values()) - - # Stop all the periodic threads that are still running, without executing - # the shutdown methods, if any. This ensures that we can stop the threads - # more promptly. - for thread in _threads_to_restart_after_fork: - log.debug("Stopping thread %s before fork", thread.name) - thread._before_fork() - - # Join all the threads to ensure they are stopped before the fork. - for thread in _threads_to_restart_after_fork: - log.debug("Joining thread %s before fork", thread.name) - thread.join() +from ddtrace.internal.utils.threads import * # noqa diff --git a/ddtrace/internal/utils/_exceptions.py b/ddtrace/internal/utils/_exceptions.py new file mode 100644 index 00000000000..e2836e036be --- /dev/null +++ b/ddtrace/internal/utils/_exceptions.py @@ -0,0 +1,37 @@ +from typing import Optional +from typing import TypeVar + + +class DDBlockException(BaseException): + """ + Base class for any in-tree decision to abort the current operation + (web request blocking, AI Guard policy abort, future product blocks). + Inherits from BaseException so a generic ``except Exception:`` handler in + user code does not accidentally swallow a blocking decision. + """ + + +class BlockingException(DDBlockException): + """ + Exception raised when a request is blocked by ASM + It derives from BaseException (via DDBlockException) to avoid being caught by the general Exception handler + """ + + +E = TypeVar("E", bound=BaseException) + + +def find_exception( + exc: BaseException, + exception_type: type[E], +) -> Optional[E]: + """Traverse an exception and its children to find the first occurrence of a specific exception type.""" + if isinstance(exc, exception_type): + return exc + # The check matches both native Python3.11+ and `exceptiongroup` compatibility package versions of ExceptionGroup + if exc.__class__.__name__ in ("BaseExceptionGroup", "ExceptionGroup") and hasattr(exc, "exceptions"): + for sub_exc in exc.exceptions: + found = find_exception(sub_exc, exception_type) + if found: + return found + return None diff --git a/ddtrace/internal/utils/compat.py b/ddtrace/internal/utils/compat.py new file mode 100644 index 00000000000..a0a122bd980 --- /dev/null +++ b/ddtrace/internal/utils/compat.py @@ -0,0 +1,125 @@ +import ipaddress +import sys +from types import TracebackType +from typing import Any +from typing import Optional # noqa:F401 +from typing import Text # noqa:F401 +from typing import Union # noqa:F401 + +import wrapt + + +__all__ = [ + "maybe_stringify", +] + +PYTHON_VERSION_INFO = sys.version_info + + +def ensure_text(s, encoding="utf-8", errors="ignore") -> str: + if isinstance(s, str): + return s + if isinstance(s, bytes): + return s.decode(encoding, errors) + raise TypeError("Expected str or bytes but received %r" % (s.__class__)) + + +def ensure_binary(s, encoding="utf-8", errors="ignore") -> bytes: + if isinstance(s, bytes): + return s + if not isinstance(s, str): + raise TypeError("Expected str or bytes but received %r" % (s.__class__)) + return s.encode(encoding, errors) + + +NumericType = Union[int, float] + + +def is_integer(obj: Any) -> bool: + """Helper to determine if the provided ``obj`` is an integer type or not""" + # DEV: We have to make sure it is an integer and not a boolean + # >>> type(True) + # + # >>> isinstance(True, int) + # True + return isinstance(obj, int) and not isinstance(obj, bool) + + +def maybe_stringify(obj: Any) -> Optional[str]: + if obj is not None: + return str(obj) + return None + + +ExcInfoType = Union[tuple[type[BaseException], BaseException, Optional[TracebackType]], tuple[None, None, None]] + +# Sentinel value to represent "no exception" +NO_EXCEPTION: ExcInfoType = (None, None, None) + + +def is_valid_ip(ip: str) -> bool: + try: + # try parsing the IP address + ipaddress.ip_address(str(ip)) + return True + except Exception: + return False + + +def ip_is_global(ip: str) -> bool: + """ + is_global is Python 3+ only. This could raise a ValueError if the IP is not valid. + """ + parsed_ip = ipaddress.ip_address(str(ip)) + + return parsed_ip.is_global + + +# This fix was implemented in 3.9.8 +# https://github.com/python/cpython/issues/83860 +if PYTHON_VERSION_INFO >= (3, 9, 8): + from functools import singledispatchmethod +else: + from functools import singledispatchmethod + + def _register(self, cls, method=None): + if hasattr(cls, "__func__"): + setattr(cls, "__annotations__", cls.__func__.__annotations__) + return self.dispatcher.register(cls, func=method) + + singledispatchmethod.register = _register # type: ignore[method-assign] + + +def get_mp_context(): + import multiprocessing + + return multiprocessing.get_context("fork" if sys.platform != "win32" else "spawn") + + +def __getattr__(name: str) -> Any: + # These attributes are expensive to pre-compute, so we make them lazy + if name == "PYTHON_VERSION": + from platform import python_version + + globals()[name] = python_version() + + elif name == "PYTHON_INTERPRETER": + from platform import python_implementation + + globals()[name] = python_implementation() + + try: + return globals()[name] + except KeyError: + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + +if hasattr(wrapt, "BaseObjectProxy"): + # This must be used for wrapt version >= 2.0.0 + wrapt_class: type = wrapt.BaseObjectProxy +else: + wrapt_class = wrapt.ObjectProxy + + +def is_wrapted(o: object) -> bool: + return isinstance(o, wrapt_class) diff --git a/ddtrace/internal/utils/constants.py b/ddtrace/internal/utils/constants.py new file mode 100644 index 00000000000..1167f75f9e6 --- /dev/null +++ b/ddtrace/internal/utils/constants.py @@ -0,0 +1,159 @@ +from ddtrace.constants import AUTO_KEEP +from ddtrace.constants import AUTO_REJECT +from ddtrace.constants import USER_KEEP +from ddtrace.constants import USER_REJECT + + +PROPAGATION_STYLE_DATADOG = "datadog" +PROPAGATION_STYLE_B3_MULTI = "b3multi" +PROPAGATION_STYLE_B3_SINGLE = "b3" +_PROPAGATION_STYLE_W3C_TRACECONTEXT = "tracecontext" +_PROPAGATION_STYLE_NONE = "none" +_PROPAGATION_STYLE_DEFAULT = "datadog,tracecontext,baggage" +_PROPAGATION_STYLE_BAGGAGE = "baggage" +PROPAGATION_STYLE_ALL = ( + _PROPAGATION_STYLE_W3C_TRACECONTEXT, + PROPAGATION_STYLE_DATADOG, + PROPAGATION_STYLE_B3_MULTI, + PROPAGATION_STYLE_B3_SINGLE, + _PROPAGATION_STYLE_NONE, + _PROPAGATION_STYLE_BAGGAGE, +) +_PROPAGATION_BEHAVIOR_CONTINUE = "continue" +_PROPAGATION_BEHAVIOR_IGNORE = "ignore" +_PROPAGATION_BEHAVIOR_RESTART = "restart" +_PROPAGATION_BEHAVIOR_DEFAULT = _PROPAGATION_BEHAVIOR_CONTINUE +W3C_TRACESTATE_KEY = "tracestate" +W3C_TRACEPARENT_KEY = "traceparent" +W3C_TRACESTATE_PARENT_ID_KEY = "p" +W3C_TRACESTATE_ORIGIN_KEY = "o" +W3C_TRACESTATE_SAMPLING_PRIORITY_KEY = "s" +DEFAULT_SAMPLING_RATE_LIMIT = 100 +SAMPLING_HASH_MODULO = 1 << 64 +# Big prime number to make hashing better distributed, it has to be the same factor as the Agent +# and other tracers to allow chained sampling +SAMPLING_KNUTH_FACTOR = 1111111111111111111 +SAMPLING_DECISION_TRACE_TAG_KEY = "_dd.p.dm" +LAST_DD_PARENT_ID_KEY = "_dd.parent_id" +DEFAULT_SERVICE_NAME = "unnamed-python-service" +# Used to set the name of an integration on a span +COMPONENT = "component" +HIGHER_ORDER_TRACE_ID_BITS = "_dd.p.tid" +MAX_UINT_64BITS = (1 << 64) - 1 +MIN_INT_64BITS = -(2**63) +MAX_INT_64BITS = 2**63 - 1 +SAMPLING_DECISION_MAKER_INHERITED = "_dd.dm.inherited" +SAMPLING_DECISION_MAKER_SERVICE = "_dd.dm.service" +SAMPLING_DECISION_MAKER_RESOURCE = "_dd.dm.resource" +SPAN_LINK_KIND = "dd.kind" +SPAN_LINKS_KEY = "_dd.span_links" +SPAN_EVENTS_KEY = "events" +SPAN_API_DATADOG = "datadog" +SPAN_API_OTEL = "otel" +SPAN_API_OPENTRACING = "opentracing" +DEFAULT_BUFFER_SIZE = 20 << 20 # 20 MB +DEFAULT_MAX_PAYLOAD_SIZE = 20 << 20 # 20 MB +DEFAULT_PROCESSING_INTERVAL = 1.0 +DEFAULT_REUSE_CONNECTIONS = False +BLOCKED_RESPONSE_HTML = """You've been blocked

Sorry, you cannot access this page. Please contact the customer service team.

Security Response ID: [security_response_id]

""" # noqa: E501 +BLOCKED_RESPONSE_JSON = """{"errors":[{"title":"You've been blocked","detail":"Sorry, you cannot access this page. Please contact the customer service team. Security provided by Datadog."}],"security_response_id":"[security_response_id]"}""" # noqa: E501 +HTTP_REQUEST_BLOCKED = "http.request.blocked" +RESPONSE_HEADERS = "http.response.headers" +HTTP_REQUEST_QUERY = "http.request.query" +HTTP_REQUEST_COOKIE_VALUE = "http.request.cookie.value" +HTTP_REQUEST_COOKIE_NAME = "http.request.cookie.name" +HTTP_REQUEST_PATH = "http.request.path" +HTTP_REQUEST_HEADER_NAME = "http.request.header.name" +HTTP_REQUEST_HEADER = "http.request.header" +HTTP_REQUEST_PARAMETER = "http.request.parameter" +HTTP_REQUEST_BODY = "http.request.body" +HTTP_REQUEST_UPGRADED = "http.upgraded" +HTTP_REQUEST_PATH_PARAMETER = "http.request.path.parameter" +REQUEST_PATH_PARAMS = "http.request.path_params" +STATUS_403_TYPE_AUTO = {"status_code": 403, "type": "auto"} +PROCESS_TAGS = "_dd.tags.process" +PROPAGATED_HASH = "_dd.propagated_hash" +_SERVICE_SOURCE = "_dd.svc_src" + +CONTAINER_ID_HEADER_NAME = "Datadog-Container-Id" +CONTAINER_TAGS_HASH = "Datadog-Container-Tags-Hash" + +ENTITY_ID_HEADER_NAME = "Datadog-Entity-ID" + +EXTERNAL_ENV_HEADER_NAME = "Datadog-External-Env" +EXTERNAL_ENV_ENVIRONMENT_VARIABLE = "DD_EXTERNAL_ENV" + +MESSAGING_BATCH_COUNT = "messaging.batch_count" +MESSAGING_DESTINATION_NAME = "messaging.destination.name" +MESSAGING_MESSAGE_ID = "messaging.message_id" +MESSAGING_OPERATION = "messaging.operation" +MESSAGING_SYSTEM = "messaging.system" + +USER_AGENT_HEADER = "user-agent" +FLASK_ENDPOINT = "flask.endpoint" +FLASK_VIEW_ARGS = "flask.view_args" +FLASK_URL_RULE = "flask.url_rule" +FLASK_RESOURCE_FULL = "flask.resource.full" + +_HTTPLIB_NO_TRACE_REQUEST = "_dd_no_trace" +DEFAULT_TIMEOUT = 2.0 + +# baggage +DD_TRACE_BAGGAGE_MAX_ITEMS = 64 +DD_TRACE_BAGGAGE_MAX_BYTES = 8192 +BAGGAGE_TAG_PREFIX = "baggage." + +# W3C Trace Context tracestate (https://www.w3.org/TR/trace-context/): +# max 32 list-members; vendors SHOULD propagate at most 512 characters (we cap parsing to that size). +DD_TRACE_TRACESTATE_MAX_ITEMS = 32 +DD_TRACE_TRACESTATE_MAX_BYTES = 512 +# Per W3C Trace Context, oversized list-members are preferred targets when truncating by size. +DD_TRACE_TRACESTATE_ITEM_MAX_CHARS = 128 + +SPAN_EVENTS_HAS_EXCEPTION = "_dd.span_events.has_exception" +COLLECTOR_MAX_SIZE_PER_SPAN = 100 + +LOG_ATTR_TRACE_ID = "dd.trace_id" +LOG_ATTR_SPAN_ID = "dd.span_id" +LOG_ATTR_ENV = "dd.env" +LOG_ATTR_VERSION = "dd.version" +LOG_ATTR_SERVICE = "dd.service" +LOG_ATTR_VALUE_ZERO = "0" +LOG_ATTR_VALUE_EMPTY = "" + + +class SamplingMechanism(object): + DEFAULT = 0 + AGENT_RATE_BY_SERVICE = 1 + REMOTE_RATE = 2 # not used, this mechanism is deprecated + LOCAL_USER_TRACE_SAMPLING_RULE = 3 + MANUAL = 4 + APPSEC = 5 + REMOTE_RATE_USER = 6 # not used, this mechanism is deprecated + REMOTE_RATE_DATADOG = 7 # not used, this mechanism is deprecated + SPAN_SAMPLING_RULE = 8 + OTLP_INGEST_PROBABILISTIC_SAMPLING = 9 # not used in ddtrace + DATA_JOBS_MONITORING = 10 # not used in ddtrace + REMOTE_USER_TRACE_SAMPLING_RULE = 11 + REMOTE_DYNAMIC_TRACE_SAMPLING_RULE = 12 + AI_GUARD = 13 + + +SAMPLING_MECHANISM_TO_PRIORITIES = { + # TODO(munir): Update mapping to include single span sampling and appsec sampling mechanisms + SamplingMechanism.AGENT_RATE_BY_SERVICE: (AUTO_KEEP, AUTO_REJECT), + SamplingMechanism.DEFAULT: (AUTO_KEEP, AUTO_REJECT), + SamplingMechanism.MANUAL: (USER_KEEP, USER_REJECT), + SamplingMechanism.APPSEC: (AUTO_KEEP, AUTO_REJECT), + SamplingMechanism.LOCAL_USER_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), + SamplingMechanism.REMOTE_USER_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), + SamplingMechanism.REMOTE_DYNAMIC_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), +} +_KEEP_PRIORITY_INDEX = 0 +_REJECT_PRIORITY_INDEX = 1 + + +# List of support values in DD_TRACE_EXPERIMENTAL_FEATURES_ENABLED +class EXPERIMENTAL_FEATURES: + # Enables submitting runtime metrics as gauges (instead of distributions) + RUNTIME_METRICS = "DD_RUNTIME_METRICS_ENABLED" diff --git a/ddtrace/internal/utils/logger.py b/ddtrace/internal/utils/logger.py new file mode 100644 index 00000000000..82d5b8dcbd4 --- /dev/null +++ b/ddtrace/internal/utils/logger.py @@ -0,0 +1,246 @@ +""" +Logging utilities for internal use. +Usage: + import ddtrace.internal.utils.logger as logger + ddlog = logger.get_logger(__name__) + + # Otherwise default is set to 1 minute or DD_TRACE_LOGGING_RATE + logger.set_tag_rate_limit("waf::init", logger.HOUR) + + # "product" is required, but other keys are optional as well as kwargs exc_info and stack_info + # supported keys are + # product: product or integration name. Required + # more_info : more information to be logged after the main tag. Default is empty string + # stack_limit: limit the stack trace to this depth for stack_info. Default is 0 (0 is no limit) + # exec_limit: limit the stack trace to this depth for exec_info. Default is 1, only the top level (0 is no limit) + + info = # format the info string + ddlog.debug('waf::init', extra={"product": "appsec", "stack_limit": 4, "more_info": info}, stack_info=True) + # This will log the message only once per hour, counting the number of skipped messages, using "waf::init" as the + # tag to keep track of the rate limit + # Different log levels can be used, the rate limit is shared between all invocations and all levels for the same tag + + # example result + DEBUG appsec::waf::init[some more info] 1 additional messages skipped + (followed by the 4 first levels of the stack trace) + + Legacy support: + if extra is not used or product is absent, the log will be treated as legacy and will be logged as is using + filename and line number of the log call + +""" + +import collections +from dataclasses import dataclass +from dataclasses import field +import logging +import time +import traceback +from typing import DefaultDict +from typing import Optional +from typing import Tuple +from typing import Union + +from ddtrace.internal.settings import env + + +SECOND = 1 +MINUTE = 60 * SECOND +HOUR = 60 * MINUTE +DAY = 24 * HOUR + + +@dataclass +class LoggerPrefix: + prefix: str + level: Optional[int] = None + children: dict[str, "LoggerPrefix"] = field(default_factory=dict) + + def lookup(self, name: str) -> Optional[int]: + """ + Lookup the log level for a given logger name in the trie. + + The name is split by '.' and each part is used to traverse the trie. + If a part is not found, it returns the level of the closest parent node. + """ + parts = name.replace("_", ".").lower().split(".") + parts.pop(0) # remove the ddtrace prefix + current = self + while parts: + if (part := parts.pop(0)) not in current.children: + return current.level + current = current.children[part] + + return current.level + + @classmethod + def build_trie(cls): + trie = cls(prefix="ddtrace", level=None, children={}) + + for logger_name, level in ((k, v) for k, v in env.items() if k.startswith("_DD_") and k.endswith("_LOG_LEVEL")): + # Remove the _DD_ prefix and _LOG_LEVEL suffix + logger_name = logger_name[4:-10] + parts = logger_name.lower().split("_") + current = trie.children + while parts: + if (part := parts.pop(0)) not in current: + current[part] = cls(prefix=part, level=getattr(logging, level, None) if not parts else None) + current = current[part].children + + return trie + + +LOG_LEVEL_TRIE = LoggerPrefix.build_trie() + + +def get_logger(name: str) -> logging.Logger: + """ + Retrieve or create a ``Logger`` instance with consistent behavior for internal use. + + Configure all loggers with a rate limiter filter to prevent excessive logging. + + """ + logger = logging.getLogger(name) + # addFilter will only add the filter if it is not already present + logger.addFilter(log_filter) + + # Set the log level from the environment variable of the closest parent + # logger. + if name.startswith("ddtrace."): # for the whole of ddtrace we have DD_TRACE_DEBUG + if (level := LOG_LEVEL_TRIE.lookup(name)) is not None: + logger.setLevel(level) + + return logger + + +_RATE_LIMITS = {} + + +def set_tag_rate_limit(tag: str, rate: int) -> None: + """ + Set the rate limit for a specific tag. + + """ + _RATE_LIMITS[tag] = rate + + +# Class used for keeping track of a log lines current time bucket and the number of log lines skipped +class LoggingBucket: + def __init__(self, bucket: float, skipped: int): + self.bucket = bucket + self.skipped = skipped + + def __repr__(self): + return f"LoggingBucket({self.bucket}, {self.skipped})" + + def is_sampled(self, record: logging.LogRecord, rate: float) -> bool: + """ + Determine if the log line should be sampled based on the rate limit. + """ + current = time.monotonic() + if current - self.bucket >= rate: + self.bucket = current + record.skipped = self.skipped + self.skipped = 0 + return True + self.skipped += 1 + return False + + +# dict to keep track of the current time bucket per name/level/pathname/lineno + +_MINF = float("-inf") + +# IMPORTANT: Do not change typing types to built-ins until minimum Python version is 3.11+ +# Module-level tuple[...] and defaultdict[...] in Python 3.10 affect import timing. See packages.py for details. +key_type = Union[Tuple[str, int, str, int], str] # noqa: UP006 +_buckets: DefaultDict[key_type, LoggingBucket] = collections.defaultdict(lambda: LoggingBucket(_MINF, 0)) # noqa: UP006 + +# Allow 1 log record per name/level/pathname/lineno every 60 seconds by default +# Allow configuring via `DD_TRACE_LOGGING_RATE` +# DEV: `DD_TRACE_LOGGING_RATE=0` means to disable all rate limiting +_rate_limit = int(env.get("DD_TRACE_LOGGING_RATE", default=60)) + + +def log_filter(record: logging.LogRecord) -> bool: + """ + Function used to determine if a log record should be outputted or not (True = output, False = skip). + + This function will: + - Rate limit log records based on the logger name, record level, filename, and line number + """ + logger = logging.getLogger(record.name) + rate_limit = _RATE_LIMITS.get(record.msg, _rate_limit) + # If rate limiting has been disabled (`DD_TRACE_LOGGING_RATE=0`) then apply no rate limit + # If the logger is set to debug, then do not apply any limits to any log + if not rate_limit or logger.getEffectiveLevel() == logging.DEBUG: + must_be_propagated = True + else: + # Allow 1 log record by pathname/lineno every X seconds or message/levelno for product logs + # This way each unique log message can get logged at least once per time period + if hasattr(record, "product"): + key: key_type = record.msg + else: + key = (record.name, record.levelno, record.pathname, record.lineno) + # If rate limiting has been disabled (`DD_TRACE_LOGGING_RATE=0`) then apply no rate limit + # If the logger is set to debug, then do not apply any limits to any log + # Only log this message if the time bucket allows it + must_be_propagated = _buckets[key].is_sampled(record, rate_limit) + if must_be_propagated: + skipped = record.__dict__.pop("skipped", 0) + if skipped: + skip_str = f" [{skipped} skipped]" + else: + skip_str = "" + product = record.__dict__.pop("product", None) + # new syntax + if product: + more_info = record.__dict__.pop("more_info", "") + stack_limit = record.__dict__.pop("stack_limit", 0) + exec_limit = record.__dict__.pop("exec_limit", 1) + # format the stacks if they are present with the right depth + if stack_limit and record.stack_info: + record.stack_info = format_stack(record.stack_info, stack_limit) + string_buffer = [f"{product}::{record.msg}{more_info}{skip_str}"] + if record.stack_info: + string_buffer.append(record.stack_info) + if record.exc_info: + string_buffer.extend(traceback.format_exception(record.exc_info[1], limit=exec_limit or None)) + record.msg = "\n".join(string_buffer) + # clean the record for any subsequent handlers + record.stack_info = None + record.exc_info = None + else: + record.msg = f"{record.msg}{skip_str}" + return must_be_propagated + + +def format_stack(stack_info, limit) -> str: + stack = stack_info.split("\n") + if len(stack) <= limit * 2 + 1: + return stack_info + stack_str = "\n".join(stack[-2 * limit :]) + return f"{stack[0]}\n{stack_str}" + + +class LogInjectionState(object): + # Log injection is disabled + DISABLED = "false" + # Log injection is enabled, but not yet configured + ENABLED = "true" + # Log injection is enabled and configured for structured logging + # This value is deprecated, but kept for backwards compatibility + STRUCTURED = "structured" + + +def get_log_injection_state(raw_config: Optional[str]) -> bool: + if raw_config: + normalized = raw_config.lower().strip() + if normalized == LogInjectionState.STRUCTURED or normalized in ("true", "1"): + return True + elif normalized not in ("false", "0"): + logging.warning( + "Invalid log injection state '%s'. Expected 'true', 'false', or 'structured'. Defaulting to 'false'.", + normalized, + ) + return False diff --git a/ddtrace/internal/utils/module.py b/ddtrace/internal/utils/module.py new file mode 100644 index 00000000000..2cfbb600d81 --- /dev/null +++ b/ddtrace/internal/utils/module.py @@ -0,0 +1,771 @@ +import abc +from collections import defaultdict +from importlib._bootstrap import _init_module_attrs +from importlib.machinery import ModuleSpec +from importlib.util import find_spec +from importlib.util import spec_from_loader +from pathlib import Path +import sys +from types import CodeType +from types import FunctionType +from types import ModuleType +import typing as t +from weakref import WeakValueDictionary as wvdict + +from ddtrace.internal.utils import get_argument_value +from ddtrace.internal.utils.logger import get_logger +from ddtrace.internal.utils.wrapping.context import WrappingContext + + +if t.TYPE_CHECKING: + from importlib.abc import Loader + + +ModuleHookType = t.Callable[[ModuleType], None] +TransformerType = t.Callable[[CodeType, ModuleType], CodeType] +PreExecHookType = t.Callable[[t.Any, ModuleType], None] +PreExecHookCond = t.Union[str, t.Callable[[str], bool]] + +ImportExceptionHookType = t.Callable[[t.Any, ModuleType], None] +ImportExceptionHookCond = t.Union[str, t.Callable[[str], bool]] + + +log = get_logger(__name__) + + +_run_code: t.Optional[t.Callable[[t.Any], t.Any]] = None +_run_module_transformers: list[TransformerType] = [] +_post_run_module_hooks: list[ModuleHookType] = [] + + +def _wrapped_run_code(*args: t.Any, **kwargs: t.Any) -> dict[str, t.Any]: + code = t.cast(CodeType, get_argument_value(args, kwargs, 0, "code")) + mod_name = t.cast(str, get_argument_value(args, kwargs, 3, "mod_name")) + + module = sys.modules[mod_name] + + for transformer in _run_module_transformers: + code = transformer(code, module) + + kwargs.pop("code", None) + new_args = (code, *args[1:]) + + try: + return _run_code(*new_args, **kwargs) # type: ignore[misc] + finally: + module = sys.modules[mod_name] + for hook in _post_run_module_hooks: + hook(module) + + +def _patch_run_code() -> None: + global _run_code + + if _run_code is None: + import runpy + + _run_code = runpy._run_code # type: ignore[attr-defined] + runpy._run_code = _wrapped_run_code # type: ignore[attr-defined] + + +def register_run_module_transformer(transformer: TransformerType) -> None: + """Register a run module transformer.""" + _run_module_transformers.append(transformer) + + +def unregister_run_module_transformer(transformer: TransformerType) -> None: + """Unregister a run module transformer. + + If the transformer was not registered, a ``ValueError`` exception is raised. + """ + _run_module_transformers.remove(transformer) + + +def register_post_run_module_hook(hook: ModuleHookType) -> None: + """Register a post run module hook. + + The hooks gets called after the module is loaded. For this to work, the + hook needs to be registered during the interpreter initialization, e.g. as + part of a sitecustomize.py script. + """ + global _run_code, _post_run_module_hooks + + _patch_run_code() + + _post_run_module_hooks.append(hook) + + +def unregister_post_run_module_hook(hook: ModuleHookType) -> None: + """Unregister a post run module hook. + + If the hook was not registered, a ``ValueError`` exception is raised. + """ + global _post_run_module_hooks + + _post_run_module_hooks.remove(hook) + + +def origin(module: ModuleType) -> t.Optional[Path]: + """Get the origin source file of the module.""" + try: + # Do not access __dd_origin__ directly to avoid force-loading lazy + # modules. + return object.__getattribute__(module, "__dd_origin__") + except AttributeError: + try: + # DEV: Use object.__getattribute__ to avoid potential side-effects. + orig = Path(object.__getattribute__(module, "__file__")).resolve() + except (AttributeError, TypeError): + # Module is probably only partially initialised, so we look at its + # spec instead + try: + # DEV: Use object.__getattribute__ to avoid potential side-effects. + orig = Path(object.__getattribute__(module, "__spec__").origin).resolve() + except (AttributeError, ValueError, TypeError): + orig = None + + if orig is not None and orig.suffix == "pyc": + orig = orig.with_suffix(".py") + + if orig is not None: + # If we failed to find a valid origin we don't cache the value and + # try again the next time. + try: + module.__dd_origin__ = orig # type: ignore[attr-defined] + except AttributeError: + pass + + return orig + + +def _resolve(path: Path) -> t.Optional[Path]: + """Resolve a (relative) path with respect to sys.path.""" + for base in (Path(_) for _ in sys.path): + if base.is_dir(): + resolved_path = (base / path.expanduser()).resolve() + if resolved_path.is_file(): + return resolved_path + return None + + +# Borrowed from the wrapt module +# https://github.com/GrahamDumpleton/wrapt/blob/df0e62c2740143cceb6cafea4c306dae1c559ef8/src/wrapt/importer.py + + +def find_loader(fullname: str) -> t.Optional["Loader"]: + return getattr(find_spec(fullname), "loader", None) + + +def is_module_installed(module_name): + return find_loader(module_name) is not None + + +def is_namespace_spec(spec: ModuleSpec) -> bool: + return spec.origin is None and spec.submodule_search_locations is not None + + +class _ImportHookChainedLoader: + def __init__(self, loader: t.Optional["Loader"], spec: t.Optional[ModuleSpec] = None) -> None: + self.loader = loader + self.spec = spec + + self.callbacks: dict[t.Any, t.Callable[[ModuleType], None]] = {} + self.import_exception_callbacks: dict[t.Any, t.Callable[[ModuleType], None]] = {} + + self.transformers: dict[t.Any, TransformerType] = {} + + # A missing loader is generally an indication of a namespace package. + if loader is None or hasattr(loader, "create_module"): + self.create_module = self._create_module + if loader is None or hasattr(loader, "exec_module"): + self.exec_module = self._exec_module + + def __getattr__(self, name): + # Proxy any other attribute access to the underlying loader. + return getattr(self.loader, name) + + def namespace_module(self, spec: ModuleSpec) -> ModuleType: + module = ModuleType(spec.name) + # Pretend that we do not have a loader (this would be self), to + # allow _init_module_attrs to create the appropriate NamespaceLoader + # for the namespace module. + spec.loader = None + + _init_module_attrs(spec, module, override=True) + + # Chain the loaders + self.loader = spec.loader + module.__loader__ = spec.loader = self # type: ignore[assignment] + + return module + + def add_callback(self, key: t.Any, callback: t.Callable[[ModuleType], None]) -> None: + self.callbacks[key] = callback + + def add_import_exception_callback(self, key: t.Any, callback: t.Callable[[ModuleType], None]) -> None: + self.import_exception_callbacks[key] = callback + + def add_transformer(self, key: t.Any, transformer: TransformerType) -> None: + self.transformers[key] = transformer + + def call_back(self, module: ModuleType) -> None: + # Restore the original loader, if possible. Some specs might be native + # and won't support attribute assignment. + try: + module.__loader__ = self.loader + except (AttributeError, TypeError): + pass + try: + object.__getattribute__(module, "spec").loader = self.loader + except (AttributeError, TypeError): + pass + + if module.__name__ == "pkg_resources": + # DEV: pkg_resources support to prevent errors such as + # NotImplementedError: Can't perform this operation for unregistered + # loader type + module.register_loader_type(_ImportHookChainedLoader, module.DefaultProvider) + + for callback in self.callbacks.values(): + callback(module) + + def load_module(self, fullname: str) -> t.Optional[ModuleType]: + if self.loader is None: + if self.spec is None: + return None + sys.modules[self.spec.name] = module = self.namespace_module(self.spec) + else: + module = self.loader.load_module(fullname) + + try: + self.call_back(module) + except Exception: + log.exception("Failed to call back on module %s", module) + + return module + + def _create_module(self, spec): + if self.loader is not None: + return self.loader.create_module(spec) + + if is_namespace_spec(spec): + return self.namespace_module(spec) + + return None + + def _find_first_hook( + self, module: ModuleType, hooks_attr: str + ) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: + for _ in sys.meta_path: + if isinstance(_, ModuleWatchdog): + try: + for ( + cond, + hook, + ) in getattr(_, hooks_attr, []): + if (isinstance(cond, str) and cond == module.__name__) or ( + callable(cond) and cond(module.__name__) + ): + return hook + except Exception: + log.debug("Exception happened while processing %s", hooks_attr, exc_info=True) + return None + + def _find_first_exception_hook(self, module: ModuleType) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: + return self._find_first_hook(module, "_import_exception_hooks") + + def _find_first_pre_exec_hook(self, module: ModuleType) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: + return self._find_first_hook(module, "_pre_exec_module_hooks") + + def _exec_module(self, module: ModuleType) -> None: + # Collect and run only the first hook that matches the module. + + _get_code = getattr(self.loader, "get_code", None) + # DEV: avoid re-wrapping the loader's get_code method (eg: in case of repeated importlib.reload() calls) + if _get_code is not None and not getattr(_get_code, "_dd_get_code", False): + + def get_code(_loader, fullname): + code = _get_code(fullname) + + for callback in self.transformers.values(): + code = callback(code, module) + + return code + + setattr(get_code, "_dd_get_code", True) + + self.loader.get_code = get_code.__get__(self.loader, type(self.loader)) # type: ignore[union-attr] + + pre_exec_hook = self._find_first_pre_exec_hook(module) + + if pre_exec_hook: + pre_exec_hook(self, module) + else: + if self.loader is None: + spec = getattr(module, "__spec__", None) + if spec is not None and is_namespace_spec(spec): + sys.modules[spec.name] = module + else: + try: + self.loader.exec_module(module) + except Exception: + exception_hook = self._find_first_exception_hook(module) + if exception_hook is not None: + exception_hook(self, module) + + # Hide the chained loader method from the traceback + _, e, tb = sys.exc_info() + if e is not None and tb is not None: + e.__traceback__ = tb.tb_next + + raise + + try: + self.call_back(module) + except Exception: + log.exception("Failed to call back on module %s", module) + + +class BaseModuleWatchdog(abc.ABC): + """Base module watchdog. + + Invokes ``after_import`` every time a new module is imported. + """ + + _instance: t.Optional["BaseModuleWatchdog"] = None + + def __init__(self) -> None: + self._finding: set[str] = set() + + # DEV: pkg_resources support to prevent errors such as + # NotImplementedError: Can't perform this operation for unregistered + pkg_resources = sys.modules.get("pkg_resources") + if pkg_resources is not None: + for _ in range(5): + # The package might be in the process of being imported by + # another thread by the time this module watchdog is created. If + # we fail to find the method we try again after a short wait. + # This is likely to happen when a module watchdog is installed + # lazily after boot, e.g. as a response to a RC enablement + # record. This means that we are unlikely to add delays to the + # main or any other use threads. + try: + pkg_resources.register_loader_type(_ImportHookChainedLoader, pkg_resources.DefaultProvider) + break + except AttributeError: + from time import sleep + + sleep(0.1) + else: + log.warning("Cannot ensure correct support with pkg_resources") + + def _add_to_meta_path(self) -> None: + sys.meta_path.insert(0, self) # type: ignore[arg-type] + + @classmethod + def _find_in_meta_path(cls) -> t.Optional[int]: + for i, meta_path in enumerate(sys.meta_path): + if type(meta_path) is cls: + return i + return None + + @classmethod + def _remove_from_meta_path(cls) -> None: + i = cls._find_in_meta_path() + + if i is None: + log.warning("%s is not installed", cls.__name__) + return + + sys.meta_path.pop(i) + + def after_import(self, module: ModuleType) -> None: + raise NotImplementedError() + + def transform(self, code: CodeType, _module: ModuleType) -> CodeType: + return code + + def find_module(self, fullname: str, path: t.Optional[str] = None) -> t.Optional["Loader"]: + if fullname in self._finding: + return None + + self._finding.add(fullname) + + try: + original_loader = find_loader(fullname) + if original_loader is not None: + loader = ( + _ImportHookChainedLoader(original_loader) + if not isinstance(original_loader, _ImportHookChainedLoader) + else original_loader + ) + + loader.add_callback(type(self), self.after_import) + loader.add_transformer(type(self), self.transform) + + return t.cast("Loader", loader) + + finally: + self._finding.remove(fullname) + + return None + + def find_spec( + self, fullname: str, path: t.Optional[str] = None, target: t.Optional[ModuleType] = None + ) -> t.Optional[ModuleSpec]: + if fullname in self._finding: + return None + + self._finding.add(fullname) + + try: + # Iterate sys.meta_path directly to avoid re-entering the import machinery + # and acquiring per-module locks (A-B deadlock risk when a background thread + # calls importlib.files() during module __init__). + spec = None + for finder in sys.meta_path: + if finder is self: + continue + _find_spec = getattr(finder, "find_spec", None) + if _find_spec is not None: + spec = _find_spec(fullname, path, target) + else: + # Fallback for legacy finders that only implement find_module() + # (deprecated in 3.4, removed in 3.12) — mirrors CPython _find_spec(). + _find_module = getattr(finder, "find_module", None) + if _find_module is None: + continue + loader = _find_module(fullname, path) + spec = spec_from_loader(fullname, loader) if loader is not None else None + if spec is not None: + break + + if spec is None: + return None + + loader = getattr(spec, "loader", None) + + if not isinstance(loader, _ImportHookChainedLoader): + spec.loader = t.cast("Loader", _ImportHookChainedLoader(loader, spec)) + + t.cast(_ImportHookChainedLoader, spec.loader).add_callback(type(self), self.after_import) + t.cast(_ImportHookChainedLoader, spec.loader).add_transformer(type(self), self.transform) + + return spec + + finally: + self._finding.remove(fullname) + + @classmethod + def install(cls) -> None: + """Install the module watchdog.""" + if cls.is_installed(): + return + cls._instance = cls() + cls._instance._add_to_meta_path() + log.debug("%s installed", cls) + + @classmethod + def is_installed(cls) -> bool: + """Check whether this module watchdog class is installed.""" + return cls._instance is not None and type(cls._instance) is cls + + @classmethod + def uninstall(cls) -> None: + """Uninstall the module watchdog. + + This will uninstall only the most recently installed instance of this + class. + """ + if not cls.is_installed(): + return + + cls._remove_from_meta_path() + + cls._instance = None + + log.debug("%s uninstalled", cls) + + +class ModuleWatchdog(BaseModuleWatchdog): + """Module watchdog. + + Hooks into the import machinery to detect when modules are loaded/unloaded. + This is also responsible for triggering any registered import hooks. + + Subclasses might customize the default behavior by overriding the + ``after_import`` method, which is triggered on every module import, once + the subclass is installed. + """ + + def __init__(self) -> None: + super().__init__() + + self._hook_map: defaultdict[str, list[ModuleHookType]] = defaultdict(list) + # DEV: It would make more sense to make this a mapping of Path to ModuleType + # but the WeakValueDictionary causes an ignored exception on shutdown + # because the pathlib module is being garbage collected. + self._om: t.Optional[t.MutableMapping[str, ModuleType]] = None + # _pre_exec_module_hooks is a set of tuples (condition, hook) instead + # of a list to ensure that no hook is duplicated + self._pre_exec_module_hooks: set[tuple[PreExecHookCond, PreExecHookType]] = set() + self._import_exception_hooks: set[tuple[ImportExceptionHookCond, ImportExceptionHookType]] = set() + + @property + def _origin_map(self) -> t.MutableMapping[str, ModuleType]: + def modules_with_origin(modules: t.Iterable[ModuleType]) -> t.MutableMapping[str, ModuleType]: + result: t.MutableMapping[str, ModuleType] = wvdict() + + for m in modules: + module_origin = origin(m) + if module_origin is None: + continue + + try: + result[str(module_origin)] = m + except TypeError: + # This can happen if the module is a special object that + # does not allow for weak references. Quite likely this is + # an object created by a native extension. We make the + # assumption that this module does not contain valuable + # information that can be used at the Python runtime level. + pass + + return result + + if self._om is None: + try: + self._om = modules_with_origin(sys.modules.values()) + except RuntimeError: + # The state of sys.modules might have been mutated by another + # thread. We try to build the full mapping at the next occasion. + # For now we take the more expensive route of building a list of + # the current values, which might be incomplete. + return modules_with_origin(list(sys.modules.values())) + return self._om + + def after_import(self, module: ModuleType) -> None: + module_path = origin(module) + path = str(module_path) if module_path is not None else None + if path is not None: + self._origin_map[path] = module + + # Collect all hooks by module origin and name + hooks = [] + if path is not None and path in self._hook_map: + hooks.extend(self._hook_map[path]) + if module.__name__ in self._hook_map: + hooks.extend(self._hook_map[module.__name__]) + + if hooks: + log.debug("Calling %d registered hooks on import of module '%s'", len(hooks), module.__name__) + for hook in hooks: + hook(module) + + @classmethod + def get_by_origin(cls, _origin: Path) -> t.Optional[ModuleType]: + """Lookup a module by its origin.""" + if not cls.is_installed(): + return None + + instance = t.cast(ModuleWatchdog, cls._instance) + + resolved_path = _resolve(_origin) + if resolved_path is not None: + path = str(resolved_path) + module = instance._origin_map.get(path) + if module is not None: + return module + + # Check if this is the __main__ module + main_module = sys.modules.get("__main__") + if main_module is not None and origin(main_module) == resolved_path: + # Register for future lookups + instance._origin_map[path] = main_module + + return main_module + + return None + + @classmethod + def register_origin_hook(cls, origin: Path, hook: ModuleHookType) -> None: + """Register a hook to be called when the module with the given origin is + imported. + + The hook will be called with the module object as argument. + """ + cls.install() + + # DEV: Under the hypothesis that this is only ever called by the probe + # poller thread, there are no further actions to take. Should this ever + # change, then thread-safety might become a concern. + resolved_path = _resolve(origin) + if resolved_path is None: + log.warning("Cannot resolve module origin %s", origin) + return + + path = str(resolved_path) + + log.debug("Registering hook '%r' on path '%s'", hook, path) + instance = t.cast(ModuleWatchdog, cls._instance) + instance._hook_map[path].append(hook) + try: + module = instance.get_by_origin(resolved_path) + if module is None: + # The path was resolved but we still haven't seen any module + # that has it as origin. Nothing more we can do for now. + return + # Sanity check: the module might have been removed from sys.modules + # but not yet garbage collected. + try: + sys.modules[module.__name__] + except KeyError: + del instance._origin_map[path] + raise + except KeyError: + # The module is not loaded yet. Nothing more we can do. + return + + # The module was already imported so we invoke the hook straight-away + log.debug("Calling hook '%r' on already imported module '%s'", hook, module.__name__) + hook(module) + + @classmethod + def unregister_origin_hook(cls, origin: Path, hook: ModuleHookType) -> None: + """Unregister the hook registered with the given module origin and + argument. + """ + if not cls.is_installed(): + return + + resolved_path = _resolve(origin) + if resolved_path is None: + log.warning("Module origin %s cannot be resolved", origin) + return + + path = str(resolved_path) + + instance = t.cast(ModuleWatchdog, cls._instance) + if path not in instance._hook_map: + log.warning("No hooks registered for origin %s", origin) + return + + try: + if path in instance._hook_map: + hooks = instance._hook_map[path] + hooks.remove(hook) + if not hooks: + del instance._hook_map[path] + except ValueError: + log.warning("Hook %r not registered for origin %s", hook, origin) + return + + @classmethod + def register_module_hook(cls, module: str, hook: ModuleHookType) -> None: + """Register a hook to be called when the module with the given name is + imported. + + The hook will be called with the module object as argument. + """ + cls.install() + + log.debug("Registering hook '%r' on module '%s'", hook, module) + instance = t.cast(ModuleWatchdog, cls._instance) + instance._hook_map[module].append(hook) + try: + module_object = sys.modules[module] + except KeyError: + # The module is not loaded yet. Nothing more we can do. + return + + # The module was already imported so we invoke the hook straight-away + log.debug("Calling hook '%r' on already imported module '%s'", hook, module) + hook(module_object) + + @classmethod + def unregister_module_hook(cls, module: str, hook: ModuleHookType) -> None: + """Unregister the hook registered with the given module name and + argument. + """ + if not cls.is_installed(): + return + + instance = t.cast(ModuleWatchdog, cls._instance) + if module not in instance._hook_map: + log.warning("No hooks registered for module %s", module) + return + + try: + if module in instance._hook_map: + hooks = instance._hook_map[module] + hooks.remove(hook) + if not hooks: + del instance._hook_map[module] + except ValueError: + log.warning("Hook %r not registered for module %r", hook, module) + return + + @classmethod + def after_module_imported(cls, module: str) -> t.Callable[[ModuleHookType], None]: + def _(hook: ModuleHookType) -> None: + cls.register_module_hook(module, hook) + + return _ + + @classmethod + def register_pre_exec_module_hook( + cls: type["ModuleWatchdog"], cond: PreExecHookCond, hook: PreExecHookType + ) -> None: + """Register a hook to execute before/instead of exec_module. + + The pre exec_module hook is executed before the module is executed + to allow for changed modules to be executed as needed. To ensure + that the hook is applied only to the modules that are required, + the condition is evaluated against the module name. + """ + cls.install() + + log.debug("Registering pre_exec module hook '%r' on condition '%s'", hook, cond) + instance = t.cast(ModuleWatchdog, cls._instance) + instance._pre_exec_module_hooks.add((cond, hook)) + + @classmethod + def remove_pre_exec_module_hook(cls: type["ModuleWatchdog"], cond: PreExecHookCond, hook: PreExecHookType) -> None: + """Register a hook to execute before/instead of exec_module. Only for testing proposes""" + instance = t.cast(ModuleWatchdog, cls._instance) + instance._pre_exec_module_hooks.remove((cond, hook)) + + @classmethod + def register_import_exception_hook( + cls: type["ModuleWatchdog"], cond: ImportExceptionHookCond, hook: ImportExceptionHookType + ): + cls.install() + + instance = t.cast(ModuleWatchdog, cls._instance) + instance._import_exception_hooks.add((cond, hook)) + + +class LazyWrappingContext(WrappingContext): + def __return__(self, value: t.Any) -> t.Any: + # Update the global (i.e. the module) scope with the local scope of the + # wrapped function. + self.__frame__.f_globals.update(self.__frame__.f_locals) + + return super().__return__(value) + + +def lazy(f: t.Callable[[], None]) -> None: + LazyWrappingContext(t.cast(FunctionType, f)).wrap() + + _globals = sys._getframe(1).f_globals + + def __getattr__(name: str) -> t.Any: + f() + try: + return _globals[name] + except KeyError: + h = AttributeError(f"module {_globals['__name__']!r} has no attribute {name!r}") + h.__suppress_context__ = True + raise h + + _globals["__getattr__"] = __getattr__ diff --git a/ddtrace/internal/schema/__init__.py b/ddtrace/internal/utils/schema/__init__.py similarity index 100% rename from ddtrace/internal/schema/__init__.py rename to ddtrace/internal/utils/schema/__init__.py diff --git a/ddtrace/internal/schema/processor.py b/ddtrace/internal/utils/schema/processor.py similarity index 100% rename from ddtrace/internal/schema/processor.py rename to ddtrace/internal/utils/schema/processor.py diff --git a/ddtrace/internal/schema/span_attribute_schema.py b/ddtrace/internal/utils/schema/span_attribute_schema.py similarity index 98% rename from ddtrace/internal/schema/span_attribute_schema.py rename to ddtrace/internal/utils/schema/span_attribute_schema.py index c9b717f6809..b42c0211b9f 100644 --- a/ddtrace/internal/schema/span_attribute_schema.py +++ b/ddtrace/internal/utils/schema/span_attribute_schema.py @@ -2,8 +2,8 @@ import sys from typing import Optional -from ddtrace.internal.constants import DEFAULT_SERVICE_NAME from ddtrace.internal.settings._inferred_base_service import detect_service +from ddtrace.internal.utils.constants import DEFAULT_SERVICE_NAME class SpanDirection(Enum): diff --git a/ddtrace/internal/utils/threads.py b/ddtrace/internal/utils/threads.py new file mode 100644 index 00000000000..faa47fdac3f --- /dev/null +++ b/ddtrace/internal/utils/threads.py @@ -0,0 +1,218 @@ +from time import monotonic_ns +import typing as t + +from ddtrace.internal import forksafe +from ddtrace.internal._threads import PeriodicThread as _PeriodicThread +from ddtrace.internal._threads import periodic_threads +from ddtrace.internal.utils.logger import get_logger + + +log = get_logger(__name__) + +# We try to import the stdlib locks from the _thread module, where they are +# implemented in C for CPython for most platforms. If that fails, we fall back +# to the threading module, which provides a pure Python implementation that +# should work on all platforms. We also make sure to grab a reference to the +# original lock classes, in case they get patched by monkey-patching libraries +# like gevent. +try: + from _thread import allocate_lock as Lock +except ImportError: + from threading import Lock + +try: + from _thread import RLock +except ImportError: + from threading import RLock + + +__all__ = [ + "Lock", + "PeriodicThread", + "RLock", +] + + +# Forking state management. This is a barrier to either prevent new threads +# from being started while forking, or to allow a thread to be started +# completely if a fork comes in the middle of it. +_forking = False +_forking_lock = Lock() + + +class BoundMethod(t.Protocol): + __self__: t.Any + + def __call__(self) -> None: ... + + +# List of threads that have requested to be started while forking. These will +# be started after the fork is complete. +_threads_to_start_after_fork: list[BoundMethod] = [] + + +def _safe_restart(start: t.Callable[[], None], name: t.Optional[str] = None) -> None: + """Invoke a post-fork thread-start callable, logging resource errors instead of raising. + + The native layer translates pthread_create failures (EAGAIN, ENOMEM) into + OSError. Post-fork restart is triggered automatically by forksafe hooks — + there is no explicit caller that can handle the error, so losing a + periodic thread to resource exhaustion must not crash the host. + Explicit start() calls let OSError propagate so the caller can react. + """ + try: + start() + except Exception as e: + log.error("failed to start periodic thread %s: %s", name, e) + + +class PeriodicThread(_PeriodicThread): + """A fork-safe periodic thread.""" + + __autorestart__ = True + + def start(self) -> None: + with _forking_lock: + # We cannot start a new thread while we are forking, because we are + # trying to stop them all. In that case, we take note of the thread + # and start it after the fork. + if not _forking: + super().start() + else: + _threads_to_start_after_fork.append(t.cast(BoundMethod, super().start)) + + +# Set of running periodic threads that need to be restarted after a fork. +_threads_to_restart_after_fork: set[_PeriodicThread] = set() + + +# A typical scenario is that of forking worker threads in a loop. For the +# parent process, this would mean having to stop and restart the threads in +# between forks, which is not ideal. Instead, we can use a timer to restart +# the threads after a certain amount of time has passed since the last fork. +# This way, we can avoid stopping and restarting the threads in between forks. +class ThreadRestartTimer(PeriodicThread): + __timeout__ = int(1e8) # nanoseconds + + _instance: t.Optional["ThreadRestartTimer"] = None + _timestamp = 0 + + def __init__(self): + super().__init__(self.__timeout__ / 1e9, self._restart_threads, name=f"{__name__}:{self.__class__.__name__}") + + def _restart_threads(self) -> None: + # Restart the threads after we have stopped calling fork for a while. + with _forking_lock: + # If we are forking, we will try again later. + if _forking: + return + + # If we haven't have calls to fork for a while, we can restart the + # threads. This way we avoid stopping and restarting the threads + # in between forks. + if monotonic_ns() >= self._timestamp: # 100ms + for thread in _threads_to_restart_after_fork.copy(): + if isinstance(thread, ThreadRestartTimer): + # Skip any ThreadRestartTimer instance, + # to avoid restarting orphaned timer instances that were + # caught in periodic_threads during a fork. + continue + log.debug("Restarting thread %s after fork", thread.name) + try: + thread._after_fork(force=True) + except Exception as e: + log.error("failed to restart periodic thread %s after fork: %s", thread.name, e) + _threads_to_restart_after_fork.clear() + + for thread_start in _threads_to_start_after_fork: + log.debug("Starting thread %s after fork", thread_start.__self__.name) + _safe_restart(thread_start, thread_start.__self__.name) + _threads_to_start_after_fork.clear() + + # We no longer need this thread so we clear it. + self.clear() + + @classmethod + def clear(cls): + """Clear the timer and stop it if it is running.""" + if cls._instance is not None: + cls._instance.stop() + cls._instance = None + + @classmethod + def touch(cls): + """Set the new expiration time for the timer.""" + cls._timestamp = monotonic_ns() + cls.__timeout__ + + @classmethod + def set(cls): + """Set the timer to restart the threads after a fork.""" + if cls._instance is None: + cls._instance = cls() + cls._instance.start() + else: + # We have already created the timer, so we let the forksafe logic + # handle the restart instead of creating a new instance. + cls._instance._after_fork() + + +@forksafe.register +def _after_fork_child(): + global _forking + + _forking = False + + # Restart the threads immediately. It is unlikely that there will be another + # call to fork here. _after_fork() (without force=True) respects + # __autorestart__: cleanup always runs, but the thread is only restarted + # when __autorestart__ is True. This is intentional in the child — threads + # with __autorestart__ = False (e.g. RemoteConfigPoller) should not run in + # forked workers. + for thread in _threads_to_restart_after_fork.copy(): + log.debug("Restarting thread %s after fork in child", thread.name) + try: + thread._after_fork() + except Exception as e: + log.error("failed to restart periodic thread %s after fork in child: %s", thread.name, e) + _threads_to_restart_after_fork.clear() + + for thread_start in _threads_to_start_after_fork.copy(): + log.debug("Starting thread %s after fork in child", thread_start.__self__.name) + _safe_restart(thread_start, thread_start.__self__.name) + _threads_to_start_after_fork.clear() + + +@forksafe.register_after_parent +def _after_fork_parent() -> None: + global _forking + + _forking = False + + if _threads_to_restart_after_fork or _threads_to_start_after_fork: + ThreadRestartTimer.set() + + +@forksafe.register_before_fork +def _before_fork() -> None: + global _threads_to_restart_after_fork, _forking_lock, _forking + + ThreadRestartTimer.touch() + + with _forking_lock: + _forking = True + + # Take note of all the periodic threads that are running and will need to be + # restarted. + _threads_to_restart_after_fork.update(periodic_threads.values()) + + # Stop all the periodic threads that are still running, without executing + # the shutdown methods, if any. This ensures that we can stop the threads + # more promptly. + for thread in _threads_to_restart_after_fork: + log.debug("Stopping thread %s before fork", thread.name) + thread._before_fork() + + # Join all the threads to ensure they are stopped before the fork. + for thread in _threads_to_restart_after_fork: + log.debug("Joining thread %s before fork", thread.name) + thread.join() diff --git a/ddtrace/internal/wrapping/__init__.py b/ddtrace/internal/utils/wrapping/__init__.py similarity index 76% rename from ddtrace/internal/wrapping/__init__.py rename to ddtrace/internal/utils/wrapping/__init__.py index b01b8c7977d..758cb916834 100644 --- a/ddtrace/internal/wrapping/__init__.py +++ b/ddtrace/internal/utils/wrapping/__init__.py @@ -4,44 +4,28 @@ from typing import Any from typing import Callable from typing import Iterator -from typing import MutableMapping from typing import Optional from typing import Protocol from typing import cast -import weakref import bytecode as bc from bytecode import Instr from ddtrace.internal.assembly import Assembly -from ddtrace.internal.threads import Lock -from ddtrace.internal.wrapping.asyncs import wrap_async -from ddtrace.internal.wrapping.generators import wrap_generator +from ddtrace.internal.utils.inspection import link_function_to_code +from ddtrace.internal.utils.wrapping.asyncs import wrap_async +from ddtrace.internal.utils.wrapping.generators import wrap_generator PY = sys.version_info[:2] -# Maps each wrapped function to its inner copy (the singly-linked list of -# wrapping layers). WeakKeyDictionary so functions are not kept alive by the -# registry alone. -_wrapped: weakref.WeakKeyDictionary[FunctionType, FunctionType] = weakref.WeakKeyDictionary() -_wrapped_lock = Lock() - -# Maps original code objects to the functions that own them. Written by -# link_function_to_code; read by functions_for_code in inspection.py. -# WeakValueDictionary so that wrapped ephemeral functions are not kept alive by -# this mapping alone — inspection.py falls back to gc.get_referrers on a miss. -_code_to_fn: MutableMapping[CodeType, FunctionType] = weakref.WeakValueDictionary() - - -def link_function_to_code(code: CodeType, function: FunctionType) -> None: - """Link a function to its original code object for fast reverse lookup.""" - _code_to_fn[code] = function - class WrappedFunction(Protocol): """A wrapped function.""" + __dd_wrapped__: Optional[FunctionType] = None + __dd_wrappers__: Optional[dict[Any, Any]] = None + def __call__(self, *args: Any, **kwargs: Any) -> Any: pass @@ -311,12 +295,11 @@ def wrap(f: FunctionType, wrapper: Wrapper) -> WrappedFunction: f.__defaults__, f.__closure__, ) - - # Carry forward the existing wrapped-function link to the new inner copy. - with _wrapped_lock: - existing_inner = _wrapped.get(f) - if existing_inner is not None: - _wrapped[wrapped] = existing_inner + try: + wf = cast(WrappedFunction, f) + cast(WrappedFunction, wrapped).__dd_wrapped__ = cast(FunctionType, wf.__dd_wrapped__) + except AttributeError: + pass wrapped.__kwdefaults__ = f.__kwdefaults__ @@ -346,54 +329,47 @@ def wrap(f: FunctionType, wrapper: Wrapper) -> WrappedFunction: # Replace the function code with the trampoline bytecode f.__code__ = wrapped_code.to_code() - # DEV: Multiple wrapping is implemented as a singly-linked list via _wrapped. - with _wrapped_lock: - _wrapped[f] = wrapped + # DEV: Multiple wrapping is implemented as a singly-linked list via the + # __dd_wrapped__ attribute. + wf = cast(WrappedFunction, f) + wf.__dd_wrapped__ = wrapped # Link the original code object to the original function link_function_to_code(code, f) - return cast(WrappedFunction, f) - - -def _unwrap_method(f: Any) -> Any: - """Return the underlying function if *f* is a bound or unbound method.""" - func = getattr(f, "__func__", f) - return func + return wf def is_wrapped(f: FunctionType) -> bool: """Check if a function is wrapped with any wrapper.""" try: - with _wrapped_lock: - inner = _wrapped.get(_unwrap_method(f)) - except TypeError: - # f is not weakly referenceable (e.g. C method_descriptor) - return False - if inner is None: + wf = cast(WrappedFunction, f) + inner = cast(FunctionType, wf.__dd_wrapped__) + + # Sanity check + assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + return True + except AttributeError: return False - assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - return True def is_wrapped_with(f: FunctionType, wrapper: Wrapper) -> bool: """Check if a function is wrapped with a specific wrapper.""" - f = cast(FunctionType, _unwrap_method(f)) try: - with _wrapped_lock: - inner = _wrapped.get(f) - except TypeError: - return False - if inner is None: - return False + wf = cast(WrappedFunction, f) + inner = cast(FunctionType, wf.__dd_wrapped__) - assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + # Sanity check + assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - if wrapper in f.__code__.co_consts: - return True + if wrapper in f.__code__.co_consts: + return True - # This is not the correct wrapping layer. Try with the next one. - return is_wrapped_with(inner, wrapper) + # This is not the correct wrapping layer. Try with the next one. + return is_wrapped_with(inner, wrapper) + + except AttributeError: + return False def unwrap(wf: WrappedFunction, wrapper: Wrapper) -> FunctionType: @@ -403,52 +379,41 @@ def unwrap(wf: WrappedFunction, wrapper: Wrapper) -> FunctionType: this will unwrap the one that uses ``wrapper``. If the function was not wrapped with ``wrapper``, it will return the first argument. """ - # DEV: Multiple wrapping layers are singly-linked via _wrapped. When we - # find the layer that needs to be removed we also have to ensure that we + # DEV: Multiple wrapping layers are singly-linked via __dd_wrapped__. When + # we find the layer that needs to be removed we also have to ensure that we # update the link at the deletion site if there is a non-empty tail. - f = cast(FunctionType, wf) - with _wrapped_lock: - inner = _wrapped.get(f) - if inner is None: + try: + inner = cast(FunctionType, wf.__dd_wrapped__) + except AttributeError: # The function is not wrapped so we return it as is. - return f + return cast(FunctionType, wf) + # Sanity check assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - if wrapper not in f.__code__.co_consts: + if wrapper not in cast(FunctionType, wf).__code__.co_consts: # This is not the correct wrapping layer. Try with the next one. return unwrap(cast(WrappedFunction, inner), wrapper) - # Remove the current wrapping layer by moving the next one over the current - # one. The code swap and registry update must be atomic: two concurrent - # unwrap calls on the same function and wrapper must not both succeed. - with _wrapped_lock: - f.__code__ = inner.__code__ - next_inner = _wrapped.get(inner) - if next_inner is not None: - _wrapped[f] = next_inner - else: - del _wrapped[f] + # Remove the current wrapping layer by moving the next one over the + # current one. + f = cast(FunctionType, wf) + f.__code__ = inner.__code__ + + try: + # Update the link to the next layer. + wf.__dd_wrapped__ = cast(WrappedFunction, inner).__dd_wrapped__ # type: ignore[assignment] + except AttributeError: + # No more wrapping layers. Restore the original function by removing + # this extra attribute. + del wf.__dd_wrapped__ return f def get_function_code(f: FunctionType) -> CodeType: - with _wrapped_lock: - inner = _wrapped.get(f) - return (inner if inner is not None else f).__code__ + return (cast(WrappedFunction, f).__dd_wrapped__ or f if is_wrapped(f) else f).__code__ def set_function_code(f: FunctionType, code: CodeType) -> None: - with _wrapped_lock: - inner = _wrapped.get(f) - (inner if inner is not None else f).__code__ = code - - -def get_wrapped(f: FunctionType) -> Optional[FunctionType]: - """Return the inner bytecode copy of *f* if it is wrapped, else None.""" - try: - with _wrapped_lock: - return _wrapped.get(cast(FunctionType, _unwrap_method(f))) - except TypeError: - return None + (cast(WrappedFunction, f).__dd_wrapped__ or f if is_wrapped(f) else f).__code__ = code # type: ignore[misc] diff --git a/ddtrace/internal/wrapping/asyncs.py b/ddtrace/internal/utils/wrapping/asyncs.py similarity index 100% rename from ddtrace/internal/wrapping/asyncs.py rename to ddtrace/internal/utils/wrapping/asyncs.py diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/utils/wrapping/context.py similarity index 50% rename from ddtrace/internal/wrapping/context.py rename to ddtrace/internal/utils/wrapping/context.py index 8333b997c36..d3efd0ef196 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/utils/wrapping/context.py @@ -8,71 +8,21 @@ from types import TracebackType import typing as t from typing import Protocol # noqa:F401 -import weakref import bytecode from bytecode import Bytecode from ddtrace.internal.assembly import Assembly -from ddtrace.internal.logger import get_logger -from ddtrace.internal.threads import Lock -from ddtrace.internal.threads import RLock -from ddtrace.internal.wrapping import WrappedFunction -from ddtrace.internal.wrapping import Wrapper -from ddtrace.internal.wrapping import get_function_code -from ddtrace.internal.wrapping import is_wrapped_with -from ddtrace.internal.wrapping import link_function_to_code -from ddtrace.internal.wrapping import set_function_code -from ddtrace.internal.wrapping import unwrap -from ddtrace.internal.wrapping import wrap - - -class _ContextRecord: - """Holds all wrapping-context metadata for a single function. - - Stores only weak references to context objects so that the WeakKeyDictionary - key (the function) is not kept alive by the registry value chain: - _registry -> _ContextRecord -> context -> context.__wrapped__ -> function - Breaking this path with weak references lets ephemeral functions be garbage - collected as soon as all external strong references drop. - """ - - __slots__ = ("_uwc_ref", "lazy_contexts") - - def __init__(self) -> None: - self._uwc_ref: t.Optional[weakref.ref["_UniversalWrappingContext"]] = None - # WeakSet so that LazyWrappingContext instances (which also hold - # __wrapped__ = f) do not prevent the function from being collected. - self.lazy_contexts: weakref.WeakSet["LazyWrappingContext"] = weakref.WeakSet() - - @property - def uwc(self) -> t.Optional["_UniversalWrappingContext"]: - ref = self._uwc_ref - return ref() if ref is not None else None - - @uwc.setter - def uwc(self, value: t.Optional["_UniversalWrappingContext"]) -> None: - self._uwc_ref = weakref.ref(value) if value is not None else None - - @classmethod - def get_or_create(cls, f: FunctionType) -> "_ContextRecord": - record = _registry.get(f) - if record is None: - with _registry_lock: - record = _registry.get(f) - if record is None: - record = cls() - _registry[f] = record - return record - - -# Per-function registry for wrapping-context machinery. WeakKeyDictionary so -# functions are not kept alive by the registry alone. Storing data here instead -# of as function attributes keeps __dict__ clean, preventing frameworks that -# copy function __dict__ (e.g. functools.wraps, -# self.__dict__.update(f.__dict__)) from capturing non-picklable objects. -_registry: weakref.WeakKeyDictionary[FunctionType, _ContextRecord] = weakref.WeakKeyDictionary() -_registry_lock = RLock() +from ddtrace.internal.utils.inspection import link_function_to_code +from ddtrace.internal.utils.logger import get_logger +from ddtrace.internal.utils.threads import Lock +from ddtrace.internal.utils.wrapping import WrappedFunction +from ddtrace.internal.utils.wrapping import Wrapper +from ddtrace.internal.utils.wrapping import get_function_code +from ddtrace.internal.utils.wrapping import is_wrapped_with +from ddtrace.internal.utils.wrapping import set_function_code +from ddtrace.internal.utils.wrapping import unwrap +from ddtrace.internal.utils.wrapping import wrap log = get_logger(__name__) @@ -343,27 +293,22 @@ class BaseWrappingContext(ABC): __priority__: int = 0 def __init__(self, f: FunctionType): - # Store a weak reference so that context objects do not keep the wrapped - # function alive. CodeType is not GC-tracked in CPython, so the cycle - # f → code.co_consts → bound_methods(uwc) → uwc.__wrapped__ → f - # cannot be broken by the cyclic GC. A weak ref here allows f's - # reference count to reach zero (and be freed) as soon as all external - # strong refs drop, without relying on the cyclic GC at all. - self._wrapped_ref: weakref.ref[FunctionType] = weakref.ref(f) + self.__wrapped__ = f self._storage: ContextVar[t.Optional[dict[str, t.Any]]] = ContextVar( f"{type(self).__name__}__storage", default=None ) - @property - def __wrapped__(self) -> FunctionType: - f = self._wrapped_ref() - if f is None: - raise RuntimeError(f"{type(self).__name__}.__wrapped__: the wrapped function has been garbage collected") - return f + def __getstate__(self) -> dict[str, t.Any]: + state = self.__dict__.copy() + state.pop("_storage", None) # remove unpicklable field + return state - @__wrapped__.setter - def __wrapped__(self, f: FunctionType) -> None: - self._wrapped_ref = weakref.ref(f) + def __setstate__(self, state: dict[str, t.Any]) -> None: + self.__dict__.update(state) + self._storage = ContextVar( + f"{type(self).__name__}__storage", + default=None, + ) def __enter__(self) -> "BaseWrappingContext": prev = self._storage.get() @@ -397,10 +342,10 @@ def set(self, key: str, value: T) -> T: @classmethod def wrapped(cls, f: FunctionType) -> "BaseWrappingContext": - try: + if cls.is_wrapped(f): context = cls.extract(f) assert isinstance(context, cls) # nosec - except ValueError: + else: context = cls(f) context.wrap() return context @@ -425,10 +370,7 @@ class WrappingContext(BaseWrappingContext): @property def __frame__(self) -> FrameType: try: - return t.cast( - FrameType, - _UniversalWrappingContext.extract(t.cast(FunctionType, self.__wrapped__)).get("__frame__"), - ) + return t.cast(FrameType, _UniversalWrappingContext.extract(self.__wrapped__).get("__frame__")) except ValueError: raise AttributeError("Wrapping context not entered") @@ -444,24 +386,31 @@ def is_wrapped(cls, f: FunctionType) -> bool: @classmethod def extract(cls, f: FunctionType) -> "WrappingContext": - try: - return _UniversalWrappingContext.extract(f).registered(cls) - except (ValueError, KeyError): - msg = f"Function is not wrapped with {cls}" - raise ValueError(msg) + if _UniversalWrappingContext.is_wrapped(f): + try: + return _UniversalWrappingContext.extract(f).registered(cls) + except KeyError: + pass + msg = f"Function is not wrapped with {cls}" + raise ValueError(msg) def wrap(self) -> None: - t.cast( - _UniversalWrappingContext, _UniversalWrappingContext.wrapped(t.cast(FunctionType, self.__wrapped__)) - ).register(self) + t.cast(_UniversalWrappingContext, _UniversalWrappingContext.wrapped(self.__wrapped__)).register(self) def unwrap(self) -> None: - f = t.cast(FunctionType, self.__wrapped__) + f = self.__wrapped__ - try: + if _UniversalWrappingContext.is_wrapped(f): _UniversalWrappingContext.extract(f).unregister(self) - except ValueError: - pass + + +class LazyWrappedFunction(Protocol): + """A lazy-wrapped function.""" + + __dd_lazy_contexts__: list[WrappingContext] + + def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: + pass class LazyWrappingContext(WrappingContext): @@ -473,11 +422,10 @@ def __init__(self, f: FunctionType): @classmethod def is_wrapped(cls, f: FunctionType) -> bool: - with _registry_lock: - record = _registry.get(f) - if record is None: - return False - return any(isinstance(c, cls) for c in record.lazy_contexts) + try: + return any(isinstance(c, cls) for c in t.cast(LazyWrappedFunction, f).__dd_lazy_contexts__) + except AttributeError: + return False def wrap(self) -> None: """Perform the bytecode wrapping on first invocation.""" @@ -485,60 +433,71 @@ def wrap(self) -> None: if self._trampoline is not None: return - # If the function is already universally wrapped it's less expensive + # If the function is already universally wrapped so it's less expensive # to do the normal wrapping. - if _UniversalWrappingContext.is_wrapped(t.cast(FunctionType, self.__wrapped__)): + if _UniversalWrappingContext.is_wrapped(self.__wrapped__): super().wrap() return def trampoline(_: t.Any, args: tuple[t.Any, ...], kwargs: dict[str, t.Any]) -> t.Any: with tl: f = t.cast(WrappedFunction, self.__wrapped__) - if is_wrapped_with(t.cast(FunctionType, self.__wrapped__), trampoline): + if is_wrapped_with(self.__wrapped__, trampoline): f = t.cast(WrappedFunction, unwrap(f, trampoline)) self._trampoline = None - inconsistent = False - with _registry_lock: - record = _registry.get(t.cast(FunctionType, f)) - if record is not None: - inconsistent = self not in record.lazy_contexts - record.lazy_contexts.discard(self) - if not record.lazy_contexts and record.uwc is None: - _registry.pop(t.cast(FunctionType, f), None) - if inconsistent: + try: + (cs := t.cast(LazyWrappedFunction, f).__dd_lazy_contexts__).remove(self) + if not cs: + del t.cast(LazyWrappedFunction, f).__dd_lazy_contexts__ + except (AttributeError, ValueError): log.warning("Inconsistent lazy wrapping context state") super(LazyWrappingContext, self).wrap() return f(*args, **kwargs) - wrap(t.cast(FunctionType, self.__wrapped__), trampoline) + wrap(self.__wrapped__, trampoline) self._trampoline = trampoline - _ContextRecord.get_or_create(t.cast(FunctionType, self.__wrapped__)).lazy_contexts.add(self) + wf = t.cast(LazyWrappedFunction, self.__wrapped__) + if not hasattr(wf, "__dd_lazy_contexts__"): + wf.__dd_lazy_contexts__ = [] + wf.__dd_lazy_contexts__.append(self) def unwrap(self) -> None: with self._trampoline_lock: - if _UniversalWrappingContext.is_wrapped(t.cast(FunctionType, self.__wrapped__)): + if _UniversalWrappingContext.is_wrapped(self.__wrapped__): assert self._trampoline is None # nosec super().unwrap() elif self._trampoline is not None: - with _registry_lock: - record = _registry.get(t.cast(FunctionType, self.__wrapped__)) - if record is not None: - record.lazy_contexts.discard(self) - if not record.lazy_contexts and record.uwc is None: - _registry.pop(t.cast(FunctionType, self.__wrapped__), None) + wf = t.cast(LazyWrappedFunction, self.__wrapped__) + if hasattr(wf, "__dd_lazy_contexts__"): + wf.__dd_lazy_contexts__.remove(self) + if not wf.__dd_lazy_contexts__: + del wf.__dd_lazy_contexts__ unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline) self._trampoline = None + def __getstate__(self) -> dict[str, t.Any]: + state = super().__getstate__() + state.pop("_trampoline_lock", None) # thread lock not picklable + state.pop("_trampoline", None) # closure not picklable + return state + + def __setstate__(self, state: dict[str, t.Any]) -> None: + super().__setstate__(state) + self._trampoline_lock = Lock() + self._trampoline = None + class ContextWrappedFunction(Protocol): """A wrapped function.""" + __dd_context_wrapped__ = None # type: t.Optional[_UniversalWrappingContext] + def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: pass @@ -614,263 +573,252 @@ def __return__(self, value: T) -> T: @classmethod def is_wrapped(cls, f: FunctionType) -> bool: try: - with _registry_lock: - record = _registry.get(f) - if record is None or record.uwc is None: - return False - # Verify the registry entry matches actual bytecode wrapping. - if sys.version_info >= (3, 11): - return record.uwc.__enter__ in get_function_code(f).co_consts - else: - return record.uwc in get_function_code(f).co_consts + # Check that we have actual bytecode wrapping. The presence of the + # __dd_context_wrapped__ attribute is not enough, as this could be + # copied over from an object state cloning. + if sys.version_info >= (3, 11): + return f.__dd_context_wrapped__.__enter__ in get_function_code(f).co_consts # type: ignore + else: + return f.__dd_context_wrapped__ in get_function_code(f).co_consts # type: ignore except AttributeError: return False @classmethod def extract(cls, f: FunctionType) -> "_UniversalWrappingContext": - with _registry_lock: - if not cls.is_wrapped(f): - raise ValueError("Function is not wrapped") - return t.cast(_UniversalWrappingContext, _registry[f].uwc) + if not cls.is_wrapped(f): + raise ValueError("Function is not wrapped") + return t.cast(_UniversalWrappingContext, t.cast(ContextWrappedFunction, f).__dd_context_wrapped__) if sys.version_info >= (3, 11): def wrap(self) -> None: f = self.__wrapped__ - with _registry_lock: - if self.is_wrapped(f): - raise ValueError("Function already wrapped") + if self.is_wrapped(f): + raise ValueError("Function already wrapped") + + bc = Bytecode.from_code(code := get_function_code(f)) + + # Prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + try: + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN.bind({"context_return": self.__return__}, lineno=instr.lineno) + elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ + return_code = CONTEXT_RETURN_CONST.bind( + {"context_return": self.__return__, "value": instr.arg}, lineno=instr.lineno + ) + else: + return_code = [] + + bc[i:i] = return_code + i += len(return_code) + except AttributeError: + # Not an instruction + pass + i += 1 + + # Search for the RESUME instruction + for i, instr in enumerate(bc, 1): + try: + if instr.name == "RESUME": + break + except AttributeError: + # Not an instruction + pass + else: + i = 0 - bc = Bytecode.from_code(code := get_function_code(f)) + bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=code.co_firstlineno) - # Prefix every return - i = 0 - while i < len(bc): - instr = bc[i] - try: - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN.bind({"context_return": self.__return__}, lineno=instr.lineno) - elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ - return_code = CONTEXT_RETURN_CONST.bind( - {"context_return": self.__return__, "value": instr.arg}, lineno=instr.lineno - ) - else: - return_code = [] + # Wrap every line outside a try block + except_label = bytecode.Label() + first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) - bc[i:i] = return_code - i += len(return_code) - except AttributeError: - # Not an instruction - pass + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None: + bc.insert(i, bytecode.TryEnd(last_try_begin)) + last_try_begin = None i += 1 - - # Search for the RESUME instruction - for i, instr in enumerate(bc, 1): - try: - if instr.name == "RESUME": + elif isinstance(instr, bytecode.TryEnd): + j = i + 1 + while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin): + if isinstance(bc[j], bytecode.Instr): + last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) + bc.insert(i + 1, last_try_begin) break - except AttributeError: - # Not an instruction - pass - else: - i = 0 + j += 1 + i += 1 + i += 1 - bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=code.co_firstlineno) + bc.insert(0, first_try_begin) - # Wrap every line outside a try block - except_label = bytecode.Label() - first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) + bc.append(bytecode.TryEnd(last_try_begin)) + bc.append(except_label) + bc.extend(CONTEXT_FOOT.bind({"context_exit": self._exit}, lineno=code.co_firstlineno)) - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None: - bc.insert(i, bytecode.TryEnd(last_try_begin)) - last_try_begin = None - i += 1 - elif isinstance(instr, bytecode.TryEnd): - j = i + 1 - while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin): - if isinstance(bc[j], bytecode.Instr): - last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) - bc.insert(i + 1, last_try_begin) - break - j += 1 - i += 1 - i += 1 + # Mark the function as wrapped by a wrapping context + t.cast(ContextWrappedFunction, f).__dd_context_wrapped__ = self - bc.insert(0, first_try_begin) + # Replace the function code with the wrapped code. We also link + # the function to its original code object so that we can retrieve + # it later if required. + link_function_to_code(code, f) - bc.append(bytecode.TryEnd(last_try_begin)) - bc.append(except_label) - bc.extend(CONTEXT_FOOT.bind({"context_exit": self._exit}, lineno=code.co_firstlineno)) - - # Register the wrapping context and write the new bytecode. - _ContextRecord.get_or_create(f).uwc = self - link_function_to_code(code, f) - set_function_code(f, bc.to_code()) + set_function_code(f, bc.to_code()) def unwrap(self) -> None: f = self.__wrapped__ - with _registry_lock: - if not self.is_wrapped(f): - return - - wc = _registry[f].uwc - - bc = Bytecode.from_code(get_function_code(f)) + if not self.is_wrapped(f): + return - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() - bc.pop() + wrapped = t.cast(ContextWrappedFunction, f) - except_label = bc.pop(0).target + bc = Bytecode.from_code(get_function_code(f)) - # Remove the try blocks - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.TryBegin) and instr.target is except_label: - bc.pop(i) - elif isinstance(instr, bytecode.TryEnd) and instr.entry.target is except_label: - bc.pop(i) - else: - i += 1 + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() + bc.pop() - # Remove the head of the try block - for i, instr in enumerate(bc): - if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: - break + except_label = bc.pop(0).target - # Search for the RESUME instruction - for i, instr in enumerate(bc, 1): - try: - if instr.name == "RESUME": - break - except AttributeError: - # Not an instruction - pass + # Remove the try blocks + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.TryBegin) and instr.target is except_label: + bc.pop(i) + elif isinstance(instr, bytecode.TryEnd) and instr.entry.target is except_label: + bc.pop(i) else: - i = 0 + i += 1 - bc[i : i + len(CONTEXT_HEAD)] = [] + # Remove the head of the try block + wc = wrapped.__dd_context_wrapped__ + for i, instr in enumerate(bc): + if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: + break - # Un-prefix every return + # Search for the RESUME instruction + for i, instr in enumerate(bc, 1): + try: + if instr.name == "RESUME": + break + except AttributeError: + # Not an instruction + pass + else: i = 0 - while i < len(bc): - instr = bc[i] - try: - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN - elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ - return_code = CONTEXT_RETURN_CONST - else: - return_code = None - - if return_code is not None: - bc[i - len(return_code) : i] = [] - i -= len(return_code) - except AttributeError: - # Not an instruction - pass - i += 1 - # Recreate the code object - set_function_code(f, bc.to_code()) + bc[i : i + len(CONTEXT_HEAD)] = [] + + # Un-prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + try: + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN + elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ + return_code = CONTEXT_RETURN_CONST + else: + return_code = None + + if return_code is not None: + bc[i - len(return_code) : i] = [] + i -= len(return_code) + except AttributeError: + # Not an instruction + pass + i += 1 + + # Recreate the code object + set_function_code(f, bc.to_code()) - # Clear the UWC from the registry; remove the record if fully empty. - record = _registry.get(f) - if record is not None: - record.uwc = None - if not record.lazy_contexts: - _registry.pop(f, None) + # Remove the wrapping context marker + del wrapped.__dd_context_wrapped__ else: def wrap(self) -> None: - f = t.cast(FunctionType, self.__wrapped__) + f = self.__wrapped__ - with _registry_lock: - if self.is_wrapped(f): - raise ValueError("Function already wrapped") + if self.is_wrapped(f): + raise ValueError("Function already wrapped") - bc = Bytecode.from_code(code := get_function_code(f)) + bc = Bytecode.from_code(code := get_function_code(f)) - # Prefix every return - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.Instr): - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN.bind({"context": self}, lineno=instr.lineno) - bc[i:i] = return_code - i += len(return_code) - i += 1 - - # Search for the GEN_START instruction, which needs to stay on top. - i = 0 - if sys.version_info >= (3, 10) and (iscoroutinefunction(f) or isgeneratorfunction(f)): - for i, instr in enumerate(bc, 1): - if isinstance(instr, bytecode.Instr) and instr.name == "GEN_START": - break + # Prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.Instr): + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN.bind({"context": self}, lineno=instr.lineno) + bc[i:i] = return_code + i += len(return_code) + i += 1 - *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": self}, lineno=code.co_firstlineno) + # Search for the GEN_START instruction, which needs to stay on top. + i = 0 + if sys.version_info >= (3, 10) and (iscoroutinefunction(f) or isgeneratorfunction(f)): + for i, instr in enumerate(bc, 1): + if isinstance(instr, bytecode.Instr) and instr.name == "GEN_START": + break - bc.append(except_label) - bc.extend(CONTEXT_FOOT.bind(lineno=code.co_firstlineno)) + *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": self}, lineno=code.co_firstlineno) - # Register the wrapping context and write the new bytecode. - _ContextRecord.get_or_create(f).uwc = self - link_function_to_code(code, f) - set_function_code(f, bc.to_code()) + bc.append(except_label) + bc.extend(CONTEXT_FOOT.bind(lineno=code.co_firstlineno)) - def unwrap(self) -> None: - f = t.cast(FunctionType, self.__wrapped__) + # Mark the function as wrapped by a wrapping context + t.cast(ContextWrappedFunction, f).__dd_context_wrapped__ = self - with _registry_lock: - if not self.is_wrapped(f): - return + # Replace the function code with the wrapped code. We also link + # the function to its original code object so that we can retrieve + # it later if required. + link_function_to_code(code, f) + set_function_code(f, bc.to_code()) - wc = _registry[f].uwc + def unwrap(self) -> None: + f = self.__wrapped__ - bc = Bytecode.from_code(get_function_code(f)) + if not self.is_wrapped(f): + return - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() + wrapped = t.cast(ContextWrappedFunction, f) - # Remove the head of the try block - for i, instr in enumerate(bc): - if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: - break + bc = Bytecode.from_code(get_function_code(f)) - bc[i : i + len(CONTEXT_HEAD) - 1] = [] + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() - # Remove all the return handlers - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.Instr) and instr.name == "RETURN_VALUE": - bc[i - len(CONTEXT_RETURN) : i] = [] - i -= len(CONTEXT_RETURN) - i += 1 + # Remove the head of the try block + wc = wrapped.__dd_context_wrapped__ + for i, instr in enumerate(bc): + if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: + break - # Recreate the code object - set_function_code(f, bc.to_code()) + bc[i : i + len(CONTEXT_HEAD) - 1] = [] - # Clear the UWC from the registry; remove the record if fully empty. - record = _registry.get(f) - if record is not None: - record.uwc = None - if not record.lazy_contexts: - _registry.pop(f, None) + # Remove all the return handlers + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.Instr) and instr.name == "RETURN_VALUE": + bc[i - len(CONTEXT_RETURN) : i] = [] + i -= len(CONTEXT_RETURN) + i += 1 + # Recreate the code object + set_function_code(f, bc.to_code()) -def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingContext]": - """Return the _UniversalWrappingContext for *f*, or None if not context-wrapped.""" - with _registry_lock: - record = _registry.get(f) - return record.uwc if record is not None else None + # Remove the wrapping context marker + del wrapped.__dd_context_wrapped__ diff --git a/ddtrace/internal/wrapping/generators.py b/ddtrace/internal/utils/wrapping/generators.py similarity index 100% rename from ddtrace/internal/wrapping/generators.py rename to ddtrace/internal/utils/wrapping/generators.py diff --git a/ddtrace/internal/wrapping.py b/ddtrace/internal/wrapping.py new file mode 100644 index 00000000000..07fd80a109e --- /dev/null +++ b/ddtrace/internal/wrapping.py @@ -0,0 +1 @@ +from ddtrace.internal.utils.wrapping import * # noqa From f2711800cff3d5d6ea80e6489dbb25b5d0ba082a Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 08:08:32 -0700 Subject: [PATCH 02/29] update from main --- ddtrace/internal/utils/schema/span_attribute_schema.py | 2 +- ddtrace/internal/utils/wrapping/__init__.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ddtrace/internal/utils/schema/span_attribute_schema.py b/ddtrace/internal/utils/schema/span_attribute_schema.py index b42c0211b9f..c9b717f6809 100644 --- a/ddtrace/internal/utils/schema/span_attribute_schema.py +++ b/ddtrace/internal/utils/schema/span_attribute_schema.py @@ -2,8 +2,8 @@ import sys from typing import Optional +from ddtrace.internal.constants import DEFAULT_SERVICE_NAME from ddtrace.internal.settings._inferred_base_service import detect_service -from ddtrace.internal.utils.constants import DEFAULT_SERVICE_NAME class SpanDirection(Enum): diff --git a/ddtrace/internal/utils/wrapping/__init__.py b/ddtrace/internal/utils/wrapping/__init__.py index 758cb916834..e3eb2424aa8 100644 --- a/ddtrace/internal/utils/wrapping/__init__.py +++ b/ddtrace/internal/utils/wrapping/__init__.py @@ -13,8 +13,8 @@ from ddtrace.internal.assembly import Assembly from ddtrace.internal.utils.inspection import link_function_to_code -from ddtrace.internal.utils.wrapping.asyncs import wrap_async -from ddtrace.internal.utils.wrapping.generators import wrap_generator +from ddtrace.internal.wrapping.asyncs import wrap_async +from ddtrace.internal.wrapping.generators import wrap_generator PY = sys.version_info[:2] From a6fa7d26102648d75a103e96c72899fea458cfbc Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 08:11:17 -0700 Subject: [PATCH 03/29] merge main --- ddtrace/internal/utils/wrapping/__init__.py | 147 ++++++++++++-------- 1 file changed, 91 insertions(+), 56 deletions(-) diff --git a/ddtrace/internal/utils/wrapping/__init__.py b/ddtrace/internal/utils/wrapping/__init__.py index e3eb2424aa8..56d0362c7ed 100644 --- a/ddtrace/internal/utils/wrapping/__init__.py +++ b/ddtrace/internal/utils/wrapping/__init__.py @@ -4,28 +4,44 @@ from typing import Any from typing import Callable from typing import Iterator +from typing import MutableMapping from typing import Optional from typing import Protocol from typing import cast +import weakref import bytecode as bc from bytecode import Instr from ddtrace.internal.assembly import Assembly -from ddtrace.internal.utils.inspection import link_function_to_code -from ddtrace.internal.wrapping.asyncs import wrap_async -from ddtrace.internal.wrapping.generators import wrap_generator +from ddtrace.internal.utils.threads import Lock +from ddtrace.internal.utils.wrapping.asyncs import wrap_async +from ddtrace.internal.utils.wrapping.generators import wrap_generator PY = sys.version_info[:2] +# Maps each wrapped function to its inner copy (the singly-linked list of +# wrapping layers). WeakKeyDictionary so functions are not kept alive by the +# registry alone. +_wrapped: weakref.WeakKeyDictionary[FunctionType, FunctionType] = weakref.WeakKeyDictionary() +_wrapped_lock = Lock() + +# Maps original code objects to the functions that own them. Written by +# link_function_to_code; read by functions_for_code in inspection.py. +# WeakValueDictionary so that wrapped ephemeral functions are not kept alive by +# this mapping alone — inspection.py falls back to gc.get_referrers on a miss. +_code_to_fn: MutableMapping[CodeType, FunctionType] = weakref.WeakValueDictionary() + + +def link_function_to_code(code: CodeType, function: FunctionType) -> None: + """Link a function to its original code object for fast reverse lookup.""" + _code_to_fn[code] = function + class WrappedFunction(Protocol): """A wrapped function.""" - __dd_wrapped__: Optional[FunctionType] = None - __dd_wrappers__: Optional[dict[Any, Any]] = None - def __call__(self, *args: Any, **kwargs: Any) -> Any: pass @@ -295,11 +311,12 @@ def wrap(f: FunctionType, wrapper: Wrapper) -> WrappedFunction: f.__defaults__, f.__closure__, ) - try: - wf = cast(WrappedFunction, f) - cast(WrappedFunction, wrapped).__dd_wrapped__ = cast(FunctionType, wf.__dd_wrapped__) - except AttributeError: - pass + + # Carry forward the existing wrapped-function link to the new inner copy. + with _wrapped_lock: + existing_inner = _wrapped.get(f) + if existing_inner is not None: + _wrapped[wrapped] = existing_inner wrapped.__kwdefaults__ = f.__kwdefaults__ @@ -329,47 +346,54 @@ def wrap(f: FunctionType, wrapper: Wrapper) -> WrappedFunction: # Replace the function code with the trampoline bytecode f.__code__ = wrapped_code.to_code() - # DEV: Multiple wrapping is implemented as a singly-linked list via the - # __dd_wrapped__ attribute. - wf = cast(WrappedFunction, f) - wf.__dd_wrapped__ = wrapped + # DEV: Multiple wrapping is implemented as a singly-linked list via _wrapped. + with _wrapped_lock: + _wrapped[f] = wrapped # Link the original code object to the original function link_function_to_code(code, f) - return wf + return cast(WrappedFunction, f) + + +def _unwrap_method(f: Any) -> Any: + """Return the underlying function if *f* is a bound or unbound method.""" + func = getattr(f, "__func__", f) + return func def is_wrapped(f: FunctionType) -> bool: """Check if a function is wrapped with any wrapper.""" try: - wf = cast(WrappedFunction, f) - inner = cast(FunctionType, wf.__dd_wrapped__) - - # Sanity check - assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - return True - except AttributeError: + with _wrapped_lock: + inner = _wrapped.get(_unwrap_method(f)) + except TypeError: + # f is not weakly referenceable (e.g. C method_descriptor) return False + if inner is None: + return False + assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + return True def is_wrapped_with(f: FunctionType, wrapper: Wrapper) -> bool: """Check if a function is wrapped with a specific wrapper.""" + f = cast(FunctionType, _unwrap_method(f)) try: - wf = cast(WrappedFunction, f) - inner = cast(FunctionType, wf.__dd_wrapped__) - - # Sanity check - assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + with _wrapped_lock: + inner = _wrapped.get(f) + except TypeError: + return False + if inner is None: + return False - if wrapper in f.__code__.co_consts: - return True + assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - # This is not the correct wrapping layer. Try with the next one. - return is_wrapped_with(inner, wrapper) + if wrapper in f.__code__.co_consts: + return True - except AttributeError: - return False + # This is not the correct wrapping layer. Try with the next one. + return is_wrapped_with(inner, wrapper) def unwrap(wf: WrappedFunction, wrapper: Wrapper) -> FunctionType: @@ -379,41 +403,52 @@ def unwrap(wf: WrappedFunction, wrapper: Wrapper) -> FunctionType: this will unwrap the one that uses ``wrapper``. If the function was not wrapped with ``wrapper``, it will return the first argument. """ - # DEV: Multiple wrapping layers are singly-linked via __dd_wrapped__. When - # we find the layer that needs to be removed we also have to ensure that we + # DEV: Multiple wrapping layers are singly-linked via _wrapped. When we + # find the layer that needs to be removed we also have to ensure that we # update the link at the deletion site if there is a non-empty tail. - try: - inner = cast(FunctionType, wf.__dd_wrapped__) - except AttributeError: + f = cast(FunctionType, wf) + with _wrapped_lock: + inner = _wrapped.get(f) + if inner is None: # The function is not wrapped so we return it as is. - return cast(FunctionType, wf) + return f - # Sanity check assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - if wrapper not in cast(FunctionType, wf).__code__.co_consts: + if wrapper not in f.__code__.co_consts: # This is not the correct wrapping layer. Try with the next one. return unwrap(cast(WrappedFunction, inner), wrapper) - # Remove the current wrapping layer by moving the next one over the - # current one. - f = cast(FunctionType, wf) - f.__code__ = inner.__code__ - - try: - # Update the link to the next layer. - wf.__dd_wrapped__ = cast(WrappedFunction, inner).__dd_wrapped__ # type: ignore[assignment] - except AttributeError: - # No more wrapping layers. Restore the original function by removing - # this extra attribute. - del wf.__dd_wrapped__ + # Remove the current wrapping layer by moving the next one over the current + # one. The code swap and registry update must be atomic: two concurrent + # unwrap calls on the same function and wrapper must not both succeed. + with _wrapped_lock: + f.__code__ = inner.__code__ + next_inner = _wrapped.get(inner) + if next_inner is not None: + _wrapped[f] = next_inner + else: + del _wrapped[f] return f def get_function_code(f: FunctionType) -> CodeType: - return (cast(WrappedFunction, f).__dd_wrapped__ or f if is_wrapped(f) else f).__code__ + with _wrapped_lock: + inner = _wrapped.get(f) + return (inner if inner is not None else f).__code__ def set_function_code(f: FunctionType, code: CodeType) -> None: - (cast(WrappedFunction, f).__dd_wrapped__ or f if is_wrapped(f) else f).__code__ = code # type: ignore[misc] + with _wrapped_lock: + inner = _wrapped.get(f) + (inner if inner is not None else f).__code__ = code + + +def get_wrapped(f: FunctionType) -> Optional[FunctionType]: + """Return the inner bytecode copy of *f* if it is wrapped, else None.""" + try: + with _wrapped_lock: + return _wrapped.get(cast(FunctionType, _unwrap_method(f))) + except TypeError: + return None From 3e4389bc3e35eacaf8e9d9ac95a3e8498e140225 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 08:13:53 -0700 Subject: [PATCH 04/29] merge main into wrapping/context --- ddtrace/internal/utils/wrapping/context.py | 574 +++++++++++---------- 1 file changed, 313 insertions(+), 261 deletions(-) diff --git a/ddtrace/internal/utils/wrapping/context.py b/ddtrace/internal/utils/wrapping/context.py index d3efd0ef196..abaad9fa4ee 100644 --- a/ddtrace/internal/utils/wrapping/context.py +++ b/ddtrace/internal/utils/wrapping/context.py @@ -8,23 +8,73 @@ from types import TracebackType import typing as t from typing import Protocol # noqa:F401 +import weakref import bytecode from bytecode import Bytecode from ddtrace.internal.assembly import Assembly -from ddtrace.internal.utils.inspection import link_function_to_code from ddtrace.internal.utils.logger import get_logger from ddtrace.internal.utils.threads import Lock +from ddtrace.internal.utils.threads import RLock from ddtrace.internal.utils.wrapping import WrappedFunction from ddtrace.internal.utils.wrapping import Wrapper from ddtrace.internal.utils.wrapping import get_function_code from ddtrace.internal.utils.wrapping import is_wrapped_with +from ddtrace.internal.utils.wrapping import link_function_to_code from ddtrace.internal.utils.wrapping import set_function_code from ddtrace.internal.utils.wrapping import unwrap from ddtrace.internal.utils.wrapping import wrap +class _ContextRecord: + """Holds all wrapping-context metadata for a single function. + + Stores only weak references to context objects so that the WeakKeyDictionary + key (the function) is not kept alive by the registry value chain: + _registry -> _ContextRecord -> context -> context.__wrapped__ -> function + Breaking this path with weak references lets ephemeral functions be garbage + collected as soon as all external strong references drop. + """ + + __slots__ = ("_uwc_ref", "lazy_contexts") + + def __init__(self) -> None: + self._uwc_ref: t.Optional[weakref.ref["_UniversalWrappingContext"]] = None + # WeakSet so that LazyWrappingContext instances (which also hold + # __wrapped__ = f) do not prevent the function from being collected. + self.lazy_contexts: weakref.WeakSet["LazyWrappingContext"] = weakref.WeakSet() + + @property + def uwc(self) -> t.Optional["_UniversalWrappingContext"]: + ref = self._uwc_ref + return ref() if ref is not None else None + + @uwc.setter + def uwc(self, value: t.Optional["_UniversalWrappingContext"]) -> None: + self._uwc_ref = weakref.ref(value) if value is not None else None + + @classmethod + def get_or_create(cls, f: FunctionType) -> "_ContextRecord": + record = _registry.get(f) + if record is None: + with _registry_lock: + record = _registry.get(f) + if record is None: + record = cls() + _registry[f] = record + return record + + +# Per-function registry for wrapping-context machinery. WeakKeyDictionary so +# functions are not kept alive by the registry alone. Storing data here instead +# of as function attributes keeps __dict__ clean, preventing frameworks that +# copy function __dict__ (e.g. functools.wraps, +# self.__dict__.update(f.__dict__)) from capturing non-picklable objects. +_registry: weakref.WeakKeyDictionary[FunctionType, _ContextRecord] = weakref.WeakKeyDictionary() +_registry_lock = RLock() + + log = get_logger(__name__) T = t.TypeVar("T") @@ -293,22 +343,27 @@ class BaseWrappingContext(ABC): __priority__: int = 0 def __init__(self, f: FunctionType): - self.__wrapped__ = f + # Store a weak reference so that context objects do not keep the wrapped + # function alive. CodeType is not GC-tracked in CPython, so the cycle + # f → code.co_consts → bound_methods(uwc) → uwc.__wrapped__ → f + # cannot be broken by the cyclic GC. A weak ref here allows f's + # reference count to reach zero (and be freed) as soon as all external + # strong refs drop, without relying on the cyclic GC at all. + self._wrapped_ref: weakref.ref[FunctionType] = weakref.ref(f) self._storage: ContextVar[t.Optional[dict[str, t.Any]]] = ContextVar( f"{type(self).__name__}__storage", default=None ) - def __getstate__(self) -> dict[str, t.Any]: - state = self.__dict__.copy() - state.pop("_storage", None) # remove unpicklable field - return state + @property + def __wrapped__(self) -> FunctionType: + f = self._wrapped_ref() + if f is None: + raise RuntimeError(f"{type(self).__name__}.__wrapped__: the wrapped function has been garbage collected") + return f - def __setstate__(self, state: dict[str, t.Any]) -> None: - self.__dict__.update(state) - self._storage = ContextVar( - f"{type(self).__name__}__storage", - default=None, - ) + @__wrapped__.setter + def __wrapped__(self, f: FunctionType) -> None: + self._wrapped_ref = weakref.ref(f) def __enter__(self) -> "BaseWrappingContext": prev = self._storage.get() @@ -342,10 +397,10 @@ def set(self, key: str, value: T) -> T: @classmethod def wrapped(cls, f: FunctionType) -> "BaseWrappingContext": - if cls.is_wrapped(f): + try: context = cls.extract(f) assert isinstance(context, cls) # nosec - else: + except ValueError: context = cls(f) context.wrap() return context @@ -370,7 +425,10 @@ class WrappingContext(BaseWrappingContext): @property def __frame__(self) -> FrameType: try: - return t.cast(FrameType, _UniversalWrappingContext.extract(self.__wrapped__).get("__frame__")) + return t.cast( + FrameType, + _UniversalWrappingContext.extract(t.cast(FunctionType, self.__wrapped__)).get("__frame__"), + ) except ValueError: raise AttributeError("Wrapping context not entered") @@ -386,31 +444,24 @@ def is_wrapped(cls, f: FunctionType) -> bool: @classmethod def extract(cls, f: FunctionType) -> "WrappingContext": - if _UniversalWrappingContext.is_wrapped(f): - try: - return _UniversalWrappingContext.extract(f).registered(cls) - except KeyError: - pass - msg = f"Function is not wrapped with {cls}" - raise ValueError(msg) + try: + return _UniversalWrappingContext.extract(f).registered(cls) + except (ValueError, KeyError): + msg = f"Function is not wrapped with {cls}" + raise ValueError(msg) def wrap(self) -> None: - t.cast(_UniversalWrappingContext, _UniversalWrappingContext.wrapped(self.__wrapped__)).register(self) + t.cast( + _UniversalWrappingContext, _UniversalWrappingContext.wrapped(t.cast(FunctionType, self.__wrapped__)) + ).register(self) def unwrap(self) -> None: - f = self.__wrapped__ + f = t.cast(FunctionType, self.__wrapped__) - if _UniversalWrappingContext.is_wrapped(f): + try: _UniversalWrappingContext.extract(f).unregister(self) - - -class LazyWrappedFunction(Protocol): - """A lazy-wrapped function.""" - - __dd_lazy_contexts__: list[WrappingContext] - - def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: - pass + except ValueError: + pass class LazyWrappingContext(WrappingContext): @@ -422,10 +473,11 @@ def __init__(self, f: FunctionType): @classmethod def is_wrapped(cls, f: FunctionType) -> bool: - try: - return any(isinstance(c, cls) for c in t.cast(LazyWrappedFunction, f).__dd_lazy_contexts__) - except AttributeError: - return False + with _registry_lock: + record = _registry.get(f) + if record is None: + return False + return any(isinstance(c, cls) for c in record.lazy_contexts) def wrap(self) -> None: """Perform the bytecode wrapping on first invocation.""" @@ -433,71 +485,60 @@ def wrap(self) -> None: if self._trampoline is not None: return - # If the function is already universally wrapped so it's less expensive + # If the function is already universally wrapped it's less expensive # to do the normal wrapping. - if _UniversalWrappingContext.is_wrapped(self.__wrapped__): + if _UniversalWrappingContext.is_wrapped(t.cast(FunctionType, self.__wrapped__)): super().wrap() return def trampoline(_: t.Any, args: tuple[t.Any, ...], kwargs: dict[str, t.Any]) -> t.Any: with tl: f = t.cast(WrappedFunction, self.__wrapped__) - if is_wrapped_with(self.__wrapped__, trampoline): + if is_wrapped_with(t.cast(FunctionType, self.__wrapped__), trampoline): f = t.cast(WrappedFunction, unwrap(f, trampoline)) self._trampoline = None - try: - (cs := t.cast(LazyWrappedFunction, f).__dd_lazy_contexts__).remove(self) - if not cs: - del t.cast(LazyWrappedFunction, f).__dd_lazy_contexts__ - except (AttributeError, ValueError): + inconsistent = False + with _registry_lock: + record = _registry.get(t.cast(FunctionType, f)) + if record is not None: + inconsistent = self not in record.lazy_contexts + record.lazy_contexts.discard(self) + if not record.lazy_contexts and record.uwc is None: + _registry.pop(t.cast(FunctionType, f), None) + if inconsistent: log.warning("Inconsistent lazy wrapping context state") super(LazyWrappingContext, self).wrap() return f(*args, **kwargs) - wrap(self.__wrapped__, trampoline) + wrap(t.cast(FunctionType, self.__wrapped__), trampoline) self._trampoline = trampoline - wf = t.cast(LazyWrappedFunction, self.__wrapped__) - if not hasattr(wf, "__dd_lazy_contexts__"): - wf.__dd_lazy_contexts__ = [] - wf.__dd_lazy_contexts__.append(self) + _ContextRecord.get_or_create(t.cast(FunctionType, self.__wrapped__)).lazy_contexts.add(self) def unwrap(self) -> None: with self._trampoline_lock: - if _UniversalWrappingContext.is_wrapped(self.__wrapped__): + if _UniversalWrappingContext.is_wrapped(t.cast(FunctionType, self.__wrapped__)): assert self._trampoline is None # nosec super().unwrap() elif self._trampoline is not None: - wf = t.cast(LazyWrappedFunction, self.__wrapped__) - if hasattr(wf, "__dd_lazy_contexts__"): - wf.__dd_lazy_contexts__.remove(self) - if not wf.__dd_lazy_contexts__: - del wf.__dd_lazy_contexts__ + with _registry_lock: + record = _registry.get(t.cast(FunctionType, self.__wrapped__)) + if record is not None: + record.lazy_contexts.discard(self) + if not record.lazy_contexts and record.uwc is None: + _registry.pop(t.cast(FunctionType, self.__wrapped__), None) unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline) self._trampoline = None - def __getstate__(self) -> dict[str, t.Any]: - state = super().__getstate__() - state.pop("_trampoline_lock", None) # thread lock not picklable - state.pop("_trampoline", None) # closure not picklable - return state - - def __setstate__(self, state: dict[str, t.Any]) -> None: - super().__setstate__(state) - self._trampoline_lock = Lock() - self._trampoline = None - class ContextWrappedFunction(Protocol): """A wrapped function.""" - __dd_context_wrapped__ = None # type: t.Optional[_UniversalWrappingContext] - def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: pass @@ -573,252 +614,263 @@ def __return__(self, value: T) -> T: @classmethod def is_wrapped(cls, f: FunctionType) -> bool: try: - # Check that we have actual bytecode wrapping. The presence of the - # __dd_context_wrapped__ attribute is not enough, as this could be - # copied over from an object state cloning. - if sys.version_info >= (3, 11): - return f.__dd_context_wrapped__.__enter__ in get_function_code(f).co_consts # type: ignore - else: - return f.__dd_context_wrapped__ in get_function_code(f).co_consts # type: ignore + with _registry_lock: + record = _registry.get(f) + if record is None or record.uwc is None: + return False + # Verify the registry entry matches actual bytecode wrapping. + if sys.version_info >= (3, 11): + return record.uwc.__enter__ in get_function_code(f).co_consts + else: + return record.uwc in get_function_code(f).co_consts except AttributeError: return False @classmethod def extract(cls, f: FunctionType) -> "_UniversalWrappingContext": - if not cls.is_wrapped(f): - raise ValueError("Function is not wrapped") - return t.cast(_UniversalWrappingContext, t.cast(ContextWrappedFunction, f).__dd_context_wrapped__) + with _registry_lock: + if not cls.is_wrapped(f): + raise ValueError("Function is not wrapped") + return t.cast(_UniversalWrappingContext, _registry[f].uwc) if sys.version_info >= (3, 11): def wrap(self) -> None: f = self.__wrapped__ - if self.is_wrapped(f): - raise ValueError("Function already wrapped") - - bc = Bytecode.from_code(code := get_function_code(f)) - - # Prefix every return - i = 0 - while i < len(bc): - instr = bc[i] - try: - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN.bind({"context_return": self.__return__}, lineno=instr.lineno) - elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ - return_code = CONTEXT_RETURN_CONST.bind( - {"context_return": self.__return__, "value": instr.arg}, lineno=instr.lineno - ) - else: - return_code = [] - - bc[i:i] = return_code - i += len(return_code) - except AttributeError: - # Not an instruction - pass - i += 1 - - # Search for the RESUME instruction - for i, instr in enumerate(bc, 1): - try: - if instr.name == "RESUME": - break - except AttributeError: - # Not an instruction - pass - else: - i = 0 + with _registry_lock: + if self.is_wrapped(f): + raise ValueError("Function already wrapped") - bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=code.co_firstlineno) + bc = Bytecode.from_code(code := get_function_code(f)) - # Wrap every line outside a try block - except_label = bytecode.Label() - first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) + # Prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + try: + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN.bind({"context_return": self.__return__}, lineno=instr.lineno) + elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ + return_code = CONTEXT_RETURN_CONST.bind( + {"context_return": self.__return__, "value": instr.arg}, lineno=instr.lineno + ) + else: + return_code = [] - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None: - bc.insert(i, bytecode.TryEnd(last_try_begin)) - last_try_begin = None + bc[i:i] = return_code + i += len(return_code) + except AttributeError: + # Not an instruction + pass i += 1 - elif isinstance(instr, bytecode.TryEnd): - j = i + 1 - while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin): - if isinstance(bc[j], bytecode.Instr): - last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) - bc.insert(i + 1, last_try_begin) + + # Search for the RESUME instruction + for i, instr in enumerate(bc, 1): + try: + if instr.name == "RESUME": break - j += 1 - i += 1 - i += 1 + except AttributeError: + # Not an instruction + pass + else: + i = 0 - bc.insert(0, first_try_begin) + bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=code.co_firstlineno) - bc.append(bytecode.TryEnd(last_try_begin)) - bc.append(except_label) - bc.extend(CONTEXT_FOOT.bind({"context_exit": self._exit}, lineno=code.co_firstlineno)) + # Wrap every line outside a try block + except_label = bytecode.Label() + first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) + + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None: + bc.insert(i, bytecode.TryEnd(last_try_begin)) + last_try_begin = None + i += 1 + elif isinstance(instr, bytecode.TryEnd): + j = i + 1 + while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin): + if isinstance(bc[j], bytecode.Instr): + last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) + bc.insert(i + 1, last_try_begin) + break + j += 1 + i += 1 + i += 1 - # Mark the function as wrapped by a wrapping context - t.cast(ContextWrappedFunction, f).__dd_context_wrapped__ = self + bc.insert(0, first_try_begin) - # Replace the function code with the wrapped code. We also link - # the function to its original code object so that we can retrieve - # it later if required. - link_function_to_code(code, f) + bc.append(bytecode.TryEnd(last_try_begin)) + bc.append(except_label) + bc.extend(CONTEXT_FOOT.bind({"context_exit": self._exit}, lineno=code.co_firstlineno)) - set_function_code(f, bc.to_code()) + # Register the wrapping context and write the new bytecode. + _ContextRecord.get_or_create(f).uwc = self + link_function_to_code(code, f) + set_function_code(f, bc.to_code()) def unwrap(self) -> None: f = self.__wrapped__ - if not self.is_wrapped(f): - return - - wrapped = t.cast(ContextWrappedFunction, f) + with _registry_lock: + if not self.is_wrapped(f): + return - bc = Bytecode.from_code(get_function_code(f)) + wc = _registry[f].uwc - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() - bc.pop() + bc = Bytecode.from_code(get_function_code(f)) - except_label = bc.pop(0).target + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() + bc.pop() - # Remove the try blocks - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.TryBegin) and instr.target is except_label: - bc.pop(i) - elif isinstance(instr, bytecode.TryEnd) and instr.entry.target is except_label: - bc.pop(i) - else: - i += 1 + except_label = bc.pop(0).target - # Remove the head of the try block - wc = wrapped.__dd_context_wrapped__ - for i, instr in enumerate(bc): - if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: - break + # Remove the try blocks + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.TryBegin) and instr.target is except_label: + bc.pop(i) + elif isinstance(instr, bytecode.TryEnd) and instr.entry.target is except_label: + bc.pop(i) + else: + i += 1 - # Search for the RESUME instruction - for i, instr in enumerate(bc, 1): - try: - if instr.name == "RESUME": + # Remove the head of the try block + for i, instr in enumerate(bc): + if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: break - except AttributeError: - # Not an instruction - pass - else: - i = 0 - bc[i : i + len(CONTEXT_HEAD)] = [] - - # Un-prefix every return - i = 0 - while i < len(bc): - instr = bc[i] - try: - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN - elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ - return_code = CONTEXT_RETURN_CONST - else: - return_code = None + # Search for the RESUME instruction + for i, instr in enumerate(bc, 1): + try: + if instr.name == "RESUME": + break + except AttributeError: + # Not an instruction + pass + else: + i = 0 - if return_code is not None: - bc[i - len(return_code) : i] = [] - i -= len(return_code) - except AttributeError: - # Not an instruction - pass - i += 1 + bc[i : i + len(CONTEXT_HEAD)] = [] - # Recreate the code object - set_function_code(f, bc.to_code()) + # Un-prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + try: + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN + elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ + return_code = CONTEXT_RETURN_CONST + else: + return_code = None + + if return_code is not None: + bc[i - len(return_code) : i] = [] + i -= len(return_code) + except AttributeError: + # Not an instruction + pass + i += 1 - # Remove the wrapping context marker - del wrapped.__dd_context_wrapped__ + # Recreate the code object + set_function_code(f, bc.to_code()) + + # Clear the UWC from the registry; remove the record if fully empty. + record = _registry.get(f) + if record is not None: + record.uwc = None + if not record.lazy_contexts: + _registry.pop(f, None) else: def wrap(self) -> None: - f = self.__wrapped__ - - if self.is_wrapped(f): - raise ValueError("Function already wrapped") + f = t.cast(FunctionType, self.__wrapped__) - bc = Bytecode.from_code(code := get_function_code(f)) + with _registry_lock: + if self.is_wrapped(f): + raise ValueError("Function already wrapped") - # Prefix every return - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.Instr): - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN.bind({"context": self}, lineno=instr.lineno) - bc[i:i] = return_code - i += len(return_code) - i += 1 + bc = Bytecode.from_code(code := get_function_code(f)) - # Search for the GEN_START instruction, which needs to stay on top. - i = 0 - if sys.version_info >= (3, 10) and (iscoroutinefunction(f) or isgeneratorfunction(f)): - for i, instr in enumerate(bc, 1): - if isinstance(instr, bytecode.Instr) and instr.name == "GEN_START": - break + # Prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.Instr): + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN.bind({"context": self}, lineno=instr.lineno) + bc[i:i] = return_code + i += len(return_code) + i += 1 - *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": self}, lineno=code.co_firstlineno) + # Search for the GEN_START instruction, which needs to stay on top. + i = 0 + if sys.version_info >= (3, 10) and (iscoroutinefunction(f) or isgeneratorfunction(f)): + for i, instr in enumerate(bc, 1): + if isinstance(instr, bytecode.Instr) and instr.name == "GEN_START": + break - bc.append(except_label) - bc.extend(CONTEXT_FOOT.bind(lineno=code.co_firstlineno)) + *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": self}, lineno=code.co_firstlineno) - # Mark the function as wrapped by a wrapping context - t.cast(ContextWrappedFunction, f).__dd_context_wrapped__ = self + bc.append(except_label) + bc.extend(CONTEXT_FOOT.bind(lineno=code.co_firstlineno)) - # Replace the function code with the wrapped code. We also link - # the function to its original code object so that we can retrieve - # it later if required. - link_function_to_code(code, f) - set_function_code(f, bc.to_code()) + # Register the wrapping context and write the new bytecode. + _ContextRecord.get_or_create(f).uwc = self + link_function_to_code(code, f) + set_function_code(f, bc.to_code()) def unwrap(self) -> None: - f = self.__wrapped__ + f = t.cast(FunctionType, self.__wrapped__) - if not self.is_wrapped(f): - return + with _registry_lock: + if not self.is_wrapped(f): + return - wrapped = t.cast(ContextWrappedFunction, f) + wc = _registry[f].uwc - bc = Bytecode.from_code(get_function_code(f)) + bc = Bytecode.from_code(get_function_code(f)) - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() - # Remove the head of the try block - wc = wrapped.__dd_context_wrapped__ - for i, instr in enumerate(bc): - if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: - break + # Remove the head of the try block + for i, instr in enumerate(bc): + if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: + break + + bc[i : i + len(CONTEXT_HEAD) - 1] = [] + + # Remove all the return handlers + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.Instr) and instr.name == "RETURN_VALUE": + bc[i - len(CONTEXT_RETURN) : i] = [] + i -= len(CONTEXT_RETURN) + i += 1 - bc[i : i + len(CONTEXT_HEAD) - 1] = [] + # Recreate the code object + set_function_code(f, bc.to_code()) - # Remove all the return handlers - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.Instr) and instr.name == "RETURN_VALUE": - bc[i - len(CONTEXT_RETURN) : i] = [] - i -= len(CONTEXT_RETURN) - i += 1 + # Clear the UWC from the registry; remove the record if fully empty. + record = _registry.get(f) + if record is not None: + record.uwc = None + if not record.lazy_contexts: + _registry.pop(f, None) - # Recreate the code object - set_function_code(f, bc.to_code()) - # Remove the wrapping context marker - del wrapped.__dd_context_wrapped__ +def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingContext]": + """Return the _UniversalWrappingContext for *f*, or None if not context-wrapped.""" + with _registry_lock: + record = _registry.get(f) + return record.uwc if record is not None else None From 1d4d4e7d45294614d35047f403f3ecb4462befa4 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 08:31:32 -0700 Subject: [PATCH 05/29] fix relative import --- ddtrace/internal/utils/formats.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ddtrace/internal/utils/formats.py b/ddtrace/internal/utils/formats.py index a40ec050234..a100a1cad91 100644 --- a/ddtrace/internal/utils/formats.py +++ b/ddtrace/internal/utils/formats.py @@ -8,8 +8,7 @@ from ddtrace.internal.constants import MAX_UINT_64BITS # noqa:F401 from ddtrace.internal.native._native import flatten_key_value # noqa: F401 from ddtrace.internal.native._native import is_sequence # noqa: F401 - -from ..compat import ensure_text +from ddtrace.internal.utils.compat import ensure_text VALUE_PLACEHOLDER = "?" From d737e82d70cd57692308157c59a060ad3ae7ea8f Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 08:40:02 -0700 Subject: [PATCH 06/29] fix relative import --- ddtrace/internal/encoding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ddtrace/internal/encoding.py b/ddtrace/internal/encoding.py index 0697d1c582d..6072c3f5410 100644 --- a/ddtrace/internal/encoding.py +++ b/ddtrace/internal/encoding.py @@ -6,6 +6,8 @@ from ddtrace.internal.settings._agent import config as agent_config # noqa:F401 from ddtrace.internal.threads import RLock +from ddtrace.internal.utils.compat import ensure_text +from ddtrace.internal.utils.logger import get_logger from ._encoding import BufferedEncoder from ._encoding import BufferFull @@ -13,8 +15,6 @@ from ._encoding import ListStringTable from ._encoding import MsgpackEncoderV04 from ._encoding import MsgpackEncoderV05 -from .compat import ensure_text -from .logger import get_logger __all__ = [ From 6c437131e4b686493cc815c9686fefb3e7da39cf Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 09:15:46 -0700 Subject: [PATCH 07/29] fancier reexporting --- ddtrace/internal/compat.py | 8 +++++++- ddtrace/internal/constants.py | 8 +++++++- ddtrace/internal/schema.py | 8 +++++++- 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/ddtrace/internal/compat.py b/ddtrace/internal/compat.py index 8463b28f23e..8c716f6b10f 100644 --- a/ddtrace/internal/compat.py +++ b/ddtrace/internal/compat.py @@ -1 +1,7 @@ -from ddtrace.internal.utils.compat import * # noqa +import importlib + + +reexport_path = "ddtrace.internal.utils.compat" +reexported_module = importlib.import_module(reexport_path) +for name in dir(reexported_module): + locals()[name] = getattr(reexported_module, name) diff --git a/ddtrace/internal/constants.py b/ddtrace/internal/constants.py index 14a047f6b24..fb4fd160fbc 100644 --- a/ddtrace/internal/constants.py +++ b/ddtrace/internal/constants.py @@ -1 +1,7 @@ -from ddtrace.internal.utils.constants import * # noqa +import importlib + + +reexport_path = "ddtrace.internal.utils.constants" +reexported_module = importlib.import_module(reexport_path) +for name in dir(reexported_module): + locals()[name] = getattr(reexported_module, name) diff --git a/ddtrace/internal/schema.py b/ddtrace/internal/schema.py index 9ba7325082f..87d2a968f4c 100644 --- a/ddtrace/internal/schema.py +++ b/ddtrace/internal/schema.py @@ -1 +1,7 @@ -from ddtrace.internal.utils.schema import * # noqa +import importlib + + +reexport_path = "ddtrace.internal.utils.schema" +reexported_module = importlib.import_module(reexport_path) +for name in dir(reexported_module): + locals()[name] = getattr(reexported_module, name) From 5b649df136a43aa50f9a98d793a4918ca244b324 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 10:00:22 -0700 Subject: [PATCH 08/29] update mypy use __all__ for reexports --- ddtrace/internal/compat.py | 8 +------- ddtrace/internal/constants.py | 8 +------- ddtrace/internal/schema.py | 7 ------- ddtrace/internal/schema/__init__.py | 1 + ddtrace/internal/schema/processor.py | 1 + ddtrace/internal/schema/span_attribute_schema.py | 1 + ddtrace/internal/utils/compat.py | 3 +++ ddtrace/internal/utils/constants.py | 3 +++ ddtrace/internal/utils/schema/__init__.py | 7 +++---- ddtrace/internal/utils/schema/processor.py | 3 +++ mypy.ini | 8 ++++---- 11 files changed, 21 insertions(+), 29 deletions(-) delete mode 100644 ddtrace/internal/schema.py create mode 100644 ddtrace/internal/schema/__init__.py create mode 100644 ddtrace/internal/schema/processor.py create mode 100644 ddtrace/internal/schema/span_attribute_schema.py diff --git a/ddtrace/internal/compat.py b/ddtrace/internal/compat.py index 8c716f6b10f..8463b28f23e 100644 --- a/ddtrace/internal/compat.py +++ b/ddtrace/internal/compat.py @@ -1,7 +1 @@ -import importlib - - -reexport_path = "ddtrace.internal.utils.compat" -reexported_module = importlib.import_module(reexport_path) -for name in dir(reexported_module): - locals()[name] = getattr(reexported_module, name) +from ddtrace.internal.utils.compat import * # noqa diff --git a/ddtrace/internal/constants.py b/ddtrace/internal/constants.py index fb4fd160fbc..14a047f6b24 100644 --- a/ddtrace/internal/constants.py +++ b/ddtrace/internal/constants.py @@ -1,7 +1 @@ -import importlib - - -reexport_path = "ddtrace.internal.utils.constants" -reexported_module = importlib.import_module(reexport_path) -for name in dir(reexported_module): - locals()[name] = getattr(reexported_module, name) +from ddtrace.internal.utils.constants import * # noqa diff --git a/ddtrace/internal/schema.py b/ddtrace/internal/schema.py deleted file mode 100644 index 87d2a968f4c..00000000000 --- a/ddtrace/internal/schema.py +++ /dev/null @@ -1,7 +0,0 @@ -import importlib - - -reexport_path = "ddtrace.internal.utils.schema" -reexported_module = importlib.import_module(reexport_path) -for name in dir(reexported_module): - locals()[name] = getattr(reexported_module, name) diff --git a/ddtrace/internal/schema/__init__.py b/ddtrace/internal/schema/__init__.py new file mode 100644 index 00000000000..9ba7325082f --- /dev/null +++ b/ddtrace/internal/schema/__init__.py @@ -0,0 +1 @@ +from ddtrace.internal.utils.schema import * # noqa diff --git a/ddtrace/internal/schema/processor.py b/ddtrace/internal/schema/processor.py new file mode 100644 index 00000000000..dacde93c163 --- /dev/null +++ b/ddtrace/internal/schema/processor.py @@ -0,0 +1 @@ +from ddtrace.internal.utils.schema.processor import * # noqa diff --git a/ddtrace/internal/schema/span_attribute_schema.py b/ddtrace/internal/schema/span_attribute_schema.py new file mode 100644 index 00000000000..85a4cc8b8c9 --- /dev/null +++ b/ddtrace/internal/schema/span_attribute_schema.py @@ -0,0 +1 @@ +from ddtrace.internal.utils.schema.span_attribute_schema import * # noqa diff --git a/ddtrace/internal/utils/compat.py b/ddtrace/internal/utils/compat.py index a0a122bd980..96d9f7b9214 100644 --- a/ddtrace/internal/utils/compat.py +++ b/ddtrace/internal/utils/compat.py @@ -123,3 +123,6 @@ def __getattr__(name: str) -> Any: def is_wrapted(o: object) -> bool: return isinstance(o, wrapt_class) + + +__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/constants.py b/ddtrace/internal/utils/constants.py index 1167f75f9e6..36472603864 100644 --- a/ddtrace/internal/utils/constants.py +++ b/ddtrace/internal/utils/constants.py @@ -157,3 +157,6 @@ class SamplingMechanism(object): class EXPERIMENTAL_FEATURES: # Enables submitting runtime metrics as gauges (instead of distributions) RUNTIME_METRICS = "DD_RUNTIME_METRICS_ENABLED" + + +__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/schema/__init__.py b/ddtrace/internal/utils/schema/__init__.py index db7226a42ae..535df404d08 100644 --- a/ddtrace/internal/utils/schema/__init__.py +++ b/ddtrace/internal/utils/schema/__init__.py @@ -2,10 +2,9 @@ from ddtrace.internal.settings import env from ddtrace.internal.utils.formats import asbool - -from .span_attribute_schema import _DEFAULT_SPAN_SERVICE_NAMES -from .span_attribute_schema import _SPAN_ATTRIBUTE_TO_FUNCTION -from .span_attribute_schema import SpanDirection +from ddtrace.internal.utils.schema.span_attribute_schema import _DEFAULT_SPAN_SERVICE_NAMES +from ddtrace.internal.utils.schema.span_attribute_schema import _SPAN_ATTRIBUTE_TO_FUNCTION +from ddtrace.internal.utils.schema.span_attribute_schema import SpanDirection log = logging.getLogger(__name__) diff --git a/ddtrace/internal/utils/schema/processor.py b/ddtrace/internal/utils/schema/processor.py index 8d5559ecae3..e5145de528f 100644 --- a/ddtrace/internal/utils/schema/processor.py +++ b/ddtrace/internal/utils/schema/processor.py @@ -23,3 +23,6 @@ def process_trace(self, trace): def _update_dd_base_service(self, span): span._set_attribute(key=_BASE_SERVICE_KEY, value=self._global_service) + + +__all__ = ["BaseServiceProcessor"] diff --git a/mypy.ini b/mypy.ini index 137791651d0..257747fed63 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1233,7 +1233,7 @@ check_untyped_defs = false disallow_any_generics = false warn_unreachable = false -[mypy-ddtrace.internal.compat] +[mypy-ddtrace.internal.utils.compat] disallow_untyped_defs = false disallow_incomplete_defs = false no_implicit_reexport = false @@ -1419,14 +1419,14 @@ disallow_untyped_defs = false disallow_incomplete_defs = false check_untyped_defs = false -[mypy-ddtrace.internal.logger] +[mypy-ddtrace.internal.utils.logger] disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false check_untyped_defs = false warn_return_any = false -[mypy-ddtrace.internal.module] +[mypy-ddtrace.internal.utils.module] disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false @@ -1778,7 +1778,7 @@ disallow_untyped_defs = false disallow_incomplete_defs = false check_untyped_defs = false -[mypy-ddtrace.internal.threads] +[mypy-ddtrace.internal.utils.threads] disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false From 15554dc7591c774423052f3a51969bc0996fbc54 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 10:03:49 -0700 Subject: [PATCH 09/29] reexport wrapping as a directory --- ddtrace/internal/utils/wrapping/__init__.py | 3 +++ ddtrace/internal/utils/wrapping/asyncs.py | 3 +++ ddtrace/internal/utils/wrapping/context.py | 3 +++ ddtrace/internal/utils/wrapping/generators.py | 3 +++ ddtrace/internal/{wrapping.py => wrapping/__init__.py} | 0 ddtrace/internal/wrapping/asyncs.py | 1 + ddtrace/internal/wrapping/context.py | 1 + ddtrace/internal/wrapping/generators.py | 1 + 8 files changed, 15 insertions(+) rename ddtrace/internal/{wrapping.py => wrapping/__init__.py} (100%) create mode 100644 ddtrace/internal/wrapping/asyncs.py create mode 100644 ddtrace/internal/wrapping/context.py create mode 100644 ddtrace/internal/wrapping/generators.py diff --git a/ddtrace/internal/utils/wrapping/__init__.py b/ddtrace/internal/utils/wrapping/__init__.py index 56d0362c7ed..530c824ee72 100644 --- a/ddtrace/internal/utils/wrapping/__init__.py +++ b/ddtrace/internal/utils/wrapping/__init__.py @@ -452,3 +452,6 @@ def get_wrapped(f: FunctionType) -> Optional[FunctionType]: return _wrapped.get(cast(FunctionType, _unwrap_method(f))) except TypeError: return None + + +__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/wrapping/asyncs.py b/ddtrace/internal/utils/wrapping/asyncs.py index a1de28a074e..d3bcd8b89cc 100644 --- a/ddtrace/internal/utils/wrapping/asyncs.py +++ b/ddtrace/internal/utils/wrapping/asyncs.py @@ -723,3 +723,6 @@ def wrap_async(instrs: bc.Bytecode, code: CodeType, lineno: int) -> None: elif bc.CompilerFlags.ASYNC_GENERATOR & code.co_flags: instrs[-1:] = ASYNC_GEN_ASSEMBLY.bind(lineno=lineno) + + +__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/wrapping/context.py b/ddtrace/internal/utils/wrapping/context.py index abaad9fa4ee..60d41002008 100644 --- a/ddtrace/internal/utils/wrapping/context.py +++ b/ddtrace/internal/utils/wrapping/context.py @@ -874,3 +874,6 @@ def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingConte with _registry_lock: record = _registry.get(f) return record.uwc if record is not None else None + + +__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/wrapping/generators.py b/ddtrace/internal/utils/wrapping/generators.py index 01c284c4686..243e3049065 100644 --- a/ddtrace/internal/utils/wrapping/generators.py +++ b/ddtrace/internal/utils/wrapping/generators.py @@ -489,3 +489,6 @@ def wrap_generator(instrs: bc.Bytecode, code: CodeType, lineno: int) -> None: instrs[0:0] = GENERATOR_HEAD_ASSEMBLY.bind(lineno=lineno) instrs[-1:] = GENERATOR_ASSEMBLY.bind(lineno=lineno) + + +__all__ = list(locals().keys()) diff --git a/ddtrace/internal/wrapping.py b/ddtrace/internal/wrapping/__init__.py similarity index 100% rename from ddtrace/internal/wrapping.py rename to ddtrace/internal/wrapping/__init__.py diff --git a/ddtrace/internal/wrapping/asyncs.py b/ddtrace/internal/wrapping/asyncs.py new file mode 100644 index 00000000000..339af72067c --- /dev/null +++ b/ddtrace/internal/wrapping/asyncs.py @@ -0,0 +1 @@ +from ddtrace.internal.utils.wrapping.asyncs import * # noqa diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py new file mode 100644 index 00000000000..1e7973a2e25 --- /dev/null +++ b/ddtrace/internal/wrapping/context.py @@ -0,0 +1 @@ +from ddtrace.internal.utils.wrapping.context import * # noqa diff --git a/ddtrace/internal/wrapping/generators.py b/ddtrace/internal/wrapping/generators.py new file mode 100644 index 00000000000..ceb1cf2fd7a --- /dev/null +++ b/ddtrace/internal/wrapping/generators.py @@ -0,0 +1 @@ +from ddtrace.internal.utils.wrapping.generators import * # noqa From 3e8383c0530f2a5ff3e656c4165ab31181f1a699 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 11:37:45 -0700 Subject: [PATCH 10/29] fix some mypy errors --- mypy.ini | 92 +++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 91 insertions(+), 1 deletion(-) diff --git a/mypy.ini b/mypy.ini index 257747fed63..521fe031665 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1,7 +1,7 @@ [mypy] files = ddtrace, docs -exclude = ddtrace/appsec/_iast/_taint_tracking/cmake-build-debug/|ddtrace/appsec/_iast/_taint_tracking/_vendor/|ddtrace/appsec/_iast/_taint_tracking/_deps/|ddtrace/internal/datadog/profiling/build|ddtrace/profiling/collector/test/build/|tests/profiling/collector/pprof_\d.*_pb2\.py +exclude = ddtrace/internal/constants\.py|ddtrace/internal/compat\.py|ddtrace/internal/module\.py|ddtrace/internal/threads\.py|ddtrace/internal/logger\.py|ddtrace/internal/_exceptions\.py|ddtrace/internal/schema/__init__\.py|ddtrace/internal/schema/processor\.py|ddtrace/internal/schema/span_attribute_schema\.py|ddtrace/internal/wrapping/__init__\.py|ddtrace/internal/wrapping/context\.py|ddtrace/internal/wrapping/generators\.py|ddtrace/internal/wrapping/asyncs\.py|ddtrace/appsec/_iast/_taint_tracking/cmake-build-debug/|ddtrace/appsec/_iast/_taint_tracking/_vendor/|ddtrace/appsec/_iast/_taint_tracking/_deps/|ddtrace/internal/datadog/profiling/build|ddtrace/profiling/collector/test/build/|tests/profiling/collector/pprof_\d.*_pb2\.py # mypy thinks .pyx files are scripts and errors out if it finds multiple scripts scripts_are_modules = true show_error_codes = true @@ -43,6 +43,40 @@ no_implicit_reexport = false ignore_errors = true no_implicit_reexport = false +# ddtrace.internal.* shims re-export from ddtrace.internal.utils.* via `import *`. +# Exclude them from the file list so mypy does not build an empty module namespace, +# and skip following imports so importers treat them as Any (avoids attr-defined). + +[mypy-ddtrace.internal._exceptions] +follow_imports = skip + +[mypy-ddtrace.internal.compat] +follow_imports = skip + +[mypy-ddtrace.internal.constants] +follow_imports = skip + +[mypy-ddtrace.internal.logger] +follow_imports = skip + +[mypy-ddtrace.internal.module] +follow_imports = skip + +[mypy-ddtrace.internal.threads] +follow_imports = skip + +[mypy-ddtrace.internal.wrapping] +follow_imports = skip + +[mypy-ddtrace.internal.wrapping.asyncs] +follow_imports = skip + +[mypy-ddtrace.internal.wrapping.context] +follow_imports = skip + +[mypy-ddtrace.internal.wrapping.generators] +follow_imports = skip + [mypy-setup] check_untyped_defs = true disallow_untyped_defs = false @@ -109,6 +143,7 @@ disallow_incomplete_defs = false [mypy-ddtrace._trace.context] disallow_untyped_calls = false +warn_return_any = false [mypy-ddtrace._trace.pin] warn_return_any = false @@ -137,6 +172,7 @@ disallow_any_generics = false disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false +warn_return_any = false [mypy-ddtrace._trace.sampling_rule] disallow_untyped_defs = false @@ -145,6 +181,8 @@ disallow_incomplete_defs = false [mypy-ddtrace._trace.span] disallow_untyped_defs = false disallow_incomplete_defs = false +warn_return_any = false +warn_unused_ignores = false warn_unreachable = false [mypy-ddtrace._trace.subscribers._base] @@ -215,6 +253,7 @@ disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false check_untyped_defs = false +warn_unused_ignores = false # ── ddtrace.appsec ────────────────────────────────────────────────────── @@ -589,6 +628,7 @@ disallow_untyped_calls = false [mypy-ddtrace.appsec._patch_utils] disallow_any_generics = false +disallow_untyped_decorators = false [mypy-ddtrace.appsec._processor] disallow_any_generics = false @@ -607,6 +647,7 @@ warn_return_any = false [mypy-ddtrace.appsec.ai_guard._api_client] disallow_any_generics = false +disallow_subclassing_any = false disallow_untyped_calls = false warn_return_any = false @@ -671,6 +712,14 @@ disallow_untyped_defs = false disallow_incomplete_defs = false +# ── ddtrace.debugging ─────────────────────────────────────────────────── + +[mypy-ddtrace.debugging.*] +disallow_subclassing_any = false +warn_return_any = false +warn_unused_ignores = false + + # ── ddtrace.errortracking ─────────────────────────────────────────────── [mypy-ddtrace.errortracking._handled_exceptions.bytecode_injector] @@ -679,6 +728,7 @@ disallow_untyped_defs = false disallow_incomplete_defs = false [mypy-ddtrace.errortracking._handled_exceptions.bytecode_reporting] +disallow_subclassing_any = false disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false @@ -696,6 +746,7 @@ disallow_untyped_defs = false disallow_incomplete_defs = false [mypy-ddtrace.errortracking._handled_exceptions.monitoring_reporting] +disallow_subclassing_any = false disallow_untyped_defs = false disallow_incomplete_defs = false check_untyped_defs = false @@ -991,6 +1042,13 @@ disallow_incomplete_defs = false # ── ddtrace.profiling ─────────────────────────────────────────────────── +[mypy-ddtrace.profiling._asyncio] +disallow_untyped_decorators = false +warn_unused_ignores = false + +[mypy-ddtrace.profiling._faulthandler] +disallow_untyped_decorators = false + [mypy-ddtrace.profiling._gevent] disallow_subclassing_any = false @@ -1015,6 +1073,7 @@ disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false check_untyped_defs = false +warn_return_any = false warn_unreachable = false @@ -1256,6 +1315,7 @@ disallow_untyped_defs = false disallow_incomplete_defs = false [mypy-ddtrace.internal.coverage.code] +disallow_subclassing_any = false disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false @@ -1617,6 +1677,7 @@ warn_return_any = false disallow_untyped_defs = false disallow_incomplete_defs = false no_implicit_reexport = false +warn_return_any = false warn_unreachable = false [mypy-ddtrace.internal.sca.product] @@ -1626,20 +1687,28 @@ disallow_untyped_defs = false disallow_incomplete_defs = false [mypy-ddtrace.internal.schema] +follow_imports = skip disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false [mypy-ddtrace.internal.schema.processor] +follow_imports = skip disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false check_untyped_defs = false [mypy-ddtrace.internal.schema.span_attribute_schema] +follow_imports = skip disallow_untyped_defs = false disallow_incomplete_defs = false +[mypy-ddtrace.internal.symbol_db.symbols] +disallow_subclassing_any = false +disallow_untyped_decorators = false +warn_return_any = false + [mypy-ddtrace.internal.serverless] warn_return_any = false @@ -1816,6 +1885,7 @@ check_untyped_defs = false warn_return_any = false no_implicit_reexport = false warn_unreachable = false +disable_error_code = arg-type [mypy-ddtrace.internal.utils.importlib] warn_return_any = false @@ -1832,6 +1902,21 @@ disallow_any_generics = false disallow_untyped_defs = false disallow_incomplete_defs = false +[mypy-ddtrace.internal.utils.schema] +disallow_untyped_calls = false +disallow_untyped_defs = false +disallow_incomplete_defs = false + +[mypy-ddtrace.internal.utils.schema.processor] +disallow_untyped_calls = false +disallow_untyped_defs = false +disallow_incomplete_defs = false +check_untyped_defs = false + +[mypy-ddtrace.internal.utils.schema.span_attribute_schema] +disallow_untyped_defs = false +disallow_incomplete_defs = false + [mypy-ddtrace.internal.utils.signals] disallow_untyped_calls = false disallow_untyped_defs = false @@ -1846,6 +1931,9 @@ warn_return_any = false [mypy-ddtrace.internal.utils.version] disallow_untyped_calls = false +[mypy-ddtrace.internal.utils.wrappers] +warn_return_any = false + [mypy-ddtrace.internal.uwsgi] disallow_any_generics = false disallow_untyped_defs = false @@ -2177,6 +2265,7 @@ disallow_untyped_calls = false disallow_untyped_calls = false disallow_untyped_defs = false disallow_incomplete_defs = false +warn_unused_ignores = false [mypy-ddtrace.contrib.internal.celery.signals] ignore_errors = true @@ -2505,6 +2594,7 @@ disallow_untyped_calls = false [mypy-ddtrace.contrib.internal.httpx.utils] disallow_untyped_defs = false disallow_incomplete_defs = false +warn_return_any = false [mypy-ddtrace.contrib.internal.jinja2.patch] disallow_untyped_calls = false From c71edbb73b0225219c199e0de6c24ca77ec3fbff Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 12:11:09 -0700 Subject: [PATCH 11/29] module __all__ --- ddtrace/internal/utils/module.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ddtrace/internal/utils/module.py b/ddtrace/internal/utils/module.py index 2cfbb600d81..20da8258f68 100644 --- a/ddtrace/internal/utils/module.py +++ b/ddtrace/internal/utils/module.py @@ -769,3 +769,6 @@ def __getattr__(name: str) -> t.Any: raise h _globals["__getattr__"] = __getattr__ + + +__all__ = list(locals().keys()) From 226ab3562298d5f52b546abc07b27055fc8bd4bb Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 12:12:16 -0700 Subject: [PATCH 12/29] more __all__ --- ddtrace/internal/utils/_exceptions.py | 2 ++ ddtrace/internal/utils/logger.py | 2 ++ ddtrace/internal/utils/threads.py | 2 ++ 3 files changed, 6 insertions(+) diff --git a/ddtrace/internal/utils/_exceptions.py b/ddtrace/internal/utils/_exceptions.py index e2836e036be..a51938c306d 100644 --- a/ddtrace/internal/utils/_exceptions.py +++ b/ddtrace/internal/utils/_exceptions.py @@ -35,3 +35,5 @@ def find_exception( if found: return found return None + +__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/logger.py b/ddtrace/internal/utils/logger.py index 82d5b8dcbd4..07415fa4b8f 100644 --- a/ddtrace/internal/utils/logger.py +++ b/ddtrace/internal/utils/logger.py @@ -244,3 +244,5 @@ def get_log_injection_state(raw_config: Optional[str]) -> bool: normalized, ) return False + +__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/threads.py b/ddtrace/internal/utils/threads.py index faa47fdac3f..6e87c9bb558 100644 --- a/ddtrace/internal/utils/threads.py +++ b/ddtrace/internal/utils/threads.py @@ -216,3 +216,5 @@ def _before_fork() -> None: for thread in _threads_to_restart_after_fork: log.debug("Joining thread %s before fork", thread.name) thread.join() + +__all__ = list(locals().keys()) From d5da0cf035e3f48b0768bcd016034316d7578811 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 12:18:43 -0700 Subject: [PATCH 13/29] fmt --- ddtrace/internal/utils/_exceptions.py | 1 + ddtrace/internal/utils/logger.py | 1 + ddtrace/internal/utils/threads.py | 1 + 3 files changed, 3 insertions(+) diff --git a/ddtrace/internal/utils/_exceptions.py b/ddtrace/internal/utils/_exceptions.py index a51938c306d..c69546158c8 100644 --- a/ddtrace/internal/utils/_exceptions.py +++ b/ddtrace/internal/utils/_exceptions.py @@ -36,4 +36,5 @@ def find_exception( return found return None + __all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/logger.py b/ddtrace/internal/utils/logger.py index 07415fa4b8f..3e1b32d321d 100644 --- a/ddtrace/internal/utils/logger.py +++ b/ddtrace/internal/utils/logger.py @@ -245,4 +245,5 @@ def get_log_injection_state(raw_config: Optional[str]) -> bool: ) return False + __all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/threads.py b/ddtrace/internal/utils/threads.py index 6e87c9bb558..80c283cedaf 100644 --- a/ddtrace/internal/utils/threads.py +++ b/ddtrace/internal/utils/threads.py @@ -217,4 +217,5 @@ def _before_fork() -> None: log.debug("Joining thread %s before fork", thread.name) thread.join() + __all__ = list(locals().keys()) From 2cbeb599a295b82a436008e513b70e58bd2200d2 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 12:34:23 -0700 Subject: [PATCH 14/29] don't check stderr --- tests/contrib/yaaredis/test_yaaredis.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/contrib/yaaredis/test_yaaredis.py b/tests/contrib/yaaredis/test_yaaredis.py index 2a2968e5f8c..84381a95906 100644 --- a/tests/contrib/yaaredis/test_yaaredis.py +++ b/tests/contrib/yaaredis/test_yaaredis.py @@ -174,7 +174,6 @@ async def test_basics(traced_yaaredis): env["PYTHONWARNINGS"] = "ignore::UserWarning:ddtrace.internal.module" out, err, status, _ = ddtrace_run_python_code_in_subprocess(code, env=env) assert status == 0, (err.decode(), out.decode()) - assert err == b"", err.decode() @pytest.mark.subprocess( From 2671256c43a43c0bb4023944870395b4c68ad2df Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 12:48:33 -0700 Subject: [PATCH 15/29] update userwarnings path --- tests/contrib/yaaredis/test_yaaredis.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/contrib/yaaredis/test_yaaredis.py b/tests/contrib/yaaredis/test_yaaredis.py index 84381a95906..80ca24e126f 100644 --- a/tests/contrib/yaaredis/test_yaaredis.py +++ b/tests/contrib/yaaredis/test_yaaredis.py @@ -171,13 +171,13 @@ async def test_basics(traced_yaaredis): env["DD_SERVICE"] = service if schema: env["DD_TRACE_SPAN_ATTRIBUTE_SCHEMA"] = schema - env["PYTHONWARNINGS"] = "ignore::UserWarning:ddtrace.internal.module" + env["PYTHONWARNINGS"] = "ignore::UserWarning:ddtrace.internal.utils.module" out, err, status, _ = ddtrace_run_python_code_in_subprocess(code, env=env) assert status == 0, (err.decode(), out.decode()) @pytest.mark.subprocess( - env=dict(DD_REDIS_RESOURCE_ONLY_COMMAND="false", PYTHONWARNINGS="ignore::UserWarning:ddtrace.internal.module") + env=dict(DD_REDIS_RESOURCE_ONLY_COMMAND="false", PYTHONWARNINGS="ignore::UserWarning:ddtrace.internal.utils.module") ) @pytest.mark.snapshot def test_full_command_in_resource_env(): From 759cb908f4e0721748e41b293beaa4c84adf7361 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 2 Jul 2026 14:16:19 -0700 Subject: [PATCH 16/29] update patching path --- tests/tracer/test_logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tracer/test_logger.py b/tests/tracer/test_logger.py index 2636eac27e0..10e6b6ba278 100644 --- a/tests/tracer/test_logger.py +++ b/tests/tracer/test_logger.py @@ -342,7 +342,7 @@ def test_logger_log_level_from_env(monkeypatch): monkeypatch.setenv("_DD_EXAMPLE_WARNING_LOG_LEVEL", "WARNING") monkeypatch.setenv("_DD_PACKAGE_WITH_UNDERSCORE_SUBMODULE_LOG_LEVEL", "ERROR") - import ddtrace.internal.logger as dd_logger + import ddtrace.internal.utils.logger as dd_logger with mock.patch.object(dd_logger, "LOG_LEVEL_TRIE", dd_logger.LoggerPrefix.build_trie()): assert get_logger("ddtrace.example.debug.foo.bar").level == logging.DEBUG From d7fcf5543219a9ee70621061ce3b997368d67059 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Fri, 3 Jul 2026 06:57:07 -0700 Subject: [PATCH 17/29] update patching path --- tests/tracer/test_logger.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/tracer/test_logger.py b/tests/tracer/test_logger.py index 10e6b6ba278..032bf9e6f0c 100644 --- a/tests/tracer/test_logger.py +++ b/tests/tracer/test_logger.py @@ -4,9 +4,9 @@ import mock import pytest -import ddtrace.internal.logger from ddtrace.internal.logger import LoggingBucket from ddtrace.internal.logger import get_logger +import ddtrace.internal.utils.logger from tests.utils import BaseTestCase @@ -20,8 +20,8 @@ def setUp(self): self.manager = logging.root.manager # Reset to default values - ddtrace.internal.logger._buckets.clear() - ddtrace.internal.logger._rate_limit = 60 + ddtrace.internal.utils.logger._buckets.clear() + ddtrace.internal.utils.logger._rate_limit = 60 def tearDown(self): # Weeee, forget all existing loggers @@ -31,8 +31,8 @@ def tearDown(self): self.manager = None # Reset to default values - ddtrace.internal.logger._buckets.clear() - ddtrace.internal.logger._rate_limit = 60 + ddtrace.internal.utils.logger._buckets.clear() + ddtrace.internal.utils.logger._rate_limit = 60 super(LoggerTestCase, self).tearDown() @@ -69,7 +69,7 @@ def test_get_logger(self): # Fetch a new logger log = get_logger("test.logger") - assert ddtrace.internal.logger.log_filter in log.filters + assert ddtrace.internal.utils.logger.log_filter in log.filters self.assertEqual(log.name, "test.logger") self.assertEqual(log.level, logging.NOTSET) @@ -106,10 +106,10 @@ def test_logger_handle_no_limit(self, call_handlers): # Configure an INFO logger with no rate limit log = get_logger("test.logger") log.setLevel(logging.INFO) - ddtrace.internal.logger._rate_limit = 0 + ddtrace.internal.utils.logger._rate_limit = 0 # Clear buckets in case any were created during logger setup - ddtrace.internal.logger._buckets.clear() + ddtrace.internal.utils.logger._buckets.clear() # Log a bunch of times very quickly (this is fast) for _ in range(1000): @@ -119,7 +119,7 @@ def test_logger_handle_no_limit(self, call_handlers): self.assertEqual(call_handlers.call_count, 1000) # Our buckets are empty (no rate limit means no buckets should be created) - self.assertEqual(ddtrace.internal.logger._buckets, dict()) + self.assertEqual(ddtrace.internal.utils.logger._buckets, dict()) @mock.patch("logging.Logger.callHandlers") def test_logger_handle_debug(self, call_handlers): @@ -132,10 +132,10 @@ def test_logger_handle_debug(self, call_handlers): log = get_logger("test.logger") log.setLevel(logging.DEBUG) assert log.getEffectiveLevel() == logging.DEBUG - assert ddtrace.internal.logger._rate_limit > 0 + assert ddtrace.internal.utils.logger._rate_limit > 0 # Clear buckets in case any were created during logger setup - ddtrace.internal.logger._buckets.clear() + ddtrace.internal.utils.logger._buckets.clear() # Log a bunch of times very quickly (this is fast) for level in ALL_LEVEL_NAMES: @@ -168,7 +168,7 @@ def test_logger_handle_bucket(self, call_handlers): # We added an bucket entry for this record key = (record.name, record.levelno, record.pathname, record.lineno) - logging_bucket = ddtrace.internal.logger._buckets.get(key) + logging_bucket = ddtrace.internal.utils.logger._buckets.get(key) self.assertIsInstance(logging_bucket, LoggingBucket) # The bucket entry is correct @@ -203,7 +203,7 @@ def test_logger_handle_bucket_limited(self, call_handlers): # We added an bucket entry for these records key = (record.name, record.levelno, record.pathname, record.lineno) - logging_bucket = ddtrace.internal.logger._buckets.get(key) + logging_bucket = ddtrace.internal.utils.logger._buckets.get(key) assert logging_bucket is not None # The bucket entry is correct @@ -229,7 +229,7 @@ def test_logger_handle_bucket_skipped_msg(self, call_handlers): key = (record.name, record.levelno, record.pathname, record.lineno) bucket = time.monotonic() # We want the time bucket to be for an older bucket - ddtrace.internal.logger._buckets[key] = LoggingBucket(bucket=bucket - 60, skipped=20) + ddtrace.internal.utils.logger._buckets[key] = LoggingBucket(bucket=bucket - 60, skipped=20) # Handle our record log.handle(record) @@ -275,7 +275,7 @@ def get_key(record): all_records = (record1, record2, record3, record4, record5, record6) [log.handle(record) for record in all_records] - buckets = ddtrace.internal.logger._buckets + buckets = ddtrace.internal.utils.logger._buckets # We have 6 records but only end up with 5 buckets self.assertEqual(len(buckets), 5) From b0213710ec3181f40571646bffa4526a007dd4c5 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Tue, 7 Jul 2026 12:02:03 -0700 Subject: [PATCH 18/29] more runners for internal tests --- tests/suitespec.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/suitespec.yml b/tests/suitespec.yml index 62fabae7f35..999c7bf922d 100644 --- a/tests/suitespec.yml +++ b/tests/suitespec.yml @@ -1,4 +1,3 @@ ---- components: $harness: - docker/* @@ -211,7 +210,7 @@ suites: - scripts/integration_registry/* internal: retry: 2 - venvs_per_job: 2 + venvs_per_job: 1 paths: - '@core' - '@remoteconfig' From bfcea88cea021f2401c414fdce08270b0e9450d2 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 07:46:19 -0700 Subject: [PATCH 19/29] experiment with older wrapt --- .../requirements/{a512f6c.txt => 1177c10.txt} | 8 ++-- .../requirements/{1b45a4d.txt => 12c807b.txt} | 6 +-- .../requirements/{c056455.txt => 16fe4e4.txt} | 10 ++--- .../requirements/{1676ee4.txt => 18ef7d5.txt} | 8 ++-- .../requirements/{15788e4.txt => 1d157d4.txt} | 8 ++-- .../requirements/{4b16172.txt => f275e43.txt} | 8 ++-- riotfile.py | 2 +- supported_versions.json | 42 +++++++++++++++++++ 8 files changed, 67 insertions(+), 25 deletions(-) rename .riot/requirements/{a512f6c.txt => 1177c10.txt} (91%) rename .riot/requirements/{1b45a4d.txt => 12c807b.txt} (94%) rename .riot/requirements/{c056455.txt => 16fe4e4.txt} (89%) rename .riot/requirements/{1676ee4.txt => 18ef7d5.txt} (91%) rename .riot/requirements/{15788e4.txt => 1d157d4.txt} (91%) rename .riot/requirements/{4b16172.txt => f275e43.txt} (91%) diff --git a/.riot/requirements/a512f6c.txt b/.riot/requirements/1177c10.txt similarity index 91% rename from .riot/requirements/a512f6c.txt rename to .riot/requirements/1177c10.txt index eb3bd311bd7..20209b99f03 100644 --- a/.riot/requirements/a512f6c.txt +++ b/.riot/requirements/1177c10.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.12 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/a512f6c.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/1177c10.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.14.3 +coverage[toml]==7.15.0 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.2 +greenlet==3.5.3 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -30,7 +30,7 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 uwsgi==2.0.31 -wrapt==2.2.2 +wrapt==2.1.1 zope-event==5.0 zope-interface==7.2 diff --git a/.riot/requirements/1b45a4d.txt b/.riot/requirements/12c807b.txt similarity index 94% rename from .riot/requirements/1b45a4d.txt rename to .riot/requirements/12c807b.txt index a9aa87cbe55..13373873e49 100644 --- a/.riot/requirements/1b45a4d.txt +++ b/.riot/requirements/12c807b.txt @@ -2,7 +2,7 @@ # This file is autogenerated by pip-compile with Python 3.9 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/1b45a4d.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/12c807b.in # attrs==26.1.0 cloudpickle==3.1.2 @@ -32,9 +32,9 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 tomli==2.4.1 -typing-extensions==4.15.0 +typing-extensions==4.16.0 uwsgi==2.0.31 -wrapt==2.2.2 +wrapt==2.1.1 zipp==3.23.1 zope-event==6.0 zope-interface==8.0.1 diff --git a/.riot/requirements/c056455.txt b/.riot/requirements/16fe4e4.txt similarity index 89% rename from .riot/requirements/c056455.txt rename to .riot/requirements/16fe4e4.txt index e3494731225..d5a54cbd18d 100644 --- a/.riot/requirements/c056455.txt +++ b/.riot/requirements/16fe4e4.txt @@ -2,15 +2,15 @@ # This file is autogenerated by pip-compile with Python 3.10 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/c056455.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/16fe4e4.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.14.3 +coverage[toml]==7.15.0 exceptiongroup==1.3.1 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.2 +greenlet==3.5.3 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -31,9 +31,9 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 tomli==2.4.1 -typing-extensions==4.15.0 +typing-extensions==4.16.0 uwsgi==2.0.31 -wrapt==2.2.2 +wrapt==2.1.1 zope-event==6.2 zope-interface==8.5 diff --git a/.riot/requirements/1676ee4.txt b/.riot/requirements/18ef7d5.txt similarity index 91% rename from .riot/requirements/1676ee4.txt rename to .riot/requirements/18ef7d5.txt index 7d95c4cb7d3..449dc6013a5 100644 --- a/.riot/requirements/1676ee4.txt +++ b/.riot/requirements/18ef7d5.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.14 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/1676ee4.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/18ef7d5.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.14.3 +coverage[toml]==7.15.0 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.2 +greenlet==3.5.3 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -30,7 +30,7 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 uwsgi==2.0.31 -wrapt==2.2.2 +wrapt==2.1.1 zope-event==5.0 zope-interface==7.2 diff --git a/.riot/requirements/15788e4.txt b/.riot/requirements/1d157d4.txt similarity index 91% rename from .riot/requirements/15788e4.txt rename to .riot/requirements/1d157d4.txt index 72b849cb343..5551dd81380 100644 --- a/.riot/requirements/15788e4.txt +++ b/.riot/requirements/1d157d4.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.11 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/15788e4.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/1d157d4.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.14.3 +coverage[toml]==7.15.0 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.2 +greenlet==3.5.3 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -30,7 +30,7 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 uwsgi==2.0.31 -wrapt==2.2.2 +wrapt==2.1.1 zope-event==6.2 zope-interface==8.5 diff --git a/.riot/requirements/4b16172.txt b/.riot/requirements/f275e43.txt similarity index 91% rename from .riot/requirements/4b16172.txt rename to .riot/requirements/f275e43.txt index b9382584ad5..ff3e4e3710c 100644 --- a/.riot/requirements/4b16172.txt +++ b/.riot/requirements/f275e43.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.13 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/4b16172.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/f275e43.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.14.3 +coverage[toml]==7.15.0 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.2 +greenlet==3.5.3 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -30,7 +30,7 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 uwsgi==2.0.31 -wrapt==2.2.2 +wrapt==2.1.1 zope-event==5.0 zope-interface==7.2 diff --git a/riotfile.py b/riotfile.py index 7686782c81b..09b1deac9bc 100644 --- a/riotfile.py +++ b/riotfile.py @@ -604,7 +604,7 @@ def select_pys(min_version: str = MIN_PYTHON_VERSION, max_version: str = MAX_PYT "python-json-logger": "==2.0.7", "pyfakefs": latest, "pytest-benchmark": latest, - "wrapt": [latest, "<2.0.0"], + "wrapt": ["==2.1.1", "<2.0.0"], "uwsgi": latest, # Ray Serve cloudpickles objects that reference the global config # (which holds forksafe locks); test_forksafe.py exercises that path. diff --git a/supported_versions.json b/supported_versions.json index a2248e4afed..bed961a86a2 100644 --- a/supported_versions.json +++ b/supported_versions.json @@ -3458,6 +3458,48 @@ } } }, + { + "dependency": "mistralai", + "integration": "mistralai", + "auto-instrumented": true, + "python_versions": { + "3.10": { + "minimum_package_version": "2.4.9", + "maximum_package_version": "2.4.9", + "tested_versions": [ + "2.4.9" + ] + }, + "3.11": { + "minimum_package_version": "2.4.9", + "maximum_package_version": "2.4.9", + "tested_versions": [ + "2.4.9" + ] + }, + "3.12": { + "minimum_package_version": "2.4.9", + "maximum_package_version": "2.4.9", + "tested_versions": [ + "2.4.9" + ] + }, + "3.13": { + "minimum_package_version": "2.4.9", + "maximum_package_version": "2.4.9", + "tested_versions": [ + "2.4.9" + ] + }, + "3.14": { + "minimum_package_version": "2.4.9", + "maximum_package_version": "2.4.9", + "tested_versions": [ + "2.4.9" + ] + } + } + }, { "dependency": "mlflow", "integration": "mlflow", From c8b040122297f3e5847c4582b684b5a0943042b1 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 08:12:58 -0700 Subject: [PATCH 20/29] Revert "experiment with older wrapt" This reverts commit bfcea88cea021f2401c414fdce08270b0e9450d2. --- .../requirements/{1d157d4.txt => 15788e4.txt} | 8 ++-- .../requirements/{18ef7d5.txt => 1676ee4.txt} | 8 ++-- .../requirements/{12c807b.txt => 1b45a4d.txt} | 6 +-- .../requirements/{f275e43.txt => 4b16172.txt} | 8 ++-- .../requirements/{1177c10.txt => a512f6c.txt} | 8 ++-- .../requirements/{16fe4e4.txt => c056455.txt} | 10 ++--- riotfile.py | 2 +- supported_versions.json | 42 ------------------- 8 files changed, 25 insertions(+), 67 deletions(-) rename .riot/requirements/{1d157d4.txt => 15788e4.txt} (91%) rename .riot/requirements/{18ef7d5.txt => 1676ee4.txt} (91%) rename .riot/requirements/{12c807b.txt => 1b45a4d.txt} (94%) rename .riot/requirements/{f275e43.txt => 4b16172.txt} (91%) rename .riot/requirements/{1177c10.txt => a512f6c.txt} (91%) rename .riot/requirements/{16fe4e4.txt => c056455.txt} (89%) diff --git a/.riot/requirements/1d157d4.txt b/.riot/requirements/15788e4.txt similarity index 91% rename from .riot/requirements/1d157d4.txt rename to .riot/requirements/15788e4.txt index 5551dd81380..72b849cb343 100644 --- a/.riot/requirements/1d157d4.txt +++ b/.riot/requirements/15788e4.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.11 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/1d157d4.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/15788e4.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.15.0 +coverage[toml]==7.14.3 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.3 +greenlet==3.5.2 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -30,7 +30,7 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 uwsgi==2.0.31 -wrapt==2.1.1 +wrapt==2.2.2 zope-event==6.2 zope-interface==8.5 diff --git a/.riot/requirements/18ef7d5.txt b/.riot/requirements/1676ee4.txt similarity index 91% rename from .riot/requirements/18ef7d5.txt rename to .riot/requirements/1676ee4.txt index 449dc6013a5..7d95c4cb7d3 100644 --- a/.riot/requirements/18ef7d5.txt +++ b/.riot/requirements/1676ee4.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.14 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/18ef7d5.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/1676ee4.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.15.0 +coverage[toml]==7.14.3 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.3 +greenlet==3.5.2 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -30,7 +30,7 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 uwsgi==2.0.31 -wrapt==2.1.1 +wrapt==2.2.2 zope-event==5.0 zope-interface==7.2 diff --git a/.riot/requirements/12c807b.txt b/.riot/requirements/1b45a4d.txt similarity index 94% rename from .riot/requirements/12c807b.txt rename to .riot/requirements/1b45a4d.txt index 13373873e49..a9aa87cbe55 100644 --- a/.riot/requirements/12c807b.txt +++ b/.riot/requirements/1b45a4d.txt @@ -2,7 +2,7 @@ # This file is autogenerated by pip-compile with Python 3.9 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/12c807b.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/1b45a4d.in # attrs==26.1.0 cloudpickle==3.1.2 @@ -32,9 +32,9 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 tomli==2.4.1 -typing-extensions==4.16.0 +typing-extensions==4.15.0 uwsgi==2.0.31 -wrapt==2.1.1 +wrapt==2.2.2 zipp==3.23.1 zope-event==6.0 zope-interface==8.0.1 diff --git a/.riot/requirements/f275e43.txt b/.riot/requirements/4b16172.txt similarity index 91% rename from .riot/requirements/f275e43.txt rename to .riot/requirements/4b16172.txt index ff3e4e3710c..b9382584ad5 100644 --- a/.riot/requirements/f275e43.txt +++ b/.riot/requirements/4b16172.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.13 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/f275e43.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/4b16172.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.15.0 +coverage[toml]==7.14.3 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.3 +greenlet==3.5.2 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -30,7 +30,7 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 uwsgi==2.0.31 -wrapt==2.1.1 +wrapt==2.2.2 zope-event==5.0 zope-interface==7.2 diff --git a/.riot/requirements/1177c10.txt b/.riot/requirements/a512f6c.txt similarity index 91% rename from .riot/requirements/1177c10.txt rename to .riot/requirements/a512f6c.txt index 20209b99f03..eb3bd311bd7 100644 --- a/.riot/requirements/1177c10.txt +++ b/.riot/requirements/a512f6c.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.12 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/1177c10.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/a512f6c.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.15.0 +coverage[toml]==7.14.3 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.3 +greenlet==3.5.2 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -30,7 +30,7 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 uwsgi==2.0.31 -wrapt==2.1.1 +wrapt==2.2.2 zope-event==5.0 zope-interface==7.2 diff --git a/.riot/requirements/16fe4e4.txt b/.riot/requirements/c056455.txt similarity index 89% rename from .riot/requirements/16fe4e4.txt rename to .riot/requirements/c056455.txt index d5a54cbd18d..e3494731225 100644 --- a/.riot/requirements/16fe4e4.txt +++ b/.riot/requirements/c056455.txt @@ -2,15 +2,15 @@ # This file is autogenerated by pip-compile with Python 3.10 # by the following command: # -# pip-compile --allow-unsafe --no-annotate .riot/requirements/16fe4e4.in +# pip-compile --allow-unsafe --no-annotate .riot/requirements/c056455.in # attrs==26.1.0 cloudpickle==3.1.2 -coverage[toml]==7.15.0 +coverage[toml]==7.14.3 exceptiongroup==1.3.1 execnet==2.1.2 gevent==26.5.0 -greenlet==3.5.3 +greenlet==3.5.2 httpretty==1.1.4 hypothesis==6.45.0 iniconfig==2.3.0 @@ -31,9 +31,9 @@ pytest-xdist==3.8.0 python-json-logger==2.0.7 sortedcontainers==2.4.0 tomli==2.4.1 -typing-extensions==4.16.0 +typing-extensions==4.15.0 uwsgi==2.0.31 -wrapt==2.1.1 +wrapt==2.2.2 zope-event==6.2 zope-interface==8.5 diff --git a/riotfile.py b/riotfile.py index 09b1deac9bc..7686782c81b 100644 --- a/riotfile.py +++ b/riotfile.py @@ -604,7 +604,7 @@ def select_pys(min_version: str = MIN_PYTHON_VERSION, max_version: str = MAX_PYT "python-json-logger": "==2.0.7", "pyfakefs": latest, "pytest-benchmark": latest, - "wrapt": ["==2.1.1", "<2.0.0"], + "wrapt": [latest, "<2.0.0"], "uwsgi": latest, # Ray Serve cloudpickles objects that reference the global config # (which holds forksafe locks); test_forksafe.py exercises that path. diff --git a/supported_versions.json b/supported_versions.json index bed961a86a2..a2248e4afed 100644 --- a/supported_versions.json +++ b/supported_versions.json @@ -3458,48 +3458,6 @@ } } }, - { - "dependency": "mistralai", - "integration": "mistralai", - "auto-instrumented": true, - "python_versions": { - "3.10": { - "minimum_package_version": "2.4.9", - "maximum_package_version": "2.4.9", - "tested_versions": [ - "2.4.9" - ] - }, - "3.11": { - "minimum_package_version": "2.4.9", - "maximum_package_version": "2.4.9", - "tested_versions": [ - "2.4.9" - ] - }, - "3.12": { - "minimum_package_version": "2.4.9", - "maximum_package_version": "2.4.9", - "tested_versions": [ - "2.4.9" - ] - }, - "3.13": { - "minimum_package_version": "2.4.9", - "maximum_package_version": "2.4.9", - "tested_versions": [ - "2.4.9" - ] - }, - "3.14": { - "minimum_package_version": "2.4.9", - "maximum_package_version": "2.4.9", - "tested_versions": [ - "2.4.9" - ] - } - } - }, { "dependency": "mlflow", "integration": "mlflow", From 55000ddb5e17e31b60f9357e643ffc0df03f1312 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 08:13:29 -0700 Subject: [PATCH 21/29] Revert "more runners for internal tests" This reverts commit b0213710ec3181f40571646bffa4526a007dd4c5. --- tests/suitespec.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/suitespec.yml b/tests/suitespec.yml index 999c7bf922d..62fabae7f35 100644 --- a/tests/suitespec.yml +++ b/tests/suitespec.yml @@ -1,3 +1,4 @@ +--- components: $harness: - docker/* @@ -210,7 +211,7 @@ suites: - scripts/integration_registry/* internal: retry: 2 - venvs_per_job: 1 + venvs_per_job: 2 paths: - '@core' - '@remoteconfig' From 8e2fd54b156df84bc6264ed2eb28e5000f7959af Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 08:35:19 -0700 Subject: [PATCH 22/29] try no-cov --- riotfile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/riotfile.py b/riotfile.py index 7686782c81b..cc8a8e0d7d6 100644 --- a/riotfile.py +++ b/riotfile.py @@ -595,7 +595,7 @@ def select_pys(min_version: str = MIN_PYTHON_VERSION, max_version: str = MAX_PYT "DD_CIVISIBILITY_ITR_ENABLED": "0", "DD_PYTEST_USE_NEW_PLUGIN": "false", }, - command="pytest -v -n auto {cmdargs} tests/internal/", + command="pytest -v --no-cov -n auto {cmdargs} tests/internal/", pkgs={ "httpretty": latest, "gevent": latest, From 4a79e3d985a95f47d81d1443f559821632c729d1 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 09:43:14 -0700 Subject: [PATCH 23/29] undo wrapping change --- ddtrace/internal/wrapping/__init__.py | 455 +++++++++++- ddtrace/internal/wrapping/asyncs.py | 726 +++++++++++++++++++- ddtrace/internal/wrapping/context.py | 877 +++++++++++++++++++++++- ddtrace/internal/wrapping/generators.py | 492 ++++++++++++- riotfile.py | 2 +- 5 files changed, 2547 insertions(+), 5 deletions(-) diff --git a/ddtrace/internal/wrapping/__init__.py b/ddtrace/internal/wrapping/__init__.py index 07fd80a109e..b01b8c7977d 100644 --- a/ddtrace/internal/wrapping/__init__.py +++ b/ddtrace/internal/wrapping/__init__.py @@ -1 +1,454 @@ -from ddtrace.internal.utils.wrapping import * # noqa +import sys +from types import CodeType +from types import FunctionType +from typing import Any +from typing import Callable +from typing import Iterator +from typing import MutableMapping +from typing import Optional +from typing import Protocol +from typing import cast +import weakref + +import bytecode as bc +from bytecode import Instr + +from ddtrace.internal.assembly import Assembly +from ddtrace.internal.threads import Lock +from ddtrace.internal.wrapping.asyncs import wrap_async +from ddtrace.internal.wrapping.generators import wrap_generator + + +PY = sys.version_info[:2] + +# Maps each wrapped function to its inner copy (the singly-linked list of +# wrapping layers). WeakKeyDictionary so functions are not kept alive by the +# registry alone. +_wrapped: weakref.WeakKeyDictionary[FunctionType, FunctionType] = weakref.WeakKeyDictionary() +_wrapped_lock = Lock() + +# Maps original code objects to the functions that own them. Written by +# link_function_to_code; read by functions_for_code in inspection.py. +# WeakValueDictionary so that wrapped ephemeral functions are not kept alive by +# this mapping alone — inspection.py falls back to gc.get_referrers on a miss. +_code_to_fn: MutableMapping[CodeType, FunctionType] = weakref.WeakValueDictionary() + + +def link_function_to_code(code: CodeType, function: FunctionType) -> None: + """Link a function to its original code object for fast reverse lookup.""" + _code_to_fn[code] = function + + +class WrappedFunction(Protocol): + """A wrapped function.""" + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + pass + + +Wrapper = Callable[[FunctionType, tuple[Any], dict[str, Any]], Any] + + +def _add(lineno: int) -> Instr: + if PY >= (3, 11): + return Instr("BINARY_OP", bc.BinaryOp.ADD, lineno=lineno) + + return Instr("INPLACE_ADD", lineno=lineno) + + +HEAD = Assembly() +if PY >= (3, 13): + HEAD.parse( + r""" + resume 0 + load_const {wrapper} + push_null + load_const {wrapped} + """ + ) + +elif PY >= (3, 11): + HEAD.parse( + r""" + resume 0 + push_null + load_const {wrapper} + load_const {wrapped} + """ + ) + +else: + HEAD.parse( + r""" + load_const {wrapper} + load_const {wrapped} + """ + ) + + +_UPDATE_MAP_FAST = Assembly() +_UPDATE_MAP_DEREF = Assembly() +if PY >= (3, 12): + _UPDATE_MAP_FAST.parse( + r""" + copy 1 + load_method $update + load_fast {varkwargsname} + call 1 + pop_top + """ + ) + _UPDATE_MAP_DEREF.parse( + r""" + copy 1 + load_method $update + load_deref {varkwargsname} + call 1 + pop_top + """ + ) +elif PY >= (3, 11): + _UPDATE_MAP_FAST.parse( + r""" + copy 1 + load_method $update + load_fast {varkwargsname} + precall 1 + call 1 + pop_top + """ + ) + _UPDATE_MAP_DEREF.parse( + r""" + copy 1 + load_method $update + load_deref {varkwargsname} + precall 1 + call 1 + pop_top + """ + ) +else: + _UPDATE_MAP_FAST.parse( + r""" + dup_top + load_attr $update + load_fast {varkwargsname} + call_function 1 + pop_top + """ + ) + _UPDATE_MAP_DEREF.parse( + r""" + dup_top + load_attr $update + load_deref {varkwargsname} + call_function 1 + pop_top + """ + ) + + +def _generate_update_map(name: str, code: CodeType, lineno: int) -> Iterator[Any]: + """Yield opcodes to call ``dict.update(name)`` where ``name`` may be a cell var.""" + if PY >= (3, 11) and name in code.co_cellvars: + yield from _UPDATE_MAP_DEREF.bind({"varkwargsname": bc.CellVar(name)}, lineno=lineno) # type: ignore[attr-defined] + else: + yield from _UPDATE_MAP_FAST.bind({"varkwargsname": name}, lineno=lineno) + + +CALL_RETURN = Assembly() +if PY >= (3, 12): + CALL_RETURN.parse( + r""" + call {arg} + return_value + """ + ) + +elif PY >= (3, 11): + CALL_RETURN.parse( + r""" + precall {arg} + call {arg} + return_value + """ + ) + +else: + CALL_RETURN.parse( + r""" + call_function {arg} + return_value + """ + ) + + +FIRSTLINENO_OFFSET = int(PY >= (3, 11)) + + +def _load_var(name: str, code: CodeType, lineno: int) -> Instr: + """Return the correct load instruction for a function parameter. + + On Python 3.11+, parameters captured by inner closures become cell + variables and must be loaded with LOAD_DEREF instead of LOAD_FAST. + """ + if PY >= (3, 11) and name in code.co_cellvars: + return Instr("LOAD_DEREF", bc.CellVar(name), lineno=lineno) # type: ignore[attr-defined] + return Instr("LOAD_FAST", name, lineno=lineno) + + +def generate_posargs(code: CodeType) -> Iterator[Any]: + """Generate the opcodes for building the positional arguments tuple.""" + varnames = code.co_varnames + lineno = code.co_firstlineno + FIRSTLINENO_OFFSET + varargs = bool(code.co_flags & bc.CompilerFlags.VARARGS) + nargs = code.co_argcount + varargsname: Optional[str] = varnames[nargs + code.co_kwonlyargcount] if varargs else None + + if nargs: # posargs [+ varargs] + yield from (_load_var(argname, code, lineno) for argname in varnames[:nargs]) + + yield Instr("BUILD_TUPLE", nargs, lineno=lineno) + if varargsname is not None: + yield _load_var(varargsname, code, lineno) + yield _add(lineno) + + elif varargsname is not None: # varargs + yield _load_var(varargsname, code, lineno) + + else: # () + yield Instr("BUILD_TUPLE", 0, lineno=lineno) + + +def generate_kwargs(code: CodeType) -> Iterator[Any]: + """Generate the opcodes for building the keyword arguments dictionary.""" + flags = code.co_flags + varnames = code.co_varnames + lineno = code.co_firstlineno + FIRSTLINENO_OFFSET + varargs = bool(flags & bc.CompilerFlags.VARARGS) + varkwargs = bool(flags & bc.CompilerFlags.VARKEYWORDS) + nargs = code.co_argcount + kwonlyargs = code.co_kwonlyargcount + varkwargsname: Optional[str] = varnames[nargs + kwonlyargs + varargs] if varkwargs else None + + if kwonlyargs: + for arg in varnames[nargs : nargs + kwonlyargs]: # kwargs [+ varkwargs] + yield Instr("LOAD_CONST", arg, lineno=lineno) + yield _load_var(arg, code, lineno) + yield Instr("BUILD_MAP", kwonlyargs, lineno=lineno) + if varkwargsname is not None: + yield from _generate_update_map(varkwargsname, code, lineno) + + elif varkwargsname is not None: # varkwargs + yield _load_var(varkwargsname, code, lineno) + + else: # {} + yield Instr("BUILD_MAP", 0, lineno=lineno) + + +def wrap_bytecode(wrapper: Wrapper, wrapped: FunctionType) -> bc.Bytecode: + """Wrap a function with a wrapper function. + + The wrapper function expects the wrapped function as the first argument, + followed by the tuple of arguments and the dictionary of keyword arguments. + The nature of the wrapped function is also honored, meaning that a generator + function will return a generator function, and a coroutine function will + return a coroutine function, and so on. The signature is also preserved to + avoid breaking, e.g., usages of the ``inspect`` module. + """ + + code = wrapped.__code__ + lineno = code.co_firstlineno + FIRSTLINENO_OFFSET + + # Push the wrapper function that is to be called and the wrapped function to + # be passed as first argument. + instrs = HEAD.bind({"wrapper": wrapper, "wrapped": wrapped}, lineno=lineno) + + # Add positional arguments + instrs.extend(generate_posargs(code)) + + # Add keyword arguments + instrs.extend(generate_kwargs(code)) + + # Call the wrapper function with the wrapped function, the positional and + # keyword arguments, and return the result. This is equivalent to + # + # >>> return wrapper(wrapped, args, kwargs) + instrs.extend(CALL_RETURN.bind({"arg": 3}, lineno=lineno)) + + # Include code for handling free/cell variables, if needed + if PY >= (3, 11): + if code.co_cellvars: + instrs[0:0] = [Instr("MAKE_CELL", bc.CellVar(_), lineno=lineno) for _ in code.co_cellvars] # type: ignore[attr-defined] + + if code.co_freevars: + instrs.insert(0, Instr("COPY_FREE_VARS", len(code.co_freevars), lineno=lineno)) + + # If the function has special flags set, like the generator, async generator + # or coroutine, inject unraveling code before the return opcode. + if (bc.CompilerFlags.GENERATOR & code.co_flags) and not (bc.CompilerFlags.COROUTINE & code.co_flags): + wrap_generator(instrs, code, lineno) + else: + wrap_async(instrs, code, lineno) + + return instrs + + +def wrap(f: FunctionType, wrapper: Wrapper) -> WrappedFunction: + """Wrap a function with a wrapper. + + The wrapper expects the function as first argument, followed by the tuple + of positional arguments and the dict of keyword arguments. + + Note that this changes the behavior of the original function with the + wrapper function, instead of creating a new function object. + """ + wrapped = FunctionType( + code := f.__code__, + f.__globals__, + "", + f.__defaults__, + f.__closure__, + ) + + # Carry forward the existing wrapped-function link to the new inner copy. + with _wrapped_lock: + existing_inner = _wrapped.get(f) + if existing_inner is not None: + _wrapped[wrapped] = existing_inner + + wrapped.__kwdefaults__ = f.__kwdefaults__ + + flags = code.co_flags + nargs = ( + (argcount := code.co_argcount) + + (kwonlycount := code.co_kwonlyargcount) + + bool(flags & bc.CompilerFlags.VARARGS) + + bool(flags & bc.CompilerFlags.VARKEYWORDS) + ) + + # Wrap the wrapped function with the wrapper + wrapped_code = wrap_bytecode(wrapper, wrapped) + + # Copy over the code attributes + wrapped_code.argcount = argcount + wrapped_code.argnames = list(code.co_varnames[:nargs]) + wrapped_code.filename = code.co_filename + wrapped_code.freevars = list(code.co_freevars) + wrapped_code.flags = bc.CompilerFlags(flags) + wrapped_code.kwonlyargcount = kwonlycount + wrapped_code.name = code.co_name + wrapped_code.posonlyargcount = code.co_posonlyargcount + if PY >= (3, 11): + wrapped_code.cellvars = list(code.co_cellvars) + + # Replace the function code with the trampoline bytecode + f.__code__ = wrapped_code.to_code() + + # DEV: Multiple wrapping is implemented as a singly-linked list via _wrapped. + with _wrapped_lock: + _wrapped[f] = wrapped + + # Link the original code object to the original function + link_function_to_code(code, f) + + return cast(WrappedFunction, f) + + +def _unwrap_method(f: Any) -> Any: + """Return the underlying function if *f* is a bound or unbound method.""" + func = getattr(f, "__func__", f) + return func + + +def is_wrapped(f: FunctionType) -> bool: + """Check if a function is wrapped with any wrapper.""" + try: + with _wrapped_lock: + inner = _wrapped.get(_unwrap_method(f)) + except TypeError: + # f is not weakly referenceable (e.g. C method_descriptor) + return False + if inner is None: + return False + assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + return True + + +def is_wrapped_with(f: FunctionType, wrapper: Wrapper) -> bool: + """Check if a function is wrapped with a specific wrapper.""" + f = cast(FunctionType, _unwrap_method(f)) + try: + with _wrapped_lock: + inner = _wrapped.get(f) + except TypeError: + return False + if inner is None: + return False + + assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + + if wrapper in f.__code__.co_consts: + return True + + # This is not the correct wrapping layer. Try with the next one. + return is_wrapped_with(inner, wrapper) + + +def unwrap(wf: WrappedFunction, wrapper: Wrapper) -> FunctionType: + """Unwrap a wrapped function. + + This is the reverse of :func:`wrap`. In case of multiple wrapping layers, + this will unwrap the one that uses ``wrapper``. If the function was not + wrapped with ``wrapper``, it will return the first argument. + """ + # DEV: Multiple wrapping layers are singly-linked via _wrapped. When we + # find the layer that needs to be removed we also have to ensure that we + # update the link at the deletion site if there is a non-empty tail. + f = cast(FunctionType, wf) + with _wrapped_lock: + inner = _wrapped.get(f) + if inner is None: + # The function is not wrapped so we return it as is. + return f + + assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + + if wrapper not in f.__code__.co_consts: + # This is not the correct wrapping layer. Try with the next one. + return unwrap(cast(WrappedFunction, inner), wrapper) + + # Remove the current wrapping layer by moving the next one over the current + # one. The code swap and registry update must be atomic: two concurrent + # unwrap calls on the same function and wrapper must not both succeed. + with _wrapped_lock: + f.__code__ = inner.__code__ + next_inner = _wrapped.get(inner) + if next_inner is not None: + _wrapped[f] = next_inner + else: + del _wrapped[f] + + return f + + +def get_function_code(f: FunctionType) -> CodeType: + with _wrapped_lock: + inner = _wrapped.get(f) + return (inner if inner is not None else f).__code__ + + +def set_function_code(f: FunctionType, code: CodeType) -> None: + with _wrapped_lock: + inner = _wrapped.get(f) + (inner if inner is not None else f).__code__ = code + + +def get_wrapped(f: FunctionType) -> Optional[FunctionType]: + """Return the inner bytecode copy of *f* if it is wrapped, else None.""" + try: + with _wrapped_lock: + return _wrapped.get(cast(FunctionType, _unwrap_method(f))) + except TypeError: + return None diff --git a/ddtrace/internal/wrapping/asyncs.py b/ddtrace/internal/wrapping/asyncs.py index 339af72067c..a1de28a074e 100644 --- a/ddtrace/internal/wrapping/asyncs.py +++ b/ddtrace/internal/wrapping/asyncs.py @@ -1 +1,725 @@ -from ddtrace.internal.utils.wrapping.asyncs import * # noqa +import sys +from types import CodeType + +import bytecode as bc + +from ddtrace.internal.assembly import Assembly + + +PY = sys.version_info[:2] + + +# ----------------------------------------------------------------------------- +# Coroutine and Async Generator Wrapping +# ----------------------------------------------------------------------------- +# DEV: The wrapping of async generators is roughly equivalent to +# +# __ddgen = wrapper(wrapped, args, kwargs) +# __ddgensend = __ddgen.asend +# try: +# value = await __ddgen.__anext__() +# while True: +# try: +# tosend = yield value +# except GeneratorExit: +# await __ddgen.aclose() +# except Exception: +# value = await __ddgen.athrow(*sys.exc_info()) +# else: +# value = await __ddgensend(tosend) +# except StopAsyncIteration: +# return +# ----------------------------------------------------------------------------- + +COROUTINE_ASSEMBLY = Assembly() +ASYNC_GEN_ASSEMBLY = Assembly() +ASYNC_HEAD_ASSEMBLY = None + +if PY >= (3, 15): + raise NotImplementedError("This version of CPython is not supported yet") + +elif PY >= (3, 14): + ASYNC_HEAD_ASSEMBLY = Assembly() + ASYNC_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + COROUTINE_ASSEMBLY.parse( + r""" + get_awaitable 0 + load_const None + + presend: + send @send + yield_value 1 + resume 3 + jump_backward_no_interrupt @presend + send: + end_send + """ + ) + + ASYNC_GEN_ASSEMBLY.parse( + r""" + try @stopiter + copy 1 + store_fast $__ddgen + load_attr (False, 'asend') + store_fast $__ddgensend + load_fast $__ddgen + load_attr (True, '__anext__') + call 0 + + loop: + get_awaitable 0 + load_const None + presend0: + send @send0 + tried + + try @genexit lasti + yield_value 0 + resume 3 + jump_backward_no_interrupt @presend0 + send0: + end_send + + yield: + call_intrinsic_1 asm.Intrinsic1Op.INTRINSIC_ASYNC_GEN_WRAP + yield_value 0 + resume 1 + push_null + load_fast $__ddgensend + swap 3 + call 1 + jump_backward @loop + tried + + genexit: + try @stopiter + push_exc_info + load_const GeneratorExit + check_exc_match + pop_jump_if_false @exc + pop_top + load_fast $__ddgen + load_attr (True, 'aclose') + call 0 + get_awaitable 0 + load_const None + + presend1: + send @send1 + yield_value 0 + resume 3 + jump_backward_no_interrupt @presend1 + send1: + end_send + pop_top + pop_except + load_const None + return_value + + exc: + pop_top + load_fast $__ddgen + load_attr (False, 'athrow') + push_null + load_const sys.exc_info + push_null + call 0 + push_null + call_function_ex + get_awaitable 0 + load_const None + + presend2: + send @send2 + yield_value 0 + resume 3 + jump_backward_no_interrupt @presend2 + send2: + end_send + swap 2 + pop_except + jump_backward @yield + tried + + stopiter: + push_exc_info + load_const StopAsyncIteration + check_exc_match + pop_jump_if_false @propagate + pop_top + pop_except + load_const None + return_value + + propagate: + reraise 0 + """ + ) + +elif PY >= (3, 13): + ASYNC_HEAD_ASSEMBLY = Assembly() + ASYNC_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + COROUTINE_ASSEMBLY.parse( + r""" + get_awaitable 0 + load_const None + + presend: + send @send + yield_value 1 + resume 3 + jump_backward_no_interrupt @presend + send: + end_send + """ + ) + + ASYNC_GEN_ASSEMBLY.parse( + r""" + try @stopiter + copy 1 + store_fast $__ddgen + load_attr (False, 'asend') + store_fast $__ddgensend + load_fast $__ddgen + load_attr (True, '__anext__') + call 0 + + loop: + get_awaitable 0 + load_const None + presend0: + send @send0 + tried + + try @genexit lasti + yield_value 0 + resume 3 + jump_backward_no_interrupt @presend0 + send0: + end_send + + yield: + call_intrinsic_1 asm.Intrinsic1Op.INTRINSIC_ASYNC_GEN_WRAP + yield_value 0 + resume 1 + push_null + load_fast $__ddgensend + swap 3 + call 1 + jump_backward @loop + tried + + genexit: + try @stopiter + push_exc_info + load_const GeneratorExit + check_exc_match + pop_jump_if_false @exc + pop_top + load_fast $__ddgen + load_attr (True, 'aclose') + call 0 + get_awaitable 0 + load_const None + + presend1: + send @send1 + yield_value 0 + resume 3 + jump_backward_no_interrupt @presend1 + send1: + end_send + pop_top + pop_except + load_const None + return_value + + exc: + pop_top + load_fast $__ddgen + load_attr (False, 'athrow') + push_null + load_const sys.exc_info + push_null + call 0 + call_function_ex 0 + get_awaitable 0 + load_const None + + presend2: + send @send2 + yield_value 0 + resume 3 + jump_backward_no_interrupt @presend2 + send2: + end_send + swap 2 + pop_except + jump_backward @yield + tried + + stopiter: + push_exc_info + load_const StopAsyncIteration + check_exc_match + pop_jump_if_false @propagate + pop_top + pop_except + load_const None + return_value + + propagate: + reraise 0 + """ + ) + +elif PY >= (3, 12): + ASYNC_HEAD_ASSEMBLY = Assembly() + ASYNC_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + COROUTINE_ASSEMBLY.parse( + r""" + get_awaitable 0 + load_const None + + presend: + send @send + yield_value 2 + resume 3 + jump_backward_no_interrupt @presend + send: + end_send + """ + ) + + ASYNC_GEN_ASSEMBLY.parse( + r""" + try @stopiter + copy 1 + store_fast $__ddgen + load_attr (False, 'asend') + store_fast $__ddgensend + load_fast $__ddgen + load_attr (True, '__anext__') + call 0 + + loop: + get_awaitable 0 + load_const None + presend0: + send @send0 + tried + + try @genexit lasti + yield_value 3 + resume 3 + jump_backward_no_interrupt @presend0 + send0: + end_send + + yield: + call_intrinsic_1 asm.Intrinsic1Op.INTRINSIC_ASYNC_GEN_WRAP + yield_value 3 + resume 1 + push_null + swap 2 + load_fast $__ddgensend + swap 2 + call 1 + jump_backward @loop + tried + + genexit: + try @stopiter + push_exc_info + load_const GeneratorExit + check_exc_match + pop_jump_if_false @exc + pop_top + load_fast $__ddgen + load_attr (True, 'aclose') + call 0 + get_awaitable 0 + load_const None + + presend1: + send @send1 + yield_value 4 + resume 3 + jump_backward_no_interrupt @presend1 + send1: + end_send + pop_top + pop_except + return_const None + + exc: + pop_top + push_null + load_fast $__ddgen + load_attr (False, 'athrow') + push_null + load_const sys.exc_info + call 0 + call_function_ex 0 + get_awaitable 0 + load_const None + + presend2: + send @send2 + yield_value 4 + resume 3 + jump_backward_no_interrupt @presend2 + send2: + end_send + swap 2 + pop_except + jump_backward @yield + tried + + stopiter: + push_exc_info + load_const StopAsyncIteration + check_exc_match + pop_jump_if_false @propagate + pop_top + pop_except + return_const None + + propagate: + reraise 0 + """ + ) + + +elif PY >= (3, 11): + ASYNC_HEAD_ASSEMBLY = Assembly() + ASYNC_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + COROUTINE_ASSEMBLY.parse( + r""" + get_awaitable 0 + load_const None + + presend: + send @send + yield_value + resume 3 + jump_backward_no_interrupt @presend + send: + """ + ) + + ASYNC_GEN_ASSEMBLY.parse( + r""" + try @stopiter + copy 1 + store_fast $__ddgen + load_attr $asend + store_fast $__ddgensend + load_fast $__ddgen + load_method $__anext__ + precall 0 + call 0 + + loop: + get_awaitable 0 + load_const None + presend0: + send @send0 + tried + + try @genexit lasti + yield_value + resume 3 + jump_backward_no_interrupt @presend0 + send0: + + yield: + async_gen_wrap + yield_value + resume 1 + push_null + swap 2 + load_fast $__ddgensend + swap 2 + precall 1 + call 1 + jump_backward @loop + tried + + genexit: + try @stopiter + push_exc_info + load_const GeneratorExit + check_exc_match + pop_jump_forward_if_false @exc + pop_top + load_fast $__ddgen + load_method $aclose + precall 0 + call 0 + get_awaitable 0 + load_const None + + presend1: + send @send1 + yield_value + resume 3 + jump_backward_no_interrupt @presend1 + send1: + pop_top + pop_except + load_const None + return_value + + exc: + pop_top + push_null + load_fast $__ddgen + load_attr $athrow + push_null + load_const sys.exc_info + precall 0 + call 0 + call_function_ex 0 + get_awaitable 0 + load_const None + + presend2: + send @send2 + yield_value + resume 3 + jump_backward_no_interrupt @presend2 + send2: + swap 2 + pop_except + jump_backward @yield + tried + + stopiter: + push_exc_info + load_const StopAsyncIteration + check_exc_match + pop_jump_forward_if_false @propagate + pop_top + pop_except + load_const None + return_value + + propagate: + reraise 0 + """ + ) + + +elif PY >= (3, 10): + COROUTINE_ASSEMBLY.parse( + r""" + get_awaitable + load_const None + yield_from + """ + ) + + ASYNC_GEN_ASSEMBLY.parse( + r""" + setup_finally @stopiter + dup_top + store_fast $__ddgen + load_attr $asend + store_fast $__ddgensend + load_fast $__ddgen + load_attr $__anext__ + call_function 0 + + loop: + get_awaitable + load_const None + yield_from + + yield: + setup_finally @genexit + yield_value + pop_block + load_fast $__ddgensend + rot_two + call_function 1 + jump_absolute @loop + + genexit: + dup_top + load_const GeneratorExit + jump_if_not_exc_match @exc + pop_top + pop_top + pop_top + pop_top + load_fast $__ddgen + load_attr $aclose + call_function 0 + get_awaitable + load_const None + yield_from + pop_except + return_value + + exc: + pop_top + pop_top + pop_top + pop_top + load_fast $__ddgen + load_attr $athrow + load_const sys.exc_info + call_function 0 + call_function_ex 0 + get_awaitable + load_const None + yield_from + rot_four + pop_except + jump_absolute @yield + + stopiter: + dup_top + load_const StopAsyncIteration + jump_if_not_exc_match @propagate + pop_top + pop_top + pop_top + pop_except + load_const None + return_value + + propagate: + reraise 0 + """ + ) + + +elif PY >= (3, 9): + COROUTINE_ASSEMBLY.parse( + r""" + get_awaitable + load_const None + yield_from + """ + ) + + ASYNC_GEN_ASSEMBLY.parse( + r""" + setup_finally @stopiter + dup_top + store_fast $__ddgen + load_attr $asend + store_fast $__ddgensend + load_fast $__ddgen + load_attr $__anext__ + call_function 0 + + loop: + get_awaitable + load_const None + yield_from + + yield: + setup_finally @genexit + yield_value + pop_block + load_fast $__ddgensend + rot_two + call_function 1 + jump_absolute @loop + + genexit: + dup_top + load_const GeneratorExit + jump_if_not_exc_match @exc + pop_top + pop_top + pop_top + pop_top + load_fast $__ddgen + load_attr $aclose + call_function 0 + get_awaitable + load_const None + yield_from + pop_except + return_value + + exc: + pop_top + pop_top + pop_top + pop_top + load_fast $__ddgen + load_attr $athrow + load_const sys.exc_info + call_function 0 + call_function_ex 0 + get_awaitable + load_const None + yield_from + rot_four + pop_except + jump_absolute @yield + + stopiter: + dup_top + load_const StopAsyncIteration + jump_if_not_exc_match @propagate + pop_top + pop_top + pop_top + pop_except + load_const None + return_value + + propagate: + reraise + """ + ) + +else: + msg = "No async wrapping support for Python %d.%d" % PY[:2] + raise RuntimeError(msg) + + +def wrap_async(instrs: bc.Bytecode, code: CodeType, lineno: int) -> None: + if (bc.CompilerFlags.ASYNC_GENERATOR | bc.CompilerFlags.COROUTINE) & code.co_flags: + if ASYNC_HEAD_ASSEMBLY is not None: + instrs[0:0] = ASYNC_HEAD_ASSEMBLY.bind(lineno=lineno) + + if bc.CompilerFlags.COROUTINE & code.co_flags: + # DEV: This is just + # >>> return await wrapper(wrapped, args, kwargs) + instrs[-1:-1] = COROUTINE_ASSEMBLY.bind(lineno=lineno) + + elif bc.CompilerFlags.ASYNC_GENERATOR & code.co_flags: + instrs[-1:] = ASYNC_GEN_ASSEMBLY.bind(lineno=lineno) diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index 1e7973a2e25..8333b997c36 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -1 +1,876 @@ -from ddtrace.internal.utils.wrapping.context import * # noqa +from abc import ABC +from contextvars import ContextVar +from inspect import iscoroutinefunction +from inspect import isgeneratorfunction +import sys +from types import FrameType +from types import FunctionType +from types import TracebackType +import typing as t +from typing import Protocol # noqa:F401 +import weakref + +import bytecode +from bytecode import Bytecode + +from ddtrace.internal.assembly import Assembly +from ddtrace.internal.logger import get_logger +from ddtrace.internal.threads import Lock +from ddtrace.internal.threads import RLock +from ddtrace.internal.wrapping import WrappedFunction +from ddtrace.internal.wrapping import Wrapper +from ddtrace.internal.wrapping import get_function_code +from ddtrace.internal.wrapping import is_wrapped_with +from ddtrace.internal.wrapping import link_function_to_code +from ddtrace.internal.wrapping import set_function_code +from ddtrace.internal.wrapping import unwrap +from ddtrace.internal.wrapping import wrap + + +class _ContextRecord: + """Holds all wrapping-context metadata for a single function. + + Stores only weak references to context objects so that the WeakKeyDictionary + key (the function) is not kept alive by the registry value chain: + _registry -> _ContextRecord -> context -> context.__wrapped__ -> function + Breaking this path with weak references lets ephemeral functions be garbage + collected as soon as all external strong references drop. + """ + + __slots__ = ("_uwc_ref", "lazy_contexts") + + def __init__(self) -> None: + self._uwc_ref: t.Optional[weakref.ref["_UniversalWrappingContext"]] = None + # WeakSet so that LazyWrappingContext instances (which also hold + # __wrapped__ = f) do not prevent the function from being collected. + self.lazy_contexts: weakref.WeakSet["LazyWrappingContext"] = weakref.WeakSet() + + @property + def uwc(self) -> t.Optional["_UniversalWrappingContext"]: + ref = self._uwc_ref + return ref() if ref is not None else None + + @uwc.setter + def uwc(self, value: t.Optional["_UniversalWrappingContext"]) -> None: + self._uwc_ref = weakref.ref(value) if value is not None else None + + @classmethod + def get_or_create(cls, f: FunctionType) -> "_ContextRecord": + record = _registry.get(f) + if record is None: + with _registry_lock: + record = _registry.get(f) + if record is None: + record = cls() + _registry[f] = record + return record + + +# Per-function registry for wrapping-context machinery. WeakKeyDictionary so +# functions are not kept alive by the registry alone. Storing data here instead +# of as function attributes keeps __dict__ clean, preventing frameworks that +# copy function __dict__ (e.g. functools.wraps, +# self.__dict__.update(f.__dict__)) from capturing non-picklable objects. +_registry: weakref.WeakKeyDictionary[FunctionType, _ContextRecord] = weakref.WeakKeyDictionary() +_registry_lock = RLock() + + +log = get_logger(__name__) + +T = t.TypeVar("T") + +# This module implements utilities for wrapping a function with a context +# manager. The rough idea is to re-write the function's bytecode to look like +# this: +# +# def foo(): +# with wrapping_context: +# # Original function code +# +# Because we also want to capture the return value, our context manager extends +# the Python one by implementing a __return__ method that will be called with +# the return value of the function. Contrary to ordinary context managers, +# though, the __exit__ method is only called if the function raises an +# exception. +# +# Because CPython 3.11 introduced zero-cost exceptions, we cannot nest try +# blocks in the function's bytecode. In this case, we call the context manager +# methods directly at the right places, and set up the appropriate exception +# handling code. For older versions of Python we rely on the with statement to +# perform entry and exit operations. Calls to __return__ are explicit in all +# cases. +# +# Some advantages of wrapping a function this way are: +# - Access to the local variables on entry and on return/exit via the frame +# object. +# - No intermediate function calls that pollute the call stack. +# - No need to call the wrapped function manually. +# +# The actual bytecode wrapping is performed once on a target function via a +# universal wrapping context. Multiple context wrapping of a function is allowed +# and it is virtually implemented on top of the concrete universal wrapping +# context. This makes multiple wrapping/unwrapping easy, as it translates to a +# single bytecode wrapping/unwrapping operation. +# +# Context wrappers should be implemented as subclasses of the WrappingContext +# class. The __priority__ attribute can be used to control the order in which +# multiple context wrappers are entered and exited. The __enter__ and __exit__ +# methods should be implemented to perform the necessary operations. The +# __exit__ method is called if the wrapped function raises an exception. The +# frame of the wrapped function can be accessed via the __frame__ property. The +# __return__ method can be implemented to capture the return value of the +# wrapped function. If implemented, its return value will be used as the wrapped +# function return value. The wrapped function can be accessed via the +# __wrapped__ attribute. Context-specific values can be stored and retrieved +# with the set and get methods. + +CONTEXT_HEAD = Assembly() +CONTEXT_RETURN = Assembly() +CONTEXT_FOOT = Assembly() + +if sys.version_info >= (3, 15): + raise NotImplementedError("Python >= 3.15 is not supported yet") +elif sys.version_info >= (3, 13): + CONTEXT_HEAD.parse( + r""" + load_const {context_enter} + push_null + call 0 + pop_top + """ + ) + CONTEXT_RETURN.parse( + r""" + push_null + load_const {context_return} + swap 3 + call 1 + """ + ) + + CONTEXT_RETURN_CONST = Assembly() + CONTEXT_RETURN_CONST.parse( + r""" + load_const {context_return} + push_null + load_const {value} + call 1 + """ + ) + + CONTEXT_FOOT.parse( + r""" + try @_except lasti + push_exc_info + load_const {context_exit} + push_null + call 0 + pop_top + reraise 2 + tried + + _except: + copy 3 + pop_except + reraise 1 + """ + ) + +elif sys.version_info >= (3, 12): + CONTEXT_HEAD.parse( + r""" + push_null + load_const {context_enter} + call 0 + pop_top + """ + ) + + CONTEXT_RETURN.parse( + r""" + load_const {context_return} + push_null + swap 3 + call 1 + """ + ) + + CONTEXT_RETURN_CONST = Assembly() + CONTEXT_RETURN_CONST.parse( + r""" + push_null + load_const {context_return} + load_const {value} + call 1 + """ + ) + + CONTEXT_FOOT.parse( + r""" + try @_except lasti + push_exc_info + push_null + load_const {context_exit} + call 0 + pop_top + reraise 2 + tried + + _except: + copy 3 + pop_except + reraise 1 + """ + ) + + +elif sys.version_info >= (3, 11): + CONTEXT_HEAD.parse( + r""" + push_null + load_const {context_enter} + precall 0 + call 0 + pop_top + """ + ) + + CONTEXT_RETURN.parse( + r""" + load_const {context_return} + push_null + swap 3 + precall 1 + call 1 + """ + ) + + CONTEXT_EXC_HEAD = Assembly() + CONTEXT_EXC_HEAD.parse( + r""" + push_null + load_const {context_exit} + precall 0 + call 0 + pop_top + """ + ) + + CONTEXT_FOOT.parse( + r""" + try @_except lasti + push_exc_info + push_null + load_const {context_exit} + precall 0 + call 0 + pop_top + reraise 2 + tried + + _except: + copy 3 + pop_except + reraise 1 + """ + ) + +elif sys.version_info >= (3, 10): + CONTEXT_HEAD.parse( + r""" + load_const {context} + setup_with @_except + pop_top + _except: + """ + ) + + CONTEXT_RETURN.parse( + r""" + pop_block + load_const {context} + load_method $__return__ + rot_three + rot_three + call_method 1 + rot_two + pop_top + """ + ) + + CONTEXT_FOOT.parse( + r""" + with_except_start + pop_top + reraise 1 + """ + ) + +elif sys.version_info >= (3, 9): + CONTEXT_HEAD.parse( + r""" + load_const {context} + setup_with @_except + pop_top + _except: + """ + ) + + CONTEXT_RETURN.parse( + r""" + pop_block + load_const {context} + load_method $__return__ + rot_three + rot_three + call_method 1 + rot_two + pop_top + """ + ) + + CONTEXT_FOOT.parse( + r""" + with_except_start + pop_top + reraise + """ + ) + + +# This is abstract and should not be used directly +class BaseWrappingContext(ABC): + __priority__: int = 0 + + def __init__(self, f: FunctionType): + # Store a weak reference so that context objects do not keep the wrapped + # function alive. CodeType is not GC-tracked in CPython, so the cycle + # f → code.co_consts → bound_methods(uwc) → uwc.__wrapped__ → f + # cannot be broken by the cyclic GC. A weak ref here allows f's + # reference count to reach zero (and be freed) as soon as all external + # strong refs drop, without relying on the cyclic GC at all. + self._wrapped_ref: weakref.ref[FunctionType] = weakref.ref(f) + self._storage: ContextVar[t.Optional[dict[str, t.Any]]] = ContextVar( + f"{type(self).__name__}__storage", default=None + ) + + @property + def __wrapped__(self) -> FunctionType: + f = self._wrapped_ref() + if f is None: + raise RuntimeError(f"{type(self).__name__}.__wrapped__: the wrapped function has been garbage collected") + return f + + @__wrapped__.setter + def __wrapped__(self, f: FunctionType) -> None: + self._wrapped_ref = weakref.ref(f) + + def __enter__(self) -> "BaseWrappingContext": + prev = self._storage.get() + self._storage.set({"__dd_wrapping_context_prev__": prev}) + + return self + + def _pop_storage(self) -> dict[str, t.Any]: + storage = t.cast(dict[str, t.Any], self._storage.get()) + self._storage.set(storage.pop("__dd_wrapping_context_prev__")) + return storage + + def __return__(self, value: T) -> T: + self._pop_storage() + return value + + def __exit__( + self, + exc_type: t.Optional[type[BaseException]], + exc_val: t.Optional[BaseException], + exc_tb: t.Optional[TracebackType], + ) -> None: + self._pop_storage() + + def get(self, key: str) -> t.Any: + return t.cast(dict[str, t.Any], self._storage.get())[key] + + def set(self, key: str, value: T) -> T: + t.cast(dict[str, t.Any], self._storage.get())[key] = value + return value + + @classmethod + def wrapped(cls, f: FunctionType) -> "BaseWrappingContext": + try: + context = cls.extract(f) + assert isinstance(context, cls) # nosec + except ValueError: + context = cls(f) + context.wrap() + return context + + @classmethod + def is_wrapped(cls, _f: FunctionType) -> bool: + raise NotImplementedError + + @classmethod + def extract(cls, _f: FunctionType) -> "BaseWrappingContext": + raise NotImplementedError + + def wrap(self) -> None: + raise NotImplementedError + + def unwrap(self) -> None: + raise NotImplementedError + + +# This is the public interface exported by this module +class WrappingContext(BaseWrappingContext): + @property + def __frame__(self) -> FrameType: + try: + return t.cast( + FrameType, + _UniversalWrappingContext.extract(t.cast(FunctionType, self.__wrapped__)).get("__frame__"), + ) + except ValueError: + raise AttributeError("Wrapping context not entered") + + def get_local(self, name: str) -> t.Any: + return self.__frame__.f_locals[name] + + @classmethod + def is_wrapped(cls, f: FunctionType) -> bool: + try: + return bool(cls.extract(f)) + except ValueError: + return False + + @classmethod + def extract(cls, f: FunctionType) -> "WrappingContext": + try: + return _UniversalWrappingContext.extract(f).registered(cls) + except (ValueError, KeyError): + msg = f"Function is not wrapped with {cls}" + raise ValueError(msg) + + def wrap(self) -> None: + t.cast( + _UniversalWrappingContext, _UniversalWrappingContext.wrapped(t.cast(FunctionType, self.__wrapped__)) + ).register(self) + + def unwrap(self) -> None: + f = t.cast(FunctionType, self.__wrapped__) + + try: + _UniversalWrappingContext.extract(f).unregister(self) + except ValueError: + pass + + +class LazyWrappingContext(WrappingContext): + def __init__(self, f: FunctionType): + super().__init__(f) + + self._trampoline: t.Optional[Wrapper] = None + self._trampoline_lock = Lock() + + @classmethod + def is_wrapped(cls, f: FunctionType) -> bool: + with _registry_lock: + record = _registry.get(f) + if record is None: + return False + return any(isinstance(c, cls) for c in record.lazy_contexts) + + def wrap(self) -> None: + """Perform the bytecode wrapping on first invocation.""" + with (tl := self._trampoline_lock): + if self._trampoline is not None: + return + + # If the function is already universally wrapped it's less expensive + # to do the normal wrapping. + if _UniversalWrappingContext.is_wrapped(t.cast(FunctionType, self.__wrapped__)): + super().wrap() + return + + def trampoline(_: t.Any, args: tuple[t.Any, ...], kwargs: dict[str, t.Any]) -> t.Any: + with tl: + f = t.cast(WrappedFunction, self.__wrapped__) + if is_wrapped_with(t.cast(FunctionType, self.__wrapped__), trampoline): + f = t.cast(WrappedFunction, unwrap(f, trampoline)) + + self._trampoline = None + + inconsistent = False + with _registry_lock: + record = _registry.get(t.cast(FunctionType, f)) + if record is not None: + inconsistent = self not in record.lazy_contexts + record.lazy_contexts.discard(self) + if not record.lazy_contexts and record.uwc is None: + _registry.pop(t.cast(FunctionType, f), None) + if inconsistent: + log.warning("Inconsistent lazy wrapping context state") + + super(LazyWrappingContext, self).wrap() + return f(*args, **kwargs) + + wrap(t.cast(FunctionType, self.__wrapped__), trampoline) + + self._trampoline = trampoline + + _ContextRecord.get_or_create(t.cast(FunctionType, self.__wrapped__)).lazy_contexts.add(self) + + def unwrap(self) -> None: + with self._trampoline_lock: + if _UniversalWrappingContext.is_wrapped(t.cast(FunctionType, self.__wrapped__)): + assert self._trampoline is None # nosec + super().unwrap() + elif self._trampoline is not None: + with _registry_lock: + record = _registry.get(t.cast(FunctionType, self.__wrapped__)) + if record is not None: + record.lazy_contexts.discard(self) + if not record.lazy_contexts and record.uwc is None: + _registry.pop(t.cast(FunctionType, self.__wrapped__), None) + + unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline) + self._trampoline = None + + +class ContextWrappedFunction(Protocol): + """A wrapped function.""" + + def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: + pass + + +# This class provides an interface between single bytecode wrapping and multiple +# logical context wrapping +class _UniversalWrappingContext(BaseWrappingContext): + def __init__(self, f: FunctionType) -> None: + super().__init__(f) + + self._contexts: list[WrappingContext] = [] + + def register(self, context: WrappingContext) -> None: + _type = type(context) + if any(isinstance(c, _type) for c in self._contexts): + raise ValueError("Context already registered") + + self._contexts.append(context) + self._contexts.sort(key=lambda c: c.__priority__) + + def unregister(self, context: WrappingContext) -> None: + try: + self._contexts.remove(context) + except ValueError: + raise ValueError("Context not registered") + + if not self._contexts: + self.unwrap() + + def is_registered(self, context: WrappingContext) -> bool: + return any(isinstance(c, type(context)) for c in self._contexts) + + def registered(self, context_type: type[WrappingContext]) -> WrappingContext: + for context in self._contexts: + if isinstance(context, context_type): + return context + raise KeyError(f"Context {context_type} not registered") + + def __enter__(self) -> "_UniversalWrappingContext": + super().__enter__() + + # Make the frame object available to the contexts + self.set("__frame__", sys._getframe(1)) + + for context in self._contexts: + context.__enter__() + + return self + + def _exit(self) -> None: + self.__exit__(*sys.exc_info()) + + def __exit__( + self, + exc_type: t.Optional[type[BaseException]], + exc_value: t.Optional[BaseException], + traceback: t.Optional[TracebackType], + ) -> None: + if exc_value is None: + return + + for context in self._contexts[::-1]: + context.__exit__(exc_type, exc_value, traceback) + + super().__exit__(exc_type, exc_value, traceback) + + def __return__(self, value: T) -> T: + for context in self._contexts[::-1]: + context.__return__(value) + + return super().__return__(value) + + @classmethod + def is_wrapped(cls, f: FunctionType) -> bool: + try: + with _registry_lock: + record = _registry.get(f) + if record is None or record.uwc is None: + return False + # Verify the registry entry matches actual bytecode wrapping. + if sys.version_info >= (3, 11): + return record.uwc.__enter__ in get_function_code(f).co_consts + else: + return record.uwc in get_function_code(f).co_consts + except AttributeError: + return False + + @classmethod + def extract(cls, f: FunctionType) -> "_UniversalWrappingContext": + with _registry_lock: + if not cls.is_wrapped(f): + raise ValueError("Function is not wrapped") + return t.cast(_UniversalWrappingContext, _registry[f].uwc) + + if sys.version_info >= (3, 11): + + def wrap(self) -> None: + f = self.__wrapped__ + + with _registry_lock: + if self.is_wrapped(f): + raise ValueError("Function already wrapped") + + bc = Bytecode.from_code(code := get_function_code(f)) + + # Prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + try: + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN.bind({"context_return": self.__return__}, lineno=instr.lineno) + elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ + return_code = CONTEXT_RETURN_CONST.bind( + {"context_return": self.__return__, "value": instr.arg}, lineno=instr.lineno + ) + else: + return_code = [] + + bc[i:i] = return_code + i += len(return_code) + except AttributeError: + # Not an instruction + pass + i += 1 + + # Search for the RESUME instruction + for i, instr in enumerate(bc, 1): + try: + if instr.name == "RESUME": + break + except AttributeError: + # Not an instruction + pass + else: + i = 0 + + bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=code.co_firstlineno) + + # Wrap every line outside a try block + except_label = bytecode.Label() + first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) + + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None: + bc.insert(i, bytecode.TryEnd(last_try_begin)) + last_try_begin = None + i += 1 + elif isinstance(instr, bytecode.TryEnd): + j = i + 1 + while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin): + if isinstance(bc[j], bytecode.Instr): + last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) + bc.insert(i + 1, last_try_begin) + break + j += 1 + i += 1 + i += 1 + + bc.insert(0, first_try_begin) + + bc.append(bytecode.TryEnd(last_try_begin)) + bc.append(except_label) + bc.extend(CONTEXT_FOOT.bind({"context_exit": self._exit}, lineno=code.co_firstlineno)) + + # Register the wrapping context and write the new bytecode. + _ContextRecord.get_or_create(f).uwc = self + link_function_to_code(code, f) + set_function_code(f, bc.to_code()) + + def unwrap(self) -> None: + f = self.__wrapped__ + + with _registry_lock: + if not self.is_wrapped(f): + return + + wc = _registry[f].uwc + + bc = Bytecode.from_code(get_function_code(f)) + + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() + bc.pop() + + except_label = bc.pop(0).target + + # Remove the try blocks + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.TryBegin) and instr.target is except_label: + bc.pop(i) + elif isinstance(instr, bytecode.TryEnd) and instr.entry.target is except_label: + bc.pop(i) + else: + i += 1 + + # Remove the head of the try block + for i, instr in enumerate(bc): + if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: + break + + # Search for the RESUME instruction + for i, instr in enumerate(bc, 1): + try: + if instr.name == "RESUME": + break + except AttributeError: + # Not an instruction + pass + else: + i = 0 + + bc[i : i + len(CONTEXT_HEAD)] = [] + + # Un-prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + try: + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN + elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ + return_code = CONTEXT_RETURN_CONST + else: + return_code = None + + if return_code is not None: + bc[i - len(return_code) : i] = [] + i -= len(return_code) + except AttributeError: + # Not an instruction + pass + i += 1 + + # Recreate the code object + set_function_code(f, bc.to_code()) + + # Clear the UWC from the registry; remove the record if fully empty. + record = _registry.get(f) + if record is not None: + record.uwc = None + if not record.lazy_contexts: + _registry.pop(f, None) + + else: + + def wrap(self) -> None: + f = t.cast(FunctionType, self.__wrapped__) + + with _registry_lock: + if self.is_wrapped(f): + raise ValueError("Function already wrapped") + + bc = Bytecode.from_code(code := get_function_code(f)) + + # Prefix every return + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.Instr): + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN.bind({"context": self}, lineno=instr.lineno) + bc[i:i] = return_code + i += len(return_code) + i += 1 + + # Search for the GEN_START instruction, which needs to stay on top. + i = 0 + if sys.version_info >= (3, 10) and (iscoroutinefunction(f) or isgeneratorfunction(f)): + for i, instr in enumerate(bc, 1): + if isinstance(instr, bytecode.Instr) and instr.name == "GEN_START": + break + + *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": self}, lineno=code.co_firstlineno) + + bc.append(except_label) + bc.extend(CONTEXT_FOOT.bind(lineno=code.co_firstlineno)) + + # Register the wrapping context and write the new bytecode. + _ContextRecord.get_or_create(f).uwc = self + link_function_to_code(code, f) + set_function_code(f, bc.to_code()) + + def unwrap(self) -> None: + f = t.cast(FunctionType, self.__wrapped__) + + with _registry_lock: + if not self.is_wrapped(f): + return + + wc = _registry[f].uwc + + bc = Bytecode.from_code(get_function_code(f)) + + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() + + # Remove the head of the try block + for i, instr in enumerate(bc): + if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: + break + + bc[i : i + len(CONTEXT_HEAD) - 1] = [] + + # Remove all the return handlers + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.Instr) and instr.name == "RETURN_VALUE": + bc[i - len(CONTEXT_RETURN) : i] = [] + i -= len(CONTEXT_RETURN) + i += 1 + + # Recreate the code object + set_function_code(f, bc.to_code()) + + # Clear the UWC from the registry; remove the record if fully empty. + record = _registry.get(f) + if record is not None: + record.uwc = None + if not record.lazy_contexts: + _registry.pop(f, None) + + +def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingContext]": + """Return the _UniversalWrappingContext for *f*, or None if not context-wrapped.""" + with _registry_lock: + record = _registry.get(f) + return record.uwc if record is not None else None diff --git a/ddtrace/internal/wrapping/generators.py b/ddtrace/internal/wrapping/generators.py index ceb1cf2fd7a..01c284c4686 100644 --- a/ddtrace/internal/wrapping/generators.py +++ b/ddtrace/internal/wrapping/generators.py @@ -1 +1,491 @@ -from ddtrace.internal.utils.wrapping.generators import * # noqa +import sys +from types import CodeType + +import bytecode as bc + +from ddtrace.internal.assembly import Assembly + + +PY = sys.version_info[:2] + + +# ----------------------------------------------------------------------------- +# Generator Wrapping +# ----------------------------------------------------------------------------- +# DEV: This is roughly equivalent to +# +# __ddgen = wrapper(wrapped, args, kwargs) +# __ddgensend = __ddgen.send +# try: +# value = next(__ddgen) +# while True: +# try: +# tosend = yield value +# except GeneratorExit: +# return __ddgen.close() +# except Exception: +# value = __ddgen.throw(*sys.exc_info()) +# else: +# value = __ddgensend(tosend) +# except StopIteration: +# return +# ----------------------------------------------------------------------------- +GENERATOR_ASSEMBLY = Assembly() +GENERATOR_HEAD_ASSEMBLY = None + +if PY >= (3, 15): + raise NotImplementedError("This version of CPython is not supported yet") + +elif PY >= (3, 14): + GENERATOR_HEAD_ASSEMBLY = Assembly() + GENERATOR_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + GENERATOR_ASSEMBLY.parse( + r""" + try @stopiter + copy 1 + store_fast $__ddgen + load_attr $send + store_fast $__ddgensend + load_const next + push_null + load_fast_borrow $__ddgen + + loop: + call 1 + tried + + yield: + try @genexit lasti + yield_value 0 + resume 1 + push_null + load_fast_borrow $__ddgensend + swap 3 + jump_backward @loop + tried + + genexit: + try @stopiter + push_exc_info + load_const GeneratorExit + check_exc_match + pop_jump_if_false @exc + pop_top + load_fast $__ddgen + load_method $close + call 0 + swap 2 + pop_except + return_value + + exc: + pop_top + load_fast $__ddgen + load_attr $throw + push_null + load_const sys.exc_info + push_null + call 0 + push_null + call_function_ex + swap 2 + pop_except + jump_backward @yield + tried + + stopiter: + push_exc_info + load_const StopIteration + check_exc_match + pop_jump_if_false @propagate + pop_top + pop_except + load_const None + return_value + + propagate: + reraise 0 + """ + ) + +elif PY >= (3, 13): + GENERATOR_HEAD_ASSEMBLY = Assembly() + GENERATOR_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + GENERATOR_ASSEMBLY.parse( + r""" + try @stopiter + copy 1 + store_fast $__ddgen + load_attr $send + store_fast $__ddgensend + load_const next + push_null + load_fast $__ddgen + + loop: + call 1 + tried + + yield: + try @genexit lasti + yield_value 0 + resume 1 + push_null + load_fast $__ddgensend + swap 3 + jump_backward @loop + tried + + genexit: + try @stopiter + push_exc_info + load_const GeneratorExit + check_exc_match + pop_jump_if_false @exc + pop_top + load_fast $__ddgen + load_method $close + call 0 + swap 2 + pop_except + return_value + + exc: + pop_top + load_fast $__ddgen + load_attr $throw + push_null + load_const sys.exc_info + push_null + call 0 + call_function_ex 0 + swap 2 + pop_except + jump_backward @yield + tried + + stopiter: + push_exc_info + load_const StopIteration + check_exc_match + pop_jump_if_false @propagate + pop_top + pop_except + load_const None + return_value + + propagate: + reraise 0 + """ + ) + +elif PY >= (3, 12): + GENERATOR_HEAD_ASSEMBLY = Assembly() + GENERATOR_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + GENERATOR_ASSEMBLY.parse( + r""" + try @stopiter + copy 1 + store_fast $__ddgen + load_attr $send + store_fast $__ddgensend + push_null + load_const next + load_fast $__ddgen + + loop: + call 1 + tried + + yield: + try @genexit lasti + yield_value 3 + resume 1 + push_null + swap 2 + load_fast $__ddgensend + swap 2 + jump_backward @loop + tried + + genexit: + try @stopiter + push_exc_info + load_const GeneratorExit + check_exc_match + pop_jump_if_false @exc + pop_top + load_fast $__ddgen + load_method $close + call 0 + swap 2 + pop_except + return_value + + exc: + pop_top + push_null + load_fast $__ddgen + load_attr $throw + push_null + load_const sys.exc_info + call 0 + call_function_ex 0 + swap 2 + pop_except + jump_backward @yield + tried + + stopiter: + push_exc_info + load_const StopIteration + check_exc_match + pop_jump_if_false @propagate + pop_top + pop_except + return_const None + + propagate: + reraise 0 + """ + ) + +elif PY >= (3, 11): + GENERATOR_HEAD_ASSEMBLY = Assembly() + GENERATOR_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + GENERATOR_ASSEMBLY.parse( + r""" + try @stopiter + copy 1 + store_fast $__ddgen + load_attr $send + store_fast $__ddgensend + push_null + load_const next + load_fast $__ddgen + + loop: + precall 1 + call 1 + tried + + yield: + try @genexit lasti + yield_value + resume 1 + push_null + swap 2 + load_fast $__ddgensend + swap 2 + jump_backward @loop + tried + + genexit: + try @stopiter + push_exc_info + load_const GeneratorExit + check_exc_match + pop_jump_forward_if_false @exc + pop_top + load_fast $__ddgen + load_method $close + precall 0 + call 0 + swap 2 + pop_except + return_value + + exc: + pop_top + push_null + load_fast $__ddgen + load_attr $throw + push_null + load_const sys.exc_info + precall 0 + call 0 + call_function_ex 0 + swap 2 + pop_except + jump_backward @yield + tried + + stopiter: + push_exc_info + load_const StopIteration + check_exc_match + pop_jump_forward_if_false @propagate + pop_top + pop_except + load_const None + return_value + + propagate: + reraise 0 + """ + ) + +elif PY >= (3, 10): + GENERATOR_ASSEMBLY.parse( + r""" + setup_finally @stopiter + dup_top + store_fast $__ddgen + load_attr $send + store_fast $__ddgensend + load_const next + load_fast $__ddgen + + loop: + call_function 1 + + yield: + setup_finally @genexit + yield_value + pop_block + load_fast $__ddgensend + rot_two + jump_absolute @loop + + genexit: + dup_top + load_const GeneratorExit + jump_if_not_exc_match @exc + pop_top + pop_top + pop_top + pop_top + load_fast $__ddgen + load_attr $close + call_function 0 + return_value + + exc: + pop_top + pop_top + pop_top + pop_top + load_fast $__ddgen + load_attr $throw + load_const sys.exc_info + call_function 0 + call_function_ex 0 + rot_four + pop_except + jump_absolute @yield + + stopiter: + dup_top + load_const StopIteration + jump_if_not_exc_match @propagate + pop_top + pop_top + pop_top + pop_except + load_const None + return_value + + propagate: + reraise 0 + """ + ) + +elif PY >= (3, 9): + GENERATOR_ASSEMBLY.parse( + r""" + setup_finally @stopiter + dup_top + store_fast $__ddgen + load_attr $send + store_fast $__ddgensend + load_const next + load_fast $__ddgen + + loop: + call_function 1 + + yield: + setup_finally @genexit + yield_value + pop_block + load_fast $__ddgensend + rot_two + jump_absolute @loop + + genexit: + dup_top + load_const GeneratorExit + jump_if_not_exc_match @exc + pop_top + pop_top + pop_top + pop_top + load_fast $__ddgen + load_attr $close + call_function 0 + return_value + + exc: + pop_top + pop_top + pop_top + pop_top + load_fast $__ddgen + load_attr $throw + load_const sys.exc_info + call_function 0 + call_function_ex 0 + rot_four + pop_except + jump_absolute @yield + + stopiter: + dup_top + load_const StopIteration + jump_if_not_exc_match @propagate + pop_top + pop_top + pop_top + pop_except + load_const None + return_value + + propagate: + reraise + """ + ) + +else: + msg = "No generator wrapping support for Python %d.%d" % PY[:2] + raise RuntimeError(msg) + + +def wrap_generator(instrs: bc.Bytecode, code: CodeType, lineno: int) -> None: + if GENERATOR_HEAD_ASSEMBLY is not None: + instrs[0:0] = GENERATOR_HEAD_ASSEMBLY.bind(lineno=lineno) + + instrs[-1:] = GENERATOR_ASSEMBLY.bind(lineno=lineno) diff --git a/riotfile.py b/riotfile.py index cc8a8e0d7d6..7686782c81b 100644 --- a/riotfile.py +++ b/riotfile.py @@ -595,7 +595,7 @@ def select_pys(min_version: str = MIN_PYTHON_VERSION, max_version: str = MAX_PYT "DD_CIVISIBILITY_ITR_ENABLED": "0", "DD_PYTEST_USE_NEW_PLUGIN": "false", }, - command="pytest -v --no-cov -n auto {cmdargs} tests/internal/", + command="pytest -v -n auto {cmdargs} tests/internal/", pkgs={ "httpretty": latest, "gevent": latest, From 5a902d044bbaae6f5b6beb72eec2b5afc4cba29d Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 10:09:53 -0700 Subject: [PATCH 24/29] undo threads change --- ddtrace/internal/threads.py | 221 +++++++++++++++++++++++++++++++++++- 1 file changed, 220 insertions(+), 1 deletion(-) diff --git a/ddtrace/internal/threads.py b/ddtrace/internal/threads.py index a6666307f10..84da0cabfa9 100644 --- a/ddtrace/internal/threads.py +++ b/ddtrace/internal/threads.py @@ -1 +1,220 @@ -from ddtrace.internal.utils.threads import * # noqa +from time import monotonic_ns +import typing as t + +from ddtrace.internal import _threads as _threads_mod +from ddtrace.internal import forksafe +from ddtrace.internal._threads import PeriodicThread as _PeriodicThread +from ddtrace.internal._threads import periodic_threads +from ddtrace.internal.logger import get_logger + + +log = get_logger(__name__) + +# We try to import the stdlib locks from the _thread module, where they are +# implemented in C for CPython for most platforms. If that fails, we fall back +# to the threading module, which provides a pure Python implementation that +# should work on all platforms. We also make sure to grab a reference to the +# original lock classes, in case they get patched by monkey-patching libraries +# like gevent. +try: + from _thread import allocate_lock as Lock +except ImportError: + from threading import Lock + +try: + from _thread import RLock +except ImportError: + from threading import RLock + + +__all__ = [ + "Lock", + "PeriodicThread", + "RLock", +] + + +# Forking state management. This is a barrier to either prevent new threads +# from being started while forking, or to allow a thread to be started +# completely if a fork comes in the middle of it. +_forking = False +_forking_lock = Lock() + + +class BoundMethod(t.Protocol): + __self__: t.Any + + def __call__(self) -> None: ... + + +# List of threads that have requested to be started while forking. These will +# be started after the fork is complete. +_threads_to_start_after_fork: list[BoundMethod] = [] + + +def _safe_restart(start: t.Callable[[], None], name: t.Optional[str] = None) -> None: + """Invoke a post-fork thread-start callable, logging resource errors instead of raising. + + The native layer translates pthread_create failures (EAGAIN, ENOMEM) into + OSError. Post-fork restart is triggered automatically by forksafe hooks — + there is no explicit caller that can handle the error, so losing a + periodic thread to resource exhaustion must not crash the host. + Explicit start() calls let OSError propagate so the caller can react. + """ + try: + start() + except Exception as e: + log.error("failed to start periodic thread %s: %s", name, e) + + +class PeriodicThread(_PeriodicThread): + """A fork-safe periodic thread.""" + + __autorestart__ = True + + def start(self) -> None: + with _forking_lock: + # We cannot start a new thread while we are forking, because we are + # trying to stop them all. In that case, we take note of the thread + # and start it after the fork. + if not _forking: + super().start() + else: + _threads_to_start_after_fork.append(t.cast(BoundMethod, super().start)) + + +# Set of running periodic threads that need to be restarted after a fork. +_threads_to_restart_after_fork: set[_PeriodicThread] = set() + + +# A typical scenario is that of forking worker threads in a loop. For the +# parent process, this would mean having to stop and restart the threads in +# between forks, which is not ideal. Instead, we can use a timer to restart +# the threads after a certain amount of time has passed since the last fork. +# This way, we can avoid stopping and restarting the threads in between forks. +class ThreadRestartTimer(PeriodicThread): + __timeout__ = int(1e8) # nanoseconds + + _instance: t.Optional["ThreadRestartTimer"] = None + _timestamp = 0 + + def __init__(self): + super().__init__(self.__timeout__ / 1e9, self._restart_threads, name=f"{__name__}:{self.__class__.__name__}") + + def _restart_threads(self) -> None: + # Restart the threads after we have stopped calling fork for a while. + with _forking_lock: + # If we are forking, we will try again later. + if _forking: + return + + # If we haven't have calls to fork for a while, we can restart the + # threads. This way we avoid stopping and restarting the threads + # in between forks. + if monotonic_ns() >= self._timestamp: # 100ms + for thread in _threads_to_restart_after_fork.copy(): + if isinstance(thread, ThreadRestartTimer): + # Skip any ThreadRestartTimer instance, + # to avoid restarting orphaned timer instances that were + # caught in periodic_threads during a fork. + continue + log.debug("Restarting thread %s after fork", thread.name) + try: + thread._after_fork(force=True) + except Exception as e: + log.error("failed to restart periodic thread %s after fork: %s", thread.name, e) + _threads_to_restart_after_fork.clear() + + for thread_start in _threads_to_start_after_fork: + log.debug("Starting thread %s after fork", thread_start.__self__.name) + _safe_restart(thread_start, thread_start.__self__.name) + _threads_to_start_after_fork.clear() + + # We no longer need this thread so we clear it. + self.clear() + + @classmethod + def clear(cls): + """Clear the timer and stop it if it is running.""" + if cls._instance is not None: + cls._instance.stop() + cls._instance = None + + @classmethod + def touch(cls): + """Set the new expiration time for the timer.""" + cls._timestamp = monotonic_ns() + cls.__timeout__ + + @classmethod + def set(cls): + """Set the timer to restart the threads after a fork.""" + if cls._instance is None: + cls._instance = cls() + cls._instance.start() + else: + # We have already created the timer, so we let the forksafe logic + # handle the restart instead of creating a new instance. + cls._instance._after_fork() + + +@forksafe.register +def _after_fork_child(): + global _forking + + _forking = False + + # Keep child at-fork work minimal: thread restarts happen asynchronously in + # the child so application code can resume immediately after fork. Parent + # process threads are still restarted in _after_fork_parent() below. + for thread in _threads_to_restart_after_fork.copy(): + log.debug("Restarting thread %s after fork in child", thread.name) + try: + thread._after_fork(force=False) + except Exception as e: + log.error("failed to restart periodic thread %s after fork in child: %s", thread.name, e) + _threads_to_restart_after_fork.clear() + + for thread_start in _threads_to_start_after_fork.copy(): + log.debug("Starting thread %s after fork in child", thread_start.__self__.name) + _safe_restart(thread_start, thread_start.__self__.name) + _threads_to_start_after_fork.clear() + + +@forksafe.register_after_parent +def _after_fork_parent() -> None: + global _forking + + _forking = False + + if _threads_to_restart_after_fork or _threads_to_start_after_fork: + ThreadRestartTimer.set() + + +@forksafe.register_before_fork +def _before_fork() -> None: + global _threads_to_restart_after_fork, _forking_lock, _forking + + ThreadRestartTimer.touch() + + with _forking_lock: + _forking = True + + # Snapshot pending restarts first so a worker moving pending -> active + # concurrently cannot be missed between the two snapshots. + pending_threads = getattr(_threads_mod, "_pending_threads", lambda: ())() + _threads_to_restart_after_fork.update(pending_threads) + # Take note of all the periodic threads that are running and will need to be + # restarted. + _threads_to_restart_after_fork.update(periodic_threads.values()) + + # Stop all the periodic threads that are still running, without executing + # the shutdown methods, if any. This ensures that we can stop the threads + # more promptly. + for thread in _threads_to_restart_after_fork: + log.debug("Stopping thread %s before fork", thread.name) + thread._before_fork() + + # Join all the threads to ensure they are stopped before the fork. + for thread in _threads_to_restart_after_fork: + log.debug("Joining thread %s before fork", thread.name) + thread.join() From 6fd698049e02e47f7479eee6ac2aa1ff8361f3d6 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 10:38:59 -0700 Subject: [PATCH 25/29] Revert "undo wrapping change" This reverts commit 4a79e3d985a95f47d81d1443f559821632c729d1. --- ddtrace/internal/wrapping/__init__.py | 455 +----------- ddtrace/internal/wrapping/asyncs.py | 726 +------------------- ddtrace/internal/wrapping/context.py | 877 +----------------------- ddtrace/internal/wrapping/generators.py | 492 +------------ riotfile.py | 2 +- 5 files changed, 5 insertions(+), 2547 deletions(-) diff --git a/ddtrace/internal/wrapping/__init__.py b/ddtrace/internal/wrapping/__init__.py index b01b8c7977d..07fd80a109e 100644 --- a/ddtrace/internal/wrapping/__init__.py +++ b/ddtrace/internal/wrapping/__init__.py @@ -1,454 +1 @@ -import sys -from types import CodeType -from types import FunctionType -from typing import Any -from typing import Callable -from typing import Iterator -from typing import MutableMapping -from typing import Optional -from typing import Protocol -from typing import cast -import weakref - -import bytecode as bc -from bytecode import Instr - -from ddtrace.internal.assembly import Assembly -from ddtrace.internal.threads import Lock -from ddtrace.internal.wrapping.asyncs import wrap_async -from ddtrace.internal.wrapping.generators import wrap_generator - - -PY = sys.version_info[:2] - -# Maps each wrapped function to its inner copy (the singly-linked list of -# wrapping layers). WeakKeyDictionary so functions are not kept alive by the -# registry alone. -_wrapped: weakref.WeakKeyDictionary[FunctionType, FunctionType] = weakref.WeakKeyDictionary() -_wrapped_lock = Lock() - -# Maps original code objects to the functions that own them. Written by -# link_function_to_code; read by functions_for_code in inspection.py. -# WeakValueDictionary so that wrapped ephemeral functions are not kept alive by -# this mapping alone — inspection.py falls back to gc.get_referrers on a miss. -_code_to_fn: MutableMapping[CodeType, FunctionType] = weakref.WeakValueDictionary() - - -def link_function_to_code(code: CodeType, function: FunctionType) -> None: - """Link a function to its original code object for fast reverse lookup.""" - _code_to_fn[code] = function - - -class WrappedFunction(Protocol): - """A wrapped function.""" - - def __call__(self, *args: Any, **kwargs: Any) -> Any: - pass - - -Wrapper = Callable[[FunctionType, tuple[Any], dict[str, Any]], Any] - - -def _add(lineno: int) -> Instr: - if PY >= (3, 11): - return Instr("BINARY_OP", bc.BinaryOp.ADD, lineno=lineno) - - return Instr("INPLACE_ADD", lineno=lineno) - - -HEAD = Assembly() -if PY >= (3, 13): - HEAD.parse( - r""" - resume 0 - load_const {wrapper} - push_null - load_const {wrapped} - """ - ) - -elif PY >= (3, 11): - HEAD.parse( - r""" - resume 0 - push_null - load_const {wrapper} - load_const {wrapped} - """ - ) - -else: - HEAD.parse( - r""" - load_const {wrapper} - load_const {wrapped} - """ - ) - - -_UPDATE_MAP_FAST = Assembly() -_UPDATE_MAP_DEREF = Assembly() -if PY >= (3, 12): - _UPDATE_MAP_FAST.parse( - r""" - copy 1 - load_method $update - load_fast {varkwargsname} - call 1 - pop_top - """ - ) - _UPDATE_MAP_DEREF.parse( - r""" - copy 1 - load_method $update - load_deref {varkwargsname} - call 1 - pop_top - """ - ) -elif PY >= (3, 11): - _UPDATE_MAP_FAST.parse( - r""" - copy 1 - load_method $update - load_fast {varkwargsname} - precall 1 - call 1 - pop_top - """ - ) - _UPDATE_MAP_DEREF.parse( - r""" - copy 1 - load_method $update - load_deref {varkwargsname} - precall 1 - call 1 - pop_top - """ - ) -else: - _UPDATE_MAP_FAST.parse( - r""" - dup_top - load_attr $update - load_fast {varkwargsname} - call_function 1 - pop_top - """ - ) - _UPDATE_MAP_DEREF.parse( - r""" - dup_top - load_attr $update - load_deref {varkwargsname} - call_function 1 - pop_top - """ - ) - - -def _generate_update_map(name: str, code: CodeType, lineno: int) -> Iterator[Any]: - """Yield opcodes to call ``dict.update(name)`` where ``name`` may be a cell var.""" - if PY >= (3, 11) and name in code.co_cellvars: - yield from _UPDATE_MAP_DEREF.bind({"varkwargsname": bc.CellVar(name)}, lineno=lineno) # type: ignore[attr-defined] - else: - yield from _UPDATE_MAP_FAST.bind({"varkwargsname": name}, lineno=lineno) - - -CALL_RETURN = Assembly() -if PY >= (3, 12): - CALL_RETURN.parse( - r""" - call {arg} - return_value - """ - ) - -elif PY >= (3, 11): - CALL_RETURN.parse( - r""" - precall {arg} - call {arg} - return_value - """ - ) - -else: - CALL_RETURN.parse( - r""" - call_function {arg} - return_value - """ - ) - - -FIRSTLINENO_OFFSET = int(PY >= (3, 11)) - - -def _load_var(name: str, code: CodeType, lineno: int) -> Instr: - """Return the correct load instruction for a function parameter. - - On Python 3.11+, parameters captured by inner closures become cell - variables and must be loaded with LOAD_DEREF instead of LOAD_FAST. - """ - if PY >= (3, 11) and name in code.co_cellvars: - return Instr("LOAD_DEREF", bc.CellVar(name), lineno=lineno) # type: ignore[attr-defined] - return Instr("LOAD_FAST", name, lineno=lineno) - - -def generate_posargs(code: CodeType) -> Iterator[Any]: - """Generate the opcodes for building the positional arguments tuple.""" - varnames = code.co_varnames - lineno = code.co_firstlineno + FIRSTLINENO_OFFSET - varargs = bool(code.co_flags & bc.CompilerFlags.VARARGS) - nargs = code.co_argcount - varargsname: Optional[str] = varnames[nargs + code.co_kwonlyargcount] if varargs else None - - if nargs: # posargs [+ varargs] - yield from (_load_var(argname, code, lineno) for argname in varnames[:nargs]) - - yield Instr("BUILD_TUPLE", nargs, lineno=lineno) - if varargsname is not None: - yield _load_var(varargsname, code, lineno) - yield _add(lineno) - - elif varargsname is not None: # varargs - yield _load_var(varargsname, code, lineno) - - else: # () - yield Instr("BUILD_TUPLE", 0, lineno=lineno) - - -def generate_kwargs(code: CodeType) -> Iterator[Any]: - """Generate the opcodes for building the keyword arguments dictionary.""" - flags = code.co_flags - varnames = code.co_varnames - lineno = code.co_firstlineno + FIRSTLINENO_OFFSET - varargs = bool(flags & bc.CompilerFlags.VARARGS) - varkwargs = bool(flags & bc.CompilerFlags.VARKEYWORDS) - nargs = code.co_argcount - kwonlyargs = code.co_kwonlyargcount - varkwargsname: Optional[str] = varnames[nargs + kwonlyargs + varargs] if varkwargs else None - - if kwonlyargs: - for arg in varnames[nargs : nargs + kwonlyargs]: # kwargs [+ varkwargs] - yield Instr("LOAD_CONST", arg, lineno=lineno) - yield _load_var(arg, code, lineno) - yield Instr("BUILD_MAP", kwonlyargs, lineno=lineno) - if varkwargsname is not None: - yield from _generate_update_map(varkwargsname, code, lineno) - - elif varkwargsname is not None: # varkwargs - yield _load_var(varkwargsname, code, lineno) - - else: # {} - yield Instr("BUILD_MAP", 0, lineno=lineno) - - -def wrap_bytecode(wrapper: Wrapper, wrapped: FunctionType) -> bc.Bytecode: - """Wrap a function with a wrapper function. - - The wrapper function expects the wrapped function as the first argument, - followed by the tuple of arguments and the dictionary of keyword arguments. - The nature of the wrapped function is also honored, meaning that a generator - function will return a generator function, and a coroutine function will - return a coroutine function, and so on. The signature is also preserved to - avoid breaking, e.g., usages of the ``inspect`` module. - """ - - code = wrapped.__code__ - lineno = code.co_firstlineno + FIRSTLINENO_OFFSET - - # Push the wrapper function that is to be called and the wrapped function to - # be passed as first argument. - instrs = HEAD.bind({"wrapper": wrapper, "wrapped": wrapped}, lineno=lineno) - - # Add positional arguments - instrs.extend(generate_posargs(code)) - - # Add keyword arguments - instrs.extend(generate_kwargs(code)) - - # Call the wrapper function with the wrapped function, the positional and - # keyword arguments, and return the result. This is equivalent to - # - # >>> return wrapper(wrapped, args, kwargs) - instrs.extend(CALL_RETURN.bind({"arg": 3}, lineno=lineno)) - - # Include code for handling free/cell variables, if needed - if PY >= (3, 11): - if code.co_cellvars: - instrs[0:0] = [Instr("MAKE_CELL", bc.CellVar(_), lineno=lineno) for _ in code.co_cellvars] # type: ignore[attr-defined] - - if code.co_freevars: - instrs.insert(0, Instr("COPY_FREE_VARS", len(code.co_freevars), lineno=lineno)) - - # If the function has special flags set, like the generator, async generator - # or coroutine, inject unraveling code before the return opcode. - if (bc.CompilerFlags.GENERATOR & code.co_flags) and not (bc.CompilerFlags.COROUTINE & code.co_flags): - wrap_generator(instrs, code, lineno) - else: - wrap_async(instrs, code, lineno) - - return instrs - - -def wrap(f: FunctionType, wrapper: Wrapper) -> WrappedFunction: - """Wrap a function with a wrapper. - - The wrapper expects the function as first argument, followed by the tuple - of positional arguments and the dict of keyword arguments. - - Note that this changes the behavior of the original function with the - wrapper function, instead of creating a new function object. - """ - wrapped = FunctionType( - code := f.__code__, - f.__globals__, - "", - f.__defaults__, - f.__closure__, - ) - - # Carry forward the existing wrapped-function link to the new inner copy. - with _wrapped_lock: - existing_inner = _wrapped.get(f) - if existing_inner is not None: - _wrapped[wrapped] = existing_inner - - wrapped.__kwdefaults__ = f.__kwdefaults__ - - flags = code.co_flags - nargs = ( - (argcount := code.co_argcount) - + (kwonlycount := code.co_kwonlyargcount) - + bool(flags & bc.CompilerFlags.VARARGS) - + bool(flags & bc.CompilerFlags.VARKEYWORDS) - ) - - # Wrap the wrapped function with the wrapper - wrapped_code = wrap_bytecode(wrapper, wrapped) - - # Copy over the code attributes - wrapped_code.argcount = argcount - wrapped_code.argnames = list(code.co_varnames[:nargs]) - wrapped_code.filename = code.co_filename - wrapped_code.freevars = list(code.co_freevars) - wrapped_code.flags = bc.CompilerFlags(flags) - wrapped_code.kwonlyargcount = kwonlycount - wrapped_code.name = code.co_name - wrapped_code.posonlyargcount = code.co_posonlyargcount - if PY >= (3, 11): - wrapped_code.cellvars = list(code.co_cellvars) - - # Replace the function code with the trampoline bytecode - f.__code__ = wrapped_code.to_code() - - # DEV: Multiple wrapping is implemented as a singly-linked list via _wrapped. - with _wrapped_lock: - _wrapped[f] = wrapped - - # Link the original code object to the original function - link_function_to_code(code, f) - - return cast(WrappedFunction, f) - - -def _unwrap_method(f: Any) -> Any: - """Return the underlying function if *f* is a bound or unbound method.""" - func = getattr(f, "__func__", f) - return func - - -def is_wrapped(f: FunctionType) -> bool: - """Check if a function is wrapped with any wrapper.""" - try: - with _wrapped_lock: - inner = _wrapped.get(_unwrap_method(f)) - except TypeError: - # f is not weakly referenceable (e.g. C method_descriptor) - return False - if inner is None: - return False - assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - return True - - -def is_wrapped_with(f: FunctionType, wrapper: Wrapper) -> bool: - """Check if a function is wrapped with a specific wrapper.""" - f = cast(FunctionType, _unwrap_method(f)) - try: - with _wrapped_lock: - inner = _wrapped.get(f) - except TypeError: - return False - if inner is None: - return False - - assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - - if wrapper in f.__code__.co_consts: - return True - - # This is not the correct wrapping layer. Try with the next one. - return is_wrapped_with(inner, wrapper) - - -def unwrap(wf: WrappedFunction, wrapper: Wrapper) -> FunctionType: - """Unwrap a wrapped function. - - This is the reverse of :func:`wrap`. In case of multiple wrapping layers, - this will unwrap the one that uses ``wrapper``. If the function was not - wrapped with ``wrapper``, it will return the first argument. - """ - # DEV: Multiple wrapping layers are singly-linked via _wrapped. When we - # find the layer that needs to be removed we also have to ensure that we - # update the link at the deletion site if there is a non-empty tail. - f = cast(FunctionType, wf) - with _wrapped_lock: - inner = _wrapped.get(f) - if inner is None: - # The function is not wrapped so we return it as is. - return f - - assert inner.__name__ == "", "Wrapper has wrapped function" # nosec - - if wrapper not in f.__code__.co_consts: - # This is not the correct wrapping layer. Try with the next one. - return unwrap(cast(WrappedFunction, inner), wrapper) - - # Remove the current wrapping layer by moving the next one over the current - # one. The code swap and registry update must be atomic: two concurrent - # unwrap calls on the same function and wrapper must not both succeed. - with _wrapped_lock: - f.__code__ = inner.__code__ - next_inner = _wrapped.get(inner) - if next_inner is not None: - _wrapped[f] = next_inner - else: - del _wrapped[f] - - return f - - -def get_function_code(f: FunctionType) -> CodeType: - with _wrapped_lock: - inner = _wrapped.get(f) - return (inner if inner is not None else f).__code__ - - -def set_function_code(f: FunctionType, code: CodeType) -> None: - with _wrapped_lock: - inner = _wrapped.get(f) - (inner if inner is not None else f).__code__ = code - - -def get_wrapped(f: FunctionType) -> Optional[FunctionType]: - """Return the inner bytecode copy of *f* if it is wrapped, else None.""" - try: - with _wrapped_lock: - return _wrapped.get(cast(FunctionType, _unwrap_method(f))) - except TypeError: - return None +from ddtrace.internal.utils.wrapping import * # noqa diff --git a/ddtrace/internal/wrapping/asyncs.py b/ddtrace/internal/wrapping/asyncs.py index a1de28a074e..339af72067c 100644 --- a/ddtrace/internal/wrapping/asyncs.py +++ b/ddtrace/internal/wrapping/asyncs.py @@ -1,725 +1 @@ -import sys -from types import CodeType - -import bytecode as bc - -from ddtrace.internal.assembly import Assembly - - -PY = sys.version_info[:2] - - -# ----------------------------------------------------------------------------- -# Coroutine and Async Generator Wrapping -# ----------------------------------------------------------------------------- -# DEV: The wrapping of async generators is roughly equivalent to -# -# __ddgen = wrapper(wrapped, args, kwargs) -# __ddgensend = __ddgen.asend -# try: -# value = await __ddgen.__anext__() -# while True: -# try: -# tosend = yield value -# except GeneratorExit: -# await __ddgen.aclose() -# except Exception: -# value = await __ddgen.athrow(*sys.exc_info()) -# else: -# value = await __ddgensend(tosend) -# except StopAsyncIteration: -# return -# ----------------------------------------------------------------------------- - -COROUTINE_ASSEMBLY = Assembly() -ASYNC_GEN_ASSEMBLY = Assembly() -ASYNC_HEAD_ASSEMBLY = None - -if PY >= (3, 15): - raise NotImplementedError("This version of CPython is not supported yet") - -elif PY >= (3, 14): - ASYNC_HEAD_ASSEMBLY = Assembly() - ASYNC_HEAD_ASSEMBLY.parse( - r""" - return_generator - pop_top - """ - ) - - COROUTINE_ASSEMBLY.parse( - r""" - get_awaitable 0 - load_const None - - presend: - send @send - yield_value 1 - resume 3 - jump_backward_no_interrupt @presend - send: - end_send - """ - ) - - ASYNC_GEN_ASSEMBLY.parse( - r""" - try @stopiter - copy 1 - store_fast $__ddgen - load_attr (False, 'asend') - store_fast $__ddgensend - load_fast $__ddgen - load_attr (True, '__anext__') - call 0 - - loop: - get_awaitable 0 - load_const None - presend0: - send @send0 - tried - - try @genexit lasti - yield_value 0 - resume 3 - jump_backward_no_interrupt @presend0 - send0: - end_send - - yield: - call_intrinsic_1 asm.Intrinsic1Op.INTRINSIC_ASYNC_GEN_WRAP - yield_value 0 - resume 1 - push_null - load_fast $__ddgensend - swap 3 - call 1 - jump_backward @loop - tried - - genexit: - try @stopiter - push_exc_info - load_const GeneratorExit - check_exc_match - pop_jump_if_false @exc - pop_top - load_fast $__ddgen - load_attr (True, 'aclose') - call 0 - get_awaitable 0 - load_const None - - presend1: - send @send1 - yield_value 0 - resume 3 - jump_backward_no_interrupt @presend1 - send1: - end_send - pop_top - pop_except - load_const None - return_value - - exc: - pop_top - load_fast $__ddgen - load_attr (False, 'athrow') - push_null - load_const sys.exc_info - push_null - call 0 - push_null - call_function_ex - get_awaitable 0 - load_const None - - presend2: - send @send2 - yield_value 0 - resume 3 - jump_backward_no_interrupt @presend2 - send2: - end_send - swap 2 - pop_except - jump_backward @yield - tried - - stopiter: - push_exc_info - load_const StopAsyncIteration - check_exc_match - pop_jump_if_false @propagate - pop_top - pop_except - load_const None - return_value - - propagate: - reraise 0 - """ - ) - -elif PY >= (3, 13): - ASYNC_HEAD_ASSEMBLY = Assembly() - ASYNC_HEAD_ASSEMBLY.parse( - r""" - return_generator - pop_top - """ - ) - - COROUTINE_ASSEMBLY.parse( - r""" - get_awaitable 0 - load_const None - - presend: - send @send - yield_value 1 - resume 3 - jump_backward_no_interrupt @presend - send: - end_send - """ - ) - - ASYNC_GEN_ASSEMBLY.parse( - r""" - try @stopiter - copy 1 - store_fast $__ddgen - load_attr (False, 'asend') - store_fast $__ddgensend - load_fast $__ddgen - load_attr (True, '__anext__') - call 0 - - loop: - get_awaitable 0 - load_const None - presend0: - send @send0 - tried - - try @genexit lasti - yield_value 0 - resume 3 - jump_backward_no_interrupt @presend0 - send0: - end_send - - yield: - call_intrinsic_1 asm.Intrinsic1Op.INTRINSIC_ASYNC_GEN_WRAP - yield_value 0 - resume 1 - push_null - load_fast $__ddgensend - swap 3 - call 1 - jump_backward @loop - tried - - genexit: - try @stopiter - push_exc_info - load_const GeneratorExit - check_exc_match - pop_jump_if_false @exc - pop_top - load_fast $__ddgen - load_attr (True, 'aclose') - call 0 - get_awaitable 0 - load_const None - - presend1: - send @send1 - yield_value 0 - resume 3 - jump_backward_no_interrupt @presend1 - send1: - end_send - pop_top - pop_except - load_const None - return_value - - exc: - pop_top - load_fast $__ddgen - load_attr (False, 'athrow') - push_null - load_const sys.exc_info - push_null - call 0 - call_function_ex 0 - get_awaitable 0 - load_const None - - presend2: - send @send2 - yield_value 0 - resume 3 - jump_backward_no_interrupt @presend2 - send2: - end_send - swap 2 - pop_except - jump_backward @yield - tried - - stopiter: - push_exc_info - load_const StopAsyncIteration - check_exc_match - pop_jump_if_false @propagate - pop_top - pop_except - load_const None - return_value - - propagate: - reraise 0 - """ - ) - -elif PY >= (3, 12): - ASYNC_HEAD_ASSEMBLY = Assembly() - ASYNC_HEAD_ASSEMBLY.parse( - r""" - return_generator - pop_top - """ - ) - - COROUTINE_ASSEMBLY.parse( - r""" - get_awaitable 0 - load_const None - - presend: - send @send - yield_value 2 - resume 3 - jump_backward_no_interrupt @presend - send: - end_send - """ - ) - - ASYNC_GEN_ASSEMBLY.parse( - r""" - try @stopiter - copy 1 - store_fast $__ddgen - load_attr (False, 'asend') - store_fast $__ddgensend - load_fast $__ddgen - load_attr (True, '__anext__') - call 0 - - loop: - get_awaitable 0 - load_const None - presend0: - send @send0 - tried - - try @genexit lasti - yield_value 3 - resume 3 - jump_backward_no_interrupt @presend0 - send0: - end_send - - yield: - call_intrinsic_1 asm.Intrinsic1Op.INTRINSIC_ASYNC_GEN_WRAP - yield_value 3 - resume 1 - push_null - swap 2 - load_fast $__ddgensend - swap 2 - call 1 - jump_backward @loop - tried - - genexit: - try @stopiter - push_exc_info - load_const GeneratorExit - check_exc_match - pop_jump_if_false @exc - pop_top - load_fast $__ddgen - load_attr (True, 'aclose') - call 0 - get_awaitable 0 - load_const None - - presend1: - send @send1 - yield_value 4 - resume 3 - jump_backward_no_interrupt @presend1 - send1: - end_send - pop_top - pop_except - return_const None - - exc: - pop_top - push_null - load_fast $__ddgen - load_attr (False, 'athrow') - push_null - load_const sys.exc_info - call 0 - call_function_ex 0 - get_awaitable 0 - load_const None - - presend2: - send @send2 - yield_value 4 - resume 3 - jump_backward_no_interrupt @presend2 - send2: - end_send - swap 2 - pop_except - jump_backward @yield - tried - - stopiter: - push_exc_info - load_const StopAsyncIteration - check_exc_match - pop_jump_if_false @propagate - pop_top - pop_except - return_const None - - propagate: - reraise 0 - """ - ) - - -elif PY >= (3, 11): - ASYNC_HEAD_ASSEMBLY = Assembly() - ASYNC_HEAD_ASSEMBLY.parse( - r""" - return_generator - pop_top - """ - ) - - COROUTINE_ASSEMBLY.parse( - r""" - get_awaitable 0 - load_const None - - presend: - send @send - yield_value - resume 3 - jump_backward_no_interrupt @presend - send: - """ - ) - - ASYNC_GEN_ASSEMBLY.parse( - r""" - try @stopiter - copy 1 - store_fast $__ddgen - load_attr $asend - store_fast $__ddgensend - load_fast $__ddgen - load_method $__anext__ - precall 0 - call 0 - - loop: - get_awaitable 0 - load_const None - presend0: - send @send0 - tried - - try @genexit lasti - yield_value - resume 3 - jump_backward_no_interrupt @presend0 - send0: - - yield: - async_gen_wrap - yield_value - resume 1 - push_null - swap 2 - load_fast $__ddgensend - swap 2 - precall 1 - call 1 - jump_backward @loop - tried - - genexit: - try @stopiter - push_exc_info - load_const GeneratorExit - check_exc_match - pop_jump_forward_if_false @exc - pop_top - load_fast $__ddgen - load_method $aclose - precall 0 - call 0 - get_awaitable 0 - load_const None - - presend1: - send @send1 - yield_value - resume 3 - jump_backward_no_interrupt @presend1 - send1: - pop_top - pop_except - load_const None - return_value - - exc: - pop_top - push_null - load_fast $__ddgen - load_attr $athrow - push_null - load_const sys.exc_info - precall 0 - call 0 - call_function_ex 0 - get_awaitable 0 - load_const None - - presend2: - send @send2 - yield_value - resume 3 - jump_backward_no_interrupt @presend2 - send2: - swap 2 - pop_except - jump_backward @yield - tried - - stopiter: - push_exc_info - load_const StopAsyncIteration - check_exc_match - pop_jump_forward_if_false @propagate - pop_top - pop_except - load_const None - return_value - - propagate: - reraise 0 - """ - ) - - -elif PY >= (3, 10): - COROUTINE_ASSEMBLY.parse( - r""" - get_awaitable - load_const None - yield_from - """ - ) - - ASYNC_GEN_ASSEMBLY.parse( - r""" - setup_finally @stopiter - dup_top - store_fast $__ddgen - load_attr $asend - store_fast $__ddgensend - load_fast $__ddgen - load_attr $__anext__ - call_function 0 - - loop: - get_awaitable - load_const None - yield_from - - yield: - setup_finally @genexit - yield_value - pop_block - load_fast $__ddgensend - rot_two - call_function 1 - jump_absolute @loop - - genexit: - dup_top - load_const GeneratorExit - jump_if_not_exc_match @exc - pop_top - pop_top - pop_top - pop_top - load_fast $__ddgen - load_attr $aclose - call_function 0 - get_awaitable - load_const None - yield_from - pop_except - return_value - - exc: - pop_top - pop_top - pop_top - pop_top - load_fast $__ddgen - load_attr $athrow - load_const sys.exc_info - call_function 0 - call_function_ex 0 - get_awaitable - load_const None - yield_from - rot_four - pop_except - jump_absolute @yield - - stopiter: - dup_top - load_const StopAsyncIteration - jump_if_not_exc_match @propagate - pop_top - pop_top - pop_top - pop_except - load_const None - return_value - - propagate: - reraise 0 - """ - ) - - -elif PY >= (3, 9): - COROUTINE_ASSEMBLY.parse( - r""" - get_awaitable - load_const None - yield_from - """ - ) - - ASYNC_GEN_ASSEMBLY.parse( - r""" - setup_finally @stopiter - dup_top - store_fast $__ddgen - load_attr $asend - store_fast $__ddgensend - load_fast $__ddgen - load_attr $__anext__ - call_function 0 - - loop: - get_awaitable - load_const None - yield_from - - yield: - setup_finally @genexit - yield_value - pop_block - load_fast $__ddgensend - rot_two - call_function 1 - jump_absolute @loop - - genexit: - dup_top - load_const GeneratorExit - jump_if_not_exc_match @exc - pop_top - pop_top - pop_top - pop_top - load_fast $__ddgen - load_attr $aclose - call_function 0 - get_awaitable - load_const None - yield_from - pop_except - return_value - - exc: - pop_top - pop_top - pop_top - pop_top - load_fast $__ddgen - load_attr $athrow - load_const sys.exc_info - call_function 0 - call_function_ex 0 - get_awaitable - load_const None - yield_from - rot_four - pop_except - jump_absolute @yield - - stopiter: - dup_top - load_const StopAsyncIteration - jump_if_not_exc_match @propagate - pop_top - pop_top - pop_top - pop_except - load_const None - return_value - - propagate: - reraise - """ - ) - -else: - msg = "No async wrapping support for Python %d.%d" % PY[:2] - raise RuntimeError(msg) - - -def wrap_async(instrs: bc.Bytecode, code: CodeType, lineno: int) -> None: - if (bc.CompilerFlags.ASYNC_GENERATOR | bc.CompilerFlags.COROUTINE) & code.co_flags: - if ASYNC_HEAD_ASSEMBLY is not None: - instrs[0:0] = ASYNC_HEAD_ASSEMBLY.bind(lineno=lineno) - - if bc.CompilerFlags.COROUTINE & code.co_flags: - # DEV: This is just - # >>> return await wrapper(wrapped, args, kwargs) - instrs[-1:-1] = COROUTINE_ASSEMBLY.bind(lineno=lineno) - - elif bc.CompilerFlags.ASYNC_GENERATOR & code.co_flags: - instrs[-1:] = ASYNC_GEN_ASSEMBLY.bind(lineno=lineno) +from ddtrace.internal.utils.wrapping.asyncs import * # noqa diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index 8333b997c36..1e7973a2e25 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -1,876 +1 @@ -from abc import ABC -from contextvars import ContextVar -from inspect import iscoroutinefunction -from inspect import isgeneratorfunction -import sys -from types import FrameType -from types import FunctionType -from types import TracebackType -import typing as t -from typing import Protocol # noqa:F401 -import weakref - -import bytecode -from bytecode import Bytecode - -from ddtrace.internal.assembly import Assembly -from ddtrace.internal.logger import get_logger -from ddtrace.internal.threads import Lock -from ddtrace.internal.threads import RLock -from ddtrace.internal.wrapping import WrappedFunction -from ddtrace.internal.wrapping import Wrapper -from ddtrace.internal.wrapping import get_function_code -from ddtrace.internal.wrapping import is_wrapped_with -from ddtrace.internal.wrapping import link_function_to_code -from ddtrace.internal.wrapping import set_function_code -from ddtrace.internal.wrapping import unwrap -from ddtrace.internal.wrapping import wrap - - -class _ContextRecord: - """Holds all wrapping-context metadata for a single function. - - Stores only weak references to context objects so that the WeakKeyDictionary - key (the function) is not kept alive by the registry value chain: - _registry -> _ContextRecord -> context -> context.__wrapped__ -> function - Breaking this path with weak references lets ephemeral functions be garbage - collected as soon as all external strong references drop. - """ - - __slots__ = ("_uwc_ref", "lazy_contexts") - - def __init__(self) -> None: - self._uwc_ref: t.Optional[weakref.ref["_UniversalWrappingContext"]] = None - # WeakSet so that LazyWrappingContext instances (which also hold - # __wrapped__ = f) do not prevent the function from being collected. - self.lazy_contexts: weakref.WeakSet["LazyWrappingContext"] = weakref.WeakSet() - - @property - def uwc(self) -> t.Optional["_UniversalWrappingContext"]: - ref = self._uwc_ref - return ref() if ref is not None else None - - @uwc.setter - def uwc(self, value: t.Optional["_UniversalWrappingContext"]) -> None: - self._uwc_ref = weakref.ref(value) if value is not None else None - - @classmethod - def get_or_create(cls, f: FunctionType) -> "_ContextRecord": - record = _registry.get(f) - if record is None: - with _registry_lock: - record = _registry.get(f) - if record is None: - record = cls() - _registry[f] = record - return record - - -# Per-function registry for wrapping-context machinery. WeakKeyDictionary so -# functions are not kept alive by the registry alone. Storing data here instead -# of as function attributes keeps __dict__ clean, preventing frameworks that -# copy function __dict__ (e.g. functools.wraps, -# self.__dict__.update(f.__dict__)) from capturing non-picklable objects. -_registry: weakref.WeakKeyDictionary[FunctionType, _ContextRecord] = weakref.WeakKeyDictionary() -_registry_lock = RLock() - - -log = get_logger(__name__) - -T = t.TypeVar("T") - -# This module implements utilities for wrapping a function with a context -# manager. The rough idea is to re-write the function's bytecode to look like -# this: -# -# def foo(): -# with wrapping_context: -# # Original function code -# -# Because we also want to capture the return value, our context manager extends -# the Python one by implementing a __return__ method that will be called with -# the return value of the function. Contrary to ordinary context managers, -# though, the __exit__ method is only called if the function raises an -# exception. -# -# Because CPython 3.11 introduced zero-cost exceptions, we cannot nest try -# blocks in the function's bytecode. In this case, we call the context manager -# methods directly at the right places, and set up the appropriate exception -# handling code. For older versions of Python we rely on the with statement to -# perform entry and exit operations. Calls to __return__ are explicit in all -# cases. -# -# Some advantages of wrapping a function this way are: -# - Access to the local variables on entry and on return/exit via the frame -# object. -# - No intermediate function calls that pollute the call stack. -# - No need to call the wrapped function manually. -# -# The actual bytecode wrapping is performed once on a target function via a -# universal wrapping context. Multiple context wrapping of a function is allowed -# and it is virtually implemented on top of the concrete universal wrapping -# context. This makes multiple wrapping/unwrapping easy, as it translates to a -# single bytecode wrapping/unwrapping operation. -# -# Context wrappers should be implemented as subclasses of the WrappingContext -# class. The __priority__ attribute can be used to control the order in which -# multiple context wrappers are entered and exited. The __enter__ and __exit__ -# methods should be implemented to perform the necessary operations. The -# __exit__ method is called if the wrapped function raises an exception. The -# frame of the wrapped function can be accessed via the __frame__ property. The -# __return__ method can be implemented to capture the return value of the -# wrapped function. If implemented, its return value will be used as the wrapped -# function return value. The wrapped function can be accessed via the -# __wrapped__ attribute. Context-specific values can be stored and retrieved -# with the set and get methods. - -CONTEXT_HEAD = Assembly() -CONTEXT_RETURN = Assembly() -CONTEXT_FOOT = Assembly() - -if sys.version_info >= (3, 15): - raise NotImplementedError("Python >= 3.15 is not supported yet") -elif sys.version_info >= (3, 13): - CONTEXT_HEAD.parse( - r""" - load_const {context_enter} - push_null - call 0 - pop_top - """ - ) - CONTEXT_RETURN.parse( - r""" - push_null - load_const {context_return} - swap 3 - call 1 - """ - ) - - CONTEXT_RETURN_CONST = Assembly() - CONTEXT_RETURN_CONST.parse( - r""" - load_const {context_return} - push_null - load_const {value} - call 1 - """ - ) - - CONTEXT_FOOT.parse( - r""" - try @_except lasti - push_exc_info - load_const {context_exit} - push_null - call 0 - pop_top - reraise 2 - tried - - _except: - copy 3 - pop_except - reraise 1 - """ - ) - -elif sys.version_info >= (3, 12): - CONTEXT_HEAD.parse( - r""" - push_null - load_const {context_enter} - call 0 - pop_top - """ - ) - - CONTEXT_RETURN.parse( - r""" - load_const {context_return} - push_null - swap 3 - call 1 - """ - ) - - CONTEXT_RETURN_CONST = Assembly() - CONTEXT_RETURN_CONST.parse( - r""" - push_null - load_const {context_return} - load_const {value} - call 1 - """ - ) - - CONTEXT_FOOT.parse( - r""" - try @_except lasti - push_exc_info - push_null - load_const {context_exit} - call 0 - pop_top - reraise 2 - tried - - _except: - copy 3 - pop_except - reraise 1 - """ - ) - - -elif sys.version_info >= (3, 11): - CONTEXT_HEAD.parse( - r""" - push_null - load_const {context_enter} - precall 0 - call 0 - pop_top - """ - ) - - CONTEXT_RETURN.parse( - r""" - load_const {context_return} - push_null - swap 3 - precall 1 - call 1 - """ - ) - - CONTEXT_EXC_HEAD = Assembly() - CONTEXT_EXC_HEAD.parse( - r""" - push_null - load_const {context_exit} - precall 0 - call 0 - pop_top - """ - ) - - CONTEXT_FOOT.parse( - r""" - try @_except lasti - push_exc_info - push_null - load_const {context_exit} - precall 0 - call 0 - pop_top - reraise 2 - tried - - _except: - copy 3 - pop_except - reraise 1 - """ - ) - -elif sys.version_info >= (3, 10): - CONTEXT_HEAD.parse( - r""" - load_const {context} - setup_with @_except - pop_top - _except: - """ - ) - - CONTEXT_RETURN.parse( - r""" - pop_block - load_const {context} - load_method $__return__ - rot_three - rot_three - call_method 1 - rot_two - pop_top - """ - ) - - CONTEXT_FOOT.parse( - r""" - with_except_start - pop_top - reraise 1 - """ - ) - -elif sys.version_info >= (3, 9): - CONTEXT_HEAD.parse( - r""" - load_const {context} - setup_with @_except - pop_top - _except: - """ - ) - - CONTEXT_RETURN.parse( - r""" - pop_block - load_const {context} - load_method $__return__ - rot_three - rot_three - call_method 1 - rot_two - pop_top - """ - ) - - CONTEXT_FOOT.parse( - r""" - with_except_start - pop_top - reraise - """ - ) - - -# This is abstract and should not be used directly -class BaseWrappingContext(ABC): - __priority__: int = 0 - - def __init__(self, f: FunctionType): - # Store a weak reference so that context objects do not keep the wrapped - # function alive. CodeType is not GC-tracked in CPython, so the cycle - # f → code.co_consts → bound_methods(uwc) → uwc.__wrapped__ → f - # cannot be broken by the cyclic GC. A weak ref here allows f's - # reference count to reach zero (and be freed) as soon as all external - # strong refs drop, without relying on the cyclic GC at all. - self._wrapped_ref: weakref.ref[FunctionType] = weakref.ref(f) - self._storage: ContextVar[t.Optional[dict[str, t.Any]]] = ContextVar( - f"{type(self).__name__}__storage", default=None - ) - - @property - def __wrapped__(self) -> FunctionType: - f = self._wrapped_ref() - if f is None: - raise RuntimeError(f"{type(self).__name__}.__wrapped__: the wrapped function has been garbage collected") - return f - - @__wrapped__.setter - def __wrapped__(self, f: FunctionType) -> None: - self._wrapped_ref = weakref.ref(f) - - def __enter__(self) -> "BaseWrappingContext": - prev = self._storage.get() - self._storage.set({"__dd_wrapping_context_prev__": prev}) - - return self - - def _pop_storage(self) -> dict[str, t.Any]: - storage = t.cast(dict[str, t.Any], self._storage.get()) - self._storage.set(storage.pop("__dd_wrapping_context_prev__")) - return storage - - def __return__(self, value: T) -> T: - self._pop_storage() - return value - - def __exit__( - self, - exc_type: t.Optional[type[BaseException]], - exc_val: t.Optional[BaseException], - exc_tb: t.Optional[TracebackType], - ) -> None: - self._pop_storage() - - def get(self, key: str) -> t.Any: - return t.cast(dict[str, t.Any], self._storage.get())[key] - - def set(self, key: str, value: T) -> T: - t.cast(dict[str, t.Any], self._storage.get())[key] = value - return value - - @classmethod - def wrapped(cls, f: FunctionType) -> "BaseWrappingContext": - try: - context = cls.extract(f) - assert isinstance(context, cls) # nosec - except ValueError: - context = cls(f) - context.wrap() - return context - - @classmethod - def is_wrapped(cls, _f: FunctionType) -> bool: - raise NotImplementedError - - @classmethod - def extract(cls, _f: FunctionType) -> "BaseWrappingContext": - raise NotImplementedError - - def wrap(self) -> None: - raise NotImplementedError - - def unwrap(self) -> None: - raise NotImplementedError - - -# This is the public interface exported by this module -class WrappingContext(BaseWrappingContext): - @property - def __frame__(self) -> FrameType: - try: - return t.cast( - FrameType, - _UniversalWrappingContext.extract(t.cast(FunctionType, self.__wrapped__)).get("__frame__"), - ) - except ValueError: - raise AttributeError("Wrapping context not entered") - - def get_local(self, name: str) -> t.Any: - return self.__frame__.f_locals[name] - - @classmethod - def is_wrapped(cls, f: FunctionType) -> bool: - try: - return bool(cls.extract(f)) - except ValueError: - return False - - @classmethod - def extract(cls, f: FunctionType) -> "WrappingContext": - try: - return _UniversalWrappingContext.extract(f).registered(cls) - except (ValueError, KeyError): - msg = f"Function is not wrapped with {cls}" - raise ValueError(msg) - - def wrap(self) -> None: - t.cast( - _UniversalWrappingContext, _UniversalWrappingContext.wrapped(t.cast(FunctionType, self.__wrapped__)) - ).register(self) - - def unwrap(self) -> None: - f = t.cast(FunctionType, self.__wrapped__) - - try: - _UniversalWrappingContext.extract(f).unregister(self) - except ValueError: - pass - - -class LazyWrappingContext(WrappingContext): - def __init__(self, f: FunctionType): - super().__init__(f) - - self._trampoline: t.Optional[Wrapper] = None - self._trampoline_lock = Lock() - - @classmethod - def is_wrapped(cls, f: FunctionType) -> bool: - with _registry_lock: - record = _registry.get(f) - if record is None: - return False - return any(isinstance(c, cls) for c in record.lazy_contexts) - - def wrap(self) -> None: - """Perform the bytecode wrapping on first invocation.""" - with (tl := self._trampoline_lock): - if self._trampoline is not None: - return - - # If the function is already universally wrapped it's less expensive - # to do the normal wrapping. - if _UniversalWrappingContext.is_wrapped(t.cast(FunctionType, self.__wrapped__)): - super().wrap() - return - - def trampoline(_: t.Any, args: tuple[t.Any, ...], kwargs: dict[str, t.Any]) -> t.Any: - with tl: - f = t.cast(WrappedFunction, self.__wrapped__) - if is_wrapped_with(t.cast(FunctionType, self.__wrapped__), trampoline): - f = t.cast(WrappedFunction, unwrap(f, trampoline)) - - self._trampoline = None - - inconsistent = False - with _registry_lock: - record = _registry.get(t.cast(FunctionType, f)) - if record is not None: - inconsistent = self not in record.lazy_contexts - record.lazy_contexts.discard(self) - if not record.lazy_contexts and record.uwc is None: - _registry.pop(t.cast(FunctionType, f), None) - if inconsistent: - log.warning("Inconsistent lazy wrapping context state") - - super(LazyWrappingContext, self).wrap() - return f(*args, **kwargs) - - wrap(t.cast(FunctionType, self.__wrapped__), trampoline) - - self._trampoline = trampoline - - _ContextRecord.get_or_create(t.cast(FunctionType, self.__wrapped__)).lazy_contexts.add(self) - - def unwrap(self) -> None: - with self._trampoline_lock: - if _UniversalWrappingContext.is_wrapped(t.cast(FunctionType, self.__wrapped__)): - assert self._trampoline is None # nosec - super().unwrap() - elif self._trampoline is not None: - with _registry_lock: - record = _registry.get(t.cast(FunctionType, self.__wrapped__)) - if record is not None: - record.lazy_contexts.discard(self) - if not record.lazy_contexts and record.uwc is None: - _registry.pop(t.cast(FunctionType, self.__wrapped__), None) - - unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline) - self._trampoline = None - - -class ContextWrappedFunction(Protocol): - """A wrapped function.""" - - def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: - pass - - -# This class provides an interface between single bytecode wrapping and multiple -# logical context wrapping -class _UniversalWrappingContext(BaseWrappingContext): - def __init__(self, f: FunctionType) -> None: - super().__init__(f) - - self._contexts: list[WrappingContext] = [] - - def register(self, context: WrappingContext) -> None: - _type = type(context) - if any(isinstance(c, _type) for c in self._contexts): - raise ValueError("Context already registered") - - self._contexts.append(context) - self._contexts.sort(key=lambda c: c.__priority__) - - def unregister(self, context: WrappingContext) -> None: - try: - self._contexts.remove(context) - except ValueError: - raise ValueError("Context not registered") - - if not self._contexts: - self.unwrap() - - def is_registered(self, context: WrappingContext) -> bool: - return any(isinstance(c, type(context)) for c in self._contexts) - - def registered(self, context_type: type[WrappingContext]) -> WrappingContext: - for context in self._contexts: - if isinstance(context, context_type): - return context - raise KeyError(f"Context {context_type} not registered") - - def __enter__(self) -> "_UniversalWrappingContext": - super().__enter__() - - # Make the frame object available to the contexts - self.set("__frame__", sys._getframe(1)) - - for context in self._contexts: - context.__enter__() - - return self - - def _exit(self) -> None: - self.__exit__(*sys.exc_info()) - - def __exit__( - self, - exc_type: t.Optional[type[BaseException]], - exc_value: t.Optional[BaseException], - traceback: t.Optional[TracebackType], - ) -> None: - if exc_value is None: - return - - for context in self._contexts[::-1]: - context.__exit__(exc_type, exc_value, traceback) - - super().__exit__(exc_type, exc_value, traceback) - - def __return__(self, value: T) -> T: - for context in self._contexts[::-1]: - context.__return__(value) - - return super().__return__(value) - - @classmethod - def is_wrapped(cls, f: FunctionType) -> bool: - try: - with _registry_lock: - record = _registry.get(f) - if record is None or record.uwc is None: - return False - # Verify the registry entry matches actual bytecode wrapping. - if sys.version_info >= (3, 11): - return record.uwc.__enter__ in get_function_code(f).co_consts - else: - return record.uwc in get_function_code(f).co_consts - except AttributeError: - return False - - @classmethod - def extract(cls, f: FunctionType) -> "_UniversalWrappingContext": - with _registry_lock: - if not cls.is_wrapped(f): - raise ValueError("Function is not wrapped") - return t.cast(_UniversalWrappingContext, _registry[f].uwc) - - if sys.version_info >= (3, 11): - - def wrap(self) -> None: - f = self.__wrapped__ - - with _registry_lock: - if self.is_wrapped(f): - raise ValueError("Function already wrapped") - - bc = Bytecode.from_code(code := get_function_code(f)) - - # Prefix every return - i = 0 - while i < len(bc): - instr = bc[i] - try: - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN.bind({"context_return": self.__return__}, lineno=instr.lineno) - elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ - return_code = CONTEXT_RETURN_CONST.bind( - {"context_return": self.__return__, "value": instr.arg}, lineno=instr.lineno - ) - else: - return_code = [] - - bc[i:i] = return_code - i += len(return_code) - except AttributeError: - # Not an instruction - pass - i += 1 - - # Search for the RESUME instruction - for i, instr in enumerate(bc, 1): - try: - if instr.name == "RESUME": - break - except AttributeError: - # Not an instruction - pass - else: - i = 0 - - bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=code.co_firstlineno) - - # Wrap every line outside a try block - except_label = bytecode.Label() - first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) - - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None: - bc.insert(i, bytecode.TryEnd(last_try_begin)) - last_try_begin = None - i += 1 - elif isinstance(instr, bytecode.TryEnd): - j = i + 1 - while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin): - if isinstance(bc[j], bytecode.Instr): - last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) - bc.insert(i + 1, last_try_begin) - break - j += 1 - i += 1 - i += 1 - - bc.insert(0, first_try_begin) - - bc.append(bytecode.TryEnd(last_try_begin)) - bc.append(except_label) - bc.extend(CONTEXT_FOOT.bind({"context_exit": self._exit}, lineno=code.co_firstlineno)) - - # Register the wrapping context and write the new bytecode. - _ContextRecord.get_or_create(f).uwc = self - link_function_to_code(code, f) - set_function_code(f, bc.to_code()) - - def unwrap(self) -> None: - f = self.__wrapped__ - - with _registry_lock: - if not self.is_wrapped(f): - return - - wc = _registry[f].uwc - - bc = Bytecode.from_code(get_function_code(f)) - - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() - bc.pop() - - except_label = bc.pop(0).target - - # Remove the try blocks - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.TryBegin) and instr.target is except_label: - bc.pop(i) - elif isinstance(instr, bytecode.TryEnd) and instr.entry.target is except_label: - bc.pop(i) - else: - i += 1 - - # Remove the head of the try block - for i, instr in enumerate(bc): - if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: - break - - # Search for the RESUME instruction - for i, instr in enumerate(bc, 1): - try: - if instr.name == "RESUME": - break - except AttributeError: - # Not an instruction - pass - else: - i = 0 - - bc[i : i + len(CONTEXT_HEAD)] = [] - - # Un-prefix every return - i = 0 - while i < len(bc): - instr = bc[i] - try: - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN - elif sys.version_info >= (3, 12) and instr.name == "RETURN_CONST": # Python 3.12+ - return_code = CONTEXT_RETURN_CONST - else: - return_code = None - - if return_code is not None: - bc[i - len(return_code) : i] = [] - i -= len(return_code) - except AttributeError: - # Not an instruction - pass - i += 1 - - # Recreate the code object - set_function_code(f, bc.to_code()) - - # Clear the UWC from the registry; remove the record if fully empty. - record = _registry.get(f) - if record is not None: - record.uwc = None - if not record.lazy_contexts: - _registry.pop(f, None) - - else: - - def wrap(self) -> None: - f = t.cast(FunctionType, self.__wrapped__) - - with _registry_lock: - if self.is_wrapped(f): - raise ValueError("Function already wrapped") - - bc = Bytecode.from_code(code := get_function_code(f)) - - # Prefix every return - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.Instr): - if instr.name == "RETURN_VALUE": - return_code = CONTEXT_RETURN.bind({"context": self}, lineno=instr.lineno) - bc[i:i] = return_code - i += len(return_code) - i += 1 - - # Search for the GEN_START instruction, which needs to stay on top. - i = 0 - if sys.version_info >= (3, 10) and (iscoroutinefunction(f) or isgeneratorfunction(f)): - for i, instr in enumerate(bc, 1): - if isinstance(instr, bytecode.Instr) and instr.name == "GEN_START": - break - - *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": self}, lineno=code.co_firstlineno) - - bc.append(except_label) - bc.extend(CONTEXT_FOOT.bind(lineno=code.co_firstlineno)) - - # Register the wrapping context and write the new bytecode. - _ContextRecord.get_or_create(f).uwc = self - link_function_to_code(code, f) - set_function_code(f, bc.to_code()) - - def unwrap(self) -> None: - f = t.cast(FunctionType, self.__wrapped__) - - with _registry_lock: - if not self.is_wrapped(f): - return - - wc = _registry[f].uwc - - bc = Bytecode.from_code(get_function_code(f)) - - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() - - # Remove the head of the try block - for i, instr in enumerate(bc): - if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: - break - - bc[i : i + len(CONTEXT_HEAD) - 1] = [] - - # Remove all the return handlers - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.Instr) and instr.name == "RETURN_VALUE": - bc[i - len(CONTEXT_RETURN) : i] = [] - i -= len(CONTEXT_RETURN) - i += 1 - - # Recreate the code object - set_function_code(f, bc.to_code()) - - # Clear the UWC from the registry; remove the record if fully empty. - record = _registry.get(f) - if record is not None: - record.uwc = None - if not record.lazy_contexts: - _registry.pop(f, None) - - -def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingContext]": - """Return the _UniversalWrappingContext for *f*, or None if not context-wrapped.""" - with _registry_lock: - record = _registry.get(f) - return record.uwc if record is not None else None +from ddtrace.internal.utils.wrapping.context import * # noqa diff --git a/ddtrace/internal/wrapping/generators.py b/ddtrace/internal/wrapping/generators.py index 01c284c4686..ceb1cf2fd7a 100644 --- a/ddtrace/internal/wrapping/generators.py +++ b/ddtrace/internal/wrapping/generators.py @@ -1,491 +1 @@ -import sys -from types import CodeType - -import bytecode as bc - -from ddtrace.internal.assembly import Assembly - - -PY = sys.version_info[:2] - - -# ----------------------------------------------------------------------------- -# Generator Wrapping -# ----------------------------------------------------------------------------- -# DEV: This is roughly equivalent to -# -# __ddgen = wrapper(wrapped, args, kwargs) -# __ddgensend = __ddgen.send -# try: -# value = next(__ddgen) -# while True: -# try: -# tosend = yield value -# except GeneratorExit: -# return __ddgen.close() -# except Exception: -# value = __ddgen.throw(*sys.exc_info()) -# else: -# value = __ddgensend(tosend) -# except StopIteration: -# return -# ----------------------------------------------------------------------------- -GENERATOR_ASSEMBLY = Assembly() -GENERATOR_HEAD_ASSEMBLY = None - -if PY >= (3, 15): - raise NotImplementedError("This version of CPython is not supported yet") - -elif PY >= (3, 14): - GENERATOR_HEAD_ASSEMBLY = Assembly() - GENERATOR_HEAD_ASSEMBLY.parse( - r""" - return_generator - pop_top - """ - ) - - GENERATOR_ASSEMBLY.parse( - r""" - try @stopiter - copy 1 - store_fast $__ddgen - load_attr $send - store_fast $__ddgensend - load_const next - push_null - load_fast_borrow $__ddgen - - loop: - call 1 - tried - - yield: - try @genexit lasti - yield_value 0 - resume 1 - push_null - load_fast_borrow $__ddgensend - swap 3 - jump_backward @loop - tried - - genexit: - try @stopiter - push_exc_info - load_const GeneratorExit - check_exc_match - pop_jump_if_false @exc - pop_top - load_fast $__ddgen - load_method $close - call 0 - swap 2 - pop_except - return_value - - exc: - pop_top - load_fast $__ddgen - load_attr $throw - push_null - load_const sys.exc_info - push_null - call 0 - push_null - call_function_ex - swap 2 - pop_except - jump_backward @yield - tried - - stopiter: - push_exc_info - load_const StopIteration - check_exc_match - pop_jump_if_false @propagate - pop_top - pop_except - load_const None - return_value - - propagate: - reraise 0 - """ - ) - -elif PY >= (3, 13): - GENERATOR_HEAD_ASSEMBLY = Assembly() - GENERATOR_HEAD_ASSEMBLY.parse( - r""" - return_generator - pop_top - """ - ) - - GENERATOR_ASSEMBLY.parse( - r""" - try @stopiter - copy 1 - store_fast $__ddgen - load_attr $send - store_fast $__ddgensend - load_const next - push_null - load_fast $__ddgen - - loop: - call 1 - tried - - yield: - try @genexit lasti - yield_value 0 - resume 1 - push_null - load_fast $__ddgensend - swap 3 - jump_backward @loop - tried - - genexit: - try @stopiter - push_exc_info - load_const GeneratorExit - check_exc_match - pop_jump_if_false @exc - pop_top - load_fast $__ddgen - load_method $close - call 0 - swap 2 - pop_except - return_value - - exc: - pop_top - load_fast $__ddgen - load_attr $throw - push_null - load_const sys.exc_info - push_null - call 0 - call_function_ex 0 - swap 2 - pop_except - jump_backward @yield - tried - - stopiter: - push_exc_info - load_const StopIteration - check_exc_match - pop_jump_if_false @propagate - pop_top - pop_except - load_const None - return_value - - propagate: - reraise 0 - """ - ) - -elif PY >= (3, 12): - GENERATOR_HEAD_ASSEMBLY = Assembly() - GENERATOR_HEAD_ASSEMBLY.parse( - r""" - return_generator - pop_top - """ - ) - - GENERATOR_ASSEMBLY.parse( - r""" - try @stopiter - copy 1 - store_fast $__ddgen - load_attr $send - store_fast $__ddgensend - push_null - load_const next - load_fast $__ddgen - - loop: - call 1 - tried - - yield: - try @genexit lasti - yield_value 3 - resume 1 - push_null - swap 2 - load_fast $__ddgensend - swap 2 - jump_backward @loop - tried - - genexit: - try @stopiter - push_exc_info - load_const GeneratorExit - check_exc_match - pop_jump_if_false @exc - pop_top - load_fast $__ddgen - load_method $close - call 0 - swap 2 - pop_except - return_value - - exc: - pop_top - push_null - load_fast $__ddgen - load_attr $throw - push_null - load_const sys.exc_info - call 0 - call_function_ex 0 - swap 2 - pop_except - jump_backward @yield - tried - - stopiter: - push_exc_info - load_const StopIteration - check_exc_match - pop_jump_if_false @propagate - pop_top - pop_except - return_const None - - propagate: - reraise 0 - """ - ) - -elif PY >= (3, 11): - GENERATOR_HEAD_ASSEMBLY = Assembly() - GENERATOR_HEAD_ASSEMBLY.parse( - r""" - return_generator - pop_top - """ - ) - - GENERATOR_ASSEMBLY.parse( - r""" - try @stopiter - copy 1 - store_fast $__ddgen - load_attr $send - store_fast $__ddgensend - push_null - load_const next - load_fast $__ddgen - - loop: - precall 1 - call 1 - tried - - yield: - try @genexit lasti - yield_value - resume 1 - push_null - swap 2 - load_fast $__ddgensend - swap 2 - jump_backward @loop - tried - - genexit: - try @stopiter - push_exc_info - load_const GeneratorExit - check_exc_match - pop_jump_forward_if_false @exc - pop_top - load_fast $__ddgen - load_method $close - precall 0 - call 0 - swap 2 - pop_except - return_value - - exc: - pop_top - push_null - load_fast $__ddgen - load_attr $throw - push_null - load_const sys.exc_info - precall 0 - call 0 - call_function_ex 0 - swap 2 - pop_except - jump_backward @yield - tried - - stopiter: - push_exc_info - load_const StopIteration - check_exc_match - pop_jump_forward_if_false @propagate - pop_top - pop_except - load_const None - return_value - - propagate: - reraise 0 - """ - ) - -elif PY >= (3, 10): - GENERATOR_ASSEMBLY.parse( - r""" - setup_finally @stopiter - dup_top - store_fast $__ddgen - load_attr $send - store_fast $__ddgensend - load_const next - load_fast $__ddgen - - loop: - call_function 1 - - yield: - setup_finally @genexit - yield_value - pop_block - load_fast $__ddgensend - rot_two - jump_absolute @loop - - genexit: - dup_top - load_const GeneratorExit - jump_if_not_exc_match @exc - pop_top - pop_top - pop_top - pop_top - load_fast $__ddgen - load_attr $close - call_function 0 - return_value - - exc: - pop_top - pop_top - pop_top - pop_top - load_fast $__ddgen - load_attr $throw - load_const sys.exc_info - call_function 0 - call_function_ex 0 - rot_four - pop_except - jump_absolute @yield - - stopiter: - dup_top - load_const StopIteration - jump_if_not_exc_match @propagate - pop_top - pop_top - pop_top - pop_except - load_const None - return_value - - propagate: - reraise 0 - """ - ) - -elif PY >= (3, 9): - GENERATOR_ASSEMBLY.parse( - r""" - setup_finally @stopiter - dup_top - store_fast $__ddgen - load_attr $send - store_fast $__ddgensend - load_const next - load_fast $__ddgen - - loop: - call_function 1 - - yield: - setup_finally @genexit - yield_value - pop_block - load_fast $__ddgensend - rot_two - jump_absolute @loop - - genexit: - dup_top - load_const GeneratorExit - jump_if_not_exc_match @exc - pop_top - pop_top - pop_top - pop_top - load_fast $__ddgen - load_attr $close - call_function 0 - return_value - - exc: - pop_top - pop_top - pop_top - pop_top - load_fast $__ddgen - load_attr $throw - load_const sys.exc_info - call_function 0 - call_function_ex 0 - rot_four - pop_except - jump_absolute @yield - - stopiter: - dup_top - load_const StopIteration - jump_if_not_exc_match @propagate - pop_top - pop_top - pop_top - pop_except - load_const None - return_value - - propagate: - reraise - """ - ) - -else: - msg = "No generator wrapping support for Python %d.%d" % PY[:2] - raise RuntimeError(msg) - - -def wrap_generator(instrs: bc.Bytecode, code: CodeType, lineno: int) -> None: - if GENERATOR_HEAD_ASSEMBLY is not None: - instrs[0:0] = GENERATOR_HEAD_ASSEMBLY.bind(lineno=lineno) - - instrs[-1:] = GENERATOR_ASSEMBLY.bind(lineno=lineno) +from ddtrace.internal.utils.wrapping.generators import * # noqa diff --git a/riotfile.py b/riotfile.py index 7686782c81b..cc8a8e0d7d6 100644 --- a/riotfile.py +++ b/riotfile.py @@ -595,7 +595,7 @@ def select_pys(min_version: str = MIN_PYTHON_VERSION, max_version: str = MAX_PYT "DD_CIVISIBILITY_ITR_ENABLED": "0", "DD_PYTEST_USE_NEW_PLUGIN": "false", }, - command="pytest -v -n auto {cmdargs} tests/internal/", + command="pytest -v --no-cov -n auto {cmdargs} tests/internal/", pkgs={ "httpretty": latest, "gevent": latest, From 2f28358d5da00eb791df84c0c09d72dbd4f96b9a Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 10:39:12 -0700 Subject: [PATCH 26/29] coverage --- riotfile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/riotfile.py b/riotfile.py index cc8a8e0d7d6..7686782c81b 100644 --- a/riotfile.py +++ b/riotfile.py @@ -595,7 +595,7 @@ def select_pys(min_version: str = MIN_PYTHON_VERSION, max_version: str = MAX_PYT "DD_CIVISIBILITY_ITR_ENABLED": "0", "DD_PYTEST_USE_NEW_PLUGIN": "false", }, - command="pytest -v --no-cov -n auto {cmdargs} tests/internal/", + command="pytest -v -n auto {cmdargs} tests/internal/", pkgs={ "httpretty": latest, "gevent": latest, From 18c76a4cc20e53a381123003c19ffa2632bc7d4e Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 10:39:32 -0700 Subject: [PATCH 27/29] deduplicate threads.py --- ddtrace/internal/utils/threads.py | 221 ------------------------------ 1 file changed, 221 deletions(-) delete mode 100644 ddtrace/internal/utils/threads.py diff --git a/ddtrace/internal/utils/threads.py b/ddtrace/internal/utils/threads.py deleted file mode 100644 index 80c283cedaf..00000000000 --- a/ddtrace/internal/utils/threads.py +++ /dev/null @@ -1,221 +0,0 @@ -from time import monotonic_ns -import typing as t - -from ddtrace.internal import forksafe -from ddtrace.internal._threads import PeriodicThread as _PeriodicThread -from ddtrace.internal._threads import periodic_threads -from ddtrace.internal.utils.logger import get_logger - - -log = get_logger(__name__) - -# We try to import the stdlib locks from the _thread module, where they are -# implemented in C for CPython for most platforms. If that fails, we fall back -# to the threading module, which provides a pure Python implementation that -# should work on all platforms. We also make sure to grab a reference to the -# original lock classes, in case they get patched by monkey-patching libraries -# like gevent. -try: - from _thread import allocate_lock as Lock -except ImportError: - from threading import Lock - -try: - from _thread import RLock -except ImportError: - from threading import RLock - - -__all__ = [ - "Lock", - "PeriodicThread", - "RLock", -] - - -# Forking state management. This is a barrier to either prevent new threads -# from being started while forking, or to allow a thread to be started -# completely if a fork comes in the middle of it. -_forking = False -_forking_lock = Lock() - - -class BoundMethod(t.Protocol): - __self__: t.Any - - def __call__(self) -> None: ... - - -# List of threads that have requested to be started while forking. These will -# be started after the fork is complete. -_threads_to_start_after_fork: list[BoundMethod] = [] - - -def _safe_restart(start: t.Callable[[], None], name: t.Optional[str] = None) -> None: - """Invoke a post-fork thread-start callable, logging resource errors instead of raising. - - The native layer translates pthread_create failures (EAGAIN, ENOMEM) into - OSError. Post-fork restart is triggered automatically by forksafe hooks — - there is no explicit caller that can handle the error, so losing a - periodic thread to resource exhaustion must not crash the host. - Explicit start() calls let OSError propagate so the caller can react. - """ - try: - start() - except Exception as e: - log.error("failed to start periodic thread %s: %s", name, e) - - -class PeriodicThread(_PeriodicThread): - """A fork-safe periodic thread.""" - - __autorestart__ = True - - def start(self) -> None: - with _forking_lock: - # We cannot start a new thread while we are forking, because we are - # trying to stop them all. In that case, we take note of the thread - # and start it after the fork. - if not _forking: - super().start() - else: - _threads_to_start_after_fork.append(t.cast(BoundMethod, super().start)) - - -# Set of running periodic threads that need to be restarted after a fork. -_threads_to_restart_after_fork: set[_PeriodicThread] = set() - - -# A typical scenario is that of forking worker threads in a loop. For the -# parent process, this would mean having to stop and restart the threads in -# between forks, which is not ideal. Instead, we can use a timer to restart -# the threads after a certain amount of time has passed since the last fork. -# This way, we can avoid stopping and restarting the threads in between forks. -class ThreadRestartTimer(PeriodicThread): - __timeout__ = int(1e8) # nanoseconds - - _instance: t.Optional["ThreadRestartTimer"] = None - _timestamp = 0 - - def __init__(self): - super().__init__(self.__timeout__ / 1e9, self._restart_threads, name=f"{__name__}:{self.__class__.__name__}") - - def _restart_threads(self) -> None: - # Restart the threads after we have stopped calling fork for a while. - with _forking_lock: - # If we are forking, we will try again later. - if _forking: - return - - # If we haven't have calls to fork for a while, we can restart the - # threads. This way we avoid stopping and restarting the threads - # in between forks. - if monotonic_ns() >= self._timestamp: # 100ms - for thread in _threads_to_restart_after_fork.copy(): - if isinstance(thread, ThreadRestartTimer): - # Skip any ThreadRestartTimer instance, - # to avoid restarting orphaned timer instances that were - # caught in periodic_threads during a fork. - continue - log.debug("Restarting thread %s after fork", thread.name) - try: - thread._after_fork(force=True) - except Exception as e: - log.error("failed to restart periodic thread %s after fork: %s", thread.name, e) - _threads_to_restart_after_fork.clear() - - for thread_start in _threads_to_start_after_fork: - log.debug("Starting thread %s after fork", thread_start.__self__.name) - _safe_restart(thread_start, thread_start.__self__.name) - _threads_to_start_after_fork.clear() - - # We no longer need this thread so we clear it. - self.clear() - - @classmethod - def clear(cls): - """Clear the timer and stop it if it is running.""" - if cls._instance is not None: - cls._instance.stop() - cls._instance = None - - @classmethod - def touch(cls): - """Set the new expiration time for the timer.""" - cls._timestamp = monotonic_ns() + cls.__timeout__ - - @classmethod - def set(cls): - """Set the timer to restart the threads after a fork.""" - if cls._instance is None: - cls._instance = cls() - cls._instance.start() - else: - # We have already created the timer, so we let the forksafe logic - # handle the restart instead of creating a new instance. - cls._instance._after_fork() - - -@forksafe.register -def _after_fork_child(): - global _forking - - _forking = False - - # Restart the threads immediately. It is unlikely that there will be another - # call to fork here. _after_fork() (without force=True) respects - # __autorestart__: cleanup always runs, but the thread is only restarted - # when __autorestart__ is True. This is intentional in the child — threads - # with __autorestart__ = False (e.g. RemoteConfigPoller) should not run in - # forked workers. - for thread in _threads_to_restart_after_fork.copy(): - log.debug("Restarting thread %s after fork in child", thread.name) - try: - thread._after_fork() - except Exception as e: - log.error("failed to restart periodic thread %s after fork in child: %s", thread.name, e) - _threads_to_restart_after_fork.clear() - - for thread_start in _threads_to_start_after_fork.copy(): - log.debug("Starting thread %s after fork in child", thread_start.__self__.name) - _safe_restart(thread_start, thread_start.__self__.name) - _threads_to_start_after_fork.clear() - - -@forksafe.register_after_parent -def _after_fork_parent() -> None: - global _forking - - _forking = False - - if _threads_to_restart_after_fork or _threads_to_start_after_fork: - ThreadRestartTimer.set() - - -@forksafe.register_before_fork -def _before_fork() -> None: - global _threads_to_restart_after_fork, _forking_lock, _forking - - ThreadRestartTimer.touch() - - with _forking_lock: - _forking = True - - # Take note of all the periodic threads that are running and will need to be - # restarted. - _threads_to_restart_after_fork.update(periodic_threads.values()) - - # Stop all the periodic threads that are still running, without executing - # the shutdown methods, if any. This ensures that we can stop the threads - # more promptly. - for thread in _threads_to_restart_after_fork: - log.debug("Stopping thread %s before fork", thread.name) - thread._before_fork() - - # Join all the threads to ensure they are stopped before the fork. - for thread in _threads_to_restart_after_fork: - log.debug("Joining thread %s before fork", thread.name) - thread.join() - - -__all__ = list(locals().keys()) From 3f8d9e08aadb77fabe96007731af7ac76c5ad5a6 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Wed, 8 Jul 2026 10:47:34 -0700 Subject: [PATCH 28/29] undo imports --- ddtrace/internal/utils/wrapping/__init__.py | 2 +- ddtrace/internal/utils/wrapping/context.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ddtrace/internal/utils/wrapping/__init__.py b/ddtrace/internal/utils/wrapping/__init__.py index 530c824ee72..7df6179ada0 100644 --- a/ddtrace/internal/utils/wrapping/__init__.py +++ b/ddtrace/internal/utils/wrapping/__init__.py @@ -14,7 +14,7 @@ from bytecode import Instr from ddtrace.internal.assembly import Assembly -from ddtrace.internal.utils.threads import Lock +from ddtrace.internal.threads import Lock from ddtrace.internal.utils.wrapping.asyncs import wrap_async from ddtrace.internal.utils.wrapping.generators import wrap_generator diff --git a/ddtrace/internal/utils/wrapping/context.py b/ddtrace/internal/utils/wrapping/context.py index 60d41002008..a20a667530b 100644 --- a/ddtrace/internal/utils/wrapping/context.py +++ b/ddtrace/internal/utils/wrapping/context.py @@ -14,9 +14,9 @@ from bytecode import Bytecode from ddtrace.internal.assembly import Assembly +from ddtrace.internal.threads import Lock +from ddtrace.internal.threads import RLock from ddtrace.internal.utils.logger import get_logger -from ddtrace.internal.utils.threads import Lock -from ddtrace.internal.utils.threads import RLock from ddtrace.internal.utils.wrapping import WrappedFunction from ddtrace.internal.utils.wrapping import Wrapper from ddtrace.internal.utils.wrapping import get_function_code From 1b7265c9c5b92e3ecc8408e52c26327ec4768a07 Mon Sep 17 00:00:00 2001 From: Emmett Butler Date: Thu, 9 Jul 2026 08:20:29 -0700 Subject: [PATCH 29/29] revert files that need more thought --- ddtrace/internal/_exceptions.py | 38 +++- ddtrace/internal/constants.py | 160 ++++++++++++++++- ddtrace/internal/schema/__init__.py | 65 ++++++- ddtrace/internal/schema/processor.py | 26 ++- .../internal/schema/span_attribute_schema.py | 123 ++++++++++++- ddtrace/internal/utils/_exceptions.py | 40 ----- ddtrace/internal/utils/constants.py | 162 ------------------ ddtrace/internal/utils/schema/__init__.py | 63 ------- ddtrace/internal/utils/schema/processor.py | 28 --- .../utils/schema/span_attribute_schema.py | 122 ------------- 10 files changed, 407 insertions(+), 420 deletions(-) delete mode 100644 ddtrace/internal/utils/_exceptions.py delete mode 100644 ddtrace/internal/utils/constants.py delete mode 100644 ddtrace/internal/utils/schema/__init__.py delete mode 100644 ddtrace/internal/utils/schema/processor.py delete mode 100644 ddtrace/internal/utils/schema/span_attribute_schema.py diff --git a/ddtrace/internal/_exceptions.py b/ddtrace/internal/_exceptions.py index 73de4a62a71..e2836e036be 100644 --- a/ddtrace/internal/_exceptions.py +++ b/ddtrace/internal/_exceptions.py @@ -1 +1,37 @@ -from ddtrace.internal.utils._exceptions import * # noqa +from typing import Optional +from typing import TypeVar + + +class DDBlockException(BaseException): + """ + Base class for any in-tree decision to abort the current operation + (web request blocking, AI Guard policy abort, future product blocks). + Inherits from BaseException so a generic ``except Exception:`` handler in + user code does not accidentally swallow a blocking decision. + """ + + +class BlockingException(DDBlockException): + """ + Exception raised when a request is blocked by ASM + It derives from BaseException (via DDBlockException) to avoid being caught by the general Exception handler + """ + + +E = TypeVar("E", bound=BaseException) + + +def find_exception( + exc: BaseException, + exception_type: type[E], +) -> Optional[E]: + """Traverse an exception and its children to find the first occurrence of a specific exception type.""" + if isinstance(exc, exception_type): + return exc + # The check matches both native Python3.11+ and `exceptiongroup` compatibility package versions of ExceptionGroup + if exc.__class__.__name__ in ("BaseExceptionGroup", "ExceptionGroup") and hasattr(exc, "exceptions"): + for sub_exc in exc.exceptions: + found = find_exception(sub_exc, exception_type) + if found: + return found + return None diff --git a/ddtrace/internal/constants.py b/ddtrace/internal/constants.py index 14a047f6b24..1167f75f9e6 100644 --- a/ddtrace/internal/constants.py +++ b/ddtrace/internal/constants.py @@ -1 +1,159 @@ -from ddtrace.internal.utils.constants import * # noqa +from ddtrace.constants import AUTO_KEEP +from ddtrace.constants import AUTO_REJECT +from ddtrace.constants import USER_KEEP +from ddtrace.constants import USER_REJECT + + +PROPAGATION_STYLE_DATADOG = "datadog" +PROPAGATION_STYLE_B3_MULTI = "b3multi" +PROPAGATION_STYLE_B3_SINGLE = "b3" +_PROPAGATION_STYLE_W3C_TRACECONTEXT = "tracecontext" +_PROPAGATION_STYLE_NONE = "none" +_PROPAGATION_STYLE_DEFAULT = "datadog,tracecontext,baggage" +_PROPAGATION_STYLE_BAGGAGE = "baggage" +PROPAGATION_STYLE_ALL = ( + _PROPAGATION_STYLE_W3C_TRACECONTEXT, + PROPAGATION_STYLE_DATADOG, + PROPAGATION_STYLE_B3_MULTI, + PROPAGATION_STYLE_B3_SINGLE, + _PROPAGATION_STYLE_NONE, + _PROPAGATION_STYLE_BAGGAGE, +) +_PROPAGATION_BEHAVIOR_CONTINUE = "continue" +_PROPAGATION_BEHAVIOR_IGNORE = "ignore" +_PROPAGATION_BEHAVIOR_RESTART = "restart" +_PROPAGATION_BEHAVIOR_DEFAULT = _PROPAGATION_BEHAVIOR_CONTINUE +W3C_TRACESTATE_KEY = "tracestate" +W3C_TRACEPARENT_KEY = "traceparent" +W3C_TRACESTATE_PARENT_ID_KEY = "p" +W3C_TRACESTATE_ORIGIN_KEY = "o" +W3C_TRACESTATE_SAMPLING_PRIORITY_KEY = "s" +DEFAULT_SAMPLING_RATE_LIMIT = 100 +SAMPLING_HASH_MODULO = 1 << 64 +# Big prime number to make hashing better distributed, it has to be the same factor as the Agent +# and other tracers to allow chained sampling +SAMPLING_KNUTH_FACTOR = 1111111111111111111 +SAMPLING_DECISION_TRACE_TAG_KEY = "_dd.p.dm" +LAST_DD_PARENT_ID_KEY = "_dd.parent_id" +DEFAULT_SERVICE_NAME = "unnamed-python-service" +# Used to set the name of an integration on a span +COMPONENT = "component" +HIGHER_ORDER_TRACE_ID_BITS = "_dd.p.tid" +MAX_UINT_64BITS = (1 << 64) - 1 +MIN_INT_64BITS = -(2**63) +MAX_INT_64BITS = 2**63 - 1 +SAMPLING_DECISION_MAKER_INHERITED = "_dd.dm.inherited" +SAMPLING_DECISION_MAKER_SERVICE = "_dd.dm.service" +SAMPLING_DECISION_MAKER_RESOURCE = "_dd.dm.resource" +SPAN_LINK_KIND = "dd.kind" +SPAN_LINKS_KEY = "_dd.span_links" +SPAN_EVENTS_KEY = "events" +SPAN_API_DATADOG = "datadog" +SPAN_API_OTEL = "otel" +SPAN_API_OPENTRACING = "opentracing" +DEFAULT_BUFFER_SIZE = 20 << 20 # 20 MB +DEFAULT_MAX_PAYLOAD_SIZE = 20 << 20 # 20 MB +DEFAULT_PROCESSING_INTERVAL = 1.0 +DEFAULT_REUSE_CONNECTIONS = False +BLOCKED_RESPONSE_HTML = """You've been blocked

Sorry, you cannot access this page. Please contact the customer service team.

Security Response ID: [security_response_id]

""" # noqa: E501 +BLOCKED_RESPONSE_JSON = """{"errors":[{"title":"You've been blocked","detail":"Sorry, you cannot access this page. Please contact the customer service team. Security provided by Datadog."}],"security_response_id":"[security_response_id]"}""" # noqa: E501 +HTTP_REQUEST_BLOCKED = "http.request.blocked" +RESPONSE_HEADERS = "http.response.headers" +HTTP_REQUEST_QUERY = "http.request.query" +HTTP_REQUEST_COOKIE_VALUE = "http.request.cookie.value" +HTTP_REQUEST_COOKIE_NAME = "http.request.cookie.name" +HTTP_REQUEST_PATH = "http.request.path" +HTTP_REQUEST_HEADER_NAME = "http.request.header.name" +HTTP_REQUEST_HEADER = "http.request.header" +HTTP_REQUEST_PARAMETER = "http.request.parameter" +HTTP_REQUEST_BODY = "http.request.body" +HTTP_REQUEST_UPGRADED = "http.upgraded" +HTTP_REQUEST_PATH_PARAMETER = "http.request.path.parameter" +REQUEST_PATH_PARAMS = "http.request.path_params" +STATUS_403_TYPE_AUTO = {"status_code": 403, "type": "auto"} +PROCESS_TAGS = "_dd.tags.process" +PROPAGATED_HASH = "_dd.propagated_hash" +_SERVICE_SOURCE = "_dd.svc_src" + +CONTAINER_ID_HEADER_NAME = "Datadog-Container-Id" +CONTAINER_TAGS_HASH = "Datadog-Container-Tags-Hash" + +ENTITY_ID_HEADER_NAME = "Datadog-Entity-ID" + +EXTERNAL_ENV_HEADER_NAME = "Datadog-External-Env" +EXTERNAL_ENV_ENVIRONMENT_VARIABLE = "DD_EXTERNAL_ENV" + +MESSAGING_BATCH_COUNT = "messaging.batch_count" +MESSAGING_DESTINATION_NAME = "messaging.destination.name" +MESSAGING_MESSAGE_ID = "messaging.message_id" +MESSAGING_OPERATION = "messaging.operation" +MESSAGING_SYSTEM = "messaging.system" + +USER_AGENT_HEADER = "user-agent" +FLASK_ENDPOINT = "flask.endpoint" +FLASK_VIEW_ARGS = "flask.view_args" +FLASK_URL_RULE = "flask.url_rule" +FLASK_RESOURCE_FULL = "flask.resource.full" + +_HTTPLIB_NO_TRACE_REQUEST = "_dd_no_trace" +DEFAULT_TIMEOUT = 2.0 + +# baggage +DD_TRACE_BAGGAGE_MAX_ITEMS = 64 +DD_TRACE_BAGGAGE_MAX_BYTES = 8192 +BAGGAGE_TAG_PREFIX = "baggage." + +# W3C Trace Context tracestate (https://www.w3.org/TR/trace-context/): +# max 32 list-members; vendors SHOULD propagate at most 512 characters (we cap parsing to that size). +DD_TRACE_TRACESTATE_MAX_ITEMS = 32 +DD_TRACE_TRACESTATE_MAX_BYTES = 512 +# Per W3C Trace Context, oversized list-members are preferred targets when truncating by size. +DD_TRACE_TRACESTATE_ITEM_MAX_CHARS = 128 + +SPAN_EVENTS_HAS_EXCEPTION = "_dd.span_events.has_exception" +COLLECTOR_MAX_SIZE_PER_SPAN = 100 + +LOG_ATTR_TRACE_ID = "dd.trace_id" +LOG_ATTR_SPAN_ID = "dd.span_id" +LOG_ATTR_ENV = "dd.env" +LOG_ATTR_VERSION = "dd.version" +LOG_ATTR_SERVICE = "dd.service" +LOG_ATTR_VALUE_ZERO = "0" +LOG_ATTR_VALUE_EMPTY = "" + + +class SamplingMechanism(object): + DEFAULT = 0 + AGENT_RATE_BY_SERVICE = 1 + REMOTE_RATE = 2 # not used, this mechanism is deprecated + LOCAL_USER_TRACE_SAMPLING_RULE = 3 + MANUAL = 4 + APPSEC = 5 + REMOTE_RATE_USER = 6 # not used, this mechanism is deprecated + REMOTE_RATE_DATADOG = 7 # not used, this mechanism is deprecated + SPAN_SAMPLING_RULE = 8 + OTLP_INGEST_PROBABILISTIC_SAMPLING = 9 # not used in ddtrace + DATA_JOBS_MONITORING = 10 # not used in ddtrace + REMOTE_USER_TRACE_SAMPLING_RULE = 11 + REMOTE_DYNAMIC_TRACE_SAMPLING_RULE = 12 + AI_GUARD = 13 + + +SAMPLING_MECHANISM_TO_PRIORITIES = { + # TODO(munir): Update mapping to include single span sampling and appsec sampling mechanisms + SamplingMechanism.AGENT_RATE_BY_SERVICE: (AUTO_KEEP, AUTO_REJECT), + SamplingMechanism.DEFAULT: (AUTO_KEEP, AUTO_REJECT), + SamplingMechanism.MANUAL: (USER_KEEP, USER_REJECT), + SamplingMechanism.APPSEC: (AUTO_KEEP, AUTO_REJECT), + SamplingMechanism.LOCAL_USER_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), + SamplingMechanism.REMOTE_USER_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), + SamplingMechanism.REMOTE_DYNAMIC_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), +} +_KEEP_PRIORITY_INDEX = 0 +_REJECT_PRIORITY_INDEX = 1 + + +# List of support values in DD_TRACE_EXPERIMENTAL_FEATURES_ENABLED +class EXPERIMENTAL_FEATURES: + # Enables submitting runtime metrics as gauges (instead of distributions) + RUNTIME_METRICS = "DD_RUNTIME_METRICS_ENABLED" diff --git a/ddtrace/internal/schema/__init__.py b/ddtrace/internal/schema/__init__.py index 9ba7325082f..db7226a42ae 100644 --- a/ddtrace/internal/schema/__init__.py +++ b/ddtrace/internal/schema/__init__.py @@ -1 +1,64 @@ -from ddtrace.internal.utils.schema import * # noqa +import logging + +from ddtrace.internal.settings import env +from ddtrace.internal.utils.formats import asbool + +from .span_attribute_schema import _DEFAULT_SPAN_SERVICE_NAMES +from .span_attribute_schema import _SPAN_ATTRIBUTE_TO_FUNCTION +from .span_attribute_schema import SpanDirection + + +log = logging.getLogger(__name__) + + +# Span attribute schema +def _validate_schema(version): + error_message = ( + "You have specified an invalid span attribute schema version: '{}'.".format(version), + "Valid options are: {}. You can change the specified value by updating".format( + _SPAN_ATTRIBUTE_TO_FUNCTION.keys() + ), + "the value exported in the 'DD_TRACE_SPAN_ATTRIBUTE_SCHEMA' environment variable.", + ) + + if version not in _SPAN_ATTRIBUTE_TO_FUNCTION.keys(): + log.warning(" ".join(error_message)) + return False + + return True + + +def _get_schema_version(): + version = env.get("DD_TRACE_SPAN_ATTRIBUTE_SCHEMA", default="v0") + if not _validate_schema(version): + version = "v0" + return version + + +SCHEMA_VERSION = _get_schema_version() +_remove_client_service_names = asbool(env.get("DD_TRACE_REMOVE_INTEGRATION_SERVICE_NAMES_ENABLED", default=False)) +_service_name_schema_version = "v0" if SCHEMA_VERSION == "v0" and not _remove_client_service_names else "v1" + +DEFAULT_SPAN_SERVICE_NAME = _DEFAULT_SPAN_SERVICE_NAMES[_service_name_schema_version] +schematize_cache_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["cache_operation"] +schematize_cloud_api_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["cloud_api_operation"] +schematize_cloud_faas_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["cloud_faas_operation"] +schematize_cloud_messaging_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["cloud_messaging_operation"] +schematize_database_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["database_operation"] +schematize_messaging_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["messaging_operation"] +schematize_service_name = _SPAN_ATTRIBUTE_TO_FUNCTION[_service_name_schema_version]["service_name"] +schematize_url_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["url_operation"] + +__all__ = [ + "DEFAULT_SPAN_SERVICE_NAME", + "SCHEMA_VERSION", + "SpanDirection", + "schematize_cache_operation", + "schematize_cloud_api_operation", + "schematize_cloud_faas_operation", + "schematize_cloud_messaging_operation", + "schematize_database_operation", + "schematize_messaging_operation", + "schematize_service_name", + "schematize_url_operation", +] diff --git a/ddtrace/internal/schema/processor.py b/ddtrace/internal/schema/processor.py index dacde93c163..8d5559ecae3 100644 --- a/ddtrace/internal/schema/processor.py +++ b/ddtrace/internal/schema/processor.py @@ -1 +1,25 @@ -from ddtrace.internal.utils.schema.processor import * # noqa +from ddtrace._trace.processor import TraceProcessor +from ddtrace.constants import _BASE_SERVICE_KEY +from ddtrace.internal.settings._config import config + +from . import schematize_service_name + + +class BaseServiceProcessor(TraceProcessor): + def __init__(self): + self._global_service = schematize_service_name((config.service or "").lower()) + + def process_trace(self, trace): + if not trace: + return trace + + traces_to_process = filter( + lambda x: x.service and x.service.lower() != self._global_service, + trace, + ) + any(map(lambda x: self._update_dd_base_service(x), traces_to_process)) + + return trace + + def _update_dd_base_service(self, span): + span._set_attribute(key=_BASE_SERVICE_KEY, value=self._global_service) diff --git a/ddtrace/internal/schema/span_attribute_schema.py b/ddtrace/internal/schema/span_attribute_schema.py index 85a4cc8b8c9..c9b717f6809 100644 --- a/ddtrace/internal/schema/span_attribute_schema.py +++ b/ddtrace/internal/schema/span_attribute_schema.py @@ -1 +1,122 @@ -from ddtrace.internal.utils.schema.span_attribute_schema import * # noqa +from enum import Enum +import sys +from typing import Optional + +from ddtrace.internal.constants import DEFAULT_SERVICE_NAME +from ddtrace.internal.settings._inferred_base_service import detect_service + + +class SpanDirection(Enum): + INBOUND = "inbound" + OUTBOUND = "outbound" + PROCESSING = "processing" + + +def service_name_v0(v0_service_name): + return v0_service_name + + +def service_name_v1(*_, **__): + from ddtrace import config as dd_config + + return dd_config.service + + +def database_operation_v0(v0_operation, database_provider=None): + return v0_operation + + +def database_operation_v1(v0_operation, database_provider=None): + operation = "query" + return "{}.{}".format(database_provider, operation) + + +def cache_operation_v0(v0_operation, cache_provider=None): + return v0_operation + + +def cache_operation_v1(v0_operation, cache_provider=None): + operation = "command" + return "{}.{}".format(cache_provider, operation) + + +def cloud_api_operation_v0(v0_operation, cloud_provider=None, cloud_service=None): + return v0_operation + + +def cloud_api_operation_v1(v0_operation, cloud_provider=None, cloud_service=None): + return "{}.{}.request".format(cloud_provider, cloud_service) + + +def cloud_faas_operation_v0(v0_operation, cloud_provider=None, cloud_service=None): + return v0_operation + + +def cloud_faas_operation_v1(v0_operation, cloud_provider=None, cloud_service=None): + return "{}.{}.invoke".format(cloud_provider, cloud_service) + + +def cloud_messaging_operation_v0(v0_operation, cloud_provider=None, cloud_service=None, direction=None): + return v0_operation + + +def cloud_messaging_operation_v1(v0_operation, cloud_provider=None, cloud_service=None, direction=None): + if direction == SpanDirection.INBOUND: + return "{}.{}.receive".format(cloud_provider, cloud_service) + elif direction == SpanDirection.OUTBOUND: + return "{}.{}.send".format(cloud_provider, cloud_service) + elif direction == SpanDirection.PROCESSING: + return "{}.{}.process".format(cloud_provider, cloud_service) + + +def messaging_operation_v0(v0_operation, provider=None, service=None, direction=None): + return v0_operation + + +def messaging_operation_v1(v0_operation, provider=None, direction=None): + if direction == SpanDirection.INBOUND: + return "{}.receive".format(provider) + elif direction == SpanDirection.OUTBOUND: + return "{}.send".format(provider) + elif direction == SpanDirection.PROCESSING: + return "{}.process".format(provider) + + +def url_operation_v0(v0_operation, protocol=None, direction=None): + return v0_operation + + +def url_operation_v1(v0_operation, protocol=None, direction=None): + server_or_client = {SpanDirection.INBOUND: "server", SpanDirection.OUTBOUND: "client"}[direction] + return "{}.{}.request".format(protocol, server_or_client) + + +_SPAN_ATTRIBUTE_TO_FUNCTION = { + "v0": { + "cache_operation": cache_operation_v0, + "cloud_api_operation": cloud_api_operation_v0, + "cloud_faas_operation": cloud_faas_operation_v0, + "cloud_messaging_operation": cloud_messaging_operation_v0, + "database_operation": database_operation_v0, + "messaging_operation": messaging_operation_v0, + "service_name": service_name_v0, + "url_operation": url_operation_v0, + }, + "v1": { + "cache_operation": cache_operation_v1, + "cloud_api_operation": cloud_api_operation_v1, + "cloud_faas_operation": cloud_faas_operation_v1, + "cloud_messaging_operation": cloud_messaging_operation_v1, + "database_operation": database_operation_v1, + "messaging_operation": messaging_operation_v1, + "service_name": service_name_v1, + "url_operation": url_operation_v1, + }, +} + +_inferred_base_service: Optional[str] = detect_service(sys.argv) + +_DEFAULT_SPAN_SERVICE_NAMES: dict[str, Optional[str]] = { + "v0": _inferred_base_service or None, + "v1": _inferred_base_service or DEFAULT_SERVICE_NAME, +} diff --git a/ddtrace/internal/utils/_exceptions.py b/ddtrace/internal/utils/_exceptions.py deleted file mode 100644 index c69546158c8..00000000000 --- a/ddtrace/internal/utils/_exceptions.py +++ /dev/null @@ -1,40 +0,0 @@ -from typing import Optional -from typing import TypeVar - - -class DDBlockException(BaseException): - """ - Base class for any in-tree decision to abort the current operation - (web request blocking, AI Guard policy abort, future product blocks). - Inherits from BaseException so a generic ``except Exception:`` handler in - user code does not accidentally swallow a blocking decision. - """ - - -class BlockingException(DDBlockException): - """ - Exception raised when a request is blocked by ASM - It derives from BaseException (via DDBlockException) to avoid being caught by the general Exception handler - """ - - -E = TypeVar("E", bound=BaseException) - - -def find_exception( - exc: BaseException, - exception_type: type[E], -) -> Optional[E]: - """Traverse an exception and its children to find the first occurrence of a specific exception type.""" - if isinstance(exc, exception_type): - return exc - # The check matches both native Python3.11+ and `exceptiongroup` compatibility package versions of ExceptionGroup - if exc.__class__.__name__ in ("BaseExceptionGroup", "ExceptionGroup") and hasattr(exc, "exceptions"): - for sub_exc in exc.exceptions: - found = find_exception(sub_exc, exception_type) - if found: - return found - return None - - -__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/constants.py b/ddtrace/internal/utils/constants.py deleted file mode 100644 index 36472603864..00000000000 --- a/ddtrace/internal/utils/constants.py +++ /dev/null @@ -1,162 +0,0 @@ -from ddtrace.constants import AUTO_KEEP -from ddtrace.constants import AUTO_REJECT -from ddtrace.constants import USER_KEEP -from ddtrace.constants import USER_REJECT - - -PROPAGATION_STYLE_DATADOG = "datadog" -PROPAGATION_STYLE_B3_MULTI = "b3multi" -PROPAGATION_STYLE_B3_SINGLE = "b3" -_PROPAGATION_STYLE_W3C_TRACECONTEXT = "tracecontext" -_PROPAGATION_STYLE_NONE = "none" -_PROPAGATION_STYLE_DEFAULT = "datadog,tracecontext,baggage" -_PROPAGATION_STYLE_BAGGAGE = "baggage" -PROPAGATION_STYLE_ALL = ( - _PROPAGATION_STYLE_W3C_TRACECONTEXT, - PROPAGATION_STYLE_DATADOG, - PROPAGATION_STYLE_B3_MULTI, - PROPAGATION_STYLE_B3_SINGLE, - _PROPAGATION_STYLE_NONE, - _PROPAGATION_STYLE_BAGGAGE, -) -_PROPAGATION_BEHAVIOR_CONTINUE = "continue" -_PROPAGATION_BEHAVIOR_IGNORE = "ignore" -_PROPAGATION_BEHAVIOR_RESTART = "restart" -_PROPAGATION_BEHAVIOR_DEFAULT = _PROPAGATION_BEHAVIOR_CONTINUE -W3C_TRACESTATE_KEY = "tracestate" -W3C_TRACEPARENT_KEY = "traceparent" -W3C_TRACESTATE_PARENT_ID_KEY = "p" -W3C_TRACESTATE_ORIGIN_KEY = "o" -W3C_TRACESTATE_SAMPLING_PRIORITY_KEY = "s" -DEFAULT_SAMPLING_RATE_LIMIT = 100 -SAMPLING_HASH_MODULO = 1 << 64 -# Big prime number to make hashing better distributed, it has to be the same factor as the Agent -# and other tracers to allow chained sampling -SAMPLING_KNUTH_FACTOR = 1111111111111111111 -SAMPLING_DECISION_TRACE_TAG_KEY = "_dd.p.dm" -LAST_DD_PARENT_ID_KEY = "_dd.parent_id" -DEFAULT_SERVICE_NAME = "unnamed-python-service" -# Used to set the name of an integration on a span -COMPONENT = "component" -HIGHER_ORDER_TRACE_ID_BITS = "_dd.p.tid" -MAX_UINT_64BITS = (1 << 64) - 1 -MIN_INT_64BITS = -(2**63) -MAX_INT_64BITS = 2**63 - 1 -SAMPLING_DECISION_MAKER_INHERITED = "_dd.dm.inherited" -SAMPLING_DECISION_MAKER_SERVICE = "_dd.dm.service" -SAMPLING_DECISION_MAKER_RESOURCE = "_dd.dm.resource" -SPAN_LINK_KIND = "dd.kind" -SPAN_LINKS_KEY = "_dd.span_links" -SPAN_EVENTS_KEY = "events" -SPAN_API_DATADOG = "datadog" -SPAN_API_OTEL = "otel" -SPAN_API_OPENTRACING = "opentracing" -DEFAULT_BUFFER_SIZE = 20 << 20 # 20 MB -DEFAULT_MAX_PAYLOAD_SIZE = 20 << 20 # 20 MB -DEFAULT_PROCESSING_INTERVAL = 1.0 -DEFAULT_REUSE_CONNECTIONS = False -BLOCKED_RESPONSE_HTML = """You've been blocked

Sorry, you cannot access this page. Please contact the customer service team.

Security Response ID: [security_response_id]

""" # noqa: E501 -BLOCKED_RESPONSE_JSON = """{"errors":[{"title":"You've been blocked","detail":"Sorry, you cannot access this page. Please contact the customer service team. Security provided by Datadog."}],"security_response_id":"[security_response_id]"}""" # noqa: E501 -HTTP_REQUEST_BLOCKED = "http.request.blocked" -RESPONSE_HEADERS = "http.response.headers" -HTTP_REQUEST_QUERY = "http.request.query" -HTTP_REQUEST_COOKIE_VALUE = "http.request.cookie.value" -HTTP_REQUEST_COOKIE_NAME = "http.request.cookie.name" -HTTP_REQUEST_PATH = "http.request.path" -HTTP_REQUEST_HEADER_NAME = "http.request.header.name" -HTTP_REQUEST_HEADER = "http.request.header" -HTTP_REQUEST_PARAMETER = "http.request.parameter" -HTTP_REQUEST_BODY = "http.request.body" -HTTP_REQUEST_UPGRADED = "http.upgraded" -HTTP_REQUEST_PATH_PARAMETER = "http.request.path.parameter" -REQUEST_PATH_PARAMS = "http.request.path_params" -STATUS_403_TYPE_AUTO = {"status_code": 403, "type": "auto"} -PROCESS_TAGS = "_dd.tags.process" -PROPAGATED_HASH = "_dd.propagated_hash" -_SERVICE_SOURCE = "_dd.svc_src" - -CONTAINER_ID_HEADER_NAME = "Datadog-Container-Id" -CONTAINER_TAGS_HASH = "Datadog-Container-Tags-Hash" - -ENTITY_ID_HEADER_NAME = "Datadog-Entity-ID" - -EXTERNAL_ENV_HEADER_NAME = "Datadog-External-Env" -EXTERNAL_ENV_ENVIRONMENT_VARIABLE = "DD_EXTERNAL_ENV" - -MESSAGING_BATCH_COUNT = "messaging.batch_count" -MESSAGING_DESTINATION_NAME = "messaging.destination.name" -MESSAGING_MESSAGE_ID = "messaging.message_id" -MESSAGING_OPERATION = "messaging.operation" -MESSAGING_SYSTEM = "messaging.system" - -USER_AGENT_HEADER = "user-agent" -FLASK_ENDPOINT = "flask.endpoint" -FLASK_VIEW_ARGS = "flask.view_args" -FLASK_URL_RULE = "flask.url_rule" -FLASK_RESOURCE_FULL = "flask.resource.full" - -_HTTPLIB_NO_TRACE_REQUEST = "_dd_no_trace" -DEFAULT_TIMEOUT = 2.0 - -# baggage -DD_TRACE_BAGGAGE_MAX_ITEMS = 64 -DD_TRACE_BAGGAGE_MAX_BYTES = 8192 -BAGGAGE_TAG_PREFIX = "baggage." - -# W3C Trace Context tracestate (https://www.w3.org/TR/trace-context/): -# max 32 list-members; vendors SHOULD propagate at most 512 characters (we cap parsing to that size). -DD_TRACE_TRACESTATE_MAX_ITEMS = 32 -DD_TRACE_TRACESTATE_MAX_BYTES = 512 -# Per W3C Trace Context, oversized list-members are preferred targets when truncating by size. -DD_TRACE_TRACESTATE_ITEM_MAX_CHARS = 128 - -SPAN_EVENTS_HAS_EXCEPTION = "_dd.span_events.has_exception" -COLLECTOR_MAX_SIZE_PER_SPAN = 100 - -LOG_ATTR_TRACE_ID = "dd.trace_id" -LOG_ATTR_SPAN_ID = "dd.span_id" -LOG_ATTR_ENV = "dd.env" -LOG_ATTR_VERSION = "dd.version" -LOG_ATTR_SERVICE = "dd.service" -LOG_ATTR_VALUE_ZERO = "0" -LOG_ATTR_VALUE_EMPTY = "" - - -class SamplingMechanism(object): - DEFAULT = 0 - AGENT_RATE_BY_SERVICE = 1 - REMOTE_RATE = 2 # not used, this mechanism is deprecated - LOCAL_USER_TRACE_SAMPLING_RULE = 3 - MANUAL = 4 - APPSEC = 5 - REMOTE_RATE_USER = 6 # not used, this mechanism is deprecated - REMOTE_RATE_DATADOG = 7 # not used, this mechanism is deprecated - SPAN_SAMPLING_RULE = 8 - OTLP_INGEST_PROBABILISTIC_SAMPLING = 9 # not used in ddtrace - DATA_JOBS_MONITORING = 10 # not used in ddtrace - REMOTE_USER_TRACE_SAMPLING_RULE = 11 - REMOTE_DYNAMIC_TRACE_SAMPLING_RULE = 12 - AI_GUARD = 13 - - -SAMPLING_MECHANISM_TO_PRIORITIES = { - # TODO(munir): Update mapping to include single span sampling and appsec sampling mechanisms - SamplingMechanism.AGENT_RATE_BY_SERVICE: (AUTO_KEEP, AUTO_REJECT), - SamplingMechanism.DEFAULT: (AUTO_KEEP, AUTO_REJECT), - SamplingMechanism.MANUAL: (USER_KEEP, USER_REJECT), - SamplingMechanism.APPSEC: (AUTO_KEEP, AUTO_REJECT), - SamplingMechanism.LOCAL_USER_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), - SamplingMechanism.REMOTE_USER_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), - SamplingMechanism.REMOTE_DYNAMIC_TRACE_SAMPLING_RULE: (USER_KEEP, USER_REJECT), -} -_KEEP_PRIORITY_INDEX = 0 -_REJECT_PRIORITY_INDEX = 1 - - -# List of support values in DD_TRACE_EXPERIMENTAL_FEATURES_ENABLED -class EXPERIMENTAL_FEATURES: - # Enables submitting runtime metrics as gauges (instead of distributions) - RUNTIME_METRICS = "DD_RUNTIME_METRICS_ENABLED" - - -__all__ = list(locals().keys()) diff --git a/ddtrace/internal/utils/schema/__init__.py b/ddtrace/internal/utils/schema/__init__.py deleted file mode 100644 index 535df404d08..00000000000 --- a/ddtrace/internal/utils/schema/__init__.py +++ /dev/null @@ -1,63 +0,0 @@ -import logging - -from ddtrace.internal.settings import env -from ddtrace.internal.utils.formats import asbool -from ddtrace.internal.utils.schema.span_attribute_schema import _DEFAULT_SPAN_SERVICE_NAMES -from ddtrace.internal.utils.schema.span_attribute_schema import _SPAN_ATTRIBUTE_TO_FUNCTION -from ddtrace.internal.utils.schema.span_attribute_schema import SpanDirection - - -log = logging.getLogger(__name__) - - -# Span attribute schema -def _validate_schema(version): - error_message = ( - "You have specified an invalid span attribute schema version: '{}'.".format(version), - "Valid options are: {}. You can change the specified value by updating".format( - _SPAN_ATTRIBUTE_TO_FUNCTION.keys() - ), - "the value exported in the 'DD_TRACE_SPAN_ATTRIBUTE_SCHEMA' environment variable.", - ) - - if version not in _SPAN_ATTRIBUTE_TO_FUNCTION.keys(): - log.warning(" ".join(error_message)) - return False - - return True - - -def _get_schema_version(): - version = env.get("DD_TRACE_SPAN_ATTRIBUTE_SCHEMA", default="v0") - if not _validate_schema(version): - version = "v0" - return version - - -SCHEMA_VERSION = _get_schema_version() -_remove_client_service_names = asbool(env.get("DD_TRACE_REMOVE_INTEGRATION_SERVICE_NAMES_ENABLED", default=False)) -_service_name_schema_version = "v0" if SCHEMA_VERSION == "v0" and not _remove_client_service_names else "v1" - -DEFAULT_SPAN_SERVICE_NAME = _DEFAULT_SPAN_SERVICE_NAMES[_service_name_schema_version] -schematize_cache_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["cache_operation"] -schematize_cloud_api_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["cloud_api_operation"] -schematize_cloud_faas_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["cloud_faas_operation"] -schematize_cloud_messaging_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["cloud_messaging_operation"] -schematize_database_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["database_operation"] -schematize_messaging_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["messaging_operation"] -schematize_service_name = _SPAN_ATTRIBUTE_TO_FUNCTION[_service_name_schema_version]["service_name"] -schematize_url_operation = _SPAN_ATTRIBUTE_TO_FUNCTION[SCHEMA_VERSION]["url_operation"] - -__all__ = [ - "DEFAULT_SPAN_SERVICE_NAME", - "SCHEMA_VERSION", - "SpanDirection", - "schematize_cache_operation", - "schematize_cloud_api_operation", - "schematize_cloud_faas_operation", - "schematize_cloud_messaging_operation", - "schematize_database_operation", - "schematize_messaging_operation", - "schematize_service_name", - "schematize_url_operation", -] diff --git a/ddtrace/internal/utils/schema/processor.py b/ddtrace/internal/utils/schema/processor.py deleted file mode 100644 index e5145de528f..00000000000 --- a/ddtrace/internal/utils/schema/processor.py +++ /dev/null @@ -1,28 +0,0 @@ -from ddtrace._trace.processor import TraceProcessor -from ddtrace.constants import _BASE_SERVICE_KEY -from ddtrace.internal.settings._config import config - -from . import schematize_service_name - - -class BaseServiceProcessor(TraceProcessor): - def __init__(self): - self._global_service = schematize_service_name((config.service or "").lower()) - - def process_trace(self, trace): - if not trace: - return trace - - traces_to_process = filter( - lambda x: x.service and x.service.lower() != self._global_service, - trace, - ) - any(map(lambda x: self._update_dd_base_service(x), traces_to_process)) - - return trace - - def _update_dd_base_service(self, span): - span._set_attribute(key=_BASE_SERVICE_KEY, value=self._global_service) - - -__all__ = ["BaseServiceProcessor"] diff --git a/ddtrace/internal/utils/schema/span_attribute_schema.py b/ddtrace/internal/utils/schema/span_attribute_schema.py deleted file mode 100644 index c9b717f6809..00000000000 --- a/ddtrace/internal/utils/schema/span_attribute_schema.py +++ /dev/null @@ -1,122 +0,0 @@ -from enum import Enum -import sys -from typing import Optional - -from ddtrace.internal.constants import DEFAULT_SERVICE_NAME -from ddtrace.internal.settings._inferred_base_service import detect_service - - -class SpanDirection(Enum): - INBOUND = "inbound" - OUTBOUND = "outbound" - PROCESSING = "processing" - - -def service_name_v0(v0_service_name): - return v0_service_name - - -def service_name_v1(*_, **__): - from ddtrace import config as dd_config - - return dd_config.service - - -def database_operation_v0(v0_operation, database_provider=None): - return v0_operation - - -def database_operation_v1(v0_operation, database_provider=None): - operation = "query" - return "{}.{}".format(database_provider, operation) - - -def cache_operation_v0(v0_operation, cache_provider=None): - return v0_operation - - -def cache_operation_v1(v0_operation, cache_provider=None): - operation = "command" - return "{}.{}".format(cache_provider, operation) - - -def cloud_api_operation_v0(v0_operation, cloud_provider=None, cloud_service=None): - return v0_operation - - -def cloud_api_operation_v1(v0_operation, cloud_provider=None, cloud_service=None): - return "{}.{}.request".format(cloud_provider, cloud_service) - - -def cloud_faas_operation_v0(v0_operation, cloud_provider=None, cloud_service=None): - return v0_operation - - -def cloud_faas_operation_v1(v0_operation, cloud_provider=None, cloud_service=None): - return "{}.{}.invoke".format(cloud_provider, cloud_service) - - -def cloud_messaging_operation_v0(v0_operation, cloud_provider=None, cloud_service=None, direction=None): - return v0_operation - - -def cloud_messaging_operation_v1(v0_operation, cloud_provider=None, cloud_service=None, direction=None): - if direction == SpanDirection.INBOUND: - return "{}.{}.receive".format(cloud_provider, cloud_service) - elif direction == SpanDirection.OUTBOUND: - return "{}.{}.send".format(cloud_provider, cloud_service) - elif direction == SpanDirection.PROCESSING: - return "{}.{}.process".format(cloud_provider, cloud_service) - - -def messaging_operation_v0(v0_operation, provider=None, service=None, direction=None): - return v0_operation - - -def messaging_operation_v1(v0_operation, provider=None, direction=None): - if direction == SpanDirection.INBOUND: - return "{}.receive".format(provider) - elif direction == SpanDirection.OUTBOUND: - return "{}.send".format(provider) - elif direction == SpanDirection.PROCESSING: - return "{}.process".format(provider) - - -def url_operation_v0(v0_operation, protocol=None, direction=None): - return v0_operation - - -def url_operation_v1(v0_operation, protocol=None, direction=None): - server_or_client = {SpanDirection.INBOUND: "server", SpanDirection.OUTBOUND: "client"}[direction] - return "{}.{}.request".format(protocol, server_or_client) - - -_SPAN_ATTRIBUTE_TO_FUNCTION = { - "v0": { - "cache_operation": cache_operation_v0, - "cloud_api_operation": cloud_api_operation_v0, - "cloud_faas_operation": cloud_faas_operation_v0, - "cloud_messaging_operation": cloud_messaging_operation_v0, - "database_operation": database_operation_v0, - "messaging_operation": messaging_operation_v0, - "service_name": service_name_v0, - "url_operation": url_operation_v0, - }, - "v1": { - "cache_operation": cache_operation_v1, - "cloud_api_operation": cloud_api_operation_v1, - "cloud_faas_operation": cloud_faas_operation_v1, - "cloud_messaging_operation": cloud_messaging_operation_v1, - "database_operation": database_operation_v1, - "messaging_operation": messaging_operation_v1, - "service_name": service_name_v1, - "url_operation": url_operation_v1, - }, -} - -_inferred_base_service: Optional[str] = detect_service(sys.argv) - -_DEFAULT_SPAN_SERVICE_NAMES: dict[str, Optional[str]] = { - "v0": _inferred_base_service or None, - "v1": _inferred_base_service or DEFAULT_SERVICE_NAME, -}