bpo-30054: Expose tracemalloc C API (#1236)

* Make PyTraceMalloc_Track() and PyTraceMalloc_Untrack() functions
  public (remove the "_" prefix)
* Remove the _PyTraceMalloc_domain_t type: use directly unsigned
  int.
* Document methods

Note: methods are already tested in test_tracemalloc.

Doc/c-api/memory.rst

@@ -429,6 +429,28 @@ Customize pymalloc Arena Allocator
    Set the arena allocator.
 
 
+tracemalloc C API
+=================
+
+.. versionadded:: 3.7
+
+.. c:function: int PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr, size_t size)
+
+   Track an allocated memory block in the :mod:`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 is already tracked, update the existing trace.
+
+.. c:function: int PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
+
+   Untrack an allocated memory block in the :mod:`tracemalloc` module.
+   Do nothing if the block was not tracked.
+
+   Return ``-2`` if tracemalloc is disabled, otherwise return ``0``.
+
+
 .. _memoryexamples:
 
 Examples

Doc/library/tracemalloc.rst

@@ -412,6 +412,9 @@ Filter
 
       Address space of a memory block (``int`` or ``None``).
 
+      tracemalloc uses the domain ``0`` to trace memory allocations made by
+      Python. C extensions can use other domains to trace other resources.
+
    .. attribute:: inclusive
 
       If *inclusive* is ``True`` (include), only match memory blocks allocated
@@ -622,6 +625,16 @@ Trace
    The :attr:`Snapshot.traces` attribute is a sequence of :class:`Trace`
    instances.
 
+   .. versionchanged:: 3.6
+      Added the :attr:`domain` attribute.
+
+   .. attribute:: domain
+
+      Address space of a memory block (``int``). Read-only property.
+
+      tracemalloc uses the domain ``0`` to trace memory allocations made by
+      Python. C extensions can use other domains to trace other resources.
+
    .. attribute:: size
 
       Size of the memory block in bytes (``int``).

Include/pymem.h

@@ -25,9 +25,6 @@ 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).
@@ -35,8 +32,8 @@ typedef unsigned int _PyTraceMalloc_domain_t;
    Return -2 if tracemalloc is disabled.
 
    If memory block is already tracked, update the existing trace. */
-PyAPI_FUNC(int) _PyTraceMalloc_Track(
-    _PyTraceMalloc_domain_t domain,
+PyAPI_FUNC(int) PyTraceMalloc_Track(
+    unsigned int domain,
     uintptr_t ptr,
     size_t size);
 
@@ -44,8 +41,8 @@ PyAPI_FUNC(int) _PyTraceMalloc_Track(
    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,
+PyAPI_FUNC(int) PyTraceMalloc_Untrack(
+    unsigned int domain,
     uintptr_t ptr);
 
 /* Get the traceback where a memory block was allocated.
@@ -57,7 +54,7 @@ PyAPI_FUNC(int) _PyTraceMalloc_Untrack(
 
    Raise an exception and return NULL on error. */
 PyAPI_FUNC(PyObject*) _PyTraceMalloc_GetTraceback(
-    _PyTraceMalloc_domain_t domain,
+    unsigned int domain,
     uintptr_t ptr);
 #endif   /* !Py_LIMITED_API */
 

Modules/_testcapimodule.c

@@ -3958,15 +3958,15 @@ tracemalloc_track(PyObject *self, PyObject *args)
 
     if (release_gil) {
         Py_BEGIN_ALLOW_THREADS
-        res = _PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
+        res = PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
         Py_END_ALLOW_THREADS
     }
     else {
-        res = _PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
+        res = PyTraceMalloc_Track(domain, (uintptr_t)ptr, size);
     }
 
     if (res < 0) {
-        PyErr_SetString(PyExc_RuntimeError, "_PyTraceMalloc_Track error");
+        PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Track error");
         return NULL;
     }
 
@@ -3987,9 +3987,9 @@ tracemalloc_untrack(PyObject *self, PyObject *args)
     if (PyErr_Occurred())
         return NULL;
 
-    res = _PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
+    res = PyTraceMalloc_Untrack(domain, (uintptr_t)ptr);
     if (res < 0) {
-        PyErr_SetString(PyExc_RuntimeError, "_PyTraceMalloc_Track error");
+        PyErr_SetString(PyExc_RuntimeError, "PyTraceMalloc_Untrack error");
         return NULL;
     }
 

Modules/_tracemalloc.c

@@ -74,7 +74,7 @@ __attribute__((packed))
 #endif
 {
     uintptr_t ptr;
-    _PyTraceMalloc_domain_t domain;
+    unsigned int domain;
 } pointer_t;
 
 /* Pack the frame_t structure to reduce the memory footprint on 64-bit
@@ -578,7 +578,7 @@ tracemalloc_use_domain(void)
 
 
 static void
-tracemalloc_remove_trace(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
+tracemalloc_remove_trace(unsigned int domain, uintptr_t ptr)
 {
     trace_t trace;
     int removed;
@@ -605,7 +605,7 @@ tracemalloc_remove_trace(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
 
 
 static int
-tracemalloc_add_trace(_PyTraceMalloc_domain_t domain, uintptr_t ptr,
+tracemalloc_add_trace(unsigned int domain, uintptr_t ptr,
                       size_t size)
 {
     pointer_t key = {ptr, domain};
@@ -1267,7 +1267,7 @@ traceback_to_pyobject(traceback_t *traceback, _Py_hashtable_t *intern_table)
 
 
 static PyObject*
-trace_to_pyobject(_PyTraceMalloc_domain_t domain, trace_t *trace,
+trace_to_pyobject(unsigned int domain, trace_t *trace,
                   _Py_hashtable_t *intern_tracebacks)
 {
     PyObject *trace_obj = NULL;
@@ -1313,7 +1313,7 @@ tracemalloc_get_traces_fill(_Py_hashtable_t *traces, _Py_hashtable_entry_t *entr
                             void *user_data)
 {
     get_traces_t *get_traces = user_data;
-    _PyTraceMalloc_domain_t domain;
+    unsigned int domain;
     trace_t trace;
     PyObject *tracemalloc_obj;
     int res;
@@ -1428,7 +1428,7 @@ _tracemalloc__get_traces_impl(PyObject *module)
 
 
 static traceback_t*
-tracemalloc_get_traceback(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
+tracemalloc_get_traceback(unsigned int domain, uintptr_t ptr)
 {
     trace_t trace;
     int found;
@@ -1783,8 +1783,8 @@ _PyTraceMalloc_Fini(void)
 }
 
 int
-_PyTraceMalloc_Track(_PyTraceMalloc_domain_t domain, uintptr_t ptr,
-                     size_t size)
+PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
+                    size_t size)
 {
     int res;
 #ifdef WITH_THREAD
@@ -1812,7 +1812,7 @@ _PyTraceMalloc_Track(_PyTraceMalloc_domain_t domain, uintptr_t ptr,
 
 
 int
-_PyTraceMalloc_Untrack(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
+PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
 {
     if (!tracemalloc_config.tracing) {
         /* tracemalloc is not tracing: do nothing */
@@ -1828,7 +1828,7 @@ _PyTraceMalloc_Untrack(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
 
 
 PyObject*
-_PyTraceMalloc_GetTraceback(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
+_PyTraceMalloc_GetTraceback(unsigned int domain, uintptr_t ptr)
 {
     traceback_t *traceback;