Merge branch 'refactor-labkit-context-to-gitlab-application-context' into 'master'

Refactor Labkit::Context to Gitlab::ApplicationContext

See merge request gitlab-org/gitlab!56217

app/controllers/application_controller.rb

@@ -462,7 +462,7 @@ class ApplicationController < ActionController::Base
       feature_category: feature_category) do
       yield
     ensure
-      @current_context = Labkit::Context.current.to_h
+      @current_context = Gitlab::ApplicationContext.current
     end
   end
 

app/graphql/mutations/admin/sidekiq_queues/delete_jobs.rb

@@ -8,7 +8,7 @@ module Mutations
 
         ADMIN_MESSAGE = 'You must be an admin to use this mutation'
 
-        Labkit::Context::KNOWN_KEYS.each do |key|
+        Gitlab::ApplicationContext::KNOWN_KEYS.each do |key|
           argument key,
                    GraphQL::STRING_TYPE,
                    required: false,

app/workers/concerns/application_worker.rb

@@ -18,7 +18,7 @@ module ApplicationWorker
     set_queue
 
     def structured_payload(payload = {})
-      context = Labkit::Context.current.to_h.merge(
+      context = Gitlab::ApplicationContext.current.merge(
         'class' => self.class.name,
         'job_status' => 'running',
         'queue' => self.class.queue,

app/workers/concerns/cronjob_queue.rb

@@ -15,7 +15,7 @@ module CronjobQueue
     # Cronjobs never get scheduled with arguments, so this is safe to
     # override
     def context_for_arguments(_args)
-      return if Gitlab::ApplicationContext.current_context_include?('meta.caller_id')
+      return if Gitlab::ApplicationContext.current_context_include?(:caller_id)
 
       Gitlab::ApplicationContext.new(caller_id: "Cronjob")
     end

ee/app/services/elastic/indexing_control_service.rb

@@ -6,7 +6,7 @@ module Elastic
   # This class should only be used with sidekiq workers which extend Elastic::IndexingControl module
   class IndexingControlService
     LIMIT = 1000
-    PROJECT_CONTEXT_KEY = "#{Labkit::Context::LOG_KEY}.project"
+    PROJECT_CONTEXT_KEY = "#{Gitlab::ApplicationContext::LOG_KEY}.project"
 
     def initialize(klass)
       raise ArgumentError, "passed class must extend Elastic::IndexingControl" unless klass.include?(Elastic::IndexingControl)
@@ -93,7 +93,7 @@ module Elastic
     end
 
     def send_to_processing_queue(job)
-      Labkit::Context.with_context(job['context']) do
+      Gitlab::ApplicationContext.with_raw_context(job['context']) do
         klass.perform_async(*job['args'])
       end
     end

ee/app/workers/concerns/elastic/indexing_control.rb

@@ -51,7 +51,7 @@ module Elastic
     end
 
     def current_context
-      Labkit::Context.current.to_h
+      Gitlab::ApplicationContext.current
     end
   end
 end

ee/spec/services/elastic/indexing_control_service_spec.rb

@@ -20,7 +20,7 @@ RSpec.describe Elastic::IndexingControlService, :clean_gitlab_redis_shared_state
   end
 
   let(:stored_context) do
-    { "#{Labkit::Context::LOG_KEY}.project" => 'gitlab-org/gitlab' }
+    { "#{Gitlab::ApplicationContext::LOG_KEY}.project" => 'gitlab-org/gitlab' }
   end
 
   let(:worker_args) { [1, 2] }
@@ -133,7 +133,7 @@ RSpec.describe Elastic::IndexingControlService, :clean_gitlab_redis_shared_state
         subject.add_to_waiting_queue!(j, worker_context)
       end
 
-      expect(Labkit::Context).to receive(:with_context).with(stored_context).exactly(jobs.count).times.and_call_original
+      expect(Gitlab::ApplicationContext).to receive(:with_raw_context).with(stored_context).exactly(jobs.count).times.and_call_original
       expect(worker_class).to receive(:perform_async).exactly(jobs.count).times
 
       expect { subject.resume_processing! }.to change { subject.has_jobs_in_waiting_queue? }.from(true).to(false)

ee/spec/workers/concerns/elastic/indexing_control_spec.rb

@@ -75,7 +75,7 @@ RSpec.describe Elastic::IndexingControl do
         expect_any_instance_of(Gitlab::Elastic::Indexer).not_to receive(:run)
         expect(Elastic::IndexingControlService).to receive(:add_to_waiting_queue!).with(worker.class, worker_args, worker_context)
 
-        Labkit::Context.with_context(worker_context) do
+        Gitlab::ApplicationContext.with_raw_context(worker_context) do
           worker.perform(*worker_args)
         end
       end

ee/spec/workers/update_all_mirrors_worker_spec.rb

@@ -27,15 +27,15 @@ RSpec.describe UpdateAllMirrorsWorker do
       worker.perform
     end
 
-    it 'removes metadata except correlation_id from the application context before scheduling mirrors' do
+    it 'removes metadata except correlation_id from the application context before scheduling mirrors', :context_aware do
       inner_context = nil
       outer_context = nil
 
       Gitlab::ApplicationContext.with_context(project: build(:project)) do
-        outer_context = Labkit::Context.current.to_h
+        outer_context = Gitlab::ApplicationContext.current
 
         expect(worker).to receive(:schedule_mirrors!) do
-          inner_context = Labkit::Context.current.to_h
+          inner_context = Gitlab::ApplicationContext.current
 
           # `schedule_mirrors!` needs to return an integer.
           0

lib/api/admin/sidekiq.rb

@@ -12,11 +12,11 @@ module API
           namespace 'queues' do
             desc 'Drop jobs matching the given metadata from the Sidekiq queue'
             params do
-              Labkit::Context::KNOWN_KEYS.each do |key|
+              Gitlab::ApplicationContext::KNOWN_KEYS.each do |key|
                 optional key, type: String, allow_blank: false
               end
 
-              at_least_one_of(*Labkit::Context::KNOWN_KEYS)
+              at_least_one_of(*Gitlab::ApplicationContext::KNOWN_KEYS)
             end
             delete ':queue_name' do
               result =

lib/gitlab/application_context.rb

@@ -8,6 +8,9 @@ module Gitlab
 
     Attribute = Struct.new(:name, :type)
 
+    LOG_KEY = Labkit::Context::LOG_KEY
+    KNOWN_KEYS = Labkit::Context::KNOWN_KEYS
+
     APPLICATION_ATTRIBUTES = [
       Attribute.new(:project, Project),
       Attribute.new(:namespace, Namespace),
@@ -24,6 +27,10 @@ module Gitlab
       application_context.use(&block)
     end
 
+    def self.with_raw_context(attributes = {}, &block)
+      Labkit::Context.with_context(attributes, &block)
+    end
+
     def self.push(args)
       application_context = new(**args)
       Labkit::Context.push(application_context.to_lazy_hash)

lib/gitlab/gitaly_client.rb

@@ -215,7 +215,7 @@ module Gitlab
         'client_name' => CLIENT_NAME
       }
 
-      context_data = Labkit::Context.current&.to_h
+      context_data = Gitlab::ApplicationContext.current
 
       feature_stack = Thread.current[:gitaly_feature_stack]
       feature = feature_stack && feature_stack[0]

lib/gitlab/grape_logging/loggers/context_logger.rb

@@ -6,7 +6,7 @@ module Gitlab
     module Loggers
       class ContextLogger < ::GrapeLogging::Loggers::Base
         def parameters(_, _)
-          Labkit::Context.current.to_h
+          Gitlab::ApplicationContext.current
         end
       end
     end

lib/gitlab/sidekiq_queue.rb

@@ -21,7 +21,7 @@ module Gitlab
       job_search_metadata =
         search_metadata
           .stringify_keys
-          .slice(*Labkit::Context::KNOWN_KEYS)
+          .slice(*Gitlab::ApplicationContext::KNOWN_KEYS)
           .transform_keys { |key| "meta.#{key}" }
           .compact
 

spec/controllers/application_controller_spec.rb

@@ -898,7 +898,7 @@ RSpec.describe ApplicationController do
       feature_category :issue_tracking
 
       def index
-        Labkit::Context.with_context do |context|
+        Gitlab::ApplicationContext.with_raw_context do |context|
           render json: context.to_h
         end
       end

spec/controllers/sessions_controller_spec.rb

@@ -524,7 +524,7 @@ RSpec.describe SessionsController do
 
       it 'sets the username and caller_id in the context' do
         expect(controller).to receive(:destroy).and_wrap_original do |m, *args|
-          expect(Labkit::Context.current.to_h)
+          expect(Gitlab::ApplicationContext.current)
             .to include('meta.user' => user.username,
                         'meta.caller_id' => 'SessionsController#destroy')
 
@@ -538,9 +538,9 @@ RSpec.describe SessionsController do
     context 'when not signed in' do
       it 'sets the caller_id in the context' do
         expect(controller).to receive(:new).and_wrap_original do |m, *args|
-          expect(Labkit::Context.current.to_h)
+          expect(Gitlab::ApplicationContext.current)
             .to include('meta.caller_id' => 'SessionsController#new')
-          expect(Labkit::Context.current.to_h)
+          expect(Gitlab::ApplicationContext.current)
             .not_to include('meta.user')
 
           m.call(*args)
@@ -557,9 +557,9 @@ RSpec.describe SessionsController do
 
       it 'sets the caller_id in the context' do
         allow_any_instance_of(User).to receive(:lock_access!).and_wrap_original do |m, *args|
-          expect(Labkit::Context.current.to_h)
+          expect(Gitlab::ApplicationContext.current)
             .to include('meta.caller_id' => 'SessionsController#create')
-          expect(Labkit::Context.current.to_h)
+          expect(Gitlab::ApplicationContext.current)
             .not_to include('meta.user')
 
           m.call(*args)

spec/lib/gitlab/application_context_spec.rb

@@ -27,6 +27,20 @@ RSpec.describe Gitlab::ApplicationContext do
     end
   end
 
+  describe '.with_raw_context' do
+    it 'yields the block' do
+      expect { |b| described_class.with_raw_context({}, &b) }.to yield_control
+    end
+
+    it 'passes the attributes unaltered on to labkit' do
+      attrs = { foo: :bar }
+
+      expect(Labkit::Context).to receive(:with_context).with(attrs)
+
+      described_class.with_raw_context(attrs) {}
+    end
+  end
+
   describe '.push' do
     it 'passes the expected context on to labkit' do
       fake_proc = duck_type(:call)
@@ -138,7 +152,7 @@ RSpec.describe Gitlab::ApplicationContext do
     it 'does not cause queries' do
       context = described_class.new(project: create(:project), namespace: create(:group, :nested), user: create(:user))
 
-      expect { context.use { Labkit::Context.current.to_h } }.not_to exceed_query_limit(0)
+      expect { context.use { Gitlab::ApplicationContext.current } }.not_to exceed_query_limit(0)
     end
   end
 end

spec/lib/gitlab/metrics/background_transaction_spec.rb

@@ -30,7 +30,7 @@ RSpec.describe Gitlab::Metrics::BackgroundTransaction do
 
   describe '#labels' do
     it 'provides labels with endpoint_id and feature_category' do
-      Labkit::Context.with_context(feature_category: 'projects', caller_id: 'TestWorker') do
+      Gitlab::ApplicationContext.with_raw_context(feature_category: 'projects', caller_id: 'TestWorker') do
         expect(transaction.labels).to eq({ endpoint_id: 'TestWorker', feature_category: 'projects' })
       end
     end
@@ -41,7 +41,7 @@ RSpec.describe Gitlab::Metrics::BackgroundTransaction do
       value = 1
       expect(prometheus_metric).to receive(metric_method).with({ endpoint_id: 'TestWorker', feature_category: 'projects' }, value)
 
-      Labkit::Context.with_context(feature_category: 'projects', caller_id: 'TestWorker') do
+      Gitlab::ApplicationContext.with_raw_context(feature_category: 'projects', caller_id: 'TestWorker') do
         transaction.send(metric_method, :test_metric, value)
       end
     end

spec/lib/gitlab/sidekiq_middleware/worker_context/server_spec.rb

@@ -18,7 +18,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Server do
       worker_context user: nil
 
       def perform(identifier, *args)
-        self.class.contexts.merge!(identifier => Labkit::Context.current.to_h)
+        self.class.contexts.merge!(identifier => Gitlab::ApplicationContext.current)
       end
     end
   end

spec/requests/api/api_spec.rb

@@ -105,7 +105,7 @@ RSpec.describe API::API do
 
     it 'logs all application context fields' do
       allow_any_instance_of(Gitlab::GrapeLogging::Loggers::ContextLogger).to receive(:parameters) do
-        Labkit::Context.current.to_h.tap do |log_context|
+        Gitlab::ApplicationContext.current.tap do |log_context|
           expect(log_context).to match('correlation_id' => an_instance_of(String),
                                        'meta.caller_id' => '/api/:version/projects/:id/issues',
                                        'meta.remote_ip' => an_instance_of(String),
@@ -122,7 +122,7 @@ RSpec.describe API::API do
 
     it 'skips fields that do not apply' do
       allow_any_instance_of(Gitlab::GrapeLogging::Loggers::ContextLogger).to receive(:parameters) do
-        Labkit::Context.current.to_h.tap do |log_context|
+        Gitlab::ApplicationContext.current.tap do |log_context|
           expect(log_context).to match('correlation_id' => an_instance_of(String),
                                        'meta.caller_id' => '/api/:version/users',
                                        'meta.remote_ip' => an_instance_of(String),

spec/spec_helper.rb

@@ -333,10 +333,20 @@ RSpec.configure do |config|
     RequestStore.clear!
   end
 
-  config.around do |example|
-    # Wrap each example in it's own context to make sure the contexts don't
-    # leak
-    Labkit::Context.with_context { example.run }
+  if ENV['SKIP_RSPEC_CONTEXT_WRAPPING']
+    config.around(:example, :context_aware) do |example|
+      # Wrap each example in it's own context to make sure the contexts don't
+      # leak
+      Gitlab::ApplicationContext.with_raw_context { example.run }
+    end
+  else
+    config.around do |example|
+      if [:controller, :request, :feature].include?(example.metadata[:type]) || example.metadata[:context_aware]
+        Gitlab::ApplicationContext.with_raw_context { example.run }
+      else
+        example.run
+      end
+    end
   end
 
   config.around do |example|

spec/support/shared_examples/lib/api/ci/runner_shared_examples.rb

 # frozen_string_literal: true
 
 RSpec.shared_examples 'API::CI::Runner application context metadata' do |api_route|
-  it 'contains correct context metadata' do
+  it 'contains correct context metadata', :context_aware do
     # Avoids popping the context from the thread so we can
     # check its content after the request.
     allow(Labkit::Context).to receive(:pop)
 
     send_request
 
-    Labkit::Context.with_context do |context|
+    Gitlab::ApplicationContext.with_raw_context do |context|
       expected_context = {
         'meta.caller_id' => api_route,
         'meta.user' => job.user.username,

spec/support/shared_examples/requests/api/logging_application_context_shared_examples.rb

 # frozen_string_literal: true
 
 RSpec.shared_examples 'storing arguments in the application context' do
-  around do |example|
-    Labkit::Context.with_context { example.run }
-  end
-
-  it 'places the expected params in the application context' do
+  it 'places the expected params in the application context', :context_aware do
     # Stub the clearing of the context so we can validate it later
-    # The `around` block above makes sure we do clean it up later
     allow(Labkit::Context).to receive(:pop)
 
     subject
 
-    Labkit::Context.with_context do |context|
-      expect(context.to_h)
-        .to include(log_hash(expected_params))
-    end
+    expect(Gitlab::ApplicationContext.current).to include(log_hash(expected_params))
   end
 
   def log_hash(hash)

spec/workers/background_migration_worker_spec.rb

@@ -101,7 +101,7 @@ RSpec.describe BackgroundMigrationWorker, :clean_gitlab_redis_shared_state do
 
     it 'sets the class that will be executed as the caller_id' do
       expect(Gitlab::BackgroundMigration).to receive(:perform) do
-        expect(Labkit::Context.current.to_h).to include('meta.caller_id' => 'Foo')
+        expect(Gitlab::ApplicationContext.current).to include('meta.caller_id' => 'Foo')
       end
 
       worker.perform('Foo', [10, 20])

spec/workers/concerns/worker_context_spec.rb

@@ -103,7 +103,7 @@ RSpec.describe WorkerContext do
   describe '#with_context' do
     it 'allows modifying context when the job is running' do
       worker.new.with_context(user: build_stubbed(:user, username: 'jane-doe')) do
-        expect(Labkit::Context.current.to_h).to include('meta.user' => 'jane-doe')
+        expect(Gitlab::ApplicationContext.current).to include('meta.user' => 'jane-doe')
       end
     end