Heap profiling for ruby 4.x - Prototype

The idea is to delay the time at which we record object ids.
Once we are outside of the allocation code path, we can request the object ID.

Author: r1viollet

ext/datadog_profiling_native_extension/collectors_cpu_and_wall_time_worker.c

@@ -16,6 +16,7 @@
 #include "private_vm_api_access.h"
 #include "setup_signal_handler.h"
 #include "time_helpers.h"
+#include "heap_recorder.h"
 
 // Used to trigger the execution of Collectors::ThreadContext, which implements all of the sampling logic
 // itself; this class only implements the "when to do it" part.
@@ -89,6 +90,9 @@ unsigned int MAX_ALLOC_WEIGHT = 10000;
   static rb_postponed_job_handle_t sample_from_postponed_job_handle;
   static rb_postponed_job_handle_t after_gc_from_postponed_job_handle;
   static rb_postponed_job_handle_t after_gvl_running_from_postponed_job_handle;
+  #ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+    static rb_postponed_job_handle_t finalize_heap_allocation_from_postponed_job_handle;
+  #endif
 #endif
 
 // Contains state for a single CpuAndWallTimeWorker instance
@@ -237,6 +241,9 @@ static VALUE rescued_after_gvl_running_from_postponed_job(VALUE self_instance);
 static VALUE _native_gvl_profiling_hook_active(DDTRACE_UNUSED VALUE self, VALUE instance);
 static inline void during_sample_enter(cpu_and_wall_time_worker_state* state);
 static inline void during_sample_exit(cpu_and_wall_time_worker_state* state);
+#ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+static void finalize_heap_allocation_from_postponed_job(DDTRACE_UNUSED void *_unused);
+#endif
 
 // We're using `on_newobj_event` function with `rb_add_event_hook2`, which requires in its public signature a function
 // with signature `rb_event_hook_func_t` which doesn't match `on_newobj_event`.
@@ -283,11 +290,17 @@ void collectors_cpu_and_wall_time_worker_init(VALUE profiling_module) {
     sample_from_postponed_job_handle = rb_postponed_job_preregister(unused_flags, sample_from_postponed_job, NULL);
     after_gc_from_postponed_job_handle = rb_postponed_job_preregister(unused_flags, after_gc_from_postponed_job, NULL);
     after_gvl_running_from_postponed_job_handle = rb_postponed_job_preregister(unused_flags, after_gvl_running_from_postponed_job, NULL);
+    #ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+    finalize_heap_allocation_from_postponed_job_handle = rb_postponed_job_preregister(unused_flags, finalize_heap_allocation_from_postponed_job, NULL);
+    #endif
 
     if (
       sample_from_postponed_job_handle == POSTPONED_JOB_HANDLE_INVALID ||
       after_gc_from_postponed_job_handle == POSTPONED_JOB_HANDLE_INVALID ||
       after_gvl_running_from_postponed_job_handle == POSTPONED_JOB_HANDLE_INVALID
+      #ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+      || finalize_heap_allocation_from_postponed_job_handle == POSTPONED_JOB_HANDLE_INVALID
+      #endif
     ) {
       rb_raise(rb_eRuntimeError, "Failed to register profiler postponed jobs (got POSTPONED_JOB_HANDLE_INVALID)");
     }
@@ -1222,6 +1235,12 @@ static void on_newobj_event(DDTRACE_UNUSED VALUE unused1, DDTRACE_UNUSED void *u
   state->stats.allocation_sampled++;
 
   during_sample_exit(state);
+
+  #ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+  // On Ruby 4+, we need to trigger a postponed job to finalize the heap allocation recording.
+  // During on_newobj_event, we can't safely call rb_obj_id(), so we defer it until the event completes.
+  rb_postponed_job_trigger(finalize_heap_allocation_from_postponed_job_handle);
+  #endif
 }
 
 static void disable_tracepoints(cpu_and_wall_time_worker_state *state) {
@@ -1404,6 +1423,27 @@ static VALUE _native_resume_signals(DDTRACE_UNUSED VALUE self) {
   }
 #endif
 
+#ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+// This postponed job callback is used to finalize heap allocation recordings on Ruby 4+.
+// During on_newobj_event, calling rb_obj_id() is unsafe because it mutates the object.
+// So we defer getting the object_id until after the event completes.
+static void finalize_heap_allocation_from_postponed_job(DDTRACE_UNUSED void *_unused) {
+  cpu_and_wall_time_worker_state *state = active_sampler_instance_state;
+
+  if (state == NULL) return;
+
+  if (!ddtrace_rb_ractor_main_p()) {
+    return;
+  }
+
+  // Get the heap_recorder from the thread_context_collector
+  heap_recorder *recorder = thread_context_collector_get_heap_recorder(state->thread_context_collector_instance);
+  if (recorder == NULL) return;
+
+  heap_recorder_finalize_pending_recordings(recorder);
+}
+#endif
+
 static inline void during_sample_enter(cpu_and_wall_time_worker_state* state) {
   // Tell the compiler it's not allowed to reorder the `during_sample` flag with anything that happens after.
   //

ext/datadog_profiling_native_extension/collectors_thread_context.c

@@ -1338,6 +1338,14 @@ VALUE enforce_thread_context_collector_instance(VALUE object) {
   return object;
 }
 
+#ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+heap_recorder* thread_context_collector_get_heap_recorder(VALUE self_instance) {
+  thread_context_collector_state *state;
+  TypedData_Get_Struct(self_instance, thread_context_collector_state, &thread_context_collector_typed_data, state);
+  return get_heap_recorder_from_stack_recorder(state->recorder_instance);
+}
+#endif
+
 // This method exists only to enable testing Datadog::Profiling::Collectors::ThreadContext behavior using RSpec.
 // It SHOULD NOT be used for other purposes.
 static VALUE _native_stats(DDTRACE_UNUSED VALUE _self, VALUE collector_instance) {

ext/datadog_profiling_native_extension/collectors_thread_context.h

@@ -4,6 +4,7 @@
 #include <stdbool.h>
 
 #include "gvl_profiling_helper.h"
+#include "heap_recorder.h"
 
 void thread_context_collector_sample(
   VALUE self_instance,
@@ -18,6 +19,12 @@ void thread_context_collector_on_gc_start(VALUE self_instance);
 __attribute__((warn_unused_result)) bool thread_context_collector_on_gc_finish(VALUE self_instance);
 VALUE enforce_thread_context_collector_instance(VALUE object);
 
+#ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+// Get the heap_recorder from the thread_context_collector instance.
+// Used to access pending recordings for deferred finalization.
+heap_recorder* thread_context_collector_get_heap_recorder(VALUE self_instance);
+#endif
+
 #ifndef NO_GVL_INSTRUMENTATION
   typedef enum {
     ON_GVL_RUNNING_UNKNOWN, // Thread is not known, it may not even be from the current Ractor

ext/datadog_profiling_native_extension/extconf.rb

@@ -144,6 +144,11 @@ def skip_building_extension!(reason)
 # On Ruby 4, we can't ask the object_id from IMEMOs (https://github.com/ruby/ruby/pull/13347)
 $defs << "-DNO_IMEMO_OBJECT_ID" unless RUBY_VERSION < "4"
 
+# On Ruby 4, we need to defer calling rb_obj_id during heap allocation recording
+# because it's not safe to mutate objects during the newobj tracepoint
+# (see https://bugs.ruby-lang.org/issues/21710)
+$defs << "-DDEFERRED_HEAP_ALLOCATION_RECORDING" unless RUBY_VERSION < "4"
+
 # This symbol is exclusively visible on certain Ruby versions: 2.6 to 3.2, as well as 3.4 (but not 4.0+)
 # It's only used to get extra information about an object when a failure happens, so it's a "very nice to have" but not
 # actually required for correct behavior of the profiler.

ext/datadog_profiling_native_extension/heap_recorder.c

@@ -79,6 +79,18 @@ static void object_record_free(heap_recorder*, object_record*);
 static VALUE object_record_inspect(heap_recorder*, object_record*);
 static object_record SKIPPED_RECORD = {0};
 
+#ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+// A pending recording is used to defer the object_id call on Ruby 4+
+// where calling rb_obj_id during on_newobj_event is unsafe.
+typedef struct {
+  VALUE object_ref;
+  heap_record *heap_record;
+  live_object_data object_data;
+} pending_recording;
+
+#define MAX_PENDING_RECORDINGS 64
+#endif
+
 struct heap_recorder {
   // Config
   // Whether the recorder should try to determine approximate sizes for tracked objects.
@@ -130,6 +142,17 @@ struct heap_recorder {
   // Data for a heap recording that was started but not yet ended
   object_record *active_recording;
 
+  #ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+  // Pending recordings that need to be finalized after on_newobj_event completes.
+  // On Ruby 4+, we can't call rb_obj_id during the newobj event, so we store the
+  // VALUE reference here and finalize it via a postponed job.
+  pending_recording pending_recordings[MAX_PENDING_RECORDINGS];
+  uint pending_recordings_count;
+  // Temporary storage for the recording in progress, used between start and end
+  VALUE active_deferred_object;
+  live_object_data active_deferred_object_data;
+  #endif
+
   // Reusable arrays, implementing a flyweight pattern for things like iteration
   #define REUSABLE_LOCATIONS_SIZE MAX_FRAMES_LIMIT
   ddog_prof_Location *reusable_locations;
@@ -294,6 +317,12 @@ void heap_recorder_after_fork(heap_recorder *heap_recorder) {
   heap_recorder->stats_lifetime = (struct stats_lifetime) {0};
 }
 
+#ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+// Sentinel marker for active_recording indicating we're using deferred recording.
+// Only used for pointer comparison, no data is read from or written to this.
+static object_record DEFERRED_RECORD = {.obj_id = -1};
+#endif
+
 void start_heap_allocation_recording(heap_recorder *heap_recorder, VALUE new_obj, unsigned int weight, ddog_CharSlice alloc_class) {
   if (heap_recorder == NULL) {
     return;
@@ -314,13 +343,29 @@ void start_heap_allocation_recording(heap_recorder *heap_recorder, VALUE new_obj
       // directly OR because the GVL got released in the middle of an update), let's skip this sample as well.
       // See notes on `heap_recorder_update` for details.
       || heap_recorder->updating
+      #ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+      // Skip if we've hit the pending recordings limit or if there's already a deferred object being recorded
+      || heap_recorder->pending_recordings_count >= MAX_PENDING_RECORDINGS
+      || heap_recorder->active_deferred_object != Qnil
+      #endif
     ) {
     heap_recorder->active_recording = &SKIPPED_RECORD;
     return;
   }
 
   heap_recorder->num_recordings_skipped = 0;
 
+  #ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+  // On Ruby 4+, we can't call rb_obj_id during on_newobj_event as it mutates the object.
+  // Instead, we store the VALUE reference and will get the object_id later via a postponed job.
+  heap_recorder->active_deferred_object = new_obj;
+  heap_recorder->active_deferred_object_data = (live_object_data) {
+    .weight = weight * heap_recorder->sample_rate,
+    .class = intern_or_raise(heap_recorder->string_storage, alloc_class),
+    .alloc_gen = rb_gc_count(),
+  };
+  heap_recorder->active_recording = &DEFERRED_RECORD;
+  #else
   VALUE ruby_obj_id = rb_obj_id(new_obj);
   if (!FIXNUM_P(ruby_obj_id)) {
     rb_raise(rb_eRuntimeError, "Detected a bignum object id. These are not supported by heap profiling.");
@@ -335,6 +380,7 @@ void start_heap_allocation_recording(heap_recorder *heap_recorder, VALUE new_obj
       .alloc_gen = rb_gc_count(),
     }
   );
+  #endif
 }
 
 // end_heap_allocation_recording_with_rb_protect gets called while the stack_recorder is holding one of the profile
@@ -383,6 +429,23 @@ static VALUE end_heap_allocation_recording(VALUE protect_args) {
     return Qnil;
   }
 
+  #ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+  if (active_recording == &DEFERRED_RECORD) {
+    // On Ruby 4+, we defer the object_id call. Store the recording in the pending list.
+    if (heap_recorder->pending_recordings_count < MAX_PENDING_RECORDINGS) {
+      heap_record *heap_record = get_or_create_heap_record(heap_recorder, locations);
+      heap_record->num_tracked_objects++; // Pre-increment since we're going to commit this later
+
+      pending_recording *pending = &heap_recorder->pending_recordings[heap_recorder->pending_recordings_count++];
+      pending->object_ref = heap_recorder->active_deferred_object;
+      pending->heap_record = heap_record;
+      pending->object_data = heap_recorder->active_deferred_object_data;
+    }
+    heap_recorder->active_deferred_object = Qnil;
+    return Qnil;
+  }
+  #endif
+
   heap_record *heap_record = get_or_create_heap_record(heap_recorder, locations);
 
   // And then commit the new allocation.
@@ -399,6 +462,89 @@ void heap_recorder_update_young_objects(heap_recorder *heap_recorder) {
   heap_recorder_update(heap_recorder, /* full_update: */ false);
 }
 
+#ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+bool heap_recorder_has_pending_recordings(heap_recorder *heap_recorder) {
+  if (heap_recorder == NULL) {
+    return false;
+  }
+  return heap_recorder->pending_recordings_count > 0;
+}
+
+bool heap_recorder_finalize_pending_recordings(heap_recorder *heap_recorder) {
+  if (heap_recorder == NULL) {
+    return false;
+  }
+
+  uint count = heap_recorder->pending_recordings_count;
+  if (count == 0) {
+    return false;
+  }
+
+  for (uint i = 0; i < count; i++) {
+    pending_recording *pending = &heap_recorder->pending_recordings[i];
+
+    VALUE obj = pending->object_ref;
+
+    // Check if the object is still valid (it might have been GC'd between
+    // the allocation and this finalization)
+    if (obj == Qnil || !RTEST(obj)) {
+      // Object is gone, decrement the pre-incremented count and cleanup
+      pending->heap_record->num_tracked_objects--;
+      cleanup_heap_record_if_unused(heap_recorder, pending->heap_record);
+      unintern_or_raise(heap_recorder, pending->object_data.class);
+      continue;
+    }
+
+    // Now it's safe to get the object_id
+    VALUE ruby_obj_id = rb_obj_id(obj);
+    if (!FIXNUM_P(ruby_obj_id)) {
+      // Bignum object id - not supported, skip this recording
+      pending->heap_record->num_tracked_objects--;
+      cleanup_heap_record_if_unused(heap_recorder, pending->heap_record);
+      unintern_or_raise(heap_recorder, pending->object_data.class);
+      continue;
+    }
+
+    long obj_id = FIX2LONG(ruby_obj_id);
+
+    // Create the object record now that we have the object_id
+    object_record *record = object_record_new(obj_id, pending->heap_record, pending->object_data);
+
+    // Commit the recording
+    int existing_error = st_update(heap_recorder->object_records, record->obj_id, update_object_record_entry, (st_data_t) record);
+    if (existing_error) {
+      // Duplicate object_id - this shouldn't happen normally, but handle it gracefully
+      pending->heap_record->num_tracked_objects--;
+      cleanup_heap_record_if_unused(heap_recorder, pending->heap_record);
+      object_record_free(heap_recorder, record);
+    }
+  }
+
+  heap_recorder->pending_recordings_count = 0;
+  return true;
+}
+
+// Mark pending recordings to prevent GC from collecting the objects
+// while they're waiting to be finalized
+void heap_recorder_mark_pending_recordings(heap_recorder *heap_recorder) {
+  if (heap_recorder == NULL) {
+    return;
+  }
+
+  for (uint i = 0; i < heap_recorder->pending_recordings_count; i++) {
+    VALUE obj = heap_recorder->pending_recordings[i].object_ref;
+    if (obj != Qnil) {
+      rb_gc_mark(obj);
+    }
+  }
+
+  // Also mark the active deferred object if it's set
+  if (heap_recorder->active_deferred_object != Qnil) {
+    rb_gc_mark(heap_recorder->active_deferred_object);
+  }
+}
+#endif
+
 // NOTE: This function needs and assumes it gets called with the GVL being held.
 //       But importantly **some of the operations inside `st_object_record_update` may cause a thread switch**,
 //       so we can't assume a single update happens in a single "atomic" step -- other threads may get some running time

ext/datadog_profiling_native_extension/heap_recorder.h

@@ -123,6 +123,20 @@ int end_heap_allocation_recording_with_rb_protect(heap_recorder *heap_recorder,
 // these objects quicker) and hopefully reduces tail latency (because there's less objects at serialization time to check).
 void heap_recorder_update_young_objects(heap_recorder *heap_recorder);
 
+#ifdef DEFERRED_HEAP_ALLOCATION_RECORDING
+// Finalize any pending heap allocation recordings by getting their object IDs.
+// This should be called via a postponed job, after the on_newobj_event has completed.
+// Returns true if there were pending recordings that were finalized.
+bool heap_recorder_finalize_pending_recordings(heap_recorder *heap_recorder);
+
+// Check if there are any pending recordings waiting to be finalized.
+bool heap_recorder_has_pending_recordings(heap_recorder *heap_recorder);
+
+// Mark pending recordings to prevent GC from collecting the objects
+// while they're waiting to be finalized.
+void heap_recorder_mark_pending_recordings(heap_recorder *heap_recorder);
+#endif
+
 // Update the heap recorder to reflect the latest state of the VM and prepare internal structures
 // for efficient iteration.
 //