Stronger check for recursive traces

Author: Fidget-Spinner

Include/internal/pycore_optimizer.h

@@ -354,7 +354,7 @@ _PyJit_TryInitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame,
     _Py_CODEUNIT *close_loop_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit,
     int oparg);
 
-void _PyJit_FinalizeTracing(PyThreadState *tstate);
+void _PyJit_ResetTracing(PyThreadState *tstate);
 
 void _PyJit_Tracer_InvalidateDependency(PyThreadState *old_tstate, void *obj);
 

Python/ceval.c

@@ -1486,7 +1486,7 @@ stop_tracing_and_jit(PyThreadState *tstate, _PyInterpreterFrame *frame)
             exit->temperature = initial_temperature_backoff_counter();
         }
     }
-    _PyJit_FinalizeTracing(tstate);
+    _PyJit_ResetTracing(tstate);
     return err;
 }
 #endif

Python/optimizer.c

@@ -138,12 +138,6 @@ _PyOptimizer_Optimize(
         // return immediately without optimization.
         return 0;
     }
-    if (_tstate->jit_tracer_state.initial_state.func == NULL) {
-        // gh-143123: It is possible for another function to finalize the current
-        // tracer's state while tracing. This might happen in a
-        // Python -> C -> Python call.
-        return 0;
-    }
     assert(!interp->compiling);
     assert(_tstate->jit_tracer_state.initial_state.stack_depth >= 0);
 #ifndef Py_GIL_DISABLED
@@ -1031,7 +1025,11 @@ _PyJit_TryInitializeTracing(
     _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
     // A recursive trace.
     // Don't trace into the inner call because it will stomp on the previous trace, causing endless retraces.
-    if (_tstate->jit_tracer_state.prev_state.code_curr_size > CODE_SIZE_EMPTY) {
+    if (_tstate->jit_tracer_state.prev_state.code_curr_size > CODE_SIZE_EMPTY ||
+        _tstate->jit_tracer_state.initial_state.func != NULL) {
+        // gh-143123: It is possible for another function to finalize the current
+        // tracer's state while tracing. This might happen in a
+        // Python -> C -> Python call.
         return 0;
     }
     if (oparg > 0xFFFF) {
@@ -1095,7 +1093,7 @@ _PyJit_TryInitializeTracing(
 }
 
 Py_NO_INLINE void
-_PyJit_FinalizeTracing(PyThreadState *tstate)
+_PyJit_ResetTracing(PyThreadState *tstate)
 {
     _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
     Py_CLEAR(_tstate->jit_tracer_state.initial_state.code);

Python/pystate.c

@@ -1526,6 +1526,7 @@ init_threadstate(_PyThreadStateImpl *_tstate,
 
 #ifdef _Py_TIER2
     _tstate->jit_tracer_state.code_buffer = NULL;
+    _PyJit_ResetTracing(tstate);
 #endif
     tstate->delete_later = NULL;
 
@@ -1846,6 +1847,7 @@ tstate_delete_common(PyThreadState *tstate, int release_gil)
         _PyObject_VirtualFree(_tstate->jit_tracer_state.code_buffer, UOP_BUFFER_SIZE);
         _tstate->jit_tracer_state.code_buffer = NULL;
     }
+    _PyJit_ResetTracing(tstate);
 #endif
 
     HEAD_UNLOCK(runtime);