refactor(profiling): unifying the allocation and live memory profilers

ddtrace/profiling/collector/_memalloc.pyi

@@ -6,10 +6,9 @@ from .. import event
 FrameType = event.DDFrame
 StackType = event.StackTraceType
 
-# (stack, nframe, thread_id)
-TracebackType = typing.Tuple[StackType, int, int]
+# (stack, thread_id)
+TracebackType = typing.Tuple[StackType, int]
 
 def start(max_nframe: int, max_events: int, heap_sample_size: int) -> None: ...
 def stop() -> None: ...
-def heap() -> typing.List[typing.Tuple[TracebackType, int]]: ...
-def iter_events() -> typing.Iterator[typing.Tuple[TracebackType, int]]: ...
+def heap() -> typing.List[typing.Tuple[TracebackType, int, int, int]]: ...

ddtrace/profiling/collector/_memalloc_heap.c

@@ -79,6 +79,8 @@ typedef struct
         memalloc_heap_map_t* allocs_m;
         ptr_array_t frees;
     } freezer;
+    /* List of freed samples that haven't been reported yet */
+    traceback_array_t unreported_samples;
 
     /* Debug guard to assert that GIL-protected critical sections are maintained
      * while accessing the profiler's state */
@@ -109,6 +111,7 @@ heap_tracker_init(heap_tracker_t* heap_tracker)
     heap_tracker->allocs_m = memalloc_heap_map_new();
     heap_tracker->freezer.allocs_m = memalloc_heap_map_new();
     ptr_array_init(&heap_tracker->freezer.frees);
+    traceback_array_init(&heap_tracker->unreported_samples);
     heap_tracker->allocated_memory = 0;
     heap_tracker->frozen = false;
     heap_tracker->sample_size = 0;
@@ -122,6 +125,7 @@ heap_tracker_wipe(heap_tracker_t* heap_tracker)
     memalloc_heap_map_delete(heap_tracker->allocs_m);
     memalloc_heap_map_delete(heap_tracker->freezer.allocs_m);
     ptr_array_wipe(&heap_tracker->freezer.frees);
+    traceback_array_wipe(&heap_tracker->unreported_samples);
 }
 
 static void
@@ -214,6 +218,12 @@ memalloc_heap_untrack_no_cpython(heap_tracker_t* heap_tracker, void* ptr)
     }
     if (!heap_tracker->frozen) {
         traceback_t* tb = memalloc_heap_map_remove(heap_tracker->allocs_m, ptr);
+        if (tb && !tb->reported) {
+            /* If the sample hasn't been reported yet, add it to the allocation list */
+            traceback_array_append(&heap_tracker->unreported_samples, tb);
+            MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard);
+            return NULL;
+        }
         MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard);
         return tb;
     }
@@ -328,7 +338,7 @@ memalloc_heap_track(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorD
        will tend to be larger for large allocations and smaller for small
        allocations, and close to the average sampling interval so that the sum
        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);
+    traceback_t* tb = memalloc_get_traceback(max_nframe, ptr, size, domain, global_heap_tracker.allocated_memory);
     if (!tb) {
         memalloc_yield_guard();
         return;
@@ -342,6 +352,36 @@ memalloc_heap_track(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorD
     memalloc_yield_guard();
 }
 
+PyObject*
+memalloc_sample_to_tuple(traceback_t* tb, bool is_live)
+{
+    PyObject* tb_and_info = PyTuple_New(4);
+    if (tb_and_info == NULL) {
+        return NULL;
+    }
+
+    size_t in_use_size;
+    size_t alloc_size;
+
+    if (is_live) {
+        /* alloc_size tracks new allocations since the last heap snapshot. Once
+         * we report it (tb->reported == true), we set the value to 0 to avoid
+         * double-counting allocations across multiple snapshots. */
+        in_use_size = tb->size;
+        alloc_size = tb->reported ? 0 : tb->size;
+    } else {
+        in_use_size = 0;
+        alloc_size = tb->size;
+    }
+
+    PyTuple_SET_ITEM(tb_and_info, 0, traceback_to_tuple(tb));
+    PyTuple_SET_ITEM(tb_and_info, 1, PyLong_FromSize_t(in_use_size));
+    PyTuple_SET_ITEM(tb_and_info, 2, PyLong_FromSize_t(alloc_size));
+    PyTuple_SET_ITEM(tb_and_info, 3, PyLong_FromSize_t(tb->count));
+
+    return tb_and_info;
+}
+
 PyObject*
 memalloc_heap(void)
 {
@@ -351,7 +391,57 @@ memalloc_heap(void)
      * New allocations will go into the secondary freezer.allocs_m map and allocations
      * tracked in allocs_m which are freed will be added to a list to be removed when
      * the profiler is thawed. */
-    PyObject* heap_list = memalloc_heap_map_export(global_heap_tracker.allocs_m);
+
+    /* Calculate total number of samples: live + freed */
+    size_t live_count = memalloc_heap_map_size(global_heap_tracker.allocs_m);
+    size_t freed_count = global_heap_tracker.unreported_samples.count;
+    size_t total_count = live_count + freed_count;
+
+    PyObject* heap_list = PyList_New(total_count);
+    if (heap_list == NULL) {
+        heap_tracker_thaw(&global_heap_tracker);
+        return NULL;
+    }
+
+    int list_index = 0;
+
+    /* First, iterate over live samples using the new iterator API */
+    memalloc_heap_map_iter_t* it = memalloc_heap_map_iter_new(global_heap_tracker.allocs_m);
+    // TODO: handle NULL return
+
+    void* key;
+    traceback_t* tb;
+
+    while (memalloc_heap_map_iter_next(it, &key, &tb)) {
+        PyObject* tb_and_info = memalloc_sample_to_tuple(tb, true);
+
+        PyList_SET_ITEM(heap_list, list_index, tb_and_info);
+        list_index++;
+
+        /* Mark as reported */
+        tb->reported = true;
+    }
+
+    memalloc_heap_map_iter_delete(it);
+
+    /* Second, iterate over freed samples from unreported_samples */
+    for (size_t i = 0; i < global_heap_tracker.unreported_samples.count; i++) {
+        traceback_t* tb = global_heap_tracker.unreported_samples.tab[i];
+
+        PyObject* tb_and_info = memalloc_sample_to_tuple(tb, false);
+
+        PyList_SET_ITEM(heap_list, list_index, tb_and_info);
+        list_index++;
+    }
+
+    /* Free all tracebacks in unreported_samples after reporting them */
+    for (size_t i = 0; i < global_heap_tracker.unreported_samples.count; i++) {
+        if (global_heap_tracker.unreported_samples.tab[i] != NULL) {
+            traceback_free(global_heap_tracker.unreported_samples.tab[i]);
+        }
+    }
+    /* Reset the count to 0 so we can reuse the memory */
+    global_heap_tracker.unreported_samples.count = 0;
 
     heap_tracker_thaw(&global_heap_tracker);
 

ddtrace/profiling/collector/_memalloc_heap_map.c

@@ -82,6 +82,11 @@ typedef struct memalloc_heap_map_t
     HeapSamples map;
 } memalloc_heap_map_t;
 
+typedef struct memalloc_heap_map_iter_t
+{
+    HeapSamples_CIter iter;
+} memalloc_heap_map_iter_t;
+
 memalloc_heap_map_t*
 memalloc_heap_map_new()
 {
@@ -178,3 +183,34 @@ memalloc_heap_map_delete(memalloc_heap_map_t* m)
     HeapSamples_destroy(&m->map);
     free(m);
 }
+
+memalloc_heap_map_iter_t*
+memalloc_heap_map_iter_new(memalloc_heap_map_t* m)
+{
+    memalloc_heap_map_iter_t* it = malloc(sizeof(memalloc_heap_map_iter_t));
+    if (it) {
+        it->iter = HeapSamples_citer(&m->map);
+    }
+    return it;
+}
+
+bool
+memalloc_heap_map_iter_next(memalloc_heap_map_iter_t* it, void** key, traceback_t** tb)
+{
+    const HeapSamples_Entry* e = HeapSamples_CIter_get(&it->iter);
+    if (!e) {
+        return false;
+    }
+    *key = e->key;
+    *tb = e->val;
+    HeapSamples_CIter_next(&it->iter);
+    return true;
+}
+
+void
+memalloc_heap_map_iter_delete(memalloc_heap_map_iter_t* it)
+{
+    if (it) {
+        free(it);
+    }
+}

ddtrace/profiling/collector/_memalloc_heap_map.h

@@ -11,6 +11,8 @@
  */
 typedef struct memalloc_heap_map_t memalloc_heap_map_t;
 
+typedef struct memalloc_heap_map_iter_t memalloc_heap_map_iter_t;
+
 /* Construct an empty map */
 memalloc_heap_map_t*
 memalloc_heap_map_new();
@@ -35,6 +37,19 @@ memalloc_heap_map_remove(memalloc_heap_map_t* m, void* key);
 PyObject*
 memalloc_heap_map_export(memalloc_heap_map_t* m);
 
+/* Create a new iterator for the heap map */
+memalloc_heap_map_iter_t*
+memalloc_heap_map_iter_new(memalloc_heap_map_t* m);
+
+/* Get the next key-value pair from the iterator. Returns true if a pair was found,
+ * false if the iterator is exhausted */
+bool
+memalloc_heap_map_iter_next(memalloc_heap_map_iter_t* it, void** key, traceback_t** tb);
+
+/* Delete the iterator */
+void
+memalloc_heap_map_iter_delete(memalloc_heap_map_iter_t* it);
+
 /* Copy the contents of src into dst, removing the items from src */
 void
 memalloc_heap_map_destructive_copy(memalloc_heap_map_t* dst, memalloc_heap_map_t* src);

ddtrace/profiling/collector/_memalloc_tb.c

@@ -160,7 +160,6 @@ traceback_free(traceback_t* tb)
         Py_DECREF(tb->frames[nframe].filename);
         Py_DECREF(tb->frames[nframe].name);
     }
-    memalloc_debug_gil_release();
     PyMem_RawFree(tb);
 }
 
@@ -250,7 +249,7 @@ memalloc_frame_to_traceback(PyFrameObject* pyframe, uint16_t max_nframe)
 }
 
 traceback_t*
-memalloc_get_traceback(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorDomain domain)
+memalloc_get_traceback(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorDomain domain, size_t weighted_size)
 {
     PyThreadState* tstate = PyThreadState_Get();
 
@@ -271,13 +270,23 @@ memalloc_get_traceback(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocat
     if (traceback == NULL)
         return NULL;
 
-    traceback->size = size;
+    traceback->size = weighted_size;
     traceback->ptr = ptr;
 
     traceback->thread_id = PyThread_get_thread_ident();
 
     traceback->domain = domain;
 
+    traceback->reported = false;
+
+    // Size 0 allocations are legal and we can hypothetically sample them,
+    // e.g. if an allocation during sampling pushes us over the next sampling threshold,
+    // but we can't sample it, so we sample the next allocation which happens to be 0
+    // bytes. Defensively make sure size isn't 0.
+    size = size > 0 ? size : 1;
+    double scaled_count = ((double)weighted_size) / ((double)size);
+    traceback->count = (size_t)scaled_count;
+
     return traceback;
 }
 
@@ -316,9 +325,8 @@ traceback_to_tuple(traceback_t* tb)
         PyTuple_SET_ITEM(stack, nframe, frame_tuple);
     }
 
-    PyObject* tuple = PyTuple_New(3);
+    PyObject* tuple = PyTuple_New(2);
     PyTuple_SET_ITEM(tuple, 0, stack);
-    PyTuple_SET_ITEM(tuple, 1, PyLong_FromUnsignedLong(tb->total_nframe));
-    PyTuple_SET_ITEM(tuple, 2, PyLong_FromUnsignedLong(tb->thread_id));
+    PyTuple_SET_ITEM(tuple, 1, PyLong_FromUnsignedLong(tb->thread_id));
     return tuple;
 }

ddtrace/profiling/collector/_memalloc_tb.h

@@ -37,6 +37,10 @@ typedef struct
     PyMemAllocatorDomain domain;
     /* Thread ID */
     unsigned long thread_id;
+    /* True if this sample has been reported previously */
+    bool reported;
+    /* Count of allocations this sample represents (for scaling) */
+    size_t count;
     /* List of frames, top frame first */
     frame_t frames[1];
 } traceback_t;
@@ -56,7 +60,7 @@ void
 traceback_free(traceback_t* tb);
 
 traceback_t*
-memalloc_get_traceback(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorDomain domain);
+memalloc_get_traceback(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorDomain domain, size_t weighted_size);
 
 PyObject*
 traceback_to_tuple(traceback_t* tb);