Track object counts using the "allocations" Gem

This allows us to track the counts of actual classes instead of "T_XXX"
nodes. This is only enabled on CRuby as it uses CRuby specific APIs.

5 files changed, 33 insertions, 3 deletions

diff --git a/Gemfile b/Gemfile
index e9e5c7df075..94660ab5217 100644
--- a/Gemfile
+++ b/Gemfile
@@ -210,6 +210,7 @@ gem 'net-ssh',            '~> 3.0.1'
 
 # Metrics
 group :metrics do
+  gem 'allocations', '~> 1.0', require: false, platform: :mri
   gem 'method_source', '~> 0.8', require: false
   gem 'influxdb', '~> 0.2', require: false
   gem 'connection_pool', '~> 2.0', require: false
diff --git a/Gemfile.lock b/Gemfile.lock
index 3f301111224..13e8168ee8a 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -49,6 +49,7 @@ GEM
     addressable (2.3.8)
     after_commit_queue (1.3.0)
       activerecord (>= 3.0)
+    allocations (1.0.1)
     annotate (2.6.10)
       activerecord (>= 3.2, <= 4.3)
       rake (~> 10.4)
@@ -818,6 +819,7 @@ DEPENDENCIES
   acts-as-taggable-on (~> 3.4)
   addressable (~> 2.3.8)
   after_commit_queue
+  allocations (~> 1.0)
   annotate (~> 2.6.0)
   asana (~> 0.4.0)
   asciidoctor (~> 1.5.2)
diff --git a/lib/gitlab/metrics.rb b/lib/gitlab/metrics.rb
index 4b92c3244fa..ce89be636d3 100644
--- a/lib/gitlab/metrics.rb
+++ b/lib/gitlab/metrics.rb
@@ -16,6 +16,10 @@ module Gitlab
       !!Settings.metrics['enabled']
     end
 
+    def self.mri?
+      RUBY_ENGINE == 'ruby'
+    end
+
     def self.method_call_threshold
       Settings.metrics['method_call_threshold'] || 10
     end
diff --git a/lib/gitlab/metrics/sampler.rb b/lib/gitlab/metrics/sampler.rb
index 03afa6324dd..828ee1f8c62 100644
--- a/lib/gitlab/metrics/sampler.rb
+++ b/lib/gitlab/metrics/sampler.rb
@@ -13,6 +13,12 @@ module Gitlab
 
         @last_minor_gc = Delta.new(GC.stat[:minor_gc_count])
         @last_major_gc = Delta.new(GC.stat[:major_gc_count])
+
+        if Gitlab::Metrics.mri?
+          require 'allocations'
+
+          Allocations.start
+        end
       end
 
       def start
@@ -52,9 +58,22 @@ module Gitlab
           new('file_descriptors', value: System.file_descriptor_count)
       end
 
-      def sample_objects
-        ObjectSpace.count_objects.each do |type, count|
-          @metrics << Metric.new('object_counts', { count: count }, type: type)
+      if Metrics.mri?
+        def sample_objects
+          sample = Allocations.to_hash
+          counts = sample.each_with_object({}) do |(klass, count), hash|
+            hash[klass.name] = count
+          end
+
+          # Symbols aren't allocated so we'll need to add those manually.
+          counts['Symbol'] = Symbol.all_symbols.length
+
+          counts.each do |name, count|
+            @metrics << Metric.new('object_counts', { count: count }, type: name)
+          end
+        end
+      else
+        def sample_objects
         end
       end
 
diff --git a/spec/lib/gitlab/metrics/sampler_spec.rb b/spec/lib/gitlab/metrics/sampler_spec.rb
index 319f287178d..69376c0b79b 100644
--- a/spec/lib/gitlab/metrics/sampler_spec.rb
+++ b/spec/lib/gitlab/metrics/sampler_spec.rb
@@ -3,6 +3,10 @@ require 'spec_helper'
 describe Gitlab::Metrics::Sampler do
   let(:sampler) { described_class.new(5) }
 
+  after do
+    Allocations.stop if Gitlab::Metrics.mri?
+  end
+
   describe '#start' do
     it 'gathers a sample at a given interval' do
       expect(sampler).to receive(:sleep).with(5)