@@ -204,6 +204,43 @@ enum perf_trampoline_type {
204204#define perf_map_file _PyRuntime.ceval.perf.map_file
205205#define persist_after_fork _PyRuntime.ceval.perf.persist_after_fork
206206#define perf_trampoline_type _PyRuntime.ceval.perf.perf_trampoline_type
207+ #define prev_eval_frame _PyRuntime.ceval.perf.prev_eval_frame
208+ #define trampoline_refcount _PyRuntime.ceval.perf.trampoline_refcount
209+ #define code_watcher_id _PyRuntime.ceval.perf.code_watcher_id
210+
211+ static void free_code_arenas (void );
212+
213+ static void
214+ perf_trampoline_reset_state (void )
215+ {
216+ free_code_arenas ();
217+ if (code_watcher_id >= 0 ) {
218+ PyCode_ClearWatcher (code_watcher_id );
219+ code_watcher_id = -1 ;
220+ }
221+ extra_code_index = -1 ;
222+ }
223+
224+ static int
225+ perf_trampoline_code_watcher (PyCodeEvent event , PyCodeObject * co )
226+ {
227+ if (event != PY_CODE_EVENT_DESTROY ) {
228+ return 0 ;
229+ }
230+ if (extra_code_index == -1 ) {
231+ return 0 ;
232+ }
233+ py_trampoline f = NULL ;
234+ int ret = _PyCode_GetExtra ((PyObject * )co , extra_code_index , (void * * )& f );
235+ if (ret != 0 || f == NULL ) {
236+ return 0 ;
237+ }
238+ trampoline_refcount -- ;
239+ if (trampoline_refcount == 0 ) {
240+ perf_trampoline_reset_state ();
241+ }
242+ return 0 ;
243+ }
207244
208245static void
209246perf_map_write_entry (void * state , const void * code_addr ,
@@ -405,6 +442,7 @@ py_trampoline_evaluator(PyThreadState *ts, _PyInterpreterFrame *frame,
405442 perf_code_arena -> code_size , co );
406443 _PyCode_SetExtra ((PyObject * )co , extra_code_index ,
407444 (void * )new_trampoline );
445+ trampoline_refcount ++ ;
408446 f = new_trampoline ;
409447 }
410448 assert (f != NULL );
@@ -428,6 +466,7 @@ int PyUnstable_PerfTrampoline_CompileCode(PyCodeObject *co)
428466 }
429467 trampoline_api .write_state (trampoline_api .state , new_trampoline ,
430468 perf_code_arena -> code_size , co );
469+ trampoline_refcount ++ ;
431470 return _PyCode_SetExtra ((PyObject * )co , extra_code_index ,
432471 (void * )new_trampoline );
433472 }
@@ -482,6 +521,10 @@ _PyPerfTrampoline_Init(int activate)
482521{
483522#ifdef PY_HAVE_PERF_TRAMPOLINE
484523 PyThreadState * tstate = _PyThreadState_GET ();
524+ if (code_watcher_id == 0 ) {
525+ // Initialize to -1 since 0 is a valid watcher ID
526+ code_watcher_id = -1 ;
527+ }
485528 if (tstate -> interp -> eval_frame &&
486529 tstate -> interp -> eval_frame != py_trampoline_evaluator ) {
487530 PyErr_SetString (PyExc_RuntimeError ,
@@ -505,6 +548,13 @@ _PyPerfTrampoline_Init(int activate)
505548 if (new_code_arena () < 0 ) {
506549 return -1 ;
507550 }
551+ code_watcher_id = PyCode_AddWatcher (perf_trampoline_code_watcher );
552+ if (code_watcher_id < 0 ) {
553+ PyErr_FormatUnraisable ("Failed to register code watcher for perf trampoline" );
554+ free_code_arenas ();
555+ return -1 ;
556+ }
557+ trampoline_refcount = 1 ; // Base refcount held by the system
508558 perf_status = PERF_STATUS_OK ;
509559 }
510560#endif
@@ -526,17 +576,19 @@ _PyPerfTrampoline_Fini(void)
526576 trampoline_api .free_state (trampoline_api .state );
527577 perf_trampoline_type = PERF_TRAMPOLINE_UNSET ;
528578 }
529- extra_code_index = -1 ;
579+
580+ // Prevent new trampolines from being created
530581 perf_status = PERF_STATUS_NO_INIT ;
531- #endif
532- return 0 ;
533- }
534582
535- void _PyPerfTrampoline_FreeArenas (void ) {
536- #ifdef PY_HAVE_PERF_TRAMPOLINE
537- free_code_arenas ();
583+ // Decrement base refcount. If refcount reaches 0, all code objects are already
584+ // dead so clean up now. Otherwise, watcher remains active to clean up when last
585+ // code object dies; extra_code_index stays valid so watcher can identify them.
586+ trampoline_refcount -- ;
587+ if (trampoline_refcount == 0 ) {
588+ perf_trampoline_reset_state ();
589+ }
538590#endif
539- return ;
591+ return 0 ;
540592}
541593
542594int
0 commit comments