Use Module#prepend for method instrumentationfix-method-instrumentation

By using Module#prepend we can define a Module containing all proxy
methods. This removes the need for setting up crazy method alias chains
and in turn prevents us from having to deal with all that madness (e.g.
methods calling each other recursively).

Fixes gitlab-org/gitlab-ce#15281

3 files changed, 65 insertions, 24 deletions

diff --git a/CHANGELOG b/CHANGELOG
index ede0c00e902..c29c207ca70 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,6 +1,7 @@
 Please view this file on the master branch, on stable branches it's out of date.
 
 v 8.7.0 (unreleased)
+  - Method instrumentation now uses Module#prepend instead of aliasing methods
   - The Projects::HousekeepingService class has extra instrumentation
   - All service classes (those residing in app/services) are now instrumented
   - Developers can now add custom tags to transactions
diff --git a/lib/gitlab/metrics/instrumentation.rb b/lib/gitlab/metrics/instrumentation.rb
index face1921d2e..708ef79f304 100644
--- a/lib/gitlab/metrics/instrumentation.rb
+++ b/lib/gitlab/metrics/instrumentation.rb
@@ -11,6 +11,8 @@ module Gitlab
     module Instrumentation
       SERIES = 'method_calls'
 
+      PROXY_IVAR = :@__gitlab_instrumentation_proxy
+
       def self.configure
         yield self
       end
@@ -91,6 +93,18 @@ module Gitlab
         end
       end
 
+      # Returns true if a module is instrumented.
+      #
+      # mod - The module to check
+      def self.instrumented?(mod)
+        mod.instance_variable_defined?(PROXY_IVAR)
+      end
+
+      # Returns the proxy module (if any) of `mod`.
+      def self.proxy_module(mod)
+        mod.instance_variable_get(PROXY_IVAR)
+      end
+
       # Instruments a method.
       #
       # type - The type (:class or :instance) of method to instrument.
@@ -99,9 +113,8 @@ module Gitlab
       def self.instrument(type, mod, name)
         return unless Metrics.enabled?
 
-        name       = name.to_sym
-        alias_name = :"_original_#{name}"
-        target     = type == :instance ? mod : mod.singleton_class
+        name   = name.to_sym
+        target = type == :instance ? mod : mod.singleton_class
 
         if type == :instance
           target = mod
@@ -113,6 +126,12 @@ module Gitlab
           method = mod.method(name)
         end
 
+        unless instrumented?(target)
+          target.instance_variable_set(PROXY_IVAR, Module.new)
+        end
+
+        proxy_module = self.proxy_module(target)
+
         # Some code out there (e.g. the "state_machine" Gem) checks the arity of
         # a method to make sure it only passes arguments when the method expects
         # any. If we were to always overwrite a method to take an `*args`
@@ -125,17 +144,13 @@ module Gitlab
           args_signature = '*args, &block'
         end
 
-        send_signature = "__send__(#{alias_name.inspect}, #{args_signature})"
-
-        target.class_eval <<-EOF, __FILE__, __LINE__ + 1
-          alias_method #{alias_name.inspect}, #{name.inspect}
-
+        proxy_module.class_eval <<-EOF, __FILE__, __LINE__ + 1
           def #{name}(#{args_signature})
             trans = Gitlab::Metrics::Instrumentation.transaction
 
             if trans
               start    = Time.now
-              retval   = #{send_signature}
+              retval   = super
               duration = (Time.now - start) * 1000.0
 
               if duration >= Gitlab::Metrics.method_call_threshold
@@ -148,10 +163,12 @@ module Gitlab
 
               retval
             else
-              #{send_signature}
+              super
             end
           end
         EOF
+
+        target.prepend(proxy_module)
       end
 
       # Small layer of indirection to make it easier to stub out the current
diff --git a/spec/lib/gitlab/metrics/instrumentation_spec.rb b/spec/lib/gitlab/metrics/instrumentation_spec.rb
index ad4290c43bb..5c885a7a982 100644
--- a/spec/lib/gitlab/metrics/instrumentation_spec.rb
+++ b/spec/lib/gitlab/metrics/instrumentation_spec.rb
@@ -33,8 +33,16 @@ describe Gitlab::Metrics::Instrumentation do
         described_class.instrument_method(@dummy, :foo)
       end
 
-      it 'renames the original method' do
-        expect(@dummy).to respond_to(:_original_foo)
+      it 'instruments the Class' do
+        target = @dummy.singleton_class
+
+        expect(described_class.instrumented?(target)).to eq(true)
+      end
+
+      it 'defines a proxy method' do
+        mod = described_class.proxy_module(@dummy.singleton_class)
+
+        expect(mod.method_defined?(:foo)).to eq(true)
       end
 
       it 'calls the instrumented method with the correct arguments' do
@@ -76,6 +84,14 @@ describe Gitlab::Metrics::Instrumentation do
 
         expect(dummy.method(:test).arity).to eq(0)
       end
+
+      describe 'when a module is instrumented multiple times' do
+        it 'calls the instrumented method with the correct arguments' do
+          described_class.instrument_method(@dummy, :foo)
+
+          expect(@dummy.foo).to eq('foo')
+        end
+      end
     end
 
     describe 'with metrics disabled' do
@@ -86,7 +102,9 @@ describe Gitlab::Metrics::Instrumentation do
       it 'does not instrument the method' do
         described_class.instrument_method(@dummy, :foo)
 
-        expect(@dummy).to_not respond_to(:_original_foo)
+        target = @dummy.singleton_class
+
+        expect(described_class.instrumented?(target)).to eq(false)
       end
     end
   end
@@ -100,8 +118,14 @@ describe Gitlab::Metrics::Instrumentation do
           instrument_instance_method(@dummy, :bar)
       end
 
-      it 'renames the original method' do
-        expect(@dummy.method_defined?(:_original_bar)).to eq(true)
+      it 'instruments instances of the Class' do
+        expect(described_class.instrumented?(@dummy)).to eq(true)
+      end
+
+      it 'defines a proxy method' do
+        mod = described_class.proxy_module(@dummy)
+
+        expect(mod.method_defined?(:bar)).to eq(true)
       end
 
       it 'calls the instrumented method with the correct arguments' do
@@ -144,7 +168,7 @@ describe Gitlab::Metrics::Instrumentation do
         described_class.
           instrument_instance_method(@dummy, :bar)
 
-        expect(@dummy.method_defined?(:_original_bar)).to eq(false)
+        expect(described_class.instrumented?(@dummy)).to eq(false)
       end
     end
   end
@@ -167,18 +191,17 @@ describe Gitlab::Metrics::Instrumentation do
     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(described_class.instrumented?(@child1.singleton_class)).to eq(true)
+      expect(described_class.instrumented?(@child2.singleton_class)).to eq(true)
 
-      expect(@child1.method_defined?(:_original_child1_bar)).to eq(true)
-      expect(@child2.method_defined?(:_original_child2_bar)).to eq(true)
+      expect(described_class.instrumented?(@child1)).to eq(true)
+      expect(described_class.instrumented?(@child2)).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)
+      expect(described_class.instrumented?(@dummy)).to eq(false)
     end
   end
 
@@ -190,7 +213,7 @@ describe Gitlab::Metrics::Instrumentation do
     it 'instruments all public class methods' do
       described_class.instrument_methods(@dummy)
 
-      expect(@dummy).to respond_to(:_original_foo)
+      expect(described_class.instrumented?(@dummy.singleton_class)).to eq(true)
     end
 
     it 'only instruments methods directly defined in the module' do
@@ -223,7 +246,7 @@ describe Gitlab::Metrics::Instrumentation do
     it 'instruments all public instance methods' do
       described_class.instrument_instance_methods(@dummy)
 
-      expect(@dummy.method_defined?(:_original_bar)).to eq(true)
+      expect(described_class.instrumented?(@dummy)).to eq(true)
     end
 
     it 'only instruments methods directly defined in the module' do