gh-154701: prevent executor self-links in JIT cold exits

Project: python/cpython | Language: C, Python | Concepts: algorithmic-optimization, error-handling | Difficulty: senior
Link: https://github.com/python/cpython/pull/155323

The problem

When the JIT (Just-In-Time) compiler in Python encounters a FOR_ITER guard failure (i.e., a loop iteration can't proceed as expected), the fallback path could incorrectly link the executor (a JIT-compiled code block) back to itself. This creates a self-loop in the executor chain, causing an infinite loop where execution can never make forward progress.

Why it's non-trivial

A naive fix might simply try to break the loop or ignore the self-link, but the executor and bytecode state must be carefully restored to ensure future executions are correct and don't get stuck or corrupt the interpreter state. Handling this at the right point in the control flow, and ensuring the stack and executor are properly detached, is subtle and requires understanding of the JIT's tiered execution and fallback paths.

The approach

1. Mark _COLD_EXIT as an escaping instruction.
This change updates the uop (micro-op) metadata to declare that _COLD_EXIT can escape, which signals to the JIT machinery that special handling is needed for state restoration during these exits.

-    [_COLD_EXIT] = HAS_SYNC_SP_FLAG,
+    [_COLD_EXIT] = HAS_ESCAPES_FLAG | HAS_SYNC_SP_FLAG,

Annotation: This makes the JIT aware that _COLD_EXIT can cause control flow to escape the optimized path, requiring stack and state validation.

2. Detect and break executor self-links in the JIT fallback path.
In the JIT's executor fallback logic, the code now checks if the executor being linked is the same as the one from the exit. If so, it detaches the executor and falls back to Tier 1 (the interpreter), preventing the infinite self-loop.

+                if (executor == _PyExecutor_FromExit(exit)) {
+                    _Py_ExecutorDetach(executor);
+                    GOTO_TIER_ONE(target);
+                }

Annotation: This check prevents the executor from being linked to itself. If detected, it detaches the executor and jumps to the interpreter's execution path.

3. Ensure stack pointer and state are validated and invalidated properly.
In the generated executor cases, stack pointer validation and invalidation are performed around the detach operation to maintain interpreter invariants.

+                if (executor == _PyExecutor_FromExit(exit)) {
+                    _PyFrame_SetStackPointer(frame, stack_pointer);
+                    _PyFrame_StackPointerValidate(frame);
+                    _Py_ExecutorDetach(executor);
+                    _PyFrame_StackPointerInvalidate(frame);
+                    SET_CURRENT_CACHED_VALUES(0);
+                    GOTO_TIER_ONE(target);
+                }

Annotation: This ensures the frame's stack pointer is correctly managed before and after detaching the executor, avoiding stack corruption.

4. Add a regression test to ensure the bug is fixed.
A new test is added to check that side exits from FOR_ITER do not create self-links, using different iterator types to trigger the JIT fallback.

+    @isolation.runInSubprocess(timeout=SHORT_TIMEOUT)
+    def test_for_iter_side_exit_does_not_self_link(self):
+        def exhaust(iterator):
+            for _ in iterator:
+                pass
+        ...
+        for iterator in warmup_iterators:
+            exhaust(iterator)
+        # A different iterator type must not link that executor to itself.
+        exhaust(map(bool, values))

Annotation: This test exhausts several iterator types to warm up the JIT, then ensures that a new iterator type does not trigger the infinite loop.

5. Make _Py_ExecutorDetach public.
The function for detaching executors is made public so it can be used in the fallback logic.

-int _Py_ExecutorInit(_PyExecutorObject *, const _PyBloomFilter *);
-void _Py_ExecutorDetach(_PyExecutorObject *);
+int _Py_ExecutorInit(_PyExecutorObject *, const _PyBloomFilter *);
+PyAPI_FUNC(void) _Py_ExecutorDetach(_PyExecutorObject *);

Annotation: This change allows the fallback path to call the detach function as needed.

6. Add a news entry.
A news file is added to document the bugfix for users and maintainers.

+Fix an infinite loop in JIT when a ``FOR_ITER`` side exit links an executor back to itself.

Annotation: Documents the root cause and fix for future reference.

From the review thread

mgorny: Thanks. I can confirm that it fixes my issue with main. Looks like it needs updating for 3.15, as there is no _PyFrame_StackPointerValidate() and _PyFrame_StackPointerUnvalidate() there.

mgorny: @cocolato, any chance you could prepare a 3.15 backport for this?

Takeaways

Always validate and restore interpreter state (such as stack pointers) when transitioning between optimized and fallback execution paths.
Preventing infinite loops in JIT or interpreter machinery often requires explicit checks for self-referential links or cycles in control flow.
Regression tests targeting subtle control flow bugs should cover multiple input types and execution paths to ensure the fix is robust.