bpo-35053: Enhance tracemalloc to trace free lists (GH-10063)

tracemalloc now tries to update the traceback when an object is
reused from a "free list" (optimization for faster object creation,
used by the builtin list type for example).

Changes:

* Add _PyTraceMalloc_NewReference() function which tries to update
  the Python traceback of a Python object.
* _Py_NewReference() now calls _PyTraceMalloc_NewReference().
* Add an unit test.

1 files changed, 84 insertions, 51 deletions

diff --git a/Modules/_tracemalloc.c b/Modules/_tracemalloc.c
index e07022c..d736b24 100644
--- a/Modules/_tracemalloc.c
+++ b/Modules/_tracemalloc.c
@@ -29,27 +29,6 @@ static struct {
     PyMemAllocatorEx obj;
 } allocators;
 
-static struct {
-    /* Module initialized?
-       Variable protected by the GIL */
-    enum {
-        TRACEMALLOC_NOT_INITIALIZED,
-        TRACEMALLOC_INITIALIZED,
-        TRACEMALLOC_FINALIZED
-    } initialized;
-
-    /* Is tracemalloc tracing memory allocations?
-       Variable protected by the GIL */
-    int tracing;
-
-    /* limit of the number of frames in a traceback, 1 by default.
-       Variable protected by the GIL. */
-    int max_nframe;
-
-    /* use domain in trace key?
-       Variable protected by the GIL. */
-    int use_domain;
-} tracemalloc_config = {TRACEMALLOC_NOT_INITIALIZED, 0, 1, 0};
 
 #if defined(TRACE_RAW_MALLOC)
 /* This lock is needed because tracemalloc_free() is called without
@@ -459,7 +438,7 @@ traceback_get_frames(traceback_t *traceback)
         tracemalloc_get_frame(pyframe, &traceback->frames[traceback->nframe]);
         assert(traceback->frames[traceback->nframe].filename != NULL);
         traceback->nframe++;
-        if (traceback->nframe == tracemalloc_config.max_nframe)
+        if (traceback->nframe == _Py_tracemalloc_config.max_nframe)
             break;
     }
 }
@@ -540,7 +519,7 @@ tracemalloc_use_domain(void)
 {
     _Py_hashtable_t *new_traces = NULL;
 
-    assert(!tracemalloc_config.use_domain);
+    assert(!_Py_tracemalloc_config.use_domain);
 
     new_traces = hashtable_new(sizeof(pointer_t),
                                sizeof(trace_t),
@@ -560,7 +539,7 @@ tracemalloc_use_domain(void)
     _Py_hashtable_destroy(tracemalloc_traces);
     tracemalloc_traces = new_traces;
 
-    tracemalloc_config.use_domain = 1;
+    _Py_tracemalloc_config.use_domain = 1;
 
     return 0;
 }
@@ -572,9 +551,9 @@ tracemalloc_remove_trace(unsigned int domain, uintptr_t ptr)
     trace_t trace;
     int removed;
 
-    assert(tracemalloc_config.tracing);
+    assert(_Py_tracemalloc_config.tracing);
 
-    if (tracemalloc_config.use_domain) {
+    if (_Py_tracemalloc_config.use_domain) {
         pointer_t key = {ptr, domain};
         removed = _Py_HASHTABLE_POP(tracemalloc_traces, key, trace);
     }
@@ -603,14 +582,14 @@ tracemalloc_add_trace(unsigned int domain, uintptr_t ptr,
     _Py_hashtable_entry_t* entry;
     int res;
 
-    assert(tracemalloc_config.tracing);
+    assert(_Py_tracemalloc_config.tracing);
 
     traceback = traceback_new();
     if (traceback == NULL) {
         return -1;
     }
 
-    if (!tracemalloc_config.use_domain && domain != DEFAULT_DOMAIN) {
+    if (!_Py_tracemalloc_config.use_domain && domain != DEFAULT_DOMAIN) {
         /* first trace using a non-zero domain whereas traces use compact
            (uintptr_t) keys: switch to pointer_t keys. */
         if (tracemalloc_use_domain() < 0) {
@@ -618,7 +597,7 @@ tracemalloc_add_trace(unsigned int domain, uintptr_t ptr,
         }
     }
 
-    if (tracemalloc_config.use_domain) {
+    if (_Py_tracemalloc_config.use_domain) {
         entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, key);
     }
     else {
@@ -639,7 +618,7 @@ tracemalloc_add_trace(unsigned int domain, uintptr_t ptr,
         trace.size = size;
         trace.traceback = traceback;
 
-        if (tracemalloc_config.use_domain) {
+        if (_Py_tracemalloc_config.use_domain) {
             res = _Py_HASHTABLE_SET(tracemalloc_traces, key, trace);
         }
         else {
@@ -956,13 +935,13 @@ tracemalloc_clear_traces(void)
 static int
 tracemalloc_init(void)
 {
-    if (tracemalloc_config.initialized == TRACEMALLOC_FINALIZED) {
+    if (_Py_tracemalloc_config.initialized == TRACEMALLOC_FINALIZED) {
         PyErr_SetString(PyExc_RuntimeError,
                         "the tracemalloc module has been unloaded");
         return -1;
     }
 
-    if (tracemalloc_config.initialized == TRACEMALLOC_INITIALIZED)
+    if (_Py_tracemalloc_config.initialized == TRACEMALLOC_INITIALIZED)
         return 0;
 
     PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &allocators.raw);
@@ -996,7 +975,7 @@ tracemalloc_init(void)
                                            hashtable_hash_traceback,
                                            hashtable_compare_traceback);
 
-    if (tracemalloc_config.use_domain) {
+    if (_Py_tracemalloc_config.use_domain) {
         tracemalloc_traces = hashtable_new(sizeof(pointer_t),
                                            sizeof(trace_t),
                                            hashtable_hash_pointer_t,
@@ -1026,7 +1005,7 @@ tracemalloc_init(void)
     tracemalloc_empty_traceback.frames[0].lineno = 0;
     tracemalloc_empty_traceback.hash = traceback_hash(&tracemalloc_empty_traceback);
 
-    tracemalloc_config.initialized = TRACEMALLOC_INITIALIZED;
+    _Py_tracemalloc_config.initialized = TRACEMALLOC_INITIALIZED;
     return 0;
 }
 
@@ -1034,9 +1013,9 @@ tracemalloc_init(void)
 static void
 tracemalloc_deinit(void)
 {
-    if (tracemalloc_config.initialized != TRACEMALLOC_INITIALIZED)
+    if (_Py_tracemalloc_config.initialized != TRACEMALLOC_INITIALIZED)
         return;
-    tracemalloc_config.initialized = TRACEMALLOC_FINALIZED;
+    _Py_tracemalloc_config.initialized = TRACEMALLOC_FINALIZED;
 
     tracemalloc_stop();
 
@@ -1077,13 +1056,13 @@ tracemalloc_start(int max_nframe)
         return -1;
     }
 
-    if (tracemalloc_config.tracing) {
+    if (_Py_tracemalloc_config.tracing) {
         /* hook already installed: do nothing */
         return 0;
     }
 
     assert(1 <= max_nframe && max_nframe <= MAX_NFRAME);
-    tracemalloc_config.max_nframe = max_nframe;
+    _Py_tracemalloc_config.max_nframe = max_nframe;
 
     /* allocate a buffer to store a new traceback */
     size = TRACEBACK_SIZE(max_nframe);
@@ -1119,7 +1098,7 @@ tracemalloc_start(int max_nframe)
     PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
 
     /* everything is ready: start tracing Python memory allocations */
-    tracemalloc_config.tracing = 1;
+    _Py_tracemalloc_config.tracing = 1;
 
     return 0;
 }
@@ -1128,11 +1107,11 @@ tracemalloc_start(int max_nframe)
 static void
 tracemalloc_stop(void)
 {
-    if (!tracemalloc_config.tracing)
+    if (!_Py_tracemalloc_config.tracing)
         return;
 
     /* stop tracing Python memory allocations */
-    tracemalloc_config.tracing = 0;
+    _Py_tracemalloc_config.tracing = 0;
 
     /* unregister the hook on memory allocators */
 #ifdef TRACE_RAW_MALLOC
@@ -1160,7 +1139,7 @@ static PyObject *
 _tracemalloc_is_tracing_impl(PyObject *module)
 /*[clinic end generated code: output=2d763b42601cd3ef input=af104b0a00192f63]*/
 {
-    return PyBool_FromLong(tracemalloc_config.tracing);
+    return PyBool_FromLong(_Py_tracemalloc_config.tracing);
 }
 
 
@@ -1174,7 +1153,7 @@ static PyObject *
 _tracemalloc_clear_traces_impl(PyObject *module)
 /*[clinic end generated code: output=a86080ee41b84197 input=0dab5b6c785183a5]*/
 {
-    if (!tracemalloc_config.tracing)
+    if (!_Py_tracemalloc_config.tracing)
         Py_RETURN_NONE;
 
     set_reentrant(1);
@@ -1299,7 +1278,7 @@ tracemalloc_get_traces_fill(_Py_hashtable_t *traces, _Py_hashtable_entry_t *entr
     PyObject *tracemalloc_obj;
     int res;
 
-    if (tracemalloc_config.use_domain) {
+    if (_Py_tracemalloc_config.use_domain) {
         pointer_t key;
         _Py_HASHTABLE_ENTRY_READ_KEY(traces, entry, key);
         domain = key.domain;
@@ -1359,7 +1338,7 @@ _tracemalloc__get_traces_impl(PyObject *module)
     if (get_traces.list == NULL)
         goto error;
 
-    if (!tracemalloc_config.tracing)
+    if (!_Py_tracemalloc_config.tracing)
         return get_traces.list;
 
     /* the traceback hash table is used temporarily to intern traceback tuple
@@ -1414,11 +1393,11 @@ tracemalloc_get_traceback(unsigned int domain, uintptr_t ptr)
     trace_t trace;
     int found;
 
-    if (!tracemalloc_config.tracing)
+    if (!_Py_tracemalloc_config.tracing)
         return NULL;
 
     TABLES_LOCK();
-    if (tracemalloc_config.use_domain) {
+    if (_Py_tracemalloc_config.use_domain) {
         pointer_t key = {ptr, domain};
         found = _Py_HASHTABLE_GET(tracemalloc_traces, key, trace);
     }
@@ -1558,7 +1537,7 @@ static PyObject *
 _tracemalloc_get_traceback_limit_impl(PyObject *module)
 /*[clinic end generated code: output=d556d9306ba95567 input=da3cd977fc68ae3b]*/
 {
-    return PyLong_FromLong(tracemalloc_config.max_nframe);
+    return PyLong_FromLong(_Py_tracemalloc_config.max_nframe);
 }
 
 
@@ -1603,7 +1582,7 @@ _tracemalloc_get_traced_memory_impl(PyObject *module)
 {
     Py_ssize_t size, peak_size;
 
-    if (!tracemalloc_config.tracing)
+    if (!_Py_tracemalloc_config.tracing)
         return Py_BuildValue("ii", 0, 0);
 
     TABLES_LOCK();
@@ -1681,7 +1660,7 @@ PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
     int res;
     PyGILState_STATE gil_state;
 
-    if (!tracemalloc_config.tracing) {
+    if (!_Py_tracemalloc_config.tracing) {
         /* tracemalloc is not tracing: do nothing */
         return -2;
     }
@@ -1700,7 +1679,7 @@ PyTraceMalloc_Track(unsigned int domain, uintptr_t ptr,
 int
 PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
 {
-    if (!tracemalloc_config.tracing) {
+    if (!_Py_tracemalloc_config.tracing) {
         /* tracemalloc is not tracing: do nothing */
         return -2;
     }
@@ -1713,6 +1692,60 @@ PyTraceMalloc_Untrack(unsigned int domain, uintptr_t ptr)
 }
 
 
+/* If the object memory block is already traced, update its trace
+   with the current Python traceback.
+
+   Do nothing if tracemalloc is not tracing memory allocations
+   or if the object memory block is not already traced. */
+int
+_PyTraceMalloc_NewReference(PyObject *op)
+{
+    assert(PyGILState_Check());
+
+    if (!_Py_tracemalloc_config.tracing) {
+        /* tracemalloc is not tracing: do nothing */
+        return -1;
+    }
+
+    uintptr_t ptr;
+    PyTypeObject *type = Py_TYPE(op);
+    if (PyType_IS_GC(type)) {
+        ptr = (uintptr_t)((char *)op - sizeof(PyGC_Head));
+    }
+    else {
+        ptr = (uintptr_t)op;
+    }
+
+    _Py_hashtable_entry_t* entry;
+    int res = -1;
+
+    TABLES_LOCK();
+    if (_Py_tracemalloc_config.use_domain) {
+        pointer_t key = {ptr, DEFAULT_DOMAIN};
+        entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, key);
+    }
+    else {
+        entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, ptr);
+    }
+
+    if (entry != NULL) {
+        /* update the traceback of the memory block */
+        traceback_t *traceback = traceback_new();
+        if (traceback != NULL) {
+            trace_t trace;
+            _Py_HASHTABLE_ENTRY_READ_DATA(tracemalloc_traces, entry, trace);
+            trace.traceback = traceback;
+            _Py_HASHTABLE_ENTRY_WRITE_DATA(tracemalloc_traces, entry, trace);
+            res = 0;
+        }
+    }
+    /* else: cannot track the object, its memory block size is unknown */
+    TABLES_UNLOCK();
+
+    return res;
+}
+
+
 PyObject*
 _PyTraceMalloc_GetTraceback(unsigned int domain, uintptr_t ptr)
 {