gh-103295: expose API for writing perf map files (#103546)

Co-authored-by: Aniket Panse <aniketpanse@fb.com>
Co-authored-by: Gregory P. Smith <greg@krypto.org>
Co-authored-by: Carl Meyer <carl@oddbird.net>

3 files changed, 96 insertions, 69 deletions

diff --git a/Python/perf_trampoline.c b/Python/perf_trampoline.c
index 3b18328..6a6f223 100644
--- a/Python/perf_trampoline.c
+++ b/Python/perf_trampoline.c
@@ -193,75 +193,33 @@ typedef struct trampoline_api_st trampoline_api_t;
 #define trampoline_api _PyRuntime.ceval.perf.trampoline_api
 #define perf_map_file _PyRuntime.ceval.perf.map_file
 
-static void *
-perf_map_get_file(void)
-{
-    if (perf_map_file) {
-        return perf_map_file;
-    }
-    char filename[100];
-    pid_t pid = getpid();
-    // Location and file name of perf map is hard-coded in perf tool.
-    // Use exclusive create flag wit nofollow to prevent symlink attacks.
-    int flags = O_WRONLY | O_CREAT | O_EXCL | O_NOFOLLOW | O_CLOEXEC;
-    snprintf(filename, sizeof(filename) - 1, "/tmp/perf-%jd.map",
-             (intmax_t)pid);
-    int fd = open(filename, flags, 0600);
-    if (fd == -1) {
-        perf_status = PERF_STATUS_FAILED;
-        PyErr_SetFromErrnoWithFilename(PyExc_OSError, filename);
-        return NULL;
-    }
-    perf_map_file = fdopen(fd, "w");
-    if (!perf_map_file) {
-        perf_status = PERF_STATUS_FAILED;
-        PyErr_SetFromErrnoWithFilename(PyExc_OSError, filename);
-        close(fd);
-        return NULL;
-    }
-    return perf_map_file;
-}
-
-static int
-perf_map_close(void *state)
-{
-    FILE *fp = (FILE *)state;
-    int ret = 0;
-    if (fp) {
-        ret = fclose(fp);
-    }
-    perf_map_file = NULL;
-    perf_status = PERF_STATUS_NO_INIT;
-    return ret;
-}
 
 static void
 perf_map_write_entry(void *state, const void *code_addr,
                          unsigned int code_size, PyCodeObject *co)
 {
-    assert(state != NULL);
-    FILE *method_file = (FILE *)state;
-    const char *entry = PyUnicode_AsUTF8(co->co_qualname);
-    if (entry == NULL) {
-        _PyErr_WriteUnraisableMsg("Failed to get qualname from code object",
-                                  NULL);
-        return;
+    const char *entry = "";
+    if (co->co_qualname != NULL) {
+        entry = PyUnicode_AsUTF8(co->co_qualname);
     }
-    const char *filename = PyUnicode_AsUTF8(co->co_filename);
-    if (filename == NULL) {
-        _PyErr_WriteUnraisableMsg("Failed to get filename from code object",
-                                  NULL);
+    const char *filename = "";
+    if (co->co_filename != NULL) {
+        filename = PyUnicode_AsUTF8(co->co_filename);
+    }
+    size_t perf_map_entry_size = snprintf(NULL, 0, "py::%s:%s", entry, filename) + 1;
+    char* perf_map_entry = (char*) PyMem_RawMalloc(perf_map_entry_size);
+    if (perf_map_entry == NULL) {
         return;
     }
-    fprintf(method_file, "%" PRIxPTR " %x py::%s:%s\n", (uintptr_t) code_addr, code_size, entry,
-            filename);
-    fflush(method_file);
+    snprintf(perf_map_entry, perf_map_entry_size, "py::%s:%s", entry, filename);
+    PyUnstable_WritePerfMapEntry(code_addr, code_size, perf_map_entry);
+    PyMem_RawFree(perf_map_entry);
 }
 
 _PyPerf_Callbacks _Py_perfmap_callbacks = {
-    &perf_map_get_file,
+    NULL,
     &perf_map_write_entry,
-    &perf_map_close
+    NULL,
 };
 
 static int
@@ -465,13 +423,6 @@ _PyPerfTrampoline_Init(int activate)
         if (new_code_arena() < 0) {
             return -1;
         }
-        if (trampoline_api.state == NULL) {
-            void *state = trampoline_api.init_state();
-            if (state == NULL) {
-                return -1;
-            }
-            trampoline_api.state = state;
-        }
         extra_code_index = _PyEval_RequestCodeExtraIndex(NULL);
         if (extra_code_index == -1) {
             return -1;
@@ -491,10 +442,6 @@ _PyPerfTrampoline_Fini(void)
         tstate->interp->eval_frame = NULL;
     }
     free_code_arenas();
-    if (trampoline_api.state != NULL) {
-        trampoline_api.free_state(trampoline_api.state);
-        trampoline_api.state = NULL;
-    }
     extra_code_index = -1;
 #endif
     return 0;
@@ -507,6 +454,7 @@ _PyPerfTrampoline_AfterFork_Child(void)
     // Restart trampoline in file in child.
     int was_active = _PyIsPerfTrampolineActive();
     _PyPerfTrampoline_Fini();
+    PyUnstable_PerfMapState_Fini();
     if (was_active) {
         _PyPerfTrampoline_Init(1);
     }
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index 93ad1fe..06c4345 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -1775,6 +1775,7 @@ Py_FinalizeEx(void)
      */
 
     _PyAtExit_Call(tstate->interp);
+    PyUnstable_PerfMapState_Fini();
 
     /* Copy the core config, PyInterpreterState_Delete() free
        the core config memory */
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
index 7510be5..33147f0 100644
--- a/Python/sysmodule.c
+++ b/Python/sysmodule.c
@@ -52,6 +52,10 @@ extern const char *PyWin_DLLVersionString;
 #include <emscripten.h>
 #endif
 
+#ifdef HAVE_FCNTL_H
+#include <fcntl.h>
+#endif
+
 /*[clinic input]
 module sys
 [clinic start generated code]*/
@@ -2144,7 +2148,7 @@ sys_activate_stack_trampoline_impl(PyObject *module, const char *backend)
     if (strcmp(backend, "perf") == 0) {
         _PyPerf_Callbacks cur_cb;
         _PyPerfTrampoline_GetCallbacks(&cur_cb);
-        if (cur_cb.init_state != _Py_perfmap_callbacks.init_state) {
+        if (cur_cb.write_state != _Py_perfmap_callbacks.write_state) {
             if (_PyPerfTrampoline_SetCallbacks(&_Py_perfmap_callbacks) < 0 ) {
                 PyErr_SetString(PyExc_ValueError, "can't activate perf trampoline");
                 return NULL;
@@ -2240,6 +2244,80 @@ sys__getframemodulename_impl(PyObject *module, int depth)
 }
 
 
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+static PerfMapState perf_map_state;
+
+PyAPI_FUNC(int) PyUnstable_PerfMapState_Init(void) {
+#ifndef MS_WINDOWS
+    char filename[100];
+    pid_t pid = getpid();
+    // Use nofollow flag to prevent symlink attacks.
+    int flags = O_WRONLY | O_CREAT | O_APPEND | O_NOFOLLOW | O_CLOEXEC;
+    snprintf(filename, sizeof(filename) - 1, "/tmp/perf-%jd.map",
+                (intmax_t)pid);
+    int fd = open(filename, flags, 0600);
+    if (fd == -1) {
+        return -1;
+    }
+    else{
+        perf_map_state.perf_map = fdopen(fd, "a");
+        if (perf_map_state.perf_map == NULL) {
+            close(fd);
+            return -1;
+        }
+    }
+    perf_map_state.map_lock = PyThread_allocate_lock();
+    if (perf_map_state.map_lock == NULL) {
+        fclose(perf_map_state.perf_map);
+        return -2;
+    }
+#endif
+    return 0;
+}
+
+PyAPI_FUNC(int) PyUnstable_WritePerfMapEntry(
+    const void *code_addr,
+    unsigned int code_size,
+    const char *entry_name
+) {
+#ifndef MS_WINDOWS
+    if (perf_map_state.perf_map == NULL) {
+        int ret = PyUnstable_PerfMapState_Init();
+        if(ret != 0){
+            return ret;
+        }
+    }
+    PyThread_acquire_lock(perf_map_state.map_lock, 1);
+    fprintf(perf_map_state.perf_map, "%" PRIxPTR " %x %s\n", (uintptr_t) code_addr, code_size, entry_name);
+    fflush(perf_map_state.perf_map);
+    PyThread_release_lock(perf_map_state.map_lock);
+#endif
+    return 0;
+}
+
+PyAPI_FUNC(void) PyUnstable_PerfMapState_Fini(void) {
+#ifndef MS_WINDOWS
+    if (perf_map_state.perf_map != NULL) {
+        // close the file
+        PyThread_acquire_lock(perf_map_state.map_lock, 1);
+        fclose(perf_map_state.perf_map);
+        PyThread_release_lock(perf_map_state.map_lock);
+
+        // clean up the lock and state
+        PyThread_free_lock(perf_map_state.map_lock);
+        perf_map_state.perf_map = NULL;
+    }
+#endif
+}
+
+#ifdef __cplusplus
+}
+#endif
+
+
 static PyMethodDef sys_methods[] = {
     /* Might as well keep this in alphabetic order */
     SYS_ADDAUDITHOOK_METHODDEF