diff --git a/ddtrace/internal/injection.py b/ddtrace/internal/injection.py index 787e0160e66..40ff6355967 100644 --- a/ddtrace/internal/injection.py +++ b/ddtrace/internal/injection.py @@ -7,6 +7,7 @@ from typing import Tuple # noqa:F401 from bytecode import Bytecode +from bytecode import Instr from ddtrace.internal.assembly import Assembly @@ -109,6 +110,11 @@ def _inject_hook(code: Bytecode, hook: HookType, lineno: int, arg: Any) -> None: raise InvalidLine("Line %d does not exist or is either blank or a comment" % lineno) for i in locs: + if isinstance(instr := code[i], Instr) and instr.name.startswith("END_"): + # This is the end of a block, e.g. a for loop. We have already + # instrumented the block on entry, so we skip instrumenting the + # end as well. + continue code[i:i] = INJECTION_ASSEMBLY.bind(dict(hook=hook, arg=arg), lineno=lineno) diff --git a/releasenotes/notes/fix-iteration-block-injection-af354b7004554dfd.yaml b/releasenotes/notes/fix-iteration-block-injection-af354b7004554dfd.yaml new file mode 100644 index 00000000000..34b06199785 --- /dev/null +++ b/releasenotes/notes/fix-iteration-block-injection-af354b7004554dfd.yaml @@ -0,0 +1,6 @@ +--- +fixes: + - | + dynamic instrumentation: fixed an issue with the instrumentation of the + first line of an iteration block (e.g. for loops) that could have caused + undefined behavior. diff --git a/tests/internal/test_injection.py b/tests/internal/test_injection.py index 871726620a8..b9366a6046f 100644 --- a/tests/internal/test_injection.py +++ b/tests/internal/test_injection.py @@ -253,3 +253,18 @@ def test_finally(): f() hook.assert_called_once_with(arg) + + +def test_for_block(): + def for_loop(): + a = [] + for i in range(10): + a.append(i) + return a + + hook, arg = mock.Mock(), mock.Mock() + + with injected_hook(for_loop, hook, arg, line=for_loop.__code__.co_firstlineno + 2): + for_loop() + + hook.assert_called_once_with(arg)