Add latest changes from gitlab-org/gitlab@13-11-stable-eev13.11.0-rc43

6 files changed, 186 insertions, 396 deletions

diff --git a/spec/lib/gitlab/sidekiq_middleware/admin_mode/client_spec.rb b/spec/lib/gitlab/sidekiq_middleware/admin_mode/client_spec.rb
index 3ba08455d01..9d5d5f28eab 100644
--- a/spec/lib/gitlab/sidekiq_middleware/admin_mode/client_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/admin_mode/client_spec.rb
@@ -74,9 +74,9 @@ RSpec.describe Gitlab::SidekiqMiddleware::AdminMode::Client, :request_store do
     end
   end
 
-  context 'admin mode feature disabled' do
+  context 'admin mode setting disabled' do
     before do
-      stub_feature_flags(user_mode_in_session: false)
+      stub_application_setting(admin_mode: false)
     end
 
     it 'yields block' do
diff --git a/spec/lib/gitlab/sidekiq_middleware/admin_mode/server_spec.rb b/spec/lib/gitlab/sidekiq_middleware/admin_mode/server_spec.rb
index e8322b11875..3ab1a9cd2f4 100644
--- a/spec/lib/gitlab/sidekiq_middleware/admin_mode/server_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/admin_mode/server_spec.rb
@@ -52,9 +52,9 @@ RSpec.describe Gitlab::SidekiqMiddleware::AdminMode::Server, :request_store do
     end
   end
 
-  context 'admin mode feature disabled' do
+  context 'admin mode setting disabled' do
     before do
-      stub_feature_flags(user_mode_in_session: false)
+      stub_application_setting(admin_mode: false)
     end
 
     it 'yields block' do
diff --git a/spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb b/spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb
index e2b36125b4e..82ca84f0697 100644
--- a/spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb
@@ -3,156 +3,33 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::SidekiqMiddleware::ClientMetrics do
-  context "with worker attribution" do
-    subject { described_class.new }
+  shared_examples "a metrics middleware" do
+    context "with mocked prometheus" do
+      let(:enqueued_jobs_metric) { double('enqueued jobs metric', increment: true) }
 
-    let(:queue) { :test }
-    let(:worker_class) { worker.class }
-    let(:job) { {} }
-    let(:default_labels) do
-      { queue: queue.to_s,
-        worker: worker_class.to_s,
-        boundary: "",
-        external_dependencies: "no",
-        feature_category: "",
-        urgency: "low" }
-    end
-
-    shared_examples "a metrics client middleware" do
-      context "with mocked prometheus" do
-        let(:enqueued_jobs_metric) { double('enqueued jobs metric', increment: true) }
-
-        before do
-          allow(Gitlab::Metrics).to receive(:counter).with(described_class::ENQUEUED, anything).and_return(enqueued_jobs_metric)
-        end
-
-        describe '#call' do
-          it 'yields block' do
-            expect { |b| subject.call(worker_class, job, :test, double, &b) }.to yield_control.once
-          end
-
-          it 'increments enqueued jobs metric with correct labels when worker is a string of the class' do
-            expect(enqueued_jobs_metric).to receive(:increment).with(labels, 1)
-
-            subject.call(worker_class.to_s, job, :test, double) { nil }
-          end
-
-          it 'increments enqueued jobs metric with correct labels' do
-            expect(enqueued_jobs_metric).to receive(:increment).with(labels, 1)
-
-            subject.call(worker_class, job, :test, double) { nil }
-          end
-        end
-      end
-    end
-
-    context "when workers are not attributed" do
       before do
-        stub_const('TestNonAttributedWorker', Class.new)
-        TestNonAttributedWorker.class_eval do
-          include Sidekiq::Worker
-        end
-      end
-
-      it_behaves_like "a metrics client middleware" do
-        let(:worker) { TestNonAttributedWorker.new }
-        let(:labels) { default_labels.merge(urgency: "") }
-      end
-    end
-
-    context "when a worker is wrapped into ActiveJob" do
-      before do
-        stub_const('TestWrappedWorker', Class.new)
-        TestWrappedWorker.class_eval do
-          include Sidekiq::Worker
-        end
-      end
-
-      it_behaves_like "a metrics client middleware" do
-        let(:job) do
-          {
-            "class" => ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper,
-            "wrapped" => TestWrappedWorker
-          }
-        end
-
-        let(:worker) { TestWrappedWorker.new }
-        let(:labels) { default_labels.merge(urgency: "") }
-      end
-    end
-
-    context "when workers are attributed" do
-      def create_attributed_worker_class(urgency, external_dependencies, resource_boundary, category)
-        klass = Class.new do
-          include Sidekiq::Worker
-          include WorkerAttributes
-
-          urgency urgency if urgency
-          worker_has_external_dependencies! if external_dependencies
-          worker_resource_boundary resource_boundary unless resource_boundary == :unknown
-          feature_category category unless category.nil?
-        end
-        stub_const("TestAttributedWorker", klass)
-      end
-
-      let(:urgency) { nil }
-      let(:external_dependencies) { false }
-      let(:resource_boundary) { :unknown }
-      let(:feature_category) { nil }
-      let(:worker_class) { create_attributed_worker_class(urgency, external_dependencies, resource_boundary, feature_category) }
-      let(:worker) { worker_class.new }
-
-      context "high urgency" do
-        it_behaves_like "a metrics client middleware" do
-          let(:urgency) { :high }
-          let(:labels) { default_labels.merge(urgency: "high") }
-        end
+        allow(Gitlab::Metrics).to receive(:counter).with(described_class::ENQUEUED, anything).and_return(enqueued_jobs_metric)
       end
 
-      context "no urgency" do
-        it_behaves_like "a metrics client middleware" do
-          let(:urgency) { :throttled }
-          let(:labels) { default_labels.merge(urgency: "throttled") }
+      describe '#call' do
+        it 'yields block' do
+          expect { |b| subject.call(worker_class, job, :test, double, &b) }.to yield_control.once
         end
-      end
 
-      context "external dependencies" do
-        it_behaves_like "a metrics client middleware" do
-          let(:external_dependencies) { true }
-          let(:labels) { default_labels.merge(external_dependencies: "yes") }
-        end
-      end
+        it 'increments enqueued jobs metric with correct labels when worker is a string of the class' do
+          expect(enqueued_jobs_metric).to receive(:increment).with(labels, 1)
 
-      context "cpu boundary" do
-        it_behaves_like "a metrics client middleware" do
-          let(:resource_boundary) { :cpu }
-          let(:labels) { default_labels.merge(boundary: "cpu") }
+          subject.call(worker_class.to_s, job, :test, double) { nil }
         end
-      end
 
-      context "memory boundary" do
-        it_behaves_like "a metrics client middleware" do
-          let(:resource_boundary) { :memory }
-          let(:labels) { default_labels.merge(boundary: "memory") }
-        end
-      end
+        it 'increments enqueued jobs metric with correct labels' do
+          expect(enqueued_jobs_metric).to receive(:increment).with(labels, 1)
 
-      context "feature category" do
-        it_behaves_like "a metrics client middleware" do
-          let(:feature_category) { :authentication }
-          let(:labels) { default_labels.merge(feature_category: "authentication") }
-        end
-      end
-
-      context "combined" do
-        it_behaves_like "a metrics client middleware" do
-          let(:urgency) { :high }
-          let(:external_dependencies) { true }
-          let(:resource_boundary) { :cpu }
-          let(:feature_category) { :authentication }
-          let(:labels) { default_labels.merge(urgency: "high", external_dependencies: "yes", boundary: "cpu", feature_category: "authentication") }
+          subject.call(worker_class, job, :test, double) { nil }
         end
       end
     end
   end
+
+  it_behaves_like 'metrics middleware with worker attribution'
 end
diff --git a/spec/lib/gitlab/sidekiq_middleware/instrumentation_logger_spec.rb b/spec/lib/gitlab/sidekiq_middleware/instrumentation_logger_spec.rb
new file mode 100644
index 00000000000..eb9ba50cdcd
--- /dev/null
+++ b/spec/lib/gitlab/sidekiq_middleware/instrumentation_logger_spec.rb
@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::SidekiqMiddleware::InstrumentationLogger do
+  let(:job) { { 'jid' => 123 } }
+  let(:queue) { 'test_queue' }
+  let(:worker) do
+    Class.new do
+      def self.name
+        'TestDWorker'
+      end
+
+      include ApplicationWorker
+
+      def perform(*args)
+      end
+    end
+  end
+
+  subject { described_class.new }
+
+  before do
+    stub_const('TestWorker', worker)
+  end
+
+  describe '.keys' do
+    it 'returns all available payload keys' do
+      expected_keys = [
+        :cpu_s,
+        :gitaly_calls,
+        :gitaly_duration_s,
+        :rugged_calls,
+        :rugged_duration_s,
+        :elasticsearch_calls,
+        :elasticsearch_duration_s,
+        :elasticsearch_timed_out_count,
+        :mem_objects,
+        :mem_bytes,
+        :mem_mallocs,
+        :redis_calls,
+        :redis_duration_s,
+        :redis_read_bytes,
+        :redis_write_bytes,
+        :redis_action_cable_calls,
+        :redis_action_cable_duration_s,
+        :redis_action_cable_read_bytes,
+        :redis_action_cable_write_bytes,
+        :redis_cache_calls,
+        :redis_cache_duration_s,
+        :redis_cache_read_bytes,
+        :redis_cache_write_bytes,
+        :redis_queues_calls,
+        :redis_queues_duration_s,
+        :redis_queues_read_bytes,
+        :redis_queues_write_bytes,
+        :redis_shared_state_calls,
+        :redis_shared_state_duration_s,
+        :redis_shared_state_read_bytes,
+        :redis_shared_state_write_bytes,
+        :db_count,
+        :db_write_count,
+        :db_cached_count,
+        :external_http_count,
+        :external_http_duration_s,
+        :rack_attack_redis_count,
+        :rack_attack_redis_duration_s
+      ]
+
+      expect(described_class.keys).to include(*expected_keys)
+    end
+  end
+
+  describe '#call', :request_store do
+    let(:instrumentation_values) do
+      {
+        cpu_s: 10,
+        unknown_attribute: 123,
+        db_count: 0,
+        db_cached_count: 0,
+        db_write_count: 0,
+        gitaly_calls: 0,
+        redis_calls: 0
+      }
+    end
+
+    before do
+      allow(::Gitlab::InstrumentationHelper).to receive(:add_instrumentation_data) do |values|
+        values.merge!(instrumentation_values)
+      end
+    end
+
+    it 'merges correct instrumentation data in the job' do
+      expect { |b| subject.call(worker, job, queue, &b) }.to yield_control
+
+      expected_values = instrumentation_values.except(:unknown_attribute)
+
+      expect(job[:instrumentation]).to eq(expected_values)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb b/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
index 71f4f2a3b64..95be76ce351 100644
--- a/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
@@ -4,296 +4,108 @@ require 'spec_helper'
 
 # rubocop: disable RSpec/MultipleMemoizedHelpers
 RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
-  context "with worker attribution" do
-    subject { described_class.new }
+  shared_examples "a metrics middleware" do
+    context "with mocked prometheus" do
+      include_context 'server metrics with mocked prometheus'
 
-    let(:queue) { :test }
-    let(:worker_class) { worker.class }
-    let(:job) { {} }
-    let(:job_status) { :done }
-    let(:labels_with_job_status) { labels.merge(job_status: job_status.to_s) }
-    let(:default_labels) do
-      { queue: queue.to_s,
-        worker: worker_class.to_s,
-        boundary: "",
-        external_dependencies: "no",
-        feature_category: "",
-        urgency: "low" }
-    end
-
-    shared_examples "a metrics middleware" do
-      context "with mocked prometheus" do
-        let(:concurrency_metric) { double('concurrency metric') }
-
-        let(:queue_duration_seconds) { double('queue duration seconds metric') }
-        let(:completion_seconds_metric) { double('completion seconds metric') }
-        let(:user_execution_seconds_metric) { double('user execution seconds metric') }
-        let(:db_seconds_metric) { double('db seconds metric') }
-        let(:gitaly_seconds_metric) { double('gitaly seconds metric') }
-        let(:failed_total_metric) { double('failed total metric') }
-        let(:retried_total_metric) { double('retried total metric') }
-        let(:redis_requests_total) { double('redis calls total metric') }
-        let(:running_jobs_metric) { double('running jobs metric') }
-        let(:redis_seconds_metric) { double('redis seconds metric') }
-        let(:elasticsearch_seconds_metric) { double('elasticsearch seconds metric') }
-        let(:elasticsearch_requests_total) { double('elasticsearch calls total metric') }
+      describe '#initialize' do
+        it 'sets concurrency metrics' do
+          expect(concurrency_metric).to receive(:set).with({}, Sidekiq.options[:concurrency].to_i)
 
-        before do
-          allow(Gitlab::Metrics).to receive(:histogram).with(:sidekiq_jobs_queue_duration_seconds, anything, anything, anything).and_return(queue_duration_seconds)
-          allow(Gitlab::Metrics).to receive(:histogram).with(:sidekiq_jobs_completion_seconds, anything, anything, anything).and_return(completion_seconds_metric)
-          allow(Gitlab::Metrics).to receive(:histogram).with(:sidekiq_jobs_cpu_seconds, anything, anything, anything).and_return(user_execution_seconds_metric)
-          allow(Gitlab::Metrics).to receive(:histogram).with(:sidekiq_jobs_db_seconds, anything, anything, anything).and_return(db_seconds_metric)
-          allow(Gitlab::Metrics).to receive(:histogram).with(:sidekiq_jobs_gitaly_seconds, anything, anything, anything).and_return(gitaly_seconds_metric)
-          allow(Gitlab::Metrics).to receive(:histogram).with(:sidekiq_redis_requests_duration_seconds, anything, anything, anything).and_return(redis_seconds_metric)
-          allow(Gitlab::Metrics).to receive(:histogram).with(:sidekiq_elasticsearch_requests_duration_seconds, anything, anything, anything).and_return(elasticsearch_seconds_metric)
-          allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_jobs_failed_total, anything).and_return(failed_total_metric)
-          allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_jobs_retried_total, anything).and_return(retried_total_metric)
-          allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_redis_requests_total, anything).and_return(redis_requests_total)
-          allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_elasticsearch_requests_total, anything).and_return(elasticsearch_requests_total)
-          allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_running_jobs, anything, {}, :all).and_return(running_jobs_metric)
-          allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_concurrency, anything, {}, :all).and_return(concurrency_metric)
-
-          allow(concurrency_metric).to receive(:set)
+          subject
         end
+      end
 
-        describe '#initialize' do
-          it 'sets concurrency metrics' do
-            expect(concurrency_metric).to receive(:set).with({}, Sidekiq.options[:concurrency].to_i)
+      describe '#call' do
+        include_context 'server metrics call'
 
-            subject
-          end
+        it 'yields block' do
+          expect { |b| subject.call(worker, job, :test, &b) }.to yield_control.once
         end
 
-        describe '#call' do
-          let(:thread_cputime_before) { 1 }
-          let(:thread_cputime_after) { 2 }
-          let(:thread_cputime_duration) { thread_cputime_after - thread_cputime_before }
-
-          let(:monotonic_time_before) { 11 }
-          let(:monotonic_time_after) { 20 }
-          let(:monotonic_time_duration) { monotonic_time_after - monotonic_time_before }
-
-          let(:queue_duration_for_job) { 0.01 }
-
-          let(:db_duration) { 3 }
-          let(:gitaly_duration) { 4 }
-
-          let(:redis_calls) { 2 }
-          let(:redis_duration) { 0.01 }
-
-          let(:elasticsearch_calls) { 8 }
-          let(:elasticsearch_duration) { 0.54 }
-
-          before do
-            allow(subject).to receive(:get_thread_cputime).and_return(thread_cputime_before, thread_cputime_after)
-            allow(Gitlab::Metrics::System).to receive(:monotonic_time).and_return(monotonic_time_before, monotonic_time_after)
-            allow(Gitlab::InstrumentationHelper).to receive(:queue_duration_for_job).with(job).and_return(queue_duration_for_job)
-            allow(ActiveRecord::LogSubscriber).to receive(:runtime).and_return(db_duration * 1000)
-
-            job[:gitaly_duration_s] = gitaly_duration
-            job[:redis_calls] = redis_calls
-            job[:redis_duration_s] = redis_duration
-
-            job[:elasticsearch_calls] = elasticsearch_calls
-            job[:elasticsearch_duration_s] = elasticsearch_duration
-
-            allow(running_jobs_metric).to receive(:increment)
-            allow(redis_requests_total).to receive(:increment)
-            allow(elasticsearch_requests_total).to receive(:increment)
-            allow(queue_duration_seconds).to receive(:observe)
-            allow(user_execution_seconds_metric).to receive(:observe)
-            allow(db_seconds_metric).to receive(:observe)
-            allow(gitaly_seconds_metric).to receive(:observe)
-            allow(completion_seconds_metric).to receive(:observe)
-            allow(redis_seconds_metric).to receive(:observe)
-            allow(elasticsearch_seconds_metric).to receive(:observe)
+        it 'calls BackgroundTransaction' do
+          expect_next_instance_of(Gitlab::Metrics::BackgroundTransaction) do |instance|
+            expect(instance).to receive(:run)
           end
 
-          it 'yields block' do
-            expect { |b| subject.call(worker, job, :test, &b) }.to yield_control.once
-          end
+          subject.call(worker, job, :test) {}
+        end
 
-          it 'calls BackgroundTransaction' do
-            expect_next_instance_of(Gitlab::Metrics::BackgroundTransaction) do |instance|
-              expect(instance).to receive(:run)
-            end
+        it 'sets queue specific metrics' do
+          expect(running_jobs_metric).to receive(:increment).with(labels, -1)
+          expect(running_jobs_metric).to receive(:increment).with(labels, 1)
+          expect(queue_duration_seconds).to receive(:observe).with(labels, queue_duration_for_job) if queue_duration_for_job
+          expect(user_execution_seconds_metric).to receive(:observe).with(labels_with_job_status, thread_cputime_duration)
+          expect(db_seconds_metric).to receive(:observe).with(labels_with_job_status, db_duration)
+          expect(gitaly_seconds_metric).to receive(:observe).with(labels_with_job_status, gitaly_duration)
+          expect(completion_seconds_metric).to receive(:observe).with(labels_with_job_status, monotonic_time_duration)
+          expect(redis_seconds_metric).to receive(:observe).with(labels_with_job_status, redis_duration)
+          expect(elasticsearch_seconds_metric).to receive(:observe).with(labels_with_job_status, elasticsearch_duration)
+          expect(redis_requests_total).to receive(:increment).with(labels_with_job_status, redis_calls)
+          expect(elasticsearch_requests_total).to receive(:increment).with(labels_with_job_status, elasticsearch_calls)
+
+          subject.call(worker, job, :test) { nil }
+        end
 
-            subject.call(worker, job, :test) {}
-          end
+        it 'sets the thread name if it was nil' do
+          allow(Thread.current).to receive(:name).and_return(nil)
+          expect(Thread.current).to receive(:name=).with(Gitlab::Metrics::Samplers::ThreadsSampler::SIDEKIQ_WORKER_THREAD_NAME)
 
-          it 'sets queue specific metrics' do
-            expect(running_jobs_metric).to receive(:increment).with(labels, -1)
-            expect(running_jobs_metric).to receive(:increment).with(labels, 1)
-            expect(queue_duration_seconds).to receive(:observe).with(labels, queue_duration_for_job) if queue_duration_for_job
-            expect(user_execution_seconds_metric).to receive(:observe).with(labels_with_job_status, thread_cputime_duration)
-            expect(db_seconds_metric).to receive(:observe).with(labels_with_job_status, db_duration)
-            expect(gitaly_seconds_metric).to receive(:observe).with(labels_with_job_status, gitaly_duration)
-            expect(completion_seconds_metric).to receive(:observe).with(labels_with_job_status, monotonic_time_duration)
-            expect(redis_seconds_metric).to receive(:observe).with(labels_with_job_status, redis_duration)
-            expect(elasticsearch_seconds_metric).to receive(:observe).with(labels_with_job_status, elasticsearch_duration)
-            expect(redis_requests_total).to receive(:increment).with(labels_with_job_status, redis_calls)
-            expect(elasticsearch_requests_total).to receive(:increment).with(labels_with_job_status, elasticsearch_calls)
+          subject.call(worker, job, :test) { nil }
+        end
 
-            subject.call(worker, job, :test) { nil }
-          end
+        context 'when job_duration is not available' do
+          let(:queue_duration_for_job) { nil }
 
-          it 'sets the thread name if it was nil' do
-            allow(Thread.current).to receive(:name).and_return(nil)
-            expect(Thread.current).to receive(:name=).with(Gitlab::Metrics::Samplers::ThreadsSampler::SIDEKIQ_WORKER_THREAD_NAME)
+          it 'does not set the queue_duration_seconds histogram' do
+            expect(queue_duration_seconds).not_to receive(:observe)
 
             subject.call(worker, job, :test) { nil }
           end
+        end
 
-          context 'when job_duration is not available' do
-            let(:queue_duration_for_job) { nil }
-
-            it 'does not set the queue_duration_seconds histogram' do
-              expect(queue_duration_seconds).not_to receive(:observe)
-
-              subject.call(worker, job, :test) { nil }
-            end
-          end
-
-          context 'when error is raised' do
-            let(:job_status) { :fail }
-
-            it 'sets sidekiq_jobs_failed_total and reraises' do
-              expect(failed_total_metric).to receive(:increment).with(labels, 1)
-
-              expect { subject.call(worker, job, :test) { raise StandardError, "Failed" } }.to raise_error(StandardError, "Failed")
-            end
-          end
-
-          context 'when job is retried' do
-            let(:job) { { 'retry_count' => 1 } }
+        context 'when error is raised' do
+          let(:job_status) { :fail }
 
-            it 'sets sidekiq_jobs_retried_total metric' do
-              expect(retried_total_metric).to receive(:increment)
+          it 'sets sidekiq_jobs_failed_total and reraises' do
+            expect(failed_total_metric).to receive(:increment).with(labels, 1)
 
-              subject.call(worker, job, :test) { nil }
-            end
+            expect { subject.call(worker, job, :test) { raise StandardError, "Failed" } }.to raise_error(StandardError, "Failed")
           end
         end
-      end
 
-      context "with prometheus integrated" do
-        describe '#call' do
-          it 'yields block' do
-            expect { |b| subject.call(worker, job, :test, &b) }.to yield_control.once
-          end
+        context 'when job is retried' do
+          let(:job) { { 'retry_count' => 1 } }
 
-          context 'when error is raised' do
-            let(:job_status) { :fail }
+          it 'sets sidekiq_jobs_retried_total metric' do
+            expect(retried_total_metric).to receive(:increment)
 
-            it 'sets sidekiq_jobs_failed_total and reraises' do
-              expect { subject.call(worker, job, :test) { raise StandardError, "Failed" } }.to raise_error(StandardError, "Failed")
-            end
+            subject.call(worker, job, :test) { nil }
           end
         end
       end
     end
 
-    context "when workers are not attributed" do
-      before do
-        stub_const('TestNonAttributedWorker', Class.new)
-        TestNonAttributedWorker.class_eval do
-          include Sidekiq::Worker
+    context "with prometheus integrated" do
+      describe '#call' do
+        it 'yields block' do
+          expect { |b| subject.call(worker, job, :test, &b) }.to yield_control.once
         end
-      end
 
-      let(:worker) { TestNonAttributedWorker.new }
-      let(:labels) { default_labels.merge(urgency: "") }
+        context 'when error is raised' do
+          let(:job_status) { :fail }
 
-      it_behaves_like "a metrics middleware"
-    end
-
-    context "when a worker is wrapped into ActiveJob" do
-      before do
-        stub_const('TestWrappedWorker', Class.new)
-        TestWrappedWorker.class_eval do
-          include Sidekiq::Worker
+          it 'sets sidekiq_jobs_failed_total and reraises' do
+            expect { subject.call(worker, job, :test) { raise StandardError, "Failed" } }.to raise_error(StandardError, "Failed")
+          end
         end
       end
-
-      let(:job) do
-        {
-          "class" => ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper,
-          "wrapped" => TestWrappedWorker
-        }
-      end
-
-      let(:worker) { TestWrappedWorker.new }
-      let(:worker_class) { TestWrappedWorker }
-      let(:labels) { default_labels.merge(urgency: "") }
-
-      it_behaves_like "a metrics middleware"
     end
+  end
 
-    context "when workers are attributed" do
-      def create_attributed_worker_class(urgency, external_dependencies, resource_boundary, category)
-        Class.new do
-          include Sidekiq::Worker
-          include WorkerAttributes
-
-          urgency urgency if urgency
-          worker_has_external_dependencies! if external_dependencies
-          worker_resource_boundary resource_boundary unless resource_boundary == :unknown
-          feature_category category unless category.nil?
-        end
-      end
-
-      let(:urgency) { nil }
-      let(:external_dependencies) { false }
-      let(:resource_boundary) { :unknown }
-      let(:feature_category) { nil }
-      let(:worker_class) { create_attributed_worker_class(urgency, external_dependencies, resource_boundary, feature_category) }
-      let(:worker) { worker_class.new }
-
-      context "high urgency" do
-        let(:urgency) { :high }
-        let(:labels) { default_labels.merge(urgency: "high") }
-
-        it_behaves_like "a metrics middleware"
-      end
-
-      context "external dependencies" do
-        let(:external_dependencies) { true }
-        let(:labels) { default_labels.merge(external_dependencies: "yes") }
-
-        it_behaves_like "a metrics middleware"
-      end
-
-      context "cpu boundary" do
-        let(:resource_boundary) { :cpu }
-        let(:labels) { default_labels.merge(boundary: "cpu") }
-
-        it_behaves_like "a metrics middleware"
-      end
-
-      context "memory boundary" do
-        let(:resource_boundary) { :memory }
-        let(:labels) { default_labels.merge(boundary: "memory") }
-
-        it_behaves_like "a metrics middleware"
-      end
-
-      context "feature category" do
-        let(:feature_category) { :authentication }
-        let(:labels) { default_labels.merge(feature_category: "authentication") }
-
-        it_behaves_like "a metrics middleware"
-      end
-
-      context "combined" do
-        let(:urgency) { :throttled }
-        let(:external_dependencies) { true }
-        let(:resource_boundary) { :cpu }
-        let(:feature_category) { :authentication }
-        let(:labels) { default_labels.merge(urgency: "throttled", external_dependencies: "yes", boundary: "cpu", feature_category: "authentication") }
-
-        it_behaves_like "a metrics middleware"
-      end
-    end
+  it_behaves_like 'metrics middleware with worker attribution' do
+    let(:job_status) { :done }
+    let(:labels_with_job_status) { labels.merge(job_status: job_status.to_s) }
   end
 end
 # rubocop: enable RSpec/MultipleMemoizedHelpers
diff --git a/spec/lib/gitlab/sidekiq_middleware/worker_context/server_spec.rb b/spec/lib/gitlab/sidekiq_middleware/worker_context/server_spec.rb
index ca473462d2e..f736a7db774 100644
--- a/spec/lib/gitlab/sidekiq_middleware/worker_context/server_spec.rb
+++ b/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