[PROF-7307] Record allocation type when sampling objects

ext/ddtrace_profiling_native_extension/collectors_thread_context.c

@@ -101,6 +101,8 @@ struct thread_context_collector_state {
   bool endpoint_collection_enabled;
   // Used to omit timestamps / timeline events from collected data
   bool timeline_enabled;
+  // Used to omit class information from collected allocation data
+  bool allocation_type_enabled;
   // Used when calling monotonic_to_system_epoch_ns
   monotonic_to_system_epoch_state time_converter_state;
   // Used to identify the main thread, to give it a fallback name
@@ -157,7 +159,8 @@ static VALUE _native_initialize(
   VALUE max_frames,
   VALUE tracer_context_key,
   VALUE endpoint_collection_enabled,
-  VALUE timeline_enabled
+  VALUE timeline_enabled,
+  VALUE allocation_type_enabled
 );
 static VALUE _native_sample(VALUE self, VALUE collector_instance, VALUE profiler_overhead_stack_thread);
 static VALUE _native_on_gc_start(VALUE self, VALUE collector_instance);
@@ -179,7 +182,8 @@ static void trigger_sample_for_thread(
   sample_values values,
   sample_type type,
   long current_monotonic_wall_time_ns,
-  ddog_CharSlice *ruby_vm_type
+  ddog_CharSlice *ruby_vm_type,
+  ddog_CharSlice *class_name
 );
 static VALUE _native_thread_list(VALUE self);
 static struct per_thread_context *get_or_create_context_for(VALUE thread, struct thread_context_collector_state *state);
@@ -219,7 +223,7 @@ void collectors_thread_context_init(VALUE profiling_module) {
   // https://bugs.ruby-lang.org/issues/18007 for a discussion around this.
   rb_define_alloc_func(collectors_thread_context_class, _native_new);
 
-  rb_define_singleton_method(collectors_thread_context_class, "_native_initialize", _native_initialize, 6);
+  rb_define_singleton_method(collectors_thread_context_class, "_native_initialize", _native_initialize, 7);
   rb_define_singleton_method(collectors_thread_context_class, "_native_inspect", _native_inspect, 1);
   rb_define_singleton_method(collectors_thread_context_class, "_native_reset_after_fork", _native_reset_after_fork, 1);
   rb_define_singleton_method(testing_module, "_native_sample", _native_sample, 2);
@@ -309,6 +313,7 @@ static VALUE _native_new(VALUE klass) {
   state->thread_list_buffer = rb_ary_new();
   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();
 
@@ -322,10 +327,12 @@ static VALUE _native_initialize(
   VALUE max_frames,
   VALUE tracer_context_key,
   VALUE endpoint_collection_enabled,
-  VALUE timeline_enabled
+  VALUE timeline_enabled,
+  VALUE allocation_type_enabled
 ) {
   ENFORCE_BOOLEAN(endpoint_collection_enabled);
   ENFORCE_BOOLEAN(timeline_enabled);
+  ENFORCE_BOOLEAN(allocation_type_enabled);
 
   struct thread_context_collector_state *state;
   TypedData_Get_Struct(collector_instance, struct thread_context_collector_state, &thread_context_collector_typed_data, state);
@@ -339,6 +346,7 @@ static VALUE _native_initialize(
   state->recorder_instance = enforce_recorder_instance(recorder_instance);
   state->endpoint_collection_enabled = (endpoint_collection_enabled == Qtrue);
   state->timeline_enabled = (timeline_enabled == Qtrue);
+  state->allocation_type_enabled = (allocation_type_enabled == Qtrue);
 
   if (RTEST(tracer_context_key)) {
     ENFORCE_TYPE(tracer_context_key, T_SYMBOL);
@@ -465,6 +473,7 @@ void update_metrics_and_sample(
     (sample_values) {.cpu_time_ns = cpu_time_elapsed_ns, .cpu_samples = 1, .wall_time_ns = wall_time_elapsed_ns},
     SAMPLE_REGULAR,
     current_monotonic_wall_time_ns,
+    NULL,
     NULL
   );
 }
@@ -609,6 +618,7 @@ VALUE thread_context_collector_sample_after_gc(VALUE self_instance) {
       (sample_values) {.cpu_time_ns = gc_cpu_time_elapsed_ns, .cpu_samples = 1, .wall_time_ns = gc_wall_time_elapsed_ns},
       SAMPLE_IN_GC,
       INVALID_TIME, // For now we're not collecting timestamps for these events
+      NULL,
       NULL
     );
 
@@ -641,14 +651,16 @@ static void trigger_sample_for_thread(
   sample_values values,
   sample_type type,
   long current_monotonic_wall_time_ns,
-  ddog_CharSlice *ruby_vm_type // Only used for allocation profiling, NULL for others; @ivoanjo: may want to refactor this at some point?
+  // These two labels are only used for allocation profiling; @ivoanjo: may want to refactor this at some point?
+  ddog_CharSlice *ruby_vm_type,
+  ddog_CharSlice *class_name
 ) {
   int max_label_count =
     1 + // thread id
     1 + // thread name
     1 + // profiler overhead
     1 + // end_timestamp_ns
-    1 + // ruby vm type
+    2 + // ruby vm type and allocation class
     2;  // local root span id and span id
   ddog_prof_Label labels[max_label_count];
   int label_pos = 0;
@@ -725,6 +737,13 @@ static void trigger_sample_for_thread(
     };
   }
 
+  if (class_name != NULL) {
+    labels[label_pos++] = (ddog_prof_Label) {
+      .key = DDOG_CHARSLICE_C("allocation class"),
+      .str = *class_name
+    };
+  }
+
   // The number of times `label_pos++` shows up in this function needs to match `max_label_count`. To avoid "oops I
   // forgot to update max_label_count" in the future, we've also added this validation.
   // @ivoanjo: I wonder if C compilers are smart enough to statically prove this check never triggers unless someone
@@ -831,6 +850,7 @@ static VALUE _native_inspect(DDTRACE_UNUSED VALUE _self, VALUE collector_instanc
   rb_str_concat(result, rb_sprintf(" stats=%"PRIsVALUE, stats_as_ruby_hash(state)));
   rb_str_concat(result, rb_sprintf(" endpoint_collection_enabled=%"PRIsVALUE, state->endpoint_collection_enabled ? Qtrue : Qfalse));
   rb_str_concat(result, rb_sprintf(" timeline_enabled=%"PRIsVALUE, state->timeline_enabled ? Qtrue : Qfalse));
+  rb_str_concat(result, rb_sprintf(" allocation_type_enabled=%"PRIsVALUE, state->allocation_type_enabled ? Qtrue : Qfalse));
   rb_str_concat(result, rb_sprintf(
     " time_converter_state={.system_epoch_ns_reference=%ld, .delta_to_epoch_ns=%ld}",
     state->time_converter_state.system_epoch_ns_reference,
@@ -1068,8 +1088,53 @@ void thread_context_collector_sample_allocation(VALUE self_instance, unsigned in
 
   VALUE current_thread = rb_thread_current();
 
+  enum ruby_value_type type = rb_type(new_object);
+
   // Tag samples with the VM internal types
-  ddog_CharSlice ruby_vm_type = ruby_value_type_to_char_slice(rb_type(new_object));
+  ddog_CharSlice ruby_vm_type = ruby_value_type_to_char_slice(type);
+
+  // Since this is stack allocated, be careful about moving it
+  ddog_CharSlice class_name;
+  ddog_CharSlice *optional_class_name = NULL;
+
+  if (state->allocation_type_enabled) {
+    optional_class_name = &class_name;
+
+    if (
+      type == RUBY_T_OBJECT   ||
+      type == RUBY_T_CLASS    ||
+      type == RUBY_T_MODULE   ||
+      type == RUBY_T_FLOAT    ||
+      type == RUBY_T_STRING   ||
+      type == RUBY_T_REGEXP   ||
+      type == RUBY_T_ARRAY    ||
+      type == RUBY_T_HASH     ||
+      type == RUBY_T_STRUCT   ||
+      type == RUBY_T_BIGNUM   ||
+      type == RUBY_T_FILE     ||
+      type == RUBY_T_DATA     ||
+      type == RUBY_T_MATCH    ||
+      type == RUBY_T_COMPLEX  ||
+      type == RUBY_T_RATIONAL ||
+      type == RUBY_T_NIL      ||
+      type == RUBY_T_TRUE     ||
+      type == RUBY_T_FALSE    ||
+      type == RUBY_T_SYMBOL   ||
+      type == RUBY_T_FIXNUM
+    ) {
+      const char *name = rb_obj_classname(new_object);
+      size_t name_length = name != NULL ? strlen(name) : 0;
+
+      if (name_length > 0) {
+        class_name = (ddog_CharSlice) {.ptr = name, .len = name_length};
+      } else {
+        class_name = DDOG_CHARSLICE_C("(Anonymous)");
+      }
+    } else {
+      class_name = ruby_vm_type; // For internal things and, we just use the VM type
+                                 // TODO: Maybe prefix them with a nice note? E.g. (VM Internal, T_IMEMO)
+    }
+  }
 
   trigger_sample_for_thread(
     state,
@@ -1079,7 +1144,8 @@ void thread_context_collector_sample_allocation(VALUE self_instance, unsigned in
     (sample_values) {.alloc_samples = sample_weight},
     SAMPLE_REGULAR,
     INVALID_TIME, // For now we're not collecting timestamps for allocation events, as per profiling team internal discussions
-    &ruby_vm_type
+    &ruby_vm_type,
+    optional_class_name
   );
 }
 

lib/datadog/profiling/collectors/thread_context.rb

@@ -14,7 +14,14 @@ module Collectors
       #
       # Methods prefixed with _native_ are implemented in `collectors_thread_context.c`
       class ThreadContext
-        def initialize(recorder:, max_frames:, tracer:, endpoint_collection_enabled:, timeline_enabled:)
+        def initialize(
+          recorder:,
+          max_frames:,
+          tracer:,
+          endpoint_collection_enabled:,
+          timeline_enabled:,
+          allocation_type_enabled: true
+        )
           tracer_context_key = safely_extract_context_key_from(tracer)
           self.class._native_initialize(
             self,
@@ -23,6 +30,7 @@ def initialize(recorder:, max_frames:, tracer:, endpoint_collection_enabled:, ti
             tracer_context_key,
             endpoint_collection_enabled,
             timeline_enabled,
+            allocation_type_enabled,
           )
         end
 

sig/datadog/profiling/collectors/thread_context.rbs

@@ -8,6 +8,7 @@ module Datadog
           tracer: Datadog::Tracing::Tracer?,
           endpoint_collection_enabled: bool,
           timeline_enabled: bool,
+          ?allocation_type_enabled: bool,
         ) -> void
 
         def self._native_initialize: (
@@ -17,6 +18,7 @@ module Datadog
           ::Symbol? tracer_context_key,
           bool endpoint_collection_enabled,
           bool timeline_enabled,
+          bool allocation_type_enabled,
         ) -> void
 
         def inspect: () -> ::String

spec/datadog/profiling/collectors/thread_context_spec.rb

@@ -39,6 +39,7 @@
   let(:tracer) { nil }
   let(:endpoint_collection_enabled) { true }
   let(:timeline_enabled) { false }
+  let(:allocation_type_enabled) { true }
 
   subject(:cpu_and_wall_time_collector) do
     described_class.new(
@@ -47,6 +48,7 @@
       tracer: tracer,
       endpoint_collection_enabled: endpoint_collection_enabled,
       timeline_enabled: timeline_enabled,
+      allocation_type_enabled: allocation_type_enabled,
     )
   end
 
@@ -981,33 +983,53 @@ def another_way_of_calling_sample(profiler_overhead_stack_thread: Thread.current
       end
     end
 
-    {
-      T_OBJECT: Object.new,
-      T_CLASS: Object,
-      T_MODULE: Kernel,
-      T_FLOAT: 1.0,
-      T_STRING: 'Hello!',
-      T_REGEXP: /Hello/,
-      T_ARRAY: [],
-      T_HASH: {},
-      T_STRUCT: Struct.new(:a).new,
-      T_BIGNUM: 2**256,
-      T_DATA: described_class.allocate, # ThreadContext is a T_DATA; we create here a dummy instance just as an example
-      T_MATCH: 'a'.match(Regexp.new('a')),
-      T_COMPLEX: Complex(1),
-      T_RATIONAL: 1/2r,
-      T_NIL: nil,
-      T_TRUE: true,
-      T_FALSE: false,
-      T_SYMBOL: :hello,
-      T_FIXNUM: 1,
-    }.each do |expected_type, object|
+    [
+      { expected_type: :T_OBJECT, object: Object.new, klass: 'Object' },
+      { expected_type: :T_CLASS, object: Object, klass: 'Class' },
+      { expected_type: :T_MODULE, object: Kernel, klass: 'Module' },
+      { expected_type: :T_FLOAT, object: 1.0, klass: 'Float' },
+      { expected_type: :T_STRING, object: 'Hello!', klass: 'String' },
+      { expected_type: :T_REGEXP, object: /Hello/, klass: 'Regexp' },
+      { expected_type: :T_ARRAY, object: [], klass: 'Array' },
+      { expected_type: :T_HASH, object: {}, klass: 'Hash' },
+      { expected_type: :T_BIGNUM, object: 2**256, klass: RUBY_VERSION < '2.4' ? 'Bignum' : 'Integer' },
+      # ThreadContext is a T_DATA; we create here a dummy instance just as an example
+      { expected_type: :T_DATA, object: described_class.allocate, klass: 'Datadog::Profiling::Collectors::ThreadContext' },
+      { expected_type: :T_MATCH, object: 'a'.match(Regexp.new('a')), klass: 'MatchData' },
+      { expected_type: :T_COMPLEX, object: Complex(1), klass: 'Complex' },
+      { expected_type: :T_RATIONAL, object: 1/2r, klass: 'Rational' },
+      { expected_type: :T_NIL, object: nil, klass: 'NilClass' },
+      { expected_type: :T_TRUE, object: true, klass: 'TrueClass' },
+      { expected_type: :T_FALSE, object: false, klass: 'FalseClass' },
+      { expected_type: :T_SYMBOL, object: :hello, klass: 'Symbol' },
+      { expected_type: :T_FIXNUM, object: 1, klass: RUBY_VERSION < '2.4' ? 'Fixnum' : 'Integer' },
+    ].each do |type|
+      expected_type = type.fetch(:expected_type)
+      object = type.fetch(:object)
+      klass = type.fetch(:klass)
+
       context "when sampling a #{expected_type}" do
         it 'includes the correct ruby vm type for the passed object' do
           sample_allocation(weight: 123, new_object: object)
 
           expect(single_sample.labels.fetch(:'ruby vm type')).to eq expected_type.to_s
         end
+
+        it 'includes the correct class for the passed object' do
+          sample_allocation(weight: 123, new_object: object)
+
+          expect(single_sample.labels.fetch(:'allocation class')).to eq klass
+        end
+
+        context 'when allocation_type_enabled is false' do
+          let(:allocation_type_enabled) { false }
+
+          it 'does not record the correct class for the passed object' do
+            sample_allocation(weight: 123, new_object: object)
+
+            expect(single_sample.labels).to_not include(:'allocation class' => anything)
+          end
+        end
       end
     end
 
@@ -1019,6 +1041,54 @@ def another_way_of_calling_sample(profiler_overhead_stack_thread: Thread.current
 
         expect(single_sample.labels.fetch(:'ruby vm type')).to eq 'T_FILE'
       end
+
+      it 'includes the correct class for the passed object' do
+        File.open(__FILE__) do |file|
+          sample_allocation(weight: 123, new_object: file)
+        end
+
+        expect(single_sample.labels.fetch(:'allocation class')).to eq 'File'
+      end
+
+      context 'when allocation_type_enabled is false' do
+        let(:allocation_type_enabled) { false }
+
+        it 'does not record the correct class for the passed object' do
+          File.open(__FILE__) do |file|
+            sample_allocation(weight: 123, new_object: file)
+          end
+
+          expect(single_sample.labels).to_not include(:'allocation class' => anything)
+        end
+      end
+    end
+
+    context 'when sampling a Struct' do
+      before do
+        stub_const('ThreadContextSpec::TestStruct', Struct.new(:a))
+      end
+
+      it 'includes the correct ruby vm type for the passed object' do
+        sample_allocation(weight: 123, new_object: ThreadContextSpec::TestStruct.new)
+
+        expect(single_sample.labels.fetch(:'ruby vm type')).to eq 'T_STRUCT'
+      end
+
+      it 'includes the correct class for the passed object' do
+        sample_allocation(weight: 123, new_object: ThreadContextSpec::TestStruct.new)
+
+        expect(single_sample.labels.fetch(:'allocation class')).to eq 'ThreadContextSpec::TestStruct'
+      end
+
+      context 'when allocation_type_enabled is false' do
+        let(:allocation_type_enabled) { false }
+
+        it 'does not record the correct class for the passed object' do
+          sample_allocation(weight: 123, new_object: ThreadContextSpec::TestStruct.new)
+
+          expect(single_sample.labels).to_not include(:'allocation class' => anything)
+        end
+      end
     end
   end