gh-93502: Add new C-API functions to trace object creation and

destruction

Author: pablogsal

Doc/c-api/init.rst

@@ -1746,6 +1746,38 @@ Python-level trace functions in previous versions.
 
 .. versionadded:: 3.12
 
+.. c:type:: int (*PyRefTracer)(PyObject *, PyRefTracerEvent event, void* data)
+
+   The type of the trace function registered using :c:func:`PyRefTracer_SetTracer` 
+   The first parameter is a Python object that has been just created (when **event**
+   is set to :c:data:`PyRefTracer_CREATE`) or about to be destroyed (when **event**
+   is set to :c:data:`PyRefTracer_CREATE`). The **data** argument is the opaque pointer
+   that was provided when :c:func:`PyRefTracer_SetTracer` was called.
+
+.. versionadded:: 3.13
+
+.. c:function:: int PyRefTracer_SetTracer(PyRefTracer tracer, void *data)
+
+   Register a reference tracer function. The function will be called when a new Python 
+   has been created or when an object is going to be destroyed. If **data** is provided
+   it must be an opaque pointer that will be provided when the tracer function is called.
+
+   Not that tracer functions **must not** create Python objects inside or otherwise the
+   call will be re-entrant.
+
+   The GIL must be held when calling this function.
+
+.. versionadded:: 3.13
+
+.. c:function:: PyRefTracer PyRefTracer_GetTracer(void** data) {
+
+   Get the registered reference tracer function and the value of the opaque data pointer that
+   was registered when :c:func:`PyRefTracer_SetTracer` was called. If no tracer was registered
+   this function will return NULL and will set the **data** pointer to NULL.
+
+   The GIL must be held when calling this function.
+
+.. versionadded:: 3.13
 
 .. _advanced-debugging:
 

Include/cpython/object.h

@@ -507,3 +507,13 @@ PyAPI_FUNC(int) PyType_Unwatch(int watcher_id, PyObject *type);
  * assigned, or 0 if a new tag could not be assigned.
  */
 PyAPI_FUNC(int) PyUnstable_Type_AssignVersionTag(PyTypeObject *type);
+
+
+typedef enum {
+    PyRefTracer_CREATE = 0,
+    PyRefTracer_DESTROY = 1,
+} PyRefTracerEvent;
+
+typedef int (*PyRefTracer)(PyObject *, PyRefTracerEvent event, void *);
+PyAPI_FUNC(int) PyRefTracer_SetTracer(PyRefTracer tracer, void *data);
+PyAPI_FUNC(PyRefTracer) PyRefTracer_GetTracer(void**);

Include/internal/pycore_object.h

@@ -252,7 +252,7 @@ extern int _PyDict_CheckConsistency(PyObject *mp, int check_content);
    when a memory block is reused from a free list.
 
    Internal function called by _Py_NewReference(). */
-extern int _PyTraceMalloc_NewReference(PyObject *op);
+extern int _PyTraceMalloc_NewReference(PyObject *op, PyRefTracerEvent event, void*);
 
 // Fast inlined version of PyType_HasFeature()
 static inline int

Include/internal/pycore_runtime.h

@@ -125,6 +125,12 @@ typedef struct _Py_DebugOffsets {
     } tuple_object;
 } _Py_DebugOffsets;
 
+/* Reference tracer state */
+struct _reftracer_runtime_state {
+    PyRefTracer tracer_func;
+    void* data;
+};
+
 /* Full Python runtime state */
 
 /* _PyRuntimeState holds the global state for the CPython runtime.
@@ -226,6 +232,7 @@ typedef struct pyruntimestate {
     struct _fileutils_state fileutils;
     struct _faulthandler_runtime_state faulthandler;
     struct _tracemalloc_runtime_state tracemalloc;
+    struct _reftracer_runtime_state reftracer;
 
     // The rwmutex is used to prevent overlapping global and per-interpreter
     // stop-the-world events. Global stop-the-world events lock the mutex

Include/internal/pycore_runtime_init.h

@@ -119,6 +119,9 @@ extern PyTypeObject _PyExc_MemoryError;
         }, \
         .faulthandler = _faulthandler_runtime_state_INIT, \
         .tracemalloc = _tracemalloc_runtime_state_INIT, \
+        .reftracer = { \
+            .tracer_func = NULL, \
+        }, \
         .stoptheworld = { \
             .is_global = 1, \
         }, \

Misc/NEWS.d/next/Core and Builtins/2024-02-26-13-14-52.gh-issue-93502.JMWRvA.rst

@@ -0,0 +1,4 @@
+Add two new functions to the C-API, :c:func:`PyRefTracer_SetTracer` and
+:c:func:`PyRefTracer_GetTracer`, that allows to track object creation and
+destruction the same way the :mod:`tracemalloc` module does. Patch by Pablo
+Galindo

Modules/_testcapimodule.c

@@ -3284,6 +3284,66 @@ test_weakref_capi(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
     _Py_COMP_DIAG_POP
 }
 
+struct simpletracer_data {
+    int count;
+    void* addresses[10];
+};
+
+static int _simpletracer(PyObject *obj, PyRefTracerEvent event, void* data) {
+    struct simpletracer_data* the_data = (struct simpletracer_data*)data;
+    the_data->count++;
+    assert(the_data->count < 10);
+    the_data->addresses[the_data->count] = obj;
+    return 0;
+}
+
+static PyObject *
+test_reftracer(PyObject *ob, PyObject *Py_UNUSED(ignored))
+{
+    struct simpletracer_data tracer_data = {0};
+    void* the_data = (void*)&tracer_data;
+    // Install a simple tracer function
+    PyRefTracer_SetTracer(_simpletracer, the_data);
+
+    // Check that the tracer was correctly installed
+    void* data;
+    if (PyRefTracer_GetTracer(&data) != _simpletracer || data != the_data) {
+        PyErr_SetString(PyExc_ValueError, "The reftracer not correctly installed");
+        PyRefTracer_SetTracer(NULL, NULL);
+        return NULL;
+    }
+    // Create a bunch of objects
+
+    PyObject* obj = PyList_New(0);
+    if (obj == NULL) {
+        return NULL;
+    }
+    PyObject* obj2 = PyDict_New();
+    if (obj2 == NULL) {
+        Py_DECREF(obj);
+        return NULL;
+    }
+
+    // Remove the tracer
+    PyRefTracer_SetTracer(NULL, NULL);
+
+    Py_DECREF(obj);
+    Py_DECREF(obj2);
+
+    // Check that the tracer was removed
+    if (PyRefTracer_GetTracer(&data) != NULL || data != NULL) {
+        PyErr_SetString(PyExc_ValueError, "The reftracer was not correctly removed");
+        return NULL;
+    }
+
+    // Check that the objects were traced and are inside the list
+    if (tracer_data.count != 2 || tracer_data.addresses[1] != obj || tracer_data.addresses[2] != obj2) {
+        PyErr_SetString(PyExc_ValueError, "The objects were not correctly traced");
+        return NULL;
+    }
+
+    Py_RETURN_NONE;
+}
 
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
@@ -3320,6 +3380,7 @@ static PyMethodDef TestMethods[] = {
     {"test_get_type_name",        test_get_type_name,            METH_NOARGS},
     {"test_get_type_qualname",    test_get_type_qualname,        METH_NOARGS},
     {"test_get_type_dict",        test_get_type_dict,            METH_NOARGS},
+    {"test_reftracer",          test_reftracer,                  METH_NOARGS},
     {"_test_thread_state",      test_thread_state,               METH_VARARGS},
 #ifndef MS_WINDOWS
     {"_spawn_pthread_waiter",   spawn_pthread_waiter,            METH_NOARGS},

Objects/object.c

@@ -2367,8 +2367,9 @@ _PyTypes_FiniTypes(PyInterpreterState *interp)
 static inline void
 new_reference(PyObject *op)
 {
-    if (_PyRuntime.tracemalloc.config.tracing) {
-        _PyTraceMalloc_NewReference(op);
+    if (_PyRuntime.reftracer.tracer_func != NULL) {
+        void* data = _PyRuntime.reftracer.data;
+        _PyRuntime.reftracer.tracer_func(op, PyRefTracer_CREATE, data);
     }
     // Skip the immortal object check in Py_SET_REFCNT; always set refcnt to 1
 #if !defined(Py_GIL_DISABLED)
@@ -2404,8 +2405,9 @@ _Py_NewReferenceNoTotal(PyObject *op)
 void
 _Py_ResurrectReference(PyObject *op)
 {
-    if (_PyRuntime.tracemalloc.config.tracing) {
-        _PyTraceMalloc_NewReference(op);
+    if (_PyRuntime.reftracer.tracer_func != NULL) {
+        void* data = _PyRuntime.reftracer.data;
+        _PyRuntime.reftracer.tracer_func(op, PyRefTracer_CREATE, data);
     }
 #ifdef Py_TRACE_REFS
     _Py_AddToAllObjects(op);
@@ -2883,6 +2885,11 @@ _Py_Dealloc(PyObject *op)
     Py_INCREF(type);
 #endif
 
+    if (_PyRuntime.reftracer.tracer_func != NULL) {
+        void* data = _PyRuntime.reftracer.data;
+        _PyRuntime.reftracer.tracer_func(op, PyRefTracer_DESTROY, data);
+    }
+
 #ifdef Py_TRACE_REFS
     _Py_ForgetReference(op);
 #endif
@@ -2970,3 +2977,17 @@ _Py_SetRefcnt(PyObject *ob, Py_ssize_t refcnt)
 {
     Py_SET_REFCNT(ob, refcnt);
 }
+
+int PyRefTracer_SetTracer(PyRefTracer tracer, void *data) {
+    _PyRuntime.reftracer.tracer_func = tracer;
+    _PyRuntime.reftracer.data = data;
+    return 0;
+}
+
+PyRefTracer PyRefTracer_GetTracer(void** data) {
+    if (data != NULL) {
+        *data = _PyRuntime.reftracer.data;
+    }
+    return _PyRuntime.reftracer.tracer_func;
+}
+

Python/tracemalloc.c

@@ -906,6 +906,8 @@ _PyTraceMalloc_Start(int max_nframe)
         return -1;
     }
 
+    _PyRuntime.reftracer.tracer_func = _PyTraceMalloc_NewReference;
+
     if (tracemalloc_config.tracing) {
         /* hook already installed: do nothing */
         return 0;
@@ -1352,8 +1354,13 @@ _PyTraceMalloc_Fini(void)
    Do nothing if tracemalloc is not tracing memory allocations
    or if the object memory block is not already traced. */
 int
-_PyTraceMalloc_NewReference(PyObject *op)
+_PyTraceMalloc_NewReference(PyObject *op, PyRefTracerEvent event, void* Py_UNUSED(ignore))
+
 {
+    if (event != PyRefTracer_CREATE) {
+        return 0;
+    }
+
     assert(PyGILState_Check());
 
     if (!tracemalloc_config.tracing) {