BigW Consortium Gitlab

instrumentation.rb 5.45 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
module Gitlab
  module Metrics
    # Module for instrumenting methods.
    #
    # This module allows instrumenting of methods without having to actually
    # alter the target code (e.g. by including modules).
    #
    # Example usage:
    #
    #     Gitlab::Metrics::Instrumentation.instrument_method(User, :by_login)
    module Instrumentation
12 13
      PROXY_IVAR = :@__gitlab_instrumentation_proxy

14 15 16 17
      def self.configure
        yield self
      end

18 19 20 21 22
      # Returns the name of the series to use for storing method calls.
      def self.series
        @series ||= "#{Metrics.series_prefix}method_calls"
      end

23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38
      # Instruments a class method.
      #
      # mod  - The module to instrument as a Module/Class.
      # name - The name of the method to instrument.
      def self.instrument_method(mod, name)
        instrument(:class, mod, name)
      end

      # Instruments an instance method.
      #
      # mod  - The module to instrument as a Module/Class.
      # name - The name of the method to instrument.
      def self.instrument_instance_method(mod, name)
        instrument(:instance, mod, name)
      end

39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61
      # 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

62
      # Instruments all public and private methods of a module.
63
      #
64 65 66 67 68
      # 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.
      #
69 70
      # mod - The module to instrument.
      def self.instrument_methods(mod)
71 72
        methods = mod.methods(false) + mod.private_methods(false)
        methods.each do |name|
73 74
          method = mod.method(name)

75 76 77 78 79
          if method.owner == mod.singleton_class
            if !block_given? || block_given? && yield(mod, method)
              instrument_method(mod, name)
            end
          end
80 81 82
        end
      end

83
      # Instruments all public and private instance methods of a module.
84
      #
85 86
      # See `instrument_methods` for more information.
      #
87 88
      # mod - The module to instrument.
      def self.instrument_instance_methods(mod)
89 90
        methods = mod.instance_methods(false) + mod.private_instance_methods(false)
        methods.each do |name|
91 92
          method = mod.instance_method(name)

93 94 95 96 97
          if method.owner == mod
            if !block_given? || block_given? && yield(mod, method)
              instrument_instance_method(mod, name)
            end
          end
98 99 100
        end
      end

101 102 103 104 105 106 107 108 109 110 111 112
      # 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

113 114 115 116 117
      # Instruments a method.
      #
      # type - The type (:class or :instance) of method to instrument.
      # mod  - The module containing the method.
      # name - The name of the method to instrument.
118 119 120
      def self.instrument(type, mod, name)
        return unless Metrics.enabled?

121 122
        name   = name.to_sym
        target = type == :instance ? mod : mod.singleton_class
123

124 125 126
        if type == :instance
          target = mod
          label  = "#{mod.name}##{name}"
127
          method = mod.instance_method(name)
128 129 130
        else
          target = mod.singleton_class
          label  = "#{mod.name}.#{name}"
131 132 133
          method = mod.method(name)
        end

134 135 136 137 138 139
        unless instrumented?(target)
          target.instance_variable_set(PROXY_IVAR, Module.new)
        end

        proxy_module = self.proxy_module(target)

140 141 142 143 144 145 146
        # 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
147
          args_signature = ''
148
        else
149
          args_signature = '*args'
150 151
        end

152
        proxy_module.class_eval <<-EOF, __FILE__, __LINE__ + 1
153
          def #{name}(#{args_signature})
154
            if trans = Gitlab::Metrics::Instrumentation.transaction
155
              trans.method_call_for(#{label.to_sym.inspect}).measure { super }
156
            else
157
              super
158
            end
159
          end
160
        EOF
161 162

        target.prepend(proxy_module)
163
      end
164 165 166 167 168 169

      # Small layer of indirection to make it easier to stub out the current
      # transaction.
      def self.transaction
        Transaction.current
      end
170 171 172
    end
  end
end