DataDog
diff --git a/‎ddtrace/profiling/collector/_memalloc_heap.c
Lines changed: 41 additions & 61 deletions b/‎ddtrace/profiling/collector/_memalloc_heap.c
Lines changed: 41 additions & 61 deletions
diff --git a/‎ddtrace/profiling/collector/_memalloc_heap_map.c
Lines changed: 177 additions & 0 deletions b/‎ddtrace/profiling/collector/_memalloc_heap_map.c
Lines changed: 177 additions & 0 deletions
@@ -4,6 +4,7 @@
 #define PY_SSIZE_T_CLEAN
 #include "_memalloc_debug.h"
 #include "_memalloc_heap.h"
+#include "_memalloc_heap_map.h"
 #include "_memalloc_reentrant.h"
 #include "_memalloc_tb.h"
 
@@ -14,15 +15,15 @@ typedef struct
     /* Current sample size of the heap profiler in bytes */
     uint64_t current_sample_size;
     /* Tracked allocations */
-    traceback_array_t allocs;
-    /* Allocated memory counter in bytes */
+    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;
@@ -88,8 +89,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;
@@ -100,8 +101,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);
 }
 
@@ -114,48 +115,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;
@@ -191,16 +170,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);
 }
@@ -231,10 +215,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;
     }
@@ -247,10 +235,11 @@ memalloc_heap_track(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorD
 
     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;
@@ -277,16 +266,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);
 
 
@@ -0,0 +1,177 @@
+#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);
+    }
+    /* Can't erase inside the loop or the iterator is invalidated */
+    HeapSamples_clear(&src->map);
+}
+
+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);
+}