Issue #26588:

* Optimize tracemalloc_add_trace(): modify hashtable entry data (trace) if the
  memory block is already tracked, rather than trying to remove the old trace
  and then add a new trace.
* Add _Py_HASHTABLE_ENTRY_WRITE_DATA() macro

3 files changed, 31 insertions, 12 deletions

diff --git a/Include/pymem.h b/Include/pymem.h
index 941e00f..431e5b6 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 0bab540..e6465a3 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 9c3fbdd..e3e8148 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 */