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.

5 files changed, 300 insertions, 11 deletions

diff --git a/Include/pymem.h b/Include/pymem.h
index b1f06ef..941e00f 100644
--- a/Include/pymem.h
+++ b/Include/pymem.h
@@ -25,6 +25,40 @@ PyAPI_FUNC(int) _PyMem_SetupAllocators(const char *opt);
 PyAPI_FUNC(int) _PyMem_PymallocEnabled(void);
 #endif
 
+/* Identifier of an address space (domain) in tracemalloc */
+typedef unsigned int _PyTraceMalloc_domain_t;
+
+/* Track an allocated memory block in the tracemalloc module.
+   Return 0 on success, return -1 on error (failed to allocate memory to store
+   the trace).
+
+   Return -2 if tracemalloc is disabled.
+
+   If memory block was already tracked, begin by removing the old trace. */
+PyAPI_FUNC(int) _PyTraceMalloc_Track(
+    _PyTraceMalloc_domain_t domain,
+    Py_uintptr_t ptr,
+    size_t size);
+
+/* Untrack an allocated memory block in the tracemalloc module.
+   Do nothing if the block was not tracked.
+
+   Return -2 if tracemalloc is disabled, otherwise return 0. */
+PyAPI_FUNC(int) _PyTraceMalloc_Untrack(
+    _PyTraceMalloc_domain_t domain,
+    Py_uintptr_t ptr);
+
+/* Get the traceback where a memory block was allocated.
+
+   Return a tuple of (filename: str, lineno: int) tuples.
+
+   Return None if the tracemalloc module is disabled or if the memory block
+   is not tracked by tracemalloc.
+
+   Raise an exception and return NULL on error. */
+PyAPI_FUNC(PyObject*) _PyTraceMalloc_GetTraceback(
+    _PyTraceMalloc_domain_t domain,
+    Py_uintptr_t ptr);
 #endif   /* !Py_LIMITED_API */
 
 
diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py
index 7b92b87..359d9c0 100644
--- a/Lib/test/test_tracemalloc.py
+++ b/Lib/test/test_tracemalloc.py
@@ -11,9 +11,15 @@ try:
     import threading
 except ImportError:
     threading = None
+try:
+    import _testcapi
+except ImportError:
+    _testcapi = None
+
 
 EMPTY_STRING_SIZE = sys.getsizeof(b'')
 
+
 def get_frames(nframe, lineno_delta):
     frames = []
     frame = sys._getframe(1)
@@ -866,12 +872,121 @@ class TestCommandLine(unittest.TestCase):
         assert_python_ok('-X', 'tracemalloc', '-c', code)
 
 
+@unittest.skipIf(_testcapi is None, 'need _testcapi')
+class TestCAPI(unittest.TestCase):
+    maxDiff = 80 * 20
+
+    def setUp(self):
+        if tracemalloc.is_tracing():
+            self.skipTest("tracemalloc must be stopped before the test")
+
+        self.domain = 5
+        self.size = 123
+        self.obj = allocate_bytes(self.size)[0]
+
+        # for the type "object", id(obj) is the address of its memory block.
+        # This type is not tracked by the garbage collector
+        self.ptr = id(self.obj)
+
+    def tearDown(self):
+        tracemalloc.stop()
+
+    def get_traceback(self):
+        frames = _testcapi.tracemalloc_get_traceback(self.domain, self.ptr)
+        if frames is not None:
+            return tracemalloc.Traceback(frames)
+        else:
+            return None
+
+    def track(self, release_gil=False, nframe=1):
+        frames = get_frames(nframe, 2)
+        _testcapi.tracemalloc_track(self.domain, self.ptr, self.size,
+                                    release_gil)
+        return frames
+
+    def untrack(self):
+        _testcapi.tracemalloc_untrack(self.domain, self.ptr)
+
+    def get_traced_memory(self):
+        # Get the traced size in the domain
+        snapshot = tracemalloc.take_snapshot()
+        domain_filter = tracemalloc.DomainFilter(True, self.domain)
+        snapshot = snapshot.filter_traces([domain_filter])
+        return sum(trace.size for trace in snapshot.traces)
+
+    def check_track(self, release_gil):
+        nframe = 5
+        tracemalloc.start(nframe)
+
+        size = tracemalloc.get_traced_memory()[0]
+
+        frames = self.track(release_gil, nframe)
+        self.assertEqual(self.get_traceback(),
+                         tracemalloc.Traceback(frames))
+
+        self.assertEqual(self.get_traced_memory(), self.size)
+
+    def test_track(self):
+        self.check_track(False)
+
+    def test_track_without_gil(self):
+        # check that calling _PyTraceMalloc_Track() without holding the GIL
+        # works too
+        self.check_track(True)
+
+    def test_track_already_tracked(self):
+        nframe = 5
+        tracemalloc.start(nframe)
+
+        # track a first time
+        self.track()
+
+        # calling _PyTraceMalloc_Track() must remove the old trace and add
+        # a new trace with the new traceback
+        frames = self.track(nframe=nframe)
+        self.assertEqual(self.get_traceback(),
+                         tracemalloc.Traceback(frames))
+
+    def test_untrack(self):
+        tracemalloc.start()
+
+        self.track()
+        self.assertIsNotNone(self.get_traceback())
+        self.assertEqual(self.get_traced_memory(), self.size)
+
+        # untrack must remove the trace
+        self.untrack()
+        self.assertIsNone(self.get_traceback())
+        self.assertEqual(self.get_traced_memory(), 0)
+
+        # calling _PyTraceMalloc_Untrack() multiple times must not crash
+        self.untrack()
+        self.untrack()
+
+    def test_stop_track(self):
+        tracemalloc.start()
+        tracemalloc.stop()
+
+        with self.assertRaises(RuntimeError):
+            self.track()
+        self.assertIsNone(self.get_traceback())
+
+    def test_stop_untrack(self):
+        tracemalloc.start()
+        self.track()
+
+        tracemalloc.stop()
+        with self.assertRaises(RuntimeError):
+            self.untrack()
+
+
 def test_main():
     support.run_unittest(
         TestTracemallocEnabled,
         TestSnapshot,
         TestFilters,
         TestCommandLine,
+        TestCAPI,
     )
 
 if __name__ == "__main__":
diff --git a/Misc/NEWS b/Misc/NEWS
index 841f8a0..29fc65d 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -232,6 +232,11 @@ Core and Builtins
 Library
 -------
 
+- Issue #26530: Add C functions :c:func:`_PyTraceMalloc_Track` and
+  :c:func:`_PyTraceMalloc_Untrack` to track memory blocks using the
+  :mod:`tracemalloc` module. Add :c:func:`_PyTraceMalloc_GetTraceback` to get
+  the traceback of an object.
+
 - Issue #26588: The _tracemalloc now supports tracing memory allocations of
   multiple address spaces (domains).
 
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);
+}