Profile requests when a header is passed

2 files changed, 90 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..8da8b754975
--- /dev/null
+++ b/lib/gitlab/request_profiler/middleware.rb
@@ -0,0 +1,47 @@
+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 = 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 = @app.call(env)
+        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
+
+        ret
+      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..f89d56903ef
--- /dev/null
+++ b/lib/gitlab/request_profiler/profile.rb
@@ -0,0 +1,43 @@
+module Gitlab
+  module RequestProfiler
+    class Profile
+      attr_reader :name, :time, :request_path
+
+      alias_method :to_param, :name
+
+      def self.all
+        Dir["#{PROFILES_DIR}/*.html"].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}"
+        return unless File.exist?(file_path)
+
+        new(name_dup)
+      end
+
+      def initialize(name)
+        @name = name
+
+        set_attributes
+      end
+
+      def content
+        File.read("#{PROFILES_DIR}/#{name}")
+      end
+
+      private
+
+      def set_attributes
+        _, path, timestamp = name.split(/(.*)_(\d+)\.html$/)
+        @request_path      = path.tr('|', '/')
+        @time              = Time.at(timestamp.to_i).utc
+      end
+    end
+  end
+end