Add latest changes from gitlab-org/gitlab@master

3 files changed, 33 insertions, 1 deletions

diff --git a/lib/gitlab/application_context.rb b/lib/gitlab/application_context.rb
index 5a9a99423c4..b950bfb0f3a 100644
--- a/lib/gitlab/application_context.rb
+++ b/lib/gitlab/application_context.rb
@@ -16,7 +16,7 @@ module Gitlab
 
     def self.with_context(args, &block)
       application_context = new(**args)
-      Labkit::Context.with_context(application_context.to_lazy_hash, &block)
+      application_context.use(&block)
     end
 
     def self.push(args)
@@ -42,6 +42,10 @@ module Gitlab
       end
     end
 
+    def use
+      Labkit::Context.with_context(to_lazy_hash) { yield }
+    end
+
     private
 
     attr_reader :set_values
diff --git a/lib/gitlab/sidekiq_middleware.rb b/lib/gitlab/sidekiq_middleware.rb
index 3dda244233f..b19853a1702 100644
--- a/lib/gitlab/sidekiq_middleware.rb
+++ b/lib/gitlab/sidekiq_middleware.rb
@@ -18,6 +18,7 @@ module Gitlab
         chain.add Labkit::Middleware::Sidekiq::Server
         chain.add Gitlab::SidekiqMiddleware::InstrumentationLogger
         chain.add Gitlab::SidekiqStatus::ServerMiddleware
+        chain.add Gitlab::SidekiqMiddleware::WorkerContext::Server
       end
     end
 
diff --git a/lib/gitlab/sidekiq_middleware/worker_context/server.rb b/lib/gitlab/sidekiq_middleware/worker_context/server.rb
new file mode 100644
index 00000000000..29d98ad16a9
--- /dev/null
+++ b/lib/gitlab/sidekiq_middleware/worker_context/server.rb
@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module SidekiqMiddleware
+    module WorkerContext
+      class Server
+        def call(worker, job, _queue, &block)
+          worker_class = worker.class
+
+          # This is not a worker we know about, perhaps from a gem
+          return yield unless worker_class.respond_to?(:get_worker_context)
+
+          # 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
+end