Merge branch 'fix/faster-sidekiq-queues-stats' into 'master'

Use a Lua script to calculate queue stats

cc/ @pcarranza

See merge request !28

Gemfile.lock

 PATH
   remote: .
   specs:
-    gitlab-monitor (0.0.16)
+    gitlab-monitor (0.0.17)
       pg (~> 0.18.4)
       quantile (~> 0.2.0)
       redis-namespace (~> 1.5.2)

lib/gitlab_monitor/sidekiq.rb

 require "sidekiq/api"
+require "digest"
 module GitLab
   module Monitor
@@ -6,25 +7,33 @@ module GitLab
     #
     # It takes the Redis URL Sidekiq is connected to
     class SidekiqProber
+      QUEUE_JOB_STATS_SCRIPT = File.read(File.expand_path("#{__FILE__}/../sidekiq_queue_job_stats.lua")).freeze
+      QUEUE_JOB_STATS_SHA    = Digest::SHA1.hexdigest(QUEUE_JOB_STATS_SCRIPT).freeze
+
       def initialize(opts, metrics: PrometheusMetrics.new)
         @opts    = opts
         @metrics = metrics
-        Sidekiq.configure_client { |config|
+        Sidekiq.configure_client do |config|
           config.redis = { url: opts[:redis_url], namespace: "resque:gitlab" }
-        }
+        end
+
+        ensure_queue_job_stats_script_loaded
       end
       def probe_queues
         return self unless connected?
-        job_stats = Hash.new(0)
+        job_stats = {}
         Sidekiq::Queue.all.each do |queue|
           @metrics.add("sidekiq_queue_size", queue.size, name: queue.name)
           @metrics.add("sidekiq_queue_latency", queue.latency, name: queue.name)
-          queue.each { |job| job_stats[job.klass] += 1 }
+          Sidekiq.redis do |conn|
+            stats = conn.evalsha(QUEUE_JOB_STATS_SHA, ["queue:#{queue.name}"])
+            job_stats.merge!(stats.to_h)
+          end
         end
         job_stats.each do |class_name, count|
@@ -84,6 +93,15 @@ module GitLab
       rescue Redis::CannotConnectError # rubocop:disable Lint/HandleExceptions
         # Maybe we're trying connecting to a slave
       end
+
+      def ensure_queue_job_stats_script_loaded
+        return unless connected?
+
+        Sidekiq.redis do |conn|
+          next if conn.script(:exists, QUEUE_JOB_STATS_SHA)
+          conn.script(:load, QUEUE_JOB_STATS_SCRIPT)
+        end
+      end
     end
   end
 end

lib/gitlab_monitor/sidekiq_queue_job_stats.lua

+--
+-- 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

lib/gitlab_monitor/version.rb

 module GitLab
   module Monitor
-    VERSION = "0.0.16".freeze
+    VERSION = "0.0.17".freeze
   end
 end