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,11 +394,29 @@ 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)
+                        try:
+                            f.__dd_context_wrapped__ = inner.__dd_context_wrapped__
+                        except AttributeError:
+                            pass
                         super(LazyWrappingContext, self).wrap()
                 return f(*args, **kwargs)
 

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,46 @@ 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):
+        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