Instrument private/protected methods

By default instrumentation will instrument public,
protected and private methods, because usually
heavy work is done on private method or at least
that’s what facts is showing

5 files changed, 63 insertions, 13 deletions

diff --git a/CHANGELOG b/CHANGELOG
index e71a154d1d5..74fb52d3aeb 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -78,6 +78,7 @@ v 8.9.0 (unreleased)
   - Remove deprecated issues_tracker and issues_tracker_id from project model
   - Allow users to create confidential issues in private projects
   - Measure CPU time for instrumented methods
+  - Instrument private methods and private instance methods by default instead just public methods
 
 v 8.8.5 (unreleased)
   - Ensure branch cleanup regardless of whether the GitHub import process succeeds
diff --git a/config/initializers/metrics.rb b/config/initializers/metrics.rb
index f6509ee43f1..989404c6a61 100644
--- a/config/initializers/metrics.rb
+++ b/config/initializers/metrics.rb
@@ -128,11 +128,6 @@ if Gitlab::Metrics.enabled?
     config.instrument_instance_methods(API::Helpers)
 
     config.instrument_instance_methods(RepositoryCheck::SingleRepositoryWorker)
-    # Iterate over each non-super private instance method to keep up to date if
-    # internals change
-    RepositoryCheck::SingleRepositoryWorker.private_instance_methods(false).each do |method|
-      config.instrument_instance_method(RepositoryCheck::SingleRepositoryWorker, method)
-    end
   end
 
   GC::Profiler.enable
diff --git a/doc/development/instrumentation.md b/doc/development/instrumentation.md
index 50d2866ca46..6cd9b274d11 100644
--- a/doc/development/instrumentation.md
+++ b/doc/development/instrumentation.md
@@ -15,8 +15,8 @@ instrument code:
 * `instrument_instance_method`: instruments a single instance method.
 * `instrument_class_hierarchy`: given a Class this method will recursively
   instrument all sub-classes (both class and instance methods).
-* `instrument_methods`: instruments all public class methods of a Module.
-* `instrument_instance_methods`: instruments all public instance methods of a
+* `instrument_methods`: instruments all public and private class methods of a Module.
+* `instrument_instance_methods`: instruments all public and private instance methods of a
   Module.
 
 To remove the need for typing the full `Gitlab::Metrics::Instrumentation`
diff --git a/lib/gitlab/metrics/instrumentation.rb b/lib/gitlab/metrics/instrumentation.rb
index ad9ce3fa442..d81d26754fe 100644
--- a/lib/gitlab/metrics/instrumentation.rb
+++ b/lib/gitlab/metrics/instrumentation.rb
@@ -56,7 +56,7 @@ module Gitlab
         end
       end
 
-      # Instruments all public methods of a module.
+      # Instruments all public and private methods of a module.
       #
       # This method optionally takes a block that can be used to determine if a
       # method should be instrumented or not. The block is passed the receiving
@@ -65,7 +65,8 @@ module Gitlab
       #
       # mod - The module to instrument.
       def self.instrument_methods(mod)
-        mod.public_methods(false).each do |name|
+        methods = mod.methods(false) + mod.private_methods(false)
+        methods.each do |name|
           method = mod.method(name)
 
           if method.owner == mod.singleton_class
@@ -76,13 +77,14 @@ module Gitlab
         end
       end
 
-      # Instruments all public instance methods of a module.
+      # Instruments all public and private instance methods of a module.
       #
       # See `instrument_methods` for more information.
       #
       # mod - The module to instrument.
       def self.instrument_instance_methods(mod)
-        mod.public_instance_methods(false).each do |name|
+        methods = mod.instance_methods(false) + mod.private_instance_methods(false)
+        methods.each do |name|
           method = mod.instance_method(name)
 
           if method.owner == mod
diff --git a/spec/lib/gitlab/metrics/instrumentation_spec.rb b/spec/lib/gitlab/metrics/instrumentation_spec.rb
index c6e979b69a4..cdf641341cb 100644
--- a/spec/lib/gitlab/metrics/instrumentation_spec.rb
+++ b/spec/lib/gitlab/metrics/instrumentation_spec.rb
@@ -9,9 +9,31 @@ describe Gitlab::Metrics::Instrumentation do
         text
       end
 
+      class << self
+        def buzz(text = 'buzz')
+          text
+        end
+        private :buzz
+
+        def flaky(text = 'flaky')
+          text
+        end
+        protected :flaky
+      end
+
       def bar(text = 'bar')
         text
       end
+
+      def wadus(text = 'wadus')
+        text
+      end
+      private :wadus
+
+      def chaf(text = 'chaf')
+        text
+      end
+      protected :chaf
     end
 
     allow(@dummy).to receive(:name).and_return('Dummy')
@@ -208,6 +230,21 @@ describe Gitlab::Metrics::Instrumentation do
       described_class.instrument_methods(@dummy)
 
       expect(described_class.instrumented?(@dummy.singleton_class)).to eq(true)
+      expect(@dummy.method(:foo).source_location.first).to match(/instrumentation\.rb/)
+    end
+
+    it 'instruments all protected class methods' do
+      described_class.instrument_methods(@dummy)
+
+      expect(described_class.instrumented?(@dummy.singleton_class)).to eq(true)
+      expect(@dummy.method(:flaky).source_location.first).to match(/instrumentation\.rb/)
+    end
+
+    it 'instruments all private instance methods' do
+      described_class.instrument_methods(@dummy)
+
+      expect(described_class.instrumented?(@dummy.singleton_class)).to eq(true)
+      expect(@dummy.method(:buzz).source_location.first).to match(/instrumentation\.rb/)
     end
 
     it 'only instruments methods directly defined in the module' do
@@ -241,6 +278,21 @@ describe Gitlab::Metrics::Instrumentation do
       described_class.instrument_instance_methods(@dummy)
 
       expect(described_class.instrumented?(@dummy)).to eq(true)
+      expect(@dummy.new.method(:bar).source_location.first).to match(/instrumentation\.rb/)
+    end
+
+    it 'instruments all protected instance methods' do
+      described_class.instrument_instance_methods(@dummy)
+
+      expect(described_class.instrumented?(@dummy)).to eq(true)
+      expect(@dummy.new.method(:chaf).source_location.first).to match(/instrumentation\.rb/)
+    end
+
+    it 'instruments all private instance methods' do
+      described_class.instrument_instance_methods(@dummy)
+
+      expect(described_class.instrumented?(@dummy)).to eq(true)
+      expect(@dummy.new.method(:wadus).source_location.first).to match(/instrumentation\.rb/)
     end
 
     it 'only instruments methods directly defined in the module' do
@@ -253,7 +305,7 @@ describe Gitlab::Metrics::Instrumentation do
 
       described_class.instrument_instance_methods(@dummy)
 
-      expect(@dummy.method_defined?(:_original_kittens)).to eq(false)
+      expect(@dummy.new.method(:kittens).source_location.first).not_to match(/instrumentation\.rb/)
     end
 
     it 'can take a block to determine if a method should be instrumented' do
@@ -261,7 +313,7 @@ describe Gitlab::Metrics::Instrumentation do
         false
       end
 
-      expect(@dummy.method_defined?(:_original_bar)).to eq(false)
+      expect(@dummy.new.method(:bar).source_location.first).not_to match(/instrumentation\.rb/)
     end
   end
 end