Make SidekiqStatus able to count number of jobs completed/runningsh-improve-sidekiq-status

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
+
+    # 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 = Sidekiq.redis do |redis|
+      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')