Add latest changes from gitlab-org/gitlab@master

3 files changed, 40 insertions, 8 deletions

diff --git a/lib/gitlab/sidekiq_middleware/worker_context.rb b/lib/gitlab/sidekiq_middleware/worker_context.rb
new file mode 100644
index 00000000000..897a9211948
--- /dev/null
+++ b/lib/gitlab/sidekiq_middleware/worker_context.rb
@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module SidekiqMiddleware
+    module WorkerContext
+      private
+
+      def wrap_in_optional_context(context_or_nil, &block)
+        return yield unless context_or_nil
+
+        context_or_nil.use(&block)
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_middleware/worker_context/client.rb b/lib/gitlab/sidekiq_middleware/worker_context/client.rb
new file mode 100644
index 00000000000..0eb52179db2
--- /dev/null
+++ b/lib/gitlab/sidekiq_middleware/worker_context/client.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module SidekiqMiddleware
+    module WorkerContext
+      class Client
+        include Gitlab::SidekiqMiddleware::WorkerContext
+
+        def call(worker_class_or_name, job, _queue, _redis_pool, &block)
+          worker_class = worker_class_or_name.to_s.safe_constantize
+
+          # Mailers can't be constantized like this
+          return yield unless worker_class
+          return yield unless worker_class.include?(::ApplicationWorker)
+
+          context_for_args = worker_class.context_for_arguments(job['args'])
+
+          wrap_in_optional_context(context_for_args, &block)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_middleware/worker_context/server.rb b/lib/gitlab/sidekiq_middleware/worker_context/server.rb
index 29d98ad16a9..d2d84742c17 100644
--- a/lib/gitlab/sidekiq_middleware/worker_context/server.rb
+++ b/lib/gitlab/sidekiq_middleware/worker_context/server.rb
@@ -4,6 +4,8 @@ module Gitlab
   module SidekiqMiddleware
     module WorkerContext
       class Server
+        include Gitlab::SidekiqMiddleware::WorkerContext
+
         def call(worker, job, _queue, &block)
           worker_class = worker.class
 
@@ -13,14 +15,6 @@ module Gitlab
           # Use the context defined on the class level as a base context
           wrap_in_optional_context(worker_class.get_worker_context, &block)
         end
-
-        private
-
-        def wrap_in_optional_context(context, &block)
-          return yield unless context
-
-          context.use(&block)
-        end
       end
     end
   end