From 0ffc2dd39fbf8a6c8e44c5355fa5289ba3da3a65 Mon Sep 17 00:00:00 2001 From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com> Date: Wed, 24 Sep 2025 20:34:23 +0100 Subject: [PATCH] Fix bug where code/func get freed halfway --- Include/internal/pycore_interp_structs.h | 4 ++-- Include/internal/pycore_optimizer.h | 2 ++ Python/ceval_macros.h | 9 ++++----- Python/generated_cases.c.h | 2 ++ Python/optimizer.c | 22 ++++++++++++++++++---- 5 files changed, 28 insertions(+), 11 deletions(-) diff --git a/Include/internal/pycore_interp_structs.h b/Include/internal/pycore_interp_structs.h index 9671325362d5..6296f09ab995 100644 --- 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 79b1931d0704..c7c304943bf8 100644 --- 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 b15b67882059..313b062bcc22 100644 --- 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 e190f651fa4c..9e855fb5afd5 100644 --- 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 382a3150793e..65bc861a1dee 100644 --- 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 -- 2.47.3