From 52144c879386efe8aafe7e43cbe622756aa7bf74 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 6 Nov 2025 14:15:22 -0500 Subject: [PATCH] perf(profiling): Only sample allocations where the pointer hashes to a given value --- ddtrace/profiling/collector/_memalloc_heap.c | 72 ++++++++++++++++---- 1 file changed, 57 insertions(+), 15 deletions(-) diff --git a/ddtrace/profiling/collector/_memalloc_heap.c b/ddtrace/profiling/collector/_memalloc_heap.c index 5616fe0c0ab..7fc63c7ae3f 100644 --- a/ddtrace/profiling/collector/_memalloc_heap.c +++ b/ddtrace/profiling/collector/_memalloc_heap.c @@ -73,6 +73,9 @@ typedef struct memalloc_heap_map_t* allocs_m; /* Bytes allocated since the last sample was collected */ uint64_t allocated_memory; + /* The factor that controls how often pointers are eligible to sample + * TODO more description here */ + uintptr_t sample_mask; /* True if we are exporting the current heap profile */ bool frozen; /* Contains the ongoing heap allocation/deallocation while frozen */ @@ -87,12 +90,20 @@ typedef struct /* Debug guard to assert that GIL-protected critical sections are maintained * while accessing the profiler's state */ memalloc_gil_debug_check_t gil_guard; + // Whether we are currently in a run where we should try to collect a sample + bool collect_sample; } heap_tracker_t; static heap_tracker_t global_heap_tracker; +static bool +is_eligible_to_sample_no_cpython(void* ptr) +{ + return (uintptr_t)ptr % global_heap_tracker.sample_mask == 0; +} + static uint32_t -heap_tracker_next_sample_size(uint32_t sample_size) +heap_tracker_next_sample_size_no_cpython(uint32_t sample_size) { /* We want to draw a sampling target from an exponential distribution with average sample_size. We use the standard technique of inverse transform @@ -114,10 +125,12 @@ heap_tracker_init(heap_tracker_t* heap_tracker) 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->sample_mask = 1; heap_tracker->allocated_memory = 0; heap_tracker->frozen = false; heap_tracker->sample_size = 0; heap_tracker->current_sample_size = 0; + heap_tracker->collect_sample = false; memalloc_gil_debug_check_init(&heap_tracker->gil_guard); } @@ -189,8 +202,11 @@ void memalloc_heap_tracker_init(uint32_t sample_size) { heap_tracker_init(&global_heap_tracker); + // TODO take as a param + // Should be prime + global_heap_tracker.sample_mask = 5; global_heap_tracker.sample_size = sample_size; - global_heap_tracker.current_sample_size = heap_tracker_next_sample_size(sample_size); + global_heap_tracker.current_sample_size = heap_tracker_next_sample_size_no_cpython(sample_size); } void @@ -218,6 +234,13 @@ memalloc_heap_untrack_no_cpython(heap_tracker_t* heap_tracker, void* ptr) MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard); return NULL; } + + // If this is not a pointer that could be sampled, then no need to do anything with it. + if (!is_eligible_to_sample_no_cpython(ptr)) { + MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard); + return NULL; + } + if (!heap_tracker->frozen) { traceback_t* tb = memalloc_heap_map_remove(heap_tracker->allocs_m, ptr); if (tb && !tb->reported) { @@ -259,7 +282,7 @@ memalloc_heap_untrack(void* ptr) * shared state, and must be called with the GIL held and without making any C * Python API calls. */ static bool -memalloc_heap_should_sample_no_cpython(heap_tracker_t* heap_tracker, size_t size) +memalloc_heap_should_sample_no_cpython(heap_tracker_t* heap_tracker, void* ptr, size_t size) { MEMALLOC_GIL_DEBUG_CHECK_ACQUIRE(&heap_tracker->gil_guard); /* Heap tracking is disabled */ @@ -270,12 +293,6 @@ memalloc_heap_should_sample_no_cpython(heap_tracker_t* heap_tracker, size_t size heap_tracker->allocated_memory += size; - /* Check if we have enough sample or not */ - if (heap_tracker->allocated_memory < heap_tracker->current_sample_size) { - MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard); - return false; - } - if (memalloc_heap_map_size(heap_tracker->allocs_m) + memalloc_heap_map_size(heap_tracker->freezer.allocs_m) > TRACEBACK_ARRAY_MAX_COUNT) { /* TODO(nick) this is vestigial from the original array-based @@ -287,10 +304,39 @@ memalloc_heap_should_sample_no_cpython(heap_tracker_t* heap_tracker, size_t size return false; } + /* Check if we have enough sample or not. + * If so, start a run of possible sample collection */ + if (heap_tracker->allocated_memory >= heap_tracker->current_sample_size) { + heap_tracker->collect_sample = true; + } + + /* If we are in a situation where we could collect samples, and the pointer is elibible, + * collect a sample */ + if (heap_tracker->collect_sample && is_eligible_to_sample_no_cpython(ptr)) { + heap_tracker->collect_sample = false; + MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard); + return true; + } + + /* Otherwise, leave collect_sample untouched, and try next time */ MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard); return true; } +/* Updates the threshold for */ +static void +memalloc_heap_update_threshold_no_cpython(heap_tracker_t* heap_tracker) +{ + MEMALLOC_GIL_DEBUG_CHECK_ACQUIRE(&heap_tracker->gil_guard); + /* Reset the counter to 0 */ + heap_tracker->allocated_memory = 0; + + /* Compute the new target sample size */ + heap_tracker->current_sample_size = heap_tracker_next_sample_size_no_cpython(heap_tracker->sample_size); + + MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard); +} + /* Track an allocation that we decided to sample. This updates shared state and * must be called with the GIL held and without making any C Python API calls. * If the allocation could not be added because the profiler was stopped, @@ -312,11 +358,7 @@ memalloc_heap_add_sample_no_cpython(heap_tracker_t* heap_tracker, traceback_t* t old = memalloc_heap_map_insert(heap_tracker->allocs_m, tb->ptr, tb); } - /* Reset the counter to 0 */ - heap_tracker->allocated_memory = 0; - - /* Compute the new target sample size */ - heap_tracker->current_sample_size = heap_tracker_next_sample_size(heap_tracker->sample_size); + memalloc_heap_update_threshold_no_cpython(heap_tracker); MEMALLOC_GIL_DEBUG_CHECK_RELEASE(&heap_tracker->gil_guard); return old; @@ -326,7 +368,7 @@ memalloc_heap_add_sample_no_cpython(heap_tracker_t* heap_tracker, traceback_t* t void memalloc_heap_track(uint16_t max_nframe, void* ptr, size_t size, PyMemAllocatorDomain domain) { - if (!memalloc_heap_should_sample_no_cpython(&global_heap_tracker, size)) { + if (!memalloc_heap_should_sample_no_cpython(&global_heap_tracker, ptr, size)) { return; }