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>

1 files changed, 27 insertions, 0 deletions

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
index 3ee3323..b91f7b6 100644
--- a/Modules/_testinternalcapi.c
+++ b/Modules/_testinternalcapi.c
@@ -760,6 +760,31 @@ clear_extension(PyObject *self, PyObject *args)
 }
 
 static PyObject *
+write_perf_map_entry(PyObject *self, PyObject *args)
+{
+    const void *code_addr;
+    unsigned int code_size;
+    const char *entry_name;
+
+    if (!PyArg_ParseTuple(args, "KIs", &code_addr, &code_size, &entry_name))
+        return NULL;
+
+    int ret = PyUnstable_WritePerfMapEntry(code_addr, code_size, entry_name);
+    if (ret == -1) {
+        PyErr_SetString(PyExc_OSError, "Failed to write performance map entry");
+        return NULL;
+    }
+    return Py_BuildValue("i", ret);
+}
+
+static PyObject *
+perf_map_state_teardown(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored))
+{
+    PyUnstable_PerfMapState_Fini();
+    Py_RETURN_NONE;
+}
+
+static PyObject *
 iframe_getcode(PyObject *self, PyObject *frame)
 {
     if (!PyFrame_Check(frame)) {
@@ -815,6 +840,8 @@ static PyMethodDef module_functions[] = {
     _TESTINTERNALCAPI_ASSEMBLE_CODE_OBJECT_METHODDEF
     {"get_interp_settings", get_interp_settings, METH_VARARGS, NULL},
     {"clear_extension", clear_extension, METH_VARARGS, NULL},
+    {"write_perf_map_entry", write_perf_map_entry, METH_VARARGS},
+    {"perf_map_state_teardown", perf_map_state_teardown, METH_NOARGS},
     {"iframe_getcode", iframe_getcode, METH_O, NULL},
     {"iframe_getline", iframe_getline, METH_O, NULL},
     {"iframe_getlasti", iframe_getlasti, METH_O, NULL},