Merge branch 'use-caller-feature-category-for-unowned-sidekiq-jobs' into 'master'

Use surrounding context's feature category for unowned workers

See merge request gitlab-org/gitlab!67355

app/workers/all_queues.yml

@@ -1890,7 +1890,7 @@
   :tags: []
 - :name: default
   :worker_name:
-  :feature_category:
+  :feature_category: :not_owned
   :has_external_dependencies:
   :urgency:
   :resource_boundary:
@@ -2215,7 +2215,7 @@
   :tags: []
 - :name: mailers
   :worker_name: ActionMailer::MailDeliveryJob
-  :feature_category: :issue_tracking
+  :feature_category: :not_owned
   :has_external_dependencies:
   :urgency: low
   :resource_boundary:

app/workers/concerns/worker_attributes.rb

@@ -46,8 +46,14 @@ module WorkerAttributes
       set_class_attribute(:feature_category, :not_owned)
     end
 
+    # Special case: if a worker is not owned, get the feature category
+    # (if present) from the calling context.
     def get_feature_category
-      get_class_attribute(:feature_category)
+      feature_category = get_class_attribute(:feature_category)
+
+      return feature_category unless feature_category == :not_owned
+
+      Gitlab::ApplicationContext.current_context_attribute('meta.feature_category') || feature_category
     end
 
     def feature_category_not_owned?

doc/development/feature_categorization/index.md

@@ -72,6 +72,11 @@ class SomeCrossCuttingConcernWorker
 end
 ```
 
+Workers marked as not owned workers will, when possible, use the
+category of their caller (worker or HTTP endpoint) in metrics and logs.
+For instance, `ReactiveCachingWorker` can have multiple feature
+categories in metrics and logs.
+
 ## Rails controllers
 
 Specifying feature categories on controller actions can be done using

ee/spec/services/elastic/indexing_control_service_spec.rb

@@ -27,6 +27,10 @@ RSpec.describe Elastic::IndexingControlService, :clean_gitlab_redis_shared_state
 
   subject { described_class.new(worker_class) }
 
+  before do
+    stub_const(worker_class.name, worker_class)
+  end
+
   describe '.initialize' do
     it 'raises an exception when passed wrong worker' do
       expect { described_class.new(Class.new) }.to raise_error(ArgumentError)
@@ -154,6 +158,10 @@ RSpec.describe Elastic::IndexingControlService, :clean_gitlab_redis_shared_state
 
     let(:other_subject) { described_class.new(second_worker_class) }
 
+    before do
+      stub_const(second_worker_class.name, second_worker_class)
+    end
+
     it 'allows to use queues independently of each other' do
       expect { subject.add_to_waiting_queue!(worker_args, worker_context) }.to change { subject.queue_size }.from(0).to(1)
 

lib/gitlab/sidekiq_config.rb

@@ -23,12 +23,12 @@ module Gitlab
     DEFAULT_WORKERS = {
       '_' => DummyWorker.new(
         queue: 'default',
-        weight: 1, tags: []
+        weight: 1,
+        tags: []
       ),
       'ActionMailer::MailDeliveryJob' => DummyWorker.new(
         name: 'ActionMailer::MailDeliveryJob',
         queue: 'mailers',
-        feature_category: :issue_tracking,
         urgency: 'low',
         weight: 2,
         tags: []

lib/gitlab/sidekiq_config/dummy_worker.rb

@@ -6,7 +6,6 @@ module Gitlab
     class DummyWorker
       ATTRIBUTE_METHODS = {
         name: :name,
-        feature_category: :get_feature_category,
         has_external_dependencies: :worker_has_external_dependencies?,
         urgency: :get_urgency,
         resource_boundary: :get_worker_resource_boundary,
@@ -27,6 +26,20 @@ module Gitlab
         nil
       end
 
+      # All dummy workers are unowned; get the feature category from the
+      # context if available.
+      def get_feature_category
+        Gitlab::ApplicationContext.current_context_attribute('meta.feature_category') || :not_owned
+      end
+
+      def get_worker_context
+        nil
+      end
+
+      def context_for_arguments(*)
+        nil
+      end
+
       ATTRIBUTE_METHODS.each do |attribute, meth|
         define_method meth do
           @attributes[attribute]

lib/gitlab/sidekiq_middleware.rb

@@ -13,6 +13,13 @@ module Gitlab
         chain.add ::Gitlab::SidekiqMiddleware::SizeLimiter::Server
         chain.add ::Gitlab::SidekiqMiddleware::Monitor
 
+        # Labkit wraps the job in the `Labkit::Context` resurrected from
+        # the job-hash. We need properties from the context for
+        # recording metrics, so this needs to be before
+        # `::Gitlab::SidekiqMiddleware::ServerMetrics` (if we're using
+        # that).
+        chain.add ::Labkit::Middleware::Sidekiq::Server
+
         if metrics
           chain.add ::Gitlab::SidekiqMiddleware::ServerMetrics
 
@@ -24,7 +31,6 @@ module Gitlab
         chain.add ::Gitlab::SidekiqMiddleware::RequestStoreMiddleware
         chain.add ::Gitlab::SidekiqMiddleware::ExtraDoneLogMetadata
         chain.add ::Gitlab::SidekiqMiddleware::BatchLoader
-        chain.add ::Labkit::Middleware::Sidekiq::Server
         chain.add ::Gitlab::SidekiqMiddleware::InstrumentationLogger
         chain.add ::Gitlab::SidekiqMiddleware::AdminMode::Server
         chain.add ::Gitlab::SidekiqVersioning::Middleware

lib/gitlab/sidekiq_middleware/client_metrics.rb

@@ -14,6 +14,7 @@ module Gitlab
       def call(worker_class, job, queue, _redis_pool)
         # worker_class can either be the string or class of the worker being enqueued.
         worker_class = worker_class.safe_constantize if worker_class.respond_to?(:safe_constantize)
+
         labels = create_labels(worker_class, queue, job)
         labels[:scheduling] = job.key?('at') ? 'delayed' : 'immediate'
 

lib/gitlab/sidekiq_middleware/metrics_helper.rb

@@ -3,14 +3,21 @@
 module Gitlab
   module SidekiqMiddleware
     module MetricsHelper
+      include ::Gitlab::SidekiqMiddleware::WorkerContext
+
       TRUE_LABEL = "yes"
       FALSE_LABEL = "no"
 
       private
 
       def create_labels(worker_class, queue, job)
-        worker_name = (job['wrapped'].presence || worker_class).to_s
-        worker = find_worker(worker_name, worker_class)
+        worker = find_worker(worker_class, job)
+
+        # This should never happen: we should always be able to find a
+        # worker class for a given Sidekiq job. But if we can't, we
+        # shouldn't blow up here, because we want to record this in our
+        # metrics.
+        worker_name = worker.try(:name) || worker.class.name
 
         labels = { queue: queue.to_s,
                    worker: worker_name,
@@ -23,9 +30,7 @@ module Gitlab
 
         labels[:urgency] = worker.get_urgency.to_s
         labels[:external_dependencies] = bool_as_label(worker.worker_has_external_dependencies?)
-
-        feature_category = worker.get_feature_category
-        labels[:feature_category] = feature_category.to_s
+        labels[:feature_category] = worker.get_feature_category.to_s
 
         resource_boundary = worker.get_worker_resource_boundary
         labels[:boundary] = resource_boundary == :unknown ? "" : resource_boundary.to_s
@@ -36,10 +41,6 @@ module Gitlab
       def bool_as_label(value)
         value ? TRUE_LABEL : FALSE_LABEL
       end
-
-      def find_worker(worker_name, worker_class)
-        Gitlab::SidekiqConfig::DEFAULT_WORKERS.fetch(worker_name, worker_class)
-      end
     end
   end
 end

lib/gitlab/sidekiq_middleware/worker_context.rb

@@ -10,6 +10,12 @@ module Gitlab
 
         context_or_nil.use(&block)
       end
+
+      def find_worker(worker_class, job)
+        worker_name = (job['wrapped'].presence || worker_class).to_s
+
+        Gitlab::SidekiqConfig::DEFAULT_WORKERS[worker_name]&.klass || worker_class
+      end
     end
   end
 end

lib/gitlab/sidekiq_middleware/worker_context/client.rb

@@ -7,20 +7,15 @@ module Gitlab
         include Gitlab::SidekiqMiddleware::WorkerContext
 
         def call(worker_class_or_name, job, _queue, _redis_pool, &block)
-          worker_class = worker_class_or_name.to_s.safe_constantize
+          worker_class = find_worker(worker_class_or_name.to_s.safe_constantize, job)
 
-          # Mailers can't be constantized like this
+          # This is not a worker we know about, perhaps from a gem
           return yield unless worker_class
-          return yield unless worker_class.include?(::ApplicationWorker)
+          return yield unless worker_class.respond_to?(:context_for_arguments)
 
           context_for_args = worker_class.context_for_arguments(job['args'])
 
-          wrap_in_optional_context(context_for_args) do
-            # This should be inside the context for the arguments so
-            # that we don't override the feature category on the worker
-            # with the one from the caller.
-            Gitlab::ApplicationContext.with_context(feature_category: worker_class.get_feature_category.to_s, &block)
-          end
+          wrap_in_optional_context(context_for_args, &block)
         end
       end
     end

lib/gitlab/sidekiq_middleware/worker_context/server.rb

@@ -7,7 +7,7 @@ module Gitlab
         include Gitlab::SidekiqMiddleware::WorkerContext
 
         def call(worker, job, _queue, &block)
-          worker_class = worker.class
+          worker_class = find_worker(worker.class, job)
 
           # This is not a worker we know about, perhaps from a gem
           return yield unless worker_class.respond_to?(:get_worker_context)

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -1631,10 +1631,19 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     let(:worker) do
       Class.new do
         include Sidekiq::Worker
+
         sidekiq_options queue: 'test'
+
+        def self.name
+          'WorkerClass'
+        end
       end
     end
 
+    before do
+      stub_const(worker.name, worker)
+    end
+
     describe '#sidekiq_queue_length' do
       context 'when queue is empty' do
         it 'returns zero' do

spec/lib/gitlab/sidekiq_middleware/client_metrics_spec.rb

@@ -63,5 +63,18 @@ RSpec.describe Gitlab::SidekiqMiddleware::ClientMetrics do
         Sidekiq::Testing.inline! { TestWorker.perform_in(1.second) }
       end
     end
+
+    context 'when the worker class cannot be found' do
+      it 'increments enqueued jobs metric with the worker labels set to NilClass' do
+        test_anonymous_worker = Class.new(TestWorker)
+
+        expect(enqueued_jobs_metric).to receive(:increment).with(a_hash_including(worker: 'NilClass'), 1)
+
+        # Sidekiq won't be able to create an instance of this class
+        expect do
+          Sidekiq::Testing.inline! { test_anonymous_worker.perform_async }
+        end.to raise_error(NameError)
+      end
+    end
   end
 end

spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb

@@ -211,6 +211,9 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
       end
     end
 
+    include_context 'server metrics with mocked prometheus'
+    include_context 'server metrics call'
+
     before do
       stub_const('TestWorker', Class.new)
       TestWorker.class_eval do
@@ -234,9 +237,6 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
       end
     end
 
-    include_context 'server metrics with mocked prometheus'
-    include_context 'server metrics call'
-
     shared_context 'worker declaring data consistency' do
       let(:worker_class) { LBTestWorker }
 
@@ -308,5 +308,67 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
       end
     end
   end
+
+  context 'feature attribution' do
+    let(:test_worker) do
+      category = worker_category
+
+      Class.new do
+        include Sidekiq::Worker
+        include WorkerAttributes
+
+        if category
+          feature_category category
+        else
+          feature_category_not_owned!
+        end
+
+        def perform
+        end
+      end
+    end
+
+    let(:context_category) { 'continuous_integration' }
+    let(:job) { { 'meta.feature_category' => 'continuous_integration' } }
+
+    before do
+      stub_const('TestWorker', test_worker)
+    end
+
+    around do |example|
+      with_sidekiq_server_middleware do |chain|
+        Gitlab::SidekiqMiddleware.server_configurator(
+          metrics: true,
+          arguments_logger: false,
+          memory_killer: false
+        ).call(chain)
+
+        Sidekiq::Testing.inline! { example.run }
+      end
+    end
+
+    include_context 'server metrics with mocked prometheus'
+    include_context 'server metrics call'
+
+    context 'when a worker has a feature category' do
+      let(:worker_category) { 'authentication_and_authorization' }
+
+      it 'uses that category for metrics' do
+        expect(completion_seconds_metric).to receive(:observe).with(a_hash_including(feature_category: worker_category), anything)
+
+        TestWorker.process_job(job)
+      end
+    end
+
+    context 'when a worker does not have a feature category' do
+      let(:worker_category) { nil }
+
+      it 'uses the category from the context for metrics' do
+        expect(completion_seconds_metric).to receive(:observe).with(a_hash_including(feature_category: context_category), anything)
+
+        TestWorker.process_job(job)
+      end
+    end
+  end
 end
 # rubocop: enable RSpec/MultipleMemoizedHelpers

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

@@ -11,8 +11,6 @@ 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
@@ -23,7 +21,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Client do
   end
 
   before do
-    stub_const('TestWithContextWorker', worker_class)
+    stub_const(worker_class.name, worker_class)
   end
 
   describe "#call" do
@@ -43,39 +41,5 @@ 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

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

@@ -3,7 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Server do
-  let(:worker_class) do
+  let(:test_worker) do
     Class.new do
       def self.name
         "TestWorker"
@@ -23,6 +23,16 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Server do
     end
   end
 
+  let(:not_owned_worker) do
+    Class.new(test_worker) do
+      def self.name
+        "NotOwnedWorker"
+      end
+
+      feature_category_not_owned!
+    end
+  end
+
   let(:other_worker) do
     Class.new do
       def self.name
@@ -37,7 +47,8 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Server do
   end
 
   before do
-    stub_const("TestWorker", worker_class)
+    stub_const("TestWorker", test_worker)
+    stub_const("NotOwnedWorker", not_owned_worker)
     stub_const("OtherWorker", other_worker)
   end
 
@@ -57,10 +68,24 @@ RSpec.describe Gitlab::SidekiqMiddleware::WorkerContext::Server do
       expect(TestWorker.contexts['identifier'].keys).not_to include('meta.user')
     end
 
-    it 'takes the feature category from the worker' do
-      TestWorker.perform_async('identifier', 1)
+    context 'feature category' do
+      it 'takes the feature category from the worker' do
+        Gitlab::ApplicationContext.with_context(feature_category: 'authentication_and_authorization') do
+          TestWorker.perform_async('identifier', 1)
+        end
+
+        expect(TestWorker.contexts['identifier']).to include('meta.feature_category' => 'foo')
+      end
 
-      expect(TestWorker.contexts['identifier']).to include('meta.feature_category' => 'foo')
+      context 'when the worker is not owned' do
+        it 'takes the feature category from the surrounding context' do
+          Gitlab::ApplicationContext.with_context(feature_category: 'authentication_and_authorization') do
+            NotOwnedWorker.perform_async('identifier', 1)
+          end
+
+          expect(NotOwnedWorker.contexts['identifier']).to include('meta.feature_category' => 'authentication_and_authorization')
+        end
+      end
     end
 
     it "doesn't fail for unknown workers" do

spec/lib/gitlab/sidekiq_middleware_spec.rb

@@ -58,13 +58,13 @@ RSpec.describe Gitlab::SidekiqMiddleware do
     let(:all_sidekiq_middlewares) do
       [
         ::Gitlab::SidekiqMiddleware::Monitor,
+        ::Labkit::Middleware::Sidekiq::Server,
         ::Gitlab::SidekiqMiddleware::ServerMetrics,
         ::Gitlab::SidekiqMiddleware::ArgumentsLogger,
         ::Gitlab::SidekiqMiddleware::MemoryKiller,
         ::Gitlab::SidekiqMiddleware::RequestStoreMiddleware,
         ::Gitlab::SidekiqMiddleware::ExtraDoneLogMetadata,
         ::Gitlab::SidekiqMiddleware::BatchLoader,
-        ::Labkit::Middleware::Sidekiq::Server,
         ::Gitlab::SidekiqMiddleware::InstrumentationLogger,
         ::Gitlab::SidekiqMiddleware::AdminMode::Server,
         ::Gitlab::SidekiqVersioning::Middleware,

spec/support/shared_contexts/lib/gitlab/sidekiq_middleware/server_metrics_shared_context.rb

@@ -17,6 +17,9 @@ RSpec.shared_context 'server metrics with mocked prometheus' do
   let(:elasticsearch_requests_total) { double('elasticsearch calls total metric') }
 
   before do
+    allow(Gitlab::Metrics).to receive(:histogram).and_call_original
+    allow(Gitlab::Metrics).to receive(:counter).and_call_original
+
     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)

spec/workers/concerns/application_worker_spec.rb

@@ -248,6 +248,10 @@ RSpec.describe ApplicationWorker do
   end
 
   describe '.perform_async' do
+    before do
+      stub_const(worker.name, worker)
+    end
+
     shared_examples_for 'worker utilizes load balancing capabilities' do |data_consistency|
       before do
         worker.data_consistency(data_consistency)
@@ -282,6 +286,10 @@ RSpec.describe ApplicationWorker do
   end
 
   describe '.bulk_perform_async' do
+    before do
+      stub_const(worker.name, worker)
+    end
+
     it 'enqueues jobs in bulk' do
       Sidekiq::Testing.fake! do
         worker.bulk_perform_async([['Foo', [1]], ['Foo', [2]]])
@@ -293,6 +301,10 @@ RSpec.describe ApplicationWorker do
   end
 
   describe '.bulk_perform_in' do
+    before do
+      stub_const(worker.name, worker)
+    end
+
     context 'when delay is valid' do
       it 'correctly schedules jobs' do
         Sidekiq::Testing.fake! do

spec/workers/concerns/worker_context_spec.rb

@@ -13,6 +13,10 @@ RSpec.describe WorkerContext do
     end
   end
 
+  before do
+    stub_const(worker.name, worker)
+  end
+
   describe '.worker_context' do
     it 'allows modifying the context for the entire worker' do
       worker.worker_context(user: build_stubbed(:user))