From 04a4f4fdb71b7d1d07e68c7e8cb54a717ebbfc29 Mon Sep 17 00:00:00 2001
From: Michal Cichra <michal@cichra.cz>
Date: Mon, 2 Oct 2017 16:29:05 +0200
Subject: [PATCH] expose job stats as prometheus metrics endpoint

show information about how many jobs of each kind are scheduled
---
 Gemfile                           |  2 ++
 Gemfile.lock                      |  4 +++
 config/application.rb             |  2 ++
 config/initializers/prometheus.rb | 17 ++++++++++
 lib/prometheus/job_stats.rb       | 53 +++++++++++++++++++++++++++++++
 5 files changed, 78 insertions(+)
 create mode 100644 config/initializers/prometheus.rb
 create mode 100644 lib/prometheus/job_stats.rb

diff --git a/Gemfile b/Gemfile
index 79db0d9c..ce219d23 100644
--- a/Gemfile
+++ b/Gemfile
@@ -39,6 +39,8 @@ gem 'message_bus' # for publishing notifications about integration status
 
 gem 'validate_url'
 
+gem 'prometheus-client', require: %w[prometheus/client prometheus/middleware/exporter]
+
 group :development, :test do
   # Call 'byebug' anywhere in the code to stop execution and get a debugger console
   gem 'pry-byebug', platforms: [:mri, :mingw, :x64_mingw]
diff --git a/Gemfile.lock b/Gemfile.lock
index 997a7f7c..6aeaa76b 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -123,6 +123,8 @@ GEM
       multi_xml (~> 0.5)
       rack (>= 1.2, < 3)
     pg (0.21.0)
+    prometheus-client (0.7.1)
+      quantile (~> 0.2.0)
     pry (0.10.4)
       coderay (~> 1.1.0)
       method_source (~> 0.8.1)
@@ -140,6 +142,7 @@ GEM
       pry (>= 0.9.11)
     public_suffix (2.0.5)
     puma (3.9.1)
+    quantile (0.2.0)
     que (0.13.1)
     rack (2.0.3)
     rack-test (0.6.3)
@@ -223,6 +226,7 @@ DEPENDENCIES
   minitest-reporters
   oauth2
   pg (>= 0.20)
+  prometheus-client
   pry-byebug
   pry-rails
   pry-rescue
diff --git a/config/application.rb b/config/application.rb
index 84d1a6eb..4fe4df5c 100644
--- a/config/application.rb
+++ b/config/application.rb
@@ -27,6 +27,8 @@ class Application < Rails::Application
     config.middleware.insert_before Rack::Sendfile,
                                     ActionDispatch::DebugLocks
 
+    config.middleware.use Prometheus::Middleware::Exporter
+
     config.que = ActiveSupport::InheritableOptions.new(config.que)
 
     config.que.worker_count = ENV.fetch('RAILS_MAX_THREADS'){ 5 }.to_i * 3
diff --git a/config/initializers/prometheus.rb b/config/initializers/prometheus.rb
new file mode 100644
index 00000000..c8ac79f6
--- /dev/null
+++ b/config/initializers/prometheus.rb
@@ -0,0 +1,17 @@
+require 'prometheus/middleware/exporter'
+require 'prometheus/job_stats'
+
+prometheus = Prometheus::Client.registry
+
+scheduled_job_stats = Prometheus::JobStats.new(:scheduled_jobs, 'Que Jobs to be executed')
+scheduled_job_stats
+prometheus.register(scheduled_job_stats)
+
+retried_job_stats = Prometheus::JobStats.new(:retried_jobs, 'Que Jobs to retried')
+retried_job_stats.filter('retries > 0')
+prometheus.register(retried_job_stats)
+
+
+pending_job_stats = Prometheus::JobStats.new(:pending_jobs, 'Que Jobs that should be already running')
+pending_job_stats.filter('run_at < now()')
+prometheus.register(pending_job_stats)
diff --git a/lib/prometheus/job_stats.rb b/lib/prometheus/job_stats.rb
new file mode 100644
index 00000000..20e3adef
--- /dev/null
+++ b/lib/prometheus/job_stats.rb
@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+require 'prometheus/client/metric'
+
+module Prometheus
+
+  # Prometheus metric to get job stats from Que.
+  class JobStats < Prometheus::Client::Metric
+
+    def initialize(*)
+      super
+      @filter = nil
+    end
+
+    CTE = <<~SQL
+      WITH jobs AS (
+          SELECT *, args::jsonb -> 0 AS options FROM que_jobs
+      ),
+       jobs_list AS (
+      SELECT options->>'job_class' AS job,
+      options->>'job_id' AS job_uuid,
+      (options->>'executions')::integer AS retries,
+      options->'arguments' AS arguments,
+      run_at, job_id FROM jobs
+      )
+    SQL
+
+    def type
+      :job_stats
+    end
+
+    def filter(sql)
+      @filter = sql
+    end
+
+    def values
+      synchronize do
+        job_stats.map do |summary|
+          [ { job: summary.fetch('job') }, summary.fetch('count') ]
+        end.to_h
+      end
+    end
+
+    protected
+
+    def job_stats
+      filter = "WHERE #{@filter}" if @filter
+      sql = <<~SQL
+          SELECT job, COUNT(*) FROM jobs_list #{filter} GROUP BY job
+      SQL
+      Que.execute(CTE + sql)
+    end
+  end
+end