Correct arity for instrumented methods w/o argsinstrumentation-signature

This ensures that an instrumented method that doesn't take arguments
reports an arity of 0, instead of -1.

If Ruby had a proper method for finding out the required arguments of a
method (e.g. Method#required_arguments) this would not have been an
issue. Sadly the only two methods we have are Method#parameters and
Method#arity, and both are equally painful to use.

Fixes gitlab-org/gitlab-ce#12450

2 files changed, 29 insertions, 3 deletions

diff --git a/lib/gitlab/metrics/instrumentation.rb b/lib/gitlab/metrics/instrumentation.rb
index d9fce2e6758..face1921d2e 100644
--- a/lib/gitlab/metrics/instrumentation.rb
+++ b/lib/gitlab/metrics/instrumentation.rb
@@ -106,20 +106,36 @@ module Gitlab
         if type == :instance
           target = mod
           label  = "#{mod.name}##{name}"
+          method = mod.instance_method(name)
         else
           target = mod.singleton_class
           label  = "#{mod.name}.#{name}"
+          method = mod.method(name)
+        end
+
+        # 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`
+        # signature this would break things. As a result we'll make sure the
+        # generated method _only_ accepts regular arguments if the underlying
+        # method also accepts them.
+        if method.arity == 0
+          args_signature = '&block'
+        else
+          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}
 
-          def #{name}(*args, &block)
+          def #{name}(#{args_signature})
             trans = Gitlab::Metrics::Instrumentation.transaction
 
             if trans
               start    = Time.now
-              retval   = __send__(#{alias_name.inspect}, *args, &block)
+              retval   = #{send_signature}
               duration = (Time.now - start) * 1000.0
 
               if duration >= Gitlab::Metrics.method_call_threshold
@@ -132,7 +148,7 @@ module Gitlab
 
               retval
             else
-              __send__(#{alias_name.inspect}, *args, &block)
+              #{send_signature}
             end
           end
         EOF
diff --git a/spec/lib/gitlab/metrics/instrumentation_spec.rb b/spec/lib/gitlab/metrics/instrumentation_spec.rb
index 2a37cd40dde..ad4290c43bb 100644
--- a/spec/lib/gitlab/metrics/instrumentation_spec.rb
+++ b/spec/lib/gitlab/metrics/instrumentation_spec.rb
@@ -66,6 +66,16 @@ describe Gitlab::Metrics::Instrumentation do
 
         @dummy.foo
       end
+
+      it 'generates a method with the correct arity when using methods without arguments' do
+        dummy = Class.new do
+          def self.test; end
+        end
+
+        described_class.instrument_method(dummy, :test)
+
+        expect(dummy.method(:test).arity).to eq(0)
+      end
     end
 
     describe 'with metrics disabled' do