Merge pull request #3103 from DataDog/ivoanjo/prof-7307-wire-up-allocation-sampling

[PROF-7307] Wire up allocation sampling into `CpuAndWallTimeWorker`

Author: ivoanjo

ext/ddtrace_profiling_native_extension/collectors_cpu_and_wall_time_worker.c

@@ -83,6 +83,7 @@ struct cpu_and_wall_time_worker_state {
   bool allocation_counting_enabled;
   bool no_signals_workaround_enabled;
   bool dynamic_sampling_rate_enabled;
+  int allocation_sample_every; // Temporarily used for development/testing of allocation profiling
   VALUE self_instance;
   VALUE thread_context_collector_instance;
   VALUE idle_sampling_helper_instance;
@@ -150,7 +151,8 @@ static VALUE _native_initialize(
   VALUE idle_sampling_helper_instance,
   VALUE allocation_counting_enabled,
   VALUE no_signals_workaround_enabled,
-  VALUE dynamic_sampling_rate_enabled
+  VALUE dynamic_sampling_rate_enabled,
+  VALUE allocation_sample_every
 );
 static void cpu_and_wall_time_worker_typed_data_mark(void *state_ptr);
 static VALUE _native_sampling_loop(VALUE self, VALUE instance);
@@ -183,9 +185,10 @@ static void grab_gvl_and_sample(void);
 static void reset_stats(struct cpu_and_wall_time_worker_state *state);
 static void sleep_for(uint64_t time_ns);
 static VALUE _native_allocation_count(DDTRACE_UNUSED VALUE self);
-static void on_newobj_event(DDTRACE_UNUSED VALUE tracepoint_data, DDTRACE_UNUSED void *unused);
+static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused);
 static void disable_tracepoints(struct cpu_and_wall_time_worker_state *state);
 static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self);
+static VALUE rescued_sample_allocation(VALUE tracepoint_data);
 
 // Note on sampler global state safety:
 //
@@ -223,7 +226,7 @@ void collectors_cpu_and_wall_time_worker_init(VALUE profiling_module) {
   // https://bugs.ruby-lang.org/issues/18007 for a discussion around this.
   rb_define_alloc_func(collectors_cpu_and_wall_time_worker_class, _native_new);
 
-  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_initialize", _native_initialize, 7);
+  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_initialize", _native_initialize, 8);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_sampling_loop", _native_sampling_loop, 1);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_stop", _native_stop, 2);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_reset_after_fork", _native_reset_after_fork, 1);
@@ -261,6 +264,7 @@ static VALUE _native_new(VALUE klass) {
   state->allocation_counting_enabled = false;
   state->no_signals_workaround_enabled = false;
   state->dynamic_sampling_rate_enabled = true;
+  state->allocation_sample_every = 0;
   state->thread_context_collector_instance = Qnil;
   state->idle_sampling_helper_instance = Qnil;
   state->owner_thread = Qnil;
@@ -287,12 +291,14 @@ static VALUE _native_initialize(
   VALUE idle_sampling_helper_instance,
   VALUE allocation_counting_enabled,
   VALUE no_signals_workaround_enabled,
-  VALUE dynamic_sampling_rate_enabled
+  VALUE dynamic_sampling_rate_enabled,
+  VALUE allocation_sample_every
 ) {
   ENFORCE_BOOLEAN(gc_profiling_enabled);
   ENFORCE_BOOLEAN(allocation_counting_enabled);
   ENFORCE_BOOLEAN(no_signals_workaround_enabled);
   ENFORCE_BOOLEAN(dynamic_sampling_rate_enabled);
+  ENFORCE_TYPE(allocation_sample_every, T_FIXNUM);
 
   struct cpu_and_wall_time_worker_state *state;
   TypedData_Get_Struct(self_instance, struct cpu_and_wall_time_worker_state, &cpu_and_wall_time_worker_typed_data, state);
@@ -301,6 +307,12 @@ static VALUE _native_initialize(
   state->allocation_counting_enabled = (allocation_counting_enabled == Qtrue);
   state->no_signals_workaround_enabled = (no_signals_workaround_enabled == Qtrue);
   state->dynamic_sampling_rate_enabled = (dynamic_sampling_rate_enabled == Qtrue);
+  state->allocation_sample_every = NUM2INT(allocation_sample_every);
+
+  if (state->allocation_sample_every < 0) {
+    rb_raise(rb_eArgError, "Unexpected value for allocation_sample_every: %d. This value must be >= 0.", state->allocation_sample_every);
+  }
+
   state->thread_context_collector_instance = enforce_thread_context_collector_instance(thread_context_collector_instance);
   state->idle_sampling_helper_instance = idle_sampling_helper_instance;
   state->gc_tracepoint = rb_tracepoint_new(Qnil, RUBY_INTERNAL_EVENT_GC_ENTER | RUBY_INTERNAL_EVENT_GC_EXIT, on_gc_event, NULL /* unused */);
@@ -880,7 +892,7 @@ static VALUE _native_allocation_count(DDTRACE_UNUSED VALUE self) {
 
 // Implements memory-related profiling events. This function is called by Ruby via the `object_allocation_tracepoint`
 // when the RUBY_INTERNAL_EVENT_NEWOBJ event is triggered.
-static void on_newobj_event(DDTRACE_UNUSED VALUE tracepoint_data, DDTRACE_UNUSED void *unused) {
+static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused) {
   // Update thread-local allocation count
   if (RB_UNLIKELY(allocation_count == UINT64_MAX)) {
     allocation_count = 0;
@@ -907,7 +919,12 @@ static void on_newobj_event(DDTRACE_UNUSED VALUE tracepoint_data, DDTRACE_UNUSED
   // defined as not being able to allocate) sets this.
   state->during_sample = true;
 
-  // TODO: Sampling goes here (calling into `thread_context_collector_sample_allocation`)
+  // TODO: This is a placeholder sampling decision strategy. We plan to replace it with a better one soon (e.g. before
+  // beta), and having something here allows us to test the rest of feature, sampling decision aside.
+  if (state->allocation_sample_every > 0 && ((allocation_count % state->allocation_sample_every) == 0)) {
+    // Rescue against any exceptions that happen during sampling
+    safely_call(rescued_sample_allocation, tracepoint_data, state->self_instance);
+  }
 
   state->during_sample = false;
 }
@@ -929,3 +946,18 @@ static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self) {
     return result;
   }
 }
+
+static VALUE rescued_sample_allocation(VALUE tracepoint_data) {
+  struct cpu_and_wall_time_worker_state *state = active_sampler_instance_state; // Read from global variable, see "sampler global state safety" note above
+
+  // This should not happen in a normal situation because on_newobj_event already checked for this, but just in case...
+  if (state == NULL) return Qnil;
+
+  rb_trace_arg_t *data = rb_tracearg_from_tracepoint(tracepoint_data);
+  VALUE new_object = rb_tracearg_object(data);
+
+  thread_context_collector_sample_allocation(state->thread_context_collector_instance, state->allocation_sample_every, new_object);
+
+  // Return a dummy VALUE because we're called from rb_rescue2 which requires it
+  return Qnil;
+}

ext/ddtrace_profiling_native_extension/collectors_thread_context.c

@@ -206,6 +206,7 @@ static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE collector
 static VALUE thread_list(struct thread_context_collector_state *state);
 static VALUE _native_sample_allocation(DDTRACE_UNUSED VALUE self, VALUE collector_instance, VALUE sample_weight, VALUE new_object);
 static VALUE _native_new_empty_thread(VALUE self);
+ddog_CharSlice ruby_value_type_to_class_name(enum ruby_value_type type);
 
 void collectors_thread_context_init(VALUE profiling_module) {
   VALUE collectors_module = rb_define_module_under(profiling_module, "Collectors");
@@ -1122,17 +1123,30 @@ void thread_context_collector_sample_allocation(VALUE self_instance, unsigned in
       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};
+      VALUE klass = rb_class_of(new_object);
+
+      // Ruby sometimes plays a bit fast and loose with some of its internal objects, e.g.
+      // `rb_str_tmp_frozen_acquire` allocates a string with no class (klass=0).
+      // 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);
+        size_t name_length = name != NULL ? strlen(name) : 0;
+
+        if (name_length > 0) {
+          class_name = (ddog_CharSlice) {.ptr = name, .len = name_length};
+        } else {
+          // @ivoanjo: I'm not sure this can ever happen, but just-in-case
+          class_name = ruby_value_type_to_class_name(type);
+        }
       } else {
-        class_name = DDOG_CHARSLICE_C("(Anonymous)");
+        // Fallback for objects with no class
+        class_name = ruby_value_type_to_class_name(type);
       }
+    } else if (type == RUBY_T_IMEMO) {
+      class_name = DDOG_CHARSLICE_C("(VM Internal, T_IMEMO)");
     } 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)
+      class_name = ruby_vm_type; // For other weird internal things we just use the VM type
     }
   }
 
@@ -1164,3 +1178,29 @@ static VALUE new_empty_thread_inner(DDTRACE_UNUSED void *arg) { return Qnil; }
 static VALUE _native_new_empty_thread(DDTRACE_UNUSED VALUE self) {
   return rb_thread_create(new_empty_thread_inner, NULL);
 }
+
+ddog_CharSlice ruby_value_type_to_class_name(enum ruby_value_type type) {
+  switch (type) {
+    case(RUBY_T_OBJECT  ): return DDOG_CHARSLICE_C("Object");
+    case(RUBY_T_CLASS   ): return DDOG_CHARSLICE_C("Class");
+    case(RUBY_T_MODULE  ): return DDOG_CHARSLICE_C("Module");
+    case(RUBY_T_FLOAT   ): return DDOG_CHARSLICE_C("Float");
+    case(RUBY_T_STRING  ): return DDOG_CHARSLICE_C("String");
+    case(RUBY_T_REGEXP  ): return DDOG_CHARSLICE_C("Regexp");
+    case(RUBY_T_ARRAY   ): return DDOG_CHARSLICE_C("Array");
+    case(RUBY_T_HASH    ): return DDOG_CHARSLICE_C("Hash");
+    case(RUBY_T_STRUCT  ): return DDOG_CHARSLICE_C("Struct");
+    case(RUBY_T_BIGNUM  ): return DDOG_CHARSLICE_C("Integer");
+    case(RUBY_T_FILE    ): return DDOG_CHARSLICE_C("File");
+    case(RUBY_T_DATA    ): return DDOG_CHARSLICE_C("(VM Internal, T_DATA)");
+    case(RUBY_T_MATCH   ): return DDOG_CHARSLICE_C("MatchData");
+    case(RUBY_T_COMPLEX ): return DDOG_CHARSLICE_C("Complex");
+    case(RUBY_T_RATIONAL): return DDOG_CHARSLICE_C("Rational");
+    case(RUBY_T_NIL     ): return DDOG_CHARSLICE_C("NilClass");
+    case(RUBY_T_TRUE    ): return DDOG_CHARSLICE_C("TrueClass");
+    case(RUBY_T_FALSE   ): return DDOG_CHARSLICE_C("FalseClass");
+    case(RUBY_T_SYMBOL  ): return DDOG_CHARSLICE_C("Symbol");
+    case(RUBY_T_FIXNUM  ): return DDOG_CHARSLICE_C("Integer");
+                  default: return DDOG_CHARSLICE_C("(VM Internal, Missing class)");
+  }
+}

lib/datadog/profiling/collectors/cpu_and_wall_time_worker.rb

@@ -21,14 +21,22 @@ def initialize(
           idle_sampling_helper: IdleSamplingHelper.new,
           # **NOTE**: This should only be used for testing; disabling the dynamic sampling rate will increase the
           # profiler overhead!
-          dynamic_sampling_rate_enabled: true
+          dynamic_sampling_rate_enabled: true,
+          allocation_sample_every: 0 # Currently only for testing; Setting this to > 0 can add a lot of overhead!
         )
           unless dynamic_sampling_rate_enabled
             Datadog.logger.warn(
               'Profiling dynamic sampling rate disabled. This should only be used for testing, and will increase overhead!'
             )
           end
 
+          if allocation_counting_enabled && allocation_sample_every > 0
+            Datadog.logger.warn(
+              "Enabled experimental allocation profiling: allocation_sample_every=#{allocation_sample_every}. This is " \
+              'experimental, not recommended, and will increase overhead!'
+            )
+          end
+
           self.class._native_initialize(
             self,
             thread_context_collector,
@@ -37,6 +45,7 @@ def initialize(
             allocation_counting_enabled,
             no_signals_workaround_enabled,
             dynamic_sampling_rate_enabled,
+            allocation_sample_every,
           )
           @worker_thread = nil
           @failure_exception = nil

lib/datadog/profiling/component.rb

@@ -91,6 +91,7 @@ def self.build_profiler_component(settings:, agent_settings:, optional_tracer:)
             allocation_counting_enabled: settings.profiling.advanced.allocation_counting_enabled,
             no_signals_workaround_enabled: no_signals_workaround_enabled,
             thread_context_collector: thread_context_collector,
+            allocation_sample_every: 0,
           )
         else
           load_pprof_support

sig/datadog/profiling/collectors/cpu_and_wall_time_worker.rbs

@@ -14,6 +14,7 @@ module Datadog
           thread_context_collector: Datadog::Profiling::Collectors::ThreadContext,
           ?idle_sampling_helper: Datadog::Profiling::Collectors::IdleSamplingHelper,
           ?dynamic_sampling_rate_enabled: bool,
+          ?allocation_sample_every: Integer,
         ) -> void
 
         def self._native_initialize: (
@@ -24,6 +25,7 @@ module Datadog
           bool allocation_counting_enabled,
           bool no_signals_workaround_enabled,
           bool dynamic_sampling_rate_enabled,
+          ::Integer allocation_sample_every,
         ) -> true
 
         def start: () -> bool?

spec/datadog/profiling/collectors/cpu_and_wall_time_worker_spec.rb

@@ -412,6 +412,69 @@
         expect(stats.fetch(:trigger_sample_attempts)).to eq(stats.fetch(:postponed_job_success))
       end
     end
+
+    context 'when allocation sampling is enabled' do
+      let(:options) { { allocation_sample_every: 1 } }
+
+      before do
+        allow(Datadog.logger).to receive(:warn)
+      end
+
+      it 'logs a warning message mentioning this is experimental' do
+        expect(Datadog.logger).to receive(:warn).with(/Enabled experimental allocation profiling/)
+
+        start
+      end
+
+      it 'records allocated objects' do
+        stub_const('CpuAndWallTimeWorkerSpec::TestStruct', Struct.new(:foo))
+
+        start
+
+        123.times { CpuAndWallTimeWorkerSpec::TestStruct.new }
+        allocation_line = __LINE__ - 1
+
+        cpu_and_wall_time_worker.stop
+
+        allocation_sample =
+          samples_for_thread(samples_from_pprof(recorder.serialize!), Thread.current)
+            .find { |s| s.labels[:'allocation class'] == 'CpuAndWallTimeWorkerSpec::TestStruct' }
+
+        expect(allocation_sample.values).to include(:'alloc-samples' => 123)
+        expect(allocation_sample.locations.first.lineno).to eq allocation_line
+      end
+
+      context 'when sampling optimized Ruby strings' do
+        # Regression test: Some internal Ruby classes use a `rb_str_tmp_frozen_acquire` function which allocates a
+        # weird "intermediate" string object that has its class pointer set to 0.
+        #
+        # When such an object gets sampled, we need to take care not to try to resolve its class name.
+        #
+        # In practice, this test is actually validating behavior of the `ThreadContext` collector, but we can only
+        # really trigger this situation when using the allocation tracepoint, which lives in the `CpuAndWallTimeWorker`.
+        it 'does not crash' do
+          start
+
+          expect(Time.new.strftime(String.new('Potato'))).to_not be nil
+        end
+      end
+    end
+
+    context 'when allocation sampling is disabled' do
+      let(:options) { { allocation_sample_every: 0 } }
+
+      it 'does not record allocations' do
+        stub_const('CpuAndWallTimeWorkerSpec::TestStruct', Struct.new(:foo))
+
+        start
+
+        123.times { CpuAndWallTimeWorkerSpec::TestStruct.new }
+
+        cpu_and_wall_time_worker.stop
+
+        expect(samples_from_pprof(recorder.serialize!).map(&:values)).to all(include(:'alloc-samples' => 0))
+      end
+    end
   end
 
   describe 'Ractor safety' do

spec/datadog/profiling/component_spec.rb

@@ -167,6 +167,7 @@
             allocation_counting_enabled: anything,
             no_signals_workaround_enabled: :no_signals_result,
             thread_context_collector: instance_of(Datadog::Profiling::Collectors::ThreadContext),
+            allocation_sample_every: 0,
           )
 
           build_profiler_component