diff --git a/Include/pymem.h b/Include/pymem.h index 941e00f6c2c..431e5b6d0f4 100644 --- a/Include/pymem.h +++ b/Include/pymem.h @@ -34,7 +34,7 @@ typedef unsigned int _PyTraceMalloc_domain_t; Return -2 if tracemalloc is disabled. - If memory block was already tracked, begin by removing the old trace. */ + If memory block is already tracked, update the existing trace. */ PyAPI_FUNC(int) _PyTraceMalloc_Track( _PyTraceMalloc_domain_t domain, Py_uintptr_t ptr, diff --git a/Modules/_tracemalloc.c b/Modules/_tracemalloc.c index 0bab540ee1a..e6465a3b749 100644 --- a/Modules/_tracemalloc.c +++ b/Modules/_tracemalloc.c @@ -552,33 +552,49 @@ static int tracemalloc_add_trace(_PyTraceMalloc_domain_t domain, Py_uintptr_t ptr, size_t size) { + pointer_t key = {ptr, domain}; traceback_t *traceback; trace_t trace; + _Py_hashtable_entry_t* entry; int res; assert(tracemalloc_config.tracing); - /* first, remove the previous trace (if any) */ - tracemalloc_remove_trace(domain, ptr); - traceback = traceback_new(); if (traceback == NULL) { return -1; } - trace.size = size; - trace.traceback = traceback; - if (tracemalloc_config.use_domain) { - pointer_t key = {ptr, domain}; - res = _Py_HASHTABLE_SET(tracemalloc_traces, key, trace); + entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, key); } else { - res = _Py_HASHTABLE_SET(tracemalloc_traces, ptr, trace); + entry = _Py_HASHTABLE_GET_ENTRY(tracemalloc_traces, ptr); } - if (res != 0) { - return res; + if (entry != NULL) { + /* the memory block is already tracked */ + _Py_HASHTABLE_ENTRY_READ_DATA(tracemalloc_traces, entry, trace); + assert(tracemalloc_traced_memory >= trace.size); + tracemalloc_traced_memory -= trace.size; + + trace.size = size; + trace.traceback = traceback; + _Py_HASHTABLE_ENTRY_WRITE_DATA(tracemalloc_traces, entry, trace); + } + else { + trace.size = size; + trace.traceback = traceback; + + if (tracemalloc_config.use_domain) { + res = _Py_HASHTABLE_SET(tracemalloc_traces, key, trace); + } + else { + res = _Py_HASHTABLE_SET(tracemalloc_traces, ptr, trace); + } + if (res != 0) { + return res; + } } assert(tracemalloc_traced_memory <= PY_SIZE_MAX - size); diff --git a/Modules/hashtable.h b/Modules/hashtable.h index 9c3fbdd6712..e3e8148fcc4 100644 --- a/Modules/hashtable.h +++ b/Modules/hashtable.h @@ -74,6 +74,9 @@ typedef struct { (PDATA), (DATA_SIZE)); \ } while (0) +#define _Py_HASHTABLE_ENTRY_WRITE_DATA(TABLE, ENTRY, DATA) \ + _Py_HASHTABLE_ENTRY_WRITE_PDATA(TABLE, ENTRY, sizeof(DATA), &(DATA)) + /* _Py_hashtable: prototypes */