fix wrapping context-lazy wrapping context compat

Author: P403n1x87

ddtrace/debugging/_function/discovery.py

@@ -139,7 +139,10 @@ def resolve(self) -> FullyNamedFunction:
             msg = "Cannot resolve pair with no code object"
             raise ValueError(msg)
 
-        if self.function is not None:
+        # Check that the function we have cached is not a wrapper layer that
+        # has been unwrapped. In this case we need to resolve the new function
+        # from the code object.
+        if (_ := self.function) is not None and _.__name__ != "<unwrapped>":
             return cast(FullyNamedFunction, self.function)
 
         code = self.code

ddtrace/internal/wrapping/__init__.py

@@ -322,6 +322,13 @@ def unwrap(wf, wrapper):
         # current one.
         f = cast(FunctionType, wf)
         f.__code__ = inner.__code__
+
+        # Mark the function as unwrapped via its name. There might be references
+        # to this function elsewhere and this would signal that the function has
+        # been unwrapped and that another function object is referencing the
+        # original code object.
+        inner.__name__ = "<unwrapped>"
+
         try:
             # Update the link to the next layer.
             inner_wf = cast(WrappedFunction, inner)

ddtrace/internal/wrapping/context.py

@@ -394,6 +394,12 @@ 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__)

tests/debugging/mocking.py

@@ -177,7 +177,7 @@ def assert_single_snapshot(self):
 
         assert len(self.test_queue) == 1
 
-        yield self.test_queue[0]
+        yield self.test_queue.pop(0)
 
 
 @contextmanager

tests/debugging/test_debugger.py

@@ -642,6 +642,48 @@ def test_debugger_line_probe_on_lazy_wrapped_function(stuff):
             assert snapshot.probe.probe_id == "line-probe-lazy-wrapping"
 
 
+def test_debugger_function_probe_on_lazy_wrapped_function(stuff):
+    from ddtrace.internal.wrapping.context import LazyWrappingContext
+
+    class LWC(LazyWrappingContext):
+        def __enter__(self):
+            self.entered = True
+            return super().__enter__()
+
+    (c := LWC(stuff.Stuff.instancestuff)).wrap()
+
+    with debugger() as d:
+        d.add_probes(
+            probe := create_snapshot_function_probe(
+                probe_id="function-probe-lazy-wrapping",
+                module="tests.submod.stuff",
+                func_qname="Stuff.instancestuff",
+                rate=float("inf"),
+            )
+        )
+
+        stuff.Stuff().instancestuff(42)
+
+        d.remove_probes(probe)
+
+        assert c.entered
+
+        with d.assert_single_snapshot() as snapshot:
+            assert snapshot.probe.probe_id == "function-probe-lazy-wrapping"
+
+        # Test that we can re-instrument the function correctly and that we
+        # don't accidentally instrument the temporary trampoline instead.
+
+        d.add_probes(probe)
+
+        stuff.Stuff().instancestuff(42)
+
+        d.remove_probes(probe)
+
+        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