Don't use Pygment,rb, use Rouge instead, and put peek-pg in the :postgres group

Signed-off-by: Rémy Coutable <remy@rymai.me>

5 files changed, 100 insertions, 6 deletions

diff --git a/Gemfile b/Gemfile
index 7a5b32a447d..c78c1f2fc4c 100644
--- a/Gemfile
+++ b/Gemfile
@@ -270,9 +270,8 @@ gem 'peek-gc', '~> 0.0.2'
 gem 'peek-host', '~> 1.0.0'
 gem 'peek-mysql2', '~> 1.1.0', group: :mysql
 gem 'peek-performance_bar', '~> 1.2.1'
-gem 'peek-pg', '~> 1.3.0'
+gem 'peek-pg', '~> 1.3.0', group: :postgres
 gem 'peek-rblineprof', '~> 0.2.0'
-gem 'pygments.rb', require: false
 gem 'peek-redis', '~> 1.2.0'
 gem 'peek-sidekiq', '~> 1.0.3'
 
diff --git a/Gemfile.lock b/Gemfile.lock
index 4c5c3db40c9..676cd977e37 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -608,8 +608,6 @@ GEM
       pry (~> 0.10)
     pry-rails (0.3.5)
       pry (>= 0.9.10)
-    pygments.rb (1.1.2)
-      multi_json (>= 1.0.0)
     pyu-ruby-sasl (0.0.3.3)
     rack (1.6.5)
     rack-accept (0.4.5)
@@ -1049,7 +1047,6 @@ DEPENDENCIES
   prometheus-client-mmap (~> 0.7.0.beta5)
   pry-byebug (~> 3.4.1)
   pry-rails (~> 0.3.4)
-  pygments.rb
   rack-attack (~> 4.4.1)
   rack-cors (~> 0.4.0)
   rack-oauth2 (~> 1.2.1)
diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index 01b1462d5ec..d52fec19a82 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -9,6 +9,7 @@ class ApplicationController < ActionController::Base
   include SentryHelper
   include WorkhorseHelper
   include EnforcesTwoFactorAuthentication
+  include Peek::Rblineprof::CustomControllerHelpers
 
   before_action :authenticate_user_from_private_token!
   before_action :authenticate_user_from_rss_token!
diff --git a/config/initializers/peek.rb b/config/initializers/peek.rb
index 6cb0e711e48..a8669ddba97 100644
--- a/config/initializers/peek.rb
+++ b/config/initializers/peek.rb
@@ -2,7 +2,13 @@ Rails.application.config.peek.adapter = :redis, { client: ::Redis.new(Gitlab::Re
 
 Peek.into Peek::Views::Host
 Peek.into Peek::Views::PerformanceBar
-Peek.into Gitlab::Database.mysql? ? Peek::Views::Mysql2 : Peek::Views::PG
+if Gitlab::Database.mysql?
+  require 'peek-mysql'
+  Peek.into Peek::Views::Mysql2
+else
+  require 'peek-pg'
+  Peek.into Peek::Views::PG
+end
 Peek.into Peek::Views::Redis
 Peek.into Peek::Views::Sidekiq
 Peek.into Peek::Views::Rblineprof
diff --git a/lib/peek/rblineprof/custom_controller_helpers.rb b/lib/peek/rblineprof/custom_controller_helpers.rb
new file mode 100644
index 00000000000..4ea4ba45e38
--- /dev/null
+++ b/lib/peek/rblineprof/custom_controller_helpers.rb
@@ -0,0 +1,91 @@
+module Peek
+  module Rblineprof
+    module CustomControllerHelpers
+      extend ActiveSupport::Concern
+
+      def pygmentize(file_name, code, lexer = nil)
+        if lexer.present?
+          Gitlab::Highlight.highlight(file_name, code)
+        else
+          "<pre>#{Rack::Utils.escape_html(code)}</pre>"
+        end
+      end
+
+      def inject_rblineprof
+        ret = nil
+        profile = lineprof(rblineprof_profiler_regex) do
+          ret = yield
+        end
+
+        if response.content_type =~ %r|text/html|
+          sort = params[:lineprofiler_sort]
+          mode = params[:lineprofiler_mode] || 'cpu'
+          min  = (params[:lineprofiler_min] || 5).to_i * 1000
+          summary = params[:lineprofiler_summary]
+
+          # Sort each file by the longest calculated time
+          per_file = profile.map do |file, lines|
+            total, child, excl, total_cpu, child_cpu, excl_cpu = lines[0]
+
+            wall = summary == 'exclusive' ? excl : total
+            cpu  = summary == 'exclusive' ? excl_cpu : total_cpu
+            idle = summary == 'exclusive' ? (excl - excl_cpu) : (total - total_cpu)
+
+            [
+              file, lines,
+              wall, cpu, idle,
+              sort == 'idle' ? idle : sort == 'cpu' ? cpu : wall
+            ]
+          end.sort_by{ |a,b,c,d,e,f| -f }
+
+          output = ''
+          per_file.each do |file_name, lines, file_wall, file_cpu, file_idle, file_sort|
+
+            output << "<div class='peek-rblineprof-file'><div class='heading'>"
+
+            show_src = file_sort > min
+            tmpl = show_src ? "<a href='#' class='js-lineprof-file'>%s</a>" : "%s"
+
+            if mode == 'cpu'
+              output << sprintf("<span class='duration'>% 8.1fms + % 8.1fms</span> #{tmpl}", file_cpu / 1000.0, file_idle / 1000.0, file_name.sub(Rails.root.to_s + '/', ''))
+            else
+              output << sprintf("<span class='duration'>% 8.1fms</span> #{tmpl}", file_wall/1000.0, file_name.sub(Rails.root.to_s + '/', ''))
+            end
+
+            output << "</div>" # .heading
+
+            next unless show_src
+
+            output << "<div class='data'>"
+            code = []
+            times = []
+            File.readlines(file_name).each_with_index do |line, i|
+              code << line
+              wall, cpu, calls = lines[i + 1]
+
+              if calls && calls > 0
+                if mode == 'cpu'
+                  idle = wall - cpu
+                  times << sprintf("% 8.1fms + % 8.1fms (% 5d)", cpu / 1000.0, idle / 1000.0, calls)
+                else
+                  times << sprintf("% 8.1fms (% 5d)", wall / 1000.0, calls)
+                end
+              else
+                times << ' '
+              end
+            end
+            output << "<pre class='duration'>#{times.join("\n")}</pre>"
+            # The following line was changed from
+            # https://github.com/peek/peek-rblineprof/blob/8d3b7a283a27de2f40abda45974516693d882258/lib/peek/rblineprof/controller_helpers.rb#L125
+            output << "<pre class='code highlight'>#{pygmentize(file_name, code.join, 'ruby')}</pre>"
+            output << "</div></div>" # .data then .peek-rblineprof-file
+          end
+
+          response.body += "<div class='peek-rblineprof-modal' id='line-profile'>#{output}</div>".html_safe
+        end
+
+        ret
+      end
+    end
+  end
+end