Add latest changes from gitlab-org/gitlab@master

2 files changed, 150 insertions, 0 deletions

diff --git a/lib/gitlab/request_profiler/middleware.rb b/lib/gitlab/request_profiler/middleware.rb
new file mode 100644
index 00000000000..acdf8d4541f
--- /dev/null
+++ b/lib/gitlab/request_profiler/middleware.rb
@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+require 'ruby-prof'
+require 'memory_profiler'
+
+module Gitlab
+  module RequestProfiler
+    class Middleware
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        if profile?(env)
+          call_with_profiling(env)
+        else
+          @app.call(env)
+        end
+      end
+
+      def profile?(env)
+        header_token = env['HTTP_X_PROFILE_TOKEN']
+        return unless header_token.present?
+
+        profile_token = Gitlab::RequestProfiler.profile_token
+        return unless profile_token.present?
+
+        header_token == profile_token
+      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
+        report = RubyProf::Profile.profile do
+          ret = catch(:warden) do # rubocop:disable Cop/BanCatchThrow
+            @app.call(env)
+          end
+        end
+
+        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 # rubocop:disable Cop/BanCatchThrow
+            @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)
+
+        begin
+          File.open(file_path, 'wb') do |file|
+            yield(file)
+          end
+        rescue StandardError
+          FileUtils.rm(file_path)
+        end
+      end
+
+      def handle_request_ret(ret)
+        if ret.is_a?(Array)
+          ret
+        else
+          throw(:warden, ret) # rubocop:disable Cop/BanCatchThrow
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/request_profiler/profile.rb b/lib/gitlab/request_profiler/profile.rb
new file mode 100644
index 00000000000..76c675658b1
--- /dev/null
+++ b/lib/gitlab/request_profiler/profile.rb
@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module RequestProfiler
+    class Profile
+      attr_reader :name, :time, :file_path, :request_path, :profile_mode, :type
+
+      alias_method :to_param, :name
+
+      def initialize(name)
+        @name = name
+        @file_path = File.join(PROFILES_DIR, name)
+
+        set_attributes
+      end
+
+      def valid?
+        @request_path.present?
+      end
+
+      def content_type
+        case type
+        when 'html'
+          'text/html'
+        when 'txt'
+          'text/plain'
+        end
+      end
+
+      private
+
+      def set_attributes
+        matches = name.match(/^(?<path>.*)_(?<timestamp>\d+)(_(?<profile_mode>\w+))?\.(?<type>html|txt)$/)
+        return unless matches
+
+        @request_path      = matches[:path].tr('|', '/')
+        @time              = Time.at(matches[:timestamp].to_i).utc
+        @profile_mode      = matches[:profile_mode] || 'unknown'
+        @type              = matches[:type]
+      end
+    end
+  end
+end