Add branch monitoring for async for loops

python · markshannon · Feb 27, 2025 · Feb 26, 2025 · Feb 26, 2025 · Feb 26, 2025
commit df3f2437a17b9b7525c2edbe44abe200f7ed246d
diff --git a/Include/internal/pycore_opcode_metadata.h b/Include/internal/pycore_opcode_metadata.h
diff --git a/Include/opcode_ids.h b/Include/opcode_ids.h
diff --git a/Lib/_opcode_metadata.py b/Lib/_opcode_metadata.py
diff --git a/Lib/test/test_monitoring.py b/Lib/test/test_monitoring.py
@@ -1657,6 +1657,29 @@ def foo(n=0):
             in_loop,
             exit_loop])
 
+    def test_async_for(self):
+
+        def func():
+            async def gen():
+                yield 2
+                yield 3
+
+            async def foo():
+                async for y in gen():
+                    2
+                pass # line 3
+
+            try:
+                foo().send(None)
+            except StopIteration:
+                pass
+
+        self.check_events(func, recorders = BRANCHES_RECORDERS, expected = [
+            ('branch left', 'foo', 1, 1),
+            ('branch left', 'foo', 1, 1),
+            ('branch right', 'foo', 1, 3),
+            ('branch left', 'func', 12, 12)])
+
 
 class TestBranchConsistency(MonitoringTestBase, unittest.TestCase):
 

@@ -1339,7 +1339,7 @@ dummy_func(
             goto exception_unwind;
         }
 
-        tier1 inst(END_ASYNC_FOR, (awaitable_st, exc_st -- )) {
+        tier1 op(_END_ASYNC_FOR, (awaitable_st, exc_st -- )) {
             PyObject *exc = PyStackRef_AsPyObjectBorrow(exc_st);
 
             assert(exc && PyExceptionInstance_Check(exc));
@@ -1355,6 +1355,16 @@ dummy_func(
             }
         }
 
+        tier1 op(_MONITOR_BRANCH_RIGHT, ( -- )) {
+            INSTRUMENTED_JUMP(prev_instr, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
+        }
+
+        macro(INSTRUMENTED_END_ASYNC_FOR) =
+            _MONITOR_BRANCH_RIGHT +
+            _END_ASYNC_FOR;
+
+        macro(END_ASYNC_FOR) = _END_ASYNC_FOR;
+
         tier1 inst(CLEANUP_THROW, (sub_iter, last_sent_val, exc_value_st -- none, value)) {
             PyObject *exc_value = PyStackRef_AsPyObjectBorrow(exc_value_st);
             #if !Py_TAIL_CALL_INTERP

@@ -2041,6 +2041,7 @@ codegen_async_for(compiler *c, stmt_ty s)
     ADDOP_LOAD_CONST(c, loc, Py_None);
     ADD_YIELD_FROM(c, loc, 1);
     ADDOP(c, loc, POP_BLOCK);  /* for SETUP_FINALLY */
+    ADDOP(c, loc, NOT_TAKEN);
 
     /* Success block for __anext__ */
     VISIT(c, expr, s->v.AsyncFor.target);

diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/instrumentation.c b/Python/instrumentation.c
@@ -106,6 +106,7 @@ static const int8_t EVENT_FOR_OPCODE[256] = {
     [INSTRUMENTED_END_SEND] = PY_MONITORING_EVENT_STOP_ITERATION,
     [NOT_TAKEN] = PY_MONITORING_EVENT_BRANCH_LEFT,
     [INSTRUMENTED_NOT_TAKEN] = PY_MONITORING_EVENT_BRANCH_LEFT,
+    [END_ASYNC_FOR] = PY_MONITORING_EVENT_BRANCH_RIGHT,
 };
 
 static const uint8_t DE_INSTRUMENT[256] = {
@@ -127,6 +128,7 @@ static const uint8_t DE_INSTRUMENT[256] = {
     [INSTRUMENTED_END_SEND] = END_SEND,
     [INSTRUMENTED_LOAD_SUPER_ATTR] = LOAD_SUPER_ATTR,
     [INSTRUMENTED_NOT_TAKEN] = NOT_TAKEN,
+    [INSTRUMENTED_END_ASYNC_FOR] = END_ASYNC_FOR,
 };
 
 static const uint8_t INSTRUMENTED_OPCODES[256] = {
@@ -166,6 +168,8 @@ static const uint8_t INSTRUMENTED_OPCODES[256] = {
     [INSTRUMENTED_LOAD_SUPER_ATTR] = INSTRUMENTED_LOAD_SUPER_ATTR,
     [NOT_TAKEN] = INSTRUMENTED_NOT_TAKEN,
     [INSTRUMENTED_NOT_TAKEN] = INSTRUMENTED_NOT_TAKEN,
+    [END_ASYNC_FOR] = INSTRUMENTED_END_ASYNC_FOR,
+    [INSTRUMENTED_END_ASYNC_FOR] = INSTRUMENTED_END_ASYNC_FOR,
 
     [INSTRUMENTED_LINE] = INSTRUMENTED_LINE,
     [INSTRUMENTED_INSTRUCTION] = INSTRUMENTED_INSTRUCTION,

diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h