Fix bug where code/func get freed halfway

diff --git a/Include/internal/pycore_interp_structs.h b/Include/internal/pycore_interp_structs.h
index 9671325362d5c35171ad6efde2a3da101c8b27a3..6296f09ab995e193e025c08f6f1cf36e210203df 100644 (file)
--- a/Include/internal/pycore_interp_structs.h
+++ b/Include/internal/pycore_interp_structs.h
@@ -951,8 +951,8 @@ struct _is {
     _Py_CODEUNIT *jit_tracer_initial_instr;
     int jit_tracer_initial_stack_depth;
     int jit_tracer_initial_chain_depth;
-    PyCodeObject *jit_tracer_initial_code; // Borrowed
-    PyFunctionObject *jit_tracer_initial_func; // Borrowed
+    PyCodeObject *jit_tracer_initial_code; // Strong
+    PyFunctionObject *jit_tracer_initial_func; // Strong
     struct _PyExitData *jit_tracer_previous_exit;
     bool jit_completed_loop;
     bool jit;

diff --git a/Include/internal/pycore_optimizer.h b/Include/internal/pycore_optimizer.h
index 79b1931d070467021206f239335c67ab7afc042a..c7c304943bf87e4fe9261f824ebcf5897d78fc2d 100644 (file)
--- a/Include/internal/pycore_optimizer.h
+++ b/Include/internal/pycore_optimizer.h
@@ -372,6 +372,8 @@ _PyJIT_translate_single_bytecode_to_trace(
 
 void
 _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *next_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit);
+
+void _PyJIT_FinalizeTracing(PyThreadState *tstate);
 #ifdef __cplusplus
 }
 #endif

diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
index b15b6788205915a398e73fe71ac9bf25fb59560f..313b062bcc22d8a9fce600969ec4ab684bd83715 100644 (file)
--- a/Python/ceval_macros.h
+++ b/Python/ceval_macros.h
@@ -141,7 +141,10 @@
 #  define LEAVE_TRACING() DISPATCH_TABLE_VAR = DISPATCH_TABLE;
 #  define BAIL_TRACING_NO_DISPATCH() \
     LEAVE_TRACING(); \
+    _PyFrame_SetStackPointer(frame, stack_pointer); \
     int _err = _PyOptimizer_Optimize(frame, tstate); \
+    _PyJIT_FinalizeTracing(tstate); \
+    stack_pointer = _PyFrame_GetStackPointer(frame); \
     if (_err < 0) { \
         JUMP_TO_LABEL(error); \
     }
@@ -149,13 +152,11 @@
     BAIL_TRACING_NO_DISPATCH() \
     DISPATCH();
 #  define RECORD_TRACE() do { \
-        frame->instr_ptr = next_instr; \
         if (add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
             BAIL_TRACING(); \
         } \
     } while (0);
 #  define RECORD_TRACE_NO_DISPATCH() do { \
-        frame->instr_ptr = next_instr; \
         if (add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
             BAIL_TRACING_NO_DISPATCH(); \
         } \
@@ -413,6 +414,7 @@ _PyFrame_SetStackPointer(frame, stack_pointer)
 
 #define TIER1_TO_TIER2(EXECUTOR)                        \
 do {                                                   \
+    LEAVE_TRACING(); \
     OPT_STAT_INC(traces_executed);                     \
     next_instr = _Py_jit_entry((EXECUTOR), frame, stack_pointer, tstate); \
     frame = tstate->current_frame;                     \
@@ -428,9 +430,6 @@ do {                                                   \
         assert(tstate->interp->jit_tracer_code_curr_size == 2); \
         ENTER_TRACING(); \
     } \
-    else { \
-        LEAVE_TRACING(); \
-    } \
     DISPATCH();                                        \
 } while (0)
 

diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index e190f651fa4c845b974bb1089e1ca579e7b56a2b..9e855fb5afd5737dfb7e1b29df6c385cda83f173 100644 (file)
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -27372,6 +27372,7 @@ JUMP_TO_LABEL(error);
 
         LABEL(error)
         {
+            assert(!IS_JIT_TRACING());
             #ifdef NDEBUG
             if (!_PyErr_Occurred(tstate)) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
@@ -27401,6 +27402,7 @@ JUMP_TO_LABEL(error);
 
         LABEL(exception_unwind)
         {
+            assert(!IS_JIT_TRACING());
             int offset = INSTR_OFFSET()-1;
             int level, handler, lasti;
             int handled = get_exception_handler(_PyFrame_GetCode(frame), offset, &level, &handler, &lasti);

diff --git a/Python/optimizer.c b/Python/optimizer.c
index 382a3150793ef35e8e46bfbcfe136b822d474369..65bc861a1deed976e8c7b7574b070e8ad505597c 100644 (file)
--- a/Python/optimizer.c
+++ b/Python/optimizer.c
@@ -129,13 +129,18 @@ _PyOptimizer_Optimize(
     // side-exit chains. We can only insert the executor into the bytecode if
     // this is true, since a deopt won't infinitely re-enter the executor:
     bool progress_needed = (chain_depth % MAX_CHAIN_DEPTH) == 0;
-    PyCodeObject *code = (PyCodeObject *)tstate->interp->jit_tracer_initial_func->func_code;
-    assert(PyCode_Check(code));
+    PyCodeObject *code = (PyCodeObject *)tstate->interp->jit_tracer_initial_code;
     _Py_CODEUNIT *start = tstate->interp->jit_tracer_initial_instr;
     if (progress_needed && !has_space_for_executor(code, start)) {
         interp->compiling = false;
         return 0;
     }
+    // We are the only one still holding a reference to this code object that
+    // is practically dead.
+    if (_PyObject_IsUniquelyReferenced(code) || _PyObject_IsUniquelyReferenced(tstate->interp->jit_tracer_initial_func)) {
+        interp->compiling = false;
+        return 0;
+    }
     _PyExecutorObject *executor;
     int err = uop_optimize(frame, tstate, &executor, progress_needed);
     if (err <= 0) {
@@ -840,8 +845,8 @@ _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_
     tstate->interp->jit_tracer_code_curr_size = 2;
     tstate->interp->jit_tracer_code_max_size = UOP_MAX_TRACE_LENGTH;
     tstate->interp->jit_tracer_initial_instr = next_instr;
-    tstate->interp->jit_tracer_initial_code = code;
-    tstate->interp->jit_tracer_initial_func = _PyFrame_GetFunction(frame);
+    tstate->interp->jit_tracer_initial_code = (PyCodeObject *)Py_NewRef(code);
+    tstate->interp->jit_tracer_initial_func = (PyFunctionObject *)Py_NewRef(_PyFrame_GetFunction(frame));
     tstate->interp->jit_tracer_previous_exit = exit;
     memset(&tstate->interp->jit_tracer_dependencies.bits, 0, sizeof(tstate->interp->jit_tracer_dependencies.bits));
     tstate->interp->jit_completed_loop = false;
@@ -849,6 +854,15 @@ _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_
     tstate->interp->jit_tracer_initial_chain_depth = chain_depth;
 }
 
+void
+_PyJIT_FinalizeTracing(PyThreadState *tstate)
+{
+    Py_CLEAR(tstate->interp->jit_tracer_initial_code);
+    Py_CLEAR(tstate->interp->jit_tracer_initial_func);
+    tstate->interp->jit_tracer_code_curr_size = 2;
+}
+
+
 #undef RESERVE
 #undef RESERVE_RAW
 #undef INSTR_IP