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

Gitlab::Profiler.profile takes a URL to profile, and returns the ruby-prof
profile result.

This is currently only for interactive use (in a console) or from the script
runner. Note that when using this interactively, changes to the application
within that console session will be used, so this can be used to test proposed
performance changes.

2 files changed, 298 insertions, 0 deletions

diff --git a/lib/gitlab/profiler.rb b/lib/gitlab/profiler.rb
new file mode 100644
index 00000000000..95d94b3cc68
--- /dev/null
+++ b/lib/gitlab/profiler.rb
@@ -0,0 +1,142 @@
+# 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
diff --git a/spec/lib/gitlab/profiler_spec.rb b/spec/lib/gitlab/profiler_spec.rb
new file mode 100644
index 00000000000..4a43dbb2371
--- /dev/null
+++ b/spec/lib/gitlab/profiler_spec.rb
@@ -0,0 +1,156 @@
+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