gh-93502: Add new C-API functions to trace object creation and destruction (#115945)

1 files changed, 84 insertions, 0 deletions

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 3448291..f5892fc 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -3219,6 +3219,89 @@ test_weakref_capi(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
     _Py_COMP_DIAG_POP
 }
 
+struct simpletracer_data {
+    int create_count;
+    int destroy_count;
+    void* addresses[10];
+};
+
+static int _simpletracer(PyObject *obj, PyRefTracerEvent event, void* data) {
+    struct simpletracer_data* the_data = (struct simpletracer_data*)data;
+    assert(the_data->create_count + the_data->destroy_count < (int)Py_ARRAY_LENGTH(the_data->addresses));
+    the_data->addresses[the_data->create_count + the_data->destroy_count] = obj;
+    if (event == PyRefTracer_CREATE) {
+        the_data->create_count++;
+    } else {
+        the_data->destroy_count++;
+    }
+    return 0;
+}
+
+static PyObject *
+test_reftracer(PyObject *ob, PyObject *Py_UNUSED(ignored))
+{
+    // Save the current tracer and data to restore it later
+    void* current_data;
+    PyRefTracer current_tracer = PyRefTracer_GetTracer(&current_data);
+
+    struct simpletracer_data tracer_data = {0};
+    void* the_data = &tracer_data;
+    // Install a simple tracer function
+    if (PyRefTracer_SetTracer(_simpletracer, the_data) != 0) {
+        goto failed;
+    }
+
+    // Check that the tracer was correctly installed
+    void* data;
+    if (PyRefTracer_GetTracer(&data) != _simpletracer || data != the_data) {
+        PyErr_SetString(PyExc_AssertionError, "The reftracer not correctly installed");
+        (void)PyRefTracer_SetTracer(NULL, NULL);
+        goto failed;
+    }
+
+    // Create a bunch of objects
+    PyObject* obj = PyList_New(0);
+    if (obj == NULL) {
+        goto failed;
+    }
+    PyObject* obj2 = PyDict_New();
+    if (obj2 == NULL) {
+        Py_DECREF(obj);
+        goto failed;
+    }
+
+    // Kill all objects
+    Py_DECREF(obj);
+    Py_DECREF(obj2);
+
+    // Remove the tracer
+    (void)PyRefTracer_SetTracer(NULL, NULL);
+
+    // Check that the tracer was removed
+    if (PyRefTracer_GetTracer(&data) != NULL || data != NULL) {
+        PyErr_SetString(PyExc_ValueError, "The reftracer was not correctly removed");
+        goto failed;
+    }
+
+    if (tracer_data.create_count != 2 ||
+        tracer_data.addresses[0] != obj ||
+        tracer_data.addresses[1] != obj2) {
+        PyErr_SetString(PyExc_ValueError, "The object creation was not correctly traced");
+        goto failed;
+    }
+
+    if (tracer_data.destroy_count != 2 ||
+        tracer_data.addresses[2] != obj ||
+        tracer_data.addresses[3] != obj2) {
+        PyErr_SetString(PyExc_ValueError, "The object destruction was not correctly traced");
+        goto failed;
+    }
+    PyRefTracer_SetTracer(current_tracer, current_data);
+    Py_RETURN_NONE;
+failed:
+    PyRefTracer_SetTracer(current_tracer, current_data);
+    return NULL;
+}
 
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
@@ -3257,6 +3340,7 @@ static PyMethodDef TestMethods[] = {
     {"get_type_fullyqualname",   get_type_fullyqualname,         METH_O},
     {"get_type_module_name",     get_type_module_name,           METH_O},
     {"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},