chore: lazy context wrapping (#15295)

## Description

We implement a lazy context wrapping mechanism whereby the bytecode
context wrapping is performed on first invocation of the function. Note
that this still requires the original function to be instrumented via
bytecode manipulations.

## Additional Notes

The motivation for this change is to allow for a more lightweight
bytecode instrumentation to trigger a more expensive one only when
needed. This is typically when an instrumented function is called for
the first time.

Author: P403n1x87

ddtrace/internal/wrapping/context.py

@@ -13,7 +13,13 @@
 from bytecode import Bytecode
 
 from ddtrace.internal.assembly import Assembly
+from ddtrace.internal.forksafe import Lock
 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 is_wrapped_with
+from ddtrace.internal.wrapping import unwrap
+from ddtrace.internal.wrapping import wrap
 
 
 T = t.TypeVar("T")
@@ -406,6 +412,44 @@ def unwrap(self) -> None:
             _UniversalWrappingContext.extract(f).unregister(self)
 
 
+class LazyWrappingContext(WrappingContext):
+    def __init__(self, f: FunctionType):
+        super().__init__(f)
+
+        self._trampoline: t.Optional[Wrapper] = None
+        self._trampoline_lock = Lock()
+
+    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(_, args, kwargs):
+                with tl:
+                    f = t.cast(WrappedFunction, self.__wrapped__)
+                    if is_wrapped_with(self.__wrapped__, trampoline):
+                        f = unwrap(f, trampoline)
+                        super(LazyWrappingContext, self).wrap()
+                return f(*args, **kwargs)
+
+            wrap(self.__wrapped__, trampoline)
+
+            self._trampoline = trampoline
+
+    def unwrap(self) -> None:
+        with self._trampoline_lock:
+            if self._trampoline is None:
+                return
+
+            if self.is_wrapped(self.__wrapped__):
+                super().unwrap()
+            else:
+                unwrap(t.cast(WrappedFunction, self.__wrapped__), self._trampoline)
+
+            self._trampoline = None
+
+
 class ContextWrappedFunction(Protocol):
     """A wrapped function."""
 

tests/internal/test_wrapping.py

@@ -10,6 +10,7 @@
 from ddtrace.internal.wrapping import is_wrapped_with
 from ddtrace.internal.wrapping import unwrap
 from ddtrace.internal.wrapping import wrap
+from ddtrace.internal.wrapping.context import LazyWrappingContext
 from ddtrace.internal.wrapping.context import WrappingContext
 from ddtrace.internal.wrapping.context import _UniversalWrappingContext
 
@@ -926,3 +927,42 @@ def foo():
 
     new_method_count = len([_ for _ in gc.get_objects() if type(_).__name__ == "method"])
     assert new_method_count <= method_count + 1
+
+
+def test_wrapping_context_lazy():
+    free = 42
+
+    def foo():
+        return free
+
+    class DummyLazyWrappingContext(LazyWrappingContext):
+        def __init__(self, f):
+            super().__init__(f)
+
+            self.count = 0
+
+        def __enter__(self):
+            self.count += 1
+            return super().__enter__()
+
+    (wc := DummyLazyWrappingContext(foo)).wrap()
+
+    assert not DummyLazyWrappingContext.is_wrapped(foo)
+
+    for _ in range(n := 10):
+        assert foo() == free
+
+        assert DummyLazyWrappingContext.is_wrapped(foo)
+
+    assert wc.count == n
+
+    wc.count = 0
+
+    wc.unwrap()
+
+    for _ in range(10):
+        assert not DummyLazyWrappingContext.is_wrapped(foo)
+
+        assert foo() == free
+
+    assert wc.count == 0