datadog 2.41.0 → 2.42.0

data/ext/datadog_profiling_native_extension/private_vm_api_access.h

@@ -2,9 +2,9 @@
 
 #include <stdbool.h>
 
-// The private_vm_api_access.c includes the RUBY_MJIT_HEADER which replaces and conflicts with any other Ruby headers;
-// so we use PRIVATE_VM_API_ACCESS_SKIP_RUBY_INCLUDES to be able to include private_vm_api_access.h on that file
-// without also dragging the incompatible includes
+// The private_vm_api_access.c includes private VM headers (via datadog-ruby_core_source) which replace and conflict
+// with any other Ruby headers; so we use PRIVATE_VM_API_ACCESS_SKIP_RUBY_INCLUDES to be able to include
+// private_vm_api_access.h on that file without also dragging the incompatible includes
 #ifndef PRIVATE_VM_API_ACCESS_SKIP_RUBY_INCLUDES
   #include <ruby/thread_native.h>
   #include <ruby/vm.h>
@@ -86,8 +86,20 @@ VALUE current_fiber_for(VALUE thread);
 
 void self_test_current_fiber_for(void);
 
+// Returns the module name without allocating, or nil if the module is anonymous (i.e., Module#name is nil).
+// Never returns an empty String (returns nil instead if it the name would be empty).
+// Also consider ddtrace_permanent_mod_name().
+//
+// Use this instead of rb_class2name() or rb_mod_name() which can trigger allocations
+VALUE ddtrace_alloc_free_rb_mod_name(VALUE mod);
+// Same as `ddtrace_alloc_free_rb_mod_name` but additionally returns nil (filters out)
+// if the module does not have a permanent name (e.g. `#<Module:0x0123>::Foo`).
+VALUE ddtrace_permanent_mod_name(VALUE mod);
+
 ssize_t ddtrace_location_label(const rb_callable_method_entry_t *cme, const rb_iseq_t *iseq, char *buf, size_t buf_size);
 VALUE ddtrace_location_base_label(const rb_callable_method_entry_t *cme, const rb_iseq_t *iseq);
 void* ddtrace_cme_cfunc_func(const rb_callable_method_entry_t *cme);
 const char *ddtrace_cme_original_method_name(const rb_callable_method_entry_t *cme);
 
+// Returns true for internal objects (like T_IMEMO/T_ICLASS/etc) and "hidden" objects (rb_class_of(obj) == 0)
+bool ddtrace_is_internal_object_p(VALUE obj);

data/ext/datadog_profiling_native_extension/profiling.c

@@ -21,6 +21,7 @@ void collectors_dynamic_sampling_rate_init(VALUE profiling_module);
 void collectors_idle_sampling_helper_init(VALUE profiling_module);
 void collectors_stack_init(VALUE profiling_module);
 void collectors_thread_context_init(VALUE profiling_module);
+void collectors_heap_recorder_init(void);
 void encoded_profile_init(VALUE profiling_module);
 void http_transport_init(VALUE profiling_module);
 void stack_recorder_init(VALUE profiling_module);
@@ -73,6 +74,7 @@ void DDTRACE_EXPORT Init_datadog_profiling_native_extension(void) {
   collectors_idle_sampling_helper_init(profiling_module);
   collectors_stack_init(profiling_module);
   collectors_thread_context_init(profiling_module);
+  collectors_heap_recorder_init();
   encoded_profile_init(profiling_module);
   http_transport_init(profiling_module);
   stack_recorder_init(profiling_module);

data/ext/datadog_profiling_native_extension/ruby_helpers.c

@@ -7,16 +7,10 @@
 
 // The following global variables are initialized at startup to save expensive lookups later.
 // They are not expected to be mutated outside of init.
-static VALUE module_object_space = Qnil;
-static ID _id2ref_id = Qnil;
 static ID inspect_id = Qnil;
 static ID to_s_id = Qnil;
 
 void ruby_helpers_init(void) {
-  rb_global_variable(&module_object_space);
-
-  module_object_space = rb_const_get(rb_cObject, rb_intern("ObjectSpace"));
-  _id2ref_id = rb_intern("_id2ref");
   inspect_id = rb_intern("inspect");
   to_s_id = rb_intern("to_s");
 }
@@ -113,40 +107,6 @@ void private_raise_enforce_syserr(
   }
 }
 
-static VALUE _id2ref(VALUE obj_id) {
-  // Call ::ObjectSpace._id2ref natively. It will raise if the id is no longer valid
-  return rb_funcall(module_object_space, _id2ref_id, 1, obj_id);
-}
-
-static VALUE _id2ref_failure(DDTRACE_UNUSED VALUE _unused1, DDTRACE_UNUSED VALUE _unused2) {
-  return Qfalse;
-}
-
-// See notes on header for important details
-bool ruby_ref_from_id(VALUE obj_id, VALUE *value) {
-  // Call ::ObjectSpace._id2ref natively. It will raise if the id is no longer valid
-  // so we need to call it via rb_rescue2
-  // TODO: Benchmark rb_rescue2 vs rb_protect here
-  VALUE result = rb_rescue2(
-    _id2ref,
-    obj_id,
-    _id2ref_failure,
-    Qnil,
-    rb_eRangeError, // rb_eRangeError is the error used to flag invalid ids
-    0 // Required by API to be the last argument
-  );
-
-  if (result == Qfalse) {
-    return false;
-  }
-
-  if (value != NULL) {
-    (*value) = result;
-  }
-
-  return true;
-}
-
 // Not part of public headers but is externed from Ruby
 size_t rb_obj_memsize_of(VALUE obj);
 
@@ -162,16 +122,9 @@ size_t rb_obj_memsize_of(VALUE obj);
 size_t ruby_obj_memsize_of(VALUE obj) {
   switch (rb_type(obj)) {
     case T_OBJECT:
-    // On Ruby 4.0, computing the size of a class/module/iclass seems to not be safe: `rb_obj_memsize_of` walks the
-    // per-namespace class extensions (`rb_class_classext_foreach` -> `classext_memsize` -> `rb_id_table_memsize`)
-    // and can crash the VM when called on objects tracked by heap profiling.
-    // See https://github.com/DataDog/dd-trace-rb/issues/5936. Class objects contribute negligibly to heap size,
-    // so we skip them (fall through to the `default` branch, returning 0) rather than risk a SIGSEGV.
-    #ifndef NO_SAFE_CLASS_MEMSIZE
     case T_MODULE:
     case T_CLASS:
     case T_ICLASS:
-    #endif
     case T_STRING:
     case T_ARRAY:
     case T_HASH:
@@ -198,36 +151,6 @@ size_t ruby_obj_memsize_of(VALUE obj) {
   }
 }
 
-// Inspired by rb_class_of but without actually returning classes or potentially doing assertions
-static bool ruby_is_obj_with_class(VALUE obj) {
-  if (!RB_SPECIAL_CONST_P(obj)) {
-    return true;
-  }
-  if (obj == RUBY_Qfalse) {
-    return true;
-  }
-  else if (obj == RUBY_Qnil) {
-    return true;
-  }
-  else if (obj == RUBY_Qtrue) {
-    return true;
-  }
-  else if (RB_FIXNUM_P(obj)) {
-    return true;
-  }
-  else if (RB_STATIC_SYM_P(obj)) {
-    return true;
-  }
-  else if (RB_FLONUM_P(obj)) {
-    return true;
-  }
-
-  return false;
-}
-
-// This function is not present in the VM headers, but is a public symbol that can be invoked.
-int rb_objspace_internal_object_p(VALUE obj);
-
 #ifdef NO_RB_OBJ_INFO
   const char* safe_object_info(DDTRACE_UNUSED VALUE obj) { return "(No rb_obj_info for current Ruby)"; }
 #else
@@ -239,8 +162,7 @@ int rb_objspace_internal_object_p(VALUE obj);
 #endif
 
 VALUE ruby_safe_inspect(VALUE obj) {
-  if (!ruby_is_obj_with_class(obj))       return rb_str_new_cstr("(Not an object)");
-  if (rb_objspace_internal_object_p(obj)) return rb_sprintf("(VM Internal, %s)", safe_object_info(obj));
+  if (ddtrace_is_internal_object_p(obj))  return rb_sprintf("(VM Internal, %s)", safe_object_info(obj));
   // @ivoanjo: I saw crashes on Ruby 3.1.4 when trying to #inspect matchdata objects. I'm not entirely sure why this
   // is needed, but since we only use this method for debug purposes I put in this alternative and decided not to
   // dig deeper.

data/ext/datadog_profiling_native_extension/ruby_helpers.h

@@ -79,13 +79,6 @@ NORETURN(void private_raise_enforce_syserr(
   const char *function_name
 ));
 
-// Native wrapper to get an object ref from an id. Returns true on success and
-// writes the ref to the value pointer parameter if !NULL. False if id doesn't
-// reference a valid object (in which case value is not changed).
-//
-// Note: GVL can be released and other threads may get to run before this method returns
-bool ruby_ref_from_id(VALUE obj_id, VALUE *value);
-
 // Native wrapper to get the approximate/estimated current size of the passed
 // object.
 size_t ruby_obj_memsize_of(VALUE obj);

data/ext/datadog_profiling_native_extension/stack_recorder.c

@@ -231,6 +231,18 @@ typedef struct {
   bool serialize_ran;
 } call_serialize_without_gvl_arguments;
 
+typedef struct {
+  // Set by caller
+  stack_recorder_state *state;
+  locked_profile_slot active_slot;
+  ddog_prof_Slice_Location locations;
+  sample_values values;
+  sample_labels labels;
+
+  // Set by callee
+  ddog_prof_Profile_Result result;
+} record_sample_arguments;
+
 static VALUE _native_new(VALUE klass);
 static void initialize_slot_concurrency_control(stack_recorder_state *state);
 static void stack_recorder_typed_data_mark(void *data);
@@ -240,6 +252,8 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
 static VALUE _native_serialize(VALUE self, VALUE recorder_instance);
 static VALUE ruby_time_from(ddog_Timespec ddprof_time);
 static void *call_serialize_without_gvl(void *call_args);
+static VALUE record_sample_locked(VALUE record_sample_arguments_as_value);
+static VALUE record_sample_unlock_ensure(VALUE record_sample_arguments_as_value);
 static locked_profile_slot sampler_lock_active_profile(stack_recorder_state *state);
 static void sampler_unlock_active_profile(locked_profile_slot active_slot);
 static profile_slot* serializer_flip_active_and_inactive_slots(stack_recorder_state *state);
@@ -252,18 +266,19 @@ static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE recorder_
 static void serializer_set_start_timestamp_for_next_profile(stack_recorder_state *state, ddog_Timespec start_time);
 static VALUE _native_record_endpoint(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE local_root_span_id, VALUE endpoint);
 static void reset_profile_slot(profile_slot *slot, ddog_Timespec start_timestamp);
-static VALUE _native_track_object(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE new_obj, VALUE weight, VALUE alloc_class);
 static VALUE _native_start_fake_slow_heap_serialization(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE _native_end_fake_slow_heap_serialization(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE _native_debug_heap_recorder(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE _native_stats(DDTRACE_UNUSED VALUE self, VALUE instance);
 static VALUE build_profile_stats(profile_slot *slot, long serialization_time_ns, long heap_iteration_prep_time_ns, long heap_profile_build_time_ns);
-static VALUE _native_is_object_recorded(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE object_id);
+static VALUE _native_is_object_recorded(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE record_id);
+static VALUE _native_record_id_for(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE obj);
 static VALUE _native_heap_recorder_reset_last_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
-static VALUE _native_recorder_after_gc_step(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
+static VALUE _native_heap_recorder_exhaust_record_ids(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
+static VALUE _native_recorder_heap_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 static VALUE _native_benchmark_intern(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE string, VALUE times, VALUE use_all);
 static VALUE _native_test_managed_string_storage_produces_valid_profiles(DDTRACE_UNUSED VALUE _self);
-static VALUE _native_finalize_pending_heap_recordings(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
+static VALUE _native_commit_heap_recordings(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance);
 
 void stack_recorder_init(VALUE profiling_module) {
   VALUE stack_recorder_class = rb_define_class_under(profiling_module, "StackRecorder", rb_cObject);
@@ -288,7 +303,6 @@ void stack_recorder_init(VALUE profiling_module) {
   rb_define_singleton_method(testing_module, "_native_slot_one_mutex_locked?", _native_is_slot_one_mutex_locked, 1);
   rb_define_singleton_method(testing_module, "_native_slot_two_mutex_locked?", _native_is_slot_two_mutex_locked, 1);
   rb_define_singleton_method(testing_module, "_native_record_endpoint", _native_record_endpoint, 3);
-  rb_define_singleton_method(testing_module, "_native_track_object", _native_track_object, 4);
   rb_define_singleton_method(testing_module, "_native_start_fake_slow_heap_serialization",
       _native_start_fake_slow_heap_serialization, 1);
   rb_define_singleton_method(testing_module, "_native_end_fake_slow_heap_serialization",
@@ -296,11 +310,13 @@ void stack_recorder_init(VALUE profiling_module) {
   rb_define_singleton_method(testing_module, "_native_debug_heap_recorder",
       _native_debug_heap_recorder, 1);
   rb_define_singleton_method(testing_module, "_native_is_object_recorded?", _native_is_object_recorded, 2);
+  rb_define_singleton_method(testing_module, "_native_record_id_for", _native_record_id_for, 2);
   rb_define_singleton_method(testing_module, "_native_heap_recorder_reset_last_update", _native_heap_recorder_reset_last_update, 1);
-  rb_define_singleton_method(testing_module, "_native_recorder_after_gc_step", _native_recorder_after_gc_step, 1);
+  rb_define_singleton_method(testing_module, "_native_heap_recorder_exhaust_record_ids", _native_heap_recorder_exhaust_record_ids, 1);
+  rb_define_singleton_method(testing_module, "_native_recorder_heap_update", _native_recorder_heap_update, 1);
   rb_define_singleton_method(testing_module, "_native_benchmark_intern", _native_benchmark_intern, 4);
   rb_define_singleton_method(testing_module, "_native_test_managed_string_storage_produces_valid_profiles", _native_test_managed_string_storage_produces_valid_profiles, 0);
-  rb_define_singleton_method(testing_module, "_native_finalize_pending_heap_recordings", _native_finalize_pending_heap_recordings, 1);
+  rb_define_singleton_method(testing_module, "_native_commit_heap_recordings", _native_commit_heap_recordings, 1);
 
   ok_symbol = ID2SYM(rb_intern_const("ok"));
   error_symbol = ID2SYM(rb_intern_const("error"));
@@ -398,7 +414,7 @@ static void stack_recorder_typed_data_mark(void *state_ptr) {
   stack_recorder_state *state = (stack_recorder_state *) state_ptr;
 
   rb_gc_mark(state->thread_context_collector_instance);
-  heap_recorder_mark_pending_recordings(state->heap_recorder);
+  heap_recorder_mark(state->heap_recorder);
 }
 
 static void stack_recorder_typed_data_free(void *state_ptr) {
@@ -605,7 +621,28 @@ void record_sample(VALUE recorder_instance, ddog_prof_Slice_Location locations,
   stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
-  locked_profile_slot active_slot = sampler_lock_active_profile(state);
+  record_sample_arguments args = {
+    .state = state,
+    .active_slot = sampler_lock_active_profile(state),
+    .locations = locations,
+    .values = values,
+    .labels = labels,
+  };
+
+  // NOTE: The heap recorder is allowed to raise exceptions if something's wrong. We use `rb_ensure` to make sure we
+  // properly unlock the active slot mutex on our way out. Otherwise, this would later lead to deadlocks (since the
+  // active slot mutex is not expected to be locked forever).
+  rb_ensure(record_sample_locked, (VALUE) &args, record_sample_unlock_ensure, (VALUE) &args);
+
+  if (args.result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
+    raise_error(rb_eArgError, "Failed to record sample: %"PRIsVALUE, get_error_details_and_drop(&args.result.err));
+  }
+}
+
+static VALUE record_sample_locked(VALUE record_sample_arguments_as_value) {
+  record_sample_arguments *args = (record_sample_arguments *) record_sample_arguments_as_value;
+  stack_recorder_state *state = args->state;
+  sample_values values = args->values;
 
   // Note: We initialize this array to have ALL_VALUE_TYPES_COUNT but only tell libdatadog to use the first
   // state->enabled_values_count values. This simplifies handling disabled value types -- we still put them on the
@@ -621,50 +658,39 @@ void record_sample(VALUE recorder_instance, ddog_prof_Slice_Location locations,
   metric_values[position_for[ALLOC_SAMPLES_UNSCALED_VALUE_ID]] = values.alloc_samples_unscaled;
   metric_values[position_for[TIMELINE_VALUE_ID]]      = values.timeline_wall_time_ns;
 
-  if (values.heap_sample) {
-    // If we got an allocation sample end the heap allocation recording to commit the heap sample.
-    // FIXME: Heap sampling currently has to be done in 2 parts because the construction of locations is happening
-    //        very late in the allocation-sampling path (which is shared with the cpu sampling path). This can
-    //        be fixed with some refactoring but for now this leads to a less impactful change.
-    //
-    // NOTE: The heap recorder is allowed to raise exceptions if something's wrong. But we also need to handle it
-    // on this side to make sure we properly unlock the active slot mutex on our way out. Otherwise, this would
-    // later lead to deadlocks (since the active slot mutex is not expected to be locked forever).
-    int exception_state = end_heap_allocation_recording_with_rb_protect(state->heap_recorder, locations);
-    if (exception_state) {
-      sampler_unlock_active_profile(active_slot);
-      rb_jump_tag(exception_state);
-    }
+  if (values.heap_sample != NULL) {
+    heap_recorder_record_allocation(
+      state->heap_recorder,
+      values.heap_sample->new_object,
+      values.alloc_samples,
+      values.heap_sample->alloc_class,
+      args->locations,
+      &values.heap_sample->out_needs_commit
+    );
   }
 
-  ddog_prof_Profile_Result result = ddog_prof_Profile_add(
-    &active_slot.data->profile,
+  args->result = ddog_prof_Profile_add(
+    &args->active_slot.data->profile,
     (ddog_prof_Sample) {
-      .locations = locations,
+      .locations = args->locations,
       .values = (ddog_Slice_I64) {.ptr = metric_values, .len = state->enabled_values_count},
-      .labels = labels.labels
+      .labels = args->labels.labels
     },
     /* To disable end_timestamp_ns to get aggregated profiles in pprof, replace the below with `0` */
-    labels.end_timestamp_ns
+    args->labels.end_timestamp_ns
   );
 
-  active_slot.data->stats.recorded_samples++;
+  args->active_slot.data->stats.recorded_samples++;
 
-  sampler_unlock_active_profile(active_slot);
-
-  if (result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
-    raise_error(rb_eArgError, "Failed to record sample: %"PRIsVALUE, get_error_details_and_drop(&result.err));
-  }
+  return Qnil; // Unused
 }
 
-// Returns needs_after_allocation: true whenever an after_sample callback is required
-bool track_object(VALUE recorder_instance, VALUE new_object, unsigned int sample_weight, ddog_CharSlice alloc_class) {
-  stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
-  // FIXME: Heap sampling currently has to be done in 2 parts because the construction of locations is happening
-  //        very late in the allocation-sampling path (which is shared with the cpu sampling path). This can
-  //        be fixed with some refactoring but for now this leads to a less impactful change.
-  return start_heap_allocation_recording(state->heap_recorder, new_object, sample_weight, alloc_class);
+static VALUE record_sample_unlock_ensure(VALUE record_sample_arguments_as_value) {
+  record_sample_arguments *args = (record_sample_arguments *) record_sample_arguments_as_value;
+
+  sampler_unlock_active_profile(args->active_slot);
+
+  return Qnil; // Unused
 }
 
 void record_endpoint(VALUE recorder_instance, uint64_t local_root_span_id, ddog_CharSlice endpoint) {
@@ -682,18 +708,18 @@ void record_endpoint(VALUE recorder_instance, uint64_t local_root_span_id, ddog_
   }
 }
 
-void recorder_after_gc_step(VALUE recorder_instance) {
+void recorder_heap_update_may_lose_gvl(VALUE recorder_instance) {
   stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   if (state->heap_clean_after_gc_enabled) heap_recorder_update_young_objects(state->heap_recorder);
 }
 
-void recorder_after_sample(VALUE recorder_instance) {
+void recorder_commit_heap_recordings_may_lose_gvl(VALUE recorder_instance) {
   stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
-  heap_recorder_finalize_pending_recordings(state->heap_recorder);
+  heap_recorder_commit_recordings_may_lose_gvl(state->heap_recorder);
 }
 
 #define MAX_LEN_HEAP_ITERATION_ERROR_MSG 256
@@ -938,15 +964,6 @@ static VALUE _native_record_endpoint(DDTRACE_UNUSED VALUE _self, VALUE recorder_
   return Qtrue;
 }
 
-static VALUE _native_track_object(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE new_obj, VALUE weight, VALUE alloc_class) {
-  ENFORCE_TYPE(weight, T_FIXNUM);
-  bool needs_after_allocation = track_object(recorder_instance, new_obj, NUM2UINT(weight), char_slice_from_ruby_string(alloc_class));
-
-  // We could instead choose to automatically trigger the after allocation here; yet, it seems kinda nice to keep it manual for
-  // the tests so we can pull on each lever separately and observe "the sausage being made" in steps
-  return needs_after_allocation ? Qtrue : Qfalse;
-}
-
 static void reset_profile_slot(profile_slot *slot, ddog_Timespec start_timestamp) {
   ddog_prof_Profile_Result reset_result = ddog_prof_Profile_reset(&slot->profile);
   if (reset_result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
@@ -1024,13 +1041,28 @@ static VALUE build_profile_stats(profile_slot *slot, long serialization_time_ns,
   return stats_as_hash;
 }
 
-static VALUE _native_is_object_recorded(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE obj_id) {
-  ENFORCE_TYPE(obj_id, T_FIXNUM);
+static VALUE _native_is_object_recorded(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE record_id) {
+  ENFORCE_TYPE(record_id, T_FIXNUM);
+
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  return heap_recorder_testonly_is_object_recorded(state->heap_recorder, FIX2LONG(record_id));
+}
+
+static VALUE _native_record_id_for(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE obj) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
+  return heap_recorder_testonly_record_id_for(state->heap_recorder, obj);
+}
+
+static VALUE _native_heap_recorder_exhaust_record_ids(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
   stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
-  return heap_recorder_testonly_is_object_recorded(state->heap_recorder, obj_id);
+  heap_recorder_testonly_exhaust_record_ids(state->heap_recorder);
+  return Qtrue;
 }
 
 static VALUE _native_heap_recorder_reset_last_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
@@ -1042,8 +1074,8 @@ static VALUE _native_heap_recorder_reset_last_update(DDTRACE_UNUSED VALUE _self,
   return Qtrue;
 }
 
-static VALUE _native_recorder_after_gc_step(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
-  recorder_after_gc_step(recorder_instance);
+static VALUE _native_recorder_heap_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
+  recorder_heap_update_may_lose_gvl(recorder_instance);
   return Qtrue;
 }
 
@@ -1153,11 +1185,11 @@ static VALUE _native_test_managed_string_storage_produces_valid_profiles(DDTRACE
   return rb_ary_new_from_args(2, encoded_pprof_1, encoded_pprof_2);
 }
 
-static VALUE _native_finalize_pending_heap_recordings(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
+static VALUE _native_commit_heap_recordings(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
   stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
-  heap_recorder_finalize_pending_recordings(state->heap_recorder);
+  heap_recorder_commit_recordings_may_lose_gvl(state->heap_recorder);
 
   return Qtrue;
 }

data/ext/datadog_profiling_native_extension/stack_recorder.h

@@ -3,14 +3,23 @@
 #include <datadog/profiling.h>
 #include <ruby.h>
 
+typedef struct {
+  // In: The object being allocated, and its class. (Its weight is the `alloc_samples` in `sample_values` below.)
+  VALUE new_object;
+  ddog_CharSlice alloc_class;
+  // Out: Set by `record_sample` to whether a `recorder_commit_heap_recordings_may_lose_gvl` callback is needed
+  bool out_needs_commit;
+} heap_sample_values;
+
 typedef struct {
   int64_t cpu_time_ns;
   int64_t wall_time_ns;
   uint32_t cpu_or_wall_samples;
   uint32_t alloc_samples;
   uint32_t alloc_samples_unscaled;
-  bool heap_sample;
   int64_t timeline_wall_time_ns;
+  // When not NULL, this sample also gets recorded for heap profiling
+  heap_sample_values *heap_sample;
 } sample_values;
 
 typedef struct {
@@ -26,8 +35,7 @@ typedef struct {
 
 void record_sample(VALUE recorder_instance, ddog_prof_Slice_Location locations, sample_values values, sample_labels labels);
 void record_endpoint(VALUE recorder_instance, uint64_t local_root_span_id, ddog_CharSlice endpoint);
-__attribute__((warn_unused_result)) bool track_object(VALUE recorder_instance, VALUE new_object, unsigned int sample_weight, ddog_CharSlice alloc_class);
-void recorder_after_sample(VALUE recorder_instance);
-void recorder_after_gc_step(VALUE recorder_instance);
+void recorder_commit_heap_recordings_may_lose_gvl(VALUE recorder_instance);
+void recorder_heap_update_may_lose_gvl(VALUE recorder_instance);
 void recorder_install_on_serialize(VALUE recorder_instance, VALUE thread_context_collector_instance);
 VALUE enforce_recorder_instance(VALUE object);

data/ext/libdatadog_api/datadog_ruby_common.h

@@ -97,13 +97,3 @@ static inline VALUE get_error_details_and_drop(ddog_Error *error) {
 // Returns the amount of characters written to string (which are necessarily
 // bounded by capacity - 1 since the string will be null-terminated).
 size_t read_ddogerr_string_and_drop(ddog_Error *error, char *string, size_t capacity);
-
-#define IMEMO_MASK 0x0f
-
-// Returns the imemo type of an imemo object.
-// This mask is the same between Ruby 2.5 and 3.3-preview3. Furthermore, the intention of this method is to be used
-// to call `rb_imemo_name` which correctly handles invalid numbers so even if the mask changes in the future, at most
-// we'll get incorrect results (and never a VM crash)
-static inline int ddtrace_imemo_type(VALUE imemo) {
-  return (RBASIC(imemo)->flags >> FL_USHIFT) & IMEMO_MASK;
-}

data/ext/libdatadog_api/di.c

@@ -14,6 +14,10 @@ void rb_objspace_each_objects(
     int (*callback)(void *start, void *end, size_t stride, void *data),
     void *data);
 
+// The mask is unchanged between Ruby 2.5 and 4.1-dev, but the `enum imemo_type` values do get
+// renumbered from time to time (e.g. Ruby 4.0 moved imemo_callinfo from 11 to 10), and 4.1-dev already uses all
+// 16 values the mask allows. Thus IMEMO_TYPE_ISEQ must be re-checked when adding support for a new Ruby.
+#define IMEMO_MASK 0x0f
 #define IMEMO_TYPE_ISEQ 7
 
 // The ID value of the string "mesg" which is used in Ruby source as
@@ -27,6 +31,12 @@ static ID id_mesg;
 // method probes on Thread#[] / Thread#[]= cannot intercept guard reads/writes.
 static ID id_datadog_di_in_probe;
 
+// Returns the imemo type of an imemo object, that is `imemo_type()` from CRuby's internal/imemo.h.
+// The caller must have already checked the object is a T_IMEMO, otherwise the result is meaningless.
+static inline int ddtrace_imemo_type(VALUE imemo) {
+  return (RBASIC(imemo)->flags >> FL_USHIFT) & IMEMO_MASK;
+}
+
 // Returns whether the argument is an IMEMO of type ISEQ.
 static bool ddtrace_imemo_iseq_p(VALUE v) {
   return rb_objspace_internal_object_p(v) && RB_TYPE_P(v, T_IMEMO) && ddtrace_imemo_type(v) == IMEMO_TYPE_ISEQ;

data/ext/libdatadog_api/extconf.rb

@@ -98,6 +98,9 @@ EXTENSION_NAME = "libdatadog_api.#{RUBY_VERSION[/\d+.\d+/]}_#{RUBY_PLATFORM}".fr
 
 have_func("rb_iseq_type")
 
+$defs << "-DHAVE_RUBY_THREAD_STORAGE_API" if RUBY_VERSION >= "3.3"
+$defs << "-DHAVE_RUBY_THREAD_HAS_GVL_P" if have_func("ruby_thread_has_gvl_p")
+
 create_makefile(EXTENSION_NAME)
 
 # rubocop:enable Style/GlobalVars

data/ext/libdatadog_api/init.c

@@ -5,6 +5,7 @@
 #include "crashtracker.h"
 #include "feature_flags.h"
 #include "library_config.h"
+#include "otel_thread_context.h"
 #include "process_discovery.h"
 #include "trace_exporter.h"
 
@@ -28,4 +29,5 @@ void DDTRACE_EXPORT Init_libdatadog_api(void) {
 
   VALUE tracing_module = rb_define_module_under(datadog_module, "Tracing");
   trace_exporter_init(tracing_module);
+  otel_thread_context_init(tracing_module);
 }