datadog 2.30.0 → 2.43.0

data/ext/datadog_profiling_native_extension/stack_recorder.c

@@ -9,6 +9,7 @@
 #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
@@ -177,6 +178,10 @@ 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;
@@ -226,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);
@@ -235,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);
@@ -247,18 +266,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);
@@ -283,7 +304,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",
@@ -291,11 +311,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"));
@@ -320,6 +343,7 @@ 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};
 
@@ -391,7 +415,8 @@ 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;
 
-  heap_recorder_mark_pending_recordings(state->heap_recorder);
+  rb_gc_mark(state->thread_context_collector_instance);
+  heap_recorder_mark(state->heap_recorder);
 }
 
 static void stack_recorder_typed_data_free(void *state_ptr) {
@@ -416,20 +441,16 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   if (options == Qnil) options = rb_hash_new();
 
   VALUE recorder_instance = rb_hash_fetch(options, ID2SYM(rb_intern("self_instance")));
-  VALUE cpu_time_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("cpu_time_enabled")));
   VALUE alloc_samples_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("alloc_samples_enabled")));
   VALUE heap_samples_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_samples_enabled")));
   VALUE heap_size_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_size_enabled")));
   VALUE heap_sample_every = rb_hash_fetch(options, ID2SYM(rb_intern("heap_sample_every")));
-  VALUE timeline_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("timeline_enabled")));
   VALUE heap_clean_after_gc_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_clean_after_gc_enabled")));
 
-  ENFORCE_BOOLEAN(cpu_time_enabled);
   ENFORCE_BOOLEAN(alloc_samples_enabled);
   ENFORCE_BOOLEAN(heap_samples_enabled);
   ENFORCE_BOOLEAN(heap_size_enabled);
   ENFORCE_TYPE(heap_sample_every, T_FIXNUM);
-  ENFORCE_BOOLEAN(timeline_enabled);
   ENFORCE_BOOLEAN(heap_clean_after_gc_enabled);
 
   stack_recorder_state *state;
@@ -440,11 +461,9 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   heap_recorder_set_sample_rate(state->heap_recorder, NUM2INT(heap_sample_every));
 
   uint8_t requested_values_count = ALL_VALUE_TYPES_COUNT -
-    (cpu_time_enabled == Qtrue ? 0 : 1) -
     (alloc_samples_enabled == Qtrue? 0 : 2) -
     (heap_samples_enabled == Qtrue ? 0 : 1) -
-    (heap_size_enabled == Qtrue ? 0 : 1) -
-    (timeline_enabled == Qtrue ? 0 : 1);
+    (heap_size_enabled == Qtrue ? 0 : 1);
 
   if (requested_values_count == ALL_VALUE_TYPES_COUNT) return Qtrue; // Nothing to do, this is the default
 
@@ -466,12 +485,13 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_WALL_TIME;
   state->position_for[WALL_TIME_VALUE_ID] = next_enabled_pos++;
 
-  if (cpu_time_enabled == Qtrue) {
-    enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_CPU_TIME;
-    state->position_for[CPU_TIME_VALUE_ID] = next_enabled_pos++;
-  } else {
-    state->position_for[CPU_TIME_VALUE_ID] = next_disabled_pos++;
-  }
+  // CPU_TIME is always enabled
+  enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_CPU_TIME;
+  state->position_for[CPU_TIME_VALUE_ID] = next_enabled_pos++;
+
+  // TIMELINE is always enabled
+  enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_TIMELINE;
+  state->position_for[TIMELINE_VALUE_ID] = next_enabled_pos++;
 
   if (alloc_samples_enabled == Qtrue) {
     enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_ALLOC_SAMPLES;
@@ -506,13 +526,6 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
     state->heap_recorder = NULL;
   }
 
-  if (timeline_enabled == Qtrue) {
-    enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_TIMELINE;
-    state->position_for[TIMELINE_VALUE_ID] = next_enabled_pos++;
-  } else {
-    state->position_for[TIMELINE_VALUE_ID] = next_disabled_pos++;
-  }
-
   ddog_prof_Profile_drop(&state->profile_slot_one.profile);
   ddog_prof_Profile_drop(&state->profile_slot_two.profile);
 
@@ -527,12 +540,16 @@ static VALUE _native_serialize(DDTRACE_UNUSED VALUE _self, VALUE recorder_instan
   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 on to the Global VM Lock; see comments on method for why
+  // 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);
   // Prepare the iteration on heap recorder we'll be doing outside the GVL. The preparation needs to
-  // happen while holding on to the GVL.
+  // 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);
@@ -560,7 +577,7 @@ static VALUE _native_serialize(DDTRACE_UNUSED VALUE _self, VALUE recorder_instan
     rb_thread_call_without_gvl2(call_serialize_without_gvl, &args, NULL /* No interruption function needed in this case */, NULL /* Not needed */);
   }
 
-  // Cleanup after heap recorder iteration. This needs to happen while holding on to the GVL.
+  // Cleanup after heap recorder iteration. This needs to happen while holding the GVL.
   heap_recorder_finish_iteration(state->heap_recorder);
 
   // NOTE: We are focusing on the serialization time outside of the GVL in this stat here. This doesn't
@@ -606,7 +623,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
@@ -622,49 +660,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
     },
-    labels.end_timestamp_ns
+    /* To disable end_timestamp_ns to get aggregated profiles in pprof, replace the below with `0` */
+    args->labels.end_timestamp_ns
   );
 
-  active_slot.data->stats.recorded_samples++;
-
-  sampler_unlock_active_profile(active_slot);
+  args->active_slot.data->stats.recorded_samples++;
 
-  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 +710,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
@@ -905,6 +933,9 @@ static ddog_Timespec system_epoch_now_timespec(void) {
 //
 // Assumption: This method gets called BEFORE restarting profiling -- e.g. there are no components attempting to
 // trigger samples at the same time.
+//
+// Note that tests call this method directly in the same process without forking,
+// and in such a case non-current Threads keep running.
 static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE recorder_instance) {
   stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
@@ -935,15 +966,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) {
@@ -1021,13 +1043,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) {
@@ -1039,8 +1087,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;
 }
 
@@ -1150,11 +1198,18 @@ 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,7 +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/datadog_profiling_native_extension/time_helpers.h

@@ -4,22 +4,24 @@
 #include <errno.h>
 #include <time.h>
 
-#include "extconf.h" // Needed for HAVE_CLOCK_MONOTONIC_COARSE
+#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 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) {
   struct timespec clock_value;
@@ -32,7 +34,16 @@ static inline long retrieve_clock_as_ns(clockid_t clock_id, raise_on_failure_set
   return clock_value.tv_nsec + SECONDS_AS_NS(clock_value.tv_sec);
 }
 
-static inline long monotonic_wall_time_now_ns(raise_on_failure_setting raise_on_failure) { return retrieve_clock_as_ns(CLOCK_MONOTONIC, raise_on_failure); }
+// CLOCK_MONOTONIC on macOS only has microsecond precision, CLOCK_MONOTONIC_RAW has nanosecond precision
+#ifdef __APPLE__
+  #define CLOCK_MONOTONIC_FOR_PROFILING CLOCK_MONOTONIC_RAW
+  #define CLOCK_MONOTONIC_COARSE_FOR_PROFILING CLOCK_MONOTONIC_RAW_APPROX
+#else
+  #define CLOCK_MONOTONIC_FOR_PROFILING CLOCK_MONOTONIC
+  #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); }
 
 // Coarse instants use CLOCK_MONOTONIC_COARSE on Linux which is expected to provide resolution in the millisecond range:
@@ -46,11 +57,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) {
- #ifdef HAVE_CLOCK_MONOTONIC_COARSE // Linux
-    return to_coarse_instant(retrieve_clock_as_ns(CLOCK_MONOTONIC_COARSE, DO_NOT_RAISE_ON_FAILURE));
-  #else // macOS
-    return to_coarse_instant(retrieve_clock_as_ns(CLOCK_MONOTONIC, DO_NOT_RAISE_ON_FAILURE));
-  #endif
+  return to_coarse_instant(retrieve_clock_as_ns(CLOCK_MONOTONIC_COARSE_FOR_PROFILING, DO_NOT_RAISE_ON_FAILURE));
 }
 
 long monotonic_to_system_epoch_ns(monotonic_to_system_epoch_state *state, long monotonic_wall_time_ns);

data/ext/datadog_profiling_native_extension/unsafe_api_calls_check.h

@@ -4,7 +4,7 @@
 //
 // Specifically, when the profiler is sampling, we're never supposed to call into Ruby code (e.g. methods
 // implemented using Ruby code) or allocate Ruby objects.
-// That's because those events introduce thread switch points, and really we don't the VM switching between threads
+// That's because those events introduce thread switch points, and really we don't want the VM switching between threads
 // in the middle of the profiler sampling. This includes raising exceptions.
 //
 // Raising exceptions as the very last operation, to stop the profiler is ok, but comes a caveat: raising exceptions
@@ -18,10 +18,12 @@
 // in most (all?) thread switch points, Ruby will check for interrupts and run the postponed jobs.
 //
 // Thus, if we set a flag while we're sampling (inside_unsafe_context), trigger the postponed job, and then only unset
-// the flag after sampling, he correct thing to happen is that the postponed job should never see the flag.
+// the flag after sampling, the correct thing to happen is that the postponed job should never see the flag.
 //
 // If, however, we have a bug and there's a thread switch point, our postponed job will see the flag and immediately
 // stop the Ruby VM before further damage happens (and hopefully giving us a stack trace clearly pointing to the culprit).
+//
+// Note that this check currently does not detect Ruby object allocations, as those do not check for interrupts.
 
 void unsafe_api_calls_check_init(void);
 

data/ext/libdatadog_api/crashtracker.c

@@ -51,12 +51,8 @@ static VALUE _native_start_or_update_on_fork(int argc, VALUE *argv, DDTRACE_UNUS
 
   VALUE version = datadog_gem_version();
 
-  // Tags and endpoint are heap-allocated, so after here we can't raise exceptions otherwise we'll leak this memory
+  // Tags are heap-allocated, so after here we can't raise exceptions otherwise we'll leak this memory
   // Start of exception-free zone to prevent leaks {{
-  ddog_Endpoint *endpoint = ddog_endpoint_from_url(char_slice_from_ruby_string(agent_base_url));
-  if (endpoint == NULL) {
-    raise_error(rb_eRuntimeError, "Failed to create endpoint from agent_base_url: %"PRIsVALUE, agent_base_url);
-  }
   ddog_Vec_Tag tags = convert_tags(tags_as_array);
 
   ddog_crasht_Config config = {
@@ -73,9 +69,11 @@ static VALUE _native_start_or_update_on_fork(int argc, VALUE *argv, DDTRACE_UNUS
     // overriding what Ruby set up seems a saner default to keep anyway.
     .create_alt_stack = false,
     .use_alt_stack = true,
-    .endpoint = endpoint,
+    .endpoint = {.url = char_slice_from_ruby_string(agent_base_url)},
     .resolve_frames = DDOG_CRASHT_STACKTRACE_COLLECTION_ENABLED_WITH_SYMBOLS_IN_RECEIVER,
     .timeout_ms = FIX2INT(upload_timeout_seconds) * 1000,
+    .collect_all_threads = true,
+    .max_threads = 128,
   };
 
   ddog_crasht_Metadata metadata = {
@@ -106,11 +104,12 @@ static VALUE _native_start_or_update_on_fork(int argc, VALUE *argv, DDTRACE_UNUS
       ) :
       ddog_crasht_update_on_fork(config, receiver_config, metadata);
 
-  first_init = false;
+  // We use first_init to know which of [init, reconfigure] needs to be called. BUT if init failed we actually need
+  // to call init next time again, not reconfigure.
+  if (result.tag == DDOG_VOID_RESULT_OK) first_init = false;
 
   // Clean up before potentially raising any exceptions
   ddog_Vec_Tag_drop(tags);
-  ddog_endpoint_drop(endpoint);
   // }} End of exception-free zone to prevent leaks
 
   CHECK_VOID_RESULT("Failed to start/update the crash tracker", result);