summaryrefslogtreecommitdiff
path: root/lib/gitlab/request_profiler/middleware.rb
blob: 7615f6f443ba87c93a430143dfe8a23354c95512 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
# frozen_string_literal: true

require 'ruby-prof'

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)
        ret = nil
        result = 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"
        file_path = "#{PROFILES_DIR}/#{file_name}"

        FileUtils.mkdir_p(PROFILES_DIR)
        File.open(file_path, 'wb') do |file|
          printer.print(file)
        end

        if ret.is_a?(Array)
          ret
        else
          throw(:warden, ret)
        end
      end
    end
  end
end