4 files changed, 116 insertions, 0 deletions
diff --git a/lib/gitlab/job_waiter.rb b/lib/gitlab/job_waiter.rb
new file mode 100644
index 00000000000..8db91d25a4b
--- /dev/null
+++ b/lib/gitlab/job_waiter.rb
@@ -0,0 +1,27 @@
+module Gitlab
+  # JobWaiter can be used to wait for a number of Sidekiq jobs to complete.
+  class JobWaiter
+    # The sleep interval between checking keys, in seconds.
+    INTERVAL = 0.1
+
+    # jobs - The job IDs to wait for.
+    def initialize(jobs)
+      @jobs = jobs
+    end
+
+    # Waits for all the jobs to be completed.
+    #
+    # timeout - The maximum amount of seconds to block the caller for. This
+    #           ensures we don't indefinitely block a caller in case a job takes
+    #           long to process, or is never processed.
+    def wait(timeout = 60)
+      start = Time.current
+
+      while (Time.current - start) <= timeout
+        break if SidekiqStatus.all_completed?(@jobs)
+
+        sleep(INTERVAL) # to not overload Redis too much.
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_status.rb b/lib/gitlab/sidekiq_status.rb
new file mode 100644
index 00000000000..aadc401ff8d
--- /dev/null
+++ b/lib/gitlab/sidekiq_status.rb
@@ -0,0 +1,66 @@
+module Gitlab
+  # The SidekiqStatus module and its child classes can be used for checking if a
+  # Sidekiq job has been processed or not.
+  #
+  # To check if a job has been completed, simply pass the job ID to the
+  # `completed?` method:
+  #
+  #     job_id = SomeWorker.perform_async(...)
+  #
+  #     if Gitlab::SidekiqStatus.completed?(job_id)
+  #       ...
+  #     end
+  #
+  # For each job ID registered a separate key is stored in Redis, making lookups
+  # much faster than using Sidekiq's built-in job finding/status API. These keys
+  # expire after a certain period of time to prevent storing too many keys in
+  # Redis.
+  module SidekiqStatus
+    STATUS_KEY = 'gitlab-sidekiq-status:%s'.freeze
+
+    # The default time (in seconds) after which a status key is expired
+    # automatically. The default of 30 minutes should be more than sufficient
+    # for most jobs.
+    DEFAULT_EXPIRATION = 30.minutes.to_i
+
+    # Starts tracking of the given job.
+    #
+    # jid - The Sidekiq job ID
+    # expire - The expiration time of the Redis key.
+    def self.set(jid, expire = DEFAULT_EXPIRATION)
+      Sidekiq.redis do |redis|
+        redis.set(key_for(jid), 1, ex: expire)
+      end
+    end
+
+    # Stops the tracking of the given job.
+    #
+    # jid - The Sidekiq job ID to remove.
+    def self.unset(jid)
+      Sidekiq.redis do |redis|
+        redis.del(key_for(jid))
+      end
+    end
+
+    # Returns true if all the given job have been completed.
+    #
+    # jids - The Sidekiq job IDs to check.
+    #
+    # Returns true or false.
+    def self.all_completed?(jids)
+      keys = jids.map { |jid| key_for(jid) }
+
+      responses = 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)
+      STATUS_KEY % jid
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_status/client_middleware.rb b/lib/gitlab/sidekiq_status/client_middleware.rb
new file mode 100644
index 00000000000..779a9998b22
--- /dev/null
+++ b/lib/gitlab/sidekiq_status/client_middleware.rb
@@ -0,0 +1,10 @@
+module Gitlab
+  module SidekiqStatus
+    class ClientMiddleware
+      def call(_, job, _, _)
+        SidekiqStatus.set(job['jid'])
+        yield
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_status/server_middleware.rb b/lib/gitlab/sidekiq_status/server_middleware.rb
new file mode 100644
index 00000000000..31dfa46ff9d
--- /dev/null
+++ b/lib/gitlab/sidekiq_status/server_middleware.rb
@@ -0,0 +1,13 @@
+module Gitlab
+  module SidekiqStatus
+    class ServerMiddleware
+      def call(worker, job, queue)
+        ret = yield
+
+        SidekiqStatus.unset(job['jid'])
+
+        ret
+      end
+    end
+  end
+end