Add 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

2 files changed, 74 insertions, 13 deletions

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