Add profile mode to extend request profilingadd-profile-mode-to-extend-request-profiling

Extend the support for "X-Profile-Token: <token>" to have an additional
header that defines type of profile we are looking for, like:
  X-Profile-Mode: execution
  X-Profile-Mode: memory

10 files changed, 221 insertions, 69 deletions

diff --git a/app/controllers/admin/requests_profiles_controller.rb b/app/controllers/admin/requests_profiles_controller.rb
index 89d4c4f18d9..1d14812186a 100644
--- a/app/controllers/admin/requests_profiles_controller.rb
+++ b/app/controllers/admin/requests_profiles_controller.rb
@@ -10,10 +10,10 @@ class Admin::RequestsProfilesController < Admin::ApplicationController
     clean_name = Rack::Utils.clean_path_info(params[:name])
     profile    = Gitlab::RequestProfiler::Profile.find(clean_name)
 
-    if profile
-      render html: profile.content.html_safe
-    else
-      redirect_to admin_requests_profiles_path, alert: 'Profile not found'
+    unless profile && profile.content_type
+      return redirect_to admin_requests_profiles_path, alert: 'Profile not found'
     end
+
+    send_file profile.file_path, type: "#{profile.content_type}; charset=utf-8", disposition: 'inline'
   end
 end
diff --git a/app/views/admin/requests_profiles/index.html.haml b/app/views/admin/requests_profiles/index.html.haml
index adfc67d66d0..86bfeef580c 100644
--- a/app/views/admin/requests_profiles/index.html.haml
+++ b/app/views/admin/requests_profiles/index.html.haml
@@ -19,7 +19,8 @@
           %ul.content-list
             - profiles.each do |profile|
               %li
-                = link_to profile.time.to_s(:long), admin_requests_profile_path(profile)
+                = link_to profile.time.to_s(:long) + ' ' + profile.profile_mode.capitalize,
+                  admin_requests_profile_path(profile)
   - else
     %p
       No profiles found
diff --git a/config/routes/admin.rb b/config/routes/admin.rb
index f609739d9fd..6f9a5552564 100644
--- a/config/routes/admin.rb
+++ b/config/routes/admin.rb
@@ -73,7 +73,7 @@ namespace :admin do
   resource :background_jobs, controller: 'background_jobs', only: [:show]
 
   resource :system_info, controller: 'system_info', only: [:show]
-  resources :requests_profiles, only: [:index, :show], param: :name, constraints: { name: /.+\.html/ }
+  resources :requests_profiles, only: [:index, :show], param: :name, constraints: { name: /.+\.(html|txt)/ }
 
   resources :projects, only: [:index]
 
diff --git a/doc/administration/monitoring/performance/img/request_profile_result.png b/doc/administration/monitoring/performance/img/request_profile_result.png
index 1b06e240fa0..3b34f207974 100644
--- a/doc/administration/monitoring/performance/img/request_profile_result.png
+++ b/doc/administration/monitoring/performance/img/request_profile_result.png
diff --git a/doc/administration/monitoring/performance/request_profiling.md b/doc/administration/monitoring/performance/request_profiling.md
index 726882fbb87..9f671e0db11 100644
--- a/doc/administration/monitoring/performance/request_profiling.md
+++ b/doc/administration/monitoring/performance/request_profiling.md
@@ -5,9 +5,9 @@
 1. Grab the profiling token from **Monitoring > Requests Profiles** admin page
   (highlighted in a blue in the image below).
   ![Profile token](img/request_profiling_token.png)
-1. Pass the header `X-Profile-Token: <token>` to the request you want to profile. You can use:
+1. Pass the header `X-Profile-Token: <token>` and `X-Profile-Mode: <mode>`(where <mode> can be `execution` or `memory`) to the request you want to profile. You can use:
     - Browser extensions. For example, [ModHeader](https://chrome.google.com/webstore/detail/modheader/idgpnmonknjnojddfkpgkljpfnnfcklj) Chrome extension.
-    - `curl`. For example, `curl --header 'X-Profile-Token: <token>' https://gitlab.example.com/group/project`.
+    - `curl`. For example, `curl --header 'X-Profile-Token: <token>' --header 'X-Profile-Mode: <mode>' https://gitlab.example.com/group/project`.
 1. Once request is finished (which will take a little longer than usual), you can
   view the profiling output from **Monitoring > Requests Profiles** admin page.
   ![Profiling output](img/request_profile_result.png)
diff --git a/lib/gitlab/request_profiler/middleware.rb b/lib/gitlab/request_profiler/middleware.rb
index 7615f6f443b..99958d7a211 100644
--- a/lib/gitlab/request_profiler/middleware.rb
+++ b/lib/gitlab/request_profiler/middleware.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'ruby-prof'
+require 'memory_profiler'
 
 module Gitlab
   module RequestProfiler
@@ -28,22 +29,73 @@ module Gitlab
       end
 
       def call_with_profiling(env)
+        case env['HTTP_X_PROFILE_MODE']
+        when 'execution', nil
+          call_with_call_stack_profiling(env)
+        when 'memory'
+          call_with_memory_profiling(env)
+        else
+          raise ActionController::BadRequest, invalid_profile_mode(env)
+        end
+      end
+
+      def invalid_profile_mode(env)
+        <<~HEREDOC
+          Invalid X-Profile-Mode: #{env['HTTP_X_PROFILE_MODE']}.
+          Supported profile mode request header:
+            - X-Profile-Mode: execution
+            - X-Profile-Mode: memory
+        HEREDOC
+      end
+
+      def call_with_call_stack_profiling(env)
         ret = nil
-        result = RubyProf::Profile.profile do
+        report = RubyProf::Profile.profile do
           ret = catch(:warden) do
             @app.call(env)
           end
         end
 
-        printer   = RubyProf::CallStackPrinter.new(result)
-        file_name = "#{env['PATH_INFO'].tr('/', '|')}_#{Time.current.to_i}.html"
+        generate_report(env, 'execution', 'html') do |file|
+          printer = RubyProf::CallStackPrinter.new(report)
+          printer.print(file)
+        end
+
+        handle_request_ret(ret)
+      end
+
+      def call_with_memory_profiling(env)
+        ret = nil
+        report = MemoryProfiler.report do
+          ret = catch(:warden) do
+            @app.call(env)
+          end
+        end
+
+        generate_report(env, 'memory', 'txt') do |file|
+          report.pretty_print(to_file: file)
+        end
+
+        handle_request_ret(ret)
+      end
+
+      def generate_report(env, report_type, extension)
+        file_name = "#{env['PATH_INFO'].tr('/', '|')}_#{Time.current.to_i}"\
+                    "_#{report_type}.#{extension}"
         file_path = "#{PROFILES_DIR}/#{file_name}"
 
         FileUtils.mkdir_p(PROFILES_DIR)
-        File.open(file_path, 'wb') do |file|
-          printer.print(file)
+
+        begin
+          File.open(file_path, 'wb') do |file|
+            yield(file)
+          end
+        rescue
+          FileUtils.rm(file_path)
         end
+      end
 
+      def handle_request_ret(ret)
         if ret.is_a?(Array)
           ret
         else
diff --git a/lib/gitlab/request_profiler/profile.rb b/lib/gitlab/request_profiler/profile.rb
index 46996ef8c51..74f2ec1d083 100644
--- a/lib/gitlab/request_profiler/profile.rb
+++ b/lib/gitlab/request_profiler/profile.rb
@@ -3,28 +3,26 @@
 module Gitlab
   module RequestProfiler
     class Profile
-      attr_reader :name, :time, :request_path
+      attr_reader :name, :time, :file_path, :request_path, :profile_mode, :type
 
       alias_method :to_param, :name
 
       def self.all
-        Dir["#{PROFILES_DIR}/*.html"].map do |path|
+        Dir["#{PROFILES_DIR}/*.{html,txt}"].map do |path|
           new(File.basename(path))
         end
       end
 
       def self.find(name)
-        name_dup = name.dup
-        name_dup << '.html' unless name.end_with?('.html')
-
-        file_path = "#{PROFILES_DIR}/#{name_dup}"
+        file_path = File.join(PROFILES_DIR, name)
         return unless File.exist?(file_path)
 
-        new(name_dup)
+        new(name)
       end
 
       def initialize(name)
         @name = name
+        @file_path = File.join(PROFILES_DIR, name)
 
         set_attributes
       end
@@ -33,12 +31,23 @@ module Gitlab
         File.read("#{PROFILES_DIR}/#{name}")
       end
 
+      def content_type
+        case type
+        when 'html'
+          'text/html'
+        when 'txt'
+          'text/plain'
+        end
+      end
+
       private
 
       def set_attributes
-        _, path, timestamp = name.split(/(.*)_(\d+)\.html$/)
+        _, path, timestamp, profile_mode, type = name.split(/(.*)_(\d+)_(.*)\.(html|txt)$/)
         @request_path      = path.tr('|', '/')
         @time              = Time.at(timestamp.to_i).utc
+        @profile_mode      = profile_mode
+        @type              = type
       end
     end
   end
diff --git a/spec/controllers/admin/requests_profiles_controller_spec.rb b/spec/controllers/admin/requests_profiles_controller_spec.rb
index 10850cb4603..289bb58c5a8 100644
--- a/spec/controllers/admin/requests_profiles_controller_spec.rb
+++ b/spec/controllers/admin/requests_profiles_controller_spec.rb
@@ -10,38 +10,63 @@ describe Admin::RequestsProfilesController do
   end
 
   describe '#show' do
-    let(:basename) { "profile_#{Time.now.to_i}.html" }
     let(:tmpdir) { Dir.mktmpdir('profiler-test') }
     let(:test_file) { File.join(tmpdir, basename) }
-    let(:profile) { Gitlab::RequestProfiler::Profile.new(basename) }
-    let(:sample_data) do
-      <<~HTML
-        <!DOCTYPE html>
-        <html>
-        <body>
-        <h1>My First Heading</h1>
-        <p>My first paragraph.</p>
-        </body>
-        </html>
-      HTML
+
+    subject do
+      get :show, params: { name: basename }
     end
 
     before do
       stub_const('Gitlab::RequestProfiler::PROFILES_DIR', tmpdir)
-      output = File.open(test_file, 'w')
-      output.write(sample_data)
-      output.close
+      File.write(test_file, sample_data)
     end
 
     after do
       File.unlink(test_file)
     end
 
-    it 'loads an HTML profile' do
-      get :show, params: { name: basename }
+    context 'when loading HTML profile' do
+      let(:basename) { "profile_#{Time.now.to_i}_execution.html" }
+
+      let(:sample_data) do
+        '<html> <body> <h1>Heading</h1> <p>paragraph.</p> </body> </html>'
+      end
+
+      it 'renders the data' do
+        subject
+
+        expect(response).to have_gitlab_http_status(200)
+        expect(response.body).to eq(sample_data)
+      end
+    end
+
+    context 'when loading TXT profile' do
+      let(:basename) { "profile_#{Time.now.to_i}_memory.txt" }
+
+      let(:sample_data) do
+        <<~TXT
+          Total allocated: 112096396 bytes (1080431 objects)
+          Total retained:  10312598 bytes (53567 objects)
+        TXT
+      end
+
+      it 'renders the data' do
+        subject
+
+        expect(response).to have_gitlab_http_status(200)
+        expect(response.body).to eq(sample_data)
+      end
+    end
+
+    context 'when loading PDF profile' do
+      let(:basename) { "profile_#{Time.now.to_i}_anything.pdf" }
+
+      let(:sample_data) { 'mocked pdf content' }
 
-      expect(response).to have_gitlab_http_status(200)
-      expect(response.body).to eq(sample_data)
+      it 'fails to render the data' do
+        expect { subject }.to raise_error(ActionController::UrlGenerationError, /No route matches.*unmatched constraints:/)
+      end
     end
   end
 end
diff --git a/spec/features/admin/admin_requests_profiles_spec.rb b/spec/features/admin/admin_requests_profiles_spec.rb
index 2503fc9067d..a962af4952b 100644
--- a/spec/features/admin/admin_requests_profiles_spec.rb
+++ b/spec/features/admin/admin_requests_profiles_spec.rb
@@ -19,42 +19,97 @@ describe 'Admin::RequestsProfilesController' do
       expect(page).to have_content("X-Profile-Token: #{Gitlab::RequestProfiler.profile_token}")
     end
 
-    it 'lists all available profiles' do
-      time1 = 1.hour.ago
-      time2 = 2.hours.ago
-      time3 = 3.hours.ago
-      profile1 = "|gitlab-org|gitlab-ce_#{time1.to_i}.html"
-      profile2 = "|gitlab-org|gitlab-ce_#{time2.to_i}.html"
-      profile3 = "|gitlab-com|infrastructure_#{time3.to_i}.html"
-
-      FileUtils.touch("#{Gitlab::RequestProfiler::PROFILES_DIR}/#{profile1}")
-      FileUtils.touch("#{Gitlab::RequestProfiler::PROFILES_DIR}/#{profile2}")
-      FileUtils.touch("#{Gitlab::RequestProfiler::PROFILES_DIR}/#{profile3}")
-
-      visit admin_requests_profiles_path
+    context 'when having multiple profiles' do
+      let(:time1) { 1.hour.ago }
+      let(:time2) { 2.hours.ago }
+
+      let(:profiles) do
+        [
+          {
+            request_path: '/gitlab-org/gitlab-ce',
+            name: "|gitlab-org|gitlab-ce_#{time1.to_i}_execution.html",
+            created: time1,
+            profile_mode: 'Execution'
+          },
+          {
+            request_path: '/gitlab-org/gitlab-ce',
+            name: "|gitlab-org|gitlab-ce_#{time2.to_i}_execution.html",
+            created: time2,
+            profile_mode: 'Execution'
+          },
+          {
+            request_path: '/gitlab-org/gitlab-ce',
+            name: "|gitlab-org|gitlab-ce_#{time1.to_i}_memory.html",
+            created: time1,
+            profile_mode: 'Memory'
+          },
+          {
+            request_path: '/gitlab-org/gitlab-ce',
+            name: "|gitlab-org|gitlab-ce_#{time2.to_i}_memory.html",
+            created: time2,
+            profile_mode: 'Memory'
+          },
+          {
+            request_path: '/gitlab-org/infrastructure',
+            name: "|gitlab-org|infrastructure_#{time1.to_i}_execution.html",
+            created: time1,
+            profile_mode: 'Execution'
+          },
+          {
+            request_path: '/gitlab-org/infrastructure',
+            name: "|gitlab-org|infrastructure_#{time2.to_i}_memory.html",
+            created: time2,
+            profile_mode: 'Memory'
+          }
+        ]
+      end
 
-      within('.card', text: '/gitlab-org/gitlab-ce') do
-        expect(page).to have_selector("a[href='#{admin_requests_profile_path(profile1)}']", text: time1.to_s(:long))
-        expect(page).to have_selector("a[href='#{admin_requests_profile_path(profile2)}']", text: time2.to_s(:long))
+      before do
+        profiles.each do |profile|
+          FileUtils.touch(File.join(Gitlab::RequestProfiler::PROFILES_DIR, profile[:name]))
+        end
       end
 
-      within('.card', text: '/gitlab-com/infrastructure') do
-        expect(page).to have_selector("a[href='#{admin_requests_profile_path(profile3)}']", text: time3.to_s(:long))
+      it 'lists all available profiles' do
+        visit admin_requests_profiles_path
+
+        profiles.each do |profile|
+          within('.card', text: profile[:request_path]) do
+            expect(page).to have_selector(
+              "a[href='#{admin_requests_profile_path(profile[:name])}']",
+              text: "#{profile[:created].to_s(:long)} #{profile[:profile_mode]}")
+          end
+        end
       end
     end
   end
 
   describe 'GET /admin/requests_profiles/:profile' do
     context 'when a profile exists' do
-      it 'displays the content of the profile' do
-        content = 'This is a request profile'
-        profile = "|gitlab-org|gitlab-ce_#{Time.now.to_i}.html"
-
+      before do
         File.write("#{Gitlab::RequestProfiler::PROFILES_DIR}/#{profile}", content)
+      end
+
+      context 'when is valid call stack profile' do
+        let(:content) { 'This is a call stack request profile' }
+        let(:profile) { "|gitlab-org|gitlab-ce_#{Time.now.to_i}_execution.html" }
+
+        it 'displays the content' do
+          visit admin_requests_profile_path(profile)
+
+          expect(page).to have_content(content)
+        end
+      end
+
+      context 'when is valid memory profile' do
+        let(:content) { 'This is a memory request profile' }
+        let(:profile) { "|gitlab-org|gitlab-ce_#{Time.now.to_i}_memory.txt" }
 
-        visit admin_requests_profile_path(profile)
+        it 'displays the content' do
+          visit admin_requests_profile_path(profile)
 
-        expect(page).to have_content(content)
+          expect(page).to have_content(content)
+        end
       end
     end
 
diff --git a/spec/requests/request_profiler_spec.rb b/spec/requests/request_profiler_spec.rb
index 75b22b1879b..851affbcf88 100644
--- a/spec/requests/request_profiler_spec.rb
+++ b/spec/requests/request_profiler_spec.rb
@@ -3,13 +3,18 @@ require 'spec_helper'
 describe 'Request Profiler' do
   let(:user) { create(:user) }
 
-  shared_examples 'profiling a request' do
+  shared_examples 'profiling a request' do |profile_type, extension|
     before do
       allow(Rails).to receive(:cache).and_return(ActiveSupport::Cache::MemoryStore.new)
       allow(RubyProf::Profile).to receive(:profile) do |&blk|
         blk.call
         RubyProf::Profile.new
       end
+      allow(MemoryProfiler).to receive(:report) do |&blk|
+        blk.call
+        MemoryProfiler.start
+        MemoryProfiler.stop
+      end
     end
 
     it 'creates a profile of the request' do
@@ -18,10 +23,11 @@ describe 'Request Profiler' do
       path    = "/#{project.full_path}"
 
       Timecop.freeze(time) do
-        get path, params: {}, headers: { 'X-Profile-Token' => Gitlab::RequestProfiler.profile_token }
+        get path, params: {}, headers: { 'X-Profile-Token' => Gitlab::RequestProfiler.profile_token, 'X-Profile-Mode' => profile_type }
       end
 
-      profile_path = "#{Gitlab.config.shared.path}/tmp/requests_profiles/#{path.tr('/', '|')}_#{time.to_i}.html"
+      profile_type = 'execution' if profile_type.nil?
+      profile_path = "#{Gitlab.config.shared.path}/tmp/requests_profiles/#{path.tr('/', '|')}_#{time.to_i}_#{profile_type}.#{extension}"
       expect(File.exist?(profile_path)).to be true
     end
 
@@ -35,10 +41,14 @@ describe 'Request Profiler' do
       login_as(user)
     end
 
-    include_examples 'profiling a request'
+    include_examples 'profiling a request', 'execution', 'html'
+    include_examples 'profiling a request', nil, 'html'
+    include_examples 'profiling a request', 'memory', 'txt'
   end
 
   context "when user is not logged-in" do
-    include_examples 'profiling a request'
+    include_examples 'profiling a request', 'execution', 'html'
+    include_examples 'profiling a request', nil, 'html'
+    include_examples 'profiling a request', 'memory', 'txt'
   end
 end