Improve pprof encoding performance by replacing MessageSet with Ruby hash

One of the things I've learned since I started doing some performance
work in the profiler, is how expensive is to hash arrays of items,
as `MessageSet` (`ObjectSet`) does.

While looking for improvements to pprof encoding, I realized that the
locations `MessageSet` was implementing a hashmap of
`Profiling::BacktraceLocation`s to pprof `Location` objects in a
rather roundabout way -- it used the full contents of the
`BacktraceLocation` objects as a key, rather than the object itself.

By replacing this arrangement with a regular Ruby hash that used
the `BacktraceLocation` as a key, I was able to quite improve the
performance of pprof encoding.

In the <benchmarks/profiler_submission.rb> benchmark
(Ruby 2.7.3/macOS 10.15.7/i7-1068NG7) I'm getting a > 2x speedup
when compared with current master (or 3.4x speedup over dd-trace-rb 0.49.0):

```
exporter locations-hashmap
                          9.241  (± 0.0%) i/s -     93.000  in  10.081019s

Comparison:
exporter locations-hashmap:        9.2 i/s
exporter cached-build-location-v2 (current master):        4.4 i/s - 2.12x  (± 0.00) slower
exporter baseline (before any work):        2.7 i/s - 3.42x  (± 0.00) slower
```

Author: ivoanjo

benchmarks/profiler_submission.rb

@@ -60,13 +60,16 @@ def run_benchmark
         run_once
       end
 
-      x.save! 'profiler-submission-results.ipsbench'
+      x.save! 'profiler-submission-results.json'
       x.compare!
     end
   end
 
   def run_forever
-    run_once while true
+    while true
+      run_once
+      print '.'
+    end
   end
 
   def run_once

lib/ddtrace/profiling/pprof/builder.rb

@@ -21,16 +21,25 @@ class Builder
 
         def initialize
           @functions = MessageSet.new(1)
-          @locations = MessageSet.new(1)
+          @locations = initialize_locations_hash
           @mappings = MessageSet.new(1)
           @sample_types = MessageSet.new
           @samples = []
           @string_table = StringTable.new
-          # Cache these procs, since it's pretty expensive to keep recreating them
-          @build_location = method(:build_location).to_proc
+
+          # Cache this proc, since it's pretty expensive to keep recreating it
           @build_function = method(:build_function).to_proc
         end
 
+        # The locations hash maps unique BacktraceLocation instances to their corresponding pprof Location objects;
+        # there's a 1:1 correspondence, since BacktraceLocations were already deduped
+        def initialize_locations_hash
+          sequence = Utils::Sequence.new(1)
+          Hash.new do |locations_hash, backtrace_location|
+            locations_hash[backtrace_location] = build_location(sequence.next, backtrace_location)
+          end
+        end
+
         def encode_profile(profile)
           Perftools::Profiles::Profile.encode(profile).force_encoding(DEFAULT_ENCODING)
         end
@@ -40,7 +49,7 @@ def build_profile
             sample_type: @sample_types.messages,
             sample: @samples,
             mapping: @mappings.messages,
-            location: @locations.messages,
+            location: @locations.values,
             function: @functions.messages,
             string_table: @string_table.strings
           )
@@ -54,45 +63,29 @@ def build_value_type(type, unit)
         end
 
         def build_locations(backtrace_locations, length)
-          locations = backtrace_locations.collect do |backtrace_location|
-            @locations.fetch(
-              # Filename
-              backtrace_location.path,
-              # Line number
-              backtrace_location.lineno,
-              # Function name
-              backtrace_location.base_label,
-              # Build function
-              &@build_location
-            )
-          end
+          locations = backtrace_locations.collect { |backtrace_location| @locations[backtrace_location] }
 
           omitted = length - backtrace_locations.length
 
           # Add placeholder stack frame if frames were truncated
           if omitted > 0
             desc = omitted == 1 ? DESC_FRAME_OMITTED : DESC_FRAMES_OMITTED
-            locations << @locations.fetch(
-              ''.freeze,
-              0,
-              "#{omitted} #{desc}",
-              &@build_location
-            )
+            locations << @locations[Profiling::BacktraceLocation.new(''.freeze, 0, "#{omitted} #{desc}")]
           end
 
           locations
         end
 
-        def build_location(id, filename, line_number, function_name = nil)
+        def build_location(id, backtrace_location)
           Perftools::Profiles::Location.new(
             id: id,
             line: [build_line(
               @functions.fetch(
-                filename,
-                function_name,
+                backtrace_location.path,
+                backtrace_location.base_label,
                 &@build_function
               ).id,
-              line_number
+              backtrace_location.lineno
             )]
           )
         end

spec/ddtrace/profiling/integration_spec.rb

@@ -168,14 +168,7 @@ def string_id_for(string)
 
     include_context 'StackSample events' do
       def stack_frame_to_location_id(backtrace_location)
-        template.builder.locations.fetch(
-          # Filename
-          backtrace_location.path,
-          # Line number
-          backtrace_location.lineno,
-          # Function name
-          backtrace_location.base_label
-        ) { raise 'Unknown stack frame!' }.id
+        template.builder.locations.fetch(backtrace_location) { raise 'Unknown stack frame!' }.id
       end
 
       def stack_frame_to_function_id(backtrace_location)

spec/ddtrace/profiling/pprof/builder_spec.rb

@@ -21,19 +21,6 @@ def string_id_for(string)
     builder.string_table.fetch(string)
   end
 
-  describe '#initialize' do
-    it do
-      is_expected.to have_attributes(
-        functions: kind_of(Datadog::Profiling::Pprof::MessageSet),
-        locations: kind_of(Datadog::Profiling::Pprof::MessageSet),
-        mappings: kind_of(Datadog::Profiling::Pprof::MessageSet),
-        sample_types: kind_of(Datadog::Profiling::Pprof::MessageSet),
-        samples: [],
-        string_table: kind_of(Datadog::Profiling::Pprof::StringTable)
-      )
-    end
-  end
-
   describe '#encode_profile' do
     subject(:build_profile) { builder.encode_profile(profile) }
 
@@ -66,7 +53,7 @@ def string_id_for(string)
           sample_type: builder.sample_types.messages,
           sample: builder.samples,
           mapping: builder.mappings.messages,
-          location: builder.locations.messages,
+          location: builder.locations.values,
           function: builder.functions.messages,
           string_table: builder.string_table.strings
         )
@@ -98,27 +85,13 @@ def string_id_for(string)
     let(:backtrace_locations) { Thread.current.backtrace_locations.first(3) }
     let(:length) { backtrace_locations.length }
 
-    let(:expected_locations) do
-      backtrace_locations.each_with_object({}) do |loc, map|
-        key = [loc.path, loc.lineno, loc.base_label]
-        # Use double instead of instance_double because protobuf doesn't define verifiable methods
-        map[key] = double('Perftools::Profiles::Location')
-      end
-    end
-
-    before do
-      expect(builder.locations).to receive(:fetch).at_least(backtrace_locations.length).times do |*args, &block|
-        expect(expected_locations).to include(args)
-        expect(block.source_location).to eq(builder.method(:build_location).source_location)
-        expected_locations[args]
-      end
-    end
-
     context 'given backtrace locations matching length' do
-      it do
-        is_expected.to be_a_kind_of(Array)
-        is_expected.to have(backtrace_locations.length).items
-        is_expected.to include(*expected_locations.values)
+      it { is_expected.to be_a_kind_of(Array) }
+      it { is_expected.to have(backtrace_locations.length).items }
+
+      it 'converts the BacktraceLocations to matching Perftools::Profiles::Location objects' do
+        # Lines are the simplest to compare, since they aren't converted to ids
+        expect(build_locations.map { |location| location.to_h[:line].first[:line] }).to eq backtrace_locations.map(&:lineno)
       end
     end
 
@@ -128,56 +101,47 @@ def string_id_for(string)
       let(:omitted_location) { double('Perftools::Profiles::Location') }
 
       before do
-        expected_locations[['', 0, "#{omitted} #{described_class::DESC_FRAMES_OMITTED}"]] = omitted_location
+        omitted_backtrace_location =
+          Datadog::Profiling::BacktraceLocation.new('', 0, "#{omitted} #{described_class::DESC_FRAMES_OMITTED}")
+
+        builder.locations[omitted_backtrace_location] = omitted_location
       end
 
-      it do
-        is_expected.to be_a_kind_of(Array)
-        is_expected.to have(backtrace_locations.length + 1).items
-        is_expected.to include(*expected_locations.values)
-        expect(build_locations.last).to be(omitted_location)
+      it { is_expected.to have(backtrace_locations.length + 1).items }
+
+      it 'converts the BacktraceLocations to matching Perftools::Profiles::Location objects' do
+        expect(build_locations[0..-2].map { |location| location.to_h[:line].first[:line] })
+          .to eq backtrace_locations.map(&:lineno)
+      end
+
+      it 'adds a placeholder frame as the last element to indicate the omitted frames' do
+        expect(build_locations.last).to be omitted_location
       end
     end
   end
 
   describe '#build_location' do
-    subject(:build_location) { builder.build_location(id, filename, line_number) }
+    subject(:build_location) do
+      builder.build_location(location_id, Datadog::Profiling::BacktraceLocation.new(function_name, line_number, filename))
+    end
 
-    let(:id) { id_sequence.next }
-    let(:filename) { double('filename') }
+    let(:location_id) { rand_int }
     let(:line_number) { rand_int }
+    let(:function_name) { 'the_function_name' }
+    let(:filename) { 'the_file_name.rb' }
 
-    # Use double instead of instance_double because protobuf doesn't define verifiable methods
-    let(:function) { double('Perftools::Profiles::Function', id: id_sequence.next) }
-
-    before do
-      expect(builder.functions).to receive(:fetch) do |*args, &block|
-        expect(args).to eq([filename, nil])
-        expect(block.source_location).to eq(builder.method(:build_function).source_location)
-        function
-      end
-    end
+    it 'creates a new Perftools::Profiles::Location object with the contents of the BacktraceLocation' do
+      function = double('Function', id: rand_int)
 
-    context 'given no function name' do
-      it do
-        is_expected.to be_a_kind_of(Perftools::Profiles::Location)
-        is_expected.to have_attributes(
-          id: id,
-          line: array_including(kind_of(Perftools::Profiles::Line))
-        )
-        expect(build_location.line).to have(1).items
-      end
-
-      describe 'returns a Location with Line that' do
-        subject(:line) { build_location.line.first }
+      expect(Perftools::Profiles::Function)
+        .to receive(:new).with(hash_including(filename: string_id_for(filename), name: string_id_for(function_name)))
+                         .and_return(function)
 
-        it do
-          is_expected.to have_attributes(
-            function_id: function.id,
-            line: line_number
-          )
-        end
-      end
+      expect(build_location).to be_a_kind_of(Perftools::Profiles::Location)
+      expect(build_location.to_h).to match(hash_including(id: location_id,
+                                                          line: [{
+                                                            function_id: function.id, line: line_number
+                                                          }]))
     end
   end
 

spec/ddtrace/profiling/pprof/stack_sample_spec.rb

@@ -314,7 +314,7 @@ def string_id_for(string)
 
         it 'each map to a Location on the profile' do
           locations.each do |id|
-            expect(builder.locations.messages[id - 1])
+            expect(builder.locations.values[id - 1])
               .to be_kind_of(Perftools::Profiles::Location)
           end
         end

spec/ddtrace/profiling/spec_helper.rb

@@ -75,7 +75,9 @@ def build_stack_sample(
 
     Datadog::Profiling::Events::StackSample.new(
       nil,
-      locations,
+      locations.map do |location|
+        Datadog::Profiling::BacktraceLocation.new(location.base_label, location.lineno, location.path)
+      end,
       locations.length,
       thread_id || rand(1e9),
       trace_id || rand(1e9),