Add C functions _PyTraceMalloc_Track()

Issue #26530:

* Add C functions _PyTraceMalloc_Track() and _PyTraceMalloc_Untrack() to track
  memory blocks using the tracemalloc module.
* Add _PyTraceMalloc_GetTraceback() to get the traceback of an object.

2 files changed, 146 insertions, 11 deletions

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 0fc7cbc..8c79485 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -3675,6 +3675,78 @@ pyobject_malloc_without_gil(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+tracemalloc_track(PyObject *self, PyObject *args)
+{
+    unsigned int domain;
+    PyObject *ptr_obj;
+    void *ptr;
+    Py_ssize_t size;
+    int release_gil = 0;
+    int res;
+
+    if (!PyArg_ParseTuple(args, "IOn|i", &domain, &ptr_obj, &size, &release_gil))
+        return NULL;
+    ptr = PyLong_AsVoidPtr(ptr_obj);
+    if (PyErr_Occurred())
+        return NULL;
+
+    if (release_gil) {
+        Py_BEGIN_ALLOW_THREADS
+        res = _PyTraceMalloc_Track(domain, (Py_uintptr_t)ptr, size);
+        Py_END_ALLOW_THREADS
+    }
+    else {
+        res = _PyTraceMalloc_Track(domain, (Py_uintptr_t)ptr, size);
+    }
+
+    if (res < 0) {
+        PyErr_SetString(PyExc_RuntimeError, "_PyTraceMalloc_Track error");
+        return NULL;
+    }
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+tracemalloc_untrack(PyObject *self, PyObject *args)
+{
+    unsigned int domain;
+    PyObject *ptr_obj;
+    void *ptr;
+    int res;
+
+    if (!PyArg_ParseTuple(args, "IO", &domain, &ptr_obj))
+        return NULL;
+    ptr = PyLong_AsVoidPtr(ptr_obj);
+    if (PyErr_Occurred())
+        return NULL;
+
+    res = _PyTraceMalloc_Untrack(domain, (Py_uintptr_t)ptr);
+    if (res < 0) {
+        PyErr_SetString(PyExc_RuntimeError, "_PyTraceMalloc_Track error");
+        return NULL;
+    }
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+tracemalloc_get_traceback(PyObject *self, PyObject *args)
+{
+    unsigned int domain;
+    PyObject *ptr_obj;
+    void *ptr;
+
+    if (!PyArg_ParseTuple(args, "IO", &domain, &ptr_obj))
+        return NULL;
+    ptr = PyLong_AsVoidPtr(ptr_obj);
+    if (PyErr_Occurred())
+        return NULL;
+
+    return _PyTraceMalloc_GetTraceback(domain, (Py_uintptr_t)ptr);
+}
+
 
 static PyMethodDef TestMethods[] = {
     {"raise_exception",         raise_exception,                 METH_VARARGS},
@@ -3861,6 +3933,9 @@ static PyMethodDef TestMethods[] = {
     {"pymem_api_misuse", pymem_api_misuse, METH_NOARGS},
     {"pymem_malloc_without_gil", pymem_malloc_without_gil, METH_NOARGS},
     {"pyobject_malloc_without_gil", pyobject_malloc_without_gil, METH_NOARGS},
+    {"tracemalloc_track", tracemalloc_track, METH_VARARGS},
+    {"tracemalloc_untrack", tracemalloc_untrack, METH_VARARGS},
+    {"tracemalloc_get_traceback", tracemalloc_get_traceback, METH_VARARGS},
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Modules/_tracemalloc.c b/Modules/_tracemalloc.c
index 784157f..5ff1f84 100644
--- a/Modules/_tracemalloc.c
+++ b/Modules/_tracemalloc.c
@@ -61,8 +61,6 @@ static PyThread_type_lock tables_lock;
 
 #define DEFAULT_DOMAIN 0
 
-typedef unsigned int domain_t;
-
 /* Pack the frame_t structure to reduce the memory footprint. */
 typedef struct
 #ifdef __GNUC__
@@ -70,7 +68,7 @@ __attribute__((packed))
 #endif
 {
     Py_uintptr_t ptr;
-    domain_t domain;
+    _PyTraceMalloc_domain_t domain;
 } pointer_t;
 
 /* Pack the frame_t structure to reduce the memory footprint on 64-bit
@@ -519,11 +517,13 @@ traceback_new(void)
 
 
 static void
-tracemalloc_remove_trace(domain_t domain, Py_uintptr_t ptr)
+tracemalloc_remove_trace(_PyTraceMalloc_domain_t domain, Py_uintptr_t ptr)
 {
     trace_t trace;
     int removed;
 
+    assert(tracemalloc_config.tracing);
+
     if (tracemalloc_config.use_domain) {
         pointer_t key = {ptr, domain};
         removed = _Py_HASHTABLE_POP(tracemalloc_traces, key, trace);
@@ -544,12 +544,15 @@ tracemalloc_remove_trace(domain_t domain, Py_uintptr_t ptr)
 
 
 static int
-tracemalloc_add_trace(domain_t domain, Py_uintptr_t ptr, size_t size)
+tracemalloc_add_trace(_PyTraceMalloc_domain_t domain, Py_uintptr_t ptr,
+                      size_t size)
 {
     traceback_t *traceback;
     trace_t trace;
     int res;
 
+    assert(tracemalloc_config.tracing);
+
     /* first, remove the previous trace (if any) */
     tracemalloc_remove_trace(domain, ptr);
 
@@ -1183,7 +1186,7 @@ traceback_to_pyobject(traceback_t *traceback, _Py_hashtable_t *intern_table)
 
 
 static PyObject*
-trace_to_pyobject(domain_t domain, trace_t *trace,
+trace_to_pyobject(_PyTraceMalloc_domain_t domain, trace_t *trace,
                   _Py_hashtable_t *intern_tracebacks)
 {
     PyObject *trace_obj = NULL;
@@ -1229,7 +1232,7 @@ tracemalloc_get_traces_fill(_Py_hashtable_t *traces, _Py_hashtable_entry_t *entr
                             void *user_data)
 {
     get_traces_t *get_traces = user_data;
-    domain_t domain;
+    _PyTraceMalloc_domain_t domain;
     trace_t *trace;
     PyObject *tracemalloc_obj;
     int res;
@@ -1340,7 +1343,7 @@ finally:
 
 
 static traceback_t*
-tracemalloc_get_traceback(domain_t domain, const void *ptr)
+tracemalloc_get_traceback(_PyTraceMalloc_domain_t domain, Py_uintptr_t ptr)
 {
     trace_t trace;
     int found;
@@ -1350,7 +1353,7 @@ tracemalloc_get_traceback(domain_t domain, const void *ptr)
 
     TABLES_LOCK();
     if (tracemalloc_config.use_domain) {
-        pointer_t key = {(Py_uintptr_t)ptr, domain};
+        pointer_t key = {ptr, domain};
         found = _Py_HASHTABLE_GET(tracemalloc_traces, key, trace);
     }
     else {
@@ -1387,7 +1390,7 @@ py_tracemalloc_get_object_traceback(PyObject *self, PyObject *obj)
     else
         ptr = (void *)obj;
 
-    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, ptr);
+    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, (Py_uintptr_t)ptr);
     if (traceback == NULL)
         Py_RETURN_NONE;
 
@@ -1415,7 +1418,7 @@ _PyMem_DumpTraceback(int fd, const void *ptr)
     traceback_t *traceback;
     int i;
 
-    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, ptr);
+    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, (Py_uintptr_t)ptr);
     if (traceback == NULL)
         return;
 
@@ -1686,3 +1689,60 @@ _PyTraceMalloc_Fini(void)
 #endif
     tracemalloc_deinit();
 }
+
+int
+_PyTraceMalloc_Track(_PyTraceMalloc_domain_t domain, Py_uintptr_t ptr,
+                     size_t size)
+{
+    int res;
+#ifdef WITH_THREAD
+    PyGILState_STATE gil_state;
+#endif
+
+    if (!tracemalloc_config.tracing) {
+        /* tracemalloc is not tracing: do nothing */
+        return -2;
+    }
+
+#ifdef WITH_THREAD
+    gil_state = PyGILState_Ensure();
+#endif
+
+    TABLES_LOCK();
+    res = tracemalloc_add_trace(domain, ptr, size);
+    TABLES_UNLOCK();
+
+#ifdef WITH_THREAD
+    PyGILState_Release(gil_state);
+#endif
+    return res;
+}
+
+
+int
+_PyTraceMalloc_Untrack(_PyTraceMalloc_domain_t domain, Py_uintptr_t ptr)
+{
+    if (!tracemalloc_config.tracing) {
+        /* tracemalloc is not tracing: do nothing */
+        return -2;
+    }
+
+    TABLES_LOCK();
+    tracemalloc_remove_trace(domain, ptr);
+    TABLES_UNLOCK();
+
+    return 0;
+}
+
+
+PyObject*
+_PyTraceMalloc_GetTraceback(_PyTraceMalloc_domain_t domain, Py_uintptr_t ptr)
+{
+    traceback_t *traceback;
+
+    traceback = tracemalloc_get_traceback(domain, ptr);
+    if (traceback == NULL)
+        Py_RETURN_NONE;
+
+    return traceback_to_pyobject(traceback, NULL);
+}