perf(debugger): reduce code origin startup time

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]
 

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.

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

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
 

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.

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__ == "<wrapped>", "Wrapper has wrapped function"  # nosec
+    # Sanity check
+    assert inner.__name__ == "<wrapped>", "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]

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__