Log running queries when a certain CPU load is reached

Gemfile

 source "https://rubygems.org"
+gemspec
+
 group :test do
   gem "rspec", "~>3.5"
   gem "rubocop", "~>0.42"

Gemfile.lock

+PATH
+  remote: .
+  specs:
+    gitlab-monitor (0.0.1)
+      pg (~> 0.18.4)
+      terminal-table (~> 1.6.0)
+      vmstat (~> 2.1.1)
+
 GEM
   remote: https://rubygems.org/
   specs:
@@ -5,6 +13,7 @@ GEM
     diff-lcs (1.2.5)
     parser (2.3.1.2)
       ast (~> 2.2)
+    pg (0.18.4)
     powerpack (0.1.1)
     rainbow (2.1.0)
     rspec (3.5.0)
@@ -27,12 +36,15 @@ GEM
       ruby-progressbar (~> 1.7)
       unicode-display_width (~> 1.0, >= 1.0.1)
     ruby-progressbar (1.8.1)
+    terminal-table (1.6.0)
     unicode-display_width (1.1.0)
+    vmstat (2.1.1)
 PLATFORMS
   ruby
 DEPENDENCIES
+  gitlab-monitor!
   rspec (~> 3.5)
   rubocop (~> 0.42)

bin/gitlab-mon

 #!/usr/bin/env ruby
+$LOAD_PATH.unshift File.expand_path("../../lib", __FILE__)
+
 require "optparse"
 require "gitlab_monitor"

gitlab-monitor.gemspec

@@ -20,5 +20,9 @@ Gem::Specification.new do |s|
   s.homepage      = "http://gitlab.com"
   s.license       = "MIT"
+  s.add_runtime_dependency "vmstat", "~> 2.1.1"
+  s.add_runtime_dependency "pg", "~> 0.18.4"
+  s.add_runtime_dependency "terminal-table", "~> 1.6.0"
+
   s.add_development_dependency "rspec", "~> 3.3"
 end

lib/gitlab_monitor.rb

@@ -6,5 +6,6 @@ module GitLab
     autoload :PrometheusMetrics, "gitlab_monitor/prometheus"
     autoload :Git,               "gitlab_monitor/git"
     autoload :GitProber,         "gitlab_monitor/git"
+    autoload :DatabaseActivityCollector, "gitlab_monitor/database_activity_collector"
   end
 end

lib/gitlab_monitor/cli.rb

+require "vmstat"
+
 module GitLab
   module Monitor
     # Stores runner classes in a single place
@@ -74,8 +76,78 @@ module GitLab
         end
       end
+      class DatabaseHeavyLoad
+        COMMAND_NAME = "db-heavy-load".freeze
+
+        def initialize(args)
+          @options = options(args)
+          @options.parse!
+
+          @collector        = ::GitLab::Monitor::DatabaseActivityCollector.new(connection_string: @db_connection_string)
+          @warning_queries  = File.open("/tmp/warning-queries", "a")
+          @critical_queries = File.open("/tmp/critical-queries", "a")
+
+          @critical_queries.sync = @warning_queries.sync = true
+        end
+
+        def options(args)
+          args.options do |opts|
+            opts.banner = "Usage: #{EXECUTABLE_NAME} #{COMMAND_NAME} [options]"
+            opts.on("-w", "--cpu-warning-level=5.7", "CPU warning level at which to collect running queries") do |val|
+              @cpu_warning_level = val.to_f
+            end
+            opts.on(
+              "-c", "--cpu-critical-level=15.5", "CPU critical level at  which to collect running queries"
+            ) do |val|
+              @cpu_critical_level = val.to_f
+            end
+            opts.on("--db-conn=\"dbname=test port=5432\"", "Database connection string") do |val|
+              @db_connection_string = val
+            end
+          end
+        end
+
+        def help
+          @options.help
+        end
+
+        def run
+          validate!
+          run_loop
+        end
+
+        private
+
+        def validate!
+          fail InvalidCLICommand.new(help) unless @cpu_warning_level && @cpu_critical_level
+        end
+
+        def run_loop
+          loop do
+            load_level = Vmstat.load_average.fifteen_minutes
+
+            if load_level < @cpu_warning_level
+              # noop
+            elsif load_level >= @cpu_warning_level && load_level < @cpu_critical_level
+              write_snapshot_to_file(@warning_queries)
+            else
+              write_snapshot_to_file(@critical_queries)
+            end
+
+            sleep 60
+          end
+        end
+
+        def write_snapshot_to_file(file)
+          @collector.run
+          file.write(Time.now.utc.to_s + "\n")
+          @collector.store(file)
+          file.write("\n")
+        end
+      end
+
       def self.commands
-        [GIT].each_with_object({}) do |command_class, commands|
+        [GIT, DatabaseHeavyLoad].each_with_object({}) do |command_class, commands|
           commands[command_class::COMMAND_NAME] = command_class
           commands
         end

lib/gitlab_monitor/database_activity_collector.rb

+require "pg"
+require "terminal-table"
+
+module GitLab
+  module Monitor
+    class DatabaseActivityCollector
+      attr_reader :result
+
+      def initialize(args)
+        @connection_string = args[:connection_string]
+      end
+
+      def run
+        @result = connection.exec <<-SQL
+          SELECT pid, application_name, client_addr, state, age(clock_timestamp(), query_start) as duration, query
+          FROM pg_stat_activity
+          WHERE query != '<IDLE>' AND query NOT ILIKE '%pg_stat_activity%' AND state != 'idle'
+          ORDER BY age(clock_timestamp(), query_start) DESC
+        SQL
+      end
+
+      def store(io)
+        return unless result
+
+        rows_count = @result.ntuples
+
+        table = Terminal::Table.new do |t|
+          @result.each.with_index do |row, index|
+            row.each do |field|
+              t.add_row(field)
+            end
+
+            t.add_separator unless rows_count == index + 1
+          end
+        end
+
+        io.write(table.to_s)
+      end
+
+      private
+
+      def connection
+        @connection ||= PG.connect(@connection_string)
+      end
+    end
+  end
+end