fix TOS line number from log function probes

ddtrace/debugging/_origin/span.py

@@ -161,8 +161,6 @@ def __enter__(self):
             self.set("context", context)
             self.set("start_time", compat.monotonic_ns())
 
-            print("snapshot created")
-
     def _close_context(self, retval=None, exc_info=(None, None, None)):
         try:
             context: SignalContext = self.get("context")

ddtrace/debugging/_signal/snapshot.py

@@ -87,6 +87,8 @@ class Snapshot(LogSignal):
     return_capture = attr.ib(type=Optional[dict], default=None)
     line_capture = attr.ib(type=Optional[dict], default=None)
 
+    _stack = attr.ib(type=Optional[list], default=None)
+
     _message = attr.ib(type=Optional[str], default=None)
     duration = attr.ib(type=Optional[int], default=None)  # nanoseconds
 
@@ -159,7 +161,7 @@ def exit(self, retval, exc_info, duration):
             return
 
         _locals = list(_safety.get_locals(self.frame))
-        _, exc, _ = exc_info
+        _, exc, tb = exc_info
         if exc is None:
             _locals.append(("@return", retval))
         else:
@@ -174,6 +176,19 @@ def exit(self, retval, exc_info, duration):
         if probe.evaluate_at != ProbeEvaluateTimingForMethod.ENTER:
             self._eval_message(dict(_args))
 
+        stack = utils.capture_stack(self.frame)
+
+        # Fix the line number of the top frame. This might have been mangled by
+        # the instrumented exception handling of function probes.
+        while tb is not None:
+            frame = tb.tb_frame
+            if frame == self.frame:
+                stack[0]["lineNumber"] = tb.tb_lineno
+                break
+            tb = tb.tb_next
+
+        self._stack = stack
+
     def line(self):
         if not isinstance(self.probe, LogLineProbe):
             return
@@ -198,6 +213,9 @@ def line(self):
             )
 
         self._eval_message(frame.f_locals)
+
+        self._stack = utils.capture_stack(frame)
+
         self.state = SignalState.DONE
 
     @property
@@ -209,7 +227,6 @@ def has_message(self) -> bool:
 
     @property
     def data(self):
-        frame = self.frame
         probe = self.probe
 
         captures = None
@@ -223,7 +240,7 @@ def data(self):
                 }
 
         return {
-            "stack": utils.capture_stack(frame),
+            "stack": self._stack,
             "captures": captures,
             "duration": self.duration,
         }

ddtrace/debugging/_signal/utils.py

@@ -1,6 +1,7 @@
 from itertools import islice
 from itertools import takewhile
 from types import FrameType
+from types import TracebackType
 from typing import Any
 from typing import Callable
 from typing import Dict
@@ -125,20 +126,34 @@ def capture_stack(top_frame: FrameType, max_height: int = 4096) -> List[dict]:
     return stack
 
 
+def capture_traceback(tb: TracebackType, max_height: int = 4096) -> List[dict]:
+    stack = []
+    h = 0
+    _tb: Optional[TracebackType] = tb
+    while _tb is not None and h < max_height:
+        frame = _tb.tb_frame
+        code = frame.f_code
+        stack.append(
+            {
+                "fileName": code.co_filename,
+                "function": code.co_name,
+                "lineNumber": _tb.tb_lineno,
+            }
+        )
+        _tb = _tb.tb_next
+        h += 1
+    return stack
+
+
 def capture_exc_info(exc_info: ExcInfoType) -> Optional[Dict[str, Any]]:
     _type, value, tb = exc_info
     if _type is None or value is None:
         return None
 
-    top_tb = tb
-    if top_tb is not None:
-        while top_tb.tb_next is not None:
-            top_tb = top_tb.tb_next
-
     return {
         "type": _type.__name__,
         "message": ", ".join([serialize(v) for v in value.args]),
-        "stacktrace": capture_stack(top_tb.tb_frame) if top_tb is not None else None,
+        "stacktrace": capture_traceback(tb) if tb is not None else None,
     }