Allow filtering of what methods to instrument

This makes it possible to determine if a method should be instrumented
or not using a block.

2 files changed, 33 insertions, 2 deletions

diff --git a/lib/gitlab/metrics/instrumentation.rb b/lib/gitlab/metrics/instrumentation.rb
index cf22aa93cdd..91e09694cd8 100644
--- a/lib/gitlab/metrics/instrumentation.rb
+++ b/lib/gitlab/metrics/instrumentation.rb
@@ -33,23 +33,38 @@ module Gitlab
 
       # Instruments all public 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
+      # module and an UnboundMethod. If the block returns a non truthy value the
+      # method is not instrumented.
+      #
       # mod - The module to instrument.
       def self.instrument_methods(mod)
         mod.public_methods(false).each do |name|
           method = mod.method(name)
 
-          instrument_method(mod, name) if method.owner == mod.singleton_class
+          if method.owner == mod.singleton_class
+            if !block_given? || block_given? && yield(mod, method)
+              instrument_method(mod, name)
+            end
+          end
         end
       end
 
       # Instruments all public 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|
           method = mod.instance_method(name)
 
-          instrument_instance_method(mod, name) if method.owner == mod
+          if method.owner == mod
+            if !block_given? || block_given? && yield(mod, method)
+              instrument_instance_method(mod, name)
+            end
+          end
         end
       end
 
diff --git a/spec/lib/gitlab/metrics/instrumentation_spec.rb b/spec/lib/gitlab/metrics/instrumentation_spec.rb
index 80dc160ebd2..5fe7a369cba 100644
--- a/spec/lib/gitlab/metrics/instrumentation_spec.rb
+++ b/spec/lib/gitlab/metrics/instrumentation_spec.rb
@@ -132,6 +132,14 @@ describe Gitlab::Metrics::Instrumentation do
 
       expect(@dummy).to_not respond_to(:_original_kittens)
     end
+
+    it 'can take a block to determine if a method should be instrumented' do
+      described_class.instrument_methods(@dummy) do
+        false
+      end
+
+      expect(@dummy).to_not respond_to(:_original_foo)
+    end
   end
 
   describe '.instrument_instance_methods' do
@@ -157,5 +165,13 @@ describe Gitlab::Metrics::Instrumentation do
 
       expect(@dummy.method_defined?(:_original_kittens)).to eq(false)
     end
+
+    it 'can take a block to determine if a method should be instrumented' do
+      described_class.instrument_instance_methods(@dummy) do
+        false
+      end
+
+      expect(@dummy.method_defined?(:_original_bar)).to eq(false)
+    end
   end
 end