fix(debugging): handle multiple injection points [backport #6759 to 1.18] (#6761)

Backport of #6759 to 1.18

Some code constructs might be translated into non-linear bytecode,
meaning that some line numbers might appear multiple time, as code
blocks might be copied in multiple place (e.g. finally blocks).

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
- [x] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [x] This PR doesn't touch any of that.

Author: P403n1x87

ddtrace/internal/injection.py

@@ -1,6 +1,8 @@
+from collections import deque
 from types import FunctionType
 from typing import Any
 from typing import Callable
+from typing import Deque
 from typing import List
 from typing import Tuple
 
@@ -30,52 +32,63 @@ def _inject_hook(code, hook, lineno, arg):
     identifier for the hook itself. This should be kept in case the hook needs
     to be removed.
     """
+    # DEV: In general there are no guarantees for bytecode to be "linear",
+    # meaning that a line number can occur multiple times. We need to find all
+    # occurrences and inject the hook at each of them. An example of when this
+    # happens is with finally blocks, which are duplicated at the end of the
+    # bytecode.
+    locs = deque()  # type: Deque[int]
+    last_lineno = None
     for i, instr in enumerate(code):
         try:
+            if instr.lineno == last_lineno:
+                continue
+            last_lineno = instr.lineno
             if instr.lineno == lineno:
-                # gotcha!
-                break
+                locs.appendleft(i)
         except AttributeError:
             # pseudo-instruction (e.g. label)
             pass
-    else:
+
+    if not locs:
         raise InvalidLine("Line %d does not exist or is either blank or a comment" % lineno)
 
     # DEV: This is the bytecode equivalent of
     # >>> hook(arg)
     # Additionally, we must discard the return value (top of the stack) to
     # restore the stack to the state prior to the call.
-    if PY < (3, 11):
-        code[i:i] = Bytecode(
-            [
-                Instr("LOAD_CONST", hook, lineno=lineno),
-                Instr("LOAD_CONST", arg, lineno=lineno),
-                Instr("CALL_FUNCTION", 1, lineno=lineno),
-                Instr("POP_TOP", lineno=lineno),
-            ]
-        )
-    elif PY >= (3, 12):
-        code[i:i] = Bytecode(
-            [
-                Instr("PUSH_NULL", lineno=lineno),
-                Instr("LOAD_CONST", hook, lineno=lineno),
-                Instr("LOAD_CONST", arg, lineno=lineno),
-                Instr("CALL", 1, lineno=lineno),
-                Instr("POP_TOP", lineno=lineno),
-            ]
-        )
-    else:
-        # Python 3.11
-        code[i:i] = Bytecode(
-            [
-                Instr("PUSH_NULL", lineno=lineno),
-                Instr("LOAD_CONST", hook, lineno=lineno),
-                Instr("LOAD_CONST", arg, lineno=lineno),
-                Instr("PRECALL", 1, lineno=lineno),
-                Instr("CALL", 1, lineno=lineno),
-                Instr("POP_TOP", lineno=lineno),
-            ]
-        )
+    for i in locs:
+        if PY < (3, 11):
+            code[i:i] = Bytecode(
+                [
+                    Instr("LOAD_CONST", hook, lineno=lineno),
+                    Instr("LOAD_CONST", arg, lineno=lineno),
+                    Instr("CALL_FUNCTION", 1, lineno=lineno),
+                    Instr("POP_TOP", lineno=lineno),
+                ]
+            )
+        elif PY >= (3, 12):
+            code[i:i] = Bytecode(
+                [
+                    Instr("PUSH_NULL", lineno=lineno),
+                    Instr("LOAD_CONST", hook, lineno=lineno),
+                    Instr("LOAD_CONST", arg, lineno=lineno),
+                    Instr("CALL", 1, lineno=lineno),
+                    Instr("POP_TOP", lineno=lineno),
+                ]
+            )
+        else:
+            # Python 3.11
+            code[i:i] = Bytecode(
+                [
+                    Instr("PUSH_NULL", lineno=lineno),
+                    Instr("LOAD_CONST", hook, lineno=lineno),
+                    Instr("LOAD_CONST", arg, lineno=lineno),
+                    Instr("PRECALL", 1, lineno=lineno),
+                    Instr("CALL", 1, lineno=lineno),
+                    Instr("POP_TOP", lineno=lineno),
+                ]
+            )
 
 
 # Default to Python 3.11 opcodes
@@ -96,6 +109,7 @@ def _eject_hook(code, hook, line, arg):
     The hook is identified by its argument. This ensures that only the right
     hook is ejected.
     """
+    locs = deque()  # type: Deque[int]
     for i, instr in enumerate(code):
         try:
             # DEV: We look at the expected opcode pattern to match the injected
@@ -106,17 +120,18 @@ def _eject_hook(code, hook, line, arg):
                 and code[i + _INJECT_ARG_OPCODE_POS].arg is arg
                 and [code[_].name for _ in range(i, i + len(_INJECT_HOOK_OPCODES))] == _INJECT_HOOK_OPCODES
             ):
-                # gotcha!
-                break
+                locs.appendleft(i)
         except AttributeError:
             # pseudo-instruction (e.g. label)
             pass
         except IndexError:
             pass
-    else:
+
+    if not locs:
         raise InvalidLine("Line %d does not contain a hook" % line)
 
-    del code[i : i + len(_INJECT_HOOK_OPCODES)]
+    for i in locs:
+        del code[i : i + len(_INJECT_HOOK_OPCODES)]
 
 
 def _function_with_new_code(f, abstract_code):

ddtrace/internal/utils/inspection.py

@@ -9,7 +9,11 @@
 def linenos(f):
     # type: (FunctionType) -> Set[int]
     """Get the line numbers of a function."""
-    ls = {instr.lineno for instr in Bytecode.from_code(f.__code__) if hasattr(instr, "lineno")}
+    ls = {
+        _
+        for _ in (instr.lineno for instr in Bytecode.from_code(f.__code__) if hasattr(instr, "lineno"))
+        if _ is not None
+    }
     if PY >= (3, 11):
         # Python 3.11 has introduced some no-op instructions that are used as
         # part of the specialisation process. These new opcodes appear on a

releasenotes/notes/fix-debugger-multiple-inject-locs-2dd1d33810cf6405.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    dynamic instrumentation: fixed an issue that prevented line probes from
+    being injected in some finally blocks.

tests/internal/test_injection.py

@@ -1,11 +1,11 @@
 from contextlib import contextmanager
 from random import shuffle
-import sys
 
 import mock
 import pytest
 from six import PY2
 
+from ddtrace.debugging._function.discovery import UnboundMethodType
 from ddtrace.internal.injection import InvalidLine
 from ddtrace.internal.injection import eject_hook
 from ddtrace.internal.injection import eject_hooks
@@ -15,18 +15,21 @@
 
 
 @contextmanager
-def injected_hook(f, hook, arg):
+def injected_hook(f, hook, arg, line=None):
+    if PY2 and isinstance(f, UnboundMethodType):
+        f = f.__func__
+
     code = f.__code__
-    line = min(linenos(f))
+
+    if line is None:
+        line = min(linenos(f))
 
     inject_hook(f, hook, line, arg)
 
     yield f
 
     eject_hook(f, hook, line, arg)
 
-    if sys.version_info[:2] not in {(3, 5), (3, 6)} and sys.version_info < (3, 11):
-        assert f.__code__ == code
     assert f.__code__ is not code
 
 
@@ -234,7 +237,7 @@ def test_inject_in_multiline():
 
 
 def test_property():
-    stuff = sys.modules["tests.submod.stuff"]
+    import tests.submod.stuff as stuff
 
     f = stuff.Stuff.propertystuff.fget
 
@@ -245,3 +248,17 @@ def test_property():
     stuff.Stuff().propertystuff
 
     hook.assert_called_once_with(arg)
+
+
+def test_finally():
+    import tests.submod.stuff as stuff
+
+    f = stuff.finallystuff
+
+    hook, arg = mock.Mock(), mock.Mock()
+
+    with injected_hook(f, hook, arg, line=157):
+        f()
+    f()
+
+    hook.assert_called_once_with(arg)

tests/submod/stuff.py

@@ -144,3 +144,15 @@ def age_checker(people, age, name=None):
 
 def caller(f, *args, **kwargs):
     return f(*args, **kwargs)
+
+
+def finallystuff():
+    a = 0
+    try:
+        if a == 0:
+            Exception("Hello", "world!", 42)
+    except Exception:
+        return a
+    finally:
+        a = 42
+    return a