datadog 2.38.0 → 2.47.0

data/ext/datadog_profiling_native_extension/stack_recorder.c

@@ -9,7 +9,6 @@
 #include "time_helpers.h"
 #include "heap_recorder.h"
 #include "encoded_profile.h"
-#include "collectors_thread_context.h"
 
 // Used to wrap a ddog_prof_Profile in a Ruby object and expose Ruby-level serialization APIs
 // This file implements the native bits of the Datadog::Profiling::StackRecorder class
@@ -178,10 +177,6 @@ typedef struct {
   heap_recorder *heap_recorder;
   bool heap_clean_after_gc_enabled;
 
-  // When set, _native_serialize will call thread_context_collector_on_serialize on this instance
-  // before serializing, so that threads suspended across the whole profile period still get sampled.
-  VALUE thread_context_collector_instance;
-
   pthread_mutex_t mutex_slot_one;
   profile_slot profile_slot_one;
   pthread_mutex_t mutex_slot_two;
@@ -231,6 +226,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 +247,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 +261,20 @@ 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_heap_recorder_set_object_alloc_gen(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE record_id, VALUE alloc_gen);
+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 +299,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 +306,14 @@ 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_heap_recorder_set_object_alloc_gen", _native_heap_recorder_set_object_alloc_gen, 3);
+  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"));
@@ -325,7 +338,6 @@ static VALUE _native_new(VALUE klass) {
   // being leaked.
 
   state->heap_clean_after_gc_enabled = false;
-  state->thread_context_collector_instance = Qnil;
 
   ddog_prof_Slice_SampleType sample_types = {.ptr = all_sample_types, .len = ALL_VALUE_TYPES_COUNT};
 
@@ -397,8 +409,7 @@ static void initialize_profiles(stack_recorder_state *state, ddog_prof_Slice_Sam
 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) {
@@ -521,21 +532,17 @@ static VALUE _native_serialize(DDTRACE_UNUSED VALUE _self, VALUE recorder_instan
   stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
-  ddog_Timespec finish_timestamp = system_epoch_now_timespec();
-  // Need to do this while still holding the Global VM Lock; see comments on method for why
-  serializer_set_start_timestamp_for_next_profile(state, finish_timestamp);
-
-  if (state->thread_context_collector_instance != Qnil) {
-    thread_context_collector_on_serialize(state->thread_context_collector_instance);
-  }
-
-  long heap_iteration_prep_start_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
+  long heap_iteration_prep_start_time_ns = monotonic_wall_time_now_ns();
   // Prepare the iteration on heap recorder we'll be doing outside the GVL. The preparation needs to
   // happen while holding the GVL.
   // NOTE: While rare, it's possible for the GVL to be released inside this function (see comments on `heap_recorder_update`)
   // and thus don't assume this is an "atomic" step -- other threads may get some running time in the meanwhile.
   heap_recorder_prepare_iteration(state->heap_recorder);
-  long heap_iteration_prep_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE) - heap_iteration_prep_start_time_ns;
+  long heap_iteration_prep_time_ns = monotonic_wall_time_now_ns() - heap_iteration_prep_start_time_ns;
+
+  ddog_Timespec finish_timestamp = system_epoch_now_timespec();
+  // Need to do this while still holding the Global VM Lock; see comments on method for why
+  serializer_set_start_timestamp_for_next_profile(state, finish_timestamp);
 
   // We'll release the Global VM Lock while we're calling serialize, so that the Ruby VM can continue to work while this
   // is pending
@@ -605,7 +612,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 +649,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 +699,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
@@ -783,7 +800,7 @@ static void build_heap_profile_without_gvl(stack_recorder_state *state, profile_
 static void *call_serialize_without_gvl(void *call_args) {
   call_serialize_without_gvl_arguments *args = (call_serialize_without_gvl_arguments *) call_args;
 
-  long serialize_no_gvl_start_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
+  long serialize_no_gvl_start_time_ns = monotonic_wall_time_now_ns();
 
   profile_slot *slot_now_inactive = serializer_flip_active_and_inactive_slots(args->state);
   args->slot = slot_now_inactive;
@@ -791,13 +808,13 @@ static void *call_serialize_without_gvl(void *call_args) {
   // Now that we have the inactive profile with all but heap samples, lets fill it with heap data
   // without needing to race with the active sampler
   build_heap_profile_without_gvl(args->state, args->slot);
-  args->heap_profile_build_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE) - serialize_no_gvl_start_time_ns;
+  args->heap_profile_build_time_ns = monotonic_wall_time_now_ns() - serialize_no_gvl_start_time_ns;
 
   // Note: The profile gets reset by the serialize call
   args->result = ddog_prof_Profile_serialize(&args->slot->profile, &args->slot->start_timestamp, &args->finish_timestamp);
   args->advance_gen_result = ddog_prof_ManagedStringStorage_advance_gen(args->state->string_storage);
   args->serialize_ran = true;
-  args->serialize_no_gvl_time_ns = long_max_of(0, monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE) - serialize_no_gvl_start_time_ns);
+  args->serialize_no_gvl_time_ns = long_max_of(0, monotonic_wall_time_now_ns() - serialize_no_gvl_start_time_ns);
 
   return NULL; // Unused
 }
@@ -897,7 +914,7 @@ static VALUE test_slot_mutex_state(VALUE recorder_instance, int slot) {
 }
 
 static ddog_Timespec system_epoch_now_timespec(void) {
-  long now_ns = system_epoch_time_now_ns(RAISE_ON_FAILURE);
+  long now_ns = system_epoch_time_now_ns();
   return (ddog_Timespec) {.seconds = now_ns / SECONDS_AS_NS(1), .nanoseconds = now_ns % SECONDS_AS_NS(1)};
 }
 
@@ -938,15 +955,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 +1032,39 @@ 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, obj_id);
+  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);
+
+  heap_recorder_testonly_exhaust_record_ids(state->heap_recorder);
+  return Qtrue;
+}
+
+static VALUE _native_heap_recorder_set_object_alloc_gen(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance, VALUE record_id, VALUE alloc_gen) {
+  ENFORCE_TYPE(record_id, T_FIXNUM);
+  ENFORCE_TYPE(alloc_gen, T_FIXNUM);
+
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  heap_recorder_testonly_set_object_alloc_gen(state->heap_recorder, FIX2LONG(record_id), NUM2SIZET(alloc_gen));
+  return Qtrue;
 }
 
 static VALUE _native_heap_recorder_reset_last_update(DDTRACE_UNUSED VALUE _self, VALUE recorder_instance) {
@@ -1042,8 +1076,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,18 +1187,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;
 }
-
-void recorder_install_on_serialize(VALUE recorder_instance, VALUE thread_context_collector_instance) {
-  stack_recorder_state *state;
-  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
-
-  state->thread_context_collector_instance = enforce_thread_context_collector_instance(thread_context_collector_instance);
-}

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,6 @@ 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_install_on_serialize(VALUE recorder_instance, VALUE thread_context_collector_instance);
+void recorder_commit_heap_recordings_may_lose_gvl(VALUE recorder_instance);
+void recorder_heap_update_may_lose_gvl(VALUE recorder_instance);
 VALUE enforce_recorder_instance(VALUE object);

data/ext/datadog_profiling_native_extension/time_helpers.c

@@ -28,10 +28,8 @@ long monotonic_to_system_epoch_ns(monotonic_to_system_epoch_state *state, long m
     (state->delta_to_epoch_ns + monotonic_wall_time_ns > state->system_epoch_ns_reference + SECONDS_AS_NS(60));
 
   if (reference_needs_update) {
-    state->system_epoch_ns_reference = system_epoch_time_now_ns(RAISE_ON_FAILURE);
-    long current_monotonic_wall_time_ns = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
-
-    state->delta_to_epoch_ns = state->system_epoch_ns_reference - current_monotonic_wall_time_ns;
+    state->system_epoch_ns_reference = system_epoch_time_now_ns();
+    state->delta_to_epoch_ns = state->system_epoch_ns_reference - monotonic_wall_time_now_ns();
   }
 
   return state->delta_to_epoch_ns + monotonic_wall_time_ns;

data/ext/datadog_profiling_native_extension/time_helpers.h

@@ -1,33 +1,33 @@
 #pragma once
 
-#include <stdbool.h>
+#include <ruby.h>
 #include <errno.h>
+#include <stdbool.h>
 #include <time.h>
 
 #include "extconf.h"
-#include "ruby_helpers.h"
-
-#define SECONDS_AS_NS(value) (value * 1000 * 1000 * 1000L)
-#define MILLIS_AS_NS(value) (value * 1000 * 1000L)
-#define MICROS_AS_NS(value) (value * 1000L)
 
-typedef enum { RAISE_ON_FAILURE, DO_NOT_RAISE_ON_FAILURE } raise_on_failure_setting;
+#define SECONDS_AS_NS(value) ((value) * 1000L * 1000L * 1000L)
+#define MILLIS_AS_NS(value)  ((value) * 1000L * 1000L)
+#define MICROS_AS_NS(value)  ((value) * 1000L)
 
 #define INVALID_TIME -1
 
 typedef struct {
+  // INVALID_TIME until the first conversion, and the only field used to decide if the state needs to be (re)computed
   long system_epoch_ns_reference;
   long delta_to_epoch_ns;
 } monotonic_to_system_epoch_state;
 
-#define MONOTONIC_TO_SYSTEM_EPOCH_INITIALIZER {.system_epoch_ns_reference = INVALID_TIME, .delta_to_epoch_ns = INVALID_TIME}
+#define MONOTONIC_TO_SYSTEM_EPOCH_INITIALIZER {.system_epoch_ns_reference = INVALID_TIME}
 
-static inline long retrieve_clock_as_ns(clockid_t clock_id, raise_on_failure_setting raise_on_failure) {
+// clock_id MUST be one of CLOCK_REALTIME/CLOCK_MONOTONIC_FOR_PROFILING/CLOCK_MONOTONIC_COARSE_FOR_PROFILING
+// as for those we're guaranteed to never see failures here unless something is really off
+static inline long retrieve_clock_as_ns(clockid_t clock_id) {
   struct timespec clock_value;
 
-  if (clock_gettime(clock_id, &clock_value) != 0) {
-    if (raise_on_failure == RAISE_ON_FAILURE) ENFORCE_SUCCESS_GVL(errno);
-    return 0;
+  if (RB_UNLIKELY(clock_gettime(clock_id, &clock_value) != 0)) {
+    rb_bug("[ddtrace] clock_gettime(%d) unexpectedly failed (errno=%d) but it should always succeed, libc/kernel bug?", (int) clock_id, errno);
   }
 
   return clock_value.tv_nsec + SECONDS_AS_NS(clock_value.tv_sec);
@@ -42,8 +42,8 @@ static inline long retrieve_clock_as_ns(clockid_t clock_id, raise_on_failure_set
   #define CLOCK_MONOTONIC_COARSE_FOR_PROFILING CLOCK_MONOTONIC_COARSE
 #endif
 
-static inline long monotonic_wall_time_now_ns(raise_on_failure_setting raise_on_failure) { return retrieve_clock_as_ns(CLOCK_MONOTONIC_FOR_PROFILING, raise_on_failure); }
-static inline long system_epoch_time_now_ns(raise_on_failure_setting raise_on_failure)   { return retrieve_clock_as_ns(CLOCK_REALTIME,  raise_on_failure); }
+static inline long monotonic_wall_time_now_ns(void) { return retrieve_clock_as_ns(CLOCK_MONOTONIC_FOR_PROFILING); }
+static inline long system_epoch_time_now_ns(void)   { return retrieve_clock_as_ns(CLOCK_REALTIME); }
 
 // Coarse instants use CLOCK_MONOTONIC_COARSE on Linux which is expected to provide resolution in the millisecond range:
 // https://docs.redhat.com/en/documentation/red_hat_enterprise_linux_for_real_time/7/html/reference_guide/sect-posix_clocks#Using_clock_getres_to_compare_clock_resolution
@@ -56,7 +56,7 @@ typedef struct {
 static inline coarse_instant to_coarse_instant(long timestamp_ns) { return (coarse_instant) {.timestamp_ns = timestamp_ns}; }
 
 static inline coarse_instant monotonic_coarse_wall_time_now_ns(void) {
-  return to_coarse_instant(retrieve_clock_as_ns(CLOCK_MONOTONIC_COARSE_FOR_PROFILING, DO_NOT_RAISE_ON_FAILURE));
+  return to_coarse_instant(retrieve_clock_as_ns(CLOCK_MONOTONIC_COARSE_FOR_PROFILING));
 }
 
 long monotonic_to_system_epoch_ns(monotonic_to_system_epoch_state *state, long monotonic_wall_time_ns);

data/ext/libdatadog_api/datadog_ruby_common.c

@@ -6,19 +6,27 @@
 static ID telemetry_message_id;
 
 void raise_unexpected_type(VALUE value, const char *value_name, const char *type_name, const char *file, int line, const char *function_name) {
-  rb_exc_raise(
-    rb_exc_new_str(
-      rb_eTypeError,
-      rb_sprintf("wrong argument %"PRIsVALUE" for '%s' (expected a %s) at %s:%d:in `%s'",
-        rb_inspect(value),
-        value_name,
-        type_name,
-        file,
-        line,
-        function_name
-      )
+  // All of the below come from `ENFORCE_TYPE` and are constant for their call-site
+  VALUE simple_message = rb_sprintf("wrong argument for '%s' (expected a %s) at %s:%d:in `%s'",
+    value_name,
+    type_name,
+    file,
+    line,
+    function_name
+  );
+  VALUE exception = rb_exc_new_str(
+    rb_eTypeError,
+    rb_sprintf("wrong argument %"PRIsVALUE" for '%s' (expected a %s) at %s:%d:in `%s'",
+      rb_inspect(value),
+      value_name,
+      type_name,
+      file,
+      line,
+      function_name
     )
   );
+  rb_ivar_set(exception, telemetry_message_id, simple_message);
+  rb_exc_raise(exception);
 }
 
 // Raises an exception with separate telemetry-safe and detailed messages.

data/ext/libdatadog_api/datadog_ruby_common.h

@@ -5,6 +5,13 @@
 #include <ruby.h>
 #include <datadog/common.h>
 
+// All recommended backward compatible macros for keyword argument functions can be found here:
+// https://github.com/ruby/ruby/blob/master/doc/extension.rdoc#defining-backward-compatible-macros-for-keyword-argument-functions
+#ifndef RB_PASS_KEYWORDS
+/* Only define macros on Ruby <2.7 */
+#define rb_funcallv_kw(o, m, c, v, kw) rb_funcallv(o, m, c, v)
+#endif
+
 // Must be called once during initialization
 void datadog_ruby_common_init(void);
 
@@ -90,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

@@ -1,7 +1,11 @@
 #include <stdbool.h>
+#include <string.h>
 
 #include "datadog_ruby_common.h"
 
+// rb_iseq is CRuby-specific
+#ifndef TRUFFLERUBY
+
 // Prototypes for Ruby functions declared in internal Ruby headers.
 // rb_iseqw_new wraps an internal iseq pointer into a Ruby-visible
 // RubyVM::InstructionSequence object.
@@ -9,11 +13,19 @@ VALUE rb_iseqw_new(const void *iseq);
 // rb_iseqw_to_iseq unwraps a RubyVM::InstructionSequence object back
 // to its internal iseq pointer.
 const void *rb_iseqw_to_iseq(VALUE iseqw);
+#ifdef HAVE_RB_IMEMO_NAME
+const char *rb_imemo_name(int imemo_type);
+#endif
 int rb_objspace_internal_object_p(VALUE obj);
 void rb_objspace_each_objects(
     int (*callback)(void *start, void *end, size_t stride, void *data),
     void *data);
 
+#ifdef HAVE_SHIFTED_IMEMO_MASK // Ruby 4.1+
+  #define IMEMO_MASK (FL_USER0 | FL_USER1 | FL_USER2 | FL_USER3 | FL_USER4)
+#else
+  #define IMEMO_MASK 0x0f
+#endif
 #define IMEMO_TYPE_ISEQ 7
 
 // The ID value of the string "mesg" which is used in Ruby source as
@@ -27,6 +39,16 @@ 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) {
+  #ifdef HAVE_SHIFTED_IMEMO_MASK // Ruby 4.1+
+    return (RBASIC(imemo)->flags & IMEMO_MASK) >> FL_USHIFT;
+  #else
+    return (RBASIC(imemo)->flags >> FL_USHIFT) & IMEMO_MASK;
+  #endif
+}
+
 // 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;
@@ -186,7 +208,20 @@ static VALUE iseq_type(DDTRACE_UNUSED VALUE _self, VALUE iseq_val) {
 }
 #endif
 
+static VALUE _native_imemo_mask(DDTRACE_UNUSED VALUE _self) {
+  return INT2NUM(IMEMO_MASK);
+}
+
 void di_init(VALUE datadog_module) {
+  #ifdef HAVE_RB_IMEMO_NAME
+    const char *iseq_name = rb_imemo_name(IMEMO_TYPE_ISEQ);
+    if (strcmp(iseq_name, "iseq") != 0) {
+      rb_raise(rb_eRuntimeError,
+        "IMEMO_TYPE_ISEQ self-test failed: rb_imemo_name(%d) returned '%s' instead of 'iseq'",
+        IMEMO_TYPE_ISEQ, iseq_name);
+    }
+  #endif
+
   id_mesg = rb_intern("mesg");
   id_datadog_di_in_probe = rb_intern("datadog_di_in_probe");
 
@@ -197,7 +232,12 @@ void di_init(VALUE datadog_module) {
   rb_define_singleton_method(di_module, "enter_probe", enter_probe, 0);
   rb_define_singleton_method(di_module, "leave_probe", leave_probe, 0);
   rb_define_singleton_method(di_module, "hash?", is_hash, 1);
-#ifdef HAVE_RB_ISEQ_TYPE
-  rb_define_singleton_method(di_module, "iseq_type", iseq_type, 1);
-#endif
+  #ifdef HAVE_RB_ISEQ_TYPE
+    rb_define_singleton_method(di_module, "iseq_type", iseq_type, 1);
+  #endif
+
+  VALUE testing_module = rb_define_module_under(di_module, "Testing");
+  rb_define_singleton_method(testing_module, "_native_imemo_mask", _native_imemo_mask, 0);
 }
+
+#endif // TRUFFLERUBY