From 3c5be56e2cf4f0742822c3110f8149d0a4c738e9 Mon Sep 17 00:00:00 2001
From: Bob Van Landuyt <bob@vanlanduyt.co>
Date: Tue, 3 Dec 2019 17:18:20 +0100
Subject: [PATCH] Wrap requests in an ApplicationContext

This provides context to all requests made to Rails controllers or
grape endpoints.

Doing this starts a new `Labkit::Context`, to which we can provide a
namespace, project and user.

We're currently setting the following values:

- Web requests: In the ApplicationController we wrap the entire
  request in a `with_context`.
  - user: based on the `auth_user` if there is one
  - project: We try to read the @project instance variable of the
    controller.
  - namespace: We try to read the @group instance variable of the
    controller. If there was none, but the project was set, we'll use
    that path to set the namespace

- API requests: The application context is pushed in a before block
  setting the following values:
  - user: to `current_user` if there is one
  - project: to `@project`
  - namespace: to `@group`

- Internal API requests: the application context is pushed in a before
  block:
  - user: When to the user set in `Api::Support::GitAccessActor`
  - project: to @project if it was available

The 3 supported attributes for a context are read lazily when
required.

This also replaces the existing correlation middlewares with the new
Labkit::Context middlewares.

The rack middleware wraps each rack request in an overarching context
that adds the correlation id. The context is cleaned up after the
request, so we're sure all child contexts are cleaned up as well.

The sidekiq client middleware will write the context into the job that
 goes into redis when a job is scheduled.

The sidekiq server middleware will then re-instantiate this context so
the job gets executed with the same context that was alive when it was
scheduled. This means that any new job scheduled from sidekiq would
also have this context.
---
 Gemfile                                       |  2 +-
 Gemfile.lock                                  |  8 +-
 app/controllers/application_controller.rb     |  9 ++
 config/initializers/correlation_id.rb         |  3 -
 config/initializers/labkit_middleware.rb      |  3 +
 config/initializers/tracing.rb                |  2 +-
 lib/api/api.rb                                |  8 ++
 lib/api/helpers/internal_helpers.rb           |  4 +-
 lib/api/internal/base.rb                      | 12 +++
 lib/gitlab/application_context.rb             | 50 +++++++++++
 lib/gitlab/middleware/correlation_id.rb       | 31 -------
 lib/gitlab/sidekiq_middleware.rb              |  4 +-
 .../correlation_injector.rb                   | 14 ----
 .../sidekiq_middleware/correlation_logger.rb  | 15 ----
 lib/gitlab/utils/lazy_attributes.rb           | 45 ++++++++++
 .../application_controller_spec.rb            | 46 +++++++++++
 spec/lib/gitlab/application_context_spec.rb   | 82 +++++++++++++++++++
 .../project_tree_restorer_spec.rb             |  5 +-
 spec/lib/gitlab/profiler_spec.rb              |  2 +-
 .../correlation_injector_spec.rb              | 49 -----------
 .../correlation_logger_spec.rb                | 35 --------
 spec/lib/gitlab/sidekiq_middleware_spec.rb    |  4 +-
 spec/lib/gitlab/utils/lazy_attributes_spec.rb | 70 ++++++++++++++++
 spec/requests/api/internal/base_spec.rb       | 12 +++
 spec/requests/api/projects_spec.rb            |  8 ++
 spec/spec_helper.rb                           |  6 ++
 ...ing_application_context_shared_examples.rb | 24 ++++++
 27 files changed, 390 insertions(+), 163 deletions(-)
 delete mode 100644 config/initializers/correlation_id.rb
 create mode 100644 config/initializers/labkit_middleware.rb
 create mode 100644 lib/gitlab/application_context.rb
 delete mode 100644 lib/gitlab/middleware/correlation_id.rb
 delete mode 100644 lib/gitlab/sidekiq_middleware/correlation_injector.rb
 delete mode 100644 lib/gitlab/sidekiq_middleware/correlation_logger.rb
 create mode 100644 lib/gitlab/utils/lazy_attributes.rb
 create mode 100644 spec/lib/gitlab/application_context_spec.rb
 delete mode 100644 spec/lib/gitlab/sidekiq_middleware/correlation_injector_spec.rb
 delete mode 100644 spec/lib/gitlab/sidekiq_middleware/correlation_logger_spec.rb
 create mode 100644 spec/lib/gitlab/utils/lazy_attributes_spec.rb
 create mode 100644 spec/support/shared_examples/logging_application_context_shared_examples.rb

diff --git a/Gemfile b/Gemfile
index 652cc04ecbc1d..0c5fffe73ea4b 100644
--- a/Gemfile
+++ b/Gemfile
@@ -301,7 +301,7 @@ gem 'sentry-raven', '~> 2.9'
 gem 'premailer-rails', '~> 1.10.3'
 
 # LabKit: Tracing and Correlation
-gem 'gitlab-labkit', '~> 0.5'
+gem 'gitlab-labkit', '0.8.0'
 
 # I18n
 gem 'ruby_parser', '~> 3.8', require: false
diff --git a/Gemfile.lock b/Gemfile.lock
index 9ca82644a9722..15ceff6e86814 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -117,7 +117,7 @@ GEM
       activemodel (>= 5.0)
     brakeman (4.2.1)
     browser (2.5.3)
-    builder (3.2.3)
+    builder (3.2.4)
     bullet (6.0.2)
       activesupport (>= 3.0.0)
       uniform_notifier (~> 1.11)
@@ -362,7 +362,7 @@ GEM
     github-markup (1.7.0)
     gitlab-chronic (0.10.5)
       numerizer (~> 0.2)
-    gitlab-labkit (0.7.0)
+    gitlab-labkit (0.8.0)
       actionpack (>= 5.0.0, < 6.1.0)
       activesupport (>= 5.0.0, < 6.1.0)
       grpc (~> 1.19)
@@ -589,7 +589,7 @@ GEM
       activesupport (>= 4)
       railties (>= 4)
       request_store (~> 1.0)
-    loofah (2.3.1)
+    loofah (2.4.0)
       crass (~> 1.0.2)
       nokogiri (>= 1.5.9)
     lumberjack (1.0.13)
@@ -1204,7 +1204,7 @@ DEPENDENCIES
   gitaly (~> 1.73.0)
   github-markup (~> 1.7.0)
   gitlab-chronic (~> 0.10.5)
-  gitlab-labkit (~> 0.5)
+  gitlab-labkit (= 0.8.0)
   gitlab-license (~> 1.0)
   gitlab-markup (~> 1.7.0)
   gitlab-net-dns (~> 0.9.1)
diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index f5306801c049f..60b5d9b6da835 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -33,6 +33,7 @@ class ApplicationController < ActionController::Base
   before_action :check_impersonation_availability
   before_action :required_signup_info
 
+  around_action :set_current_context
   around_action :set_locale
   around_action :set_session_storage
 
@@ -448,6 +449,14 @@ def u2f_app_id
     request.base_url
   end
 
+  def set_current_context(&block)
+    Gitlab::ApplicationContext.with_context(
+      user: -> { auth_user },
+      project: -> { @project },
+      namespace: -> { @group },
+      &block)
+  end
+
   def set_locale(&block)
     Gitlab::I18n.with_user_locale(current_user, &block)
   end
diff --git a/config/initializers/correlation_id.rb b/config/initializers/correlation_id.rb
deleted file mode 100644
index 2a7c138dc40fa..0000000000000
--- a/config/initializers/correlation_id.rb
+++ /dev/null
@@ -1,3 +0,0 @@
-# frozen_string_literal: true
-
-Rails.application.config.middleware.use(Gitlab::Middleware::CorrelationId)
diff --git a/config/initializers/labkit_middleware.rb b/config/initializers/labkit_middleware.rb
new file mode 100644
index 0000000000000..ea4103f052fd9
--- /dev/null
+++ b/config/initializers/labkit_middleware.rb
@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+Rails.application.config.middleware.use(Labkit::Middleware::Rack)
diff --git a/config/initializers/tracing.rb b/config/initializers/tracing.rb
index 0ae57021fcf39..aaf74eb4cd304 100644
--- a/config/initializers/tracing.rb
+++ b/config/initializers/tracing.rb
@@ -2,7 +2,7 @@
 
 if Labkit::Tracing.enabled?
   Rails.application.configure do |config|
-    config.middleware.insert_after Gitlab::Middleware::CorrelationId, ::Labkit::Tracing::RackMiddleware
+    config.middleware.insert_after Labkit::Middleware::Rack, ::Labkit::Tracing::RackMiddleware
   end
 
   # Instrument the Sidekiq client
diff --git a/lib/api/api.rb b/lib/api/api.rb
index 56eccb036b6f1..eae10738f32b2 100644
--- a/lib/api/api.rb
+++ b/lib/api/api.rb
@@ -43,6 +43,14 @@ class API < Grape::API
       header['X-Content-Type-Options'] = 'nosniff'
     end
 
+    before do
+      Gitlab::ApplicationContext.push(
+        user: -> { current_user },
+        project: -> { @project },
+        namespace: -> { @group }
+      )
+    end
+
     # The locale is set to the current user's locale when `current_user` is loaded
     after { Gitlab::I18n.use_default_locale }
 
diff --git a/lib/api/helpers/internal_helpers.rb b/lib/api/helpers/internal_helpers.rb
index 607e078441539..b719e5c0886f9 100644
--- a/lib/api/helpers/internal_helpers.rb
+++ b/lib/api/helpers/internal_helpers.rb
@@ -107,8 +107,10 @@ def set_project
         if params[:gl_repository]
           @project, @repo_type = Gitlab::GlRepository.parse(params[:gl_repository])
           @redirected_path = nil
-        else
+        elsif params[:project]
           @project, @repo_type, @redirected_path = Gitlab::RepoPath.parse(params[:project])
+        else
+          @project, @repo_type, @redirected_path = nil, nil, nil
         end
       end
       # rubocop:enable Gitlab/ModuleWithInstanceVariables
diff --git a/lib/api/internal/base.rb b/lib/api/internal/base.rb
index 11f2a2ea1c027..d64de2bb46512 100644
--- a/lib/api/internal/base.rb
+++ b/lib/api/internal/base.rb
@@ -6,6 +6,13 @@ module Internal
     class Base < Grape::API
       before { authenticate_by_gitlab_shell_token! }
 
+      before do
+        Gitlab::ApplicationContext.push(
+          user: -> { actor&.user },
+          project: -> { project }
+        )
+      end
+
       helpers ::API::Helpers::InternalHelpers
 
       UNKNOWN_CHECK_RESULT_ERROR = 'Unknown check result'.freeze
@@ -205,7 +212,12 @@ def check_allowed(params)
           status 200
 
           response = Gitlab::InternalPostReceive::Response.new
+
+          # Try to load the project and users so we have the application context
+          # available for logging before we schedule any jobs.
           user = actor.user
+          project
+
           push_options = Gitlab::PushOptions.new(params[:push_options])
 
           response.reference_counter_decreased = Gitlab::ReferenceCounter.new(params[:gl_repository]).decrease
diff --git a/lib/gitlab/application_context.rb b/lib/gitlab/application_context.rb
new file mode 100644
index 0000000000000..b9190b519a0cf
--- /dev/null
+++ b/lib/gitlab/application_context.rb
@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Gitlab
+  # A GitLab-rails specific accessor for `Labkit::Logging::ApplicationContext`
+  class ApplicationContext
+    include Gitlab::Utils::LazyAttributes
+
+    def self.with_context(args, &block)
+      application_context = new(**args)
+      Labkit::Context.with_context(application_context.to_lazy_hash, &block)
+    end
+
+    def self.push(args)
+      application_context = new(**args)
+      Labkit::Context.push(application_context.to_lazy_hash)
+    end
+
+    def initialize(user: nil, project: nil, namespace: nil)
+      @user, @project, @namespace = user, project, namespace
+    end
+
+    def to_lazy_hash
+      { user: -> { username },
+        project: -> { project_path },
+        root_namespace: -> { root_namespace_path } }
+    end
+
+    private
+
+    lazy_attr_reader :user, type: User
+    lazy_attr_reader :project, type: Project
+    lazy_attr_reader :namespace, type: Namespace
+
+    def project_path
+      project&.full_path
+    end
+
+    def username
+      user&.username
+    end
+
+    def root_namespace_path
+      if namespace
+        namespace.full_path_components.first
+      else
+        project&.full_path_components&.first
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/middleware/correlation_id.rb b/lib/gitlab/middleware/correlation_id.rb
deleted file mode 100644
index fffd5da827fd8..0000000000000
--- a/lib/gitlab/middleware/correlation_id.rb
+++ /dev/null
@@ -1,31 +0,0 @@
-# frozen_string_literal: true
-
-# A dumb middleware that steals correlation id
-# and sets it as a global context for the request
-module Gitlab
-  module Middleware
-    class CorrelationId
-      include ActionView::Helpers::TagHelper
-
-      def initialize(app)
-        @app = app
-      end
-
-      def call(env)
-        ::Labkit::Correlation::CorrelationId.use_id(correlation_id(env)) do
-          @app.call(env)
-        end
-      end
-
-      private
-
-      def correlation_id(env)
-        request(env).request_id
-      end
-
-      def request(env)
-        ActionDispatch::Request.new(env)
-      end
-    end
-  end
-end
diff --git a/lib/gitlab/sidekiq_middleware.rb b/lib/gitlab/sidekiq_middleware.rb
index c6726dcfa675c..4893cbc1f45f7 100644
--- a/lib/gitlab/sidekiq_middleware.rb
+++ b/lib/gitlab/sidekiq_middleware.rb
@@ -15,7 +15,7 @@ def self.server_configurator(metrics: true, arguments_logger: true, memory_kille
         chain.add Gitlab::SidekiqMiddleware::MemoryKiller if memory_killer
         chain.add Gitlab::SidekiqMiddleware::RequestStoreMiddleware if request_store
         chain.add Gitlab::SidekiqMiddleware::BatchLoader
-        chain.add Gitlab::SidekiqMiddleware::CorrelationLogger
+        chain.add Labkit::Middleware::Sidekiq::Server
         chain.add Gitlab::SidekiqMiddleware::InstrumentationLogger
         chain.add Gitlab::SidekiqStatus::ServerMiddleware
       end
@@ -27,7 +27,7 @@ def self.server_configurator(metrics: true, arguments_logger: true, memory_kille
     def self.client_configurator
       lambda do |chain|
         chain.add Gitlab::SidekiqStatus::ClientMiddleware
-        chain.add Gitlab::SidekiqMiddleware::CorrelationInjector
+        chain.add Labkit::Middleware::Sidekiq::Client
       end
     end
   end
diff --git a/lib/gitlab/sidekiq_middleware/correlation_injector.rb b/lib/gitlab/sidekiq_middleware/correlation_injector.rb
deleted file mode 100644
index 1539fd706abe6..0000000000000
--- a/lib/gitlab/sidekiq_middleware/correlation_injector.rb
+++ /dev/null
@@ -1,14 +0,0 @@
-# frozen_string_literal: true
-
-module Gitlab
-  module SidekiqMiddleware
-    class CorrelationInjector
-      def call(worker_class, job, queue, redis_pool)
-        job[Labkit::Correlation::CorrelationId::LOG_KEY] ||=
-          Labkit::Correlation::CorrelationId.current_or_new_id
-
-        yield
-      end
-    end
-  end
-end
diff --git a/lib/gitlab/sidekiq_middleware/correlation_logger.rb b/lib/gitlab/sidekiq_middleware/correlation_logger.rb
deleted file mode 100644
index cffc448357382..0000000000000
--- a/lib/gitlab/sidekiq_middleware/correlation_logger.rb
+++ /dev/null
@@ -1,15 +0,0 @@
-# frozen_string_literal: true
-
-module Gitlab
-  module SidekiqMiddleware
-    class CorrelationLogger
-      def call(worker, job, queue)
-        correlation_id = job[Labkit::Correlation::CorrelationId::LOG_KEY]
-
-        Labkit::Correlation::CorrelationId.use_id(correlation_id) do
-          yield
-        end
-      end
-    end
-  end
-end
diff --git a/lib/gitlab/utils/lazy_attributes.rb b/lib/gitlab/utils/lazy_attributes.rb
new file mode 100644
index 0000000000000..79f3a7dcb532c
--- /dev/null
+++ b/lib/gitlab/utils/lazy_attributes.rb
@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Utils
+    module LazyAttributes
+      extend ActiveSupport::Concern
+      include Gitlab::Utils::StrongMemoize
+
+      class_methods do
+        def lazy_attr_reader(*one_or_more_names, type: nil)
+          names = Array.wrap(one_or_more_names)
+          names.each { |name| define_lazy_reader(name, type: type) }
+        end
+
+        def lazy_attr_accessor(*one_or_more_names, type: nil)
+          names = Array.wrap(one_or_more_names)
+          names.each do |name|
+            define_lazy_reader(name, type: type)
+            define_lazy_writer(name)
+          end
+        end
+
+        private
+
+        def define_lazy_reader(name, type:)
+          define_method(name) do
+            strong_memoize("#{name}_lazy_loaded") do
+              value = instance_variable_get("@#{name}")
+              value = value.call if value.respond_to?(:call)
+              value = nil if type && !value.is_a?(type)
+              value
+            end
+          end
+        end
+
+        def define_lazy_writer(name)
+          define_method("#{name}=") do |value|
+            clear_memoization("#{name}_lazy_loaded")
+            instance_variable_set("@#{name}", value)
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/spec/controllers/application_controller_spec.rb b/spec/controllers/application_controller_spec.rb
index e72ab16f62a7b..0c299dcda3484 100644
--- a/spec/controllers/application_controller_spec.rb
+++ b/spec/controllers/application_controller_spec.rb
@@ -895,4 +895,50 @@ def index
       end
     end
   end
+
+  context '#set_current_context' do
+    controller(described_class) do
+      def index
+        Labkit::Context.with_context do |context|
+          render json: context.to_h
+        end
+      end
+    end
+
+    let_it_be(:user) { create(:user) }
+
+    before do
+      sign_in(user)
+    end
+
+    it 'does not break anything when no group or project method is defined' do
+      get :index
+
+      expect(response).to have_gitlab_http_status(:success)
+    end
+
+    it 'sets the username in the context when signed in' do
+      get :index
+
+      expect(json_response['meta.user']).to eq(user.username)
+    end
+
+    it 'sets the group if it was available' do
+      group = build_stubbed(:group)
+      controller.instance_variable_set(:@group, group)
+
+      get :index, format: :json
+
+      expect(json_response['meta.root_namespace']).to eq(group.path)
+    end
+
+    it 'sets the project if one was available' do
+      project = build_stubbed(:project)
+      controller.instance_variable_set(:@project, project)
+
+      get :index, format: :json
+
+      expect(json_response['meta.project']).to eq(project.full_path)
+    end
+  end
 end
diff --git a/spec/lib/gitlab/application_context_spec.rb b/spec/lib/gitlab/application_context_spec.rb
new file mode 100644
index 0000000000000..1c8fcb8d737d6
--- /dev/null
+++ b/spec/lib/gitlab/application_context_spec.rb
@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::ApplicationContext do
+  describe '.with_context' do
+    it 'yields the block' do
+      expect { |b| described_class.with_context({}, &b) }.to yield_control
+    end
+
+    it 'passes the expected context on to labkit' do
+      fake_proc = duck_type(:call)
+      expected_context = hash_including(user: fake_proc, project: fake_proc, root_namespace: fake_proc)
+
+      expect(Labkit::Context).to receive(:with_context).with(expected_context)
+
+      described_class.with_context(
+        user: build(:user),
+        project: build(:project),
+        namespace: build(:namespace)) {}
+    end
+
+    it 'raises an error when passing invalid options' do
+      expect { described_class.with_context(no: 'option') {} }.to raise_error(ArgumentError)
+    end
+  end
+
+  describe '.push' do
+    it 'passes the expected context on to labkit' do
+      fake_proc = duck_type(:call)
+      expected_context = hash_including(user: fake_proc, project: fake_proc, root_namespace: fake_proc)
+
+      expect(Labkit::Context).to receive(:push).with(expected_context)
+
+      described_class.push(user: build(:user))
+    end
+
+    it 'raises an error when passing invalid options' do
+      expect { described_class.push(no: 'option')}.to raise_error(ArgumentError)
+    end
+  end
+
+  describe '#to_lazy_hash' do
+    let(:user) { build(:user) }
+    let(:project) { build(:project) }
+    let(:namespace) { build(:group) }
+    let(:subgroup) { build(:group, parent: namespace) }
+
+    def result(context)
+      context.to_lazy_hash.transform_values { |v| v.call }
+    end
+
+    it 'does not call the attributes until needed' do
+      fake_proc = double('Proc')
+
+      expect(fake_proc).not_to receive(:call)
+
+      described_class.new(user: fake_proc, project: fake_proc, namespace: fake_proc).to_lazy_hash
+    end
+
+    it 'correctly loads the expected values when they are wrapped in a block' do
+      context = described_class.new(user: -> { user }, project: -> { project }, namespace: -> { subgroup })
+
+      expect(result(context))
+        .to include(user: user.username, project: project.full_path, root_namespace: namespace.full_path)
+    end
+
+    it 'correctly loads the expected values when passed directly' do
+      context = described_class.new(user: user, project: project, namespace: subgroup)
+
+      expect(result(context))
+        .to include(user: user.username, project: project.full_path, root_namespace: namespace.full_path)
+    end
+
+    it 'falls back to a projects namespace when a project is passed but no namespace' do
+      context = described_class.new(project: project)
+
+      expect(result(context))
+        .to include(project: project.full_path, root_namespace: project.full_path_components.first)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/import_export/project_tree_restorer_spec.rb b/spec/lib/gitlab/import_export/project_tree_restorer_spec.rb
index ec1b935ad6309..dae5b028c1593 100644
--- a/spec/lib/gitlab/import_export/project_tree_restorer_spec.rb
+++ b/spec/lib/gitlab/import_export/project_tree_restorer_spec.rb
@@ -654,13 +654,10 @@
     let(:user) { create(:user) }
     let!(:project) { create(:project, :builds_disabled, :issues_disabled, name: 'project', path: 'project') }
     let(:project_tree_restorer) { described_class.new(user: user, shared: shared, project: project) }
-    let(:correlation_id) { 'my-correlation-id' }
 
     before do
       setup_import_export_config('with_invalid_records')
 
-      # Import is running from the rake task, `correlation_id` is not assigned
-      expect(Labkit::Correlation::CorrelationId).to receive(:new_id).and_return(correlation_id)
       subject
     end
 
@@ -682,7 +679,7 @@
         expect(import_failure.relation_index).to be_present
         expect(import_failure.exception_class).to eq('ActiveRecord::RecordInvalid')
         expect(import_failure.exception_message).to be_present
-        expect(import_failure.correlation_id_value).to eq('my-correlation-id')
+        expect(import_failure.correlation_id_value).not_to be_empty
         expect(import_failure.created_at).to be_present
       end
     end
diff --git a/spec/lib/gitlab/profiler_spec.rb b/spec/lib/gitlab/profiler_spec.rb
index c27e3ca7ace24..8f6fb6eda653d 100644
--- a/spec/lib/gitlab/profiler_spec.rb
+++ b/spec/lib/gitlab/profiler_spec.rb
@@ -84,7 +84,7 @@
           expect(severity).to eq(Logger::DEBUG)
           expect(message).to include('public').and include(described_class::FILTERED_STRING)
           expect(message).not_to include(private_token)
-        end.twice
+        end.at_least(1) # This spec could be wrapped in more blocks in the future
 
         custom_logger.debug("public #{private_token}")
       end
diff --git a/spec/lib/gitlab/sidekiq_middleware/correlation_injector_spec.rb b/spec/lib/gitlab/sidekiq_middleware/correlation_injector_spec.rb
deleted file mode 100644
index d5ed939485a48..0000000000000
--- a/spec/lib/gitlab/sidekiq_middleware/correlation_injector_spec.rb
+++ /dev/null
@@ -1,49 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-describe Gitlab::SidekiqMiddleware::CorrelationInjector do
-  class TestWorker
-    include ApplicationWorker
-  end
-
-  before do |example|
-    Sidekiq.client_middleware do |chain|
-      chain.add described_class
-    end
-  end
-
-  after do |example|
-    Sidekiq.client_middleware do |chain|
-      chain.remove described_class
-    end
-
-    Sidekiq::Queues.clear_all
-  end
-
-  around do |example|
-    Sidekiq::Testing.fake! do
-      example.run
-    end
-  end
-
-  it 'injects into payload the correlation id' do
-    expect_next_instance_of(described_class) do |instance|
-      expect(instance).to receive(:call).and_call_original
-    end
-
-    Labkit::Correlation::CorrelationId.use_id('new-correlation-id') do
-      TestWorker.perform_async(1234)
-    end
-
-    expected_job_params = {
-      "class" => "TestWorker",
-      "args" => [1234],
-      "correlation_id" => "new-correlation-id"
-    }
-
-    expect(Sidekiq::Queues.jobs_by_worker).to a_hash_including(
-      "TestWorker" => a_collection_containing_exactly(
-        a_hash_including(expected_job_params)))
-  end
-end
diff --git a/spec/lib/gitlab/sidekiq_middleware/correlation_logger_spec.rb b/spec/lib/gitlab/sidekiq_middleware/correlation_logger_spec.rb
deleted file mode 100644
index 27eea96340238..0000000000000
--- a/spec/lib/gitlab/sidekiq_middleware/correlation_logger_spec.rb
+++ /dev/null
@@ -1,35 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-describe Gitlab::SidekiqMiddleware::CorrelationLogger do
-  class TestWorker
-    include ApplicationWorker
-  end
-
-  before do |example|
-    Sidekiq::Testing.server_middleware do |chain|
-      chain.add described_class
-    end
-  end
-
-  after do |example|
-    Sidekiq::Testing.server_middleware do |chain|
-      chain.remove described_class
-    end
-  end
-
-  it 'injects into payload the correlation id', :sidekiq_might_not_need_inline do
-    expect_any_instance_of(described_class).to receive(:call).and_call_original
-
-    expect_any_instance_of(TestWorker).to receive(:perform).with(1234) do
-      expect(Labkit::Correlation::CorrelationId.current_id).to eq('new-correlation-id')
-    end
-
-    Sidekiq::Client.push(
-      'queue' => 'test',
-      'class' => TestWorker,
-      'args' => [1234],
-      'correlation_id' => 'new-correlation-id')
-  end
-end
diff --git a/spec/lib/gitlab/sidekiq_middleware_spec.rb b/spec/lib/gitlab/sidekiq_middleware_spec.rb
index aef472e0648c1..ef4a898bdb6b1 100644
--- a/spec/lib/gitlab/sidekiq_middleware_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware_spec.rb
@@ -38,7 +38,7 @@ def perform(_arg)
       [
        Gitlab::SidekiqMiddleware::Monitor,
        Gitlab::SidekiqMiddleware::BatchLoader,
-       Gitlab::SidekiqMiddleware::CorrelationLogger,
+       Labkit::Middleware::Sidekiq::Server,
        Gitlab::SidekiqMiddleware::InstrumentationLogger,
        Gitlab::SidekiqStatus::ServerMiddleware,
        Gitlab::SidekiqMiddleware::Metrics,
@@ -120,7 +120,7 @@ def perform(_arg)
       # This test ensures that this does not happen
       it "invokes the chain" do
         expect_any_instance_of(Gitlab::SidekiqStatus::ClientMiddleware).to receive(:call).with(*middleware_expected_args).once.and_call_original
-        expect_any_instance_of(Gitlab::SidekiqMiddleware::CorrelationInjector).to receive(:call).with(*middleware_expected_args).once.and_call_original
+        expect_any_instance_of(Labkit::Middleware::Sidekiq::Client).to receive(:call).with(*middleware_expected_args).once.and_call_original
 
         expect { |b| chain.invoke(worker_class_arg, job, queue, redis_pool, &b) }.to yield_control.once
       end
diff --git a/spec/lib/gitlab/utils/lazy_attributes_spec.rb b/spec/lib/gitlab/utils/lazy_attributes_spec.rb
new file mode 100644
index 0000000000000..c0005c194c4c3
--- /dev/null
+++ b/spec/lib/gitlab/utils/lazy_attributes_spec.rb
@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+require 'fast_spec_helper'
+require 'active_support/concern'
+
+describe Gitlab::Utils::LazyAttributes do
+  subject(:klass) do
+    Class.new do
+      include Gitlab::Utils::LazyAttributes
+
+      lazy_attr_reader :number, type: Numeric
+      lazy_attr_reader :reader_1, :reader_2
+      lazy_attr_accessor :incorrect_type, :string_attribute, :accessor_2, type: String
+
+      def initialize
+        @number = -> { 1 }
+        @reader_1, @reader_2 = 'reader_1', -> { 'reader_2' }
+        @incorrect_type, @accessor_2 = -> { :incorrect_type }, -> { 'accessor_2' }
+      end
+    end
+  end
+
+  context 'class methods' do
+    it { is_expected.to respond_to(:lazy_attr_reader, :lazy_attr_accessor) }
+    it { is_expected.not_to respond_to(:define_lazy_reader, :define_lazy_writer) }
+  end
+
+  context 'instance methods' do
+    subject(:instance) { klass.new }
+
+    it do
+      is_expected.to respond_to(:number, :reader_1, :reader_2, :incorrect_type,
+                                :incorrect_type=, :accessor_2, :accessor_2=,
+                                :string_attribute, :string_attribute=)
+    end
+
+    context 'reading attributes' do
+      it 'returns the correct values for procs', :aggregate_failures do
+        expect(instance.number).to eq(1)
+        expect(instance.reader_2).to eq('reader_2')
+        expect(instance.accessor_2).to eq('accessor_2')
+      end
+
+      it 'does not return the value if the type did not match what was specified' do
+        expect(instance.incorrect_type).to be_nil
+      end
+
+      it 'only calls the block once even if it returned `nil`', :aggregate_failures do
+        expect(instance.instance_variable_get('@number')).to receive(:call).once.and_call_original
+        expect(instance.instance_variable_get('@accessor_2')).to receive(:call).once.and_call_original
+        expect(instance.instance_variable_get('@incorrect_type')).to receive(:call).once.and_call_original
+
+        2.times do
+          instance.number
+          instance.incorrect_type
+          instance.accessor_2
+        end
+      end
+    end
+
+    context 'writing attributes' do
+      it 'sets the correct values', :aggregate_failures do
+        instance.string_attribute = -> { 'updated 1' }
+        instance.accessor_2 = nil
+
+        expect(instance.string_attribute).to eq('updated 1')
+        expect(instance.accessor_2).to be_nil
+      end
+    end
+  end
+end
diff --git a/spec/requests/api/internal/base_spec.rb b/spec/requests/api/internal/base_spec.rb
index 4fe0d4a59832b..12e6e7c7a093b 100644
--- a/spec/requests/api/internal/base_spec.rb
+++ b/spec/requests/api/internal/base_spec.rb
@@ -389,6 +389,12 @@
           end
         end
       end
+
+      it_behaves_like 'storing arguments in the application context' do
+        let(:expected_params) { { user: key.user.username, project: project.full_path } }
+
+        subject { push(key, project) }
+      end
     end
 
     context "access denied" do
@@ -885,6 +891,12 @@
       post api('/internal/post_receive'), params: valid_params
     end
 
+    it_behaves_like 'storing arguments in the application context' do
+      let(:expected_params) { { user: user.username, project: project.full_path } }
+
+      subject { post api('/internal/post_receive'), params: valid_params }
+    end
+
     context 'when there are merge_request push options' do
       before do
         valid_params[:push_options] = ['merge_request.create']
diff --git a/spec/requests/api/projects_spec.rb b/spec/requests/api/projects_spec.rb
index 9af4f484f9993..975e59346b859 100644
--- a/spec/requests/api/projects_spec.rb
+++ b/spec/requests/api/projects_spec.rb
@@ -1626,6 +1626,14 @@
         end
       end
     end
+
+    it_behaves_like 'storing arguments in the application context' do
+      let_it_be(:user) { create(:user) }
+      let_it_be(:project) { create(:project, :public) }
+      let(:expected_params) { { user: user.username, project: project.full_path } }
+
+      subject { get api("/projects/#{project.id}", user) }
+    end
   end
 
   describe 'GET /projects/:id/users' do
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 1f0119108a850..6393e48290453 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -245,6 +245,12 @@
     Rails.cache = caching_store
   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 }
+  end
+
   config.around(:each, :clean_gitlab_redis_cache) do |example|
     redis_cache_cleanup!
 
diff --git a/spec/support/shared_examples/logging_application_context_shared_examples.rb b/spec/support/shared_examples/logging_application_context_shared_examples.rb
new file mode 100644
index 0000000000000..038ede884c88c
--- /dev/null
+++ b/spec/support/shared_examples/logging_application_context_shared_examples.rb
@@ -0,0 +1,24 @@
+# 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
+    # 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
+  end
+
+  def log_hash(hash)
+    hash.transform_keys! { |key| "meta.#{key}" }
+  end
+end
-- 
GitLab