diff --git a/ddtrace/debugging/_function/store.py b/ddtrace/debugging/_function/store.py index 8b2e4e98146..e50baff8b00 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 @@ -89,5 +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(): + # 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 5bcbbc978d8..9e2b6d6aea9 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,369 @@ 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 + 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 + + # 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]: + """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: int | None = None + instrs: set[str] = 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: tuple[int, str] = 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: 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: list[str] = [] + 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 = Bytecode.from_code(get_function_code(f)) -_INJECT_HOOK_OPCODES = [_.name for _ in INJECTION_ASSEMBLY] + failed: list[HookInfoType] = [] + 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 = Bytecode.from_code(f.__code__) - return failed + failed: list[HookInfoType] = [] + 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 = 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 = 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/monitoring.py b/ddtrace/internal/monitoring.py index f63994cd711..883bcae640a 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 @@ -81,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) @@ -91,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() @@ -105,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: @@ -217,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 @@ -230,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 @@ -243,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/__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/asyncs.py b/ddtrace/internal/wrapping/asyncs.py index a1de28a074e..396e9c20577 100644 --- a/ddtrace/internal/wrapping/asyncs.py +++ b/ddtrace/internal/wrapping/asyncs.py @@ -1,5 +1,7 @@ +import enum import sys from types import CodeType +from typing import Optional import bytecode as bc @@ -9,6 +11,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 # ----------------------------------------------------------------------------- @@ -33,10 +54,137 @@ COROUTINE_ASSEMBLY = Assembly() ASYNC_GEN_ASSEMBLY = Assembly() -ASYNC_HEAD_ASSEMBLY = None +ASYNC_HEAD_ASSEMBLY: Optional[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 039a71ebb2e..aa128fa4add 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -9,7 +9,7 @@ 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 @@ -49,7 +49,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 @@ -58,7 +58,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) @@ -118,6 +118,12 @@ def get_or_create(cls, f: FunctionType) -> "_ContextRecord": _STORAGE_PREV = "__dd_wrapping_context_prev__" _STORAGE_OWNER = "__dd_wrapping_context_owner__" +# 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. # @@ -214,8 +220,12 @@ def _release_storage_var(name: str, var: StorageVar) -> None: 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""" @@ -424,6 +434,32 @@ def _release_storage_var(name: str, var: StorageVar) -> None: ) +# 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. + # 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 + # 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 @@ -473,7 +509,9 @@ def __enter__(self) -> "BaseWrappingContext": return self def _pop_storage(self) -> dict[str, t.Any]: - storage = t.cast(WrappingContextStorage, self._storage.get()) + storage = t.cast(t.Optional[WrappingContextStorage], self._storage.get()) + if storage is None: + return {} self._storage.set(storage.pop(_STORAGE_PREV)) del storage[_STORAGE_OWNER] return storage @@ -566,88 +604,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 wrap(self) -> None: + """Perform the bytecode wrapping on first invocation.""" + with (tl := self._trampoline_lock): + if self._trampoline is not None: + return - def trampoline(_: t.Any, args: tuple[t.Any, ...], kwargs: dict[str, t.Any]) -> t.Any: - with tl: - f = t.cast(WrappedFunction, self.__wrapped__) - if is_wrapped_with(t.cast(FunctionType, self.__wrapped__), trampoline): - f = t.cast(WrappedFunction, unwrap(f, trampoline)) + # 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 - self._trampoline = None + 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)) - 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") + self._trampoline = None - super(LazyWrappingContext, self).wrap() - return f(*args, **kwargs) + 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") - wrap(t.cast(FunctionType, self.__wrapped__), trampoline) + super(LazyWrappingContext, self).wrap() + return f(*args, **kwargs) - self._trampoline = trampoline + wrap(t.cast(FunctionType, self.__wrapped__), trampoline) - _ContextRecord.get_or_create(t.cast(FunctionType, self.__wrapped__)).lazy_contexts.add(self) + self._trampoline = 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) + _ContextRecord.get_or_create(t.cast(FunctionType, self.__wrapped__)).lazy_contexts.add(self) - unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline) - self._trampoline = None + 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[type, ...] = (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) @@ -685,7 +747,7 @@ def __enter__(self) -> "_UniversalWrappingContext": storage = t.cast(WrappingContextStorage, self._storage.get()) # Make the frame object available to the contexts - storage["__frame__"] = sys._getframe(1) + storage["__frame__"] = sys._getframe(_ENTER_FRAME_DEPTH) # Freeze the list of contexts so that we know exactly which ones to # exit, in case new contexts are registered during the execution of @@ -705,6 +767,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__( @@ -728,184 +797,236 @@ def __exit__( super().__exit__(exc_type, exc_value, traceback) def __return__(self, value: T) -> T: + storage = t.cast(WrappingContextStorage, self._storage.get()) try: - contexts = t.cast(WrappingContextStorage, self._storage.get())["__contexts__"] + contexts = storage["__contexts__"] except (TypeError, KeyError): log.debug("Universal wrapping context returned without entering") return super().__return__(value) - for context in contexts[::-1]: - context.__return__(value) - - return super().__return__(value) - - @classmethod - def is_wrapped(cls, f: FunctionType) -> bool: try: - with _registry_lock: - record = _registry.get(f) - if record is None or record.uwc is None: + for context in 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[_SKIP_EXIT_KEY] = True + raise + + return t.cast(T, super().__return__(value)) + + if sys.version_info >= (3, 15): + # 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_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. + + def on_py_start(self, code: t.Any, instruction_offset: int) -> None: + try: + self.__enter__() + except BaseException: + storage = self._storage.get() + if storage is not None: + storage[_SKIP_EXIT_KEY] = True + raise + + 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: + storage = self._storage.get() + if storage is not None and storage.pop(_SKIP_EXIT_KEY, False): + return + 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) - - def wrap(self) -> None: - f = t.cast(FunctionType, self.__wrapped__) - - with _registry_lock: - if self.is_wrapped(f): - raise ValueError("Function already wrapped") - - code = get_function_code(f) - - # Closures created from repeated calls to the same factory share - # the same code object: _build_template is memoized so the - # expensive decompile/instrument/recompile step is reused across - # them. - template = self._build_template(code) - - # Register the wrapping context and link the function to the new - # code object. - _ContextRecord.get_or_create(f).uwc = self - link_function_to_code(code, f) - - # Substitute the template's placeholder consts with the real, - # instance-specific values. - replacements = self._template_replacements() - set_function_code(f, template.replace(co_consts=tuple(replacements.get(c, c) for c in template.co_consts))) - - if sys.version_info >= (3, 11): - - @staticmethod - @lru_cache(maxsize=_TEMPLATE_CACHE_MAX_SIZE) - def _build_template(code: "CodeType") -> "CodeType": - """Build a cacheable, instance-agnostic instrumented copy of *code*. - - The instance-specific context_enter/context_return/context_exit - bound methods are replaced with placeholders so that the result - can be shared across multiple closures backed by the same code - object; see wrap(). Memoized via lru_cache, keyed on the code - object, so repeated wraps of closures sharing the same underlying - code skip the decompile/instrument/recompile. - """ - bc = Bytecode.from_code(code) - - # 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": _RETURN_PLACEHOLDER}, 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": _RETURN_PLACEHOLDER, "value": instr.arg}, lineno=instr.lineno - ) + # 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: FunctionType = self.__wrapped__ + original_code: CodeType = get_function_code(f) + with _ctx_registry_lock: + if original_code in _ctx_registry: + 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: - return_code = [] - - bc[i:i] = return_code - i += len(return_code) - except AttributeError: - # Not an instruction - pass - i += 1 - - # Search for the RESUME instruction - for i, instr in enumerate(bc, 1): - try: - if instr.name == "RESUME": - break - except AttributeError: - # Not an instruction - pass - else: - i = 0 - - bc[i:i] = CONTEXT_HEAD.bind({"context_enter": _ENTER_PLACEHOLDER}, lineno=code.co_firstlineno) - - # Wrap every line outside a try block - except_label = bytecode.Label() - first_try_begin = last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) - - i = 0 - while i < len(bc): - instr = bc[i] - if isinstance(instr, bytecode.TryBegin) and last_try_begin is not None: - bc.insert(i, bytecode.TryEnd(last_try_begin)) - last_try_begin = None - i += 1 - elif isinstance(instr, bytecode.TryEnd): - j = i + 1 - while j < len(bc) and not isinstance(bc[j], bytecode.TryBegin): - if isinstance(bc[j], bytecode.Instr): - last_try_begin = bytecode.TryBegin(except_label, push_lasti=True) - bc.insert(i + 1, last_try_begin) - break - j += 1 - i += 1 - i += 1 - - bc.insert(0, first_try_begin) - - bc.append(bytecode.TryEnd(last_try_begin)) - bc.append(except_label) - bc.extend(CONTEXT_FOOT.bind({"context_exit": _EXIT_PLACEHOLDER}, lineno=code.co_firstlineno)) - - return bc.to_code() - - def _template_replacements(self) -> dict[object, object]: - return { - _ENTER_PLACEHOLDER: self.__enter__, - _RETURN_PLACEHOLDER: self.__return__, - _EXIT_PLACEHOLDER: self._exit, - } + raise ValueError("Function already wrapped") + + # 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) + + _ctx_registry[monitor_code] = self + _fn_registry[f] = self + self._finalize = weakref.finalize( + f, + _finalize_monitoring_wrap, + weakref.ref(self), + monitor_code, + ) + self._finalize.atexit = False + # 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 = self.__wrapped__ - - with _registry_lock: - if not self.is_wrapped(f): + 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 - wc = _registry[f].uwc + migrate_line_hooks(code, original_code) + set_function_code(f, original_code) + del self._original_code - bc = Bytecode.from_code(get_function_code(f)) + else: - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() - bc.pop() + @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 - except_label = bc.pop(0).target + @classmethod + def extract(cls, f: FunctionType) -> "_UniversalWrappingContext": + with _registry_lock: + if not cls.is_wrapped(f): + raise ValueError("Function is not wrapped") + return t.cast(_UniversalWrappingContext, _registry[f].uwc) - # Remove the try blocks + def wrap(self) -> None: + f = t.cast(FunctionType, self.__wrapped__) + + with _registry_lock: + if self.is_wrapped(f): + raise ValueError("Function already wrapped") + + code = get_function_code(f) + + # Closures created from repeated calls to the same factory share + # the same code object: _build_template is memoized so the + # expensive decompile/instrument/recompile step is reused across + # them. + template = self._build_template(code) + + # Register the wrapping context and link the function to the new + # code object. + _ContextRecord.get_or_create(f).uwc = self + link_function_to_code(code, f) + + # Substitute the template's placeholder consts with the real, + # instance-specific values. + replacements = self._template_replacements() + set_function_code( + f, template.replace(co_consts=tuple(replacements.get(c, c) for c in template.co_consts)) + ) + + if sys.version_info >= (3, 11): + + @staticmethod + @lru_cache(maxsize=_TEMPLATE_CACHE_MAX_SIZE) + def _build_template(code: "CodeType") -> "CodeType": + """Build a cacheable, instance-agnostic instrumented copy of *code*. + + The instance-specific context_enter/context_return/context_exit + bound methods are replaced with placeholders so that the result + can be shared across multiple closures backed by the same code + object; see wrap(). Memoized via lru_cache, keyed on the code + object, so repeated wraps of closures sharing the same underlying + code skip the decompile/instrument/recompile. + """ + bc = Bytecode.from_code(code) + + # Prefix every return 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 + try: + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN.bind( + {"context_return": _RETURN_PLACEHOLDER}, 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": _RETURN_PLACEHOLDER, "value": instr.arg}, lineno=instr.lineno + ) + else: + return_code = [] - # 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] = 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): @@ -918,126 +1039,245 @@ def unwrap(self) -> None: else: i = 0 - bc[i : i + len(CONTEXT_HEAD)] = [] + bc[i:i] = CONTEXT_HEAD.bind({"context_enter": _ENTER_PLACEHOLDER}, 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) - # 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 + 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 - # Recreate the code object - set_function_code(f, bc.to_code()) + bc.insert(0, first_try_begin) - # 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) + bc.append(bytecode.TryEnd(last_try_begin)) + bc.append(except_label) + bc.extend(CONTEXT_FOOT.bind({"context_exit": _EXIT_PLACEHOLDER}, lineno=code.co_firstlineno)) - else: + return bc.to_code() - @staticmethod - @lru_cache(maxsize=_TEMPLATE_CACHE_MAX_SIZE) - def _build_template(code: "CodeType") -> "CodeType": - """Build a cacheable, instance-agnostic instrumented copy of *code*. - - The instance-specific context object is replaced with a placeholder - so that the result can be shared across multiple closures backed by - the same code object; see wrap(). Memoized via lru_cache, keyed on - the code object, so repeated wraps of closures sharing the same - underlying code skip the decompile/instrument/ - recompile. - """ - bc = Bytecode.from_code(code) - - # 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": _CONTEXT_PLACEHOLDER}, lineno=instr.lineno) - bc[i:i] = return_code - i += len(return_code) - i += 1 + def _template_replacements(self) -> dict[object, object]: + return { + _ENTER_PLACEHOLDER: self.__enter__, + _RETURN_PLACEHOLDER: self.__return__, + _EXIT_PLACEHOLDER: self._exit, + } - # Search for the GEN_START instruction, which needs to stay on top. - i = 0 - if sys.version_info >= (3, 10) and (code.co_flags & (CO_GENERATOR | CO_COROUTINE)): - for i, instr in enumerate(bc, 1): - if isinstance(instr, bytecode.Instr) and instr.name == "GEN_START": - break + def unwrap(self) -> None: + f = self.__wrapped__ - *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": _CONTEXT_PLACEHOLDER}, lineno=code.co_firstlineno) + with _registry_lock: + if not self.is_wrapped(f): + return - bc.append(except_label) - bc.extend(CONTEXT_FOOT.bind(lineno=code.co_firstlineno)) + wc = _registry[f].uwc - return bc.to_code() + bc = Bytecode.from_code(get_function_code(f)) - def _template_replacements(self) -> dict[object, object]: - return {_CONTEXT_PLACEHOLDER: self} + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() + bc.pop() - def unwrap(self) -> None: - f = t.cast(FunctionType, self.__wrapped__) + except_label = bc.pop(0).target - with _registry_lock: - if not self.is_wrapped(f): - return + # 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 - wc = _registry[f].uwc + # Remove the head of the try block + for i, instr in enumerate(bc): + if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: + break - bc = Bytecode.from_code(get_function_code(f)) + # 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 - # Remove the exception handling code - bc[-len(CONTEXT_FOOT) :] = [] - bc.pop() + bc[i : i + len(CONTEXT_HEAD)] = [] - # 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 + # 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) - 1] = [] + # Recreate the code object + set_function_code(f, bc.to_code()) - # Remove all the return handlers + # 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) + + else: + + @staticmethod + @lru_cache(maxsize=_TEMPLATE_CACHE_MAX_SIZE) + def _build_template(code: "CodeType") -> "CodeType": + """Build a cacheable, instance-agnostic instrumented copy of *code*. + + The instance-specific context object is replaced with a placeholder + so that the result can be shared across multiple closures backed by + the same code object; see wrap(). Memoized via lru_cache, keyed on + the code object, so repeated wraps of closures sharing the same + underlying code skip the decompile/instrument/recompile. + """ + bc = Bytecode.from_code(code) + + # Prefix every return 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) + if isinstance(instr, bytecode.Instr): + if instr.name == "RETURN_VALUE": + return_code = CONTEXT_RETURN.bind({"context": _CONTEXT_PLACEHOLDER}, lineno=instr.lineno) + bc[i:i] = return_code + i += len(return_code) i += 1 - # Recreate the code object - set_function_code(f, bc.to_code()) + # Search for the GEN_START instruction, which needs to stay on top. + i = 0 + if sys.version_info >= (3, 10) and (code.co_flags & (CO_GENERATOR | CO_COROUTINE)): + 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": _CONTEXT_PLACEHOLDER}, lineno=code.co_firstlineno + ) - # 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) + bc.append(except_label) + bc.extend(CONTEXT_FOOT.bind(lineno=code.co_firstlineno)) + + return bc.to_code() + + def _template_replacements(self) -> dict[object, object]: + return {_CONTEXT_PLACEHOLDER: self} + def unwrap(self) -> None: + f = t.cast(FunctionType, self.__wrapped__) -def wrapping_context_for(f: FunctionType) -> "t.Optional[_UniversalWrappingContext]": - """Return the _UniversalWrappingContext for *f*, or None if not context-wrapped.""" - with _registry_lock: - record = _registry.get(f) - return record.uwc if record is not None else None + with _registry_lock: + if not self.is_wrapped(f): + return + + wc = _registry[f].uwc + + bc = Bytecode.from_code(get_function_code(f)) + + # Remove the exception handling code + bc[-len(CONTEXT_FOOT) :] = [] + bc.pop() + + # Remove the head of the try block + for i, instr in enumerate(bc): + if isinstance(instr, bytecode.Instr) and instr.name == "LOAD_CONST" and instr.arg is wc: + break + + bc[i : i + len(CONTEXT_HEAD) - 1] = [] + + # Remove all the return handlers + i = 0 + while i < len(bc): + instr = bc[i] + if isinstance(instr, bytecode.Instr) and instr.name == "RETURN_VALUE": + bc[i - len(CONTEXT_RETURN) : i] = [] + i -= len(CONTEXT_RETURN) + i += 1 + + # Recreate the code object + set_function_code(f, bc.to_code()) + + # Clear the UWC from the registry; remove the record if fully empty. + record: t.Optional[_ContextRecord] = _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]", + code: CodeType, + ) -> None: + """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() + if self is None: + return + try: + 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 cleanup of monitoring wrap") + + 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..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,10 +32,85 @@ # return # ----------------------------------------------------------------------------- GENERATOR_ASSEMBLY = Assembly() -GENERATOR_HEAD_ASSEMBLY = None +GENERATOR_HEAD_ASSEMBLY: Optional[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/mypy.ini b/mypy.ini index 18ae19a7fdb..77d88d65134 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1333,6 +1333,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/bytecode_injection/test_injection.py b/tests/internal/bytecode_injection/test_injection.py index a60f4fe5bce..9388bef8bd7 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 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: @@ -31,7 +41,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 +308,58 @@ 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) + + +@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 06e610f7aff..003ba43f077 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 @@ -864,6 +865,58 @@ 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. + + __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): + 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 + assert not wc.exited, "__exit__ must not run when __return__ itself is the failure" + + def test_wrapping_context_mid_call_registration(): """Registering a new context on a function while a call is in flight must not crash that call. @@ -1285,6 +1338,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 @@ -1362,6 +1416,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 @@ -1732,63 +1787,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() -> None: + """Storage set by one thread must not bleed into a concurrent thread.""" + import threading + + results: dict[int, int] = {} + errors: list[str] = [] + + class ThreadIsolationContext(DummyWrappingContext): + def __enter__(self) -> "ThreadIsolationContext": + super().__enter__() + self.set("tid", threading.get_ident()) + return self + + def __return__(self, value: Any) -> Any: + 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() -> None: + barrier.wait() + foo() + threads: list[threading.Thread] = [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() -> None: + """Each recursive coroutine call must have its own isolated storage slot.""" + values: list[tuple[str, int]] = [] - CaptureOnExit(foo).wrap() + class AsyncRecursiveContext(DummyWrappingContext): + def __enter__(self) -> "AsyncRecursiveContext": + super().__enter__() + n: int = self.__frame__.f_locals["n"] + self.set("n", n) + values.append(("enter", n)) + return self - with pytest.raises(_Sentinel): - foo() + 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) - 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: int) -> int: + 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..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 @@ -53,6 +56,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 @@ -68,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. @@ -80,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 f6de0b55f10..0a4a74f6938 100644 --- a/tests/wrapping/test_unwrap.py +++ b/tests/wrapping/test_unwrap.py @@ -10,15 +10,23 @@ 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 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 collections.abc import Callable import inspect +import sys +from types import FunctionType +from typing import Any +from typing import Union +from typing import cast import pytest @@ -30,11 +38,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) @@ -45,44 +53,48 @@ 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() +_Restore = Union[_InternalRestore, _ContextRestore] + + @pytest.fixture(params=[_InternalRestore, _ContextRestore], ids=["internal_wrap", "wrapping_context"]) -def restore(request): - return request.param() +def restore(request: pytest.FixtureRequest) -> _Restore: + return cast(_Restore, 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: _Restore) -> 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(): @@ -114,6 +126,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", )