[prof-heap] Heap size and sample rate configurations

benchmarks/profiler_sample_loop_v2.rb

@@ -21,6 +21,7 @@ def create_profiler
       alloc_samples_enabled: false,
       heap_samples_enabled: false,
       heap_size_enabled: false,
+      heap_sample_every: 1,
       timeline_enabled: false,
     )
     @collector = Datadog::Profiling::Collectors::ThreadContext.new(
@@ -42,8 +43,11 @@ def thread_with_very_deep_stack(depth: 500)
 
   def run_benchmark
     Benchmark.ips do |x|
-      benchmark_time = VALIDATE_BENCHMARK_MODE ? {time: 0.01, warmup: 0} : {time: 10, warmup: 2}
-      x.config(**benchmark_time, suite: report_to_dogstatsd_if_enabled_via_environment_variable(benchmark_name: 'profiler_sample_loop_v2'))
+      benchmark_time = VALIDATE_BENCHMARK_MODE ? { time: 0.01, warmup: 0 } : { time: 10, warmup: 2 }
+      x.config(
+        **benchmark_time,
+        suite: report_to_dogstatsd_if_enabled_via_environment_variable(benchmark_name: 'profiler_sample_loop_v2')
+      )
 
       x.report("stack collector #{ENV['CONFIG']}") do
         Datadog::Profiling::Collectors::ThreadContext::Testing._native_sample(@collector, PROFILER_OVERHEAD_STACK_THREAD)
@@ -58,7 +62,9 @@ def run_benchmark
 
   def run_forever
     while true
-      1000.times { Datadog::Profiling::Collectors::ThreadContext::Testing._native_sample(@collector, PROFILER_OVERHEAD_STACK_THREAD) }
+      1000.times do
+        Datadog::Profiling::Collectors::ThreadContext::Testing._native_sample(@collector, PROFILER_OVERHEAD_STACK_THREAD)
+      end
       @recorder.serialize
       print '.'
     end

ext/ddtrace_profiling_native_extension/heap_recorder.c

@@ -93,6 +93,7 @@ typedef struct {
 static object_record* object_record_new(long, heap_record*, live_object_data);
 static void object_record_free(object_record*);
 static VALUE object_record_inspect(object_record*);
+static object_record SKIPPED_RECORD = {0};
 
 struct heap_recorder {
   // Map[key: heap_record_key*, record: heap_record*]
@@ -122,6 +123,10 @@ struct heap_recorder {
 
   // Reusable location array, implementing a flyweight pattern for things like iteration.
   ddog_prof_Location *reusable_locations;
+
+  // Sampling controls
+  uint sample_rate;
+  uint num_recordings_skipped;
 };
 static heap_record* get_or_create_heap_record(heap_recorder*, ddog_prof_Slice_Location);
 static void cleanup_heap_record_if_unused(heap_recorder*, heap_record*);
@@ -151,6 +156,7 @@ heap_recorder* heap_recorder_new(void) {
   recorder->object_records_snapshot = NULL;
   recorder->reusable_locations = ruby_xcalloc(MAX_FRAMES_LIMIT, sizeof(ddog_prof_Location));
   recorder->partial_object_record = NULL;
+  recorder->sample_rate = 1; // By default do no sampling
 
   return recorder;
 }
@@ -184,6 +190,19 @@ void heap_recorder_free(heap_recorder *heap_recorder) {
   ruby_xfree(heap_recorder);
 }
 
+void heap_recorder_set_sample_rate(heap_recorder *heap_recorder, int sample_rate) {
+  if (heap_recorder == NULL) {
+    return;
+  }
+
+  if (sample_rate < 0) {
+    rb_raise(rb_eArgError, "Heap sample rate must be a positive integer value but was %d", sample_rate);
+  }
+
+  heap_recorder->sample_rate = sample_rate;
+  heap_recorder->num_recordings_skipped = 0;
+}
+
 // WARN: Assumes this gets called before profiler is reinitialized on the fork
 void heap_recorder_after_fork(heap_recorder *heap_recorder) {
   if (heap_recorder == NULL) {
@@ -214,6 +233,14 @@ void start_heap_allocation_recording(heap_recorder *heap_recorder, VALUE new_obj
     return;
   }
 
+  if (heap_recorder->num_recordings_skipped + 1 < heap_recorder->sample_rate) {
+    heap_recorder->partial_object_record = &SKIPPED_RECORD;
+    heap_recorder->num_recordings_skipped++;
+    return;
+  }
+
+  heap_recorder->num_recordings_skipped = 0;
+
   VALUE ruby_obj_id = rb_obj_id(new_obj);
   if (!FIXNUM_P(ruby_obj_id)) {
     rb_raise(rb_eRuntimeError, "Detected a bignum object id. These are not supported by heap profiling.");
@@ -224,7 +251,7 @@ void start_heap_allocation_recording(heap_recorder *heap_recorder, VALUE new_obj
   }
 
   heap_recorder->partial_object_record = object_record_new(FIX2LONG(ruby_obj_id), NULL, (live_object_data) {
-    .weight =  weight,
+    .weight =  weight * heap_recorder->sample_rate,
     .class = alloc_class != NULL ? string_from_char_slice(*alloc_class) : NULL,
     .alloc_gen = rb_gc_count(),
   });
@@ -241,11 +268,15 @@ void end_heap_allocation_recording(struct heap_recorder *heap_recorder, ddog_pro
     // Recording ended without having been started?
     rb_raise(rb_eRuntimeError, "Ended a heap recording that was not started");
   }
-
   // From now on, mark active recording as invalid so we can short-circuit at any point and
   // not end up with a still active recording. partial_object_record still holds the object for this recording
   heap_recorder->partial_object_record = NULL;
 
+  if (partial_object_record == &SKIPPED_RECORD) {
+    // special marker when we decided to skip due to sampling
+    return;
+  }
+
   heap_record *heap_record = get_or_create_heap_record(heap_recorder, locations);
 
   // And then commit the new allocation.

ext/ddtrace_profiling_native_extension/heap_recorder.h

@@ -53,6 +53,21 @@ heap_recorder* heap_recorder_new(void);
 // Free a previously initialized heap recorder.
 void heap_recorder_free(heap_recorder *heap_recorder);
 
+// Set sample rate used by this heap recorder.
+//
+// Controls how many recordings will be ignored before committing a heap allocation and
+// the weight of the committed heap allocation.
+//
+// A value of 1 will effectively track all objects that are passed through
+// start/end_heap_allocation_recording pairs. A value of 10 will only track every 10th
+// object passed through such calls and its effective weight for the purposes of heap
+// profiling will be multiplied by 10.
+//
+// NOTE: Default is 1, i.e., track all heap allocation recordings.
+//
+// WARN: Non-positive values will lead to an exception being thrown.
+void heap_recorder_set_sample_rate(heap_recorder *heap_recorder, int sample_rate);
+
 // Do any cleanup needed after forking.
 // WARN: Assumes this gets called before profiler is reinitialized on the fork
 void heap_recorder_after_fork(heap_recorder *heap_recorder);

ext/ddtrace_profiling_native_extension/stack_recorder.c

@@ -218,6 +218,7 @@ static VALUE _native_initialize(
   VALUE alloc_samples_enabled,
   VALUE heap_samples_enabled,
   VALUE heap_sizes_enabled,
+  VALUE heap_sample_every,
   VALUE timeline_enabled
 );
 static VALUE _native_serialize(VALUE self, VALUE recorder_instance);
@@ -257,7 +258,7 @@ void stack_recorder_init(VALUE profiling_module) {
   // https://bugs.ruby-lang.org/issues/18007 for a discussion around this.
   rb_define_alloc_func(stack_recorder_class, _native_new);
 
-  rb_define_singleton_method(stack_recorder_class, "_native_initialize", _native_initialize, 6);
+  rb_define_singleton_method(stack_recorder_class, "_native_initialize", _native_initialize, 7);
   rb_define_singleton_method(stack_recorder_class, "_native_serialize",  _native_serialize, 1);
   rb_define_singleton_method(stack_recorder_class, "_native_reset_after_fork", _native_reset_after_fork, 1);
   rb_define_singleton_method(testing_module, "_native_active_slot", _native_active_slot, 1);
@@ -373,17 +374,21 @@ static VALUE _native_initialize(
   VALUE alloc_samples_enabled,
   VALUE heap_samples_enabled,
   VALUE heap_size_enabled,
+  VALUE heap_sample_every,
   VALUE timeline_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);
 
   struct stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, struct stack_recorder_state, &stack_recorder_typed_data, state);
 
+  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 : 1) -

lib/datadog/core/configuration/settings.rb

@@ -356,6 +356,21 @@ def initialize(*_)
               o.default false
             end
 
+            # Can be used to enable/disable the collection of heap size profiles.
+            #
+            # This feature is alpha and enabled by default when heap profiling is enabled.
+            #
+            # @warn To enable heap size profiling you are required to also enable allocation and heap profiling.
+            # @note Heap profiles are not yet GA in the Datadog UI, get in touch if you want to help us test it.
+            #
+            # @default `DD_PROFILING_EXPERIMENTAL_HEAP_SIZE_ENABLED` environment variable as a boolean, otherwise
+            # whatever the value of DD_PROFILING_EXPERIMENTAL_HEAP_ENABLED is.
+            option :experimental_heap_size_enabled do |o|
+              o.type :bool
+              o.env 'DD_PROFILING_EXPERIMENTAL_HEAP_SIZE_ENABLED'
+              o.default true # This gets ANDed with experimental_heap_enabled in the profiler component.
+            end
+
             # Can be used to configure the allocation sampling rate: a sample will be collected every x allocations.
             #
             # The lower the value, the more accuracy in allocation and heap tracking but the bigger the overhead. In
@@ -368,6 +383,22 @@ def initialize(*_)
               o.default 50
             end
 
+            # Can be used to configure the heap sampling rate: a heap sample will be collected for every x allocation
+            # samples.
+            #
+            # The lower the value, the more accuracy in heap tracking but the bigger the overhead. In particular, a
+            # value of 1 will track ALL allocations samples for heap profiles.
+            #
+            # The effective heap sampling rate in terms of allocations (not allocation samples) can be calculated via
+            # effective_heap_sample_rate = allocation_sample_rate * heap_sample_rate.
+            #
+            # @default `DD_PROFILING_EXPERIMENTAL_HEAP_SAMPLE_RATE` environment variable, otherwise `10`.
+            option :experimental_heap_sample_rate do |o|
+              o.type :int
+              o.env 'DD_PROFILING_EXPERIMENTAL_HEAP_SAMPLE_RATE'
+              o.default 10
+            end
+
             # Can be used to disable checking which version of `libmysqlclient` is being used by the `mysql2` gem.
             #
             # This setting is only used when the `mysql2` gem is installed.

lib/datadog/profiling/component.rb

@@ -7,7 +7,7 @@ module Component
       # Passing in a `nil` tracer is supported and will disable the following profiling features:
       # * Code Hotspots panel in the trace viewer, as well as scoping a profile down to a span
       # * Endpoint aggregation in the profiler UX, including normalization (resource per endpoint call)
-      def self.build_profiler_component(settings:, agent_settings:, optional_tracer:)
+      def self.build_profiler_component(settings:, agent_settings:, optional_tracer:) # rubocop:disable Metrics/MethodLength
         require_relative '../profiling/diagnostics/environment_logger'
 
         Profiling::Diagnostics::EnvironmentLogger.collect_and_log!
@@ -43,14 +43,18 @@ def self.build_profiler_component(settings:, agent_settings:, optional_tracer:)
         timeline_enabled = settings.profiling.advanced.experimental_timeline_enabled
         allocation_sample_every = get_allocation_sample_every(settings)
         allocation_profiling_enabled = enable_allocation_profiling?(settings, allocation_sample_every)
-        heap_profiling_enabled = enable_heap_profiling?(settings, allocation_profiling_enabled)
+        heap_sample_every = get_heap_sample_every(settings)
+        heap_profiling_enabled = enable_heap_profiling?(settings, allocation_profiling_enabled, heap_sample_every)
+        heap_size_profiling_enabled = enable_heap_size_profiling?(settings, heap_profiling_enabled)
 
         overhead_target_percentage = valid_overhead_target(settings.profiling.advanced.overhead_target_percentage)
         upload_period_seconds = [60, settings.profiling.advanced.upload_period_seconds].max
 
         recorder = build_recorder(
           allocation_profiling_enabled: allocation_profiling_enabled,
           heap_profiling_enabled: heap_profiling_enabled,
+          heap_size_profiling_enabled: heap_size_profiling_enabled,
+          heap_sample_every: heap_sample_every,
           timeline_enabled: timeline_enabled,
         )
         thread_context_collector = build_thread_context_collector(settings, recorder, optional_tracer, timeline_enabled)
@@ -67,6 +71,7 @@ def self.build_profiler_component(settings:, agent_settings:, optional_tracer:)
           no_signals_workaround_enabled: no_signals_workaround_enabled,
           timeline_enabled: timeline_enabled,
           allocation_sample_every: allocation_sample_every,
+          heap_sample_every: heap_sample_every,
         }.freeze
 
         exporter = build_profiler_exporter(settings, recorder, internal_metadata: internal_metadata)
@@ -79,13 +84,16 @@ def self.build_profiler_component(settings:, agent_settings:, optional_tracer:)
       private_class_method def self.build_recorder(
         allocation_profiling_enabled:,
         heap_profiling_enabled:,
+        heap_size_profiling_enabled:,
+        heap_sample_every:,
         timeline_enabled:
       )
         Datadog::Profiling::StackRecorder.new(
           cpu_time_enabled: RUBY_PLATFORM.include?('linux'), # Only supported on Linux currently
           alloc_samples_enabled: allocation_profiling_enabled,
           heap_samples_enabled: heap_profiling_enabled,
-          heap_size_enabled: heap_profiling_enabled,
+          heap_size_enabled: heap_size_profiling_enabled,
+          heap_sample_every: heap_sample_every,
           timeline_enabled: timeline_enabled,
         )
       end
@@ -147,6 +155,14 @@ def self.build_profiler_component(settings:, agent_settings:, optional_tracer:)
         allocation_sample_rate
       end
 
+      private_class_method def self.get_heap_sample_every(settings)
+        heap_sample_rate = settings.profiling.advanced.experimental_heap_sample_rate
+
+        raise ArgumentError, "Heap sample rate must be a positive integer. Was #{heap_sample_rate}" if heap_sample_rate <= 0
+
+        heap_sample_rate
+      end
+
       private_class_method def self.enable_allocation_profiling?(settings, allocation_sample_every)
         unless settings.profiling.advanced.experimental_allocation_enabled
           # Allocation profiling disabled, short-circuit out
@@ -200,7 +216,7 @@ def self.build_profiler_component(settings:, agent_settings:, optional_tracer:)
         true
       end
 
-      private_class_method def self.enable_heap_profiling?(settings, allocation_profiling_enabled)
+      private_class_method def self.enable_heap_profiling?(settings, allocation_profiling_enabled, heap_sample_rate)
         heap_profiling_enabled = settings.profiling.advanced.experimental_heap_enabled
 
         return false unless heap_profiling_enabled
@@ -227,7 +243,20 @@ def self.build_profiler_component(settings:, agent_settings:, optional_tracer:)
         end
 
         Datadog.logger.warn(
-          'Enabled experimental heap profiling. This is experimental, not recommended, and will increase overhead!'
+          "Enabled experimental heap profiling: heap_sample_rate=#{heap_sample_rate}. This is experimental, not " \
+          'recommended, and will increase overhead!'
+        )
+
+        true
+      end
+
+      private_class_method def self.enable_heap_size_profiling?(settings, heap_profiling_enabled)
+        heap_size_profiling_enabled = settings.profiling.advanced.experimental_heap_size_enabled
+
+        return false unless heap_profiling_enabled && heap_size_profiling_enabled
+
+        Datadog.logger.warn(
+          'Enabled experimental heap size profiling. This is experimental, not recommended, and will increase overhead!'
         )
 
         true

lib/datadog/profiling/stack_recorder.rb

@@ -6,7 +6,7 @@ module Profiling
     class StackRecorder
       def initialize(
         cpu_time_enabled:, alloc_samples_enabled:, heap_samples_enabled:, heap_size_enabled:,
-        timeline_enabled:
+        heap_sample_every:, timeline_enabled:
       )
         # This mutex works in addition to the fancy C-level mutexes we have in the native side (see the docs there).
         # It prevents multiple Ruby threads calling serialize at the same time -- something like
@@ -22,6 +22,7 @@ def initialize(
           alloc_samples_enabled,
           heap_samples_enabled,
           heap_size_enabled,
+          heap_sample_every,
           timeline_enabled,
         )
       end