gitlab-monitor 4.0.0

data/lib/gitlab_monitor/sidekiq_queue_job_stats.lua

@@ -0,0 +1,42 @@
+--
+-- Adapted from https://github.com/mperham/sidekiq/blob/2f9258e4fe77991c526f7a65c92bcf792eef8338/lib/sidekiq/api.rb#L231
+--
+local queue_name      = KEYS[1]
+local initial_size    = redis.call('llen', queue_name)
+local deleted_size    = 0
+local page            = 0
+local page_size       = 2000
+local temp_job_stats  = {}
+local final_job_stats = {}
+
+while true do
+  local range_start = page * page_size - deleted_size
+  local range_end   = range_start + page_size - 1
+  local entries     = redis.call('lrange', queue_name, range_start, range_end)
+
+  if #entries == 0 then
+    break
+  end
+
+  page = page + 1
+
+  for index, entry in next, entries do
+    local class = cjson.decode(entry)['class']
+    if class ~= nil then
+      if temp_job_stats[class] ~= nil then
+        temp_job_stats[class] = temp_job_stats[class] + 1
+      else
+        temp_job_stats[class] = 1
+      end
+    end
+  end
+
+  deleted_size = initial_size - redis.call('llen', queue_name)
+end
+
+for class, count in next, temp_job_stats do
+  local stat_entry = {class, count}
+  table.insert(final_job_stats, stat_entry)
+end
+
+return final_job_stats

data/lib/gitlab_monitor/util.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+module GitLab
+  module Monitor
+    # Simple time wrapper that provides a to_i and wraps the execution result
+    TrackedResult = Struct.new(:result, :time) do
+      def to_i
+        time
+      end
+    end
+
+    # Time tracking object
+    #
+    # Provides a simple time tracking, and returns back the result plus the tracked time
+    # wraped in a TrackedResult struct
+    class TimeTracker
+      def track
+        @start = Time.now.to_f
+        result = yield
+        TrackedResult.new(result, Time.now.to_f - @start)
+      end
+    end
+
+    # Helper methods, some stuff was copied from ActiveSupport
+    module Utils
+      def camel_case_string(str)
+        str.gsub(/(?:_|^)([a-z\d]*)/i) { $1.capitalize } # rubocop:disable PerlBackrefs
+      end
+      module_function :camel_case_string
+
+      def deep_symbolize_hash_keys(hash)
+        deep_transform_keys_in_object(hash, &:to_sym)
+      end
+      module_function :deep_symbolize_hash_keys
+
+      def deep_transform_keys_in_object(object, &block)
+        case object
+        when Hash
+          object.keys.each do |key|
+            value = object.delete(key)
+            object[yield(key)] = deep_transform_keys_in_object(value, &block)
+          end
+          object
+        when Array
+          object.map! { |e| deep_transform_keys_in_object(e, &block) }
+        else
+          object
+        end
+      end
+      module_function :deep_transform_keys_in_object
+
+      def pgrep(pattern)
+        # pgrep will include the PID of the shell, so strip that out
+        exec_pgrep(pattern).split("\n").each_with_object([]) do |line, arr|
+          pid, name = line.split(" ")
+          arr << pid if name != "sh"
+        end
+      end
+      module_function :pgrep
+
+      def exec_pgrep(pattern)
+        `pgrep -fl "#{pattern}"`
+      end
+      module_function :exec_pgrep
+
+      def system_uptime
+        File.read("/proc/uptime").split(" ")[0].to_f
+      end
+      module_function :system_uptime
+
+      def wrap_in_array(object)
+        if object.nil?
+          []
+        elsif object.respond_to?(:to_ary)
+          object.to_ary || [object]
+        else
+          [object]
+        end
+      end
+      module_function :wrap_in_array
+    end
+  end
+end

data/lib/gitlab_monitor/version.rb

@@ -0,0 +1,5 @@
+module GitLab
+  module Monitor
+    VERSION = "4.0.0".freeze
+  end
+end

data/lib/gitlab_monitor/web_exporter.rb

@@ -0,0 +1,77 @@
+require "sinatra/base"
+require "English"
+
+module GitLab
+  module Monitor
+    # Metrics web exporter
+    class WebExporter < Sinatra::Base
+      # A middleware to kill the process if we exceeded a certain threshold
+      class MemoryKillerMiddleware
+        def initialize(app, memory_threshold)
+          @app = app
+          @memory_threshold = memory_threshold.to_i * 1024
+        end
+
+        def call(env)
+          if memory_usage > @memory_threshold
+            puts "Memory usage of #{memory_usage} exceeded threshold of #{@memory_threshold}, signalling KILL"
+            Process.kill("KILL", $PID)
+          end
+
+          @app.call(env)
+        end
+
+        private
+
+        def memory_usage
+          io = IO.popen(%W(ps -o rss= -p #{$PID}))
+
+          mem = io.read
+          io.close
+
+          return 0 unless $CHILD_STATUS.to_i.zero?
+
+          mem.to_i
+        end
+      end
+
+      class << self
+        def setup(config)
+          setup_server(config[:server])
+          setup_probes(config[:probes])
+
+          memory_threshold = (config[:server] && config[:server][:memory_threshold]) || 1024
+          use MemoryKillerMiddleware, memory_threshold
+        end
+
+        def setup_server(config)
+          config ||= {}
+
+          set(:bind, config.fetch(:listen_address, "0.0.0.0"))
+          set(:port, config.fetch(:listen_port, 9168))
+        end
+
+        def setup_probes(config)
+          (config || {}).each do |probe_name, params|
+            opts =
+              if params.delete(:multiple)
+                params
+              else
+                { probe_name => params }
+              end
+
+            get "/#{probe_name}" do
+              content_type "text/plain; version=0.0.4"
+              prober = Prober.new(opts, metrics: PrometheusMetrics.new(include_timestamp: false))
+
+              prober.probe_all
+              prober.write_to(response)
+
+              response
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/spec/cli_spec.rb

@@ -0,0 +1,31 @@
+require "spec_helper"
+require "gitlab_monitor/cli"
+
+context "With valid pair of repositories" do
+  let(:repos) { GitRepoBuilder.new }
+
+  after do
+    repos.cleanup
+  end
+
+  describe GitLab::Monitor::CLI do
+    it "returns the rigth parser" do
+      expect(GitLab::Monitor::CLI.for("git")).to be(GitLab::Monitor::CLI::GIT)
+    end
+
+    it "returns a null parser if it is not found" do
+      expect(GitLab::Monitor::CLI.for("invalid")).to be(GitLab::Monitor::CLI::NullRunner)
+    end
+  end
+
+  describe GitLab::Monitor::CLI::GIT do
+    let(:output) { StringIO.new }
+    it "works end to end" do
+      args = CLIArgs.new([repos.cloned_repo, output])
+      ssh = GitLab::Monitor::CLI::GIT.new(args)
+      ssh.run
+      output.rewind
+      expect(output.read).to match(/git_push_time_milliseconds \d+ \d+/)
+    end
+  end
+end

data/spec/database/bloat_spec.rb

@@ -0,0 +1,99 @@
+require "spec_helper"
+require "gitlab_monitor/database/bloat"
+
+describe GitLab::Monitor::Database::BloatCollector do
+  let(:connection_pool) { double("connection pool") }
+  let(:connection) { double("connection") }
+
+  subject { described_class.new(connection_string: "").run(type) }
+  let(:type) { :btree }
+  let(:query) { "select something" }
+
+  before do
+    allow_any_instance_of(described_class).to receive(:connection_pool).and_return(connection_pool)
+    allow(connection_pool).to receive(:with).and_yield(connection)
+  end
+
+  it "converts query results into a hash" do
+    row1 = { "object_name" => "o", "more_stuff" => 1 }
+    row2 = { "object_name" => "a", "more_stuff" => 2 }
+
+    expect(File).to receive(:read).and_return(query)
+    expect(connection).to receive(:exec).with(query).and_return([row1, row2])
+
+    expect(subject).to eq("o" => row1, "a" => row2)
+  end
+end
+
+describe GitLab::Monitor::Database::BloatProber do
+  let(:opts) { { bloat_types: %i(btree table) } }
+  let(:metrics) { double("PrometheusMetrics", add: nil) }
+  let(:collector) { double("BloatCollector", run: data) }
+
+  let(:data) do
+    {
+      "object" => {
+        "object_name" => "object",
+        "bloat_ratio" => 1,
+        "bloat_size" => 2,
+        "extra_size" => 3,
+        "real_size" => 4
+      }
+    }
+  end
+
+  describe "#probe_db" do
+    subject { described_class.new(opts, metrics: metrics, collector: collector).probe_db }
+
+    it "invokes the collector for each bloat type" do
+      expect(collector).to receive(:run).with(:btree)
+      expect(collector).to receive(:run).with(:table)
+
+      subject
+    end
+
+    it "adds bloat_ratio metric" do
+      opts[:bloat_types].each do |type|
+        expect(metrics).to receive(:add).with("gitlab_database_bloat_#{type}_bloat_ratio", 1, query_name: "object")
+      end
+
+      subject
+    end
+
+    it "adds bloat_size metric" do
+      opts[:bloat_types].each do |type|
+        expect(metrics).to receive(:add).with("gitlab_database_bloat_#{type}_bloat_size", 2, query_name: "object")
+      end
+
+      subject
+    end
+
+    it "adds extra_size metric" do
+      opts[:bloat_types].each do |type|
+        expect(metrics).to receive(:add).with("gitlab_database_bloat_#{type}_extra_size", 3, query_name: "object")
+      end
+
+      subject
+    end
+
+    it "adds real_size metric" do
+      opts[:bloat_types].each do |type|
+        expect(metrics).to receive(:add).with("gitlab_database_bloat_#{type}_real_size", 4, query_name: "object")
+      end
+
+      subject
+    end
+  end
+
+  describe "#write_to" do
+    let(:target) { double }
+    let(:metrics) { double("PrometheusMetrics", to_s: double) }
+    subject { described_class.new(opts, metrics: metrics).write_to(target) }
+
+    it "writes to given target" do
+      expect(target).to receive(:write).with(metrics.to_s)
+
+      subject
+    end
+  end
+end

data/spec/database/ci_builds_spec.rb

@@ -0,0 +1,421 @@
+require "spec_helper"
+require "gitlab_monitor/database/ci_builds"
+
+# rubocop:disable Metrics/LineLength
+describe GitLab::Monitor::Database do
+  let(:set_random_page_cost_query) { "SET random_page_cost" }
+  let(:builds_query_ee) { "SELECT BUILDS EE" }
+  let(:builds_query_ce) { "SELECT BUILDS CE" }
+  let(:stale_builds_query) { "SELECT NOT UPDATED RUNNING" }
+  let(:per_runner_query_ee) { "SELECT ALL RUNNING PER RUNNER EE" }
+  let(:per_runner_query_ce) { "SELECT ALL RUNNING PER RUNNER CE" }
+  let(:mirror_column_query) { "SELECT DOES MIRROR COLUMN EXISTS" }
+  let(:repeated_commands_query_ee) { "SELECT EE REPEATED COMNANDS %d" }
+  let(:repeated_commands_query_ce) { "SELECT CE REPEATED COMNANDS %d" }
+  let(:unarchived_traces_query) { "SELECT UNARCHIVED TRACES %s LIST" }
+  let(:connection_pool) { double("connection pool") }
+  let(:connection) { double("connection") }
+  let(:allowed_repeated_commands_count) { 5 }
+  let(:created_builds_counting_disabled) { true }
+  let(:time_now) { Time.new(2019, 4, 9, 6, 30, 0) }
+  let(:unarchived_traces_query_time) { "2019-04-09 05:30:00" }
+  let(:unarchived_traces_offset_minutes) { 60 }
+
+  def stub_ee
+    allow(connection).to receive(:exec).with(mirror_column_query).and_return([{ "exists" => "t" }])
+  end
+
+  def stub_ce
+    allow(connection).to receive(:exec).with(mirror_column_query).and_return([{ "exists" => "f" }])
+  end
+
+  def builds_query_row_ee(shared_runners_enabled, status, namespace_id, has_minutes, count)
+    row = builds_query_row_ce(shared_runners_enabled, status, namespace_id, count)
+    row["has_minutes"] = has_minutes
+    row
+  end
+
+  def builds_query_row_ce(shared_runners_enabled, status, namespace_id, count)
+    { "shared_runners_enabled" => shared_runners_enabled,
+      "status" => status,
+      "namespace_id" => namespace_id,
+      "count" => count }
+  end
+
+  # rubocop:disable Metrics/ParameterLists
+  def per_runner_query_row_ee(runner_id, is_shared, namespace_id, mirror, mirror_trigger_builds, pipeline_schedule_id, trigger_request_id, has_minutes, count)
+    row = per_runner_query_row_ce(runner_id, is_shared, namespace_id, pipeline_schedule_id, trigger_request_id, count)
+    row["mirror"] = mirror
+    row["mirror_trigger_builds"] = mirror_trigger_builds
+    row["has_minutes"] = has_minutes
+    row
+  end
+  # rubocop:enable Metrics/ParameterLists
+
+  # rubocop:disable Metrics/ParameterLists
+  def per_runner_query_row_ce(runner_id, is_shared, namespace_id, pipeline_schedule_id, trigger_request_id, count)
+    { "runner_id" => runner_id,
+      "is_shared" => is_shared,
+      "namespace_id" => namespace_id,
+      "pipeline_schedule_id" => pipeline_schedule_id,
+      "trigger_request_id" => trigger_request_id,
+      "count" => count }
+  end
+  # rubocop:enable Metrics/ParameterLists
+
+  # rubocop:disable Metrics/ParameterLists
+  def repeated_commands_query_row_ee(namespace_id, shared_runners_enabled, project_id, status, has_minutes, count)
+    row = repeated_commands_query_row_ce(namespace_id, shared_runners_enabled, project_id, status, count)
+    row["has_minutes"] = has_minutes
+    row
+  end
+  # rubocop:enable Metrics/ParameterLists
+
+  def repeated_commands_query_row_ce(namespace_id, shared_runners_enabled, project_id, status, count)
+    { "namespace_id" => namespace_id,
+      "shared_runners_enabled" => shared_runners_enabled,
+      "project_id" => project_id,
+      "status" => status,
+      "count" => count }
+  end
+
+  before do
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::SET_RANDOM_PAGE_COST", set_random_page_cost_query)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::BUILDS_QUERY_EE", builds_query_ee)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::BUILDS_QUERY_CE", builds_query_ce)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::STALE_BUILDS_QUERY", stale_builds_query)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::PER_RUNNER_QUERY_EE", per_runner_query_ee)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::PER_RUNNER_QUERY_CE", per_runner_query_ce)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::MIRROR_COLUMN_QUERY", mirror_column_query)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::REPEATED_COMMANDS_QUERY_EE", repeated_commands_query_ee)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::REPEATED_COMMANDS_QUERY_CE", repeated_commands_query_ce)
+    stub_const("GitLab::Monitor::Database::CiBuildsCollector::UNARCHIVED_TRACES_QUERY", unarchived_traces_query)
+
+    allow_any_instance_of(GitLab::Monitor::Database::CiBuildsCollector).to receive(:connection_pool).and_return(connection_pool)
+    allow(connection_pool).to receive(:with).and_yield(connection)
+
+    allow(connection).to receive(:transaction).and_yield(connection)
+    allow(connection).to receive(:exec).with(set_random_page_cost_query)
+
+    allow(Time).to receive(:now).and_return(time_now)
+
+    allow(connection).to receive(:exec).with(builds_query_ee)
+      .and_return([builds_query_row_ee("f", "created", "1", "f", 10),
+                   builds_query_row_ee("t", "pending", "1", "t", 30),
+                   builds_query_row_ee("f", "created", "2", "f", 20),
+                   builds_query_row_ee("t", "pending", "2", "t", 50),
+                   builds_query_row_ee("t", "pending", "3", "f", 1),
+                   builds_query_row_ee("t", "pending", "4", "t", 2),
+                   builds_query_row_ee("f", "pending", "5", "f", 2)])
+    allow(connection).to receive(:exec).with(builds_query_ce)
+      .and_return([builds_query_row_ce("f", "created", "1", 10),
+                   builds_query_row_ce("t", "pending", "1", 30),
+                   builds_query_row_ce("f", "created", "2", 20),
+                   builds_query_row_ce("t", "pending", "2", 50),
+                   builds_query_row_ce("t", "pending", "3", 1),
+                   builds_query_row_ce("t", "pending", "4", 2),
+                   builds_query_row_ce("f", "pending", "5", 2)])
+
+    allow(connection).to receive(:exec).with(stale_builds_query).and_return([{ "count" => 2 }])
+
+    allow(connection).to receive(:exec).with(per_runner_query_ee)
+      .and_return([per_runner_query_row_ee(1, "t", 1, "f", "f", 1, nil, "t", 15),
+                   per_runner_query_row_ee(2, "f", 2, "t", "t", nil, 3, "f", 5),
+                   per_runner_query_row_ee(2, "f", 3, "t", "t", nil, 3, "t", 5),
+                   per_runner_query_row_ee(3, "f", 4, "t", "t", nil, 3, "f", 5)])
+
+    allow(connection).to receive(:exec).with(per_runner_query_ce)
+      .and_return([per_runner_query_row_ce(1, "t", 1, 1, nil, 15),
+                   per_runner_query_row_ce(2, "f", 2, nil, 3, 5),
+                   per_runner_query_row_ce(2, "f", 3, nil, 3, 5),
+                   per_runner_query_row_ce(3, "f", 4, nil, 3, 5)])
+
+    # rubocop:disable Style/FormatString
+    repeated_commands_query_ee_with_limit = repeated_commands_query_ee % [allowed_repeated_commands_count]
+    repeated_commands_query_ce_with_limit = repeated_commands_query_ce % [allowed_repeated_commands_count]
+    # rubocop:enable Style/FormatString
+
+    allow(connection).to receive(:exec)
+      .with(repeated_commands_query_ee_with_limit)
+      .and_return([repeated_commands_query_row_ee(1, "t", 1, "pending", "t", 10),
+                   repeated_commands_query_row_ee(2, "f", 2, "running", "f", 20),
+                   repeated_commands_query_row_ee(1, "f", 3, "pending", "t", 30),
+                   repeated_commands_query_row_ee(2, "t", 4, "running", "f", 40)])
+    allow(connection).to receive(:exec)
+      .with(repeated_commands_query_ce_with_limit)
+      .and_return([repeated_commands_query_row_ce(1, "t", 1, "pending", 10),
+                   repeated_commands_query_row_ce(2, "f", 2, "running", 20),
+                   repeated_commands_query_row_ce(1, "f", 3, "pending", 30),
+                   repeated_commands_query_row_ce(2, "t", 4, "running", 40)])
+
+    unarchived_traces_query_with_time = unarchived_traces_query % [unarchived_traces_query_time] # rubocop:disable Style/FormatString
+
+    allow(connection).to receive(:exec).with(unarchived_traces_query_with_time).and_return([{ "count" => 10 }])
+  end
+
+  describe GitLab::Monitor::Database::CiBuildsCollector do
+    let(:collector) do
+      described_class.new(connection_string: "host=localhost",
+                          allowed_repeated_commands_count: allowed_repeated_commands_count,
+                          created_builds_counting_disabled: created_builds_counting_disabled,
+                          unarchived_traces_offset_minutes: unarchived_traces_offset_minutes)
+    end
+    let(:expected_stale_builds) { 2 }
+    let(:expected_unarchived_traces) { 10 }
+
+    shared_examples "data collector" do
+      subject { collector.run }
+
+      it "returns raw per_runner data" do
+        expect(subject[:per_runner]).to include(*expected_per_runner)
+      end
+
+      it "returns raw pending_builds data" do
+        expect(subject[:pending_builds]).to include(*expected_pending_builds)
+      end
+
+      context "when created_builds_counting_disabled is set to false" do
+        let(:created_builds_counting_disabled) { false }
+
+        it "returns raw created_builds data" do
+          expect(subject).to have_key(:created_builds)
+          expect(subject[:created_builds]).to include(*expected_created_builds)
+        end
+      end
+
+      context "when created_builds_counting_disabled is set to true" do
+        let(:created_builds_counting_disabled) { true }
+
+        it "doesn't return raw created_builds data" do
+          expect(subject).not_to have_key(:created_builds)
+        end
+      end
+
+      it "returns raw stale_builds data" do
+        expect(subject[:stale_builds]).to eq(expected_stale_builds)
+      end
+
+      it "returns raw repeated_commands data" do
+        expect(subject[:repeated_commands]).to include(*expected_repeated_commands)
+      end
+
+      it "returns raw unarchived_traces data" do
+        expect(subject[:unarchived_traces]).to eq(expected_unarchived_traces)
+      end
+    end
+
+    context "when executed on EE" do
+      let(:expected_pending_builds) do
+        [{ namespace: "1", shared_runners: "yes", has_minutes: "yes", value: 30 },
+         { namespace: "2", shared_runners: "yes", has_minutes: "yes", value: 50 },
+         { namespace: "3", shared_runners: "yes", has_minutes: "no", value: 1 },
+         { namespace: "4", shared_runners: "yes", has_minutes: "yes", value: 2 },
+         { namespace: "5", shared_runners: "no", has_minutes: "no", value: 2 }]
+      end
+      let(:expected_created_builds) do
+        [{ namespace: "1", shared_runners: "no", has_minutes: "no", value: 10 },
+         { namespace: "2", shared_runners: "no", has_minutes: "no", value: 20 }]
+      end
+      let(:expected_per_runner) do
+        [{ runner: "1", shared_runner: "yes", namespace: "1", mirror: "no", mirror_trigger_builds: "no", scheduled: "yes", triggered: "no", has_minutes: "yes", value: 15 },
+         { runner: "2", shared_runner: "no", namespace: "2", mirror: "yes", mirror_trigger_builds: "yes", scheduled: "no", triggered: "yes", has_minutes: "no", value: 5 },
+         { runner: "2", shared_runner: "no", namespace: "3", mirror: "yes", mirror_trigger_builds: "yes", scheduled: "no", triggered: "yes", has_minutes: "yes", value: 5 },
+         { runner: "3", shared_runner: "no", namespace: "4", mirror: "yes", mirror_trigger_builds: "yes", scheduled: "no", triggered: "yes", has_minutes: "no", value: 5 }]
+      end
+      let(:expected_repeated_commands) do
+        [{ namespace: "1", project: "1", shared_runners: "yes", status: "pending", has_minutes: "yes", value: 10 },
+         { namespace: "2", project: "2", shared_runners: "no", status: "running", has_minutes: "no", value: 20 },
+         { namespace: "1", project: "3", shared_runners: "no", status: "pending", has_minutes: "yes", value: 30 },
+         { namespace: "2", project: "4", shared_runners: "yes", status: "running", has_minutes: "no", value: 40 }]
+      end
+
+      before do
+        stub_ee
+      end
+
+      it_behaves_like "data collector"
+    end
+
+    context "when executed on CE" do
+      let(:expected_pending_builds) do
+        [{ namespace: "1", shared_runners: "yes", value: 30 },
+         { namespace: "2", shared_runners: "yes", value: 50 },
+         { namespace: "3", shared_runners: "yes", value: 1 },
+         { namespace: "4", shared_runners: "yes", value: 2 },
+         { namespace: "5", shared_runners: "no", value: 2 }]
+      end
+      let(:expected_created_builds) do
+        [{ namespace: "1", shared_runners: "no", value: 10 },
+         { namespace: "2", shared_runners: "no", value: 20 }]
+      end
+      let(:expected_per_runner) do
+        [{ runner: "1", shared_runner: "yes", namespace: "1", scheduled: "yes", triggered: "no", value: 15 },
+         { runner: "2", shared_runner: "no", namespace: "2", scheduled: "no", triggered: "yes", value: 5 },
+         { runner: "2", shared_runner: "no", namespace: "3", scheduled: "no", triggered: "yes", value: 5 },
+         { runner: "3", shared_runner: "no", namespace: "4", scheduled: "no", triggered: "yes", value: 5 }]
+      end
+      let(:expected_repeated_commands) do
+        [{ namespace: "1", project: "1", shared_runners: "yes", status: "pending", value: 10 },
+         { namespace: "2", project: "2", shared_runners: "no", status: "running", value: 20 },
+         { namespace: "1", project: "3", shared_runners: "no", status: "pending", value: 30 },
+         { namespace: "2", project: "4", shared_runners: "yes", status: "running", value: 40 }]
+      end
+
+      before do
+        stub_ce
+      end
+
+      it_behaves_like "data collector"
+    end
+  end
+
+  describe GitLab::Monitor::Database::CiBuildsProber do
+    let(:writer) { StringIO.new }
+    let(:prober) do
+      opts = { connection_string: "host=localhost",
+               allowed_repeated_commands_count: allowed_repeated_commands_count,
+               created_builds_counting_disabled: created_builds_counting_disabled,
+               unarchived_traces_offset_minutes: unarchived_traces_offset_minutes }
+      described_class.new(opts,
+                          metrics: GitLab::Monitor::PrometheusMetrics.new(include_timestamp: false))
+    end
+
+    before do
+      allow_any_instance_of(GitLab::Monitor::Database::CiBuildsCollector).to receive(:connected?).and_return(true)
+    end
+
+    shared_examples "metrics server" do
+      subject do
+        prober.probe_db
+        prober.write_to(writer)
+        writer.string
+      end
+
+      context "when PG exceptions aren't raised" do
+        context "when created_builds_counting_disabled is set to false" do
+          let(:created_builds_counting_disabled) { false }
+
+          it "responds with created builds Prometheus metrics" do
+            ci_created_builds_expected_lines.each do |expected_line|
+              expect(subject).to match(Regexp.new("^#{expected_line}$", Regexp::MULTILINE))
+            end
+          end
+        end
+
+        context "when created_builds_counting_disabled is set to true" do
+          let(:created_builds_counting_disabled) { true }
+
+          it "doesn't respond with created builds Prometheus metrics" do
+            ci_created_builds_expected_lines.each do |expected_line|
+              expect(subject).not_to match(Regexp.new("^#{expected_line}$", Regexp::MULTILINE))
+            end
+          end
+        end
+
+        it "responds with pending builds Prometheus metrics" do
+          ci_pending_builds_expected_lines.each do |expected_line|
+            expect(subject).to match(Regexp.new("^#{expected_line}$", Regexp::MULTILINE))
+          end
+        end
+
+        it "responds with running builds Prometheus metrics" do
+          ci_running_builds_expected_lines.each do |expected_line|
+            expect(subject).to match(Regexp.new("^#{expected_line}$", Regexp::MULTILINE))
+          end
+        end
+
+        it "responds with repeated commands Prometheus metrics" do
+          ci_repeated_commands_builds_lines.each do |expected_line|
+            expect(subject).to match(Regexp.new("^#{expected_line}$", Regexp::MULTILINE))
+          end
+        end
+
+        it "responds with stale builds Prometheus metrics" do
+          expect(subject).to match(/^ci_stale_builds 2$/m)
+        end
+
+        it "responds with unarchived traces Prometheus metrics" do
+          expect(subject).to match(/^ci_unarchived_traces 10$/m)
+        end
+      end
+
+      context "when PG exceptions are raised" do
+        before do
+          allow(connection).to receive(:exec).and_raise(PG::UndefinedColumn)
+        end
+
+        it "responds with Prometheus metrics" do
+          prober.probe_db
+          prober.write_to(writer)
+          output = writer.string
+
+          expect(output).to match(/^ci_stale_builds 0$/m)
+        end
+      end
+    end
+
+    context "when executed on EE" do
+      let(:ci_created_builds_expected_lines) do
+        ['ci_created_builds\{has_minutes="no",namespace="1",shared_runners="no"\} 10',
+         'ci_created_builds\{has_minutes="no",namespace="2",shared_runners="no"\} 20']
+      end
+      let(:ci_pending_builds_expected_lines) do
+        ['ci_pending_builds\{has_minutes="yes",namespace="1",shared_runners="yes"\} 30',
+         'ci_pending_builds\{has_minutes="yes",namespace="2",shared_runners="yes"\} 50',
+         'ci_pending_builds\{has_minutes="no",namespace="",shared_runners="yes"\} 1',
+         'ci_pending_builds\{has_minutes="yes",namespace="",shared_runners="yes"\} 2',
+         'ci_pending_builds\{has_minutes="no",namespace="",shared_runners="no"\} 2']
+      end
+      let(:ci_running_builds_expected_lines) do
+        ['ci_running_builds\{has_minutes="yes",mirror="no",mirror_trigger_builds="no",namespace="1",runner="1",scheduled="yes",shared_runner="yes",triggered="no"\} 15',
+         'ci_running_builds\{has_minutes="no",mirror="yes",mirror_trigger_builds="yes",namespace="",runner="2",scheduled="no",shared_runner="no",triggered="yes"\} 5',
+         'ci_running_builds\{has_minutes="yes",mirror="yes",mirror_trigger_builds="yes",namespace="",runner="2",scheduled="no",shared_runner="no",triggered="yes"\} 5',
+         'ci_running_builds\{has_minutes="no",mirror="yes",mirror_trigger_builds="yes",namespace="",runner="3",scheduled="no",shared_runner="no",triggered="yes"\} 5']
+      end
+      let(:ci_repeated_commands_builds_lines) do
+        ['ci_repeated_commands_builds\{namespace="1",project="1",shared_runners="yes",status="pending",has_minutes="yes"\} 10',
+         'ci_repeated_commands_builds\{namespace="2",project="2",shared_runners="no",status="running",has_minutes="no"\} 20',
+         'ci_repeated_commands_builds\{namespace="1",project="3",shared_runners="no",status="pending",has_minutes="yes"\} 30',
+         'ci_repeated_commands_builds\{namespace="2",project="4",shared_runners="yes",status="running",has_minutes="no"\} 40']
+      end
+      let(:namespace_out_of_limit) { 2 }
+
+      before do
+        stub_ee
+      end
+
+      it_behaves_like "metrics server"
+    end
+
+    context "when executed on CE" do
+      let(:ci_created_builds_expected_lines) do
+        ['ci_created_builds\{namespace="1",shared_runners="no"\} 10',
+         'ci_created_builds\{namespace="2",shared_runners="no"\} 20']
+      end
+      let(:ci_pending_builds_expected_lines) do
+        ['ci_pending_builds\{namespace="1",shared_runners="yes"\} 30',
+         'ci_pending_builds\{namespace="2",shared_runners="yes"\} 50',
+         'ci_pending_builds\{namespace="",shared_runners="yes"\} 3',
+         'ci_pending_builds\{namespace="",shared_runners="no"\} 2']
+      end
+      let(:ci_running_builds_expected_lines) do
+        ['ci_running_builds\{namespace="1",runner="1",scheduled="yes",shared_runner="yes",triggered="no"\} 15',
+         'ci_running_builds\{namespace="",runner="2",scheduled="no",shared_runner="no",triggered="yes"\} 10',
+         'ci_running_builds\{namespace="",runner="3",scheduled="no",shared_runner="no",triggered="yes"\} 5']
+      end
+      let(:ci_repeated_commands_builds_lines) do
+        ['ci_repeated_commands_builds\{namespace="1",project="1",shared_runners="yes",status="pending"\} 10',
+         'ci_repeated_commands_builds\{namespace="2",project="2",shared_runners="no",status="running"\} 20',
+         'ci_repeated_commands_builds\{namespace="1",project="3",shared_runners="no",status="pending"\} 30',
+         'ci_repeated_commands_builds\{namespace="2",project="4",shared_runners="yes",status="running"\} 40']
+      end
+      let(:namespace_out_of_limit) { 0 }
+
+      before do
+        stub_ce
+      end
+
+      it_behaves_like "metrics server"
+    end
+  end
+end