Add latest changes from gitlab-org/gitlab@14-1-stable-eev14.1.0-rc42

5 files changed, 239 insertions, 65 deletions

diff --git a/spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb b/spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb
index 82ca84f0697..698758a13fd 100644
--- a/spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb
@@ -3,11 +3,12 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::SidekiqMiddleware::ClientMetrics do
+  let(:enqueued_jobs_metric) { double('enqueued jobs metric', increment: true) }
+
   shared_examples "a metrics middleware" do
     context "with mocked prometheus" do
-      let(:enqueued_jobs_metric) { double('enqueued jobs metric', increment: true) }
-
       before do
+        labels[:scheduling] = 'immediate'
         allow(Gitlab::Metrics).to receive(:counter).with(described_class::ENQUEUED, anything).and_return(enqueued_jobs_metric)
       end
 
@@ -32,4 +33,35 @@ RSpec.describe Gitlab::SidekiqMiddleware::ClientMetrics do
   end
 
   it_behaves_like 'metrics middleware with worker attribution'
+
+  context 'when mounted' do
+    before do
+      stub_const('TestWorker', Class.new)
+      TestWorker.class_eval do
+        include Sidekiq::Worker
+
+        def perform(*args)
+        end
+      end
+
+      allow(Gitlab::Metrics).to receive(:counter).and_return(Gitlab::Metrics::NullMetric.instance)
+      allow(Gitlab::Metrics).to receive(:counter).with(described_class::ENQUEUED, anything).and_return(enqueued_jobs_metric)
+    end
+
+    context 'when scheduling jobs for immediate execution' do
+      it 'increments enqueued jobs metric with scheduling label set to immediate' do
+        expect(enqueued_jobs_metric).to receive(:increment).with(a_hash_including(scheduling: 'immediate'), 1)
+
+        Sidekiq::Testing.inline! { TestWorker.perform_async }
+      end
+    end
+
+    context 'when scheduling jobs for future execution' do
+      it 'increments enqueued jobs metric with scheduling label set to delayed' do
+        expect(enqueued_jobs_metric).to receive(:increment).with(a_hash_including(scheduling: 'delayed'), 1)
+
+        Sidekiq::Testing.inline! { TestWorker.perform_in(1.second) }
+      end
+    end
+  end
 end
diff --git a/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job_spec.rb b/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job_spec.rb
index a10a8883591..d67cb95f483 100644
--- a/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job_spec.rb
@@ -13,7 +13,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
   let(:queue) { 'authorized_projects' }
 
   let(:idempotency_key) do
-    hash = Digest::SHA256.hexdigest("#{job['class']}:#{job['args'].join('-')}")
+    hash = Digest::SHA256.hexdigest("#{job['class']}:#{Sidekiq.dump_json(job['args'])}")
     "#{Gitlab::Redis::Queues::SIDEKIQ_NAMESPACE}:duplicate:#{queue}:#{hash}"
   end
 
diff --git a/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb b/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
index 34b4541f339..3ec8d404bf0 100644
--- a/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
@@ -8,11 +8,77 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
     context "with mocked prometheus" do
       include_context 'server metrics with mocked prometheus'
 
-      describe '#initialize' do
+      describe '.initialize_process_metrics' do
         it 'sets concurrency metrics' do
           expect(concurrency_metric).to receive(:set).with({}, Sidekiq.options[:concurrency].to_i)
 
-          subject
+          described_class.initialize_process_metrics
+        end
+
+        it 'initializes sidekiq_jobs_completion_seconds for the workers in the current Sidekiq process' do
+          allow(Gitlab::SidekiqConfig)
+            .to receive(:current_worker_queue_mappings)
+                  .and_return('MergeWorker' => 'merge', 'Ci::BuildFinishedWorker' => 'default')
+
+          expect(completion_seconds_metric)
+            .to receive(:get).with(queue: 'merge',
+                                   worker: 'MergeWorker',
+                                   urgency: 'high',
+                                   external_dependencies: 'no',
+                                   feature_category: 'source_code_management',
+                                   boundary: '',
+                                   job_status: 'done')
+
+          expect(completion_seconds_metric)
+            .to receive(:get).with(queue: 'merge',
+                                   worker: 'MergeWorker',
+                                   urgency: 'high',
+                                   external_dependencies: 'no',
+                                   feature_category: 'source_code_management',
+                                   boundary: '',
+                                   job_status: 'fail')
+
+          expect(completion_seconds_metric)
+            .to receive(:get).with(queue: 'default',
+                                   worker: 'Ci::BuildFinishedWorker',
+                                   urgency: 'high',
+                                   external_dependencies: 'no',
+                                   feature_category: 'continuous_integration',
+                                   boundary: 'cpu',
+                                   job_status: 'done')
+
+          expect(completion_seconds_metric)
+            .to receive(:get).with(queue: 'default',
+                                   worker: 'Ci::BuildFinishedWorker',
+                                   urgency: 'high',
+                                   external_dependencies: 'no',
+                                   feature_category: 'continuous_integration',
+                                   boundary: 'cpu',
+                                   job_status: 'fail')
+
+          described_class.initialize_process_metrics
+        end
+
+        context 'when the sidekiq_job_completion_metric_initialize feature flag is disabled' do
+          before do
+            stub_feature_flags(sidekiq_job_completion_metric_initialize: false)
+          end
+
+          it 'sets the concurrency metric' do
+            expect(concurrency_metric).to receive(:set).with({}, Sidekiq.options[:concurrency].to_i)
+
+            described_class.initialize_process_metrics
+          end
+
+          it 'does not initialize sidekiq_jobs_completion_seconds' do
+            allow(Gitlab::SidekiqConfig)
+              .to receive(:current_worker_queue_mappings)
+                    .and_return('MergeWorker' => 'merge', 'Ci::BuildFinishedWorker' => 'default')
+
+            expect(completion_seconds_metric).not_to receive(:get)
+
+            described_class.initialize_process_metrics
+          end
         end
       end
 
@@ -47,6 +113,26 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
           subject.call(worker, job, :test) { nil }
         end
 
+        it 'sets sidekiq_jobs_completion_seconds values that are compatible with those from .initialize_process_metrics' do
+          label_validator = Prometheus::Client::LabelSetValidator.new([:le])
+
+          allow(Gitlab::SidekiqConfig)
+            .to receive(:current_worker_queue_mappings)
+                  .and_return('MergeWorker' => 'merge', 'Ci::BuildFinishedWorker' => 'default')
+
+          allow(completion_seconds_metric).to receive(:get) do |labels|
+            expect { label_validator.validate(labels) }.not_to raise_error
+          end
+
+          allow(completion_seconds_metric).to receive(:observe) do |labels, _duration|
+            expect { label_validator.validate(labels) }.not_to raise_error
+          end
+
+          described_class.initialize_process_metrics
+
+          subject.call(worker, job, :test) { nil }
+        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)
@@ -109,22 +195,20 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
   end
 
   context 'DB load balancing' do
-    using RSpec::Parameterized::TableSyntax
-
     subject { described_class.new }
 
     let(:queue) { :test }
     let(:worker_class) { worker.class }
-    let(:job) { {} }
-    let(:job_status) { :done }
-    let(:labels_with_job_status) { default_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" }
+    let(:worker) { TestWorker.new }
+    let(:client_middleware) { Gitlab::Database::LoadBalancing::SidekiqClientMiddleware.new }
+    let(:load_balancer) { double.as_null_object }
+    let(:load_balancing_metric) { double('load balancing metric') }
+    let(:job) { { "retry" => 3, "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e" } }
+
+    def process_job
+      client_middleware.call(worker_class, job, queue, double) do
+        worker_class.process_job(job)
+      end
     end
 
     before do
@@ -132,84 +216,97 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
       TestWorker.class_eval do
         include Sidekiq::Worker
         include WorkerAttributes
+
+        def perform(*args)
+        end
       end
+
+      allow(::Gitlab::Database::LoadBalancing).to receive_message_chain(:proxy, :load_balancer).and_return(load_balancer)
+      allow(load_balancing_metric).to receive(:increment)
+      allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_load_balancing_count, anything).and_return(load_balancing_metric)
     end
 
-    let(:worker) { TestWorker.new }
+    around do |example|
+      with_sidekiq_server_middleware do |chain|
+        chain.add Gitlab::Database::LoadBalancing::SidekiqServerMiddleware
+        chain.add described_class
+        Sidekiq::Testing.inline! { example.run }
+      end
+    end
 
     include_context 'server metrics with mocked prometheus'
+    include_context 'server metrics call'
+    include_context 'clear DB Load Balancing configuration'
 
-    context 'when load_balancing is enabled' do
-      let(:load_balancing_metric) { double('load balancing metric') }
-
-      include_context 'clear DB Load Balancing configuration'
+    shared_context 'worker declaring data consistency' do
+      let(:worker_class) { LBTestWorker }
 
       before do
-        allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        allow(Gitlab::Metrics).to receive(:counter).with(:sidekiq_load_balancing_count, anything).and_return(load_balancing_metric)
-      end
-
-      describe '#initialize' do
-        it 'sets load_balancing metrics' do
-          expect(Gitlab::Metrics).to receive(:counter).with(:sidekiq_load_balancing_count, anything).and_return(load_balancing_metric)
+        stub_const('LBTestWorker', Class.new(TestWorker))
+        LBTestWorker.class_eval do
+          include ApplicationWorker
 
-          subject
+          data_consistency :delayed
         end
       end
+    end
 
-      describe '#call' do
-        include_context 'server metrics call'
-
-        context 'when :database_chosen is provided' do
-          where(:database_chosen) do
-            %w[primary retry replica]
-          end
-
-          with_them do
-            context "when #{params[:database_chosen]} is used" do
-              let(:labels_with_load_balancing) do
-                labels_with_job_status.merge(database_chosen: database_chosen, data_consistency: 'delayed')
-              end
+    context 'when load_balancing is enabled' do
+      before do
+        allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+      end
 
-              before do
-                job[:database_chosen] = database_chosen
-                job[:data_consistency] = 'delayed'
-                allow(load_balancing_metric).to receive(:increment)
-              end
+      describe '#call' do
+        context 'when worker declares data consistency' do
+          include_context 'worker declaring data consistency'
 
-              it 'increment sidekiq_load_balancing_count' do
-                expect(load_balancing_metric).to receive(:increment).with(labels_with_load_balancing, 1)
+          it 'increments load balancing counter with defined data consistency' do
+            process_job
 
-                described_class.new.call(worker, job, :test) { nil }
-              end
-            end
+            expect(load_balancing_metric).to have_received(:increment).with(
+              a_hash_including(
+                data_consistency: :delayed,
+                load_balancing_strategy: 'replica'
+              ), 1)
           end
         end
 
-        context 'when :database_chosen is not provided' do
-          it 'does not increment sidekiq_load_balancing_count' do
-            expect(load_balancing_metric).not_to receive(:increment)
+        context 'when worker does not declare data consistency' do
+          it 'increments load balancing counter with default data consistency' do
+            process_job
 
-            described_class.new.call(worker, job, :test) { nil }
+            expect(load_balancing_metric).to have_received(:increment).with(
+              a_hash_including(
+                data_consistency: :always,
+                load_balancing_strategy: 'primary'
+              ), 1)
           end
         end
       end
     end
 
     context 'when load_balancing is disabled' do
-      include_context 'clear DB Load Balancing configuration'
+      include_context 'worker declaring data consistency'
 
       before do
         allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
       end
 
       describe '#initialize' do
-        it 'doesnt set load_balancing metrics' do
+        it 'does not set load_balancing metrics' do
           expect(Gitlab::Metrics).not_to receive(:counter).with(:sidekiq_load_balancing_count, anything)
 
           subject
         end
       end
+
+      describe '#call' do
+        it 'does not increment load balancing counter' do
+          process_job
+
+          expect(load_balancing_metric).not_to have_received(:increment)
+        end
+      end
     end
   end
 end
diff --git a/spec/lib/gitlab/sidekiq_middleware/size_limiter/validator_spec.rb b/spec/lib/gitlab/sidekiq_middleware/size_limiter/validator_spec.rb
index 4fbe59c3c27..440eca10a88 100644
--- a/spec/lib/gitlab/sidekiq_middleware/size_limiter/validator_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/size_limiter/validator_spec.rb
@@ -230,11 +230,11 @@ RSpec.describe Gitlab::SidekiqMiddleware::SizeLimiter::Validator do
     end
 
     context 'in compress mode' do
+      let(:size_limit) { 50 }
+      let(:compression_threshold) { 30 }
       let(:mode) { 'compress' }
 
       context 'when job size is less than compression threshold' do
-        let(:size_limit) { 50 }
-        let(:compression_threshold) { 30 }
         let(:job) { job_payload(a: 'a' * 10) }
 
         it 'does not raise an exception' do
@@ -244,8 +244,6 @@ RSpec.describe Gitlab::SidekiqMiddleware::SizeLimiter::Validator do
       end
 
       context 'when job size is bigger than compression threshold and less than size limit after compressed' do
-        let(:size_limit) { 50 }
-        let(:compression_threshold) { 30 }
         let(:args) { { a: 'a' * 300 } }
         let(:job) { job_payload(args) }
 
@@ -260,9 +258,20 @@ RSpec.describe Gitlab::SidekiqMiddleware::SizeLimiter::Validator do
         end
       end
 
+      context 'when the job was already compressed' do
+        let(:job) do
+          job_payload({ a: 'a' * 10 })
+            .merge(Gitlab::SidekiqMiddleware::SizeLimiter::Compressor::COMPRESSED_KEY => true)
+        end
+
+        it 'does not compress the arguments again' do
+          expect(Gitlab::SidekiqMiddleware::SizeLimiter::Compressor).not_to receive(:compress)
+
+          expect { validate.call(TestSizeLimiterWorker, job) }.not_to raise_error
+        end
+      end
+
       context 'when job size is bigger than compression threshold and bigger than size limit after compressed' do
-        let(:size_limit) { 50 }
-        let(:compression_threshold) { 30 }
         let(:args) { { a: 'a' * 3000 } }
         let(:job) { job_payload(args) }
 
diff --git a/spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb b/spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb
index fff925f8532..d6cc787f53d 100644
--- a/spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb
@@ -11,6 +11,8 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Client do
 
       include ApplicationWorker
 
+      feature_category :issue_tracking
+
       def self.job_for_args(args)
         jobs.find { |job| job['args'] == args }
       end
@@ -41,5 +43,39 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Client do
       expect(job1['meta.user']).to eq(user_per_job['job1'].username)
       expect(job2['meta.user']).to eq(user_per_job['job2'].username)
     end
+
+    context 'when the feature category is set in the context_proc' do
+      it 'takes the feature category from the worker, not the caller' do
+        TestWithContextWorker.bulk_perform_async_with_contexts(
+          %w(job1 job2),
+          arguments_proc: -> (name) { [name, 1, 2, 3] },
+          context_proc: -> (_) { { feature_category: 'code_review' } }
+        )
+
+        job1 = TestWithContextWorker.job_for_args(['job1', 1, 2, 3])
+        job2 = TestWithContextWorker.job_for_args(['job2', 1, 2, 3])
+
+        expect(job1['meta.feature_category']).to eq('issue_tracking')
+        expect(job2['meta.feature_category']).to eq('issue_tracking')
+      end
+    end
+
+    context 'when the feature category is already set in the surrounding block' do
+      it 'takes the feature category from the worker, not the caller' do
+        Gitlab::ApplicationContext.with_context(feature_category: 'authentication_and_authorization') do
+          TestWithContextWorker.bulk_perform_async_with_contexts(
+            %w(job1 job2),
+            arguments_proc: -> (name) { [name, 1, 2, 3] },
+            context_proc: -> (_) { {} }
+          )
+        end
+
+        job1 = TestWithContextWorker.job_for_args(['job1', 1, 2, 3])
+        job2 = TestWithContextWorker.job_for_args(['job2', 1, 2, 3])
+
+        expect(job1['meta.feature_category']).to eq('issue_tracking')
+        expect(job2['meta.feature_category']).to eq('issue_tracking')
+      end
+    end
   end
 end