Support for instrumenting class hierarchies

This will be used to (for example) instrument all ActiveRecord models.

2 files changed, 56 insertions, 0 deletions

diff --git a/lib/gitlab/metrics/instrumentation.rb b/lib/gitlab/metrics/instrumentation.rb
index ca2dffbc46a..06fc2f25948 100644
--- a/lib/gitlab/metrics/instrumentation.rb
+++ b/lib/gitlab/metrics/instrumentation.rb
@@ -31,6 +31,29 @@ module Gitlab
         instrument(:instance, mod, name)
       end
 
+      # Recursively instruments all subclasses of the given root module.
+      #
+      # This can be used to for example instrument all ActiveRecord models (as
+      # these all inherit from ActiveRecord::Base).
+      #
+      # This method can optionally take a block to pass to `instrument_methods`
+      # and `instrument_instance_methods`.
+      #
+      # root - The root module for which to instrument subclasses. The root
+      #        module itself is not instrumented.
+      def self.instrument_class_hierarchy(root, &block)
+        visit = root.subclasses
+
+        until visit.empty?
+          klass = visit.pop
+
+          instrument_methods(klass, &block)
+          instrument_instance_methods(klass, &block)
+
+          klass.subclasses.each { |c| visit << c }
+        end
+      end
+
       # Instruments all public methods of a module.
       #
       # This method optionally takes a block that can be used to determine if a
diff --git a/spec/lib/gitlab/metrics/instrumentation_spec.rb b/spec/lib/gitlab/metrics/instrumentation_spec.rb
index 71d7209db0f..a7eab9d11cc 100644
--- a/spec/lib/gitlab/metrics/instrumentation_spec.rb
+++ b/spec/lib/gitlab/metrics/instrumentation_spec.rb
@@ -133,6 +133,39 @@ describe Gitlab::Metrics::Instrumentation do
     end
   end
 
+  describe '.instrument_class_hierarchy' do
+    before do
+      allow(Gitlab::Metrics).to receive(:enabled?).and_return(true)
+
+      @child1 = Class.new(@dummy) do
+        def self.child1_foo; end
+        def child1_bar; end
+      end
+
+      @child2 = Class.new(@child1) do
+        def self.child2_foo; end
+        def child2_bar; end
+      end
+    end
+
+    it 'recursively instruments a class hierarchy' do
+      described_class.instrument_class_hierarchy(@dummy)
+
+      expect(@child1).to respond_to(:_original_child1_foo)
+      expect(@child2).to respond_to(:_original_child2_foo)
+
+      expect(@child1.method_defined?(:_original_child1_bar)).to eq(true)
+      expect(@child2.method_defined?(:_original_child2_bar)).to eq(true)
+    end
+
+    it 'does not instrument the root module' do
+      described_class.instrument_class_hierarchy(@dummy)
+
+      expect(@dummy).to_not respond_to(:_original_foo)
+      expect(@dummy.method_defined?(:_original_bar)).to eq(false)
+    end
+  end
+
   describe '.instrument_methods' do
     before do
       allow(Gitlab::Metrics).to receive(:enabled?).and_return(true)