merged main

Author: danielsn

ddtrace/internal/datadog/profiling/dd_wrapper/include/sample.hpp

@@ -82,7 +82,8 @@ class Sample
     std::vector<int64_t> values = {};
 
     // Additional metadata
-    int64_t endtime_ns = 0; // end of the event
+    int64_t endtime_ns = 0;         // end of the event
+    bool reverse_locations = false; // whether to reverse locations when exporting/flushing
 
     // Backing memory for string copies
     internal::StringArena string_storage{};
@@ -101,6 +102,8 @@ class Sample
     bool push_release(int64_t lock_time, int64_t count);
     bool push_alloc(int64_t size, int64_t count);
     bool push_heap(int64_t size);
+    void reset_alloc();
+    void reset_heap();
     bool push_gpu_gputime(int64_t time, int64_t count);
     bool push_gpu_memory(int64_t size, int64_t count);
     bool push_gpu_flops(int64_t flops, int64_t count);
@@ -132,8 +135,15 @@ class Sample
                     int64_t line               // for ddog_prof_Location
     );
 
+    // Set whether to reverse locations when exporting/flushing
+    void set_reverse_locations(bool reverse) { reverse_locations = reverse; }
+
     // Flushes the current buffer, clearing it
-    bool flush_sample(bool reverse_locations = false);
+    bool flush_sample();
+
+    // Exports the sample data to the profile without clearing buffers
+    // This is useful when the Sample object is embedded and will be destroyed later
+    bool export_sample();
 
     static ProfileBorrow profile_borrow();
     static void postfork_child();

ddtrace/internal/datadog/profiling/dd_wrapper/src/ddup_interface.cpp

@@ -332,7 +332,8 @@ ddup_flush_sample(Datadog::Sample* sample) // cppcheck-suppress unusedFunction
 void
 ddup_flush_sample_v2(Datadog::Sample* sample) // cppcheck-suppress unusedFunction
 {
-    sample->flush_sample(/*reverse_locations*/ true);
+    sample->set_reverse_locations(true);
+    sample->flush_sample();
 }
 
 void

ddtrace/internal/datadog/profiling/dd_wrapper/src/sample.cpp

@@ -132,12 +132,25 @@ Datadog::Sample::clear_buffers()
     labels.clear();
     locations.clear();
     dropped_frames = 0;
+    reverse_locations = false;
     string_storage.reset();
 }
 
 bool
-Datadog::Sample::flush_sample(bool reverse_locations)
+Datadog::Sample::flush_sample()
 {
+    // Export the sample data (export_sample handles dropped frames)
+    auto ret = export_sample();
+
+    // Clear buffers after exporting
+    clear_buffers();
+    return ret;
+}
+
+bool
+Datadog::Sample::export_sample()
+{
+    // Handle dropped frames by adding them to locations if needed
     if (dropped_frames > 0) {
         const std::string name =
           "<" + std::to_string(dropped_frames) + " frame" + (1 == dropped_frames ? "" : "s") + " omitted>";
@@ -146,6 +159,7 @@ Datadog::Sample::flush_sample(bool reverse_locations)
 
     if (reverse_locations) {
         std::reverse(locations.begin(), locations.end());
+        reverse_locations = false; // Reset after reversing
     }
 
     const ddog_prof_Sample sample = {
@@ -154,9 +168,8 @@ Datadog::Sample::flush_sample(bool reverse_locations)
         .labels = { labels.data(), labels.size() },
     };
 
-    const bool ret = profile_state.collect(sample, endtime_ns);
-    clear_buffers();
-    return ret;
+    // Export to profile without clearing buffers
+    return profile_state.collect(sample, endtime_ns);
 }
 
 bool
@@ -292,6 +305,30 @@ Datadog::Sample::push_heap(int64_t size)
     return false;
 }
 
+void
+Datadog::Sample::reset_alloc()
+{
+    if (0U != (type_mask & SampleType::Allocation)) {
+        const size_t alloc_space_idx = profile_state.val().alloc_space;
+        const size_t alloc_count_idx = profile_state.val().alloc_count;
+        if (alloc_space_idx < values.size() && alloc_count_idx < values.size()) {
+            values[alloc_space_idx] = 0;
+            values[alloc_count_idx] = 0;
+        }
+    }
+}
+
+void
+Datadog::Sample::reset_heap()
+{
+    if (0U != (type_mask & SampleType::Heap)) {
+        const size_t heap_space_idx = profile_state.val().heap_space;
+        if (heap_space_idx < values.size()) {
+            values[heap_space_idx] = 0;
+        }
+    }
+}
+
 bool
 Datadog::Sample::push_gpu_gputime(int64_t time, int64_t count)
 {

ddtrace/profiling/collector/CMakeLists.txt

@@ -55,7 +55,7 @@ if(NOT TARGET Python3::Python)
 endif()
 
 # Source files for the extension
-set(SOURCE_FILES _memalloc.cpp _memalloc_tb.cpp _memalloc_heap.cpp _memalloc_reentrant.cpp _memalloc_heap_map.cpp)
+set(SOURCE_FILES _memalloc.cpp _memalloc_tb.cpp _memalloc_heap.cpp _memalloc_reentrant.cpp)
 
 # Get the extension name from setup.py or use default Note: EXTENSION_NAME from setup.py already includes the full
 # suffix
@@ -78,8 +78,10 @@ endif()
 
 # Include directories
 target_include_directories(
-    ${FULL_EXTENSION_NAME} PRIVATE ${Python3_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}
-                                   ${CMAKE_CURRENT_SOURCE_DIR}/../../internal/datadog/profiling/dd_wrapper/include)
+    ${FULL_EXTENSION_NAME}
+    PRIVATE ${Python3_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}
+            ${CMAKE_CURRENT_SOURCE_DIR}/../../internal/datadog/profiling/dd_wrapper/include
+            ${RUST_GENERATED_HEADERS_DIR})
 
 # Link libraries Python3::Python target might not exist in all build environments (e.g., manylinux) Python modules
 # should use -undefined dynamic_lookup on macOS and not link to libpython on Linux

ddtrace/profiling/collector/_memalloc.cpp

@@ -11,6 +11,9 @@
 #include "_memalloc_tb.h"
 #include "_pymacro.h"
 
+// Ensure profile_state is initialized before creating Sample objects
+#include "../../internal/datadog/profiling/dd_wrapper/include/ddup_interface.hpp"
+
 typedef struct
 {
     PyMemAllocatorEx pymem_allocator_obj;
@@ -40,7 +43,7 @@ memalloc_free(void* ctx, void* ptr)
     if (ptr == NULL)
         return;
 
-    memalloc_heap_untrack(ptr);
+    memalloc_heap_untrack_no_cpython(ptr);
 
     alloc->free(alloc->ctx, ptr);
 }
@@ -80,9 +83,9 @@ memalloc_realloc(void* ctx, void* ptr, size_t new_size)
 {
     memalloc_context_t* memalloc_ctx = (memalloc_context_t*)ctx;
     void* ptr2 = memalloc_ctx->pymem_allocator_obj.realloc(memalloc_ctx->pymem_allocator_obj.ctx, ptr, new_size);
-
+    // TODO(dsn) is the GIL held here? Can there be a race between the realloc and the untrack of the old value?
     if (ptr2) {
-        memalloc_heap_untrack(ptr);
+        memalloc_heap_untrack_no_cpython(ptr);
         memalloc_heap_track(memalloc_ctx->max_nframe, ptr2, new_size, memalloc_ctx->domain);
     }
 
@@ -108,6 +111,11 @@ memalloc_start(PyObject* Py_UNUSED(module), PyObject* args)
         return NULL;
     }
 
+    // Ensure profile_state is initialized before creating Sample objects
+    // This initializes the Sample::profile_state which is required for Sample objects to work correctly
+    // ddup_start() uses std::call_once, so it's safe to call multiple times
+    ddup_start();
+
     char* val = getenv("_DD_MEMALLOC_DEBUG_RNG_SEED");
     if (val) {
         /* NB: we don't bother checking whether val is actually a valid integer.
@@ -144,7 +152,7 @@ memalloc_start(PyObject* Py_UNUSED(module), PyObject* args)
         PyUnicode_InternInPlace(&object_string);
     }
 
-    if (!memalloc_heap_tracker_init((uint32_t)heap_sample_size))
+    if (!memalloc_heap_tracker_init_no_cpython((uint32_t)heap_sample_size))
         return NULL;
 
     PyMemAllocatorEx alloc;
@@ -187,7 +195,7 @@ memalloc_stop(PyObject* Py_UNUSED(module), PyObject* Py_UNUSED(args))
      * or memalloc_heap. The higher-level collector deals with this. */
     PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &global_memalloc_ctx.pymem_allocator_obj);
 
-    memalloc_heap_tracker_deinit();
+    memalloc_heap_tracker_deinit_no_cpython();
 
     /* Finally, we know in-progress sampling won't use the buffer pool, so clear it out */
     traceback_t::deinit();
@@ -201,7 +209,7 @@ PyDoc_STRVAR(memalloc_heap_py__doc__,
              "heap($module, /)\n"
              "--\n"
              "\n"
-             "Get the sampled heap representation.\n");
+             "Export sampled heap allocations to the pprof profile.\n");
 static PyObject*
 memalloc_heap_py(PyObject* Py_UNUSED(module), PyObject* Py_UNUSED(args))
 {
@@ -210,7 +218,8 @@ memalloc_heap_py(PyObject* Py_UNUSED(module), PyObject* Py_UNUSED(args))
         return NULL;
     }
 
-    return memalloc_heap();
+    memalloc_heap_no_cpython();
+    Py_RETURN_NONE;
 }
 
 static PyMethodDef module_methods[] = { { "start", (PyCFunction)memalloc_start, METH_VARARGS, memalloc_start__doc__ },
@@ -234,11 +243,5 @@ PyInit__memalloc(void)
     if (m == NULL)
         return NULL;
 
-    // Initialize the DDFrame namedtuple class
-    // Do this early so we don't have complicated cleanup
-    if (!memalloc_ddframe_class_init()) {
-        return NULL;
-    }
-
     return m;
 }