gh-143183: Link trace to side exits, rather than stop (GH-143268)

Author: Fidget-Spinner

Lib/test/test_capi/test_opt.py

@@ -60,6 +60,13 @@ def iter_opnames(ex):
 def get_opnames(ex):
     return list(iter_opnames(ex))
 
+def iter_ops(ex):
+    for item in ex:
+        yield item
+
+def get_ops(ex):
+    return list(iter_ops(ex))
+
 
 @requires_specialization
 @unittest.skipIf(Py_GIL_DISABLED, "optimizer not yet supported in free-threaded builds")
@@ -3003,14 +3010,25 @@ def f():
         # Outer loop warms up later, linking to the inner one.
         # Therefore, we have at least two executors.
         self.assertGreaterEqual(len(all_executors), 2)
+        executor_ids = [id(e) for e in all_executors]
         for executor in all_executors:
-            opnames = list(get_opnames(executor))
+            ops = get_ops(executor)
             # Assert all executors first terminator ends in
             # _EXIT_TRACE or _JUMP_TO_TOP, not _DEOPT
-            for idx, op in enumerate(opnames):
-                if op == "_EXIT_TRACE" or op == "_JUMP_TO_TOP":
+            for idx, op in enumerate(ops):
+                opname = op[0]
+                if opname == "_EXIT_TRACE":
+                    # As this is a link outer executor to inner
+                    # executor problem, all executors exits should point to
+                    # another valid executor. In this case, none of them
+                    # should be the cold executor.
+                    exit = op[3]
+                    link_to = _testinternalcapi.get_exit_executor(exit)
+                    self.assertIn(id(link_to), executor_ids)
+                    break
+                elif opname == "_JUMP_TO_TOP":
                     break
-                elif op == "_DEOPT":
+                elif opname == "_DEOPT":
                     self.fail(f"_DEOPT encountered first at executor"
                               f" {executor} at offset {idx} rather"
                               f" than expected _EXIT_TRACE")

Modules/_testinternalcapi.c

@@ -1245,6 +1245,22 @@ invalidate_executors(PyObject *self, PyObject *obj)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+get_exit_executor(PyObject *self, PyObject *arg)
+{
+    if (!PyLong_CheckExact(arg)) {
+        PyErr_SetString(PyExc_TypeError, "argument must be an ID to an _PyExitData");
+        return NULL;
+    }
+    uint64_t ptr;
+    if (PyLong_AsUInt64(arg, &ptr) < 0) {
+        // Error set by PyLong API
+        return NULL;
+    }
+    _PyExitData *exit = (_PyExitData *)ptr;
+    return Py_NewRef(exit->executor);
+}
+
 #endif
 
 static int _pending_callback(void *arg)
@@ -2546,6 +2562,7 @@ static PyMethodDef module_functions[] = {
 #ifdef _Py_TIER2
     {"add_executor_dependency", add_executor_dependency, METH_VARARGS, NULL},
     {"invalidate_executors", invalidate_executors, METH_O, NULL},
+    {"get_exit_executor", get_exit_executor, METH_O, NULL},
 #endif
     {"pending_threadfunc", _PyCFunction_CAST(pending_threadfunc),
      METH_VARARGS | METH_KEYWORDS},

Python/optimizer.c

@@ -625,7 +625,6 @@ _PyJit_translate_single_bytecode_to_trace(
     int trace_length = _tstate->jit_tracer_state.prev_state.code_curr_size;
     _PyUOpInstruction *trace = _tstate->jit_tracer_state.code_buffer;
     int max_length = _tstate->jit_tracer_state.prev_state.code_max_size;
-    int exit_op = stop_tracing_opcode == 0 ? _EXIT_TRACE : stop_tracing_opcode;
 
     _Py_CODEUNIT *this_instr =  _tstate->jit_tracer_state.prev_state.instr;
     _Py_CODEUNIT *target_instr = this_instr;
@@ -691,13 +690,18 @@ _PyJit_translate_single_bytecode_to_trace(
         goto full;
     }
 
-    if (stop_tracing_opcode != 0) {
+    if (stop_tracing_opcode == _DEOPT) {
         // gh-143183: It's important we rewind to the last known proper target.
         // The current target might be garbage as stop tracing usually indicates
         // we are in something that we can't trace.
         DPRINTF(2, "Told to stop tracing\n");
         goto unsupported;
     }
+    else if (stop_tracing_opcode != 0) {
+        assert(stop_tracing_opcode == _EXIT_TRACE);
+        ADD_TO_TRACE(stop_tracing_opcode, 0, 0, target);
+        goto done;
+    }
 
     DPRINTF(2, "%p %d: %s(%d) %d %d\n", old_code, target, _PyOpcode_OpName[opcode], oparg, needs_guard_ip, old_stack_level);
 
@@ -733,7 +737,7 @@ _PyJit_translate_single_bytecode_to_trace(
                 int32_t old_target = (int32_t)uop_get_target(curr);
                 curr++;
                 trace_length++;
-                curr->opcode = exit_op;
+                curr->opcode = _DEOPT;
                 curr->format = UOP_FORMAT_TARGET;
                 curr->target = old_target;
             }