Merge branch 'master' of gitlab.com:gitlab-org/gitlab-ce

1 files changed, 48 insertions, 9 deletions

diff --git a/lib/gitlab/job_waiter.rb b/lib/gitlab/job_waiter.rb
index 208f0e1bbea..4d6bbda15f3 100644
--- a/lib/gitlab/job_waiter.rb
+++ b/lib/gitlab/job_waiter.rb
@@ -1,12 +1,31 @@
 module Gitlab
   # JobWaiter can be used to wait for a number of Sidekiq jobs to complete.
+  #
+  # Its use requires the cooperation of the sidekiq jobs themselves. Set up the
+  # waiter, then start the jobs, passing them its `key`. Their `perform` methods
+  # should look like:
+  #
+  #     def perform(args, notify_key)
+  #       # do work
+  #     ensure
+  #       ::Gitlab::JobWaiter.notify(notify_key, jid)
+  #     end
+  #
+  # The JobWaiter blocks popping items from a Redis array. All the sidekiq jobs
+  # push to that array when done. Once the waiter has popped `count` items, it
+  # knows all the jobs are done.
   class JobWaiter
-    # The sleep interval between checking keys, in seconds.
-    INTERVAL = 0.1
+    def self.notify(key, jid)
+      Gitlab::Redis::SharedState.with { |redis| redis.lpush(key, jid) }
+    end
+
+    attr_reader :key, :jobs_remaining, :finished
 
-    # jobs - The job IDs to wait for.
-    def initialize(jobs)
-      @jobs = jobs
+    # jobs_remaining - the number of jobs left to wait for
+    def initialize(jobs_remaining)
+      @key = "gitlab:job_waiter:#{SecureRandom.uuid}"
+      @jobs_remaining = jobs_remaining
+      @finished = []
     end
 
     # Waits for all the jobs to be completed.
@@ -15,13 +34,33 @@ module Gitlab
     #           ensures we don't indefinitely block a caller in case a job takes
     #           long to process, or is never processed.
     def wait(timeout = 10)
-      start = Time.current
+      deadline = Time.now.utc + timeout
+
+      Gitlab::Redis::SharedState.with do |redis|
+        # Fallback key expiry: allow a long grace period to reduce the chance of
+        # a job pushing to an expired key and recreating it
+        redis.expire(key, [timeout * 2, 10.minutes.to_i].max)
+
+        while jobs_remaining > 0
+          # Redis will not take fractional seconds. Prefer waiting too long over
+          # not waiting long enough
+          seconds_left = (deadline - Time.now.utc).ceil
 
-      while (Time.current - start) <= timeout
-        break if SidekiqStatus.all_completed?(@jobs)
+          # Redis interprets 0 as "wait forever", so skip the final `blpop` call
+          break if seconds_left <= 0
 
-        sleep(INTERVAL) # to not overload Redis too much.
+          list, jid = redis.blpop(key, timeout: seconds_left)
+          break unless list && jid # timed out
+
+          @finished << jid
+          @jobs_remaining -= 1
+        end
+
+        # All jobs have finished, so expire the key immediately
+        redis.expire(key, 0) if jobs_remaining == 0
       end
+
+      finished
     end
   end
 end