Merge pull request #3786 from DataDog/ivoanjo/small-profiler-tweaks

[NO-TICKET] Fixes for rare issues in the profiler + small optimization

ext/datadog_profiling_native_extension/collectors_cpu_and_wall_time_worker.c

@@ -349,6 +349,10 @@ static VALUE _native_new(VALUE klass) {
 
   discrete_dynamic_sampler_init(&state->allocation_sampler, "allocation", now);
 
+  // Note: As of this writing, no new Ruby objects get created and stored in the state. If that ever changes, remember
+  // to keep them on the stack and mark them with RB_GC_GUARD -- otherwise it's possible for a GC to run and
+  // since the instance representing the state does not yet exist, such objects will not get marked.
+
   return state->self_instance = TypedData_Wrap_Struct(klass, &cpu_and_wall_time_worker_typed_data, state);
 }
 

ext/datadog_profiling_native_extension/collectors_discrete_dynamic_sampler.c

@@ -359,6 +359,10 @@ static VALUE _native_new(VALUE klass) {
   }
   discrete_dynamic_sampler_init(&state->sampler, "test sampler", now_ns);
 
+  // Note: As of this writing, no new Ruby objects get created and stored in the state. If that ever changes, remember
+  // to keep them on the stack and mark them with RB_GC_GUARD -- otherwise it's possible for a GC to run and
+  // since the instance representing the state does not yet exist, such objects will not get marked.
+
   return TypedData_Wrap_Struct(klass, &sampler_typed_data, state);
 }
 

ext/datadog_profiling_native_extension/collectors_idle_sampling_helper.c

@@ -83,6 +83,10 @@ static VALUE _native_new(VALUE klass) {
 
   reset_state(state);
 
+  // Note: As of this writing, no new Ruby objects get created and stored in the state. If that ever changes, remember
+  // to keep them on the stack and mark them with RB_GC_GUARD -- otherwise it's possible for a GC to run and
+  // since the instance representing the state does not yet exist, such objects will not get marked.
+
   return TypedData_Wrap_Struct(klass, &idle_sampling_helper_typed_data, state);
 }
 

ext/datadog_profiling_native_extension/collectors_thread_context.c

@@ -345,17 +345,28 @@ static VALUE _native_new(VALUE klass) {
     st_init_numtable();
   state->recorder_instance = Qnil;
   state->tracer_context_key = MISSING_TRACER_CONTEXT_KEY;
-  state->thread_list_buffer = rb_ary_new();
+  VALUE thread_list_buffer = rb_ary_new();
+  state->thread_list_buffer = thread_list_buffer;
   state->endpoint_collection_enabled = true;
   state->timeline_enabled = true;
   state->allocation_type_enabled = true;
   state->time_converter_state = (monotonic_to_system_epoch_state) MONOTONIC_TO_SYSTEM_EPOCH_INITIALIZER;
-  state->main_thread = rb_thread_main();
+  VALUE main_thread = rb_thread_main();
+  state->main_thread = main_thread;
   state->otel_current_span_key = Qnil;
   state->gc_tracking.wall_time_at_previous_gc_ns = INVALID_TIME;
   state->gc_tracking.wall_time_at_last_flushed_gc_event_ns = 0;
 
-  return TypedData_Wrap_Struct(klass, &thread_context_collector_typed_data, state);
+  // Note: Remember to keep any new allocated objects that get stored in the state also on the stack + mark them with
+  // RB_GC_GUARD -- otherwise it's possible for a GC to run and
+  // since the instance representing the state does not yet exist, such objects will not get marked.
+
+  VALUE instance = TypedData_Wrap_Struct(klass, &thread_context_collector_typed_data, state);
+
+  RB_GC_GUARD(thread_list_buffer);
+  RB_GC_GUARD(main_thread); // Arguably not needed, but perhaps can be move in some future Ruby release?
+
+  return instance;
 }
 
 static VALUE _native_initialize(
@@ -1257,7 +1268,7 @@ void thread_context_collector_sample_allocation(VALUE self_instance, unsigned in
       // Thus, we need to make sure there's actually a class before getting its name.
 
       if (klass != 0) {
-        const char *name = rb_obj_classname(new_object);
+        const char *name = rb_class2name(klass);
         size_t name_length = name != NULL ? strlen(name) : 0;
 
         if (name_length > 0) {

ext/datadog_profiling_native_extension/heap_recorder.c

@@ -219,7 +219,7 @@ void heap_recorder_free(heap_recorder *heap_recorder) {
   st_foreach(heap_recorder->heap_records, st_heap_record_entry_free, 0);
   st_free_table(heap_recorder->heap_records);
 
-  if (heap_recorder->active_recording.object_record != NULL) {
+  if (heap_recorder->active_recording.object_record != NULL && heap_recorder->active_recording.object_record != &SKIPPED_RECORD) {
     // If there's a partial object record, clean it up as well
     object_record_free(heap_recorder->active_recording.object_record);
   }

ext/datadog_profiling_native_extension/stack_recorder.c

@@ -342,6 +342,10 @@ static VALUE _native_new(VALUE klass) {
   // Note: At this point, slot_one_profile and slot_two_profile contain null pointers. Libdatadog validates pointers
   // before using them so it's ok for us to go ahead and create the StackRecorder object.
 
+  // Note: As of this writing, no new Ruby objects get created and stored in the state. If that ever changes, remember
+  // to keep them on the stack and mark them with RB_GC_GUARD -- otherwise it's possible for a GC to run and
+  // since the instance representing the state does not yet exist, such objects will not get marked.
+
   VALUE stack_recorder = TypedData_Wrap_Struct(klass, &stack_recorder_typed_data, state);
 
   // NOTE: We initialize this because we want a new recorder to be operational even without initialization and our