From ffd970d97dc9faf82752e6494aab8ba6fdce759a Mon Sep 17 00:00:00 2001
From: Stan Hu <stanhu@gmail.com>
Date: Fri, 3 Mar 2017 14:12:35 -0800
Subject: [PATCH] Make SidekiqStatus able to count number of jobs
 completed/running

---
 lib/gitlab/sidekiq_status.rb           | 35 +++++++++++++++++++++-----
 spec/lib/gitlab/sidekiq_status_spec.rb | 26 +++++++++++++++++++
 2 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/lib/gitlab/sidekiq_status.rb b/lib/gitlab/sidekiq_status.rb
index aadc401ff8d..11e5f1b645c 100644
--- a/lib/gitlab/sidekiq_status.rb
+++ b/lib/gitlab/sidekiq_status.rb
@@ -44,19 +44,42 @@ module Gitlab
 
     # Returns true if all the given job have been completed.
     #
-    # jids - The Sidekiq job IDs to check.
+    # job_ids - The Sidekiq job IDs to check.
     #
     # Returns true or false.
-    def self.all_completed?(jids)
-      keys = jids.map { |jid| key_for(jid) }
+    def self.all_completed?(job_ids)
+      self.num_running(job_ids).zero?
+    end
 
-      responses = Sidekiq.redis do |redis|
+    # Returns the number of jobs that are running.
+    #
+    # job_ids - The Sidekiq job IDs to check.
+    def self.num_running(job_ids)
+      responses = self.job_status(job_ids)
+
+      responses.select(&:present?).count
+    end
+
+    # Returns the number of jobs that have completed.
+    #
+    # job_ids - The Sidekiq job IDs to check.
+    def self.num_completed(job_ids)
+      job_ids.size - self.num_running(job_ids)
+    end
+
+    # Returns the job status for each of the given job IDs.
+    #
+    # job_ids - The Sidekiq job IDs to check.
+    #
+    # Returns an array of true or false indicating job completion.
+    def self.job_status(job_ids)
+      keys = job_ids.map { |jid| key_for(jid) }
+
+      Sidekiq.redis do |redis|
         redis.pipelined do
           keys.each { |key| redis.exists(key) }
         end
       end
-
-      responses.all? { |value| !value }
     end
 
     def self.key_for(jid)
diff --git a/spec/lib/gitlab/sidekiq_status_spec.rb b/spec/lib/gitlab/sidekiq_status_spec.rb
index 0aa36a3416b..56f06b61afb 100644
--- a/spec/lib/gitlab/sidekiq_status_spec.rb
+++ b/spec/lib/gitlab/sidekiq_status_spec.rb
@@ -39,6 +39,32 @@ describe Gitlab::SidekiqStatus do
     end
   end
 
+  describe '.num_running', :redis do
+    it 'returns 0 if all jobs have been completed' do
+      expect(described_class.num_running(%w(123))).to eq(0)
+    end
+
+    it 'returns 2 if two jobs are still running' do
+      described_class.set('123')
+      described_class.set('456')
+
+      expect(described_class.num_running(%w(123 456 789))).to eq(2)
+    end
+  end
+
+  describe '.num_completed', :redis do
+    it 'returns 1 if all jobs have been completed' do
+      expect(described_class.num_completed(%w(123))).to eq(1)
+    end
+
+    it 'returns 1 if a job has not yet been completed' do
+      described_class.set('123')
+      described_class.set('456')
+
+      expect(described_class.num_completed(%w(123 456 789))).to eq(1)
+    end
+  end
+
   describe '.key_for' do
     it 'returns the key for a job ID' do
       key = described_class.key_for('123')