Use surrounding context's feature category for unowned workers

When we're in Sidekiq, in most cases we want to use the feature category
of the current worker class in logs and metrics. If a worker with the
'users' feature category is called by an endpoint that has the
'issue tracking' feature category, then 'users' is the more specific
category.

However, for workers without an owning feature category, this is worse:
we're effectively throwing away the information from the surrounding
context to give the least-specific information we could possibly give.
This changes that so that in logs and in metrics, we'll use the feature
category from the calling context for workers that don't have an owner.

app/workers/all_queues.yml

@@ -1881,7 +1881,7 @@
   :tags: []
 - :name: default
   :worker_name:
-  :feature_category:
+  :feature_category: :not_owned
   :has_external_dependencies:
   :urgency:
   :resource_boundary:
@@ -2206,7 +2206,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))