Add Gitlab::Profiler for profiling from a console or runner

lib/gitlab/profiler.rb

+# coding: utf-8
+module Gitlab
+  module Profiler
+    FILTERED_STRING = '[FILTERED]'.freeze
+
+    IGNORE_BACKTRACES = %w[
+      lib/gitlab/i18n.rb
+      lib/gitlab/request_context.rb
+      config/initializers
+      lib/gitlab/database/load_balancing/
+      lib/gitlab/etag_caching/
+      lib/gitlab/metrics/
+      lib/gitlab/middleware/
+      lib/gitlab/performance_bar/
+      lib/gitlab/request_profiler/
+      lib/gitlab/profiler.rb
+    ].freeze
+
+    # Takes a URL to profile (can be a fully-qualified URL, or an absolute path)
+    # and returns the ruby-prof profile result. Formatting that result is the
+    # caller's responsibility. Requests are GET requests unless post_data is
+    # passed.
+    #
+    # Optional arguments:
+    # - logger: will be used for SQL logging, including a summary at the end of
+    #   the log file of the total time spent per model class.
+    #
+    # - post_data: a string of raw POST data to use. Changes the HTTP verb to
+    #   POST.
+    #
+    # - user: a user to authenticate as. Only works if the user has a valid
+    #   personal access token.
+    #
+    # - private_token: instead of providing a user instance, the token can be
+    #   given as a string. Takes precedence over the user option.
+    def self.profile(url, logger: nil, post_data: nil, user: nil, private_token: nil)
+      app = ActionDispatch::Integration::Session.new(Rails.application)
+      verb = :get
+      headers = {}
+
+      if post_data
+        verb = :post
+        headers['Content-Type'] = 'application/json'
+      end
+
+      if user
+        private_token ||= user.personal_access_tokens.active.pluck(:token).first
+      end
+
+      headers['Private-Token'] = private_token if private_token
+      logger = create_custom_logger(logger, private_token: private_token)
+
+      RequestStore.begin!
+
+      # Make an initial call for an asset path in development mode to avoid
+      # sprockets dominating the profiler output.
+      ActionController::Base.helpers.asset_path('katex.css') if Rails.env.development?
+
+      # Rails loads internationalization files lazily the first time a
+      # translation is needed. Running this prevents this overhead from showing
+      # up in profiles.
+      ::I18n.t('.')[:test_string]
+
+      # Remove API route mounting from the profile.
+      app.get('/api/v4/users')
+
+      result = with_custom_logger(logger) do
+        RubyProf.profile { app.public_send(verb, url, post_data, headers) } # rubocop:disable GitlabSecurity/PublicSend
+      end
+
+      RequestStore.end!
+
+      log_load_times_by_model(logger)
+
+      result
+    end
+
+    def self.create_custom_logger(logger, private_token: nil)
+      return unless logger
+
+      logger.dup.tap do |new_logger|
+        new_logger.instance_variable_set(:@private_token, private_token)
+
+        class << new_logger
+          attr_reader :load_times_by_model, :private_token
+
+          def debug(message, *)
+            message.gsub!(private_token, FILTERED_STRING) if private_token
+
+            _, type, time = *message.match(/(\w+) Load \(([0-9.]+)ms\)/)
+
+            if type && time
+              @load_times_by_model ||= {}
+              @load_times_by_model[type] ||= 0
+              @load_times_by_model[type] += time.to_f
+            end
+
+            super
+
+            backtrace = Rails.backtrace_cleaner.clean(caller)
+
+            backtrace.each do |caller_line|
+              next if caller_line.match(Regexp.union(IGNORE_BACKTRACES))
+
+              stripped_caller_line = caller_line.sub("#{Rails.root}/", '')
+
+              super("  ↳ #{stripped_caller_line}")
+            end
+          end
+        end
+      end
+    end
+
+    def self.with_custom_logger(logger)
+      original_colorize_logging = ActiveSupport::LogSubscriber.colorize_logging
+      original_activerecord_logger = ActiveRecord::Base.logger
+      original_actioncontroller_logger = ActionController::Base.logger
+
+      if logger
+        ActiveSupport::LogSubscriber.colorize_logging = false
+        ActiveRecord::Base.logger = logger
+        ActionController::Base.logger = logger
+      end
+
+      result = yield
+
+      ActiveSupport::LogSubscriber.colorize_logging = original_colorize_logging
+      ActiveRecord::Base.logger = original_activerecord_logger
+      ActionController::Base.logger = original_actioncontroller_logger
+
+      result
+    end
+
+    def self.log_load_times_by_model(logger)
+      return unless logger.respond_to?(:load_times_by_model)
+
+      logger.load_times_by_model.to_a.sort_by(&:last).reverse.each do |(model, time)|
+        logger.info("#{model} total: #{time.round(2)}ms")
+      end
+    end
+  end
+end

spec/lib/gitlab/profiler_spec.rb

+require 'spec_helper'
+
+describe Gitlab::Profiler do
+  RSpec::Matchers.define_negated_matcher :not_change, :change
+
+  let(:null_logger) { Logger.new('/dev/null') }
+  let(:private_token) { 'private' }
+
+  describe '.profile' do
+    let(:app) { double(:app) }
+
+    before do
+      allow(ActionDispatch::Integration::Session).to receive(:new).and_return(app)
+      allow(app).to receive(:get)
+    end
+
+    it 'returns a profile result' do
+      expect(described_class.profile('/')).to be_an_instance_of(RubyProf::Profile)
+    end
+
+    it 'uses the custom logger given' do
+      expect(described_class).to receive(:create_custom_logger)
+                                   .with(null_logger, private_token: anything)
+                                   .and_call_original
+
+      described_class.profile('/', logger: null_logger)
+    end
+
+    it 'sends a POST request when data is passed' do
+      post_data = '{"a":1}'
+
+      expect(app).to receive(:post).with(anything, post_data, anything)
+
+      described_class.profile('/', post_data: post_data)
+    end
+
+    it 'uses the private_token for auth if given' do
+      expect(app).to receive(:get).with('/', nil, 'Private-Token' => private_token)
+      expect(app).to receive(:get).with('/api/v4/users')
+
+      described_class.profile('/', private_token: private_token)
+    end
+
+    it 'uses the user for auth if given' do
+      user = double(:user)
+      user_token = 'user'
+
+      allow(user).to receive_message_chain(:personal_access_tokens, :active, :pluck, :first).and_return(user_token)
+
+      expect(app).to receive(:get).with('/', nil, 'Private-Token' => user_token)
+      expect(app).to receive(:get).with('/api/v4/users')
+
+      described_class.profile('/', user: user)
+    end
+
+    it 'uses the private_token for auth if both it and user are set' do
+      user = double(:user)
+      user_token = 'user'
+
+      allow(user).to receive_message_chain(:personal_access_tokens, :active, :pluck, :first).and_return(user_token)
+
+      expect(app).to receive(:get).with('/', nil, 'Private-Token' => private_token)
+      expect(app).to receive(:get).with('/api/v4/users')
+
+      described_class.profile('/', user: user, private_token: private_token)
+    end
+  end
+
+  describe '.create_custom_logger' do
+    it 'does nothing when nil is passed' do
+      expect(described_class.create_custom_logger(nil)).to be_nil
+    end
+
+    context 'the new logger' do
+      let(:custom_logger) do
+        described_class.create_custom_logger(null_logger, private_token: private_token)
+      end
+
+      it 'does not affect the existing logger' do
+        expect(null_logger).not_to receive(:debug)
+        expect(custom_logger).to receive(:debug).and_call_original
+
+        custom_logger.debug('Foo')
+      end
+
+      it 'strips out the private token' do
+        expect(custom_logger).to receive(:add) do |severity, _progname, message|
+          expect(severity).to eq(Logger::DEBUG)
+          expect(message).to include('public').and include(described_class::FILTERED_STRING)
+          expect(message).not_to include(private_token)
+        end
+
+        custom_logger.debug("public #{private_token}")
+      end
+
+      it 'tracks model load times by model' do
+        custom_logger.debug('This is not a model load')
+        custom_logger.debug('User Load (1.2ms)')
+        custom_logger.debug('User Load (1.3ms)')
+        custom_logger.debug('Project Load (10.4ms)')
+
+        expect(custom_logger.load_times_by_model).to eq('User' => 2.5,
+                                                        'Project' => 10.4)
+      end
+
+      it 'logs the backtrace, ignoring lines as appropriate' do
+        # Skip Rails's backtrace cleaning.
+        allow(Rails.backtrace_cleaner).to receive(:clean, &:itself)
+
+        expect(custom_logger).to receive(:add)
+                                   .with(Logger::DEBUG,
+                                         anything,
+                                         a_string_matching(File.basename(__FILE__)))
+                                   .twice
+
+        expect(custom_logger).not_to receive(:add).with(Logger::DEBUG,
+                                                        anything,
+                                                        a_string_matching('lib/gitlab/profiler.rb'))
+
+        # Force a part of the backtrace to be in the (ignored) profiler source
+        # file.
+        described_class.with_custom_logger(nil) { custom_logger.debug('Foo') }
+      end
+    end
+  end
+
+  describe '.with_custom_logger' do
+    context 'when the logger is set' do
+      it 'uses the replacement logger for the duration of the block' do
+        expect(null_logger).to receive(:debug).and_call_original
+
+        expect { described_class.with_custom_logger(null_logger) { ActiveRecord::Base.logger.debug('foo') } }
+          .to not_change { ActiveRecord::Base.logger }
+          .and not_change { ActionController::Base.logger }
+          .and not_change { ActiveSupport::LogSubscriber.colorize_logging }
+      end
+
+      it 'returns the result of the block' do
+        expect(described_class.with_custom_logger(null_logger) { 2 }).to eq(2)
+      end
+    end
+
+    context 'when the logger is nil' do
+      it 'returns the result of the block' do
+        expect(described_class.with_custom_logger(nil) { 2 }).to eq(2)
+      end
+
+      it 'does not modify the standard Rails loggers' do
+        expect { described_class.with_custom_logger(nil) { } }
+          .to not_change { ActiveRecord::Base.logger }
+          .and not_change { ActionController::Base.logger }
+          .and not_change { ActiveSupport::LogSubscriber.colorize_logging }
+      end
+    end
+  end
+end