Add latest changes from gitlab-org/gitlab@master

1 files changed, 38 insertions, 0 deletions

diff --git a/app/workers/concerns/worker_context.rb b/app/workers/concerns/worker_context.rb
index d85565e3446..ca006eaad5d 100644
--- a/app/workers/concerns/worker_context.rb
+++ b/app/workers/concerns/worker_context.rb
@@ -12,8 +12,46 @@ module WorkerContext
       @worker_context || superclass_context
     end
 
+    def bulk_perform_async_with_contexts(objects, arguments_proc:, context_proc:)
+      with_batch_contexts(objects, arguments_proc, context_proc) do |arguments|
+        bulk_perform_async(arguments)
+      end
+    end
+
+    def bulk_perform_in_with_contexts(delay, objects, arguments_proc:, context_proc:)
+      with_batch_contexts(objects, arguments_proc, context_proc) do |arguments|
+        bulk_perform_in(delay, arguments)
+      end
+    end
+
+    def context_for_arguments(args)
+      batch_context&.context_for(args)
+    end
+
     private
 
+    BATCH_CONTEXT_KEY = "#{name}_batch_context"
+
+    def batch_context
+      Thread.current[BATCH_CONTEXT_KEY]
+    end
+
+    def batch_context=(value)
+      Thread.current[BATCH_CONTEXT_KEY] = value
+    end
+
+    def with_batch_contexts(objects, arguments_proc, context_proc)
+      self.batch_context = Gitlab::BatchWorkerContext.new(
+        objects,
+        arguments_proc: arguments_proc,
+        context_proc: context_proc
+      )
+
+      yield(batch_context.arguments)
+    ensure
+      self.batch_context = nil
+    end
+
     def superclass_context
       return unless superclass.include?(WorkerContext)