perf(profiling): improve scaling of memory profiler for large heaps (#13317)

The heap profiler component of memalloc uses a plain array to track
allocations. It performs a linear scan of this array on every free to
see if the freed allocation is tracked. This means the cost of freeing
an allocation increases substantially for programs with large heaps.

Switch to a hash table for tracking allocations. This should have roughly
constant overhead even for larger heaps, where we track more allocations.

This PR uses the C implementation of the [Abseil
SwissTable](https://abseil.io/blog/20180927-swisstables)
from https://github.com/google/cwisstable. This hash table is designed
for memory efficiency, both in terms of memory used (1 byte per entry of
metadata) and in terms of access (cache-friendly layout). The table
is designed in particular for fast lookup and insertion. We do a lookup
for every single free, so it's important that it's fast.

This PR previously tried to use a C++ `std::unordered_map`, but it
proved far too difficult to integrate with the existing C `memalloc`
implementation.

This PR includes the single `cwisstable.h` header with the following
modifications, which can be found in the header by searching for
`BEGIN MODIFICATION`:

- Fix `CWISS_Mul128`: for 64-bit windows we need a different intrinsic,
   and for 32-bit systems we can't use it at all
- Fix a multiple declaration and invalid C++ syntax in
`CWISS_LeadingZeroes64`
   for 32-bit Windows

The PR also supplies a 32-bit compatible hash function since the
one in `cwisstable.h` couldn't be easily made to work. We use `xxHash`
instead. Note that `cwisstable.h` requires the hash function to return
a `size_t`, which is 32 bits on 32-bit platforms, but the SwissTable design
expects 64-bit hashes. I don't know if this is a huge problem in practice
or if people even use our profiler on 32-bit systems, but it's something to
watch out for.

Note that the `cwisstable.h` hash table implementation basically never
gets rid of backing memory unless we completely clear the table. This
is not a big deal in practice. Each entry in the table is an 8 byte `void*` key
and an 8 byte `traceback_t*`, plus 1 byte of control metadata. If we assume
a target load factor of 50%, and we keep the existing `2^16` element cap,
then a completely full table is only about 2MiB. Most of the heap profiler
memory usage comes from the tracebacks themselves, which we _do_
free when they're removed from the profiler structures.

This PR also fixes a potential memory leak. Since we use try-locks, it theoretically
possible to fail to un-track an allocation if we can't acquire the lock.
When this happens, the previous implementation would just leak an entry.
With the hash table, if we find an existing entry for a given address, we can free the
old traceback and replace it with the new one. The real fix will be to fix how we
handle locking, but this PR improves the situation.

I've tested this locally and see very good performance even for larger heaps,
and this has also been tested in #13307.

Fixes #13307

Author: nsrip-dd

ddtrace/profiling/collector/_memalloc_heap.c

@@ -3,6 +3,7 @@
 
 #define PY_SSIZE_T_CLEAN
 #include "_memalloc_heap.h"
+#include "_memalloc_heap_map.h"
 #include "_memalloc_reentrant.h"
 #include "_memalloc_tb.h"
 
@@ -66,15 +67,15 @@ typedef struct
     /* Next heap sample target, in bytes allocated */
     uint64_t current_sample_size;
     /* Tracked allocations */
-    traceback_array_t allocs;
+    memalloc_heap_map_t* allocs_m;
     /* Bytes allocated since the last sample was collected */
     uint64_t allocated_memory;
     /* True if the heap tracker is frozen */
     bool frozen;
     /* Contains the ongoing heap allocation/deallocation while frozen */
     struct
     {
-        traceback_array_t allocs;
+        memalloc_heap_map_t* allocs_m;
         ptr_array_t frees;
     } freezer;
 } heap_tracker_t;
@@ -146,8 +147,8 @@ heap_tracker_next_sample_size(uint32_t sample_size)
 static void
 heap_tracker_init(heap_tracker_t* heap_tracker)
 {
-    traceback_array_init(&heap_tracker->allocs);
-    traceback_array_init(&heap_tracker->freezer.allocs);
+    heap_tracker->allocs_m = memalloc_heap_map_new();
+    heap_tracker->freezer.allocs_m = memalloc_heap_map_new();
     ptr_array_init(&heap_tracker->freezer.frees);
     heap_tracker->allocated_memory = 0;
     heap_tracker->frozen = false;
@@ -158,8 +159,8 @@ heap_tracker_init(heap_tracker_t* heap_tracker)
 static void
 heap_tracker_wipe(heap_tracker_t* heap_tracker)
 {
-    traceback_array_wipe(&heap_tracker->allocs);
-    traceback_array_wipe(&heap_tracker->freezer.allocs);
+    memalloc_heap_map_delete(heap_tracker->allocs_m);
+    memalloc_heap_map_delete(heap_tracker->freezer.allocs_m);
     ptr_array_wipe(&heap_tracker->freezer.frees);
 }
 
@@ -172,48 +173,26 @@ heap_tracker_freeze(heap_tracker_t* heap_tracker)
 static void
 heap_tracker_untrack_thawed(heap_tracker_t* heap_tracker, void* ptr)
 {
-    /* This search is O(n) where `n` is the number of tracked traceback,
-       which is linearly linked to the heap size. This search could probably be
-       optimized in a couple of ways:
-
-       - sort the traceback in allocs by ptr so we can find the ptr in O(log2 n)
-       - use a Bloom filter?
-
-       That being said, we start iterating at the end of the array because most
-       of the time this is where the untracked ptr is (the most recent object
-       get de-allocated first usually). This might be a good enough
-       trade-off. */
-    for (TRACEBACK_ARRAY_COUNT_TYPE i = heap_tracker->allocs.count; i > 0; i--) {
-        traceback_t** tb = &heap_tracker->allocs.tab[i - 1];
-
-        if (ptr == (*tb)->ptr) {
-            /* Free the traceback */
-            traceback_free(*tb);
-            traceback_array_remove(&heap_tracker->allocs, tb);
-            break;
-        }
+    traceback_t* tb = memalloc_heap_map_remove(heap_tracker->allocs_m, ptr);
+    if (tb) {
+        traceback_free(tb);
     }
 }
 
 static void
 heap_tracker_thaw(heap_tracker_t* heap_tracker)
 {
-    /* Add the frozen allocs at the end */
-    traceback_array_splice(&heap_tracker->allocs,
-                           heap_tracker->allocs.count,
-                           0,
-                           heap_tracker->freezer.allocs.tab,
-                           heap_tracker->freezer.allocs.count);
+    memalloc_heap_map_destructive_copy(heap_tracker->allocs_m, heap_tracker->freezer.allocs_m);
 
     /* Handle the frees: we need to handle the frees after we merge the allocs
        array together to be sure that there's no free in the freezer matching
        an alloc that is also in the freezer; heap_tracker_untrack_thawed does
        not care about the freezer, by definition. */
-    for (MEMALLOC_HEAP_PTR_ARRAY_COUNT_TYPE i = 0; i < heap_tracker->freezer.frees.count; i++)
+    for (MEMALLOC_HEAP_PTR_ARRAY_COUNT_TYPE i = 0; i < heap_tracker->freezer.frees.count; i++) {
         heap_tracker_untrack_thawed(heap_tracker, heap_tracker->freezer.frees.tab[i]);
+    }
 
     /* Reset the count to zero so we can reused the array and overwrite previous values */
-    heap_tracker->freezer.allocs.count = 0;
     heap_tracker->freezer.frees.count = 0;
 
     heap_tracker->frozen = false;
@@ -249,16 +228,21 @@ memalloc_heap_untrack(void* ptr)
         return;
     }
     if (global_heap_tracker.frozen) {
-        /* Check that we still have space to store the free. If we don't have
-           enough space, we ignore the untrack. That's sad as there is a change
-           the heap profile won't be valid anymore. However, that's the best we
-           can do since reporting an error is not an option here. What's gonna
-           free more than 2^64 pointers anyway?!
-        */
-        if (global_heap_tracker.freezer.frees.count < MEMALLOC_HEAP_PTR_ARRAY_MAX_COUNT)
+        /* NB: because we also take the lock in memalloc_heap, we're not going to
+         * take this branch. We will, however, if we start using the GIL as our
+         * lock instead.
+         */
+        traceback_t* tb = memalloc_heap_map_remove(global_heap_tracker.freezer.allocs_m, ptr);
+        if (tb) {
+            traceback_free(tb);
+        } else if (memalloc_heap_map_contains(global_heap_tracker.allocs_m, ptr)) {
+            /* We're tracking this pointer but can't remove it right now because
+             * we're iterating over the map. Save the pointer to remove later */
             ptr_array_append(&global_heap_tracker.freezer.frees, ptr);
-    } else
+        }
+    } else {
         heap_tracker_untrack_thawed(&global_heap_tracker, ptr);
+    }
 
     memlock_unlock(&g_memheap_lock);
 }
@@ -289,10 +273,14 @@ memalloc_heap_track(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorD
         return false;
     }
 
-    /* Check if we can add more samples: the sum of the freezer + alloc tracker
-     cannot be greater than what the alloc tracker can handle: when the alloc
-     tracker is thawed, all the allocs in the freezer will be moved there!*/
-    if (global_heap_tracker.freezer.allocs.count + global_heap_tracker.allocs.count >= TRACEBACK_ARRAY_MAX_COUNT) {
+    if (memalloc_heap_map_size(global_heap_tracker.allocs_m) +
+          memalloc_heap_map_size(global_heap_tracker.freezer.allocs_m) >
+        TRACEBACK_ARRAY_MAX_COUNT) {
+        /* TODO(nick) this is vestigial from the original array-based
+         * implementation. Do we actually want this? It gives us bounded memory
+         * use, but the size limit is arbitrary and once we hit the arbitrary
+         * limit our reported numbers will be inaccurate.
+         */
         memlock_unlock(&g_memheap_lock);
         return false;
     }
@@ -310,10 +298,11 @@ memalloc_heap_track(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorD
        of sample live allocations stays close to the actual heap size */
     traceback_t* tb = memalloc_get_traceback(max_nframe, ptr, global_heap_tracker.allocated_memory, domain);
     if (tb) {
-        if (global_heap_tracker.frozen)
-            traceback_array_append(&global_heap_tracker.freezer.allocs, tb);
-        else
-            traceback_array_append(&global_heap_tracker.allocs, tb);
+        if (global_heap_tracker.frozen) {
+            memalloc_heap_map_insert(global_heap_tracker.freezer.allocs_m, tb->ptr, tb);
+        } else {
+            memalloc_heap_map_insert(global_heap_tracker.allocs_m, tb->ptr, tb);
+        }
 
         /* Reset the counter to 0 */
         global_heap_tracker.allocated_memory = 0;
@@ -340,16 +329,7 @@ memalloc_heap()
 
     heap_tracker_freeze(&global_heap_tracker);
 
-    PyObject* heap_list = PyList_New(global_heap_tracker.allocs.count);
-
-    for (TRACEBACK_ARRAY_COUNT_TYPE i = 0; i < global_heap_tracker.allocs.count; i++) {
-        traceback_t* tb = global_heap_tracker.allocs.tab[i];
-
-        PyObject* tb_and_size = PyTuple_New(2);
-        PyTuple_SET_ITEM(tb_and_size, 0, traceback_to_tuple(tb));
-        PyTuple_SET_ITEM(tb_and_size, 1, PyLong_FromSize_t(tb->size));
-        PyList_SET_ITEM(heap_list, i, tb_and_size);
-    }
+    PyObject* heap_list = memalloc_heap_map_export(global_heap_tracker.allocs_m);
 
     heap_tracker_thaw(&global_heap_tracker);
 

ddtrace/profiling/collector/_memalloc_heap_map.c

@@ -0,0 +1,180 @@
+#include <stdlib.h>
+
+#include <Python.h>
+
+#include "_memalloc_tb.h"
+#include "vendor/cwisstable.h"
+
+/* cwisstable.h provides a C implementation of SwissTables hash maps, originally
+ * implemented in the Abseil C++ library.
+ *
+ * This header is was generated from https://github.com/google/cwisstable
+ * at commit 6de0e5f2e55f90017534a3366198ce7d3e3b7fef
+ * and lightly modified to compile for Windows and 32-bit platforms we support.
+ * See "BEGIN MODIFICATION" and "END MODIFICATION" in the header.
+ *
+ * The following macro will expand to a type-safe implementation with void* keys
+ * and traceback_t* values for use by the heap profiler. We encapsulate this
+ * implementation in a wrapper specialized for use by the heap profiler, both to
+ * keep compilation fast (the cwisstables header is big) and to allow us to swap
+ * out the implementation if we want.
+ *
+ * Note that the HeapSample tables will, in general, never free their backing
+ * memory unless we completely clear them. The table takes 17 bytes per entry: 8
+ * for the void* keys, 8 for the traceback* values, and 1 byte per entry for
+ * control metadata. Assuming a load factor target of ~50%, meaning our table
+ * has roughly twice as many slots as actual entries, then for our default
+ * maximum of 2^16 entries the table will be about 2MiB. A table this large
+ * would correspond to a program with a ~65GiB live heap with a 1MiB default
+ * sampling interval. Most of the memory usage of the profiler will come from
+ * the tracebacks themselves, which we _do_ free when we're done with them.
+ */
+#if defined(_WIN_64) || defined(__x86_64__) || defined(__aarch_64__)
+CWISS_DECLARE_FLAT_HASHMAP(HeapSamples, void*, traceback_t*);
+#else
+/* The default cwisstable hash relies on full-width 64-bit
+ * multiplication, which is really slow on 32-bit.
+ * For 32-bit, we define a custom hash function with reasonable quality.
+ * Derived from:
+ * https://github.com/Cyan4973/xxHash/blob/dev/doc/xxhash_spec.md#xxh32-algorithm-description.
+ *
+ * NOTE: cwisstable.h requires the hash function to return a size_t.
+ * On 32-bit platforms this is 32 bits, while the SwissTable design
+ * expects 64-bit hashes, with 7 of the bits are used for metadata.
+ * So we get much lower entropy on 32-bit platforms.
+ */
+static size_t
+void_ptr_hash(const void* value)
+{
+#define PRIME32_1 0x9E3779B1U
+#define PRIME32_2 0x85EBCA77U
+#define PRIME32_3 0xC2B2AE3DU
+#define PRIME32_4 0x27D4EB2FU
+#define PRIME32_5 0x165667B1U
+
+    /* "Special case: input is less than 16 bytes".
+     * Here our seed is fixed at 0 so we elide it */
+    uint32_t acc = PRIME32_5;
+
+    /* "Input length" is the size of a pointer. */
+    acc = acc + sizeof(void*);
+
+    /* "Consume remaining input".
+     * Here we know our input is just 4 bytes, the size of a pointer */
+    uint32_t lane = *((uint32_t*)value);
+    acc += lane * PRIME32_3;
+    acc = (acc << 17) * PRIME32_4;
+
+    acc ^= (acc >> 15);
+    acc *= PRIME32_2;
+    acc ^= (acc >> 13);
+    acc *= PRIME32_3;
+    acc ^= (acc >> 16);
+    return acc;
+}
+CWISS_DECLARE_FLAT_MAP_POLICY(HeapSamples_policy32, void*, traceback_t*, (key_hash, void_ptr_hash));
+CWISS_DECLARE_HASHMAP_WITH(HeapSamples, void*, traceback_t*, HeapSamples_policy32);
+#endif
+
+typedef struct memalloc_heap_map_t
+{
+    HeapSamples map;
+} memalloc_heap_map_t;
+
+memalloc_heap_map_t*
+memalloc_heap_map_new()
+{
+    memalloc_heap_map_t* m = calloc(sizeof(memalloc_heap_map_t), 1);
+    m->map = HeapSamples_new(0);
+    return m;
+}
+
+size_t
+memalloc_heap_map_size(memalloc_heap_map_t* m)
+{
+    return HeapSamples_size(&m->map);
+}
+
+void
+memalloc_heap_map_insert(memalloc_heap_map_t* m, void* key, traceback_t* value)
+{
+    HeapSamples_Entry k = { key = key, value = value };
+    HeapSamples_Insert res = HeapSamples_insert(&m->map, &k);
+    if (!res.inserted) {
+        /* This shouldn't happen, but due to the current try-locking implementation
+         * we can hypothetically fail to remove a sample, leading to two entries
+         * with the same pointer. If we see this key already, then the allocation
+         * for the existing entry was freed. Replace it with the new one
+         */
+        HeapSamples_Entry* e = HeapSamples_Iter_get(&res.iter);
+        traceback_free(e->val);
+        e->val = value;
+    }
+}
+
+bool
+memalloc_heap_map_contains(memalloc_heap_map_t* m, void* key)
+{
+    return HeapSamples_contains(&m->map, &key);
+}
+
+traceback_t*
+memalloc_heap_map_remove(memalloc_heap_map_t* m, void* key)
+{
+    traceback_t* res = NULL;
+    HeapSamples_Iter it = HeapSamples_find(&m->map, &key);
+    HeapSamples_Entry* e = HeapSamples_Iter_get(&it);
+    if (e != NULL) {
+        res = e->val;
+        /* This erases the entry but won't shrink the table. */
+        HeapSamples_erase_at(it);
+    }
+    return res;
+}
+
+PyObject*
+memalloc_heap_map_export(memalloc_heap_map_t* m)
+{
+    PyObject* heap_list = PyList_New(HeapSamples_size(&m->map));
+    if (heap_list == NULL) {
+        return NULL;
+    }
+
+    int i = 0;
+    HeapSamples_CIter it = HeapSamples_citer(&m->map);
+    for (const HeapSamples_Entry* e = HeapSamples_CIter_get(&it); e != NULL; e = HeapSamples_CIter_next(&it)) {
+        traceback_t* tb = e->val;
+
+        PyObject* tb_and_size = PyTuple_New(2);
+        PyTuple_SET_ITEM(tb_and_size, 0, traceback_to_tuple(tb));
+        PyTuple_SET_ITEM(tb_and_size, 1, PyLong_FromSize_t(tb->size));
+        PyList_SET_ITEM(heap_list, i, tb_and_size);
+        i++;
+    }
+    return heap_list;
+}
+
+void
+memalloc_heap_map_destructive_copy(memalloc_heap_map_t* dst, memalloc_heap_map_t* src)
+{
+    HeapSamples_Iter it = HeapSamples_iter(&src->map);
+    for (const HeapSamples_Entry* e = HeapSamples_Iter_get(&it); e != NULL; e = HeapSamples_Iter_next(&it)) {
+        HeapSamples_insert(&dst->map, e);
+        /* Erasing the element doesn't free up the backing storage. This is
+         * probably fine; even at full capacity the table will probably only be
+         * ~1-2MiB
+         */
+        HeapSamples_erase_at(it);
+    }
+}
+
+void
+memalloc_heap_map_delete(memalloc_heap_map_t* m)
+{
+    HeapSamples_CIter it = HeapSamples_citer(&m->map);
+    for (const HeapSamples_Entry* e = HeapSamples_CIter_get(&it); e != NULL; e = HeapSamples_CIter_next(&it)) {
+        traceback_free(e->val);
+    }
+    HeapSamples_destroy(&m->map);
+    free(m);
+}