Extract WaitableWorker out of AuthorizedProjectsWorker

2 files changed, 46 insertions, 34 deletions

diff --git a/app/workers/authorized_projects_worker.rb b/app/workers/authorized_projects_worker.rb
index 09559e3b696..d7e24491516 100644
--- a/app/workers/authorized_projects_worker.rb
+++ b/app/workers/authorized_projects_worker.rb
@@ -1,42 +1,10 @@
 class AuthorizedProjectsWorker
   include ApplicationWorker
+  prepend WaitableWorker
 
-  # Schedules multiple jobs and waits for them to be completed.
-  def self.bulk_perform_and_wait(args_list)
-    # Short-circuit: it's more efficient to do small numbers of jobs inline
-    return bulk_perform_inline(args_list) if args_list.size <= 3
-
-    waiter = Gitlab::JobWaiter.new(args_list.size)
-
-    # Point all the bulk jobs at the same JobWaiter. Converts, [[1], [2], [3]]
-    # into [[1, "key"], [2, "key"], [3, "key"]]
-    waiting_args_list = args_list.map { |args| [*args, waiter.key] }
-    bulk_perform_async(waiting_args_list)
-
-    waiter.wait
-  end
-
-  # Performs multiple jobs directly. Failed jobs will be put into sidekiq so
-  # they can benefit from retries
-  def self.bulk_perform_inline(args_list)
-    failed = []
-
-    args_list.each do |args|
-      begin
-        new.perform(*args)
-      rescue
-        failed << args
-      end
-    end
-
-    bulk_perform_async(failed) if failed.present?
-  end
-
-  def perform(user_id, notify_key = nil)
+  def perform(user_id)
     user = User.find_by(id: user_id)
 
     user&.refresh_authorized_projects
-  ensure
-    Gitlab::JobWaiter.notify(notify_key, jid) if notify_key
   end
 end
diff --git a/app/workers/concerns/waitable_worker.rb b/app/workers/concerns/waitable_worker.rb
new file mode 100644
index 00000000000..7d2266aea3c
--- /dev/null
+++ b/app/workers/concerns/waitable_worker.rb
@@ -0,0 +1,44 @@
+module WaitableWorker
+  extend ActiveSupport::Concern
+
+  module ClassMethods
+    # Schedules multiple jobs and waits for them to be completed.
+    def bulk_perform_and_wait(args_list)
+      # Short-circuit: it's more efficient to do small numbers of jobs inline
+      return bulk_perform_inline(args_list) if args_list.size <= 3
+
+      waiter = Gitlab::JobWaiter.new(args_list.size)
+
+      # Point all the bulk jobs at the same JobWaiter. Converts, [[1], [2], [3]]
+      # into [[1, "key"], [2, "key"], [3, "key"]]
+      waiting_args_list = args_list.map { |args| [*args, waiter.key] }
+      bulk_perform_async(waiting_args_list)
+
+      waiter.wait
+    end
+
+    # Performs multiple jobs directly. Failed jobs will be put into sidekiq so
+    # they can benefit from retries
+    def bulk_perform_inline(args_list)
+      failed = []
+
+      args_list.each do |args|
+        begin
+          new.perform(*args)
+        rescue
+          failed << args
+        end
+      end
+
+      bulk_perform_async(failed) if failed.present?
+    end
+  end
+
+  def perform(*args)
+    notify_key = args.pop if Gitlab::JobWaiter.key?(args.last)
+
+    super(*args)
+  ensure
+    Gitlab::JobWaiter.notify(notify_key, jid) if notify_key
+  end
+end