diff --git a/ddtrace/debugging/_function/discovery.py b/ddtrace/debugging/_function/discovery.py index 7fc219f6ef1..711f861b115 100644 --- a/ddtrace/debugging/_function/discovery.py +++ b/ddtrace/debugging/_function/discovery.py @@ -25,6 +25,7 @@ from ddtrace.internal.utils.inspection import linenos from ddtrace.internal.utils.inspection import resolved_code_origin from ddtrace.internal.utils.inspection import undecorated +from ddtrace.internal.wrapping import get_function_code log = get_logger(__name__) @@ -135,13 +136,13 @@ def __init__(self, code: Optional[CodeType] = None, function: Optional[FunctionT self.code = function.__code__ if function is not None else code def resolve(self) -> FullyNamedFunction: + if self.function is not None: + return cast(FullyNamedFunction, self.function) + if self.code is None: msg = "Cannot resolve pair with no code object" raise ValueError(msg) - if self.function is not None: - return cast(FullyNamedFunction, self.function) - code = self.code functions = functions_for_code(code) n = len(functions) @@ -285,11 +286,11 @@ def __init__(self, module: ModuleType) -> None: if ( function not in seen_functions - and resolved_code_origin(cast(FunctionType, function).__code__) == module_path + and resolved_code_origin(code := get_function_code(cast(FunctionType, function))) == module_path ): # We only map line numbers for functions that actually belong to # the module. - for lineno in linenos(cast(FunctionType, function)): + for lineno in linenos(cast(FunctionType, code)): self[lineno].append(_FunctionCodePair(function=cast(FunctionType, function))) seen_functions.add(function) @@ -299,6 +300,8 @@ def at_line(self, line: int) -> List[FullyNamedFunction]: Note that, in general, there can be multiple copies of the same functions. This can happen as a result, e.g., of using decorators. """ + fcp: _FunctionCodePair + if line in self._cached: return self._cached[line] diff --git a/ddtrace/debugging/_function/store.py b/ddtrace/debugging/_function/store.py index e52ac3cc0d5..c35f7946f0c 100644 --- a/ddtrace/debugging/_function/store.py +++ b/ddtrace/debugging/_function/store.py @@ -13,6 +13,7 @@ from ddtrace.internal.bytecode_injection import HookType from ddtrace.internal.bytecode_injection import eject_hooks from ddtrace.internal.bytecode_injection import inject_hooks +from ddtrace.internal.wrapping import get_function_code from ddtrace.internal.wrapping.context import ContextWrappedFunction from ddtrace.internal.wrapping.context import WrappingContext @@ -52,7 +53,7 @@ def __exit__(self, *exc): def _store(self, function: FunctionType) -> None: if function not in self._code_map: - self._code_map[function] = function.__code__ + self._code_map[function] = get_function_code(function) def inject_hooks(self, function: FullyNamedContextWrappedFunction, hooks: List[HookInfoType]) -> Set[str]: """Bulk-inject hooks into a function. diff --git a/ddtrace/debugging/_origin/span.py b/ddtrace/debugging/_origin/span.py index c4b9c182f45..c3491400060 100644 --- a/ddtrace/debugging/_origin/span.py +++ b/ddtrace/debugging/_origin/span.py @@ -20,7 +20,7 @@ from ddtrace.internal.forksafe import Lock from ddtrace.internal.logger import get_logger from ddtrace.internal.safety import _isinstance -from ddtrace.internal.wrapping.context import WrappingContext +from ddtrace.internal.wrapping.context import LazyWrappingContext log = get_logger(__name__) @@ -67,7 +67,14 @@ class EntrySpanLocation: probe: EntrySpanProbe -class EntrySpanWrappingContext(WrappingContext): +class EntrySpanWrappingContext(LazyWrappingContext): + """Entry span wrapping context. + + This context is lazy to avoid paid any upfront instrumentation costs for + large functions that might not get invoked. Instead, the actual wrapping + will be performed on the first invocation. + """ + __enabled__ = False __priority__ = 199 @@ -183,7 +190,7 @@ def instrument_view(cls, f: t.Union[FunctionType, MethodType]) -> None: _f = t.cast(FunctionType, f) if not EntrySpanWrappingContext.is_wrapped(_f): - log.debug("Patching entrypoint %r for code origin", f) + log.debug("Lazy wrapping entrypoint %r for code origin", f) EntrySpanWrappingContext(cls.__uploader__, _f).wrap() @classmethod diff --git a/ddtrace/internal/bytecode_injection/__init__.py b/ddtrace/internal/bytecode_injection/__init__.py index a151fbea654..73b05aed9c4 100644 --- a/ddtrace/internal/bytecode_injection/__init__.py +++ b/ddtrace/internal/bytecode_injection/__init__.py @@ -10,6 +10,8 @@ from ddtrace.internal.assembly import Assembly from ddtrace.internal.compat import PYTHON_VERSION_INFO as PY +from ddtrace.internal.wrapping import get_function_code +from ddtrace.internal.wrapping import set_function_code HookType = Callable[[Any], Any] @@ -172,7 +174,7 @@ 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(f.__code__) + abstract_code = Bytecode.from_code(get_function_code(f)) failed = [] for hook, line, arg in hooks: @@ -182,7 +184,7 @@ def inject_hooks(f: FunctionType, hooks: List[HookInfoType]) -> List[HookInfoTyp failed.append((hook, line, arg)) if len(failed) < len(hooks): - f.__code__ = abstract_code.to_code() + set_function_code(f, abstract_code.to_code()) return failed diff --git a/ddtrace/internal/utils/inspection.py b/ddtrace/internal/utils/inspection.py index a8a3ca28e82..db5c8fd7bb2 100644 --- a/ddtrace/internal/utils/inspection.py +++ b/ddtrace/internal/utils/inspection.py @@ -50,9 +50,6 @@ def undecorated(f: FunctionType, name: str, path: Path) -> FunctionType: def match(g): return g.__code__.co_name == name and resolved_code_origin(g.__code__) == path - if _isinstance(f, FunctionType) and match(f): - return f - seen_functions = {f} q = deque([f]) # FIFO: use popleft and append @@ -101,6 +98,15 @@ def match(g): q.append(c) seen_functions.add(c) + # If the function has bytecode wrapping we return the function itself. + # We don't want to return the temporary wrapped function from the + # __dd_wrapped__ attribute. + try: + object.__getattribute__(g, "__dd_wrapped__") + return g + except AttributeError: + pass + # Look for a function attribute (method decoration) # DEV: We don't recurse over arbitrary objects. We stop at the first # depth level. diff --git a/ddtrace/internal/wrapping/__init__.py b/ddtrace/internal/wrapping/__init__.py index 53a0b6c35b1..b9563e4eb34 100644 --- a/ddtrace/internal/wrapping/__init__.py +++ b/ddtrace/internal/wrapping/__init__.py @@ -14,6 +14,7 @@ from bytecode import Instr from ddtrace.internal.assembly import Assembly +from ddtrace.internal.utils.inspection import link_function_to_code from ddtrace.internal.wrapping.asyncs import wrap_async from ddtrace.internal.wrapping.generators import wrap_generator @@ -301,6 +302,9 @@ def wrap(f: FunctionType, wrapper: Wrapper) -> WrappedFunction: wf = cast(WrappedFunction, f) wf.__dd_wrapped__ = wrapped + # Link the original code object to the original function + link_function_to_code(code, f) + return wf @@ -348,30 +352,36 @@ def unwrap(wf: WrappedFunction, wrapper: Wrapper) -> FunctionType: # update the link at the deletion site if there is a non-empty tail. try: inner = cast(FunctionType, wf.__dd_wrapped__) + except AttributeError: + # The function is not wrapped so we return it as is. + return cast(FunctionType, wf) - # Sanity check - assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + # Sanity check + assert inner.__name__ == "", "Wrapper has wrapped function" # nosec + + if wrapper not in cast(FunctionType, wf).__code__.co_consts: + # This is not the correct wrapping layer. Try with the next one. + return unwrap(cast(WrappedFunction, inner), wrapper) - if wrapper not in cast(FunctionType, wf).__code__.co_consts: - # This is not the correct wrapping layer. Try with the next one. - inner_wf = cast(WrappedFunction, inner) - return unwrap(inner_wf, wrapper) - - # Remove the current wrapping layer by moving the next one over the - # current one. - f = cast(FunctionType, wf) - f.__code__ = inner.__code__ - try: - # Update the link to the next layer. - inner_wf = cast(WrappedFunction, inner) - wf.__dd_wrapped__ = inner_wf.__dd_wrapped__ # type: ignore[assignment] - except AttributeError: - # No more wrapping layers. Restore the original function by removing - # this extra attribute. - del wf.__dd_wrapped__ - - return f + # Remove the current wrapping layer by moving the next one over the + # current one. + f = cast(FunctionType, wf) + f.__code__ = inner.__code__ + try: + # Update the link to the next layer. + wf.__dd_wrapped__ = cast(WrappedFunction, inner).__dd_wrapped__ # type: ignore[assignment] except AttributeError: - # The function is not wrapped so we return it as is. - return cast(FunctionType, wf) + # No more wrapping layers. Restore the original function by removing + # this extra attribute. + del wf.__dd_wrapped__ + + return f + + +def get_function_code(f: FunctionType) -> CodeType: + return (cast(WrappedFunction, f).__dd_wrapped__ or f if is_wrapped(f) else f).__code__ + + +def set_function_code(f: FunctionType, code: CodeType) -> None: + (cast(WrappedFunction, f).__dd_wrapped__ or f if is_wrapped(f) else f).__code__ = code # type: ignore[misc] diff --git a/ddtrace/internal/wrapping/context.py b/ddtrace/internal/wrapping/context.py index c18070a3843..756f552ceb7 100644 --- a/ddtrace/internal/wrapping/context.py +++ b/ddtrace/internal/wrapping/context.py @@ -17,7 +17,9 @@ from ddtrace.internal.utils.inspection import link_function_to_code from ddtrace.internal.wrapping import WrappedFunction from ddtrace.internal.wrapping import Wrapper +from ddtrace.internal.wrapping import get_function_code from ddtrace.internal.wrapping import is_wrapped_with +from ddtrace.internal.wrapping import set_function_code from ddtrace.internal.wrapping import unwrap from ddtrace.internal.wrapping import wrap @@ -394,11 +396,30 @@ def wrap(self) -> None: if self._trampoline is not None: return + # If the function is already universally wrapped so it's less expensive + # to do the normal wrapping. + if _UniversalWrappingContext.is_wrapped(self.__wrapped__): + super().wrap() + return + def trampoline(_, args, kwargs): with tl: f = t.cast(WrappedFunction, self.__wrapped__) if is_wrapped_with(self.__wrapped__, trampoline): + # If the wrapped function was instrumented with a + # wrapping context before the first invocation we need + # to carry that over to the original function when we + # remove the trampoline. + try: + inner = f.__dd_wrapped__ + except AttributeError: + inner = None f = unwrap(f, trampoline) + self._trampoline = None + try: + f.__dd_context_wrapped__ = inner.__dd_context_wrapped__ + except AttributeError: + pass super(LazyWrappingContext, self).wrap() return f(*args, **kwargs) @@ -408,15 +429,12 @@ def trampoline(_, args, kwargs): def unwrap(self) -> None: with self._trampoline_lock: - if self._trampoline is None: - return - if self.is_wrapped(self.__wrapped__): + assert self._trampoline is None # nosec super().unwrap() - else: + elif self._trampoline is not None: unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline) - - self._trampoline = None + self._trampoline = None class ContextWrappedFunction(Protocol): @@ -498,9 +516,9 @@ def is_wrapped(cls, f: FunctionType) -> bool: # __dd_context_wrapped__ attribute is not enough, as this could be # copied over from an object state cloning. if sys.version_info >= (3, 11): - return f.__dd_context_wrapped__.__enter__ in f.__code__.co_consts # type: ignore + return f.__dd_context_wrapped__.__enter__ in get_function_code(f).co_consts # type: ignore else: - return f.__dd_context_wrapped__ in f.__code__.co_consts # type: ignore + return f.__dd_context_wrapped__ in get_function_code(f).co_consts # type: ignore except AttributeError: return False @@ -518,7 +536,7 @@ def wrap(self) -> None: if self.is_wrapped(f): raise ValueError("Function already wrapped") - bc = Bytecode.from_code(f.__code__) + bc = Bytecode.from_code(code := get_function_code(f)) # Prefix every return i = 0 @@ -552,7 +570,7 @@ def wrap(self) -> None: else: i = 0 - bc[i:i] = CONTEXT_HEAD.bind({"context_enter": self.__enter__}, lineno=f.__code__.co_firstlineno) + 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() @@ -588,8 +606,9 @@ def wrap(self) -> None: # Replace the function code with the wrapped code. We also link # the function to its original code object so that we can retrieve # it later if required. - link_function_to_code(f.__code__, f) - f.__code__ = bc.to_code() + link_function_to_code(code, f) + + set_function_code(f, bc.to_code()) def unwrap(self) -> None: f = self.__wrapped__ @@ -599,7 +618,7 @@ def unwrap(self) -> None: wrapped = t.cast(ContextWrappedFunction, f) - bc = Bytecode.from_code(f.__code__) + bc = Bytecode.from_code(get_function_code(f)) # Remove the exception handling code bc[-len(CONTEXT_FOOT) :] = [] @@ -663,7 +682,7 @@ def unwrap(self) -> None: i += 1 # Recreate the code object - f.__code__ = bc.to_code() + set_function_code(f, bc.to_code()) # Remove the wrapping context marker del wrapped.__dd_context_wrapped__ @@ -676,7 +695,7 @@ def wrap(self) -> None: if self.is_wrapped(f): raise ValueError("Function already wrapped") - bc = Bytecode.from_code(f.__code__) + bc = Bytecode.from_code(code := get_function_code(f)) # Prefix every return i = 0 @@ -706,7 +725,7 @@ def wrap(self) -> None: # Not an instruction pass - *bc[i:i], except_label = CONTEXT_HEAD.bind({"context": self}, lineno=f.__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()) @@ -717,8 +736,8 @@ def wrap(self) -> None: # Replace the function code with the wrapped code. We also link # the function to its original code object so that we can retrieve # it later if required. - link_function_to_code(f.__code__, f) - f.__code__ = bc.to_code() + link_function_to_code(code, f) + set_function_code(f, bc.to_code()) def unwrap(self) -> None: f = self.__wrapped__ @@ -728,7 +747,7 @@ def unwrap(self) -> None: wrapped = t.cast(ContextWrappedFunction, f) - bc = Bytecode.from_code(f.__code__) + bc = Bytecode.from_code(get_function_code(f)) # Remove the exception handling code bc[-len(CONTEXT_FOOT) :] = [] @@ -760,7 +779,7 @@ def unwrap(self) -> None: i += 1 # Recreate the code object - f.__code__ = bc.to_code() + set_function_code(f, bc.to_code()) # Remove the wrapping context marker del wrapped.__dd_context_wrapped__ diff --git a/tests/debugging/mocking.py b/tests/debugging/mocking.py index 746f9cd2691..59b4a444629 100644 --- a/tests/debugging/mocking.py +++ b/tests/debugging/mocking.py @@ -110,7 +110,10 @@ def wait_for_payloads(self, cond=lambda _: bool(_), timeout=1.0): raise PayloadWaitTimeout(_cond, timeout) sleep(0.05) - return self.payloads + try: + return self.payloads + finally: + self.flush() def flush(self) -> None: self.queue.clear() @@ -177,7 +180,7 @@ def assert_single_snapshot(self): assert len(self.test_queue) == 1 - yield self.test_queue[0] + yield self.test_queue.pop(0) @contextmanager diff --git a/tests/debugging/origin/test_span.py b/tests/debugging/origin/test_span.py index fee287f3652..ea34e0af51d 100644 --- a/tests/debugging/origin/test_span.py +++ b/tests/debugging/origin/test_span.py @@ -174,3 +174,44 @@ def entry_call(self): entry.get_tag("_dd.code_origin.frames.0.method") == "SpanProbeTestCase.test_span_origin_entry_method..App.entry_call" ) + + +def test_instrument_view_benchmark(benchmark): + """Benchmark instrument_view performance when wrapping functions.""" + MockSpanCodeOriginProcessorEntry.enable() + + try: + + def setup(): + """Create a unique function to wrap for each iteration.""" + + # Create a more realistic view function similar to Flask views + # with decorators, imports, and more complex code + def realistic_view(request_arg, *args, **kwargs): + """A realistic view function with actual code.""" + import json + import os # noqa + + data = {"status": "ok", "items": []} + for i in range(10): + item = { + "id": i, + "name": f"item_{i}", + "value": i * 100, + } + data["items"].append(item) + + result = json.dumps(data) + return result + + return (realistic_view,), {} + + # Benchmark the wrapping operation + benchmark.pedantic( + MockSpanCodeOriginProcessorEntry.instrument_view, + setup=setup, + rounds=100, + ) + + finally: + MockSpanCodeOriginProcessorEntry.disable() diff --git a/tests/debugging/test_debugger.py b/tests/debugging/test_debugger.py index 23a8b295be2..7cf952048bd 100644 --- a/tests/debugging/test_debugger.py +++ b/tests/debugging/test_debugger.py @@ -621,6 +621,101 @@ def test_debugger_line_probe_on_wrapped_function(stuff): assert snapshot.probe.probe_id == "line-probe-wrapped-method" +def test_debugger_function_and_line_probse_on_lazy_wrapped_function(stuff): + """Test that we can correctly instrument view functions with line and function + probes. + """ + from ddtrace.internal.wrapping.context import LazyWrappingContext + + class CounterWC(LazyWrappingContext): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.count = 0 + + def __enter__(self): + self.count += 1 + return super().__enter__() + + (wc := CounterWC(stuff.durationstuff)).wrap() + + line_probes = [ + create_snapshot_line_probe( + probe_id=f"line-probe-lazy-wrapping-{n}", + source_file="tests/submod/stuff.py", + line=132, + condition=None, + rate=float("inf"), + ) + for n in range(10) + ] + + function_probes = [ + create_snapshot_function_probe( + probe_id=f"function-probe-lazy-wrapping-{n}", + module="tests.submod.stuff", + func_qname="durationstuff", + rate=float("inf"), + ) + for n in range(10) + ] + + with debugger() as d: + for r in range(10): + d.add_probes(*function_probes, *line_probes) + + stuff.durationstuff(0) + + d.remove_probes(*function_probes, *line_probes) + + snapshots = d.uploader.wait_for_payloads(len(function_probes) + len(line_probes)) + + assert {s["debugger"]["snapshot"]["probe"]["id"] for s in snapshots} == { + f"line-probe-lazy-wrapping-{n}" for n in range(len(line_probes)) + } | {f"function-probe-lazy-wrapping-{n}" for n in range(len(function_probes))} + + assert wc.count == r + 1 + + +def test_debugger_function_probe_on_lazy_wrapped_function(stuff): + from ddtrace.internal.wrapping.context import LazyWrappingContext + + class LWC(LazyWrappingContext): + entered = False + + def __enter__(self): + self.entered = True + return super().__enter__() + + (c := LWC(stuff.throwexcstuff)).wrap() + + probe = create_snapshot_function_probe( + probe_id="function-probe-lazy-wrapping", + module="tests.submod.stuff", + func_qname="throwexcstuff", + rate=float("inf"), + ) + + with debugger() as d: + # Test that we can re-instrument the function correctly and that we + # don't accidentally instrument the temporary trampoline instead. + for _ in range(10): + d.add_probes(probe) + + try: + stuff.throwexcstuff() + except Exception: + pass + + d.remove_probes(probe) + + assert c.entered + + c.entered = False + + with d.assert_single_snapshot() as snapshot: + assert snapshot.probe.probe_id == "function-probe-lazy-wrapping" + + def test_probe_status_logging(remote_config_worker, stuff): assert remoteconfig_poller.status == ServiceStatus.STOPPED @@ -1083,7 +1178,7 @@ def test_debugger_modified_probe(stuff): stuff.Stuff().instancestuff() - _, msg = d.uploader.wait_for_payloads(2) + (msg,) = d.uploader.wait_for_payloads(1) assert "hello brave new world" == msg["message"], msg assert msg["debugger"]["snapshot"]["probe"]["version"] == 2, msg