From a73aa48e6bec900be7edd3431deaa5fc1d809e6f Mon Sep 17 00:00:00 2001
From: Pablo Galindo Salgado <Pablogsal@gmail.com>
Date: Fri, 1 Dec 2023 13:20:51 +0000
Subject: [PATCH] gh-112367: Only free perf trampoline arenas at shutdown
 (#112368)

Signed-off-by: Pablo Galindo <pablogsal@gmail.com>
---
 Include/internal/pycore_ceval.h               |  1 +
 ...-11-24-14-10-57.gh-issue-112367.9z1IDp.rst |  2 +
 Python/perf_trampoline.c                      | 40 ++++++++++++++++---
 Python/pylifecycle.c                          |  2 +-
 4 files changed, 39 insertions(+), 6 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2023-11-24-14-10-57.gh-issue-112367.9z1IDp.rst

diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h
index c372b7224fb..3f7ac922bdf 100644
--- a/Include/internal/pycore_ceval.h
+++ b/Include/internal/pycore_ceval.h
@@ -101,6 +101,7 @@ extern int _PyPerfTrampoline_SetCallbacks(_PyPerf_Callbacks *);
 extern void _PyPerfTrampoline_GetCallbacks(_PyPerf_Callbacks *);
 extern int _PyPerfTrampoline_Init(int activate);
 extern int _PyPerfTrampoline_Fini(void);
+extern void _PyPerfTrampoline_FreeArenas(void);
 extern int _PyIsPerfTrampolineActive(void);
 extern PyStatus _PyPerfTrampoline_AfterFork_Child(void);
 #ifdef PY_HAVE_PERF_TRAMPOLINE
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-11-24-14-10-57.gh-issue-112367.9z1IDp.rst b/Misc/NEWS.d/next/Core and Builtins/2023-11-24-14-10-57.gh-issue-112367.9z1IDp.rst
new file mode 100644
index 00000000000..991e45ad47f
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2023-11-24-14-10-57.gh-issue-112367.9z1IDp.rst	
@@ -0,0 +1,2 @@
+Avoid undefined behaviour when using the perf trampolines by not freeing the
+code arenas until shutdown. Patch by Pablo Galindo
diff --git a/Python/perf_trampoline.c b/Python/perf_trampoline.c
index 208ced6c101..540b650192e 100644
--- a/Python/perf_trampoline.c
+++ b/Python/perf_trampoline.c
@@ -216,10 +216,24 @@ perf_map_write_entry(void *state, const void *code_addr,
     PyMem_RawFree(perf_map_entry);
 }
 
+static void*
+perf_map_init_state(void)
+{
+    PyUnstable_PerfMapState_Init();
+    return NULL;
+}
+
+static int
+perf_map_free_state(void *state)
+{
+    PyUnstable_PerfMapState_Fini();
+    return 0;
+}
+
 _PyPerf_Callbacks _Py_perfmap_callbacks = {
-    NULL,
+    &perf_map_init_state,
     &perf_map_write_entry,
-    NULL,
+    &perf_map_free_state,
 };
 
 static int
@@ -415,7 +429,6 @@ _PyPerfTrampoline_SetCallbacks(_PyPerf_Callbacks *callbacks)
     trampoline_api.write_state = callbacks->write_state;
     trampoline_api.free_state = callbacks->free_state;
     trampoline_api.state = NULL;
-    perf_status = PERF_STATUS_OK;
 #endif
     return 0;
 }
@@ -434,6 +447,7 @@ _PyPerfTrampoline_Init(int activate)
     }
     if (!activate) {
         tstate->interp->eval_frame = NULL;
+        perf_status = PERF_STATUS_NO_INIT;
     }
     else {
         tstate->interp->eval_frame = py_trampoline_evaluator;
@@ -444,6 +458,9 @@ _PyPerfTrampoline_Init(int activate)
         if (extra_code_index == -1) {
             return -1;
         }
+        if (trampoline_api.state == NULL && trampoline_api.init_state != NULL) {
+            trampoline_api.state = trampoline_api.init_state();
+        }
         perf_status = PERF_STATUS_OK;
     }
 #endif
@@ -454,16 +471,29 @@ int
 _PyPerfTrampoline_Fini(void)
 {
 #ifdef PY_HAVE_PERF_TRAMPOLINE
+    if (perf_status != PERF_STATUS_OK) {
+        return 0;
+    }
     PyThreadState *tstate = _PyThreadState_GET();
     if (tstate->interp->eval_frame == py_trampoline_evaluator) {
         tstate->interp->eval_frame = NULL;
     }
-    free_code_arenas();
+    if (perf_status == PERF_STATUS_OK) {
+        trampoline_api.free_state(trampoline_api.state);
+    }
     extra_code_index = -1;
+    perf_status = PERF_STATUS_NO_INIT;
 #endif
     return 0;
 }
 
+void _PyPerfTrampoline_FreeArenas(void) {
+#ifdef PY_HAVE_PERF_TRAMPOLINE
+    free_code_arenas();
+#endif
+    return;
+}
+
 int
 PyUnstable_PerfTrampoline_SetPersistAfterFork(int enable){
 #ifdef PY_HAVE_PERF_TRAMPOLINE
@@ -477,8 +507,8 @@ PyStatus
 _PyPerfTrampoline_AfterFork_Child(void)
 {
 #ifdef PY_HAVE_PERF_TRAMPOLINE
-    PyUnstable_PerfMapState_Fini();
     if (persist_after_fork) {
+        _PyPerfTrampoline_Fini();
         char filename[256];
         pid_t parent_pid = getppid();
         snprintf(filename, sizeof(filename), "/tmp/perf-%d.map", parent_pid);
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index ac8d5208322..aff67d7a835 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -1797,6 +1797,7 @@ finalize_interp_clear(PyThreadState *tstate)
         _PyArg_Fini();
         _Py_ClearFileSystemEncoding();
         _PyPerfTrampoline_Fini();
+        _PyPerfTrampoline_FreeArenas();
     }
 
     finalize_interp_types(tstate->interp);
@@ -1854,7 +1855,6 @@ Py_FinalizeEx(void)
      */
 
     _PyAtExit_Call(tstate->interp);
-    PyUnstable_PerfMapState_Fini();
 
     /* Copy the core config, PyInterpreterState_Delete() free
        the core config memory */