From feb748d975244b1e38dda63d3ff34b18b1e8d7b6 Mon Sep 17 00:00:00 2001 From: "Gabriele N. Tornetta" Date: Wed, 22 Jul 2026 17:31:20 -0400 Subject: [PATCH 01/13] feat(internal): add sys.monitoring multiplexer for Python 3.15 Introduces ddtrace.internal.monitoring for Python 3.15+. Part of the #17849 split (PR 1/6). --- ddtrace/internal/monitoring.py | 275 ++++++++++++++++++++++++++++++ tests/internal/test_monitoring.py | 177 +++++++++++++++++++ 2 files changed, 452 insertions(+) create mode 100644 ddtrace/internal/monitoring.py create mode 100644 tests/internal/test_monitoring.py diff --git a/ddtrace/internal/monitoring.py b/ddtrace/internal/monitoring.py new file mode 100644 index 00000000000..768a7aebb84 --- /dev/null +++ b/ddtrace/internal/monitoring.py @@ -0,0 +1,275 @@ +"""Multiplexed sys.monitoring interface for ddtrace internal use. + +A single sys.monitoring tool ID is shared across all ddtrace sub-systems. +Sub-systems implement :class:`MonitoringEventHandler` and register instances +via :func:`register`; the multiplexer dispatches each monitoring event to all +handlers registered for that code object. + +Only the events corresponding to overridden handler methods are enabled, +so a handler that only overrides ``on_py_start`` pays no cost for the other +two events. + +The handler instance itself serves as the registration key: pass the same +object to :func:`unregister` to remove it. +""" + +from abc import ABC +import sys +from types import CodeType +from typing import NamedTuple +from typing import Optional +import weakref + +from ddtrace.internal.threads import Lock + + +if sys.version_info < (3, 15): + raise ImportError("ddtrace.internal.monitoring requires Python 3.15+") + +_E = sys.monitoring.events # type: ignore[unreachable] +DISABLE = sys.monitoring.DISABLE +_DISABLE = DISABLE + +# sys.monitoring distinguishes "local" events, which can be enabled per code +# object via set_local_events, from events that can only be enabled globally +# via set_events. PY_START, PY_RETURN and LINE are local; PY_UNWIND is not and +# must be enabled globally, otherwise set_local_events raises +# "ValueError: invalid local event set". +_LOCAL_EVENTS = _E.PY_START | _E.PY_RETURN | _E.LINE +_GLOBAL_EVENTS = _E.PY_UNWIND + +_tool_id: Optional[int] = None +_tool_lock = Lock() + +# The set of global-only events currently enabled via sys.monitoring.set_events. +_active_global_events: int = 0 + +_registry: "weakref.WeakKeyDictionary[CodeType, dict[int, _Entry]]" = weakref.WeakKeyDictionary() +_registry_lock = Lock() + + +class MonitoringEventHandler(ABC): + """Base class for sys.monitoring event handlers. + + Override only the methods for events you need. The multiplexer enables + only those events, so un-overridden methods incur no monitoring overhead. + + .. warning:: + Do not call :func:`register` or :func:`unregister` from inside an + event handler method. Doing so mutates the handler list while it is + being iterated, which produces undefined behavior. + """ + + def on_py_start(self, code: CodeType, instruction_offset: int) -> None: + pass + + def on_py_return(self, code: CodeType, instruction_offset: int, retval: object) -> None: + pass + + def on_py_unwind(self, code: CodeType, instruction_offset: int, exception: BaseException) -> None: + pass + + def on_py_line(self, code: CodeType, line_number: int) -> Optional[object]: + """Return ``sys.monitoring.DISABLE`` to stop future events on this line.""" + return None + + +def _events_for_handler(handler: MonitoringEventHandler) -> int: + """Return the OR of events corresponding to overridden handler methods.""" + cls: type[MonitoringEventHandler] = type(handler) + base: type[MonitoringEventHandler] = MonitoringEventHandler + events: int = 0 + if cls.on_py_start is not base.on_py_start: + events |= _E.PY_START + if cls.on_py_return is not base.on_py_return: + events |= _E.PY_RETURN + if cls.on_py_unwind is not base.on_py_unwind: + events |= _E.PY_UNWIND + if cls.on_py_line is not base.on_py_line: + events |= _E.LINE + return events + + +def _events_for(entries: "dict[int, _Entry]") -> int: + events: int = 0 + for e in list(entries.values()): + events |= e.events + return events + + +class _Entry(NamedTuple): + handler: MonitoringEventHandler + events: int # pre-computed from _events_for_handler + + +def _setup() -> int: + """Claim a free tool ID and install the global callbacks (idempotent).""" + global _tool_id + + if _tool_id is not None: + return _tool_id + + with _tool_lock: + if _tool_id is not None: + return _tool_id + + for tid in range(5, -1, -1): + try: + sys.monitoring.use_tool_id(tid, "ddtrace") + _tool_id = tid + break + except ValueError: + continue + else: + raise RuntimeError("No free sys.monitoring tool ID available for ddtrace") + + sys.monitoring.register_callback(_tool_id, _E.PY_START, _on_py_start) + sys.monitoring.register_callback(_tool_id, _E.PY_RETURN, _on_py_return) + sys.monitoring.register_callback(_tool_id, _E.PY_UNWIND, _on_py_unwind) + sys.monitoring.register_callback(_tool_id, _E.LINE, _on_py_line) + + return _tool_id + + +# --------------------------------------------------------------------------- +# Hot-path callbacks — no lock, no allocation +# --------------------------------------------------------------------------- + + +def _on_py_start(code: CodeType, instruction_offset: int) -> Optional[object]: + entries: Optional[dict[int, _Entry]] = _registry.get(code) + if not entries: + return _DISABLE + for e in list(entries.values()): + if e.events & _E.PY_START: + e.handler.on_py_start(code, instruction_offset) + return None + + +def _on_py_return(code: CodeType, instruction_offset: int, retval: object) -> Optional[object]: + entries: Optional[dict[int, _Entry]] = _registry.get(code) + if not entries: + return _DISABLE + for e in list(entries.values()): + if e.events & _E.PY_RETURN: + e.handler.on_py_return(code, instruction_offset, retval) + return None + + +def _on_py_unwind(code: CodeType, instruction_offset: int, exception: BaseException) -> Optional[object]: + entries: Optional[dict[int, _Entry]] = _registry.get(code) + # PY_UNWIND is a global event, so this callback fires for every unwinding + # frame regardless of registration. We must not return DISABLE for + # unregistered code: doing so would permanently disable the event for that + # code location, and a later register() would not re-arm it (we never call + # restart_events). Unwinding only happens on exceptions, so the extra lookup + # cost on this already-slow path is negligible. + if not entries: + return None + for e in list(entries.values()): + if e.events & _E.PY_UNWIND: + e.handler.on_py_unwind(code, instruction_offset, exception) + return None + + +def _on_py_line(code: CodeType, line_number: int) -> Optional[object]: + entries: Optional[dict[int, _Entry]] = _registry.get(code) + if not entries: + return _DISABLE + disable: bool = True + for e in list(entries.values()): + if e.events & _E.LINE: + if e.handler.on_py_line(code, line_number) is not _DISABLE: + disable = False + return _DISABLE if disable else None + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def _enable_global_events(events: int) -> None: + """Ensure the given global-only *events* are enabled (additive).""" + global _active_global_events + if events & ~_active_global_events: + _active_global_events |= events + assert _tool_id is not None # nosec + sys.monitoring.set_events(_tool_id, _active_global_events) + + +def _recompute_global_events() -> None: + """Re-derive the set of global-only events from the current registry.""" + global _active_global_events + if _tool_id is None: + return + needed: int = 0 + for entries in _registry.values(): + needed |= _events_for(entries) + needed &= _GLOBAL_EVENTS + if needed != _active_global_events: + _active_global_events = needed + sys.monitoring.set_events(_tool_id, needed) + + +def register(code: CodeType, handler: MonitoringEventHandler) -> None: + """Register a monitoring event handler for *code*. + + The handler instance itself is the registration key; pass the same object + to :func:`unregister` to remove it. + """ + handler_events: int = _events_for_handler(handler) + if not handler_events: + raise ValueError("Handler overrides no MonitoringEventHandler methods") + + tool_id: int = _setup() + entry: _Entry = _Entry(handler, handler_events) + + with _registry_lock: + entries: Optional[dict[int, _Entry]] = _registry.get(code) + if entries is None: + _registry[code] = entries = {} + entries[id(handler)] = entry + all_events: int = _events_for(entries) + sys.monitoring.set_local_events(tool_id, code, all_events & _LOCAL_EVENTS) + _enable_global_events(all_events & _GLOBAL_EVENTS) + + +def refresh(code: CodeType) -> None: + """Re-apply local events for *code*, resetting any per-line DISABLE state. + + Call this after adding a new hook for a line that may have been previously + disabled via a ``DISABLE`` return from :meth:`MonitoringEventHandler.on_py_line`. + """ + with _registry_lock: + entries: Optional[dict[int, _Entry]] = _registry.get(code) + if entries and _tool_id is not None: + events: int = _events_for(entries) & _LOCAL_EVENTS + # A DISABLE returned from a per-line callback is sticky until the + # monitored event set changes or restart_events() is called. + # Re-applying the same local events does not clear it; toggling + # local events off and back on re-arms only this tool's DISABLE + # marks for *code* without the global restart_events() call that + # would reset other tools' disabled-event bookkeeping (coverage.py). + sys.monitoring.set_local_events(_tool_id, code, 0) + sys.monitoring.set_local_events(_tool_id, code, events) + + +def unregister(code: CodeType, handler: MonitoringEventHandler) -> None: + """Remove *handler* from the handlers registered for *code*.""" + with _registry_lock: + existing: Optional[dict[int, _Entry]] = _registry.get(code) + if existing is None: + return + + existing.pop(id(handler), None) + + if not existing: + del _registry[code] + if _tool_id is not None: + sys.monitoring.set_local_events(_tool_id, code, 0) + else: + assert _tool_id is not None # nosec + sys.monitoring.set_local_events(_tool_id, code, _events_for(existing) & _LOCAL_EVENTS) + + _recompute_global_events() diff --git a/tests/internal/test_monitoring.py b/tests/internal/test_monitoring.py new file mode 100644 index 00000000000..77d4ee85e60 --- /dev/null +++ b/tests/internal/test_monitoring.py @@ -0,0 +1,177 @@ +"""Tests for ddtrace.internal.monitoring, the multiplexed sys.monitoring layer. + +These focus on the local-vs-global event split for PY_UNWIND. PY_UNWIND is a +global-only sys.monitoring event: passing it to ``set_local_events`` raises +``ValueError: invalid local event set``. The module therefore enables PY_UNWIND +via ``set_events`` (global) and keeps PY_START/PY_RETURN/LINE per-code via +``set_local_events`` (local). +""" + +import sys +from typing import TYPE_CHECKING +from typing import Any + +import pytest + + +# The module only imports on 3.15+ (it raises ImportError below that). On older +# interpreters importorskip skips the whole module at collection time. Under a +# type checker we import it directly so member/base-class references resolve. +if TYPE_CHECKING: + from ddtrace.internal import monitoring +else: + monitoring = pytest.importorskip("ddtrace.internal.monitoring") + +# TODO(py-315): the full 3.15 monitoring/profiling stack (PR #17624 and its +# dependencies) is not yet enabled on this branch. Skip on 3.15 for now; the +# stacked PRs remove this mark and these tests run and pass on 3.15. +pytestmark = pytest.mark.skipif( + sys.version_info >= (3, 15), + reason="TODO(py-315): enable once the 3.15 monitoring stack lands (PR #17624 + deps)", +) + +# Fetched via getattr so the type checker treats it as Any: the source module's +# `_E = sys.monitoring.events` has an indeterminate type when mypy analyzes it +# under a pre-3.15 Python version. +_E: Any = getattr(monitoring, "_E") + + +class UnwindHandler(monitoring.MonitoringEventHandler): + def __init__(self): + self.unwinds = [] + + def on_py_unwind(self, code, instruction_offset, exception): + self.unwinds.append((code, exception)) + + +class StartAndUnwindHandler(monitoring.MonitoringEventHandler): + def __init__(self): + self.started = False + self.unwound = False + + def on_py_start(self, code, instruction_offset): + self.started = True + + def on_py_unwind(self, code, instruction_offset, exception): + self.unwound = True + + +@pytest.fixture +def registered(): + """Register a handler for a code object and always unregister afterwards.""" + registrations = [] + + def _register(code, handler): + monitoring.register(code, handler) + registrations.append((code, handler)) + return handler + + yield _register + + for code, handler in registrations: + monitoring.unregister(code, handler) + + +def test_register_unwind_handler_does_not_raise(registered): + """Regression: registering a PY_UNWIND-only handler must not raise. + + Before the fix, ``register`` passed PY_UNWIND to ``set_local_events`` which + raised ``ValueError: invalid local event set``. + """ + + def boom(): + raise ValueError("boom") + + registered(boom.__code__, UnwindHandler()) + + +def test_unwind_enabled_globally_not_locally(registered): + """PY_UNWIND must be a global event; it must not appear in local events.""" + + def boom(): + raise ValueError("boom") + + registered(boom.__code__, UnwindHandler()) + + tool_id = monitoring._tool_id + assert tool_id is not None + + local_events = sys.monitoring.get_local_events(tool_id, boom.__code__) + global_events = sys.monitoring.get_events(tool_id) + + assert not (local_events & _E.PY_UNWIND), "PY_UNWIND must not be a local event" + assert global_events & _E.PY_UNWIND, "PY_UNWIND must be enabled globally" + + +def test_on_py_unwind_does_not_disable_unregistered_code(): + """Regression: the unwind callback must return None (never DISABLE). + + PY_UNWIND fires for every unwinding frame, including code with no handler. + Returning DISABLE would permanently disarm the global event for that code + location with no re-arm path, so unregistered code must yield None. + """ + + def unrelated(): + pass + + result = monitoring._on_py_unwind(unrelated.__code__, 0, ValueError("x")) + assert result is None, "unregistered code must not be disabled (must return None, not DISABLE)" + + +def test_unwind_callback_fires_on_exception(registered): + """A registered handler receives on_py_unwind when its code unwinds.""" + + def boom(): + raise ValueError("kaboom") + + handler = registered(boom.__code__, UnwindHandler()) + + with pytest.raises(ValueError): + boom() + + assert any(exc.args == ("kaboom",) for _, exc in handler.unwinds), ( + "on_py_unwind was not called for the unwinding frame" + ) + + +def test_unregister_disables_global_unwind(): + """Unregistering the last unwind handler clears the global PY_UNWIND event.""" + + def boom(): + raise ValueError("boom") + + handler = UnwindHandler() + monitoring.register(boom.__code__, handler) + + tool_id = monitoring._tool_id + assert tool_id is not None + assert sys.monitoring.get_events(tool_id) & _E.PY_UNWIND + + monitoring.unregister(boom.__code__, handler) + + assert not (sys.monitoring.get_events(tool_id) & _E.PY_UNWIND), ( + "global PY_UNWIND should be disabled once no handlers need it" + ) + + +def test_mixed_local_and_global_events(registered): + """A handler overriding both PY_START and PY_UNWIND gets each at its scope.""" + + def fn(): + raise ValueError("mixed") + + handler = registered(fn.__code__, StartAndUnwindHandler()) + + tool_id = monitoring._tool_id + assert tool_id is not None + + local_events = sys.monitoring.get_local_events(tool_id, fn.__code__) + assert local_events & _E.PY_START, "PY_START must be a local event" + assert not (local_events & _E.PY_UNWIND), "PY_UNWIND must not be local" + assert sys.monitoring.get_events(tool_id) & _E.PY_UNWIND, "PY_UNWIND must be global" + + with pytest.raises(ValueError): + fn() + + assert handler.started, "on_py_start did not fire" + assert handler.unwound, "on_py_unwind did not fire" From 74e27bca4b4eb15ae8f28a323db55e383218c8e3 Mon Sep 17 00:00:00 2001 From: "Gabriele N. Tornetta" Date: Wed, 22 Jul 2026 17:31:23 -0400 Subject: [PATCH 02/13] chore(wrapping): add Python 3.15 wrapping context support Part of the #17849 split (PR 2/6). --- .../internal/bytecode_injection/__init__.py | 532 ++++++++----- ddtrace/internal/wrapping/asyncs.py | 149 +++- ddtrace/internal/wrapping/context.py | 743 +++++++++++------- ddtrace/internal/wrapping/generators.py | 77 +- .../bytecode_injection/test_injection.py | 19 +- tests/internal/test_wrapping.py | 118 +-- tests/wrapping/conftest.py | 15 + tests/wrapping/test_generators.py | 7 - tests/wrapping/test_unwrap.py | 9 +- 9 files changed, 1137 insertions(+), 532 deletions(-) diff --git a/ddtrace/internal/bytecode_injection/__init__.py b/ddtrace/internal/bytecode_injection/__init__.py index 5bcbbc978d8..6ee36859165 100644 --- a/ddtrace/internal/bytecode_injection/__init__.py +++ b/ddtrace/internal/bytecode_injection/__init__.py @@ -1,4 +1,5 @@ from collections import deque +from types import CodeType from types import FunctionType from typing import Any # noqa:F401 from typing import Callable # noqa:F401 @@ -24,223 +25,368 @@ class InvalidLine(Exception): """ -# DEV: This is the bytecode equivalent of -# >>> hook(arg) -# Additionally, we must discard the return value (top of the stack) to restore -# the stack to the state prior to the call. - -INJECTION_ASSEMBLY = Assembly() if PY >= (3, 15): - raise NotImplementedError("Python >= 3.15 is not supported yet") -elif PY >= (3, 13): - INJECTION_ASSEMBLY.parse( - r""" - load_const {hook} - push_null - load_const {arg} - call 1 - pop_top + import weakref + + from ddtrace.internal import monitoring as _monitoring + from ddtrace.internal.threads import Lock + from ddtrace.internal.utils.inspection import linenos + + class _LineHookHandler(_monitoring.MonitoringEventHandler): + """Per-code-object handler that dispatches line hooks via sys.monitoring.""" + + def __init__(self) -> None: + # lineno -> list of (hook, arg) pairs in registration order + self._hooks: dict[int, list[tuple[HookType, Any]]] = {} + + def on_py_line(self, code: Any, line_number: int) -> Any: + hooks: "list[tuple[HookType, Any]] | None" = self._hooks.get(line_number) + if not hooks: + return _monitoring.DISABLE # type: ignore[has-type] + for hook, arg in hooks: + hook(arg) + return None + + def add(self, line: int, hook: HookType, arg: Any) -> None: + self._hooks.setdefault(line, []).append((hook, arg)) + + def remove(self, line: int, hook: HookType, arg: Any) -> None: + hooks: "list[tuple[HookType, Any]] | None" = self._hooks.get(line) + if hooks is not None: + try: + hooks.remove((hook, arg)) + except ValueError: + pass + if not hooks: + del self._hooks[line] + + @property + def is_empty(self) -> bool: + return not self._hooks + + # WeakKeyDictionary: code object -> _LineHookHandler + _line_hook_registry: "weakref.WeakKeyDictionary[CodeType, _LineHookHandler]" = weakref.WeakKeyDictionary() + _line_hook_lock = Lock() + + def inject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoType]: + """Bulk-inject a list of hooks into a function. + + Hooks are specified via a list of tuples, where each tuple contains the hook + itself, the line number and the identifying argument passed to the hook. + + Returns the list of hooks that failed to be injected. """ - ) -elif PY >= (3, 12): - INJECTION_ASSEMBLY.parse( - r""" - push_null - load_const {hook} - load_const {arg} - call 1 - pop_top + code: CodeType = get_function_code(f) + valid_lines: set[int] = linenos(code) + failed: list[HookInfoType] = [] + + with _line_hook_lock: + handler: _LineHookHandler | None = _line_hook_registry.get(code) + if handler is None: + handler = _LineHookHandler() + new_handler: bool = True + else: + new_handler = False + + for hook, line, arg in hooks: + if line not in valid_lines: + failed.append((hook, line, arg)) + continue + handler.add(line, hook, arg) + + if not handler.is_empty: + if new_handler: + _line_hook_registry[code] = handler + _monitoring.register(code, handler) + else: + # Reset any lines that were DISABLE'd so newly added hooks fire. + _monitoring.refresh(code) + + return failed + + def migrate_line_hooks(src: CodeType, dst: CodeType) -> None: + """Move line-hook registration from *src* to *dst*. + + Wrapping on 3.15 clones the function's code object; line probes installed + on the pre-wrap code must follow the clone or they stop firing. """ - ) -elif PY >= (3, 11): - INJECTION_ASSEMBLY.parse( - r""" - push_null - load_const {hook} - load_const {arg} - precall 1 - call 1 - pop_top + with _line_hook_lock: + handler: _LineHookHandler | None = _line_hook_registry.get(src) + if handler is None: + return + del _line_hook_registry[src] + _monitoring.unregister(src, handler) + _line_hook_registry[dst] = handler + _monitoring.register(dst, handler) + + def eject_all_hooks(f: FunctionType) -> None: + """Remove every line hook registered for *f*.""" + code: CodeType = get_function_code(f) + with _line_hook_lock: + handler: _LineHookHandler | None = _line_hook_registry.get(code) + if handler is None: + return + del _line_hook_registry[code] + _monitoring.unregister(code, handler) + + def eject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoType]: + """Bulk-eject a list of hooks from a function. + + The hooks are specified via a list of tuples, where each tuple contains the + hook line number and the identifying argument. + + Returns the list of hooks that failed to be ejected. """ - ) + code: CodeType = get_function_code(f) + failed: list[HookInfoType] = [] + + with _line_hook_lock: + handler: _LineHookHandler | None = _line_hook_registry.get(code) + if handler is None: + return list(hooks) + + for hook, line, arg in hooks: + before: int = len(handler._hooks.get(line, ())) + handler.remove(line, hook, arg) + if len(handler._hooks.get(line, ())) == before: + failed.append((hook, line, arg)) + + if handler.is_empty: + del _line_hook_registry[code] + _monitoring.unregister(code, handler) + + return failed + + def inject_hook(f: FunctionType, hook: HookType, line: int, arg: Any) -> FunctionType: + """Inject a hook into a function at the given line number.""" + failed: list[HookInfoType] = inject_hooks(f, [(hook, line, arg)]) + if failed: + raise InvalidLine("Line %d does not exist or is either blank or a comment" % line) + return f + + def eject_hook(f: FunctionType, hook: HookType, line: int, arg: Any) -> FunctionType: + """Eject a hook from a function at the given line number.""" + failed: list[HookInfoType] = eject_hooks(f, [(hook, line, arg)]) + if failed: + raise InvalidLine("Line %d does not contain a hook" % line) + return f + else: - INJECTION_ASSEMBLY.parse( - r""" - load_const {hook} - load_const {arg} - call_function 1 - pop_top + # DEV: This is the bytecode equivalent of + # >>> hook(arg) + # Additionally, we must discard the return value (top of the stack) to restore + # the stack to the state prior to the call. + + INJECTION_ASSEMBLY = Assembly() + if PY >= (3, 13): + INJECTION_ASSEMBLY.parse( + r""" + load_const {hook} + push_null + load_const {arg} + call 1 + pop_top + """ + ) + elif PY >= (3, 12): + INJECTION_ASSEMBLY.parse( + r""" + push_null + load_const {hook} + load_const {arg} + call 1 + pop_top + """ + ) + elif PY >= (3, 11): + INJECTION_ASSEMBLY.parse( + r""" + push_null + load_const {hook} + load_const {arg} + precall 1 + call 1 + pop_top + """ + ) + else: + INJECTION_ASSEMBLY.parse( + r""" + load_const {hook} + load_const {arg} + call_function 1 + pop_top + """ + ) + + _INJECT_HOOK_OPCODES = [_.name for _ in INJECTION_ASSEMBLY] + + def _inject_hook(code: Bytecode, hook: HookType, lineno: int, arg: Any) -> None: + """Inject a hook at the given line number inside an abstract code object. + + The hook is called with the given argument, which is also used as an + identifier for the hook itself. This should be kept in case the hook needs + to be removed. + """ + # DEV: In general there are no guarantees for bytecode to be "linear", + # meaning that a line number can occur multiple times. We need to find all + # occurrences and inject the hook at each of them. An example of when this + # happens is with finally blocks, which are duplicated at the end of the + # bytecode. + locs: deque[tuple[int, str]] = deque() + last_lineno = None + instrs = set() + for i, item in enumerate(code): + if not isinstance(item, Instr): + continue + if item.lineno == last_lineno: + continue + last_lineno = item.lineno + # Some lines might be implemented across multiple instruction + # offsets, and sometimes a NOP is used as a placeholder. We skip + # those to avoid duplicate injections. + if item.lineno == lineno: + locs.appendleft((i, item.name)) + instrs.add(item.name) + + if not locs: + raise InvalidLine("Line %d does not exist or is either blank or a comment" % lineno) + + if instrs == {"NOP"}: + # If the line occurs on NOPs only, we instrument only the first one + last_instr = locs.pop() + locs.clear() + locs.append(last_instr) + elif "NOP" in instrs: + # If the line occurs on NOPs and other instructions, we remove the NOPs + # to avoid injecting the hook multiple times. The NOP in this case is + # just a placeholder. + locs = deque((i, instr) for i, instr in locs if instr != "NOP") + + for i, instr in locs: + if instr.startswith("END_"): + # This is the end of a block, e.g. a for loop. We have already + # instrumented the block on entry, so we skip instrumenting the + # end as well. + continue + code[i:i] = INJECTION_ASSEMBLY.bind(dict(hook=hook, arg=arg), lineno=lineno) + + _INJECT_HOOK_OPCODE_POS = 1 if (3, 11) <= PY < (3, 13) else 0 + _INJECT_ARG_OPCODE_POS = 1 if PY < (3, 11) else 2 + + def _eject_hook(code: Bytecode, hook: HookType, line: int, arg: Any) -> None: + """Eject a hook from the abstract code object at the given line number. + + The hook is identified by its argument. This ensures that only the right + hook is ejected. """ - ) + locs: deque[int] = deque() + for i, item in enumerate(code): + if not isinstance(item, Instr): + continue + try: + hook_op = code[i + _INJECT_HOOK_OPCODE_POS] + arg_op = code[i + _INJECT_ARG_OPCODE_POS] + if not isinstance(hook_op, Instr) or not isinstance(arg_op, Instr): + continue + opcodes = [] + for j in range(i, i + len(_INJECT_HOOK_OPCODES)): + op = code[j] + if not isinstance(op, Instr): + break + opcodes.append(op.name) + else: + # DEV: We look at the expected opcode pattern to match the injected + # hook and we also test for the expected opcode arguments + if ( + item.lineno == line + and hook_op.arg == hook # bound methods don't like identity comparisons + and arg_op.arg is arg + and opcodes == _INJECT_HOOK_OPCODES + ): + locs.appendleft(i) + except IndexError: + pass + + if not locs: + raise InvalidLine("Line %d does not contain a hook" % line) + + for i in locs: + del code[i : i + len(_INJECT_HOOK_OPCODES)] + + def inject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoType]: + """Bulk-inject a list of hooks into a function. + + Hooks are specified via a list of tuples, where each tuple contains the hook + itself, the line number and the identifying argument passed to the hook. + + Returns the list of hooks that failed to be injected. + """ + abstract_code = Bytecode.from_code(get_function_code(f)) -_INJECT_HOOK_OPCODES = [_.name for _ in INJECTION_ASSEMBLY] + failed = [] + for hook, line, arg in hooks: + try: + _inject_hook(abstract_code, hook, line, arg) + except InvalidLine: + failed.append((hook, line, arg)) + if len(failed) < len(hooks): + set_function_code(f, abstract_code.to_code()) -def _inject_hook(code: Bytecode, hook: HookType, lineno: int, arg: Any) -> None: - """Inject a hook at the given line number inside an abstract code object. + return failed - The hook is called with the given argument, which is also used as an - identifier for the hook itself. This should be kept in case the hook needs - to be removed. - """ - # DEV: In general there are no guarantees for bytecode to be "linear", - # meaning that a line number can occur multiple times. We need to find all - # occurrences and inject the hook at each of them. An example of when this - # happens is with finally blocks, which are duplicated at the end of the - # bytecode. - locs: deque[tuple[int, str]] = deque() - last_lineno = None - instrs = set() - for i, instr in enumerate(code): - if not isinstance(instr, Instr): - # pseudo-instruction (e.g. label) - continue - if instr.lineno == last_lineno: - continue - last_lineno = instr.lineno - # Some lines might be implemented across multiple instruction - # offsets, and sometimes a NOP is used as a placeholder. We skip - # those to avoid duplicate injections. - if instr.lineno == lineno: - locs.appendleft((i, instr.name)) - instrs.add(instr.name) - - if not locs: - raise InvalidLine("Line %d does not exist or is either blank or a comment" % lineno) - - if instrs == {"NOP"}: - # If the line occurs on NOPs only, we instrument only the first one - last_instr = locs.pop() - locs.clear() - locs.append(last_instr) - elif "NOP" in instrs: - # If the line occurs on NOPs and other instructions, we remove the NOPs - # to avoid injecting the hook multiple times. The NOP in this case is - # just a placeholder. - locs = deque((i, instr) for i, instr in locs if instr != "NOP") - - for i, opname in locs: - if opname.startswith("END_"): - # This is the end of a block, e.g. a for loop. We have already - # instrumented the block on entry, so we skip instrumenting the - # end as well. - continue - code[i:i] = INJECTION_ASSEMBLY.bind(dict(hook=hook, arg=arg), lineno=lineno) - - -_INJECT_HOOK_OPCODE_POS = 1 if (3, 11) <= PY < (3, 13) else 0 -_INJECT_ARG_OPCODE_POS = 1 if PY < (3, 11) else 2 - - -def _eject_hook(code: Bytecode, hook: HookType, line: int, arg: Any) -> None: - """Eject a hook from the abstract code object at the given line number. - - The hook is identified by its argument. This ensures that only the right - hook is ejected. - """ - locs: deque[int] = deque() - for i, instr in enumerate(code): - if not isinstance(instr, Instr): - # pseudo-instruction (e.g. label) - continue - try: - # DEV: We look at the expected opcode pattern to match the injected - # hook and we also test for the expected opcode arguments - _hook_instr = code[i + _INJECT_HOOK_OPCODE_POS] - _arg_instr = code[i + _INJECT_ARG_OPCODE_POS] - _window = [code[_] for _ in range(i, i + len(_INJECT_HOOK_OPCODES))] - if ( - instr.lineno == line - and isinstance(_hook_instr, Instr) - and _hook_instr.arg == hook # bound methods don't like identity comparisons - and isinstance(_arg_instr, Instr) - and _arg_instr.arg is arg - and all(isinstance(_c, Instr) for _c in _window) - and [_c.name for _c in _window if isinstance(_c, Instr)] == _INJECT_HOOK_OPCODES - ): - locs.appendleft(i) - except (AttributeError, IndexError): - pass - - if not locs: - raise InvalidLine("Line %d does not contain a hook" % line) - - for i in locs: - del code[i : i + len(_INJECT_HOOK_OPCODES)] - - -def inject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoType]: - """Bulk-inject a list of hooks into a function. - - Hooks are specified via a list of tuples, where each tuple contains the hook - itself, the line number and the identifying argument passed to the hook. - - Returns the list of hooks that failed to be injected. - """ - abstract_code = Bytecode.from_code(get_function_code(f)) + def eject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoType]: + """Bulk-eject a list of hooks from a function. - failed = [] - for hook, line, arg in hooks: - try: - _inject_hook(abstract_code, hook, line, arg) - except InvalidLine: - failed.append((hook, line, arg)) + The hooks are specified via a list of tuples, where each tuple contains the + hook line number and the identifying argument. - if len(failed) < len(hooks): - set_function_code(f, abstract_code.to_code()) + Returns the list of hooks that failed to be ejected. + """ + abstract_code = Bytecode.from_code(f.__code__) - return failed + failed = [] + for hook, line, arg in hooks: + try: + _eject_hook(abstract_code, hook, line, arg) + except InvalidLine: + failed.append((hook, line, arg)) + if len(failed) < len(hooks): + f.__code__ = abstract_code.to_code() -def eject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoType]: - """Bulk-eject a list of hooks from a function. + return failed - The hooks are specified via a list of tuples, where each tuple contains the - hook line number and the identifying argument. + def inject_hook(f: FunctionType, hook: HookType, line: int, arg: Any) -> FunctionType: + """Inject a hook into a function. - Returns the list of hooks that failed to be ejected. - """ - abstract_code = Bytecode.from_code(f.__code__) + The hook is injected at the given line number and called with the given + argument. The latter is also used as an identifier for the hook. This should + be kept in case the hook needs to be removed. + """ + abstract_code = Bytecode.from_code(f.__code__) - failed = [] - for hook, line, arg in hooks: - try: - _eject_hook(abstract_code, hook, line, arg) - except InvalidLine: - failed.append((hook, line, arg)) + _inject_hook(abstract_code, hook, line, arg) - if len(failed) < len(hooks): f.__code__ = abstract_code.to_code() - return failed + return f + def eject_hook(f: FunctionType, hook: HookType, line: int, arg: Any) -> FunctionType: + """Eject a hook from a function. -def inject_hook(f: FunctionType, hook: HookType, line: int, arg: Any) -> FunctionType: - """Inject a hook into a function. - - The hook is injected at the given line number and called with the given - argument. The latter is also used as an identifier for the hook. This should - be kept in case the hook needs to be removed. - """ - abstract_code = Bytecode.from_code(f.__code__) - - _inject_hook(abstract_code, hook, line, arg) - - f.__code__ = abstract_code.to_code() - - return f - + The hook is identified by its line number and the argument passed to the + hook. + """ + abstract_code = Bytecode.from_code(f.__code__) -def eject_hook(f: FunctionType, hook: HookType, line: int, arg: Any) -> FunctionType: - """Eject a hook from a function. + _eject_hook(abstract_code, hook, line, arg) - The hook is identified by its line number and the argument passed to the - hook. - """ - abstract_code = Bytecode.from_code(f.__code__) - - _eject_hook(abstract_code, hook, line, arg) + f.__code__ = abstract_code.to_code() - f.__code__ = abstract_code.to_code() + return f - return f + def eject_all_hooks(f: FunctionType) -> None: + """No-op on Python <3.15: line hooks are removed when __code__ is restored.""" + return None diff --git a/ddtrace/internal/wrapping/asyncs.py b/ddtrace/internal/wrapping/asyncs.py index a1de28a074e..10989b91f5c 100644 --- a/ddtrace/internal/wrapping/asyncs.py +++ b/ddtrace/internal/wrapping/asyncs.py @@ -1,3 +1,4 @@ +import enum import sys from types import CodeType @@ -9,6 +10,25 @@ PY = sys.version_info[:2] +def _ensure_common_constant_none() -> None: + """Extend bytecode.CommonConstant with CONSTANT_NONE when missing. + + bytecode 0.18.1 only defines ASSERTION_ERROR..BUILTIN_ANY. CPython 3.15's + LOAD_COMMON_CONSTANT uses index 7 for None (see CONSTANT_NONE in + Include/internal/pycore_opcode_utils.h). + """ + import bytecode.instr as instr_mod + + cc = instr_mod.CommonConstant + if hasattr(cc, "CONSTANT_NONE"): + return + + members: list[tuple[str, int]] = [(m.name, m.value) for m in cc] + members.append(("CONSTANT_NONE", 7)) + extended = enum.IntEnum("CommonConstant", members) # type: ignore[misc] + instr_mod.CommonConstant = extended # type: ignore[misc, assignment] + + # ----------------------------------------------------------------------------- # Coroutine and Async Generator Wrapping # ----------------------------------------------------------------------------- @@ -36,7 +56,134 @@ ASYNC_HEAD_ASSEMBLY = None if PY >= (3, 15): - raise NotImplementedError("This version of CPython is not supported yet") + _ensure_common_constant_none() + ASYNC_HEAD_ASSEMBLY = Assembly() + ASYNC_HEAD_ASSEMBLY.parse( + r""" + return_generator + pop_top + """ + ) + + COROUTINE_ASSEMBLY.parse( + r""" + get_awaitable 0 + push_null + load_common_constant asm.instr.CommonConstant.CONSTANT_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 + push_null + load_common_constant asm.instr.CommonConstant.CONSTANT_NONE + presend0: + send @send0 + tried + + try @genexit lasti + yield_value 1 + 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 + push_null + load_common_constant asm.instr.CommonConstant.CONSTANT_NONE + + presend1: + send @send1 + yield_value 1 + 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 + push_null + load_common_constant asm.instr.CommonConstant.CONSTANT_NONE + + presend2: + send @send2 + yield_value 1 + 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, 14): ASYNC_HEAD_ASSEMBLY = Assembly() diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index 8333b997c36..f211e4c4e34 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -3,11 +3,12 @@ from inspect import iscoroutinefunction from inspect import isgeneratorfunction import sys +from types import CodeType from types import FrameType from types import FunctionType from types import TracebackType import typing as t -from typing import Protocol # noqa:F401 +from typing import Protocol import weakref import bytecode @@ -47,7 +48,7 @@ def __init__(self) -> None: @property def uwc(self) -> t.Optional["_UniversalWrappingContext"]: - ref = self._uwc_ref + ref: t.Optional[weakref.ref["_UniversalWrappingContext"]] = self._uwc_ref return ref() if ref is not None else None @uwc.setter @@ -56,7 +57,7 @@ def uwc(self, value: t.Optional["_UniversalWrappingContext"]) -> None: @classmethod def get_or_create(cls, f: FunctionType) -> "_ContextRecord": - record = _registry.get(f) + record: t.Optional["_ContextRecord"] = _registry.get(f) if record is None: with _registry_lock: record = _registry.get(f) @@ -128,8 +129,12 @@ def get_or_create(cls, f: FunctionType) -> "_ContextRecord": CONTEXT_RETURN = Assembly() CONTEXT_FOOT = Assembly() -if sys.version_info >= (3, 15): - raise NotImplementedError("Python >= 3.15 is not supported yet") +if sys.version_info >= (3, 16): + raise NotImplementedError("This version of Python is not supported yet") +elif sys.version_info >= (3, 15): + # We rely on sys.monitoring for wrapping, so no bytecode manipulation is + # needed. + pass elif sys.version_info >= (3, 13): CONTEXT_HEAD.parse( r""" @@ -338,6 +343,27 @@ def get_or_create(cls, f: FunctionType) -> "_ContextRecord": ) +# On the bytecode path __enter__ is invoked directly from inside the wrapped +# function, so the monitored frame is one level up. On the monitoring path +# (3.15+) the stack is: +# monitored function → monitoring._on_py_start → uwc.on_py_start → __enter__ +# so the monitored frame is three levels up. +_ENTER_FRAME_DEPTH = 3 if sys.version_info >= (3, 15) else 1 + +if sys.version_info >= (3, 15): + from ddtrace.internal import monitoring as _monitoring + + # Keyed by code object: drives sys.monitoring dispatch and is_wrapped/extract lookup. + _ctx_registry: "weakref.WeakKeyDictionary[CodeType, _UniversalWrappingContext]" = weakref.WeakKeyDictionary() + # Keyed by function instance: distinguishes functions that share a code object + # (e.g. closures re-created in a loop) from one another. Kept off the function's + # __dict__ (unlike a plain attribute) so functools.wraps does not propagate + # wrapping metadata and deepcopy of decorated functions does not traverse a + # context holding unpicklable state (locks, ContextVars). See issue #16443. + _fn_registry: "weakref.WeakKeyDictionary[FunctionType, _UniversalWrappingContext]" = weakref.WeakKeyDictionary() + _ctx_registry_lock = Lock() + + # This is abstract and should not be used directly class BaseWrappingContext(ABC): __priority__: int = 0 @@ -372,7 +398,9 @@ def __enter__(self) -> "BaseWrappingContext": return self def _pop_storage(self) -> dict[str, t.Any]: - storage = t.cast(dict[str, t.Any], self._storage.get()) + storage = self._storage.get() + if storage is None: + return {} self._storage.set(storage.pop("__dd_wrapping_context_prev__")) return storage @@ -464,88 +492,112 @@ def unwrap(self) -> None: pass -class LazyWrappingContext(WrappingContext): - def __init__(self, f: FunctionType): - super().__init__(f) +if sys.version_info >= (3, 15): + # Monitoring-based instrumentation has negligible per-function overhead, so + # there is no benefit to deferring wrapping until first call. On Python 3.15+ + # this is a transparent alias for WrappingContext kept only for API compatibility. + LazyWrappingContext = WrappingContext - self._trampoline: t.Optional[Wrapper] = None - self._trampoline_lock = Lock() +else: - @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) + class LazyWrappingContext(WrappingContext): + def __init__(self, f: FunctionType): + super().__init__(f) - def wrap(self) -> None: - """Perform the bytecode wrapping on first invocation.""" - with (tl := self._trampoline_lock): - if self._trampoline is not None: - return + self._trampoline: t.Optional[Wrapper] = None + self._trampoline_lock = Lock() - # 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 + @classmethod + def is_wrapped(cls, f: FunctionType) -> bool: + with _registry_lock: + record: t.Optional[_ContextRecord] = _registry.get(f) + if record is None: + return False + return any(isinstance(c, cls) for c in record.lazy_contexts) - 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)) + def wrap(self) -> None: + """Perform the bytecode wrapping on first invocation.""" + with (tl := self._trampoline_lock): + if self._trampoline is not None: + return - self._trampoline = None + # 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 - 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") + 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)) - super(LazyWrappingContext, self).wrap() - return f(*args, **kwargs) + self._trampoline = None - wrap(t.cast(FunctionType, self.__wrapped__), trampoline) + inconsistent: bool = False + with _registry_lock: + record: t.Optional[_ContextRecord] = _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") - self._trampoline = trampoline + super(LazyWrappingContext, self).wrap() + return f(*args, **kwargs) - _ContextRecord.get_or_create(t.cast(FunctionType, self.__wrapped__)).lazy_contexts.add(self) + wrap(t.cast(FunctionType, self.__wrapped__), trampoline) - 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) + self._trampoline = trampoline - unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline) - self._trampoline = None + _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: t.Optional[_ContextRecord] = _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.""" + """A function that is (or can be) wrapped with a WrappingContext. + + Used purely as a structural type marker for call sites that operate on + wrapped functions. Per-function wrapping state is tracked off the function + object (see _fn_registry on 3.15+ / the registry on older versions), so this + protocol intentionally carries no wrapping-metadata attribute. + """ def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: pass +# On 3.15+ _UniversalWrappingContext also implements MonitoringEventHandler so +# it can be registered directly with the multiplexer via register(code, self). +if sys.version_info >= (3, 15): + from ddtrace.internal.monitoring import MonitoringEventHandler as _MonitoringEventHandler + + _UWC_BASES: tuple = (BaseWrappingContext, _MonitoringEventHandler) +else: + _UWC_BASES = (BaseWrappingContext,) + + # This class provides an interface between single bytecode wrapping and multiple # logical context wrapping -class _UniversalWrappingContext(BaseWrappingContext): +class _UniversalWrappingContext(*_UWC_BASES): # type: ignore[misc] def __init__(self, f: FunctionType) -> None: super().__init__(f) @@ -581,7 +633,7 @@ def __enter__(self) -> "_UniversalWrappingContext": super().__enter__() # Make the frame object available to the contexts - self.set("__frame__", sys._getframe(1)) + self.set("__frame__", sys._getframe(_ENTER_FRAME_DEPTH)) for context in self._contexts: context.__enter__() @@ -609,268 +661,399 @@ def __return__(self, value: T) -> T: for context in self._contexts[::-1]: context.__return__(value) - return super().__return__(value) + return t.cast(T, 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: + if sys.version_info >= (3, 15): + + def on_py_start(self, code: t.Any, instruction_offset: int) -> None: + self.__enter__() + + def on_py_return(self, code: t.Any, instruction_offset: int, retval: t.Any) -> None: + self.__return__(retval) + + def on_py_unwind(self, code: t.Any, instruction_offset: int, exception: BaseException) -> None: + self.__exit__(type(exception), exception, exception.__traceback__) + + @classmethod + def is_wrapped(cls, f: FunctionType) -> bool: + try: + code: CodeType = get_function_code(f) + if code not in _ctx_registry: 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 + # Also verify that THIS function instance is wrapped, not just some + # other function that shares the same code object (e.g. closures + # re-created in a loop). + return _fn_registry.get(f) is _ctx_registry[code] + except Exception: + return False - @classmethod - def extract(cls, f: FunctionType) -> "_UniversalWrappingContext": - with _registry_lock: - if not cls.is_wrapped(f): + @classmethod + def extract(cls, f: FunctionType) -> "_UniversalWrappingContext": + ctx: t.Optional["_UniversalWrappingContext"] = _ctx_registry.get(get_function_code(f)) + if ctx is None: raise ValueError("Function is not wrapped") - return t.cast(_UniversalWrappingContext, _registry[f].uwc) - - if sys.version_info >= (3, 11): + # Monitoring dispatches per code object, so a fresh function instance + # that merely shares a code object (e.g. a closure re-created in a + # loop) maps to the same registry entry without being wrapped itself. + # Mirror is_wrapped()'s per-instance check so callers such as wrapped() + # replace the stale registration via wrap() instead of double- + # registering on a context that belongs to another (often dead) function. + if _fn_registry.get(f) is not ctx: + raise ValueError("Function is not wrapped") + return ctx def wrap(self) -> None: - f = self.__wrapped__ + f: FunctionType = self.__wrapped__ + original_code: CodeType = get_function_code(f) + with _ctx_registry_lock: + if original_code in _ctx_registry: + # Allow wrapping a new function instance that shares the same + # code object as an already-wrapped (but orphaned) function. + # This happens when closures are re-created in a loop: each + # iteration produces a new function object but the same code + # object. If the new function is not itself registered, the old + # monitoring registration is stale and should be replaced. + if _fn_registry.get(f) is _ctx_registry[original_code]: + raise ValueError("Function already wrapped") + # Stale entry: clean up old monitoring before re-registering. + old: "_UniversalWrappingContext" = _ctx_registry.pop(original_code) + _monitoring.unregister(original_code, old) + + # sys.monitoring dispatches per code object. Clone the code so + # unwrapped siblings that share the same CodeType are not affected. + from ddtrace.internal.bytecode_injection import migrate_line_hooks + + link_function_to_code(original_code, f) + monitor_code: CodeType = original_code.replace() + migrate_line_hooks(original_code, monitor_code) + set_function_code(f, monitor_code) + self._original_code = original_code + + _ctx_registry[monitor_code] = self + _fn_registry[f] = self + self._finalize = weakref.finalize( + f, + _finalize_monitoring_wrap, + weakref.ref(self), + weakref.ref(f), + ) + self._finalize.atexit = False + _monitoring.register(monitor_code, self) + def unwrap(self) -> None: + f: FunctionType = self.__wrapped__ + finalize: t.Optional[weakref.finalize] = getattr(self, "_finalize", None) + if finalize is not None: + finalize.detach() + del self._finalize + code: CodeType = get_function_code(f) + with _ctx_registry_lock: + if code not in _ctx_registry: + return + del _ctx_registry[code] + _fn_registry.pop(f, None) + _monitoring.unregister(code, self) + original_code: t.Optional[CodeType] = getattr(self, "_original_code", None) + if original_code is not None: + from ddtrace.internal.bytecode_injection import migrate_line_hooks + + migrate_line_hooks(code, original_code) + set_function_code(f, original_code) + del self._original_code + + else: + + @classmethod + def is_wrapped(cls, f: FunctionType) -> bool: + try: + with _registry_lock: + record: t.Optional[_ContextRecord] = _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 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 + if not cls.is_wrapped(f): + raise ValueError("Function is not wrapped") + return t.cast(_UniversalWrappingContext, _registry[f].uwc) - bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=code.co_firstlineno) + 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") - # 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 = Bytecode.from_code(code := get_function_code(f)) - 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 + # 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 - 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 + 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 - i += 1 - bc.insert(0, first_try_begin) + 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)) + 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()) + # 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__ + def unwrap(self) -> None: + f = self.__wrapped__ - with _registry_lock: - if not self.is_wrapped(f): - return + with _registry_lock: + if not self.is_wrapped(f): + return - wc = _registry[f].uwc + 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() - bc.pop() + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() + bc.pop() - except_label = bc.pop(0).target + 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) + # 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 += 1 + i = 0 - # 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)] = [] - # 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: + # 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 - 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 + # Recreate the code object + set_function_code(f, bc.to_code()) - if return_code is not None: - bc[i - len(return_code) : i] = [] - i -= len(return_code) - except AttributeError: - # Not an instruction - pass - i += 1 + # Clear the UWC from the registry; remove the record if fully empty. + record: t.Optional[_ContextRecord] = _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()) + else: - # 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 wrap(self) -> None: + f = t.cast(FunctionType, self.__wrapped__) - else: + with _registry_lock: + if self.is_wrapped(f): + raise ValueError("Function already wrapped") - def wrap(self) -> None: - 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") - - 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 + # 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 + # 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[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)) + 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()) + # 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__) + def unwrap(self) -> None: + f = t.cast(FunctionType, self.__wrapped__) - with _registry_lock: - if not self.is_wrapped(f): - return + with _registry_lock: + if not self.is_wrapped(f): + return + + wc = _registry[f].uwc + + bc = Bytecode.from_code(get_function_code(f)) - wc = _registry[f].uwc + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() - bc = Bytecode.from_code(get_function_code(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 - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() + bc[i : i + len(CONTEXT_HEAD) - 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 + # 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: t.Optional[_ContextRecord] = _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()) - # 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) +if sys.version_info >= (3, 15): + def _finalize_monitoring_wrap( + self_ref: "weakref.ref[_UniversalWrappingContext]", + f_ref: weakref.ref[FunctionType], + ) -> None: + """Unregister sys.monitoring when a wrapped function is collected without unwrap().""" + self: t.Optional["_UniversalWrappingContext"] = self_ref() + f: t.Optional[FunctionType] = f_ref() + if self is None or f is None: + return + try: + if _fn_registry.get(f) is self: + self.unwrap() + except Exception: + log.exception( + "ddtrace: error during finalizer unwrap of %s", + getattr(f, "__qualname__", "?"), + ) -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 + def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingContext]": + """Return the _UniversalWrappingContext for *f*, or None if not context-wrapped.""" + try: + return _UniversalWrappingContext.extract(f) + except ValueError: + return None + +else: + + def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingContext]": + """Return the _UniversalWrappingContext for *f*, or None if not context-wrapped.""" + with _registry_lock: + record: t.Optional[_ContextRecord] = _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 01c284c4686..d88341c818b 100644 --- a/ddtrace/internal/wrapping/generators.py +++ b/ddtrace/internal/wrapping/generators.py @@ -34,7 +34,82 @@ GENERATOR_HEAD_ASSEMBLY = None if PY >= (3, 15): - raise NotImplementedError("This version of CPython is not supported yet") + 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, 14): GENERATOR_HEAD_ASSEMBLY = Assembly() diff --git a/tests/internal/bytecode_injection/test_injection.py b/tests/internal/bytecode_injection/test_injection.py index a60f4fe5bce..82b6bb51ed6 100644 --- a/tests/internal/bytecode_injection/test_injection.py +++ b/tests/internal/bytecode_injection/test_injection.py @@ -1,7 +1,7 @@ from contextlib import contextmanager import sys +from unittest import mock -import mock import pytest from ddtrace.internal.bytecode_injection import InvalidLine @@ -31,7 +31,13 @@ def injected_hook(f, hook, arg, line=None): eject_hook(f, hook, line, arg) - assert f.__code__ is not code + if sys.version_info >= (3, 15): + # The 3.15+ monitoring-based injection path attaches hooks via + # sys.monitoring rather than rewriting bytecode, so the code object is + # intentionally left unchanged across inject/eject. + assert f.__code__ is code + else: + assert f.__code__ is not code def injection_target(a, b): @@ -292,4 +298,11 @@ def for_loop(): with injected_hook(for_loop, hook, arg, line=for_loop.__code__.co_firstlineno + 2): for_loop() - hook.assert_called_once_with(arg) + if sys.version_info >= (3, 15): + # The monitoring-based path fires a LINE event every time the loop + # header line is (re-)entered, i.e. once per iteration, rather than + # once at loop setup as the bytecode-rewriting path does. + assert hook.call_count >= 10 + hook.assert_called_with(arg) + else: + hook.assert_called_once_with(arg) diff --git a/tests/internal/test_wrapping.py b/tests/internal/test_wrapping.py index 749b829b4a5..d484261e944 100644 --- a/tests/internal/test_wrapping.py +++ b/tests/internal/test_wrapping.py @@ -1134,6 +1134,7 @@ def __enter__(self): return super().__enter__() +@pytest.mark.skipif(sys.version_info >= (3, 15), reason="LazyWrappingContext is eager on 3.15+") def test_wrapping_context_lazy(): free = 42 @@ -1211,6 +1212,7 @@ def __enter__(self): assert c1.count == c2.count == 0 +@pytest.mark.skipif(sys.version_info >= (3, 15), reason="LazyWrappingContext is eager on 3.15+") def test_wrapping_context_lazy_unwrap_before_call(): free = 42 @@ -1377,63 +1379,91 @@ async def inner(): ) -def test_wrapping_context_foot_has_valid_linenos(): - """Regression test: CONTEXT_FOOT.bind() must pass lineno to avoid None line numbers. +# Thread-based concurrency: ContextVar storage must be isolated per thread. +# --------------------------------------------------------------------------- - The CONTEXT_FOOT assembly (exception-handler epilogue injected by - _UniversalWrappingContext.wrap) was previously emitted without line-number - information. inspect.stack() / inspect.getframeinfo() raises TypeError when - it encounters a frame whose f_lineno is None, so any tool that inspects the - call stack from inside a WrappingContext-wrapped function (e.g. IAST's - report_stack) would crash. - """ - stack_from_inside: list = [] + +def test_wrapping_context_thread_concurrent(): + """Storage set by one thread must not bleed into a concurrent thread.""" + import threading + + results = {} + errors = [] + + class ThreadIsolationContext(DummyWrappingContext): + def __enter__(self): + super().__enter__() + self.set("tid", threading.get_ident()) + return self + + def __return__(self, value): + stored = self.get("tid") + current = threading.get_ident() + if stored != current: + errors.append(f"tid mismatch: stored={stored} current={current}") + results[current] = stored + return super().__return__(value) def foo(): - stack_from_inside.extend(inspect.stack()) - return 42 + return threading.get_ident() - wc = DummyWrappingContext(foo) + wc = ThreadIsolationContext(foo) wc.wrap() - result = foo() - assert result == 42 + barrier = threading.Barrier(10) - for frame_info in stack_from_inside: - lineno = frame_info.lineno - assert lineno is not None, ( - f"Frame {frame_info.filename}:{frame_info.function} has lineno=None — " - "CONTEXT_FOOT was bound without line-number information" - ) + def run(): + barrier.wait() + foo() + threads = [threading.Thread(target=run) for _ in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() -def test_wrapping_context_foot_has_valid_linenos_on_exception(): - """Same lineno regression check when the wrapped function raises an exception. + assert not errors, errors + # Each thread saw its own tid in storage. + for tid, stored in results.items(): + assert tid == stored - The CONTEXT_FOOT assembly is only executed when an exception propagates, so - the lineno fix must also cover the exception path. - """ - stack_from_handler: list = [] - class _Sentinel(Exception): - pass +# --------------------------------------------------------------------------- +# Async recursion: ContextVar stack must be correct for recursive coroutines. +# --------------------------------------------------------------------------- - class CaptureOnExit(WrappingContext): - def __exit__(self, exc_type, exc_value, traceback): - stack_from_handler.extend(inspect.stack()) - return super().__exit__(exc_type, exc_value, traceback) - def foo(): - raise _Sentinel("boom") +@pytest.mark.asyncio +async def test_wrapping_context_async_recursive(): + """Each recursive coroutine call must have its own isolated storage slot.""" + values = [] - CaptureOnExit(foo).wrap() + class AsyncRecursiveContext(DummyWrappingContext): + def __enter__(self): + super().__enter__() + n = self.__frame__.f_locals["n"] + self.set("n", n) + values.append(("enter", n)) + return self - with pytest.raises(_Sentinel): - foo() + def __return__(self, value): + n = self.__frame__.f_locals["n"] + assert self.get("n") == n, f"storage mismatch: expected {n}, got {self.get('n')}" + values.append(("return", n)) + return super().__return__(value) - for frame_info in stack_from_handler: - lineno = frame_info.lineno - assert lineno is not None, ( - f"Frame {frame_info.filename}:{frame_info.function} has lineno=None — " - "CONTEXT_FOOT exception path was bound without line-number information" - ) + async def afactorial(n): + if n == 0: + return 1 + return n * await afactorial(n - 1) + + wc = AsyncRecursiveContext(afactorial) + wc.wrap() + + result = await afactorial(5) + assert result == 120 + + entered = [n for ev, n in values if ev == "enter"] + returned = [n for ev, n in values if ev == "return"] + assert entered == [5, 4, 3, 2, 1, 0] + assert returned == [0, 1, 2, 3, 4, 5] diff --git a/tests/wrapping/conftest.py b/tests/wrapping/conftest.py index dbef918f9ef..d91992831b1 100644 --- a/tests/wrapping/conftest.py +++ b/tests/wrapping/conftest.py @@ -53,6 +53,21 @@ def pytest_collection_modifyitems(items): this directory -- otherwise a broader run (e.g. ``pytest tests/``) would abort on every test that legitimately has no ``mech``. """ + # On 3.15+ WrappingContext uses sys.monitoring instead of bytecode rewriting, + # so the strict xfails for wrapping_context t-string cases are obsolete. The + # markers live in test_tstrings_py314.py, which cannot be edited with a + # version-gated xfail condition (ruff rejects t-string syntax when the file + # is passed directly to the pre-commit hook). + if sys.version_info >= (3, 15): + for item in items: + if not str(getattr(item, "path", "")).startswith(_HERE + os.sep): + continue + if not item.name.startswith("test_tstring"): + continue + if "wrapping_context" not in item.nodeid: + continue + item.own_markers = [m for m in item.own_markers if m.name != "xfail"] + missing = [ item.nodeid for item in items diff --git a/tests/wrapping/test_generators.py b/tests/wrapping/test_generators.py index 20dffb99af6..cf1bc1a7921 100644 --- a/tests/wrapping/test_generators.py +++ b/tests/wrapping/test_generators.py @@ -4,8 +4,6 @@ wrapped generator must produce. """ -import sys - import pytest from tests.wrapping.mechanisms import xfail_mechanism @@ -74,11 +72,6 @@ def g(): assert log == ["cleanup"] -@xfail_mechanism( - "wrapping_context", - reason="WrappingContext.throw() on an unstarted generator crashes on 3.11+ (internal AttributeError)", - condition=sys.version_info >= (3, 11), -) def test_throw_on_unstarted_generator(mech): def g(): yield 1 diff --git a/tests/wrapping/test_unwrap.py b/tests/wrapping/test_unwrap.py index f6de0b55f10..2df8352bdd5 100644 --- a/tests/wrapping/test_unwrap.py +++ b/tests/wrapping/test_unwrap.py @@ -10,15 +10,17 @@ These tests assert a wrap->unwrap round-trip restores call behaviour and the signature, and pin a divergence in how completely the original is reinstated: ``internal_wrap`` puts back the exact original ``__code__`` object (a true -inverse, even when layers are nested); ``WrappingContext.unwrap`` restores -behaviour but rebuilds the code object rather than reinstating the original, so -``__code__`` identity is not restored (codified with a strict xfail). +inverse, even when layers are nested). On Python < 3.15, ``WrappingContext.unwrap`` +restores behaviour but rebuilds the code object rather than reinstating the +original, so ``__code__`` identity is not restored (codified with a strict xfail). +On 3.15+ the monitoring path never mutates ``__code__``, so identity is preserved. Mechanism-specific (the matrix's other two mechanisms have no in-place unwrap), so these opt out of the all-mechanisms ``mech`` guardrail. """ import inspect +import sys import pytest @@ -114,6 +116,7 @@ def f(a): @pytest.mark.xfail( strict=True, + condition=sys.version_info < (3, 15), reason="WrappingContext.unwrap restores behaviour but rebuilds the code object instead of " "reinstating the original, so __code__ identity is not restored after a wrap->unwrap round-trip", ) From bed77fe28c34105acd603eb4b295cd01c3a5459d Mon Sep 17 00:00:00 2001 From: Vlad Scherbich Date: Thu, 23 Jul 2026 10:30:28 -0400 Subject: [PATCH 03/13] typing: annotate monitoring tests --- tests/internal/test_monitoring.py | 86 ++++++++++++++++++------------- 1 file changed, 50 insertions(+), 36 deletions(-) diff --git a/tests/internal/test_monitoring.py b/tests/internal/test_monitoring.py index 77d4ee85e60..816ffd09478 100644 --- a/tests/internal/test_monitoring.py +++ b/tests/internal/test_monitoring.py @@ -8,8 +8,11 @@ """ import sys +from types import CodeType from typing import TYPE_CHECKING from typing import Any +from typing import Callable +from typing import Iterator import pytest @@ -34,34 +37,37 @@ # `_E = sys.monitoring.events` has an indeterminate type when mypy analyzes it # under a pre-3.15 Python version. _E: Any = getattr(monitoring, "_E") +_sys_monitoring: Any = getattr(sys, "monitoring", None) class UnwindHandler(monitoring.MonitoringEventHandler): - def __init__(self): - self.unwinds = [] + def __init__(self) -> None: + self.unwinds: list[tuple[CodeType, BaseException]] = [] - def on_py_unwind(self, code, instruction_offset, exception): + def on_py_unwind(self, code: CodeType, instruction_offset: int, exception: BaseException) -> None: self.unwinds.append((code, exception)) class StartAndUnwindHandler(monitoring.MonitoringEventHandler): - def __init__(self): - self.started = False - self.unwound = False + def __init__(self) -> None: + self.started: bool = False + self.unwound: bool = False - def on_py_start(self, code, instruction_offset): + def on_py_start(self, code: CodeType, instruction_offset: int) -> None: self.started = True - def on_py_unwind(self, code, instruction_offset, exception): + def on_py_unwind(self, code: CodeType, instruction_offset: int, exception: BaseException) -> None: self.unwound = True @pytest.fixture -def registered(): +def registered() -> Iterator[ + Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler] +]: """Register a handler for a code object and always unregister afterwards.""" - registrations = [] + registrations: list[tuple[CodeType, monitoring.MonitoringEventHandler]] = [] - def _register(code, handler): + def _register(code: CodeType, handler: monitoring.MonitoringEventHandler) -> monitoring.MonitoringEventHandler: monitoring.register(code, handler) registrations.append((code, handler)) return handler @@ -72,38 +78,42 @@ def _register(code, handler): monitoring.unregister(code, handler) -def test_register_unwind_handler_does_not_raise(registered): +def test_register_unwind_handler_does_not_raise( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: """Regression: registering a PY_UNWIND-only handler must not raise. Before the fix, ``register`` passed PY_UNWIND to ``set_local_events`` which raised ``ValueError: invalid local event set``. """ - def boom(): + def boom() -> None: raise ValueError("boom") registered(boom.__code__, UnwindHandler()) -def test_unwind_enabled_globally_not_locally(registered): +def test_unwind_enabled_globally_not_locally( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: """PY_UNWIND must be a global event; it must not appear in local events.""" - def boom(): + def boom() -> None: raise ValueError("boom") registered(boom.__code__, UnwindHandler()) - tool_id = monitoring._tool_id + tool_id: int | None = monitoring._tool_id assert tool_id is not None - local_events = sys.monitoring.get_local_events(tool_id, boom.__code__) - global_events = sys.monitoring.get_events(tool_id) + local_events: int = _sys_monitoring.get_local_events(tool_id, boom.__code__) + global_events: int = _sys_monitoring.get_events(tool_id) assert not (local_events & _E.PY_UNWIND), "PY_UNWIND must not be a local event" assert global_events & _E.PY_UNWIND, "PY_UNWIND must be enabled globally" -def test_on_py_unwind_does_not_disable_unregistered_code(): +def test_on_py_unwind_does_not_disable_unregistered_code() -> None: """Regression: the unwind callback must return None (never DISABLE). PY_UNWIND fires for every unwinding frame, including code with no handler. @@ -111,20 +121,22 @@ def test_on_py_unwind_does_not_disable_unregistered_code(): location with no re-arm path, so unregistered code must yield None. """ - def unrelated(): + def unrelated() -> None: pass - result = monitoring._on_py_unwind(unrelated.__code__, 0, ValueError("x")) + result: object | None = monitoring._on_py_unwind(unrelated.__code__, 0, ValueError("x")) assert result is None, "unregistered code must not be disabled (must return None, not DISABLE)" -def test_unwind_callback_fires_on_exception(registered): +def test_unwind_callback_fires_on_exception( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: """A registered handler receives on_py_unwind when its code unwinds.""" - def boom(): + def boom() -> None: raise ValueError("kaboom") - handler = registered(boom.__code__, UnwindHandler()) + handler: UnwindHandler = registered(boom.__code__, UnwindHandler()) # type: ignore[assignment] with pytest.raises(ValueError): boom() @@ -134,41 +146,43 @@ def boom(): ) -def test_unregister_disables_global_unwind(): +def test_unregister_disables_global_unwind() -> None: """Unregistering the last unwind handler clears the global PY_UNWIND event.""" - def boom(): + def boom() -> None: raise ValueError("boom") - handler = UnwindHandler() + handler: UnwindHandler = UnwindHandler() monitoring.register(boom.__code__, handler) - tool_id = monitoring._tool_id + tool_id: int | None = monitoring._tool_id assert tool_id is not None - assert sys.monitoring.get_events(tool_id) & _E.PY_UNWIND + assert _sys_monitoring.get_events(tool_id) & _E.PY_UNWIND monitoring.unregister(boom.__code__, handler) - assert not (sys.monitoring.get_events(tool_id) & _E.PY_UNWIND), ( + assert not (_sys_monitoring.get_events(tool_id) & _E.PY_UNWIND), ( "global PY_UNWIND should be disabled once no handlers need it" ) -def test_mixed_local_and_global_events(registered): +def test_mixed_local_and_global_events( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: """A handler overriding both PY_START and PY_UNWIND gets each at its scope.""" - def fn(): + def fn() -> None: raise ValueError("mixed") - handler = registered(fn.__code__, StartAndUnwindHandler()) + handler: StartAndUnwindHandler = registered(fn.__code__, StartAndUnwindHandler()) # type: ignore[assignment] - tool_id = monitoring._tool_id + tool_id: int | None = monitoring._tool_id assert tool_id is not None - local_events = sys.monitoring.get_local_events(tool_id, fn.__code__) + local_events: int = _sys_monitoring.get_local_events(tool_id, fn.__code__) assert local_events & _E.PY_START, "PY_START must be a local event" assert not (local_events & _E.PY_UNWIND), "PY_UNWIND must not be local" - assert sys.monitoring.get_events(tool_id) & _E.PY_UNWIND, "PY_UNWIND must be global" + assert _sys_monitoring.get_events(tool_id) & _E.PY_UNWIND, "PY_UNWIND must be global" with pytest.raises(ValueError): fn() From 3cf7dcc2295477030f63b9865a950b367ecacdd1 Mon Sep 17 00:00:00 2001 From: Vlad Scherbich Date: Thu, 23 Jul 2026 10:26:00 -0400 Subject: [PATCH 04/13] typing: annotate wrapping context for py3.15 --- .../internal/bytecode_injection/__init__.py | 24 +++++------ ddtrace/internal/wrapping/asyncs.py | 2 +- ddtrace/internal/wrapping/context.py | 2 +- ddtrace/internal/wrapping/generators.py | 2 +- .../bytecode_injection/test_injection.py | 12 +++++- tests/internal/test_wrapping.py | 29 +++++++------ tests/wrapping/conftest.py | 11 +++-- tests/wrapping/test_unwrap.py | 43 +++++++++++-------- 8 files changed, 72 insertions(+), 53 deletions(-) diff --git a/ddtrace/internal/bytecode_injection/__init__.py b/ddtrace/internal/bytecode_injection/__init__.py index 6ee36859165..c2ba59f3230 100644 --- a/ddtrace/internal/bytecode_injection/__init__.py +++ b/ddtrace/internal/bytecode_injection/__init__.py @@ -234,8 +234,8 @@ def _inject_hook(code: Bytecode, hook: HookType, lineno: int, arg: Any) -> None: # happens is with finally blocks, which are duplicated at the end of the # bytecode. locs: deque[tuple[int, str]] = deque() - last_lineno = None - instrs = set() + last_lineno: int | None = None + instrs: set[str] = set() for i, item in enumerate(code): if not isinstance(item, Instr): continue @@ -254,7 +254,7 @@ def _inject_hook(code: Bytecode, hook: HookType, lineno: int, arg: Any) -> None: if instrs == {"NOP"}: # If the line occurs on NOPs only, we instrument only the first one - last_instr = locs.pop() + last_instr: tuple[int, str] = locs.pop() locs.clear() locs.append(last_instr) elif "NOP" in instrs: @@ -285,11 +285,11 @@ def _eject_hook(code: Bytecode, hook: HookType, line: int, arg: Any) -> None: if not isinstance(item, Instr): continue try: - hook_op = code[i + _INJECT_HOOK_OPCODE_POS] - arg_op = code[i + _INJECT_ARG_OPCODE_POS] + hook_op: object = code[i + _INJECT_HOOK_OPCODE_POS] + arg_op: object = code[i + _INJECT_ARG_OPCODE_POS] if not isinstance(hook_op, Instr) or not isinstance(arg_op, Instr): continue - opcodes = [] + opcodes: list[str] = [] for j in range(i, i + len(_INJECT_HOOK_OPCODES)): op = code[j] if not isinstance(op, Instr): @@ -322,9 +322,9 @@ def inject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoTyp Returns the list of hooks that failed to be injected. """ - abstract_code = Bytecode.from_code(get_function_code(f)) + abstract_code: Bytecode = Bytecode.from_code(get_function_code(f)) - failed = [] + failed: list[HookInfoType] = [] for hook, line, arg in hooks: try: _inject_hook(abstract_code, hook, line, arg) @@ -344,9 +344,9 @@ def eject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoType Returns the list of hooks that failed to be ejected. """ - abstract_code = Bytecode.from_code(f.__code__) + abstract_code: Bytecode = Bytecode.from_code(f.__code__) - failed = [] + failed: list[HookInfoType] = [] for hook, line, arg in hooks: try: _eject_hook(abstract_code, hook, line, arg) @@ -365,7 +365,7 @@ def inject_hook(f: FunctionType, hook: HookType, line: int, arg: Any) -> Functio argument. The latter is also used as an identifier for the hook. This should be kept in case the hook needs to be removed. """ - abstract_code = Bytecode.from_code(f.__code__) + abstract_code: Bytecode = Bytecode.from_code(f.__code__) _inject_hook(abstract_code, hook, line, arg) @@ -379,7 +379,7 @@ def eject_hook(f: FunctionType, hook: HookType, line: int, arg: Any) -> Function The hook is identified by its line number and the argument passed to the hook. """ - abstract_code = Bytecode.from_code(f.__code__) + abstract_code: Bytecode = Bytecode.from_code(f.__code__) _eject_hook(abstract_code, hook, line, arg) diff --git a/ddtrace/internal/wrapping/asyncs.py b/ddtrace/internal/wrapping/asyncs.py index 10989b91f5c..58264f0865d 100644 --- a/ddtrace/internal/wrapping/asyncs.py +++ b/ddtrace/internal/wrapping/asyncs.py @@ -53,7 +53,7 @@ def _ensure_common_constant_none() -> None: COROUTINE_ASSEMBLY = Assembly() ASYNC_GEN_ASSEMBLY = Assembly() -ASYNC_HEAD_ASSEMBLY = None +ASYNC_HEAD_ASSEMBLY: Assembly | None = None if PY >= (3, 15): _ensure_common_constant_none() diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index f211e4c4e34..7b7b0e3ed33 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -590,7 +590,7 @@ def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: if sys.version_info >= (3, 15): from ddtrace.internal.monitoring import MonitoringEventHandler as _MonitoringEventHandler - _UWC_BASES: tuple = (BaseWrappingContext, _MonitoringEventHandler) + _UWC_BASES: tuple[type, ...] = (BaseWrappingContext, _MonitoringEventHandler) else: _UWC_BASES = (BaseWrappingContext,) diff --git a/ddtrace/internal/wrapping/generators.py b/ddtrace/internal/wrapping/generators.py index d88341c818b..d036e37953c 100644 --- a/ddtrace/internal/wrapping/generators.py +++ b/ddtrace/internal/wrapping/generators.py @@ -31,7 +31,7 @@ # return # ----------------------------------------------------------------------------- GENERATOR_ASSEMBLY = Assembly() -GENERATOR_HEAD_ASSEMBLY = None +GENERATOR_HEAD_ASSEMBLY: Assembly | None = None if PY >= (3, 15): GENERATOR_HEAD_ASSEMBLY = Assembly() diff --git a/tests/internal/bytecode_injection/test_injection.py b/tests/internal/bytecode_injection/test_injection.py index 82b6bb51ed6..949754a2864 100644 --- a/tests/internal/bytecode_injection/test_injection.py +++ b/tests/internal/bytecode_injection/test_injection.py @@ -1,9 +1,14 @@ +from collections.abc import Iterator from contextlib import contextmanager import sys +from types import FunctionType +from typing import Any +from typing import Optional from unittest import mock import pytest +from ddtrace.internal.bytecode_injection import HookType from ddtrace.internal.bytecode_injection import InvalidLine from ddtrace.internal.bytecode_injection import eject_hook from ddtrace.internal.bytecode_injection import eject_hooks @@ -13,7 +18,12 @@ @contextmanager -def injected_hook(f, hook, arg, line=None): +def injected_hook( + f: FunctionType, + hook: HookType, + arg: Any, + line: Optional[int] = None, +) -> Iterator[FunctionType]: code = f.__code__ if line is None: diff --git a/tests/internal/test_wrapping.py b/tests/internal/test_wrapping.py index d484261e944..bc2653bda61 100644 --- a/tests/internal/test_wrapping.py +++ b/tests/internal/test_wrapping.py @@ -5,6 +5,7 @@ import sys from types import CoroutineType from types import FunctionType +from typing import Any from typing import cast import pytest @@ -1383,20 +1384,20 @@ async def inner(): # --------------------------------------------------------------------------- -def test_wrapping_context_thread_concurrent(): +def test_wrapping_context_thread_concurrent() -> None: """Storage set by one thread must not bleed into a concurrent thread.""" import threading - results = {} - errors = [] + results: dict[int, int] = {} + errors: list[str] = [] class ThreadIsolationContext(DummyWrappingContext): - def __enter__(self): + def __enter__(self) -> "ThreadIsolationContext": super().__enter__() self.set("tid", threading.get_ident()) return self - def __return__(self, value): + def __return__(self, value: Any) -> Any: stored = self.get("tid") current = threading.get_ident() if stored != current: @@ -1412,11 +1413,11 @@ def foo(): barrier = threading.Barrier(10) - def run(): + def run() -> None: barrier.wait() foo() - threads = [threading.Thread(target=run) for _ in range(10)] + threads: list[threading.Thread] = [threading.Thread(target=run) for _ in range(10)] for t in threads: t.start() for t in threads: @@ -1434,25 +1435,25 @@ def run(): @pytest.mark.asyncio -async def test_wrapping_context_async_recursive(): +async def test_wrapping_context_async_recursive() -> None: """Each recursive coroutine call must have its own isolated storage slot.""" - values = [] + values: list[tuple[str, int]] = [] class AsyncRecursiveContext(DummyWrappingContext): - def __enter__(self): + def __enter__(self) -> "AsyncRecursiveContext": super().__enter__() - n = self.__frame__.f_locals["n"] + n: int = self.__frame__.f_locals["n"] self.set("n", n) values.append(("enter", n)) return self - def __return__(self, value): - n = self.__frame__.f_locals["n"] + def __return__(self, value: Any) -> Any: + n: int = self.__frame__.f_locals["n"] assert self.get("n") == n, f"storage mismatch: expected {n}, got {self.get('n')}" values.append(("return", n)) return super().__return__(value) - async def afactorial(n): + async def afactorial(n: int) -> int: if n == 0: return 1 return n * await afactorial(n - 1) diff --git a/tests/wrapping/conftest.py b/tests/wrapping/conftest.py index d91992831b1..a9c4fefc44f 100644 --- a/tests/wrapping/conftest.py +++ b/tests/wrapping/conftest.py @@ -15,7 +15,10 @@ import os import re import sys +from typing import Any +import _pytest.config +import _pytest.nodes import pytest from tests.wrapping.mechanisms import ALL_MECHANISMS @@ -31,7 +34,7 @@ collect_ignore.append(_name) -def pytest_configure(config): +def pytest_configure(config: _pytest.config.Config) -> None: config.addinivalue_line( "markers", "mechanism_specific: this test targets one wrapping mechanism by name (e.g. tracer.wrap() " @@ -39,7 +42,7 @@ def pytest_configure(config): ) -def pytest_collection_modifyitems(items): +def pytest_collection_modifyitems(items: list[_pytest.nodes.Item]) -> None: """Guardrail: every matrix test must run over all wrapping mechanisms. A test that forgets the ``mech`` argument (and does not parametrize it via @@ -83,7 +86,7 @@ def pytest_collection_modifyitems(items): @pytest.fixture(autouse=True) -def _dummy_tracer(tracer): +def _dummy_tracer(tracer: Any) -> None: """Ensure every test in this suite uses a DummyWriter, not the NativeWriter. ``tracer_wrap`` cases call ``ddtrace.tracer.wrap()`` which creates real spans. @@ -95,7 +98,7 @@ def _dummy_tracer(tracer): @pytest.fixture(params=list(ALL_MECHANISMS.values()), ids=list(ALL_MECHANISMS)) -def mech(request): +def mech(request: pytest.FixtureRequest) -> Any: """The wrapping mechanism under test. Every test taking a ``mech`` argument is automatically run once per mechanism (internal_wrap, tracer_wrap, wrapt, wrapping_context). diff --git a/tests/wrapping/test_unwrap.py b/tests/wrapping/test_unwrap.py index 2df8352bdd5..0bc529bdc32 100644 --- a/tests/wrapping/test_unwrap.py +++ b/tests/wrapping/test_unwrap.py @@ -19,8 +19,12 @@ these opt out of the all-mechanisms ``mech`` guardrail. """ +from collections.abc import Callable import inspect import sys +from types import FunctionType +from typing import Any +from typing import cast import pytest @@ -32,11 +36,11 @@ pytestmark = pytest.mark.mechanism_specific -def _noop(wrapped, args, kwargs): +def _noop(wrapped: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: return wrapped(*args, **kwargs) -def _noop2(wrapped, args, kwargs): +def _noop2(wrapped: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: return wrapped(*args, **kwargs) @@ -47,44 +51,45 @@ class _NoopContext(WrappingContext): class _InternalRestore: """internal wrap()/unwrap() round-trip on a function (mutates in place).""" - def wrap(self, fn): + def wrap(self, fn: FunctionType) -> None: _internal_wrap(fn, _noop) - def unwrap(self, fn): + def unwrap(self, fn: FunctionType) -> None: _internal_unwrap(fn, _noop) class _ContextRestore: """WrappingContext wrap()/unwrap(); the instance is retained to unwrap with.""" - def __init__(self): - self._ctx = {} + def __init__(self) -> None: + self._ctx: dict[FunctionType, WrappingContext] = {} - def wrap(self, fn): + def wrap(self, fn: FunctionType) -> None: ctx = _NoopContext(fn) ctx.wrap() self._ctx[fn] = ctx - def unwrap(self, fn): + def unwrap(self, fn: FunctionType) -> None: self._ctx.pop(fn).unwrap() @pytest.fixture(params=[_InternalRestore, _ContextRestore], ids=["internal_wrap", "wrapping_context"]) -def restore(request): - return request.param() +def restore(request: pytest.FixtureRequest) -> _InternalRestore | _ContextRestore: + return cast(_InternalRestore | _ContextRestore, request.param()) -def test_roundtrip_restores_behavior_and_signature(restore): - def f(a, b=2, *, k=3): +def test_roundtrip_restores_behavior_and_signature(restore: _InternalRestore | _ContextRestore) -> None: + def f(a: Any, b: int = 2, *, k: int = 3) -> tuple[Any, int, int]: return (a, b, k) - sig = str(inspect.signature(f)) - assert f(1) == (1, 2, 3) - restore.wrap(f) - assert f(1) == (1, 2, 3) # transparent while wrapped - restore.unwrap(f) - assert f(1) == (1, 2, 3) # behaviour restored after unwrap - assert str(inspect.signature(f)) == sig + fn = cast(FunctionType, f) + sig = str(inspect.signature(fn)) + assert fn(1) == (1, 2, 3) + restore.wrap(fn) + assert fn(1) == (1, 2, 3) # transparent while wrapped + restore.unwrap(fn) + assert fn(1) == (1, 2, 3) # behaviour restored after unwrap + assert str(inspect.signature(fn)) == sig def test_internal_wrap_unwrap_reinstates_original_code(): From be123d3538a04e57d16ab04a02dacdf929fcf767 Mon Sep 17 00:00:00 2001 From: Vlad Scherbich Date: Thu, 23 Jul 2026 11:24:39 -0400 Subject: [PATCH 05/13] fix(wrapping): use Optional for 3.9-compatible module annotations --- ddtrace/internal/wrapping/asyncs.py | 3 ++- ddtrace/internal/wrapping/generators.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/ddtrace/internal/wrapping/asyncs.py b/ddtrace/internal/wrapping/asyncs.py index 58264f0865d..396e9c20577 100644 --- a/ddtrace/internal/wrapping/asyncs.py +++ b/ddtrace/internal/wrapping/asyncs.py @@ -1,6 +1,7 @@ import enum import sys from types import CodeType +from typing import Optional import bytecode as bc @@ -53,7 +54,7 @@ def _ensure_common_constant_none() -> None: COROUTINE_ASSEMBLY = Assembly() ASYNC_GEN_ASSEMBLY = Assembly() -ASYNC_HEAD_ASSEMBLY: Assembly | None = None +ASYNC_HEAD_ASSEMBLY: Optional[Assembly] = None if PY >= (3, 15): _ensure_common_constant_none() diff --git a/ddtrace/internal/wrapping/generators.py b/ddtrace/internal/wrapping/generators.py index d036e37953c..fc80c3d37d7 100644 --- a/ddtrace/internal/wrapping/generators.py +++ b/ddtrace/internal/wrapping/generators.py @@ -1,5 +1,6 @@ import sys from types import CodeType +from typing import Optional import bytecode as bc @@ -31,7 +32,7 @@ # return # ----------------------------------------------------------------------------- GENERATOR_ASSEMBLY = Assembly() -GENERATOR_HEAD_ASSEMBLY: Assembly | None = None +GENERATOR_HEAD_ASSEMBLY: Optional[Assembly] = None if PY >= (3, 15): GENERATOR_HEAD_ASSEMBLY = Assembly() From 53a0f7ceaef00b5cfb309130a834b2998189827b Mon Sep 17 00:00:00 2001 From: Vlad Scherbich Date: Thu, 23 Jul 2026 16:39:54 -0400 Subject: [PATCH 06/13] first round of bot comments --- ddtrace/internal/monitoring.py | 241 +++++++++++++++++++++++------- mypy.ini | 4 + tests/internal/test_monitoring.py | 8 - 3 files changed, 190 insertions(+), 63 deletions(-) diff --git a/ddtrace/internal/monitoring.py b/ddtrace/internal/monitoring.py index 768a7aebb84..5e4b705c4cf 100644 --- a/ddtrace/internal/monitoring.py +++ b/ddtrace/internal/monitoring.py @@ -7,7 +7,7 @@ Only the events corresponding to overridden handler methods are enabled, so a handler that only overrides ``on_py_start`` pays no cost for the other -two events. +events. The handler instance itself serves as the registration key: pass the same object to :func:`unregister` to remove it. @@ -16,17 +16,22 @@ from abc import ABC import sys from types import CodeType +from typing import Any +from typing import Iterator from typing import NamedTuple from typing import Optional import weakref +from ddtrace.internal.logger import get_logger from ddtrace.internal.threads import Lock if sys.version_info < (3, 15): raise ImportError("ddtrace.internal.monitoring requires Python 3.15+") -_E = sys.monitoring.events # type: ignore[unreachable] +log = get_logger(__name__) + +_E = sys.monitoring.events DISABLE = sys.monitoring.DISABLE _DISABLE = DISABLE @@ -38,16 +43,78 @@ _LOCAL_EVENTS = _E.PY_START | _E.PY_RETURN | _E.LINE _GLOBAL_EVENTS = _E.PY_UNWIND +# CPython tool IDs (see ddtrace/profiling/collector/_exception.pyx): +# 0 DEBUGGER_ID, 1 COVERAGE_ID, 2 PROFILER_ID, 3 handled exceptions, 5 OPTIMIZER_ID. +# Never claim reserved/unmigrated slots; prefer 4 then 3. +_HANDLED_EXCEPTIONS_TOOL_ID = 3 +_CANDIDATE_TOOL_IDS = (4, _HANDLED_EXCEPTIONS_TOOL_ID) + _tool_id: Optional[int] = None _tool_lock = Lock() # The set of global-only events currently enabled via sys.monitoring.set_events. _active_global_events: int = 0 -_registry: "weakref.WeakKeyDictionary[CodeType, dict[int, _Entry]]" = weakref.WeakKeyDictionary() _registry_lock = Lock() +class _IdentityWeakKeyDictionary: + """Weak mapping keyed by object identity (not equality). + + Unlike ``weakref.WeakKeyDictionary``, lookups use ``is`` rather than + ``CodeType.__eq__``, so distinct code objects for the same source remain + separate entries. + """ + + __slots__ = ("_data", "_on_remove") + + def __init__(self, on_remove: Optional[Any] = None) -> None: + self._data: dict[int, tuple[weakref.ref[Any], Any]] = {} + self._on_remove = on_remove + + def _make_remove(self, key_id: int) -> Any: + def remove(_ref: weakref.ref[Any], selfref: weakref.ref[Any] = weakref.ref(self)) -> None: + self = selfref() + if self is None: + return + if self._data.pop(key_id, None) is not None and self._on_remove is not None: + self._on_remove() + + return remove + + def get(self, key: CodeType, default: Any = None) -> Any: + item = self._data.get(id(key)) + if item is None: + return default + ref, value = item + if ref() is key: + return value + return default + + def __setitem__(self, key: CodeType, value: Any) -> None: + key_id = id(key) + self._data[key_id] = (weakref.ref(key, self._make_remove(key_id)), value) + + def __delitem__(self, key: CodeType) -> None: + key_id = id(key) + if key_id not in self._data: + raise KeyError(key) + del self._data[key_id] + + def values(self) -> Iterator[Any]: + for ref, value in self._data.values(): + if ref() is not None: + yield value + + +def _on_registry_entry_removed() -> None: + with _registry_lock: + _recompute_global_events() + + +_registry: _IdentityWeakKeyDictionary = _IdentityWeakKeyDictionary(on_remove=_on_registry_entry_removed) + + class MonitoringEventHandler(ABC): """Base class for sys.monitoring event handlers. @@ -90,18 +157,39 @@ def _events_for_handler(handler: MonitoringEventHandler) -> int: return events -def _events_for(entries: "dict[int, _Entry]") -> int: - events: int = 0 - for e in list(entries.values()): - events |= e.events - return events - - class _Entry(NamedTuple): handler: MonitoringEventHandler events: int # pre-computed from _events_for_handler +class _CodeHandlers: + """Per-code handler table with a pre-built snapshot for hot-path dispatch.""" + + __slots__ = ("_by_handler", "snapshot") + + def __init__(self) -> None: + self._by_handler: dict[int, _Entry] = {} + self.snapshot: tuple[_Entry, ...] = () + + def __len__(self) -> int: + return len(self._by_handler) + + def set_handler(self, handler_id: int, entry: _Entry) -> None: + self._by_handler[handler_id] = entry + self.snapshot = tuple(self._by_handler.values()) + + def pop_handler(self, handler_id: int) -> None: + self._by_handler.pop(handler_id, None) + self.snapshot = tuple(self._by_handler.values()) + + +def _events_for(handlers: _CodeHandlers) -> int: + events: int = 0 + for e in handlers.snapshot: + events |= e.events + return events + + def _setup() -> int: """Claim a free tool ID and install the global callbacks (idempotent).""" global _tool_id @@ -113,7 +201,7 @@ def _setup() -> int: if _tool_id is not None: return _tool_id - for tid in range(5, -1, -1): + for tid in _CANDIDATE_TOOL_IDS: try: sys.monitoring.use_tool_id(tid, "ddtrace") _tool_id = tid @@ -132,55 +220,83 @@ def _setup() -> int: # --------------------------------------------------------------------------- -# Hot-path callbacks — no lock, no allocation +# Hot-path callbacks — no lock; iterate a pre-built handler snapshot tuple # --------------------------------------------------------------------------- +def _dispatch_start(code: CodeType, instruction_offset: int, entry: _Entry) -> None: + entry.handler.on_py_start(code, instruction_offset) + + +def _dispatch_return(code: CodeType, instruction_offset: int, retval: object, entry: _Entry) -> None: + entry.handler.on_py_return(code, instruction_offset, retval) + + +def _dispatch_unwind(code: CodeType, instruction_offset: int, exception: BaseException, entry: _Entry) -> None: + entry.handler.on_py_unwind(code, instruction_offset, exception) + + +def _dispatch_line(code: CodeType, line_number: int, entry: _Entry) -> Optional[object]: + return entry.handler.on_py_line(code, line_number) + + def _on_py_start(code: CodeType, instruction_offset: int) -> Optional[object]: - entries: Optional[dict[int, _Entry]] = _registry.get(code) - if not entries: + handlers: Optional[_CodeHandlers] = _registry.get(code) + if not handlers or not handlers.snapshot: return _DISABLE - for e in list(entries.values()): + for e in handlers.snapshot: if e.events & _E.PY_START: - e.handler.on_py_start(code, instruction_offset) + try: + _dispatch_start(code, instruction_offset, e) + except Exception: + log.warning("monitoring PY_START handler failed", exc_info=True) return None def _on_py_return(code: CodeType, instruction_offset: int, retval: object) -> Optional[object]: - entries: Optional[dict[int, _Entry]] = _registry.get(code) - if not entries: + handlers: Optional[_CodeHandlers] = _registry.get(code) + if not handlers or not handlers.snapshot: return _DISABLE - for e in list(entries.values()): + for e in handlers.snapshot: if e.events & _E.PY_RETURN: - e.handler.on_py_return(code, instruction_offset, retval) + try: + _dispatch_return(code, instruction_offset, retval, e) + except Exception: + log.warning("monitoring PY_RETURN handler failed", exc_info=True) return None def _on_py_unwind(code: CodeType, instruction_offset: int, exception: BaseException) -> Optional[object]: - entries: Optional[dict[int, _Entry]] = _registry.get(code) + handlers: Optional[_CodeHandlers] = _registry.get(code) # PY_UNWIND is a global event, so this callback fires for every unwinding # frame regardless of registration. We must not return DISABLE for # unregistered code: doing so would permanently disable the event for that # code location, and a later register() would not re-arm it (we never call # restart_events). Unwinding only happens on exceptions, so the extra lookup # cost on this already-slow path is negligible. - if not entries: + if not handlers or not handlers.snapshot: return None - for e in list(entries.values()): + for e in handlers.snapshot: if e.events & _E.PY_UNWIND: - e.handler.on_py_unwind(code, instruction_offset, exception) + try: + _dispatch_unwind(code, instruction_offset, exception, e) + except Exception: + log.warning("monitoring PY_UNWIND handler failed", exc_info=True) return None def _on_py_line(code: CodeType, line_number: int) -> Optional[object]: - entries: Optional[dict[int, _Entry]] = _registry.get(code) - if not entries: + handlers: Optional[_CodeHandlers] = _registry.get(code) + if not handlers or not handlers.snapshot: return _DISABLE disable: bool = True - for e in list(entries.values()): + for e in handlers.snapshot: if e.events & _E.LINE: - if e.handler.on_py_line(code, line_number) is not _DISABLE: - disable = False + try: + if _dispatch_line(code, line_number, e) is not _DISABLE: + disable = False + except Exception: + log.warning("monitoring LINE handler failed", exc_info=True) return _DISABLE if disable else None @@ -204,14 +320,29 @@ def _recompute_global_events() -> None: if _tool_id is None: return needed: int = 0 - for entries in _registry.values(): - needed |= _events_for(entries) + for handlers in _registry.values(): + needed |= _events_for(handlers) needed &= _GLOBAL_EVENTS if needed != _active_global_events: _active_global_events = needed sys.monitoring.set_events(_tool_id, needed) +def _set_local_events(tool_id: int, code: CodeType, events: int) -> None: + sys.monitoring.set_local_events(tool_id, code, events) + + +def _rearm_local_events(tool_id: int, code: CodeType, events: int) -> None: + # A DISABLE returned from a per-line callback is sticky until the monitored + # event set changes or restart_events() is called. Re-applying the same + # local events does not clear it; toggling local events off and back on + # re-arms only this tool's DISABLE marks for *code* without the global + # restart_events() call that would reset other tools' disabled-event + # bookkeeping (coverage.py). + _set_local_events(tool_id, code, 0) + _set_local_events(tool_id, code, events) + + def register(code: CodeType, handler: MonitoringEventHandler) -> None: """Register a monitoring event handler for *code*. @@ -226,13 +357,20 @@ def register(code: CodeType, handler: MonitoringEventHandler) -> None: entry: _Entry = _Entry(handler, handler_events) with _registry_lock: - entries: Optional[dict[int, _Entry]] = _registry.get(code) - if entries is None: - _registry[code] = entries = {} - entries[id(handler)] = entry - all_events: int = _events_for(entries) - sys.monitoring.set_local_events(tool_id, code, all_events & _LOCAL_EVENTS) - _enable_global_events(all_events & _GLOBAL_EVENTS) + handlers: Optional[_CodeHandlers] = _registry.get(code) + if handlers is None: + _registry[code] = handlers = _CodeHandlers() + + had_line: bool = any(e.events & _E.LINE for e in handlers.snapshot) + handlers.set_handler(id(handler), entry) + local_events: int = _events_for(handlers) & _LOCAL_EVENTS + + if (handler_events & _E.LINE) and had_line: + _rearm_local_events(tool_id, code, local_events) + else: + _set_local_events(tool_id, code, local_events) + + _enable_global_events(_events_for(handlers) & _GLOBAL_EVENTS) def refresh(code: CodeType) -> None: @@ -242,34 +380,27 @@ def refresh(code: CodeType) -> None: disabled via a ``DISABLE`` return from :meth:`MonitoringEventHandler.on_py_line`. """ with _registry_lock: - entries: Optional[dict[int, _Entry]] = _registry.get(code) - if entries and _tool_id is not None: - events: int = _events_for(entries) & _LOCAL_EVENTS - # A DISABLE returned from a per-line callback is sticky until the - # monitored event set changes or restart_events() is called. - # Re-applying the same local events does not clear it; toggling - # local events off and back on re-arms only this tool's DISABLE - # marks for *code* without the global restart_events() call that - # would reset other tools' disabled-event bookkeeping (coverage.py). - sys.monitoring.set_local_events(_tool_id, code, 0) - sys.monitoring.set_local_events(_tool_id, code, events) + handlers: Optional[_CodeHandlers] = _registry.get(code) + if handlers and _tool_id is not None: + events: int = _events_for(handlers) & _LOCAL_EVENTS + _rearm_local_events(_tool_id, code, events) def unregister(code: CodeType, handler: MonitoringEventHandler) -> None: """Remove *handler* from the handlers registered for *code*.""" with _registry_lock: - existing: Optional[dict[int, _Entry]] = _registry.get(code) - if existing is None: + handlers: Optional[_CodeHandlers] = _registry.get(code) + if handlers is None: return - existing.pop(id(handler), None) + handlers.pop_handler(id(handler)) - if not existing: + if not handlers: del _registry[code] if _tool_id is not None: - sys.monitoring.set_local_events(_tool_id, code, 0) + _set_local_events(_tool_id, code, 0) else: assert _tool_id is not None # nosec - sys.monitoring.set_local_events(_tool_id, code, _events_for(existing) & _LOCAL_EVENTS) + _set_local_events(_tool_id, code, _events_for(handlers) & _LOCAL_EVENTS) _recompute_global_events() diff --git a/mypy.ini b/mypy.ini index 719ebad5e27..d982d1590af 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1390,6 +1390,10 @@ warn_return_any = false strict_equality = false warn_unreachable = false +[mypy-ddtrace.internal.monitoring] +disallow_any_generics = false +warn_unreachable = false + [mypy-ddtrace.internal.native] no_implicit_reexport = false diff --git a/tests/internal/test_monitoring.py b/tests/internal/test_monitoring.py index 816ffd09478..0ecbd24b074 100644 --- a/tests/internal/test_monitoring.py +++ b/tests/internal/test_monitoring.py @@ -25,14 +25,6 @@ else: monitoring = pytest.importorskip("ddtrace.internal.monitoring") -# TODO(py-315): the full 3.15 monitoring/profiling stack (PR #17624 and its -# dependencies) is not yet enabled on this branch. Skip on 3.15 for now; the -# stacked PRs remove this mark and these tests run and pass on 3.15. -pytestmark = pytest.mark.skipif( - sys.version_info >= (3, 15), - reason="TODO(py-315): enable once the 3.15 monitoring stack lands (PR #17624 + deps)", -) - # Fetched via getattr so the type checker treats it as Any: the source module's # `_E = sys.monitoring.events` has an indeterminate type when mypy analyzes it # under a pre-3.15 Python version. From 3d5230424ff15942ebb1eab6cd4c4d4a35a45e26 Mon Sep 17 00:00:00 2001 From: Vlad Scherbich Date: Thu, 23 Jul 2026 12:02:36 -0400 Subject: [PATCH 07/13] fix(tests): use Union in test_unwrap for Python 3.9 --- tests/wrapping/test_unwrap.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/wrapping/test_unwrap.py b/tests/wrapping/test_unwrap.py index 0bc529bdc32..12a3e05da8d 100644 --- a/tests/wrapping/test_unwrap.py +++ b/tests/wrapping/test_unwrap.py @@ -24,6 +24,7 @@ import sys from types import FunctionType from typing import Any +from typing import Union from typing import cast import pytest @@ -73,12 +74,15 @@ def unwrap(self, fn: FunctionType) -> None: self._ctx.pop(fn).unwrap() +_Restore = Union[_InternalRestore, _ContextRestore] + + @pytest.fixture(params=[_InternalRestore, _ContextRestore], ids=["internal_wrap", "wrapping_context"]) -def restore(request: pytest.FixtureRequest) -> _InternalRestore | _ContextRestore: - return cast(_InternalRestore | _ContextRestore, request.param()) +def restore(request: pytest.FixtureRequest) -> _Restore: + return cast(_Restore, request.param()) -def test_roundtrip_restores_behavior_and_signature(restore: _InternalRestore | _ContextRestore) -> None: +def test_roundtrip_restores_behavior_and_signature(restore: _Restore) -> None: def f(a: Any, b: int = 2, *, k: int = 3) -> tuple[Any, int, int]: return (a, b, k) From bc9465df8119f0e0656410eb46db167b0f1624ac Mon Sep 17 00:00:00 2001 From: Vlad Scherbich Date: Thu, 23 Jul 2026 16:54:00 -0400 Subject: [PATCH 08/13] second round of bot comments --- ddtrace/internal/monitoring.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ddtrace/internal/monitoring.py b/ddtrace/internal/monitoring.py index 5e4b705c4cf..78abb67346b 100644 --- a/ddtrace/internal/monitoring.py +++ b/ddtrace/internal/monitoring.py @@ -102,7 +102,7 @@ def __delitem__(self, key: CodeType) -> None: del self._data[key_id] def values(self) -> Iterator[Any]: - for ref, value in self._data.values(): + for ref, value in list(self._data.values()): if ref() is not None: yield value @@ -137,7 +137,12 @@ def on_py_unwind(self, code: CodeType, instruction_offset: int, exception: BaseE pass def on_py_line(self, code: CodeType, line_number: int) -> Optional[object]: - """Return ``sys.monitoring.DISABLE`` to stop future events on this line.""" + """Return ``sys.monitoring.DISABLE`` to request disabling future LINE events. + + The multiplexer forwards ``DISABLE`` to CPython only when every registered + LINE handler for this code object returns it. If any handler returns a + different value, LINE events continue for that location. + """ return None From c611c201e6962c2059420d093ff5de9b9d9fceb3 Mon Sep 17 00:00:00 2001 From: Vlad Scherbich Date: Thu, 23 Jul 2026 17:09:45 -0400 Subject: [PATCH 09/13] final round of bot comments --- ddtrace/internal/monitoring.py | 88 +++++++------------------------ tests/internal/test_monitoring.py | 51 +++++++----------- 2 files changed, 39 insertions(+), 100 deletions(-) diff --git a/ddtrace/internal/monitoring.py b/ddtrace/internal/monitoring.py index 78abb67346b..7bd63d8e28c 100644 --- a/ddtrace/internal/monitoring.py +++ b/ddtrace/internal/monitoring.py @@ -17,7 +17,6 @@ import sys from types import CodeType from typing import Any -from typing import Iterator from typing import NamedTuple from typing import Optional import weakref @@ -35,26 +34,21 @@ DISABLE = sys.monitoring.DISABLE _DISABLE = DISABLE -# sys.monitoring distinguishes "local" events, which can be enabled per code -# object via set_local_events, from events that can only be enabled globally -# via set_events. PY_START, PY_RETURN and LINE are local; PY_UNWIND is not and -# must be enabled globally, otherwise set_local_events raises -# "ValueError: invalid local event set". -_LOCAL_EVENTS = _E.PY_START | _E.PY_RETURN | _E.LINE -_GLOBAL_EVENTS = _E.PY_UNWIND +# On Python 3.15+, PY_UNWIND is a per-code "other" event and can be enabled via +# set_local_events alongside PY_START/PY_RETURN/LINE. +_LOCAL_EVENTS = _E.PY_START | _E.PY_RETURN | _E.LINE | _E.PY_UNWIND # CPython tool IDs (see ddtrace/profiling/collector/_exception.pyx): # 0 DEBUGGER_ID, 1 COVERAGE_ID, 2 PROFILER_ID, 3 handled exceptions, 5 OPTIMIZER_ID. -# Never claim reserved/unmigrated slots; prefer 4 then 3. -_HANDLED_EXCEPTIONS_TOOL_ID = 3 -_CANDIDATE_TOOL_IDS = (4, _HANDLED_EXCEPTIONS_TOOL_ID) +# Slot 4 is today's exception-profiler slot; the 3.15 stack migrates that +# collector onto this multiplexer. Never claim 0/1/2/3/5 or a slot owned by +# another tool name. +_MULTIPLEXER_TOOL_NAME = "ddtrace" +_CANDIDATE_TOOL_IDS = (4,) _tool_id: Optional[int] = None _tool_lock = Lock() -# The set of global-only events currently enabled via sys.monitoring.set_events. -_active_global_events: int = 0 - _registry_lock = Lock() @@ -66,19 +60,14 @@ class _IdentityWeakKeyDictionary: separate entries. """ - __slots__ = ("_data", "_on_remove") + __slots__ = ("_data",) - def __init__(self, on_remove: Optional[Any] = None) -> None: + def __init__(self) -> None: self._data: dict[int, tuple[weakref.ref[Any], Any]] = {} - self._on_remove = on_remove def _make_remove(self, key_id: int) -> Any: - def remove(_ref: weakref.ref[Any], selfref: weakref.ref[Any] = weakref.ref(self)) -> None: - self = selfref() - if self is None: - return - if self._data.pop(key_id, None) is not None and self._on_remove is not None: - self._on_remove() + def remove(_ref: weakref.ref[Any]) -> None: + self._data.pop(key_id, None) return remove @@ -101,18 +90,8 @@ def __delitem__(self, key: CodeType) -> None: raise KeyError(key) del self._data[key_id] - def values(self) -> Iterator[Any]: - for ref, value in list(self._data.values()): - if ref() is not None: - yield value - - -def _on_registry_entry_removed() -> None: - with _registry_lock: - _recompute_global_events() - -_registry: _IdentityWeakKeyDictionary = _IdentityWeakKeyDictionary(on_remove=_on_registry_entry_removed) +_registry: _IdentityWeakKeyDictionary = _IdentityWeakKeyDictionary() class MonitoringEventHandler(ABC): @@ -207,8 +186,12 @@ def _setup() -> int: return _tool_id for tid in _CANDIDATE_TOOL_IDS: + existing = sys.monitoring.get_tool(tid) + if existing is not None and existing != _MULTIPLEXER_TOOL_NAME: + continue try: - sys.monitoring.use_tool_id(tid, "ddtrace") + if existing is None: + sys.monitoring.use_tool_id(tid, _MULTIPLEXER_TOOL_NAME) _tool_id = tid break except ValueError: @@ -273,14 +256,8 @@ def _on_py_return(code: CodeType, instruction_offset: int, retval: object) -> Op def _on_py_unwind(code: CodeType, instruction_offset: int, exception: BaseException) -> Optional[object]: handlers: Optional[_CodeHandlers] = _registry.get(code) - # PY_UNWIND is a global event, so this callback fires for every unwinding - # frame regardless of registration. We must not return DISABLE for - # unregistered code: doing so would permanently disable the event for that - # code location, and a later register() would not re-arm it (we never call - # restart_events). Unwinding only happens on exceptions, so the extra lookup - # cost on this already-slow path is negligible. if not handlers or not handlers.snapshot: - return None + return _DISABLE for e in handlers.snapshot: if e.events & _E.PY_UNWIND: try: @@ -310,29 +287,6 @@ def _on_py_line(code: CodeType, line_number: int) -> Optional[object]: # --------------------------------------------------------------------------- -def _enable_global_events(events: int) -> None: - """Ensure the given global-only *events* are enabled (additive).""" - global _active_global_events - if events & ~_active_global_events: - _active_global_events |= events - assert _tool_id is not None # nosec - sys.monitoring.set_events(_tool_id, _active_global_events) - - -def _recompute_global_events() -> None: - """Re-derive the set of global-only events from the current registry.""" - global _active_global_events - if _tool_id is None: - return - needed: int = 0 - for handlers in _registry.values(): - needed |= _events_for(handlers) - needed &= _GLOBAL_EVENTS - if needed != _active_global_events: - _active_global_events = needed - sys.monitoring.set_events(_tool_id, needed) - - def _set_local_events(tool_id: int, code: CodeType, events: int) -> None: sys.monitoring.set_local_events(tool_id, code, events) @@ -375,8 +329,6 @@ def register(code: CodeType, handler: MonitoringEventHandler) -> None: else: _set_local_events(tool_id, code, local_events) - _enable_global_events(_events_for(handlers) & _GLOBAL_EVENTS) - def refresh(code: CodeType) -> None: """Re-apply local events for *code*, resetting any per-line DISABLE state. @@ -407,5 +359,3 @@ def unregister(code: CodeType, handler: MonitoringEventHandler) -> None: else: assert _tool_id is not None # nosec _set_local_events(_tool_id, code, _events_for(handlers) & _LOCAL_EVENTS) - - _recompute_global_events() diff --git a/tests/internal/test_monitoring.py b/tests/internal/test_monitoring.py index 0ecbd24b074..11c05e106b7 100644 --- a/tests/internal/test_monitoring.py +++ b/tests/internal/test_monitoring.py @@ -1,10 +1,7 @@ """Tests for ddtrace.internal.monitoring, the multiplexed sys.monitoring layer. -These focus on the local-vs-global event split for PY_UNWIND. PY_UNWIND is a -global-only sys.monitoring event: passing it to ``set_local_events`` raises -``ValueError: invalid local event set``. The module therefore enables PY_UNWIND -via ``set_events`` (global) and keeps PY_START/PY_RETURN/LINE per-code via -``set_local_events`` (local). +On Python 3.15+, PY_UNWIND is a per-code event and is enabled via +``set_local_events`` together with PY_START/PY_RETURN/LINE. """ import sys @@ -29,6 +26,7 @@ # `_E = sys.monitoring.events` has an indeterminate type when mypy analyzes it # under a pre-3.15 Python version. _E: Any = getattr(monitoring, "_E") +_DISABLE: Any = getattr(monitoring, "_DISABLE") _sys_monitoring: Any = getattr(sys, "monitoring", None) @@ -73,11 +71,7 @@ def _register(code: CodeType, handler: monitoring.MonitoringEventHandler) -> mon def test_register_unwind_handler_does_not_raise( registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], ) -> None: - """Regression: registering a PY_UNWIND-only handler must not raise. - - Before the fix, ``register`` passed PY_UNWIND to ``set_local_events`` which - raised ``ValueError: invalid local event set``. - """ + """Registering a PY_UNWIND-only handler enables the per-code unwind event.""" def boom() -> None: raise ValueError("boom") @@ -85,10 +79,10 @@ def boom() -> None: registered(boom.__code__, UnwindHandler()) -def test_unwind_enabled_globally_not_locally( +def test_unwind_enabled_locally( registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], ) -> None: - """PY_UNWIND must be a global event; it must not appear in local events.""" + """PY_UNWIND is enabled per code object on Python 3.15+.""" def boom() -> None: raise ValueError("boom") @@ -101,23 +95,18 @@ def boom() -> None: local_events: int = _sys_monitoring.get_local_events(tool_id, boom.__code__) global_events: int = _sys_monitoring.get_events(tool_id) - assert not (local_events & _E.PY_UNWIND), "PY_UNWIND must not be a local event" - assert global_events & _E.PY_UNWIND, "PY_UNWIND must be enabled globally" - + assert local_events & _E.PY_UNWIND, "PY_UNWIND must be a local event on 3.15+" + assert not (global_events & _E.PY_UNWIND), "PY_UNWIND must not be enabled globally" -def test_on_py_unwind_does_not_disable_unregistered_code() -> None: - """Regression: the unwind callback must return None (never DISABLE). - PY_UNWIND fires for every unwinding frame, including code with no handler. - Returning DISABLE would permanently disarm the global event for that code - location with no re-arm path, so unregistered code must yield None. - """ +def test_on_py_unwind_disables_unregistered_code() -> None: + """The unwind callback returns DISABLE when no handler is registered.""" def unrelated() -> None: pass result: object | None = monitoring._on_py_unwind(unrelated.__code__, 0, ValueError("x")) - assert result is None, "unregistered code must not be disabled (must return None, not DISABLE)" + assert result is _DISABLE def test_unwind_callback_fires_on_exception( @@ -138,8 +127,8 @@ def boom() -> None: ) -def test_unregister_disables_global_unwind() -> None: - """Unregistering the last unwind handler clears the global PY_UNWIND event.""" +def test_unregister_clears_local_unwind() -> None: + """Unregistering the last unwind handler clears the per-code PY_UNWIND event.""" def boom() -> None: raise ValueError("boom") @@ -149,19 +138,19 @@ def boom() -> None: tool_id: int | None = monitoring._tool_id assert tool_id is not None - assert _sys_monitoring.get_events(tool_id) & _E.PY_UNWIND + assert _sys_monitoring.get_local_events(tool_id, boom.__code__) & _E.PY_UNWIND monitoring.unregister(boom.__code__, handler) - assert not (_sys_monitoring.get_events(tool_id) & _E.PY_UNWIND), ( - "global PY_UNWIND should be disabled once no handlers need it" + assert not (_sys_monitoring.get_local_events(tool_id, boom.__code__) & _E.PY_UNWIND), ( + "local PY_UNWIND should be disabled once no handlers need it" ) -def test_mixed_local_and_global_events( +def test_mixed_local_events( registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], ) -> None: - """A handler overriding both PY_START and PY_UNWIND gets each at its scope.""" + """A handler overriding both PY_START and PY_UNWIND gets each as local events.""" def fn() -> None: raise ValueError("mixed") @@ -173,8 +162,8 @@ def fn() -> None: local_events: int = _sys_monitoring.get_local_events(tool_id, fn.__code__) assert local_events & _E.PY_START, "PY_START must be a local event" - assert not (local_events & _E.PY_UNWIND), "PY_UNWIND must not be local" - assert _sys_monitoring.get_events(tool_id) & _E.PY_UNWIND, "PY_UNWIND must be global" + assert local_events & _E.PY_UNWIND, "PY_UNWIND must be a local event on 3.15+" + assert not (_sys_monitoring.get_events(tool_id) & _E.PY_UNWIND), "PY_UNWIND must not be global" with pytest.raises(ValueError): fn() From 74b6d3df67b1f198c086d67ace6780f44b003063 Mon Sep 17 00:00:00 2001 From: Vlad Scherbich Date: Thu, 23 Jul 2026 17:15:09 -0400 Subject: [PATCH 10/13] fix(wrapping): harden 3.15 monitoring wrap registration Only evict stale _ctx_registry entries when the prior wrapped function was collected, and register sys.monitoring handlers before swapping __code__ to close the race where calls could run without callbacks. --- ddtrace/internal/wrapping/context.py | 30 ++++++++++++++++------------ 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index 7b7b0e3ed33..d1e1c306c00 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -707,17 +707,19 @@ def wrap(self) -> None: original_code: CodeType = get_function_code(f) with _ctx_registry_lock: if original_code in _ctx_registry: - # Allow wrapping a new function instance that shares the same - # code object as an already-wrapped (but orphaned) function. - # This happens when closures are re-created in a loop: each - # iteration produces a new function object but the same code - # object. If the new function is not itself registered, the old - # monitoring registration is stale and should be replaced. - if _fn_registry.get(f) is _ctx_registry[original_code]: + existing: "_UniversalWrappingContext" = _ctx_registry[original_code] + if _fn_registry.get(f) is existing: + raise ValueError("Function already wrapped") + # Only replace a registry entry when the prior wrapped function + # has been collected. A live function sharing this code object is + # still actively wrapped and must not be unregistered. + try: + existing.__wrapped__ + except RuntimeError: + _ctx_registry.pop(original_code) + _monitoring.unregister(original_code, existing) + else: raise ValueError("Function already wrapped") - # Stale entry: clean up old monitoring before re-registering. - old: "_UniversalWrappingContext" = _ctx_registry.pop(original_code) - _monitoring.unregister(original_code, old) # sys.monitoring dispatches per code object. Clone the code so # unwrapped siblings that share the same CodeType are not affected. @@ -726,8 +728,6 @@ def wrap(self) -> None: link_function_to_code(original_code, f) monitor_code: CodeType = original_code.replace() migrate_line_hooks(original_code, monitor_code) - set_function_code(f, monitor_code) - self._original_code = original_code _ctx_registry[monitor_code] = self _fn_registry[f] = self @@ -738,7 +738,11 @@ def wrap(self) -> None: weakref.ref(f), ) self._finalize.atexit = False - _monitoring.register(monitor_code, self) + # Register monitoring before swapping __code__ so no thread can + # observe monitor_code without an active handler. + _monitoring.register(monitor_code, self) + set_function_code(f, monitor_code) + self._original_code = original_code def unwrap(self) -> None: f: FunctionType = self.__wrapped__ From 8d1d128659c3aac5eb1fca512c9a52286d5ba9c9 Mon Sep 17 00:00:00 2001 From: "Gabriele N. Tornetta" Date: Fri, 14 Aug 2026 16:04:43 +0100 Subject: [PATCH 11/13] address review comments --- ddtrace/debugging/_function/store.py | 7 + .../internal/bytecode_injection/__init__.py | 2 +- ddtrace/internal/monitoring.py | 64 ++++----- ddtrace/internal/wrapping/__init__.py | 49 ++++++- ddtrace/internal/wrapping/context.py | 128 +++++++++++++++--- tests/wrapping/test_unwrap.py | 3 +- 6 files changed, 198 insertions(+), 55 deletions(-) diff --git a/ddtrace/debugging/_function/store.py b/ddtrace/debugging/_function/store.py index 8b2e4e98146..31bbb8d9eec 100644 --- a/ddtrace/debugging/_function/store.py +++ b/ddtrace/debugging/_function/store.py @@ -7,6 +7,7 @@ from ddtrace.debugging._function.discovery import FullyNamed from ddtrace.internal.bytecode_injection import HookInfoType from ddtrace.internal.bytecode_injection import HookType +from ddtrace.internal.bytecode_injection import eject_all_hooks from ddtrace.internal.bytecode_injection import eject_hooks from ddtrace.internal.bytecode_injection import inject_hooks from ddtrace.internal.wrapping import get_function_code @@ -90,4 +91,10 @@ def unwrap(self, function: FullyNamedContextWrappedFunction) -> None: def restore_all(self) -> None: """Restore all the patched functions to their original form.""" for function, code in self._code_map.items(): + # On 3.15+, line hooks are sys.monitoring registrations keyed by + # code object rather than injected bytecode, so restoring + # __code__ alone leaves them firing; eject_all_hooks is a no-op + # on older versions, where the registration lives in the bytecode + # this restores over. + eject_all_hooks(function) function.__code__ = code diff --git a/ddtrace/internal/bytecode_injection/__init__.py b/ddtrace/internal/bytecode_injection/__init__.py index c2ba59f3230..a077b366f1c 100644 --- a/ddtrace/internal/bytecode_injection/__init__.py +++ b/ddtrace/internal/bytecode_injection/__init__.py @@ -42,7 +42,7 @@ def __init__(self) -> None: def on_py_line(self, code: Any, line_number: int) -> Any: hooks: "list[tuple[HookType, Any]] | None" = self._hooks.get(line_number) if not hooks: - return _monitoring.DISABLE # type: ignore[has-type] + return _monitoring._DISABLE # type: ignore[has-type] for hook, arg in hooks: hook(arg) return None diff --git a/ddtrace/internal/monitoring.py b/ddtrace/internal/monitoring.py index cd25c271e72..676800cc836 100644 --- a/ddtrace/internal/monitoring.py +++ b/ddtrace/internal/monitoring.py @@ -28,10 +28,7 @@ if sys.version_info < (3, 15): raise ImportError("ddtrace.internal.monitoring requires Python 3.15+") -# mypy's configured python_version is below 3.15, so it considers everything -# past the guard above unreachable; that's expected here since the module -# raises ImportError before this point on unsupported versions. -log = get_logger(__name__) # type: ignore[unreachable] +log = get_logger(__name__) _E = sys.monitoring.events _DISABLE = sys.monitoring.DISABLE @@ -40,12 +37,11 @@ # set_local_events alongside PY_START/PY_RETURN/LINE. _LOCAL_EVENTS = _E.PY_START | _E.PY_RETURN | _E.LINE | _E.PY_UNWIND -# CPython tool IDs (see ddtrace/profiling/collector/_exception.pyx): -# 0 DEBUGGER_ID, 1 COVERAGE_ID, 2 PROFILER_ID, 3 handled exceptions, 5 OPTIMIZER_ID. -# Slot 4 is today's exception-profiler slot; the 3.15 stack migrates that -# collector onto this multiplexer. Never claim 0/1/2/3/5 or a slot owned by -# another tool name. Fall back to 3 if 4 is claimed by another tool. _MULTIPLEXER_TOOL_NAME = "ddtrace" +# sys.monitoring exposes six tool IDs (0–5). 0/1/2/5 are conventionally reserved +# for debugger/coverage/profiler/optimizer; 3 and 4 are the only undefined +# slots for custom tools (see CPython docs). Prefer 4 first, consistent with +# coverage's _DD_CANDIDATE_SLOTS, and fall back to 3 if another tool claimed it. _CANDIDATE_TOOL_IDS = (4, 3) _tool_id: Optional[int] = None @@ -82,6 +78,16 @@ def get(self, key: CodeType, default: Any = None) -> Any: return value return default + def __contains__(self, key: CodeType) -> bool: + item = self._data.get(id(key)) + return item is not None and item[0]() is key + + def __getitem__(self, key: CodeType) -> Any: + item = self._data.get(id(key)) + if item is None or item[0]() is not key: + raise KeyError(key) + return item[1] + def __setitem__(self, key: CodeType, value: Any) -> None: key_id = id(key) self._data[key_id] = (weakref.ref(key, self._make_remove(key_id)), value) @@ -92,6 +98,16 @@ def __delitem__(self, key: CodeType) -> None: raise KeyError(key) del self._data[key_id] + def pop(self, key: CodeType, *default: Any) -> Any: + try: + value = self[key] + except KeyError: + if default: + return default[0] + raise + del self[key] + return value + _registry: _IdentityWeakKeyDictionary = _IdentityWeakKeyDictionary() @@ -214,22 +230,6 @@ def _setup() -> int: # --------------------------------------------------------------------------- -def _dispatch_start(code: CodeType, instruction_offset: int, entry: _Entry) -> None: - entry.handler.on_py_start(code, instruction_offset) - - -def _dispatch_return(code: CodeType, instruction_offset: int, retval: object, entry: _Entry) -> None: - entry.handler.on_py_return(code, instruction_offset, retval) - - -def _dispatch_unwind(code: CodeType, instruction_offset: int, exception: BaseException, entry: _Entry) -> None: - entry.handler.on_py_unwind(code, instruction_offset, exception) - - -def _dispatch_line(code: CodeType, line_number: int, entry: _Entry) -> Optional[object]: - return entry.handler.on_py_line(code, line_number) - - def _on_py_start(code: CodeType, instruction_offset: int) -> Optional[object]: handlers: Optional[_CodeHandlers] = _registry.get(code) if not handlers or not handlers.snapshot: @@ -237,7 +237,7 @@ def _on_py_start(code: CodeType, instruction_offset: int) -> Optional[object]: for e in handlers.snapshot: if e.events & _E.PY_START: try: - _dispatch_start(code, instruction_offset, e) + e.handler.on_py_start(code, instruction_offset) except Exception: log.warning("monitoring PY_START handler failed", exc_info=True) return None @@ -250,7 +250,7 @@ def _on_py_return(code: CodeType, instruction_offset: int, retval: object) -> Op for e in handlers.snapshot: if e.events & _E.PY_RETURN: try: - _dispatch_return(code, instruction_offset, retval, e) + e.handler.on_py_return(code, instruction_offset, retval) except Exception: log.warning("monitoring PY_RETURN handler failed", exc_info=True) return None @@ -263,7 +263,7 @@ def _on_py_unwind(code: CodeType, instruction_offset: int, exception: BaseExcept for e in handlers.snapshot: if e.events & _E.PY_UNWIND: try: - _dispatch_unwind(code, instruction_offset, exception, e) + e.handler.on_py_unwind(code, instruction_offset, exception) except Exception: log.warning("monitoring PY_UNWIND handler failed", exc_info=True) return None @@ -277,7 +277,7 @@ def _on_py_line(code: CodeType, line_number: int) -> Optional[object]: for e in handlers.snapshot: if e.events & _E.LINE: try: - if _dispatch_line(code, line_number, e) is not _DISABLE: + if e.handler.on_py_line(code, line_number) is not _DISABLE: disable = False except Exception: log.warning("monitoring LINE handler failed", exc_info=True) @@ -298,9 +298,9 @@ def _rearm_local_events(tool_id: int, code: CodeType, events: int) -> None: # A DISABLE returned from a per-line callback is sticky until the monitored # event set changes or restart_events() is called. Re-applying the same # local events does not clear it; toggling local events off and back on - # re-arms only this tool's DISABLE marks for *code* without the global + # re-arms only this tool's DISABLE marks for code without the global # restart_events() call that would reset other tools' disabled-event - # bookkeeping (coverage.py). + # bookkeeping. _set_local_events(tool_id, code, 0) _set_local_events(tool_id, code, events) @@ -337,7 +337,7 @@ def refresh(code: CodeType) -> None: """Re-apply local events for *code*, resetting any per-line DISABLE state. Call this after adding a new hook for a line that may have been previously - disabled via a ``DISABLE`` return from :meth:`MonitoringEventHandler.on_py_line`. + disabled via a DISABLE return from :meth:`MonitoringEventHandler.on_py_line`. """ with _registry_lock: handlers: Optional[_CodeHandlers] = _registry.get(code) diff --git a/ddtrace/internal/wrapping/__init__.py b/ddtrace/internal/wrapping/__init__.py index b01b8c7977d..7e3d2b12ef7 100644 --- a/ddtrace/internal/wrapping/__init__.py +++ b/ddtrace/internal/wrapping/__init__.py @@ -4,7 +4,6 @@ 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 @@ -27,11 +26,53 @@ _wrapped: weakref.WeakKeyDictionary[FunctionType, FunctionType] = weakref.WeakKeyDictionary() _wrapped_lock = Lock() + +class _IdentityWeakValueDictionary: + """Maps code objects to functions by code-object identity, not equality. + + CodeType overrides __eq__/__hash__ (equal for structurally-identical code, + e.g. two CodeType.replace() clones of the same original, or two exec()'d + copies of identical source). A plain weakref.WeakValueDictionary keyed by + CodeType would conflate such distinct objects, returning the wrong + function for a given code object. This keys on id(code) instead, holding + the code object strongly (it may otherwise be referenced only here, once + the owning function's __code__ is replaced by wrapping) and the function + weakly, so wrapped ephemeral functions are not kept alive by this mapping + alone -- inspection.py falls back to gc.get_referrers on a miss. + """ + + __slots__ = ("_data",) + + def __init__(self) -> None: + self._data: dict[int, tuple[CodeType, "weakref.ref[FunctionType]"]] = {} + + def _make_remove(self, code_id: int) -> Any: + def remove(_ref: "weakref.ref[FunctionType]") -> None: + self._data.pop(code_id, None) + + return remove + + def __setitem__(self, code: CodeType, function: FunctionType) -> None: + code_id = id(code) + self._data[code_id] = (code, weakref.ref(function, self._make_remove(code_id))) + + def __getitem__(self, code: CodeType) -> FunctionType: + item = self._data.get(id(code)) + if item is not None: + stored_code, ref = item + if stored_code is code: + function = ref() + if function is not None: + return function + raise KeyError(code) + + def clear(self) -> None: + self._data.clear() + + # 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() +_code_to_fn: _IdentityWeakValueDictionary = _IdentityWeakValueDictionary() def link_function_to_code(code: CodeType, function: FunctionType) -> None: diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index 5c32d701e31..9e71a647631 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -110,6 +110,57 @@ def get_or_create(cls, f: FunctionType) -> "_ContextRecord": T = t.TypeVar("T") +StorageVar = ContextVar[t.Optional[dict[str, t.Any]]] + +_STORAGE_PREV = "__dd_wrapping_context_prev__" +_STORAGE_OWNER = "__dd_wrapping_context_owner__" + +# Free lists of storage context variables, keyed by variable name. +# +# Once a ContextVar has been set, the running Context holds a strong reference +# to it for the lifetime of the thread, and there is no way to drop that entry +# again (ContextVar.reset is not usable here, see the note on _pop_storage). +# Since wrapping contexts are created per function object, code that decorates +# ephemeral functions on every call would otherwise pin one variable per +# invocation. Recycling the variables of collected wrapping contexts caps the +# number of live ones at the peak number of concurrently live wrapping +# contexts. The pool is deliberately never trimmed: dropping a variable from it +# would not release the Context entries it already holds, and would only force +# the allocation of a new variable, adding entries instead of reusing them. That +# peak is therefore retained for the lifetime of the process, but it no longer +# grows with the number of functions that get wrapped. +# +# A recycled variable may still be set in some Context when it is handed out: a +# context that is collected without exiting leaves its storage behind, and the +# finalizer cannot reset it because it runs in an unrelated Context. Storage +# dicts are tagged with an owner token so that the new owner can tell a leftover +# value apart from one of its own; see __enter__. +_storage_var_pools: dict[str, list[StorageVar]] = {} + +# Reentrant because the release happens from a finalizer, which can run at +# any point, including in the middle of an acquisition on the same thread. +_storage_var_pools_lock = RLock() + + +def _acquire_storage_var(name: str) -> StorageVar: + with _storage_var_pools_lock: + pool = _storage_var_pools.get(name) + if pool: + var = pool.pop() + if not pool: + # Drop exhausted pools so that the names of wrapping contexts + # that are no longer in use don't accumulate either. + del _storage_var_pools[name] + return var + + return ContextVar(name, default=None) + + +def _release_storage_var(name: str, var: StorageVar) -> None: + with _storage_var_pools_lock: + _storage_var_pools.setdefault(name, []).append(var) + + # 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: @@ -384,7 +435,12 @@ def get_or_create(cls, f: FunctionType) -> "_ContextRecord": from ddtrace.internal import monitoring as _monitoring # Keyed by code object: drives sys.monitoring dispatch and is_wrapped/extract lookup. - _ctx_registry: "weakref.WeakKeyDictionary[CodeType, _UniversalWrappingContext]" = weakref.WeakKeyDictionary() + # CodeType overrides __eq__/__hash__ (equal for structurally-identical code), and + # monitor_code is produced by original_code.replace(), which compares equal to + # original_code. A plain weakref.WeakKeyDictionary would therefore conflate the + # clone with the original -- and, via that, conflate closures that share one + # original code object -- so this uses the identity-keyed mapping instead. + _ctx_registry: "_monitoring._IdentityWeakKeyDictionary" = _monitoring._IdentityWeakKeyDictionary() # Keyed by function instance: distinguishes functions that share a code object # (e.g. closures re-created in a loop) from one another. Kept off the function's # __dict__ (unlike a plain attribute) so functools.wraps does not propagate @@ -406,9 +462,17 @@ def __init__(self, f: FunctionType): # 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 - ) + + # Identifies the storage dicts written by this context. A dedicated token + # is used rather than self so that the storage dict cannot keep the + # context, and therefore the wrapped function, alive. + self._storage_owner = object() + + # Qualified so that same-named context types (e.g. the two + # LazyWrappingContext classes in this package) do not share a pool. + name = f"{type(self).__module__}.{type(self).__qualname__}__storage" + self._storage: StorageVar = _acquire_storage_var(name) + weakref.finalize(self, _release_storage_var, name, self._storage).atexit = False @property def __wrapped__(self) -> FunctionType: @@ -423,7 +487,14 @@ def __wrapped__(self, f: FunctionType) -> None: def __enter__(self) -> "BaseWrappingContext": prev = self._storage.get() - self._storage.set({"__dd_wrapping_context_prev__": prev}) + if prev is not None and prev.get(_STORAGE_OWNER) is not self._storage_owner: + # Storage left behind by a previous owner of this recycled variable. + # Chaining it into our own prev would restore it on every exit from + # now on, pinning it (and the frame it holds, for a universal + # wrapping context) for the lifetime of the thread. Dropping it here + # instead frees it as soon as we overwrite the variable below. + prev = None + self._storage.set({_STORAGE_PREV: prev, _STORAGE_OWNER: self._storage_owner}) return self @@ -431,7 +502,8 @@ def _pop_storage(self) -> dict[str, t.Any]: storage = self._storage.get() if storage is None: return {} - self._storage.set(storage.pop("__dd_wrapping_context_prev__")) + self._storage.set(storage.pop(_STORAGE_PREV)) + del storage[_STORAGE_OWNER] return storage def __return__(self, value: T) -> T: @@ -694,6 +766,23 @@ def __return__(self, value: T) -> T: return t.cast(T, super().__return__(value)) if sys.version_info >= (3, 15): + # AIDEV-NOTE: sys.monitoring itself does propagate a PY_START/PY_UNWIND + # callback's exception into the monitored function's frame (verified + # against CPython 3.15), which would give exceptions raised here true + # with-statement semantics matching the bytecode path (an __enter__ + # failure aborts the call before the body runs). But ddtrace.internal. + # monitoring multiplexes one sys.monitoring tool ID across independent + # subsystems (this wrapping context, DI line hooks, the exception + # profiler); its dispatch loop (_on_py_start et al.) therefore catches + # and logs each handler's exceptions instead of propagating them, so + # that a bug in one subsystem's handler cannot break monitored + # execution -- or a sibling handler registered on the same code + # object -- for the others. As a side effect, an exception raised by a + # registered WrappingContext's __enter__/__exit__ here is swallowed + # (logged, not raised), so unlike the bytecode path it does NOT abort + # the wrapped call. Fixing this would mean adding an opt-in + # propagate-exceptions mode to the shared multiplexer, which has + # implications for its other consumers; not done here. def on_py_start(self, code: t.Any, instruction_offset: int) -> None: self.__enter__() @@ -765,7 +854,7 @@ def wrap(self) -> None: f, _finalize_monitoring_wrap, weakref.ref(self), - weakref.ref(f), + monitor_code, ) self._finalize.atexit = False # Register monitoring before swapping __code__ so no thread can @@ -1103,21 +1192,26 @@ def unwrap(self) -> None: def _finalize_monitoring_wrap( self_ref: "weakref.ref[_UniversalWrappingContext]", - f_ref: weakref.ref[FunctionType], + code: CodeType, ) -> None: - """Unregister sys.monitoring when a wrapped function is collected without unwrap().""" + """Unregister sys.monitoring when a wrapped function is collected without unwrap(). + + weakref.finalize fires only after the wrapped function is unreachable, so by the + time this runs, self.__wrapped__ is already gone; unwrap() (and self.unwrap's use + of self.__wrapped__) cannot be used here. Clean up via the cloned monitor code + object instead, which the finalizer callback captures directly. + """ self: t.Optional["_UniversalWrappingContext"] = self_ref() - f: t.Optional[FunctionType] = f_ref() - if self is None or f is None: + if self is None: return try: - if _fn_registry.get(f) is self: - self.unwrap() + with _ctx_registry_lock: + if _ctx_registry.get(code) is not self: + return + del _ctx_registry[code] + _monitoring.unregister(code, self) except Exception: - log.exception( - "ddtrace: error during finalizer unwrap of %s", - getattr(f, "__qualname__", "?"), - ) + log.exception("ddtrace: error during finalizer cleanup of monitoring wrap") def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingContext]": """Return the _UniversalWrappingContext for *f*, or None if not context-wrapped.""" diff --git a/tests/wrapping/test_unwrap.py b/tests/wrapping/test_unwrap.py index 12a3e05da8d..0a4a74f6938 100644 --- a/tests/wrapping/test_unwrap.py +++ b/tests/wrapping/test_unwrap.py @@ -13,7 +13,8 @@ inverse, even when layers are nested). On Python < 3.15, ``WrappingContext.unwrap`` restores behaviour but rebuilds the code object rather than reinstating the original, so ``__code__`` identity is not restored (codified with a strict xfail). -On 3.15+ the monitoring path never mutates ``__code__``, so identity is preserved. +On 3.15+ the monitoring path swaps in a cloned code object while wrapped but restores the +exact original ``__code__`` object on unwrap, so identity is preserved. Mechanism-specific (the matrix's other two mechanisms have no in-place unwrap), so these opt out of the all-mechanisms ``mech`` guardrail. From a6165fda9a68a9f8119e471651ef98810aafcf73 Mon Sep 17 00:00:00 2001 From: "Gabriele N. Tornetta" Date: Tue, 18 Aug 2026 14:10:25 +0100 Subject: [PATCH 12/13] address review comments --- ddtrace/debugging/_function/store.py | 13 +- .../internal/bytecode_injection/__init__.py | 9 +- ddtrace/internal/monitoring.py | 27 +++-- ddtrace/internal/wrapping/context.py | 50 +++++--- .../bytecode_injection/test_injection.py | 47 ++++++++ tests/internal/test_monitoring.py | 112 ++++++++++++++++++ tests/internal/test_wrapping.py | 56 +++++++++ 7 files changed, 274 insertions(+), 40 deletions(-) diff --git a/ddtrace/debugging/_function/store.py b/ddtrace/debugging/_function/store.py index 31bbb8d9eec..e50baff8b00 100644 --- a/ddtrace/debugging/_function/store.py +++ b/ddtrace/debugging/_function/store.py @@ -90,11 +90,14 @@ def unwrap(self, function: FullyNamedContextWrappedFunction) -> None: def restore_all(self) -> None: """Restore all the patched functions to their original form.""" + for function, wrapping_context in list(self._wrapper_map.items()): + wrapping_context.unwrap() + for function, code in self._code_map.items(): - # On 3.15+, line hooks are sys.monitoring registrations keyed by - # code object rather than injected bytecode, so restoring - # __code__ alone leaves them firing; eject_all_hooks is a no-op - # on older versions, where the registration lives in the bytecode - # this restores over. + # Restoring __code__ alone would leave 3.15+ line hooks (keyed by + # code object) still firing. eject_all_hooks(function) function.__code__ = code + + self._code_map.clear() + self._wrapper_map.clear() diff --git a/ddtrace/internal/bytecode_injection/__init__.py b/ddtrace/internal/bytecode_injection/__init__.py index a077b366f1c..9e2b6d6aea9 100644 --- a/ddtrace/internal/bytecode_injection/__init__.py +++ b/ddtrace/internal/bytecode_injection/__init__.py @@ -26,8 +26,6 @@ class InvalidLine(Exception): if PY >= (3, 15): - import weakref - from ddtrace.internal import monitoring as _monitoring from ddtrace.internal.threads import Lock from ddtrace.internal.utils.inspection import linenos @@ -64,8 +62,11 @@ def remove(self, line: int, hook: HookType, arg: Any) -> None: def is_empty(self) -> bool: return not self._hooks - # WeakKeyDictionary: code object -> _LineHookHandler - _line_hook_registry: "weakref.WeakKeyDictionary[CodeType, _LineHookHandler]" = weakref.WeakKeyDictionary() + # Identity-keyed (not CodeType.__eq__) weak mapping: code object -> _LineHookHandler. + # Distinct code objects can compare structurally equal (e.g. repeated identical + # compiles, or CodeType.replace() clones), so a plain WeakKeyDictionary would let + # a hook registered for one code object be looked up under another. + _line_hook_registry: "_monitoring._IdentityWeakKeyDictionary" = _monitoring._IdentityWeakKeyDictionary() _line_hook_lock = Lock() def inject_hooks(f: FunctionType, hooks: list[HookInfoType]) -> list[HookInfoType]: diff --git a/ddtrace/internal/monitoring.py b/ddtrace/internal/monitoring.py index 676800cc836..883bcae640a 100644 --- a/ddtrace/internal/monitoring.py +++ b/ddtrace/internal/monitoring.py @@ -122,6 +122,15 @@ class MonitoringEventHandler(ABC): Do not call :func:`register` or :func:`unregister` from inside an event handler method. Doing so mutates the handler list while it is being iterated, which produces undefined behavior. + + .. warning:: + Exceptions from ``on_py_start``/``on_py_return``/``on_py_unwind`` are + not caught -- they propagate into the monitored frame and skip any + later handler for the same event, exactly as sys.monitoring itself + would deliver a callback failure. Catch your own exceptions if a + handler must not affect the monitored function's behavior. + ``on_py_line`` is caught and logged instead, since independent + handlers commonly share one code object's LINE registration. """ def on_py_start(self, code: CodeType, instruction_offset: int) -> None: @@ -234,12 +243,10 @@ def _on_py_start(code: CodeType, instruction_offset: int) -> Optional[object]: handlers: Optional[_CodeHandlers] = _registry.get(code) if not handlers or not handlers.snapshot: return _DISABLE + # Deliberately uncaught: see the propagation warning on MonitoringEventHandler. for e in handlers.snapshot: if e.events & _E.PY_START: - try: - e.handler.on_py_start(code, instruction_offset) - except Exception: - log.warning("monitoring PY_START handler failed", exc_info=True) + e.handler.on_py_start(code, instruction_offset) return None @@ -247,12 +254,10 @@ def _on_py_return(code: CodeType, instruction_offset: int, retval: object) -> Op handlers: Optional[_CodeHandlers] = _registry.get(code) if not handlers or not handlers.snapshot: return _DISABLE + # Deliberately uncaught: see the propagation warning on MonitoringEventHandler. for e in handlers.snapshot: if e.events & _E.PY_RETURN: - try: - e.handler.on_py_return(code, instruction_offset, retval) - except Exception: - log.warning("monitoring PY_RETURN handler failed", exc_info=True) + e.handler.on_py_return(code, instruction_offset, retval) return None @@ -260,12 +265,10 @@ def _on_py_unwind(code: CodeType, instruction_offset: int, exception: BaseExcept handlers: Optional[_CodeHandlers] = _registry.get(code) if not handlers or not handlers.snapshot: return _DISABLE + # Deliberately uncaught: see the propagation warning on MonitoringEventHandler. for e in handlers.snapshot: if e.events & _E.PY_UNWIND: - try: - e.handler.on_py_unwind(code, instruction_offset, exception) - except Exception: - log.warning("monitoring PY_UNWIND handler failed", exc_info=True) + e.handler.on_py_unwind(code, instruction_offset, exception) return None diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index 9e71a647631..c1d83f2b422 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -114,6 +114,10 @@ def get_or_create(cls, f: FunctionType) -> "_ContextRecord": _STORAGE_PREV = "__dd_wrapping_context_prev__" _STORAGE_OWNER = "__dd_wrapping_context_owner__" +# 3.15+ only: set in per-call storage when on_py_start/on_py_return fails, so the +# PY_UNWIND event sys.monitoring raises for that same synthetic exception does not +# also trigger __exit__. See _UniversalWrappingContext.on_py_unwind. +_SKIP_UNWIND_KEY = "__dd_wrapping_context_skip_unwind__" # Free lists of storage context variables, keyed by variable name. # @@ -766,31 +770,39 @@ def __return__(self, value: T) -> T: return t.cast(T, super().__return__(value)) if sys.version_info >= (3, 15): - # AIDEV-NOTE: sys.monitoring itself does propagate a PY_START/PY_UNWIND - # callback's exception into the monitored function's frame (verified - # against CPython 3.15), which would give exceptions raised here true - # with-statement semantics matching the bytecode path (an __enter__ - # failure aborts the call before the body runs). But ddtrace.internal. - # monitoring multiplexes one sys.monitoring tool ID across independent - # subsystems (this wrapping context, DI line hooks, the exception - # profiler); its dispatch loop (_on_py_start et al.) therefore catches - # and logs each handler's exceptions instead of propagating them, so - # that a bug in one subsystem's handler cannot break monitored - # execution -- or a sibling handler registered on the same code - # object -- for the others. As a side effect, an exception raised by a - # registered WrappingContext's __enter__/__exit__ here is swallowed - # (logged, not raised), so unlike the bytecode path it does NOT abort - # the wrapped call. Fixing this would mean adding an opt-in - # propagate-exceptions mode to the shared multiplexer, which has - # implications for its other consumers; not done here. + # Exceptions here are deliberately left uncaught (see the propagation + # warning on MonitoringEventHandler), which matches bytecode-path + # with-statement semantics -- safe because this is the only handler + # ddtrace registers for these events on a given code object. + # + # CPython also fires a synthetic PY_UNWIND after a failing PY_START/ + # PY_RETURN; _SKIP_UNWIND_KEY suppresses the resulting __exit__ call so + # it only fires for a real exception from the wrapped function body. + # It lives in per-call storage (a ContextVar), not a plain attribute, + # because this same instance is shared across concurrent calls. def on_py_start(self, code: t.Any, instruction_offset: int) -> None: - self.__enter__() + try: + self.__enter__() + except BaseException: + storage = self._storage.get() + if storage is not None: + storage[_SKIP_UNWIND_KEY] = True + raise def on_py_return(self, code: t.Any, instruction_offset: int, retval: t.Any) -> None: - self.__return__(retval) + try: + self.__return__(retval) + except BaseException: + storage = self._storage.get() + if storage is not None: + storage[_SKIP_UNWIND_KEY] = True + raise def on_py_unwind(self, code: t.Any, instruction_offset: int, exception: BaseException) -> None: + storage = self._storage.get() + if storage is not None and storage.pop(_SKIP_UNWIND_KEY, False): + return self.__exit__(type(exception), exception, exception.__traceback__) @classmethod diff --git a/tests/internal/bytecode_injection/test_injection.py b/tests/internal/bytecode_injection/test_injection.py index 949754a2864..9388bef8bd7 100644 --- a/tests/internal/bytecode_injection/test_injection.py +++ b/tests/internal/bytecode_injection/test_injection.py @@ -316,3 +316,50 @@ def for_loop(): hook.assert_called_with(arg) else: hook.assert_called_once_with(arg) + + +@pytest.mark.skipif(sys.version_info < (3, 15), reason="line hook registry is only keyed by code identity on 3.15+") +def test_line_hooks_isolated_across_structurally_equal_code_objects(): + """Two distinct code objects that compare equal must not share a line hook registration.""" + src = "def target(x):\n return x + 1\n" + ns_a: dict = {} + ns_b: dict = {} + exec(src, ns_a) + exec(src, ns_b) + f_a = ns_a["target"] + f_b = ns_b["target"] + + assert f_a.__code__ is not f_b.__code__ + assert f_a.__code__ == f_b.__code__, "the two code objects must be structurally equal for this test to matter" + + hook = mock.Mock() + lo = min(linenos(f_a)) + inject_hook(f_a, hook, lo, 42) + + f_b(1) + + hook.assert_not_called() + + eject_hook(f_a, hook, lo, 42) + + +@pytest.mark.skipif(sys.version_info < (3, 15), reason="line hook registry is only keyed by code identity on 3.15+") +def test_line_hooks_isolated_across_code_replace_clone(): + """A hook registered against the original code object must not fire for a code.replace() clone.""" + + def target(x): + return x + 1 + + hook = mock.Mock() + lo = min(linenos(target)) + inject_hook(target, hook, lo, 42) + + cloned = FunctionType(target.__code__.replace(), target.__globals__, "cloned") + assert cloned.__code__ is not target.__code__ + assert cloned.__code__ == target.__code__ + + cloned(1) + + hook.assert_not_called() + + eject_hook(target, hook, lo, 42) diff --git a/tests/internal/test_monitoring.py b/tests/internal/test_monitoring.py index 9ae9c7f1c6d..37e48f6df91 100644 --- a/tests/internal/test_monitoring.py +++ b/tests/internal/test_monitoring.py @@ -74,6 +74,33 @@ def on_py_unwind(self, code: CodeType, instruction_offset: int, exception: BaseE self.unwound = True +class RaisingStartHandler(monitoring.MonitoringEventHandler): + def __init__(self) -> None: + self.called: bool = False + + def on_py_start(self, code: CodeType, instruction_offset: int) -> None: + self.called = True + raise RuntimeError("start handler exploded") + + +class RaisingReturnHandler(monitoring.MonitoringEventHandler): + def __init__(self) -> None: + self.called: bool = False + + def on_py_return(self, code: CodeType, instruction_offset: int, retval: object) -> None: + self.called = True + raise RuntimeError("return handler exploded") + + +class RaisingUnwindHandler(monitoring.MonitoringEventHandler): + def __init__(self) -> None: + self.called: bool = False + + def on_py_unwind(self, code: CodeType, instruction_offset: int, exception: BaseException) -> None: + self.called = True + raise RuntimeError("unwind handler exploded") + + @pytest.fixture def registered() -> Iterator[ Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler] @@ -237,3 +264,88 @@ def fn() -> None: result: object | None = monitoring._on_py_line(fn.__code__, fn.__code__.co_firstlineno) assert result is not _DISABLE + + +def test_on_py_start_propagates_exception( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: + """A PY_START handler's exception is never caught -- it always reaches the caller.""" + + def fn() -> None: + pass + + handler: RaisingStartHandler = registered(fn.__code__, RaisingStartHandler()) # type: ignore[assignment] + + with pytest.raises(RuntimeError, match="start handler exploded"): + monitoring._on_py_start(fn.__code__, 0) + + assert handler.called + + +def test_on_py_start_propagation_aborts_the_monitored_call( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: + """A propagating PY_START failure must prevent the monitored function body from running.""" + + ran: bool = False + + def fn() -> None: + nonlocal ran + ran = True + + registered(fn.__code__, RaisingStartHandler()) + + with pytest.raises(RuntimeError, match="start handler exploded"): + fn() + + assert not ran, "the function body must not run when a propagating PY_START handler raises" + + +def test_on_py_return_propagates_exception( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: + """A PY_RETURN handler's exception is never caught -- it always reaches the caller.""" + + def fn() -> None: + pass + + handler: RaisingReturnHandler = registered(fn.__code__, RaisingReturnHandler()) # type: ignore[assignment] + + with pytest.raises(RuntimeError, match="return handler exploded"): + monitoring._on_py_return(fn.__code__, 0, None) + + assert handler.called + + +def test_on_py_unwind_propagates_exception( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: + """A PY_UNWIND handler's exception is never caught -- it always reaches the caller.""" + + def fn() -> None: + pass + + handler: RaisingUnwindHandler = registered(fn.__code__, RaisingUnwindHandler()) # type: ignore[assignment] + + with pytest.raises(RuntimeError, match="unwind handler exploded"): + monitoring._on_py_unwind(fn.__code__, 0, ValueError("original")) + + assert handler.called + + +def test_propagating_handler_skips_later_handlers_for_same_event( + registered: Callable[[CodeType, monitoring.MonitoringEventHandler], monitoring.MonitoringEventHandler], +) -> None: + """A propagating handler's exception skips any sibling handler registered after it.""" + + def fn() -> None: + pass + + raiser: RaisingStartHandler = registered(fn.__code__, RaisingStartHandler()) # type: ignore[assignment] + sibling: StartAndUnwindHandler = registered(fn.__code__, StartAndUnwindHandler()) # type: ignore[assignment] + + with pytest.raises(RuntimeError, match="start handler exploded"): + monitoring._on_py_start(fn.__code__, 0) + + assert raiser.called + assert not sibling.started, "a sibling handler after a propagating raiser must not run" diff --git a/tests/internal/test_wrapping.py b/tests/internal/test_wrapping.py index b954802c980..3b66e5b6860 100644 --- a/tests/internal/test_wrapping.py +++ b/tests/internal/test_wrapping.py @@ -865,6 +865,62 @@ def foo(): assert exc.args == ("foo",) +def test_wrapping_context_exc_on_enter(): + """A raising __enter__ must abort the call before the wrapped body runs.""" + + ran = False + + class BrokenEnterWrappingContext(DummyWrappingContext): + def __enter__(self): + super().__enter__() + raise RuntimeError("broken enter") + + def foo(): + nonlocal ran + ran = True + return 42 + + wc = BrokenEnterWrappingContext(foo) + wc.wrap() + + with pytest.raises(RuntimeError): + foo() + + assert not ran, "the wrapped function body must not run when __enter__ raises" + assert wc.entered + + +def test_wrapping_context_exc_on_return(): + """A raising __return__ must override the wrapped function's return value. + + Whether __exit__ also runs differs by path: on <3.15 it's part of the same + try/except as the rest of the call, so it does; on 3.15+ _SKIP_UNWIND_KEY + (see context.py) suppresses the synthetic PY_UNWIND this triggers, so it + doesn't. + """ + + class BrokenReturnWrappingContext(DummyWrappingContext): + def __return__(self, value): + super().__return__(value) + raise RuntimeError("broken return") + + def foo(): + return 42 + + wc = BrokenReturnWrappingContext(foo) + wc.wrap() + + with pytest.raises(RuntimeError): + foo() + + assert wc.entered + assert wc.return_value == 42 + if sys.version_info >= (3, 15): + assert not wc.exited, "__exit__ must not run when __return__ itself is the failure" + else: + assert wc.exited + + def test_wrapping_context_priority(): class HighPriorityWrappingContext(DummyWrappingContext): def __enter__(self): From 0c8becf2bf1c07a76a2489d02021f35d275cd79d Mon Sep 17 00:00:00 2001 From: "Gabriele N. Tornetta" Date: Tue, 18 Aug 2026 16:05:32 +0100 Subject: [PATCH 13/13] uniform exc-on-return behaviour --- ddtrace/internal/wrapping/context.py | 45 ++++++++++++++++++---------- tests/internal/test_wrapping.py | 12 +++----- 2 files changed, 33 insertions(+), 24 deletions(-) diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index c1d83f2b422..27aed1dc49d 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -114,10 +114,12 @@ def get_or_create(cls, f: FunctionType) -> "_ContextRecord": _STORAGE_PREV = "__dd_wrapping_context_prev__" _STORAGE_OWNER = "__dd_wrapping_context_owner__" -# 3.15+ only: set in per-call storage when on_py_start/on_py_return fails, so the -# PY_UNWIND event sys.monitoring raises for that same synthetic exception does not -# also trigger __exit__. See _UniversalWrappingContext.on_py_unwind. -_SKIP_UNWIND_KEY = "__dd_wrapping_context_skip_unwind__" +# Set in per-call storage when a raise originates from context machinery itself +# (__return__, or on 3.15+ on_py_start) rather than from the wrapped function +# body, so the resulting exception does not also trigger __exit__. Consumed by +# _UniversalWrappingContext._exit (bytecode path, >=3.11) and on_py_unwind +# (sys.monitoring path, >=3.15). +_SKIP_EXIT_KEY = "__dd_wrapping_context_skip_exit__" # Free lists of storage context variables, keyed by variable name. # @@ -747,6 +749,13 @@ def __enter__(self) -> "_UniversalWrappingContext": return self def _exit(self) -> None: + # Only reached on the bytecode path for Python >= 3.11, where the + # injected __return__ call sits inside the same try/except as the rest + # of the wrapped function body (see CONTEXT_FOOT). Skip the exit here + # too, so a failing __return__ behaves the same on every version. + storage = self._storage.get() + if storage is not None and storage.pop(_SKIP_EXIT_KEY, False): + return self.__exit__(*sys.exc_info()) def __exit__( @@ -764,8 +773,18 @@ def __exit__( super().__exit__(exc_type, exc_value, traceback) def __return__(self, value: T) -> T: - for context in self._contexts[::-1]: - context.__return__(value) + try: + for context in self._contexts[::-1]: + context.__return__(value) + except BaseException: + # A failing __return__ must not be treated as an exception from the + # wrapped function body -- it never gets to run, so __exit__ must + # not run for it either. See _exit and on_py_unwind, which consume + # this flag on the bytecode and sys.monitoring paths respectively. + storage = self._storage.get() + if storage is not None: + storage[_SKIP_EXIT_KEY] = True + raise return t.cast(T, super().__return__(value)) @@ -776,7 +795,7 @@ def __return__(self, value: T) -> T: # ddtrace registers for these events on a given code object. # # CPython also fires a synthetic PY_UNWIND after a failing PY_START/ - # PY_RETURN; _SKIP_UNWIND_KEY suppresses the resulting __exit__ call so + # PY_RETURN; _SKIP_EXIT_KEY suppresses the resulting __exit__ call so # it only fires for a real exception from the wrapped function body. # It lives in per-call storage (a ContextVar), not a plain attribute, # because this same instance is shared across concurrent calls. @@ -787,21 +806,15 @@ def on_py_start(self, code: t.Any, instruction_offset: int) -> None: except BaseException: storage = self._storage.get() if storage is not None: - storage[_SKIP_UNWIND_KEY] = True + storage[_SKIP_EXIT_KEY] = True raise def on_py_return(self, code: t.Any, instruction_offset: int, retval: t.Any) -> None: - try: - self.__return__(retval) - except BaseException: - storage = self._storage.get() - if storage is not None: - storage[_SKIP_UNWIND_KEY] = True - raise + self.__return__(retval) def on_py_unwind(self, code: t.Any, instruction_offset: int, exception: BaseException) -> None: storage = self._storage.get() - if storage is not None and storage.pop(_SKIP_UNWIND_KEY, False): + if storage is not None and storage.pop(_SKIP_EXIT_KEY, False): return self.__exit__(type(exception), exception, exception.__traceback__) diff --git a/tests/internal/test_wrapping.py b/tests/internal/test_wrapping.py index 3b66e5b6860..181e7b5c822 100644 --- a/tests/internal/test_wrapping.py +++ b/tests/internal/test_wrapping.py @@ -893,10 +893,9 @@ def foo(): def test_wrapping_context_exc_on_return(): """A raising __return__ must override the wrapped function's return value. - Whether __exit__ also runs differs by path: on <3.15 it's part of the same - try/except as the rest of the call, so it does; on 3.15+ _SKIP_UNWIND_KEY - (see context.py) suppresses the synthetic PY_UNWIND this triggers, so it - doesn't. + __exit__ must not run either, on any version: __return__ failing means the + wrapped function's own body never got a chance to raise, so this is not the + exception __exit__ exists to observe. See _SKIP_EXIT_KEY in context.py. """ class BrokenReturnWrappingContext(DummyWrappingContext): @@ -915,10 +914,7 @@ def foo(): assert wc.entered assert wc.return_value == 42 - if sys.version_info >= (3, 15): - assert not wc.exited, "__exit__ must not run when __return__ itself is the failure" - else: - assert wc.exited + assert not wc.exited, "__exit__ must not run when __return__ itself is the failure" def test_wrapping_context_priority():